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: A BSTRACT
Revised: A BSTRACT

Original: Sorting input objects is an important step in 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 with respect to 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
Revised: In this work, we propose

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

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

Original: 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.
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, including a fully differentiable, parameterized extension of the k -nearest neighbors algorithm
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: 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).
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: Such algorithms typically require learning informative representations of complex, high-dimensional data, such as images, before 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 (Weinberger et al., 2006).
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: 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.
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 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.
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 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.
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: Instead of a permutation matrix,
Revised: Instead of a permutation matrix,

Original: NeuralSort returns a unimodal row-stochastic matrix.
Revised: NeuralSort returns a unimodal row-stochastic matrix.

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

Original: All permutation matrices are unimodal row-stochastic matrices.
Revised: All permutation matrices are unimodal row-stochastic matrices.

Original: 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.
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: 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.
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, NeuralSort is also suitable for efﬁcient straight-through gradient optimization (Bengio et al., 2013), which requires “exact” permutation matrices to evaluate learning objectives.
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: As the second primary contribution, we consider the use of NeuralSort for 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 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.,
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.,

Original: ; Chierichetti et al., 2018).
Revised: ; Chierichetti et al., 2018).

Original: By learning distributions over unobserved permutations, we can account for the uncertainty in these 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 with stochastic optimization over discrete distributions lies in gradient estimation with respect to the distribution parameters.
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).
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).
Revised: In this regard, we consider the Plackett-Luce (PL) family of distributions over permutations (Plackett, 1975; Luce, 1959).

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

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

Original: 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.
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: By using NeuralSort instead, we can approximate the objective and obtain well-deﬁned reparameterized gradient estimates for stochastic optimization.
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 NeuralSort to tasks that require us to learn semantic orderings of complex, highdimensional 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 sorting 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 quantile regression, where we want to estimate the median ( 50 -th percentile) 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
Revised: Because the choice of the k nearest neighbors requires a non-differentiable sorting, we use

Original: NeuralSort to obtain an approximate, differentiable surrogate.
Revised: NeuralSort to obtain an approximate, differentiable surrogate.

Original: On all tasks, we observe signiﬁcant empirical improvements due to NeuralSort over the relevant baselines and competing relaxations to permutation matrices.
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).

Original: [Figure]
Revised: [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 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.
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
Revised: The proofs of all theoretical results in this work are deferred to

Original: Appendix B.
Revised: Appendix B.

Original:  N EURAL S ORT : 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
Revised: E.g. , in straight-through gradient estimation (Bengio et al., 2013; Jang et al., 2017), the

Original: /9 / 16 5 / 16 1 /  
Revised: /9 / 16 5 / 16 1 /  

Original: 
Revised: 

Original: P D
Revised: P D

Original: R
Revised: R

Original: 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.
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.
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: If we remove 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) nj =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
Revised: We denote U n as the set of

Original: n
Revised: n

Original: ×
Revised: ×

Original: n
Revised: n

Original: unimodal
Revised: unimodal

Original: row
Revised: row

Original: stochastic matrices.
Revised: stochastic matrices.

Original: All row stochastic matrices satisfy the ﬁrst two conditions.
Revised: 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 matrix via a row-wise arg max operation.
Revised: The condition provides a straightforward mechanism for extracting a permutation from a unimodal row stochastic matrix via a row-wise arg max operation.

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

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

Original: [Figure]
Revised: [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: 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 τ deﬁned in Eq.
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.t.
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: 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:
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 NeuralSort relaxation to the sort operator:
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 relevance 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, 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 traditional 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: 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.
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.

Original: the set of doubly-stochastic matrices.
Revised: the set of doubly-stochastic matrices.

Original: A doubly-stochastic matrix is a permutation matrix
Revised: A doubly-stochastic matrix is a permutation matrix

Original: 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: 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;

Original: Fiori et al., 2013; Lim & Wright, 2014).
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 near the
Revised: (2018) propose a rounding procedure that uses the Sinkhorn operator to directly sample matrices near the

Original: 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, NeuralSort maps permutation matrices to the set of unimodal rowstochastic 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 (Kim et al., 2016; Mussmann et al., 2017), 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 NeuralSort and Stochastic
Revised: We refer to the two approaches proposed in Sections 3, 4 as Deterministic NeuralSort and Stochastic

Original: NeuralSort, 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
Revised: Every sequence is this dataset consists of

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

Original: A
Revised: A

Original: [Figure]
Revised: [Figure]

Original: =
Revised: =

Original: [Figure]
Revised: [Figure]

