Instructions to use Jiarui-Wang/MemSFT-Qwen3-Bio-Memory-4B with libraries, inference providers, notebooks, and local apps. Follow these links to get started.

Libraries

How to use Jiarui-Wang/MemSFT-Qwen3-Bio-Memory-4B with Transformers:

# Use a pipeline as a high-level helper
from transformers import pipeline

pipe = pipeline("text-generation", model="Jiarui-Wang/MemSFT-Qwen3-Bio-Memory-4B")
messages = [
    {"role": "user", "content": "Who are you?"},
]
pipe(messages)

# Load model directly
from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("Jiarui-Wang/MemSFT-Qwen3-Bio-Memory-4B")
model = AutoModelForCausalLM.from_pretrained("Jiarui-Wang/MemSFT-Qwen3-Bio-Memory-4B", device_map="auto")
messages = [
    {"role": "user", "content": "Who are you?"},
]
inputs = tokenizer.apply_chat_template(
	messages,
	add_generation_prompt=True,
	tokenize=True,
	return_dict=True,
	return_tensors="pt",
).to(model.device)

outputs = model.generate(**inputs, max_new_tokens=40)
print(tokenizer.decode(outputs[0][inputs["input_ids"].shape[-1]:]))

Notebooks
Google Colab
Kaggle
Local Apps Settings

vLLM

How to use Jiarui-Wang/MemSFT-Qwen3-Bio-Memory-4B with vLLM:

Install from pip and serve model

# Install vLLM from pip:
pip install vllm
# Start the vLLM server:
vllm serve "Jiarui-Wang/MemSFT-Qwen3-Bio-Memory-4B"
# Call the server using curl (OpenAI-compatible API):
curl -X POST "http://localhost:8000/v1/chat/completions" \
	-H "Content-Type: application/json" \
	--data '{
		"model": "Jiarui-Wang/MemSFT-Qwen3-Bio-Memory-4B",
		"messages": [
			{
				"role": "user",
				"content": "What is the capital of France?"
			}
		]
	}'

Use Docker

docker model run hf.co/Jiarui-Wang/MemSFT-Qwen3-Bio-Memory-4B

SGLang

How to use Jiarui-Wang/MemSFT-Qwen3-Bio-Memory-4B with SGLang:

Install from pip and serve model

# Install SGLang from pip:
pip install sglang
# Start the SGLang server:
python3 -m sglang.launch_server \
    --model-path "Jiarui-Wang/MemSFT-Qwen3-Bio-Memory-4B" \
    --host 0.0.0.0 \
    --port 30000
# Call the server using curl (OpenAI-compatible API):
curl -X POST "http://localhost:30000/v1/chat/completions" \
	-H "Content-Type: application/json" \
	--data '{
		"model": "Jiarui-Wang/MemSFT-Qwen3-Bio-Memory-4B",
		"messages": [
			{
				"role": "user",
				"content": "What is the capital of France?"
			}
		]
	}'

Use Docker images

docker run --gpus all \
    --shm-size 32g \
    -p 30000:30000 \
    -v ~/.cache/huggingface:/root/.cache/huggingface \
    --env "HF_TOKEN=<secret>" \
    --ipc=host \
    lmsysorg/sglang:latest \
    python3 -m sglang.launch_server \
        --model-path "Jiarui-Wang/MemSFT-Qwen3-Bio-Memory-4B" \
        --host 0.0.0.0 \
        --port 30000
# Call the server using curl (OpenAI-compatible API):
curl -X POST "http://localhost:30000/v1/chat/completions" \
	-H "Content-Type: application/json" \
	--data '{
		"model": "Jiarui-Wang/MemSFT-Qwen3-Bio-Memory-4B",
		"messages": [
			{
				"role": "user",
				"content": "What is the capital of France?"
			}
		]
	}'

Docker Model Runner
How to use Jiarui-Wang/MemSFT-Qwen3-Bio-Memory-4B with Docker Model Runner:
```
docker model run hf.co/Jiarui-Wang/MemSFT-Qwen3-Bio-Memory-4B
```

MemSFT-Qwen3-Bio-Memory-4B

📄 Paper • 💻 GitHub • 🤗 HF Collection

Introduction

MemSFT specializes modern large language models with an external parametric memory. This checkpoint contains the Qwen3-4B memory trained on Biology-Instructions. The memory learns to approximate retrieval-based teacher distributions over domain SFT data. At each decoding step, a learned token-level router combines the next-token distributions of the frozen base model and memory. This checkpoint is an auxiliary memory, not a standalone chat model. Its key advantages are:

Plug-and-Play: Attaches to a frozen backbone without modifying its parameters or architecture.
Strong Specialization: Improves domain performance with negligible degradation in general capabilities.
Cross-Scale Reuse: The MemSFT design supports reuse across compatible Qwen3 backbones without retraining the memory for each backbone.

Quick Start

This example pairs the 4B memory with the Qwen3-14B backbone used in the paper. It requires a CUDA GPU with sufficient memory to load both models in BF16.

1. Install

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

conda create -n memsft-generate python=3.10 pip -y
conda activate memsft-generate
python -m pip install -e .
python -m pip install \
  "torch>=2.4,<2.7" \
  "transformers==4.51.3" \
  "huggingface-hub==0.35.3" \
  "accelerate>=0.34,<2"

