Differentiable Ranks and Sorting using Optimal Transport

Introduction

On (not) learning with sorting and ranking. Operators $R$ and $S$ play an important role across statistics and machine learning. For instance, $R$ is the main workhorse behind order statistics , but also appears prominently in $k$-NN rules, in which $R$ is applied on a vector of distances to select the closest neighbors to a query point. Ranking is also used to assess the performance of an algorithm: either at test time, such as $0/1$ and top-$k$ classification accuracies and NDCG metrics when learning-to-rank , or at train time, by selecting pairs and triplets of points of interest. The sorting operator $S$ is of no less importance, and can be used to handle outliers in robust statistics, as in trimmed and least-quantile regression or median-of-means estimators . Yet, and although examples of using $R$ and $S$ abound in ML, neither $R$ nor $S$ are actively used in end-to-end learning approaches: while $S$ is not differentiable everywhere, $R$ is outright pathological, since it is piecewise constant and has therefore a Jacobian $\partial R/\partial\mathbf{x}$ that is almost everywhere zero.

Outline. We recall first the link between the $R$ and $S$ operators and OT between 1D measures, to define then generalized Kantorovich rank and sort operators in §2. We turn them into differentiable operators using entropic regularization, and discuss in §3 the several parameters that can shape this smoothness. Using these smooth operators, we propose in §4 alternatives to cross-entropy and least-quantile losses to learn classifiers and regression functions.

Ranking and Sorting as an Optimal Transport Problem

The fact that solving the OT problem between two discrete univariate measures boils down to sorting is well known [33, §2]. The usual narrative states that the Wasserstein distance between two univariate measures reduces to comparing their quantile functions, which can be obtained by inverting CDFs, which are themselves computed by considering the sorted values of the supports of these measures. This downstream connection from OT to quantiles, CDFs and finally sorting has been exploited in several works, notably because the $n\log n$ price for sorting is far cheaper than the order $n^{3}\log n$ one has to pay to solve generic OT problems. This is evidenced by the recent surge in interest for sliced Wasserstein distances . We propose in this section to go instead upstream, that is to redefine ranking and sorting functions as byproducts of the resolution of an optimal assignment problem between measures supported on the reals. We then propose in Def.1 generalized rank and sort operators using the Kantorovich formulation of OT.

We make in what follows the additional assumption that $h$ is convex. A fundamental result [33, Theorem 2.9] states that in that case (see also for the more involved case where $h$ is concave) $\operatorname{OT}_{h}(\xi,\upsilon)$ can be computed in closed form using the quantile functions $Q_{\xi},Q_{\upsilon}$ of $\xi,\upsilon$:

Therefore, to compute OT between $\xi$ and $\upsilon$, one only needs to integrate the difference in their quantile functions, which can be done by inverting the empirical distribution functions for $\xi,\upsilon$, which itself only requires sorting the entries in $\mathbf{x}$ and $\mathbf{y}$ to obtain their sorting permutations $\sigma$ and $\tau$. Additionally, Eq. (2) allows us not only to recover the value of $\operatorname{OT}_{h}$ as defined in Eq. (1), but it can also be used to recover the corresponding optimal solution $P_{\star}$ in $n+m$ operations, using the permutations $\sigma$ and $\tau$ to build a so-called north-west corner solution [28, §3.4.2]:

Let $\sigma$ and $\tau$ be sorting permutations for $\mathbf{x}$ and $\mathbf{y}$. Define $N$ to be the north-west corner solution using permuted weights $\mathbf{a}_{\sigma},\mathbf{b}_{\tau}$. Then $N_{\sigma^{-1},\tau^{-1}}$ is optimal for (1).

Such a permuted north-western corner solution is illustrated in Figure 1(b). It is indeed easy to check that in that case $(P_{\star})_{\sigma,\tau}$ runs from the top-left (north-west) to the bottom right corner. In the simple case where $n=m$ and $\mathbf{a}=\mathbf{b}=\mathbf{1}_{n}/n$, the solution $N_{\sigma^{-1},\tau^{-1}}$ is a permutation matrix divided by $n$, namely a matrix equal to everywhere except for its entries indexed by $(i,\tau\circ\sigma^{-1})_{i}$ which are all equal to $1/n$. That solution is a vertex of the Birkhoff polytope, namely, an optimal assignment which to the $i$-th value in $\mathbf{x}$ associates the $(\tau\circ\sigma^{-1})_{i}$-th value in $\mathbf{y}$; informally, this solution assigns the $i$-th smallest entry in $\mathbf{x}$ to the $i$-th smallest entry in $\mathbf{y}$.

