TopK Activation in Deep Learning

• TopK Activation is a hard sparsifying operator that retains only the k largest activations, ensuring strict control over neuron activity.
• It is applied in sparse autoencoders, transformer models, and mixture-of-experts systems to manage latent cardinality and computational efficiency.
• Recent advancements include batch-level selection, differentiable relaxations, and hierarchical methods that enhance scalability and interpretability.

Searching arXiv for papers on TopK activation across transformers, sparse autoencoders, differentiable top-k, and MoE routing. arxiv_search(query="TopK activation sparse autoencoder transformer differentiable top-k mixture of experts", max_results=10, sort_by="relevance") TopK activation is a hard sparsifying operator that retains only a selected subset of activations and sets the remainder to zero. Across recent arXiv literature, it appears in several distinct roles: as the latent nonlinearity of sparse autoencoders, as a replacement for dense MLP activations in transformers, as a pre-softmax approximation over attention logits, and as a routing rule in mixture-of-experts systems. Although the common motif is “keep the $k$ largest entries,” the precise semantics vary by domain: some formulations select the largest raw pre-activations, others select post-ReLU values, some operate per sample, and others relax the budget to the batch or sequence level (Gao et al., 2024, Takahashi et al., 26 Jun 2025, Dong et al., 2024, Wen et al., 9 Nov 2025).

1. Formal operator and selection semantics

A canonical formulation defines TopK on a vector $x \in \mathbb{R}^d$ by thresholding at the $k$-th largest coordinate. In "TopK LLMs" (Takahashi et al., 26 Jun 2025), letting $\tau_k(x)$ denote the $k$-th largest value among the coordinates of $x$, the activation is

$$y_i = f(x_i)\,\mathbf{1}_{\{\,x_i \ge \tau_k(x)\}},$$

with selection performed on the largest raw values of $x_i$, not on $|x_i|$, and not after the nonlinearity. In that formulation, non-selected activations are explicitly set to zero, so the output has at most $k$ nonzero entries, modulo ties; because the rule is written with $x \in \mathbb{R}^d$0, ties can preserve more than $x \in \mathbb{R}^d$1 coordinates (Takahashi et al., 26 Jun 2025).

This general definition does not fix a single operational meaning across architectures. In sparse autoencoders, TopK often acts on encoder activations and enforces exact or near-exact latent cardinality per sample (Gao et al., 2024). In other settings it follows ReLU, so the competition is effectively among nonnegative post-ReLU values rather than signed pre-activations (Haque et al., 5 Dec 2025). In "Topkima-Former" (Dong et al., 2024), the selected object is neither a latent code nor a token subset but the attention-score activation before softmax: only the $x \in \mathbb{R}^d$2 largest logits in a row of $x \in \mathbb{R}^d$3 are sent to the softmax calculation block. The paper is explicit that this is not top-$x \in \mathbb{R}^d$4 token pruning, MoE gating, or generic activation sparsification; it is top-$x \in \mathbb{R}^d$5 selection on attention-score activations before softmax (Dong et al., 2024).

2. TopK as the sparsity mechanism in sparse autoencoders

The sparse-autoencoder literature has made TopK activation a central alternative to ReLU plus an explicit sparsity penalty. In "Scaling and evaluating sparse autoencoders" (Gao et al., 2024), a TopK SAE replaces the standard

$x \in \mathbb{R}^d$6

with

$x \in \mathbb{R}^d$7

and trains with pure reconstruction loss

$x \in \mathbb{R}^d$8

The stated motivation is direct control of the number of active latents, avoidance of $x \in \mathbb{R}^d$9-induced activation shrinkage, and cleaner scaling-law analysis. That work reports clean scaling laws with respect to autoencoder size and sparsity, and trains a 16 million latent autoencoder on GPT-4 activations for 40 billion tokens (Gao et al., 2024).

Subsequent work relaxes the fixed per-sample budget without discarding TopK selection itself. "BatchTopK Sparse Autoencoders" (Bussmann et al., 2024) defines BatchTopK by selecting the top $k$0 activations across an entire mini-batch of $k$1 samples, rather than the top $k$2 within each sample. The stated effect is variable numbers of active latents per sample while preserving the same average sparsity. That paper reports that BatchTopK SAEs consistently outperform TopK SAEs in reconstructing activations from GPT-2 Small and Gemma 2 2B, achieve comparable performance to JumpReLU SAEs, and allow the average number of latents to be specified directly rather than approximately tuned through a sparsity coefficient sweep (Bussmann et al., 2024).

A further extension is "HierarchicalTopK" (Balagansky et al., 30 May 2025), which trains one SAE to work across multiple sparsity budgets up to a maximum $k$3. Instead of optimizing only the reconstruction using the top $k$4 features, it optimizes all prefix reconstructions

$k$5

with objective

$k$6

This imposes a nested TopK prefix structure and is reported to yield Pareto-optimal trade-offs between sparsity and explained variance on Gemma-2 2B, while preserving high interpretability scores even at higher sparsity (Balagansky et al., 30 May 2025).

