TopK Activation in Deep Learning

• TopK activation function is a neural technique that retains only the top k highest activations, promoting sparsity and efficient computation in models.
• Different variants, including differentiable relaxations such as ASH, Smooth TopK, and convex formulations, enable feasible gradient flow and end-to-end learning.
• Its practical integration in architectures like Transformers and CNNs enhances interpretability and performance in language and vision tasks.

The TopK activation function promotes sparsity and selectivity in neural representations by retaining only the k highest (or most "informative") activations in a given vector or tensor, zeroing out the rest. Originally employed in sparse coding and winner-take-all circuits, TopK and its differentiable analogues are now critical for engineered sparsity, interpretable models, and computational efficiency in deep learning architectures, including LLMs and vision systems.

1. Mathematical Definitions and Variants

The canonical TopK activation function processes an input vector $z = (z_1, ..., z_d) \in \mathbb{R}^d$, a nonlinearity $f(\cdot)$ (e.g., ReLU), and a sparsity budget $k$ $(1 \leq k \leq d)$. Defining the $k$-th order statistic $\tau_k(z)$ as the $k$-th largest element of $z$, the hard TopK activation $T_k$ is given by: $$y_i = \begin{cases} f(z_i) & \text{if } z_i \geq \tau_k(z) \ 0 & \text{otherwise} \end{cases}$$ or equivalently, in vector form: $y = f(z) \odot 1_{z \geq \tau_k(z)}$ where $1_{z \geq \tau_k(z)}$ is the elementwise indicator mask (Takahashi et al., 26 Jun 2025).

Variants in the literature include:

• ASH (Adaptive SwisH): Instead of a hard threshold, ASH computes a learnable, feature-wise Z-score threshold $\theta = \mu_X + z_k \sigma_X$ (with $\mu_X, \sigma_X$ the mean and standard deviation over the activation tensor, and $z_k$ a trainable scalar), and applies a smooth sigmoid sigmoid-based thresholding:

$$A(x^{(i)}) \approx x^{(i)} S(2\alpha(x^{(i)}-\mu_X-z_k\sigma_X))$$

leading to a form $A(x) = x S(a x + b)$, which subsumes Swish and approaches TopK in the limit of high steepness and appropriate $z_k$ (Lee et al., 2022).

• Smooth TopK (Optimal Transport): The SOFT TopK operator provides a continuous relaxation by casting selection as an Entropic Optimal Transport (EOT) problem. For $x \in \mathbb{R}^n$, it seeks an assignment matrix $\Gamma^{*}$ minimizing

$$\langle C, \Gamma \rangle + \epsilon\sum_{i,j}\Gamma_{ij}\ln \Gamma_{ij}$$

subject to marginal constraints, where the smoothed mask is $A^\epsilon = n \Gamma_{:,1}$ and $\epsilon$ is a smoothness hyperparameter (Xie et al., 2020).

• Convex/Isotonic TopK: By posing TopK as a convex program over the permutahedron with a $p$-norm regularization, one attains a differentiable, sparse operator. Computationally, this reduces to isotonic regression, solvable via the Pool-Adjacent-Violators (PAV) or vectorized Dykstra algorithms (Sander et al., 2023).

2. Differentiable Relaxations and Gradient Flow

Naive (hard) TopK is non-differentiable due to discontinuous index selection; gradients with respect to $x$ are zero almost everywhere. To enable end-to-end learning:

• Sigmoid Relaxations: Replace the step function or hard threshold with a steep sigmoid, yielding nonzero gradients with respect to both the threshold and $x$. Used in ASH, this approach ensures backpropagation can tune context-dependent thresholds per layer or feature map (Lee et al., 2022).
• Optimal Transport/Entropic Smoothing: SOFT TopK leverages entropic regularization, using forward–backward Sinkhorn–Knopp iterations, and computes gradients via implicit differentiation of KKT conditions, resulting in $O(n)$ complexity per pass and practical end-to-end differentiability (Xie et al., 2020).
• Convex Analysis: The permutahedron and $p$-norm smoothed formulations allow for isotonic regression solvers with closed-form Jacobians for efficient backpropagation. The choice of $p$ determines smoothness and sparsity: as $p \rightarrow 1^+$ one approaches hard TopK but risks ill-conditioned gradients; $p=2$ yields smoother transitions (Sander et al., 2023).

3. Algorithmic and Implementation Considerations

Efficient implementation is critical for scaling TopK activations in modern architectures.

• Sorting/Selection: Hard TopK requires selection of the $k$ largest elements per input vector/tensor, implemented via batched primitives (e.g., torch.topk) on GPU. Average complexity is $O(d)$, and memory overhead for masks is usually negligible compared to activations (Takahashi et al., 26 Jun 2025).
• Threshold Annealing: In practice, a linear annealing schedule:

