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

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

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

代碼倉(cāng)庫(kù)：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)練語(yǔ)言模型（LLMs）近幾年在眾多任務(wù)上取得了突破性的進(jìn)展，成為眾多應(yīng)用的基礎(chǔ)模型。

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

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

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

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

方法介紹

InfLLM示意圖

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

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

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

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

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

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

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

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

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

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

此外，語(yǔ)義塊級(jí)別的記憶單元確保了連續(xù)的內(nèi)存訪問(wèn)，并減少了內(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. 將歷史上下文切分成語(yǔ)義塊，構(gòu)成上下文記憶模塊中的記憶單元。每個(gè)記憶單元通過(guò)其在之前注意力計(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)行比較，在長(zhǎng)文本數(shù)據(jù) Infinite-Bench 和 Longbench 上取得了顯著的效果提升。

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

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

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

總結(jié)

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

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