BatchTopK is a method for training Sparse Autoencoders (SAEs); its distinguishing characteristic lies in achieving a computationally efficient balance between sparsity and reconstruction performance, an inherent tradeoff in SAEs. This makes it a recommended choice for mechanistic interpretability researchers.

Its conceptual innovation lies in allowing the number of latents ‘invested’ in linearising a sample of superimposed features to fluctuate between samples while remaining constant across each batch of samples. This forced economisation renders the number of latents sensitive, in relative terms, to the respective complexity encoded in the sampled activations. The resulting ‘weighting’ of the SAE distribution may therefore offer a more accurate account of the latent dynamics learned by the model.

Background

SAEs emerged as a viable mechanistic interpretability technique following Anthropic Inc.’s 2023 “Towards Monosemanticity” study. This suggested that the apparently uninterpretable, polysemantic features learned by neural networks’ individual neurons were actually superpositions of interpretable, monosemantic, features. In this way, a given neural network can learn more features than it has neurons; it ‘crams’ unrelated features into orthogonal subspaces.¹ SAEs allow researchers to recover a list of interpretable, monosemantic features from the activations of the target neural network.

These findings had an enormous impact on the mechanistic interpretability research community. On the one hand, they offered a tantalising suggestion of an elegant and powerful framework for illuminating the black box. Thus, they sparked a wave of enthusiasm in the field (especially vis-à-vis LLMs) extending far beyond both corporate, “frontier” AI labs and academic institutions alike. Its participants took on, inter alia, the other side of that development: ascertaining the shape and limits—in terms of completeness and veracity—of the outcomes gleaned from this new tool.

look man it is a chart idk — On Gemma 2 2B activations, BatchTopK outperforms TopK SAEs across different values of k. Although BatchTopK has a better reconstruction performance (left), it only outperforms JumpReLU in terms of downstream CE degradation in the setting where k=16 (right). (Bussmanm, Leask & Nanda, 2024)

BatchTopK contibutes to both dimensions. At a given compute budget, it gleans a more informative distribution of learned latent features from the target model, making it an invaluable tool for the democraticisation of mechanistic interpretability research. Not only does it help mitigate the need for compute-intensive hyperparameter sweeps, the particular quantities it optimises between began hinting at certain contours to the SAE programme. BatchTopK ‘adaptively’ allocates a greater quantity of latents to interpret the more complex of the samples pulled from the target AI’s activations.² It does so while maintaining the same average sparsity as, and either improving or roughly matching the overall accuracy of, alternative methods³.

Use

The BatchTopK architecture has been integrated into various SAE-focused tools, such as SAEBench, dictionary_learning, and SAE Lens.

To take up an example offered by Anthropic, consider a neuron that fires on academic citation components, http requests, snippets of Hangul unicode, and English dialogue. Taken all together at the neuron level, this combination appears incomprehensible—yet it is precisely because these (separately reasonably comprehensible) features are so alien from one another that the model ‘is able’ to arrange them ‘on top of each other’; it exploits their non-overlapping geometry to fulfill the compression demand (Bricken et al, 2023). ↩︎
From among the data presented in their BatchTopK SAEs paper, Bussmann, Leask, and Nanda (2024) highlight a particularly striking example of this complexity-sensitive-(re)calibration. At the extreme ends of the distribution, a mere 2 features suffice for the simplest samples—while the most complex few register active feature counts in the mid-80s all the way up to 99 (mean:32, mode:31, sd.:10.42; see fig. 3). ↩︎
Metrics follow the prevailing SAE literature: accuracy is quantified as reconstruction loss (mean squared error and cross entropy). SAEs fundamentally trade-off between this quantity and sparsity.
Models, specifically the accurate-but-expensive JumpReLU as well as the TopK SAEs that inspired Bussmann et al.—both of which already improve on the vanilla L1 SAEs. Meanwhile there exist also further variations, such as Matching Pursuit-, Gated-, Matryoshka-, and tanh-JumpReLU-SAEs… For a precise and up-to-date accounting of techniques and their respective tradeoffs, see e.g. SAE Lens’s “Training Sparse Autoencoders” documentation. ↩︎