Stochastic Optimization of Sorting Networks via Continuous Relaxations

By Aditya Grover, Eric Wang, Aaron Zweig, Stefano Ermon. 2019

220 citations recorded by the source corpus.

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

Subfields: Artificial Intelligence

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

Citations in the corpus

Graphite: Iterative Generative Modeling of Graphs (2018, background): “This is an exciting challenge for future work, which could potentially be resolved by incorporate graph representations robust to permutation invariances (Verma & Zhang, 2017) or modeling distributions over permutations of node orderings via recent approaches such as NeuralSort (Grover et al., 2019)”
Stochastic Beams and Where to Find Them: The Gumbel-Top-k Trick for Sampling Sequences Without Replacement (2019, background): “Our Stochastic Beam Search is a principled alternative that shares some of the beneﬁts of these heuristic variants, such as the ability to control diversity or produce randomized output.”
Learning Fast Algorithms for Linear Transforms Using Butterfly Factorizations (2019, background): “, 2018), sorting (Grover et al., 2019), and traveling salesman (Bello et al.”
Gradient Estimation with Stochastic Softmax Tricks (2020, methodology): “Previously, [57] considered U ∼ Gumbel( θ ) and [31] uses correlated Gumbel-based utilities that induce a Plackett-Luce model [51][65].”
Git Re-Basin: Merging Models modulo Permutation Symmetries (2022, background): “Akin to differentiable permutation learning, many prior works have studied differentiable sorting (Grover et al., 2019; Prillo & Eisenschlos, 2020; Cuturi et al., 2019; Petersen et al., 2022; 2021; Mena et al., 2018).”
Fast Differentiable Sorting and Ranking (2020, methodology): “This method is reﬁned by Grover et al. (2019) using unimodal row-stochastic matrices.”
Unbiased Gradient Estimation with Balanced Assignments for Mixtures of Experts (2021, background): “Sampling and optimization of balanced assignments is closely related to the sampling and optimization of permutations (n × n assignments) [24, 1, 27, 30, 29, 13, 32], which relate to estimation of the matrix permanent [16, 17, 22].”
FSPool: Learning Set Representations with Featurewise Sort Pooling (2019, methodology): “To avoid this issue, we use the recently proposed sorting networks [12], which provide a continuous relaxation of numerical sorting.”
Reparameterizable Subset Sampling via Continuous Relaxations (2019, methodology): “Table 3: Forward pass average runtimes (100 trials) for a small CNN selecting k = 5 neighbors from m candidates using the NeuralSort relaxation [Grover et al., 2019] and our method (RelaxSubSample).”
Differentiable Top-k with Optimal Transport (2020, methodology): “The settings of the neural networks, the training procedure, and the number of neighbors k, and the tuning procedures are similar to Grover et al. (2019). The tuning o ranging from 10−6 to 10−2.”