Learning Latent Permutations with Gumbel-Sinkhorn Networks

Introduction

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 (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šnjak et al., 2017), and first order logic (Rocktäschel & 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-to-end 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 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 (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: 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 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 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.

The Sinkhorn 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 defined as $\text{softmax}_{\tau}(x)_{i}=\exp(x_{i}/\tau)/\sum_{j=1}\exp(x_{j}/\tau)$. For positive values of $\tau$, $\text{softmax}_{\tau}(x)_{i}$ is a point in the probability simplex. Also, in the limit $\tau\rightarrow 0$, $\text{softmax}_{\tau}(x)_{i}$ converges to a vertex of the simplex, a one-hot vector corresponding to the largest $x_{i}$ With the exception of the degenerate case of ties.. This approximation is a key ingredient in the successful implementations by Jang et al. (2016); Maddison et al. (2016), and here we extend it to permutations.

To do so, we first 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. Specifically, following Adams & Zemel (2011), we define the Sinkhorn operator $S(X)$ over an $N$ dimensional square matrix $X$ as:

where $\mathcal{T}_{r}(X)=X\oslash(X\mathbf{1}_{N}\mathbf{1}_{N}^{\top}),$ and $\mathcal{T}_{c}(X)=X\oslash(\mathbf{1}_{N}\mathbf{1}_{N}^{\top}X$) as the row and column-wise normalization operators of a matrix, with $\oslash$ denoting the element-wise division and $\mathbf{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 $\mathcal{B}_{N}$ This theorem requires certain technical conditions which are trivially satisfied if $X$ has positive entries, motivating the use of the component-wise exponential $\exp(\cdot)$ in the first line of equation 1..

We call $M(\cdot)$ 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/\tau)$, meaning that one can approximate $M(X)\approx S(X/\tau)$ with a small $\tau$. Theorem 1 summarizes our finding. We provide a rigorous proof in appendix A; briefly, it is based on showing that $S(X/\tau)$ solves a certain entropy-regularized problem in $\mathcal{B}_{n}$, which in the limit converges to the matching problem in equation 2.

For a doubly-stochastic matrix $P$, define its entropy as $h(P)=-\sum_{i,j}P_{i,j}\log\left(P_{i,j}\right)$. Then, one has,

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 $\tau$ and $L$.

Sinkhorn Networks

Now we show how to apply the approximation in Theorem 1 in the context of artificial 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.

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. Specifically, we require networks to satisfy:

2 Summary

Probabilistic aspects: the Gumbel-Sinkhorn and Gumbel-Matching distributions

Recently, in Jang et al. (2016) and Maddison et al. (2016), the Gumbel-Softmax or Concrete distributions were defined 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.

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 difficult 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 $\gamma(y)$; i.e., $\operatorname*{arg\,max}_{y\in\mathcal{Y}}\{X(y)+\gamma(y)\}\sim p(\cdot|X)$. Therefore, one can re-parameterize any categorical distribution (corresponding to (6) with $X(y)=\langle X,y\rangle$) by the choice of a category, after injecting noise.

However, the above scheme is unfeasible in our context, as $|\mathcal{Y}|=N!$. Nonetheless, we appeal to an interesting result: in cases where $\mathcal{Y}$ factorizes, $\mathcal{Y}=\prod_{i=1}^{N}\mathcal{Y}_{i}$ It suffices that $\mathcal{Y}$ is a subset of the product space, which here is true as $\mathcal{Y}=\mathcal{P}_{n}\subseteq\{1,\ldots,N\}^{N}$., the use of rank-one perturbations $\gamma(y)=\sum_{i=1}^{N}\gamma_{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\sim\mathcal{G.M.}(X)$, if it has the distribution arising by the rank-one perturbation of (6) on permutations, with the linear potential $X(P)=\left<X,P\right\rangle_{F}$ (replacing $y$ with $P$). One can verify, in a similar line as in Li et al. (2013), that $M(X+\varepsilon)\sim\mathcal{G.M.}(X)$, if $\varepsilon$ 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 define its relaxation for doubly stochastic matrices as follows: we say $P$ follows the Gumbel-Sinkhorn distribution with parameter $X$ and temperature $\tau$ , denoted $P\sim\mathcal{G.S.}(X,\tau)$, if it has the distribution of $S((X+\varepsilon)/\tau)$. Samples of $\mathcal{G.S.}(X,\tau)$ 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 defined 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ár, 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.

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. (2017). Specifically, we aim to maximize the ELBO, the r.h.s. of (7):

We assume that both the prior and variational posteriors decompose as products (mean-field). 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$; $\mathcal{G.M.}(X)$. To enable differentiability, we replace them by $\mathcal{G.S.}(X,\tau)$ distributions, leading to a surrogate ELBO that uses relaxed (continuous) variables. In more detail, for our uniform prior over permutations we use the isotropic $\mathcal{G.S.}(X=0,\tau_{prior})$ distribution, while for the variational posterior we consider the more generic $\mathcal{G.S.}(X,\tau)$.

Unfortunately, the term $KL(q(P|Y)\;\|\;p(P))=KL(\mathcal{G.S.}(X,\tau)\;\|\;\mathcal{G.S.}(X=0,\tau_{prior}))$ in equation (7) is intractable as there is not closed form expression for the density of $\mathcal{G.S.}$ random variables. As a solution, we use that our prior and posterior are re-parameterizable in terms of matrices $\varepsilon$ of Gumbel i.i.d variables: we have $S((X+\varepsilon)/\tau)\sim\mathcal{G.S.}(X,\tau)$ and $S(\varepsilon/\tau_{prior})\sim\mathcal{G.S.}(X=0,\tau_{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+\varepsilon)/\tau$ instead. Then, the KL term substantially simplifies to $KL((X+\varepsilon)/\tau\;\|\;\varepsilon/\tau_{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}\;\|\;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(\cdot)$, one has $KL(Z_{1}\;\|\;Z_{2})\geq KL(g(Z_{1})\;\|\;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(\cdot)$) 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.

Experiments

In this section we perform several experiments comparing to existing methods. In the first 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.

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, the network learns to sort numbers even when test examples are not sampled from $U(0,1)$, but on a considerably different interval. This indicates the network is not overfitting. These results can be compared with those from Vinyals et al. (2015), where a much more complex (recurrent) network was used, 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\approx 100$ for most test intervals.

2 Jigsaw Puzzles

For evaluation on test data, we report several measures: first, 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 $l1$ and $l2$ (train) losses and the Kendall tau, a “correlation coefficient” 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 $l1$ and $l2$ 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 unidentifiable 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 findings of Maddison et al. (2016); Jang et al. (2016). This might not be obvious a-priori, as one could reason that temperature over-parameterizes the network. However, results confirm 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 .72 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 finally (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. Specifically, 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\times 4096\times 9\approx 19,000,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.

3 Assembly 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 sufficiently 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 identified by the naked eye. Moreover, this judgement is supported by the assessment of a neural network. Specifically, we trained a two-layer CNN Specifically, we used the one described in the Deep MNIST for experts tutorial. 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 specific results, regarding specific 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.

4 Posterior inference over permutations with the Gumbel-Sinkhorn estimator

We illustrate how the $\mathcal{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 identification problem”, originally introduced in Linderman et al. (2017). We refer the reader to (Linderman et al., 2017) for an in-depth introduction, but briefly, C. elegans is a nematode (worm) whose biological neural configuration – 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 identified 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 profit from prior information that may constrain the possibilities. Specifically, one states a linear dynamical system $Y_{t}=PWP^{\mathsf{T}}Y_{t-1}+\nu_{t}$, where $\nu_{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 flat 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 difficulty. 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 confirm 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 defined on the entire space and may not property act as an efficient regularizer.

Related 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 re-parameterization trick (Kingma & Welling, 2013) was proposed as a competitive alternative to REINFORCE for the marginal inference scenario. Specifically, Maddison et al. (2016); 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 first 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 first proposed in Genevay et al. (2017) in the context of generative modeling. Although this field 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. (2018); Genevay et al. (2017) in that the OT cost function (here, the matrix $X$) is learned using an artificial 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)\approx 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)\approx E(P)$. We comment on this in appendix D.1. Note that Sinkhorn iteration can be interpreted as mean-field 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 $\tau$, but their network still successfully learns, as $\tau=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. (2017), and we hope our more general theory; particularly, Theorem 1 and the notion of equivariance, may aid further developments in that direction.

Discussion

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 Bello et al. (2016). Either sequential improvement would help solve the “Order Matters” problem (Vinyals et al., 2015), 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 .

References

Appendix A Proof of Theorem 1

In this section we give a rigorous proof of Theorem 1. Also, in A.2 we briefly comment on how Theorem 1 extend a perhaps more intuitive results, in the probability simplex.

Before stating Theorem 1 we need some preliminary definitions. We start by recalling a well-known result in matrix theory, the Sinkhorn theorem.

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$.

See Sinkhorn (1964); Sinkhorn & Knopp (1967); Knight (2008). ∎

For our purposes, it is useful to define the Sinkhorn operator $S(\cdot)$ as follows:

Let $X$ be an arbitrary matrix with dimension $N$. Denote $\mathcal{T}_{r}(X)=X\oslash(X1_{N}1_{N}^{\top}),\;\mathcal{T}_{c}(X)=X\oslash(1_{N}1_{N}^{\top}X$) (with $\oslash$ representing the element-wise division and $1_{n}$ the $n$ dimensional vector of ones) the row and column-wise normalization operators, respectively. Then, we define the Sinkhorn operator applied to $X$; $S(X)$, as follows:

Here, the $\exp(\cdot)$ 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 define a relevant geometric object.

We denote by $\mathcal{B}_{N}$ the $N$-Birkhoff polytope, i.e., the set of doubly stochastic matrices of dimension $N$. Likewise, we denote $\mathcal{P}_{n}$ be the set of permutation matrices of size $N$. Alternatively,

$\mathcal{P}_{N}$ is the set of extremal points of $\mathcal{B}_{N}$. In other words, the convex hull of $\mathcal{B}_{N}$ equals $\mathcal{P}_{N}$.

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 define the matching operator $M(\cdot)$ as the one that returns the solution of the assignment problem:

Now we state the main theorem of this work:

For a doubly stochastic matrix $P$ define its entropy as $h(P)=-\sum_{i,j}P_{i,j}\log\left(P_{i,j}\right)$. Then, one has,

Now, assume also the entries of $X$ are drawn independently from a distribution that is absolutely continuous with respect to the Lebesgue measure in $\mathcal{R}$. Then, almost surely the following convergence holds:

We divide the proof of Theorem 1 in three steps. First, in Lemma 1 we state a relation between $S(X/\tau)$ 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.

We first notice that the solution $P_{\tau}$ 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

It is easy to see, by stating the equality $\partial{\mathcal{L}}/\partial{P}=0$ that one must have for each $i,j$,

in other words, $P_{\tau}=D_{1}\exp(X_{i,j}/\tau)D_{2}$ for certain diagonal matrices $D_{1},D_{2}$, with positive diagonals. By Sinkhorn’s theorem, and our definition of the Sinkhorn operator, we must have that $S(X/\tau)=P_{\tau}$. ∎

This is a known result from sensibility analysis on linear programming which we prove for completeness. Notice first 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 $\mathcal{F}$ be a face of $\mathcal{B}_{N}$ of dimension $\geq 1$, and take $P_{1},P_{2}\in\mathcal{F}$, $P_{1}\neq P_{2}$. If $\mathcal{F}$ is an optimal face for a certain $X_{\mathcal{F}}$, then $X_{\mathcal{F}}\in\{X\ :\ \left<P_{1},X\right>_{F}=\left<P_{2},X\right>_{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 $\geq 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. ∎

Call $P_{\tau}$ the solution to the problem in equation 10, i.e. $P_{\tau}=P_{\tau}(X)=S(X/\tau)$. Under the assumptions of Lemma 2, $P_{\tau}\rightarrow P_{0}$ when if $\tau\rightarrow 0^{+}$.

Proof Notice that by Lemmas 1 and 2, $P_{\tau}$ is well defined and unique for each $\tau\geq 0$. Moreover, at $\tau=0$, $P_{0}=M(X)$ is the unique solution of a linear program. Now, let’s define $f_{\tau}(\cdot)=\left<\cdot,X\right>_{F}+\tau h(\cdot)$. We observe that $f_{0}(P_{\tau})\rightarrow f_{0}(P_{0})$. Indeed, one has:

From which convergence follows trivially. Moreover, in this case convergence of the values implies the converge of $P_{\tau}$: suppose $P_{\tau}$ does not converge to $P_{0}$. Then, there would exist a certain $\delta$ and sequence $\tau_{n}\rightarrow 0$ such that $\|P_{\tau_{n}}-P_{0}\|>\delta$. On the other hand, since $P_{0}$ is the unique maximizer of an LP, there exists $\varepsilon>0$ such that $f_{0}(P_{0})-f_{0}(P)>\varepsilon$ whenever $\|P-P_{0}\|>\delta$, $P\in\mathcal{B}_{N}$. This contradicts the convergence of $f_{0}(P_{\tau_{n}})$. ∎

A.1.2 Proof of Theorem 1

The first statement is Lemma 1. Convergence (equation 11) is a direct consequence of Lemma 3, after noticing $P_{\tau}=S(X/\tau)$ and $P_{0}=M(X)$. We note that an alternative approach for the limiting argument is presented in Cominetti & San Martín (1994).

A.2 Relation to softmax

Finally, we notice that all of the above results can be understood as a generalization of the well-known approximation result $\operatorname*{arg\,max}_{i}x_{i}=\lim_{\tau\rightarrow 0^{+}}softmax(x/\tau)$. To see this, treat a category as a one-hot vector. Then, one has

where $\mathcal{S}_{n}$ is the probability simplex, the convex hull of the one-hot vectors (denoted $\mathcal{H}_{n}$). Again, by the fundamental theorem of linear algebra, the following holds:

On the other hand, by a similar (but simpler) argument than of the proof of theorem 4 one can easily show that

where the entropy $h(\cdot)$ is not defined as $h(e)=-\sum_{i=1}^{n}e_{i}\log(e_{i})$

A.3 Illustrating theorem 1

Appendix B Supplemental Methods

All experiments were run on a cluster using Tensorflow Abadi et al. (2016), using several GPU (Tesla K20, K40, K80 and P100) in parallel to enable an efficient 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 specifications stated in Linderman et al. (2017)

B.2 Number of parameters on Sinkhorn Networks

For images, the first layer is a convolution, composed by $n_{f}$ convolutional filters of receptive field 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 first layer is given by $K_{s}^{2}\times n_{c}\times n_{f}+n_{f}$. The second layers connects the output of a convolution, i.e., the stacked convolved $l\times l$ images by each of the filters (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})\times n_{f}\times p^{2}=l^{2}/s^{2}\times n_{f}$, up to rounding and padding subtleties. Then, the total number of parameters is $l^{2}/s^{2}\times n_{f}+K_{s}^{2}\times n_{c}\times 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}\times l^{2}/s^{2}\times n_{f}+K_{s}^{2}\times n_{c}\times n_{f}+n_{f}$, where $n_{l}$ is the number of labels (here, $n_{l}=10$).

B.3 Inference with the implicit Gumbel-Sinkhorn distribution

Here we show how to compute $KL((X+\varepsilon)/\tau\;\|\;\varepsilon/\tau_{prior})$, as defined in 4.1. We first 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:

Therefore, the log density ratio $LR(z)$ between each component of $h_{1}=(x_{i,j}+\varepsilon_{i,j})/\tau$ and $h_{2}=\varepsilon_{i,j}/{\tau_{prior}}$ is (suppressing indexing for simplicity)

We need to take expectations with respect to the distribution of $h_{1}$. To compute this expectation, we first express the above ratio in terms of $\varepsilon$

Now we appeal to the law of the unconscious statistician, and take the expectation with respect to $\varepsilon$. Using the identities

$E(\varepsilon)=\gamma\approx 0.5772$ (the Euler-Mascheroni constant)

Moment generating function $E(\exp(t\varepsilon))=\Gamma(1-t)$; implying $E(\exp(-\varepsilon))=1$ and $E(\exp\left(-\tau_{prior}/\tau\varepsilon\right))=\Gamma(1+\tau_{prior}/\tau)$)

From this, it easily follows (adding all the $N^{2}$ components) that

where $S_{1}=\tau_{prior}/\tau\sum_{i,j}x_{i,j}$ and $S_{2}=\sum_{i,j}\exp\left(-x_{i,j}\tau_{prior}/\tau\right)$.

Appendix C Supplemental Results

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.2 Transformations 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 figure 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’.

C.3 Results on categorial VAE in MNIST

In general, for arbitrary random variables $Z_{1},Z_{2}$ and a function $g$, one has

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, define

Under this definition, one can verify that

But $KL((q(y|z)\;\|\;p(y|z))=0$, as $y$ is a deterministic function of $z$. Therefore, $KL((q(z)\;\|\;p(z))=KL(q(y)\;\|\;p(y))+KL(q(z|y)\;\|\;p(z|y))$, and since the second term is positive (a KL divergence) we conclude $KL(q(z)\;\|\;p(z))\geq KL(q(y)\;\|\;p(y))$.

C.4 Supplementary 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. (2017). Likewise, Table 3 correspond to the setting of Figure 4(b) in Linderman et al. (2017).

Appendix D Supplementary discussion

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_{\theta}(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 $\mathcal{M}$ the marginal polytope, the convex hull of the set of realizable sufficient statistics, that here coincides with $\mathcal{B}_{n}$. Also, lets call $\mathcal{H}(\mu)$ the entropy of (6) for the parameter $\theta(\mu)$ such that $\mu=E_{\theta(\mu)}(P)$. Then,

Notice the only difference between the optimization problems in (17) and (10) is the entropy term, after identifying $X$ with $\theta$. Therefore, one may understand the Sinkhorn operator as providing approximations for the partition function and the marginals, which will be accurate insofar as $h(\mu)$ is a good approximation for $\mathcal{H}(\mu)$. In this way, one can understand $S(X)$ as an approximation for $E_{\theta}(P)$, that may complement more classical ones, as the Bethe and Kituchani’s approximations for $\mathcal{H}(\mu)$, and the corresponding approximate inference algorithms that they give rise to (Yedidia et al., 2001; Vilnis et al., 2015).

D.2 Summary of extensions