全新檢索模式：在無限長token下，大語言模型自身或能檢索信息！

受大語言模型（LLM）上下文窗口大小的限制，處理輸入token數超過上限的各種任務頗具挑戰性，無論是簡單的直接檢索任務，還是復雜的多跳推理任務。

盡管新提出的各種方法用來增強大語言模型的長上下文處理能力，但這些方法痛點突出：

要么會產生高昂的訓練后成本，

要么需要額外的工具模塊（如檢索增強生成RAG），

要么在實際任務中顯示出改進，并不明顯。

研究團隊觀察了各層注意力分布與生成答案之間的相關性，通過實驗證實了注意力分配與檢索增強能力是一致的。

基于上述見解，研究團隊提出了一種全新的方法InfiniRetri，該方法利用大語言模型自身的注意力信息，實現對任意長度輸入的精確檢索。

實驗表明，在100萬個token的「大海撈針」（Needle-In-a-Haystack，NIH）測試中，InfiniRetri將5億參數的模型從44.6%提升到了100%的準確率。

沒有InfiniRetri的NIH測試只有44.6%的準確率

要知道GPT-4在NIH測試中都做不到100%準確率。

InfiniRetri一舉超過了其他方法或更大的模型，創造了當前最佳（SOTA）結果。

值得注意的是，某7B模型在HotpotQA任務上的得分，超越了其他同等參數規模的模型。

類似地，Mistral-7B-Instruct v0.2作為擅長短文本推理的模型，在長文本任務中的表現也得到了顯著提升。

此外，新方法在實際基準測試中也取得了顯著的性能提升，最大提升幅度達到288%。

另外，無需額外訓練，InfiniRetri就可應用于任何基于Transformer的大語言模型，并且能大幅降低長文本推理延遲和計算開銷。

論文鏈接：https://arxiv.org/abs/2502.12962

項目地址：https://github.com/CapitalCode2020/InfiniRetri2

文章主要貢獻如下：

• 創新性提出「注意力分配與檢索增強對齊」概念，并成功利用這一特性提升LLM處理長文本的能力。
• 無需額外訓練，InfiniRetri可直接應用于任何基于Transformer的LLM，使其具備處理無限長上下文的能力。
• 不同于RAG依賴外部嵌入模型，提出了「注意力中的檢索」這一新穎視角，充分利用LLM內在能力增強長文本處理能力。
• 顯著降低推理延遲和計算開銷，在大規模數據集上的檢索與問答任務中表現優異，展現出在極長文本處理場景中的實際應用價值。
• 不僅僅是擴展上下文窗口，證明了可以通過多種方式提升LLM處理長文本能力。未來的改進可通過在較小上下文窗口內增強模型內部能力，從而獲得更優的長文本處理效果。

鞭辟入里，靈魂三問

近年來，許多主流的LLM都在致力于擴展上下文窗口的規模。

例如，GPT-4、Llama3和DeepSeekV3支持最長128K的上下文長度，Claude-3可達200K，而Gemini-Pro-1.5和Qwen2.5-1M甚至支持1M的上下文長度。

然而，實際效果并未完全達到這些模型所宣稱的長度。

LLM在處理超長上下文方面仍然面臨重大挑戰。

那自然考慮：是否能不斷延長上下文窗口？

這是第一個問題，答案是不能，有3大原因：

1. 增加計算成本，產生明顯的延遲。
2. 輸入長度具有長尾效應，簡單地擴展上下文窗口帶來的提升注定會帶來越來越低。
3. 持續和階段性訓練成本極高，絕大多數研究人員無法承受。

接下來，進一步考慮第二個問題：要打破處理長上下文時不同窗口之間的信息壁壘，是否存在一種低成本的方法？

事實上，業界已經通過檢索增強生成（Retrieval-Augmented Generation，RAG）提供了答案。

RAG由兩大核心組件組成：檢索模塊和生成模塊。其中，檢索模塊依賴外部嵌入模型，根據輸入查詢從長文本中檢索相關內容。

然而，RAG系統普遍難以建立檢索信息之間的關聯。

相比之下，在推理過程中，LLM的注意力機制，能夠高效地建立不同信息片段之間的聯系。

由此引出了第三個問題：為何不直接利用LLM自身的檢索能力，來處理長文本上下文？

為了讓LLM以低成本的方式壓縮并存儲過去的鍵（key）和值（value），研究團隊認為，只有打破不同上下文窗口之間的信息壁壘，LLM才能真正提升其處理長文本的能力。

此外，通過觀察LLM在推理過程中回答問題時的注意力分配模式，研究團隊提出：這種注意力分配模式與檢索增強能力（RAG）高度契合。

研究團隊測試了基于所有層的注意力分布來檢索答案的準確性。

