---
abstract: 'While sorting is an important procedure in computer science, the `argsort` operator - which takes as input a vector and returns its sorting permutation - has a discrete image and thus zero gradients almost everywhere. This prohibits end-to-end, gradient-based learning of models that rely on the `argsort` operator. A natural way to overcome this problem is to replace the `argsort` operator with a continuous relaxation. Recent work has shown a number of ways to do this, but the relaxations proposed so far are computationally complex. In this work we propose a simple continuous relaxation for the `argsort` operator which has the following qualities: it can be implemented in three lines of code, achieves state-of-the-art performance, is easy to reason about mathematically - substantially simplifying proofs - and is faster than competing approaches. We open source the code to reproduce all of the experiments and results.'
bibliography:
- 'references.bib'
---

Introduction
============

(b1) [$v_1$]{}; (b2) \[below=of b1\] [$v_2$]{}; (b3) \[below=of b2\] [$v_3$]{};

;

; ; ;

(bb2) \[right=2cm of b1\] [$v_2$]{}; (bb3) \[right=2cm of b2\] [$v_3$]{}; (bb1) \[right=2cm of b3\] [$v_1$]{};

(b1.east) to\[out=0,in=180\] (bb1.west); (b2.east) to\[out=0,in=180\] (bb2.west); (b3.east) to\[out=0,in=180\] (bb3.west);

(v1) \[right=6cm of b1\][$v_1$]{}; (v2) \[below=of v1\] [$v_2$]{}; (v3) \[below=of v2\] [$v_3$]{};

;

; ; ;

(vv11) \[right=2cm of v1\] [$v_1$]{}; ; ; (vv12) \[right=4cm of v1\] [$v_2$]{}; ; ; (vv13) \[right=6cm of v1\] [$v_3$]{}; ; ;

(vv21) \[right=2cm of v2\] [$v_1$]{}; ; ; (vv22) \[right=4cm of v2\] [$v_2$]{}; ; ; (vv23) \[right=6cm of v2\] [$v_3$]{}; ; ;

(vv31) \[right=2cm of v3\] [$v_1$]{}; ; ; (vv32) \[right=4cm of v3\] [$v_2$]{}; ; ; (vv33) \[right=6cm of v3\] [$v_3$]{}; ; ;

(v1.east) to\[out=-30,in=150\] (vv31.west); (v2.east) to\[out=45,in=-135\] (vv12.south); (v3.east) to\[out=45,in=-135\] (vv23.south);

-0.2in

Gradient-based optimization lies at the core of the success of deep learning. However, many common operators have discrete images and thus exhibit zero gradients almost everywhere, making them unsuitable for gradient-based optimization. Examples include the *Heaviside step* function, the `argmax` operator, and - the central object of this work - the `argsort` operator. To overcome this limitation, continuous relaxations for these operators can be constructed. For example, the sigmoid function serves as a continuous relaxation for the *Heaviside step* function, and the `softmax` operator serves as a continuous relaxation for the `argmax`. These continuous relaxations have the crucial property that they provide meaningful gradients that can drive optimization. Because of this, operators such as the `softmax` are ubiquitous in deep learning. In this work we are concerned with continuous relaxations for the `argsort` operator.

Formally, we define the `argsort` operator as the mapping $\texttt{argsort} : \mathbb R^n \rightarrow \mathcal S_n$ from $n$-dimensional real vectors $s \in \mathbb R^n$ to the set of permutations over $n$ elements $\mathcal S_n \subseteq \{1,2,\dots,n\}^n$, where $\texttt{argsort}(s)$ returns the permutation that sorts $s$ in decreasing order[^1]. For example, for the input vector $s = [9,1,5,2]^T$ we have $\texttt{argsort}(s) = [1, 3, 4, 2]^T$. If we let $\mathcal P_n \subseteq \{0, 1\}^{n \times n} \subset \mathbb R^{n \times n}$ denote the set of permutation matrices of dimension $n$, following [@2018_grover] we can define, for a permutation $\pi \in \mathcal S_n$, the permutation matrix $P_\pi \in \mathcal P_n$ as: $$P_\pi[i, j] =
\begin{cases}
1\text{ if }j = \pi_i,\\
0\text{ otherwise}
\end{cases}$$ This is simply the one-hot representation of $\pi$. Note that with these definitions, the mapping $\texttt{sort} : \mathbb R^n \rightarrow \mathbb R^n$ that sorts $s$ in decreasing order is $\texttt{sort}(s) = P_{\texttt{argsort}(s)}s$. Also, if we let $\bar{1}_n = [1,2,\dots,n]^T$, then the `argsort` operator can be recovered from $P_{\texttt{argsort}(\cdot)} : \mathbb R^n \rightarrow \mathcal P_n$ by a simple matrix multiplication via $$\texttt{argsort}(s) = P_{\texttt{argsort}(s)}\bar{1}_n$$ Because of this reduction from the `argsort` operator to the $P_{\texttt{argsort}(\cdot)}$ operator, in this work, as in previous works [@2018_mena; @2018_grover; @2019_cuturi], our strategy to derive a continuous relaxation for the `argsort` operator is to instead derive a continuous relaxation for the $P_{\texttt{argsort}(\cdot)}$ operator. This is analogous to the way that the `softmax` operator relaxes the `argmax` operator.

