VIA CONTINUOUS RELAXATIONS
Year: 2019

Abstract:
Sorting input objects is an important step in many machine learning pipelines. However, the sorting operator is non-differentiable with respect to its inputs, which prohibits end-to-end gradient-based optimization. In this work, we propose NeuralSort, a general-purpose continuous relaxation of the output of the sorting operator from permutation matrices to the set of unimodal row-stochastic matrices, where every row sums to one and has a distinct argmax. This relaxation permits straight-through optimization of any computational graph involve a sorting operation. Further, we use this relaxation to enable gradient-based stochastic optimization over the combinatorially large space of permutations by deriving a reparameterized gradient estimator for the Plackett-Luce family of distributions over permutations. We demonstrate the usefulness of our framework on three tasks that require learning semantic orderings of high-dimensional objects, including a fully differentiable, parameterized extension of the k-nearest neighbors algorithm.

heading: 1 INTRODUCTION
text: Learning to automatically sort objects is useful in many machine learning applications, such as topk multi-class classification (Berrada et al., 2018), ranking documents for information retrieval (Liu et al., 2009), and multi-object target tracking in computer vision (Bar-Shalom & Li, 1995). Such algorithms typically require learning informative representations of complex, high-dimensional data, such as images, before sorting and subsequent downstream processing. For instance, the k-nearest neighbors image classification algorithm, which orders the neighbors based on distances in the canonical pixel basis, can be highly suboptimal for classification (Weinberger et al., 2006). Deep neural networks can instead be used to learn representations, but these representations cannot be optimized end-to-end for a downstream sorting-based objective, since the sorting operator is not differentiable with respect to its input.
In this work, we seek to remedy this shortcoming by proposing NeuralSort, a continuous relaxation to the sorting operator that is differentiable almost everywhere with respect to the inputs. The output of any sorting algorithm can be viewed as a permutation matrix, which is a square matrix with entries in {0, 1} such that every row and every column sums to 1. Instead of a permutation matrix, NeuralSort returns a unimodal row-stochastic matrix. A unimodal row-stochastic matrix is defined as a square matrix with positive real entries, where each row sums to 1 and has a distinct argmax. All permutation matrices are unimodal row-stochastic matrices. NeuralSort has a temperature knob that controls the degree of approximation, such that in the limit of zero temperature, we recover a permutation matrix that sorts the inputs. Even for a non-zero temperature, we can efficiently project any unimodal matrix to the desired permutation matrix via a simple row-wise argmax operation. Hence, NeuralSort is also suitable for efficient straight-through gradient optimization (Bengio et al., 2013), which requires “exact” permutation matrices to evaluate learning objectives.
As the second primary contribution, we consider the use of NeuralSort for stochastic optimization over permutations. In many cases, such as latent variable models, the permutations may be latent but directly influence observed behavior, e.g., utility and choice models are often expressed as distributions over permutations which govern the observed decisions of agents (Regenwetter et al.,
∗Equal contribution
2006; Chierichetti et al., 2018). By learning distributions over unobserved permutations, we can account for the uncertainty in these permutations in a principled manner. However, the challenge with stochastic optimization over discrete distributions lies in gradient estimation with respect to the distribution parameters. Vanilla REINFORCE estimators are impractical for most cases, or necessitate custom control variates for low-variance gradient estimation (Glasserman, 2013).
In this regard, we consider the Plackett-Luce (PL) family of distributions over permutations (Plackett, 1975; Luce, 1959). A common modeling choice for ranking models, the PL distribution is parameterized by n scores, with its support defined over the symmetric group consisting of n! permutations. We derive a reparameterizable sampler for stochastic optimization with respect to this distribution, based on Gumbel perturbations to the n (log-)scores. However, the reparameterized sampler requires sorting these perturbed scores, and hence the gradients of a downstream learning objective with respect to the scores are not defined. By using NeuralSort instead, we can approximate the objective and obtain well-defined reparameterized gradient estimates for stochastic optimization.
Finally, we apply NeuralSort to tasks that require us to learn semantic orderings of complex, highdimensional input data. First, we consider sorting images of handwritten digits, where the goal is to learn to sort images by their unobserved labels. Our second task extends the first one to quantile regression, where we want to estimate the median (50-th percentile) of a set of handwritten numbers. In addition to identifying the index of the median image in the sequence, we need to learn to map the inferred median digit to its scalar representation. In the third task, we propose an algorithm that learns a basis representation for the k-nearest neighbors (kNN) classifier in an end-to-end procedure. Because the choice of the k nearest neighbors requires a non-differentiable sorting, we use NeuralSort to obtain an approximate, differentiable surrogate. On all tasks, we observe significant empirical improvements due to NeuralSort over the relevant baselines and competing relaxations to permutation matrices.