正如下圖3所示，越接近輸出層，這種模式的增強效果就越明顯。

此外，在第14層和第15層，檢索準確率達到一個局部峰值。

因此，InfiniRetri引入了全新的策略，利用LLM自身的注意力信息，而非依賴外部嵌入模型，去提升其長文本處理能力。

無需對基于Transformer的LLM進行額外訓練，InfiniRetri可以開箱即用。得益于此，研究團隊在多個模型上進行了全面的對比實驗。

在跨上下文長度的事實檢索任務「大海撈針」（Needle In A Haystack，NIH）中，InfiniRetri僅使用5億個額外參數，就將模型的上下文長度從原始的32K擴展至100萬個token（如圖1所示）。

原始方法（上）與InfiniRetri方法（下）對比，準確檢索的最大文本長度從原始32K提升至超過1M tokens。

更值得注意的是，InfiniRetri在NIH任務上實現了無限長度范圍內的精準檢索，不僅超越了當前主流方法，還有效解決了NIH任務的挑戰。

此外，在LongBench提供的9個真實數據集上，InfiniRetri在基于KV Cache甚至Full KV的主流方法中均取得了超越性的表現，尤其是在多文檔問答（Multi-Document QA）任務（如HotpotQA、2WikiMQA和Musique）中，Qwen2-7B-Instruct采用InfiniRetri后，平均性能提升高達369.6%。

InfiniRetri方法：帶著問題閱讀

那么，如何應用這一模式來處理超出上下文窗口的長文本，并真正提升LLM的長文本處理能力呢？

在本節中，將InfiniRetri方法拆分為三個小節，分別介紹該模式的應用方式，以提升LLM的長文本處理能力。

如圖4所示，該方法的完整工作流程包括五個主要步驟，詳細介紹這些步驟：

• 步驟1（切分，Chunk）、步驟2（合并，Merge）和步驟3（推理，Inference）。
• 步驟4（檢索，Retrieval）是方法的核心部分。
• 步驟5（緩存，Cache）。

圖4：InfiniRetri方法在增強LLM長上下文處理能力中的完整工作流程

文本分割

新方法受到人類閱讀書籍過程的啟發，專門針對LLM在處理超出上下文窗口的文本時所面臨的挑戰。

盡管人類的視野有限，一次只能看到一頁內容，但仍然可以逐頁閱讀并理解整本書。在這個過程中，大腦就像一個緩存（cache），通過記憶保留并整合每一頁的信息，從而掌握整本書的內容。

類似地，InfiniRetri將整篇文本拆分為連續的文本塊（chunk）。

這種切分方式雖然與RAG相似，但不同之處在于，InfiniRetri不是并行處理每個文本塊，而是按照順序逐塊迭代地處理每個文檔。

這種方法能夠保留文本的順序信息，更符合人類的閱讀習慣。

具體而言，如圖4所示，在步驟1（切分，Chunk）中，研究團隊根據句子邊界將整個長文本劃分為長度大致相等的文檔塊，其長度由方法參數ChunkSize3決定。

然后，這些文檔塊依次與緩存（Cache）中保留的token進行合并，形成完整的輸入序列，稱為MergeToken，并將其輸入LLM進行處理。

InfiniRetri采用了類似滑動窗口注意力（Slide Window Attention，SWA）的迭代方式，按順序處理每個文本片段。

然而，InfiniRetri對緩存的處理方式與傳統方法有本質區別。

• 傳統緩存通常在每一層存儲過去的鍵值（Key-Value）狀態，而新方法則重新定義了緩存概念，改為存儲過去的token ID。
• 如圖4，步驟2（合并，Merge）所示，新方法在輸入LLM之前，將緩存的token ID與當前文本片段的token進行合并，從而取代了推理過程中對歷史鍵值狀態的合并需求。
• 因此，在步驟3（推理，Inference）階段，LLM仍然使用標準注意力機制，而非SWA。對于第h層的注意力得分，其計算公式如下：

其中，A^h∈R^{n×m}表示查詢（query）和鍵（key）組成的注意力矩陣，n是查詢的數量，m是鍵的數量。

在注意力中檢索（Retrieval In Attention）

在單個上下文窗口內，根據問題token準確定位相關的上下文token，注意力分配模式能夠幫助LLM找到正確答案。

如果在滑動窗口框架內的每次推理過程中持續應用這一模式，理論上，LLM就可以在保持查詢不變的情況下，對整個長文本進行推理。

這一過程與人類的閱讀方式高度相似，類似于公認的「帶著問題閱讀」學習策略，即通過問題作為錨點，在LLM可處理的范圍內逐步整合相關信息。

因此，LLM能否精準檢索與問題最相關的文本，是本方法有效性的核心。

