# TopK Language Models

arXiv: https://arxiv.org/abs/2506.21468 · HTML: https://mlsanity.com/paper/2506.21468/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in SAE features and their applications, one of 13 directions in Mechanistic interpretability.
- Builds on 32 of its 43 references that are in the map.
- First posted on arXiv 26 Jun 2025.

## Placed in

- Mechanistic interpretability / [SAE features and their applications](https://mlsanity.com/mechinterp/sae-features-and-their-applications/)

## Builds on (most cited first, 30 of 32)

- [LLaMA: Open and Efficient Foundation Language Models](https://mlsanity.com/paper/2302.13971/) (2023)
- [Think you have Solved Question Answering? Try ARC, the AI2 Reasoning Challenge](https://mlsanity.com/paper/1803.05457/) (2018)
- [HellaSwag: Can a Machine Really Finish Your Sentence?](https://mlsanity.com/paper/1905.07830/) (2019)
- [Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer](https://mlsanity.com/paper/1701.06538/) (2017)
- [Pointer Sentinel Mixture Models](https://mlsanity.com/paper/1609.07843/) (2016)
- [Can a Suit of Armor Conduct Electricity? A New Dataset for Open Book Question Answering](https://mlsanity.com/paper/1809.02789/) (2018)
- [Generating Long Sequences with Sparse Transformers](https://mlsanity.com/paper/1904.10509/) (2019)
- [Sparse Autoencoders Find Highly Interpretable Features in Language Models](https://mlsanity.com/paper/2309.08600/) (2023)
- [The FineWeb Datasets: Decanting the Web for the Finest Text Data at Scale](https://mlsanity.com/paper/2406.17557/) (2024)
- [The LAMBADA dataset: Word prediction requiring a broad discourse context](https://mlsanity.com/paper/1606.06031/) (2016)
- [Semantic Uncertainty: Linguistic Invariances for Uncertainty Estimation in Natural Language Generation](https://mlsanity.com/paper/2302.09664/) (2023)
- [PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel](https://mlsanity.com/paper/2304.11277/) (2023)
- [Learning Sparse Neural Networks through L0 Regularization](https://mlsanity.com/paper/1712.01312/) (2017)
- [Rigging the Lottery: Making All Tickets Winners](https://mlsanity.com/paper/1911.11134/) (2019)
- [Causal Abstractions of Neural Networks](https://mlsanity.com/paper/2106.02997/) (2021)
- [Scalable Training of Artificial Neural Networks with Adaptive Sparse Connectivity inspired by Network Science](https://mlsanity.com/paper/1707.04780/) (2017)
- [Improving Dictionary Learning with Gated Sparse Autoencoders](https://mlsanity.com/paper/2404.16014/) (2024)
- [Soft Threshold Weight Reparameterization for Learnable Sparsity](https://mlsanity.com/paper/2002.03231/) (2020)
- [Llama Scope: Extracting Millions of Features from Llama-3.1-8B with Sparse Autoencoders](https://mlsanity.com/paper/2410.20526/) (2024)
- [Automatically Interpreting Millions of Features in Large Language Models](https://mlsanity.com/paper/2410.13928/) (2024)
- [Sparse Autoencoders Trained on the Same Data Learn Different Features](https://mlsanity.com/paper/2501.16615/) (2025)
- [Sparse is Enough in Scaling Transformers](https://mlsanity.com/paper/2111.12763/) (2021)
- [Top-KAST: Top-K Always Sparse Training](https://mlsanity.com/paper/2106.03517/) (2021)
- [A Survey on Sparse Autoencoders: Interpreting the Internal Mechanisms of Large Language Models](https://mlsanity.com/paper/2503.05613/) (2025)
- [The State of Sparse Training in Deep Reinforcement Learning](https://mlsanity.com/paper/2206.10369/) (2022)
- [Projecting Assumptions: The Duality Between Sparse Autoencoders and Concept Geometry](https://mlsanity.com/paper/2503.01822/) (2025)
- [Memory-efficient Transformers via Top-k Attention](https://mlsanity.com/paper/2106.06899/) (2021)
- [Codebook Features: Sparse and Discrete Interpretability for Neural Networks](https://mlsanity.com/paper/2310.17230/) (2023)
- [How Can We Be So Dense? The Benefits of Using Highly Sparse Representations](https://mlsanity.com/paper/1903.11257/) (2019)
- [Tracking the Feature Dynamics in LLM Training: A Mechanistic Study](https://mlsanity.com/paper/2412.17626/) (2024)

## Abstract (arXiv)

Sparse autoencoders (SAEs) have become an important tool for analyzing and interpreting the activation space of transformer-based language models (LMs). However, SAEs suffer several shortcomings that diminish their utility and internal validity. Since SAEs are trained post-hoc, it is unclear if the failure to discover a particular concept is a failure on the SAE's side or due to the underlying LM not representing this concept. This problem is exacerbated by training conditions and architecture choices affecting which features an SAE learns. When tracing how LMs learn concepts during training, the lack of feature stability also makes it difficult to compare SAEs features across different checkpoints. To address these limitations, we introduce a modification to the transformer architecture that incorporates a TopK activation function at chosen layers, making the model's hidden states equivalent to the latent features of a TopK SAE. This approach eliminates the need for post-hoc training while providing interpretability comparable to SAEs. The resulting TopK LMs offer a favorable trade-off between model size, computational efficiency, and interpretability. Despite this simple architectural change, TopK LMs maintain their original capabilities while providing robust interpretability benefits. Our experiments demonstrate that the sparse representations learned by TopK LMs enable successful steering through targeted neuron interventions and facilitate detailed analysis of neuron formation processes across checkpoints and layers. These features make TopK LMs stable and reliable tools for understanding how language models learn and represent concepts, which we believe will significantly advance future research on model interpretability and controllability.