EPIC Lab團隊 投稿量子位 | 公眾號 QbitAI

首個用于加速擴散式大語言模型（diffusion-based Large Language Models, 簡稱 dLLMs）推理過程的免訓練方法。

上海交通大學EPIC Lab團隊提出了一種無需訓練、即插即用的高效推理緩存機制：dLLM-Cache。

其核心思想在于，在一個多步去噪過程中，復用相鄰時間步上變化較小的特征，僅更新那些變化較大的特征，從而實現(xiàn)了計算量的大幅降低，并保持了原有的生成質量。

圖1 不同dLLMs使用dLLM–Cache和不使用dLLM–Cache在速度和質量上的對比

dLLM-Cache具有幾個重要的亮點：

1. 訓練無關，即插即用。dLLM-Cache完全在推理過程中工作，無需修改模型參數(shù)或重訓練。dLLM-Cache可以在完全不損失模型輸出質量的前提下，帶來最高9.1倍的推理速度提升 。

2.通用于主流dLLM架構，如LLaDA、Dream以及LLaDA-V、MMaDA、Dimple等多模態(tài)模型。

3. 在推理過程中，首次識別出了prompt部分的Transformer中間層特征（Key、Value、Attention output、FFN output）長期穩(wěn)定，而response部分僅有一小部分tokens的特征變化較大，為緩存特征并后續(xù)復用提供了理論基礎。

4. 獨創(chuàng)了以V-verify機制為核心的選擇更新策略。以Value向量的變化為選擇基準，成功識別出了response部分變化較大的那些tokens，通過僅更新這些特征，摒棄了高達75%的冗余計算。

本論文共同第一作者劉知遠和楊奕存是哈爾濱工業(yè)大學2022級本科生，目前在上海交通大學EPIC Lab進行科研實習，師從張林峰助理教授，主要研究方向為高效深度學習，此前曾在CVPR2025上收獲滿分論文。

接下來，我們一起來看看該研究的細節(jié)。

研究動機

基于擴散的大語言模型正成為語言生成領域最受關注的新范式之一。隨著模型架構的發(fā)展、去噪算法的優(yōu)化以及Masked Diffusion在語言建模中逐步展現(xiàn)出與自回歸模型不同的建模能力，這類模型正在逐步成為挑戰(zhàn) GPT 等主流模型的重要力量。

以LLaDA、Dream為代表的擴散語言模型，基于迭代去噪的生成過程，不再依賴嚴格的自回歸因果結構，天然支持雙向建模、全局依賴和反向推理等能力，已經(jīng)在“逆轉詛咒”、數(shù)學推理等任務上展現(xiàn)出領先性能。

然而，這種范式的優(yōu)勢也伴隨著巨大的代價。為了確保生成的質量，dLLMs在推理過程中通常需要執(zhí)行長達數(shù)百步的去噪迭代，每一步都需重新計算attention、FFN等所有層的特征，計算量相當于多次完整前向傳播。這為dLLMs的推理效率帶來了嚴重的瓶頸，制約了其實際部署。更重要的是，主流的加速手段如用于自回歸模型的KV Cache，由于不兼容雙向注意力架構，在dLLMs中完全失效。

與傳統(tǒng)的自回歸語言模型不同，dLLMs不再依賴順序生成下一個token，而是采用隨機遮蔽(mask) + 逐步還原的方式建模token分布，這種機制使得模型具備天然的雙向建模能力，理論上能夠更好地處理逆向邏輯、長距離依賴等任務。

LLaDA 等模型已經(jīng)在多個基準任務中超越主流ARMs，尤其在“逆轉詛咒”上明顯勝出。

然而，這種擴散式推理帶來一個嚴重的挑戰(zhàn)：為了確保生成質量，dLLMs通常需要上百步的去噪迭代，每一步都需全量計算Attention、FFN等模塊，導致其推理速度相比ARMs慢一個數(shù)量級，落地成本高。同時，ARMs 通用的加速方法如KV-Cache因dLLMs的雙向注意力設計而無法兼容。這些造成了dLLMs在推理時既慢又缺乏加速手段的現(xiàn)象。這正是 dLLM-Cache所要破解的核心問題。

方法簡介

