大模型只能夠記憶與理解有限的上下文已經成為大模型在真實應用中的能力瓶頸，例如對話式AI系統往往無法記憶你前一天對話的內容，利用大模型構建智能體會產生前后不一致的行為與記憶。

為了讓大模型能夠記憶并處理更長的上下文，來自清華大學、麻省理工學院和人民大學的研究人員聯合提出無需額外訓練的大模型長文本理解方法 InfLLM，利用少量計算和顯存開銷實現了 LLM的超長文本處理。

論文地址：https://arxiv.org/abs/2402.04617

代碼倉庫：https://github.com/thunlp/InfLLM

實驗結果表明，InfLLM能夠有效地擴展Mistral、LLaMA的上下文處理窗口，并在1024K上下文的海底撈針任務中實現100%召回。

研究背景

大規模預訓練語言模型（LLMs）近幾年在眾多任務上取得了突破性的進展，成為眾多應用的基礎模型。

這些真實應用也給LLMs處理超長序列的能力提出了更高的要求，例如LLM驅動的智能體需要基于所有歷史記憶持續處理從外部環境接收的信息，對話式AI需要更好地記憶與用戶的對話內容從而產生個性化回答。

然而，現有的大模型往往僅在長度只有數千個Token的序列上進行預訓練，將他們應用到超長文本中存在兩大挑戰：

1. 分布外長度：直接將LLMs應用到更長長度的文本中，往往需要LLMs處理超過訓練范圍的位置編碼，從而造成Out-of-Distribution問題，無法泛化；

2. 注意力干擾：過長的上下文將使模型注意力被過度分散到無關的信息中，從而無法有效建模上下文中遠距離語義依賴。

方法介紹

InfLLM示意圖

為了高效地實現大模型的長度泛化能力，作者提出了一種無需訓練的記憶增強方法，InfLLM，用于流式地處理超長序列。

InfLLM旨在激發LLMs的內在能力，以有限的計算成本捕獲超長上下文中的長距離語義依賴關系，從而實現高效的長文本理解。

整體框架：考慮到長文本注意力的稀疏性，處理每個Token通常只需要其上下文的一小部分。

作者構建了一個外部記憶模塊，用于存儲超長上下文信息；采用滑動窗口機制，每個計算步驟，只有與當前Token距離相近的Tokens（Local Tokens）和外部記憶模塊中的少量相關信息參與到注意力層的計算中，而忽略其他不相關的噪聲。

因此，LLMs可以使用有限的窗口大小來理解整個長序列，并避免引入噪聲。

然而，超長序列中的海量上下文對于記憶模塊中有效的相關信息定位和記憶查找效率帶來了重大挑戰。

為了應對這些挑戰，上下文記憶模塊中每個記憶單元由一個語義塊構成，一個語義塊由連續的若干Token構成。

具體而言， (1)為了有效定位相關記憶單元，每個語義塊的連貫語義比碎片化的Token更能有效滿足相關信息查詢的需求。

此外，作者從每個語義塊中選擇語義上最重要的Token，即接收到注意力分數最高的Token，作為語義塊的表示，這種方法有助于避免在相關性計算中不重要Token的干擾。

(2)為了高效的內存查找，語義塊級別的記憶單元避免了逐Token，逐注意力的相關性計算，降低了計算復雜性。

此外，語義塊級別的記憶單元確保了連續的內存訪問，并減少了內存加載成本。

得益于此，作者設計了一種針對上下文記憶模塊的高效卸載機制（Offloading）。

考慮到大多數記憶單元的使用頻率不高，InfLLM將所有記憶單元卸載到CPU內存上，并動態保留頻繁使用的記憶單元放在GPU顯存中，從而顯著減少了顯存使用量。

可以將InfLLM總結為：

1. 在滑動窗口的基礎上，加入遠距離的上下文記憶模塊。

2. 將歷史上下文切分成語義塊，構成上下文記憶模塊中的記憶單元。每個記憶單元通過其在之前注意力計算中的注意力分數確定代表性Token，作為記憶單元的表示。從而避免上下文中的噪音干擾，并降低記憶查詢復雜度

實驗分析

作者在 Mistral-7b-Inst-v0.2（32K） 和 Vicuna-7b-v1.5（4K）模型上應用 InfLLM，分別使用4K和2K的局部窗口大小。

與原始模型、位置編碼內插、Infinite-LM以及StreamingLLM進行比較，在長文本數據 Infinite-Bench 和 Longbench 上取得了顯著的效果提升。

超長文本實驗

此外，作者繼續探索了 InfLLM 在更長文本上的泛化能力，在 1024K 長度的「海底撈針」任務中仍能保持 100% 的召回率。

海底撈針實驗結果

總結

在本文中，團隊提出了 InfLLM，無需訓練即可實現 LLM 的超長文本處理拓展，并可以捕捉到長距離的語義信息。

InfLLM 在滑動窗口的基礎上，增加了包含長距離上下文信息的記憶模塊，并使用緩存和offload 機制實現了少量計算和顯存消耗的流式長文本推理。