Source-linked AI summary
Learning Latent Permutations with Gumbel-Sinkhorn Networks
Gonzalo Mena, David Belanger, Scott Linderman, Jasper Snoek
TL;DR
Latent matching models are difficult to train because exact marginalization requires intractable structured partition functions. The paper replaces hard matching with Sinkhorn-based differentiable relaxations, including Gumbel-Sinkhorn, and reports stronger performance than competitive baselines across sorting, jigsaw, and worm-neural-signal tasks.
Problem
Learning latent matchings is difficult because computing their probabilities requires an intractable partition function, despite their use in alignment, canonicalization, sorting, and neuroscience.
Method
The paper approximates maximum-weight permutation matching with the continuous Sinkhorn operator and uses Gumbel-Sinkhorn to enable reparametrized optimization.
Results
The methods outperform strong neural-network baselines on sorting numbers, solving jigsaw puzzles, and identifying neural signals from C. elegans worms.
Takeaways & Limitations
Sinkhorn networks can learn permutations for reconstruction and matching tasks, including sorting, jigsaw puzzles, and neural-signal identification.
Takeaways & Limitations
The approach has clear limits on the more challenging ImageNet setting, motivating sequential improvement or richer permutation parameterizations.
Abstract
from arXiv · showhide
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.
1 INTRODUCTION
Matchings and permutations support alignment, canonicalization, and sorting, but learning latent matchings is difficult because their marginal likelihoods involve intractable partition functions. The paper introduces differentiable relaxations and evaluates them across several tasks.
- Latent matchings are relevant when matching annotations are unavailable, including identifying neurons in C. elegans.
- Marginal likelihood optimization is challenging because structured matching distributions have intractable partition functions.This prevents unbiased score-function gradients for a given matching.
- The paper introduces Sinkhorn-based differentiable relaxations and Gumbel-Sinkhorn for optimizing latent permutation models with reparametrization.
- The methods outperform strong neural-network baselines on sorting numbers, solving jigsaw puzzles, and identifying neural signals from C. elegans worms.
2 THE SINKHORN OPERATOR: AN ANALOG OF THE SOFTMAX FOR PERMUTATIONS
The Sinkhorn operator provides a continuous analog of hard permutation selection by iteratively normalizing matrices, while the matching operator selects a maximum-weight permutation. The paper shows that the continuous relaxation approaches the discrete solution as temperature decreases, though practical computation requires truncated iterations.
- As τ approaches zero, softmax selects a one-hot vector, motivating the analogous extension from categorical variables to permutations.
- Sinkhorn normalization iteratively normalizes matrix rows and columns, producing a doubly stochastic matrix in the Birkhoff polytope.
- The matching operator parameterizes a permutation as the solution to a linear assignment problem over permutation matrices.
- M(X) can be approximated by S(X/τ) with small τ because the entropy-regularized problem converges to the matching problem.
- The theoretical limit requires infinitely many Sinkhorn iterations, so practical implementations truncate the iteration count to L.
3 SINKHORN NETWORKS
Sinkhorn networks predict permutations from scrambled inputs by replacing a nondifferentiable matching operator with a differentiable Sinkhorn relaxation. They use permutation-equivariant architectures to produce consistent approximate assignments for reconstruction and matching tasks.
- A neural network maps scrambled inputs to a matrix whose matching solution represents the predicted permutation.
- The nondifferentiable matching operator is replaced by S(g(˜X, θ)/τ) so the network can be trained with gradients.
- Temperature requires care: values that are too small can make gradients vanish, while values that are too large can yield nonsensical reconstructions.
- 3.1 PERMUTATION EQUIVARIANCE: Permutation equivariance is implemented by processing each scrambled piece with the same network before applying Sinkhorn normalization.
- 3 SINKHORN NETWORKS: Sinkhorn networks can reconstruct scrambled objects and learn matchings between objects from two equally sized sets.
4 PROBABILISTIC ASPECTS: THE GUMBEL-SINKHORN AND GUMBEL-MATCHING DISTRIBUTIONS
The paper extends Gumbel-Softmax-style reparameterization to latent permutations by combining Gumbel perturbations with differentiable Sinkhorn relaxation. These constructions support surrogate variational inference, while introducing limitations from intractable densities and potentially looser bounds.
- Computational considerations: The approach avoids enumerating all permutations by exploiting factorized perturbations and tractable optimization alternatives.Direct perturb-and-MAP sampling is infeasible because the permutation space has size N!, while factorized settings motivate more tractable approximations.
- Gumbel-Matching: Gumbel-Matching samples permutations by maximizing independently perturbed linear assignment potentials.The permutation potential is the Frobenius inner product ⟨X, P⟩F, and adding i.i.d. Gumbel noise to X yields the Gumbel-Matching distribution.
- Gumbel-Sinkhorn: Gumbel-Sinkhorn replaces nondifferentiable Gumbel-Matching samples with Sinkhorn-relaxed doubly stochastic matrices.Samples are formed as S((X + ε)/τ) and converge almost surely to Gumbel-Matching samples as the relaxation approaches the discrete regime.
- Limitations: Gumbel-Matching and Gumbel-Sinkhorn lack tractable densities, so the method computes a surrogate divergence in the underlying Gumbel space.The same Sinkhorn transformation applied to prior and posterior enables tractable computation, but non-invertibility can make the resulting lower bound less tight.
- Variational inference: The resulting relaxed distributions provide reparameterizable latent permutation variables for differentiable variational inference.The surrogate ELBO replaces Gumbel-Matching prior and posterior distributions with Gumbel-Sinkhorn distributions.
- Variational inference: Mean-field factorization separates the permutation variable from other latent variables, allowing the discrete part to be treated independently.The assumed decompositions are q(P, W|Y) = q(P)q(W) and p(P, W) = p(P)p(W).
5 EXPERIMENTS
The experiments evaluate Sinkhorn-based models across sorting, jigsaw reconstruction, arbitrary MNIST digit assembly, and posterior inference over latent permutations. Results show strong performance across these settings, while also exposing sensitivity to temperature, ambiguity in reconstruction losses, and dataset difficulty.
- Experimental scope: The experiments compare Sinkhorn architectures and probabilistic constructions across sorting, jigsaw puzzles, MNIST digit assembly, and latent-permutation inference.The first three experiments study increasingly complex Sinkhorn network architectures, while the fourth addresses marginal inference over an unobserved permutation.
- Sorting numbers: The network learns to sort up to N = 120 numbers and generalizes beyond the U(0, 1) training interval.Performance starts to degrade only after N ≈100 for most test intervals, compared with guarantees up to N = 15 for a prior recurrent approach.
- Jigsaw puzzles: Celeba puzzles are solved up to 5x5 pieces with only 21% of face pieces incorrectly ordered.Performance decreases substantially at extreme temperatures, whereas one Sinkhorn iteration remains reasonable when temperature is appropriately chosen.
- Jigsaw puzzles: On Imagenet, the network ties the .72 Kendall tau score reported by DeepPermNet while using a much simpler architecture.The best model has around 1,050,000 parameters, whereas a single DeepPermNet layer has approximately 19,000,000 parameters before counting AlexNet parameters.
- MNIST digit assembly: For arbitrary MNIST digit transformations from 10x10 scrambled pieces, the evaluation CNN judged reconstructions appropriate 85.1% of the time.The same network also produced meaningful assemblies from strongly mixed pieces, though unsupervised training for this alternative generative use was beyond the paper’s scope.
- Posterior inference over permutations: In the C. elegans task, variational inference with Gumbel-Sinkhorn outperforms alternative methods in most experimental conditions.The comparisons include local-swipe MCMC, a rounding method, and a no-regularization variant; poor MCMC mixing is identified as one reason for its weaker performance.
6 RELATED WORK
Prior matching work mainly addressed supervised structured prediction, while this paper targets marginal inference with latent matchings. It connects Sinkhorn-based relaxations to optimal transport and extends Gumbel-Softmax ideas to permutations.
- Prior work: Prior matching applications mostly concern structured prediction, whereas this work studies probabilistic marginal inference with latent matchings.The latent setting is harder because standard score-function methods are not applicable when matching distributions have intractable partition functions.
- Continuous relaxations: Gumbel-Softmax reparameterizes categorical variables and replaces them with relaxed surrogates to enable gradient-based learning.The paper parallels this strategy for probabilistic latent permutations.
- Sinkhorn connections: The paper interprets its Sinkhorn approximation through entropy-regularized optimal transport, building on earlier permutation–Sinkhorn connections.Earlier work discussed related limits in statistical physics and economics, whereas this approach emphasizes optimal transport.
- Related architectures: Compared with ranking networks, this work uses a temperature parameter and a limit argument because Sinkhorn approximates the hard permutation choice.The distinction is important for latent matching rather than merely approximating marginal matrices.
- Related architectures: Cruz et al. introduced Sinkhorn networks concurrently for CNN representation, while this paper focuses on their more fundamental properties.Their networks omit an explicit temperature parameter, with τ = 1 remaining within a reasonable range on the reported jigsaw task.
7 DISCUSSION
The discussion reports that Sinkhorn networks can learn permutations in elementary sequential-structure tasks, while ImageNet exposes clear limits of the formulation. The section also states the core Sinkhorn definitions and their relationship to doubly stochastic and permutation matrices.
- Discussion: Sinkhorn networks learned the correct permutation in elementary tasks where training samples share a sequential structure, including sorted numbers and face pieces.The discussion presents this as evidence that neural networks can learn to solve the linear assignment problem.
- Discussion: ImageNet was substantially more challenging, indicating clear limits to the formulation.The authors suggest sequential improvement, richer permutation parameterizations, or reinforcement learning as possible extensions.
- Implementation: The paper provides TensorFlow code implementing the number-sorting experiment.The code is made available through the repository named in the paper.
- Definitions: The Sinkhorn theorem states that a positive square matrix can be diagonally rescaled into a doubly stochastic matrix by alternating row and column normalization.The resulting diagonal factors are unique up to a scalar factor.
- Definitions: The Sinkhorn operator is defined by iterated row and column normalization and produces a doubly stochastic matrix.The operator is denoted S(X) and is obtained as the limit of repeated normalization steps.
- Definitions: The Birkhoff polytope contains doubly stochastic matrices, whose extremal points are the permutation matrices.The paper separately defines these matrix sets through row and column-sum constraints.
A.1 AN APPROXIMATION THEOREM FOR THE MATCHING PROBLEM
The matching problem maximizes a linear score over permutation matrices, while its relaxed counterpart optimizes over the Birkhoff polytope. The main theorem connects the hard matching solution to a differentiable, entropy-regularized Sinkhorn relaxation as temperature vanishes.
- Matching formulation: The matching operator M(X) returns a permutation maximizing the Frobenius inner product with X.The relaxed operator changes the feasible set from permutation matrices to the Birkhoff polytope.
- Matching formulation: The hard matching solution is contained in the relaxed solution set because permutation matrices are extremal points of the Birkhoff polytope.Non-unique solutions may form a set of permutations or a face of the Birkhoff polytope.
- Theorem 1: Theorem 1 is the paper’s main result connecting permutation matching to a differentiable Sinkhorn relaxation.The theorem defines entropy for doubly stochastic matrices and establishes the approximation framework.
- Theorem 1: With independently drawn absolutely continuous entries in X, the theorem’s convergence statement holds almost surely.The distributional assumption supplies the stochastic regime used for the limiting argument.
- Proof strategy: The proof proceeds by relating S(X/τ) to an entropy-regularized optimization problem, establishing uniqueness, and then showing convergence.This decomposes the theorem into intermediate analytical steps.
A.1.1 INTERMEDIATE RESULTS FOR THEOREM 1
The intermediate results identify the Sinkhorn output with the entropy-regularized assignment solution and establish convergence to the hard matching solution under an almost-sure uniqueness condition. This extends the familiar softmax-to-argmax limit to permutations.
- Intermediate results: Strict concavity of the entropy-regularized objective makes its solution Pτ exist and be unique.The argument uses entropy’s strict concavity on the constrained optimization problem.
- Intermediate results: The constrained optimization is expressed with a Lagrangian enforcing row and column sums equal to one.Its Lagrangian combines the Frobenius score, entropy regularization, and multiplier terms.
- Intermediate results: The optimizer has the form Pτ = D1 exp(X/τ)D2, so Sinkhorn’s theorem identifies it with S(X/τ).The diagonal matrices have positive diagonals and provide the required rescaling.
- Intermediate results: For independently drawn absolutely continuous entries of X, the relaxed and hard matching solutions almost surely coincide as a unique permutation matrix.This is the uniqueness result used by the convergence argument.
- Intermediate results: Under the same assumptions, Pτ = S(X/τ) converges to P0 as τ approaches 0 from above.The proof uses convergence of objective values and uniqueness of the linear-program solution.
- Interpretation: The permutation result generalizes the categorical identity arg max_i x_i = lim_{τ→0+} softmax(x/τ).The analogy treats categories as one-hot vectors and permutations as structured discrete choices.
A.3 ILLUSTRATING THEOREM 1
Figure 3 visualizes how Sinkhorn normalization approximates hard matching and how the corresponding relaxed distributions behave across temperatures.
- The matching operator produces a hard permutation matrix from parameter matrix X.
- Sinkhorn normalization approximates matching while varying temperature τ and iteration count L.
- At low temperature, Matching-Sinkhorn and Gumbel-Sinkhorn samples are indistinguishable.
B.1 EXPERIMENTAL PROTOCOLS
The supplementary protocols specify implementation settings and evaluate sensitivity to Sinkhorn hyperparameters across puzzle and digit-transformation experiments.
- Experiments used TensorFlow with multiple GPUs to explore temperature, learning rate, and network dimensions.
- Training used L = 20 Sinkhorn iterations and 10 Gumbel-perturbed reconstructions for each sample in a 10x10 batch.
- The Hungarian algorithm computed M(X) for predicted-match evaluation.
- The image-puzzle architecture used convolutional filters, ReLU, max pooling, and a final assignment representation.
- Extreme temperature settings and a single Sinkhorn iteration produced worse puzzle results, although one iteration remained reasonable in an appropriate temperature range.
- A two-layer CNN identified transformed digits, with the most troublesome transformation being transformation into one.
C.3 RESULTS ON CATEGORIAL VAE IN MNIST
The section relates deterministic transformations of latent variables to KL-divergence bounds and discusses their implications for categorical and permutation relaxations.
- A deterministic mapping yields KL(q(z) ∥ p(z)) = KL(q(y) ∥ p(y)) + KL(q(z|y) ∥ p(z|y)).
- Because the conditional KL term is nonnegative, transforming variables cannot increase the retained KL divergence.
- The resulting lower ELBO is reported to have minimal performance impact in the categorical MNIST experiment.
- Figure 4 shows arbitrary digits reconstructed from pieces that may be strokes or dust rather than digit-like structures.
- Sinkhorn outputs are interpreted as approximations to matching marginals and the partition function.