TopK activation has also been combined with expert routing rather than replaced by it. "Efficient Dictionary Learning with Switch Sparse Autoencoders" (Mudide et al., 2024) keeps each expert as a TopK SAE and introduces single-expert routing across many smaller dictionaries. In that design, TopK still determines which latent features fire within an expert, while routing determines which expert dictionary is evaluated. The paper presents this as a compute-scaling extension of TopK SAEs rather than a new latent activation rule, and reports a substantial Pareto improvement in the reconstruction-versus-sparsity frontier for a fixed training compute budget (Mudide et al., 2024).

3. TopK as a native transformer computation

TopK activation has moved from post-hoc analysis tools into the forward pass of transformer architectures. In "TopK LLMs" (Takahashi et al., 26 Jun 2025), selected transformer layers replace the usual dense activation in the MLP pathway with TopK, making the model’s hidden states equivalent to the latent features of a TopK SAE. The paper trains decoder-only Llama-style transformers with TopK activation in the first $k$7 layers and leaves the final $k$8 layers dense. It uses an annealed activation

$k$9

with $\tau_k(x)$0 decayed linearly from 1 to 0 over the first $\tau_k(x)$1 of training steps, and a main setting of $\tau_k(x)$2. The reported trade-off is modest perplexity degradation with preserved downstream capability and much lower token and semantic entropy of neurons; for example, in the 24-layer, $\tau_k(x)$3 setting, validation perplexity rises from $\tau_k(x)$4 to $\tau_k(x)$5, while LAMBADA accuracy rises from $\tau_k(x)$6 to $\tau_k(x)$7 (Takahashi et al., 26 Jun 2025).

A distinct transformer use appears in "Topkima-Former" (Dong et al., 2024), where TopK is applied to attention-score activations before softmax rather than to MLP hidden states. The design keeps only the $\tau_k(x)$8 largest logits in each attention row and ignores the rest on the grounds that softmax amplifies large logits exponentially. The paper combines three levels of co-design: a training scheme called top-$\tau_k(x)$9 forward-complete backward propagation, a scale-free attention implementation to remove explicit division by $k$0, and a circuit-level topkima in-memory ADC mechanism that finds top-$k$1 logits without sorting latency. With $k$2, the reported accuracy reduction is only $k$3 to $k$4 across ViT, distilBERT, and BERT-base on CIFAR-10, CIFAR-100, and SQuAD, while the full Topkima-Former system provides $k$5–$k$6 speedup and $k$7–$k$8 energy efficiency over prior IMC accelerators (Dong et al., 2024).

These two lines of work instantiate two non-equivalent meanings of TopK activation inside transformers: one treats TopK as a sparse hidden-state nonlinearity, the other as an approximation to softmax support over attention logits.

4. Routing, expert selection, and budget granularity

In mixture-of-experts systems, TopK usually denotes router-based expert selection. "Route Experts by Sequence, not by Token" (Wen et al., 9 Nov 2025) formalizes the conventional token-level rule as choosing exactly $k$9 experts per token, then replaces it with sequence-level TopK. Given a length-$x$0 sequence and expert-score matrix $x$1, SeqTopK selects the top $x$2 token-expert pairs across the entire sequence,

$x$3

This preserves the same overall budget while allowing some tokens to receive more than $x$4 experts and others fewer. The paper states that SeqTopK is parameter-free, requires only a few lines of code, adds less than $x$5 overhead, remains fully compatible with pretrained MoE models, and produces gains that become substantially larger under higher sparsity, up to $x$6 (Wen et al., 9 Nov 2025).

The same literature also contains direct critiques of TopK routing. "Routing-Free Mixture-of-Experts" (Liu et al., 1 Apr 2026) argues that Top-K and Softmax impose rigid centralized inductive biases: fixed per-token compute, discontinuous selection, and external-router information bottlenecks. It replaces centralized Top-K routing with expert-local thresholded self-activation,

$x$7

and controls compute through global density regulation rather than fixed per-token cardinality. In that framing, TopK activation is no longer approximated or softened; it is abandoned as the organizing principle of sparse expert usage (Liu et al., 1 Apr 2026).

Taken together, these MoE results show that TopK activation is not only a sparsity operator but also a budgeting rule. Recent work either moves that budget from token level to sequence level or removes fixed-$x$8 routing entirely.

5. Differentiable and systems realizations

Hard TopK is discontinuous, and this has driven a parallel line of work on differentiable sparse top-$x$9 operators. "Successive Halving Top-k Operator" (Pietruszka et al., 2020) proposes a tournament-style relaxation in which candidates are repeatedly sorted, paired strongest-to-weakest, and merged via sharpened pairwise softmax weights. The method avoids iterative softmax over the full score vector, is faster than an earlier iterative baseline, and achieves higher normalized Chamfer Cosine Similarity to exact top-$$y_i = f(x_i)\,\mathbf{1}_{\{\,x_i \ge \tau_k(x)\}},$$0, but it remains only partially differentiable because the sorting permutation is discrete (Pietruszka et al., 2020).

