我們介紹一致性擴散語言模型（CDLM），透過結合基於一致性的多詞元最終確定與區塊式 KV 快取，將擴散語言模型推理速度提升高達 14.5 倍，適用於數學和程式碼任務。

擴散語言模型（DLM）正成為自迴歸（AR）語言模型的有潛力替代方案。DLM 並非一次產生一個詞元，而是透過多個採樣步驟迭代地精煉部分遮蔽的序列，逐步將完全遮蔽的序列轉換為乾淨的文本。這個精煉過程創造了一個引人注目的機會：它支援平行生成，讓模型每次迭代可以最終確定多個詞元，潛在地實現比 AR 解碼更高的吞吐量。同時，它還可以利用雙向上下文來解鎖文本填補和精煉等新功能。

然而，在實際應用中，標準 DLM 存在兩項主要效率低下的問題。

CDLM 透過一種後訓練方法來解決這兩個瓶頸，該方法使較少步驟的推理變得可靠，同時實現精確的區塊式 KV 快取。

DLM 生成是一個在 N 個離散採樣步驟上的迭代精煉過程。它將時間 t=1 時的完全遮蔽序列轉換為 t=0 時的乾淨序列。在每個步驟，模型根據當前雜訊序列 xt 和提示 c，預測一個乾淨序列分佈 x0：

$p_{ heta}(\mathbf{X}_0 \mid \mathbf{X}_t, c)$

一種常見的確定性實例是低信賴度重新遮蔽：模型貪婪地取消遮蔽詞元（通常在區塊內），最終確定最高信賴度的遮蔽位置，同時保持其他位置的遮蔽狀態。這導致了解碼軌跡：

$\mathcal{T}_{\mathbf{x}} = \left(\mathbf{x}_{t_0}, \mathbf{x}_{t_1}, \ldots, \mathbf{x}_{t_N} ight), \quad t_k = 1 - rac{k}{N}$

該軌跡記錄了部分精煉序列如何逐步演變。這個軌跡成為 CDLM 訓練的核心對象。

我們透過在特定領域的提示上運行 DLM 推理來離線收集軌跡。對於每個提示 x，我們記錄詞元級解碼軌跡 T_x、包含詞元最終確定時刻最後一層隱藏狀態的緊湊隱藏狀態緩衝區 H_x，以及真實文本 ŷ。具體來說，我們採用區塊式解碼，生成長度 L_g = 256，區塊大小 B = 32，以及總共 N = L_g 個步驟（即在當前區塊內每步最終確定一個詞元）。這種保守的設定產生了更高品質的軌跡用於蒸餾。

在軌跡提取過程中，我們使用完整的雙向注意力遮罩。相比之下，在訓練 CDLM 時，我們採用區塊式的因果遮罩，該遮罩關注提示、先前已完成的區塊以及當前解碼區塊。此設計使模型能夠從完全雙向切換到區塊式擴散模型（如 [2]），從而為已最終確定的區塊實現精確的區塊式 KV 快取。

CDLM 聯合最小化三個目標：

(i) 蒸餾損失（新取消遮蔽的位置）

對於在中間狀態 y 和其區塊完成 y* 之間新取消遮蔽的位置，我們將學生的預測分佈與從儲存的隱藏狀態獲得的教師重建分佈進行匹配。

直覺：此目標作為主要錨點，教導學生在區塊因果約束下最終確定區塊內的詞元。

(ii) 一致性損失（仍被遮蔽的位置）

我們透過對學生在狀態 y 的預測與其在更具資訊的狀態 y* 下對仍被遮蔽位置的預測進行對齊來強制執行區塊內的時間一致性，並使用 stop-gradient 目標。

直覺：此目標鼓勵了解碼軌跡上的穩定多步驟轉換。

(iii) 輔助 DLM 遮蔽去噪損失

我們包含一個標準的遮蔽去噪目標，應用於隨機遮蔽的真實文本。

直覺：此目標保留了模型通用的遮蔽詞元預測能力，並有助於保留推理行為，尤其是在數學任務上。

在推理時，CDLM 以區塊式的自迴歸方式進行解碼，重複使用提示和所有先前已最終確定區塊的 KV 快取。在每個區塊內，我們應用信賴度閾值化的平行最終確定。[3] 我們還採用提前停止策略，一旦在當前區塊中出現結束文本詞元。

我們故意避免引入額外超參數的額外啟發式方法（例如，具有任務相關設定的區塊間平行處理），而是專注於基於精確 KV 快取和可靠步驟減少的穩健預設解碼流程。

有效的步驟減少：為什麼訓練很重要

naively 截斷步驟數會導致明顯的準確度下降，而 CDLM 在類似的步驟預算下保持了品質（並且由於快取，延遲約減半）。這凸顯了核心要點：穩定的多詞元精煉並非免費；它需要強制執行軌跡一致行為的訓練。

系統級分析：為什麼區塊式 DLM 處於最佳位置

為了理解硬體利用率，我們分析了算術強度（AI），即移動每位元組的 FLOPs，隨著批次大小的增加，比較了：AR 解碼、標準（全注意力）DLM、區塊式 DLM（CDLM）與 B∈{4,16,32}。

總體而言，該分析解釋了為什麼 CDLM 類的區塊式擴散可以在小批次大小下提供強大的效率：它利用平行處理來分攤記憶體訪問，同時保持在仍然受益於實際擴展的範圍內。

DLM 中的全雙向注意力需要在每次去噪步驟中重新計算 O(L^2) 的注意力，使得推理計算量非常大。CDLM 在保留每個區塊內雙向上下文的同時，實現了精確的 KV 快取，保留了局部精煉能力（例如，在當前區塊內進行填補）。

CDLM 是一種後訓練方法，可以應用於任何區塊擴散模型，並且隨著更強大的 DLM 的出現，其優勢應該會不斷增長。一個有前景的方向是從更大、更強的 DLM 教師那裡收集軌跡，並使用 CDLM 訓練中等規模的學生。

我們提出了 CDLM，這是一種基於訓練的加速方案，將一致性建模引入 DLM。透過強制執行區塊內的時間一致性並微調區塊式的因果學生模型，CDLM 減少了精煉步驟並實現了精確的 KV 快取。在數學和程式碼任務上，CDLM 實現了更快的推理、更少的步驟、更低的延遲和更高的吞吐量，同時保持了具有競爭力的準確度。