Research map / Architectures and efficiency / RNNs and normalization
Identifying Functionally Important Features with End-to-End Sparse Dictionary Learning
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 end-to-end sparse dictionary learning, training SAEs by minimizing KL divergence of model outputs (optionally plus downstream reconstruction) with L1 sparsity, yielding a Pareto improvement over local MSE-trained SAEs on GPT2-small and TinyStories-1M.
Problem
Standard sparse autoencoders are trained to minimize activation reconstruction MSE, which may not prioritize features that are functionally important for explaining network behavior, wasting capacity on less important features.
Approach
Replace the original activations with SAE outputs and train the SAE end-to-end by minimizing KL divergence between the original and SAE-inserted output distributions plus an L1 sparsity penalty (e2e), with an optional downstream-layer MSE term (e2e+downstream) to keep computational pathways close to the original model; the base model is frozen.
Key results
- For a given L0 both e2e and e2e+downstream SAEs have CE loss increase below 45% of the local SAE's CE loss increase (GPT2-small, layer 6).
- e2e+downstream SAEs have interpretability scores not significantly different from local SAEs at similar L0, and are more interpretable at similar CE loss increase in layers 2 (p=0.0053) and 6 (p=0.0005).
- Local SAEs show higher feature splitting: mean max cosine similarity is 0.04 higher than e2e+downstream and 0.166 higher than e2e at layer 6.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Replace local activation-reconstruction training with output-distribution KL and sparsity training | SAEs on GPT2-small · GPT2-small residual stream before attention layer 6; comparisons across sparsity coefficients | CE loss increase at a given L0: less than 45% of the CE loss increase of Local SAE | Helped |
| Replace local activation-reconstruction training with output-distribution KL and sparsity training | SAEs on GPT2-small · GPT2-small residual stream before attention layer 6; comparisons at approximately equivalent CE loss increase | L0: 69.4 -> 27.5 | Helped |
| Replace local activation-reconstruction training with output-distribution KL and sparsity training | SAEs on GPT2-small · GPT2-small residual stream before attention layer 6; comparisons at approximately equivalent CE loss increase | alive dictionary elements: 26k -> 22k | Helped |
| Add downstream-layer reconstruction loss to end-to-end KL and sparsity training (ablation) | End-to-end SAE · GPT2-small residual stream before attention layer 6 | downstream reconstruction loss | Helped |
| Add downstream-layer reconstruction loss to end-to-end KL and sparsity training (ablation) | End-to-end SAE · GPT2-small residual stream before attention layer 6; approximately matched CE loss increase | CE loss increase: 0.144 -> 0.125 | Helped |
| Increase learning rate above 0.0005 (ablation) | Local SAE, End-to-end SAE, and E2E + downstream SAE · GPT2-small layers 2, 6, and 10 | alive dictionary elements at the same CE loss increase | Helped |
| Increase learning rate above 0.0005 (ablation) | Local SAE, End-to-end SAE, and E2E + downstream SAE · GPT2-small layers 2, 6, and 10 | L0 versus CE loss increase | Mixed |
| Increase initial dictionary size (ablation) | Local SAE, End-to-end SAE, and E2E + downstream SAE · GPT2-small residual stream before attention layer 6 | L0 versus CE loss increase | Helped |
| Increase initial dictionary size (ablation) | Local SAE, End-to-end SAE, and E2E + downstream SAE · GPT2-small residual stream before attention layer 6 | alive dictionary elements | Hurt |
| Increase Local SAE training samples from 50k to 200k, 400k, or 800k (ablation) | Local SAE · GPT2-small residual stream before attention layer 6 | L0 versus CE loss increase | No clear effect |
| Increase End-to-end SAE training samples (ablation) | End-to-end SAE · GPT2-small residual stream before attention layer 6 | L0 versus CE loss increase | Helped |
| Increase E2E + downstream training samples up to 400k (ablation) | E2E + downstream SAE · GPT2-small residual stream before attention layer 6 | L0 and alive dictionary elements versus CE loss increase | Helped |
Surprising findings
- e2e SAEs have much worse per-layer activation reconstruction error than local SAEs despite explaining more performance, with e2e SAE features not robust across random seeds while local and downstream SAEs are robust.
- Increasing dictionary size to 100x or training samples to 800k does not noticeably improve the local SAE Pareto frontier, so e2e gains are not due to extra compute.
Limitations
- e2e training is 2-3.5x slower in wall-clock time than local SAE training.
- e2e SAEs suffer feature suppression and poor reconstruction of activation directions at subsequent layers.
- Interpretability is assessed with an imperfect automated score and limited qualitative analysis; e2e+downstream loss coefficients were not extensively tuned.
Experimental setup
Datasets: OpenWebText TinyStories
Benchmarks: OpenWebText evaluation split
Metrics: CE loss increase L0 alive dictionary elements downstream reconstruction MSE automated interpretability score cosine similarity of dictionary features
Baselines: local SAE (MSE reconstruction) Gated SAE Towards Monosemanticity SAEs
Builds on: sparse autoencoders sparse dictionary learning Towards Monosemanticity
Closest papers by meaning
- Efficient Dictionary Learning with Switch Sparse Autoencoders (ICLR 2025 · similarity 0.91)
- Learning Multi-Level Features with Matryoshka Sparse Autoencoders (ICML 2025 · similarity 0.89)
- Sanity Checks for Sparse Autoencoders: Do SAEs Beat Random Baselines? (NeurIPS 2026 · similarity 0.89)
- Towards Principled Evaluations of Sparse Autoencoders for Interpretability and Control (ICLR 2025 · similarity 0.89)
- Compute Optimal Inference and Provable Amortisation Gap in Sparse Autoencoders (ICML 2025 · similarity 0.88)
- Sparse Autoencoders Do Not Find Canonical Units of Analysis (ICLR 2025 · similarity 0.88)
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.