# Distribution-Aware Feature Selection for SAEs

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Sparse autoencoder training and evaluation, one of 13 directions in Mechanistic interpretability.
- Builds on 17 of its 29 references that are in the map.
- First posted on arXiv 29 Aug 2025.

## Placed in

- Mechanistic interpretability / [Sparse autoencoder training and evaluation](https://mlsanity.com/mechinterp/sparse-autoencoder-training-and-evaluation/)

## Builds on (most cited first, 17 of 17)

- [The Pile: An 800GB Dataset of Diverse Text for Language Modeling](https://mlsanity.com/paper/2101.00027/) (2020)
- [Pythia: A Suite for Analyzing Large Language Models Across Training and Scaling](https://mlsanity.com/paper/2304.01373/) (2023)
- [Sparse Autoencoders Find Highly Interpretable Features in Language Models](https://mlsanity.com/paper/2309.08600/) (2023)
- [Finding Neurons in a Haystack: Case Studies with Sparse Probing](https://mlsanity.com/paper/2305.01610/) (2023)
- [Sketching as a Tool for Numerical Linear Algebra](https://mlsanity.com/paper/1411.4357/) (2014)
- [Improving Dictionary Learning with Gated Sparse Autoencoders](https://mlsanity.com/paper/2404.16014/) (2024)
- [A is for Absorption: Studying Feature Splitting and Absorption in Sparse Autoencoders](https://mlsanity.com/paper/2409.14507/) (2024)
- [SAEBench: A Comprehensive Benchmark for Sparse Autoencoders in Language Model Interpretability](https://mlsanity.com/paper/2503.09532/) (2025)
- [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)
- [Input Sparsity Time Low-Rank Approximation via Ridge Leverage Score Sampling](https://mlsanity.com/paper/1511.07263/) (2015)
- [Frequent Directions : Simple and Deterministic Matrix Sketching](https://mlsanity.com/paper/1501.01711/) (2015)
- [Interpretability in Parameter Space: Minimizing Mechanistic Description Length with Attribution-based Parameter Decomposition](https://mlsanity.com/paper/2501.14926/) (2025)
- [Online Row Sampling](https://mlsanity.com/paper/1604.05448/) (2016)
- [ℓp Row Sampling by Lewis Weights](https://mlsanity.com/paper/1412.0588/) (2014)
- [Interpreting Learned Feedback Patterns in Large Language Models](https://mlsanity.com/paper/2310.08164/) (2023)
- [Online Matrix Completion and Online Robust PCA](https://mlsanity.com/paper/1503.03525/) (2015)

## Abstract (arXiv)

Sparse autoencoders (SAEs) decompose neural activations into interpretable features. A widely adopted variant, the TopK SAE, reconstructs each token from its K most active latents. However, this approach is inefficient, as some tokens carry more information than others. BatchTopK addresses this limitation by selecting top activations across a batch of tokens. This improves average reconstruction but risks an "activation lottery," where rare high-magnitude features crowd out more informative but lower-magnitude ones. To address this issue, we introduce Sampled-SAE: we score the columns (representing features) of the batch activation matrix (via L2 norm or entropy), forming a candidate pool of size Kl, and then apply Top-K to select tokens across the batch from the restricted pool of features. Varying l traces a spectrum between batch-level and token-specific selection. At l=1, tokens draw only from K globally influential features, while larger l expands the pool toward standard BatchTopK and more token-specific features across the batch. Small l thus enforces global consistency; large l favors fine-grained reconstruction. On Pythia-160M, no single value optimizes l across all metrics: the best choice depends on the trade-off between shared structure, reconstruction fidelity, and downstream performance. Sampled-SAE thus reframes BatchTopK as a tunable, distribution-aware family.