作者熊璟，香港大學(xué)一年級(jí)博士生，師從黃毅教授和孔令鵬教授。已在 ICLR、ICML、NeurIPS、ACL、EMNLP、TMLR等頂級(jí)會(huì)議/期刊發(fā)表論文，研究方向?yàn)楦咝Т笳Z(yǔ)言模型推理與自動(dòng)定理證明。擔(dān)任NAACL、EMNLP、ACL、ICML、ICLR、NeurIPS、COLING等會(huì)議審稿人。個(gè)人主頁(yè)： https://menik1126.github.io/

引言：大模型長(zhǎng)文本推理的瓶頸與突破

隨著大語(yǔ)言模型（LLMs）能力日益提升，AI 對(duì)超長(zhǎng)文本的理解和處理需求也變得前所未有地重要。然而，目前主流 LLM 雖然依賴旋轉(zhuǎn)位置編碼（RoPE）等機(jī)制，在訓(xùn)練階段能高效處理 4K-8K tokens 級(jí)別的上下文，但一旦推理階段外推遇到如 128K 以上長(zhǎng)度的長(zhǎng)文本時(shí)，模型往往受到顯存瓶頸的限制和注意力下沉 (attention sink) 等問(wèn)題影響，采用常規(guī)的文本截?cái)喾桨溉菀壮霈F(xiàn)信息遺失，這極大限制了大模型在實(shí)際場(chǎng)景中的應(yīng)用拓展。

業(yè)界目前嘗試的處理長(zhǎng)文本的高效推理主要的瓶頸有兩個(gè),  一個(gè)是位置編碼的長(zhǎng)度外推, 再一個(gè)是長(zhǎng)度外推中的內(nèi)存瓶頸。     

目前的位置編碼包括兩類(lèi)：一是基于頻率區(qū)分的 NTK 插值方法，為不同頻段位置編碼設(shè)計(jì)專(zhuān)屬策略以試圖拓展長(zhǎng)度上限；二是各種分塊（chunking）方法，將文本切分、塊內(nèi)復(fù)用位置編碼，避免重新訓(xùn)練的高昂成本。 在處理超長(zhǎng)上下文(>128K)外推時(shí), 兩種方案的優(yōu)劣仍是未知的。      

在解決顯存瓶頸時(shí) , 普遍采用了KV cache壓縮的方案 , 并且一些方案發(fā)現(xiàn)在傳統(tǒng)自回歸場(chǎng)景下的注意力普遍遭遇了 “注意力下沉”（Attention Sink）現(xiàn)象影響 —— 模型的注意力極易集中在文本首尾  , 如果加以合理利用可以提升模型性能，但在并行注意力下的注意力下沉的機(jī)制是仍待探索的。

圖 1: 標(biāo)準(zhǔn) Attention 和并行 Attention 機(jī)制下的注意力分布雖然有所不同，但都容易陷入這種 “塌縮” 狀態(tài)。而并行注意力機(jī)制在極長(zhǎng)文本中產(chǎn)生的多峰新型 “sink” 現(xiàn)象，尚缺乏系統(tǒng)剖析與有效解決之道。

我們的創(chuàng)新：ParallelComp，高效超長(zhǎng)文本推理新范式

針對(duì)上述問(wèn)題，我們提出了一種全新的訓(xùn)練免調(diào)（Training-Free）長(zhǎng)文本處理方案 ——ParallelComp，其核心包括并行 Attention 分塊、KV 緩存智能淘汰與注意力偏差校準(zhǔn)三大技術(shù)創(chuàng)新。

• 論文標(biāo)題：ParallelComp: Parallel Long-Context Compressor for Length Extrapolation
• 論文鏈接: https://arxiv.org/abs/2502.14317
• 代碼鏈接: https://github.com/menik1126/ParallelComp

圖 2:ParallelComp 整體流程示意圖

1. 并行分塊注意力，讓超長(zhǎng)文本 “輕松分段消化”

首先，ParallelComp 借鑒并拓展了分塊思想，將輸入長(zhǎng)文本按模型最大處理長(zhǎng)度自動(dòng)切分為若干塊（chunk），并與查詢一起并行輸入模型進(jìn)行本地 Attention 計(jì)算。這樣不僅可以獲得塊內(nèi)的注意力分布，還可以通過(guò)塊的位置編碼重用機(jī)制可以讓模型靈活外推至 128K 級(jí)別的序列長(zhǎng)度而性能無(wú)明顯損失。（見(jiàn)圖 2）

2. 雙重淘汰機(jī)制，極致壓縮顯存消耗

• 分塊淘汰（Chunk Eviction）： 基于每個(gè)塊內(nèi)部 token 對(duì)查詢的 “自信息量” 進(jìn)行在線打分，僅保留信息量最高、對(duì)當(dāng)前任務(wù)最相關(guān)的一部分文本塊，極大縮減內(nèi)存開(kāi)銷(xiāo)。
• KV 緩存淘汰（Parallel KV Cache Eviction）： 采用高效的 FlashAttention 推斷算法，動(dòng)態(tài)評(píng)估每個(gè) token 在 Attention 計(jì)算中的全局重要性，自動(dòng)驅(qū)逐影響較小的 token，實(shí)現(xiàn)了分塊粒度的 KV 緩存智能壓縮。

