作者熊璟，香港大學一年級博士生，師從黃毅教授和孔令鵬教授。已在 ICLR、ICML、NeurIPS、ACL、EMNLP、TMLR等頂級會議/期刊發表論文，研究方向為高效大語言模型推理與自動定理證明。擔任NAACL、EMNLP、ACL、ICML、ICLR、NeurIPS、COLING等會議審稿人。個人主頁： https://menik1126.github.io/

引言：大模型長文本推理的瓶頸與突破

隨著大語言模型（LLMs）能力日益提升，AI 對超長文本的理解和處理需求也變得前所未有地重要。然而，目前主流 LLM 雖然依賴旋轉位置編碼（RoPE）等機制，在訓練階段能高效處理 4K-8K tokens 級別的上下文，但一旦推理階段外推遇到如 128K 以上長度的長文本時，模型往往受到顯存瓶頸的限制和注意力下沉 (attention sink) 等問題影響，采用常規的文本截斷方案容易出現信息遺失，這極大限制了大模型在實際場景中的應用拓展。

業界目前嘗試的處理長文本的高效推理主要的瓶頸有兩個,  一個是位置編碼的長度外推, 再一個是長度外推中的內存瓶頸。     

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

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

圖 1: 標準 Attention 和并行 Attention 機制下的注意力分布雖然有所不同，但都容易陷入這種 “塌縮” 狀態。而并行注意力機制在極長文本中產生的多峰新型 “sink” 現象，尚缺乏系統剖析與有效解決之道。

我們的創新：ParallelComp，高效超長文本推理新范式

針對上述問題，我們提出了一種全新的訓練免調（Training-Free）長文本處理方案 ——ParallelComp，其核心包括并行 Attention 分塊、KV 緩存智能淘汰與注意力偏差校準三大技術創新。

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

圖 2:ParallelComp 整體流程示意圖

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

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

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

• 分塊淘汰（Chunk Eviction）： 基于每個塊內部 token 對查詢的 “自信息量” 進行在線打分，僅保留信息量最高、對當前任務最相關的一部分文本塊，極大縮減內存開銷。
• KV 緩存淘汰（Parallel KV Cache Eviction）： 采用高效的 FlashAttention 推斷算法，動態評估每個 token 在 Attention 計算中的全局重要性，自動驅逐影響較小的 token，實現了分塊粒度的 KV 緩存智能壓縮。

通過這兩項機制，全流程推理所需顯存可以壓縮到 64G，且基于并行預填充過程即可完成高達 128K 長度推理任務，顯著提升 Batch inference 和多 GPU 設置下的推理效率。

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

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

3. 注意力偏差校準，攻克并行下的多峰 “sink” 現象

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

進一步的分層實驗揭示，不同類型注意力偏差（“首段偏置”“尾段偏置”“中部塌縮”）可分別通過不同深度層的 token 淘汰策略加以緩解，為長文本推理設計提供了細致化的實操指引。

一個有趣的觀察是有些特定層的注意力偏差對模型的上下文學習能力和推理能力至關重要。

表 1 : 不同層的 bias token 對特定任務的影響

i) 淺層（第 1-8 層）的首段偏置對于檢索類任務至關重要。去除這些 token 會顯著削弱模型性能。

ii) 深層（第 9-16 層）的尾段偏置在模型推理能力中起到關鍵作用。淘汰這些 token 會導致編碼和數學任務的性能下降。

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

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

理論與實驗分析

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

表 2: 在 longbench 上的性能。

表 3: 在 infinitebench 上的性能。

表 4: PPL 的性能

表 5: 加速比和性能分析

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

結論與展望

ParallelComp 為長文本推理時代的 LLM 結構創新帶來了 “多快好省” 的系統級躍升 —— 無需新訓練、僅借助智能分塊和淘汰機制，即可極大提升處理長度、效率，并擺脫原有的注意力失衡難題。我們期待該方法未來在法律文檔、醫學文獻、長篇小說等需求復雜的產業應用中進一步開花結果。