Skip to yearly menu bar Skip to main content


Poster

Stochastic Optimization of Sorting Networks via Continuous Relaxations

Aditya Grover · Eric J. Wang · Aaron Zweig · Stefano Ermon

Great Hall BC #74

Keywords: [ plackett-luce ] [ stochastic computation graphs ] [ continuous relaxations ] [ sorting ] [ permutation ]


Abstract:

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

Live content is unavailable. Log in and register to view live content