首页 > 生活

扩散语言模型九倍推理加速！上海交大：KV Cache并非自回归模型的专属技巧

首个用于加速扩散式大语言模型（diffusion-based Large Language Models, 简称 dLLMs）推理过程的免训练方法。

上海交通大学 EPIC Lab 团队提出了一种无需训练、即插即用的高效推理缓存机制：dLLM-Cache。

其核心思想在于，在一个多步去噪过程中，复用相邻时间步上变化较小的特征，仅更新那些变化较大的特征，从而实现了计算量的大幅降低，并保持了原有的生成质量。

图 1 不同 dLLMs 使用 dLLM – Cache 和不使用 dLLM – Cache 在速度和质量上的对比

dLLM-Cache 具有几个重要的亮点：

1. 训练无关，即插即用。dLLM-Cache 完全在推理过程中工作，无需修改模型参数或重训练。dLLM-Cache 可以在完全不损失模型输出质量的前提下，带来最高 9.1 倍的推理速度提升。

2. 通用于主流 dLLM 架构，如 LLaDA、Dream 以及 LLaDA-V、MMaDA、Dimple 等多模态模型。

3. 在推理过程中，首次识别出了 prompt 部分的 Transformer 中间层特征（Key、Value、Attention output、FFN output）长期稳定，而 response 部分仅有一小部分 tokens 的特征变化较大，为缓存特征并后续复用提供了理论基础。

4. 独创了以 V-verify 机制为核心的选择更新策略。以 Value 向量的变化为选择基准，成功识别出了 response 部分变化较大的那些 tokens，通过仅更新这些特征，摒弃了高达 75% 的冗余计算。

本论文共同第一作者刘知远和杨奕存是哈尔滨工业大学 2022 级本科生，目前在上海交通大学 EPIC Lab 进行科研实习，师从张林峰助理教授，主要研究方向为高效深度学习，此前曾在 CVPR2025 上收获满分论文。

接下来，我们一起来看看该研究的细节。

研究动机

基于扩散的大语言模型正成为语言生成领域最受关注的新范式之一。随着模型架构的发展、去噪算法的优化以及 Masked Diffusion 在语言建模中逐步展现出与自回归模型不同的建模能力，这类模型正在逐步成为挑战 GPT 等主流模型的重要力量。

以 LLaDA、Dream 为代表的扩散语言模型，基于迭代去噪的生成过程，不再依赖严格的自回归因果结构，天然支持双向建模、全局依赖和反向推理等能力，已经在 " 逆转诅咒 "、数学推理等任务上展现出领先性能。

然而，这种范式的优势也伴随着巨大的代价。为了确保生成的质量，dLLMs 在推理过程中通常需要执行长达数百步的去噪迭代，每一步都需重新计算 attention、FFN 等所有层的特征，计算量相当于多次完整前向传播。这为 dLLMs 的推理效率带来了严重的瓶颈，制约了其实际部署。更重要的是，主流的加速手段如用于自回归模型的 KV Cache，由于不兼容双向注意力架构，在 dLLMs 中完全失效。

与传统的自回归语言模型不同，dLLMs 不再依赖顺序生成下一个 token，而是采用随机遮蔽 ( mask ) + 逐步还原的方式建模 token 分布，这种机制使得模型具备天然的双向建模能力，理论上能够更好地处理逆向逻辑、长距离依赖等任务。

LLaDA 等模型已经在多个基准任务中超越主流 ARMs，尤其在 " 逆转诅咒 " 上明显胜出。

然而，这种扩散式推理带来一个严重的挑战：为了确保生成质量，dLLMs 通常需要上百步的去噪迭代，每一步都需全量计算 Attention、FFN 等模块，导致其推理速度相比 ARMs 慢一个数量级，落地成本高。同时，ARMs 通用的加速方法如 KV-Cache 因 dLLMs 的双向注意力设计而无法兼容。这些造成了 dLLMs 在推理时既慢又缺乏加速手段的现象。这正是 dLLM-Cache 所要破解的核心问题。

方法简介

本文作者仔细研究了 dLLMs 推理的中间特征变化过程，发现如下关键现象：

图 2 dLLM 中两个相邻去噪步骤之间的 Key、Value、Attention Output 和 FFN Output 的余弦相似度

Prompt tokens 的特征在整个去噪过程中基本保持稳定，每一步都重新计算这些特征是完全不必要且浪费计算资源的；

Response tokens 多数变化很小，仅少部分变化剧烈，全量计算所有 response tokens 存在冗余。

