DepCap: Adaptive Block-Wise Parallel Decoding for Efficient Diffusion LM Inference

Xiang Xia, Wuyang Zhang, Jiazheng Liu, Cheng Yan, Yanyong Zhang

View Original ↗
AI 導讀 technology AI 重要性 4/5

DepCap 免訓練框架動態劃分區塊並偵測運算衝突,為 LLaDA-1.5 模型創下 5.63 倍加速並提升 7.4% 準確率。

  • 捨棄固定區塊大小,以跨步 KL 散度與預測熵動態決定 8 到 128 個 token 的區塊邊界。
  • 計算 token 之間的聯合對數機率互斥分數,精準排除會破壞語意的衝突平行運算。
  • 無需重新訓練即可部署,與現存 Prefix Cache 等快取技術完全相容並能疊加吞吐量效能。

擴散語言模型(DLMs)具備全域修改與平行生成的優勢,卻深陷解碼速度與品質的拉扯。來自中國科學技術大學的研究團隊提出免訓練推理框架 DepCap,不僅在 LLaDA 系列模型平均達成 3.57 倍速度躍升,更在 MBPP 程式碼基準測試中,為 LLaDA-1.5 帶來 5.63 倍的極致加速,同時將相對準確率逆勢推升 7.4%

擴散語言模型推理面臨的區塊邊界與平行解碼挑戰

擴散語言模型透過反覆的迭代降噪程序產生文字,有別於傳統自迴歸模型嚴格由左至右的限制,能夠運用雙向注意力機制在推理過程中平行更新多個位置的 token(文字處理的基本單位)。這種非序列式的生成範式為複雜推論與程式碼撰寫任務帶來極大的靈活性,卻也引發了龐大的運算開銷,使得如何兼顧生成品質與解碼速度成為核心難題。

為縮減無謂的運算負擔,近期的研究轉向採用區塊式解碼策略。系統會將整段待生成的文字序列分割成多個不重疊的區塊,逐塊完成擴散降噪,避免在每一步都對整條超長序列進行無效的全域微調。這種做法不僅提升了推理的實用性,更替後續的記憶體快取技術打開了大門。

然而,決定下一個區塊的邊界位置與挑選區塊內可平行處理的 token 成為兩大技術障礙。現行方法多半依賴固定的區塊排程,或是只觀察當前步驟的語意結構(例如遇到換行符號即截斷)。在挑選平行解碼的 token 時,傳統策略又過度依賴單一位置的預測信心度,為了防範生成品質崩壞而設定極度保守的門檻,導致系統平白錯失大量平行運算的空間。

DepGA-Block 引入跨步訊號動態劃分 8 到 128 區塊大小

針對區塊邊界的決策,研究團隊設計了依賴引導的自適應區塊劃分策略(DepGA-Block)。這項機制徹底捨棄死板的固定大小與局部啟發式規則,轉而關注擴散降噪過程中的跨步訊號。具體而言,系統會動態評估「剛解碼完成的最後一個區塊」對尚未生成的未來位置產生了多大的實質支撐力。

在運算層面上,DepGA-Block 會捕捉加入上一個區塊前後的邊際預測分佈變化,並使用 KL 散度(衡量兩種機率分佈差異的統計指標)來量化這份影響力。當影響力極大時,代表前文提供了強烈的跨步資訊支撐;同時,系統也會計算預測分佈的 Shannon 熵(衡量不確定性的數學工具),若熵值過高,則代表該位置的預測依然模糊不清。

框架將這兩個指標結合成專屬的「依賴分數」。演算法會從當前解碼前緣向後掃描,直到發現第一個不確定性超越前區塊支撐力的位置,便在該處設下明確的區塊邊界。為了確保推論過程的穩定性,DepCap 設有 8 個 token 的最小冷啟動長度,並將單一區塊的最大長度動態限制在 128 個 token。

資訊理論支撐條件互訊息的局部重疊假設

DepGA-Block 的依賴分數設計並非僅憑直覺,其背後有著嚴謹的資訊理論基礎。研究分析指出,對於較短的局部候選區塊,前一個已解碼區塊提供給各個未來位置的資訊量,其重疊部分其實相當有限,這在理論上被稱為局部重疊假設。

在這種合理的數學假設下,前一個解碼區塊對整個候選區塊的累積影響力,可以被近似為個別 token 影響力的直接加總。此累積影響力在數學公式中等同於條件互訊息的期望值,反映了已知上下文對未知內容的約束程度。

