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. In principle, deep networks can learn arbitrarily sophisticated mappings from inputs to outputs. However, in practice we must encode specific inductive biases in order to learn accurate models from limit data. In a variety of recent research efforts, practitioners have provided models with the ability to explicitly manipulate latent combinatorial objects such as stacks BID12 BID23 , memory slots BID16 BID45 , mathematical expressions BID33 , program traces BID13 BID6 , and first order logic . Operations on these discrete objects can be approximated using differentiable operations on continuous relaxations of the objects. As such, these operations can be included as modules in neural network models that can be trained end-toend by gradient descent.Matchings and permutations are a fundamental building block in a variety of applications, as they can be used to align, canonicalize, and sort data. Prior work has developed learning algorithms for supervised learning where the training data includes annotated matchings BID36 BID46 . However, we would like to learn models with latent matchings, where the matching is not provided to us as supervision. This is a common and relevant setting. For example, BID30 showed a problem from neuroscience involving the identification of neurons from the worm C. elegans can be cast as the inference of latent permutation on a larger hierarchical structure.Unfortunately, maximizing the marginal likelihood for problems with latent matchings is very challenging. Unlike for problems with categorical latent variables, we cannot obtain unbiased stochastic gradients of the marginal likelihood using the score function estimator BID54 , as computing the probability of a given matching requires computing an intractable partition function for a structured distribution. Instead, we draw on recent work that obtains biased stochastic gradients by relaxing the discrete latent variables into continuous random variables that support the reparametrization trick BID22 BID31 .Our contributions are the following: first, in Section 2 we present a theoretical result showing that the non-differentiable parameterization of a permutation can be approximated in terms of a differentiable relaxation, the so-called Sinkhorn operator. Based on this result, in Section 3 we introduce Sinkhorn networks, which generalize the work of method of BID1 for predicting rankings, and complements the concurrent work by BID9 , by focusing on more fundamental aspects. Further , in Section 4 we introduce the Gumbel-Sinkhorn, an analog of the Gumbel Softmax distribution BID22 BID31 for permutations. This enables optimization of the marginal likelihood by the reparametrization trick. Finally, in Section 5 we demonstrate that our methods outperform strong neural network baselines on the tasks of sorting numbers, solving jigsaw puzzles, and identifying neural signals from C. elegans worms. We have demonstrated Sinkhorn networks are able to learn to find the right permutation in the most elementary cases; where all training samples obey the same sequential structure; e.g., in sorted number and in pieces of faces, as we expect parts of faces occupy similar positions from sample to sample. This is already non-trivial, as indicates one can train a neural network to solve the linear assignment problem.However, the fact that Imagenet represented a much more challenging scenario indicates there are clear limits to our formulation. As the most obvious extension we propose to introduce a sequential stage, in which current solutions are kept on a memory buffer, and improved. One way to achieve this would be by exploring more complex parameterizations for permutations; i.e. replacing M (X) by a quadratic operator that may parameterize a notion of local distance between pieces. Alternatively, one may resort to reinforcement learning techniques, as suggested in BID3 . Either sequential improvement would help solve the "Order Matters" problem BID52 , and we deem our elementary work as a significant step in that direction.We have made available Tensorflow code for Gumbel-Sinkhorn networks featuring an implementation of the number sorting experiment at http://github.com/google/gumbel sinkhorn .
A new method for gradient-descent inference of permutations, with applications to latent matching inference and supervised learning of permutations with neural networks The paper utilizes finite approximation of the Sinkhorn operator to describe how one can construct a neural network for learning from permutation valued training data.  The paper proposes a new method that approximates the discrete max-weight for learning latent permutations