---
title: TopK Activation in Deep Learning
url: https://www.emergentmind.com/topics/topk-activation
type: topic
---

# TopK Activation in Deep Learning

Searching arXiv for recent 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 [2406.04093] [2506.21468] [2411.13050] [2511.06494].

## 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 Language Models" [2506.21468], 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_i \ge \tau_k(x)\), ties can preserve more than \(k\) coordinates [2506.21468].

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 [2406.04093]. In other settings it follows ReLU, so the competition is effectively among nonnegative post-ReLU values rather than signed pre-activations [2512.05794]. In "Topkima-Former" [2411.13050], the selected object is neither a latent code nor a token subset but the attention-score activation before softmax: only the \(k\) largest logits in a row of \(QK^T\) are sent to the softmax calculation block. The paper is explicit that this is not top-\(k\) token pruning, MoE gating, or generic activation sparsification; it is top-\(k\) selection on attention-score activations before softmax [2411.13050].

## 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" [2406.04093], a TopK SAE replaces the standard
\[
z = \mathrm{ReLU}(W_{\mathrm{enc}}(x-b_{\mathrm{pre}})+b_{\mathrm{enc}})
\]
with
\[
z = \mathrm{TopK}(W_{\mathrm{enc}}(x-b_{\mathrm{pre}})),
\qquad
\hat{x}=W_{\mathrm{dec}}z+b_{\mathrm{pre}},
\]
and trains with pure reconstruction loss
\[
\mathcal{L} = \|x-\hat{x}\|_2^2.
\]
The stated motivation is direct control of the number of active latents, avoidance of \(L_1\)-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 [2406.04093].

Subsequent work relaxes the fixed per-sample budget without discarding TopK selection itself. "BatchTopK Sparse Autoencoders" [2412.06410] defines BatchTopK by selecting the top \(n \times k\) activations across an entire mini-batch of \(n\) samples, rather than the top \(k\) 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 [2412.06410].

A further extension is "HierarchicalTopK" [2505.24473], which trains one SAE to work across multiple sparsity budgets up to a maximum \(K\). Instead of optimizing only the reconstruction using the top \(k\) features, it optimizes all prefix reconstructions
\[
\hat{\mathbf{x}}_j = \sum_{i\in\operatorname{top}_j}\mathbf{l}_i(\mathbf{x})\,\mathbf{e}_i + \mathbf{b}_{\mathrm{dec}},
\]
with objective
\[
\mathcal{L}_{\mathrm{hierarchical}}=
\frac{1}{|\mathcal{J}|}\sum_{j\in\mathcal{J}}\mathcal{L}_{\mathrm{rec}}^j.
\]
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 [2505.24473].

TopK activation has also been combined with expert routing rather than replaced by it. "Efficient Dictionary Learning with Switch Sparse Autoencoders" [2410.08201] 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 [2410.08201].

## 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 Language Models" [2506.21468], 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 \(L-n_{\mathrm{nontopk}}\) layers and leaves the final \(n_{\mathrm{nontopk}}=2\) layers dense. It uses an annealed activation
\[
y = \alpha f(x) + (1-\alpha)\bigl(f(x)\odot \mathbf{1}_{\{x \ge \tau_k(x)\}}\bigr),
\]
with \(\alpha\) decayed linearly from 1 to 0 over the first \(20\%\) of training steps, and a main setting of \(k=64\). 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, \(D=1024\) setting, validation perplexity rises from \(11.76\) to \(14.96\), while LAMBADA accuracy rises from \(31.61\%\) to \(33.32\%\) [2506.21468].

