Original: S TOCHASTIC O PTIMIZATION OF S ORTING N ETWORKS VIA C ONTINUOUS R ELAXATIONS
Revised: S TOCHASTIC O PTIMIZATION OF S ORTING N ETWORKS VIA C ONTINUOUS R ELAXATIONS

Original: Anonymous authors

Original: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Sorting input objects is an important step within many machine learning pipelines.
Revised: Sorting input objects is an important step in many machine learning pipelines.

Original: However, the sorting operator is non-differentiable w.r.t. its inputs, which prohibits end-to-end gradient-based optimization.
Revised: However, the sorting operator is non-differentiable with respect to its inputs, which prohibits end-to-end gradient-based optimization.

Original: In this work, we propose  a generalpurpose continuous relaxation of the output of the sorting operator from permutation matrices to the set of unimodal matrices .
Revised: 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 arg max .

Revised: This relaxation permits straight-through optimization of any computational graph involve a sorting operation.

Original: Further, we use this relaxation to enable more efﬁcient stochastic optimization over the combinatorially large space of permutations. In particular, we derive a reparameterized gradient estimator for thewidely used Plackett-Luce family of distributions.
Revised: 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.

Original: We demonstrate the usefulness of our framework on three tasks that require learning semantic orderings of high-dimensional objects.
Revised: 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

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: In machine learning, learning to automatically sort objects is useful in many  applications, such as top- k losses for classiﬁcation (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).
Revised: Learning to automatically sort objects is useful in many machine learning applications, such as topk multi-class classiﬁcation (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).

Original: Many such algorithms involve learning permutations of complex, high-dimensional objects, e.g. , images. Such algorithms typically require learning informative representations of data prior to sorting and subsequent downstream processing.
Revised: Such algorithms typically require learning informative representations of complex, high-dimensional data, such as images, before sorting and subsequent downstream processing.

Original: For instance, the k -nearest neighbors image classiﬁcation algorithm, which orders the neighbors based on distances in the canonical pixel basis, can be highly suboptimal for classiﬁcation.
Revised: For instance, the k -nearest neighbors image classiﬁcation algorithm, which orders the neighbors based on distances in the canonical pixel basis, can be highly suboptimal for classiﬁcation (Weinberger et al., 2006).

Original: While gradient-based optimization methods for deep neural networks excel at representation learning, the non-differentiability of sorting w.r.t. the input representations makes it challenging to use the sort operator in conjunction with automatic differentiation tools.
Revised: 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.

Original: In this work, we propose a  relaxation to the sorting operator that iscontinuous everywhere and differentiable almost everywhere w.r.t. the inputs.
Revised: 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.

Original: The outcome of any sorting algorithm can be viewed as a permutation matrix, which is a square matrix with0/1 entries  such that every row and every column sums to 1.
Revised: 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.

Original: We propose a continuous relaxation of  permutation matrices onto the set of unimodal  row stochastic matrices.
Revised: Instead of a permutation matrix, NeuralSort returns a unimodal row-stochastic matrix. A unimodal row-stochastic matrix is deﬁned as a square matrix with positive real entries, where each row sums to 1 and has a distinct arg max . All permutation matrices are unimodal row-stochastic matrices.

Original: The relaxation includes a temperature knob that controls the degree of approximation, recovering a permutation matrix in the limit of zero temperature.
Revised: 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.

Original: Besides every row summing to 1 , unimodal matrices additionally can be transformed to a permutation matrix that sorts the inputs by simply taking the arg max within each row.
Revised: Even for a non-zero temperature, we can efﬁciently project any unimodal matrix to the desired permutation matrix via a simple row-wise arg max operation.

Original: Hence, such matrices are also suitable for efﬁcient straight-through gradient optimization (Bengio et al., 2013) that require “hard” permutation matrices for evaluating learning objectives during the forward pass.
Revised: Hence, NeuralSort is also suitable for efﬁcient straight-through gradient optimization (Bengio et al., 2013), which requires “exact” permutation matrices to evaluate learning objectives.

Original: Our second contribution extends the use of the proposed deterministic sorting relaxation to stochastic optimization over permutations.
Revised: As the second primary contribution, we consider the use of NeuralSort for stochastic optimization over permutations.

Original: In many cases, such as latent variable models, the permutations maynot be  directly observed.
Revised: In many cases, such as latent variable models, the permutations may be latent but directly inﬂuence 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., ; Chierichetti et al., 2018).

Original: By allowing for distributions over  permutations, we can account for the uncertainty due to the unobserved permutations in a principled manner.
Revised: By learning distributions over unobserved permutations, we can account for the uncertainty in these permutations in a principled manner.

Original: However, the challenge withincluding stochastic nodes over discrete objects lies in gradient estimation with respect to the parameters of the discrete distribution.
Revised: However, the challenge with stochastic optimization over discrete distributions lies in gradient estimation with respect to the distribution parameters.

Original: Vanilla REINFORCE estimators are impractical for most cases, or necessitate custom control variates for low-variance gradient estimation (Glasserman, 2013; Schulman et al., 2015).
Revised: Vanilla REINFORCE estimators are impractical for most cases, or necessitate custom control variates for low-variance gradient estimation (Glasserman, 2013).

Original: 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.
Revised: 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 deﬁned over the symmetric group consisting of n !

Original: Appealingly, we can also derive a reparameterizable sampler for stochastic optimization with respect to this distribution, based on Gumbel perturbations to the distribution parameters.
Revised: permutations. We derive a reparameterizable sampler for stochastic optimization with respect to this distribution, based on Gumbel perturbations to the n (log-)scores.

Original: The reparameterized sampler, however requiresa sorting operation which makes the  objective non-differentiable w.r.t. these parameters.
Revised: 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 deﬁned.

Original: Our proposed continuous relaxation to sorting allows us to approximate the objective and obtain welldeﬁned reparameterized gradient estimates.
Revised: By using NeuralSort instead, we can approximate the objective and obtain well-deﬁned reparameterized gradient estimates for stochastic optimization.

Original: Finally, we apply this framework to tasks that require us to learn semantic orderings ofthe complex, high-dimensional input data.
Revised: Finally, we apply NeuralSort to tasks that require us to learn semantic orderings of complex, highdimensional input data.

Original: First, we consider sortingof images of handwritten digits, where the goal is to learn to sort images by their unobserved labels.
Revised: First, we consider sorting images of handwritten digits, where the goal is to learn to sort images by their unobserved labels.

Original: Our second task extends the ﬁrst one to identifying the  50 -th quantile (a.k.a. the median) of a set of handwritten numbers.
Revised: Our second task extends the ﬁrst one to quantile regression, where we want to estimate the median ( 50 -th percentile) of a set of handwritten numbers.

Original: 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.
Revised: 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.

Original: In the third task, we propose an algorithm that learns a basis representation for the k -nearest neighbors (kNN) classiﬁer in an end-to-end procedure.
Revised: In the third task, we propose an algorithm that learns a basis representation for the k -nearest neighbors (kNN) classiﬁer in an end-to-end procedure.

Original: Because the choice of the k nearest neighbors requires a non-differentiable sorting, we use our relaxation to obtain an approximate, differentiable surrogate.
Revised: Because the choice of the k nearest neighbors requires a non-differentiable sorting, we use NeuralSort to obtain an approximate, differentiable surrogate.

Original: On all tasks, we observe signiﬁcant empirical improvements  over the  competing benchmarks.
Revised: On all tasks, we observe signiﬁcant empirical improvements due to NeuralSort over the relevant baselines and competing relaxations to permutation matrices.

Original:  P RELIMINARIES
Revised:  P RELIMINARIES

Original: An n -dimensional permutation z =
Revised: An n -dimensional permutation z =

Original: [ z 1 , z 2 , . . .
Revised: [ z 1 , z 2 , . . .

Original: , z n ] T is a list of unique indices { 1 , 2 , . . . , n } .
Revised: , z n ] T is a list of unique indices { 1 , 2 , . . . , n } .

Original: Every permutation z is associated with a permutation matrix P z ∈ { 0 , 1 } n × n with entries given as:
Revised: Every permutation z is associated with a permutation matrix P z ∈ { 0 , 1 } n × n with entries given as:

Original: P z [ i, j ] = (cid:26) 1 if j = z i 0 otherwise .
Revised: P z [ i, j ] = (cid:26) 1 if j = z i 0 otherwise .

Original: Let Z n denote the set of all n ! possible permutations in the symmetric group.
Revised: Let Z n denote the set of all n ! possible permutations in the symmetric group.

Original: We deﬁne the sort : R n → Z n operator as a mapping of n real-valued inputs to a permutation corresponding to a descending ordering of these inputs.
Revised: We deﬁne the sort : R n → Z n operator as a mapping of n real-valued inputs to a permutation corresponding to a descending ordering of these inputs.

Original: E.g. , if the input vector s =
Revised: E.g. , if the input vector s =

Original: [9 , 1 , 5 , 2] T , then sort ( s ) =
Revised: [9 , 1 , 5 , 2] T , then sort ( s ) =

Original: [1 , 3 , 4 , 2] T since the largest element is at the ﬁrst index, second largest element is at the third index and so on.
Revised: [1 , 3 , 4 , 2] T since the largest element is at the ﬁrst index, second largest element is at the third index and so on.

Original: In case of ties, elements are assigned indices in the order they appear.
Revised: In case of ties, elements are assigned indices in the order they appear.

Original: We can obtain the sorted vector simply via P sort ( s ) s .
Revised: We can obtain the sorted vector simply via P sort ( s ) s .

Original: 2. P LACKETT -L UCE DISTRIBUTIONS
Revised: 2. P LACKETT -L UCE DISTRIBUTIONS

Original: 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 .
Revised: 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 .

Original: A common assumption in ranking models is that the underlying generating process for any observed permutation of n items satisﬁes Luce’s choice
Revised: A common assumption in ranking models is that the underlying generating process for any observed permutation of n items satisﬁes Luce’s choice

Original: axiom (Luce, 1959).
Revised: axiom (Luce, 1959).

Original: Mathematically, this axiom deﬁnes the ‘choice’ probability of an item with index i as: q ( i ) ∝ s i where s i > 0 is interpreted as the score of item with index i .
Revised: Mathematically, this axiom deﬁnes the ‘choice’ probability of an item with index i as: q ( i ) ∝ s i where s i > 0 is interpreted as the score of item with index i .

Original: The normalization constant is given by Z = (cid:80) i ∈{ 1 , 2 ,...,n } s i .
Revised: The normalization constant is given by Z = (cid:80) i ∈{ 1 , 2 ,...,n } s i .

Original: 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.
Revised: 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.

Original: This distribution is referred to as the Plackett-Luce (PL) distribution, and its probability mass function for any z ∈ Z n is given by:
Revised: This distribution is referred to as the Plackett-Luce (PL) distribution, and its probability mass function for any z ∈ Z n is given by:

Original: [Equation]
Revised: [Equation]

Original: where s = { s 1 , s 2 , . .
Revised: where s = { s 1 , s 2 , . .

Original: , s n } is the vector of scores parameterizing this distribution (Plackett, 1975). 
Revised: , s n } is the vector of scores parameterizing this distribution (Plackett, 1975). [Figure]

Original: 2. S TOCHASTIC COMPUTATION GRAPHS
Revised: 2. S TOCHASTIC COMPUTATION GRAPHS

Original: The abstraction of stochastic computation graphs (SCG) compactly speciﬁes the forward value and the backward gradient computation for computational circuits.
Revised: The abstraction of stochastic computation graphs (SCG) compactly speciﬁes the forward value and the backward gradient computation for computational circuits.

Original: 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.
Revised: 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.

Original: See Schulman et al.
Revised: See Schulman et al.

Original: (2015) for a review.
Revised: (2015) for a review.

Original: To deﬁne gradients of an objective function w.r.t. any node in the graph, the chain rule necessitates that the gradients w.r.t. the intermediate nodes are well-deﬁned.
Revised: To deﬁne 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-deﬁned.

Original: This is not the case for the sort operator.
Revised: This is not the case for the sort operator.

Original: In Section 3, we propose to extend stochastic computation graphs with nodes corresponding to a relaxation of the deterministic sort operator.
Revised: In Section 3, we propose to extend stochastic computation graphs with nodes corresponding to a relaxation of the deterministic sort operator.