通過(guò)這兩項(xiàng)機(jī)制，全流程推理所需顯存可以壓縮到 64G，且基于并行預(yù)填充過(guò)程即可完成高達(dá) 128K 長(zhǎng)度推理任務(wù)，顯著提升 Batch inference 和多 GPU 設(shè)置下的推理效率。

圖 3：三類(lèi)典型注意力偏置分布案例（U 型、山峰型、均勻型）

圖 4: 幾種 KV cache 驅(qū)逐策略后的模型的分布

3. 注意力偏差校準(zhǔn)，攻克并行下的多峰 “sink” 現(xiàn)象

通過(guò)理論分析與大量可視化實(shí)證，我們發(fā)現(xiàn)并行分塊機(jī)制下，不同于經(jīng)典的 U 型 Attention Sink，容易出現(xiàn)多峰、區(qū)域性異常（見(jiàn)圖 3、圖 4）。為此，我們提出在 Attention 分布中對(duì)異常高分 token 實(shí)施 “偏差校準(zhǔn)”—— 分層驅(qū)逐被極端關(guān)注的 token，使 Attention 分布趨于平滑，極大降低了模型關(guān)注 “死角” 和信息遺漏風(fēng)險(xiǎn)。

進(jìn)一步的分層實(shí)驗(yàn)揭示，不同類(lèi)型注意力偏差（“首段偏置”“尾段偏置”“中部塌縮”）可分別通過(guò)不同深度層的 token 淘汰策略加以緩解，為長(zhǎng)文本推理設(shè)計(jì)提供了細(xì)致化的實(shí)操指引。

一個(gè)有趣的觀察是有些特定層的注意力偏差對(duì)模型的上下文學(xué)習(xí)能力和推理能力至關(guān)重要。

表 1 : 不同層的 bias token 對(duì)特定任務(wù)的影響

i) 淺層（第 1-8 層）的首段偏置對(duì)于檢索類(lèi)任務(wù)至關(guān)重要。去除這些 token 會(huì)顯著削弱模型性能。

ii) 深層（第 9-16 層）的尾段偏置在模型推理能力中起到關(guān)鍵作用。淘汰這些 token 會(huì)導(dǎo)致編碼和數(shù)學(xué)任務(wù)的性能下降。

iii) 淺層的中部偏置（第 1-8 層）會(huì)損害模型的理解能力，淘汰這些 token 反而能夠提升模型表現(xiàn)。而深層的中部偏置（第 24-31 層）有助于模型在閱讀理解任務(wù)（如 En.MC）中的能力，去除它們會(huì)降低模型表現(xiàn)。

iv) 早期層（第 1-8 層）的尾段偏置對(duì)于模型的上下文學(xué)習(xí)能力同樣非常重要。

理論與實(shí)驗(yàn)分析

我們的理論分析表明，并行 Attention 下 “稀疏性” 與 “有效信息量” 之間存在量化可控的門(mén)檻。隨著 chunk 數(shù)量增多和長(zhǎng)度變長(zhǎng)，Attention 大概率只會(huì)專(zhuān)注于極少數(shù) token。合理選擇稀疏參數(shù)、分層校準(zhǔn)策略，可在計(jì)算效率與信息保真度之間取得最優(yōu)平衡。

表 2: 在 longbench 上的性能。

表 3: 在 infinitebench 上的性能。

表 4: PPL 的性能

表 5: 加速比和性能分析

大規(guī)模實(shí)驗(yàn)顯示，ParallelComp 在 A100-80GB 單卡環(huán)境中，模型推理長(zhǎng)度從 8K 無(wú)縫外推至 128K，prefilling 階段加速高達(dá) 23.5 倍；使用僅 8B 參數(shù)、且僅在 8K 上下文訓(xùn)練的小模型即可在超長(zhǎng)文本任務(wù)中整體性能達(dá)到 GPT-4o 性能的 91.17% , 在特定任務(wù)下可以超過(guò)GPT-4o的性能，甚至超過(guò) Claude-2 和 Kimi-Chat 等高參數(shù)閉源大模型。這一成果充分證明，我們的方法不僅適用于算力受限場(chǎng)景，更為 LLM 實(shí)際落地部署帶來(lái)了嶄新范式。

結(jié)論與展望

ParallelComp 為長(zhǎng)文本推理時(shí)代的 LLM 結(jié)構(gòu)創(chuàng)新帶來(lái)了 “多快好省” 的系統(tǒng)級(jí)躍升 —— 無(wú)需新訓(xùn)練、僅借助智能分塊和淘汰機(jī)制，即可極大提升處理長(zhǎng)度、效率，并擺脫原有的注意力失衡難題。我們期待該方法未來(lái)在法律文檔、醫(yī)學(xué)文獻(xiàn)、長(zhǎng)篇小說(shuō)等需求復(fù)雜的產(chǎn)業(yè)應(yīng)用中進(jìn)一步開(kāi)花結(jié)果。