關鍵在于設計基于注意力得分分布的token檢索策略和算法，以確保模型能夠在長文本中高效提取關鍵信息。

借鑒實驗結果，研究團隊選取了多頭注意力（Multi-Head Attention）的最后一層，并對所有注意力頭的得分進行求和聚合（如公式2所示），探索了一種能夠準確判斷模型關注重點的方法。

通過可視化注意力得分，研究團隊觀察到：與答案相關的信息通常由連續的token組成。

也就是說，它們以短語級別的粒度存在。

這一發現與在圖2c的實驗結果一致，進一步確認了LLM在token級別上具備較高的注意力精度。

因此，希望設計的操作是：計算每個token及其相鄰token在2D注意力得分矩陣中的累積注意力得分。

計算結果將作為新的特征，在后續的檢索過程中用于排序。經過深入分析，發現此操作等效于使用一個填充了1的卷積核（kernel）進1D卷積。

對于查詢i和鍵j，其特征重要性計算如下：

1. 聚合所有注意力頭的得分（公式2）：

2. 基于1D卷積計算每個token及其相鄰token的重要性（公式3）：

其中，k是1D卷積核的大小，對應于方法中的參數Phrase Token Num。

3. 沿矩陣的列方向求和，計算每個上下文token的總重要性分數（公式4）：

其中，s_i代表第i個上下文token的綜合重要性分數。

4. 最后，選取重要性分數最高的前K個上下文token，并將其所在句子的所有token寫入緩存（cache）。這個過程可表示為：

即，從所有token的重要性分數v中，選擇排名前K的token，并將它們所在的完整句子存入緩存，以便后續推理時使用。

緩存句子Token（Cache Sentence Token）

在推理過程中對緩存（cache）的使用方式，InfiniRetri與傳統方法存在本質性區別。

相比于直接使用緩存，InfiniRetri將其用于存儲過去的上下文信息。

具體而言，主要有以下兩點不同：

1 新方法在模型外部緩存token ID，而不是每層的歷史鍵值（Key-Value）狀態。具體來說，在推理過程中不使用傳統的Key-Value緩存，而是在每次推理前，將過去的上下文信息與當前輸入合并后再進行推理。

2 新方法基于短語級特征進行檢索，并在緩存中存儲包含Top-K token的句子級token。也就是說，存儲的是完整的句子，而不是單獨的token，從而確保檢索到的信息更具上下文完整性。

事實上，正是這兩項創新性改進，使得新方法在無需微調的情況下，就能比傳統的KV緩存方法更有效地提升LLM處理長文本的能力。

新方法并不試圖壓縮緩存中的token，而是保留句子級別的相關上下文信息。這是因為，句子是最小的完整語義單元，相比于單個token，更能確保LLM對上下文的理解。

在LLM逐步推理每個文本片段的過程中，緩存中保留的中間結果是動態變化的，它們由先前存儲的token和當前輸入片段的組合決定。因此，在整個過程中，這些中間結果會相對調整和更新，以適應模型的理解需求。

AI界的「大海撈針」

大海撈針（Needle-in-a-Haystack，NIH）任務要求模型在一篇「超長文檔」（「大海撈針」之海）中，精準檢索出一個特定的目標句子（「針」），該句子可以被隨機插入到文檔的任意位置。

通過調整「針」的放置位置（文檔深度）和上下文長度，反復測試以衡量模型的表現。

為了直觀分析模型的檢索能力，采用了熱力圖（heatmap）可視化實驗：

• 綠色代表完美檢索（準確找到目標句），
• 其他顏色表示檢索錯誤。

這種可視化方法可以直觀展示LLM處理長文本的能力上限，因此被廣泛用于評估。

實驗1：Llama3-8B-Instruct的對比

如圖6所示，在Llama3-8B-Instruct模型上測試NIH任務，并與以下方法進行了對比：Full KV（完整KV緩存）

StreamingLLM、H2O、SnapKV、PyramidKV以及InfiniRetri（新方法）。

在最長32K token的輸入文本上進行實驗，結果表明：

1. 傳統KV緩存壓縮方法雖有所改進，但沒有超越Full KV的性能。
2. InfiniRetri方法的表現優于FullKV，顯著增強了Llama3-8B-Instruct處理NIH任務的能力，甚至突破了原始8K token上下文窗口的限制。

實驗2：Mistral-7B-Instruct的對比

為了進一步驗證InfiniRetri的有效性，在Mistral-7B-Instruct上擴展了輸入長度。

Mistral-7B-Instruct官方支持的上下文窗口為32K token，對比了FullKV和InfiniRetri的表現，如圖7所示：

