MemSFT: Mitigating Alignment Tax with an External Parametric Memory

Adapting Large Language Models (LLMs) to specialized domains like biology, law, or geoscience often leads to the 'alignment tax'—a phenomenon where fine-tuning improves domain-specific performance but causes catastrophic forgetting of the model's original general reasoning and instruction-following abilities. This paper investigates whether this trade-off can be bypassed by decoupling domain knowledge from the backbone model parameters.

The authors propose MemSFT, a modular architecture that keeps the backbone LLM frozen while offloading domain-specific knowledge to an external, trainable parametric memory. The memory is trained to mimic a non-parametric retriever that operates over domain-specific data, effectively distilling retrieval-based knowledge into a compact module. During inference, a learned router dynamically fuses the output distributions of the frozen backbone and the memory at each token, allowing the model to selectively invoke domain expertise only when necessary. A key feature of this design is its modularity: once a memory is trained for a specific domain, it can be reused across different sizes of the same LLM family without retraining.

Experiments across biology (BioIns), geoscience (OpenSWI), and law (LawBench) demonstrate that MemSFT consistently outperforms traditional fine-tuning methods. While full supervised fine-tuning (SFT) and LoRA show significant degradation in general capabilities as domain performance increases, MemSFT achieves substantial domain-specific gains with negligible impact on general benchmarks like MMLU-Redux and MATH-500. Furthermore, the authors show that a single 8B memory module can be successfully plugged into various backbone sizes, ranging from 8B to 235B parameters, significantly reducing the computational cost of domain adaptation compared to full-parameter fine-tuning.