Learning Permutations with Sinkhorn Policy Gradient

Abstract

Many problems at the intersection of combinatorics and computer science require solving for a permutation that optimally matches, ranks, or sorts some data. These problems usually have a task-specific, often non-differentiable objective function that data-driven algorithms can use as a learning signal. In this paper, we propose the Sinkhorn Policy Gradient (SPG) algorithm for learning policies on permutation matrices. The actor-critic neural network architecture we introduce for SPG uniquely decouples representation learning of the state space from the highly-structured action space of permutations with a temperature-controlled Sinkhorn layer. The Sinkhorn layer produces continuous relaxations of permutation matrices so that the actor-critic architecture can be trained end-to-end. Our empirical results show that agents trained with SPG can perform competitively on sorting, the Euclidean TSP, and matching tasks. We also observe that SPG is significantly more data efficient at the matching task than the baseline methods, which indicates that SPG is conducive to learning representations that are useful for reasoning about permutations.

arXiv comment16 pages, under review for NIPS 2018 conference

References in corpus (14)

• Adam: A Method for Stochastic Optimization
• Continuous control with deep reinforcement learning
• Inductive Representation Learning on Large Graphs
• Categorical Reparameterization with Gumbel-Softmax
• Prioritized Experience Replay
• Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation
• Learning Combinatorial Optimization Algorithms over Graphs
• The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables
• Discriminative Embeddings of Latent Variable Models for Structured Data
• Parameter Space Noise for Exploration
• Neural Combinatorial Optimization with Reinforcement Learning
• Order Matters: Sequence to sequence for sets
• Learning Latent Permutations with Gumbel-Sinkhorn Networks
• Ranking via Sinkhorn Propagation

Cited by in corpus (10)

• GOT: An Optimal Transport framework for Graph comparison
• Seq2Slate: Re-ranking and Slate Optimization with RNNs
• Generalize a Small Pre-trained Model to Arbitrarily Large TSP Instances
• Provably Good Solutions to the Knapsack Problem via Neural Networks of Bounded Size
• Set-to-Sequence Methods in Machine Learning: a Review
• Learning Index Selection with Structured Action Spaces
• TauRieL: Targeting Traveling Salesman Problem with a deep reinforcement learning inspired architecture
• Learning Revenue-Maximizing Auctions With Differentiable Matching
• Learning Combined Set Covering and Traveling Salesman Problem
• Learning to Reconstruct and Segment 3D Objects