Stochastic Optimization of Sorting Networks via Continuous Relaxations

Abstract

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.

ICLR 2019

References in corpus (9)

• Auto-Encoding Variational Bayes
• Categorical Reparameterization with Gumbel-Softmax
• Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation
• The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables
• On the Properties of the Softmax Function with Application in Game Theory and Reinforcement Learning
• Gradient Estimation Using Stochastic Computation Graphs
• Ranking via Sinkhorn Propagation
• Robust Multimodal Graph Matching: Sparse Coding Meets Graph Matching
• Sorting Network Relaxations for Vector Permutation Problems

Cited by in corpus (24)

• Learning Object Bounding Boxes for 3D Instance Segmentation on Point Clouds
• Single Image Depth Estimation: An Overview
• FSPool: Learning Set Representations with Featurewise Sort Pooling
• Graphite: Iterative Generative Modeling of Graphs
• Gradient Estimation with Stochastic Softmax Tricks
• NeuralNDCG: Direct Optimisation of a Ranking Metric via Differentiable Relaxation of Sorting
• Differentiable Top-k Operator with Optimal Transport
• Neural Proof Nets
• Fast Differentiable Sorting and Ranking
• Differentiable Ranks and Sorting using Optimal Transport
• A Deep Generative Model for Reordering Adjacency Matrices
• Set-to-Sequence Methods in Machine Learning: a Review
• Quantile Regularization: Towards Implicit Calibration of Regression Models
• Abductive Knowledge Induction From Raw Data
• Learning Randomly Perturbed Structured Predictors for Direct Loss Minimization
• Learning Revenue-Maximizing Auctions With Differentiable Matching
• Permutation-Invariant Variational Autoencoder for Graph-Level Representation Learning
• CARMS: Categorical-Antithetic-REINFORCE Multi-Sample Gradient Estimator
• Differentiable Sorting Networks for Scalable Sorting and Ranking Supervision
• A Universal End-to-End Approach to Portfolio Optimization via Deep Learning
• Learning to Solve Combinatorial Optimization under Positive Linear Constraints via Non-Autoregressive Neural Networks
• Modeling Orders of User Behaviors via Differentiable Sorting: A Multi-task Framework to Predicting User Post-click Conversion
• Unbiased Gradient Estimation with Balanced Assignments for Mixtures of Experts
• Learning to Reconstruct and Segment 3D Objects