BatchTopK: Batch-Level Hard Sparsity Operator

• BatchTopK is a batch-level hard-sparsity selection operator that enforces an exact average sparsity budget across a batch, improving adaptive feature allocation and reconstruction quality.
• It operates by flattening activation matrices and efficiently selecting the top n×k values via algorithms like radix selection, enabling scalable GPU implementations in sparse autoencoders and crosscoders.
• The method enhances mechanistic interpretability and model diffing in transformers by providing precise control over latent features, facilitating tasks such as language model probing and text-to-speech steering.

BatchTopK is a batch-level hard-sparsity selection operator and optimization target widely used in modern sparse autoencoders (SAEs) and crosscoders to interpret, decompose, and control the internal representations of neural networks—particularly transformers. By enforcing an exact sparsity budget at the batch level rather than per sample, BatchTopK delivers improved reconstruction accuracy, adaptive feature allocation, and direct control over average sparsity. It has become a methodological foundation in mechanistic interpretability, model diffing, and high-throughput inference, with applications ranging from LLM probing to text-to-speech steering and dense-vs-MoE comparisons (Bussmann et al., 2024, Chaudhari et al., 6 Mar 2026, Oozeer et al., 29 Aug 2025, Minder et al., 3 Apr 2025, Koriagin et al., 8 Jun 2026, Kassem et al., 16 Feb 2026).

1. Mathematical Formulation and Operator Definition

Let $X \in \mathbb{R}^{n \times d}$ denote a minibatch of $n$ input vectors (each of dimension $d$). The encoder computes latent preactivations $$Z = W_{\mathrm{enc}} X + b_{\mathrm{enc}} \in \mathbb{R}^{n \times m}$$, with $m$ the dictionary size. Traditional TopK SAEs enforce a fixed $k$-sparsity within each sample by keeping the $k$ largest entries per row of $Z$. In contrast, BatchTopK aggregates all preactivations across the batch and selects the top $B = n \cdot k$ entries overall (by value or optionally by scaled score).

The explicit thresholding is:

• Flatten $Z$ to $n$0.
• Find the threshold $n$1 that is the $n$2-th largest element in $n$3.
• Define a binary mask: $n$4 if $n$5, else $n$6.
• Apply: $n$7.

This yields exactly $n$8 nonzero activations per batch, i.e., an average of $n$9 active features per sample, but any sample may be allocated more or fewer based on the distribution of activations. The operator extends directly to signed activations by selection by magnitude $d$0.

In crosscoder applications (model-diffing), batch-level top-$d$1 is further often imposed on a scaled activation matrix that includes decoder norm weights, aligning the competitive budget with feature salience (Chaudhari et al., 6 Mar 2026, Minder et al., 3 Apr 2025).

2. Training Objectives and Loss Function

BatchTopK regularizes via an exact batch-level $d$2 constraint, obviating the need for explicit per-activation $d$3- or $d$4-penalties. The prototypical SAE objective is: $d$5 where the reconstruction proceeds as $d$6 (with $d$7 obtained via BatchTopK thresholding), and $d$8 is an auxiliary loss (e.g., for “dead” features). In crosscoders, losses may include multi-model reconstructions, decoder norm penalties, and specialized contrastive or delta terms to highlight model-specific latent directions (Bussmann et al., 2024, Kassem et al., 16 Feb 2026).

BatchTopK guarantees exact average sparsity per batch, and the tradeoff between sparsity and reconstruction quality is governed solely by $d$9—the target active features per sample—directly set by the practitioner without expensive hyperparameter sweeps.

3. Algorithmic Workflow and GPU Implementation

The central algorithmic step in BatchTopK is efficient selection of the $$Z = W_{\mathrm{enc}} X + b_{\mathrm{enc}} \in \mathbb{R}^{n \times m}$$0 largest (by value or score) entries from a possibly high-dimensional batch matrix. This is achieved via:

• Flattening the batch activation matrix.
• O($$Z = W_{\mathrm{enc}} X + b_{\mathrm{enc}} \in \mathbb{R}^{n \times m}$$1) partial selection algorithms (e.g., quickselect or radix-based selection suites) to determine the top $$Z = W_{\mathrm{enc}} X + b_{\mathrm{enc}} \in \mathbb{R}^{n \times m}$$2 threshold without full sorting (Li et al., 24 Jan 2025).
• Masking and backpropagation, where the mask is considered constant per batch (straight-through gradient estimators may be employed for partial differentiability in some settings).

The computational cost per step is O($$Z = W_{\mathrm{enc}} X + b_{\mathrm{enc}} \in \mathbb{R}^{n \times m}$$3), dominated by encoder and decoder matrix multiplications in SAEs and crosscoders. GPU-parallelized BatchTopK kernels, e.g., in the RadiK framework, achieve scalable evaluation and selection for very large $$Z = W_{\mathrm{enc}} X + b_{\mathrm{enc}} \in \mathbb{R}^{n \times m}$$4 and batch sizes via radix selection and adaptive scaling for pathological input distributions, with up to 4.8× speedup over merge-based or serial alternatives (Li et al., 24 Jan 2025).

