TOKENSWIFT團隊 投稿
量子位 | 公眾號 QbitAI
大語言模型長序列文字生成效率新突破——
生成10萬Token的文字，傳統自迴歸模型需要近5個小時，現在僅需90分鐘！
最新研究提出了一個名為TOKENSWIFT框架，從模型載入、KV快取管理到Token生成策略進行了全方位的最佳化。
實驗結果證明，該方法不僅能大幅提升生成效率，更在保證生成質量和多樣性上實現了無損加速。
而且支援R1-Distill，團隊釋出經過微調的DeepSeek-R1-Distill-Qwen-32B模型，同樣具備3倍加速效果。
來看demo展示：
本研究由來自北京通用人工智慧研究院的團隊完成，以下是更多細節。
TOKENSWIFT框架長啥樣？
隨著LLMs長上下文視窗能力的不斷提升，複雜任務對超長文字生成的需求越來越高。傳統的自迴歸（AR）生成方式雖然在短文字上表現良好，但在長文字生成中存在明顯瓶頸，主要體現在以下三個方面：
-
模型頻繁載入問題
由於自迴歸生成每生成一個Token都需要從GPU儲存中重新載入模型權重，導致I/O操作頻繁、延遲高。在生成10萬Token時，模型需要重複載入上萬次，嚴重拖慢整體生成速度。
-
KV緩存的動態管理
生成超長文字過程中，模型內部的鍵值對（KV Cache）不斷增長，若直接使用全量KV快取，不僅超出記憶體預算，還會大幅增加計算時間。如何在保證關鍵資訊不丟失的前提下，實現KV快取的高效更新成為一大難題。
-
重複性生成
長序列生成易出現重複和冗餘問題，影響文字的多樣性和質量。雖然重複問題並非論文的主要聚焦點，但在超長文字生成中依然需要有效抑制。
為解決上述難題，論文提出了TOKENSWIFT——一個全新的框架，旨在實現無損加速超長序列生成，其主要創新點體現在以下幾個方面：
1）多Token並行生成與Token複用
論文借鑑了Medusa等方法，透過引入額外的線性層，使模型在一次前向傳播中能夠同時生成多個草稿Token。
更重要的是，基於生成文字中的n-gram頻率資訊，系統會自動檢索並複用高頻短語，從而進一步減少模型重新載入的次數，提升整體效率。
2）動態KV快取更新策略
在KV快取管理上，TOKENSWIFT採用動態更新策略。系統在生成過程中將初始KV快取保留，同時根據Token的重要性對後續快取進行有序替換。
這種方式不僅有效控制了快取的規模，還確保了關鍵資訊始終被儲存，大幅降低了因快取載入帶來的延遲。
3）基於樹結構的多候選Token驗證
為保證生成結果與目標模型預測的一致性，TOKENSWIFT引入了樹形注意力機制。
透過構建包含多個候選Token組合的樹形結構，並採用並行驗證的方式，從中隨機選擇最長且有效的n-gram作為最終輸出，確保生成過程無損且多樣性得到提升。
4）上下文懲罰策略
為了進一步抑制重複生成問題，論文設計了一種上下文懲罰方法。該方法在生成過程中為近期生成的Token施加懲罰，使得模型在選擇下一Token時更傾向於多樣化輸出，從而有效減少重複現象。
TOKENSWIFT效果如何？
實驗部分，論文在多種模型架構（包括MHA和GQA）及不同規模（1.5B、7B、8B、14B）上進行了充分測試。
結果表明，TOKENSWIFT在生成10萬Token長序列時，相較於傳統自迴歸方法，平均實現了3倍以上的加速，且生成結果在準確性和多樣性上基本保持無損。
1）加速效果
實驗資料顯示，在LLaMA3.1-8B模型下，傳統AR生成10萬Token約需4.9小時，而使用TOKENSWIFT後僅需90分鐘，大幅節省時間。在Qwen2.5-14B時，傳統AR生成10萬Token更是達到了7.9小時，加速後僅需142分鐘。這一成果對於實際應用中需要即時或高效長文字生成的場景具有重要意義。
2）驗證率與接受率
論文設計了多項指標來評估生成質量，包括Token接受率和Distinct-n指標。結果表明，TOKENSWIFT不僅在速度上顯著領先，還能在保持無損生成的前提下，有效提升文字的多樣性。
消融實驗與案例分析
在深入理解TOKENSWIFT各模組貢獻的過程中，論文還進行了全面的消融實驗和案例分析，為最佳化方案提供了充分依據。
消融實驗：關鍵元件的作用
-
Token複用消融：實驗中將Token複用引數設為0（即不復用），結果顯示，接受率和生成速度均顯著下降，表明複用機制在減少模型載入次數、提升整體效率方面發揮了關鍵作用。
-
KV快取更新策略對比：研究對比了全量快取、一次性更新和動態更新三種方案。結果表明，全量快取雖然在接受率上略佔優勢，但其高計算開銷使得整體加速效果不理想；而一次性更新則因快取膨脹導致效能下降。動態更新策略則在保持高接受率的同時，實現了最佳的速度與資源平衡。
-
上下文懲罰效果：在不同取樣方法下，加入上下文懲罰後，生成文字的Distinct-n指標明顯提高。
例如，在min-p取樣場景下，Distinct-n平均得分從0.12提升至0.69，僅帶來約8%的速度損失，充分驗證了該策略在抑制重複生成方面的有效性。
案例分析：真實生成對比
論文還對比了在有無上下文懲罰條件下生成文字的差異，案例分析結果令人印象深刻：
-
重複現象的抑制：在未使用上下文懲罰時，生成文字在大約5K Token處就出現明顯重複，且多為逐字重複；而採用上下文懲罰後，重複問題明顯延遲至約60K Token，且重複部分多表現為語義層次上的相似，而非直接複製，整體文字連貫性和可讀性大幅改善。
-
文字質量的提升：案例對比顯示，使用完整TOKENSWIFT流程的生成結果在邏輯連貫性、表達多樣性和創新性方面均優於傳統方法，為實際應用提供了更高質量的文字輸出。
透過這些消融實驗和案例分析，論文不僅證明了各關鍵技術模組的重要性，也為後續最佳化指明瞭方向，充分體現了TOKENSWIFT在超長文字生成領域的先進性和實用性。
Arxiv：https://arxiv.org/abs/2502.18890
Github：https://github.com/bigai-nlco/TokenSwift
Blog：https://bigai-nlco.github.io/TokenSwift/
Github：https://github.com/bigai-nlco/TokenSwift
Blog：https://bigai-nlco.github.io/TokenSwift/
學術投稿請於工作日發郵件到：
標題註明【投稿】，告訴我們：
你是誰，從哪來，投稿內容
附上論文/專案主頁連結，以及聯絡方式哦
我們會（儘量）及時回覆你
一鍵關注 👇 點亮星標
一鍵三連「點贊」「轉發」「小心心」
歡迎在評論區留下你的想法！