PaperlayerView exact revision on arXiv

Source-linked AI summary

MemSFT: Mitigating Alignment Tax with an External Parametric Memory

Jiarui Wang, Xiang Shi, Jiaqi Cao, Rubin Wei, Xiquan Wang, Hao Sun, Jingzhi Wang, Zhiqi Yang, Qipeng Guo, Bowen Zhou, Zhouhan Lin

arXiv:2607.25614v1cs.LGcs.CL

TL;DR

Domain-specific fine-tuning can cause catastrophic forgetting and erode broadly useful post-training capabilities. MemSFT instead stores domain expertise in a reusable external parametric memory while keeping the backbone frozen, improving domain performance with negligible degradation in general capabilities across biology, geoscience, and law.

  • Problem

    Domain SFT can incur an alignment tax, with specialized gains coming at the expense of broadly useful post-training capabilities.

  • Method

    MemSFT trains a separate parametric memory to imitate retrieval-based teacher distributions and dynamically fuses its outputs with a frozen backbone through token-level routing.

  • Results

    MemSFT consistently improves domain performance with negligible degradation in general capabilities across biology, geoscience, and law, while the same memory transfers across Qwen3-8B to Qwen3-235B-A22B.

  • Takeaways & Limitations

    MemSFT offers a modular path to equip modern LLMs with specialized capabilities without repeated backbone fine-tuning.

  • Takeaways & Limitations

    The same memory currently transfers only across backbones sharing a compatible tokenizer and output vocabulary, while cross-vocabulary adaptation remains future work.

Abstract

from arXiv · show

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.

1. Introduction

MemSFT addresses the alignment tax by storing domain expertise in an external parametric memory while keeping the backbone frozen, enabling selective specialization with minimal general-capability degradation. Across biology, geoscience, and law, it improves domain performance, supports memory reuse across Qwen3 backbones, and reduces adaptation compute.

  • Motivation: Parameter updates for domain SFT can cause catastrophic forgetting of previously acquired capabilities, including instruction following, reasoning, and aligned response behavior.Post-alignment models can shift toward the broader pretraining distribution after subsequent fine-tuning.
  • Method: MemSFT externalizes domain knowledge into a separate parametric memory while keeping the backbone LLM frozen.The memory imitates a non-parametric retriever built over target-domain SFT data to internalize instruction-level expertise.
  • Results: MemSFT improves domain performance with negligible general-capability degradation, whereas SFT and LoRA suffer severe forgetting.The comparison spans biology, geoscience, and law with Qwen3 backbones at multiple scales.
  • Results: The same domain-specific memory can be reused as a plug-and-play module across Qwen3-8B to Qwen3-235B-A22B without retraining for each backbone.This reuse decouples domain specialization from backbone size.
  • Results: MemSFT requires far less adaptation compute than adapting all four Qwen3 backbones independently.The supplied passage reports 9.23… for adapting all four backbones, but the excerpt truncates the comparison.

2. Method

MemSFT decouples domain specialization from backbone parameter updates using an external parametric memory trained from retrieval-based teacher distributions. A token-level router then selectively fuses the frozen backbone and memory outputs during generation, enabling reuse across compatible backbones.

  • Architecture: MemSFT combines specialized memory training with token-level dynamic inference that augments different backbone LLMs through interpolation.The backbone parameters remain frozen while specialization-specific updates are confined to the external memory and router.
  • Datastore construction: The QA-oriented datastore retrieves similar domain contexts and converts their next-token values into non-parametric teacher distributions for answer-side training.Each teacher-forced answer position is queried using a hidden representation of the decoding context.
  • Memory training: The memory LM learns retrieval-derived domain behavior by combining KL divergence to the teacher distribution with cross-entropy on the gold SFT answer.The coefficient β ∈ [0, 1] controls the balance between the retrieval-based teacher and the gold answer.
  • Router training: A lightweight two-layer MLP router predicts λ_t from frozen-model representations and output confidence or entropy features to control memory contribution at each token.The fused distribution is p_fused(· | c_t) = (1 − λ_t) · p_base(· | c_t) + λ_t · p_mem(· | c_t).
  • Router training: Signed regularization pushes λ_t upward on domain examples and downward on general examples while cross-entropy optimizes the final next-token distribution.The router is trained on mixed general and domain instruction data, with s_t < 0 for domain examples and s_t > 0 for general examples.
  • Inference and modularity: During inference, the memory and frozen backbone process the same input in parallel, and the trained memory can be reused across frozen LLMs that share a tokenizer.The backbone supplies general reasoning and instruction-following, while the router selects domain-specialized memory predictions.

