在大語言模型(LLM)領域,推理效率是制約其實際應用的關鍵因素之一。谷歌 DeepMind 的 Gemini diffusion 曾以1400 tokens / 秒的生成速度震驚學界,展現了擴散模型在并行生成上的潛力。然而,開源擴散 LLM 卻因缺乏 KV 緩存機制和并行解碼質量衰退,實際推理速度長期被自回歸模型壓制.
近日,NVIDIA 聯合香港大學、MIT 等機構重磅推出Fast-dLLM,以無需訓練的即插即用加速方案,實現了推理速度的突破!
- 論文:Fast-dLLM: Training-free Acceleration of Diffusion LLM by Enabling KV Cache and Parallel Decoding
- 項目地址:https://nvlabs.github.io/Fast-dLLM
- 論文鏈接:http://arxiv.org/abs/2505.22618
- GitHub 鏈接:https://github.com/NVlabs/Fast-dLLM
通過創新的技術組合,在不依賴重新訓練模型的前提下,該工作為擴散模型的推理加速帶來了突破性進展。本文將結合具體技術細節與實驗數據,解析其核心優勢。
一、 核心技術
分塊 KV 緩存與置信度感知并行解碼
1. 分塊 KV 緩存(Block-Wise KV Cache):激活重用率超 90% 的雙向加速
傳統擴散模型因雙向注意力機制難以直接復用計算結果,導致長序列推理效率低下。Fast-dLLM 提出分塊 KV 緩存機制,通過以下設計實現高效計算:
- 雙向緩存策略:采用 DualCache 同時緩存前綴(Prompt)和后綴(Masked Tokens)的注意力激活值(KV Cache),如圖 1 (a)(b) 所示。在分塊生成時,前序塊的 KV 激活可直接復用于后續塊,減少重復計算。
- 高相似度驗證:實驗表明,相鄰推理步驟的 KV 激活余弦相似度接近 1(圖 2),證明緩存復用的可行性。例如,在 LLaDA 模型中,通過緩存可實現 90% 以上的激活重用,單步計算量顯著降低。
2. 置信度感知并行解碼(Confidence-Aware Parallel Decoding)
并行解碼雖能提升速度,但條件獨立假設易破壞 token 依賴關系,比方說這個例子 The list of poker hands that consist of two English words are: _ _.。后續兩個單詞可以是 “high card,” “two pair,” “full house,” 或者是 “straight flush.”。值得注意的是,這兩個單詞之間存在關聯。
然而,MDMs 中的多令牌預測過程首先為每個令牌生成一個概率分布,然后從這些分布中獨立采樣。這種獨立采樣可能導致不理想的組合(如生成 “high house” 等無效組合)。Fast-dLLM 通過動態置信度篩選解決這一問題(所謂置信度,是指模型給 token 賦予的概率大小):
- 閾值激活策略:僅對置信度超過閾值(如≥0.9)的 token 進行并行解碼,低置信度 token 留待后續步驟處理。如圖 3 所示,該策略可在保證生成質量的前提下,并行輸出多個 token。
- 理論證明:當 (n+1)?≤1 時(n 為并行解碼 token 數,并且并行解碼的 n 個 token 的置信度都大于 1-?),貪婪解碼策略下并行解碼與順序解碼結果一致,從數學層面確保了生成邏輯的連貫性。
3. 偽代碼:分塊 KV 緩存與置信度感知并行解碼流程
以下是 Fast-dLLM 算法的核心偽代碼,結合了分塊 KV 緩存以及置信度感知并行解碼,無需訓練就可以在現有的開源 Diffusion LLM(如 LLaDA、Dream)上即插即用進行推理加速。
二、 性能突破
速度與精度的均衡優化
1. 長文本生成:27.6 倍端到端加速
在 LLaDA 模型上,針對 1024 token 的長文本生成任務,Fast-dLLM 將單步延遲從 0.26 秒降至 0.09 秒,整體耗時從 266 秒壓縮至 12 秒,實現 27.6 倍端到端加速。這一提升在代碼生成、數學推理等長序列場景中尤為顯著,例如 8-shot 提示的 GSM8K 任務中,加速后仍能保持 76% 的準確率。
2. 精度保持:損失 < 2% 的基準測試表現
在主流基準測試中,Fast-dLLM 的準確率損失控制在 2% 以內:
- GSM8K(5-shot):LLaDA+Fast-dLLM 準確率為 78.5%,僅比基線低 0.8%,但吞吐量提升 8.1 倍(圖 5)。
- HumanEval(代碼生成):準確率達 44.5%,較基線提升 1.2%,同時吞吐量提升 3.7 倍。
- 多模型兼容:在 LLaDA、Dream 等模型上均實現高效加速,驗證了技術的通用性。
三、 應用價值
無需訓練的即插即用方案
Fast-dLLM 的零訓練成本特性使其成為理想的推理優化工具,能夠快速集成到現有的系統中。對于那些已經在使用擴散模型的企業和開發者來說,可以在不改變模型架構和訓練流程的基礎上,直接利用 Fast-dLLM 提升推理效率,縮短長文本生成耗時,為實際部署提供更可行的方案。
四、 總結與展望
Fast-dLLM 通過分塊 KV 緩存與置信度感知并行解碼的創新組合,實現了擴散模型推理效率的跨越式提升,同時保持了生成質量的穩定性。其技術設計為非自回歸生成模型提供了新的優化范式,有望推動擴散模型在實時交互、長文本生成等場景中的廣泛應用。未來,隨著模型規模的擴大和蒸餾技術的深化,Fast-dLLM 或將進一步縮小與自回歸模型的性能差距,成為 LLM 推理加速的核心方案之一。
特別聲明:以上內容(如有圖片或視頻亦包括在內)為自媒體平臺“網易號”用戶上傳并發布,本平臺僅提供信息存儲服務。
Notice: The content above (including the pictures and videos if any) is uploaded and posted by a user of NetEase Hao, which is a social media platform and only provides information storage services.