Research map / Architectures and efficiency / RNNs and normalization
Mechanistic Permutability: Match Features Across Layers
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 SAE Match, a data-free method that aligns Sparse Autoencoder features across layers of Gemma 2 by minimizing MSE between threshold-folded SAE weights solved via linear assignment.
Problem
Aligning interpretable SAE features across different layers of a language model is an open problem, since features at each layer are learned independently and their correspondence is unknown.
Approach
Match features across layers as a permutation problem: fold JumpReLU thresholds into encoder/decoder weights so that MSE between decoder rows reflects hidden-state norm scales, then solve a linear assignment problem; compose permutations to approximate distant layer matches and use them to skip layers.
Key results
- Folded matching yields better LLM-judged feature match quality than unfolded matching across layers of Gemma 2 2B.
- Starting from layer 10 onward, pruning a layer via matched-feature encode-permute-decode causes minimal loss increase, per the paper's layer-pruning experiments.
- Feature similarity gradually declines with path length but remains significant for about five layers.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Fold activation thresholds into SAE weights before matching (ablation) | SAE Match · Gemma 2 feature matching across layers | feature matching quality | Helped |
| Use encoder-only weights instead of encoder-decoder weights for matching (ablation) | SAE Match · Gemma 2 layers 19–20 | Matching Score | Hurt |
| Use decoder-only weights instead of encoder-decoder weights for matching (ablation) | SAE Match · Gemma 2 layers 19–20 | Matching Score | No clear effect |
| Compose adjacent-layer permutations instead of computing a full permutation (ablation) | SAE Match · Gemma 2 feature paths across layers | semantic matching quality | Mixed |
| Use SAEs with similar activation norms across layers (ablation) | SAE Match · Gemma 2 initial layers | Explained Variance | Helped |
| Increase SAE mean l0-norm (ablation) | SAE Match · Gemma 2 layers 20–21 | change in cross-entropy loss (ΔL) | Helped |
| Vary SAE sparsity; use mean l0-norm approximately 70 (ablation) | SAE Match · Gemma 2 layers 20–21 | Explained Variance | Helped |
| Use more than 100 active features when matching (ablation) | SAE Match · Gemma 2 layers 20–21 | Explained Variance | Hurt |
| Approximate a subsequent layer from earlier-layer features rather than reconstructing the current hidden state with its own SAE (ablation) | SAE Match layer pruning · Gemma 2 layers 20–21 | change in cross-entropy loss (ΔL) | Hurt |
| Prune a layer using matched features | SAE Match layer pruning · OpenWebText, Code, and WikiText | change in cross-entropy loss (ΔL) | Mixed |
| Use SAE Match rather than plain encode-decode at lower layers (ablation) | SAE Match layer pruning · Gemma 2 lower layers | Explained Variance | Hurt |
Surprising findings
- Early-layer matching fails (near-zero SAME labels), which the authors attribute to large l0 variance across canonical SAEs at those layers.
- Explained variance peaks at mean l0 around 70 rather than monotonically improving with lower sparsity.
Limitations
- Only bijective (one-to-one) matching is considered, so some features remain unmatched.
- Early-layer matching quality is poor, attributed to differing l0 norms across canonical SAEs.
- Matching uses all SAE weights jointly; the contribution of individual weight types is left to future work.
Experimental setup
Datasets: OpenWebText GitHub code subset (github-small-near-dedup) WikiText Neuronpedia feature descriptions
Benchmarks: OpenWebText Code WikiText
Metrics: MSE of matched decoder/encoder weights LLM-judged description similarity (SAME/MAYBE/DIFFERENT) matching score change in cross-entropy loss explained variance
Baselines: encoder-only matching unfolded matching naive matching encode-decode previous hidden layer drop
Builds on: Gemma Scope JumpReLU sparse autoencoders Git Re-Basin (weight permutation alignment)
Closest papers by meaning
- Semantic Optimal Transport for Sparse Autoencoder Feature Matching and Circuit Compression (NeurIPS 2026 · similarity 0.88)
- From Flat to Hierarchical: Extracting Sparse Representations with Matching Pursuit (NeurIPS 2025 · similarity 0.87)
- Universal Sparse Autoencoders: Interpretable Cross-Model Concept Alignment (ICML 2025 · similarity 0.85)
- Sanity Checks for Sparse Autoencoders: Do SAEs Beat Random Baselines? (NeurIPS 2026 · similarity 0.85)
- Sparse Autoencoders Do Not Find Canonical Units of Analysis (ICLR 2025 · similarity 0.84)
- HyperDAS: Towards Automating Mechanistic Interpretability with Hypernetworks (ICLR 2025 · similarity 0.84)
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.