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

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

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

要么會產(chǎn)生高昂的訓(xùn)練后成本，

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

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

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

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

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

沒有InfiniRetri的NIH測試只有44.6%的準(zhǔn)確率

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

InfiniRetri一舉超過了其他方法或更大的模型，創(chuàng)造了當(dāng)前最佳（SOTA）結(jié)果。

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

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

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

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

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

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

文章主要貢獻如下：

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

鞭辟入里，靈魂三問

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

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

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

LLM在處理超長上下文方面仍然面臨重大挑戰(zhàn)。

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

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

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

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

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

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

然而，RAG系統(tǒng)普遍難以建立檢索信息之間的關(guān)聯(lián)。

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

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

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

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

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

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

此外，在第14層和第15層，檢索準(zhǔn)確率達(dá)到一個局部峰值。

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

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

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

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

更值得注意的是，InfiniRetri在NIH任務(wù)上實現(xiàn)了無限長度范圍內(nèi)的精準(zhǔn)檢索，不僅超越了當(dāng)前主流方法，還有效解決了NIH任務(wù)的挑戰(zhàn)。

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

InfiniRetri方法：帶著問題閱讀

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

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

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

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

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

文本分割

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

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

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

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

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

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

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

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

然而，InfiniRetri對緩存的處理方式與傳統(tǒng)方法有本質(zhì)區(qū)別。

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

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

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

在單個上下文窗口內(nèi)，根據(jù)問題token準(zhǔn)確定位相關(guān)的上下文token，注意力分配模式能夠幫助LLM找到正確答案。

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

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

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

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

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

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

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

這一發(fā)現(xiàn)與在圖2c的實驗結(jié)果一致，進一步確認(rèn)了LLM在token級別上具備較高的注意力精度。

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

計算結(jié)果將作為新的特征，在后續(xù)的檢索過程中用于排序。經(jīng)過深入分析，發(fā)現(xiàn)此操作等效于使用一個填充了1的卷積核（kernel）進1D卷積。

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

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

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

其中，k是1D卷積核的大小，對應(yīng)于方法中的參數(shù)Phrase Token Num。

3. 沿矩陣的列方向求和，計算每個上下文token的總重要性分?jǐn)?shù)（公式4）：

其中，s_i代表第i個上下文token的綜合重要性分?jǐn)?shù)。

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

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

緩存句子Token（Cache Sentence Token）

在推理過程中對緩存（cache）的使用方式，InfiniRetri與傳統(tǒng)方法存在本質(zhì)性區(qū)別。

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

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

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

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

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

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

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

AI界的「大海撈針」

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

通過調(diào)整「針」的放置位置（文檔深度）和上下文長度，反復(fù)測試以衡量模型的表現(xiàn)。

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

• 綠色代表完美檢索（準(zhǔn)確找到目標(biāo)句），
• 其他顏色表示檢索錯誤。

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

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

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

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

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

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

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

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

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

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

關(guān)鍵發(fā)現(xiàn)與額外實驗

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

基于這一發(fā)現(xiàn)，在更小的開源模型上的額外實驗，結(jié)果符合預(yù)期：

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

LongBench實驗

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

主要實驗結(jié)果：

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任務(wù)上的表現(xiàn)提升最為顯著，最大增幅達(dá)到288%（14.8→57.52）。

關(guān)鍵發(fā)現(xiàn)

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

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

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

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

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

因此，新方法可能需要進一步的優(yōu)化和改進，以更好地應(yīng)對這些摘要任務(wù)。

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

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

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

降低延遲與計算開銷

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

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

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

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

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

這一發(fā)現(xiàn)表明，提升LLM的長文本處理能力不僅可以通過擴展上下文窗口，還可以通過優(yōu)化模型在小窗口內(nèi)的推理能力，再結(jié)合新方法機制，實現(xiàn)高效處理長文本的目標(biāo)。

評論

這篇論文強調(diào)了一個本應(yīng)顯而易見的事實：預(yù)測和檢索是同一枚硬幣的兩面。

要有效地進行預(yù)測，首先必須確定什么是相關(guān)的。令人驚訝的是，適當(dāng)?shù)乩?/strong>注意力模式，擁有5億參數(shù)的模型可以在100萬個token上執(zhí)行完美的檢索。

這引發(fā)了一個有趣的問題：如果圍繞檢索能力明確設(shè)計架構(gòu)會怎樣？

Transformer架構(gòu)是為預(yù)測而設(shè)計的，檢索是作為副產(chǎn)品出現(xiàn)的。那么，專門為檢索優(yōu)化的架構(gòu)會是什么樣子呢？

許多資金已經(jīng)投入到構(gòu)建大規(guī)模RAG（檢索增強生成）系統(tǒng)中。

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