Stochastic Optimization of Sorting Networks via Continuous Relaxations

1 INTRODUCTION . Learning to automatically sort objects is useful in many machine learning applications , such as topk multi-class classification ( Berrada et al. , 2018 ) , ranking documents for information retrieval ( Liu et al. , 2009 ) , and multi-object target tracking in computer vision ( Bar-Shalom & Li , 1995 ) . Such algorithms typically require learning informative representations of complex , high-dimensional data , such as images , before sorting and subsequent downstream processing . For instance , the k-nearest neighbors image classification algorithm , which orders the neighbors based on distances in the canonical pixel basis , can be highly suboptimal for classification ( Weinberger et al. , 2006 ) . Deep neural networks can instead be used to learn representations , but these representations can not be optimized end-to-end for a downstream sorting-based objective , since the sorting operator is not differentiable with respect to its input . In this work , we seek to remedy this shortcoming by proposing NeuralSort , a continuous relaxation to the sorting operator that is differentiable almost everywhere with respect to the inputs . The output of any sorting algorithm can be viewed as a permutation matrix , which is a square matrix with entries in { 0 , 1 } such that every row and every column sums to 1 . Instead of a permutation matrix , NeuralSort returns a unimodal row-stochastic matrix . A unimodal row-stochastic matrix is defined as a square matrix with positive real entries , where each row sums to 1 and has a distinct argmax . All permutation matrices are unimodal row-stochastic matrices . NeuralSort has a temperature knob that controls the degree of approximation , such that in the limit of zero temperature , we recover a permutation matrix that sorts the inputs . Even for a non-zero temperature , we can efficiently project any unimodal matrix to the desired permutation matrix via a simple row-wise argmax operation . Hence , NeuralSort is also suitable for efficient straight-through gradient optimization ( Bengio et al. , 2013 ) , which requires “ exact ” permutation matrices to evaluate learning objectives . As the second primary contribution , we consider the use of NeuralSort for stochastic optimization over permutations . In many cases , such as latent variable models , the permutations may be latent but directly influence observed behavior , e.g. , utility and choice models are often expressed as distributions over permutations which govern the observed decisions of agents ( Regenwetter et al. , ∗Equal contribution 2006 ; Chierichetti et al. , 2018 ) . By learning distributions over unobserved permutations , we can account for the uncertainty in these permutations in a principled manner . However , the challenge with stochastic optimization over discrete distributions lies in gradient estimation with respect to the distribution parameters . Vanilla REINFORCE estimators are impractical for most cases , or necessitate custom control variates for low-variance gradient estimation ( Glasserman , 2013 ) . In this regard , we consider the Plackett-Luce ( PL ) family of distributions over permutations ( Plackett , 1975 ; Luce , 1959 ) . A common modeling choice for ranking models , the PL distribution is parameterized by n scores , with its support defined over the symmetric group consisting of n ! permutations . We derive a reparameterizable sampler for stochastic optimization with respect to this distribution , based on Gumbel perturbations to the n ( log- ) scores . However , the reparameterized sampler requires sorting these perturbed scores , and hence the gradients of a downstream learning objective with respect to the scores are not defined . By using NeuralSort instead , we can approximate the objective and obtain well-defined reparameterized gradient estimates for stochastic optimization . Finally , we apply NeuralSort to tasks that require us to learn semantic orderings of complex , highdimensional input data . First , we consider sorting images of handwritten digits , where the goal is to learn to sort images by their unobserved labels . Our second task extends the first one to quantile regression , where we want to estimate the median ( 50-th percentile ) of a set of handwritten numbers . In addition to identifying the index of the median image in the sequence , we need to learn to map the inferred median digit to its scalar representation . In the third task , we propose an algorithm that learns a basis representation for the k-nearest neighbors ( kNN ) classifier in an end-to-end procedure . Because the choice of the k nearest neighbors requires a non-differentiable sorting , we use NeuralSort to obtain an approximate , differentiable surrogate . On all tasks , we observe significant empirical improvements due to NeuralSort over the relevant baselines and competing relaxations to permutation matrices . 2 PRELIMINARIES . An n-dimensional permutation z = [ z1 , z2 , . . . , zn ] T is a list of unique indices { 1 , 2 , . . . , n } . Every permutation z is associated with a permutation matrix Pz ∈ { 0 , 1 } n×n with entries given as : Pz [ i , j ] = { 1 if j = zi 0 otherwise . Let Zn denote the set of all n ! possible permutations in the symmetric group . We define the sort : Rn → Zn operator as a mapping of n real-valued inputs to a permutation corresponding to a descending ordering of these inputs . E.g. , if the input vector s = [ 9 , 1 , 5 , 2 ] T , then sort ( s ) = [ 1 , 3 , 4 , 2 ] T since the largest element is at the first index , second largest element is at the third index and so on . In case of ties , elements are assigned indices in the order they appear . We can obtain the sorted vector simply via Psort ( s ) s . 2.1 PLACKETT-LUCE DISTRIBUTIONS . The family of Plackett-Luce distributions over permutations is best described via a generative process : Consider a sequence of n items , each associated with a canonical index i = 1 , 2 , . . . , n. A common assumption in ranking models is that the underlying generating process for any observed permutation of n items satisfies Luce ’ s choice axiom ( Luce , 1959 ) . Mathematically , this axiom defines the ‘ choice ’ probability of an item with index i as : q ( i ) ∝ si where si > 0 is interpreted as the score of item with index i . The normalization constant is given by Z = ∑ i∈ { 1,2 , ... , n } si . If we choose the n items one at a time ( without replacement ) based on these choice probabilities , we obtain a discrete distribution over all possible permutations . This distribution is referred to as the Plackett-Luce ( PL ) distribution , and its probability mass function for any z ∈ Zn is given by : q ( z|s ) = sz1 Z sz2 Z − sz1 · · · szn Z − ∑n−1 i=1 szi ( 1 ) where s = { s1 , s2 , . . . , sn } is the vector of scores parameterizing this distribution ( Plackett , 1975 ) . 2.2 STOCHASTIC COMPUTATION GRAPHS . The abstraction of stochastic computation graphs ( SCG ) compactly specifies the forward value and the backward gradient computation for computational circuits . An SCG is a directed acyclic graph that consists of three kinds of nodes : input nodes which specify external inputs ( including parameters ) , deterministic nodes which are deterministic functions of their parents , and stochastic nodes which are distributed conditionally on their parents . See Schulman et al . ( 2015 ) for a review . To define gradients of an objective function with respect to any node in the graph , the chain rule necessitates that the gradients with respect to the intermediate nodes are well-defined . This is not the case for the sort operator . In Section 3 , we propose to extend stochastic computation graphs with nodes corresponding to a relaxation of the deterministic sort operator . In Section 4 , we further use this relaxation to extend computation graphs to include stochastic nodes corresponding to distributions over permutations . The proofs of all theoretical results in this work are deferred to Appendix B . 3 NEURALSORT : THE RELAXED SORTING OPERATOR . Our goal is to optimize training objectives involving a sort operator with gradient-based methods . Consider the optimization of objectives written in the following form : L ( θ , s ) = f ( Pz ; θ ) ( 2 ) where z = sort ( s ) . Here , s ∈ Rn denotes a vector of n real-valued scores , z is the permutation that ( deterministically ) sorts the scores s , and f ( · ) is an arbitrary function of interest assumed to be differentiable w.r.t a set of parameters θ and z . For example , in a ranking application , these scores could correspond to the inferred relevances of n webpages and f ( · ) could be a ranking loss . Figure 1 shows the stochastic computation graph corresponding to the objective in Eq . 2 . We note that this could represent part of a more complex computation graph , which we skip for ease of presentation while maintaining the generality of the scope of this work . While the gradient of the above objective w.r.t . θ is well-defined and can be computed via standard backpropogation , the gradient w.r.t . the scores s is not defined since the sort operator is not differentiable w.r.t . s. Our solution is to derive a relaxation to the sort operator that leads to a surrogate objective with well-defined gradients . In particular , we seek to use such a relaxation to replace the permutation matrix Pz in Eq . 2 with an approximation P̂z such that the surrogate objective f ( P̂z ; θ ) is differentiable w.r.t . the scores s. The general recipe to relax non-differentiable operators with discrete codomains N is to consider differentiable alternatives that map the input to a larger continuous codomain M with desirable properties . For gradient-based optimization , we are interested in two key properties : 1 . The relaxation is continuous everywhere and differentiable ( almost- ) everywhere with respect to elements in the input domain . 2 . There exists a computationally efficient projection fromM back to N . Relaxations satisfying the first requirement are amenable to automatic differentiation for optimizing stochastic computational graphs . The second requirement is useful for evaluating metrics and losses that necessarily require a discrete output akin to the one obtained from the original , non-relaxed operator . E.g. , in straight-through gradient estimation ( Bengio et al. , 2013 ; Jang et al. , 2017 ) , the non-relaxed operator is used for evaluating the learning objective in the forward pass and the relaxed operator is used in the backward pass for gradient estimation . The canonical example is the 0/1 loss used for binary classification . While the 0/1 loss is discontinuous w.r.t . its inputs ( real-valued predictions from a model ) , surrogates such as the logistic and hinge losses are continuous everywhere and differentiable almost-everywhere ( property 1 ) , and can give hard binary predictions via thresholding ( property 2 ) . Note : For brevity , we assume that the argmax operator is applied over a set of elements with a unique maximizer and hence , the operator has well-defined semantics . With some additional bookkeeping for resolving ties , the results in this section hold even if the elements to be sorted are not unique . See Appendix C. Unimodal Row Stochastic Matrices . The sort operator maps the input vector to a permutation , or equivalently a permutation matrix . Our relaxation to sort is motivated by the geometric structure of permutation matrices . The set of permutation matrices is a subset of doubly-stochastic matrices , i.e. , a non-negative matrix such that the every row and column sums to one . If we remove the requirement that every column should sum to one , we obtain a larger set of row stochastic matrices . In this work , we propose a relaxation to sort that maps inputs to an alternate subset of row stochastic matrices , which we refer to as the unimodal row stochastic matrices . Definition 1 ( Unimodal Row Stochastic Matrices ) . An n× n matrix is Unimodal Row Stochastic if it satisfies the following conditions : 1 . Non-negativity : U [ i , j ] ≥ 0 ∀i , j ∈ { 1 , 2 , . . . , n } . 2 . Row Affinity : ∑n j=1 U [ i , j ] = 1 ∀i ∈ { 1 , 2 , . . . , n } . 3 . Argmax Permutation : Let u denote an n-dimensional vector with entries such that ui = argmaxj U [ i , j ] ∀i ∈ { 1 , 2 , . . . , n } . Then , u ∈ Zn , i.e. , it is a valid permuation . We denote Un as the set of n× n unimodal row stochastic matrices . All row stochastic matrices satisfy the first two conditions . The third condition is useful for gradient based optimization involving sorting-based losses . The condition provides a straightforward mechanism for extracting a permutation from a unimodal row stochastic matrix via a row-wise argmax operation . Figure 2 shows the relationships between the different subsets of square matrices . NeuralSort . Our relaxation to the sort operator is based on a standard identity for evaluating the sum of the k largest elements in any input vector . Lemma 2 . [ Lemma 1 in Ogryczak & Tamir ( 2003 ) ] For an input vector s = [ s1 , s2 , . . . , sn ] T that is sorted as s [ 1 ] ≥ s [ 2 ] ≥ . . . ≥ s [ n ] , we have the sum of the k-largest elements given as : k∑ i=1 s [ i ] = min λ∈ { s1 , s2 , ... , sn } λk + n∑ i=1 max ( si − λ , 0 ) . ( 3 ) The identity in Lemma 2 outputs the sum of the top-k elements . The k-th largest element itself can be recovered by taking the difference of the sum of top-k elements and the top- ( k − 1 ) elements . Corollary 3 . Let s = [ s1 , s2 , . . . , sn ] T be a real-valued vector of length n. LetAs denote the matrix of absolute pairwise differences of the elements of s such that As [ i , j ] = |si − sj | . The permutation matrix Psort ( s ) corresponding to sort ( s ) is given by : Psort ( s ) [ i , j ] = { 1 if j = argmax [ ( n+ 1− 2i ) s−As1 ] 0 otherwise ( 4 ) where 1 denotes the column vector of all ones . E.g. , if we set i = b ( n + 1 ) /2c then the non-zero entry in the i-th row Psort ( s ) [ i , : ] corresponds to the element with the minimum sum of ( absolute ) distance to the other elements . As desired , this corresponds to the median element . The relaxation requires O ( n2 ) operations to compute As , as opposed to the O ( n log n ) overall complexity for the best known sorting algorithms . In practice however , it is highly parallelizable and can be implemented efficiently on GPU hardware . The argmax operator is non-differentiable which prohibits the direct use of Corollary 3 for gradient computation . Instead , we propose to replace the argmax operator with softmax to obtain a continuous relaxation P̂sort ( s ) ( τ ) . In particular , the i-th row of P̂sort ( s ) ( τ ) is given by : P̂sort ( s ) [ i , : ] ( τ ) = softmax [ ( ( n+ 1− 2i ) s−As1 ) /τ ] ( 5 ) where τ > 0 is a temperature parameter . Our relaxation is continuous everywhere and differentiable almost everywhere with respect to the elements of s. Furthermore , we have the following result . Theorem 4 . Let P̂sort ( s ) denote the continuous relaxation to the permutation matrix Psort ( s ) for an arbitrary input vector s and temperature τ defined in Eq . 5 . Then , we have : 1 . Unimodality : ∀τ > 0 , P̂sort ( s ) is a unimodal row stochastic matrix . Further , let u denote the permutation obtained by applying argmax row-wise to P̂sort ( s ) . Then , u = sort ( s ) . 2 . Limiting behavior : If we assume that the entries of s are drawn independently from a distribution that is absolutely continuous w.r.t . the Lebesgue measure in R , then the following convergence holds almost surely : lim τ→0+ P̂sort ( s ) [ i , : ] ( τ ) = Psort ( s ) [ i , : ] ∀i ∈ { 1 , 2 , . . . , n } . ( 6 ) Unimodality allows for efficient projection of the relaxed permutation matrix P̂sort ( s ) to the hard matrix Psort ( s ) via a row-wise argmax , e.g. , for straight-through gradients . For analyzing limiting behavior , independent draws ensure that the elements of s are distinct almost surely . The temperature τ controls the degree of smoothness of our approximation . At one extreme , the approximation becomes tighter as the temperature is reduced . In practice however , the trade-off is in the variance of these estimates , which is typically lower for larger temperatures .

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