3. Experimental Setup

The experiments test MemSFT across biology, geoscience, and law using Qwen3 backbones and several adaptation baselines. Evaluation combines domain-specific metrics with five general-capability benchmarks to assess specialization and retention.

  • Domain Datasets: Experiments cover BioIns biology, OpenSWI geoscience, and DISC-Law professional-domain evaluation, spanning sequence understanding, surface-wave inversion, and legal tasks.BioIns contains 21 multi-omics tasks across DNA, RNA, protein, and multi-sequence inputs.
  • Backbones and Main Baselines: Qwen3-8B, Qwen3-14B, Qwen3-32B, and Qwen3-235B-A22B are compared with the original backbone, full-parameter SFT, and LoRA.LawBench uses a single Qwen3-14B configuration; BioIns additionally includes LoRA + MixTraining(1:1) and Wise-FT-style interpolation.
  • Training Details: BioIns models train for one epoch on 500K samples, OpenSWI models for three epochs on 30K samples, and DISC-Law SFT and LoRA use 55K samples.Maximum sequence length is 2048, except legal SFT and LoRA runs at 3072; LoRA uses rank 8, alpha 16, and dropout 0.05.
  • Evaluation Metrics: BioIns reports average task score across 21 tasks, while OpenSWI reports shallow-setting RMSE, with lower values better.BioIns aggregates classification, correlation, regression, and sequence-function scores; OpenSWI evaluation follows OpenCompass generation-style evaluation.
  • Evaluation Metrics: General capability retention is measured with lm-evaluation-harness on MATH-500, C-Eval, IFEval, MMLU-Redux, and INCLUDE.The evaluation follows the Qwen3 Technical Report settings for consistency and comparability.

4. Results

MemSFT improves specialized-domain performance while largely preserving general capabilities across biology, geoscience, and law. Its external memory scales across backbone sizes and families without rewriting the backbone.

  • BioIns results: MemSFT delivers large, stable BioIns gains across Qwen3 sizes, while full SFT’s stronger or comparable domain gains incur substantial general-capability degradation.Full SFT’s largest general-task failures occur on MATH-500 and IFEval, partly because BioIns contains sequence-centric tasks unlike natural-language instructions.
  • Scaling across Qwen3 backbones: The same 8B memory can be reused across frozen Qwen3 backbones up to Qwen3-235B-A22B, maintaining domain gains and consistent general retention.This supports MemSFT as a scalable plug-in memory module for modern LLMs.
  • Forgetting-mitigation baselines: On Qwen3-14B, targeted forgetting-mitigation baselines are evaluated against MemSFT using mixed general-domain training and Wise-FT-style weight interpolation.The passage defines LoRA + MixTraining(1:1) with 500K BioIns and 500K general instruction examples, alongside interpolation between the original and BioIns-LoRA backbones.
  • Adaptation compute: MemSFT’s adaptation-cost analysis includes a Qwen3-8B datastore-construction pass and memory-training stages, while full SFT and LoRA use 6P_actT and 4P_actT estimates.For Qwen3-235B-A22B, the estimate uses 22B activated parameters per token rather than 235B total parameters.
  • Backbone-family validation: With LLaMA2-13B and a 7B memory, MemSFT substantially improves BioIns while keeping the general average nearly unchanged, unlike full SFT.This validates the approach beyond the Qwen3 backbone family.
  • Additional domain evaluations: MemSFT achieves the lowest OpenSWI RMSE on 8B, 14B, and 32B while keeping general averages nearly unchanged, and reaches LawBench performance comparable to the strongest trainable baseline without similar forgetting.SFT and LoRA reduce OpenSWI error but degrade general capabilities; on LawBench, SFT and LoRA especially harm MATH-500 and IFEval.

