Research map / Label-efficient and robust learning / Recognition and long-tail
Transferring Linear Features Across Language Models With Model Stitching
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Learns affine stitches between residual streams of same-family LMs of different sizes and uses them to transfer SAEs, probes, and steering vectors, cutting SAE training FLOPs by about 50% via transferred initialization.
Problem
Whether linear representation spaces of language models of different sizes in the same family are similar enough to transfer interpretable features (SAEs, probes, steering vectors) between them, and whether this can reduce SAE training cost.
Approach
Train two affine mappings (up and down) between residual streams of two models with an MSE reconstruction loss plus an inversion penalty, then analytically reparameterize a trained SAE from the small model into an SAE for the large model and use it as initialization.
Key results
- Transferred SAE initialization reaches ~90% explained variance threshold in roughly 30-50% fewer FLOPs than training from scratch on Pythia-160m, including stitch training cost.
- Transferred probes without retraining beat random features on average across 8 SAEBench datasets; retraining on transferred features nearly recovers ground-truth probe performance.
- Transferred response-language steering vectors from gemma-2-9b to gemma-2-2b work for some languages (e.g., de-en 0.79 vs 0.81 ground truth) but poorly for others; stitch next-token loss is near ground truth in the A-to-B-to-A inversion.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Initialize the larger SAE with the transferred smaller-model SAE instead of random initialization | SAE training on Pythia-160m · Pythia-160m, explained-variance thresholds around 90% | FLOPs savings: 30%-50% less FLOPs | Helped |
| Transfer a small-model SAE without further training | Pythia-160m, GPT2-medium, and Gemma-2-9B · Zero-shot SAE transfer | FUV: 0.17 / 0.52; 0.11 / 0.39; 0.21 / 0.42 | Hurt |
| Transfer a steering vector between model sizes | Gemma-2 language steering · EuroParl target-language generation | target-language response accuracy: 0.79; 0.69 | Mixed |
| Retrain the probe on transferred SAE features | Pythia-70m to Pythia-160m sparse probing · 8 SAEBench datasets | test accuracy | Helped |
| Use transferred vectors for IFEval instruction-following steering rather than ground-truth vectors | Gemma-2-2B and Gemma-2-9B instruction-tuned models · IFEval steering | instruction-following performance | Hurt |
| Add an inversion penalty to the affine stitch objective (ablation) | Model stitching · Downstream next-token fidelity | next-token CE loss | Helped |
| Set inversion-penalty strength to α = 1.0 (ablation) | Model stitching · Inversion ablation | downstream next-token CE loss | Helped |
| Use stitch initialization instead of scratch initialization (ablation) | GPT2-medium SAE training · Approximate SAE scaling law | A coefficient: 42.3 → 5.0 | Helped |
| Use stitch initialization instead of scratch initialization (ablation) | Pythia-160m SAE training · Approximate SAE scaling law | A coefficient: 41.2 → 18.5 | Helped |
| Transfer attention deactivation feature from GPT2-small to GPT2-medium | GPT2 attention feature · Attention deactivation and effective null-space composition | effective null-space composition: 0.27 → 0.77 | Mixed |
Surprising findings
- Structural features transfer more consistently than semantic features, while semantic features are polarized between transferring well and not transferring.
- Transferred zero-shot SAEs are imperfect (FUV 0.21 to 0.42 for Gemma), yet the stitched initialization still speeds up SAE training.
Limitations
- Stitches trained only on general internet text within same family and tokenizer; cross-family transfer not tested.
- Scaling-law fits lack irreducible loss term due to limited high-compute data.
- Steering on instruction following with chat templates is weak because stitches are not trained on chat-template tokens.
Experimental setup
Datasets: OpenWebText EuroParl IFEval SAEBench probing datasets
Benchmarks: OpenWebText SAEBench sparse probing EuroParl language steering IFEval
Metrics: next-token cross-entropy loss fraction unexplained variance delta loss explained variance probe accuracy percent responses in target language attribution correlation FLOPs to threshold
Baselines: SAE trained from scratch random features probing dense probe ground truth steering
Builds on: model stitching SVCCA TopK SAE JumpReLU SAE SAELens
Closest papers by meaning
- Low-Rank Adapting Models for Sparse Autoencoders (ICML 2025 · similarity 0.90)
- Scaling and evaluating sparse autoencoders (ICLR 2025 · similarity 0.89)
- Improving Sparse Decomposition of Language Model Activations with Gated Sparse Autoencoders (NeurIPS 2024 · similarity 0.89)
- SAEBench: A Comprehensive Benchmark for Sparse Autoencoders in Language Model Interpretability (ICML 2025 · similarity 0.89)
- Semantic Optimal Transport for Sparse Autoencoder Feature Matching and Circuit Compression (NeurIPS 2026 · similarity 0.89)
- A is for Absorption: Studying Feature Splitting and Absorption in Sparse Autoencoders (NeurIPS 2025 · similarity 0.89)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.