擴散語言模型九倍推理加速！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的自適應緩存策略

通過這種“長間隔”與“自適應”相結合的緩存策略，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