L EARNING L ATENT P ERMUTATIONS WITH G UMBELS INKHORN N ETWORKS
A BSTRACT
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 difﬁcult, 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 operator is attractive because it functions as a simple, easy-to-implement analog of the softmax operator.
With this, we can deﬁne the Gumbel-Sinkhorn method, an extension of the Gumbel-Softmax method (Jang et al., 2016; Maddison 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.
I NTRODUCTION
In principle, deep networks can learn arbitrarily sophisticated mappings from inputs to outputs.
However, in practice we must encode speciﬁc 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 (Dyer et al., 2015;
Joulin & Mikolov, 2015), memory slots (Graves et al., 2014; Sukhbaatar et al., 2015), mathematical expressions (Neelakantan et al., 2015), program traces (Gaunt et al., 2016;
Bo ˇ snjak et al., 2017), and ﬁrst order logic (Rockt ¨ aschel & Riedel, 2017).
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 (Caetano et al., 2009;
Petterson et al., 2009; Tang et al., 2016).
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, Linderman et al.
(2017) showed a problem from neuroscience involving the identiﬁcation 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 (Williams, 1992), 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 (Jang et al., 2016; Maddison et al., 2016).
Our contributions are the following: ﬁrst, 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 Adams & Zemel (2011) for predicting rankings, and complements the concurrent work by Cruz et al.
(2017), by focusing on more fundamental aspects.
Further, in Section 4 we introduce the Gumbel-Sinkhorn , an analog of the
Gumbel Softmax distribution (Jang et al., 2016; Maddison et al., 2016) for permutations.
This enables optimization of the marginal likelihood by the reparametrization trick.
Finally, in Sectionwe 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.
T HE S INKHORN OPERATOR : AN ANALOG OF THE SOFTMAX FOR PERMUTATIONS
One sensible way to approximate a discrete category by continuous values is by using a temperature-dependent softmax function, component-wise deﬁned as softmax τ ( x ) i =
exp
( x i /τ ) / (cid:80) j =1 exp( x j /τ ) .
For positive values of τ , softmax τ ( x ) i is a point in the probability simplex.
Also, in the limit τ → 0 , softmax τ ( x ) i converges to a vertex of the simplex, a one-hot vector corresponding to the largest x i 1 .
This approximation is a key ingredient in the successful implementations by Jang et al.
; Maddison et al.
and here we extend it to permutations.
To do so, we ﬁrst state an analog of the normalization implemented by the softmax.
This is achieved through the Sinkhorn operator (or Sinkhorn normalization, or Sinkhorn balancing), which iteratively normalizes rows and columns of a matrix.
Speciﬁcally, following Adams & Zemel (2011), we deﬁne the Sinkhorn operator S ( X ) over an N dimensional square matrix X as:
[Equation]
where T r ( X ) =
X (cid:11) (
X 1 N 1 (cid:62) N ) , and T c ( X ) =
X (cid:11) ( 1 N 1 (cid:62) N X ) as the row and column-wise normalization operators of a matrix, with (cid:11) denoting the element-wise division and 1 N a column vector of ones.
Sinkhorn (1964) proved that S ( X ) must belong to the Birkhoff polytope, the set of doubly stochastic matrices, that we denote B N 2 .
Building on our analogy with categories, notice that choosing a category can always be cast as a maximization problem: the choice arg max i x
i is the one that maximizes the function (cid:104) x, v (cid:105) (with v being a one-hot vector), i.e. the maximizing v ∗ indexes the largest x i .
Similarly, one may parameterize the choice of a permutation P through a square matrix X , as the solution to the linear assignment problem (Kuhn, 1955), with P N denoting the set of permutation matrices and (cid:104) A, B (cid:105) F = trace( A (cid:62) B ) the (Frobenius) inner product of matrices:
[Equation]
We call M ( · ) the matching operator, through which we parameterize the hard choice of a permutation (see Figure 3a for an example).
Our theoretical contribution is to show that M ( X ) can be obtained as the limit of S ( X/τ ) , meaning that one can approximate M ( X )
≈ S ( X/τ ) with a small τ .
Theorem 1 summarizes our ﬁnding.
We provide a rigorous proof in appendix A; brieﬂy, it is based on showing that S ( X/τ ) solves a certain entropy-regularized problem in B n , which in the limit converges to the matching problem in equation 2.
Theorem 1.
For a doubly-stochastic matrix P , deﬁne its entropy as h ( P ) =
− (cid:80) i,j P i,j log ( P i,j ) .
Then, one has,
[Equation]
Now, assume also the entries of X are drawn independently from a distribution that is absolutely continuous with respect to the Lebesgue measure in R .
Then, almost surely, the following convergence holds:
[Equation]
Finally, we note that Theorem 1 cannot be realized in practice, as it involves a limit on the Sinkhorn iterations l .
Instead, we’ll always consider the incomplete version of the Sinkhorn operator (Adams & Zemel, 2011), where we truncate l in (1) to L .
Figure 3b in appendix A.3 illustrates the dependence of the approximation in τ and L .
S INKHORN N ETWORKS
Now we show how to apply the approximation in Theorem 1 in the context of artiﬁcial neural networks.
We construct a layer that encodes the representation of a permutation, and show how to train networks containing such layers as intermediate representations.
We deﬁne the components of this network through a minimal example: consider the supervised task of learning a mapping from scrambled objects
˜ X
to actual, non-scrambled X .
Data, then, are M pairs ( X i , ˜ X i ) where ˜ X i can be constructed by randomly permuting pieces of X i .
We state this problem as a permutation-valued regression X i =
P − 1 θ, ˜ X i ˜
X i + ε i , where ε i is a noise term, and
P θ, ˜ X
i is the permutation matrix mapping
X
i to ˜ X
i , which depends on ˜ X
i and parameters θ .
We are concerned with minimization of the reconstruction error 3 :
[Equation]
One way to express a complex parameterization of this kind is through a neural network: this network receives ˜ X i as input, which is then passed through some intermediate, feed-forward computations of the type g h
( W h x h + b h ) , where g h are nonlinear activation functions, x h is the output of a previous layer, and θ = { ( W h , b h ) } h are the network parameters.
To make the ﬁnal network output be a permutation, we appeal to constructions developed in Section 2: by assuming that the ﬁnal network output P θ, ˜ X can be parameterized as the solution of the assignments problem; i.e.,
P θ, ˜ X
= M ( g ( ˜ X, θ )) , where g ( · , θ ) represents the outcome of all operations involving g h .
Unfortunately, the above construction involves a non-differentiable f (in θ ).
We use Theorem 1 as a justiﬁcation for replacing M ( g ( ˜ X, θ )) by the differentiable S ( g ( ˜ X, θ ) /τ ) in the computational graph.
The value of τ must be chosen with caution: if τ is too small, gradients vanishes almost everywhere, as S ( g ( ˜ X, θ ) /τ ) approaches the non-differentiable M ( g ( ˜ X, θ )) .
Conversely, if τ is too large, S ( X/τ ) may be far from the vertices of the Birkhoff polytope, and reconstructions P − 1 θ, ˜ X
˜ X may be nonsensical (see Figure 2a).
Importantly, we will always add noise to the output layer g ( ˜ X, θ ) as a regularization device: by doing so we ensure uniqueness of M ( g ( ˜ X, θ )) , which is required for convergence in Theorem 1.
3. P ERMUTATION EQUIVARIANCE
Among all possible architectures that respect the aforementioned parameterization, we will only consider networks that are permutation equivariant , the natural kind of symmetry arising in this context.
Speciﬁcally, we require networks to satisfy:
[Equation]
where P (cid:48) is an arbitrary permutation.
The underlying intuition is simple: reconstructions of objects should not depend on how pieces were scrambled, but only on the pieces themselves.
We achieve permutation equivariance by using the same network to process each piece of ˜ X , throwing an N dimensional output.
Then, these N outputs (each with N components) are used to create the rows of the matrix g ( ˜ X, θ ) , to which we ﬁnally apply the (differentiable) Sinkhorn operator (i.e. g stacks the composition of the g h acting locally on each piece).
One can interpret each row as representing a vector of local likelihoods of assignment, but they might be inconsistent.
The Sinkhorn operator, then, mixes those separate representations, and ensures that consistent (approximate) assignment are produced.
With permutation equivariance, the only consideration left to the practitioner is the choice of the particular architecture, which will depend on the particular kind of data.
In Sectionwe illustrate the uses of Sinkhorn networks with three examples, each of them using a different architecture.
Also, in ﬁgure 1 we illustrate a network architecture used in one of our examples.
3. S UMMARY
Sinkhorn network is a supervised method for learning to reconstruct a scrambled object ˜ X (input) given several training examples ( X i , ˜ X i ) .
By applying some non-linear transformations, a Sinkhorn network richly parameterizes the mapping between ˜ X and the permutation P that once applied to ˜
X , will allow to reconstruct the original object as X rec =
P (cid:62) ˜
X (the output).
We note that Sinkhorn networks may be similarly used not only to learn permutations, but also to learn matchings between objects of two sets of the same size.
P θ, ˜ X that resolves inconsistencies in g ( ˜ X, θ ) .
P θ, ˜ X is then used to recover the actual X at training, although at test time one may use the actual M ( g ( ˜ X, θ )) .
P ROBABILISTIC ASPECTS : THE G UMBEL -S INKHORN AND G UMBEL -M ATCHING DISTRIBUTIONS
Recently, in Jang et al.
(2016) and Maddison et al.
(2016), the Gumbel-Softmax or Concrete distributions were deﬁned for computational graphs with stochastic nodes; i.e, latent probabilistic representations.
Their choice is guided by the following i) they seek re-parameterizable distributions to enable the re-parameterization trick (Kingma & Welling, 2013), and note that via the Gumbel trick (see below) any categorical distribution is re-parameterizable, ii) since the re-parameterization in i) is not differentiable, they consider instead sampling under the softmax approximation.
This gives rise to the Gumbel-Softmax distribution.
Here we parallel these choices to enable learning of a probabilistic latent representation of permutations.
To this aim, we start by considering a generic distribution on the discrete set Y , with potential function X : Y → R :
[Equation]
Regarding i), the Gumbel trick arises in the context of Perturb and MAP methods (Papandreou & Yuille, 2011) for sampling in discrete graphical models.
This has recently received renewed interest (Balog et al., 2017), as it recasts the a difﬁcult sampling problem as an easier optimization problem.
In detail, sampling from (6), can be achieved by the maximization of random perturbations of each potential X ( y ) , with Gumbel i.i.d.
noise γ ( y ) ; i.e., arg max y ∈Y { X ( y ) + γ ( y ) } ∼ p ( ·| X ) .
Therefore, one can re-parameterize any categorical distribution (corresponding to (6) with X ( y ) =
(cid:104) X, y (cid:105) ) by the choice of a category, after injecting noise.
However, the above scheme is unfeasible in our context, as |Y| = N ! .
Nonetheless, we appeal to an interesting result:
in cases where Y factorizes, Y = (cid:81) Ni =
Y i 4 , the use of rank-one perturbations γ ( y )
=
(cid:80) N i =1 γ
i
( y i ) is proposed as a more tractable alternative.
Although ultimately heuristic, they lead to bounds in the partition function (Hazan & Jaakkola, 2012;
Balog et al., 2017), and can also be understood as providing approximate or unbiased samples from the true density (Hazan et al., 2013; Tomczak, 2016).
Guided by this, we say the random permutation P follows the Gumbel-Matching distribution with parameter X , denoted P ∼ G .
M .
( X ) , if it has the distribution arising by the rank-one perturbation of (6) on permutations, with the linear potential X ( P ) =
(cid:104) X, P (cid:105) F (replacing y with P ).
One can verify, in a similar line as in Li et al.
, that M ( X + ε ) ∼ G .
M .
( X ) , if ε is a matrix of standard i.i.d.
Gumbel noise.
Unfortunately, as ii) with the categorical case, Gumbel-Matching distribution samples are not differentiable in X , but by appealing to Theorem 1, we deﬁne its relaxation for doubly stochastic matrices as follows:
we say P follows the Gumbel-Sinkhorn distribution with parameter X and temperature τ , denoted P ∼ G .
S .
( X, τ ) , if it has the distribution of S (( X + ε ) /τ ) .
Samples of G .
S .
( X, τ ) converge almost surely to samples of the Gumbel-Matching distribution (see Fig 3c in appendix
A.3).
Unlike for the categorical case, neither the Gumbel-Matching nor Gumbel-Sinkhorn distributions have tractable densities.
However, this does not preclude inference: likelihood-free methods have recently been developed to enable learning in such implicitly deﬁned distributions (Ranganath et al., 2016; Tran et al., 2017).
These methods avoid evaluating the likelihood based on the observation that in many cases inference can be cast as the estimation of a likelihood ratio, which can be obtained from samples (Husz ´ ar, 2017).
Regardless of these useful advances, in the following we develop a solution based on using the likelihoods of random variables whose densities are available .
4. A PPROXIMATE P OSTERIOR I NFERENCE
Consider a latent variable model probabilistic model with observed data Y , and latent Z = { P, W } where P is a permutation and W are other variables.
Here we illustrate how to approximate the posterior probability p ( { P, W }| Y ) using variational inference Blei et al.
Speciﬁcally, we aim to maximize the ELBO, the r.h.s.
of (7):
[Equation]
We assume that both the prior and variational posteriors decompose as products (mean-ﬁeld).
That is, q ( { P, W }| Y ) =
q ( P ) q ( W ) , p ( P, W ) =
p ( P ) p ( W ) .
With this assumption, we may focus only on the discrete part of the problem, i.e. without loss of generality we can assume Z = P .
We parameterize our variational prior and posteriors on P using the Gumbel-Matching distributions with some parameter X ; G .
M .
( X ) .
To enable differentiability, we replace them by G . S .
( X, τ ) distributions, leading to a surrogate ELBO that uses relaxed (continuous) variables.
In more detail,
[Table]
Unfortunately, the term KL ( q ( P | Y ) (cid:107) p ( P )) =
KL ( G . S .
( X, τ ) (cid:107) G .
S .
( X = 0 , τ prior )) in equation (7) is intractable as there is not closed form expression for the density of G .
S .
random variables.
As a solution
[Table]
we use that our prior and posterior are re-parameterizable in terms of matrices ε of Gumbel i.i.d variables: we have S (( X + ε ) /τ ) ∼ G .
S .
( X, τ ) and S ( ε/τ prior ) ∼ G .
S .
( X = 0 , τ prior ) , for the posterior and prior, respectively.
To obtain a tractable expression, we propose to use as ‘code’ or stochastic node Z , the variable ( X + ε ) /τ instead.
Then, the KL term substantially simpliﬁes to KL (( X + ε ) /τ
(cid:107) ε/τ prior ) .
This term can be computed explicitly, as shown in appendix B.3.
This ‘trick’, however, comes at a cost: the divergence KL ( Z 1 (cid:107) Z 2 ) would certainly remain unchanged by applying the same invertible transformation g to both variables Z 1 and Z 2 , but in the general case, for non-invertible transformations, such as S ( · ) , one has KL ( Z 1 (cid:107) Z 2 ) ≥
KL ( g ( Z 1 ) (cid:107) g
( Z 2 )) .
This implies that working in the ‘Gumbel space’ might entail the optimization of a less tight lower bound.
Nonetheless, through categorical experiments on MNIST (see appendix
C.3)
we observe this loss of tightness is minimal, suggesting the suitability of our approach on permutations.
Finally, we note that key to to our treatment of the problem is the fact that both the prior and posterior were the same function ( S ( · ) ) of a simpler distribution.
This may not be the case in more general models.
To conclude this section, we refer the reader to table 8 in appendix D.2 for a summary of all the constructions on permutations developed in this work.
E XPERIMENTS
In this section we perform several experiments comparing to existing methods.
In the ﬁrst three experiments we explore different Sinkhorn network architectures of increasing complexity, and therefore, they mostly implements section 3.
The fourth experiment relates to the probabilistic constructions described in section 4, and addresses a problem involving marginal inferences over a latent, unobserved permutation.
All experimental details not stated here are in appendix B.
5. S ORTING NUMBERS
To illustrate the capabilities of Sinkhorn Networks in a simple scenario, we consider the task of sorting numbers using artiﬁcial neural networks as in Vinyals et al.
Speciﬁcally, we sample uniform random numbers ˜ X in the [0 , 1] interval and we train our network with pairs ( ˜ X, X ) where
X are the same ˜ X but in sorted order.
The network has a ﬁrst fully connected layer that links a number with an intermediate representation (with 32 units), and a second (also fully connected) layer that turns that representation into a row of the matrix g ( ˜ X, θ ) .
Table 1 shows our network learns to sort up to N = 120 numbers.
As an evaluation measure, we report the proportion of sequences where there was at least one error (Prop. any wrong).
Surprisingly,
[Table] the network learns to sort numbers even when test examples are not sampled from U (0 , 1) [Table] but on a considerably different interval.
This indicates the network is not overﬁtting.
These results can be compared with those from Vinyals et al.
, where a much more complex (recurrent) network was used
[Table]
but performance guarantees were obtained only with at most N = 15 numbers.
In that case, the reported error rate is 0.9, whereas ours starts to degrade only after N ≈ 100 for most test intervals.
5. J IGSAW P UZZLES
A more complex scenario for learning permutations arises in the reconstruction of an image X from a collection of scrambled “jigsaw” pieces ˜ X (Noroozi & Favaro, 2016; Cruz et al., 2017).
In this example, our network differs from the one in 5.1 in the ﬁrst layer is a simple CNN (convolution + max pooling), which maps the puzzle pieces to an intermediate representation (see ﬁgure 1 for details).
For evaluation on test data, we report several measures: ﬁrst, in addition to Prop.
any wrong we also consider Prop.
wrong, the overall proportion of scrambled pieces that were wrongly assigned to their actual position.
Also, we use l 1 and l 2 (train) losses and the Kendall tau, a “correlation coefﬁcient” for ranked data.
In Table 2, we benchmark results for the MNIST, Celeba and Imagenet datasets, with puzzles between 2x2 and 6x6 pieces.
In MNIST we achieve very low l 1 and l 2 on up to 6x6 puzzles but a high proportion of errors.
This is a consequence of our loss being agnostic to particular permutations, but only caring about reconstruction errors: as the number of black pieces increases with the number of puzzle pieces, many become unidentiﬁable under this loss.
In Celeba, we are able to solve puzzles of up to 5x5 pieces with only 21% of pieces of faces being incorrectly ordered (see Figure 2a for examples of reconstructions).
For this dataset, we provide additional baselines in Table 4 of appendix C.1:
there, we show that performance substantially decreases if the temperature is too small or large, but only slightly decreases if only one Sinkhorn iterations is made.
We observe that temperature does play a relevant role, consistent with the ﬁndings of Maddison et al.
; Jang et al.
This might not be obvious a-priori, as one could reason that temperature over-parameterizes the network.
However, results conﬁrm this is not the case.
We hypothesize that different temperatures result in parameter convergence in different phases or regions.
Also, the minor difference for a single iteration suggest that only a few might be necessary, implying potential savings in the memory needed to unroll computations in the graph, during training.
Learning in the Imagenet dataset is much more challenging, as there isn’t a sequential structure that generalizes among images, unlike Celeba and MNIST.
In this dataset, our network ties with the .
Kendall tau score reported in (Cruz et al., 2017).
Their network, named DeepPermNet, is based on the stacking of up to the sixth fully connected layer fc6 of AlexNet (Krizhevsky et al., 2012), which ﬁnally (fully) connects to a Sinkhorn layer through intermediate fc7 and fc8 .
We note, however, our network is much simpler, with only two layers and far fewer parameters.
Speciﬁcally, the network that produced our best results had around 1,050,000 parameters (see appendix B for a derivation), while in DeepPermNet, the layer connecting fc6 with fc7 has 512 × 4096 × 9 ≈ 19 , 000 ,parameters, let alone the AlexNet parameters (also to be learned).
Indeed, we believe there is no reason to consider a complex stacking of convolutions: as the number of pieces increases, each piece is smaller and the convolutional layer eventually becomes fully connected.
In the following experiment we explore this phenomenon in more detail.
5. A SSEMBLY OF ARBITRARY MNIST DIGITS FROM PIECES
We also consider an original application, motivated by the observation that the Jigsaw Puzzle task becomes ill-posed if a puzzle contains too many pieces.
Indeed, consider the binarized MNIST dataset: there, reconstructions are not unique if pieces are sufﬁciently atomic, and in the limit case of pieces of size 1x1 squared pixels, for a given scrambled MNIST digit there are as many valid reconstructions as there are MNIST digits with the same number of white pixels.
In other words, reconstructions stop being probabilistic and become a multimodal distribution over permutations.
We exploit this intuition to ask whether a neural network can be trained to achieve arbitrary digit reconstructions, given their loose atomic pieces.
To address this question, we slightly changed the network in 5.2, this time stacking several second layers linking an intermediate representation to the output.
We trained the network to reconstruct a particular digit with each layer, by using digit identity to indicate which layer should activate with a particular training example.
Our results demonstrate a positive answer: Figure 2b shows reconstructions of arbitrary digits given 10x10 scrambled pieces.
In general, they can be unambiguously identiﬁed by the naked eye.
Moreover, this judgement is supported by the assessment of a neural network.
Speciﬁcally, we trained a two-layer CNN 5 on MNIST (achieving a 99.2% accuracy on test set) and evaluated its performance on the test set generated by arbitrary transformations of each digit of the original test set into any other digit.
We found the CNN made an appropriate judgement in 85.1% of the time.
More speciﬁc results, regarding speciﬁc transformations are presented in Table 5 of appendix C.2.
Finally, we note that meaningful assemblies are possible regardless of the original digit: in Figure 4 of appendix C.2 we show arbitrary reconstructions, by this same network, of “digits” from a ‘strongly mixed’ MNIST dataset.
In detail, these “digits” were crafted by sampling, without replacement, from a bag containing all the small pieces from all original digits.
These reconstructions suggest the possibility of an alternative to generative modeling, based on the (random) assembly of small pieces of noise, instead of the processing of noise through a neural network.
However, this would require training the network without supervision, which is beyond the scope of this work. [Table]
5. P OSTERIOR INFERENCE OVER PERMUTATIONS WITH THE G UMBEL -S INKHORN ESTIMATOR
We illustrate how the G .
S .
distribution can be used as a continuous relaxation for stochastic nodes in a computational graph.
To this end, we revisit the “C. elegans neural identiﬁcation problem”, originally introduced in Linderman et al.
We refer the reader to (Linderman et al., 2017) for an in-depth introduction, but brieﬂy, C. elegans is a nematode (worm) whose biological neural conﬁguration – the connectome – is stereotypical; i.e. specimens always posses the same number of somatic neurons (282) (Varshney et al., 2011), and the ways those neurons connect and interact changes little from worm to worm.
Therefore, its brain can be thought of as a canonical object, and its neurons can unequivocally be identiﬁed with names.
The task, then, consists of matching traces from the observed neural dynamics Y to identities (neuron names) in the canonical brain.
This problem is stated in terms of a Bayesian hierarchical model, in order to proﬁt from prior information that may constrain the possibilities.
Speciﬁcally, one states a linear dynamical system Y
[Equation] 1 + ν t , where ν t is a noise term and W and P are latent variables with respective prior distributions.
W encodes the dynamics, with a prior p ( W ) to represent the sparseness of the connectome, etc., and P is a permutation matrix representing the matching between indexes of observed neurons and their canonical counterparts, where we place a ﬂat prior p ( P ) over permutations.
Notably, within the framework it is possible to model the simultaneous problem with many worms sharing the same dynamical system, but here we avoid explicit references to individuals for notational ease.
Given this model, we seek the posterior distribution p ( { P, W }| Y ) , a problem that we address with variational inference (Blei et al., 2017) using the constructions developed in 4.1.
In Table 3 (and also in Table 7 of appendix C.4) we show results for this task, using accuracy in matching as the performance measure.
These are broken down by relevant experimental covariates (Linderman et al., 2017): different proportion of neurons known beforehand, and by task difﬁculty.
As baselines, we include i) a simple MCMC sampler that proposes local swipes on permutations ii) the rounding method presented in Linderman et al.
(2017), iii) our method, where we also consider the absence of regularization.
Results show our method outperforms the alternatives in most cases.
MCMC fails because mixing is poor, but differences are much subtler with the other baselines.
With them, we see that clear differences with the no-regularization case conﬁrm the stochastic nature of this problem, i.e., that it is truly necessary to represent a latent probabilistic permutation.
We believe our method outperforms the one in Linderman et al.
(2017) because theirs, although it provides a explicit density, is a less tight relaxation, in the sense that points can be anywhere in the space, and not only on the Birkhoff polytope.
Therefore, their prior also needs to be deﬁned on the entire space and may not property act as an efﬁcient regularizer.
R ELATED WORK
Learning with matchings has been extensively been studied in the machine learning community; but current applications mostly relate to structured prediction (Petterson et al., 2009; Tang et al., 2016).
However, our probabilistic treatment focuses on marginal inference in a model with a latent matching.
This is a more challenging scenario, as standard learning techniques, i.e. the score function estimator or REINFORCE (Williams, 1992), are not applicable due to the partition function for non-trivial distributions over matchings.
In the case of latent categories, a recent technique that combines a relaxation and the reparameterization trick (Kingma & Welling, 2013) was proposed as a competitive alternative to REINFORCE for the marginal inference scenario.
Speciﬁcally, Maddison et al.
; Jang et al.
(2016) use the Gumbel-trick to re-parameterize a discrete density, and then replace it with a relaxed surrogate, the Gumbel Softmax distribution, to enable gradient-descent.
Our work, like the simultaneous work of Linderman et al.
(2017), aims to extends the scope of this technique to latent permutations.
We deem our Gumbel Sinkhorn distributions as the most natural tractable extension of the Gumbel Softmax to permutations, as we clearly parallel each of the steps leading to its construction.
A parallel is also presented in Linderman et al.
(2017); and notably, unlike ours, their framework produces tractable densities.
However, it is less clear how their constructions extend each of the features of the Gumbel Softmax: for example, their rounding-based relaxation also utilizes the Sinkhorn operator, but the limit they consider does not make use of the non-trivial statement of
Theorem 1, which naturally extends the categorical case (see appendix A.2 for details).
In practice, we see our results favor the Gumbel Sinkhorn distribution, since it is a tighter relaxation.
Connections between permutations and the Sinkhorn operator have been known for at least twenty years.
Indeed, the limit in Theorem 1 was ﬁrst presented in Kosowsky & Yuille (1994), but their interpretation and motivation were more linked to statistical physics and economics.
However, our approach is different and links to recent developments in optimal transport (OT) (Villani, 2003):
Theorem 1 draws on the entropy-regularization for OT technique developed inCuturi (2013), where the entropy-regularized transportation problem is referred to as a ‘Sinkhorn distance’.
The extension is sensible as in the case of transportation between two discrete measures (here) the Birkhoff polytope appears naturally as the optimization set (Villani, 2003).
Entropy regularization as means to achieve a differentiable version of a loss was ﬁrst proposed in Genevay et al.
(2017) in the context of generative modeling.
Although this ﬁeld may appear separate, recent work (Salimans et al., 2018) makes explicit the connection to permutations: to compute a (Wasserstein) distance between a batch of dataset samples and one of generative samples of the same size, one needs to solve the matching problem so that the distance between matched samples is minimized.
Finally, we note our work shares with Salimans et al.
; Genevay et al.
(2017) in that the OT cost function (here, the matrix X ) is learned using an artiﬁcial neural network.
We understand our work as extending Adams & Zemel (2011), which developed neural networks to learn a permutation-like structure; a ranking.
However, there, as in Helmbold & Warmuth (2009), the objective function was linear and the Sinkhorn operator was instead used as an approximation of a matrix of the marginals, i.e., S ( P ) ≈ E ( P ) .
In consequence, there was no need to introduce a temperature parameter and consider a limit argument, which is critical to our case.
Interestingly, equation (10) can be understood in terms of approximate marginal inference, justifying the approximation S ( P )
≈ E ( P ) .
We comment on this in appendix D.1. Note that Sinkhorn iteration can be interpreted as mean-ﬁeld inference in an associated Gibbs distribution over matchings.
With this in mind, backpropagation through Sinkhorn is an end-to-end learning in an unrolled inference algorithm Stoyanov et al.
(2011); Domke (2013).
In future work, it may be fruitful to unroll alternative algorithms for marginal inference over matchings, such as belief propagation (Huang & Jebara, 2009).
Sinkhorn networks were also very recently introduced in Cruz et al.
(2017), although their work substantially differs from ours.
While their interest lies in the representational aspects of CNN’s, we are more concerned with the more fundamental properties.
In their work, they don’t consider a temperature parameter τ , but their network still successfully learns, as τ = 1 happens to fall within the range of reasonable values.
On the Jigsaw puzzle task, we showed that we achieve equivalent performance with a much simpler network having several times fewer parameters and layers.
Nonetheless, we recognize the need for more complex architectures for the tasks considered in Cruz et al.
and we hope our more general theory; particularly, Theorem 1 and the notion of equivariance, may aid further developments in that direction.
D ISCUSSION
We have demonstrated Sinkhorn networks are able to learn to ﬁnd 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 Bello et al.
Either sequential improvement would help solve the “Order Matters” problem (Vinyals et al., 2015) and we deem our elementary work as a signiﬁcant step in that direction.
We have made available Tensorﬂow code for Gumbel-Sinkhorn networks featuring an implementation of the number sorting experiment at http://github.com/google/gumbel sinkhorn .
A P ROOF OF T HEOREM
In this section we give a rigorous proof of Theorem 1.
Also, in A.2 we brieﬂy comment on how
Theorem 1 extend a perhaps more intuitive results, in the probability simplex.
Before stating Theorem 1 we need some preliminary deﬁnitions.
We start by recalling a well-known result in matrix theory, the Sinkhorn theorem.
Theorem (Sinkhorn).
Let A be an N dimensional square matrix with positive entries.
Then, there exists two diagonal matrices D 1 , D 2 , with positive diagonals, so that P = D 1 AD 2 is a doubly stochastic matrix.
These D 1 , D 2 are unique up to a scalar factor.
Also, P can be obtained through the iterative process of alternatively normalizing the rows and columns of A .
Proof.
See Sinkhorn (1964);
Sinkhorn & Knopp (1967); Knight (2008).
For our purposes, it is useful to deﬁne the Sinkhorn operator S ( · ) as follows:
Deﬁnition 1.
Let X be an arbitrary matrix with dimension N .
Denote T r ( X ) =
X (cid:11) (
X 1 N 1 (cid:62) N ) , T c ( X ) =
X (cid:11) (1 N 1 (cid:62) N X ) (with (cid:11) representing the element-wise division and 1 n the n dimensional vector of ones) the row and column-wise normalization operators, respectively.
Then, we deﬁne the Sinkhorn operator applied to X ; S ( X ) , as follows:
[Equation]
Here, the exp( · ) operator is interpreted as the component-wise exponential.
By Sinkhorn’s theorem,
S ( X ) is a doubly stochastic matrix.
Finally, we review some key properties related to the space of doubly stochastic matrices.
First, we need to deﬁne a relevant geometric object.
Deﬁnition 2.
We denote by B N the N -Birkhoff polytope, i.e., the set of doubly stochastic matrices of dimension N .
Likewise, we denote P n be the set of permutation matrices of size N .
Alternatively,
[Equation]
Theorem (Birkhoff).
P N is the set of extremal points of B N .
In other words, the convex hull of B N equals P N .
Proof.
See Birkhoff (1946).
A. A N APPROXIMATION THEOREM FOR THE MATCHING PROBLEM
Let’s now focus on the standard combinatorial assignment (or matching) problem, for an arbitrary
N dimensional matrix X .
We aim to maximize a linear functional (in the sense of the Frobenius norm) in the space of permutation matrices.
In this context, let’s deﬁne the matching operator M ( · ) as the one that returns the solution of the assignment problem:
[Equation]
Likewise, we deﬁne ˜ M ( · ) as a related operator, but changing the feasible space by the Birkhoff polytope:
[Equation]
Notice that in general ˜ M ( X ) , M ( X ) might not be unique matrices, but a face of the Birkhoff polytope, or a set of permutations, respectively (see Lemma 2 for details).
In any case, the relation
M ( X ) ⊆
˜ M ( X ) holds by virtue of Birkhoff’s theorem, and the fundamental theorem of linear programming.
Now we state the main theorem of this work:
Theorem 1.
For a doubly stochastic matrix P deﬁne its entropy as h ( P ) =
− (cid:80) i,j P i,j log ( P i,j ) .
Then, one has,
[Equation]
Now, assume also the entries of X are drawn independently from a distribution that is absolutely continuous with respect to the Lebesgue measure in R .
Then, almost surely the following convergence holds:
[Equation]
We divide the proof of Theorem 1 in three steps.
First, in Lemma 1 we state a relation between
S ( X/τ ) and the entropy regularized problem in equation (10).
Then, in Lemma 2 we show that under our stochastic regime, uniqueness of solutions holds.
Finally, in Lemma 3 we show that in this well-behaved regime, convergence of solutions holds.
states that and Lemma 2b endows us with the tools to make a limit argument.
A.1. I NTERMEDIATE RESULTS FOR T HEOREM
Lemma 1.
[Equation]
Proof.
We ﬁrst notice that the solution P τ of the above problem exists, and it is unique.
This is a simple consequence of the strict concavity of the objective (recall the entropy is strictly concave Rao (1984)).
Now, let’s state the Lagrangian of this constrained problem
[Equation]
It is easy to see, by stating the equality ∂ L /∂P =
that one must have for each i, j ,
[Equation]
in other words, P τ = D 1 exp( X i,j /τ )
D 2 for certain diagonal matrices D 1 , D 2 , with positive diagonals.
By Sinkhorn’s theorem, and our deﬁnition of the Sinkhorn operator, we must have that
S ( X/τ ) =
P τ .
Lemma 2.
Suppose the entries of X are drawn independently from a distribution that is absolutely continuous with respect to the Lebesgue measure in R .
Then, almost surely, ˜ M
( X ) =
M ( X ) is a unique permutation matrix.
Proof.
This is a known result from sensibility analysis on linear programming which we prove for completeness.
Notice ﬁrst that the problem in (2) is a linear program on a polytope.
As such, by the fundamental theorem of linear program, the optimal solution set must correspond to a face of the polytope.
Let F be a face of B N of dimension ≥ 1 , and take P 1 , P 2 ∈ F , P 1 (cid:54) =
P 2 .
If F is an optimal face for a certain X F , then X F ∈ { X : (cid:104) P 1 , X (cid:105) F = (cid:104) P 2 , X (cid:105) F } .
Nonetheless, the latter set does not have full dimension, and consequently has measure zero, given our distributional assumption on X .
Repeating the argument for every face of dimension ≥ 1 and taking a union bound we conclude that, almost surely, the optimal solution lies on a face of dimension 0, i.e, a vertex.
From here uniqueness follows.
Lemma 3.
Call P τ the solution to the problem in equation 10, i.e. P τ = P τ ( X ) =
S ( X/τ ) .
Under the assumptions of Lemma 2, P τ → P 0 when if τ → 0 + .
Proof.
Proof Notice that by Lemmas 1 and 2, P τ is well deﬁned and unique for each τ ≥ 0 .
Moreover, at τ = 0 , P 0 = M ( X ) is the unique solution of a linear program.
Now, let’s deﬁne f τ ( · ) =
(cid:104)· , X (cid:105) F + τh ( · ) .
We observe that f 0 ( P τ ) →
f 0 ( P 0 ) .
Indeed, one has:
[Equation]
From which convergence follows trivially.
Moreover, in this case convergence of the values implies the converge of P τ : suppose P τ does not converge to P 0 .
Then, there would exist a certain δ and sequence τ n → 0 such that (cid:107) P τ n − P
(cid:107) > δ .
On the other hand, since P 0 is the unique maximizer of an LP, there exists ε > 0 such that f 0 ( P 0 ) − f 0 ( P ) >
ε whenever (cid:107) P − P 0 (cid:107) > δ , P ∈ B N .
This contradicts the convergence of f
( P τ n ) .
A.1. P ROOF OF T HEOREM
The ﬁrst statement is Lemma 1.
Convergence (equation 11) is a direct consequence of Lemma 3, after noticing P τ = S
( X/τ ) and P 0 = M ( X ) .
We note that an alternative approach for the limiting argument is presented in Cominetti & San Mart ´ ın (1994).
A. R ELATION TO SOFTMAX
Finally, we notice that all of the above results can be understood as a generalization of the wellknown approximation result arg max i x i =
lim τ → 0 + softmax
( x/τ ) .
To see this, treat a category as a one-hot vector.
Then, one has
[Equation]
12) where S n is the probability simplex, the convex hull of the one-hot vectors (denoted H n ).
Again, by the fundamental theorem of linear algebra, the following holds:
[Equation]
On the other hand, by a similar (but simpler) argument than of the proof of theorem 4 one can easily show that
[Equation]
where the entropy h ( · ) is not deﬁned as h ( e ) = − (cid:80)
ni =1 e i log( e i )
A. I LLUSTRATING THEOREM
[Figure]
B S UPPLEMENTAL M ETHODS B. E XPERIMENTAL PROTOCOLS
All experiments were run on a cluster using Tensorﬂow Abadi et al.
(2016), using several GPU (Tesla K20, K40, K80 and P100) in parallel to enable an efﬁcient exploration of the hyperparameter space: temperature, learning rate, and neural network parameters (dimensions).
In all cases, we used L = 20 Sinkhorn Operator Iterations, and a 10x10 batch size: for each sample in the batch we used Gumbel perturbations to generate 10 different reconstructions.
For evaluation, we used the Hungarian Algorithm Munkres (1957) to compute M ( X ) required to infer the predicted matching.
Finally, experiments of section 5.4 were done consistent with model speciﬁcations stated in Linderman et al.
B. N UMBER OF PARAMETERS ON S INKHORN N ETWORKS
In the simplest network, the one that sorts number, the number of parameters is given by n u + N × n u :
Indeed, each number is connected with the hidden layer with n u (here, 32) units.
This layer connects with another layer with N units, representing a row of g ( ˜ X, θ ) .
For images, the ﬁrst layer is a convolution, composed by n f convolutional ﬁlters of receptive ﬁeld size K s with n c channels (one or three) followed by a ReLU + max-pooling (with stride s ) operations.
Then, the number of parameters in the ﬁrst layer is given by K 2 s × n c × n f + n f .
The second layers connects the output of a convolution, i.e., the stacked convolved l × l images by each of the ﬁlters (after max-pooling) and p 2 units, where p is the number of pieces each side was divided by.
Therefore, the number of parameters is given by l 2 / ( p 2 s 2 ) × n f × p 2 = l 2 /s 2 × n f , up to rounding and padding subtleties.
Then, the total number of parameters is l 2 /s 2 × n f + K 2 s × n c × n f + n f .
For the 3x3 puzzle on Imagenet, l = 256 , p = 3 , n c = 3 and the optimal network was such that n f = 64 , s = 2 , K s = 5 .
Then, it had 1,053,440 parameters.
Finally, for arbitrary assembly experiments, as one includes additional fully connected second layers, the total number of parameters is n l × l 2 /s 2 × n f + K 2 s × n c × n f + n f , where n l is the number of labels (here, n l = 10 ).
B. I NFERENCE WITH THE IMPLICIT G UMBEL -S INKHORN DISTRIBUTION
Here we show how to compute KL (( X + ε ) /τ (cid:107) ε/τ prior ) , as deﬁned in 4.1.
We ﬁrst notice that the density of the variable h
= ( a + g ) /b , where g has a Gumbel distribution and a, b are constants is given by:
[Equation]
Therefore
[Equation]
the log density ratio LR ( z ) between each component of h 1 = ( x i,j + ε i,j ) /τ and h 2 = ε i,j /τ prior is (suppressing indexing for simplicity)
[Equation]
We need to take expectations with respect to the distribution of h 1 .
To compute this expectation
[Equation]
we ﬁrst express the above ratio in terms of ε
[Equation]
Now we appeal to the law of the unconscious statistician
[Equation]
and take the expectation with respect to ε .
Using the identities
• E ( ε ) =
γ ≈ 0 .
5772 (the Euler-Mascheroni constant) • Moment generating function E (exp( tε ))
= Γ(1 − t ) ; implying E (exp( − ε ))
= 1 and E (exp ( − τ prior /τε ))
= Γ(
+ τ prior /τ ) )
we have:
[Equation]
(1 + τ prior /τ )) .
From this, it easily follows (adding all the N 2 components)
that
KL (( X + ε ) /τ (cid:107) ε/τ prior ) =
(cid:88) i,j
E g 1 ( LR ( z i,j )) =
N 2 (log( τ/τ prior ) − 1 + γ
( τ prior /τ − 1))
+ S 1 + Γ(1 + τ prior /τ )
S 2 , where S 1 = τ prior /τ (cid:80) i,j x i,j and S
= (cid:80) i,j exp ( − x i,
j τ prior /τ ) .
C S UPPLEMENTAL R ESULTS C. P UZZLES
In table 4 we provide further performance measures for the Jigsaw puzzle task on Celeba, for extreme hyper-parameter values: small temperature, large temperature, and a single Sinkhorn iteration These are worse than the ones in table 2, although surprisingly, one Sinkhorn iteration already provides reasonable performance, as long temperature is chosen in an appropriate range.
C. T RANSFORMATIONS INTO ARBITRARY DIGITS
In table 5 we show performance of a 2-layer CNN in detecting transformed digits as the ones they are intended to be.
From this we see the most troublesome transformation was to one, as this network most of the times categorized it as a different number.
Also, in ﬁgure 4 we show transformations, showing that to reconstruct to arbitrary digits it is not required that the original ones have an actual digit-like structure, but they can be only pieces of ‘strokes’ or ‘dust’.
[Table]
[Figure]
[Table]
g it .
In general, for arbitrary random variables Z 1 , Z 2 and a function g , one has
[Equation]
We prove this in the discrete case, for simplicity: call q ( z ) and p ( z ) the densities of Z 1 , Z 2 , and call y =
g ( z ) .
This induces two joint distributions, p ( z, y ) and q ( z, y ) .
Now, deﬁne
[Equation]
Under this deﬁnition, one can verify that
[Equation]
But KL (( q ( y | z ) (cid:107) p
( y | z ))
= 0 , as y is a deterministic function of z .
Therefore,
KL (( q ( z ) (cid:107) p ( z )) =
KL ( q ( y ) (cid:107) p
( y ))
+ KL ( q ( z | y ) (cid:107) p ( z | y )) , and since the second term is positive (a KL divergence)
we conclude KL ( q ( z ) (cid:107) p ( z )) ≥ KL ( q ( y ) (cid:107) p
( y )) .
This implies a lower (or less tight) ELBO if using Z 1 , Z 2 instead of g ( Z 1 ) , g ( Z 2 ) .
However, we note that in the categorical case this has a minimal impact in performance.
Indeed, we replicated the density estimation on MNIST task described in Jang et al.
; Maddison et al.
(2016), and
as alternative method we considered the concrete distribution, but using as stochastic node ( ε + x ) /τ (with prior ε/τ prior instead of two concrete distributions.
In other words, for us g ( x ) = softmax τ ( x ) and Z 1 = ( ε + x ) /τ, Z 2 = ( ε ) /τ prior (in law).
Results are shown in Table 6.
We ﬁrst see that
Concrete distribution does worse than Gumbel-Softmax, which we attribute to a sub-optimal parameter search.
However, we see that working in the Gumbel space has little impact on log p ( x ) : the difference was smaller than .
5 nats.
[Figure]
Metaphorically, one is able to reconstruct pieces out of ‘dust’.
[Table] [Table]
et
[Table]
[Table]
C. S UPPLEMENTARY RESULTS ON C. ELEGANS
Finally, in Table 7 we show additional results for the C.elegans experiment.
The setting is the same as in Figure 4(a) in Linderman et al.
Likewise, Table 3 correspond to the setting of Figure 4(b) in Linderman et al.
D S UPPLEMENTARY DISCUSSION D. S INKHORN OPERATOR FOR APPROXIMATE MARGINAL INFERENCE
A second connection between the distribution in (6) (and therefore, the Matching Gumbel distribution) and the Sinkhorn operator arises as a consequence of Theorem 1.
This relates to the estimation of the marginals E θ ( P i,j ) , known to be a #P hard problem.
A well known result (Globerson &
Jaakkola, 2007; Wainwright et al., 2008), consequence of Fenchel (conjugate) duality (Rockafellar, 1970) applied to exponential families, links this problem to optimization in the following way: lets denote by M the marginal polytope, the convex hull of the set of realizable sufﬁcient statistics, that here coincides with B n .
Also, lets call H ( µ ) the entropy of (6) for the parameter θ ( µ ) such that µ =
E θ ( µ ) ( P ) .
Then,
[Equation]
Notice the only difference between the optimization problems in (17) and (10) is the entropy term, after identifying X with θ .
Therefore, one may understand the Sinkhorn operator as providing approximations for the partition function and the marginals, which will be accurate insofar as h ( µ ) is a good approximation for H ( µ ) .
In this way, one can understand S ( X ) as an approximation for
E θ ( P ) , that may complement more classical ones, as the Bethe and Kituchani’s approximations for
H ( µ ) , and the corresponding approximate inference algorithms that they give rise to (Yedidia et al., 2001; Vilnis et al., 2015). D. S UMMARY OF EXTENSIONS [Table]
arg max x i = arg max s ∈S (cid:104) x, s (cid:105)
M ( X ) =
arg max P ∈B (cid:104) P, X (cid:105) F
[Table]
arg max i x i
=
lim τ → 0 + softmax
( x/τ )
M ( X ) = lim τ → 0 + S ( X/τ )
[Table]
h ( s ) =
(cid:80) i −
s i log
s i
h ( P ) =
(cid:80) i,j − P i,j log ( P i,j )
[Table]
softmax ( x/τ ) =
arg max s ∈S (cid:104) x, s (cid:105) + τh ( s )
S ( X/τ ) =
arg max P ∈B (cid:104) P, X (cid:105) F + τh ( P )
[Table]
arg max i ( x i + (cid:15) i )
[Table]
Matching G M ( X )
M
[Equation]
X
[Equation]
[Table]
softmax (( x + (cid:15) ) /τ )
Gumbel-Sinkhorn G S ( X, τ )
S (( X + (cid:15) ) /τ )