為了讓前文提供的支援能夠穩壓未來的累積不確定性,DepCap 實作中提取了單次採樣的 KL 散度減去帶權重的預測熵,形成一個計算輕量且高度貼合機率分佈變化的線上決策指標,確保了區塊動態擴展具備紮實的數學合理性。

CAP-Decoding 透過互斥分數篩選安全的平行生成子集

選定動態區塊後,系統還必須決定該區塊內有哪些位置可以安全地同步降噪。傳統基於信心度的過濾機制,雖然能篩掉低機率的候選者,卻無法察覺 token 之間的潛在衝突。強行同步解碼彼此衝突的 token,會導致整段語意嚴重碎裂甚至失去邏輯。

為解決此問題,框架內建了衝突感知的平行解碼策略(CAP-Decoding)。系統會替候選池中任意兩個 token 計算互斥分數(基於聯合對數機率)。當分數超過預設的 -16.0 門檻時,代表這兩個位置的機率分佈高度糾纏、極易互相干擾,系統便會將其標記為互斥對象,嚴格禁止它們在同一回合被運算。

CAP-Decoding 的篩選作業分為兩大階段:首先是「高信心優先」階段,直接挑選機率大於 0.95 的絕對安全候選者,並剔除與其衝突的對象;接著進入「貪婪補全」階段,依據機率由高至低依序遞補剩下的候選者,每選中一個就移除周邊的衝突者。透過顯式的衝突排查,模型能在不傷害語意品質的前提下,大幅增加每一回合被平行消化的文字數量。

LLaDA-1.5 在 MBPP 測試達成 5.63 倍加速與準確率提升

研究團隊在 LLaDA-8B-InstructDream-v0-base-7B 以及最新的 LLaDA-1.5 等三款主流模型上部署 DepCap,並橫跨數學推理(GSM8K、Math-500)與程式碼生成(MBPP、HumanEval)四大基準測試進行驗證。數據顯示,與固定的區塊排程相比,自適應機制能更加穩定地兼顧高文字吞吐量與低運算步數。

在實際硬體表現上,結合兩大策略的完整 DepCap 展現出極致的推論效能。在 MBPP 基準測試中,搭載於 LLaDA-1.5 的框架飆出 5.63 倍的推理加速,由於更合理的區塊邊界分配,模型的相對準確率不僅沒有下降,反而逆勢推升了 7.4%。在 LLaDA 系列的八組跨任務測試平均中,也取得 3.57 倍的優異加速成績。

除了機制本身的優越性,DepCap 的隨插即用特性更是一大亮點。在整合 Fast-dLLM 所使用的 Prefix Cache(快取模型內部運算狀態以加速推理的技術)與 Dual Cache 後,吞吐量獲得進一步的疊加放大,在 GSM8K 測試中達到每秒 51.2 個 token 的生成速度與 79.3 的高準確率,證明該框架能完美相容於現存的軟體基礎架構。

透過動態評估跨步訊號與精確的衝突迴避過濾,免訓練框架 DepCap 成功釋放了擴散語言模型的平行解碼潛力,為非序列式生成技術的大規模部署鋪平了道路。

Abstract

Diffusion language models (DLMs) have emerged as a promising alternative to autoregressive language generation due to their potential for parallel decoding and global refinement of the entire sequence. To unlock this potential, DLM inference must carefully balance generation quality and decoding speed. Recent block-wise DLM decoding methods improve this trade-off by performing diffusion-based decoding sequentially in blocks. However, existing methods typically rely on fixed block schedules or current-step local signals to determine block boundaries, and use conservative confidence-based parallel decoding to avoid conflicts, limiting the quality-speed trade-off. In this paper, we argue that block-wise DLM inference requires more suitable signals for its two core decisions: cross-step signals for determining block boundaries, and token-level conflict signals for parallel decoding. Based on this view, we propose DepCap, a training-free framework for efficient block-wise DLM inference. Specifically, DepCap instantiates the cross-step signal as the influence of the last decoded block and uses it to adaptively determine how far the next block should extend, while identifying a conflict-free subset of tokens for safe parallel decoding within each block, enabling substantial inference acceleration with negligible quality degradation. DepCap is a plug-and-play method applicable to various DLMs, and compatible with existing KV-cache strategies for block-wise DLM. An information-theoretic analysis further suggests that the cumulative last-block influence on a candidate block is approximately additive across tokens, supporting the proposed block-partitioning criterion. Experimental results show that DepCap achieves favorable speed-quality trade-offs across multiple DLM backbones and reasoning and coding benchmarks, with up to 5.63$\times$ speedup without significant performance degradation.