Research map / Architectures and efficiency / Pruning and efficient tuning
Low-Rank Adapting Models for Sparse Autoencoders
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.
Proposes fine-tuning the language model itself with LoRA (KL objective) around a frozen pretrained SAE to reduce the CE loss gap caused by SAE insertion, which is faster than end-to-end SAE training.
Problem
Sparse autoencoder reconstructions inserted into a language model substantially increase cross entropy loss, and prior improvements require expensive SAE training, including gradient-based or end-to-end methods.
Approach
Freeze both the SAE and the base model, and train only low-rank adapters on the layers after the SAE so the SAE-inserted model's next-token distribution matches the original model via KL divergence; this optimizes the model rather than the SAE.
Key results
- Across SAE sparsity, width, layer and model size, LoRA closes 30% to 55% of the CE loss gap with Gemma Scope SAEs inserted.
- TopK + LoRA reaches the same CE loss as e2e SAEs 3x to 20x faster on Gemma-2-2B and 2x to 10x faster on Llama-3.2-1B.
- On Llama-3.1-8B with 7 Llama Scope SAEs, compound CE loss drops from 7.83 to 2.78 nats with LoRA.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Add rank 64 LoRA to the model with the TopK SAE | Gemma-2-2B · SAEBench SCR (max) | SCR: 0.448 -> 0.526 | Helped |
| Add rank 64 LoRA to the model with the TopK SAE | Gemma-2-2B · SAEBench TPP (max) | TPP: 0.372 -> 0.412 | Helped |
| Add rank 64 LoRA to the model with the TopK SAE | Gemma-2-2B · SAEBench Sparse Probing (top 1) | Sparse Probing: 0.732 -> 0.760 | Helped |
| Add rank 64 LoRA to the model with the TopK SAE | Gemma-2-2B · SAEBench Autointerp | Autointerp: 0.832 -> 0.830 | Hurt |
| Add rank 64 LoRA to the model with the TopK SAE | Gemma-2-2B · SAEBench Absorption | Absorption: 0.205 -> 0.210 | Hurt |
| Add LoRA while the SAE is inserted | Gemma-2-2B · MMLU | MMLU: 44.2±0.4 -> 45.8±0.4 | Helped |
| Add LoRA while the SAE is inserted | Gemma-2-2B · HellaSwag | HellaSwag: 50.9±0.5 -> 52.1±0.5 | Helped |
| Add LoRA while the SAE is inserted | Gemma-2-27B · BLEU | BLEU: 40.9± 1.7 -> 38.9 ± 1.7 | Hurt |
| Use TopK + LoRA instead of e2e SAE training | Gemma-2-2B · CE loss threshold 2.57 | speedup: 20.01x | Helped |
| Use TopK + LoRA instead of e2e SAE training | Llama-3.2-1B · CE loss threshold 2.73 | speedup: 10.38x | Helped |
| Add LoRA with 7 SAEs inserted | Llama-3.1-8B · 7 SAE compound CE loss | compound cross entropy loss: 7.83 -> 2.78 | Helped |
| Add LoRA with 3 SAEs inserted | Llama-3.1-8B · 3 SAE compound CE loss | compound cross entropy loss: 3.55 -> 2.45 | Helped |
Surprising findings
- Training LoRA only on the layer right after the SAE achieves 88.14% of the loss reduction of training on all layers after the SAE.
- The LoRA-adapted model without any SAE often outperforms the original model on MMLU, HellaSwag and other benchmarks.
Limitations
- Changes the model being interpreted, which may be unacceptable for some mechanistic interpretability applications.
- e2e SAE baselines may not have converged; comparisons of converged accuracy are not possible, and learning rate was not tuned.
- Gains are on a 15M-token training budget and adapters did not converge, so unclear how far they scale.
Experimental setup
Datasets: The Pile (uncopyrighted) RedPajama MMLU HellaSwag TruthfulQA SAEBench
Benchmarks: SAEBench MMLU HellaSwag TruthfulQA Gemma Scope cross entropy loss gap
Metrics: cross entropy loss gap SCR TPP sparse probing autointerp absorption steering log-likelihood change BLEU ROUGE MC1
Baselines: TopK SAE end-to-end SAE JumpReLU SAE (Gemma Scope)
Builds on: LoRA TopK SAE end-to-end SAE JumpReLU SAE
Closest papers by meaning
- Improving Sparse Decomposition of Language Model Activations with Gated Sparse Autoencoders (NeurIPS 2024 · similarity 0.93)
- SAEBench: A Comprehensive Benchmark for Sparse Autoencoders in Language Model Interpretability (ICML 2025 · similarity 0.92)
- Scaling and evaluating sparse autoencoders (ICLR 2025 · similarity 0.91)
- Sparse Autoencoders, Again? (ICML 2025 · similarity 0.90)
- Transferring Linear Features Across Language Models With Model Stitching (NeurIPS 2025 · similarity 0.90)
- Dynamic Low-Rank Sparse Adaptation for Large Language Models (ICLR 2025 · similarity 0.90)
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.