$$y = \alpha f(z) + (1-\alpha)[f(z)\odot 1_{z\geq\tau_k(z)}]$$

is used, with $\alpha$ decayed from 1 to 0 during an initial period (e.g., first 20% of training steps), easing optimization and reducing early sparsity's destabilizing effects (Takahashi et al., 26 Jun 2025).

• GPU/TPU Optimization: Smooth TopK variants using isotonic regression (e.g., via Dykstra's method) are vectorizable and can be parallelized efficiently. PAV is $O(n)$ but not easily vectorized (Sander et al., 2023).
• Forward Passes in ASH: ASH avoids sorting with purely elementwise operations, relying on contextual (dynamic) thresholds using per-feature-map statistics, leading to memory-coalesced, GPU-friendly code (Lee et al., 2022).

4. Integration into Neural Network Architectures

TopK activation was recently integrated directly into Transformer-based LLMs ("TopK LMs") for intrinsic sparsity and interpretability:

• Transformer Modification: In TopK LMs, the activation in the feed-forward (MLP) block of the first $L-n_{\textrm{nontopk}}$ layers is replaced by TopK, while the last $n_{\textrm{nontopk}}$ layers retain dense activations (e.g., ReLU or GELU) for expressivity (Takahashi et al., 26 Jun 2025).
• ASH in CNNs: ASH is incorporated as the main nonlinearity after convolutional layers, allowing for per-layer and per-feature-map adaptive thresholds (Lee et al., 2022).
• Sparse Mixture of Experts and Routers: The convex/differentiable TopK serves as a router in sparse MoE architectures, and as a masking mechanism in pruning and feature selection (Sander et al., 2023).

5. Hyperparameters and Trade-Offs

Key hyperparameters across TopK variants:

Parameter	Effect	Empirical Setting
$k$	Sparsity ratio ($k/d$)	$k/d \approx 5$–$10\%$
$\alpha$	Anneal dense$\to$sparse schedule	Linearly to 0 in 10–30% steps
$p$	Smoothness of convex TopK	$p=2$ (stable), $p=4/3$ (smooth)
$\epsilon$	Smoothing in SOFT TopK	Small for near-discrete; not too small for stable gradients
$n_{\textrm{nontopk}}$	Dense final layers in Transformer	$n_{\textrm{nontopk}}=2$ for $L\leq24$

Smaller $k$ increases sparsity and selectivity but may degrade expressiveness; larger $k$ recovers dense models. Excessively sharp thresholds ($p \approx 1$ or very small $\epsilon$) risk gradient instability. Final performance depends on architectural context and dataset (Takahashi et al., 26 Jun 2025, Lee et al., 2022, Sander et al., 2023).

6. Empirical Properties and Interpretability

Empirical studies across modalities establish the following:

• Vision Tasks (ASH): On ImageNet, CIFAR-10/100, ADE20K, and COCO, ASH surpasses ReLU, GELU, and even Swish by 0.5–1.5% in top-1 accuracy or mean IoU/mAP, and enables 10–20% faster convergence. Grad-CAM reveals that ASH produces more sharply localized feature activations (Lee et al., 2022).
• Language Modeling (TopK LM): Slight increase in validation perplexity (e.g., from 11.76 to 14.96 for $d=1024, L=24$) is observed, but zero-shot accuracies remain comparable to the dense baseline. Neuron activations become highly monosemantic and can be causally manipulated during generation for steerability and targeted interventions (e.g., amplifying neuron 22:894 yields text about “work”) (Takahashi et al., 26 Jun 2025).
• Sparse Attention / kNN / Beam Search (SOFT TopK): Replacing hard selection with differentiable TopK provides SOTA or improved classification and decoding accuracy, enhances interpretability in attention mechanisms, and closes the exposure bias gap in beam search (Xie et al., 2020).
• Checkpoint Stability: Fixing neuron indices with TopK restores feature traceability across training, unlike post-hoc sparse autoencoders where feature permutations undermine interpretability (Takahashi et al., 26 Jun 2025).

7. Theoretical and Practical Significance

TopK activations represent a principled approach to embedding explicit, controlled sparsity and nonlinearity in neural function. Differentiable TopK mechanisms unify the strengths of thresholding, selectivity, and learned context, enabling new regimes of model pruning, mixture-of-experts routing, and interpretable representation learning. The generalization to smooth, convex, or adaptive versions preserves the essential sparsity while opening gradient-based optimization, establishing TopK as a foundational building block in modern deep learning systems (Lee et al., 2022, Takahashi et al., 26 Jun 2025, Xie et al., 2020, Sander et al., 2023).