機構(gòu): 北京大學(xué)人工智能研究院 北京通用人工智能研究院

作者: 毛彥升 徐宇飛 李佳琪 孟繁續(xù) 楊昊桐 鄭子隆 王希元 張牧涵

長文本任務(wù)是當下大模型研究的重點之一。在實際場景和應(yīng)用中，普遍存在大量長序列（文本、語音、視頻等），有些甚至長達百萬級 tokens。擴充模型的長文本能力不僅意味著可以在上下文窗口中裝入更長的文本，更是能夠更好地建模文本段落間信息的長程依賴關(guān)系，增強對長文的閱讀理解和推理。

現(xiàn)有大模型解決長文本任務(wù)的難點之一是傳統(tǒng)的 dot-product attention 對輸入長度呈平方復(fù)雜度，且存儲 KV cache 的開銷隨輸入長度增加，時間和空間開銷都較高。

此外，模型難以真正理解散落在長文本各處信息間的長程依賴。主流的長文本解決方法包括 Retrieval-Augmented Generation（RAG）[1]、long-context adaption 等。

RAG 從長文本中抽取與問題相關(guān)的信息放入 context window 進行推理，但它依賴準確的檢索方法，大量的噪聲和無關(guān)信息會進一步引起模型幻覺。

long-context adaption 通過在大量長文本的數(shù)據(jù)集上后訓(xùn)練[2]擴展模型的 context window，但其推理復(fù)雜度隨文本長度平方增長、顯存占用高，且 context window 仍然有限。

為了應(yīng)對長文本開銷大、難以建立長程依賴的挑戰(zhàn)，北京大學(xué)張牧涵團隊提出全新的框架 Long Input Fine-Tuning（LIFT）。通過將長輸入文本訓(xùn)練進模型參數(shù)中，LIFT 可以使任意短上下文窗口模型獲得長文本能力。

• 題目: LIFT: Improving Long Context Understanding of Large Language Models through Long Input Fine-Tuning
• 文章鏈接: https://arxiv.org/abs/2502.14644

表 1 是 LIFT 和現(xiàn)有常見方法的對比。

表 1. LIFT 與傳統(tǒng)長文本處理方法的對比

LIFT 首次提出將長文本知識存儲在模型參數(shù)中，而不是外部數(shù)據(jù)庫或上下文窗口中，類比人類將 working memory 轉(zhuǎn)成 long-term memory，實現(xiàn)知識的內(nèi)化。

與此相比，我們認為無限地擴充 context window 無法真正解決長文本、長歷史的挑戰(zhàn)，因為無論再長的 context window 仍然有耗盡的一天，而只有將上下文持續(xù)地轉(zhuǎn)變成 parametric knowledge，才能實現(xiàn)無限地學(xué)習(xí)。

研究創(chuàng)新

我們的方案具有以下優(yōu)勢：

• 動態(tài)高效的長輸入訓(xùn)練。LIFT 能夠通過調(diào)整模型參數(shù)，動態(tài)適應(yīng)新的長輸入文本，將其作為新的知識源，無需進行資源密集型的 long-context adaptation。針對每一篇需要處理的長文本，LIFT 通過分段的 language modeling 以及精心設(shè)計的輔助任務(wù)來微調(diào)模型，實現(xiàn)用模型參數(shù)來記憶和理解長文本，從而避免過長的 context 造成的推理復(fù)雜度提升和長程依賴丟失。
• 平衡模型參數(shù)知識和原有能力。由于模型原有參數(shù)（比如 Llama 3 8B）通常顯著大于記憶長文本所需的參數(shù)量，全參數(shù)微調(diào)面臨過擬合長文本而損失模型基礎(chǔ)能力的風(fēng)險。為了在模型原有能力和微調(diào)后新的參數(shù)內(nèi)知識之間找到平衡，我們提出了一種專門的參數(shù)高效微調(diào)模塊——門控記憶適配器（Gated Memory Adapter），它能平衡原始模型的 In-Context Learning（ICL）能力和 LIFT 訓(xùn)練后對長輸入的記憶理解能力。
• 在流行的長上下文任務(wù)上取得了巨大提升。在幾個廣泛認可的長上下文基準集（例如 LooGLE [3]、Longbench [4]）上的評估表明，不同 LLM 始終能通過 LIFT 在常見的長/短依賴問答和摘要等通用任務(wù)上受益。例如，在非常具有挑戰(zhàn)性的 LooGLE 長依賴問答上，相較僅通過 ICL，LIFT 過后的 Llama 3 8B 的正確率從 15.44% 提升至 29.97%。在 LooGLE 短依賴問答上，LIFT 將 Gemma 2 9B 的正確率從 37.37% 提升至 50.33%。

圖 1.LIFT 框架以及和現(xiàn)有方法對比

LIFT 方法

長文本切段訓(xùn)練

