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

生成10萬Token的文本，傳統(tǒng)自回歸模型需要近5個(gè)小時(shí)，現(xiàn)在僅需90分鐘！

最新研究提出了一個(gè)名為TOKENSWIFT框架，從模型加載、KV緩存管理到Token生成策略進(jìn)行了全方位的優(yōu)化。

實(shí)驗(yàn)結(jié)果證明，該方法不僅能大幅提升生成效率，更在保證生成質(zhì)量和多樣性上實(shí)現(xiàn)了無損加速。

而且支持R1-Distill，團(tuán)隊(duì)發(fā)布經(jīng)過微調(diào)的DeepSeek-R1-Distill-Qwen-32B模型，同樣具備3倍加速效果。

來看demo展示：

本研究由來自北京通用人工智能研究院的團(tuán)隊(duì)完成，以下是更多細(xì)節(jié)。

TOKENSWIFT框架長啥樣？

隨著LLMs長上下文窗口能力的不斷提升，復(fù)雜任務(wù)對超長文本生成的需求越來越高。傳統(tǒng)的自回歸（AR）生成方式雖然在短文本上表現(xiàn)良好，但在長文本生成中存在明顯瓶頸，主要體現(xiàn)在以下三個(gè)方面：

• 模型頻繁加載問題

由于自回歸生成每生成一個(gè)Token都需要從GPU存儲中重新加載模型權(quán)重，導(dǎo)致I/O操作頻繁、延遲高。在生成10萬Token時(shí)，模型需要重復(fù)加載上萬次，嚴(yán)重拖慢整體生成速度。

• KV緩存的動態(tài)管理

生成超長文本過程中，模型內(nèi)部的鍵值對（KV Cache）不斷增長，若直接使用全量KV緩存，不僅超出內(nèi)存預(yù)算，還會大幅增加計(jì)算時(shí)間。如何在保證關(guān)鍵信息不丟失的前提下，實(shí)現(xiàn)KV緩存的高效更新成為一大難題。

• 重復(fù)性生成

長序列生成易出現(xiàn)重復(fù)和冗余問題，影響文本的多樣性和質(zhì)量。雖然重復(fù)問題并非論文的主要聚焦點(diǎn)，但在超長文本生成中依然需要有效抑制。

為解決上述難題，論文提出了TOKENSWIFT——一個(gè)全新的框架，旨在實(shí)現(xiàn)無損加速超長序列生成，其主要創(chuàng)新點(diǎn)體現(xiàn)在以下幾個(gè)方面：

1）多Token并行生成與Token復(fù)用

論文借鑒了Medusa等方法，通過引入額外的線性層，使模型在一次前向傳播中能夠同時(shí)生成多個(gè)草稿Token。

更重要的是，基于生成文本中的n-gram頻率信息，系統(tǒng)會自動檢索并復(fù)用高頻短語，從而進(jìn)一步減少模型重新加載的次數(shù)，提升整體效率。

2）動態(tài)KV緩存更新策略

在KV緩存管理上，TOKENSWIFT采用動態(tài)更新策略。系統(tǒng)在生成過程中將初始KV緩存保留，同時(shí)根據(jù)Token的重要性對后續(xù)緩存進(jìn)行有序替換。

這種方式不僅有效控制了緩存的規(guī)模，還確保了關(guān)鍵信息始終被保存，大幅降低了因緩存加載帶來的延遲。

3）基于樹結(jié)構(gòu)的多候選Token驗(yàn)證

為保證生成結(jié)果與目標(biāo)模型預(yù)測的一致性，TOKENSWIFT引入了樹形注意力機(jī)制。

通過構(gòu)建包含多個(gè)候選Token組合的樹形結(jié)構(gòu)，并采用并行驗(yàn)證的方式，從中隨機(jī)選擇最長且有效的n-gram作為最終輸出，確保生成過程無損且多樣性得到提升。

4）上下文懲罰策略

為了進(jìn)一步抑制重復(fù)生成問題，論文設(shè)計(jì)了一種上下文懲罰方法。該方法在生成過程中為近期生成的Token施加懲罰，使得模型在選擇下一Token時(shí)更傾向于多樣化輸出，從而有效減少重復(fù)現(xiàn)象。