Original: ( x
Revised: ( x

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: , z )
Revised: , z )

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: [Table]
Revised: [Table]

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

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: [Table]
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 NeuralSort 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 NeuralSort approach is the consistent best performer on
Revised: The stochastic NeuralSort approach is the consistent best performer on

Original: MSE, while the deterministic NeuralSort is slightly better on the R 2 metric.
Revised: 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.

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

Original: The differentiable kNN algorithm consists of two hyperparameters: the number of training neighbors n , the number of top candidates k , and the sorting temperature τ .
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 x and a randomly sampled subset of n candidate nearest neighbors from the training set, say { x 1 , x 2 , . . .
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 , . . .

Original: , x n } .
Revised: , 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 h φ ( · ) ( e.g. , CNN) such that the top- k candidate points with the minimum
Revised: The loss function optimizes for a representation space h φ ( · ) ( e.g. , CNN) such that the top- k candidate points with the minimum

Original: Euclidean distance to the query point in the representation space have the same label as the query point.
Revised: Euclidean distance to the query point in the representation space have the same label as the query point.

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

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

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

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: where { y 1 , y 2 , . . . ,
Revised: where { y 1 , y 2 , . . . ,

Original: y n } are the labels for the candidate points.
Revised: y n } are the labels for the candidate points.

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

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

Original: Deterministic:
Revised: Deterministic:

Original: [Equation]
Revised: [Equation]

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

Original: [Equation]
Revised: [Equation]

Original: x

Original: [Equation]

Original:  − h [Equation] x [Equation]
Revised: (15) where each entry of s is given as s j = −(cid:107) h φ ( x ) − h φ ( x j ) (cid:107) 22 .

Original: Datasets.
Revised: Datasets.

Original: 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.
Revised: 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), an autonencoder (AE).
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 CNN with the same architecture as the one used for NeuralSort (except the ﬁnal layer) using the cross-entropy loss.
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 proposed NeuralSort, 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 any computation graph involving a sort operator.
Revised: Our relaxation facilitates gradient estimation on any computation graph involving a sort operator.

Original: Further, we derived a reparameterized gradient estimator for the
Revised: Further, we derived a reparameterized gradient estimator for the

Original: Plackett-Luce distribution for efﬁcient stochastic optimization over permutations.
Revised: Plackett-Luce distribution for efﬁcient stochastic optimization over permutations.

Original: 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.
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 deterministic and stochastic NeuralSort 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
Revised: We provide reference implementations in Tensorﬂow (Abadi et al., 2016) and PyTorch (Paszke et al., 2017) in

Original: Appendix A.
Revised: Appendix A.

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

Original: A CKNOWLEDGEMENTS
Revised: A CKNOWLEDGEMENTS

