MemSFT: Mitigating Alignment Tax with an External Parametric Memory

Overview

MemSFT mitigates alignment tax by extending a frozen language model with an external parametric memory. The memory is trained to approximate the output distribution of a non-parametric retriever built over domain SFT data. It can be integrated with modern LLMs as a plug-and-play domain memory, effectively enhancing their capabilities in evolving or newly emerging domains.

At each decoding step, the frozen base and memory models process the same context in parallel. A lightweight token-level router combines their next-token distributions, increasing the memory contribution when domain knowledge is useful while retaining the base model's general capabilities. One memory can be reused across multiple compatible Qwen3 base model sizes.

This repository provides code and checkpoints for Biology-Instructions, OpenSWI, and Law. We also provide a complete OpenSWI reproduction pipeline using a Qwen3-14B base model and a Qwen3-8B memory.

MemSFT first trains the memory to approximate retrieval-based teacher distributions, then learns a token-level router that adaptively combines the frozen base and memory distributions.

Performance

Biology-Instructions

OpenSWI

OpenSWI reports RMSE, where lower is better.

Quick Start

1. Environment Setup

Clone the repository and create the evaluation environment:

git clone https://github.com/LUMIA-Group/MemSFT.git
cd MemSFT

conda create -n memsft-eval python=3.10 pip -y
conda activate memsft-eval
export PYTHONNOUSERSITE=1

python -m pip install -e .
python -m pip install -r requirements-eval.txt

MemSFT's evaluation adapters target lm-evaluation-harness==v0.4.11:

export PROJECT_ROOT="$PWD"
export LM_EVAL_ROOT=/path/to/lm-evaluation-harness

git clone --depth 1 --branch v0.4.11 \
  https://github.com/EleutherAI/lm-evaluation-harness.git \
  "$LM_EVAL_ROOT"
python -m pip install -e "${LM_EVAL_ROOT}[math]"
python scripts/utils/install_lm_eval_adapters.py \
  --lm-eval-root "$LM_EVAL_ROOT"
python scripts/utils/check_install.py --eval-only \
  --lm-eval-root "$LM_EVAL_ROOT" --openswi-tasks

For memory training, datastore construction, and router training, create the separate validated Python 3.9 environment:

conda env create -f environment.yml
conda activate memsft
export PYTHONNOUSERSITE=1

conda install -y -c pytorch -c nvidia \
  pytorch=2.4.0 pytorch-cuda=12.1 faiss-gpu=1.8.0
python -m pip install -r requirements.txt
python -m pip install -e .
python scripts/utils/check_install.py

The complete training recipe and hardware requirements are documented in MemSFT Reproduction. Keep the evaluation and training environments separate because vLLM and the training stack use different PyTorch/CUDA combinations. The pinned evaluation stack uses vLLM 0.8.5.post1, Transformers 4.51.3, and its PyTorch 2.6/CUDA 12.4 dependency set; run it in a compatible NVIDIA CUDA runtime.

2. Evaluate

Download a Qwen3 base model, the OpenSWI memory, and its router bundle:

hf download Qwen/Qwen3-14B \
  --revision 40c069824f4251a91eefaf281ebe4c544efd3e18 \
  --local-dir /path/to/Qwen3-14B

hf download Jiarui-Wang/MemSFT-Qwen3-OpenSWI-Memory-8B \
  --revision v1.0.0 \
  --local-dir /path/to/MemSFT-Qwen3-OpenSWI-Memory-8B

hf download Jiarui-Wang/MemSFT-Qwen3-Routers \
  --revision v1.0.0 \
  --include "Qwen3-14B-OpenSWI-M8B-Router/*" \
  --local-dir /path/to/MemSFT-Qwen3-Routers

Set the model, router, and local OpenSWI shallow-1K evaluation-data paths:

export QWEN3_14B=/path/to/Qwen3-14B
export MEMORY_PATH=/path/to/MemSFT-Qwen3-OpenSWI-Memory-8B
export ROUTER_PATH=/path/to/MemSFT-Qwen3-Routers/Qwen3-14B-OpenSWI-M8B-Router
export OPENSWI_DATASET_ROOT="$PROJECT_ROOT/data/openswi_shallow_1k"
export RESULTS_ROOT=/path/to/results/openswi

The repository includes the exact ordered OpenSWI shallow-1K evaluation split used in the paper. Its source, CC BY 4.0 attribution, transformations, and SHA256 are documented in data/openswi_shallow_1k/README.md.

Run one evaluation cell:

BENCHMARK=openswi SEED=42 \
  bash scripts/openswi/eval_seed.sh

Run the complete frozen evaluation protocol and aggregate the OpenSWI seeds:

bash scripts/openswi/eval_5x.sh
python scripts/openswi/aggregate_openswi_5x.py \
  --results_root "$RESULTS_ROOT"

The full data construction, memory training, router training, frozen hashes, and five-seed evaluation protocol are in docs/MEMSFT_REPRODUCTION.md.

3. Generate

The following command runs the released OpenSWI example with Qwen3-14B, the Qwen3-8B OpenSWI memory, and their matching router. It reads id_ddm=707 from the included shallow-1K split and uses the same sampling configuration as the evaluation workflow:

python scripts/generate_openswi_example.py

Example output:

[1.3487,1.3487,1.3487,1.3487,1.3487,1.3487,1.3487,1.3487,1.8542,1.8542,1.8542,1.8542,1.8542,1.8542,1.8542,1.8542,1.8542,1.8542,1.8542,1.8542,1.8542,1.8542,1.8542,1.8542,1.8542,1.8542,1.8542,1.8542,1.8542,1.8542,1.8542,1.8542,1.8542,1.8542,1.8542,1.8542,1.8542,1.8542,1.8542,1.8542,1.8542,1.8542,1.8542,1.8542,1.8542,1.8542,1.8542,1.8542,1.8542,1.8542,2.2210,2.2210,2.2210,2.2210,2.2210,2.2210,2.2210,2.2210,2.2210,2.2210,2.2210,2.2210,2.2210,2.2210,2.2210,2.2210,2.2210,2.2210,2.2210,2.1837]

This output contains the required 70 values and has an RMSE of 0.1404. OpenSWI targets a layered S-wave velocity profile, so identical values across consecutive depth layers are expected. The example requires a CUDA GPU with enough memory to load the 14B base and 8B memory in BF16.

MemSFT prefers tensor-only .safetensors router checkpoints. Legacy .pt router checkpoints are loaded with PyTorch's restricted weights_only=True mode, but checkpoints should still be obtained only from trusted sources.

Citation

If you find MemSFT helpful in your research, please consider citing:

@misc{wang2026memsftmitigatingalignmenttax,
      title={MemSFT: Mitigating Alignment Tax with an External Parametric Memory},
      author={Jiarui Wang and Xiang Shi and Jiaqi Cao and Rubin Wei and Xiquan Wang and Hao Sun and Jingzhi Wang and Zhiqi Yang and Qipeng Guo and Bowen Zhou and Zhouhan Lin},
      year={2026},
      eprint={2607.25614},
      archivePrefix={arXiv},
      primaryClass={cs.LG},
      url={https://arxiv.org/abs/2607.25614},
}

Acknowledgements

The MemSFT implementation uses code and utilities from Memory Decoder, Hugging Face Transformers, and lm-evaluation-harness. We thank the authors for making their work publicly available. See Third-Party Software and Data for attribution and dataset terms.

This work is supported by the Shanghai General AI Foundation Models Program (Grant No. 2025SHZDZX025G09). Technical collaboration: Shanghai Artificial Intelligence Laboratory.

Contact

For questions and discussions, feel free to email: wangjiarui1@sjtu.edu.cn.

License

MemSFT source code and released Qwen3-derived MemSFT checkpoints are provided under the Apache License 2.0. Third-party models, software, and datasets remain subject to their respective licenses and terms.