Smooth Top-K Relaxations: Differentiable Approaches

• Smooth Top-K relaxations are differentiable approximations of the non-differentiable Top-K operator, enabling gradient-based learning in classification, ranking, and structured prediction tasks.
• They employ methodologies such as log-sum-exp smoothing, convex regularization, dynamic programming, and optimal transport, each balancing computational efficiency, sparsity, and approximation accuracy.
• These techniques enhance model performance in applications like multiclass classification, attention mechanisms, and information retrieval by providing stable dense or sparse gradients and faster convergence.

Smooth Top-K Relaxations enable the integration of inherently discrete Top-K selection operations into differentiable optimization frameworks. They approximate the non-differentiable Top-K operator—critical in classification, recognition, information retrieval, and structured prediction tasks—with continuous, differentiable surrogates, facilitating end-to-end gradient-based learning. These relaxations span a spectrum from log-sum-exp smoothings, convex regularization, optimal transport, dynamic programming, to differentiable sorting/ranking operators. Each approach trades off exactness, computational complexity, gradient structure, sparsity, and practical impact.

1. Mathematical Principles of Top-K Relaxations

The hard Top-$K$ operator selects the $k$ largest elements of a score vector $x\in\mathbb{R}^n$, returning an indicator mask $z\in\{0,1\}^n$ s.t. $\sum_i z_i=k$, with $z_i=1$ if $i$ is among the top $k$. Its discontinuities preclude direct use in neural optimization. Smooth Top-K relaxations construct soft masks $\hat{z}\in[0,1]^n$ (with $\sum_i\hat{z}_i=k$ or $k$0), which approximate the top-$k$1 set in a differentiable manner.

Several core methodologies have been advanced:

• Entropic and log-sum-exp smoothings: Replace hard selection with softmax or log-sum-exp approximations over scores or subsets (Berrada et al., 2018, Pietruszka et al., 2020, Lapin et al., 2016). The temperature parameter controls smoothness; as it vanishes, the relaxation becomes sharp but gradients degenerate.
• Convex and regularized optimization: Frame Top-K as an LP over the permutahedron or simplex and apply $k$2-norm or entropy regularization to yield smooth, often sparse masks (Sander et al., 2023, Lapin et al., 2016).
• Dynamic programming with softmax gates: Express Top-K as a discrete maximization (knapsack DP) and smooth recurrences via differentiable gates and log-sum-exp (Vivier-Ardisson et al., 29 Jan 2026).
• Entropic optimal transport: Formulate Top-K as marginal-constrained optimal transport with entropic regularization, solved via Sinkhorn iterations (Xie et al., 2020).
• Differentiable sorting/ranking operators: Relax permutation matrices to row- or doubly-stochastic matrices admitting continuous dependence on scores, enabling differentiable "top-$k$3" via soft permutations or sorting networks (Petersen et al., 2022).
• Tournament-style successive halving: Cascade pairwise differentiable comparisons to simulate tournament rounds, drastically reducing global softmax cost and yielding faithful soft Top-K masks (Pietruszka et al., 2020).

2. Algorithmic Constructions

Distinct algorithmic constructions offer varying trade-offs between computational cost, approximation quality, and sparsity:

Method	Core Mechanism	Complexity
Iterative Softmax	$k$4 steps of global softmax	$k$5
Successive Halving	$k$6 pairwise softmax tournaments	$k$7
Convex/Permutahedron	$k$8-norm regularized LPs, isotonic	$k$9
Entropic OT (Sinkhorn)	Entropy-regularized transport	$x\in\mathbb{R}^n$0 per iter
Dynamic Prog (Soft-gate)	Smoothed knapsack recursion	$x\in\mathbb{R}^n$1
DiffSort/NeuralSort	Smoothed permutation matrices	$x\in\mathbb{R}^n$2 / $x\in\mathbb{R}^n$3

