Optimal Transport SOFT Top-k

• Optimal Transport (OT) is a framework that couples probability measures under cost constraints, widely applied for smooth relaxations of discrete selection.
• The SOFT top-k operator employs entropic and quadratic regularizations to approximate non-differentiable top-k selection, enabling gradient-based learning.
• Efficient algorithms like Sinkhorn iterations and implicit differentiation compute sparsity and gradients, making OT methods scalable for tasks like quantile regression and sparse attention.

Optimal transport (OT) seeks a coupling between two probability measures that matches prescribed marginals and optimizes a given cost. In many classical and modern applications—such as machine learning layers, regression, or selection routines—OT underpins differentiable relaxations of sorting, ranking, quantile estimation, and top-k selection, otherwise non-smooth or combinatorial. The SOFT (Scalable Optimal transport-based diFferenTiable) top-k operator and its related differentiable OT-based approximations constitute a core methodology for end-to-end learning in such settings, leveraging entropic regularization, stochastic root-finding, and sparsity constraints.

1. Discrete Top-k Selection, Non-differentiability, and OT Relaxations

The standard top-k operator maps a score vector $s\in\mathbb{R}^n$ to the set of the $k$ largest elements. Formally, it produces a binary indicator $a\in\{0,1\}^n$ with $\sum_{i=1}^n a_i = k$, where $a_i = 1$ if $s_i$ is among the top-k. This mapping is piecewise constant and discontinuous: its gradient is undefined at ties and vanishes almost everywhere, precluding gradient-based optimization. This motivates relaxation via OT.

Discrete sorting and top-k can be posed as finite OT problems: each data point is a source, targets correspond to the $k$ selection bins, and the goal is to maximize inner product subject to selection constraints:

$$a^* = \arg\max_{a \in \{0,1\}^n, \sum_i a_i = k} a^\top s$$

This combinatorial optimization aligns precisely with an OT problem where mass must be transported from uniform source marginals to targets encoding the selection cardinalities (Xie et al., 2020, Cuturi et al., 2019).

2. Entropic OT and the SOFT Top-k Operator

Introducing an entropy penalty renders the OT problem strictly convex and smooths the solution. Given scores $s\in\mathbb{R}^n$:

• Let $\mu \in \mathbb{R}^n$, $k$0 (uniform source).
• Target measure $k$1: $k$2.
• Cost matrix $k$3:

$k$4

The entropic OT (EOT) problem is

$k$5

subject to $k$6, $k$7, $k$8.

The solution $k$9 is unique, smooth in $a\in\{0,1\}^n$0, and computed efficiently by Sinkhorn iterations:

$a\in\{0,1\}^n$1

updating scaling vectors $a\in\{0,1\}^n$2, $a\in\{0,1\}^n$3 alternately to enforce marginal constraints (Xie et al., 2020, Cuturi et al., 2019).

The output soft top-k indicator is $a\in\{0,1\}^n$4, which converges in the $a\in\{0,1\}^n$5 limit to the discrete indicator.

3. Differentiability and Gradient Computation

The entropic OT framework ensures the mapping $a\in\{0,1\}^n$6 is differentiable almost everywhere. The solution $a\in\{0,1\}^n$7 admits dual variables, and gradients can be computed via implicit differentiation of the Sinkhorn fixed-point equations or through KKT conditions:

$a\in\{0,1\}^n$8

The derivatives $a\in\{0,1\}^n$9 and ultimately $\sum_{i=1}^n a_i = k$0 are computable in closed form due to smoothness of the marginal constraint system. Full block-matrix gradient expressions appear in the respective appendices (Xie et al., 2020). This property is central to enabling backpropagation in complex models using SOFT top-k routines as differentiable layers.

4. Sparsity-Constrained OT and Soft Top-k via Quadratic Regularization

In contrast to entropic regularization, quadratic regularization enables explicit control over sparsity via cardinality constraints. The primal formulation imposes:

$\sum_{i=1}^n a_i = k$1

where $\sum_{i=1}^n a_i = k$2 is the $\sum_{i=1}^n a_i = k$3 ball, and $\sum_{i=1}^n a_i = k$4 is the corresponding indicator.

