Sinkhorn Networks: Using Optimal Transport Techniques to Learn Permutations

Recently, Optimal Transport (OT) has received significant attention in the Machine Learning community. It has been shown to be useful as a tool for generative modeling, in which the density estimation problem is cast as the minimization of a linear function on the transportation polytope. Entropy regularization of this problem (Cuturi, 2013) has been demonstrated to be particularly useful, as its solution can be characterized in terms of the Sinkhorn operator, which i) can be computed more efficiently than the original problem and ii) enables efficient automatic differentiation (AD). We show that this technique extends to the Birkhoff polytope, and we use it to understand the solution of the linear assignment problem as a limit of the Sinkhorn operator. This observation justifies and enables the use of AD in computation graphs containing permutations as intermediate representations. As a result, we are able to introduce Sinkhorn networks for learning permutations, extending the work of Adams & Zemel (2011), and apply them to a variety of tasks. The success of our extension suggests entropy regularization might be used in other polytopes as well, enabling AD in other discrete structures.