Learning Latent Permutations with Gumbel-Sinkhorn Networks

Abstract

Permutations and matchings are core building blocks in a variety of latent variable models, as they allow us to align, canonicalize, and sort data. Learning in such models is difficult, however, because exact marginalization over these combinatorial objects is intractable. In response, this paper introduces a collection of new methods for end-to-end learning in such models that approximate discrete maximum-weight matching using the continuous Sinkhorn operator. Sinkhorn iteration is attractive because it functions as a simple, easy-to-implement analog of the softmax operator. With this, we can define the Gumbel-Sinkhorn method, an extension of the Gumbel-Softmax method (Jang et al. 2016, Maddison2016 et al. 2016) to distributions over latent matchings. We demonstrate the effectiveness of our method by outperforming competitive baselines on a range of qualitatively different tasks: sorting numbers, solving jigsaw puzzles, and identifying neural signals in worms.

Cited by in corpus (38)

• Self-supervised learning methods and applications in medical imaging analysis: A survey
• Rethinking Search: Making Domain Experts out of Dilettantes
• Sparse Sinkhorn Attention
• Deep Set Prediction Networks
• FSPool: Learning Set Representations with Featurewise Sort Pooling
• Seq2Slate: Re-ranking and Slate Optimization with RNNs
• Learning Permutations with Sinkhorn Policy Gradient
• Investigating Object Compositionality in Generative Adversarial Networks
• Insertion-based Decoding with automatically Inferred Generation Order
• Neural Graph Matching Network: Learning Lawler's Quadratic Assignment Problem with Extension to Hypergraph and Multiple-graph Matching
• Gradient Estimation with Stochastic Softmax Tricks
• Scalable Optimal Transport Methods in Machine Learning: A Contemporary Survey
• A probabilistic atlas for cell identification
• The Limited Multi-Label Projection Layer
• Neural Proof Nets
• Deep Learning Based Resource Assignment for Wireless Networks
• Differentiable Ranks and Sorting using Optimal Transport
• Towards Listening to 10 People Simultaneously: An Efficient Permutation Invariant Training of Audio Source Separation Using Sinkhorn's Algorithm
• A Deep Generative Model for Reordering Adjacency Matrices
• r-local sensing: Improved algorithm and applications
• Learning Index Selection with Structured Action Spaces
• Set-to-Sequence Methods in Machine Learning: a Review
• Direct Optimization through $\arg \max$ for Discrete Variational Auto-Encoder
• Learning Functions over Sets via Permutation Adversarial Networks
• Searching Dense Point Correspondences via Permutation Matrix Learning
• Learning Randomly Perturbed Structured Predictors for Direct Loss Minimization
• Show, Control and Tell: A Framework for Generating Controllable and Grounded Captions
• COMET: Learning Cardinality Constrained Mixture of Experts with Trees and Local Search
• The Virtual Patch Clamp: Imputing C. elegans Membrane Potentials from Calcium Imaging
• Learning to rank for censored survival data
• Differentiable Sorting Networks for Scalable Sorting and Ranking Supervision
• Modeling Orders of User Behaviors via Differentiable Sorting: A Multi-task Framework to Predicting User Post-click Conversion
• Learning data association without data association: An EM approach to neural assignment prediction
• Learning to Solve Combinatorial Optimization under Positive Linear Constraints via Non-Autoregressive Neural Networks
• A Differentiable Relaxation of Graph Segmentation and Alignment for AMR Parsing
• Unbiased Gradient Estimation with Balanced Assignments for Mixtures of Experts
• Set Cross Entropy: Likelihood-based Permutation Invariant Loss Function for Probability Distributions
• Learning sparse transformations through backpropagation