本文作者仔細研究了dLLMs推理的中間特征變化過程，發(fā)現(xiàn)如下關鍵現(xiàn)象：

圖2 dLLM中兩個相鄰去噪步驟之間的Key、Value、Attention Output和FFN Output的余弦相似度

Prompt tokens的特征在整個去噪過程中基本保持穩(wěn)定，每一步都重新計算這些特征是完全不必要且浪費計算資源的；

Response tokens多數(shù)變化很小，僅少部分變化劇烈，全量計算所有response tokens存在冗余。

由此，問題轉化為了如何高效識別出這些變化劇烈的response tokens。

圖3 Response tokens的K或V變化與其他特征變化的相關性

本文作者首創(chuàng)性得提出了V-verify機制。它的提出源于另一項重要的發(fā)現(xiàn)：作者量化了response tokens的底層特征（Key, Value向量）的變化與其上層復雜特征（Attention Output, FFN Output）的變化之間的關系，結果顯示它們存在著極強的正相關性，皮爾遜相關系數(shù)最高可達0.944。

這意味著，一個token底層的Value向量是否發(fā)生變化，是其整體狀態(tài)是否發(fā)生改變的一個極佳的、且計算成本極低的“指示器”。

基于以上這些關鍵的觀察，本文作者提出了dLLM-Cache ，具體的框架設計如下：

圖4 dLLM-Cache方法整體pipeline

Prompt緩存：長間隔重用

對于prompt部分，作者設計了長間隔Prompt緩存，每隔Kp步（在實驗中一般設置為100）更新一次prompt的Key、Value、Attention Output、FFN Output，其余步驟全部復用先前結果。這樣避免了對穩(wěn)定不變的特征的重復計算，大幅減少了計算量

Response緩存：自適應部分更新

對生成目標response區(qū)域，由于response tokens的特征并不是一直保持穩(wěn)定不變的，作者設計了較短間隔的Response緩存，每隔Kr步（在實驗中一般設置為8左右）全量更新一次response的Key、Value、Attention Output、FFN Output，在其余的步驟，作者提出了基于V-verify的自適應緩存策略

1. 在每個去噪步驟，首先計算所有response tokens最新的Value向量。
2. 然后，通過計算新Value向量與緩存中舊Value向量的余弦相似度，將余弦相似度作為每個response tokens的一個“變化分”。
3. 選出“變化分”最高（即相似度最低）的極少數(shù)tokens（例如，變化最劇烈的25%），將它們標記為“待更新” 。
4. 最后，模型只對這些被標記的“待更新”tokens，進行完整的特征重計算。而其余75%的“穩(wěn)定”tokens，則繼續(xù)高效地從緩存中復用其特征

通過這種“長間隔”與“自適應”相結合的緩存策略，dLLM-Cache在Transformer的每一層都實現(xiàn)了計算量的極致優(yōu)化，且整個過程無需任何額外訓練，做到了真正的即插即用

3 實驗結果

本文在 LLaDA 8B和Dream 7B兩大代表性的開源dLLM的基礎版與指令微調版上，針對數(shù)學與科學、通用任務、代碼生成三大領域的8個主流基準測試，對dLLM-Cache的有效性進行了嚴苛的檢驗 。評估維度不僅包括推理速度（TPS）計算效率（FLOPs），更核心的是模型性能得分（Score），以確保加速不是以犧牲模型能力為代價

本文在LLaDA 8B的基礎版和指令微調版上都部署了dLLM-Cache，下圖的實驗結果充分展示了其強大的加速能力和卓越的生成質量保持。在幾乎所有的基準測試中，達到了5倍以上的加速效果，且在絕大部分情況下，生成質量都沒有降低，甚至有輕微的提升。特別是當面對LongBench任務時，prompt的穩(wěn)定性帶來了更顯著的加速效果，在HotpotQA上實現(xiàn)了高達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個點的巨大優(yōu)勢領先于LLaMA3，但在推理速度上卻遠不及。然而，在使用了本文的dLLM-Cache之后，LLaDA的推理速度獲得了超過3.3倍的提升，首次超過了LLaMA3的推理速度。這一結果有力地證明，本文提出的dLLM-Cache能夠讓dLLMs在保持其顯著準確率優(yōu)勢的同時，獲得與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