The main contribution of this paper is the proposal of a family of simple continuous relaxation for the $P_{\texttt{argsort}(\cdot)}$ operator, which we call $\SimpleSort{}$, and define as follows: $$\label{eq:simplesort}
\SimpleSort{}^d_{\tau}(s) = \texttt{softmax}\left(\frac{-d\left(\texttt{sort}(s) \mathds 1^T, \mathds 1 s^T\right)}{\tau}\right)$$ where the softmax operator is applied row-wise, $d$ is any differentiable almost everywhere, semi–metric function of $\mathbb R$ applied pointwise, and $\tau$ is a temperature parameter that controls the degree of the approximation. In simple words: *the $r$-th row of the operator is the `softmax` of the negative distances to the $r$-th largest element*.

Throughout this work we will predominantly use $d(x, y) = |x - y|$ (the $L_1$ norm), but our theoretical results hold for any such $d$, making the approach flexible. The operator is trivial to implement in automatic differentiation libraries such as TensorFlow [@2016_abadi] and PyTorch [@2017_paszke], and we show that:

-   achieves state-of-the-art performance on multiple benchmarks that involve reordering elements.

-   shares the same desirable properties as the `NeuralSort` operator [@2018_grover]. Namely, it is row-stochastic, converges to $P_{\texttt{argsort}(\cdot)}$ in the limit of the temperature, and can be projected onto a permutation matrix using the row-wise `argmax` operation. However, is significantly easier to reason about mathematically, which substantially simplifies the proof of these properties.

-   The operator is faster than the `NeuralSort` operator of [@2018_grover], the fastest competing approach, and empirically just as easy to optimize in terms of the number of gradient steps required for the training objective to converge.

Therefore, the operator advances the state of the art in differentiable sorting by significantly simplifying previous approaches. To better illustrate the usefulness of the mapping defined by , we show in Figure \[fig:simplesort\_plots\] the result of applying the operator to soft-sort a sequence of vectors $v_i$ $(1\leq i \leq n)$ according to the order given by respective scores $s_i \in \mathbb R$. Soft-sorting the $v_i$ is achieved by multiplying to the left by $\SimpleSort(s)$.

The code and experiments are available at <https://github.com/sprillo/softsort>

Related Work
============

Relaxed operators for sorting procedures were first proposed in the context of Learning to Rank with the end goal of enabling direct optimization of Information Retrieval (IR) metrics. Many IR metrics, such as the Normalized Discounted Cumulative Gain (NDCG) [@2002_kalervo], are defined in terms of *ranks*. Formally, the `rank` operator is defined as the function $\texttt{rank} : \mathbb R^n \rightarrow \mathcal S_n$ that returns the inverse of the `argsort` permutation: $\texttt{rank}(s) = \texttt{argsort}(s)^{-1}$, or equivalently $P_{\texttt{rank}(s)} = P_{\texttt{argsort}(s)}^T$. The $\texttt{rank}$ operator is thus intimately related to the $\texttt{argsort}$ operator; in fact, a relaxation for the $P_{\texttt{rank}(\cdot)}$ operator can be obtained by transposing a relaxation for the $P_{\texttt{argsort}(\cdot)}$ operator, and vice-versa; this duality is apparent in the work of [@2019_cuturi].

We begin by discussing previous work on relaxed `rank` operators in section \[sec:relaxed\_rank\_operators\]. Next, we discuss more recent work, which deals with relaxations for the $P_{\texttt{argsort}(\cdot)}$ operator.