受 LLM 預(yù)訓(xùn)練的啟發(fā)，LIFT 將「記憶長文本」的任務(wù)建模為語言建模（Language Modeling）任務(wù)。但在整篇長文本上進行語言建模訓(xùn)練開銷過大，且短上下文模型不具備直接在長文本上訓(xùn)練的能力。為此，LIFT 將長文本切分為固定長度的片段，對所有片段并行進行語言建模訓(xùn)練。

如果將長文本切分為互不相交的片段（如圖 2 中 Trivial segmentation 所示），模型將丟失片段間的正確順序，而順序?qū)τ陂L文本中的長程依賴和總體理解非常重要。因此，LIFT 要求相鄰片段有一定重疊（如圖 2 中的 Our segmentation 所示）——每個片段的末尾就是下一個片段的開頭。

這樣，如果模型能夠記憶某個片段，那么它就能夠續(xù)寫出下一個片段，直到按順序續(xù)寫出全文。在實驗中，我們?nèi)≈丿B的長度為片段長度的 5/8，因此訓(xùn)練的復(fù)雜度對長文本的長度呈線性。

圖 2. LIFT 的文章切段方法

輔助任務(wù)訓(xùn)練

在特定任務(wù)上微調(diào) LLM 通常會導(dǎo)致其在其他任務(wù)上能力下降。同理，長文本切段訓(xùn)練可能導(dǎo)致 LLM 的 reasoning、instruction-following 等能力損失。

研究團隊提出在合成的輔助任務(wù)上訓(xùn)練，一方面彌補模型的能力損失，另一方面幫助模型學(xué)會應(yīng)用長文本中的信息回答問題。具體而言，研究團隊用預(yù)訓(xùn)練的 LLM 基于長文本片段自動生成幾十個問答類型的輔助任務(wù)。

于是 LIFT 訓(xùn)練分為兩個階段，第一個階段只在長文本切段任務(wù)上進行語言建模訓(xùn)練，第二個階段在輔助任務(wù)上訓(xùn)練模型基于長文本回答問題的能力。

Gated Memory 架構(gòu)

盡管 LIFT 可以任意使用全參數(shù)微調(diào)或 LoRA/PiSSA 等參數(shù)高效微調(diào)方法來訓(xùn)練模型，我們提出了一個專用的 Gated Memory Adapter 來平衡長文本記憶和能力。其核心在于用短窗口模型模擬假設(shè)長文本在上下文窗口中時模型的行為和內(nèi)部表示。

為此我們將假設(shè)的「全上下文」分為窗口外（out-of-context）和窗口中（in-context）兩部分——窗口外放置的是預(yù)計將通過微調(diào)放入?yún)?shù)中的長文本，而窗口中放置的是關(guān)于長文本的問題。

我們的目的是設(shè)計一個模型，用 LIFT 參數(shù) + 窗口外內(nèi)容（短上下文）去模擬全上下文的行為，以此達成只用短上下文模型實現(xiàn)長上下文的注意力機制。

為此，我們設(shè)計了一個門控記憶模塊（Gated Memory）（見圖 3）。該模塊為每個注意力層增加了兩個特殊的 MLP（圖 3 中的 Memory MLP 和 Gate MLP），均以每個位置的 query vector 為輸入，分別用于學(xué)習(xí)「窗口外部分的權(quán)重」（gate）和「窗口外部分的記憶提取內(nèi)容」（memory）。

這樣，當一個新的 query 進入，模型可以動態(tài)地調(diào)控其使用多少 LIFT 記憶的窗口外內(nèi)容：當 gate=0，模型將恢復(fù)為純 ICL，不用任何 LIFT 記憶的信息；當 gate=1，模型將完全依賴 LIFT 知識而忽略當前窗口中的上下文。

這種動態(tài)分配機制可以有效地平衡對長文本的記憶和模型原本的 ICL 能力。LIFT 訓(xùn)練過程中，我們將只微調(diào) Gated Memory 中的參數(shù)，實現(xiàn)了模型在微調(diào)較小參數(shù)量的情況下，有效地記憶長文本內(nèi)容并用于下游任務(wù)。

實驗證明了這一結(jié)構(gòu)的有效性（見下文表 4）。

圖 3.Gated Memory 模塊

實驗測評

為了評估 LIFT 的有效性，研究團隊在 Llama 3 8B 和 Gemma 2 9B 兩個短文本開源模型（上下文窗口為 8k）上和 GPT 3.5 商用模型（上下文窗口為 16k）上比較了 LIFT 方法和使用截斷 ICL 的 baselines。

baselines 使用原模型，盡可能將長文本填入模型的上下文窗口（優(yōu)先填入開頭和末尾 tokens，其余截斷），并保證問題 prompt 全部填入。LIFT 在測試時的輸入與 baseline 相同，但使用的模型為經(jīng)過 LIFT 訓(xùn)練的模型，并默認使用 Gated Memory 適配器。