Generalizing sorting, CDFs and quantiles using optimal transport. From now on in this paper, we make the crucial assumption that $\mathbf{y}$ is already sorted, that is, $y_{1}<\dots<y_{m}$. $\tau$ is therefore the identity permutation. When in addition $n=m$, the $i$-th value in $\mathbf{x}$ is simply assigned to the $\sigma_{i}^{-1}$-th value in $\mathbf{y}$. Conversely, and as illustrated in Figure 1(a), the rank $i$ value in $\mathbf{x}$ is assigned to the $i$-th value $y_{i}$. Because of this, $R$ and $S$ can be rewritten using the optimal assignment matrix $P_{\star}$:

These identities stem from the fact that $nP_{\star}$ is a permutation matrix, which can be applied to the vector $n\overline{\mathbf{b}}=(1,\dots,n)$ to recover the rank of each entry in $\mathbf{x}$, or transposed and applied to $\mathbf{x}$ to recover the sorted values of $\mathbf{x}$. The former expression can be equivalently interpreted as $n$ times the CDF of $\xi$ evaluated elementwise to $\mathbf{x}$, the latter as the quantiles of $\xi$ at levels $\overline{\mathbf{b}}$. The identities in Prop. 2 are valid when the input measures $\xi,\upsilon$ are uniform and of the same size. The first contribution of this paper is to consider more general scenarios, in which $m$, the size of $\mathbf{y}$, can be smaller than $n$, and where weights $\mathbf{a},\mathbf{b}$ need not be uniform. This is a major departure from previous references , which all require pairwise comparisons between the entries in $\mathbf{x}$. We show in our applications that $m$ can be as small as 3 when trying to recover a quantile, as in Figs. 1, 3.

The K-ranking operator computes convex combinations of rank values (as described in the entries $n\overline{\mathbf{b}}$) while the K-sorting operator computes convex combinations of values contained in $\mathbf{x}$ directly. Note that we consider here convex combinations (weighted averages) of these ranks/values, according to the Euclidean geometry. Extending more generally these combinations to Fréchet means using alternative geometries (KL, hyperbolic, etc) on these ranks/values is left for future work. Because these quantities are only defined pointwisely (we output vectors and not functions) and depend on the ordering of $\mathbf{a},\mathbf{x},\mathbf{b},\mathbf{y}$, we drop our reference to measure $\xi$ in notations.

The K-rank vector map $\widetilde{R}$ outputs a vector of size $n$ containing a continuous rank for each entry for $\mathbf{x}$ (these entries can be alternatively interpreted as $n$ times a “synthetic” CDF value in $$, itself a convex mixture of the CDF values$\overline{\mathbf{b}}_{j}$of the$y_{j}$onto which each$x_{i}$is transported).$\widetilde{S}$is a split-quantile operator outputting$m$increasing values which are each, respectively, barycenters of some of the entries in$\mathbf{x}$. The fact that these values are increasing can be obtained by a simple argument in which$\xi$and$\upsilon$are cast again as uniform measures of the same size using duplicated supports$x_{i}$and$y_{j}$, and then use the monotonicity given by the third identity of Prop. 2.

Computations and Non-differentiability The generalized ranking and sorting operators presented in Def. 1 are interesting in their own right, but have very little practical appeal. For one, their computation relies on solving an OT problem at a cost of $O(nm(n+m)\log(nm))$ and remains therefore far more costly than regular sorting, even when $m$ is very small. Furthermore, these operators remain fundamentally not differentiable. This can be hinted by the simple fact that it is difficult to guarantee in general that a solution $P_{\star}$ to (1) is unique. Most importantly, the Jacobian $\partial P_{\star}/\partial\mathbf{x}$ is, very much like $R$, null almost everywhere. This can be visualized by looking at Figure 1(b) to notice that an infinitesimal change in $\mathbf{x}$ would not change $P_{\star}$ (notice however that an infinitesimal change in weights $\mathbf{a}$ would; that Jacobian would involve North-west corner type mass transfers). All of these pathologies — computational cost, non-uniqueness of optimal solution and non-differentiability — can be avoided by using regularized OT .

