# Fast, Differentiable and Sparse Top-k: a Convex Analysis Perspective

arXiv: https://arxiv.org/abs/2302.01425 · HTML: https://mlsanity.com/paper/2302.01425/

## Key facts

- Cited by 26 papers in the map (data through 14 Oct 2026).
- Builds on 36 of its 61 references that are in the map.
- First posted on arXiv 2 Feb 2023.

## Builds on (most cited first, 30 of 36)

- [Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer](https://mlsanity.com/paper/1701.06538/) (2017)
- [Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity](https://mlsanity.com/paper/2101.03961/) (2021)
- [Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distances](https://mlsanity.com/paper/1306.0895/) (2013)
- [The Lottery Ticket Hypothesis: Finding Sparse, Trainable Neural Networks](https://mlsanity.com/paper/1803.03635/) (2018)
- [Learning both Weights and Connections for Efficient Neural Networks](https://mlsanity.com/paper/1506.02626/) (2015)
- [Revisiting Unreasonable Effectiveness of Data in Deep Learning Era](https://mlsanity.com/paper/1707.02968/) (2017)
- [Scaling Vision with Sparse Mixture of Experts](https://mlsanity.com/paper/2106.05974/) (2021)
- [Mixture-of-Experts with Expert Choice Routing](https://mlsanity.com/paper/2202.09368/) (2022)
- [What is the State of Neural Network Pruning?](https://mlsanity.com/paper/2003.03033/) (2020)
- [The State of Sparsity in Deep Neural Networks](https://mlsanity.com/paper/1902.09574/) (2019)
- [BASE Layers: Simplifying Training of Large, Sparse Models](https://mlsanity.com/paper/2103.16716/) (2021)
- [Efficient and Modular Implicit Differentiation](https://mlsanity.com/paper/2105.15183/) (2021)
- [Fast Differentiable Sorting and Ranking](https://mlsanity.com/paper/2002.08871/) (2020)
- [A Review of Sparse Expert Models in Deep Learning](https://mlsanity.com/paper/2209.01667/) (2022)
- [Sequence-to-Sequence Learning as Beam-Search Optimization](https://mlsanity.com/paper/1606.02960/) (2016)
- [Learning with Fenchel-Young Losses](https://mlsanity.com/paper/1901.02324/) (2019)
- [DSelect-k: Differentiable Selection in the Mixture of Experts with Applications to Multi-Task Learning](https://mlsanity.com/paper/2106.03760/) (2021)
- [Differentiable Ranks and Sorting using Optimal Transport](https://mlsanity.com/paper/1905.11885/) (2019)
- [Stochastic Optimization of Sorting Networks via Continuous Relaxations](https://mlsanity.com/paper/1903.08850/) (2019)
- [Implicit MLE: Backpropagating Through Discrete Exponential Family Distributions](https://mlsanity.com/paper/2106.01798/) (2021)
- [Differentiable Patch Selection for Image Recognition](https://mlsanity.com/paper/2104.03059/) (2021)
- [SoftSort: A Continuous Relaxation for the argsort Operator](https://mlsanity.com/paper/2006.16038/) (2020)
- [Optimizing Rank-based Metrics with Blackbox Differentiation](https://mlsanity.com/paper/1912.03500/) (2019)
- [Smooth Loss Functions for Deep Top-k Classification](https://mlsanity.com/paper/1802.07595/) (2018)
- [Loss Functions for Top-k Error: Analysis and Insights](https://mlsanity.com/paper/1512.00486/) (2015)
- [Sparsity-Constrained Optimal Transport](https://mlsanity.com/paper/2209.15466/) (2022)
- [Differentiable Sorting Networks for Scalable Sorting and Ranking Supervision](https://mlsanity.com/paper/2105.04019/) (2021)
- [Top-k Multiclass SVM](https://mlsanity.com/paper/1511.06683/) (2015)
- [The Limited Multi-Label Projection Layer](https://mlsanity.com/paper/1906.08707/) (2019)
- [Structured Prediction with Projection Oracles](https://mlsanity.com/paper/1910.11369/) (2019)

## Built on by (most cited first, 26 of 26)

- [Routers in Vision Mixture of Experts: An Empirical Study](https://mlsanity.com/paper/2401.15969/) (2024)
- [Differentiable Clustering with Perturbed Spanning Forests](https://mlsanity.com/paper/2305.16358/) (2023)
- [LapSum -- One Method to Differentiate Them All: Ranking, Sorting and Top-k Selection](https://mlsanity.com/paper/2503.06242/) (2025)
- [High-Similarity-Pass Attention for Single Image Super-Resolution](https://mlsanity.com/paper/2305.15768/) (2023)
- [A path-norm toolkit for modern networks: consequences, promises and challenges](https://mlsanity.com/paper/2310.01225/) (2023)
- [Foundations of Top-k Decoding For Language Models](https://mlsanity.com/paper/2505.19371/) (2025)
- [Universal generalization guarantees for Wasserstein distributionally robust models](https://mlsanity.com/paper/2402.11981/) (2024)
- [Generalizing Stochastic Smoothing for Differentiation and Gradient Estimation](https://mlsanity.com/paper/2410.08125/) (2024)
- [Dynamic Sparse Training of Diagonally Sparse Networks](https://mlsanity.com/paper/2506.11449/) (2025)
- [Geometric Algorithms for Neural Combinatorial Optimization with Constraints](https://mlsanity.com/paper/2510.24039/) (2025)
- [How Many Tokens Do 3D Point Cloud Transformer Architectures Really Need?](https://mlsanity.com/paper/2511.05449/) (2025)
- [Fibottention: Inceptive Visual Representation Learning with Diverse Attention Across Heads](https://mlsanity.com/paper/2406.19391/) (2024)
- [Soft Condorcet Optimization for Ranking of General Agents](https://mlsanity.com/paper/2411.00119/) (2024)
- [Joint Learning of Energy-based Models and their Partition Function](https://mlsanity.com/paper/2501.18528/) (2025)
- [Differentiable Fast Top-K Selection for Large-Scale Recommendation](https://mlsanity.com/paper/2510.11472/) (2025)
- [Learning Cascade Ranking as One Network](https://mlsanity.com/paper/2503.09492/) (2025)
- [Deep greedy unfolding: Sorting out argsorting in greedy sparse recovery algorithms](https://mlsanity.com/paper/2505.15661/) (2025)
- [SoftJAX & SoftTorch: Empowering Automatic Differentiation Libraries with Informative Gradients](https://mlsanity.com/paper/2603.08824/) (2026)
- [Separate, Dynamic and Differentiable (SMART) Pruner for Block/Output Channel Pruning on Computer Vision Tasks](https://mlsanity.com/paper/2403.19969/) (2024)
- [Decision-aware training of spatiotemporal forecasting models to select a top K subset of sites for intervention](https://mlsanity.com/paper/2503.05622/) (2025)
- [Weaver: Kronecker Product Approximations of Spatiotemporal Attention for Traffic Network Forecasting](https://mlsanity.com/paper/2511.08888/) (2025)
- [Differentiable Knapsack and Top-k Operators via Dynamic Programming](https://mlsanity.com/paper/2601.21775/) (2026)
- [Accurate and Efficient Multi-Channel Time Series Forecasting via Sparse Attention Mechanism](https://mlsanity.com/paper/2603.18712/) (2026)
- [Teacher-Guided Routing for Sparse Vision Mixture-of-Experts](https://mlsanity.com/paper/2604.21330/) (2026)
- [LaPrune: Controllable Differentiable Sparsity at Million Scale](https://mlsanity.com/paper/2608.04057/) (2026)
- [Fast LapSum: Exact Differentiable Top-k at Million Scale](https://mlsanity.com/paper/2608.06912/) (2026)

## Abstract (arXiv)

The top-k operator returns a sparse vector, where the non-zero values correspond to the k largest values of the input. Unfortunately, because it is a discontinuous function, it is difficult to incorporate in neural networks trained end-to-end with backpropagation. Recent works have considered differentiable relaxations, based either on regularization or perturbation techniques. However, to date, no approach is fully differentiable and sparse. In this paper, we propose new differentiable and sparse top-k operators. We view the top-k operator as a linear program over the permutahedron, the convex hull of permutations. We then introduce a p-norm regularization term to smooth out the operator, and show that its computation can be reduced to isotonic optimization. Our framework is significantly more general than the existing one and allows for example to express top-k operators that select values in magnitude. On the algorithmic side, in addition to pool adjacent violator (PAV) algorithms, we propose a new GPU/TPU-friendly Dykstra algorithm to solve isotonic optimization problems. We successfully use our operators to prune weights in neural networks, to fine-tune vision transformers, and as a router in sparse mixture of experts.