對于 GPT3.5，我們直接調(diào)用 GPT 3.5 的訓(xùn)練 API。我們主要在兩個代表性的長文本評測集 LooGLE 和 LongBench 上評測，其中 LooGLE 包含大量人工標注的極具挑戰(zhàn)性的長依賴問答（LongQA）和 LLM 自動生成的短依賴問答（ShortQA），LongBench 包含問答、摘要等多種任務(wù)。

結(jié)果如表 2、表 3 所示，實驗表明：

• LIFT 極大提升了短文本模型在 LooGLE 上的表現(xiàn)。LIFT 穩(wěn)定提升了被測模型在 ShortQA 和 LongQA 中的平均指標。值得注意的是，Llama 3 在 LongQA 上的指標從 15.44% 提升至 29.97%，Gemma 2 在 ShortQA 上的指標從 37.37% 提升至 50.33%。
• LIFT 提升了短文本模型在 Longbench 的大多數(shù)子任務(wù)上的表現(xiàn)。研究團隊從 LongBench 中選取了 5 個具有代表性的子任務(wù)進行測試，任務(wù)包括多篇文章間的多跳推理、閱讀理解和概括、檢索召回等，Llama 3 通過 LIFT 在其中 4 個子任務(wù)上均有提升。
• LIFT 的效果與模型的原有能力以及測試任務(wù)有關(guān)。LIFT 雖然普遍提升了模型的長文本能力，但在部分子任務(wù)上仍有改進空間。通過分析各個子任務(wù)，研究團隊認為與測試問題相似的輔助任務(wù)可以促進模型關(guān)注對測試任務(wù)有用的長上下文信息，有助于下游任務(wù)表現(xiàn)。

表 2. LIFT 在 LooGLE 上的 GPT4_score 指標

表 3. LIFT 在 LongBench 上的表現(xiàn)（評測指標和原數(shù)據(jù)集一致）

表 4. LIFT Gated Memory 架構(gòu)的消融實驗

此外，我們通過消融實驗驗證了 Gated Memory 適配器的作用。如表 4 所示，在 LooGLE ShortQA 數(shù)據(jù)集上，Gated Memory 架構(gòu)相比于使用 PiSSA[5]（一種 LoRA 的改進版方法）微調(diào)的原模型，GPT-4 score 提升了 5.48%。

總結(jié)、展望和討論

在本文中，我們提出了一種新穎的框架——LIFT，以增強 LLMs 的長上下文理解能力。LIFT 通過高效微調(diào)模型參數(shù)，利用參數(shù)內(nèi)知識（in-parameter knowledge）來動態(tài)適應(yīng)長輸入，從而提升長上下文任務(wù)的能力。實驗結(jié)果表明，在 LooGLE 和 LongBench 等流行基準測試中，LIFT 顯著提升了短上下文 LLMs 在長上下文任務(wù)中的表現(xiàn)。

然而，LIFT 仍然存在一定局限性。首先，在 context window 不夠的情況下，我們經(jīng)常需要截斷上下文來做長文本推理，但對于需要精確信息提取的任務(wù)，如「大海撈針任務(wù)」（Needle in a Haystack），該方法仍然性能欠佳。

其次，LIFT 通過將長文本輸入注入模型參數(shù)，增強了模型對數(shù)據(jù)的熟悉度，但下游任務(wù)的效果仍然依賴于模型能否自主提取和利用 LIFT 過程中獲得的參數(shù)化知識。分析表明，模型在「in-context」和「out-of-context」問題上的表現(xiàn)存在顯著差距，表明 LIFT 后的參數(shù)化知識提取能力仍需進一步優(yōu)化。

此外，我們發(fā)現(xiàn)在 LIFT 過程中引入輔助任務(wù)并不能總是顯著提高模型能力，其性能嚴重依賴下游測試任務(wù)和輔助任務(wù)的相似程度，甚至可能因過擬合而導(dǎo)致性能下降。因此，如何設(shè)計更通用的輔助任務(wù)是未來的研究重點。

最后，盡管 Gated Memory 架構(gòu)顯著提升了長文本記憶和 ICL 能力的平衡，我們發(fā)現(xiàn) LIFT 后的模型仍存在對原有能力的破壞，如何設(shè)計更好的適配器來平衡記憶和能力，也留作未來工作。

LIFT 的理念非常有趣，因為人類的短期記憶也會轉(zhuǎn)化為長期記憶，這一過程類似于 LIFT 將上下文中的知識轉(zhuǎn)換為參數(shù)化知識。雖然距離徹底解決 LLMs 的長上下文挑戰(zhàn)仍然任重道遠，但我們的初步結(jié)果表明，LIFT 提供了一個極具潛力和前景的研究方向。

我們鼓勵社區(qū)一同探索 LIFT 在更廣泛的訓(xùn)練數(shù)據(jù)、更豐富的模型、更先進的輔助任務(wù)設(shè)計以及更強計算資源支持下的潛在能力。