The Sinkhorn Ranking and Sorting Operators

Both K-rank $\widetilde{R}$ and K-sort $\widetilde{S}$ operators are expressed using the optimal solution $P_{\star}$ to the linear program in (1). However, $P_{\star}$ is not differentiable w.r.t inputs $\mathbf{a},\mathbf{x}$ nor parameters $\mathbf{b},\mathbf{y}$ [2, §5]. We propose instead to rely on a differentiable variant of the OT problem that uses entropic regularization , as detailed in [28, §4]. This differentiability is reflected in the fact that the optimal regularized transport plan is a dense matrix (yielding more arrows in Fig. 1(c)), which ensures differentiability everywhere w.r.t. both $\mathbf{a}$ and $\mathbf{x}$.

Consider first a regularization strength $\varepsilon>0$ to define the solution to the regularized OT problem:

Given a regularization strength $\varepsilon>0$, run Alg.1 to define

Sensitivity to $\varepsilon$. Parameter $\varepsilon$ plays the same role as other temperature parameters in previously proposed smoothed sorting operators : the smaller $\varepsilon$ is, the closer the Sinkhorn operator’s output is to the original vectors of ranks and sorted values; The bigger $\varepsilon$, the closer $P_{\star}^{\varepsilon}$ to matrix $\mathbf{a}\mathbf{b}^{T}$, and therefore all entries of $\widetilde{R}_{\varepsilon}$ collapse to the average of $n\bar{\mathbf{b}}$, while all entries of $\widetilde{S}_{\varepsilon}$ collapse to the weigted average (using $\mathbf{a}$) of $\mathbf{x}$, as illustrated in Fig. 2. Although choosing a small value for $\varepsilon$ might seem natural, in the sense that $\widetilde{R}_{\varepsilon},\widetilde{S}_{\varepsilon}$ approximate more faithfully $R,S$, one should not forget that this would result in recovering the deficiencies of $R,S$ in terms of differentiability. When learning with such operators, it may therefore be desirable to use a value for $\varepsilon$ that is large enough to ensure $\partial P^{\varepsilon}_{\star}/\partial\mathbf{x}$ has non-null entries. We usually set $\varepsilon=10^{-2}$ or $10^{-3}$ when $\mathbf{x},\mathbf{y}$ lie in as in Fig. 2. We have kept $\varepsilon$ fixed throughout Alg. 1, but we do notice some speedups using scheduling as advocated by .

Numerical Stability. When using small regularization strengths, we recommend to cast Sinkhorn iterations in the log-domain by considering the following stabilized iterations for each pair of vectors $\mathbf{x}_{s},\mathbf{y}$, resulting in the following updates (with $\boldsymbol{\alpha}$ and $\boldsymbol{\beta}$ initialized to $\mathbf{0}_{n}$ and $\mathbf{0}_{m}$),

where $\min_{\varepsilon}$ is the soft-minimum operator applied linewise to a matrix to output a vector, namely for $M\in\mathbf{R}^{n\times m}$, $\min_{\varepsilon}(M)\in\mathbf{R}^{n}$ and is such that $[\min_{\varepsilon}(M)]_{i}=-\varepsilon(\log\sum_{j}e^{-M_{ij}/\varepsilon})$. The rationale behind the substractions/additions of $\boldsymbol{\alpha}$ and $\boldsymbol{\beta}$ above is that once a Sinkhorn iteration is carried out, the terms inside the parenthesis above are normalized, in the sense that once divided by $\varepsilon$, their exponentials sum to one (they can be used to recover a coupling). Therefore, they must be negative, which improves the stability of summing exponentials [28, §4.4].

The choices that we have made are summarized in Alg. 2, but we believe there are opportunities to perfect them depending on the task.

Learning with Smoothed Ranks and Sorts

