# Differentiable Top-k Operator with Optimal Transport

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

## Key facts

- Cited by 17 papers in the map (data through 14 Oct 2026).
- First posted on arXiv 16 Feb 2020.

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

- [Differentiable Patch Selection for Image Recognition](https://mlsanity.com/paper/2104.03059/) (2021)
- [A Unified Framework for Implicit Sinkhorn Differentiation](https://mlsanity.com/paper/2205.06688/) (2022)
- [Dodo: Dynamic Contextual Compression for Decoder-only LMs](https://mlsanity.com/paper/2310.02409/) (2023)
- [LapSum -- One Method to Differentiate Them All: Ranking, Sorting and Top-k Selection](https://mlsanity.com/paper/2503.06242/) (2025)
- [Scalable Decision-Focused Learning in Restless Multi-Armed Bandits with Application to Maternal and Child Health](https://mlsanity.com/paper/2202.00916/) (2022)
- [Combinatorial Optimization for Panoptic Segmentation: A Fully Differentiable Approach](https://mlsanity.com/paper/2106.03188/) (2021)
- [Leveraging Recursive Gumbel-Max Trick for Approximate Inference in Combinatorial Spaces](https://mlsanity.com/paper/2110.15072/) (2021)
- [TUSK: Task-Agnostic Unsupervised Keypoints](https://mlsanity.com/paper/2206.08460/) (2022)
- [Differentiable Spline Approximations](https://mlsanity.com/paper/2110.01532/) (2021)
- [Auto-Compressing Subset Pruning for Semantic Image Segmentation](https://mlsanity.com/paper/2201.11103/) (2022)
- [CS-Net:Contribution-based Sampling Network for Point Cloud Simplification](https://mlsanity.com/paper/2501.10789/) (2025)
- [What is the Right Notion of Distance between Predict-then-Optimize Tasks?](https://mlsanity.com/paper/2409.06997/) (2024)
- [GSINA: Improving Subgraph Extraction for Graph Invariant Learning via Graph Sinkhorn Attention](https://mlsanity.com/paper/2402.07191/) (2024)
- [Differentiable Knapsack and Top-k Operators via Dynamic Programming](https://mlsanity.com/paper/2601.21775/) (2026)
- [DiffPrune: differentiable information throttling for token pruning in vision-language models](https://mlsanity.com/paper/2608.01985/) (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 operation, i.e., finding the k largest or smallest elements from a collection of scores, is an important model component, which is widely used in information retrieval, machine learning, and data mining. However, if the top-k operation is implemented in an algorithmic way, e.g., using bubble algorithm, the resulting model cannot be trained in an end-to-end way using prevalent gradient descent algorithms. This is because these implementations typically involve swapping indices, whose gradient cannot be computed. Moreover, the corresponding mapping from the input scores to the indicator vector of whether this element belongs to the top-k set is essentially discontinuous. To address the issue, we propose a smoothed approximation, namely the SOFT (Scalable Optimal transport-based diFferenTiable) top-k operator. Specifically, our SOFT top-k operator approximates the output of the top-k operation as the solution of an Entropic Optimal Transport (EOT) problem. The gradient of the SOFT operator can then be efficiently approximated based on the optimality conditions of EOT problem. We apply the proposed operator to the k-nearest neighbors and beam search algorithms, and demonstrate improved performance.