Original: In Section 4, we further use this relaxation to extend computation graphs to include stochastic nodes corresponding to distributions over permutations.
Revised: In Section 4, we further use this relaxation to extend computation graphs to include stochastic nodes corresponding to distributions over permutations.

Original: The proofs of all theoretical results in this work are deferred to Appendix B.
Revised: The proofs of all theoretical results in this work are deferred to Appendix B.

Original:  [Figure] T HE R ELAXED S ORTING O PERATOR
Revised:  N EURAL S ORT : T HE R ELAXED S ORTING O PERATOR

Original: Our goal is to optimize training objectives involving a sort operator with gradient-based methods.
Revised: Our goal is to optimize training objectives involving a sort operator with gradient-based methods.

Original: Consider the optimization of objectives written in the following form:
Revised: Consider the optimization of objectives written in the following form:

Original: [Equation]
Revised: [Equation]

Original: where z = sort ( s ) .
Revised: where z = sort ( s ) .

Original: Here, s ∈ R n 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 .
Revised: Here, s ∈ R n 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 .

Original: For example, in a ranking application, these scores could correspond to the inferred relevances of n webpages and f ( · ) could be a ranking loss.
Revised: For example, in a ranking application, these scores could correspond to the inferred relevances of n webpages and f ( · ) could be a ranking loss.

Original: Figure 1 shows the stochastic computation graph corresponding to the objective in Eq.
Revised: Figure 1 shows the stochastic computation graph corresponding to the objective in Eq.

Original: 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.
Revised: 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.

Original: While the gradient of the above objective w.r.t.
Revised: While the gradient of the above objective w.r.t.

Original: θ is well-deﬁned and can be computed via standard backpropogation, the gradient w.r.t.
Revised: θ is well-deﬁned and can be computed via standard backpropogation, the gradient w.r.t.

Original: the scores s is not deﬁned since the sort operator is not differentiable w.r.t.
Revised: the scores s is not deﬁned since the sort operator is not differentiable w.r.t.

Original: s .
Revised: s .

Original: Our solution is to derive a relaxation to the sort operator that leads to a surrogate objective with well-deﬁned gradients.
Revised: Our solution is to derive a relaxation to the sort operator that leads to a surrogate objective with well-deﬁned gradients.

Original: In particular, we seek to use such a relaxation to replace the permutation matrix P z in Eq.
Revised: In particular, we seek to use such a relaxation to replace the permutation matrix P z in Eq.

Original: 2 with an approximation (cid:98) P
Revised: 2 with an approximation (cid:98) P

Original: z such that the surrogate objective f
Revised: z such that the surrogate objective f

Original: ( (cid:98) P z ; θ ) is differentiable w.r.t.
Revised: ( (cid:98) P z ; θ ) is differentiable w.r.t.

Original: the scores s .
Revised: the scores s .

Original: 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.
Revised: 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.

Original: For gradient-based optimization, we are interested in two key properties:
Revised: For gradient-based optimization, we are interested in two key properties:

Original: The relaxation is continuous everywhere and differentiable (almost-)everywhere with respect to elements in the input domain.
Revised: The relaxation is continuous everywhere and differentiable (almost-)everywhere with respect to elements in the input domain.

Original: There exists a computationally efﬁcient projection from M back to N .
Revised: There exists a computationally efﬁcient projection from M back to N .

Original: Relaxations satisfying the ﬁrst requirement are amenable to automatic differentiation for optimizing stochastic computational graphs.
Revised: Relaxations satisfying the ﬁrst requirement are amenable to automatic differentiation for optimizing stochastic computational graphs.

Original: 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.
Revised: 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.

Original: 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.
Revised: E.g. , in straight-through gradient estimation (Bengio et al., 2013; Jang et al., 2017), the /9 / 16 5 / 16 1 /    P D

Revised: R

Revised: U 3 / 8 1 / 8 1 /3 / 4 1 / 41 / 4 1 / 2 1 /    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.

Original: The canonical example is the 0 / 1 loss used for binary classiﬁcation.
Revised: The canonical example is the 0 / 1 loss used for binary classiﬁcation.

Original: While the 0 / 1 loss is discontinuous w.r.t.
Revised: While the 0 / 1 loss is discontinuous w.r.t.

Original: 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).
Revised: 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).

Original: Note: For brevity, we assume that the arg max operator is applied over a set of elements with a unique maximizer and hence, the operator has well-deﬁned semantics.
Revised: Note: For brevity, we assume that the arg max operator is applied over a set of elements with a unique maximizer and hence, the operator has well-deﬁned semantics.

Original: With some additional bookkeeping for resolving ties, the results in this section hold even if the elements to be sorted are not unique.
Revised: With some additional bookkeeping for resolving ties, the results in this section hold even if the elements to be sorted are not unique.

Original: See Appendix C.
Revised: See Appendix C.

Original: Unimodal Row Stochastic Matrices.
Revised: Unimodal Row Stochastic Matrices.

Original: The sort operator maps the input vector to a permutation, or equivalently a permutation matrix.
Revised: The sort operator maps the input vector to a permutation, or equivalently a permutation matrix.

Original: Our relaxation to sort is motivated by the geometric structure of permutation matrices.
Revised: Our relaxation to sort is motivated by the geometric structure of permutation matrices.

Original: 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 /9 / 16 5 / 16 1 /   
Revised: 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.

Original: P D R U 3 / 8 1 / 8 1 /3 / 4 1 / 41 / 4 1 / 2 1 /    the requirement that every column should sum to one, we obtain a larger set of row stochastic matrices .
Revised: If we remove the requirement that every column should sum to one, we obtain a larger set of row stochastic matrices .

Original: 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 .
Revised: 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 .

Original: Deﬁnition 1 (Unimodal Row Stochastic Matrices) .
Revised: Deﬁnition 1 (Unimodal Row Stochastic Matrices) .

Original: An n × n matrix is Unimodal Row Stochastic if it satisﬁes the following conditions:
Revised: An n × n matrix is Unimodal Row Stochastic if it satisﬁes the following conditions:

Original: Non-negativity: U
Revised: Non-negativity: U

Original: [ i, j ] ≥ 0 ∀ i, j ∈ { 1 , 2 , . . . , n } .
Revised: [ i, j ] ≥ 0 ∀ i, j ∈ { 1 , 2 , . . . , n } .

Original: Row Afﬁnity: (cid:80) n j =1 U [ i, j ] = 1 ∀ i ∈ { 1 , 2 , . . . , n } . 3.
Revised: Row Afﬁnity: (cid:80) nj =1 U [ i, j ] = 1 ∀ i ∈ { 1 , 2 , . . . , n } . 3.

Original: Argmax Permutation: Let u denote an n -dimensional vector with entries such that u i = arg max j U [ i, j ] ∀ i ∈ { 1 , 2 , . . . , n } .
Revised: Argmax Permutation: Let u denote an n -dimensional vector with entries such that u i = arg max j U [ i, j ] ∀ i ∈ { 1 , 2 , . . . , n } .

Original: Then, u ∈ Z n , i e., it is a valid permuation.
Revised: Then, u ∈ Z n , i.e., it is a valid permuation.

Original: We denote U n as the set of n × nunimodal row stochastic matrices.
Revised: We denote U n as the set of n × n

Original:  All row stochastic matrices satisfy the ﬁrst two conditions.
Revised: unimodal row stochastic matrices. All row stochastic matrices satisfy the ﬁrst two conditions.

Original: The third condition is useful for gradient based optimization involving sorting-based losses.
Revised: The third condition is useful for gradient based optimization involving sorting-based losses.

Original: The condition provides a straightforward mechanism for extracting a permutation from a unimodal row stochastic 
Revised: The condition provides a straightforward mechanism for extracting a permutation from a unimodal row stochastic matrix via a row-wise arg max operation.

Revised: Figure 2 shows the relationships between the different subsets of square matrices.

Original: Sorting Relaxation.
Revised: NeuralSort.

Original: 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.
Revised: 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.

Original: Lemma 2.
Revised: Lemma 2.

Original: [Lemma 1 in Ogryczak & Tamir (2003)]
Revised: [Lemma 1 in Ogryczak & Tamir (2003)]

Original: For an input vector s =
Revised: For an input vector s =

Original: [ s 1 , s 2 , . . .
Revised: [ s 1 , s 2 , . . .

Original: , s n ] T that is sorted as s [1] ≥ s [2] ≥ . .
Revised: , s n ] T that is sorted as s [1] ≥ s [2] ≥ . .

Original: ≥ s [ n ] , we have the sum of the k -largest elements given as:
Revised: ≥ s [ n ] , we have the sum of the k -largest elements given as:

Original: [Equation]
Revised: [Equation]

Original: The identity in Lemma 2 outputs the sum of the top- k elements.
Revised: The identity in Lemma 2 outputs the sum of the top- k elements.

Original: 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. 
Revised: 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. [Figure]

Original: Corollary 3.
Revised: Corollary 3.

Original: Let s =
Revised: Let s =

Original: [ s 1 , s 2 , . . .
Revised: [ s 1 , s 2 , . . .

Original: , s n ] T be a real-valued vector of length n .
Revised: , s n ] T be a real-valued vector of length n .

Original: Let A s denote the matrix of absolute pairwise differences of the elements of s such that A s [ i, j ]
Revised: Let A s denote the matrix of absolute pairwise differences of the elements of s such that A s [ i, j ]

Original: = | s i − s j | .
Revised: = | s i − s j | .

Original: The permutation matrix P sort ( s ) corresponding to sort ( s ) is given by:
Revised: The permutation matrix P sort ( s ) corresponding to sort ( s ) is given by:

Original: [Equation]
Revised: [Equation]

Original: where 1 denotes the column vector of all ones.
Revised: where 1 denotes the column vector of all ones.

Original: E.g. , if we set i = (cid:98) ( n + 1) / 2 (cid:99) then the non-zero entry in the i -th row P sort ( s ) [ i, :] corresponds to the element with the minimum sum of (absolute) distance to the other elements.
Revised: E.g. , if we set i = (cid:98) ( n + 1) / 2 (cid:99) then the non-zero entry in the i -th row P sort ( s ) [ i, :] corresponds to the element with the minimum sum of (absolute) distance to the other elements.

Original: As desired, this corresponds to the median element.
Revised: As desired, this corresponds to the median element.

Original: The relaxation requires O ( n 2 ) operations to compute A s , as opposed to the O ( n log n ) overall complexity for the best known sorting algorithms.
Revised: The relaxation requires O ( n 2 ) operations to compute A s , as opposed to the O ( n log n ) overall complexity for the best known sorting algorithms.

Original: In practice however, it is highly parallelizable and can be implemented efﬁciently on GPU hardware.
Revised: In practice however, it is highly parallelizable and can be implemented efﬁciently on GPU hardware.

Original: The arg max operator is non-differentiable which prohibits the direct use of Corollary 3 for gradient computation.
Revised: The arg max operator is non-differentiable which prohibits the direct use of Corollary 3 for gradient computation.

Original: Instead, we propose to replace the arg max operator with soft max to obtain a continuous relaxation (cid:98) P sort ( s ) ( τ ) .
Revised: Instead, we propose to replace the arg max operator with soft max to obtain a continuous relaxation (cid:98) P sort ( s ) ( τ ) .

Original: In particular, the i -th row of (cid:98) P sort ( s ) ( τ ) is given by:
Revised: In particular, the i -th row of (cid:98) P sort ( s ) ( τ ) is given by:

Original: [Equation]
Revised: [Equation]

Original: [Figure]

Original: where τ > 0 is a temperature parameter.
Revised: where τ > 0 is a temperature parameter.

Original: Our relaxation is continuous everywhere and differentiable almost everywhere with respect to the elements of s .
Revised: Our relaxation is continuous everywhere and differentiable almost everywhere with respect to the elements of s .

Original: Furthermore, we have the following result.
Revised: Furthermore, we have the following result.

Original: Theorem 4.
Revised: Theorem 4.

Original: Let (cid:98) P sort ( s ) denote the continuous relaxation to the permutation matrix
Revised: Let (cid:98) P sort ( s ) denote the continuous relaxation to the permutation matrix

Original: P sort ( s ) for an arbitrary input vector s and temperature τ .
Revised: P sort ( s ) for an arbitrary input vector s and temperature τ deﬁned in Eq.

Original: Then, we have:
Revised: Then, we have:

Original: Unimodality: ∀ τ > 0 , (cid:98) P sort ( s ) is a unimodal row stochastic matrix.
Revised: Unimodality: ∀ τ > 0 , (cid:98) P sort ( s ) is a unimodal row stochastic matrix.

Original: Further, let u denote the permutation obtained by applying arg max row-wise to (cid:98) P sort ( s ) .
Revised: Further, let u denote the permutation obtained by applying arg max row-wise to (cid:98) P sort ( s ) .

Original: Then, u = sort ( s ) .
Revised: Then, u = sort ( s ) .

Original: Limiting behavior: If we assume that the entries of s are drawn independently from a distribution that is absolutely continuous w.r
Revised: Limiting behavior: If we assume that the entries of s are drawn independently from a distribution that is absolutely continuous w.r.t.

Original: t. the Lebesgue measure in R , then the following convergence holds almost surely:
Revised: the Lebesgue measure in R , then the following convergence holds almost surely:

Original: [Equation]
Revised: [Equation]

Original: Unimodality allows for efﬁcient projection of the relaxed permutation matrix (cid:98) P sort ( s ) to the hard matrix P sort ( s ) via a row-wise arg max , e.g. , for straight-through gradients.
Revised: Unimodality allows for efﬁcient projection of the relaxed permutation matrix (cid:98) P sort ( s ) to the hard matrix P sort ( s ) via a row-wise arg max , e.g. , for straight-through gradients.

Original: For analyzing limiting behavior, independent draws ensure that the elements of s are distinct almost surely.
Revised: For analyzing limiting behavior, independent draws ensure that the elements of s are distinct almost surely.

Original: The temperature τ controls the degree of smoothness of our approximation.
Revised: The temperature τ controls the degree of smoothness of our approximation.

Original: At one extreme, the approximation becomes tighter as the temperature is reduced.
Revised: At one extreme, the approximation becomes tighter as the temperature is reduced.

Original: In practice however, the trade-off is in the variance of these estimates, which is typically lower for larger temperatures.
Revised: In practice however, the trade-off is in the variance of these estimates, which is typically lower for larger temperatures.

Original:  S TOCHASTIC O PTIMIZATION OVER P ERMUTATIONS
Revised:  S TOCHASTIC O PTIMIZATION OVER P ERMUTATIONS

Original: 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.
Revised: 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.

Original: Random variables that assume values corresponding to permutations can be represented via stochastic nodes in the stochastic computation graph.
Revised: Random variables that assume values corresponding to permutations can be represented via stochastic nodes in the stochastic computation graph.

Original: For optimizing the parameters of such a graph, consider the following class of objectives:
Revised: For optimizing the parameters of such a graph, consider the following class of objectives:

Original: [Equation]
Revised: [Equation]

Original: where θ and s denote sets of parameters, P z is the permutation matrix corresponding to the permutation z , q ( · ) is a parameterized distribution over the elements of the symmetric group Z n , and f ( · ) is an arbitrary function of interest assumed to be differentiable in θ and z .
Revised: where θ and s denote sets of parameters, P z is the permutation matrix corresponding to the permutation z , q ( · ) is a parameterized distribution over the elements of the symmetric group Z n , and f ( · ) is an arbitrary function of interest assumed to be differentiable in θ and z .

Original: The SCG is shown in
Revised: The SCG is shown in

Original: Figure 3a.
Revised: Figure 3a.

Original: 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.
Revised: 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.

Original: While such objectives are typically intractable to evaluate exactly since they require summing over a combinatorially large set, we can obtain unbiased estimates efﬁciently via Monte Carlo.
Revised: While such objectives are typically intractable to evaluate exactly since they require summing over a combinatorially large set, we can obtain unbiased estimates efﬁciently via Monte Carlo.

Original: Monte
Revised: Monte

Original: Carlo estimates of gradients w.r.t.
Revised: Carlo estimates of gradients w.r.t.

Original: θ can be derived simply via linearity of expectation.
Revised: θ can be derived simply via linearity of expectation.

Original: However, the gradient estimates w.r.t.
Revised: However, the gradient estimates w.r.t.

Original: s cannot be obtained directly since the sampling distribution depends on s .
Revised: s cannot be obtained directly since the sampling distribution depends on s .

Original: The REINFORCE gradient estimator (Glynn, 1990; Williams, 1992; Fu, 2006) uses the fact that ∇ s q ( z | s ) =
Revised: The REINFORCE gradient estimator (Glynn, 1990; Williams, 1992; Fu, 2006) uses the fact that ∇ s q ( z | s ) =

Original: q ( z | s ) ∇ s log q ( z | s ) to derive the following Monte Carlo gradient estimates:[Equation]
Revised: q ( z | s ) ∇ s log q ( z | s ) to derive the following Monte Carlo gradient estimates:

Original: [Equation]
Revised: [Equation]

Original: 4. R EPARAMETERIZED GRADIENT ESTIMATORS FOR PL DISTRIBUTIONS
Revised: 4. R EPARAMETERIZED GRADIENT ESTIMATORS FOR PL DISTRIBUTIONS

Original: REINFORCE gradient estimators typically suffer from high variance (Schulman et al., 2015;
Revised: REINFORCE gradient estimators typically suffer from high variance (Schulman et al., 2015;

Original: Glasserman, 2013).
Revised: Glasserman, 2013).

Original: Reparameterized samplers provide an alternate gradient estimator by expressing samples from a distribution as a deterministic function of its parameters and a ﬁxed source of randomness (Kingma & Welling, 2014; Rezende et al., 2014; Titsias & L´azaro-Gredilla, 2014).
Revised: Reparameterized samplers provide an alternate gradient estimator by expressing samples from a distribution as a deterministic function of its parameters and a ﬁxed source of randomness (Kingma & Welling, 2014; Rezende et al., 2014; Titsias & L ´ azaro-Gredilla, 2014).

Original: Since the randomness is from a ﬁxed distribution, Monte Carlo gradient estimates can be derived by pushing the gradient operator inside the expectation (via linearity).
Revised: Since the randomness is from a ﬁxed distribution, Monte Carlo gradient estimates can be derived by pushing the gradient operator inside the expectation (via linearity).

Original: In this section, we will derive a reparameterized sampler and gradient estimator for the Plackett-Luce (PL) family of distributions.
Revised: In this section, we will derive a reparameterized sampler and gradient estimator for the Plackett-Luce (PL) family of distributions.

Original: Let the score s i for an item i ∈ { 1 , 2 , . . . , n } be an unobserved random variable drawn from some underlying score distribution (Thurstone, 1927).
Revised: Let the score s i for an item i ∈ { 1 , 2 , . . . , n } be an unobserved random variable drawn from some underlying score distribution (Thurstone, 1927).

Original: Now for each item, we draw a score from its corresponding score distribution.
Revised: Now for each item, we draw a score from its corresponding score distribution.

Original: Next, we generate a permutation by applying the deterministic sort operator to these n randomly sampled scores.
Revised: Next, we generate a permutation by applying the deterministic sort operator to these n randomly sampled scores.

Original: 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.
Revised: 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.

Original: Proposition 5.
Revised: Proposition 5.

Original: [adapted from Yellott Jr (1977)]
Revised: [adapted from Yellott Jr (1977)]

Original: Let s be a vector of scores for the n items.
Revised: Let s be a vector of scores for the n items.

Original: For each item i , sample g i ∼ Gumbel (0 , β ) independently with zero mean and a ﬁxed scale β .
Revised: For each item i , sample g i ∼ Gumbel (0 , β ) independently with zero mean and a ﬁxed scale β .

Original: Let ˜s denote the vector of Gumbel perturbed log-scores with entries such that ˜ s i =
Revised: Let ˜s denote the vector of Gumbel perturbed log-scores with entries such that ˜ s i =

Original: β log s i
Revised: β log s i

Original: + g i .
Revised: + g i .

Original: Then:
Revised: Then:

Original: [Equation]
Revised: [Equation]

Original: For ease of presentation, we assume β = 1 in the rest of this work.
Revised: For ease of presentation, we assume β = 1 in the rest of this work.

Original: 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.
Revised: 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.

Original: 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 ﬁxed source of randomness (Figure 3b).
Revised: 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 ﬁxed source of randomness (Figure 3b).

Original: Letting g denote the vector of i.i.d.
Revised: Letting g denote the vector of i.i.d.

Original: Gumbel perturbations, we can express the objective in Eq.
Revised: Gumbel perturbations, we can express the objective in Eq.

Original: 7 as:
Revised: 7 as:

Original: [Equation]
Revised: [Equation]

Original: 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 .
Revised: 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 .

Original: In order to obtain a differentiable surrogate, we approximate the objective based on the continuous relaxation to the sort operator proposed previously:
Revised: In order to obtain a differentiable surrogate, we approximate the objective based on the NeuralSort relaxation to the sort operator:

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: Accordingly, we get the following reparameterized gradient estimates for the approximation:
Revised: Accordingly, we get the following reparameterized gradient estimates for the approximation:

Original: [Equation]
Revised: [Equation]

Original: 12) which can be estimated efﬁciently via Monte Carlo because the expectation is with respect to a distribution that does not depend on s .
Revised: (12) which can be estimated efﬁciently via Monte Carlo because the expectation is with respect to a distribution that does not depend on s .

Original:  D ISCUSSION AND R ELATED W ORK
Revised:  D ISCUSSION AND R ELATED W ORK

Original: The problem of learning to rank documents based on relevances has been studied extensively in the context of information retrieval.
Revised: The problem of learning to rank documents based on relevance has been studied extensively in the context of information retrieval.

Original: In particular,the listwise approaches learn functions that map objects to scores.
Revised: In particular, listwise approaches learn functions that map objects to scores.

Original: 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.
Revised: 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.

Original: (2008) extends this with a loss that maximizes the PL likelihood of ground-truth permutations directly.
Revised: (2008) extends this with a loss that maximizes the PL likelihood of ground-truth permutations directly.

Original: The differentiable pairwise approaches to ranking, such as Rigutini et al.
Revised: The differentiable pairwise approaches to ranking, such as Rigutini et al.

Original: (2011), learn to approximate the comparator between pairs of objects.
Revised: (2011), learn to approximate the comparator between pairs of objects.

Original: Our work considers a generalized setting where sorting based operators can be inserted anywhere in computation graphs to extend traditionalalgorithmic pipelines e.g. , kNN.
Revised: Our work considers a generalized setting where sorting based operators can be inserted anywhere in computation graphs to extend traditional pipelines e.g. , kNN.

Original: Recent works have proposed relaxations of permutation matrices to the Birkhoff polytope, which is deﬁned as the  set of doubly stochastic matrices (Adams & Zemel, 2011; Mena et al., 2018; Linderman et al., 2018).
Revised: Prior works have proposed relaxations of permutation matrices to the Birkhoff polytope, which is deﬁned 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;

Revised: Fiori et al., 2013; Lim & Wright, 2014).

Original: Adams & Zemel (2011) proposed the use of the Sinkhorn operator to map any square matrix to the Birkhoff polytope.
Revised: Adams & Zemel (2011) proposed the use of the Sinkhorn operator to map any square matrix to the Birkhoff polytope.

Original: They interpret the resulting doubly-stochastic matrix as the marginals of a distribution over permutations.
Revised: They interpret the resulting doubly-stochastic matrix as the marginals of a distribution over permutations.

Original: Mena et al.
Revised: Mena et al.

Original: (2018) propose an alternate method where the square matrix deﬁnes a latent distribution over the doubly-stochastic matrices themselves.
Revised: (2018) propose an alternate method where the square matrix deﬁnes a latent distribution over the doubly-stochastic matrices themselves.

Original: These distributions can be sampled from by adding elementwise Gumbel perturbations.
Revised: These distributions can be sampled from by adding elementwise Gumbel perturbations.

Original: Linderman et al.
Revised: Linderman et al.

Original: (2018) propose a rounding procedure that uses the Sinkhorn operator to directly sample matrices xxxxxs i =
Revised: (2018) propose a rounding procedure that uses the Sinkhorn operator to directly sample matrices near the

Original: CNN ( x i ) sssssz =

Original: Sortnet ( s )

Original: Task 1: Sorting Loss( [3 , 5 , 1 , 4 , 2] T , z )

Original: ˆ y

Original: = CNN (ˆ x ) Task 2: Median Regression Loss( 2960 , ˆ y ) ˆ x = x z [3] near the Birkhoff polytope.
Revised: Birkhoff polytope.