A more general convex formulation appears in "Fast, Differentiable and Sparse Top-k: a Convex Analysis Perspective" (Sander et al., 2023). That paper casts top-$$y_i = f(x_i)\,\mathbf{1}_{\{\,x_i \ge \tau_k(x)\}},$$1 as linear optimization over the permutahedron, adds a $$y_i = f(x_i)\,\mathbf{1}_{\{\,x_i \ge \tau_k(x)\}},$$2-norm regularization term, reduces computation to isotonic optimization, and derives sparse operators that are differentiable almost everywhere for $$y_i = f(x_i)\,\mathbf{1}_{\{\,x_i \ge \tau_k(x)\}},$$3 and differentiable everywhere for $$y_i = f(x_i)\,\mathbf{1}_{\{\,x_i \ge \tau_k(x)\}},$$4. It also introduces a GPU/TPU-friendly Dykstra algorithm in addition to pool-adjacent-violators solvers, and demonstrates the resulting operators in pruning, vision-transformer fine-tuning, and MoE routing (Sander et al., 2023).

TopK also appears as a systems primitive outside model architecture. "Activations and Gradients Compression for Model-Parallel Training" (Rudakov et al., 2024) studies TopK compression of communicated activations and activation gradients between model-parallel stages, defining TopK as selecting the largest $$y_i = f(x_i)\,\mathbf{1}_{\{\,x_i \ge \tau_k(x)\}},$$5 values in absolute value and zeroing the rest. The paper reports that $$y_i = f(x_i)\,\mathbf{1}_{\{\,x_i \ge \tau_k(x)\}},$$6 is the lowest TopK compression level that does not harm convergence severely in its CNN experiments, but also that models trained with TopK perform well only when compression is also applied during inference. Error feedback does not improve training quality relative to plain compression, although it allows inference without compression with almost no quality drop (Rudakov et al., 2024).

6. Interpretability, pathologies, and unresolved limitations

TopK activation has been especially influential in mechanistic interpretability, but the literature distinguishes between feature discovery, causal control, and circuit sparsity. In "TopK LLMs" (Takahashi et al., 26 Jun 2025), sparse hidden states enable successful steering through targeted neuron interventions and support direct analysis of feature formation across checkpoints and layers. In "Mechanistic Interpretability of Antibody LLMs Using SAEs" (Haque et al., 5 Dec 2025), TopK SAEs reveal biologically meaningful latent features in p-IgGen, but high feature concept correlation does not guarantee causal control over generation: TopK features identify concepts effectively, whereas Ordered SAEs are reported to be preferable when precise generative steering is required (Haque et al., 5 Dec 2025).

A separate limitation is that within-layer sparsity does not imply sparse cross-layer interactions. "SCALAR: Benchmarking SAE Interaction Sparsity in Toy LLMs" (Fillingham et al., 10 Nov 2025) uses TopK SAEs as the baseline and argues that independently trained SAEs can reconstruct activations well while still inducing dense inter-feature connectivity across layers. Under its SCALAR benchmark, Staircase SAEs improve relative sparsity over TopK SAEs by $$y_i = f(x_i)\,\mathbf{1}_{\{\,x_i \ge \tau_k(x)\}},$$7 on feedforward blocks and $$y_i = f(x_i)\,\mathbf{1}_{\{\,x_i \ge \tau_k(x)\}},$$8 on transformer blocks, while JSAEs improve over TopK by $$y_i = f(x_i)\,\mathbf{1}_{\{\,x_i \ge \tau_k(x)\}},$$9 on feedforward layers but cannot train effectively across transformer blocks (Fillingham et al., 10 Nov 2025).

The most explicit failure-mode analysis concerns dead features in TopK SAEs. "On the Relationship Between Activation Outliers and Feature Death in Sparse Autoencoders" (Simon et al., 29 May 2026) decomposes initialization-time pre-activations into a constant shift and a varying signal,

$x_i$0

and formalizes outlier severity as

$x_i$1

The paper reports Spearman $x_i$2 between $x_i$3 and dead-by-TopK rates across 454 model-layer combinations, and argues that high-$x_i$4 activation means cause the same aligned features to win the TopK competition repeatedly. Mean-centering the inputs eliminates outlier-induced death across all tested models, and the paper presents this as a principled basis for when that preprocessing step is necessary (Simon et al., 29 May 2026).

Even when dead-feature pathologies are addressed, fixed-$x_i$5 itself remains a known limitation. "Scaling and evaluating sparse autoencoders" (Gao et al., 2024) states that forcing every token to use exactly $x_i$6 latents is likely suboptimal and notes that ordinary TopK SAEs overfit to the training sparsity level. Its Multi-TopK variant trains across more than one budget to obtain a more progressive code. This line of work, together with BatchTopK, HierarchicalTopK, and SeqTopK, indicates that the main open tension is no longer whether TopK can induce useful sparsity, but how rigidly that sparsity budget should be enforced (Gao et al., 2024).