The semi-dual problem is

$\sum_{i=1}^n a_i = k$5

with explicit expressions for the conjugate and its gradient:

$\sum_{i=1}^n a_i = k$6

This yields the "soft top-k" operator: take the $\sum_{i=1}^n a_i = k$7 largest entries of $\sum_{i=1}^n a_i = k$8, subtract a normalization scalar $\sum_{i=1}^n a_i = k$9 (found by simplex projection), and threshold. This operator interpolates between hard top-k for $a_i = 1$0 and fully dense mapping as $a_i = 1$1 increases (Liu et al., 2022).

Gradient-based optimization of this sparse OT is tractable and efficient, scaling as $a_i = 1$2 per update.

5. Quantile Optimization and Soft Top-k in Semidiscrete OT

Recent work studies minimizing quantiles—not means—of the cost in semidiscrete OT. Here, one measure is continuous ($a_i = 1$3 with density $a_i = 1$4), the other discrete ($a_i = 1$5 with probabilities $a_i = 1$6), and the cost $a_i = 1$7.

The quantile OT problem is:

$a_i = 1$8

where $a_i = 1$9 is the $s_i$0-quantile of cost.

Optimal solutions reduce to a finite-dimensional convex program in binary variables $s_i$1. The "tie-breaking" necessary to preserve prescribed marginals, when the assignment cell $s_i$2 of $s_i$3 comprises multiple indices, is realized by entropy-regularized softmax distributions on the active set.

The softmax tie-breaking rule is formally:

$s_i$4

where $s_i$5 solves a strictly convex, unconstrained dual. Stochastic approximation algorithms suffice for finding $s_i$6. Provably, induced randomized kernels converge at rate $s_i$7, as does the estimation of optimal quantile threshold $s_i$8 (Zhu et al., 11 Feb 2026).

This framework exhibits a "soft power diagram" geometry: for each class $s_i$9, the cell $k$0 is defined by sublevel sets; where these overlap, probabilistic mixing occurs based on $k$1. The geometry interpolates between deterministic power diagrams at the mean and randomized, partially overlapping cells at higher quantiles.

6. Algorithmic Implementation and Scalability

Algorithmic implementation is streamlined by the underlying regularized OT structure. For the SOFT top-k:

• Sinkhorn iterations for small $k$2 ($k$3 for $k$4 iterations)
• Memory requirements are low (storage of kernel $k$5, scaling vectors)
• For sparse OT, projections onto top-$k$6+simplex constraints are efficiently computable by sorting ($k$7)
• Closed-form expressions for gradients are available for both entropic and quadratic regimes

Root-finding for quantile thresholds in the semidiscrete setting exploits monotonicity and empirical process concentration ($k$8 rates). Each top-k or quantile evaluation reduces to a linear program or explicit projection step (Zhu et al., 11 Feb 2026, Liu et al., 2022).

7. Applications and Extensions

The SOFT top-k paradigm is central to differentiable $k$9-nearest neighbors, beam search, sparse mixture-of-experts layers, sparse attention modules, soft cumulative ranking/statistics, and quantile regression. In neural architectures, soft top-k layers allow for end-to-end training—replacing discontinuous selection with differentiable relaxations—while maintaining interpretability and efficient memory usage (Xie et al., 2020, Cuturi et al., 2019, Liu et al., 2022).

Extensions encompass:

• Sorted SOFT top-k with explicit rank preservation
• Unbalanced OT relaxations for adaptive cardinality
• Geometric partitioning with soft power diagrams driven by quantile or tail cost objectives
• Sparse assignment constraints for computational savings in large-scale or expert-model architectures

Empirical analyses indicate comparable or improved performance relative to cross-entropy or standard softmax models across diverse benchmarks.

---

The use of optimal transport, especially its entropic and sparsity-constrained regularizations, has enabled robust, efficient, and provably accurate soft top-k and quantile-based operations for core machine learning tasks (Zhu et al., 11 Feb 2026, Liu et al., 2022, Xie et al., 2020, Cuturi et al., 2019).