Jiarui-Wang

MemSFT-Qwen3-OpenSWI-Memory-8B

README

License: apache-2.0

Introduction

MemSFT specializes modern large language models with an external parametric memory. This checkpoint contains the Qwen3-8B memory trained on OpenSWI. 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: Works with Qwen3 backbones from 8B to 235B-A22B without retraining the memory for each backbone.

Quick Start

The 14B + 8B example is intended for a CUDA GPU with sufficient memory to load both models in BF16. It reads one record from the OpenSWI shallow-1K file included in the MemSFT repository.

1. Install

bash
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

python
from pathlib import Path

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

from memsft.router.adaptive_memdec import AdaptiveMemoryDecoder

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

router_root = snapshot_download(
    repo_id=router_repo,
    revision="v1.0.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.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

python
import json

example_path = Path("data/openswi_shallow_1k/shallow/test.jsonl")
with example_path.open("r", encoding="utf-8") as handle:
    example = next(
        record
        for record in map(json.loads, handle)
        if record["id_ddm"] == 707
    )

system_prompt = (
    "You are a professional geophysical inversion expert, proficient in "
    "utilizing surface wave dispersion data to infer subsurface S-wave "
    "velocity structures. Based on the provided surface wave dispersion "
    "data, perform nonlinear inversion to obtain an S-wave velocity (Vs) "
    "sequence at specified depth points."
)
format_instruction = (
    'Strict format requirements:\n'
    '1) Return exactly one valid Python list of floats.\n'
    '2) The list length must be exactly M (M = "layers in total" in the input prompt).\n'
    '3) Output must start with "[" and end with "]".\n'
    '4) Output only one line, no prefix/suffix text, no reasoning, no markdown, no code block.\n'
    'Final answer format example: [0.3123, 0.4210, ...]'
)
prompt = f"{system_prompt}\n\n{example['prompt']}\n{format_instruction}"

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)

set_seed(42)
with torch.inference_mode():
    output_ids = model.generate(
        **inputs,
        do_sample=True,
        temperature=0.6,
        top_p=0.95,
        top_k=20,
        max_new_tokens=8192,
        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:

text
[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]

For OpenSWI record id_ddm=707, the raw MemSFT 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. With the same prompt, seed, and generation settings, Qwen3-14B alone produces a 96-value list instead of the required 70-value sequence. The router's mean memory mixing weight is 0.9900, with the memory route receiving the larger weight on 99.8% of generation steps. The outputs were reproduced in BF16 on an NVIDIA A800 80GB GPU.

The example is drawn from the released OpenSWI shallow-1K evaluation file. Its source, selection, modifications, and CC BY 4.0 attribution are documented in the dataset README.

Performance

The same Qwen3-8B OpenSWI memory is reused with each base model; each pairing uses its corresponding router. OpenSWI is evaluated by RMSE, where lower is better.

Table with columns: Base model, OpenSWI RMSE ↓, General ↑
Base model	OpenSWI RMSE ↓	General ↑
Qwen3-8B + MemSFT	0.47	81.13
Qwen3-14B + MemSFT	0.47	83.80
Qwen3-32B + MemSFT	0.47	85.41
Qwen3-235B-A22B + MemSFT	0.47	87.22

The complete OpenSWI training and five-seed evaluation workflow is provided in the MemSFT reproduction guide.

Compatible Pairing

The example above uses:

base: Qwen/Qwen3-14B
memory: Jiarui-Wang/MemSFT-Qwen3-OpenSWI-Memory-8B
router: Jiarui-Wang/MemSFT-Qwen3-Routers/Qwen3-14B-OpenSWI-M8B-Router

MemSFT prefers tensor-only .safetensors router checkpoints. Legacy .pt checkpoints should be loaded only from trusted sources; the MemSFT loader uses PyTorch's restricted weights_only=True mode for compatibility.

Intended Use and Limitations

This checkpoint is intended for reproducing MemSFT and for augmenting compatible Qwen3 base models on OpenSWI surface-wave inversion tasks. It should be used with the router matching the selected base/memory pair. Performance outside the evaluated model combinations and domain has not been established. Model outputs should be reviewed by qualified domain experts before use in consequential geophysical applications.

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. OpenSWI-derived evaluation data included in the MemSFT repository is provided with attribution under the Creative Commons Attribution 4.0 International license.

Citation

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

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

Contact

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

Available on FriendliAI

Dedicated Endpoints

Run this model inference on single tenant GPU with unmatched speed and reliability at scale.

Learn more

Container

Run this model inference with full control and performance in your environment.

Learn more

Model Details

Model Provider

Jiarui-Wang

Model Tree

Base

Qwen/Qwen3-8B

Fine-tuned

this model

Input Modalities

Text

Output Modalities