Research map / Architectures and efficiency / Attention and transformers
Residual Stream Analysis with Multi-Layer SAEs
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.
Trains a single k-sparse autoencoder on residual-stream activations from all layers of Pythia (and other) transformers to study how latents are distributed across layers, quantifying multi-layer activity via variance of the layer index.
Problem
Standard sparse autoencoders are trained separately per layer, making it hard to study how information flows across residual-stream layers and to match features across layers.
Approach
Train one SAE (single hidden layer, TopK+ReLU, tied parameters) on residual stream activation vectors from every layer, treating each layer's activation as a separate training example; quantify how latents spread across layers using the variance of a layer-index distribution, with and without tuned-lens transformations.
Key results
- Aggregating over 10 million tokens, the fraction of total variance over layers explained by individual latents (Var(L|J)/Var(L)) is between 54% and 88%, broadly increasing with model size.
- Within a single latent, variance explained by individual tokens (E[Var(L|J,T)]/E[Var(L|J)]) is on the order of 1 to 10%.
- Multi-layer SAEs achieve reconstruction and delta CE loss comparable to single-layer SAEs, performing more consistently across layers; tuned-lens transformations decrease FVU and delta CE for Pythia-70m but make the cross-layer variance ratio drop to 37-41%.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Train one SAE on activations from every layer rather than separate single-layer SAEs | Pythia-70m, Pythia-160m, and Pythia-410m · test activations across layers | FVU and delta cross-entropy loss | No clear effect |
| Increase model size | Pythia MLSAEs with fixed hyperparameters · 10 million test tokens | latent layer-variance ratio: between 54 and 88% | Helped |
| Increase model size | Pythia MLSAEs · 10 million test tokens | normalized mean entropy of latent activations over layers | Helped |
| Apply pre-trained tuned-lens transformations to activations (ablation) | Pythia-70m MLSAE · test set | FVU | Helped |
| Apply pre-trained tuned-lens transformations to activations (ablation) | Pythia-70m MLSAE · test set | delta cross-entropy loss | Helped |
| Apply pre-trained tuned-lens transformations to activations (ablation) | Pythia-70m MLSAE · test set | KL divergence | No clear effect |
| Apply pre-trained tuned-lens transformations to activations (ablation) | MLSAE · 10 million test tokens | latent layer-variance ratio: between 37% and 41% | Hurt |
| Increase expansion factor R (ablation) | Pythia MLSAE with fixed sparsity k=32 · test activations | FVU and MSE | Helped |
| Increase expansion factor R (ablation) | Pythia MLSAE with fixed sparsity k=32 · last-layer activations | FVU and MSE | Hurt |
| Increase expansion factor R (ablation) | Pythia MLSAE with fixed sparsity k=32 · test activations | delta CE loss and KL divergence | Helped |
| Increase sparsity k (ablation) | Pythia MLSAE with fixed expansion factor R=64 · test activations | FVU and MSE | Helped |
| Increase sparsity k (ablation) | Pythia MLSAE with fixed expansion factor R=64 · test activations | delta CE loss and KL divergence | Helped |
Surprising findings
- Tuned-lens transformations do not obviously increase multi-layer latent activity; the aggregate variance ratio actually decreases, contrary to the authors' expectation.
- Applying tuned-lens transformations degrades single-layer SAEs on data from other layers, unlike multi-layer SAEs.
Limitations
- Tuned-lens pre-trained transforms not available for Pythia-1b, leaving gaps in tuned-lens results.
- Last-layer activations excluded from some comparisons because they represent next-token predictions rather than intermediate computations.
- Variance of layer index treats ordinal layer data as interval; threshold-based alternative depends strongly on threshold choice.
Experimental setup
Datasets: The Pile (uncopyrighted subset)
Benchmarks: The Pile test set (1M tokens for FVU/CE, 10M tokens for layer distributions)
Metrics: FVU MSE L1 norm delta cross-entropy loss KL divergence variance of layer index ratio mean max cosine similarity entropy over layers
Baselines: single-layer SAEs
Builds on: k-sparse autoencoder tuned lens AuxK auxiliary loss Pythia
Closest papers by meaning
- Bigram Subnetworks: Mapping to Next Tokens in Transformer Language Models (NeurIPS 2025 · similarity 0.88)
- Sparse Autoencoders, Again? (ICML 2025 · similarity 0.88)
- Semantic Optimal Transport for Sparse Autoencoder Feature Matching and Circuit Compression (NeurIPS 2026 · similarity 0.87)
- Jacobian Sparse Autoencoders: Sparsify Computations, Not Just Activations (ICML 2025 · similarity 0.87)
- One-Step is Enough: Sparse Autoencoders for Text-to-Image Diffusion Models (NeurIPS 2025 · similarity 0.87)
- Dense SAE Latents Are Features, Not Bugs (NeurIPS 2025 · similarity 0.87)
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.