Original: Unlike Mena et al.
Revised: Unlike Mena et al.

Original: (2018), the resulting distribution over matrices has a tractable density.
Revised: (2018), the resulting distribution over matrices has a tractable density.

Original: In practice, however, the approach of Mena et al.
Revised: In practice, however, the approach of Mena et al.

Original: (2018) performs better and will be the main baseline we will be comparing against in our experiments in Section 6.
Revised: (2018) performs better and will be the main baseline we will be comparing against in our experiments in Section 6.

Original: As discussed in Section 3, our relaxation maps permutation matrices to the set of unimodal row stochastic matrices.
Revised: As discussed in Section 3, NeuralSort maps permutation matrices to the set of unimodal rowstochastic matrices.

Original: For the stochastic setting, the PL distribution permits efﬁcient sampling, exact and tractable density estimation, making it an attractive choice for several applications, e.g. , variational inference over latent permutations.
Revised: For the stochastic setting, the PL distribution permits efﬁcient sampling, exact and tractable density estimation, making it an attractive choice for several applications, e.g. , variational inference over latent permutations.

Original: Our reparameterizable sampler, while also making use of the Gumbel distribution, is based on a result unique to the PL distribution (Proposition 5).
Revised: Our reparameterizable sampler, while also making use of the Gumbel distribution, is based on a result unique to the PL distribution (Proposition 5).

Original: The use of the Gumbel distribution for deﬁning continuous relaxations to discrete distributions was ﬁrst proposed concurrently by Jang et al.
Revised: The use of the Gumbel distribution for deﬁning continuous relaxations to discrete distributions was ﬁrst proposed concurrently by Jang et al.

Original: (2017) and Maddison et al.
Revised: (2017) and Maddison et al.

Original: (2017) for categorical variables, referred to as Gumbel-Softmax.
Revised: (2017) for categorical variables, referred to as Gumbel-Softmax.

Original: 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 !
Revised: 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 !

Original: categories.
Revised: categories.

Original: Gumbel-softmax does not scale to a combinatorially large number of categories, necessitating the use of alternate relaxations, such as the one considered in this work.
Revised: 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.

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: We refer to the two approaches proposed in Sections 3, 4 as Deterministic Sortnet and Stochastic
Revised: We refer to the two approaches proposed in Sections 3, 4 as Deterministic NeuralSort and Stochastic

Original: Sortnet, respectively.
Revised: NeuralSort, respectively.

Original: For additional hyperparameter details and analysis, see Appendix D.
Revised: For additional hyperparameter details and analysis, see Appendix D.

Original: 6. S ORTING HANDWRITTEN NUMBERS
Revised: 6. S ORTING HANDWRITTEN NUMBERS

Original: Dataset.
Revised: Dataset.

Original: We ﬁrst create the large-MNIST dataset, which extends the MNIST dataset of handwritten digits.
Revised: We ﬁrst create the large-MNIST dataset, which extends the MNIST dataset of handwritten digits.

Original: 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.
Revised: 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.

Original: Each image is associated with a real-valued label, which corresponds to its concatenated MNIST labels, e.g. , the label of is 1810 .
Revised: Each image is associated with a real-valued label, which corresponds to its concatenated MNIST labels, e.g. , the label of is 1810 .

Original: Using the large-MNIST dataset, we ﬁnally create a dataset of sequences.
Revised: Using the large-MNIST dataset, we ﬁnally create a dataset of sequences.

Original: Every sequence is this dataset consists of n randomly sampled large-MNIST images.
Revised: Every sequence is this dataset consists of [Table]

Original: Setup.
Revised: Setup.

Original: 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.
Revised: 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.

Original: Figure 4 (Task 1) illustrates this task on an example sequence of n = 5 large-MNIST images.
Revised: Figure 4 (Task 1) illustrates this task on an example sequence of n = 5 large-MNIST images.

Original: This task is a challenging extension of the one considered by Mena et al.
Revised: This task is a challenging extension of the one considered by Mena et al.

Original: (2018) in sorting scalars, since it involves learning the semantics of high-dimensional objects prior to sorting. 
Revised: (2018) in sorting scalars, since it involves learning the semantics of high-dimensional objects prior to sorting. A [Figure] =

Revised: [Figure]

Revised: ( x

Original: A good model needs to learn to dissect the individual digits in an image, rank these digits, and ﬁnally, compose such rankings based on the digit positions within an image.
Revised: [Figure] [Figure] [Figure] [Figure]

Revised: , z )

Revised: [Figure]

Revised: [Figure]

Revised: [Table]

Revised: good model needs to learn to dissect the individual digits in an image, rank these digits, and ﬁnally, compose such rankings based on the digit positions within an image.

Original: The available supervision, in the form of the ground-truth permutation, is very weak compared to a classiﬁcation setting that gives direct access to the image labels.
Revised: The available supervision, in the form of the ground-truth permutation, is very weak compared to a classiﬁcation setting that gives direct access to the image labels.

Original: Baselines.
Revised: Baselines.

Original: All baselines use a CNN that is shared across all images in a sequence to map each large-MNIST image to a feature space.
Revised: All baselines use a CNN that is shared across all images in a sequence to map each large-MNIST image to a feature space.

Original: The vanilla row stochastic (RS) baseline concatenates the
Revised: The vanilla row-stochastic (RS) baseline concatenates the

Original: CNN representations for n images into a single vector that is fed into a multilayer perceptron that outputs n multiclass predictions of the image probabilities for each rank.
Revised: CNN representations for n images into a single vector that is fed into a multilayer perceptron that outputs n multiclass predictions of the image probabilities for each rank.

Original: The Sinkhorn and GumbelSinkhorn baselines, as discussed in Section 5, use the Sinkhorn operator to map the stacked CNN[Table] [Table] representations for the n objects into a doubly stochastic matrix.
Revised: 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.

Original: For all methods, we minimized the cross-entropy loss between the predicted matrix and the ground-truth permutation matrix.
Revised: For all methods, we minimized the cross-entropy loss between the predicted matrix and the ground-truth permutation matrix.

Original: Results.
Revised: Results.

Original: Following Mena et al.
Revised: Following Mena et al.

Original: (2018), our evaluation metric is the the proportion of correctly predicted permutations on a test set of sequences.
Revised: (2018), our evaluation metric is the the proportion of correctly predicted permutations on a test set of sequences.

Original: Additionally, we evaluate the proportion of individual elements ranked correctly.
Revised: Additionally, we evaluate the proportion of individual elements ranked correctly.

Original: Table 1 demonstrates that the approaches based on the proposed sorting relaxation signiﬁcantly outperform the baseline approaches for all n considered.
Revised: Table 1 demonstrates that the approaches based on the proposed sorting relaxation signiﬁcantly outperform the baseline approaches for all n considered.

Original: The performance of the deterministic and stochastic variants are comparable.
Revised: The performance of the deterministic and stochastic variants are comparable.

Original: The vanilla RS baseline performs well in ranking individual elements, but is not good at recovering the overall square matrix.
Revised: The vanilla RS baseline performs well in ranking individual elements, but is not good at recovering the overall square matrix.

Original: We believe the poor performance of the Sinkhorn baselines is partly because these methods were designed and evaluated for matchings .
Revised: We believe the poor performance of the Sinkhorn baselines is partly because these methods were designed and evaluated for matchings .

Original: Like the output of sort , matchings can also be represented as permutation matrices.
Revised: Like the output of sort , matchings can also be represented as permutation matrices.

Original: 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.
Revised: 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.

Original: 6. Q UANTILE REGRESSION
Revised: 6. Q UANTILE REGRESSION

Original: Setup.
Revised: Setup.

Original: In this experiment, we extend the sorting task to regression.
Revised: In this experiment, we extend the sorting task to regression.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: 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).
Revised: 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).

Original: 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. 
Revised: 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. [Table] [Figure] [Table]

Revised: [Figure]

Revised: [Table]

Original: Baselines.
Revised: Baselines.

Original: In addition to Sinkhorn and Gumbel-Sinkhorn , we design two more baselines.
Revised: In addition to Sinkhorn and Gumbel-Sinkhorn , we design two more baselines.

Original: The
Revised: The

Original: Constant baseline always returns the median of the full range of possible outputs, ignoring the input sequence.
Revised: Constant baseline always returns the median of the full range of possible outputs, ignoring the input sequence.

Original: This corresponds to 4999 .
Revised: This corresponds to 4999 .

Original: 5 since we are sampling large-MNIST images uniformly in the range of four-digit numbers.
Revised: 5 since we are sampling large-MNIST images uniformly in the range of four-digit numbers.

Original: The vanilla neural net (NN) baseline directly maps the input sequence of images to a real-valued prediction for the median.
Revised: The vanilla neural net (NN) baseline directly maps the input sequence of images to a real-valued prediction for the median.

Original: Results.
Revised: Results.

Original: Our evaluation metric is the mean squared error (MSE) and R 2 on a test set of sequences.
Revised: Our evaluation metric is the mean squared error (MSE) and R 2 on a test set of sequences.

Original: Results for n = { 5 , 9 , 15 } images are shown in Table 2.
Revised: Results for n = { 5 , 9 , 15 } images are shown in Table 2.

Original: The Vanilla NN baseline while incurring a large MSE, is competitive on the R 2 metric.
Revised: The Vanilla NN baseline while incurring a large MSE, is competitive on the R 2 metric.

Original: The other baselines give comparable performance on the MSE metric.
Revised: The other baselines give comparable performance on the MSE metric.

Original: The proposed sortnet approaches outperform the competing methods on both the metrics considered.
Revised: The proposed NeuralSort approaches outperform the competing methods on both the metrics considered.

Original: The stochastic sortnet approach is the consistent best performer on MSE, while the deterministic sortnet is slightly better on the R 2 metric.[Table]
Revised: The stochastic NeuralSort approach is the consistent best performer on MSE, while the deterministic NeuralSort is slightly better on the R 2 metric.

Original: 6. E ND - TO - END , DIFFERENTIABLE k -N EAREST N EIGHBORS
Revised: 6. E ND - TO - END , DIFFERENTIABLE k -N EAREST N EIGHBORS

Original: Setup.
Revised: Setup.

Original: In this experiment, we design a fully differentiable, end-to-end k -nearest neighbors (kNN) classiﬁer.
Revised: In this experiment, we design a fully differentiable, end-to-end k -nearest neighbors (kNN) classiﬁer.

Original: Unlike a standard kNN classiﬁer which computes distances between points in a predeﬁned space, we learn a representation of the data points before evaluating the k -nearest neighbors.
Revised: Unlike a standard kNN classiﬁer which computes distances between points in a predeﬁned space, we learn a representation of the data points before evaluating the k -nearest neighbors.

Revised: We are given access to a dataset D of ( x , y ) pairs of standard input data and their class labels respectively.

Revised: The differentiable kNN algorithm consists of two hyperparameters: the number of training neighbors n , the number of top candidates k , and the sorting temperature τ .

Original: Every sequence of items here consists of a query point  and a randomly sampled subset of n candidate nearest neighbors from the training set.
Revised: 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 { x 1 , x 2 , . . . , x n } .

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: The loss function optimizes for a representation space  such that the top- k candidate points (with the minimum Euclidean distance to the query point) have the same label as the query point.
Revised: 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.

Revised: 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 classiﬁer.

Original: Figure 5 illustrates the proposed algorithm.
Revised: Figure 5 illustrates the proposed algorithm.

Revised: Formally, for any datapoint x , let z denote a permutation of the n candidate points.

Revised: The uniformlyweighted kNN loss, denoted as (cid:96) kNN ( · ) , can be written as follows:

Original: Datasets.
Revised: [Equation] [Table] where { y 1 , y 2 , . . . , y n } are the labels for the candidate points.

Revised: Note that when (cid:98) 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 .

Revised: Using Eq. 13, the training objectives for Deterministic and Stochastic NeuralSort are given as:

Revised: Deterministic:

Revised: [Equation]

Revised: 15) where each entry of s is given as s j = −(cid:107)

Revised: [Equation]

Revised: x

Revised: [Equation]

Revised: − h

Revised: [Equation]

Original:  We consider the benchmark  MNIST dataset of handwritten digits and the CIFARdataset of natural images (no data augmentation) with the canonical splits for training and testing.
Revised: x [Equation] 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.

Original: Baselines.
Revised: Baselines.

