Source-linked AI summary
Differentiable Ranks and Sorting using Optimal Transport
Marco Cuturi, Olivier Teboul, Jean-Philippe Vert
TL;DR
Sorting and ranking are widely used in machine learning but provide little gradient information for end-to-end training. The paper connects sorting to one-dimensional optimal transport, generalizes the resulting operators, and makes them differentiable with entropic regularization and Sinkhorn iterations. It applies these operators to learning settings including quantile regression and top-k losses.
Problem
Sorting and ranking support many machine-learning operations, but sorting is nondifferentiable and ranking has a Jacobian that is almost everywhere zero.
Method
The paper models sorting through one-dimensional optimal transport, generalizes target measures and weights, and uses entropic regularization with Sinkhorn iterations to obtain differentiable operators.
Results
The proposed Sinkhorn operators are used for smooth sorting-based losses and applications including quantile regression and differentiable top-k classification losses.
Takeaways & Limitations
Sorting-based metrics and quantile procedures can be incorporated into end-to-end learning through automatically differentiable approximations.
Takeaways & Limitations
The unregularized generalized operators are computationally expensive and remain nondifferentiable, while the K-operators are defined pointwise and depend on the ordering of inputs and target measures.
Abstract
from arXiv · showhide
Sorting an array is a fundamental routine in machine learning, one that is used to compute rank-based statistics, cumulative distribution functions (CDFs), quantiles, or to select closest neighbors and labels. The sorting function is however piece-wise constant (the sorting permutation of a vector does not change if the entries of that vector are infinitesimally perturbed) and therefore has no gradient information to back-propagate. We propose a framework to sort elements that is algorithmically differentiable. We leverage the fact that sorting can be seen as a particular instance of the optimal transport (OT) problem on $\mathbb{R}$, from input values to a predefined array of sorted values (e.g. $1,2,\dots,n$ if the input array has $n$ elements). Building upon this link , we propose generalized CDFs and quantile operators by varying the size and weights of the target presorted array. Because this amounts to using the so-called Kantorovich formulation of OT, we call these quantities K-sorts, K-CDFs and K-quantiles. We recover differentiable algorithms by adding to the OT problem an entropic regularization, and approximate it using a few Sinkhorn iterations. We call these operators S-sorts, S-CDFs and S-quantiles, and use them in various learning settings: we benchmark them against the recently proposed neuralsort [Grover et al. 2019], propose applications to quantile regression and introduce differentiable formulations of the top-k accuracy that deliver state-of-the art performance.
1 Introduction
Sorting and ranking support many machine-learning and statistical operations, but their nondifferentiability limits direct end-to-end learning. The paper frames differentiable proxies through optimal transport and Sinkhorn iterations.
- Ranking and sorting underpin order statistics, nearest-neighbor selection, evaluation metrics, pair or triplet selection, and robust statistical procedures.
- The rank operator R is piecewise constant, so its Jacobian with respect to the input is almost everywhere zero.
- Differentiable proxies could turn sorting-based test metrics into training losses and support rank-based constraints such as quantile-based fairness criteria.
- Prior smoothed-ranking approaches use Gaussian perturbation expectations, logistic pairwise comparisons, or unimodal row-stochastic matrices.
- The paper recovers sorting permutations through optimal assignment and develops generalized operators using flexible target measures, then regularizes OT for differentiable Sinkhorn computation.
- The introduction positions the framework as defining Kantorovich operators before making them differentiable and applying them to smooth alternatives for classification and regression losses.
2 Ranking and Sorting as an Optimal Transport Problem
Optimal transport recasts one-dimensional sorting as matching input values to an increasing target measure. Kantorovich transport generalizes ranks and sorted values to unequal sizes and weights, while entropic regularization produces differentiable operators.
- OT formulation: In one dimension, optimal transport between discrete measures can be computed from their quantile functions, which require sorting the supports.The corresponding optimal transport plan can be constructed using the sorted supports and a permuted north-west corner solution.
- OT formulation: When target values are increasing and source and target measures have equal uniform size, the optimal assignment matches the i-th smallest input to the i-th target value.The assignment matrix is a scaled permutation matrix, revealing both the sorting permutation and the sorted vector.
- Kantorovich generalization: Kantorovich transport allows target arrays with fewer points and non-uniform weights by splitting source mass across multiple target points.The resulting K-rank operator mixes target rank values, while K-sort mixes input values to produce continuous ranks and synthetic quantiles.
- Kantorovich generalization: K-ranks and K-sorts extend ordinary ranks and sorting by computing convex combinations of ranks or input values from the optimal transport plan.K-ranks output continuous rank values for each input, interpretable as scaled synthetic CDF values; K-sorts output m increasing barycentric values.
- Differentiable extension: Entropic regularization replaces sparse optimal plans with dense transport plans, yielding differentiable S-operations with respect to inputs and weights.The paper approximates these regularized plans using the Sinkhorn algorithm.
- Differentiable extension: The generalized operators remain costly and non-differentiable because optimal plans may be non-unique and locally unchanged by perturbations of input values.The stated computation cost is O(nm(n + m) log(nm)), motivating regularized optimal transport.
3 The Sinkhorn Ranking and Sorting Operators
The Sinkhorn operators replace nondifferentiable optimal-transport sorting plans with entropically regularized, dense plans that are differentiable in the inputs. Their smoothness is controlled by ε, while target measures can also define soft quantiles and other generalized operators.
- Differentiable operators: Entropic regularization replaces the nondifferentiable OT plan with a dense transport matrix differentiable with respect to a and x.The regularized plan is approximated through Sinkhorn iterations and differentiated algorithmically after a finite number of iterations.
- Sinkhorn computation: The Sinkhorn algorithm alternates updates of u and v until the transported column masses are within tolerance η of b.The kernel is K = exp(−Cxy/ε), and the iterations use v ← b/(K^T u) and u ← a/(Kv).
- Smoothness control: As ε decreases, Sinkhorn ranks and sorts approach the original integer ranks and sorted values; as ε increases, outputs become smoother.At large ε, ranks collapse toward the average of n̄b and sorted values toward the a-weighted average of x.
- Design choices: The ground cost may use any nonnegative convex h, including h(u) = |u|^p with p = 1 or 2.The paper also rescales inputs to [0,1] and uses a regular grid for y to make ε choices more consistent across value ranges.
- Soft quantiles: Non-uniform target weights yield differentiable soft τ-quantiles by placing a small filler mass between masses approximately τ and 1 − τ.With m = 3, weights [τ − t/2, t, 1 − τ − t/2] and locations [0, 1/2, 1], the second soft-sort entry approximates the τ-quantile.
4 Learning with Smoothed Ranks and Sorts
The paper applies smoothed rank and sorting operators to classification, sorting, and quantile regression, replacing discontinuous objectives with differentiable formulations. Experiments report stronger training quantile errors, comparable test/MSE errors, and competitive or improved sorting performance.
- Learning to sort: Table 1 reports better exact and partial sorting precision than the method in [18] across all sorting tasks, using the same network architecture.The results are averaged over 10 runs.
- Differentiable top-k classification: The proposed differentiable top-k loss combines Sinkhorn ranks with a smoothed Heaviside-like function, replacing discontinuous 0/1 and top-k losses.The formulation is designed as an alternative to cross-entropy and generalized top-k cross-entropy losses.
- Differentiable top-k classification: The method evaluates rank-based learning on CIFAR-10 and CIFAR-100 using a CNN and Resnet18 with ReLU losses J_k(u) = max(0, u − k + 1).Experiments focus on k = 1 and report test-set classification accuracies across epochs.
- Learning to sort: On the MNIST CNN sorting task, S-sort performs on par with NeuralSort using the same experimental setup.The task concatenates four MNIST images representing numbers, with labels given by their ranks; ε = 0.005.
- Least quantile regression: Least quantile regression uses a soft quantile operator with ε = 10^-2 and filler weight t = 1/512 instead of selecting one empirical quantile-error point.The baseline uses ε = 0 and takes gradient steps with respect to the point identified as the desired quantile.
- Least quantile regression: Overall, the method achieves better training quantile errors but comparable test/MSE errors in least quantile regression.Table 2 reports results at the 50% and 90% quantiles, with different behaviors on test sets for quantile loss and MSE.