由此，问题转化为了如何高效识别出这些变化剧烈的 response tokens。

图 3 Response tokens 的 K 或 V 变化与其他特征变化的相关性

本文作者首创性得提出了 V-verify 机制。它的提出源于另一项重要的发现：作者量化了 response tokens 的底层特征（Key, Value 向量）的变化与其上层复杂特征（Attention Output, FFN Output）的变化之间的关系，结果显示它们存在着极强的正相关性，皮尔逊相关系数最高可达 0.944。

这意味着，一个 token 底层的 Value 向量是否发生变化，是其整体状态是否发生改变的一个极佳的、且计算成本极低的 " 指示器 "。

基于以上这些关键的观察，本文作者提出了 dLLM-Cache ，具体的框架设计如下：

图 4 dLLM-Cache 方法整体 pipeline

Prompt 缓存：长间隔重用

对于 prompt 部分，作者设计了长间隔 Prompt 缓存，每隔 Kp 步（在实验中一般设置为 100）更新一次 prompt 的 Key、Value、Attention Output、FFN Output，其余步骤全部复用先前结果。这样避免了对稳定不变的特征的重复计算，大幅减少了计算量。

Response 缓存：自适应部分更新

对生成目标 response 区域，由于 response tokens 的特征并不是一直保持稳定不变的，作者设计了较短间隔的 Response 缓存，每隔 Kr 步（在实验中一般设置为 8 左右）全量更新一次 response 的 Key、Value、Attention Output、FFN Output，在其余的步骤，作者提出了基于 V-verify 的自适应缓存策略：

在每个去噪步骤，首先计算所有 response tokens 最新的 Value 向量。

然后，通过计算新 Value 向量与缓存中旧 Value 向量的余弦相似度，将余弦相似度作为每个 response tokens 的一个 " 变化分 "。

只选出 " 变化分 " 最高（即相似度最低）的极少数tokens（例如，变化最剧烈的 25%），将它们标记为 " 待更新 " 。

最后，模型只对这些被标记的 " 待更新 "tokens，进行完整的特征重计算。而其余75%的 " 稳定 "tokens，则继续高效地从缓存中复用其特征。

通过这种 " 长间隔 " 与 " 自适应 " 相结合的缓存策略，dLLM-Cache 在 Transformer 的每一层都实现了计算量的极致优化，且整个过程无需任何额外训练，做到了真正的即插即用。

3 实验结果

本文在 LLaDA 8B 和 Dream 7B 两大代表性的开源 dLLM 的基础版与指令微调版上，针对数学与科学、通用任务、代码生成三大领域的8个主流基准测试，对 dLLM-Cache 的有效性进行了严苛的检验。评估维度不仅包括推理速度（TPS）和计算效率（FLOPs），更核心的是模型性能得分（Score），以确保加速不是以牺牲模型能力为代价。

本文在 LLaDA 8B 的基础版和指令微调版上都部署了 dLLM-Cache，下图的实验结果充分展示了其强大的加速能力和卓越的生成质量保持。在几乎所有的基准测试中，都达到了5 倍以上的加速效果，且在绝大部分情况下，生成质量都没有降低，甚至有轻微的提升。特别是当面对 LongBench 任务时，prompt 的稳定性带来了更显著的加速效果，在 HotpotQA 上实现了高达9.1 倍的无损加速。

图 5 dLLM-Cache 在 LLaDA 模型上的效果

为了进一步证明 dLLM-Cache 的通用性和鲁棒性，作者将其无缝迁移至另一款架构略有不同的 dLLM —— Dream 7B 上。下图的实验结果再次印证了 dLLM-Cache 方法的有效性，充分说明了其通用于主流 dLLM 架构。

图 6 dLLM-Cache 在 Dream 模型上的效果

作者还将 dLLM 和主流的基于 ARM 的 LLM 进行了对比，下图展示了 LLaDA 8B 与 LLaMA3 8B 在 GSM8K 任务上的比较。结果显示，原始的 LLaDA 在准确率上以近 20 个点的巨大优势领先于 LLaMA3，但在推理速度上却远不及。然而，在使用了本文的 dLLM-Cache 之后，LLaDA 的推理速度获得了超过 3.3 倍的提升，首次超过了 LLaMA3 的推理速度。这一结果有力地证明，本文提出的 dLLM-Cache 能够让 dLLMs 在保持其显著准确率优势的同时，获得与 ARMs 相当竞争力的推理速度。