heading: 2 PRELIMINARIES
text: An n-dimensional permutation z = [z1, z2, . . . , zn]T is a list of unique indices {1, 2, . . . , n}. Every permutation z is associated with a permutation matrix Pz ∈ {0, 1}n×n with entries given as:
Pz[i, j] = { 1 if j = zi 0 otherwise.
Let Zn denote the set of all n! possible permutations in the symmetric group. We define the sort : Rn → Zn operator as a mapping of n real-valued inputs to a permutation corresponding to a descending ordering of these inputs. E.g., if the input vector s = [9, 1, 5, 2]T , then sort(s) = [1, 3, 4, 2]T since the largest element is at the first index, second largest element is at the third index and so on. In case of ties, elements are assigned indices in the order they appear. We can obtain the sorted vector simply via Psort(s)s.

heading: 2.1 PLACKETT-LUCE DISTRIBUTIONS
text: The family of Plackett-Luce distributions over permutations is best described via a generative process: Consider a sequence of n items, each associated with a canonical index i = 1, 2, . . . , n. A common assumption in ranking models is that the underlying generating process for any observed permutation of n items satisfies Luce’s choice axiom (Luce, 1959). Mathematically, this axiom defines the ‘choice’ probability of an item with index i as: q(i) ∝ si where si > 0 is interpreted as the score of item with index i. The normalization constant is given by Z = ∑ i∈{1,2,...,n} si.
If we choose the n items one at a time (without replacement) based on these choice probabilities, we obtain a discrete distribution over all possible permutations. This distribution is referred to as the Plackett-Luce (PL) distribution, and its probability mass function for any z ∈ Zn is given by:
q(z|s) = sz1 Z sz2 Z − sz1 · · · szn Z − ∑n−1 i=1 szi
(1)
where s = {s1, s2, . . . , sn} is the vector of scores parameterizing this distribution (Plackett, 1975).

heading: 2.2 STOCHASTIC COMPUTATION GRAPHS
text: The abstraction of stochastic computation graphs (SCG) compactly specifies the forward value and the backward gradient computation for computational circuits. An SCG is a directed acyclic graph that consists of three kinds of nodes: input nodes which specify external inputs (including parameters), deterministic nodes which are deterministic functions of their parents, and stochastic nodes which are distributed conditionally on their parents. See Schulman et al. (2015) for a review.
To define gradients of an objective function with respect to any node in the graph, the chain rule necessitates that the gradients with respect to the intermediate nodes are well-defined. This is not the case for the sort operator. In Section 3, we propose to extend stochastic computation graphs with nodes corresponding to a relaxation of the deterministic sort operator. In Section 4, we further use this relaxation to extend computation graphs to include stochastic nodes corresponding to distributions over permutations. The proofs of all theoretical results in this work are deferred to Appendix B.

heading: 3 NEURALSORT: THE RELAXED SORTING OPERATOR
text: Our goal is to optimize training objectives involving a sort operator with gradient-based methods. Consider the optimization of objectives written in the following form:
L(θ, s) = f(Pz; θ) (2) where z = sort(s).
Here, s ∈ Rn denotes a vector of n real-valued scores, z is the permutation that (deterministically) sorts the scores s, and f(·) is an arbitrary function of interest assumed to be differentiable w.r.t a set of parameters θ and z. For example, in a ranking application, these scores could correspond to the inferred relevances of n webpages and f(·) could be a ranking loss. Figure 1 shows the stochastic computation graph corresponding to the objective in Eq. 2. We note that this could represent part of a more complex computation graph, which we skip for ease of presentation while maintaining the generality of the scope of this work.
While the gradient of the above objective w.r.t. θ is well-defined and can be computed via standard backpropogation, the gradient w.r.t. the scores s is not defined since the sort operator is not differentiable w.r.t. s. Our solution is to derive a relaxation to the sort operator that leads to a surrogate objective with well-defined gradients. In particular, we seek to use such a relaxation to replace the permutation matrix Pz in Eq. 2 with an approximation P̂z such that the surrogate objective f(P̂z; θ) is differentiable w.r.t. the scores s.
The general recipe to relax non-differentiable operators with discrete codomains N is to consider differentiable alternatives that map the input to a larger continuous codomain M with desirable properties. For gradient-based optimization, we are interested in two key properties:
1. The relaxation is continuous everywhere and differentiable (almost-)everywhere with respect to elements in the input domain.
2. There exists a computationally efficient projection fromM back to N .
Relaxations satisfying the first requirement are amenable to automatic differentiation for optimizing stochastic computational graphs. The second requirement is useful for evaluating metrics and losses that necessarily require a discrete output akin to the one obtained from the original, non-relaxed operator. E.g., in straight-through gradient estimation (Bengio et al., 2013; Jang et al., 2017), the
non-relaxed operator is used for evaluating the learning objective in the forward pass and the relaxed operator is used in the backward pass for gradient estimation.
The canonical example is the 0/1 loss used for binary classification. While the 0/1 loss is discontinuous w.r.t. its inputs (real-valued predictions from a model), surrogates such as the logistic and hinge losses are continuous everywhere and differentiable almost-everywhere (property 1), and can give hard binary predictions via thresholding (property 2).
Note: For brevity, we assume that the argmax operator is applied over a set of elements with a unique maximizer and hence, the operator has well-defined semantics. With some additional bookkeeping for resolving ties, the results in this section hold even if the elements to be sorted are not unique. See Appendix C.
Unimodal Row Stochastic Matrices. The sort operator maps the input vector to a permutation, or equivalently a permutation matrix. Our relaxation to sort is motivated by the geometric structure of permutation matrices. The set of permutation matrices is a subset of doubly-stochastic matrices, i.e., a non-negative matrix such that the every row and column sums to one. If we remove the requirement that every column should sum to one, we obtain a larger set of row stochastic matrices. In this work, we propose a relaxation to sort that maps inputs to an alternate subset of row stochastic matrices, which we refer to as the unimodal row stochastic matrices. Definition 1 (Unimodal Row Stochastic Matrices). An n× n matrix is Unimodal Row Stochastic if it satisfies the following conditions:
1. Non-negativity: U [i, j] ≥ 0 ∀i, j ∈ {1, 2, . . . , n}. 2. Row Affinity: ∑n j=1 U [i, j] = 1 ∀i ∈ {1, 2, . . . , n}.
3. Argmax Permutation: Let u denote an n-dimensional vector with entries such that ui = argmaxj U [i, j] ∀i ∈ {1, 2, . . . , n}. Then, u ∈ Zn, i.e., it is a valid permuation.
We denote Un as the set of n× n unimodal row stochastic matrices.
All row stochastic matrices satisfy the first two conditions. The third condition is useful for gradient based optimization involving sorting-based losses. The condition provides a straightforward mechanism for extracting a permutation from a unimodal row stochastic matrix via a row-wise argmax operation. Figure 2 shows the relationships between the different subsets of square matrices.
NeuralSort. Our relaxation to the sort operator is based on a standard identity for evaluating the sum of the k largest elements in any input vector. Lemma 2. [Lemma 1 in Ogryczak & Tamir (2003)] For an input vector s = [s1, s2, . . . , sn]T that is sorted as s[1] ≥ s[2] ≥ . . . ≥ s[n], we have the sum of the k-largest elements given as:
k∑ i=1 s[i] = min λ∈{s1,s2,...,sn} λk + n∑ i=1 max(si − λ, 0). (3)
The identity in Lemma 2 outputs the sum of the top-k elements. The k-th largest element itself can be recovered by taking the difference of the sum of top-k elements and the top-(k − 1) elements.
Corollary 3. Let s = [s1, s2, . . . , sn]T be a real-valued vector of length n. LetAs denote the matrix of absolute pairwise differences of the elements of s such that As[i, j] = |si − sj |. The permutation matrix Psort(s) corresponding to sort(s) is given by:
Psort(s)[i, j] = { 1 if j = argmax[(n+ 1− 2i)s−As1] 0 otherwise
(4)
where 1 denotes the column vector of all ones.
E.g., if we set i = b(n + 1)/2c then the non-zero entry in the i-th row Psort(s)[i, :] corresponds to the element with the minimum sum of (absolute) distance to the other elements. As desired, this corresponds to the median element. The relaxation requires O(n2) operations to compute As, as opposed to the O(n log n) overall complexity for the best known sorting algorithms. In practice however, it is highly parallelizable and can be implemented efficiently on GPU hardware.
The argmax operator is non-differentiable which prohibits the direct use of Corollary 3 for gradient computation. Instead, we propose to replace the argmax operator with softmax to obtain a continuous relaxation P̂sort(s)(τ). In particular, the i-th row of P̂sort(s)(τ) is given by:
P̂sort(s)[i, :](τ) = softmax [((n+ 1− 2i)s−As1)/τ ] (5)
where τ > 0 is a temperature parameter. Our relaxation is continuous everywhere and differentiable almost everywhere with respect to the elements of s. Furthermore, we have the following result.
Theorem 4. Let P̂sort(s) denote the continuous relaxation to the permutation matrix Psort(s) for an arbitrary input vector s and temperature τ defined in Eq. 5. Then, we have:
1. Unimodality: ∀τ > 0, P̂sort(s) is a unimodal row stochastic matrix. Further, let u denote the permutation obtained by applying argmax row-wise to P̂sort(s). Then, u = sort(s).
2. Limiting behavior: If we assume that the entries of s are drawn independently from a distribution that is absolutely continuous w.r.t. the Lebesgue measure in R, then the following convergence holds almost surely:
lim τ→0+
P̂sort(s)[i, :](τ) = Psort(s)[i, :] ∀i ∈ {1, 2, . . . , n}. (6)
Unimodality allows for efficient projection of the relaxed permutation matrix P̂sort(s) to the hard matrix Psort(s) via a row-wise argmax, e.g., for straight-through gradients. For analyzing limiting behavior, independent draws ensure that the elements of s are distinct almost surely. The temperature τ controls the degree of smoothness of our approximation. At one extreme, the approximation becomes tighter as the temperature is reduced. In practice however, the trade-off is in the variance of these estimates, which is typically lower for larger temperatures.

heading: 4 STOCHASTIC OPTIMIZATION OVER PERMUTATIONS
text: In many scenarios, we would like the ability to express our uncertainty in inferring a permutation e.g., latent variable models with latent nodes corresponding to permutations. Random variables that assume values corresponding to permutations can be represented via stochastic nodes in the
stochastic computation graph. For optimizing the parameters of such a graph, consider the following class of objectives:
L(θ, s) = Eq(z|s) [f(Pz; θ)] (7) where θ and s denote sets of parameters, Pz is the permutation matrix corresponding to the permutation z, q(·) is a parameterized distribution over the elements of the symmetric group Zn, and f(·) is an arbitrary function of interest assumed to be differentiable in θ and z. The SCG is shown in Figure 3a. In contrast to the SCG considered in the previous section (Figure 1), here we are dealing with a distribution over permutations as opposed to a single (deterministically computed) one.
While such objectives are typically intractable to evaluate exactly since they require summing over a combinatorially large set, we can obtain unbiased estimates efficiently via Monte Carlo. Monte Carlo estimates of gradients w.r.t. θ can be derived simply via linearity of expectation. However, the gradient estimates w.r.t. s cannot be obtained directly since the sampling distribution depends on s. The REINFORCE gradient estimator (Glynn, 1990; Williams, 1992; Fu, 2006) uses the fact that∇sq(z|s) = q(z|s)∇s log q(z|s) to derive the following Monte Carlo gradient estimates:
∇sL(θ, s) = Eq(z|s) [f(Pz; θ)∇s log q(z|s)] + Eq(z|s) [∇sf(Pz; θ)] . (8)

heading: 4.1 REPARAMETERIZED GRADIENT ESTIMATORS FOR PL DISTRIBUTIONS
text: REINFORCE gradient estimators typically suffer from high variance (Schulman et al., 2015; Glasserman, 2013). Reparameterized samplers provide an alternate gradient estimator by expressing samples from a distribution as a deterministic function of its parameters and a fixed source of randomness (Kingma & Welling, 2014; Rezende et al., 2014; Titsias & Lázaro-Gredilla, 2014). Since the randomness is from a fixed distribution, Monte Carlo gradient estimates can be derived by pushing the gradient operator inside the expectation (via linearity). In this section, we will derive a reparameterized sampler and gradient estimator for the Plackett-Luce (PL) family of distributions.
Let the score si for an item i ∈ {1, 2, . . . , n} be an unobserved random variable drawn from some underlying score distribution (Thurstone, 1927). Now for each item, we draw a score from its corresponding score distribution. Next, we generate a permutation by applying the deterministic sort operator to these n randomly sampled scores. Interestingly, prior work has shown that the resulting distribution over permutations corresponds to a PL distribution if and only if the scores are sampled independently from Gumbel distributions with identical scales. Proposition 5. [adapted from Yellott Jr (1977)] Let s be a vector of scores for the n items. For each item i, sample gi ∼ Gumbel(0, β) independently with zero mean and a fixed scale β. Let s̃ denote the vector of Gumbel perturbed log-scores with entries such that s̃i = β log si + gi. Then:
q(s̃z1 ≥ · · · ≥ s̃zn) = sz1 Z sz2 Z − sz1 · · · szn Z − ∑n−1 i=1 szi . (9)
For ease of presentation, we assume β = 1 in the rest of this work. Proposition 5 provides a method for sampling from PL distributions with parameters s by adding Gumbel perturbations to the logscores and applying the sort operator to the perturbed log-scores. This procedure can be seen as a reparameterization trick that expresses a sample from the PL distribution as a deterministic function of the scores and a fixed source of randomness (Figure 3b). Letting g denote the vector of i.i.d. Gumbel perturbations, we can express the objective in Eq. 7 as:
L(θ, s) = Eg [ f(Psort(log s+g); θ) ] . (10)
While the reparameterized sampler removes the dependence of the expectation on the parameters s, it introduces a sort operator in the computation graph such that the overall objective is nondifferentiable in s. In order to obtain a differentiable surrogate, we approximate the objective based on the NeuralSort relaxation to the sort operator:
Eg [ f(Psort(log s+g); θ) ] ≈ Eg [ f(P̂sort(log s+g); θ) ] := L̂(θ, s). (11)
Accordingly, we get the following reparameterized gradient estimates for the approximation: ∇sL̂(θ, s) = Eg [ ∇sf(P̂sort(log s+g); θ) ] (12)
which can be estimated efficiently via Monte Carlo because the expectation is with respect to a distribution that does not depend on s.

heading: 5 DISCUSSION AND RELATED WORK
text: The problem of learning to rank documents based on relevance has been studied extensively in the context of information retrieval. In particular, listwise approaches learn functions that map objects to scores. Much of this work concerns the PL distribution: the RankNet algorithm (Burges et al., 2005) can be interpreted as maximizing the PL likelihood of pairwise comparisons between items, while the ListMLE ranking algorithm in Xia et al. (2008) extends this with a loss that maximizes the PL likelihood of ground-truth permutations directly. The differentiable pairwise approaches to ranking, such as Rigutini et al. (2011), learn to approximate the comparator between pairs of objects. Our work considers a generalized setting where sorting based operators can be inserted anywhere in computation graphs to extend traditional pipelines e.g., kNN.
Prior works have proposed relaxations of permutation matrices to the Birkhoff polytope, which is defined as the convex hull of the set of permutation matrices a.k.a. the set of doubly-stochastic matrices. A doubly-stochastic matrix is a permutation matrix iff it is orthogonal and continuous relaxations based on these matrices have been used previously for solving NP-complete problems such as seriation and graph matching (Fogel et al., 2013; Fiori et al., 2013; Lim & Wright, 2014). Adams & Zemel (2011) proposed the use of the Sinkhorn operator to map any square matrix to the Birkhoff polytope. They interpret the resulting doubly-stochastic matrix as the marginals of a distribution over permutations. Mena et al. (2018) propose an alternate method where the square matrix defines a latent distribution over the doubly-stochastic matrices themselves. These distributions can be sampled from by adding elementwise Gumbel perturbations. Linderman et al. (2018) propose a rounding procedure that uses the Sinkhorn operator to directly sample matrices near the Birkhoff polytope. Unlike Mena et al. (2018), the resulting distribution over matrices has a tractable density. In practice, however, the approach of Mena et al. (2018) performs better and will be the main baseline we will be comparing against in our experiments in Section 6.
As discussed in Section 3, NeuralSort maps permutation matrices to the set of unimodal rowstochastic matrices. For the stochastic setting, the PL distribution permits efficient sampling, exact and tractable density estimation, making it an attractive choice for several applications, e.g., variational inference over latent permutations. Our reparameterizable sampler, while also making use of the Gumbel distribution, is based on a result unique to the PL distribution (Proposition 5).
The use of the Gumbel distribution for defining continuous relaxations to discrete distributions was first proposed concurrently by Jang et al. (2017) and Maddison et al. (2017) for categorical variables, referred to as Gumbel-Softmax. The number of possible permutations grow factorially with the dimension, and thus any distribution over n-dimensional permutations can be equivalently seen as a distribution over n! categories. Gumbel-softmax does not scale to a combinatorially large number of categories (Kim et al., 2016; Mussmann et al., 2017), necessitating the use of alternate relaxations, such as the one considered in this work.

heading: 6 EXPERIMENTS
text: We refer to the two approaches proposed in Sections 3, 4 as Deterministic NeuralSort and Stochastic NeuralSort, respectively. For additional hyperparameter details and analysis, see Appendix D.

heading: 6.1 SORTING HANDWRITTEN NUMBERS
text: Dataset. We first create the large-MNIST dataset, which extends the MNIST dataset of handwritten digits. The dataset consists of multi-digit images, each a concatenation of 4 randomly selected individual images from MNIST, e.g., is one such image in this dataset. Each image is associated with a real-valued label, which corresponds to its concatenated MNIST labels, e.g., the label of is 1810. Using the large-MNIST dataset, we finally create a dataset of sequences. Every sequence is this dataset consists of n randomly sampled large-MNIST images.
Setup. Given a dataset of sequences of large-MNIST images, our goal is to learn to predict the permutation that sorts the labels of the sequence of images, given a training set of ground-truth permutations. Figure 4 (Task 1) illustrates this task on an example sequence of n = 5 large-MNIST images. This task is a challenging extension of the one considered by Mena et al. (2018) in sorting scalars, since it involves learning the semantics of high-dimensional objects prior to sorting. A
good model needs to learn to dissect the individual digits in an image, rank these digits, and finally, compose such rankings based on the digit positions within an image. The available supervision, in the form of the ground-truth permutation, is very weak compared to a classification setting that gives direct access to the image labels.
Baselines. All baselines use a CNN that is shared across all images in a sequence to map each large-MNIST image to a feature space. The vanilla row-stochastic (RS) baseline concatenates the CNN representations for n images into a single vector that is fed into a multilayer perceptron that outputs nmulticlass predictions of the image probabilities for each rank. The Sinkhorn and GumbelSinkhorn baselines, as discussed in Section 5, use the Sinkhorn operator to map the stacked CNN representations for the n objects into a doubly-stochastic matrix. For all methods, we minimized the cross-entropy loss between the predicted matrix and the ground-truth permutation matrix.
Results. Following Mena et al. (2018), our evaluation metric is the the proportion of correctly predicted permutations on a test set of sequences. Additionally, we evaluate the proportion of individual elements ranked correctly. Table 1 demonstrates that the approaches based on the proposed sorting relaxation significantly outperform the baseline approaches for all n considered. The performance of the deterministic and stochastic variants are comparable. The vanilla RS baseline performs well in ranking individual elements, but is not good at recovering the overall square matrix.
We believe the poor performance of the Sinkhorn baselines is partly because these methods were designed and evaluated for matchings. Like the output of sort, matchings can also be represented as permutation matrices. However, distributions over matchings need not satisfy Luce’s choice axiom or imply a total ordering, which could explain the poor performance on the tasks considered.

heading: 6.2 QUANTILE REGRESSION
text: Setup. In this experiment, we extend the sorting task to regression. Again, each sequence contains n large-MNIST images, and the regression target for each sequence is the 50-th quantile (i.e., the median) of the n labels of the images in the sequence. Figure 4 (Task 2) illustrates this task on an example sequence of n = 5 large-MNIST images, where the goal is to output the third largest label. The design of this task highlights two key challenges since it explicitly requires learning both a suitable representation for sorting high-dimensional inputs and a secondary function that approximates the label itself (regression). Again, the supervision available in the form of the label of only a single image at an arbitrary and unknown location in the sequence is weak.
Baselines. In addition to Sinkhorn and Gumbel-Sinkhorn, we design two more baselines. The Constant baseline always returns the median of the full range of possible outputs, ignoring the input sequence. This corresponds to 4999.5 since we are sampling large-MNIST images uniformly in the range of four-digit numbers. The vanilla neural net (NN) baseline directly maps the input sequence of images to a real-valued prediction for the median.
Results. Our evaluation metric is the mean squared error (MSE) and R2 on a test set of sequences. Results for n = {5, 9, 15} images are shown in Table 2. The Vanilla NN baseline while incurring a large MSE, is competitive on the R2 metric. The other baselines give comparable performance on the MSE metric. The proposed NeuralSort approaches outperform the competing methods on both the metrics considered. The stochastic NeuralSort approach is the consistent best performer on MSE, while the deterministic NeuralSort is slightly better on the R2 metric.
6.3 END-TO-END, DIFFERENTIABLE k-NEAREST NEIGHBORS
Setup. In this experiment, we design a fully differentiable, end-to-end k-nearest neighbors (kNN) classifier. Unlike a standard kNN classifier which computes distances between points in a predefined space, we learn a representation of the data points before evaluating the k-nearest neighbors.
We are given access to a datasetD of (x, y) pairs of standard input data and their class labels respectively. The differentiable kNN algorithm consists of two hyperparameters: the number of training neighbors n, the number of top candidates k, and the sorting temperature τ . Every sequence of items here consists of a query point x and a randomly sampled subset of n candidate nearest neighbors from the training set, say {x1,x2, . . . ,xn}. In principle, we could use the entire training set (excluding the query point) as candidate points, but this can hurt the learning both computationally and statistically. The query points are randomly sampled from the train/validation/test sets as appropriate but the nearest neighbors are always sampled from the training set. The loss function optimizes for a representation space hφ(·) (e.g., CNN) such that the top-k candidate points with the minimum Euclidean distance to the query point in the representation space have the same label as the query point. Note that at test time, once the representation space hφ is learned, we can use the entire training set as the set of candidate points, akin to a standard kNN classifier. Figure 5 illustrates the proposed algorithm.
Formally, for any datapoint x, let z denote a permutation of the n candidate points. The uniformlyweighted kNN loss, denoted as `kNN(·), can be written as follows:
`kNN(P̂z, y, y1, y2, . . . , yn) = − 1
k k∑ j=1 n∑ i=1 1(yi = y)P̂z[i, j] (13)
where {y1, y2, . . . , yn} are the labels for the candidate points. Note that when P̂z is an exact permutation matrix (i.e., temperature τ → 0), this expression is exactly the negative of the fraction of k nearest neighbors that have the same label as x.
Using Eq. 13, the training objectives for Deterministic and Stochastic NeuralSort are given as:
Deterministic: min φ
1 |D| ∑
(x,y)∈D
`kNN(P̂sort(s), y, y1, . . . , yn) (14)
Stochastic: min φ
1 |D| ∑
(x,y)∈D
Ez∼q(z|s) [ `kNN(P̂z, y, y1, y2, . . . , yn) ] (15)
where each entry of s is given as sj = −‖hφ(x)− hφ(xj)‖22.
Datasets. We consider three benchmark datasetes: MNIST dataset of handwritten digits, FashionMNIST dataset of fashion apparel, and the CIFAR-10 dataset of natural images (no data augmentation) with the canonical splits for training and testing.
Baselines. We consider kNN baselines that operate in three standard representation spaces: the canonical pixel basis, the basis specified by the top 50 principal components (PCA), an autonencoder (AE). Additionally, we experimented with k = 1, 3, 5, 9 nearest neighbors and across two distance metrics: uniform weighting of all k-nearest neighbors and weighting nearest neighbors by the inverse of their distance. For completeness, we trained a CNN with the same architecture as the one used for NeuralSort (except the final layer) using the cross-entropy loss.
Results. We report the classification accuracies on the standard test sets in Table 3. On both datasets, the differentiable kNN classifier outperforms all the baseline kNN variants including the convolutional autoencoder approach. The performance is much closer to the accuracy of a standard CNN.

heading: 7 CONCLUSION
text: In this paper, we proposed NeuralSort, a continuous relaxation of the sorting operator to the set of unimodal row-stochastic matrices. Our relaxation facilitates gradient estimation on any computation graph involving a sort operator. Further, we derived a reparameterized gradient estimator for the Plackett-Luce distribution for efficient stochastic optimization over permutations. On three illustrative tasks including a fully differentiable k-nearest neighbors, our proposed relaxations outperform prior work in end-to-end learning of semantic orderings of high-dimensional objects.
In the future, we would like to explore alternate relaxations to sorting as well as applications that extend widely-used algorithms such as beam search (Goyal et al., 2018). Both deterministic and stochastic NeuralSort are easy to implement. We provide reference implementations in Tensorflow (Abadi et al., 2016) and PyTorch (Paszke et al., 2017) in Appendix A. The full codebase for this work is open-sourced at https://github.com/ermongroup/neuralsort.

heading: ACKNOWLEDGEMENTS
text: This research was supported by NSF (#1651565, #1522054, #1733686), ONR, AFOSR (FA955019-1-0024), FLI, and Amazon AWS. AG is supported by MSR fellowship and Stanford Data Science scholarship. We are thankful to Jordan Alexander, Kristy Choi, Adithya Ganesh, Karan Goel, Neal Jean, Daniel Levy, Jiaming Song, Yang Song, Serena Yeung, and Hugh Zhang for feedback.

heading: APPENDICES

heading: A SORTING OPERATOR

heading: A.1 TENSORFLOW
text: Sorting Relaxation for Deterministic NeuralSort:
import tensorflow as tf
def deterministic_NeuralSort(s, tau): """ s: input elements to be sorted. Shape: batch_size x n x 1 tau: temperature for relaxation. Scalar. """
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
Reparameterized Sampler for Stochastic NeuralSort:
def sample_gumbel(samples_shape, eps = 1e-10):
U = tf.random_uniform(samples_shape, minval=0, maxval=1) return -tf.log(-tf.log(U + eps) + eps)
def stochastic_NeuralSort(s, n_samples, tau): """ s: parameters of the PL distribution. Shape: batch_size x n x 1. n_samples: number of samples from the PL distribution. Scalar. tau: temperature for the relaxation. Scalar. """
batch_size = tf.shape(s)[0] n = tf.shape(s)[1] log_s_perturb = s + sample_gumbel([n_samples, batch_size, n, 1]) log_s_perturb = tf.reshape(log_s_perturb, [n_samples * batch_size, n,
1])
P_hat = deterministic_NeuralSort(log_s_perturb, tau) P_hat = tf.reshape(P_hat, [n_samples, batch_size, n, n])
return P_hat

heading: A.2 PYTORCH
text: Sorting Relaxation for Deterministic NeuralSort:
import torch
def deterministic_NeuralSort(s, tau): """ s: input elements to be sorted. Shape: batch_size x n x 1 tau: temperature for relaxation. Scalar. """
n = s.size()[1] one = torch.ones((n, 1), dtype = torch.float32)
A_s = torch.abs(s - s.permute(0, 2, 1)) B = torch.matmul(A_s, torch.matmul(one, torch.transpose(one, 0, 1))) scaling = (n + 1 - 2 * (torch.arange(n) + 1)).type(torch.float32) C = torch.matmul(s, scaling.unsqueeze(0))
P_max = (C-B).permute(0, 2, 1) sm = torch.nn.Softmax(-1) P_hat = sm(P_max / tau)
return P_hat
Reparamterized Sampler for Stochastic NeuralSort:
def sample_gumbel(samples_shape, eps = 1e-10):
U = torch.rand(samples_shape) return -torch.log(-torch.log(U + eps) + eps)
def stochastic_NeuralSort(s, n_samples, tau): """ s: parameters of the PL distribution. Shape: batch_size x n x 1. n_samples: number of samples from the PL distribution. Scalar. tau: temperature for the relaxation. Scalar. """
batch_size = s.size()[0] n = s.size()[1] log_s_perturb = torch.log(s) + sample_gumbel([n_samples, batch_size, n,
1]) log_s_perturb = log_s_perturb.view(n_samples * batch_size, n, 1)
P_hat = deterministic_NeuralSort(log_s_perturb, tau) P_hat = P_hat.view(n_samples, batch_size, n, n)
return P_hat

heading: B PROOFS OF THEORETICAL RESULTS

heading: B.1 LEMMA 2
text: Proof. For any value of λ, the following inequalities hold:
k∑ i=1 s[i] = λk + k∑ i=1 (s[i] − λ)
≤ λk + k∑ i=1 max(s[i] − λ, 0)
≤ λk + n∑ i=1 max(si − λ, 0).
Furthermore, for λ = s[k]:
λk + n∑ i=1 max(si − λ, 0) = s[k]k + n∑ i=1 max(si − s[k], 0)
= s[k]k + k∑ i=1 (s[i] − s[k])
= k∑ i=1 s[i].
This finishes the proof.

heading: B.2 COROLLARY 3
text: Proof. We first consider at exactly what values of λ the sum in Lemma 2 is minimized. For simplicity we will only prove the case where all values of s are distinct.
The equality ∑k i=1 s[i] = λk + ∑n i=1 max(si − λ, 0) holds only when s[k] ≤ λ ≤ s[k+1]. By Lemma 2, these values of λ also minimize the RHS of the equality.
Symmetrically, if one considers the score vector t = −s, then λ(n− k+1)+ ∑n i=1 max(ti−λ, 0) is minimized at t[n−k+1] ≤ λ ≤ t[n−k+2].
Replacing λ by −λ and using the definition of t implies that λ(k − 1− n) + ∑n i=1 max(λ− si, 0) is minimized at s[k−1] ≤ λ ≤ s[k]. It follows that:
s[k] = argmin λ∈s
( λk +
n∑ i=1 max(si − λ, 0)
) + ( λ(k − 1− n) +
n∑ i=1 max(λ− si, 0)
)
= argmin λ∈s λ(2k − 1− n) + n∑ i=1 |si − λ|.
Thus, if si = s[k], then i = argmin(2k − 1− n)s+As1. This finishes the proof.

heading: B.3 THEOREM 4
text: We prove the two properties in the statement of the theorem independently:
1. Unimodality
Proof. By definition of the softmax function, the entries ef P̂ are positive and sum to 1. To show that P̂ satisfies the argmax permutation property, . Formally, for any given row i, we construct the argmax permutation vector u as:
ui = argmax[softmax((n+ 1− 2i)s−As1)] = argmax[(n+ 1− 2i)s−As1] = [i]
where the square notation [i] denotes the index of the i-th largest element. The first step follows from the fact that the softmax function is monotonically increasing and hence, it preserves the argmax. The second equality directly follows from Corollary 3. By definition, sort(s) = {[1], [2], . . . , [n]}, finishing the proof.
2. Limiting behavior
Proof. As shown in Gao & Pavel (2017), the softmax function may be equivalently defined as softmax(z/τ) = argmaxx∈∆n−1〈x, z〉 − τ ∑n i=1 xi log xi. In particular, limτ→0 softmax(z/τ) = argmaxx. The distributional assumptions ensure that the elements of s are distinct a.s., so plugging in z = (n + 1 − 2k)s − As1 completes the proof.

heading: B.4 PROPOSITION 5
text: This result follows from an earlier result by Yellott Jr (1977). We give the proof sketch below and refer the reader to Yellott Jr (1977) for more details.
Sketch. Consider random variables {Xi}ni=1 such that Xi ∼ Exp(szi)). We may prove by induction a generalization of the memoryless property:
q(X1 ≤ · · · ≤ Xn|x ≤ min i Xi)
= ∫ ∞ 0 q(x ≤ X1 ≤ x+ t|x ≤ min i Xi)q(X2 ≤ · · · ≤ Xn|x+ t ≤ min i≥2 Xi)dt
= ∫ ∞ 0 q(0 ≤ X1 ≤ t)q(X2 ≤ · · · ≤ Xn|x+ t ≤ min i≥2 Xi)dt.
If we assume as inductive hypothesis that q(X2 ≤ · · · ≤ Xn|x+ t ≤ mini≥2Xi) = q(X2 ≤ · · · ≤ Xn|t ≤ mini≥2Xi), we complete the induction as:
q(X1 ≤ · · · ≤ Xn|x ≤ min i Xi)
= ∫ ∞ 0 q(0 ≤ X1 ≤ t)q(X2 ≤ · · · ≤ Xn|t ≤ min i≥2 Xi)dt = q(X1 ≤ X2 ≤ · · · ≤ Xn|0 ≤ min i Xi).
It follows from a familiar property of argmin of exponential distributions that: q(X1 ≤ X2 ≤ · · · ≤ Xn) = q(X1 ≤ min
i Xi)q(X2 ≤ · · · ≤ Xn|X1 ≤ min i Xi)
= sz1 Z q(X2 ≤ · · · ≤ Xn|X1 ≤ min i Xi)
= sz1 Z ∫ ∞ 0 q(X1 = x)q(X2 ≤ · · · ≤ Xn|x ≤ min i≥2 Xi)dx = sz1 Z q(X2 ≤ · · · ≤ Xn),
and by another induction, we have q(X1 ≤ · · · ≤ Xn) = ∏n i=1 szi Z− ∑i−1 k=1 szk .
Finally, following the argument of Balog et al. (2017), we apply the strictly decreasing function g(x) = −β log x to this identity, which from the definition of the Gumbel distribution implies:
q(s̃z1 ≥ · · · ≥ s̃zn) = n∏ i=1
szi Z − ∑i−1 k=1 szk .

heading: C ARG MAX SEMANTICS FOR TIED MAX ELEMENTS
text: While applying the argmax operator to a vector with duplicate entries attaining the max value, we need to define the operator semantics for argmax to handle ties in the context of the proposed relaxation.
Definition 6. For any vector with ties, let argmax set denote the operator that returns the set of all indices containing the max element. We define the argmax of the i-th in a matrix M recursively:
1. If there exists an index j ∈ {1, 2, . . . , n} that is a member of argmax set(M [i, :]) and has not been assigned as an argmax of any row k < i, then the argmax is the smallest such index.
2. Otherwise, the argmax is the smallest index that is a member of the argmax set(M [i, :]).
This function is efficiently computable with additional bookkeeping.
Lemma 7. For an input vector s with the sort permutation matrix given as Psort(s), we have sj1 = sj2 if and only if there exists a row i such that P̂ [i, j1] = P̂ [i, j2] for all j1, j2 ∈ {1, 2, . . . , n}.
Proof. From Eq. 5, we have the i-th row of P̂ [i, :] given as:
P̂ [i, :] = softmax [((n+ 1− 2i)s−As1)/τ ]
. Therefore, we have the equations:
P̂ [i, j1] = exp(((n+ 1− 2i)sj1 − (As1)i)/τ)
Z
= P̂ [i, j2] = exp(((n+ 1− 2i)sj2 − (As1)i)/τ)
Z
for some fixed normalization constant Z. As the function f(x) = exp(((n+1−2i)x−(As1)i)/τ)Z is invertible, both directions of the lemma follow immediately.
Lemma 8. If argmax set(P̂ [i1, :]) and argmax set(P̂ [i2, :]) have a non-zero intersection, then argmax set(P̂ [i1, :]) = argmax set(P̂ [i2, :]).
Proof. Assume without loss of generality that | argmax set(P̂ [i1, :])| > 1 for some i. Let j1, j2 be two members of | argmax set(P̂ [i1, :])|. By Lemma 7, sj1 = sj2 , and therefore P̂ [i2, j1] = P̂ [i2, j2]. Hence if j1 ∈ argmax set(P̂ [i2, :]), then j2 is also an element. A symmetric argument implies that if j2 ∈ argmax set(P̂ [i2, :]), then j1 is also an element for arbitrary j1, j2 ∈ | argmax set(P̂ [i1, :])|. This completes the proof.
Proposition 9. (Argmax Permutation with Ties) For s = [s1, s2, . . . , sn]T ∈ Rn, the vector z defined by zi = argmaxj P̂sort(s)[i, j] is such that z ∈ Zn.
Proof. From Corollary 3, we know that the row P̂sort(s)[i, :] attains its maximum (perhaps nonuniquely) at some P̂sort(s)[i, j] where sj = s[i]. Note that s[i] is well-defined even in the case of ties.
Consider an arbitrary row P̂sort(s)[i, :] and let argmax set(P̂sort(s)[i, :]) = {j1, . . . , jm}. It follows from Lemma 7 that there are exactly m scores in s that are equal to s[i]: sj1 , . . . , sjm . These scores corresponds to m values of s[i′] such that s[i′] = s[i], and consequently to m rows P [i′, :] that are maximized with values s[i′] = s[i] and consequently (by Lemma 8) at indices j1, . . . , jm.
Suppose we now chose an i′ such that s[i′] 6= s[i]. Then P̂ [i′, :] attains its maximum at some P̂sort(s)[i′, j′] where sj′ = s[i′]. Because sj′ = s[i′] 6= s[i] = sj , Lemma 7 tells us that P̂ [i′, :] does not attain its maximum at any of j1, . . . , jm. Therefore, only m rows have a non-zero argmax set intersection with argmax set(P [i, :]).
Because P [i, :] is one of these rows, there can be up to m − 1 such rows above it. Because each row above only has one argmax assigned via the tie-breaking protocol, it is only possible for up to m − 1 elements of argmax set(P [i, :]) to have been an argmax of a previous row k < i. As | argmax set(P [i, :])| = m, there exists at least one element that has not been specified as the argmax of a previous row (pigeon-hole principle). Thus, the argmax of each row are distinct. Because each argmax is also an element of {1, . . . , n}, it follows that z ∈ Zn.

heading: D EXPERIMENTAL DETAILS AND ANALYSIS
text: We used Tensorflow (Abadi et al., 2016) and PyTorch (Paszke et al., 2017) for our experiments. In Appendix A, we provide “plug-in” snippets for implementing our proposed relaxations in both Tensorflow and PyTorch. The full codebase for reproducing the experiments can be found at https://github.com/ermongroup/neuralsort.
For the sorting and quantile regression experiments, we used standard training/validation/test splits of 50, 000/10, 000/10, 000 images of MNIST for constructing the large-MNIST dataset. We ensure that only digits in the standard training/validation/test sets of the MNIST dataset are composed together to generate the corresponding sets of the large-MNIST dataset. For CIFAR-10, we used a split of 45, 000/5000/10, 000 examples for training/validation/test. With regards to the baselines considered, we note that the REINFORCE based estimators were empirically observed to be worse than almost all baselines for all our experiments.

heading: D.1 SORTING HANDWRITTEN NUMBERS
text: Architectures. We control for the choice of computer vision models by using the same convolutional network architecture for each sorting method. This architecture is as follows:
Conv[Kernel: 5x5, Stride: 1, Output: 140x28x32, Activation: Relu] → Pool[Stride: 2, Output: 70x14x32] → Conv[Kernel: 5x5, Stride: 1, Output: 70x14x64, Activation: Relu] → Pool[Stride: 2, Output: 35x7x64] → FC[Units: 64, Activation: Relu]
Note that the dimension of a 5-digit large-MNIST image is 140×28. The primary difference between our methods is how we combine the scores to output a row-stochastic prediction matrix.
For NeuralSort-based methods, we use another fully-connected layer of dimension 1 to map the image representations to n scalar scores. In the case of Stochastic NeuralSort, we then sample from the PL distribution by perturbing the scores multiple times with Gumbel noise. Finally, we use the NeuralSort operator to map the set of n scores (or each set of n perturbed scores) to its corresponding unimodal row-stochastic matrix.
For Sinkhorn-based methods, we use a fully-connected layer of dimension n to map each image to an n-dimensional vector. These vectors are then stacked into an n × n matrix. We then either map this matrix to a corresponding doubly-stochastic matrix (Sinkhorn) or sample directly from a distribution over permutation matrices via Gumbel perturbations (Gumbel-Sinkhorn). We implemented the Sinkhorn operator based on code snippets obtained from the open source implementation of Mena et al. (2018) available at https://github.com/google/gumbel sinkhorn.
For the Vanilla RS baseline, we ran each element through a fully-connected n dimensional layer, concatenated the representations of each element and then fed the results through three fullyconnected n2-unit layers to output multiclass predictions for each rank.
All our methods yield row-stochastic n × n matrices as their final output. Our loss is the row-wise cross-entropy loss between the true permutation matrix and the row-stochastic output.
Hyperparameters. For this experiment, we used an Adam optimizer with an initial learning rate of 10−4 and a batch size of 20. Continuous relaxations to sorting also introduce another hyperparameter: the temperature τ for the Sinkhorn-based and NeuralSort-based approaches. We tuned this hyperparameter on the set {1, 2, 4, 8, 16} by picking the model with the best validation accuracy on predicting entire permutations (as opposed to predicting individual maps between elements and ranks).
Effect of temperature. In Figure 7, we report the log-variance in gradient estimates as a function of the temperature τ . Similar to the effect of temperature observed for other continuous relaxations to discrete objects such as Gumbel-softmax (Jang et al., 2017; Maddison et al., 2017), we note that higher temperatures lead to lower variance in gradient estimates. The element-wise mean squared
difference between unimodal approximations P̂sort(s) and the projected hard permutation matrices Psort(s) for the best τ on the test set is shown in Table 4.

heading: D.2 QUANTILE REGRESSION
text: Architectures. Due to resource constraints, we ran the quantile regression experiment on 4-digit numbers instead of 5-digit numbers. We use the same neural network architecture as previously used in the sorting experiment.
Conv[Kernel: 5x5, Stride: 1, Output: 112x28x32, Activation: Relu] → Pool[Stride: 2, Output: 56x14x32] → Conv[Kernel: 5x5, Stride: 1, Output: 56x14x64, Activation: Relu] → Pool[Stride: 2, Output: 28x7x64] → FC[Units: 64, Activation: Relu]
The vanilla NN baseline for quantile regression was generated by feeding the CNN representations into a series of three fully-connected layers of ten units each, the last of which mapped to a singleunit estimate of the median. In the other experiments, one copy of this network was used to estimate each element’s rank through a method like Gumbel-Sinkhorn or NeuralSort that produces a rowstochastic matrix, while another copy was used to estimate each element’s value directly. Point predictions are obtained by multiplying the center row of the matrix with the column vector of estimated values, and we minimize the `2 loss between these point predictions and the true median, learning information about ordering and value simultaneously.
Hyperparameters. We used the Adam optimizer with an initial learning rate of 10−4 and a batch size of 5. The temperature τ was tuned on the set {1, 2, 4, 8, 16} based on the validation loss.
Further Analysis. In Figure 8, we show the scatter plots for the true vs. predicted medians on 2000 test points from the large-MNIST dataset as we vary n. For stochastic NeuralSort, we average the predictions across 5 samples. As we increase n, the distribution of true medians concentrates, leading to an easier prediction problem (at an absolute scale) and hence, we observe lower MSE for larger n in Table 2. However, the relatively difficulty of the problem increases with increasing n, as the model is trying to learn a semantic sorting across a larger set of elements. This is reflected in the R2 values in Table 2 which show a slight dip as n increases.
D.3 END-TO-END, DIFFERENTIABLE k-NEAREST NEIGHBORS
Architectures. The baseline kNN implementation for the pixel basis, PCA basis and the autoencoder basis was done using sklearn. For the autoencoder baselines for kNN, we used the following standard architectures.
MNIST and Fashion-MNIST: The dimension of the encoding used for distance computation in kNN is 50.
FC[Units: 500, Activation: Relu] → FC[Units: 500, Activation: Relu] → FC[Units: 50, Activation: Relu] = (embedding) → FC[Units: 500, Activation: Relu] → FC[Units: 500, Activation: Relu] → FC[Units: 784, Activation: Sigmoid]
CIFAR-10: The dimension of the encoding used for distance computation in kNN is 256. The architecture and training procedure follows the one available at https://github.com/shibuiwilliam/Keras Autoencoder.
Conv[Kernel: 3x3, Stride: 1, Output: 32x32x64, Activation: Relu] → Pool[Stride: 2, Output: 16x16x64] → Conv[Kernel: 3x3, Stride: 1, Output: 16x16x32, Normalization: BatchNorm, Activation: Relu] → Pool[Stride: 2, Output: 8x8x32] → Conv[Stride: 3, Output: 8x8x16, Normalization: BatchNorm, Activation: Relu] → MaxPool[Stride: 2, Output: 4x4x16] = (embedding) → Conv[Kernel: 3x3, Stride: 1, Output: 4x4x16, Normalization: BatchNorm, Activation: Relu] → UpSampling[Size: 2x2, Output: 8x8x16] → Conv[Kernel: 3x3, Stride: 1, Output: 8x8x32, Normalization: BatchNorm, Activation: Relu] → UpSampling[Size: 2x2, Output: 16x16x32] → Conv[Kernel: 3x3, Output: 16x16x64, Normalization: BatchNorm, Activation: Relu] → UpSampling[Size: 2x2, Output: 32x32x64] → Conv[Kernel: 3x3, Stride: 1, Output: 32x32x3, Normalization: BatchNorm, Activation: Sigmoid]
For the MNIST experiments with NeuralSort, we used a network similar to the large-MNIST network used in the previous experiments:
Conv[Kernel: 5x5, Stride: 1, Output: 24x24x20, Activation: Relu] → Pool[Stride: 2, Output: 12x12x20] → Conv[Kernel: 5x5, Stride: 1, Output: 8x8x50, Activation: Relu] → Pool[Stride: 2, Output: 4x4x50] → FC[Units: 500, Activation: Relu]
For the Fashion-MNIST and CIFAR experiments with NeuralSort, we use the ResNet18 architecture as described in https://github.com/kuangliu/pytorch-cifar.
Hyperparameters. For this experiment, we used an SGD optimizer with a momentum parameter of 0.9, with a batch size of 100 queries and 100 neighbor candidates at a time. We chose the temperature hyperparameter from the set {1, 16, 64}, the constant learning rate from {10−4, 10−5}, and the number of nearest neighbors k from the set {1, 3, 5, 9}. The model with the best evaluation loss was evaluated on the test set. We suspect that accuracy improvements can be made by a more expensive hyperparameter search and a more fine-grained learning rate schedule.
Accuracy for different k. In Table 5, we show the performance of Deterministic and Stochastic NeuralSort for different choice of the hyperparameter k for the differentiable k-nearest neighbors algorithm.

heading: E LOSS FUNCTIONS
text: For each of the experiments in this work, we assume we have access to a finite dataset D = {(x(1),y(1)), (x(2),y(2)), . . .}. Our goal is to learn a predictor for y given x, as in a standard supervised learning (classification/regression) setting. Below, we state and elucidate the semantics of the training objective optimized by Deterministic and Stochastic NeuralSort for the sorting and quantile regression experiments.

heading: E.1 SORTING HANDWRITTEN NUMBERS
text: We are given a dataset D of sequences of large-MNIST images and the permutations that sort the sequences. That is, every datapoint in D consists of an input x, which corresponds to a sequence containing n images, and the desired output label y, which corresponds to the permutation that sorts this sequence (as per the numerical values of the images in the input sequence). For example, Figure 4 shows one input sequence of n = 5 images, and the permutation y = [3, 5, 1, 4, 2] that sorts this sequence.
For any datapoint x, let `CE(·) denote the average multiclass cross entropy (CE) error between the rows of the true permutation matrix Py and a permutation matrix Pŷ corresponding to a predicted permutation, say ŷ.
`CE(Py, Pŷ) = 1
n n∑ i=1 n∑ j=1 1(Py[i, j] = 1) logPŷ[i, j]
where 1(·) denotes the indicator function. Now, we state the training objective functions for the Deterministic and Stochastic NeuralSort approaches respectively.
1. Deterministic NeuralSort
min φ
1 |D| ∑
(x,y)∈D
`CE(Py, P̂sort(s)) (16)
where each entry of s is given as sj = hφ(xj).
2. Stochastic NeuralSort
min φ
1 |D| ∑
(x,y)∈D
Ez∼q(z|s) [ `CE(Py, P̂z)) ] (17)
where each entry of s is given as sj = hφ(xj).
To ground this in our experimental setup, the score sj for each large-MNIST image xj in any input sequence x of n = 5 images is obtained via a CNN hφ() with parameters φ. Note that the CNN parameters φ are shared across the different images x1,x2, . . . ,xn in the sequence for efficient learning.

heading: E.2 QUANTILE REGRESSION
text: In contrast to the previous experiment, here we are given a dataset D of sequences of large-MNIST images and only the numerical value of the median element for each sequence. For example, the desired label corresponds to y = 2960 (a real-valued scalar) for the input sequence of n = 5 images in Figure 4.
For any datapoint x, let `MSE(·) denote the mean-squared error between the true median y and the prediction, say ŷ.
`MSE(y, ŷ) = ‖y − ŷ‖22
For the NeuralSort approaches, we optimize the following objective functions.
1. Deterministic NeuralSort
min φ,θ
1 |D| ∑
(x,y)∈D
`MSE(y, gθ(P̂sort(s)x)) (18)
where each entry of s is given as sj = hφ(xj).
2. Stochastic NeuralSort
min φ,θ
1 |D| ∑
(x,y)∈D
Ez∼q(z|s) [ `MSE(y, gθ(P̂zx)) ] (19)
where each entry of s is given as sj = hφ(xj).
As before, the score sj for each large-MNIST image xj in any input sequence x of n images is obtained via a CNN hφ() with parameters φ. Once we have a predicted permutation matrix P̂sort(s) (or P̂z) for deterministic (or stochastic) approaches, we extract the median image via P̂sort(s)x (or P̂zx). Finally, we use a neural network gθ(·)with parameters θ to regress this image to a scalar prediction for the median.

Review summaries:
In many machine learning applications, sorting is an important step such as ranking. However, the sorting operator is not differentiable with respect to its inputs. The main idea of the paper is to introduce a continuous relaxation of the sorting operator in order to construct an end-to-end gradient-based optimization. This relaxation is introduced as \hat{P}_{sort(s)} (see Equation 4). The paper also introduces a stochastic extension of its method