英偉達、港大等發布創新KV緩存,實現擴散模型無訓練加速
擴散語言模型與傳統的自回歸模型(如GPT系列)逐個生成單詞或字符的方式不同,主要通過逐步去除文本中的噪聲來生成高質量的文本,能在單次迭代中生成多個單詞,生成效率也就更高。
但在實際應用中多數開源的擴散語言模型在推理效率上往往不如自回歸模型,主要是由缺乏鍵值(KV)緩存支持和并行解碼時生成質量下降兩大原因造成。
所以,英偉達、香港大學和麻省理工的研究人員聯合提出了Fast-dLLM來解決這些難題。

KV緩存是自回歸模型中用于加速推理的關鍵技術。通過存儲和重用先前計算的注意力狀態,減少了重復計算,從而顯著提高了生成速度。而擴散語言模型由于是雙向注意力機制,直接應用KV緩存并不容易。
在Fast-dLLM架構中,文本生成過程被劃分為多個塊,每個塊包含一定數量的token。這種塊狀生成方式的核心在于,允許模型在生成一個塊之前,預先計算并存儲其他塊的KV緩存。
在生成過程中,模型會重用這些緩存的KV激活,從而避免了重復計算。這種近似KV緩存機制的關鍵在于,它利用了KV激活在相鄰推理步驟中的高度相似性。

例如,在生成一個塊之前,模型會針對其他塊展開計算,并存儲對應的 KV Cache,以便后續步驟能夠進行重用。當完成一個塊的生成后,再重新對所有塊的KV Cache 予以計算。這種以塊為單位的生成方式,在保障模型性能基本穩定的同時,成功減少了大量冗余計算。
盡管KV緩存機制顯著提升了擴散語言模型的推理速度,但在并行解碼時,生成質量往往會下降。這是因為擴散型模型在解碼時假設了條件獨立性,而實際上,標記之間可能存在復雜的依賴關系。這種依賴關系的破壞會導致生成文本的連貫性和準確性下降。

為了解決這一問題,Fast-dLLM提出了一種基于置信度的平行解碼策略。Fast-dLLM在每個解碼步驟中計算每個標記的置信度,并選擇置信度超過閾值的標記進行解碼。如果沒有標記的置信度超過閾值,則模型會選擇置信度最高的標記進行解碼,以確保解碼過程能夠繼續進行。
這種基于置信度的解碼策略的關鍵在于,它能夠在高置信度的情況下,安全地進行并行解碼,而不會引入過多的錯誤。
為了測試dLLM性能,研究人員使用NVIDIA A100 80GB GPU對LLaDA和Dream兩種擴散語言模型進行全面評估,測試基準包括GSM8K、MATH、HumanEval和MBPP四個數據集,涵蓋數學推理與代碼生成等任務。
在KV緩存機制測試中,塊大小設為4至32。結果顯示,塊大小32時吞吐量與準確性最佳,吞吐量達54.4 tokens/s,準確率78.5%;而塊大小8時吞吐量降至49.2 tokens/s,準確率略升至79.3%,但資源消耗較大。塊大小64則因上下文不匹配導致準確率下降。
并行解碼測試采用置信度閾值0.5至1.0,默認值0.9。

動態閾值策略優于固定token數基線。以GSM8K(5-shot)為例,閾值0.9時每步處理2個token,準確率78.5%;固定解碼每步2 token時準確率79.2%。降低閾值至0.7時,每步處理4 token,準確率仍達79.3%,顯示策略靈活性強。
整體吞吐量方面,LLaDA模型GSM8K(5-shot,生成長度256)任務中,僅用KV Cache加速3.2倍至21.2 tokens/s,并行解碼加速2.5倍至16.5 tokens/s,二者結合加速8.1倍至54.4 tokens/s。生成長度1024時端到端加速達27.6倍。

Dream模型MBPP(3-shot,生成長度512)任務中,結合策略加速7.8倍至73.6 tokens/s;GSM8K(5-shot,生成長度512)加速5.6倍至42.9 tokens/s。
準確性方面,LLaDA任務中準確率僅下降0.8個百分點至78.5%;Dream的HumanEval任務準確率反升至54.3%。所有測試準確率波動均在1-2個百分點內,表明Fast-dLLM在加速同時有效保持生成質量。




