Original: We consider kNN baselines that operate in three standard representation spaces: the canonical pixel basis, the basis speciﬁed by the top 50 principal components (PCA), a convolutional autonencoder (CAE).
Revised: We consider kNN baselines that operate in three standard representation spaces: the canonical pixel basis, the basis speciﬁed by the top 50 principal components (PCA), an autonencoder (AE).

Original: 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.
Revised: 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.

Original: For completeness, we trained a CNNusing the crossentropy loss with the same architecture as the one used for sortnet (except the ﬁnal layer).
Revised: For completeness, we trained a CNN with the same architecture as the one used for NeuralSort (except the ﬁnal layer) using the cross-entropy loss.

Original: Results.
Revised: Results.

Original: We report the classiﬁcation accuracies on the standard test sets in Table 3.
Revised: We report the classiﬁcation accuracies on the standard test sets in Table 3.

Original: On both datasets, the differentiable kNN classiﬁer outperforms all the baseline kNN variants including the convolutional autoencoder approach.
Revised: On both datasets, the differentiable kNN classiﬁer outperforms all the baseline kNN variants including the convolutional autoencoder approach.

Original: The performance is much closer to the accuracy of a standard CNN.
Revised: The performance is much closer to the accuracy of a standard CNN.

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: In this paper, we propose a continuous relaxation of the sorting operator to the set of unimodal row stochastic matrices.
Revised: In this paper, we proposed NeuralSort, a continuous relaxation of the sorting operator to the set of unimodal row-stochastic matrices.

Original: Our relaxation facilitates gradient estimation on stochastic computation graphs and can be extended to include stochastic nodes corresponding to distributions over permutations.
Revised: Our relaxation facilitates gradient estimation on any computation graph involving a sort operator.

Original: Further, we derived a reparameterized gradient estimator for the Plackett-Luce distribution for efﬁcient stochastic optimization.
Revised: Further, we derived a reparameterized gradient estimator for the Plackett-Luce distribution for efﬁcient stochastic optimization over permutations.

Original: On three illustrative tasks, our proposed relaxations outperform prior work in end-to-end learning of semantic orderings of high-dimensional objects.
Revised: 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.

Original: 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).
Revised: 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).

Original: Both our relaxed sort operator and the reparameterizable sampler are easy to implement.
Revised: Both deterministic and stochastic NeuralSort are easy to implement.

Original: We provide reference implementations in Tensorﬂow (Abadi et al., 2016) and PyTorch (Paszke et al., 2017) in Appendix A.
Revised: We provide reference implementations in Tensorﬂow (Abadi et al., 2016) and PyTorch (Paszke et al., 2017) in Appendix A.

Revised: The full codebase for this work is open-sourced at https://github.com/ermongroup/neuralsort .

Original: A PPENDICES A S ORTING O PERATOR A. T ENSORFLOW
Revised: A CKNOWLEDGEMENTS

