北大團隊提出LIFT：將長上下文知識注入模型引數，提升大模型長文字能力
机器之心發表於2025-03-17
機構: 北京大學人工智慧研究院 北京通用人工智慧研究院作者: 毛彥升 徐宇飛 李佳琪 孟繁續 楊昊桐 鄭子隆 王希元 張牧涵長文字任務是當下大模型研究的重點之一。在實際場景和應用中，普遍存在大量長序列（文字、語音、影片等），有些甚至長達百萬級 tokens。擴充模型的長文字能力不僅意味著可以在上下文視窗中裝入更長的文字，更是能夠更好地建模文字段落間資訊的長程依賴關係，增強對長文的閱讀理解和推理。現有大模型解決長文字任務的難點之一是傳統的 dot-product attention 對輸入長度呈平方複雜度，且儲存 KV cache 的開銷隨輸入長度增加，時間和空間開銷都較高。此外，模型難以真正理解散落在長文字各處資訊間的長程依賴。主流的長文字解決方法包括 Retrieval-Augmented Generation（RAG）[1]、long-context adaption 等。RAG 從長文字中抽取與問題相關的資訊放入 context window 進行推理，但它依賴準確的檢索方法，大量的噪聲和無關資訊會進一步引起模型幻覺。long-context adaption 透過在大量長文字的資料集上後訓練[2]擴充套件模型的 context window，但其推理複雜度隨文字長度平方增長、視訊記憶體佔用高，且 context window 仍然有限。為了應對長文字開銷大、難以建立長程依賴的挑戰，北京大學張牧涵團隊提出全新的框架 Long Input Fine-Tuning（LIFT）。透過將長輸入文字訓練進模型引數中，LIFT 可以使任意短上下文視窗模型獲得長文字能力。- 題目: LIFT: Improving Long Context Understanding of Large Language Models through Long Input Fine-Tuning
- 文章連結: https://arxiv.org/abs/2502.14644
LIFT 首次提出將長文字知識儲存在模型引數中，而不是外部資料庫或上下文視窗中，類比人類將 working memory 轉成 long-term memory，實現知識的內化。與此相比，我們認為無限地擴充 context window 無法真正解決長文字、長曆史的挑戰，因為無論再長的 context window 仍然有耗盡的一天，而只有將上下文持續地轉變成 parametric knowledge，才能實現無限地學習。- 動態高效的長輸入訓練。LIFT 能夠透過調整模型引數，動態適應新的長輸入文字，將其作為新的知識源，無需進行資源密集型的 long-context adaptation。針對每一篇需要處理的長文字，LIFT 透過分段的 language modeling 以及精心設計的輔助任務來微調模型，實現用模型引數來記憶和理解長文字，從而避免過長的 context 造成的推理複雜度提升和長程依賴丟失。
- 平衡模型引數知識和原有能力。由於模型原有引數（比如 Llama 3 8B）通常顯著大於記憶長文字所需的引數量，全引數微調面臨過擬合長文字而損失模型基礎能力的風險。為了在模型原有能力和微調後新的引數內知識之間找到平衡，我們提出了一種專門的引數高效微調模組——門控記憶介面卡（Gated Memory Adapter），它能平衡原始模型的 In-Context Learning（ICL）能力和 LIFT 訓練後對長輸入的記憶理解能力。
- 在流行的長上下文任務上取得了巨大提升。在幾個廣泛認可的長上下文基準集（例如 LooGLE [3]、Longbench [4]）上的評估表明，不同 LLM 始終能透過 LIFT 在常見的長/短依賴問答和摘要等通用任務上受益。例如，在非常具有挑戰性的 LooGLE 長依賴問答上，相較僅透過 ICL，LIFT 過後的 Llama 3 8B 的正確率從 15.44% 提升至 29.97%。在 LooGLE 短依賴問答上，LIFT 將 Gemma 2 9B 的正確率從 37.37% 提升至 50.33%。
受 LLM 預訓練的啟發，LIFT 將「記憶長文字」的任務建模為語言建模（Language Modeling）任務。但在整篇長文字上進行語言建模訓練開銷過大，且短上下文模型不具備直接在長文字上訓練的能力。為此，LIFT 將長文字切分為固定長度的片段，對所有片段並行進行語言建模訓練。如果將長文字切分為互不相交的片段（如圖 2 中 Trivial segmentation 所示），模型將丟失片段間的正確順序，而順序對於長文字中的長程依賴和總體理解非常重要。因此，LIFT 要求相鄰片段有一定重疊（如圖 2 中的 Our segmentation 所示）——每個片段的末尾就是下一個片段的開頭。這樣，如果模型能夠記憶某個片段，那麼它就能夠續寫出下一個片段，直到按順序續寫出全文。在實驗中，我們取重疊的長度為片段長度的 5/8，因此訓練的複雜度對長文字的長度呈線性。在特定任務上微調 LLM 通常會導致其在其他任務上能力下降。同理，長文字切段訓練可能導致 LLM 的 reasoning、instruction-following 等能力損失。研究團隊提出在合成的輔助任務上訓練，一方面彌補模型的能力損失，另一方面幫助模型學會應用長文字中的資訊回答問題。具體而言，研究團隊用預訓練的 LLM 基於長文字片段自動生成幾十個問答型別的輔助任務。於是 LIFT 訓練分為兩個階段，第一個階段只在長文字切段任務上進行語言建模訓練，第二個階段在輔助任務上訓練模型基於長文字回答問題的能力。儘管 LIFT 可以任意使用全引數微調或 LoRA/PiSSA 等引數高效微調方法來訓練模型，我們提出了一個專用的 Gated Memory Adapter 來平衡長文字記憶和能力。其核心在於用短視窗模型模擬假設長文字在上下文視窗中時模型的行為和內部表示。為此我們將假設的「全上下文」分為視窗外（out-of-context）和視窗中（in-context）兩部分——視窗外放置的是預計將透過微調放入引數中的長文字，而視窗中放置的是關於長文字的問題。我們的目的是設計一個模型，用 LIFT 引數 + 視窗外內容（短上下文）去模擬全上下文的行為，以此達成只用短上下文模型實現長上下文的注意力機制。為此，我們設計了一個門控記憶模組（Gated Memory）（見圖 3）。該模組為每個注意力層增加了兩個特殊的 MLP（圖 3 中的 Memory MLP 和 Gate MLP），均以每個位置的 query vector 為輸入，分別用於學習「視窗外部分的權重」（gate）和「視窗外部分的記憶提取內容」（memory）。這樣，當一個新的 query 進入，模型可以動態地調控其使用多少 LIFT 記憶的視窗外內容：當 gate=0，模型將恢復為純 ICL，不用任何 LIFT 記憶的資訊；當 gate=1，模型將完全依賴 LIFT 知識而忽略當前視窗中的上下文。這種動態分配機制可以有效地平衡對長文字的記憶和模型原本的 ICL 能力。LIFT 訓練過程中，我們將只微調 Gated Memory 中的引數，實現了模型在微調較小引數量的情況下，有效地記憶長文字內容並用於下游任務。為了評估 LIFT 的有效性，研究團隊在 Llama 3 8B 和 Gemma 2 9B 兩個短文字開源模型（上下文視窗為 8k）上和 GPT 3.5 商用模型（上下文視窗為 16k）上比較了 LIFT 方法和使用截斷 ICL 的 baselines。baselines 使用原模型，儘可能將長文字填入模型的上下文視窗（優先填入開頭和末尾 tokens，其餘截斷），並保證問題 prompt 全部填入。LIFT 在測試時的輸入與 baseline 相同，但使用的模型為經過 LIFT 訓練的模型，並預設使用 Gated Memory 介面卡。對於 GPT3.5，我們直接呼叫 GPT 3.5 的訓練 API。我們主要在兩個代表性的長文字評測集 LooGLE 和 LongBench 上評測，其中 LooGLE 包含大量人工標註的極具挑戰性的長依賴問答（LongQA）和 LLM 自動生成的短依賴問答（ShortQA），LongBench 包含問答、摘要等多種任務。- LIFT 極大提升了短文字模型在 LooGLE 上的表現。LIFT 穩定提升了被測模型在 ShortQA 和 LongQA 中的平均指標。值得注意的是，Llama 3 在 LongQA 上的指標從 15.44% 提升至 29.97%，Gemma 2 在 ShortQA 上的指標從 37.37% 提升至 50.33%。
- LIFT 提升了短文字模型在 Longbench 的大多數子任務上的表現。研究團隊從 LongBench 中選取了 5 個具有代表性的子任務進行測試，任務包括多篇文章間的多跳推理、閱讀理解和概括、檢索召回等，Llama 3 透過 LIFT 在其中 4 個子任務上均有提升。
- LIFT 的效果與模型的原有能力以及測試任務有關。LIFT 雖然普遍提升了模型的長文字能力，但在部分子任務上仍有改進空間。透過分析各個子任務，研究團隊認為與測試問題相似的輔助任務可以促進模型關注對測試任務有用的長上下文資訊，有助於下游任務表現。
表 2. LIFT 在 LooGLE 上的 GPT4_score 指標 表 3. LIFT 在 LongBench 上的表現（評測指標和原資料集一致） 表 4. LIFT Gated Memory 架構的消融實驗此外，我們透過消融實驗驗證了 Gated Memory 介面卡的作用。如表 4 所示，在 LooGLE ShortQA 資料集上，Gated Memory 架構相比於使用 PiSSA[5]（一種 LoRA 的改進版方法）微調的原模型，GPT-4 score 提升了 5.48%。在本文中，我們提出了一種新穎的框架——LIFT，以增強 LLMs 的長上下文理解能力。LIFT 透過高效微調模型引數，利用引數內知識（in-parameter knowledge）來動態適應長輸入，從而提升長上下文任務的能力。實驗結果表明，在 LooGLE 和 LongBench 等流行基準測試中，LIFT 顯著提升了短上下文 LLMs 在長上下文任務中的表現。然而，LIFT 仍然存在一定侷限性。首先，在 context window 不夠的情況下，我們經常需要截斷上下文來做長文字推理，但對於需要精確資訊提取的任務，如「大海撈針任務」（Needle in a Haystack），該方法仍然效能欠佳。其次，LIFT 透過將長文字輸入注入模型引數，增強了模型對資料的熟悉度，但下游任務的效果仍然依賴於模型能否自主提取和利用 LIFT 過程中獲得的引數化知識。分析表明，模型在「in-context」和「out-of-context」問題上的表現存在顯著差距，表明 LIFT 後的引數化知識提取能力仍需進一步最佳化。此外，我們發現在 LIFT 過程中引入輔助任務並不能總是顯著提高模型能力，其效能嚴重依賴下游測試任務和輔助任務的相似程度，甚至可能因過擬合而導致效能下降。因此，如何設計更通用的輔助任務是未來的研究重點。最後，儘管 Gated Memory 架構顯著提升了長文字記憶和 ICL 能力的平衡，我們發現 LIFT 後的模型仍存在對原有能力的破壞，如何設計更好的介面卡來平衡記憶和能力，也留作未來工作。LIFT 的理念非常有趣，因為人類的短期記憶也會轉化為長期記憶，這一過程類似於 LIFT 將上下文中的知識轉換為引數化知識。雖然距離徹底解決 LLMs 的長上下文挑戰仍然任重道遠，但我們的初步結果表明，LIFT 提供了一個極具潛力和前景的研究方向。我們鼓勵社群一同探索 LIFT 在更廣泛的訓練資料、更豐富的模型、更先進的輔助任務設計以及更強計算資源支援下的潛在能力。[1] Jiang, Ziyan, Xueguang Ma, and Wenhu Chen. "Longrag: Enhancing retrieval-augmented generation with long-context llms." arXiv preprint arXiv:2406.15319 (2024).[2] Chen, Yukang, Shengju Qian, Haotian Tang, Xin Lai, Zhijian Liu, Song Han, and Jiaya Jia. "LongLoRA: Efficient Fine-tuning of Long-Context Large Language Models." In The Twelfth International Conference on Learning Representations.[3] Li, Jiaqi, Mengmeng Wang, Zilong Zheng, and Muhan Zhang. "LooGLE: Can Long-Context Language Models Understand Long Contexts?." In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 16304-16333. 2024.[4] Bai, Yushi, Xin Lv, Jiajie Zhang, Hongchang Lyu, Jiankai Tang, Zhidian Huang, Zhengxiao Du et al. "LongBench: A Bilingual, Multitask Benchmark for Long Context Understanding." In ACL (1). 2024.[5] Meng, Fanxu, Zhaohui Wang, and Muhan Zhang. "PiSSA: Principal singular values and singular vectors adaptation of large language models." Advances in Neural Information Processing Systems 37 (2024): 121038-121072.[6] Hong, Junyuan, Lingjuan Lyu, Jiayu Zhou, and Michael Spranger. "Mecta: Memory-economic continual test-time model adaptation." In 2023 International Conference on Learning Representations. 2023. 
