---
title: Optimal Transport SOFT Top-k
url: https://www.emergentmind.com/topics/optimal-transport-soft-top-k
type: topic
---

# Optimal Transport SOFT Top-k

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 [2002.06504, 1905.11885].

## 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$, $\mu_i = 1/n$ (uniform source).
- Target measure $\nu \in \mathbb{R}^2$: $\nu = [k/n, (n-k)/n]^\top$.
- Cost matrix $C\in\mathbb{R}^{n\times 2}$:
  $$
  C_{i1} = (s_i - 0)^2,\quad C_{i2} = (s_i - 1)^2
  $$

The entropic OT (EOT) problem is
$$
T^\varepsilon = \arg\min_{T\in U(s,k)} \langle T, C \rangle + \varepsilon \sum_{i,j} T_{ij}(\log T_{ij} - 1),
$$
subject to $T\mathbf{1}_2 = \mu$, $T^\top \mathbf{1}_n = \nu$, $T\ge 0$.

The solution $T^\varepsilon$ is unique, smooth in $s$, and computed efficiently by Sinkhorn iterations:
$$
K = \exp(-C/\varepsilon),\quad \text{with normalization steps}
$$
updating scaling vectors $u \in \mathbb{R}^n$, $v \in \mathbb{R}^2$ alternately to enforce marginal constraints [2002.06504, 1905.11885].

The output soft top-k indicator is $p = n \cdot T^\varepsilon_{:,2} \in [0,1]^n$, which converges in the $\varepsilon \to 0$ limit to the discrete indicator.

## 3. Differentiability and Gradient Computation

The entropic OT framework ensures the mapping $s\mapsto p$ is differentiable almost everywhere. The solution $T^\varepsilon$ admits dual variables, and gradients can be computed via implicit differentiation of the Sinkhorn fixed-point equations or through KKT conditions:
$$
T^\varepsilon = \operatorname{diag}(\exp(\xi/\varepsilon))\, \exp(-C/\varepsilon)\, \operatorname{diag}(\exp(\zeta/\varepsilon))
$$
The derivatives $\partial T^\varepsilon / \partial C$ and ultimately $\partial p / \partial s$ are computable in closed form due to smoothness of the marginal constraint system. Full block-matrix gradient expressions appear in the respective appendices [2002.06504]. 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:
$$
\min_{T\in\mathbb{R}^{m\times n}} \langle T, C \rangle + \frac{\gamma}{2}\|T\|^2_F + \sum_{j=1}^n \delta_{B_k}(t_j)
$$
where $B_k = \{ t \in \mathbb{R}^m : \|t\|_0 \le k \}$ is the $\ell_0$ ball, and $\delta_{B_k}$ is the corresponding indicator.

The semi-dual problem is
$$
S_k(\mu, b; C) = \max_{\alpha \in \mathbb{R}^m} \langle \alpha, \mu \rangle - \sum_{j=1}^n \Omega^*_{b_j}(\alpha - c_j)
$$
with explicit expressions for the conjugate and its gradient:
$$
\nabla \Omega^*_b(s) = \operatorname{proj}_{b\Delta^m \cap B_k}(s) = [\mathrm{top}_k(s) - \tau 1_m]_+
$$
This yields the "soft top-k" operator: take the $k$ largest entries of $s$, subtract a normalization scalar $\tau$ (found by simplex projection), and threshold. This operator interpolates between hard top-k for $k=1$ and fully dense mapping as $k$ increases [2209.15466].

Gradient-based optimization of this sparse OT is tractable and efficient, scaling as $O(m n \log k)$ 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 ($X$ with density $f$), the other discrete ($Y$ with probabilities $p$), and the cost $c: \mathcal{X}\times \{1,\ldots,K\}\to [0,\infty)$.

The quantile OT problem is:
$$
\inf_{\pi \ge 0} Q^\alpha_\pi[c(X,Y)] \quad \text{subject to} \quad \int_{\mathcal{X}} \pi(x,k)\,dx = p_k;\ \sum_{k=1}^K \pi(x,k) = f(x)
$$
where $Q^\alpha_\pi$ is the $\alpha$-quantile of cost.

Optimal solutions reduce to a finite-dimensional convex program in binary variables $\psi \in \{0,1\}^K$. The "tie-breaking" necessary to preserve prescribed marginals, when the assignment cell $S(x)$ of $x$ comprises multiple indices, is realized by entropy-regularized softmax distributions on the active set.

The softmax tie-breaking rule is formally:
$$
v_{S,k} = \frac{\exp(\theta^*_k)}{\sum_{j\in S} \exp(\theta^*_j)},\quad k\in S
$$
where $\theta^*$ solves a strictly convex, unconstrained dual. Stochastic approximation algorithms suffice for finding $\theta^*$. Provably, induced randomized kernels converge at rate $O(1/\sqrt{m})$, as does the estimation of optimal quantile threshold $t^*(\alpha)$ [2602.10515].

This framework exhibits a "soft power diagram" geometry: for each class $k$, the cell $E_k = \{ x : c(x,k) \le t^* \}$ is defined by sublevel sets; where these overlap, probabilistic mixing occurs based on $\psi^*_k$. 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 $n$ ($O(nkL)$ for $L$ iterations)
- Memory requirements are low (storage of kernel $K$, scaling vectors)
- For sparse OT, projections onto top-$k$+simplex constraints are efficiently computable by sorting ($O(m\log k)$)
- 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 ($O_p(N^{-1/2})$ rates). Each top-k or quantile evaluation reduces to a linear program or explicit projection step [2602.10515, 2209.15466].

## 7. Applications and Extensions

The SOFT top-k paradigm is central to differentiable $k$-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 [2002.06504, 1905.11885, 2209.15466].

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 [2602.10515, 2209.15466, 2002.06504, 1905.11885].

Source: https://www.emergentmind.com/topics/optimal-transport-soft-top-k