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

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

• 論文標(biāo)題：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/

重新定義超長生成：為什么傳統(tǒng)方法「慢」？

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

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

主要原因有三：

• 模型重復(fù)重載：每生成一個 Token，都會觸發(fā)一次完整的前向推理過程；在多進程或流水線執(zhí)行時，模型需要不斷讀取參數(shù)，造成 I/O 瓶頸。
• KV 緩存無限膨脹：Transformer 架構(gòu)要求保留所有歷史 Token 的 Key/Value 信息，用于后續(xù) Token 的注意力計算。隨著生成進程推進，KV 緩存占用不斷增加，導(dǎo)致計算與內(nèi)存開銷雪上加霜。
• 語義重復(fù)堆疊：生成越長，模型越容易陷入句式與主題的復(fù)讀循環(huán)，降低輸出多樣性與用戶體驗。

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

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

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

讓我們來逐步拆解 TokenSwift 的關(guān)鍵技術(shù)理念：

 多 Token 并行草擬：告別一次一 Token 的低效時代

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

更關(guān)鍵的是，草擬階段并非 「胡亂猜測」。引入了上下文引導(dǎo)機制，使草擬結(jié)果具備較高的語義相關(guān)性與語法一致性，隨后通過結(jié)構(gòu)化的驗證機制確保其與標(biāo)準(zhǔn) AR 路徑一致。

 n-gram 啟發(fā)式補全：巧用歷史片段，精確草擬結(jié)構(gòu)

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

此外，TokenSwift 還引入了 「隨機篩選器」，在多個 n-gram 片段中選擇語義最優(yōu)的一組進行結(jié)構(gòu)構(gòu)建。這不僅提升了草擬的準(zhǔn)確性，也保證了后續(xù)驗證的成功率。

 樹結(jié)構(gòu)驗證機制：確保與 AR 一致，輸出 「無損」 保障

有了草擬機制，還需要驗證其 「正當(dāng)性」。TokenSwift 的另一大亮點在于提出了樹結(jié)構(gòu)的并行驗證模塊，通過構(gòu)建多個驗證路徑，對草擬的多個 Token 并行進行預(yù)測評分判斷，并篩選出與標(biāo)準(zhǔn) AR 路徑一致的候選。

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

動態(tài) KV 管理 + 重復(fù)懲罰：越長越快，越長越優(yōu)

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

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

實驗評估：全面剖析 TokenSwift 的性能與質(zhì)量

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

• 加速比普遍在 3 倍以上
• 生成質(zhì)量與原模型一致
• Distinct-n 指標(biāo)顯著優(yōu)于原始 AR 路徑

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

Token 重用的決定性作用

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

動態(tài) KV 更新的巧妙平衡

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

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

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

小結(jié)

TokenSwift 擺脫了傳統(tǒng)自回歸生成的性能枷鎖，讓我們在面對超長文本推理任務(wù)時，不再被時間與內(nèi)存所拖累。

它并非一個 「另起爐灶」 的新模型，而是一種可直接嵌入現(xiàn)有主流模型（LLaMA、Qwen 等）的通用加速策略，具備極強的兼容性與部署便利性。

更重要的是，TokenSwift 對推理質(zhì)量的 「無損」 保證，讓我們在享受速度提升的同時，不犧牲任何精度與語義一致性。我們相信，這一框架將為未來多輪推理、代碼生成、Agent 計劃編排等長文本場景提供堅實的技術(shù)支撐。

歡迎閱讀 ICML 2025 原論文，了解更多技術(shù)細(xì)節(jié)！