• Successive Halving (Pietruszka et al., 2020): Arranges $x\in\mathbb{R}^n$4 candidates into a succession of $x\in\mathbb{R}^n$5 rounds of pairwise softmaxes with high "boost" ($x\in\mathbb{R}^n$6); each element's soft Top-K inclusion is the product of its softmax victories along the unique tournament path.
• Permutahedron-based Relaxations (Sander et al., 2023): Formulate Top-K as a linear program over $x\in\mathbb{R}^n$7, relax with $x\in\mathbb{R}^n$8-norm regularization, and solve via isotonic regression (PAV or Dykstra algorithm) for $x\in\mathbb{R}^n$9 cost.
• Soft Dynamic Programming (Vivier-Ardisson et al., 29 Jan 2026): Classical DP recurrences are replaced by soft (log-sum-exp) recursions. The recursion gates are differentiated to obtain the soft mask, with explicit parallelizable forward and backward passes.
• Optimal Transport (Xie et al., 2020): The Top-K mask is the normalized marginal of an OT plan that minimizes a cost plus an entropy term under row/column constraints. Sinkhorn iterations yield the plan; gradients are computed via the KKT implicit function.
• Differentiable Sorting (Petersen et al., 2022): Top-K picks are relaxed via soft permutation matrices (SoftSort, NeuralSort, SinkhornSort, or differentiable sorting networks), enabling smooth estimation of $z\in\{0,1\}^n$0.
• Smooth Top-K Loss via Log-Sum-Exp Subset Sums (Berrada et al., 2018): Top-k SVM or entropy-based loss is regularized with log-sum-exp over all $z\in\{0,1\}^n$1-tuples, with polynomial algebra for efficient computation.

3. Gradient Structure and Optimization Properties

The choice of relaxation strongly shapes the gradient structure, sparsity, and statistical behavior:

• Smooth dense gradients (e.g., softmax, OT, DP relaxations) enable stable SGD, critical for deep learning (Berrada et al., 2018, Xie et al., 2020, Pietruszka et al., 2020).
• Sparsity control: $z\in\{0,1\}^n$2-norm and isotonic relaxations with $z\in\{0,1\}^n$3 can be exactly $z\in\{0,1\}^n$4-sparse; Shannon-entropy and log-sum-exp relaxations yield strictly dense masks for all $z\in\{0,1\}^n$5 (Sander et al., 2023, Vivier-Ardisson et al., 29 Jan 2026).
• Convexity and calibration: Many relaxations are convex, e.g., Moreau–Yosida-smoothed SVM and top-$z\in\{0,1\}^n$6 entropy (Lapin et al., 2016). Convexity facilitates global convergence and closed-form gradients.
• Permutation equivariance: Within the dynamic programming framework, the Shannon entropy is uniquely determined by permutation equivariance; other regularizers can induce bias or lose symmetry (Vivier-Ardisson et al., 29 Jan 2026).
• Gradient analytic formulas: For isotonic ($z\in\{0,1\}^n$7-norm) and dynamic programming approaches, explicit closed-form or blockwise Jacobian formulas (via implicit differentiation or chain rule) enable efficient backpropagation (Sander et al., 2023, Vivier-Ardisson et al., 29 Jan 2026, Xie et al., 2020).

4. Application Domains and Adaptations

Smooth Top-K relaxations are fundamental for:

• Multiclass/top-$z\in\{0,1\}^n$8 classification: Smooth Top-K SVM and entropy losses, generalizing softmax, improve both top-1 and top-k accuracies, with calibration guarantees (Lapin et al., 2016, Berrada et al., 2018, Petersen et al., 2022).
• Ranking and information retrieval: The need for differentiable NDCG@K or recall@K metrics has led to quantile- and softmax-based upper-bound relaxations, enabling the direct optimization of ranking objectives (Yang et al., 4 Aug 2025).
• Sparse and mixture-of-experts routing: Sparse Top-K masks (isotonic and Dykstra variants) allow efficient routing in large parameter models (Vision MoE), achieving superior throughput and accuracy (Sander et al., 2023).
• Structured and decision-focused learning: Smoothed Top-K enables neural architectures to incorporate greedy or combinatorial selection (e.g., differentiable beam search, dynamic assortment RL) within gradient-based learning (Xie et al., 2020, Vivier-Ardisson et al., 29 Jan 2026).
• Attention and selection networks: Soft Top-K used for enforcing sparsity in attention (e.g., Top-K attention) and for robust, trainable neighbor selection in k-nearest neighbor modules (Xie et al., 2020).

5. Computational and Empirical Comparisons

Computation scales as follows:

• Classical iterative softmax methods incur $z\in\{0,1\}^n$9 cost, with entangled gradients and significant runtime issues for large $\sum_i z_i=k$0 or $\sum_i z_i=k$1 (Pietruszka et al., 2020, Lapin et al., 2016).
• Successive halving reduces both the number of softmax operations and the chain length for backpropagation, achieving 2–10× faster runtimes than iterative baselines and higher normalized Chamfer–Cosine Similarity (nCCS) to hard masks (Pietruszka et al., 2020).
• Isotonic/permutahedron relaxations yield order-of-magnitude runtime improvements for sparse Top-K, with exact sparsity when $\sum_i z_i=k$2 (Sander et al., 2023).
• OT-based and DP methods provide explicit bias-variance accounting and allow parallel hardware acceleration (Xie et al., 2020, Vivier-Ardisson et al., 29 Jan 2026).
• Empirical studies confirm faster convergence (e.g., 20–50% fewer epochs for the successive-halving layer), superior robustness to label noise, and better stability in loss landscapes compared to both naïve surrogates and non-smooth approaches (Berrada et al., 2018, Pietruszka et al., 2020, Yang et al., 4 Aug 2025).

6. Extensions, Theoretical Insights, and Open Questions

Smooth Top-K frameworks generalize in several directions:

• Magnitude-based and signed Top-K: Isotonic methods admit smooth selection based on absolute or signed score magnitude, with links to OWL and $\sum_i z_i=k$3-support norms (Sander et al., 2023).
• Truncated and adaptive-$\sum_i z_i=k$4 relaxations: Mixtures over $\sum_i z_i=k$5 (e.g., top-1 and top-5) in differentiable sort-based objectives enhance performance on both metrics (Petersen et al., 2022).
• Scaling and implementation: Parallel algorithms for Dykstra, DP recursions, and sorting networks make large-scale, GPU/TPU deployment efficient (Sander et al., 2023, Vivier-Ardisson et al., 29 Jan 2026, Petersen et al., 2022).
• Fundamental limits: Only Shannon entropy guarantees permutation-equivariant, fully dense smoothings; other regularizers trade-off sparsity against symmetry and approximation (Vivier-Ardisson et al., 29 Jan 2026). Sparse relaxations may lose permutation symmetry or fail to provide everywhere-differentiable masks.
• Calibration: Standard softmax and smooth SVM objectives attain uniform top-k calibration, whereas some truncated or hybrid losses do not, impacting risk consistency (Lapin et al., 2016).

7. Representative Empirical Results

Method	Setting	Metric	Empirical Gain
Successive Halving	Synthetic, CIFAR-10	nCCS	0.98–0.995 vs. 0.93–0.98 (softmax)
Dykstra Top-K (ViT MoE, JFT)	MoE routing	precision@1	+ improvement over discrete
SoftmaxLoss@$\sum_i z_i=k$6 (RS)	RecSys (NDCG@K)	NDCG@20	+6.03% avg. over prior best
Smooth Top-$\sum_i z_i=k$7 Loss (Berrada et al., 2018)	CIFAR-100, ImageNet	Acc@5	+2–5% under label noise
DiffSort Top-$\sum_i z_i=k$8 (Petersen et al., 2022)	ImageNet-1K	Acc@1,5	+0.2% Acc@1, +0.17% Acc@5

These results underpin the practical utility of Smooth Top-K relaxations in both accuracy and efficiency across domains (Pietruszka et al., 2020, Sander et al., 2023, Yang et al., 4 Aug 2025, Petersen et al., 2022, Berrada et al., 2018).

---

References

• "Successive Halving Top-k Operator" (Pietruszka et al., 2020)
• "Fast, Differentiable and Sparse Top-k: a Convex Analysis Perspective" (Sander et al., 2023)
• "Differentiable Knapsack and Top-k Operators via Dynamic Programming" (Vivier-Ardisson et al., 29 Jan 2026)
• "Breaking the Top-$\sum_i z_i=k$9 Barrier: Advancing Top-$z_i=1$0 Ranking Metrics Optimization in Recommender Systems" (Yang et al., 4 Aug 2025)
• "Analysis and Optimization of Loss Functions for Multiclass, Top-k, and Multilabel Classification" (Lapin et al., 2016)
• "Differentiable Top-k Operator with Optimal Transport" (Xie et al., 2020)
• "Smooth Loss Functions for Deep Top-k Classification" (Berrada et al., 2018)
• "Differentiable Top-k Classification Learning" (Petersen et al., 2022)