ICML 2025｜如何憑「自動補全」實現100K生成3×加速？

在當前大模型推理愈發複雜的時代，如何快速、高效地產生超長文本，成為了模型部署與優化中的一大核心挑戰。隨著 GPT-o3, DeepSeek R1 等具備「超級上下文窗口」能力的大模型持續刷新業界記錄，百萬甚至千萬 Token 級別的推理任務已從研究話題邁入現實場景。然而，生成這些超長文本的背後，卻隱藏著令人咋舌的計算成本 —— 長時間的等待、巨大的內存負擔以及偶爾重覆乏味的輸出，嚴重製約了這些模型的真正潛力。

面對這一挑戰，BIGAI NLCo 團隊提出了一項全新的推理加速框架 —— TokenSwift，該工作已成功被 ICML 2025 正式接收！在這項研究中提出了一套可插拔、無損、高效的生成加速策略，專為 100K Token 級別的長文本推理而設計。在保持原始模型輸出一致性的前提下，加速比達到 3 倍以上，極大提升了推理效率。

論文標題：TokenSwift: Lossless Acceleration of Ultra Long Sequence Generation

Arxiv: https://arxiv.org/abs/2502.18890
Github: https://github.com/bigai-nlco/TokenSwift
Blog: https://bigai-nlco.github.io/TokenSwift/

🧠 重新定義超長生成：為什麼傳統方法「慢」？

為了更好地理解 TokenSwift 的意義，我們先看一下目前主流大模型（如 LLaMA、Qwen 等）在長文本生成中的瓶頸所在。

儘管這些模型具備了強大的生成長上下文的能力，但大多數依然採用傳統的自回歸（Autoregressive）生成方式：每次僅生成一個新的 Token，並以其作為輸入接著生成下一個。這種方式本身在短文本生成中問題不大，但當序列長度擴展至 10 萬甚至更多時，性能就會急劇下降。

主要原因有三：

模型重覆重載：每生成一個 Token，都會觸發一次完整的前向推理過程；在多進程或流水線執行時，模型需要不斷讀取參數，造成 I/O 瓶頸。
KV 緩存無限膨脹：Transformer 架構要求保留所有歷史 Token 的 Key/Value 信息，用於後續 Token 的注意力計算。隨著生成進程推進，KV 緩存佔用不斷增加，導致計算與內存開銷雪上加霜。
語義重覆堆疊：生成越長，模型越容易陷入句式與主題的複讀循環，降低輸出多樣性與用戶體驗。

尤其在當前日益增長的多輪對話、大模型代理（Agent）、逐步推理等任務中，一個 Query 可能會觸發幾千甚至上萬的推理 Token 輸出。傳統自回歸的效率顯然已經難以滿足需求。

🚀 TokenSwift：擁抱並行的超長推理時代

TokenSwift 的提出，正是為瞭解決上述超長生成中的三大瓶頸。它通過一個極為輕量且高效的框架，對傳統自回歸推理進行了「重構」，提出了以「多 Token 草擬 + 並行驗證 + 動態緩存更新」為核心的全新機制。

讓我們來逐步拆解 TokenSwift 的關鍵技術理念：

✳️ 多 Token 並行草擬：告別一次一 Token 的低效時代

在 TokenSwift 中，不再堅持「一步一 Token」的生成模式，而是通過對已有模型的最小化修改（添加極少量的線性層），實現了「一次性草擬多個 Token」。這意味著，模型每一次前向傳播，都可以並行生成 γ 個候選 Token，大幅降低模型重載頻率，顯著節省 I/O 時間。

更關鍵的是，草擬階段並非「胡亂猜測」。引入了上下文引導機制，使草擬結果具備較高的語義相關性與語法一致性，隨後通過結構化的驗證機制確保其與標準 AR 路徑一致。

🧩 n-gram 啟髮式補全：巧用歷史片段，精確草擬結構

為了避免粗略草擬帶來的語義偏離，TokenSwift 設計了基於歷史生成內容的 n-gram 片段緩存機制。會定期保存頻率較高的 n-gram 序列，並在草擬新 Token 時，借助這些高頻片段進行「自動補全」。

此外，TokenSwift 還引入了「隨機篩選器」，在多個 n-gram 片段中選擇語義最優的一組進行結構構建。這不僅提升了草擬的準確性，也保證了後續驗證的成功率。

🧠 樹結構驗證機制：確保與 AR 一致，輸出「無損」保障

有了草擬機制，還需要驗證其「正當性」。TokenSwift 的另一大亮點在於提出了樹結構的並行驗證模塊，通過構建多個驗證路徑，對草擬的多個 Token 並行進行預測評分判斷，並篩選出與標準 AR 路徑一致的候選。

換句話說，TokenSwift 不僅快，而且不犧牲任何輸出質量，生成內容與原始模型保持一致，是一套真正「無損」的加速方案。

💡 動態 KV 管理 + 重覆懲罰：越長越快，越長越優

為瞭解決 KV 緩存膨脹的問題，TokenSwift 實現了動態的 KV 裁剪機制。模型會根據 Token 重要度與時間衰減策略，對 KV 對進行「主動淘汰」，在保證上下文保留質量的同時，顯著減輕緩存負擔。

與此同時，為了緩解長文本生成過程中的重覆問題，設計了重覆懲罰機制，在生成過程中動態降低重覆 n-gram 的概率，確保最終輸出具備更高的多樣性和可讀性。

📈 實驗評估：全面剖析 TokenSwift 的性能與質量

在多個主流模型上，包括 YaRN-LLaMA2-7b-128k、LLaMA3.1-8b、Qwen2.5-1.5B, 7B, 14B 等，進行了大規模實驗，序列長度涵蓋從 20K 到 100K，TokenSwift 表現均極其亮眼：

加速比普遍在 3 倍以上
生成質量與原模型一致
Distinct-n 指標顯著優於原始 AR 路徑

更令人振奮的是，隨著序列越長，TokenSwift 的加速效果越顯著。在 100K Token 生成任務中，LLaMA3.1-8B 的生成時間從近 5 小時縮短至 1.5 小時，極大降低了實際使用成本。

Token 重用的決定性作用

我們對比了禁用（k=0）與啟用狀態下的接受率和加速比。結果顯示，未啟用重用時，隨著生成長度增長，接受率和加速比均出現明顯下滑；而在 k=20 時，接受率不僅維持在 70–90% 的高水平，還隨序列越長而略有提升，加速比亦從～2.1× 提升至～3.1×，凸顯 Token 重用在減少模型重載和保持驗證高效中的關鍵作用。

動態 KV 更新的巧妙平衡

針對 KV 緩存管理，我們進一步實驗了「全量緩存」、「僅預填後一次更新」與 TokenSwift 中的「動態部分緩存」三種策略。在關閉 Token 重用的前提下，「部分緩存」因接受率低而加速有限，「全量緩存」雖保證了較高接受率，卻因管理成本抵消了速度收益；而 TokenSwift 的動態更新策略恰好在接受率與計算開銷之間取得平衡，使得在 100K Token 任務上依然能保持近 2× 以上的加速效果。

上下文懲罰對多樣性的提升

為了抑制超長生成中的機械複讀，TokenSwift 採用了僅對最近 W Token 應用懲罰值 θ 的「上下文懲罰」機制。在與 top-p、η-sampling、min-p 等多種采樣方法結合的實驗中，未啟用懲罰時 Distinct-n 平均值僅約 0.12；引入 θ=1.2、W=1024 後，平均多樣性飆升至 0.43–0.69，同時加速比僅輕微下降～2–8% 左右，證明了該機制在保持質量一致性與提升文本多樣性上的高效性。