A distinct transformer use appears in "Topkima-Former" [2411.13050], where TopK is applied to attention-score activations before softmax rather than to MLP hidden states. The design keeps only the \(k\) 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-\(k\) forward-complete backward propagation, a scale-free attention implementation to remove explicit division by \(\sqrt{d_k}\), and a circuit-level topkima in-memory ADC mechanism that finds top-\(k\) logits without sorting latency. With \(k=5\), the reported accuracy reduction is only \(0.4\%\) to \(1.2\%\) across ViT, distilBERT, and BERT-base on CIFAR-10, CIFAR-100, and SQuAD, while the full Topkima-Former system provides \(1.8\times\)–\(84\times\) speedup and \(1.3\times\)–\(35\times\) energy efficiency over prior IMC accelerators [2411.13050].

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" [2511.06494] formalizes the conventional token-level rule as choosing exactly \(K\) experts per token, then replaces it with sequence-level TopK. Given a length-\(T\) sequence and expert-score matrix \(S \in \mathbb{R}^{T\times N}\), SeqTopK selects the top \(T\cdot K\) token-expert pairs across the entire sequence,
\[
R_i(h_t)=
\begin{cases}
1, & \text{if } (t,i)\in \argtopk_{t\in[T],\,i\in[N]} s_{t,i},\\
0, & \text{otherwise}.
\end{cases}
\]
This preserves the same overall budget while allowing some tokens to receive more than \(K\) experts and others fewer. The paper states that SeqTopK is parameter-free, requires only a few lines of code, adds less than \(1\%\) overhead, remains fully compatible with pretrained MoE models, and produces gains that become substantially larger under higher sparsity, up to \(16.9\%\) [2511.06494].

The same literature also contains direct critiques of TopK routing. "Routing-Free Mixture-of-Experts" [2604.00801] 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,
\[
G_i(\mathbf{x})=\mathrm{ReLU}(\|\mathbf{xA}_{\mathrm{gate},i}\|_2-b_i),\qquad
f_i(\mathbf{x})=\mathds{1}\{G_i(\mathbf{x})-\theta\ge 0\},
\]
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 [2604.00801].

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-\(K\) routing entirely.

## 5. Differentiable and systems realizations

Hard TopK is discontinuous, and this has driven a parallel line of work on differentiable sparse top-\(k\) operators. "Successive Halving Top-k Operator" [2010.15552] 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-\(k\), but it remains only partially differentiable because the sorting permutation is discrete [2010.15552].

A more general convex formulation appears in "Fast, Differentiable and Sparse Top-k: a Convex Analysis Perspective" [2302.01425]. That paper casts top-\(k\) as linear optimization over the permutahedron, adds a \(p\)-norm regularization term, reduces computation to isotonic optimization, and derives sparse operators that are differentiable almost everywhere for \(p=2\) and differentiable everywhere for \(1<p<2\). 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 [2302.01425].

TopK also appears as a systems primitive outside model architecture. "Activations and Gradients Compression for Model-Parallel Training" [2401.07788] studies TopK compression of communicated activations and activation gradients between model-parallel stages, defining TopK as selecting the largest \(K\%\) values in absolute value and zeroing the rest. The paper reports that \(K=10\%\) 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 [2401.07788].

## 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 Language Models" [2506.21468], 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 Language Models Using SAEs" [2512.05794], 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 [2512.05794].

A separate limitation is that within-layer sparsity does not imply sparse cross-layer interactions. "SCALAR: Benchmarking SAE Interaction Sparsity in Toy LLMs" [2511.07572] 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 \(59.67\% \pm 1.83\%\) on feedforward blocks and \(63.15\% \pm 1.35\%\) on transformer blocks, while JSAEs improve over TopK by \(8.54\% \pm 0.38\%\) on feedforward layers but cannot train effectively across transformer blocks [2511.07572].

The most explicit failure-mode analysis concerns dead features in TopK SAEs. "On the Relationship Between Activation Outliers and Feature Death in Sparse Autoencoders" [2605.31518] decomposes initialization-time pre-activations into a constant shift and a varying signal,
\[
z_i(\mathbf{x}) = \mathbf{w}_i\cdot \boldsymbol{\mu} + \mathbf{w}_i\cdot(\mathbf{x}-\boldsymbol{\mu}),
\]
and formalizes outlier severity as
\[
\gamma = \frac{\|\boldsymbol{\mu}\|}{\|\boldsymbol{\sigma}\|}.
\]
The paper reports Spearman \(\rho=0.89\) between \(\gamma\) and dead-by-TopK rates across 454 model-layer combinations, and argues that high-\(\gamma\) 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 [2605.31518].

Even when dead-feature pathologies are addressed, fixed-\(k\) itself remains a known limitation. "Scaling and evaluating sparse autoencoders" [2406.04093] states that forcing every token to use exactly \(k\) 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 [2406.04093].

Source: https://www.emergentmind.com/topics/topk-activation