英偉達(dá)、港大等發(fā)布創(chuàng)新KV緩存，實現(xiàn)擴(kuò)散模型無訓(xùn)練加速
擴(kuò)散語言模型與傳統(tǒng)的自回歸模型（如GPT系列）逐個生成單詞或字符的方式不同，主要通過逐步去除文本中的噪聲來生成高質(zhì)量的文本，能在單次迭代中生成多個單詞，生成效率也就更高。
但在實際應(yīng)用中多數(shù)開源的擴(kuò)散語言模型在推理效率上往往不如自回歸模型，主要是由缺乏鍵值（KV）緩存支持和并行解碼時生成質(zhì)量下降兩大原因造成。
所以，英偉達(dá)、香港大學(xué)和麻省理工的研究人員聯(lián)合提出了Fast-dLLM來解決這些難題。
KV緩存是自回歸模型中用于加速推理的關(guān)鍵技術(shù)。通過存儲和重用先前計算的注意力狀態(tài)，減少了重復(fù)計算，從而顯著提高了生成速度。而擴(kuò)散語言模型由于是雙向注意力機(jī)制，直接應(yīng)用KV緩存并不容易。
在Fast-dLLM架構(gòu)中，文本生成過程被劃分為多個塊，每個塊包含一定數(shù)量的token。這種塊狀生成方式的核心在于，允許模型在生成一個塊之前，預(yù)先計算并存儲其他塊的KV緩存。
在生成過程中，模型會重用這些緩存的KV激活，從而避免了重復(fù)計算。這種近似KV緩存機(jī)制的關(guān)鍵在于，它利用了KV激活在相鄰?fù)评聿襟E中的高度相似性。
例如，在生成一個塊之前，模型會針對其他塊展開計算，并存儲對應(yīng)的 KV Cache，以便后續(xù)步驟能夠進(jìn)行重用。當(dāng)完成一個塊的生成后，再重新對所有塊的KV Cache 予以計算。這種以塊為單位的生成方式，在保障模型性能基本穩(wěn)定的同時，成功減少了大量冗余計算。
盡管KV緩存機(jī)制顯著提升了擴(kuò)散語言模型的推理速度，但在并行解碼時，生成質(zhì)量往往會下降。這是因為擴(kuò)散型模型在解碼時假設(shè)了條件獨立性，而實際上，標(biāo)記之間可能存在復(fù)雜的依賴關(guān)系。這種依賴關(guān)系的破壞會導(dǎo)致生成文本的連貫性和準(zhǔn)確性下降。
為了解決這一問題，F(xiàn)ast-dLLM提出了一種基于置信度的平行解碼策略。Fast-dLLM在每個解碼步驟中計算每個標(biāo)記的置信度，并選擇置信度超過閾值的標(biāo)記進(jìn)行解碼。如果沒有標(biāo)記的置信度超過閾值，則模型會選擇置信度最高的標(biāo)記進(jìn)行解碼，以確保解碼過程能夠繼續(xù)進(jìn)行。
這種基于置信度的解碼策略的關(guān)鍵在于，它能夠在高置信度的情況下，安全地進(jìn)行并行解碼，而不會引入過多的錯誤。
為了測試dLLM性能，研究人員使用NVIDIA A100 80GB GPU對LLaDA和Dream兩種擴(kuò)散語言模型進(jìn)行全面評估，測試基準(zhǔn)包括GSM8K、MATH、HumanEval和MBPP四個數(shù)據(jù)集，涵蓋數(shù)學(xué)推理與代碼生成等任務(wù)。
在KV緩存機(jī)制測試中，塊大小設(shè)為4至32。結(jié)果顯示，塊大小32時吞吐量與準(zhǔn)確性最佳，吞吐量達(dá)54.4 tokens/s，準(zhǔn)確率78.5%；而塊大小8時吞吐量降至49.2 tokens/s，準(zhǔn)確率略升至79.3%，但資源消耗較大。塊大小64則因上下文不匹配導(dǎo)致準(zhǔn)確率下降。
并行解碼測試采用置信度閾值0.5至1.0，默認(rèn)值0.9。
動態(tài)閾值策略優(yōu)于固定token數(shù)基線。以GSM8K（5-shot）為例，閾值0.9時每步處理2個token，準(zhǔn)確率78.5%；固定解碼每步2 token時準(zhǔn)確率79.2%。降低閾值至0.7時，每步處理4 token，準(zhǔn)確率仍達(dá)79.3%，顯示策略靈活性強(qiáng)。
整體吞吐量方面，LLaDA模型GSM8K（5-shot，生成長度256）任務(wù)中，僅用KV Cache加速3.2倍至21.2 tokens/s，并行解碼加速2.5倍至16.5 tokens/s，二者結(jié)合加速8.1倍至54.4 tokens/s。生成長度1024時端到端加速達(dá)27.6倍。
Dream模型MBPP（3-shot，生成長度512）任務(wù)中，結(jié)合策略加速7.8倍至73.6 tokens/s；GSM8K（5-shot，生成長度512）加速5.6倍至42.9 tokens/s。
準(zhǔn)確性方面，LLaDA任務(wù)中準(zhǔn)確率僅下降0.8個百分點至78.5%；Dream的HumanEval任務(wù)準(zhǔn)確率反升至54.3%。所有測試準(zhǔn)確率波動均在1-2個百分點內(nèi)，表明Fast-dLLM在加速同時有效保持生成質(zhì)量。
