MemSFT: Mitigating Alignment Tax with an External Parametric Memory

📅 2026-07-28
📈 Citations: 0
Influential: 0
📄 PDF
🤖 AI Summary
This work addresses the degradation of general capabilities in large language models during domain-specific fine-tuning—a phenomenon often attributed to “alignment tax.” To mitigate this, the authors propose a plug-in parametric memory module that decouples domain knowledge acquisition from the backbone model. The module is trained via behavioral cloning to emulate a non-parametric retriever and is dynamically fused with the backbone’s output during inference through a learned router. This approach achieves, for the first time, parameter-level disentanglement between general-purpose abilities and domain-specific knowledge, enabling memory modules to be reused across different model scales without catastrophic forgetting. Experiments in biomedical, geoscience, and legal domains demonstrate substantial gains in domain performance while preserving near-original general-task proficiency, outperforming standard supervised fine-tuning baselines.
📝 Abstract
Adapting Large Language Models (LLMs) to specialized domains often incurs an alignment tax, as fine-tuning on domain-specific tasks can cause catastrophic forgetting and substantially degrade performance on general tasks. We propose MemSFT, which mitigates the alignment tax by decoupling domain specialization from backbone parameter updates through a plug-and-play parametric memory. The memory is trained to imitate the behavior of a non-parametric retriever operating over domain data, thereby memorizing knowledge and patterns that would otherwise be accessed through retrieval. Once trained on a specific domain, the memory can be reused across LLMs of different sizes. During generation, a learned router dynamically fuses the output distributions of the memory and backbone at each decoding step, allowing domain expertise to be invoked selectively. Across biology, geoscience, and law, evaluations with models ranging from Qwen3-8B to Qwen3-235B-A22B show that MemSFT consistently improves domain performance with negligible degradation in general performance, whereas full SFT suffers severe forgetting on general tasks. Overall, our results demonstrate a practical path to decoupling general model capabilities from domain-specific knowledge at the parameter level, thereby equipping LLMs with new specialized capabilities without compromising their general capabilities.
Problem

Research questions and friction points this paper is trying to address.

alignment tax
catastrophic forgetting
domain adaptation
large language models
specialized domains
Innovation

Methods, ideas, or system contributions that make the work stand out.

parametric memory
alignment tax
catastrophic forgetting
domain adaptation
plug-and-play
🔎 Similar Papers
2024-05-10arXiv.orgCitations: 2
Jiarui Wang
Jiarui Wang
Shanghai Jiao Tong University
Natural Language ProcessingMemory-Augmented LLM
X
Xiang Shi
LUMIA Lab, School of Artificial Intelligence, Shanghai Jiao Tong University
Jiaqi Cao
Jiaqi Cao
Shanghai Jiao Tong University
Natural Language ProcessingLong-term Memory
Rubin Wei
Rubin Wei
Shanghai Jiao Tong University
LLMMemory-Augmented LLM
X
Xiquan Wang
LUMIA Lab, School of Artificial Intelligence, Shanghai Jiao Tong University; Shanghai AI Laboratory
H
Hao Sun
LUMIA Lab, School of Artificial Intelligence, Shanghai Jiao Tong University; Shanghai AI Laboratory
J
Jingzhi Wang
LUMIA Lab, School of Artificial Intelligence, Shanghai Jiao Tong University
Z
Zhiqi Yang
LUMIA Lab, School of Artificial Intelligence, Shanghai Jiao Tong University
Qipeng Guo
Qipeng Guo
Fudan University
Bowen Zhou
Bowen Zhou
Chair Professor, Department of Electrical Engineering, Tsinghua University; Founder of Frontis.ai
Machine LearningNatural Language ProcessingRepresentation Learning and ReasoningConversational
Z
Zhouhan Lin
LUMIA Lab, School of Artificial Intelligence, Shanghai Jiao Tong University; Shanghai AI Laboratory