We propose a differentiable variant of the 0/1 and more generally top $k$ losses that bypasses combinatorial consideration nor builds upon non-differentiable surrogates . Ignoring the degenerate case in which $l^{\star}$ is not unique, given a query $\omega$, stating that the the label $l^{\star}$ has been selected is equivalent to stating that the entry indexed at $l^{\star}$ of the vector of ranks $R(f_{\theta}(\omega))$ is $L$. Given a labelled pair $(\omega,l)$, the 0/1 loss of the classifier for that pair is therefore,

where $H$ is the heaviside function: $H(u)=1$ if $u>0$ and $H(u)=0$ for $u\leq 0$. More generally, if for some labelled input $\omega$, the entry $[R(f_{\theta})]_{l_{o}}$ is bigger than $L-k+1$, then that labelled example has a top-$k$ error of . Conversely, if $[R(f_{\theta})]_{l}$ is smaller than $L-k+1$, then the top-$k$ error is $1$. The top-$k$ error can be therefore formulated as in (6), where the argument $L-\left[R(f_{\theta}(\omega)\right]_{l}$ within the Heaviside function is replaced by $L-\left[R(f_{\theta}(\omega)\right]_{l}-k+1$.

We propose the simple family of ReLU losses $J_{k}(u)=\max(0,u-k+1)$, and have focused our experiments on the case $k=1$. We train a vanilla CNN (4 Conv2D with 2 max-pooling layers, ReLU activation, 2 fully connected layers, batchnorm on each) and a Resnet18 on CIFAR-10 and CIFAR-100. Fig. 4 and 5 report test-set classification accuracies / epochs. We used $\varepsilon=10^{-3}$, $\eta=10^{-3}$, a squared distance cost $h(u)=u^{2}$ and a stepsize of $10^{-4}$ with the ADAM optimizer.

Learning CNNs by sorting handwritten numbers. We use the MNIST experiment setup in , in which a CNN is given $n$ numbers between between 0 and 9999 given as 4 concatenated MNIST images. The labels are the ranks (within $n$ pairs) of each of these $n$ numbers. We use the code kindly made available by the authors. We use $100$ epochs, and confirm experimentally that S-sort performs on par with their neural-sort function. We set $\varepsilon=0.005$.

For each quantile/dataset pair, we report the original (not-regularized) $\tau$ quantile of the errors evaluated on the entire training set, on an entire held-out test set, and the MSE on the test set of the function that is recovered. We notice that our algorithm reaches overall better quantile errors on the training set—this is our main goal—but comparable test/MSE errors.

Conclusion. We have proposed in this paper differentiable proxies to the ranking and sorting operations. These proxies build upon the existing connection between sorting and the computation of OT in 1D. By generalizing sorting using OT, and then introducing a regularized form that can be solved using Sinkhorn iterations, we recover the simple benefit that all of its steps can be easily automatically differentiated. We have shown that, with a focus on numerical stability, one can use there operators in various settings, including smooth extensions of test-time metrics that rely on sort, and which can be now used as training losses. For instance, we have used the Sinkhorn sort operator to provide a smooth approximation of quantiles to solve least-quantile regression problems, and the Sinkhorn rank operator to formulate an alternative to the cross-entropy that can mimic the $0/1$ loss in multiclass classification. This smooth approximation to the rank, and the resulting gradient flow that we obtain is strongly reminiscent of rank based dynamics, in which players in a given game produce an effort (a gradient) that is a direct function of their rank (or standing) within the game, as introduced by . Our use of the Sinkhorn algorithm can therefore be interpreted as a smooth mechanism to enact such dynamics. Several open questions remain: although the choice of a cost function $h$, target vector $\mathbf{y}$ and squashing function $g$ (used to form vector $\widetilde{x}$ in Alg. 1, using Eq. 4) have in principle no influence on the vector of Sinkhorn ranks or sorted values in the limit when $\varepsilon$ goes to (they all converge to $R$ and $S$), these choices strongly shape the differentiability of $\widetilde{R}_{\varepsilon}$ and $\widetilde{S}_{\varepsilon}$ when $\varepsilon>0$. Our empirical findings suggest that whitening and squashing all entries within $$is crucial to obtain stable numerically, but more generally to retain consistent gradients across iterations, without having to re-define$\varepsilon$ at each iteration.

References