大模型只能夠記憶與理解有限的上下文已經(jīng)成為大模型在真實(shí)應(yīng)用中的能力瓶頸，例如對(duì)話式AI系統(tǒng)往往無法記憶你前一天對(duì)話的內(nèi)容，利用大模型構(gòu)建智能體會(huì)產(chǎn)生前后不一致的行為與記憶。

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

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

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

實(shí)驗(yàn)結(jié)果表明，InfLLM能夠有效地?cái)U(kuò)展Mistral、LLaMA的上下文處理窗口，并在1024K上下文的海底撈針任務(wù)中實(shí)現(xiàn)100%召回。

研究背景

大規(guī)模預(yù)訓(xùn)練語言模型（LLMs）近幾年在眾多任務(wù)上取得了突破性的進(jìn)展，成為眾多應(yīng)用的基礎(chǔ)模型。

這些真實(shí)應(yīng)用也給LLMs處理超長序列的能力提出了更高的要求，例如LLM驅(qū)動(dòng)的智能體需要基于所有歷史記憶持續(xù)處理從外部環(huán)境接收的信息，對(duì)話式AI需要更好地記憶與用戶的對(duì)話內(nèi)容從而產(chǎn)生個(gè)性化回答。

然而，現(xiàn)有的大模型往往僅在長度只有數(shù)千個(gè)Token的序列上進(jìn)行預(yù)訓(xùn)練，將他們應(yīng)用到超長文本中存在兩大挑戰(zhàn)：

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

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

方法介紹

InfLLM示意圖

為了高效地實(shí)現(xiàn)大模型的長度泛化能力，作者提出了一種無需訓(xùn)練的記憶增強(qiáng)方法，InfLLM，用于流式地處理超長序列。

InfLLM旨在激發(fā)LLMs的內(nèi)在能力，以有限的計(jì)算成本捕獲超長上下文中的長距離語義依賴關(guān)系，從而實(shí)現(xiàn)高效的長文本理解。

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

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

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

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

為了應(yīng)對(duì)這些挑戰(zhàn)，上下文記憶模塊中每個(gè)記憶單元由一個(gè)語義塊構(gòu)成，一個(gè)語義塊由連續(xù)的若干Token構(gòu)成。

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

此外，作者從每個(gè)語義塊中選擇語義上最重要的Token，即接收到注意力分?jǐn)?shù)最高的Token，作為語義塊的表示，這種方法有助于避免在相關(guān)性計(jì)算中不重要Token的干擾。

(2)為了高效的內(nèi)存查找，語義塊級(jí)別的記憶單元避免了逐Token，逐注意力的相關(guān)性計(jì)算，降低了計(jì)算復(fù)雜性。

此外，語義塊級(jí)別的記憶單元確保了連續(xù)的內(nèi)存訪問，并減少了內(nèi)存加載成本。

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

考慮到大多數(shù)記憶單元的使用頻率不高，InfLLM將所有記憶單元卸載到CPU內(nèi)存上，并動(dòng)態(tài)保留頻繁使用的記憶單元放在GPU顯存中，從而顯著減少了顯存使用量。

可以將InfLLM總結(jié)為：

1. 在滑動(dòng)窗口的基礎(chǔ)上，加入遠(yuǎn)距離的上下文記憶模塊。

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

實(shí)驗(yàn)分析

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

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

超長文本實(shí)驗(yàn)

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

海底撈針實(shí)驗(yàn)結(jié)果

總結(jié)

在本文中，團(tuán)隊(duì)提出了 InfLLM，無需訓(xùn)練即可實(shí)現(xiàn) LLM 的超長文本處理拓展，并可以捕捉到長距離的語義信息。

InfLLM 在滑動(dòng)窗口的基礎(chǔ)上，增加了包含長距離上下文信息的記憶模塊，并使用緩存和offload 機(jī)制實(shí)現(xiàn)了少量計(jì)算和顯存消耗的流式長文本推理。