Original: This research was supported by NSF (#1651565, #1522054, #1733686), ONR, AFOSR (FA955019-1-0024), and FLI.
Revised: This research was supported by NSF (#1651565, #1522054, #1733686), ONR, AFOSR (FA955019-1-0024), and FLI.

Original: AG is supported by MSR Ph.D. fellowship and Stanford Data Science scholarship.
Revised: AG is supported by MSR Ph.D. fellowship and Stanford Data Science scholarship.

Original: We are thankful to Jordan Alexander, Kristy Choi, Adithya Ganesh, Karan Goel, Neal Jean,
Revised: We are thankful to Jordan Alexander, Kristy Choi, Adithya Ganesh, Karan Goel, Neal Jean,

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

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

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

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

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

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

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

Original: return P hat
Revised: hat

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

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

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

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

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

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

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

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

Original: A.
Revised: A

Original: P

Original: [Table]
Revised: [Table]

Original: ORCH

Original: Sorting

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

Original: [Table]
Revised: [Table]

Original: import torch
Revised: Reparamterized Sampler

Original: [Table]
Revised: [Table]

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

Original: [Table]
Revised: [Table]

Original: 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)
Revised: batch_size = s [Table] size()[0] n = s

Original: P_
Revised: [Table] size( [Table]

Original: [Table]
Revised: [Table]

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:  = deterministic_\name{}(log
Revised: P_hat = deterministic_NeuralSort(log_s_

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 ] :
Revised: Furthermore, for λ = s [ k ] :

Original: [Equation]
Revised: [Equation]

Original: +
Revised: +

Original: [Equation]
Revised: [Equation]

Original: i = max
Revised: i = max

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

Original: =
Revised: =

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

Original: ] k + n (cid:88)
Revised: ] k + n (cid:88)

Original: i = max( s i − s [ k ]
Revised: i = max( s i − s [ k ]

Original: [Equation]
Revised: [Equation]

Original: =
Revised: =

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

Original: i =
Revised: i =

Original: ( s [ i ]
Revised: ( s [ i ]

Original: [Equation]
Revised: [Equation]

Original: s [ k ] )
Revised: 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 ? , the softmax function may be equivalently deﬁned as 
Revised: As shown in Gao & Pavel (2017), the softmax function may be equivalently deﬁned as soft max( z/τ ) =

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

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

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

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

Original: The distributional assumptions ensure that the elements of s are distinct a.s., so plugging in z = [Equation] n + 1 − 2 k [Equation]
Revised: 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.

Original: 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) ni =
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: soft max
Revised: soft max

Original: [Equation]
Revised: [Equation]

Original: n
Revised: n

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

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

Original: [Equation]
Revised: [Equation]

Original: Z 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) and PyTorch (Paszke et al., 2017) 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 “plug-in” snippets for implementing our proposed relaxations in both
Revised: In Appendix A, we provide “plug-in” snippets for implementing our proposed relaxations in both

Original: Tensorﬂow and PyTorch.
Revised: Tensorﬂow and PyTorch.

Original: The full codebase for reproducing the experiments can be found at https://github.com/ermongroup/neuralsort .
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: 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 NeuralSort-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 NeuralSort, 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
Revised: Finally, we use the

Original: NeuralSort operator to map the set of n scores (or each set of n perturbed scores) to its corresponding unimodal row-stochastic matrix.
Revised: 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 relaxations to sorting also introduce another hyperparameter: the temperature τ for the Sinkhorn-based and NeuralSort-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 squared
Revised: The element-wise mean squared

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: difference between unimodal approximations (cid:98) P sort ( s ) and the projected hard permutation matrices
Revised: difference between unimodal approximations (cid:98) P sort ( s ) and the projected hard permutation matrices

Original: P sort ( s ) for the best τ on the test set is shown in Table 4.
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 NeuralSort that produces a rowstochastic 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 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: We used the 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 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 .
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 NeuralSort, 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 the autoencoder basis 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 the autoencoder baselines for kNN, we used the following standard architectures.
Revised: For the autoencoder baselines for kNN, we used the following standard architectures.

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

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

Original: FC[Units: 500, Activation: Relu] → FC[Units: 500, Activation: Relu]
Revised: FC[Units: 500, Activation: Relu] → FC[Units: 500, Activation: Relu]

Original: → FC[Units: 50, Activation: Relu]
Revised: → FC[Units: 50, Activation: Relu]

Original: = (embedding) → FC[Units: 500, Activation: Relu] → FC[Units: 500, Activation: Relu] → FC[Units: 784, 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: Sigmoid]
Revised: → UpSampling[Size: 2x2, Output: 32x32x64] → Conv[Kernel: 3x3, Stride: 1, Output: 32x32x3, Normalization: BatchNorm, Activation: Sigmoid]

Original: [Table]
Revised: [Table]

Original: For the MNIST experiments with NeuralSort, 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: 24x24x20, Activation: Relu]
Revised: Conv[Kernel: 5x5, Stride: 1, Output: 24x24x20, Activation: Relu]

Original: → Pool[Stride: 2, Output: 12x12x20] → Conv[Kernel: 5x5, Stride: 1, Output: 8x8x50, Activation: Relu]
Revised: → Pool[Stride: 2, Output: 12x12x20] → Conv[Kernel: 5x5, Stride: 1, Output: 8x8x50, Activation: Relu]

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

Original: For the Fashion-MNIST and CIFAR experiments with NeuralSort, we use the ResNet18 architecture as described in https://github.com/kuangliu/pytorch-cifar .
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 SGD optimizer with a momentum parameter of 0 .
Revised: For this experiment, we used an SGD optimizer with a momentum parameter of 0 .

Original: 9 , with a batch size of 100 queries and 100 neighbor candidates at a time.
Revised: 9 , with a batch size of 100 queries and 100 neighbor candidates at a time.

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

Original: We suspect that accuracy improvements can be made by a more expensive hyperparameter search and a more ﬁne-grained learning rate schedule.
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 5, we show the performance of Deterministic and Stochastic
Revised: In Table 5, we show the performance of Deterministic and Stochastic

Original: NeuralSort for different choice of the hyperparameter k for the differentiable k -nearest neighbors algorithm.
Revised: 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 NeuralSort for the sorting and quantile regression 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 NeuralSort approaches respectively.
Revised: Deterministic and Stochastic NeuralSort approaches respectively.

Original: 1. Deterministic NeuralSort
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: 2. Stochastic NeuralSort
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 NeuralSort approaches, we optimize the following objective functions.
Revised: For the NeuralSort approaches, we optimize the following objective functions.

Original: Deterministic NeuralSort
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: Stochastic NeuralSort
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.