• Mistral-7B-Instruct+FullKV（圖7a）：最多可在43K token長度范圍內正確完成NIH任務。
• Mistral-7B-Instruct+InfiniRetri（圖7b）：在相同的參數設置下，新方法不僅超越了Llama3-8B-Instruct，而且在NIH任務上達到了100%的檢索準確率，并將可處理的輸入長度擴展至1M token，且沒有損失準確率。

關鍵發現與額外實驗

進一步觀察到：只要LLM在有限上下文窗口內具備足夠的檢索能力，新方法就可以賦能模型處理超長文本的檢索任務，理論上可支持無限長輸入。

基于這一發現，在更小的開源模型上的額外實驗，結果符合預期：

新方法將該模型的有效上下文token長度從32K擴展至1M+，從而使其在NIH任務上具備了近乎無限的長文本處理能力（如圖1所示）。

LongBench實驗

從表1的整體實驗結果來看，新方法是唯一一個在所有模型上全面超越Full KV的方法，其中在文檔問答（DocumentQA）任務中的提升最為顯著。

主要實驗結果：

1. LLaMA3-8B-Instruct：相對提升4.9%（32.92→34.56）
2. Qwen2-7B-Instruct：相對提升70.5%（25.11→42.82）
3. Mistral-7B-Instruct-v0.3：相對提升55.8%（24.17→37.68）

其中，Qwen2-7B-Instruct在HotpotQA任務上的表現提升最為顯著，最大增幅達到288%（14.8→57.52）。

關鍵發現

• 值得注意的是，Qwen2-7B-Instruct在HotpotQA任務上的得分，超越了其他同等參數規模的模型，表明其在短文本推理方面的優勢。這進一步說明，Qwen2-7B-Instruct通過新方法，能夠有效提升其長文本推理能力。
• 類似地，Mistral-7B-Instruct v0.2作為擅長短文本推理的模型，在長文本任務中的表現也得到了顯著提升。

隨后，雖然新方法在長文檔問答任務中取得了顯著改進，但在文檔摘要任務上的表現相對較差。

這種差異可能源于摘要任務的性質，這些任務通常需要更豐富的上下文信息來生成高質量的輸出。

新方法無法一次性訪問所有相關信息，這在一定程度上限制了它在這些任務中的有效性。

與問答和檢索任務不同，在這些任務中，答案通常依賴于長上下文的一小部分，摘要任務則高度依賴于對整個上下文的全面理解。

因此，新方法可能需要進一步的優化和改進，以更好地應對這些摘要任務。

為了進一步評估InfiniRetri的有效性，使用最新的Qwen2.5-7B-Instruct模型在LongBenchV2上進行了額外的實驗。

正如表2所示的結果，在應用了新方法InfiniRetri后，在處理LongBenchV2上的長文本和中等長度文本方面，Qwen2.5-7B整體性能與72B的對比模型相當，表現出顯著的改進。

這一結果進一步驗證了，只要LLMs在短上下文場景中表現出色，新方法就可以有效地提高其處理更長上下文文本的能力。

降低延遲與計算開銷

如前所述，新方法采用分段滑動窗口機制+迭代處理，在確保LLM推理長度保持在方法參數范圍內的同時，僅在緩存中保留最相關的token。

這種機制使得LLM僅需處理長文本中的少部分關鍵信息，從而顯著降低長文本處理時的推理延遲和計算開銷。

正如表4所示，即使未對方法參數進行精細調整，新方法仍能在LongBench文檔問答（QA）任務中，大幅降低推理成本，適用于Llama3-8B-Instruct、Mistral-7B-InstructV0.2和Qwen2-7B-Instruct等模型。

例如，在NtvQA任務中，僅保留4.5%的原始輸入文本（18409->834）；在HotpotQA任務中，Qwen2-7B-Instruct僅需處理8.7%的原始文本（9152->795），但推理性能提升高達288%。

這些實驗結果進一步證明，新方法能夠通過優化LLM在較小上下文窗口內的能力，有效提升其長文本處理能力。

這一發現表明，提升LLM的長文本處理能力不僅可以通過擴展上下文窗口，還可以通過優化模型在小窗口內的推理能力，再結合新方法機制，實現高效處理長文本的目標。

評論

這篇論文強調了一個本應顯而易見的事實：預測和檢索是同一枚硬幣的兩面。

要有效地進行預測，首先必須確定什么是相關的。令人驚訝的是，適當地利用注意力模式，擁有5億參數的模型可以在100萬個token上執行完美的檢索。

這引發了一個有趣的問題：如果圍繞檢索能力明確設計架構會怎樣？

Transformer架構是為預測而設計的，檢索是作為副產品出現的。那么，專門為檢索優化的架構會是什么樣子呢？

許多資金已經投入到構建大規模RAG（檢索增強生成）系統中。

如果這篇論文所承諾的性能改進是真實的，其影響將是巨大的。