Source-linked AI summary
Fast Differentiable Sorting and Ranking
Mathieu Blondel, Olivier Teboul, Quentin Berthet, Josip Djolonga
TL;DR
Sorting and ranking are difficult to differentiate, and existing differentiable proxies do not attain sorting-like complexity. The paper constructs exact differentiable operators through regularized permutahedron projections and isotonic optimization, achieving O(n log n) computation and an order-of-magnitude speedup while enabling two applications.
Problem
Sorting has nondifferentiable kinks, ranking has null or undefined derivatives, and existing differentiable ranking approaches use worse-than-O(n log n) complexity.
Method
The paper regularizes permutahedron-based linear programs and uses isotonic optimization to compute and differentiate soft sorting and ranking operators.
Results
The operators achieve O(n log n) computation and O(n) differentiation, with an order-of-magnitude speedup over existing approaches.
Takeaways & Limitations
The operators serve as drop-in replacements for O(n^2) methods and support differentiable Spearman rank correlation and soft least trimmed squares.
Takeaways & Limitations
For top-k classification, scores must be squashed to [0, 1]^n and ε tuned; in soft least trimmed squares, existing O(n^2) operators can be prohibitive when n is the number of training samples.
Abstract
from arXiv · showhide
The sorting operation is one of the most commonly used building blocks in computer programming. In machine learning, it is often used for robust statistics. However, seen as a function, it is piecewise linear and as a result includes many kinks where it is non-differentiable. More problematic is the related ranking operator, often used for order statistics and ranking metrics. It is a piecewise constant function, meaning that its derivatives are null or undefined. While numerous works have proposed differentiable proxies to sorting and ranking, they do not achieve the $O(n \log n)$ time complexity one would expect from sorting and ranking operations. In this paper, we propose the first differentiable sorting and ranking operators with $O(n \log n)$ time and $O(n)$ space complexity. Our proposal in addition enjoys exact computation and differentiation. We achieve this feat by constructing differentiable operators as projections onto the permutahedron, the convex hull of permutations, and using a reduction to isotonic optimization. Empirically, we confirm that our approach is an order of magnitude faster than existing approaches and showcase two novel applications: differentiable Spearman's rank correlation coefficient and least trimmed squares.
1. Introduction
Sorting and ranking are useful but poorly differentiated: sorting has many kinks, while ranking is piecewise constant with null or undefined derivatives. The paper introduces exact, efficient differentiable operators based on permutahedron projections and isotonic optimization, and demonstrates applications and speedups.
- Sorting is piecewise linear with many nondifferentiable kinks and can induce non-convexity when composed with other functions.
- Ranking is piecewise constant, so its derivatives are null or undefined and cannot support gradient backpropagation.
- The proposed differentiable sorting and ranking operators have O(n log n) time and O(n) memory complexity with exact computation and differentiation.
- The approach casts sorting and ranking as linear programs over the permutahedron and regularizes them into differentiable projections.
- O(n log n) computation and O(n) differentiation are achieved through a reduction to isotonic optimization.
- An order-of-magnitude speedup is reported alongside differentiable Spearman’s rank coefficient and soft least trimmed squares applications.
2. Preliminaries
The paper defines score vectors, permutations, sorting, argsort, and ranking using descending order. These operations can be computed in O(n log n) time.
- The input θ is an n-dimensional vector of model scores or logits, potentially containing label scores.
- A permutation reorders vector coordinates, while ρ denotes the reversing permutation (n, n −1, . . . , 1).
- Argsort returns indices σ such that θσ1 ≥ · · · ≥ θσn, with ties broken arbitrarily.
- Sorting returns the input values in descending order, whereas ranking assigns each coordinate its position in that descending sort.
- Sorting, argsort, and ranking can all be computed in O(n log n) time.
3. Sorting and ranking as linear programs
Sorting and ranking can be formulated as linear programs over the permutahedron, whose vertices are permutations. This formulation exposes sorting’s almost-everywhere differentiability but ranking’s unusable Jacobian.
- Argsort and ranking are first formulated as optimization problems over the set of permutations.
- The permutahedron is the convex hull of permutations of a vector, forming a convex polytope whose vertices correspond to those permutations.
- Linear programs over the permutahedron recover sorting and ranking because optimal solutions are almost surely attained at permutation vertices.
- Sorting is piecewise linear and differentiable almost everywhere, with a permutation-matrix Jacobian when the ordering is unique.
- Ranking is discontinuous and piecewise constant, so small perturbations can change its selected permutation and leave derivatives null or undefined.
4. Differentiable sorting and ranking
The paper regularizes permutahedron-based linear programs to obtain differentiable soft sorting and ranking operators. Their behavior is controlled by ε, with fast projection algorithms and a trade-off between fidelity to hard operators and smoothness.
- Strongly convex regularization turns the sorting and ranking linear programs into differentiable projection operators.
- As ε →0, soft sorting and ranking converge to their hard counterparts; as ε →∞, they collapse to constants or means.
- The operators use quadratic or entropic regularization, with entropic regularization interpreted through a log KL projection.
- Larger ε reduces sorting kinks and makes objectives increasingly convex, but increases departure from hard sorting or ranking.
- Soft operators are differentiable almost everywhere and preserve descending sort order and the corresponding rank order for finite ε.
- For top-k classification, squashing scores to [0, 1]^n and tuning ε is important, whereas tuning ε was unnecessary in the label-ranking experiment.
- The approach regularizes permutahedron projections directly, enabling O(n log n) time and O(n) space instead of approximate O(Tn^2) Sinkhorn computation.
5. Fast computation and differentiation
The paper reduces permutahedron projections to isotonic optimization, enabling exact soft sorting and ranking with O(n log n) computation and O(n) differentiation. Its Jacobian products are linear-time, while PAV provides exact optimization without iterative precision choices.
- Reduction to isotonic optimization: O(n log n) forward computation and O(n) backward differentiation are obtained by reducing permutahedron projections to isotonic optimization.The reduction uses simple chain constraints and underlies the fast computation and differentiation of the soft operators.
- Computation: PAV returns exact solutions for both quadratic and entropic isotonic problems in O(n) time.Unlike Sinkhorn, PAV requires neither a chosen iteration count nor a precision level.
- Computation: The total projection complexity is O(n log n) because sorting precedes the O(n)-time isotonic computation.Computing the projections requires obtaining the sorted scores first.
- Differentiation: Jacobian-vector products for isotonic optimization take O(n) time in both quadratic and entropic regularization cases.Quadratic regularization averages gradients uniformly within blocks, whereas entropic regularization uses softmax-weighted averages.
- Differentiation: Permutation of Jacobian rows and columns preserves linear-time multiplication for the full projection Jacobian.The identity (Jπ)z = (Jzπ−1)π reuses the O(n) isotonic-Jacobian multiplication.
- Differentiation: The resulting soft sorting and ranking operators can multiply their Jacobians in O(n) time and space.Differentiation of the operators follows by applying the chain rule to the projection formulas.
6. Experiments
Experiments show that the proposed soft sorting and ranking operators retain competitive accuracy while improving scalability, and support label ranking and robust regression applications.
- Top-k classification: On CIFAR-10 and CIFAR-100, the proposed soft rank formulations achieve comparable accuracy to OT while being significantly faster.The comparison averages results over 12 runs.
- Top-k classification: 29 hours (OT), 21 hours (rQ), 23 hours (rE) and 16 hours (All-pairs) were required for 600 CIFAR-100 training epochs.The proposed operators were faster than OT but slower than All-pairs in this GPU comparison.
- Runtime comparison: effect of input dimension: OT and All-pairs scale quadratically with n, whereas the proposed operators use O(n) memory and remain scalable as dimensionality increases.OT and All-pairs ran out of memory from n = 2000 and n = 3000, respectively, while backpropagation caused out-of-memory at n = 1000 and n = 2500.
- Label ranking via soft Spearman’s rank correlation coefficient: 15 of 21 label-ranking datasets improved with the soft rank layer, 4 were similar, and 2 were worse.The experiment used Spearman’s rank correlation coefficient and compared models with and without the soft ranking layer.
- Robust regression via soft least trimmed squares: Soft least trimmed squares interpolates between least trimmed squares as ε →0 and least squares as ε →∞.This interpolation was empirically validated, and ε can provide adaptivity to the unknown outlier percentage.
- Robust regression via soft least trimmed squares: On all datasets, ridge-regression accuracy deteriorated significantly as outliers increased, while trimmed-squares methods varied relative to Huber loss across datasets.Hard least trimmed squares worked well across the evaluated datasets, with regularization useful in some cases.
7. Conclusion
The paper constructs differentiable sorting and ranking operators from permutahedron projections, achieving efficient exact computation and differentiation and demonstrating applications in ranking and robust regression.
- Conclusion: The operators provide exact O(n log n) computation and O(n) differentiation as drop-in replacements for existing O(n^2) operators.The paper reports an order-of-magnitude speed-up.
- Conclusion: The paper applies the soft operators to differentiable Spearman’s rank correlation and robust regression through soft least trimmed squares.These applications use the proposed operators for label ranking and outlier-robust regression.
A. Additional empirical results
The label-ranking experiment evaluates differentiable ranking across 21 datasets, with results summarized by Spearman rank correlation and detailed comparisons against no projection.
- 21 datasets are included in the detailed label-ranking evaluation.
- Spearman’s rank correlation coefficient is averaged over 5 runs for each method.
- Table 1 marks methods with better or worse Spearman rank correlation than using no projection.
B. Proofs
The proofs establish claims about the operators by analyzing their behavior for ordered inputs and relating permutations to permutahedron vertices.
- The first claim is considered for vectors w satisfying w1 > w2 > · · · > wn.
- Setting w = ρ and using the stated permutation identities proves the claim.
- The proof uses the fact that the permutahedron equals the convex hull of permutations of its defining vector.
B.3. Proof of Proposition 2 (Properties of soft sorting and ranking operators)
The proof characterizes differentiability, order preservation, and regularization limits for soft sorting and ranking through regularized projections and isotonic optimization.
- Strong convexity makes the projection solution Lipschitz continuous and differentiable almost everywhere under twice-differentiable regularization.
- Soft sorting follows the descending order of ρ, while soft ranking follows the ordering of −θ.
- The regularization strength ε controls the trade-off between approximating the original operator and increasing smoothness.
- For sufficiently small ε, the vector s/ε − w is already feasible and gives the optimal isotonic solution.
- For ε > εmax(s, w), the PAV algorithm merges violating neighboring blocks until the full index set is optimal.
B.4. Proof of Proposition 3 (Reduction to isotonic optimization)
The proof reduces regularized projections onto permutahedra to isotonic optimization through a dual formulation and separable convex objectives.
- A strongly convex regularized linear program has a dual minimizing Ψ*(z − u) plus the support function of the constraint set.
- For cardinality-based submodular base polytopes, the dual variables can be represented in descending order and reduced to isotonic optimization.
- The permutahedron is a special case of a base polytope, with its ordered defining vector determining the corresponding greedy representation.
- Quadratic regularization yields a squared-distance term between the sorted input and isotonic variables, with the entropic case treated similarly.
- When the regularizer’s conjugate is separable, the resulting isotonic problem can be solved using a generalized PAV algorithm.
- The exposition focuses on quadratic and entropic regularization, although other regularizers may also be possible.
B.7. Proof of Proposition 4 (Jacobian of isotonic optimization)
The Jacobian of isotonic optimization has a block structure induced by the optimal partition, with derivatives determined blockwise under quadratic or entropic regularization. At nondifferentiable points, the paper uses Clarke generalized Jacobians.
- Block structure: The optimal isotonic solution partitions [n] into blocks, and each coordinate shares the value associated with its unique block.This block value is written as γΨ(Bl; s, w) for coordinate i in block Bl.
- Jacobian structure: The Jacobian is block diagonal because derivatives are nonzero only when input and output indices belong to the same block.Within each block, the nonzero partial derivatives form a block matrix whose columns correspond to input coordinates.
- Quadratic regularization: For quadratic regularization, each block’s derivative entries are constant within columns and equal 1/|Bl| for indices in the same block.Entries outside the block are zero.
- Entropic regularization: For entropic regularization, within-block derivatives are given by softmax(sBl)j, while derivatives outside the block are zero.The softmax expression applies when both indices i and j belong to block Bl.
- Nondifferentiability: At points where an optimal block can split into equal-valued blocks, the solution is nondifferentiable and an arbitrary Clarke generalized Jacobian may be used.The generalized Jacobian is the convex hull of limiting Jacobians; directional derivatives may disagree at these points.