90分鍾生成10萬Token，新框架實現3倍無損加速超長文本生成，支持DeepSeek-R1和QwQ！

今天小編分享的科學經驗：90分鍾生成10萬Token，新框架實現3倍無損加速超長文本生成，支持DeepSeek-R1和QwQ！，歡迎閲讀。

大語言模型長序列文本生成效率新突破——

生成 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 不僅在速度上顯著領先，還能在保持無損生成的前提下，有效提升文本的多樣性。