Revised: This research was supported by NSF (#1651565, #1522054, #1733686), ONR, AFOSR (FA955019-1-0024), and FLI. AG is supported by MSR Ph.D. fellowship and Stanford Data Science scholarship. We are thankful to Jordan Alexander, Kristy Choi, Adithya Ganesh, Karan Goel, Neal Jean,

Revised: Daniel Levy, Jiaming Song, Yang Song, Serena Yeung, and Hugh Zhang for helpful comments.

Revised: A PPENDICES A S ORTING O PERATOR A. T ENSORFLOW

Original: Sorting Relaxation for Deterministic Sortnet: import tensorflow as tf
Revised: Sorting Relaxation for Deterministic NeuralSort: import tensorflow as tf

Original: [Table]
Revised: [Table]

Original: [Table] [Table][Table]
Revised: [Table] [Table]

Original: [Table] [Table][Table]
Revised: [Table] [Table]

Original: [Table] [Table]
Revised: [Table] [Table]

Original: return P_hat
Revised: return P hat

Original: Reparameterized Sampler for Stochastic Sortnet:
Revised: Reparameterized Sampler for Stochastic NeuralSort:

Original: [Table]
Revised: [Table]

Original: """ batch_
Revised: [Table] [Table] [Table] [Table]

Revised: [Table]

Original: [Table]
Revised: [Table]

Original: = tf

Original: [Table]
Revised: [Table]

Original: [0] n

Original: [Table]
Revised: [Table]

Original: s
Revised: [Table]

Original: [Table]
Revised: [Table]

Original: [1] log_s_perturb = s + sample_gumbel([n_samples, batch_size, n, 1])
Revised: log_s_perturb = s + sample_gumbel([n_samples, batch_ return P hat

Original: A
Revised: A.

Original: P
Revised: P

Original: [Table]
Revised: [Table]

Original: ORCH
Revised: ORCH

Original: Sorting Relaxation for Deterministic 
Revised: Sorting Relaxation for Deterministic NeuralSort

Original: Sortnet
Revised: [Table]

Original: : import torch
Revised: import torch

Original: [Table]
Revised: [Table]

Original: Reparamterized Sampler for Stochastic Sortnet:
Revised: Reparamterized Sampler for Stochastic NeuralSort:

Original: [Table]
Revised: [Table]

Original: _hat P hat.view( samples batch
Revised: 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)

Original: ) return P hat
Revised: P_ [Table] = deterministic_\name{}(log

Original: B P ROOFS OF T HEORETICAL R ESULTS B. L EMMA
Revised: B P ROOFS OF T HEORETICAL R ESULTS B. L EMMA

Original: Proof.
Revised: Proof.

Original: For any value of λ , the following inequalities hold: k (cid:88) i = s
Revised: For any value of λ , the following inequalities hold: k (cid:88) i = s

Original: [ i ] = λk + k (cid:88)
Revised: [ i ] = λk + k (cid:88)

Original: i =
Revised: i =

Original: ( s [ i ] − λ ) ≤ λk + k (cid:88)
Revised: ( s [ i ] − λ ) ≤ λk + k (cid:88)

Original: i = max( s [ i ] − λ, 0) ≤ λk + n (cid:88) i
Revised: i = max( s [ i ] − λ, 0) ≤ λk + n (cid:88) i

Original: = max( s i − λ, 0) .
Revised: = max( s i − λ, 0) .

Original: Furthermore, for λ = s [ k ] : λk + n (cid:88)
Revised: Furthermore, for λ = s [ k ] : [Equation] + [Equation]

Original: i = max( s i − λ, 0) =
Revised: i = max [Equation] [Equation] =

Revised: [Equation]

Original: s [ k ] k + n (cid:88)
Revised: [Equation] ] k + n (cid:88)

Original: i = max( s i − s [ k ] , 0)
Revised: i = max( s i − s [ k ] [Equation]

Original: =
Revised: =

Original: s [ k ] k + k (cid:88)
Revised: s [ k ] k + k (cid:88)

Original: i =
Revised: i =

Original: ( s [ i ] − s [ k ] )
Revised: ( s [ i ] [Equation] s [ k ] )

Original: =
Revised: =

Original: k (cid:88) i = s
Revised: k (cid:88) i = s

Original: [ i ] .
Revised: [ i ] .

Original: This ﬁnishes the proof.
Revised: This ﬁnishes the proof.

Original: B. C OROLLARY
Revised: B. C OROLLARY

Original: Proof.
Revised: Proof.

Original: We ﬁrst consider at exactly what values of λ the sum in Lemma 2 is minimized.
Revised: We ﬁrst consider at exactly what values of λ the sum in Lemma 2 is minimized.

Original: For simplicity we will only prove the case where all values of s are distinct.
Revised: For simplicity we will only prove the case where all values of s are distinct.

Original: The equality (cid:80) ki =1 s [ i ] =
Revised: The equality (cid:80) ki =1 s [ i ] =

Original: λk + (cid:80) ni
Revised: λk + (cid:80) ni

Original: =1 max( s i − λ, 0) holds only when s [ k ] ≤ λ ≤ s [ k +1] .
Revised: =1 max( s i − λ, 0) holds only when s [ k ] ≤ λ ≤ s [ k +1] .

Original: By
Revised: By

Original: Lemma 2, these values of λ also minimize the RHS of the equality.
Revised: Lemma 2, these values of λ also minimize the RHS of the equality.

Original: Symmetrically, if one considers the score vector t = − s , then λ ( n − k + 1) + (cid:80) ni =1 max( t i − λ, 0) is minimized at t [ n − k +1] ≤ λ ≤ t [ n − k +2] .
Revised: Symmetrically, if one considers the score vector t = − s , then λ ( n − k + 1) + (cid:80) ni =1 max( t i − λ, 0) is minimized at t [ n − k +1] ≤ λ ≤ t [ n − k +2] .

Original: Replacing λ by − λ and using the deﬁnition of t implies that λ ( k − 1 − n ) + (cid:80) n i =1 max( λ − s i , 0) is minimized at s [ k − 1] ≤ λ ≤ s [ k ] .
Revised: Replacing λ by − λ and using the deﬁnition of t implies that λ ( k − 1 − n ) + (cid:80) n i =1 max( λ − s i , 0) is minimized at s [ k − 1] ≤ λ ≤ s [ k ] .

Original: It follows that:
Revised: It follows that:

Original: [Equation]
Revised: [Equation]

Original: Thus, if s i
Revised: Thus, if s i

Original: = s [ k ] , then i =
Revised: = s [ k ] , then i =

Original: arg min(2 k − 1 − n ) s
Revised: arg min(2 k − 1 − n ) s

Original: + A s 1 .
Revised: + A s 1 .

Original: This ﬁnishes the proof.
Revised: This ﬁnishes the proof.

Original: B. T HEOREM
Revised: B. T HEOREM

Original: We prove the two properties in the statement of the theorem independently:
Revised: We prove the two properties in the statement of the theorem independently:

Original: 1. Unimodality
Revised: 1. Unimodality

Original: Proof.
Revised: Proof.

Original: By deﬁnition of the softmax function, the entries ef (cid:98) P are positive and sum to 1.
Revised: By deﬁnition of the softmax function, the entries ef (cid:98) P are positive and sum to 1.

Original: To show that (cid:98) P satisﬁes the argmax permutation property, .
Revised: To show that (cid:98) P satisﬁes the argmax permutation property, .

Original: Formally, for any given row i , we construct the argmax permutation vector u as:
Revised: Formally, for any given row i , we construct the argmax permutation vector u as:

Original: [Equation]
Revised: [Equation]

Original: where the square notation [ i ] denotes the index of the i -th largest element.
Revised: where the square notation [ i ] denotes the index of the i -th largest element.

Original: The ﬁrst step follows from the fact that the softmax function is monotonically increasing and hence, it preserves the argmax.
Revised: The ﬁrst step follows from the fact that the softmax function is monotonically increasing and hence, it preserves the argmax.

Original: The second equality directly follows from Corollary 3.
Revised: The second equality directly follows from Corollary 3.

Original: By deﬁnition, sort ( s ) =
Revised: By deﬁnition, sort ( s ) =

Original: , [ n ] } , ﬁnishing the proof.
Revised: , [ n ] } , ﬁnishing the proof.

Original: 2. Limiting behavior
Revised: 2. Limiting behavior

Original: Proof.
Revised: Proof.

Original: As shown in Gao & Pavel (2017), the softmax function may be equivalently deﬁned assoft max( z/τ ) =
Revised: As shown in ? , the softmax function may be equivalently deﬁned as

Original: arg max x ∈ ∆ n − 1 (cid:104) x, z (cid:105) − τ (cid:80) ni

Original: =1 x i log x i .
Revised: [Equation]

Original: In particular,lim τ → 0 soft max( z/τ ) =
Revised: In particular,

Original: arg max x .
Revised: [Equation] x .

Original: The distributional assumptions ensure that the elements of s are distinct a.s., so plugging in z = ( n + 1 − 2 k ) s − A s 1 completes the proof.
Revised: The distributional assumptions ensure that the elements of s are distinct a.s., so plugging in z = [Equation] n + 1 − 2 k [Equation]

Revised: s − A s 1 completes the proof.

Original: B. P ROPOSITION
Revised: B. P ROPOSITION

Original: This result follows from an earlier result by Yellott Jr (1977).
Revised: This result follows from an earlier result by Yellott Jr (1977).

Original: We give the proof sketch below and refer the reader to Yellott Jr (1977) for more details.
Revised: We give the proof sketch below and refer the reader to Yellott Jr (1977) for more details.

Original: Sketch.
Revised: Sketch.

Original: Consider random variables { X i } ni =1 such that X i ∼
Revised: Consider random variables { X i } ni =1 such that X i ∼

Original: Exp ( s z i )) .
Revised: Exp ( s z i )) .

Original: We may prove by induction a generalization of the memoryless property:
Revised: We may prove by induction a generalization of the memoryless property:

Original: [Equation]
Revised: [Equation]

Original: If we assume as inductive hypothesis that q ( X 2 ≤ · · · ≤ X n | x + t ≤ min i ≥ 2 X i )
Revised: If we assume as inductive hypothesis that q ( X 2 ≤ · · · ≤ X n | x + t ≤ min i ≥ 2 X i )

Original: =
Revised: =

Original: q ( X 2 ≤ · · · ≤
Revised: q ( X 2 ≤ · · · ≤

Original: X n | t ≤ min i ≥ 2 X i )
Revised: X n | t ≤ min i ≥ 2 X i )

Original: [Equation]
Revised: [Equation]

Original: we complete the induction as:
Revised: we complete the induction as:

Original: [Equation]
Revised: [Equation]

Original: It follows from a familiar property of argmin of exponential distributions that:
Revised: It follows from a familiar property of argmin of exponential distributions that:

Original: [Equation]
Revised: [Equation]

Original: and by another induction, we have q ( X 1 ≤ · · · ≤ X n ) =
Revised: and by another induction, we have q ( X 1 ≤ · · · ≤ X n ) =

Original: (cid:81) [Equation]
Revised: (cid:81) ni =

Original: [Equation]
Revised: [Equation]

Original: Finally, following the argument of Balog et al.
Revised: Finally, following the argument of Balog et al.

Original: , we apply the strictly decreasing function g ( x ) = − β log x to this identity, which from the deﬁnition of the Gumbel distribution implies:
Revised: , we apply the strictly decreasing function g ( x ) = − β log x to this identity, which from the deﬁnition of the Gumbel distribution implies:

Original: [Equation]
Revised: [Equation]

Original: C A RG M AX SEMANTICS FOR T IED M AX E LEMENTS
Revised: C A RG M AX SEMANTICS FOR T IED M AX E LEMENTS

Original: While applying the arg max operator to a vector with duplicate entries attaining the max value, we need to deﬁne the operator semantics for arg max to handle ties in the context of the proposed relaxation.
Revised: While applying the arg max operator to a vector with duplicate entries attaining the max value, we need to deﬁne the operator semantics for arg max to handle ties in the context of the proposed relaxation.

Original: Deﬁnition 6.
Revised: Deﬁnition 6.

Original: For any vector with ties, let arg max set denote the operator that returns the set of all indices containing the max element.
Revised: For any vector with ties, let arg max set denote the operator that returns the set of all indices containing the max element.

Original: We deﬁne the arg max of the i -th in a matrix M recursively:
Revised: We deﬁne the arg max of the i -th in a matrix M recursively:

Original: If there exists an index j ∈ { 1 , 2 , . .
Revised: If there exists an index j ∈ { 1 , 2 , . .

Original: . , n } that is a member of arg max set( M [ i, :]) and has not been assigned as an arg max of any row k < i , then the arg max is the smallest such index.
Revised: . , n } that is a member of arg max set( M [ i, :]) and has not been assigned as an arg max of any row k < i , then the arg max is the smallest such index.

Original: Otherwise, the arg max is the smallest index that is a member of the arg max set( M [ i, :]) .
Revised: Otherwise, the arg max is the smallest index that is a member of the arg max set( M [ i, :]) .

Original: This function is efﬁciently computable with additional bookkeeping.
Revised: This function is efﬁciently computable with additional bookkeeping.

Original: Lemma 7.
Revised: Lemma 7.

Original: For an input vector s with the sort permutation matrix given as P sort ( s ) , we have s j 

Revised: For an input vector s with the sort permutation matrix given as P sort ( s ) , we have s j

Original: =
Revised: =

Original: s
Revised: s

Original: j 2 if
Revised: j 2 if

Original: and only if there exists a row i such that (cid:98) P [ i, j 1 ]
Revised: and only if there exists a row i such that (cid:98) P [ i, j 1 ]

Original: = (cid:98) P
Revised: = (cid:98) P

Original: [ i, j 2 ] for all j 1 , j 2 ∈ { 1 , 2 , . . . , n } .
Revised: [ i, j 2 ] for all j 1 , j 2 ∈ { 1 , 2 , . . . , n } .

Original: Proof.
Revised: Proof.

Original: From Eq.
Revised: From Eq.

Original: 5, we have the i -th row of (cid:98) P [ i, :] given as:
Revised: 5, we have the i -th row of (cid:98) P [ i, :] given as:

Original: [Equation]
Revised: [Equation]

Original: [ i, :] = soft max[
Revised: soft max

Original: [Equation]
Revised: [Equation]

Original: n
Revised: n

Original: [Equation]
Revised: [Equation]

Original: [Equation]

Original: [Equation]
Revised: [Equation]

Original: /τ ] .

Original: Therefore, we have the equations:
Revised: Therefore, we have the equations:

Original: [Equation]
Revised: [Equation]

Original:  for some ﬁxed normalization constant Z .
Revised: Z for some ﬁxed normalization constant Z .

Original: As the function f
Revised: As the function f

Original: ( x ) =
Revised: ( x ) =

Original: exp((( n +1 − 2 i ) x − ( A s 1 ) i ) /τ )
Revised: exp((( n +1 − 2 i ) x − ( A s 1 ) i ) /τ )

Original: Z is invertible, both directions of the lemma follow immediately.
Revised: Z is invertible, both directions of the lemma follow immediately.

Original: Lemma 8.
Revised: Lemma 8.

Original: If arg max set( (cid:98) P [ i 1 , :]) and arg max set( (cid:98) P [ i 2 , :]) have a non-zero intersection, then arg max set( (cid:98) P [ i 1 , :]) =
Revised: If arg max set( (cid:98) P [ i 1 , :]) and arg max set( (cid:98) P [ i 2 , :]) have a non-zero intersection, then arg max set( (cid:98) P [ i 1 , :]) =

Original: arg max set( (cid:98) P [ i 2 , :]) .
Revised: arg max set( (cid:98) P [ i 2 , :]) .

Original: Proof.
Revised: Proof.

Original: Assume without loss of generality that | arg max set( (cid:98) P [ i 1 , :]) | > 1 for some i .
Revised: Assume without loss of generality that | arg max set( (cid:98) P [ i 1 , :]) | > 1 for some i .

Original: Let j 1 , j 2 be two members of | arg max set( (cid:98) P [ i 1 , :]) | .
Revised: Let j 1 , j 2 be two members of | arg max set( (cid:98) P [ i 1 , :]) | .

Original: By Lemma 7, s j 

Revised: By Lemma 7, s j

Original: = s j 2 , and therefore (cid:98) P
Revised: = s j 2 , and therefore (cid:98) P

Original: [ i 2 , j 1 ] = (cid:98) P [ i 2 , j 2 ] .
Revised: [ i 2 , j 1 ] = (cid:98) P [ i 2 , j 2 ] .

Original: Hence if j 1 ∈ arg max set( (cid:98) P [ i 2 , :]) , then j 2 is also an element.
Revised: Hence if j 1 ∈ arg max set( (cid:98) P [ i 2 , :]) , then j 2 is also an element.

Original: A symmetric argument implies that if j 2 ∈ arg max set( (cid:98) P [ i 2 , :]) , then j 1 is also an element for arbitrary j 1 , j 2 ∈ | arg max set( (cid:98) P [ i 1 , :]) | .
Revised: A symmetric argument implies that if j 2 ∈ arg max set( (cid:98) P [ i 2 , :]) , then j 1 is also an element for arbitrary j 1 , j 2 ∈ | arg max set( (cid:98) P [ i 1 , :]) | .

Original: This completes the proof.
Revised: This completes the proof.

Original: Proposition 9.
Revised: Proposition 9.

Original: (Argmax Permutation with Ties)
Revised: (Argmax Permutation with Ties)

Original: For s =
Revised: For s =

Original: [ s 1 , s 2 , . .
Revised: [ s 1 , s 2 , . .

Original: , s n ] T ∈ R n , the vector z deﬁned by z i = arg max j (cid:98) P sort ( s )
Revised: , s n ] T ∈ R n , the vector z deﬁned by z i = arg max j (cid:98) P sort ( s )

Original: [ i, j ] is such that z ∈ Z n .
Revised: [ i, j ] is such that z ∈ Z n .

Original: [Figure]
Revised: [Figure]

Original: Proof.
Revised: Proof.

Original: From Corollary 3, we know that the row (cid:98) P sort ( s )
Revised: From Corollary 3, we know that the row (cid:98) P sort ( s )

Original: [ i, :] attains its maximum (perhaps nonuniquely) at some (cid:98) P sort ( s )
Revised: [ i, :] attains its maximum (perhaps nonuniquely) at some (cid:98) P sort ( s )

Original: [ i, j ] where s j = s
Revised: [ i, j ] where s j = s

Original: [ i ] .
Revised: [ i ] .

Original: Note that s [ i ] is well-deﬁned even in the case of ties.
Revised: Note that s [ i ] is well-deﬁned even in the case of ties.

Original: Consider an arbitrary row (cid:98) P sort ( s )
Revised: Consider an arbitrary row (cid:98) P sort ( s )

Original: [ i, :] and let arg max set( (cid:98) P sort ( s )
Revised: [ i, :] and let arg max set( (cid:98) P sort ( s )

Original: [ i, :]) =
Revised: [ i, :]) =

Original: { j 1 , . . .
Revised: { j 1 , . . .

Original: j m } .
Revised: j m } .

Original: It follows from Lemma 7 that there are exactly m scores in s that are equal to s [ i ] : s j 1 , . . .
Revised: It follows from Lemma 7 that there are exactly m scores in s that are equal to s [ i ] : s j 1 , . . .

Original: , s j m .
Revised: , s j m .

Original: These scores corresponds to m values of s [ i (cid:48) ] such that s [ i (cid:48) ]
Revised: These scores corresponds to m values of s [ i (cid:48) ] such that s [ i (cid:48) ]

Original: = s [ i ] , and consequently to m rows P [ i (cid:48) , :] that are maximized with values s [ i (cid:48) ]
Revised: = s [ i ] , and consequently to m rows P [ i (cid:48) , :] that are maximized with values s [ i (cid:48) ]

Original: = s [ i ] and consequently (by Lemma 8) at indices j 1 , . . .
Revised: = s [ i ] and consequently (by Lemma 8) at indices j 1 , . . .

Original: , j m .
Revised: , j m .

Original: Suppose we now chose an i (cid:48) such that s [ i (cid:48) ] (cid:54) =
Revised: Suppose we now chose an i (cid:48) such that s [ i (cid:48) ] (cid:54) =

Original: s [ i ] .
Revised: s [ i ] .

Original: Then (cid:98) P [ i (cid:48) , :] attains its maximum at some (cid:98) P sort ( s )
Revised: Then (cid:98) P [ i (cid:48) , :] attains its maximum at some (cid:98) P sort ( s )

Original: [ i (cid:48) , j (cid:48) ] where s j (cid:48)
Revised: [ i (cid:48) , j (cid:48) ] where s j (cid:48)

Original: =
Revised: =

Original: s [ i (cid:48) ] .
Revised: s [ i (cid:48) ] .

Original: Because s j (cid:48)
Revised: Because s j (cid:48)

Original: =
Revised: =

Original: s [ i (cid:48) ] (cid:54) =
Revised: s [ i (cid:48) ] (cid:54) =

Original: s [ i ] =
Revised: s [ i ] =

Original: s j , Lemma 7 tells us that (cid:98) P [ i (cid:48) , :] does not attain its maximum at any of j 1 , . . .
Revised: s j , Lemma 7 tells us that (cid:98) P [ i (cid:48) , :] does not attain its maximum at any of j 1 , . . .

Original: , j m .
Revised: , j m .

Original: Therefore, only m rows have a non-zero arg max set intersection with arg max set( P [ i, :]) .
Revised: Therefore, only m rows have a non-zero arg max set intersection with arg max set( P [ i, :]) .

Original: Because P [ i, :] is one of these rows, there can be up to m − 1 such rows above it.
Revised: Because P [ i, :] is one of these rows, there can be up to m − 1 such rows above it.

Original: Because each row above only has one arg max assigned via the tie-breaking protocol, it is only possible for up to m − 1 elements of arg max set( P [ i, :]) to have been an arg max of a previous row k < i .
Revised: Because each row above only has one arg max assigned via the tie-breaking protocol, it is only possible for up to m − 1 elements of arg max set( P [ i, :]) to have been an arg max of a previous row k < i .

Original: As | arg max set( P [ i, :]) | = m , there exists at least one element that has not been speciﬁed as the arg max of a previous row (pigeon-hole principle).
Revised: As | arg max set( P [ i, :]) | = m , there exists at least one element that has not been speciﬁed as the arg max of a previous row (pigeon-hole principle).

Original: Thus, the arg max of each row are distinct.
Revised: Thus, the arg max of each row are distinct.

Original: Because each argmax is also an element of { 1 , . . . , n } , it follows that z ∈ Z n .
Revised: Because each argmax is also an element of { 1 , . . . , n } , it follows that z ∈ Z n .

Original: D E XPERIMENTAL D ETAILS AND A NALYSIS
Revised: D E XPERIMENTAL D ETAILS AND A NALYSIS

Original: We used Tensorﬂow (Abadi et al., 2016)  for our experiments.
Revised: We used Tensorﬂow (Abadi et al., 2016) and PyTorch (Paszke et al., 2017) for our experiments.

Original: In Appendix A, we provide “plugin” snippets for implementing our proposed relaxations in bothTensorﬂow (Abadi et al., 2016) and
Revised: In Appendix A, we provide “plug-in” snippets for implementing our proposed relaxations in both

Original: PyTorch (Paszke et al., 2017).
Revised: Tensorﬂow and PyTorch.

Original: We will make the full codebaseused for reproducing the experiments open source upon the completion of the review process.
Revised: The full codebase for reproducing the experiments can be found at https://github.com/ermongroup/neuralsort .

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: For CIFAR-10, we used a split of 45 , 000 / 5000 / 10 , 000 examples for training/validation/test.
Revised: For CIFAR-10, we used a split of 45 , 000 / 5000 / 10 , 000 examples for training/validation/test.

Original: 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.
Revised: 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.

Original: We now report experimental details with regards to the individual experiments.

Original: D. S ORTING HANDWRITTEN NUMBERS
Revised: D. S ORTING HANDWRITTEN NUMBERS

Original: Architectures.
Revised: Architectures.

Original: We control for the choice of computer vision models by using the same convolutional network architecture for each sorting method.
Revised: We control for the choice of computer vision models by using the same convolutional network architecture for each sorting method.

Original: This architecture is as follows:
Revised: This architecture is as follows:

Original: Conv[Kernel: 5x5, Stride: 1, Output: 140x28x32, Activation: Relu]
Revised: Conv[Kernel: 5x5, Stride: 1, Output: 140x28x32, Activation: Relu]

Original: → Pool[Stride: 2, Output: 70x14x32] → Conv[Kernel: 5x5, Stride: 1, Output: 70x14x64, Activation: Relu]
Revised: → Pool[Stride: 2, Output: 70x14x32] → Conv[Kernel: 5x5, Stride: 1, Output: 70x14x64, Activation: Relu]

Original: → Pool[Stride: 2, Output: 35x7x64] → FC[Units: 64, Activation: Relu]
Revised: → Pool[Stride: 2, Output: 35x7x64] → FC[Units: 64, Activation: Relu]

Original: [Figure]
Revised: [Figure]

Original: Note that the dimension of a 5-digit large-MNIST image is 140 × 28 .
Revised: Note that the dimension of a 5-digit large-MNIST image is 140 × 28 .

Original: The primary difference between our methods is how we combine the scores to output a row stochastic prediction matrix.
Revised: The primary difference between our methods is how we combine the scores to output a row-stochastic prediction matrix.

Original: For Sortnet-based methods, we use another fully-connected layer of dimension 1 to map the image representations to n scalar scores.
Revised: For NeuralSort-based methods, we use another fully-connected layer of dimension 1 to map the image representations to n scalar scores.

Original: In the case of Stochastic Sortnet, we then sample from the PL distribution by perturbing the scores multiple times with Gumbel noise.
Revised: In the case of Stochastic NeuralSort, we then sample from the PL distribution by perturbing the scores multiple times with Gumbel noise.

Original: Finally, we use the Sortnet operator to map the set of n scores (or each set of n perturbed scores) to its corresponding unimodal row stochastic matrix.
Revised: 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.

Original: For Sinkhorn-based methods, we use a fully-connected layer of dimension n to map each image to an n -dimensional vector.
Revised: For Sinkhorn-based methods, we use a fully-connected layer of dimension n to map each image to an n -dimensional vector.

Original: These vectors are then stacked into an n × n matrix.
Revised: These vectors are then stacked into an n × n matrix.

Original: 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).
Revised: 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).

Original: We implemented the Sinkhorn operator based on code snippets obtained from the open source implementation of
Revised: We implemented the Sinkhorn operator based on code snippets obtained from the open source implementation of

Original: Mena et al.
Revised: Mena et al.

Original: (2018) available at https://github.com/google/gumbel sinkhorn .
Revised: (2018) available at https://github.com/google/gumbel sinkhorn .

Original: 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 n 2 -unit layers to output multiclass predictions for each rank.
Revised: 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 n 2 -unit layers to output multiclass predictions for each rank.

Original: All our methods yield row stochastic n × n matrices as their ﬁnal output.
Revised: All our methods yield row-stochastic n × n matrices as their ﬁnal output.

Original: Our loss is the row-wise cross-entropy loss between the true permutation matrix and the row stochastic output.
Revised: Our loss is the row-wise cross-entropy loss between the true permutation matrix and the row-stochastic output.

Original: Hyperparameters.
Revised: Hyperparameters.

Original: For this experiment, we used an Adam optimizer with an initial learning rate of 10 − 4 and a batch size of 20 .
Revised: For this experiment, we used an Adam optimizer with an initial learning rate of 10 − 4 and a batch size of 20 .

Original: Continuous relaxation to sorting also introduce another hyperparameter: the temperature τ for the Sinkhorn-based and sortnet-based approaches.
Revised: Continuous relaxations to sorting also introduce another hyperparameter: the temperature τ for the Sinkhorn-based and NeuralSort-based approaches.

Original: 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).
Revised: 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).

Original: Effect of temperature.
Revised: Effect of temperature.

Original: In Figure 7, we report the log-variance in gradient estimates as a function of the temperature τ .
Revised: In Figure 7, we report the log-variance in gradient estimates as a function of the temperature τ .

Original: 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.
Revised: 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.

Original: The element-wise mean squareddifference between unimodal approximations (cid:98) P sort ( s ) and the projected hard permutation matrices P sort ( s ) for the best τ on the test set is shown in Table 4.
Revised: The element-wise mean squared

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: Signal-to-Noise Ratio for the Stochastic Optimization over Permutations. To get a better understanding of the behavior of the Gumbel-Sinkhorn and the proposed reparameterized gradient estimates for Stochastic Sortnet, we plot the Signal-to-Noise (SNR) ratio for the gradient estimates in Figure 8 for 5 Gumbel samples and τ = 1 . Here, SNR is deﬁned as the ratio of the absolute value of the mean in gradient estimates to the standard deviation. We observe a much higher SNR for the
Revised: difference between unimodal approximations (cid:98) P sort ( s ) and the projected hard permutation matrices

Original: Stochastic Sortnet approach in line with the overall gains we see on the underlying task.
Revised: P sort ( s ) for the best τ on the test set is shown in Table 4.

Original: D. Q UANTILE R EGRESSION
Revised: D. Q UANTILE R EGRESSION

Original: Architectures.
Revised: Architectures.

Original: Due to resource constraints, we ran the quantile regression experiment on 4-digit numbers instead of 5-digit numbers.
Revised: Due to resource constraints, we ran the quantile regression experiment on 4-digit numbers instead of 5-digit numbers.

Original: We use the same neural network architecture as previously used in the sorting experiment.
Revised: We use the same neural network architecture as previously used in the sorting experiment.

Original: Conv[Kernel: 5x5, Stride: 1, Output: 112x28x32, Activation: Relu]
Revised: Conv[Kernel: 5x5, Stride: 1, Output: 112x28x32, Activation: Relu]

Original: → Pool[Stride: 2, Output: 56x14x32] → Conv[Kernel: 5x5, Stride: 1, Output: 56x14x64, Activation: Relu]
Revised: → Pool[Stride: 2, Output: 56x14x32] → Conv[Kernel: 5x5, Stride: 1, Output: 56x14x64, Activation: Relu]

Original: → Pool[Stride: 2, Output: 28x7x64] → FC[Units: 64, Activation: Relu]
Revised: → Pool[Stride: 2, Output: 28x7x64] → FC[Units: 64, Activation: Relu]

Original: 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.
Revised: 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.

Original: In the other experiments, one copy of this network was used to estimate each element’s rank through a method like Gumbel-Sinkhorn or Sortnet that produces a row stochastic matrix, while another copy was used to estimate each element’s value directly.
Revised: 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.

Original: Point predictions[Figure] are obtained by multiplying the center row of the matrix with the column vector of estimated values, and we minimize the (cid:96) 2 loss between these point predictions and the true median, learning information about ordering and value simultaneously.
Revised: Point predictions are obtained by multiplying the center row of the matrix with the column vector of estimated values, and we minimize the (cid:96) 2 loss between these point predictions and the true median, learning information about ordering and value simultaneously.

Original: Hyperparameters.
Revised: Hyperparameters.

Original: For this experiment, we used an Adam optimizer with an initial learning rate of 10 − 4 and a batch size of 5.
Revised: We used the Adam optimizer with an initial learning rate of 10 − 4 and a batch size of 5.

Original: The temperature τ was tuned on the set { 1 , 2 , 4 , 8 , 16 } based on the validation loss.
Revised: The temperature τ was tuned on the set { 1 , 2 , 4 , 8 , 16 } based on the validation loss.

Original: Further Analysis.
Revised: Further Analysis.

Original: In Figure 9, we show the scatter plots for the true vs. predicted medians on 2000 test points from the large-MNIST dataset as we vary n .
Revised: 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 .

Original: For stochastic sortnet, we average the predictions across 5 samples.
Revised: For stochastic NeuralSort, we average the predictions across 5 samples.

Original: 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.
Revised: 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.

Original: However, the relatively difﬁculty of the problem increases with increasing n , as the model is trying to learn a semantic sorting across a larger set of elements.
Revised: However, the relatively difﬁculty of the problem increases with increasing n , as the model is trying to learn a semantic sorting across a larger set of elements.

Original: This is reﬂected in the
Revised: This is reﬂected in the

Original: R 2 values in Table 2 which show a slight dip as n increases.
Revised: R 2 values in Table 2 which show a slight dip as n increases.

Original: D. E ND - TO - END , DIFFERENTIABLE k - NEAREST NEIGHBORS
Revised: D. E ND - TO - END , DIFFERENTIABLE k - NEAREST NEIGHBORS

Original: Architectures.
Revised: Architectures.

Original: The baseline kNN implementation for the pixel basis, PCA basis and theconvolutional autoencoder  was done using sklearn .
Revised: The baseline kNN implementation for the pixel basis, PCA basis and the autoencoder basis was done using sklearn .

Original: For theconvolutional autoencoder baseline for kNN, we used the following standard architectures.
Revised: For the autoencoder baselines for kNN, we used the following standard architectures.

Original: MNIST:
Revised: MNIST and Fashion-MNIST:

Original: The dimension of the encoding used for distance computation in kNN is 256 .
Revised: The dimension of the encoding used for distance computation in kNN is 50 .

Original: The architecture and training procedure follows the one available at https://blog.keras.io Conv[Kernel: 3x3, Stride: 1, Output: 28x28x16, Activation: Relu] → Pool[Stride: 2, Output: 14x14x16] → Conv[Kernel: 3x3, Stride: 1, Output: 14x14x8, Activation: Relu]
Revised: FC[Units: 500, Activation: Relu] → FC[Units: 500, Activation: Relu]

Original: → Pool[Stride: 2, Output: 7x7x8] → Conv[Stride: 3, Stride: 1, Output: 7x7x8, Activation: Relu] → MaxPool[Stride: 2, Output: 4x4x8] = (embedding) → Conv[Kernel: 3x3, Stride: 1, Output: 4x4x8, Activation: Relu] → UpSampling[Size: 2x2, Output: 8x8x8] → Conv[Kernel: 3x3, Stride: 1, Output: 8x8x8, Activation: Relu] → UpSampling[Size: 2x2, Output: 16x16x8] → Conv[Kernel: 3x3, Output: 14x14x16, Activation: Relu]
Revised: → FC[Units: 50, Activation: Relu]

Original:  → UpSampling[Size: 2x2, Output: 28x28x16] → Conv[Kernel: 3x3, Stride: 1, Output: 28x28x1, Activation: Sigmoid]
Revised: = (embedding) → FC[Units: 500, Activation: Relu] → FC[Units: 500, Activation: Relu] → FC[Units: 784, Activation: Sigmoid]

Original: CIFAR-10:
Revised: CIFAR-10:

Original: The dimension of the encoding used for distance computation in kNN is 256 .
Revised: The dimension of the encoding used for distance computation in kNN is 256 .

Original: The architecture and training procedure follows the one available at https://github.com/shibuiwilliam/Keras Autoencoder .
Revised: The architecture and training procedure follows the one available at https://github.com/shibuiwilliam/Keras Autoencoder .

Original: 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]
Revised: 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]

Original: → Pool[Stride: 2, Output: 8x8x32] → Conv[Stride: 3, Output: 8x8x16, Normalization: BatchNorm, Activation: Relu]
Revised: → Pool[Stride: 2, Output: 8x8x32] → Conv[Stride: 3, Output: 8x8x16, Normalization: BatchNorm, Activation: Relu]

Original: → 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]
Revised: → 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]

Original: → UpSampling[Size: 2x2, Output: 16x16x32] → Conv[Kernel: 3x3, Output: 16x16x64, Normalization: BatchNorm, Activation: Relu]
Revised: → UpSampling[Size: 2x2, Output: 16x16x32] → Conv[Kernel: 3x3, Output: 16x16x64, Normalization: BatchNorm, Activation: Relu]

Original: → UpSampling[Size: 2x2, Output: 32x32x64] → Conv[Kernel: 3x3, Stride: 1, Output: 32x32x3, Normalization: BatchNorm, Activation: 
Revised: → UpSampling[Size: 2x2, Output: 32x32x64] → Conv[Kernel: 3x3, Stride: 1, Output: 32x32x3, Normalization: BatchNorm, Activation: Sigmoid]

Original: Sigmoid]
Revised: [Table]

Original: For the MNIST experiments with Sortnet, we used a network similar to the large-MNIST network used in the previous experiments:
Revised: For the MNIST experiments with NeuralSort, we used a network similar to the large-MNIST network used in the previous experiments:

Original: Conv[Kernel: 5x5, Stride: 1, Output: 28x28x32, Activation: Relu] → Pool[Stride: 2, Output: 14x14x32] → Conv[Kernel: 5x5, Stride: 1, Output: 14x14x64, Activation: Relu]
Revised: Conv[Kernel: 5x5, Stride: 1, Output: 24x24x20, Activation: Relu] → Pool[Stride: 2, Output: 12x12x20] → Conv[Kernel: 5x5, Stride: 1, Output: 8x8x50, Activation: Relu]

Original: → Pool[Stride: 2, Output: 7x7x64] → FC[Units: 64, Activation: Relu]
Revised: → Pool[Stride: 2, Output: 4x4x50] → FC[Units: 500, Activation: Relu]

Original: For the  CIFAR experiments with Sortnet, wealso make use of local response normalization (LRN) layers and a second fully-connected layer: [Table]
Revised: For the Fashion-MNIST and CIFAR experiments with NeuralSort, we use the ResNet18 architecture as described in https://github.com/kuangliu/pytorch-cifar .

Original: Hyperparameters.
Revised: Hyperparameters.

Original: For this experiment, we used an Adam optimizer with an initial learning rate of 10 − 5 and a batch size of 10 queries and 100 neighbor candidates.
Revised: 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.

Original: We tuned the temperature parameter on the set { 1 , 2 , 4 , 8 , 16 } . The number of nearest neighbors kwas selected from the set { 1, 3, 5, 9 } .
Revised: 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 } .

Original: The model with the best evaluation loss was evaluated on the test set.
Revised: The model with the best evaluation loss was evaluated on the test set.

Revised: We suspect that accuracy improvements can be made by a more expensive hyperparameter search and a more ﬁne-grained learning rate schedule.

Original: Accuracy for different k .
Revised: Accuracy for different k .

Original: In Table

Original: [Table] we show the performance of Deterministic and Stochastic Sortnet for different choice of the hyperparameter k for the differentiable k -nearest neighbors algorithm [Table]
Revised: 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.

Original: E L OSS FUNCTIONS
Revised: E L OSS FUNCTIONS

Original: For each of the experiments in this work, we assume we have access to a ﬁnite dataset D
Revised: For each of the experiments in this work, we assume we have access to a ﬁnite dataset D

Original: = { ( x (1) , y (1) ) , ( x (2) , y (2) ) , . . . } .
Revised: = { ( x (1) , y (1) ) , ( x (2) , y (2) ) , . . . } .

Original: Our goal is to learn a predictor for y given x , as in a standard supervised learning (classiﬁcation/regression) setting.
Revised: Our goal is to learn a predictor for y given x , as in a standard supervised learning (classiﬁcation/regression) setting.

Original: Below, we state and elucidate the semantics of the training objective optimized by Deterministic and Stochastic Sortnet foreach of the  experiments.
Revised: Below, we state and elucidate the semantics of the training objective optimized by Deterministic and Stochastic NeuralSort for the sorting and quantile regression experiments.

Original: E. S ORTING HANDWRITTEN NUMBERS
Revised: E. S ORTING HANDWRITTEN NUMBERS

Original: We are given a dataset D of sequences of large-MNIST images and the permutations that sort the sequences.
Revised: We are given a dataset D of sequences of large-MNIST images and the permutations that sort the sequences.

Original: 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).
Revised: 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).

Original: For example,
Revised: For example,

Original: Figure 4 shows one input sequence of n = 5 images, and the permutation y =
Revised: Figure 4 shows one input sequence of n = 5 images, and the permutation y =

Original: [3 , 5 , 1 , 4 , 2] that sorts this sequence.
Revised: [3 , 5 , 1 , 4 , 2] that sorts this sequence.

Original: For any datapoint x , let (cid:96) CE ( · ) denote the average multiclass cross entropy (CE) error between the rows of the true permutation matrix P y and a permutation matrix P (cid:98)
Revised: For any datapoint x , let (cid:96) CE ( · ) denote the average multiclass cross entropy (CE) error between the rows of the true permutation matrix P y and a permutation matrix P (cid:98)

Original: y corresponding to a predicted permutation, say (cid:98) y .
Revised: y corresponding to a predicted permutation, say (cid:98) y .

Original: [Equation]
Revised: [Equation]

Original: where 1 ( · ) denotes the indicator function.
Revised: where 1 ( · ) denotes the indicator function.

Original: Now, we state the training objective functions for the
Revised: Now, we state the training objective functions for the

Original: Deterministic and Stochastic Sortnet approaches respectively.
Revised: Deterministic and Stochastic NeuralSort approaches respectively.

Original: 1. Deterministic Sortnet
Revised: 1. Deterministic NeuralSort

Original: [Equation]
Revised: [Equation]

Original: where each entry of s is given as s j = h φ ( x j ) .
Revised: where each entry of s is given as s j = h φ ( x j ) .

Original:  Stochastic Sortnet
Revised: 2. Stochastic NeuralSort

Original: [Equation]
Revised: [Equation]

Original:  where each entry of s is given as s j = h φ ( x j ) .
Revised: ) where each entry of s is given as s j = h φ ( x j ) .

Original: To ground this in our experimental setup, the score s j for each large-MNIST image x j in any input sequence x of n = 5 images is obtained via a CNN
Revised: To ground this in our experimental setup, the score s j for each large-MNIST image x j in any input sequence x of n = 5 images is obtained via a CNN

Original: h φ () with parameters φ .
Revised: h φ () with parameters φ .

Original: Note that the CNN parameters φ are shared across the different images x 1 , x 2 , . . .
Revised: Note that the CNN parameters φ are shared across the different images x 1 , x 2 , . . .

Original: , x n in the sequence for efﬁcient learning.
Revised: , x n in the sequence for efﬁcient learning.

Original: E. Q UANTILE REGRESSION
Revised: E. Q UANTILE REGRESSION

Original: 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.
Revised: 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.

Original: For example, the desired label corresponds to y = 

Revised: For example, the desired label corresponds to y =

Original: (a real-valued scalar) for the input sequence of n = 5 images in Figure 4.
Revised: (a real-valued scalar) for the input sequence of n = 5 images in Figure 4.

Original: For any datapoint x , let (cid:96) MSE ( · ) denote the mean-squared error between the true median y and the prediction, say (cid:98) y .
Revised: For any datapoint x , let (cid:96) MSE ( · ) denote the mean-squared error between the true median y and the prediction, say (cid:98) y .

Original: [Equation]
Revised: [Equation]

Original: For the Sortnet approaches, we optimize the following objective functions.
Revised: For the NeuralSort approaches, we optimize the following objective functions.

Original: Deterministic Sortnet
Revised: Deterministic NeuralSort

Original: [Equation]
Revised: [Equation]

Original: where each entry of s is given as s j = h φ ( x j ) .
Revised: where each entry of s is given as s j = h φ ( x j ) .

Original: 2. Stochastic Sortnet
Revised: Stochastic NeuralSort

Original: [Equation]
Revised: [Equation]

Original: where each entry of s is given as s j = h φ ( x j ) .
Revised: where each entry of s is given as s j = h φ ( x j ) .

Original: As before, the score s j for each large-MNIST image x j in any input sequence x of n images is obtained via a CNN h φ () with parameters φ .
Revised: As before, the score s j for each large-MNIST image x j in any input sequence x of n images is obtained via a CNN h φ () with parameters φ .

Original: Once we have a predicted permutation matrix (cid:98) P sort ( s ) (or (cid:98) P z ) for deterministic (or stochastic) approaches, we extract the median image via (cid:98) P sort ( s ) x
Revised: Once we have a predicted permutation matrix (cid:98) P sort ( s ) (or (cid:98) P z ) for deterministic (or stochastic) approaches, we extract the median image via (cid:98) P sort ( s ) x

Original: (or (cid:98) P z x ).
Revised: (or (cid:98) P z x ).

Original: Finally, we use a neural network g θ ( · ) with parameters θ to regress this image to a scalar prediction for the median.
Revised: Finally, we use a neural network g θ ( · ) with parameters θ to regress this image to a scalar prediction for the median.

Original: E. E ND - TO - END , DIFFERENTIABLE k -N EAREST N EIGHBORS

Original: In this experiment, we are given access to a dataset D of ( x , y ) pairs of standard MNIST/CIFARimages and their class labels.

Original: The goal is to perform classiﬁcation on the image labels.

Original: For any datapoint x , let (cid:96) kNN ( · ) denote the uniformly-weighted kNN loss of the permutation matrix

Original: P (cid:98) y given n training datapoints, which we write as follows:

Original: [Equation]

Original: Note that when P (cid:98) y is an exact permutation matrix, this expression is exactly the negative of the number of the k nearest neighbors that have the same label as x .

Original: The training objective function for the Sortnet approaches are as follows.

Original: 1. Deterministic Sortnet

Original: [Equation]

Original: where

Original: each entry of

Original: s is

Original: given

Original: as s j = −(cid:107)

Original: [Equation]

Original: 2. Stochastic Sortnet

Original: [Equation]

Original: where each

Original: entry of s

Original: is given as s j = −(cid:107)

Original: h

Original: φ

Original: ( x ) − h

Original: φ

Original: ( x j ) (cid:107) 22 .

Original: At each gradient update, a set of n = 100 candidate neighbors ( x 1 , y 1 ) , ..., ( x 100 , y 100 ) is sampled from the training set.

Original: A neural network h φ ( · ) maps a query point x

Original: and the candidate neighbors to representations h φ ( x ) , h φ ( x 1 ) , h φ ( x 2 ) , . .

Original: , h φ ( x 100 ) .

Original: Finally, the score vector s is computed based on the Euclidean distances of the query point representation to the representations of the candidate neighbors as

Original: [Equation]