Research map / Architectures and efficiency / RNNs and normalization
Transformer Key-Value Memories Are Nearly as Interpretable as 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.
Compares feed-forward key-value memory neurons (FF-KV, with TopK and normalized variants) against SAEs and Transcoders on SAEBench metrics and human annotation, finding comparable interpretability and low feature alignment between FF-KV and Transcoder features.
Problem
Whether SAE-learned features are more interpretable than the feature vectors already present in transformer feed-forward layers viewed as key-value memories, which is rarely compared systematically.
Approach
Treat FF layer neuron activations and value vectors directly as an SAE-like dictionary (FF-KV), optionally with TopK sparsity and row-norm normalization, and evaluate them with the same SAEBench metrics, human annotation, and cosine-based feature alignment against Transcoders.
Key results
- On Gemma-2-2B layer 13, FF-KV reaches explained variance 1.000 and absorption 0.000 versus SAE 0.699 and 0.087, while sparse probing is similar (0.827 vs 0.846).
- Autointerp is similar across methods (FF-KV 0.710 vs SAE 0.782 on Gemma-2-2B), with SAEs slightly ahead on autointerp and SCR.
- Only 5.7% of FF-KV features are aligned with Transcoder features (cosine above 0.9), and 41% of FF-KV features have no aligned Transcoder counterpart.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| use FF-KV features instead of SAE features | Gemma-2-2B layer 13 · SAEBench | Absorption: 0.087±0.173 -> 0.000±0.001 | Helped |
| use FF-KV features instead of SAE features | Llama-3.1-8B layer 16 · SAEBench | Sparse Probing: 0.879±0.123 -> 0.876±0.098 | No clear effect |
| apply TopK activation sparsity to FF-KV (ablation) | Gemma-2-2B layer 13 · SAEBench | Auto-Interpretation: 0.710±0.246 -> 0.772±0.276 | Helped |
| apply TopK activation sparsity to FF-KV (ablation) | Gemma-2-2B layer 13 · SAEBench | Sparse Probing: 0.827±0.158 -> 0.768±0.168 | Hurt |
| normalize FF-KV value vectors (ablation) | Gemma-2-2B layer 13 FF-KV · SAEBench | Auto-Interpretation: 0.710±0.246 -> 0.706±0.255 | No clear effect |
| increase TopK sparsity (ablation) | TopK FF-KV · Gemma-2-2B layer 13 | Interpretability scores | Mixed |
| increase FF hidden dimension size (ablation) | FF-KV · Pythia model family | RAVEL interpretability | No clear effect |
| increase FF hidden dimension size (ablation) | TopK FF-KV · Pythia model family | Explained Variance: 0.227±0.000 -> 0.129±0.000 -> 0.082±0.000 -> 0.277±0.000 -> 0.090±0.000 -> 0.047±0.000 -> 0.195±0.000 | Mixed |
| use FF-KV instead of SAE features | Gemma-2-2B layer 13 · SAEBench | Auto-Interpretation: 0.782±0.274 -> 0.710±0.246 | Hurt |
| use FF-KV instead of SAE features | Gemma-2-2B layer 13 · SAEBench | Sparse Probing: 0.846±0.161 -> 0.827±0.158 | No clear effect |
| apply TopK activation sparsity to FF-KV (ablation) | Gemma-2-2B layer 13 · human feature-origin judgment | Judging accuracy: 0.86 -> 0.28 | Helped |
| use Transcoder features rather than FF-KV features | Gemma-2-2B layer 13 · FF/Transcoder feature alignment | Unaligned feature proportion: 41% -> 66% | Hurt |
Surprising findings
- Vanilla FF key-value features match or beat SAEs and Transcoders on several SAEBench metrics, including absorption and sparse probing.
- Human annotators could identify FF-KV features as originating from the FF module (86% accuracy) but could not distinguish SAE, Transcoder, and TopK FF-KV features from one another.
Limitations
- Only fixed-width pretrained SAEs and Transcoders were compared; no broader configuration sweep.
- Transcoders were not available for all models, and training them at scale was judged too costly.
- SCR and TPP scores are highly unstable across the choice of K, and the faithfulness of proxy features remains unresolved.
Experimental setup
Datasets: OpenWebText monology/pile-uncopyrighted SHIFT RAVEL
Benchmarks: SAEBench RAVEL SCR TPP sparse probing absorption
Metrics: feature alive rate explained variance absorption score sparse probing accuracy autointerpretation score SCR TPP RAVEL isolation RAVEL causality max-cosine score alignment
Baselines: SAE Transcoder Random Transformer
Builds on: SAEBench Sparse autoencoders Transcoders Geva et al. FF key-value memories Gemma Scope
Closest papers by meaning
- Towards Principled Evaluations of Sparse Autoencoders for Interpretability and Control (ICLR 2025 · 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)
- Rethinking Evaluation of Sparse Autoencoders through the Representation of Polysemous Words (ICLR 2025 · similarity 0.89)
- Are Sparse Autoencoders Useful? A Case Study in Sparse Probing (ICML 2025 · similarity 0.88)
- A is for Absorption: Studying Feature Splitting and Absorption in Sparse Autoencoders (NeurIPS 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.