扩散语言模型九倍推理加速！KV Cache并非自回归模型专属

扩散语言模型九倍推理加速！KV Cache并非自回归模型专属

EPIC Lab团队投稿量子位 | 民众号 QbitAI

首个用于加速扩散式大语言模型（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相当竞争力的推理速度

图7 使用dLLM-Cache的dLLM vs 使用KV-Cache的ARM

论文链接： https://github.com/maomaocun/dLLM-cache/blob/main/asset/paper.pdf代码已开源： https://github.com/maomaocun/dLLM-Cache

分享： 2025-06-02 17:42:31 共81款