Zyphra釋出ZAYA1-8B-Diffusion-Preview：首個從自迴歸LLM轉換的MoE擴散模型，速度提升達7.7倍

Zyphra，這家總部位於舊金山的AI實驗室，釋出了ZAYA1-8B-Diffusion-Preview，這是其擴散語言模型早期工作的預覽版。該釋出表明，現有的自迴歸語言模型可以轉換為離散擴散模型，而不會造成評估效能的系統性損失，同時在AMD硬體上實現了顯著的推理加速。

自迴歸解碼的問題 要理解這一突破的重要性，首先需要了解大多數語言模型目前生成文本的方式。標準大型語言模型是自迴歸的：它們按順序逐個解碼token。對於每個新token，注意力機制必須回顧所有之前生成的token，並從GPU記憶體載入它們的儲存表示（即KV快取）。關鍵問題是，由於批次中每個使用者的token歷史不同，每個使用者的KV快取必須單獨載入，無法在請求之間共享。這造成了瓶頸。當GPU花費更多時間從記憶體行動數據而非執行實際計算時，系統便受限於記憶體頻寬而非計算能力。這限制了現代GPU硬體（其計算FLOPs增長速度超過記憶體頻寬）在推理過程中的高效利用。

擴散模型提供了另一種方案。擴散模型不是逐個生成token，而是同時生成N個token的多個草稿，並多次迭代這一草稿過程。由於塊中的所有N個token共享相同的KV快取，操作從記憶體頻寬受限轉變為計算受限，這意味著GPU可以更高效地被利用。在ZAYA1-8B-Diffusion-Preview中，模型對塊中的每個token執行從掩碼到token的單步變換——即直接預測未被掩碼的token，而非迭代去噪。

從自迴歸到擴散的轉換：無需從頭訓練 從頭訓練擴散語言模型在技術上很困難，且已知的成熟方法很少。Zyphra團隊給出了選擇轉換而非從頭訓練的兩個理由：首先，從頭訓練本身就很難，已知配方極少；其次，以擴散模式訓練並無優勢，因為訓練本身已是計算受限——擴散解決的頻寬瓶頸僅出現在推理階段。這意味著擴散的所有優勢都是推理時的優勢，現有的預訓練堆疊可以原樣複用。

基於TiDAR配方，Zyphra以ZAYA1-8B-base檢查點為起點，額外進行了6000億token的擴散轉換中訓練（上下文長度32k），隨後進行了5000億token的原生上下文擴充套件至128k，最後進行了擴散監督微調（SFT）階段。

ZAYA1-8B-Diffusion-Preview是首個從自迴歸LLM轉換的MoE擴散模型，也是首個在AMD GPU上訓練的擴散語言模型。Zyphra報告稱，與基礎自迴歸檢查點相比，評估效能幾乎沒有下降，在某些基準（如LCB-v6）上甚至有所提升。他們將此部分歸因於改進的中訓練資料集，以及擴散式塊內非因果推理相比因果自迴歸具有更強的表達能力。

擴散取樣器的工作原理 推理過程中，ZAYA1-8B-Diffusion-Preview同時生成16個token的草稿。基於從推測解碼借鑑的取樣標準，接受其中一部分token。關鍵在於，同一個模型在前向傳播中同時扮演推測器和驗證器的角色，這消除了傳統方法（如EAGLE或dFlash）中執行兩個獨立模型的開銷。在高度記憶體頻寬受限的情況下，幾乎所有被接受的token都代表了對自迴歸解碼的免費加速——GPU已經載入完畢，額外token的計算成本極低。

Zyphra團隊報告了兩種具有不同速度-質量權衡的取樣器：

• 無損擴散取樣器：使用標準推測解碼接受標準min(1, p(x)/q(x))，其中p是自迴歸模型的logit分佈，q是擴散模型的分佈。拒絕時，下一個token從p(x)-q(x)的殘差分佈中取樣。該取樣器實現4.6倍加速，且沒有系統性的評估效能下降。
• Logit-mixing取樣器：首先混合擴散推測器和自迴歸模型的logit，然後使用平均分佈進行驗證。這提高了接受率，因為驗證logit更接近擴散logit，但對質量有一定影響。該取樣器實現7.7倍加速。速度與質量的權衡可在執行時選擇。

需要注意的是，由於ZAYA1-8B-Diffusion-Preview是未經RL訓練的基礎中訓練檢查點，Zyphra使用pass@評估而非標準準確率基準，以更好地代表模型在RL訓練後的最終潛力。讀者在將這些數字與其他模型報告的基準進行比較時應牢記這一點。

Zyphra團隊還指出，擴散帶來的加速高於多token預測（MTP）和EAGLE3等推測解碼策略的替代方法。由於TiDAR風格的擴散模型僅使用單次前向傳播，即使接受率與dFlash相當，也能實現顯著的加速。

架構細節 ZAYA1-8B-Diffusion-Preview是一種單步推測擴散模型，採用順序約束生成，即擴散模型僅能生成從字首開始的連續子序列。與無約束掩碼擴散目標或集合塊解碼相比，這一約束極大地提高了訓練穩定性，這也是Zyphra基於TiDAR配方的主要原因。

該模型使用Zyphra的CCA注意力變體。CCA大幅減少了注意力中的預填充FLOPs，這對擴散直接有利，因為擴散將解碼轉變為類似預填充的操作。這意味著CCA允許模型在達到計算限制之前並行處理更多token。

具體來說，架構採用CCGQA，查詢頭與鍵頭的比例為4:1。設計選擇之一是故意避免MLA（多頭潛在注意力），因為其高算術強度與CCGQA不匹配。由於塊擴散訪問相同的快取，算術強度隨塊大小和每次前向傳播的塊數量而擴充套件。在AMD MI300x硬體上（bf16），系統每次前向傳播大約支援三個塊大小的提議；在MI355x上，這一數字上升到約五個。CCGQA還以2倍壓縮執行，這使得Zyphra能夠負擔與TiDAR中訓練相關的額外訓練FLOPs。AMD GPU硬體更大的VRAM容量進一步使得擴散訓練整體上更高效。

在實踐中，實現理論加速更具挑戰性，因為擴散帶來了額外的操作開銷，且擴散模型的推理堆疊遠不如自迴歸推理的成熟工具那樣最佳化。