2. Load the base, memory, and router

from pathlib import Path

import torch
from huggingface_hub import snapshot_download
from transformers import AutoModelForCausalLM, AutoTokenizer

from router.adaptive_memdec import AdaptiveMemoryDecoder

device = torch.device("cuda:0")
base_id = "Qwen/Qwen3-14B"
memory_id = "Jiarui-Wang/MemSFT-Qwen3-Bio-Memory-4B"
router_repo = "Jiarui-Wang/MemSFT-Qwen3-Routers"
router_subdir = "Qwen3-14B-Bio-M4B-Router"

router_root = snapshot_download(
    repo_id=router_repo,
    revision="v1.0",
    allow_patterns=[f"{router_subdir}/*"],
)
router_path = str(Path(router_root) / router_subdir)

tokenizer = AutoTokenizer.from_pretrained(
    base_id,
    revision="40c069824f4251a91eefaf281ebe4c544efd3e18",
)
base = AutoModelForCausalLM.from_pretrained(
    base_id,
    revision="40c069824f4251a91eefaf281ebe4c544efd3e18",
    torch_dtype=torch.bfloat16,
    low_cpu_mem_usage=True,
).to(device).eval()
memory = AutoModelForCausalLM.from_pretrained(
    memory_id,
    revision="v1.0",
    torch_dtype=torch.bfloat16,
    low_cpu_mem_usage=True,
).to(device).eval()

vocab_size = len(tokenizer)
base.resize_token_embeddings(vocab_size)
memory.resize_token_embeddings(vocab_size)
base.requires_grad_(False)
memory.requires_grad_(False)
model = AdaptiveMemoryDecoder(
    base_lm=base,
    knn_generator=memory,
    router_path=router_path,
    router_device=device,
).eval()
model.set_tokenizer(tokenizer)

3. Generate

sequence = (
    "MKSILIEKPNQLAIVEREIPTPSAGEVRVKVKLAGICGSDSHIYRGHNPFAKYPRVIGHEFFGVIDAV"
    "GEGVESARVGERVAVDPVVSCGHCYPCSIGKPNVCTTLAVLGVHADGGFSEYAVVPAKNAWKIPEAVA"
    "DQYAVMIEPFTIAANVTGHGQPTENDTVLVYGAGPIGLTIVQVLKGVYNVKNVIVADRIDERLEKAKE"
    "SGADWAINNSQTPLGEIFTEKGIKPTLIIDAACHPSILKEAVTLASPAARIVLMGFSSEPSEVIQQGI"
    "TGKELSIFSSRLNANKFPIVIDWLSKGLIKPEKLITHTFDFQHVADAISLFEQDQKHCCKVLLTFSE"
)
prompt = (
    r"<PROTEIN> "
    + sequence
    + r" </PROTEIN> What is the EC number associated with the enzymatic "
    r"function of this protein? Please put the final enzyme within \boxed{} "
    r"using an EC number such as ECx.x.x.x, and separate multiple entries "
    r"with commas."
)

messages = [{"role": "user", "content": prompt}]
prompt_text = tokenizer.apply_chat_template(
    messages,
    tokenize=False,
    add_generation_prompt=True,
    enable_thinking=False,
)
inputs = tokenizer(prompt_text, return_tensors="pt").to(device)

with torch.inference_mode():
    output_ids = model.generate(
        **inputs,
        do_sample=False,
        max_new_tokens=32,
        eos_token_id=tokenizer.eos_token_id,
        pad_token_id=tokenizer.eos_token_id,
    )

answer = tokenizer.decode(
    output_ids[0, inputs["input_ids"].shape[1]:],
    skip_special_tokens=True,
)
print(answer)

Example output:

\boxed{EC1.1.1.-}

UniProtKB annotates Escherichia coli K-12 RspB (P38105) with EC 1.1.1.-. For comparison, using the same prompt and deterministic generation configuration, Qwen3-14B alone predicts EC 4.2.1.22. The outputs were reproduced in BF16 on NVIDIA A800 80GB GPUs.

Performance

The memory-scale experiment uses the same Biology-Instructions training data, Qwen3-8B retrieval teacher, and frozen Qwen3-14B backbone across memory sizes.

Configuration	Memory size	Biology-Instructions ↑
Qwen3-14B	—	6.64
Qwen3-14B + MemSFT	4B	37.12

The paper reports a General average range of 83.23–83.62 across the 1.7B, 4B, and 8B memory configurations, rather than a separate General average for this checkpoint.

Compatible Pairing

base: Qwen/Qwen3-14B
memory: Jiarui-Wang/MemSFT-Qwen3-Bio-Memory-4B
router: Jiarui-Wang/MemSFT-Qwen3-Routers/Qwen3-14B-Bio-M4B-Router

Intended Use and Limitations

This checkpoint is intended for reproducing the MemSFT memory-scale experiment and for augmenting Qwen3-14B on Biology-Instructions tasks. It should be used with the matching Qwen3-14B/Bio/M4B router. This 4B memory was not evaluated with the other backbone sizes in the paper. Performance outside the evaluated pairing and domain has not been established.

License

This MemSFT checkpoint is released under the Apache License 2.0. Upstream models, software, and datasets remain subject to their respective licenses and terms.

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}, 
}