TOKENSWIFT效果如何？

實(shí)驗(yàn)部分，論文在多種模型架構(gòu)（包括MHA和GQA）及不同規(guī)模（1.5B、7B、8B、14B）上進(jìn)行了充分測試。

結(jié)果表明，TOKENSWIFT在生成10萬Token長序列時(shí)，相較于傳統(tǒng)自回歸方法，平均實(shí)現(xiàn)了3倍以上的加速，且生成結(jié)果在準(zhǔn)確性和多樣性上基本保持無損。

1）加速效果

實(shí)驗(yàn)數(shù)據(jù)顯示，在LLaMA3.1-8B模型下，傳統(tǒng)AR生成10萬Token約需4.9小時(shí)，而使用TOKENSWIFT后僅需90分鐘，大幅節(jié)省時(shí)間。在Qwen2.5-14B時(shí)，傳統(tǒng)AR生成10萬Token更是達(dá)到了7.9小時(shí)，加速后僅需142分鐘。這一成果對于實(shí)際應(yīng)用中需要實(shí)時(shí)或高效長文本生成的場景具有重要意義。

2）驗(yàn)證率與接受率

論文設(shè)計(jì)了多項(xiàng)指標(biāo)來評估生成質(zhì)量，包括Token接受率和Distinct-n指標(biāo)。結(jié)果表明，TOKENSWIFT不僅在速度上顯著領(lǐng)先，還能在保持無損生成的前提下，有效提升文本的多樣性。

消融實(shí)驗(yàn)與案例分析

在深入理解TOKENSWIFT各模塊貢獻(xiàn)的過程中，論文還進(jìn)行了全面的消融實(shí)驗(yàn)和案例分析，為優(yōu)化方案提供了充分依據(jù)。

消融實(shí)驗(yàn)：關(guān)鍵組件的作用

• Token復(fù)用消融：實(shí)驗(yàn)中將Token復(fù)用參數(shù)設(shè)為0（即不復(fù)用），結(jié)果顯示，接受率和生成速度均顯著下降，表明復(fù)用機(jī)制在減少模型加載次數(shù)、提升整體效率方面發(fā)揮了關(guān)鍵作用。
• KV緩存更新策略對比：研究對比了全量緩存、一次性更新和動態(tài)更新三種方案。結(jié)果表明，全量緩存雖然在接受率上略占優(yōu)勢，但其高計(jì)算開銷使得整體加速效果不理想；而一次性更新則因緩存膨脹導(dǎo)致性能下降。動態(tài)更新策略則在保持高接受率的同時(shí)，實(shí)現(xiàn)了最佳的速度與資源平衡。
• 上下文懲罰效果：在不同采樣方法下，加入上下文懲罰后，生成文本的Distinct-n指標(biāo)明顯提高。

例如，在min-p采樣場景下，Distinct-n平均得分從0.12提升至0.69，僅帶來約8%的速度損失，充分驗(yàn)證了該策略在抑制重復(fù)生成方面的有效性。

案例分析：真實(shí)生成對比

論文還對比了在有無上下文懲罰條件下生成文本的差異，案例分析結(jié)果令人印象深刻：

• 重復(fù)現(xiàn)象的抑制：在未使用上下文懲罰時(shí)，生成文本在大約5K Token處就出現(xiàn)明顯重復(fù)，且多為逐字重復(fù)；而采用上下文懲罰后，重復(fù)問題明顯延遲至約60K Token，且重復(fù)部分多表現(xiàn)為語義層次上的相似，而非直接復(fù)制，整體文本連貫性和可讀性大幅改善。
• 文本質(zhì)量的提升：案例對比顯示，使用完整TOKENSWIFT流程的生成結(jié)果在邏輯連貫性、表達(dá)多樣性和創(chuàng)新性方面均優(yōu)于傳統(tǒng)方法，為實(shí)際應(yīng)用提供了更高質(zhì)量的文本輸出。

通過這些消融實(shí)驗(yàn)和案例分析，論文不僅證明了各關(guān)鍵技術(shù)模塊的重要性，也為后續(xù)優(yōu)化指明了方向，充分體現(xiàn)了TOKENSWIFT在超長文本生成領(lǐng)域的先進(jìn)性和實(shí)用性。

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