5. Analysis

The analysis shows that token-wise routing avoids the domain–general capability trade-off of fixed interpolation, while larger memories improve specialization without materially reducing general performance. Routing is selective within responses, and MemSFT substantially outperforms direct BM25 retrieval on BioIns.

  • Router and interpolation: 42.92 BioIns with an 83.62 general average shows that learned token-wise routing surpasses the fixed-interpolation trade-off.Fixed interpolation reaches only 6.28 BioIns at λ=0.1 while preserving an 83.21 general average.
  • Memory size: Larger memories improve BioIns performance from 30.38 with 1.7B to 37.12 with 4B and 42.92 with 8B, while preserving general performance.The 1.7B and 4B memories improve the backbone by 23.74 and 30.48 points, respectively, with comparable general-capability retention.
  • Token-level routing: Token routing varies by role: function tokens average 0.20 memory weight, domain-bearing terms 0.43, and numerical profile tokens 0.99.This preserves backbone reasoning while exposing memory to spans that benefit from specialization.
  • Token-level routing: Across additional cases, the router responds to task role: BioIns labels are strongly memory-routed, ordinary mathematical numbers receive modest weights, and MMLU answers remain almost entirely backbone-routed.The comparison separates domain relevance from surface form rather than relying on fixed lexical categories.
  • Retrieval comparison: 42.92 MemSFT BioIns performance remains far above BM25 RAG’s 12.23 at top-50, despite both using the same 500K-example BioIns subset.Increasing top-k improves the original backbone from 6.64 to 12.23, but direct retrieval remains substantially weaker than MemSFT.

6. Related Work

Related work frames domain specialization as effective but potentially harmful to general capabilities, and distinguishes short-term from long-term memory. MemSFT extends parametric memory with a fixed, plug-and-play model for providing domain expertise.

  • Domain specialization: Domain specialization commonly uses domain-specific instruction tuning or task-oriented fine-tuning, but full-parameter SFT can erode general capabilities.The passage identifies further optimization on target-domain supervision as a practical specialization route and notes the associated risk from parameter updates.
  • Short-term memory: Short-term memory retains input-derived representations to extend accessible context, including cross-segment hidden-state reuse, compressed activations, and retrieved past representations.Examples include Transformer-XL, Compressive Transformer, Memorizing Transformers, and LongMem.
  • Long-term memory: Long-term memory preserves information beyond the current context through external datastores or learned parametric states, encompassing non-parametric retrieval and parametric memory.Retrieval-augmented language models exemplify non-parametric memory by retrieving textual evidence from external corpora at inference time.
  • MemSFT: MemSFT extends the parametric memory paradigm to modern post-trained LLMs through a plug-and-play memory model that remains fixed at inference.This memory model provides domain expertise without requiring inference-time updates to the memory itself.

7. Conclusion

MemSFT mitigates the alignment tax by decoupling domain specialization from backbone parameter updates through a plug-and-play parametric memory. Across biology, geoscience, and law, it improves domain performance with negligible general-performance degradation while reusing one domain memory across Qwen3 backbones.

  • Conclusion: MemSFT decouples domain specialization from backbone parameter updates through a plug-and-play parametric memory.The approach is modular and targets the alignment tax.
  • Conclusion: Across biology, geoscience, and law, MemSFT improves domain performance with negligible degradation in general performance.The conclusion reports consistent gains across all three domains.
  • Conclusion: The same domain memory can be reused across Qwen3-8B to Qwen3-235B-A22B backbones, avoiding repeated backbone fine-tuning.This reuse enables specialized capabilities across differently sized backbones.

8. Limitations

MemSFT’s current limitations concern portability across model families and the scope of the memory-training setting. The same domain memory is currently reusable only across backbones with compatible tokenizers and output vocabularies, while broader transfer and cross-vocabulary adaptation remain future work.

  • Portability: MemSFT currently reuses a domain memory only across backbones with compatible tokenizers and output vocabularies.This compatibility condition matches the Qwen3 experiments.
  • Portability: Transferring the same memory across unrelated model families would require additional vocabulary alignment.
  • Future work: Cross-vocabulary adaptation through brief continued training is proposed as a promising direction but left for future work.
  • Scope: The paper focuses on supervised memory.