Relaxed Rank Operators {#sec:relaxed_rank_operators}
----------------------

The first work to propose a relaxed `rank` operator is that of [@2008_taylor]. The authors introduce the relaxation $\texttt{SoftRank}_\tau : \mathbb R^n \rightarrow \mathbb R^n$ given by $\texttt{SoftRank}_\tau(s) = \mathbb E[\texttt{rank}(s + z)]$ where $z \sim \mathcal N_n(0, \tau I_n)$, and show that this relaxation, as well as its gradients, can be computed exactly in time $\mathcal O(n^3)$. Note that as $\tau \rightarrow 0$, $\texttt{SoftRank}_\tau(s) \rightarrow \texttt{rank}(s)$[^2]. This relaxation is used in turn to define a surrogate for NDCG which can be optimized directly.

In [@2010_qin], another relaxation for the `rank` operator $\widehat{\pi}_\tau : \mathbb R^n \rightarrow \mathbb R^n$ is proposed, defined as: $$\label{eqn:2010_qin}
\widehat{\pi}_\tau(s)[i] = 1 + \sum_{j \neq i} \sigma\left(\frac{s_i - s_j}{\tau}\right)$$ where $\sigma(x) = (1 + \exp\{-x\})^{-1}$ is the sigmoid function. Again, $\widehat{\pi}_\tau(s) \rightarrow \texttt{rank}(s)$ as $\tau \rightarrow 0$. This operator can be computed in time $\mathcal{O}(n^2)$, which is faster than the $\mathcal O(n^3)$ approach of [@2008_taylor]. Note that the above two approaches cannot be used to define a relaxation for the `argsort` operator. Indeed, $\texttt{SoftRank}_\tau(s)$ and $\widehat{\pi}_\tau(s)$ are not relaxations for the $P_{\texttt{rank}(\cdot)}$ operator. Instead, they directly relax the $\texttt{rank}$ operator, and there is no easy way to obtain a relaxed `argsort` or $P_{\texttt{argsort}(\cdot)}$ operator from them.

Sorting via Bipartite Maximum Matchings
---------------------------------------

The work of [@2018_mena] draws an analogy between the `argmax` operator and the `argsort` operator by means of bipartite maximum matchings: the `argmax` operator applied to an $n$-dimensional vector $s$ can be viewed as a maximum matching on a bipartite graph with $n$ vertices in one component and $1$ vertex in the other component, the edge weights equal to the given vector $s$; similarly, a permutation matrix can be seen as a maximum matching on a bipartite graph between two groups of $n$ vertices with edge weights given by a matrix $X \in \mathbb R^{n \times n}$. This induces a mapping $M$ (for ‘matching’) from the set of square matrices $X \in \mathbb R^{n \times n}$ to the set $\mathcal P_n$. Note that this mapping has arity $M : \mathbb R^{n \times n} \rightarrow \mathcal P_n$, unlike the $P_{\texttt{argsort}(\cdot)}$ operator which has arity $P_{\texttt{argsort}(\cdot)} : \mathbb R^n \rightarrow \mathcal P_n$.

Like the $P_{\texttt{argsort}(\cdot)}$ operator, the $M$ operator has discrete image $\mathcal P_n$, so to enable end-to-end gradient-based optimization, [@2018_mena] propose to relax the matching operator $M(X)$ by means of the Sinkhorn operator $S(X / \tau)$; $\tau$ is a temperature parameter that controls the degree of the approximation; as $\tau \rightarrow 0$ they show that $S(X / \tau) \rightarrow M(X)$. The Sinkhorn operator $S$ maps the square matrix $X / \tau$ to the Birkhoff polytope $\mathcal B_n$, which is defined as the set of doubly stochastic matrices (i.e. rows and columns summing to $1$).

The computational complexity of the [@2018_mena] approach to differentiable sorting is thus $\mathcal O(Ln^2)$ where $L$ is the number of Sinkhorn iterates used to approximate $S(X / \tau)$; the authors use $L = 20$ for all experiments.

Sorting via Optimal Transport
-----------------------------

The recent work of [@2019_cuturi] also makes use of the Sinkhorn operator to derive a continuous relaxation for the $P_{\texttt{argsort}(\cdot)}$ operator. This time, the authors are motivated by the observation that a sorting permutation for $s \in \mathbb R^n$ can be recovered from an optimal transport plan between two discrete measures defined on the real line, one of which is supported on the elements of $s$ and the other of which is supported on arbitrary values $y_1 < \dots < y_n$. Indeed, the optimal transport plan between the probability measures $\frac{1}{n} \sum_{i = 1}^n \delta_{s_i}$ and $\frac{1}{n} \sum_{i = 1}^n \delta_{y_i}$ (where $\delta_x$ is the Dirac delta at $x$) is given by the matrix $P_{\texttt{argsort}(s)}^T$.

Notably, a variant of the optimal transport problem with entropy regularization yields instead a continuous relaxation $P^\epsilon_{\texttt{argsort}(s)}$ mapping $s$ to the Birkhoff polytope $\mathcal B_n$; $\epsilon$ plays a role similar to the temperature in [@2018_mena], with $P^\epsilon_{\texttt{argsort}(s)} \rightarrow P_{\texttt{argsort}(s)}$ as $\epsilon \rightarrow 0$. This relaxation can be computed via Sinkhorn iterates, and enables the authors to relax $P_{\texttt{argsort}(\cdot)}$ by means of $P^\epsilon_{\texttt{argsort}(s)}$. Gradients can be computed by backpropagating through the Sinkhorn operator as in [@2018_mena].

The computational complexity of this approach is again $\mathcal O(Ln^2)$. However, the authors show that a generalization of their method can be used to compute relaxed quantiles in time $\mathcal O(Ln)$, which is interesting in its own right.

Sorting via a sum-of-top-k elements identity
--------------------------------------------

Finally, a more computationally efficient approach to differentiable sorting is proposed in [@2018_grover]. The authors rely on an identity that expresses the sum of the top $k$ elements of a vector $s \in \mathbb R^n$ as a symmetric function of $s_1,\dots,s_n$ that only involves `max` and `min` operations [@2003_ogryczak Lemma 1]. Based on this identity, and denoting by $A_s$ the matrix of *absolute* pairwise differences of elements of $s$, namely $A_s[i, j] = |s_i - s_j|$, the authors prove the identity: $$\label{eq:2018_grover_eq4}
P_{\texttt{argsort}(s)}[i, j] =
\begin{cases}
1\text{ if }j = \texttt{argmax}(c_i),\\
0\text{ otherwise}
\end{cases}$$ where $c_i = (n + 1 - 2i)s - A_s \mathds 1$, and $\mathds 1$ denotes the column vector of all ones.

Therefore, by replacing the $\texttt{argmax}$ operator in Eq. \[eq:2018\_grover\_eq4\] by a row-wise `softmax`, the authors arrive at the following continuous relaxation for the $P_{\texttt{argsort}(\cdot)}$ operator, which they call `NeuralSort`: $$\label{eq:2018_grover_eq5}
\texttt{NeuralSort}_\tau(s)[i, :] = \texttt{softmax}\left(\frac{c_i}{\tau}\right)$$ Again, $\tau > 0$ is a temperature parameter that controls the degree of the approximation; as $\tau \rightarrow 0$ they show that $\texttt{NeuralSort}_\tau(s) \rightarrow P_{\texttt{argsort}(s)}$. Notably, the relaxation proposed by [@2018_grover] can be computed in time $\mathcal O(n^2)$, making it much faster than the competing approaches of [@2018_mena; @2019_cuturi].

: A simple relaxed sorting operator
===================================

$$\texttt{NeuralSort}_\tau(s) = g_\tau
\begin{pmatrix}
0                        & s_2 - s_1         & 3s_3 - s_1 - 2s_2 & 5s_4 - s_1 - 2s_2 - 2s_3 \\
s_2 - s_1                & 0                 & s_3 - s_2         & 3s_4 - s_2 - 2s_3 \\
2s_2 + s_3 - 3s_1        & s_3 - s_2         & 0                 & s_4 - s_3 \\
2s_2 + 2s_3 + s_4 - 5s_1 & 2s_3 + s_4 - 3s_2 & s_4 - s_3         & 0 \\
\end{pmatrix}$$

-0.2in

$$\SimpleSort{}^{|\cdot|}_{\tau}(s) = g_\tau
\begin{pmatrix}
0         & s_2 - s_1 & s_3 - s_1 & s_4 - s_1 \\
s_2 - s_1 & 0         & s_3 - s_2 & s_4 - s_2 \\
s_3 - s_1 & s_3 - s_2 & 0         & s_4 - s_3 \\
s_4 - s_1 & s_4 - s_2 & s_4 - s_3 & 0         \\
\end{pmatrix}$$

-0.2in

In this paper we propose , a simple continuous relaxation for the $P_{\texttt{argsort}(\cdot)}$ operator. We define as follows: $$\label{eq:simplesort2}
\SimpleSort{}^d_{\tau}(s) = \texttt{softmax}\left(\frac{-d\left(\texttt{sort}(s) \mathds 1^T, \mathds 1 s^T\right)}{\tau}\right)$$ where $\tau > 0$ is a temperature parameter that controls the degree of the approximation and $d$ is semi–metric function applied pointwise that is differentiable almost everywhere. Recall that a semi–metric has all the properties of a metric except the triangle inequality, which is not required. Examples of semi–metrics in $\mathbb R$ include any positive power of the absolute value. The operator has similar desirable properties to those of the `NeuralSort` operator, while being significantly simpler. Here we state and prove these properties. We start with the definition of *Unimodal Row Stochastic Matrices* [@2018_grover], which summarizes the properties of our relaxed operator:

(Unimodal Row Stochastic Matrices). An $n \times n$ matrix is Unimodal Row Stochastic (URS) if it satisfies the following conditions:

1.  **Non-negativity:** $U[i, j] \ge 0\quad \forall i,j \in \{1,2,\dots,n\}$.

2.  **Row Affinity:** $\sum_{j = 1}^n U[i, j] = 1\quad \forall i \in \{1,2,\dots,n\}$.

3.  **Argmax Permutation:** Let $u$ denote a vector of size $n$ such that $u_i = \arg\max_j U[i, j]\quad \forall i \in \{1,2,\dots,n\}$. Then, $u \in \mathcal S_n$, i.e., it is a valid permutation.

While `NeuralSort` and `SoftSort` yield URS matrices (we will prove this shortly), the approaches of [@2018_mena; @2019_cuturi] yield bistochastic matrices. It is natural to ask whether URS matrices should be preferred over bistochastic matrices for relaxing the $P_{\texttt{argsort}(\cdot)}$ operator. Note that URS matrices are not comparable to bistochastic matrices: they drop the column-stochasticity condition, but require that each row have a distinct `argmax`, which is not true of bistochastic matrices. This means that URS matrices can be trivially projected onto the probability simplex, which is useful for e.g. straight-through gradient optimization, or whenever hard permutation matrices are required, such as at test time. The one property URS matrices lack is column-stochasticity, but this is not central to soft sorting. Instead, this property arises in the work of [@2018_mena] because their goal is to relax the bipartite matching operator (rather than the `argsort` operator), and in this context bistochastic matrices are the natural choice. Similarly, [@2019_cuturi] yields bistochastic matrices because they are the solutions to optimal transport problems (this does, however, allow them to simultaneously relax the `argsort` and `rank` operators). Since our only goal (as in the `NeuralSort` paper) is to relax the `argsort` operator, column-stochasticity can be dropped, and URS matrices are the more natural choice.

Now on to the main Theorem, which shows that has the same desirable properties as `NeuralSort`. These are [@2018_grover Theorem 4]:

\[main\_theorem\] The operator has the following properties:

1.  Unimodality: $\forall \tau > 0$, $\SimpleSort{}^d_\tau(s)$ is a unimodal row stochastic matrix. Further, let $u$ denote the permutation obtained by applying $\texttt{argmax}$ row-wise to $\SimpleSort{}^d_\tau(s)$. Then, $u = \texttt{argsort}(s)$.

2.  Limiting behavior: If no elements of $s$ coincide, then: $$\lim_{\tau \rightarrow 0^+} \SimpleSort{}^d_\tau(s) = P_{\texttt{argsort}(s)}$$ In particular, this limit holds almost surely if the entries of $s$ are drawn from a distribution that is absolutely continuous w.r.t. the Lebesgue measure on $\mathbb R$.

**Proof**.

1.  Non-negativity and row affinity follow from the fact that every row in $\SimpleSort{}^d_\tau(s)$ is the result of a `softmax` operation. For the third property, we use that `softmax` preserves maximums and that $d(\cdot, x)$ has a unique minimum at $x$ for every $x\in\mathbb R$. Formally, let $\texttt{sort}(s) = [s_\ord{1}, \dots, s_\ord{n}]^T$, i.e. $s_\ord{1} \ge \dots \ge s_\ord{n}$ are the decreasing order statistics of $s$. Then: $$\begin{aligned}
    u_i &= \argmax_j \SimpleSort{}^d_\tau(s)[i, j] \\
    &= \argmax_j \left(\texttt{softmax}(-d(s_\ord{i}, s_j) / \tau)\right) \\
    &= \argmin_j \left(d(s_\ord{i}, s_j)\right) \\
    &= \texttt{argsort}(s)[i]\end{aligned}$$ as desired.

2.  It suffices to show that the $i$-th row of $\SimpleSort{}^d_{\tau}(s)$ converges to the one-hot representation $h$ of $\texttt{argsort}(s)[i]$. But by part 1, the $i$-th row of $\SimpleSort{}^d_{\tau}(s)$ is the softmax of $v / \tau$ where $v$ is a vector whose unique argmax is $\texttt{argsort}(s)[i]$. Since it is a well-known property of the softmax that $\lim_{\tau \rightarrow 0^+} \texttt{softmax}(v / \tau) = h$ [@1994_elfadel], we are done.

Note that the proof of unimodality of the operator is straightforward, unlike the proof for the `NeuralSort` operator, which requires proving a more technical Lemma and Corollary [@2018_grover Lemma 2, Corollary 3].

The row-stochastic property can be loosely interpreted as follows: row $r$ of and `NeuralSort` encodes a distribution over the value of the rank $r$ element, more precisely, the probability of it being equal to $s_j$ for each $j$. In particular, note that the first row of the $\SimpleSort{}^{|\cdot|}$ operator is precisely the `softmax` of the input vector. In general, the $r$-th row of the $\SimpleSort{}^d$ operator is the `softmax` of the negative distances to the $r$-th largest element.

Finally, regarding the choice of $d$ in , even though a large family of semi–metrics could be considered, in this work we experimented with the absolute value as well as the square distance and found the absolute value to be marginally better during experimentation. With this in consideration, in what follows we fix $d = |\cdot|$ the absolute value function, unless stated otherwise. We leave for future work learning the metric $d$ or exploring a larger family of such functions.

Comparing to `NeuralSort`
=========================

    def neural_sort(s, tau):
        n = tf.shape(s)[1]
        one = tf.ones((n, 1), dtype = tf.float32)
        A_s = tf.abs(s - tf.transpose(s, perm=[0, 2, 1]))
        B = tf.matmul(A_s, tf.matmul(one, tf.transpose(one)))
        scaling = tf.cast(n + 1 - 2 * (tf.range(n) + 1), dtype = tf.float32)
        C = tf.matmul(s, tf.expand_dims(scaling, 0))
        P_max = tf.transpose(C-B, perm=[0, 2, 1])
        P_hat = tf.nn.softmax(P_max / tau, -1)
        return P_hat

    def soft_sort(s, tau):
        s_sorted = tf.sort(s, direction='DESCENDING', axis=1)
        pairwise_distances = -tf.abs(tf.transpose(s, perm=[0, 2, 1]) - s_sorted)
        P_hat = tf.nn.softmax(pairwise_distances / tau, -1)
        return P_hat

-0.2in

Mathematical Simplicity {#sec:mathematical_simplicity}
-----------------------

The difference between and `NeuralSort` becomes apparent once we write down what the actual operators look like; the equations defining them (Eq. \[eq:2018\_grover\_eq5\], Eq. \[eq:simplesort2\]) are compact but do not offer much insight. Note that even though the work of [@2018_grover] shows that the `NeuralSort` operator has the desirable properties of Theorem \[main\_theorem\], the paper never gives a concrete example of what the operator instantiates to in practice, which keeps some of its complexity hidden.

Let $g_\tau : \mathbb R^{n \times n} \rightarrow \mathbb R^{n \times n}$ be the function defined as $g_\tau(X) = \texttt{softmax}(X / \tau)$, where the `softmax` is applied row-wise. Suppose that $n = 4$ and that $s$ is sorted in decreasing order $s_1 \ge s_2 \ge s_3 \ge s_4$. Then the `NeuralSort` operator is given in Figure \[eq:neuralsort\_logits\] and the $\SimpleSort{}$ operator is given in Figure \[eq:simplesort\_logits\]. Note that the diagonal of the logit matrix has been $0$-centered by subtracting a constant value from each row; this does not change the `softmax` and simplifies the expressions. The $\SimpleSort{}$ operator is straightforward, with the $i,j$-th entry of the logit matrix given by $-|s_i - s_j|$. In contrast, the $i,j$-th entry of the `NeuralSort` operator depends on all intermediate values $s_i,s_{i+1},\dots,s_j$. This is a consequence of the coupling between the `NeuralSort` operator and the complex identity used to derive it. As we show in this paper, this complexity can be completely avoided, and results in further benefits beyond aesthetic simplicity such as flexibility, speed and mathematical simplicity.

Note that for an arbitrary region of space other than $s_1 \ge s_2 \ge s_3 \ge s_4$, the `NeuralSort` and $\SimpleSort{}$ operators look just like Figures \[eq:neuralsort\_logits\] and \[eq:simplesort\_logits\] respectively except for relabelling of the $s_i$ and column permutations. Indeed, we have:

\[prop:invariance\] For both $f = \SimpleSort^d_\tau$ and $f = \texttt{NeuralSort}_\tau$, the following identity holds: $$f(s) = f(\texttt{sort}(s))P_{\texttt{argsort}(s)}$$

We defer the proof to appendix \[sec:proof\_of\_invariance\]. This proposition is interesting because it implies that the behaviour of the and `NeuralSort` operators can be completely characterized by their functional form on the region of space where $s_1 \ge s_2 \ge \dots \ge s_n$. Indeed, for any other value of $s$, we can compute the value of $\SimpleSort(s)$ or $\texttt{NeuralSort}(s)$ by first sorting $s$, then applying $\SimpleSort$ or $\texttt{NeuralSort}$, and finally sorting the columns of the resulting matrix with the inverse permutation that sorts $s$. In particular, to our point, the proposition shows that Figures \[eq:neuralsort\_logits\] and \[eq:simplesort\_logits\] are valid for *all* $s$ up to relabeling of the $s_i$ (by $s_\ord{i}$) and column permutations (by the inverse permutation that sorts $s$).

To further our comparison, in appendix \[sec:size\] we show how the $\SimpleSort{}$ and `NeuralSort` operators differ in terms of the *size* of their matrix entries.

Code Simplicity
---------------

In Figures \[code:neuralsort\] and \[code:simplesort\] we show TensorFlow implementations of the `NeuralSort` and operators respectively. has a simpler implementation than `NeuralSort`, which we shall see is reflected in its faster speed. (See section \[speed-comparison\])

Note that our implementation of is based directly off Eq. \[eq:simplesort2\], and we rely on the `sort` operator. We would like to remark that there is nothing wrong with using the `sort` operator in a stochastic computation graph. Indeed, the `sort` function is continuous, almost everywhere differentiable (with non-zero gradients) and piecewise linear, just like the `max`, `min` or `ReLU` functions.

Finally, the unimodality property (Theorem \[main\_theorem\]) implies that any algorithm that builds a relaxed permutation matrix can be used to construct the true discrete permutation matrix. This means that any relaxed sorting algorithm (in particular, `NeuralSort`) is lower bounded by the complexity of sorting, which justifies relying on sorting as a subroutine. As we show later, is faster than `NeuralSort`. Also, we believe that this modular approach is a net positive since sorting in CPU and GPU is a well studied problem [@2017_singh] and any underlying improvements will benefit ’s speed as well. For instance, the current implementation in TensorFlow relies on radix sort and heap sort depending on input size.

Experiments
===========

\[tab:run\_sort\]

\[tab:run\_median\]

  Algorithm                                     $n = 5$                             $n = 9$                            $n = 15$
  ------------------------------- ----------------------------------- ----------------------------------- -----------------------------------
  Deterministic NeuralSort         $\textbf{21.52}\ (\textbf{0.97})$   $\textbf{15.00}\ (\textbf{0.97})$       $18.81\ (\textbf{0.95})$
  Stochastic NeuralSort                $24.78\ (\textbf{0.97})$                 $17.79\ (0.96)$                     $18.10\ (0.94)$
  Deterministic SoftSort (Ours)        $23.44\ (\textbf{0.97})$                 $19.26\ (0.96)$            $\textbf{15.54}\ (\textbf{0.95})$
  Stochastic SoftSort (Ours)           $26.17\ (\textbf{0.97})$                 $19.06\ (0.96)$                     $20.65\ (0.94)$

\[tab:knn\]

                Algorithm                  MNIST     Fashion-MNIST    CIFAR-10
  ------------------------------------- ----------- --------------- ------------
               kNN+Pixels                  97.2%         85.8%         35.4%
                 kNN+PCA                   97.6%         85.9%         40.9%
                 kNN+AE                    97.6%         87.5%         44.2%
     kNN + Deterministic NeuralSort      **99.5%**       93.5%         90.7%
       kNN + Stochastic NeuralSort         99.4%         93.4%         89.5%
   kNN + Deterministic SoftSort (Ours)    99.37%        93.60 %      **92.03**%
    kNN + Stochastic SoftSort (Ours)      99.42%      **93.67**%       90.64%
              CNN (w/o kNN)                99.4%         93.4%       **95.1%**

We first compare to `NeuralSort` on the benchmarks from the `NeuralSort` paper [@2018_grover], using the code kindly open-sourced by the authors. We show that performs comparably to `NeuralSort`. Then, we devise a specific experiment to benchmark the speeds of the and `NeuralSort` operators in isolation, and show that is faster than `NeuralSort` while taking the same number of gradient steps to converge. This makes not only the simplest, but also the fastest relaxed sorting operator proposed to date.

Models
------

For both and `NeuralSort` we consider their deterministic and stochastic variants as in [@2018_grover]. The deterministic operators are those given by equations \[eq:2018\_grover\_eq5\] and \[eq:simplesort2\]. The stochastic variants are Plackett-Luce distributions reparameterized via Gumbel distributions [@2018_grover Section 4.1], where the $P_{\texttt{argsort}(\cdot)}$ operator that is applied to the samples is relaxed by means of the or `NeuralSort` operator; this is analogous to the Gumbel-Softmax trick where the `argmax` operator that is applied to the samples is relaxed via the `softmax` operator [@2017_jang; @2017_maddison].

Sorting Handwritten Numbers
---------------------------

The *large-MNIST* dataset of handwritten *numbers* is formed by concatenating $4$ randomly selected MNIST *digits*. In this task, a neural network is presented with a sequence of $n$ large-MNIST numbers and must learn the permutation that sorts these numbers. Supervision is provided only in the form of the ground-truth permutation. Performance on the task is measured by:

1.  The proportion of *permutations* that are perfectly recovered.

2.  The proportion of *permutation elements* that are correctly recovered.

Note that the first metric is always less than or equal to the second metric. We use the setup in [@2019_cuturi] to be able to compare against their Optimal-Transport-based method. They use $100$ epochs to train all models.

The results for the first metric are shown in Table \[tab:run\_sort\]. We report the mean and standard deviation over 10 runs. We see that and `NeuralSort` perform identically for all values of $n$. Moreover, our results for `NeuralSort` are better than those reported in [@2019_cuturi], to the extent that `NeuralSort` and outperform the method of [@2019_cuturi] for $n = 9, 15$, unlike reported in said paper. We found that the hyperparameter values reported in [@2018_grover] and used by [@2019_cuturi] for `NeuralSort` are far from ideal: [@2018_grover] reports using a learning rate of $10^{-4}$ and temperature values from the set $\{1,2,4,8,16\}$. However, a higher learning rate dramatically improves `NeuralSort`’s results, and higher temperatures also help. Concretely, we used a learning rate of $0.005$ for all the and `NeuralSort` models, and a value of $\tau = 1024$ for $n = 3,5,7$ and $\tau = 128$ for $n = 9, 15$. The results for the second metric are reported in appendix \[app:run\_sort\_aux\]. In the appendix we also include learning curves for and `NeuralSort`, which show that they converge at the same speed.

Quantile Regression
-------------------

As in the sorting task, a neural network is presented with a sequence of $n$ large-MNIST numbers. The task is to predict the median element from the sequence, and this is the only available form of supervision. Performance on the task is measured by mean squared error and Spearman’s rank correlation. We used $100$ iterations to train all models.

The results are shown in Table \[tab:run\_median\]. We used a learning rate of $10^{-3}$ for all models - again, higher than that reported in [@2018_grover] - and grid-searched the temperature on the set $\{128, 256, 512, 1024, 2048, 4096\}$ - again, higher than that reported in [@2018_grover]. We see that and `NeuralSort` perform similarly, with `NeuralSort` sometimes outperforming and vice versa. The results for `NeuralSort` are also much better than those reported in [@2018_grover], which we attribute to the better choice of hyperparameters, concretely, the higher learning rate. In the appendix we also include learning curves for and `NeuralSort`, which show that they converge at the same speed.

Differentiable kNN
------------------

In this setup, we explored using the operator to learn a differentiable $k$-nearest neighbours (kNN) classifier that is able to learn a representation function, used to measure the distance between the candidates.

In a supervised training framework, we have a dataset that consists of pairs $(x, y)$ of a datapoint and a label. We are interested in learning a map $\Phi$ to embed every $x$ such that we can use a kNN classifier to identify the class of $\hat{x}$ by looking at the class of its closest neighbours according to the distance $\|\Phi(x) - \Phi(\hat{x})\|$. Such a classifier would be valuable by virtue of being interpretable and robust to both noise and unseen classes.

This is achieved by constructing episodes during training that consist of one pair $\hat{x}, \hat{y}$ and $n$ candidate pairs $(x_i, y_i)$ for $i = 1\dots n$, arranged in two column vectors $X$ and $Y$. The probability $P(\hat{y}|\hat{x},X,Y)$ of class $\hat{y}$ under a kNN classifier is the average of the first $k$ entries in the vector $$P_{\texttt{argsort}(-\|\Phi(X) - \Phi(\hat{x})\|^2)} \mathbb{I}_{Y = \hat{y}}$$ where $\|\Phi(X) - \Phi(\hat{x})\|^2$ is the vector of squared distances from the candidate points and $\mathbb{I}_{Y = \hat{y}}$ is the binary vector indicating which indexes have class $\hat{y}$. Thus, if we replace $P_\texttt{argsort}$ by the operator we obtain a differentiable relaxation $\widehat{P}(\hat{y}|\hat{x},X,Y)$. To compute the loss we follow [@2018_grover] and use $-\widehat{P}(\hat{y}|\hat{x},X,Y)$. We also experimented with the cross entropy loss, but the performance went down for both methods.

When $k=1$, our method is closely related to *Matching Networks* [@2016_vinyals]. This follows from the following result: (See proof in Appendix \[sec:proof\_of\_dknn\])

\[dknn\] Let $k=1$ and $\widehat{P}$ be the differentiable kNN operator using $\SimpleSort{}^{|\cdot|}_2$. If we choose the embedding function $\Phi$ to be of norm $1$, then $$\widehat{P}(\hat{y}|\hat{x},X,Y) =
    {\sum_{i:y_i = \hat{y}} e^{\Phi(\hat{x})\cdot\Phi(x_i)}} \bigg/
    {\sum_{i=1\dots n} e^{\Phi(\hat{x})\cdot\Phi(x_i)}}$$

This suggests that our method is a generalization of *Matching Networks*, since in our experiments larger values of $k$ yielded better results consistently and we expect a kNN classifier to be more robust in a setup with noisy labels. However, *Matching Networks* use contextual embedding functions, and different networks for $\hat{x}$ and the candidates $x_i$, both of which could be incorporated to our setup. A more comprehensive study comparing both methods on a few shot classification dataset such as *Omniglot* [@2011_lake] is left for future work.

We applied this method to three benchmark datasets: MNIST, Fashion MNIST and CIFAR-10 with canonical splits. As baselines, we compare against `NeuralSort` as well as other kNN models with fixed representations coming from raw pixels, a PCA feature extractor and an auto-encoder. All the results are based on the ones reported in [@2018_grover]. We also included for comparison a standard classifier using a convolutional network.

Results are shown in Table \[tab:knn\]. In every case, we achieve comparable accuracies with `NeuralSort` implementation, either slightly outperforming or underperforming `NeuralSort`. See hyperparameters used in appendix \[sec:experimental\_details\].

Speed Comparison
----------------

We set up an experiment to compare the speed of the and `NeuralSort` operators. We are interested in exercising both their forward and backward calls. To this end, we set up a dummy learning task where the goal is to perturb an $n$-dimensional input vector $\theta$ to make it become sorted. We scale $\theta$ to $[0, 1]$ and feed it through the or `NeuralSort` operator to obtain $\widehat{P}(\theta)$, and place a loss on $\widehat{P}(\theta)$ that encourages it to become equal to the identity matrix, and thus encourages the input to become sorted.

Concretely, we place the cross-entropy loss between the true permutation matrix and the predicted soft URS matrix: $$L(\widehat{P}) = - \frac{1}{n}\sum_{i = 1}^n \log \widehat{P}[i, i]$$ This encourages the probability mass from each row of $\widehat{P}$ to concentrate on the diagonal, which drives $\theta$ to sort itself. Note that this is a trivial task, since for example a pointwise ranking loss $\frac{1}{n} \sum_{i = 1}^n (\theta_i + i)^2$ [@2008_taylor Section 2.2] leads the input to become sorted too, without any need for the or `NeuralSort` operators. However, this task is a reasonable benchmark to measure the speed of the two operators in a realistic learning setting.

We benchmark $n$ in the range $100$ to $4000$, and batch $20$ inputs $\theta$ together to exercise batching. Thus the input is a parameter tensor of shape $20 \times n$. Models are trained for $100$ epochs, which we verified is enough for the parameter vectors to become perfectly sorted by the end of training (i.e., to succeed at the task).

In Figure \[fig:benchmark\_tensorflow\] we show the results for the TensorFlow implementations of `NeuralSort` and given in Figures \[code:neuralsort\] and \[code:simplesort\] respectively. We see that on both CPU and GPU, is faster than `NeuralSort`. For $N = 4000$, is about 6 times faster than the `NeuralSort` implementation of [@2018_grover] on both CPU and GPU. We tried to speed up the `NeuralSort` implementation of [@2018_grover], and although we were able to improve it, `NeuralSort` was still slower than , concretely: $80\%$ slower on CPU and $40\%$ slower on GPU. Details of our improvements to the speed of the `NeuralSort` operator are provided in appendix \[appendix:neuralsort\_performance\_improvement\].

The performance results for PyTorch are provided in the appendix and are similar to the TensorFlow results. In the appendix we also show that the learning curves of with $d = |\cdot|^2$ and `NeuralSort` are almost identical; interestingly, we found that using $d = |\cdot|$ converges more slowly on this synthetic task.

We also investigated if relying on a sorting routine could cause slower run times in worst-case scenarios. When using sequences sorted in the opposite order we did not note any significant slowdowns. Furthermore, in applications where this could be a concern, the effect can be avoided entirely by shuffling the inputs before applying our operator.

As a final note, given that the cost of sorting is sub-quadratic, and most of the computation is payed when building and applying the $n \times n$ matrix, we also think that our algorithm could be made faster asymptotically by constructing sparse versions of the operator. For applications like differentiable nearest neighbors, evidence suggests than processing longer sequences yields better results, which motivates improvements in the asymptotic complexity. We leave this topic for future work.

Conclusion
==========

We have introduced , a simple continuous relaxation for the `argsort` operator. The $r$-th row of the operator is simply the `softmax` of the negative distances to the $r$-th largest element.

has similar properties to those of the `NeuralSort` operator of [@2018_grover]. However, due to its simplicity, is trivial to implement, more modular, faster than `NeuralSort`, and proofs regarding the operator are effortless. We also empirically find that it is just as easy to optimize.

Fundamentally, advances the state of the art in differentiable sorting by significantly simplifying previous approaches. Our code and experiments can be found at <https://github.com/sprillo/softsort>.

Acknowledgments {#acknowledgments .unnumbered}
===============

We thank Assistant Professor Jordan Boyd-Graber from University of Maryland and Visting Researcher at Google, and Thomas Müller from Google Language Research at Zurich for their feedback and comments on earlier versions of the manuscript. We would also like to thank the anonymous reviewers for their feedback that helped improve this work.

[^1]: This is called the `sort` operator in [@2018_grover]. We adopt the more conventional naming.

[^2]: Except when there are ties, which we assume is not the case. Ties are merely a technical nuisance and do not represent a problem for any of the methods (ours or other’s) discussed in this paper.
