Research map / Architectures and efficiency / RNNs and normalization
Codebook Features: Sparse and Discrete Interpretability for Neural Networks
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.
Finetunes pretrained Transformers with sparse vector-quantization codebook bottlenecks at each attention head or MLP layer, then identifies and steers behaviors by activating interpretable codes.
Problem
Dense, continuous hidden states make neural networks hard to interpret, with many features superposed across neurons and no clean unit for analysis or control.
Approach
Replace each layer's activation with the sum of the top-k most cosine-similar codes from a learned codebook (trained with straight-through estimation plus a reconstruction loss), producing sparse discrete states whose codes can be found and activated to causally steer model behavior.
Key results
- On WikiText-103 with Pythia-410M, the Attn k=8 codebook model reaches 46.68% accuracy vs 50.52% for the finetuned base (loss 2.74 vs 2.41).
- On TokFSM, a k=1 attention-only codebook model matches the original (96.39% state accuracy vs 96.77%), and state codes average 97.1% precision vs 70.5% for the best neuron.
- Steering with topic codes raised topic mention on WikiText-103 to 95% for football and 90% for movie when all codes were replaced, vs baseline frequencies of 7.5% and 27.5%.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Add attention codebooks | GPTNeoX-style 4-layer Transformer · TokFSM | State Acc: 96.77 -> 96.39 | No clear effect |
| Add attention and MLP codebooks with k=1 | GPTNeoX-style 4-layer Transformer · TokFSM | State Acc: 96.77 -> 63.65 | Hurt |
| Increase k for attention-and-MLP codebooks (ablation) | GPTNeoX-style 4-layer Transformer · TokFSM | State Acc: 63.65 -> 76.32; 63.65 -> 91.53; 63.65 -> 95.82 | Helped |
| Add attention codebooks with k=8 | TinyStories 21M · TinyStories | Acc: 59.27 -> 57.91 | Hurt |
| Increase attention codebook k from 8 to 64 (ablation) | Pythia 410M · WikiText-103 | Acc: 46.68 -> 48.44 | Helped |
| Group MLP activations into 16 independently coded chunks (ablation) | Pythia 410M · WikiText-103 | Acc: 42.47 -> 48.46 | Helped |
| Use a randomly initialized model instead of a pretrained model (ablation) | Codebook Transformer · TinyStories-1L | Next-token accuracy: 57.91 -> 55.67 | Hurt |
| Freeze the base model and train only codebooks with the original loss (ablation) | Codebook Transformer · TinyStories-1L | Next-token accuracy: 57.91 -> 47.08 | Hurt |
| Train codebooks using only the autoencoding loss (ablation) | Codebook Transformer · TinyStories-1L | Next-token accuracy: 57.91 -> 51.73 | Hurt |
| Use a randomly initialized model instead of a pretrained model (ablation) | Codebook Transformer · FSM-4L | Next-state transition accuracy: 96.39 -> 52.35 | Hurt |
| Freeze the base model and train only codebooks with the original loss (ablation) | Codebook Transformer · FSM-4L | Next-state transition accuracy: 96.39 -> 58.48 | Hurt |
| Train codebooks using only the autoencoding loss (ablation) | Codebook Transformer · FSM-4L | Next-state transition accuracy: 96.39 -> 43.44 | Hurt |
Surprising findings
- Networks operate with only k=1 active code per codebook on the FSM task with near-original accuracy.
- Training with codebooks from a randomly initialized model was far worse than finetuning a pretrained model, while codebook-only training with frozen weights also underperformed.
Limitations
- Only studied Transformers on one algorithmic dataset and two natural-language datasets; no vision or CNN experiments.
- Codebook models are slower at inference (34% to 69% slowdown in tokens per second), reduced with FAISS.
- Topic steering evaluated only with a simple word-based classifier and manually filtered topic codes.
Experimental setup
Datasets: TokFSM (synthetic FSM rollouts) TinyStories WikiText-103
Benchmarks: TokFSM TinyStories WikiText-103
Metrics: language modeling loss next-token accuracy state accuracy Jensen-Shannon divergence classifier precision topic mention frequency tokens per second
Baselines: No codebook original model Pythia 160M (finetuned) neuron-level classifiers
Builds on: VQ-VAE GPTNeoX Pythia product quantization
Closest papers by meaning
- Sparse Feature Circuits: Discovering and Editing Interpretable Causal Graphs in Language Models (ICLR 2025 · similarity 0.87)
- Explainable Neural Networks with Guarantee: A Sparse Estimation Approach (AAAI 2025 · similarity 0.87)
- Transcoders find interpretable LLM feature circuits (NeurIPS 2024 · similarity 0.86)
- Towards Interpretability Without Sacrifice: Faithful Dense Layer Decomposition with Mixture of Decoders (NeurIPS 2025 · similarity 0.85)
- Neurons to Words: A Novel Method for Automated Neural Network Interpretability and Alignment (AAAI 2025 · similarity 0.85)
- Training Neural Networks as Recognizers of Formal Languages (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.