LIFT: Improving Long Context Understanding of Large Language Models through Long Input Fine-Tuning

本文提出LIFT框架，通过长输入微调和Gated Memory适配器提升短上下文LLMs的长上下文理解能力，实验显示显著性能改进。

Large Language Model, Long Context, Fine-Tuning, In-Context Learning, Parameter-Efficient Fine-Tuning, Reasoning

Yansheng Mao, Yufei Xu, Jiaqi Li, Fanxu Meng, Haotong Yang, Zilong Zheng, Xiyuan Wang, Muhan Zhang

Peking University, BIGAI

Generated by grok-3-mini-latest

Background Problem

大型语言模型（LLMs）在处理长上下文时面临重大挑战，主要由于自注意力机制的计算复杂度为二次方，导致处理长输入时计算负担过重、硬件资源消耗巨大，且难以捕获散布在长输入中的长距离依赖关系，从而影响模型在实际应用中的整体信息理解和推理性能。现有方法如检索增强生成（RAG）依赖外部数据源的检索精度，可能引入噪声或幻觉问题，而长上下文适应则需要大量计算资源，且扩展上下文窗口后生成成本仍随输入长度二次增长。此外，上下文窗口的有限性使得模型无法泛化到无限长输入。LIFT的工作起点是直接将长输入吸收进模型参数中，动态适应参数以提升短上下文模型的长上下文理解能力，而非简单扩展上下文窗口。

Method

LIFT框架的核心思想是通过长输入微调动态适应模型参数，将长输入存储在参数中以提升长上下文性能。具体方法包括：

• 分段训练： 将长输入x分割成重叠段落（长度为ℓ，偏移为s，例如s = 3/8 ℓ），以保持序列连续性，目标函数为$\mathcal{L}_{input}(\mathbf{x}; \boldsymbol{\theta}) = \sum_{k=1}^{K} \mathcal{L}_{LM}(\mathbf{x}_{l_k:r_k}; \boldsymbol{\theta})$。
• 辅助任务： 引入基于长输入合成的问答（QA）任务，目标函数为$\mathcal{L}_{AT}((\mathbf{q}_i, \mathbf{a}_i)_{i=1}^m; \boldsymbol{\theta}) = -\sum_{i=1}^m \log \mathbb{P}(\mathbf{a}_i \mid \mathbf{q}_i; \boldsymbol{\theta})$，并联合优化$\mathcal{L}(\mathbf{x}, (\mathbf{q}_i, \mathbf{a}_i)_{i=1}^m; \theta) = \mathcal{L}_{input}(\mathbf{x}; \theta) + \mathcal{L}_{AT}((\mathbf{q}_i, \mathbf{a}_i)_{i=1}^m; \theta)$。
• 上下文化训练： 修改目标函数为监督式微调格式，提供上下文c_k和提示p，优化$\mathcal{L}_{input}(\mathbf{x};\theta) = -\sum_{k=1}^{K} \log \mathbb{P}(\mathbf{x}_{l_k:r_k} \mid \text{concat}(\mathbf{c}_k, \mathbf{p}); \theta)$ 和 $\mathcal{L}_{AT}((\mathbf{q}_i, \mathbf{a}_i)_{i=1}^m; \theta) = -\sum_{i=1}^m \log \mathbb{P}(\mathbf{a}_i \mid \text{concat}(\mathbf{c}_q, \mathbf{q}_i); \theta)$，以统一训练和测试格式。
• Gated Memory架构： 一种参数高效微调（PEFT）方法，添加门控函数g和记忆函数m，每个注意力头学习动态平衡参数内知识和上下文学习，公式为$\text{attn}(\hat{q}_L, \hat{\mathbf{k}}_{1:L}, \hat{\mathbf{v}}_{1:L}) = g(\hat{q}_L) \cdot m(\hat{q}_L) + (1 - g(\hat{q}_L)) \cdot \text{attn}(\hat{q}_L, \hat{\mathbf{k}}_{l'+1:L}, \hat{\mathbf{v}}_{l'+1:L})$，端到端训练以适应长输入。

Experiment

实验在LooGLE和LongBench等基准上评估LIFT的性能，使用数据集如长短问答任务，实验设置包括比较LIFT与截断ICL方法的表现。LIFT显著提升准确率，例如在LooGLE LongQA任务上，Llama-3的GPT4分数从15.44%提高到29.97%，在LongBench的NarrativeQA和QMSum任务上也表现出改善。实验设计合理，采用消融研究验证了上下文化训练和Gated Memory的有效性（如无上下文化训练时性能下降），效率测试显示LIFT在生成长序列时解码速度更快（输入长度超过1500 token时优于ICL）。结果符合预期，证明LIFT在保持原模型能力的同时提升了长上下文理解，但也暴露了在某些任务如精确记忆时的局限性。

Further Thoughts

LIFT的理念类似于人类认知中将短时记忆转化为长时记忆，值得探索与其他测试时训练方法（如TTT）的结合，以提升模型在动态环境中的适应性；未来可扩展到多模态数据或与RAG整合以提高检索精度；同时，需解决辅助任务设计中的计算开销和过拟合问题，并通过改进Gated Memory的训练策略（如注意力蒸馏）来增强参数知识提取能力，这可能启发更泛化的持续学习框架。