Mask the Target: A Plug-and-Play Regularizer Against LoRA Forgetting

📄 arXiv: 2605.29498v1 📥 PDF

作者: Runze Xu, Arpit Garg, Hemanth Saratchandran, Simon Lucey

分类: cs.CL, cs.CV

发布日期: 2026-05-28

备注: In Submission


💡 一句话要点

提出Mask the Target正则化方法,解决LoRA微调中的灾难性遗忘问题

🎯 匹配领域: 支柱九:具身大模型 (Embodied Foundation Models)

关键词: 低秩适应 LoRA微调 灾难性遗忘 正则化 知识保留

📋 核心要点

  1. LoRA微调在提升特定任务性能的同时,容易导致模型遗忘预训练阶段学到的通用知识,尤其是在目标分布与预训练分布差异较大时。
  2. 论文提出一种名为“Mask the Target”的输出空间正则化方法,通过在损失函数层面屏蔽ground-truth token,保留模型对其他token的相对偏好。
  3. 该方法无需重放数据或修改模型架构,可直接应用于现有LoRA变体,并在多种LoRA变体和backbone上验证了其有效性,提升了新学习和遗忘之间的平衡。

📝 摘要(中文)

低秩适应(LoRA)已成为将大型语言模型适应于新领域、任务和用户的最广泛使用的微调机制之一。然而,仅仅是适应性能可能会掩盖一个重要的失败模式:LoRA更新可能提高目标分布的性能,同时降低预训练和对齐期间学习到的先前能力。我们表明,当适应分布与模型的原始训练或对齐分布有很大差异时,这种遗忘变得尤其严重。在实际环境中,原始训练和对齐数据通常不可用,这加剧了挑战。受此约束的推动,我们研究了基于LoRA的适应如何在无重放设置中平衡新学习与遗忘,并引入了一个简单的输出空间正则化器,可以直接添加到现有的训练流程中。我们的方法从基础模型和适应模型的分布中移除ground-truth token,重新归一化剩余的概率,并仅对非目标词汇应用KL正则化。这保留了基础模型在替代token之间的相对偏好,而不会直接反对适应所需的交叉熵信号。由于正则化器仅在损失层面起作用,因此不需要重放数据、架构更改、适配器重新设计或推理时开销,并且可以直接应用于现有的LoRA变体。在我们测试的所有LoRA变体和各种backbone上,当适应分布与基础模型的原始训练或对齐分布有很大差异时,我们的方法改进了新学习和遗忘之间的平衡,这表明了一种广泛适用的、更可靠的LLM更新途径。

🔬 方法详解

问题定义:LoRA微调虽然能有效使LLM适应特定任务,但会引起灾难性遗忘,即模型在目标任务上表现提升的同时,丧失了在原始预训练任务上的能力。尤其当微调数据分布与预训练数据分布差异较大时,这种遗忘现象更为严重。现有方法通常依赖于重放原始数据,但在实际应用中,原始数据往往不可获取。

核心思路:论文的核心思路是在输出空间进行正则化,通过屏蔽ground-truth token,避免LoRA微调过度拟合目标任务,从而保留模型在预训练阶段学到的知识。具体来说,该方法不是直接对抗交叉熵损失,而是保留模型对非目标token的相对偏好,从而在学习新知识的同时,避免过度遗忘原有知识。

技术框架:该方法是一个即插即用的正则化器,可以直接添加到现有的LoRA训练流程中。其主要步骤包括:1) 从基础模型和微调模型的输出分布中移除ground-truth token;2) 对剩余的概率进行重新归一化;3) 仅对非目标词汇应用KL散度正则化。整个过程在损失函数层面进行,无需修改模型架构或引入额外的推理开销。

关键创新:该方法最重要的创新点在于其正则化策略。与传统的正则化方法不同,该方法不是直接惩罚模型对目标任务的拟合,而是通过屏蔽ground-truth token,保留模型对其他token的相对偏好。这种策略能够在学习新知识的同时,有效缓解灾难性遗忘问题。

关键设计:该方法的关键设计在于KL散度正则化的应用。KL散度用于衡量基础模型和微调模型在非目标词汇上的分布差异。通过最小化这种差异,可以促使微调模型保留基础模型在预训练阶段学到的知识。此外,该方法无需调整额外的超参数,可以直接应用于现有的LoRA变体。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明,在多种LoRA变体和backbone上,该方法均能有效提升新学习和遗忘之间的平衡。具体来说,在适应分布与原始训练分布差异较大的情况下,该方法能够显著降低模型的遗忘程度,同时保持甚至提升在目标任务上的性能。该方法无需重放数据或修改模型架构,具有很强的实用性和通用性。

🎯 应用场景

该研究成果可广泛应用于各种需要对大型语言模型进行微调的场景,例如:特定领域的文本生成、对话系统、代码生成等。通过缓解灾难性遗忘问题,可以提高微调后模型的泛化能力和鲁棒性,使其在适应新任务的同时,保持原有的知识和能力。这对于构建可靠且可持续更新的LLM至关重要。

📄 摘要(原文)

Low-Rank Adaptation (LoRA) has become one of the most widely used fine-tuning mechanisms for adapting large language models to new domains, tasks, and users. Yet adaptation performance alone can obscure an important failure mode: LoRA updates may improve performance on the target distribution while degrading prior capabilities learned during pretraining and alignment. We show that this forgetting becomes especially severe when the adaptation distribution differs substantially from the models original training or alignment distributions. The challenge is amplified in practical settings, where the original training and alignment data are typically unavailable. Motivated by this constraint, we study how LoRA based adaptation balances new learning against forgetting in a replay-free setting, and introduce a simple output space regularizer that can be added directly to existing training pipelines. Our method removes the ground-truth token from both the base and adapted model distributions, renormalizes the remaining probabilities, and applies KL regularization only over the non-target vocabulary. This preserves the base models relative preferences among alternative tokens without directly opposing the cross-entropy signal required for adaptation. As the regularizer acts only at the loss level, it requires no replay data, architectural changes, adapter redesign, or inference-time overhead, and can be applied directly to existing LoRA variants. Across all LoRA variants tested and across various backbones, our method improves the frontier between new learning and forgetting when the adaptation distribution differs substantially from the base models original training or alignment distributions, suggesting a broadly applicable route toward more reliable LLM updating.