Sparse Transformer: a Controlled Testbed for Mechanistic Interpretability
The paper introduces Sparse Transformer (ST), a 5-layer language model trained on TinyStories that enforces exact activation sparsity in its MLPs via a TopK gate. By allowing only K neurons to fire per token, the model creates a known reference basis that serves as ground truth for evaluating mechanistic interpretability methods. The authors use this testbed to compare feature-recovery techniques (SAEs vs. Transcoders) and to study how compression induces polysemanticity, releasing all 35 trained models and checkpoints to support future research.
The paper is a useful and well-timed contribution that fills a methodological gap between synthetic toy models and production LLMs. The ground-truth testbed is conceptually clean, the case studies are internally consistent, and the public release of models and checkpoints is exemplary. However, the model's tiny scale and narrow domain limit how far conclusions can generalize, and the enthusiasm for scaling laws is not fully warranted by the fitted exponents.
The TopK-sparse MLP design is elegant and empirically effective. The primary ST model (M=32, K=8) achieves an AutoInterp score of 0.866, substantially above loss-matched (0.605) and width-matched (0.667) dense baselines, showing that exact sparsity genuinely forces monosemantic specialization. Case Study #1 is methodologically rigorous because the true neuron vectors are known by construction, allowing direct cosine-similarity measurement rather than proxy metrics. The width-sparsity trade-off analysis and release of checkpoints every 10% of training further strengthen reproducibility.
First, scale and domain are severely restricted. The 5-layer, 256-hidden model trained on TinyStories learns surface-level narrative patterns, not abstract reasoning or world knowledge. The limitation section concedes this but does not fully temper the broader claims about scalability. Second, the power-law fit $\mathcal{L}(M,K)=c \cdot M^{\alpha} \cdot K^{\beta}$ yields exponents near zero ($\alpha = -0.038$, $\beta = -0.015$), meaning loss is barely sensitive to width or sparsity; extrapolating this to large models is speculative. Third, dead-neuron rates reach 22.9% at wide settings, and the selected model still requires resampling every 1,000 steps to maintain 3.5% dead neurons, meaning the ground truth is partly engineered. Finally, the high AutoInterp score relies on GPT-4o-mini, an automated evaluator whose alignment with human judgment is itself contested.
The evidence supports the central claim that ST neurons are interpretable and that Transcoders recover ground-truth vectors better than SAEs (0.709 vs. 0.465 cosine similarity, consistent across BatchTopK and Matryoshka variants). However, the SAE vs. Transcoder comparison, while fair within this testbed, may overstate the failure of SAEs: SAEs are designed without knowledge of the true basis, and the ST MLP architecture is intentionally chosen to resemble a Transcoder. The polysemanticity distillation study is compelling, but the student is trained only on MSE against the teacher, not on language modeling, so it may not reflect how dense polysemanticity arises in standard training. Related work coverage is adequate but could more directly address whether SAE evaluation on toy grounds has already produced similar conclusions.
Reproducibility is a strength. The authors release code, all 35 trained models on HuggingFace, and checkpoints every 10% of training. Hyperparameters are reported in detail: AdamW, batch size 64, cosine LR from $3 \times 10^{-4}$ to $10^{-5}$, context length 256, and neuron resampling every 1,000 steps with a 5e-6 frequency threshold. All models share the same random seed and data order. Gaps include unspecified tokenizer details, no explicit train/validation split size, and reliance on GPT-4o-mini for AutoInterp, whose API outputs may drift over time. Reproducing the exact resampling behavior also depends on implementation details of the dead-neuron detection logic that are not fully specified.
Mechanistic interpretability aims to decompose neural network computations into human-understandable features, yet progress is limited by the absence of ground-truth units to evaluate against. We address this with a Sparse Transformer (ST): a 5-layer language model trained on TinyStories whose wide, TopK-sparse MLP forces neuron specialization. We show that ST neurons are highly interpretable: they outperform a loss-matched dense transformer by 43% on AutoInterp, and circuits of a median of 12 neurons suffice to reproduce individual model predictions. This testbed allows us to directly compare feature learning methods: Transcoder, which approximates the MLP input-output mapping, recovers ST's neuron vectors with 52.6% higher cosine similarity than Sparse Autoencoder, which only reconstructs outputs. Furthermore, the ST enables a controlled study of polysemanticity: distilling its sparse MLP into a dense counterpart reveals how compression leads to feature superposition. These results show how a known reference basis sharpens evaluation and understanding, a role we hope the ST can serve for the field.
Pick a starting point or write your own. Challenges run in the background, so you can keep reading while the AI investigates.
No challenges yet. Disagree with the review? Ask the AI to revisit a specific claim.