# Distribution-Aware Feature Selection for SAEs

## TL;DR

This paper introduces Sampled-SAE, a flexible feature selection method for sparse autoencoders that balances global and token-specific features by tuning a parameter, improving interpretability and reconstruction.

## Contribution

It proposes Sampled-SAE, a novel, distribution-aware feature selection approach that generalizes BatchTopK by controlling the feature pool size, enhancing interpretability and performance.

## Key findings

- Sampled-SAE outperforms BatchTopK in various metrics.
- The optimal parameter l varies depending on the desired trade-offs.
- Sampled-SAE offers a tunable spectrum between global and token-specific features.

## Abstract

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 $L_2$ 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.