A. Datasets and Evaluation Protocols … A.3. Router Training Data

The paper evaluates domain specialization and general-capability retention across biology, geoscience, and law using disjoint domain splits and standardized general-task protocols. Router training combines domain and general examples to learn selective memory invocation while preserving backbone behavior.

  • A. Datasets and Evaluation Protocols: Domain datasets are reused across full SFT, LoRA, and memory training so methods differ in training mechanism rather than access to domain supervision.BioIns and OpenSWI use official training splits for training and corresponding disjoint official evaluation splits.
  • A.1. Domain Specialization and Evaluation Data: Biology-Instructions contributes 500K examples covering all 21 DNA, RNA, protein, and multi-sequence task categories.After Qwen3 tokenization and truncation to 2048 tokens, the subset contains 90,019,678 input tokens and 9,003,134 answer-side supervision tokens.
  • A.1. Domain Specialization and Evaluation Data: OpenSWI contributes 30K training instances for shallow generation of near-surface S-wave velocity profiles, evaluated with RMSE.Inputs contain period, phase-velocity, and group-velocity values, while targets contain depth and subsurface properties.
  • A.1. Domain Specialization and Evaluation Data: Legal specialization uses 55,295 Qwen3-formatted examples constructed from DISC-Law-SFT files and evaluates LawBench across 19 tasks.The set uses task-balanced sampling, format-targeted supplementation, and removal of overlaps with LawBench evaluation data.
  • A.2. General Capability Evaluation: General-capability retention uses the same task configuration for each original backbone and its adapted variants across MATH-500, C-Eval, IFEval, MMLU-Redux, and INCLUDE.These benchmarks measure mathematical reasoning, Chinese knowledge and reasoning, instruction following, broad academic knowledge, and multilingual and cross-cultural knowledge, respectively.
  • A.2. General Capability Evaluation: MATH-500, IFEval, and MMLU-Redux report means over five runs with decoding seeds 42–46, whereas C-Eval and INCLUDE use deterministic evaluation.INCLUDE follows reported parameter settings where available and otherwise uses a 5-shot multiple-choice likelihood configuration without a chat template.
  • A.3. Router Training Data: Router training combines domain and general examples, using 2,250, 6,000, and 5,610 domain examples for biology, OpenSWI, and legal routers, respectively.Each router also uses 15,336 general examples sampled from NVIDIA’s Nemotron-Post-Training-Dataset-v1.

B. Datastore and Retrieval-Teacher Construction

The retrieval teacher is built from answer-side supervision using hidden-state keys from a frozen datastore encoder and FAISS nearest-neighbor search. Query positions shape retrieval context but are excluded as stored values and self-neighbors are removed from teacher targets.

  • Datastore construction: Datastore key-value pairs are constructed only at answer positions, while query labels are masked and query tokens influence hidden states only as prefix context.The query itself is not stored as a datastore value.
  • Feature extraction: The feature extractor uses the input hidden state to the final decoder block’s MLP in a frozen datastore encoder.For Qwen models, this hooks the final Transformer layer before its MLP sublayer, after attention and residual computation.
  • Retrieval teacher: FAISS L2 nearest-neighbor search constructs the retrieval-based teacher distribution.The keys are not additionally L2-normalized, so the retrieval metric is not cosine similarity.
  • Retrieval teacher: 2047 non-self neighbors contribute to each teacher distribution after retrieving the top 2048 neighbors and discarding the query’s self-match.The teacher distribution uses temperature τ = 16.0 in Eq. (3).

C. Training and Implementation Details … G. BM25 Retrieval Baseline

