AI News HubLIVE
站內改寫2 分鐘閱讀

雙流Transformer:將主預填充路徑與額外解碼計算解耦

Dual-Flow Transformer 將主流程與輔助流程解耦,輔助流僅在解碼階段從最後一個提示位置開始激活,在不寫入持久 KV 緩存的情況下增加續寫計算。實驗顯示在匹配 token 數量下驗證損失更低,且在 MoE 模型中可按階段獨立分配專家計算。

來源arXiv AI作者: Liming Liu, Mingze Wang, Tuo Zhao

隨着大語言模型(LLM)在實際服務中面對越來越多的請求,推理成本正在從一次性訓練成本之外的次要因素,逐漸變成影響總體擁有成本的關鍵變量。論文指出,推理的兩個階段對底層硬件提出了差異很大的要求。提示預填充(prefill)階段需要處理整個輸入上下文,擁有很高的並行度,通常受限於計算吞吐;自迴歸解碼階段則是逐個 token 生成,本質上是一個串行過程,在 KV 緩存變大後常常受制於內存帶寬。傳統的模型擴展方式,無論是增加寬度還是深度,都會讓每一層新增的參數同時參與預填充和解碼,因此兩個階段的成本會被同步推高。

在這篇論文中,Liming Liu、Mingze Wang 和 Tuo Zhao 提出了 Dual-Flow Transformer,目的是在保持既有預填充計算和單一持久 KV 緩存的同時,把額外的可學習計算集中在續寫預測階段。該架構將模型分成兩條流。主流程(primary flow)是一個完整的因果語言模型,負責讀取整個提示並像標準 Transformer 一樣寫入 KV 緩存。輔助流程(auxiliary flow)在提示處理階段完全不參與計算,只有生成到最後一個提示位置之後才被激活。這樣,輔助流程可以在不修改持久狀態的情況下增加後續 token 預測的建模能力,也不會對主流程的表示產生干擾。

兩條流程共享大部分注意力層、MLP 和輸出矩陣,但同時使用獨立的 token 嵌入和輕量級的耦合方式。這種設計既控制了新增參數量,也帶來了執行層面的複用好處:在批量或分組解碼時,權重已經被加載到顯存/內存中,主流程 KV 緩存裏的鍵和值也可以被輔助流程重複使用。因此,增加輔助計算不一定意味着成比例地增加訪存壓力。

為了驗證思路,作者進行了匹配 token 數量的對比實驗。也就是説,在比較 Dual-Flow 與基線模型時,讓它們處理相同數量的 token,而不是簡單比較相同參數量。結果顯示,Dual-Flow 在不同架構和數據配置下都取得了更低的驗證損失,説明額外計算被更有效地分配到了續寫預測上。

論文還進一步探討了混合專家(MoE)模型中的情況。在 MoE 中,每條 token 只會激活一部分專家網絡,專家扇出(expert fan-out)決定了使用的專家數量。Dual-Flow 的分離讓主流程和輔助流程的專家扇出變成兩個相互獨立的控制旋鈕,從而可以分別調節預填充成本、續寫成本和預測質量。作者研究了兩種配置:一是在固定預填充專家計算量的前提下增加解碼階段的計算;二是在固定的解碼專家預算下,把資源在兩條流程之間重新分配。實驗揭示了預填充—解碼—質量之間的權衡關係,並展示了按推理階段單獨分配專家計算的潛力。

總的來説,Dual-Flow Transformer 為推理成本敏感的 LLM 部署提供了一個新的架構方向:不是簡單地把更多計算平攤到所有階段,而是把額外容量精確投放到最需要“思考”的續寫階段。對於長上下文、高併發服務、邊緣設備上的生成任務等場景,這種面向階段的解耦設計可能會成為模型選型和推理優化的重要參考。