Kernel	Time Complexity	Max k Supported	Remarks
Bitonic	O(k log k)	$$Z = W_{\mathrm{enc}} X + b_{\mathrm{enc}} \in \mathbb{R}^{n \times m}$$52K	Fast for small k
Radix (RadiK)	O(N)	Arbitrarily large	Suits BatchTopK

4. Empirical Findings and Practical Considerations

Empirical work has demonstrated that BatchTopK SAEs and crosscoders consistently outperform per-sample TopK in normalized mean squared error (NMSE) and downstream language-model cross-entropy (CE) degradation across LLMs (GPT-2 Small, Gemma 2 2B) and tasks (Bussmann et al., 2024). The flexibility of BatchTopK enables adaptive latent allocation: simple samples may receive only one or two active codes, while complex samples absorb a much larger share. The ability to control average $$Z = W_{\mathrm{enc}} X + b_{\mathrm{enc}} \in \mathbb{R}^{n \times m}$$6 sparsity with a single parameter $$Z = W_{\mathrm{enc}} X + b_{\mathrm{enc}} \in \mathbb{R}^{n \times m}$$7 removes the trial-and-error calibration required by $$Z = W_{\mathrm{enc}} X + b_{\mathrm{enc}} \in \mathbb{R}^{n \times m}$$8-penalized or threshold-based approaches (e.g., JumpReLU).

Hyperparameter recommendations for robust convergence are:

• Batch size $$Z = W_{\mathrm{enc}} X + b_{\mathrm{enc}} \in \mathbb{R}^{n \times m}$$9
• Learning rate $m$0 (Adam)
• Sparsity $m$1 in $m$2 for standard interpretability/reconstruction tradeoffs
• Auxiliary loss weight $m$3

BatchTopK has also demonstrated resilience to training instabilities common in alternative sparse coding regimes, and its batch-wise enforcement leads to interpretable and reliable representations (Bussmann et al., 2024, Minder et al., 3 Apr 2025).

5. Applications in Model Interpretation and Diffing

BatchTopK underpins a wide range of applications in mechanistic interpretability:

• SAE-based LLM analysis: Enables extraction of sparse, monosemantic features that are directly attributable to linguistic or structural phenomena (e.g., phonemes, speaker attributes, laughter in TTS), with causal steering in latent space (Koriagin et al., 8 Jun 2026).
• Crosscoders in model diffing: By dictating a batch-level hard budget on shared and model-specific features, BatchTopK yields cleaner splits between base and fine-tuned model concepts, mitigating artifacts such as latent decoupling and decoder shrinkage inherent to $m$4-based crosscoders (Minder et al., 3 Apr 2025, Kassem et al., 16 Feb 2026).
• Mixture-of-Expert (MoE) and dense model comparisons: Facilitates quantification of feature overlap and specialization, with fractional variance explained (FVE) in joint activation spaces exceeding 87% using explicit BatchTopK-controlled crosscoders (Chaudhari et al., 6 Mar 2026).

Concrete case studies reveal that features recovered or discovered via BatchTopK are often causal—with steering or patching interventions along individual features or small sets yielding substantial behavioral changes in downstream generative tasks (Koriagin et al., 8 Jun 2026, Kassem et al., 16 Feb 2026).

6. Generalizations, Limitations, and Variants

The rigid global competition of BatchTopK, while advantageous for adaptive allocation, can induce the “activation lottery” wherein rare, high-magnitude features crowd out mid-frequency, semantically stable activations. The Sampled-SAE framework generalizes BatchTopK by introducing a filter stage based on batch-level feature scores (e.g., $m$5 norm, entropy), creating a tunable spectrum between global and local selection controlled by a pool-multiplier $m$6. Moderate $m$7 yields substantial improvements in probing accuracy at minor cost to reconstruction fidelity, trading off between consistency and per-token resolution (Oozeer et al., 29 Aug 2025).

Unlike soft-thresholded methods, BatchTopK's non-differentiable masking means gradient flow is zero for non-selected features at the batch step; practical schemes include straight-through estimators. Limitations noted include minor residual artifacts for especially large models, and the inability to distinguish entirely new versus repurposed latents in model diffing (Minder et al., 3 Apr 2025).

Variant	Motivation	Key Tradeoff
BatchTopK	Adaptive allocation	May suffer activation lottery
Sampled-SAE	Distribution-aware	Trades FVU vs. probing
Per-sample TopK	Simpler semantics	Rigid, less adaptive

7. Impact and Methodological Recommendations

BatchTopK has reshaped best practices in sparse coding and interpretability research:

• It provides robust, interpretable, and causal latent dictionaries aligned with intended semantic units.
• Offers hyperparameter-free control of average sparsity, facilitating reproducibility and interpretability.
• Recommended for crosscoder-based model diffing to avoid artifacts such as Complete Shrinkage and Latent Decoupling; LatentScaling metrics are advised to empirically validate feature attributions (Minder et al., 3 Apr 2025).
• In high-throughput and GPU workloads, batch-wise selection exploits parallelism more efficiently than merge-based alternatives (Li et al., 24 Jan 2025).

In summary, BatchTopK is a versatile, high-fidelity sparsification and feature selection strategy that enables the next generation of mechanistic model analysis, reliable behavioral control, and efficient parallel inference. Its adoption across interpretability, model diffing, and generative control tasks underscores its broad technical impact.