The paper specifies domain-specific training and routing configurations, then illustrates selective memory use across specialized and general tasks. It also documents per-task evaluations, routing ablations, forgetting-mitigation baselines, and a BM25 retrieval baseline.

  • C. Training and Implementation Details: MemSFT uses domain-specific memory training, frozen backbone and memory LMs during router training, and token-level routing informed by hidden states and uncertainty features.Maximum sequence length is 2048 except for legal SFT and LoRA runs, which use 3072; LoRA targets attention and MLP projections but not embeddings, normalization, or the language-model head.
  • C.1.1. BioIns Case: 0.93 is the mean λ for positive BioIns answer tokens, versus 0.18 for explanation tokens and 0.40 for boxed-format tokens.Memory use therefore concentrates on the specialized BioIns decision in the displayed case.
  • C.1.2. MATH Case I: Alternating Sum; C.1.3. MATH Case II: Complex Rotation: MATH examples remain backbone-dominant: alternating-sum text and mathematical-formatting tokens average 0.001 and 0.002, while complex-rotation excerpts keep every displayed token below 0.1.Numeric tokens average 0.033 in the alternating-sum case and 0.0169 in the complex-rotation case; the latter has a maximum of 0.0942.
  • C.1.4. MMLU Case I: Narcolepsy; C.1.5. MMLU Case II: Mars Greenhouse Effect: MMLU examples also remain backbone-dominant, with final D-token weights of 0.0006 for narcolepsy and 0.0005 for Mars greenhouse effect.In the Mars example, the complete reasoning-and-answer trace remains below 0.1, with a maximum of 0.0591.
  • D. Full Per-Task Domain Results; D.1. Biology-Instructions; D.2. LawBench: Per-task result tables expand Biology-Instructions for Qwen3-8B, Qwen3-14B, Qwen3-32B, and Qwen3-235B-A22B, while LawBench reports 19 evaluated subtasks.The LawBench aggregate excludes “2-10 trigger_word_extraction” because its evaluator is order-sensitive for set-like outputs.
  • E. Ablating Routing and External Memory: Fixed interpolation cannot balance specialization and general capability, whereas learned token-level routing selectively uses retrieval-trained memory while preserving the frozen backbone.The ablation compares fixed coefficients with learned routing and contrasts a standard domain SFT external module against retrieval-based memory supervision.
  • F. Forgetting-Mitigation Baselines; G. BM25 Retrieval Baseline: The paper additionally evaluates MixTraining and Wise-FT on Qwen3-14B Biology-Instructions, and tests BM25 RAG using a global index over 500K training examples.BM25 prompts prepend retrieved examples to the original evaluation prompt; Table 12 reports OpenCompass summary metrics, with Top-50 using shorter truncation.

H. Adaptation FLOPs Accounting

The FLOPs accounting estimates the training cost of full SFT, LoRA, and MemSFT, including shared datastore and memory work plus backbone-specific routing. MemSFT requires 6.73 EFLOPs for standalone Qwen3-235B-A22B specialization and 9.23 EFLOPs when reused across four backbones, while omitted FAISS operations add small measured overhead.

  • Analytical accounting: The Biology-Instructions corpus contains T=90,019,678 post-tokenization input tokens per training epoch, with estimates of 6P_actT_model FLOPs for full SFT and 4P_actT for LoRA.The LoRA estimate is an optimistic lower bound that removes the backbone parameter-gradient contribution while retaining backbone forward and activation-gradient computation.
  • Baseline adaptation costs: 11.88 EFLOPs for SFT and 7.92 EFLOPs for LoRA are estimated for standalone Qwen3-235B-A22B adaptation.Across Qwen3-8B/14B/32B, the summed active parameter count is 54B, yielding 29.17 EFLOPs for SFT and 19.44 EFLOPs for LoRA; four-backbone totals are 41.05 and 27.37 EFLOPs.
  • MemSFT accounting: 6.73 EFLOPs are required for standalone Qwen3-235B-A22B MemSFT specialization, comprising datastore construction, memory training, and router training.The components are 1.44 EFLOPs for datastore construction, 4.32 EFLOPs for memory training, and 0.96 EFLOPs for router training; reuse across four backbones raises only router cost, producing 9.23 EFLOPs.
  • Measured system overhead: 12 minutes and 3 seconds of FAISS index construction and KNN-target materialization account for 1.61 allocated GPU-hours, or 1.55% of cumulative wall-clock time and 1.76% of allocated GPU-hours.These system operations are omitted from the analytical model-FLOPs total, and their measured overhead is small relative to memory and router training.
Loading 2607.25614v1…