Title: LEARNING DEEP-LATENT HIERARCHIES BY STACK-ING WASSERSTEIN AUTOENCODERS

Abstract: Probabilistic models with hierarchical-latent-variable structures provide state-ofthe-art results amongst non-autoregressive, unsupervised density-based models. However, the most common approach to training such models based on Variational Autoencoders often fails to leverage deep-latent hierarchies; successful approaches require complex inference and optimisation schemes. Optimal Transport is an alternative, non-likelihood-based framework for training generative models with appealing theoretical properties, in principle allowing easier training convergence between distributions. In this work we propose a novel approach to training models with deep-latent hierarchies based on Optimal Transport, without the need for highly bespoke models and inference networks. We show that our method enables the generative model to fully leverage its deep-latent hierarchy, and that in-so-doing, it is more effective than the original Wasserstein Autoencoder with Maximum Mean Discrepancy divergence.

1. INTRODUCTION:

Probabilistic latent-variable modelling is widely applicable in machine learning as a method for discovering structure directly from large, unlabelled datasets. Variational Autoencoders (VAEs) (Kingma & Welling, 2014; Rezende et al., 2014) have proven to be effective for training generative models parametrised by powerful neural networks, by mapping the data into a low-dimensional embedding space. While allowing the training of expressive models, VAEs often fail when using deeper hierarchies with several stochastic latent layers.

In fact, many of the most successful probabilistic latent-variable models use only one stochastic latent layer. Auto-regressive models (Larochelle & Murray, 2011; Van Den Oord et al., 2016b; a; Salimans et al., 2017) , for example, produce state-of-the-art samples and likelihood scores. However, auto-regressive models suffer from poor scalability to high-dimensional data. Flow-based models (Rezende & Mohamed, 2015; Kingma et al., 2016; Dinh et al., 2016; Kingma & Dhariwal, 2018) are another class of generative models providing competitive sample quality and are able to scale to higher-dimensional data, but they still lag behind auto-regressive models in terms of likelihood score.

Deep-latent-variable models are highly expressive models that aim to capture the structure of data in a hierarchical manner, and could thus potentially compete with auto-regressive models for state-ofthe-art performance. However, they remain hard to train. Many explanations have been proposed for this, from the use of dissimilarity measure directly in pixel space (Larsen et al., 2016) resulting in poor sample quality, to the lack of efficient representation in the latent (Zhao et al., 2017) , to simply the poor expressiveness of the models used (Zha et al., 2017; Maaløe et al., 2019) .

Recent works have tried to overcome the latter problem introducing latent skip connections in the generative model and inference network (Bachman, 2016) , sharing generative model and inference network parameters (Sønderby et al., 2016) , as well as bidirectional inference network (Maaløe et al., 2019) . Maaløe et al. (2019) managed to train very deep-hierarchical-latent models achieving near state-of-the-art sample generations. However, in order to leverage their latent hierarchy (working in the VAE framework), they needed both complex, tailored inference networks, and deterministic skip connections in the generative model.

Optimal Transport (OT) (Villani, 2008) is a mathematical framework to compute distances between distributions. Recently, it has been used as a training method for generative models (Genevay et al., 2018; Bousquet et al., 2017) . Tolstikhin et al. (2018) introduced Wasserstein Autoencoders (WAEs), where as with VAEs, an encoding distribution maps the data into a latent space, aiming to learn a low-dimensional representation of the data. WAE is a non-likelihood based framework with appealing topological properties (Arjovsky et al., 2017; Bousquet et al., 2017) , in theory allowing for easier training convergence between distributions. Gaujac et al. (2018) trained a two-latent-layer generative model using WAE, showing promising results in the capacity of the WAE framework to leverage a latent hierarchy in generative modelling. Following these early successes, we propose to train deep-hierarchical-latent models using the WAE framework, without the need for complex dependency paths in both the generative model and inference network. As in the works of Sønderby et al. (2016) and Maaløe et al. (2019) , we believe that a deep-latent hierarchy offers the potential to improve generative models, if they could be trained properly. In order to leverage the deep-latent hierarchy, we derive a novel objective function by stacking WAEs, introducing an inference network at each level, and encoding the information up to the deepest layer in a bottom-up way. For convenience, we refer to our method as STACKEDWAE.

Our contributions are two-fold: first, we introduce STACKEDWAE, a novel objective function based on OT, designed specifically for generative modelling with latent hierarchies, and show that it is able to fully leverage its hierarchy of latents. Second, we show that STACKEDWAE performs significantly better when training hierarchical-latent models than the original WAE framework regularising the latent with the Maximum Mean Discrepancy (MMD) divergence (Gretton et al., 2012) .

2. STACKED WAE:

STACKEDWAEs are probabilistic-latent-variable models with a deep hierarchy of latents. They can be minimalistically simple in their inference and generative models, but are trained using OT in a novel way. We start by defining the probabilistic models considered in this work, then introduce OT, and finally discuss how to train probabilistic models with deep-latent hierarchies using OT, the method that we refer to as STACKEDWAE.

Throughout this paper, we use calligraphic letters (e.g. X ) for sets, capital letters (e.g. X) for random variables, and lower case letters (e.g. x) for their values. We denote probability distributions with capital letters (e.g. P (X)) and their densities with lower case letters (e.g. p(x)).

2.1. GENERATIVE MODELS WITH DEEP HIERARCHIES OF LATENT VARIABLES:

We will consider deep-generative models with Markovian hierarchies in their latent variables. Namely, where each latent variable depends exclusively on the previous one. Denoting by P ⇥ the parametric model with N latent layers, where ⇥ = {✓ 1 , . . . , ✓ N }, we have:

p ⇥ (x) = Z Z1...Z N p ✓1 (x|z 1 ) p ✓2 (z 1 |z 2 ) . . . p ✓ N (z N 1 |z N ) p(z N ) dz 1 . . . dz N (1)

where the data is X 2 X and the latent variables are Z n 2 Z n and we chose p(z N ) = N (z N ; 0 Z N , I Z N ). The corresponding graphical model for N = 3 is given Figure 1a . We will be using variational inference through the WAE framework of Tolstikhin et al. (2018) , introducing variational distributions, q (z 1 , . . . , z N |x), to facilitate the evaluation of the integral in Eq. (1), analogous to VAEs (Kingma & Welling, 2014; Rezende et al., 2014) . It will be shown in Section 2.3 that without loss of generality, q (z 1 , . . . , z N |x) can have a Markovian latent hierarchy when following the STACKEDWAE approach. That is,

EQUATION

where each q N (z i |z i 1 ) is introduced iteratively by stacking WAEs at each latent layer. The corresponding graphical model for N = 3 is given Figure 1b .

We focus on this simple Markovian latent-variable structure for the generative model as a proof point for STACKEDWAE. This simple modelling setup is famously difficult to train, as is discussed extensively in the VAE framework (see for example Burda et al. (2015) ; Sønderby et al. (2016) ; Zhao et al. (2017) ). The difficulty in training such models comes from the Markovian latent structure of the generative model; in particular, the difficulty of learning structure in the deeper latents. This is because, to generate samples x ⇠ p ⇥ (x), only the joint p ⇥ (x, z 1 ) is needed as

p ⇥ (x) = R Z1 p ✓1 (x|z 1 )p ⇥ (z 1 )dz 1 .

Learning a smooth structure in each latent layer is not a strict requirement for learning a good generative model, however, it is sought after if the latent is to be used downstream or interpreted. We find empirically (see Section 3.1) that a better generative model is also achieved when the latent hierarchy is well learnt all the way down. Sønderby et al. (2016) sought to overcome the difficulty of learning a deep-latent hierarchy by using deterministic bottom-up inference, followed by top-down inference that shared parameters with the generative model. With additional optimisation tricks (e.g. KL annealing), their deeper-latent distributions would still go unused for sufficiently deep-latent hierarchies (discussed in Maaløe et al. (2019) ). In order to get deeper hierarchies of latents, Maaløe et al. (2019) introduced additional deterministic connections in the generative model as well as bidirectional inference network to facilitate the deep information flow needed to ensure the usage of the deeper-latent layers.

While the approach in Maaløe et al. (2019) is well motivated and achieves excellent results, we choose the OT framework for training deep-latent models due to its topological properties (see Arjovsky et al. (2017) ). Still, the standard WAE encounters the same difficulties as the VAE in learning deep-latent hierarchies. We thus modify the original WAE objective, effectively stacking WAEs, to improve the learning of both the generative model and the inference distribution throughout the latent hierarchy.

2.2. WASSERSTEIN AUTOENCODERS:

The Kantorovich formulation of the OT between the true-but-unknown data distribution P D and the model distribution P ⇥ , for a given cost function c, is defined by: OT c (P D , P ⇥ ) = inf

EQUATION

where P(P D , P ⇥ ) is the space of all couplings of P D and P ⇥ ; namely, the space of joint distributions on X ⇥ X whose densities have marginals p D and p ⇥ :

EQUATION

In the WAE framework introduced by Tolstikhin et al. (2018) , the space of couplings is constrained to joint distributions , with density , of the form:

(x, x) = Z Z1...Z N p ⇥ (x|z 1 , . . . , z N ) q 1 (z 1 , . . . , z N |x) p D (x) dz 1 . . . dz N (5)

where Q 1 (Z 1 , . . . , Z N |X) plays the same role as the variational distribution in variational inference.

Marginalising over x in Eq. ( 5) automatically gives p D (x), however the second marginal constraint in Eq. (4) (that over x giving p ⇥ ) is not guaranteed. Due to the Markovian structure of the generative model, a sufficient condition for satisfying the second marginal constraint is (see Appendix A.1):

Z X q 1 (z 1 |x)p D (x)dx = Z Z2...Z N p ✓2 (z 1 |z 2 ) . . . p ✓ N (z N 1 |z N )p(z N )dz 2 . . . dz N def = p ⇥ 2:N (z 1 )

(6) Finally, to get the WAE objective of Tolstikhin et al. (2018) , the marginal constraint of Eq. ( 6) is relaxed using a Lagrange multiplier (see Appendix A.2):

c W c (P D , P ⇥ ) = inf Q 1 (Z1|x) " 1 D 1 ✓ Z X Q 1 (Z 1 |x) p D (x) dx, P ⇥ 2:N (Z 1 ) ◆ (7) + Z X ⇥X Z Z1 c(x, x) p ✓1 (x|z 1 ) q 1 (z 1 |x) p D (x) dx dx dz 1 #

where D 1 is any divergence function and 1 a relaxation parameter. Note that the inf is taken only over q 1 (z 1 |x) instead of the full q 1 (z 1 , . . . , z N |x) because the expression does not depend on z >1 .

While Eq. ( 7) is in-principle tractable (e.g. for Gaussian q 1 (z 1 |x) and sample-based divergence function such as MMD), it only depends on the first latent Z 1 . Thus it will learn only a good approximation for p ⇥ (x, z 1 ) = p ✓1 (x|z 1 )p(z 1 ), rather than a full hierarchy with each latent living on a smooth manifold. We show empirically in Section 3.1 that Eq. ( 7) is indeed insufficient.

2.3. STACKING WAES FOR DEEP-GENERATIVE-LATENT VARIABLE MODELLING:

In theory (e.g. 1 ! 1), Eq. ( 7) does not depend on the choice of divergence D 1 . However, given the set of approximations used, a divergence that takes into account the smoothness of the full stack of latents will likely help the optimisation. We now show that by using the Wasserstein distance itself for D 1 , we can derive an objective that naturally pairs up inference and generation at every level in the deep-latent hierarchy. After all, the divergence in Eq. ( 7) is between an aggregate distribution:

Q agg 1 (Z 1 ) def = Z X Q 1 (Z 1 |x) p D (x) dx (8)

from which we can only access samples, and an analytically-known distribution P ⇥ 2:N (Z 1 ), which is analgous to where we started with the Wasserstein distance between P D and the full P ⇥ .

Specifically, we choose D 1 in Eq. ( 7) to be the relaxed Wasserstein distance c W c1 , which following the same arguments as before, requires the introduction of a new variational distribution Q 2 (Z 2 |Z 1 ):

D 1 ⇣ Q agg 1 (Z 1 ), P ⇥ 2:N (Z 1 ) ⌘ = inf Q 2 (Z2|z1) " 2 D 2 ⇣ Q agg 2 (Z 2 ), P ⇥ 3:N (Z 2 ) ⌘ (9) + Z Z1⇥Z1 Z Z2 c 1 (z 1 , z1 ) p ✓2 (z 1 |z 2 ) q 2 (z 2 |z 1 ) q agg 1 (z 1 ) dz 1 dz 1 dz 2 #

where the notation in Eq. ( 6) has been re-used for the prior P ⇥ 3:N (Z 2 ), and similarly for the aggregated posterior,

Q agg 2 (Z 2 ) def = R Z1 Q 2 (Z 2 |z 1 )q agg 1 (z 1 )

, as in Eq. ( 8). Just as before, Q 2 (Z 2 |Z 1 ) does not need to provide a distribution over the z >2 without loss of generality.

The divergence D 1 that arose in Eq. ( 7) between two distributions over Z 1 is thus mapped onto the latent at the next level in the latent hierarchy, Z 2 , via Eq. ( 9). This process can be repeated by using c W c2 again for D 2 in Eq. ( 9) to get an expression that maps to Z 3 , requiring the introduction of another variational distribution Q 3 (Z 3 |Z 2 ). Repeating this process until we get to the final layer of the hierarchical-latent-variable model gives the STACKEDWAE objective:

W STACKEDWAE (P D , P ⇥ ) = inf Q (Z1,...,Z N |x)  N Y i=1 i D N ⇣ Q agg N (Z N ), P (Z N ) ⌘ (10) + N 1 X n=0  n Y i=1 i Z Zn⇥Zn Z Zn+1 c n (z n , zn ) p ✓n+1 (z n |z n+1 ) q n+1 (z n+1 |z n ) q agg n (z n ) dz n dz n dz n+1 !

where we denote (Z 0 , Z 0 , z 0 ) = (X , X, x) and we define the empty product Q 0 i=1 i def = 1. Each p ✓n is the n th layer of the generative model given in Eq. ( 1). The q n 's are the inference models introduced each time a WAE is "stacked", which combine to make the overall STACKEDWAE Markovian inference model given in Eq. ( 2), and the aggregated posterior distributions are defined as

Q agg 0 def = P D and Q agg n (Z n ) def = Z Zn 1 Q n (Z n |z n 1 ) q agg n 1 (z n 1 ) dz n 1 (11)

Note that the STACKEDWAE objective function in Eq. ( 10) still requires the specification of a divergence at the highest latent layer D N , which we simply take to be the MMD as originally proposed by Tolstikhin et al. (2018) . Other choices can be made, as in Patrini et al. (2018) , who choose a Wasserstein distance computed using the Sinkhorn algorithm (Cuturi, 2013) . While Patrini et al. (2018) provide a theoretical justification for the minimisation of a Wasserstein distance in the prior space, we found that it did not result in significant improvement and comes at an extra efficiency cost. Similarly, one could choose different cost functions c n at each layer; for simplicity we take all cost functions to be the squared Euclidean distance in their respective spaces.

3. EXPERIMENTS:

We now show how the STACKEDWAE approach of Section 2 can be used to train deep-latent hierarchies without customizing the generative or inference models (e.g. skip connections, parameter sharing) . We also show through explicit comparison that the STACKEDWAE approach performs significantly better than the standard WAE when training deep-hierarchical-latent models. 2a , samples from the generative model in Figure 2b , and latent space interpolations in Figure 2c where digits are reconstructed from points in Z 5 taken evenly in a grid varying between ±2 standard deviations from the origin. STACKEDWAE generates compelling samples and learns a smooth manifold. Note that the choice d Z5 = 2 allows for easy visualisation of the learned manifold, rather than being the optimal dimension for capturing all of the variance in MNIST.

Figure 2c shows that STACKEDWAE manages to use all of its latent layers, capturing most of the covariance structure of the data in the deepest-latent layer, which is something that VAE methods struggle to accomplish (Sønderby et al., 2016) . Figure 3a shows the encoded input images through the latent layers, with corresponding digit labels coloured. We see through each layer that STACKEDWAE leverages the full hierarchy in its latents, with structured manifolds learnt at each stochastic layer.

An advantage of deep-latent hierarchies is their capacity to capture information at different levels, augmenting a single-layer latent space. In Figure 3b , input images, shown in the bottom row, are encoded and reconstructed for each latent layer. More specifically, the inputs are encoded up to the latent layer i and reconstructed from the encoded z i using the generative model p(x|z 1 )p(z 1 |z 2 ) . . . p(z i 1 |z i ). Row i in Figure 3b shows the reconstruction obtained from encoding up to layer i. We can see that each additional encoding layer moves slightly farther away from copying the input image as it moves towards fitting the encoding into the 2-dimensional, unit-normal prior distribution. The dimensionality of the deeper-latent layers is a modelling choice which determines how much information is preserved; this can be seen through the loss of information from deeper reconstructions in Figure 3b . Indeed, in each layer, the encoder is asked to map the incoming encodings into a lower-dimensional latent space, filtering the amount of information to keep and pass up to the deeper layer. Thus, if there is a mismatch in the dimensions between the true underlying generative process of the data and the chosen model, the encoder will have to project the encodings into lower-dimensional space, losing information along the way.

ABLATION STUDY: STACKEDWAE VERSUS STANDARD WAE

In this section, we compare STACKEDWAE with the original WAE framework for training generative models with deep-hierarchical latents. In particular, we train a WAE using the objective defined in Eq. ( 7) and an inference distribution as in Eq. ( 2); the corresponding graphical model is shown in Figure 1c . We use the same the experimental setup as outlined earlier in this section, and the same parametrised networks.

The results are shown in Figure 4 , with reconstructions in Figure 4a , generated samples in Figure 4b , and deepest-latent interpolations in Figure 4c . The latter two should be compared directly with Figure 2b and Figure 2c , respectively, for the STACKEDWAE. The quality of the generated samples from the WAE is poor in comparison with that of the STACKEDWAE.

The lack of smooth interpolations in Figure 4c shows that almost no structure has been captured in the deepest-latent layer. This is likely due to the fact that the standard WAE, with objective given in Eq. ( 7), is independent of the deeper-latent inference distributions, thus reducing the smoothness in the deeper layers. The relatively accurate reconstructions in Figure 4a indicate that the model managed to capture most of the structure of the data in the first latent layer. This behaviour is similar to that of the Markov HVAE as described in Zhao et al. (2017) . They show that, for Markov HVAEs to learn interpretable latents, one needs additional constraints beyond simply maximising the likelihood.

3.2. STREET VIEW HOUSE NUMBERS:

We now show that STACKEDWAE is able to leverage a deep-latent hierarchy on more real-world datasets, in particular Street View House Numbers (SVHN) (Netzer et al., 2011) . We trained a 6-layer STACKEDWAE with inference networks and generative models at each latent layer taken to be Gaussian distributions with mean and covariance functions parametrised by 3-layer ResNet-style (Kaiming et al., 2015) neural networks. The details are given in Appendix B.2.

As before, we choose c n (z n , zn ) = kz n zn k 2 L 2 , and compute the expectations in Eq. (10) analytically whenever possible, otherwise using Monte Carlo sampling.

For simplicity, our architecture choices constrain the dimensionality on the latent space (details are in Appendix B.2). The latent dimension is given by the size of the output feature map at that layer times the number of these feature maps, leading to latent dimensions of: 16⇥16⇥1, 8⇥8⇥2, 4⇥4⇥4, 4⇥4⇥2, 4⇥4⇥1 and 2⇥2⇥2. These make for relatively high-dimensional latent spaces for the first few latent layers. Rubenstein et al. (2018) observed that when training WAEs with high-dimensional latent space, the variance of the inference network tends to collapse to 0. The authors argue that this might come either from an optimisation issue or from the failing of the divergence used to regularise the latents. Either way, the collapse to deterministic encoders results in poor sample quality as the deterministic encoder is being asked to map the input manifold into a higher-dimensional space than its intrinsic dimension. One solution proposed in Rubenstein et al. (2018) is to include in the objective function Eq. (10) a regulariser that maintains a non-zero variance. As in Rubenstein et al. (2018) , we include a log-variance penalty term as given in Eq. ( 12):

EQUATION

We find that ⌃ i = 10 (2+i) for i = 1, . . . , 6 works well in our setting. This choice has been motivated by the fact that the bigger the latent dimension, the more likely it is that d Z > d I , where d I and d Z are the intrinsic dimension of the input and the latent dimension, respectively. Given our 6b ), the first latent layers will suffer more from the collapse to deterministic encoder than the deeper ones with lower dimension.

Our results on SVHN are given Figure 5 . Samples from the generative model are shown in Figure 5a .

Interpolations across the deepest-latent layer are shown in Figure 5b with the anchor data points in the first and last columns; the interpolations are smooth as expected, though perhaps due to the choice of low-dimensional deepest layers, the reconstructions (second and second-to-last columns) may fail to match perfectly the anchor data points (first and last columns). The encoded latent space is shown in Figure 5c for each latent layer, where the structure of SVHN can be seen; it is likely that the most salient structure in these embeddings represents the background colour, rather than the central digit, but that is to be expected from SVHN. Finally, Figure 5d shows the reconstructions of the data points (along the bottom row) at each latent layer in the hierarchy. Similarly to the results for MNIST, we can see that the deepest latent layer may not be large enough to enable high-fidelity reconstructions.

Our intention is to show that the hierarchy of latents can be learnt, which is clearly the case, rather than to model SVHN perfectly, so we do not attempt to tune to the optimal latent dimensionality.

4. CONCLUSION:

In this work we introduced a novel objective function for training generative models with deep hierarchies of latent variables using Optimal Transport. Our approach recursively applies the Wasserstein distance as the regularisation divergence, allowing the stacking of WAEs for arbitrarily deep-latent hierarchies. We showed that this approach enables the learning of smooth latent distributions even in deep latent hierarchies, which otherwise requires extensive model design and tweaking of the optimisation procedure to train. We also showed that our approach is significantly more effective at learning smooth hierarchical latents than the standard WAE.

A STACKEDWAE DERIVATION DETAILS:

A.1 MARGINAL CONSTRAINT The space of couplings 2 P(P D , P ⇥ ) that defines the OT distance in Eq. ( 3) is constrained according to Eq. ( 4). WAE assumes a joint density of the form given in Eq. ( 5), which automatically satisfies the p D marginal constraint, but requires the further sufficient condition given in Eq. ( 6) in order to satisfy the p ⇥ marginal constraint. To see that Eq. ( 6) is indeed a sufficient condition for the p ⇥ marginal constraint, note that from the Markovian assumption of the generative model (see Eq. ( 1)), we can write as 8

EQUATION

The constraint in Eq. ( 4) on the p D marginal is trivially true as the integral over the second variable can be brought inside the integral over Z 1 , after which all the integrals simply integrate to unity leaving p D .

The constraint on the second marginal is obtained by integrating Eq. ( 13) over the first variable, 8x 2 X , Z

EQUATION

Thus, to satisfy Eq. ( 4), we need:

EQUATION

def = Z Z1 p ✓1 (x|z 1 ) p ⇥ 2:N (z 1 ) dz 1

where the definition of the generative model from Eq. ( 1) was used and we introduced:

EQUATION

To satisfy Eq. ( 15), one obvious sufficient condition on the aggregated posterior distribution Q agg 1 (Z 1 ) defined in Eq. ( 8) is Eq. ( 6), namely that

EQUATION

which is what we sought out to show. However, Eq. ( 17) is a sufficient condition, not a necessary one: indeed Eq. ( 17) must only hold under R Z1 P ⇥ ( X|z 1 )dz 1 . So for example, if P ✓1 ( X|z 1 ) = P ✓1 ( X), then Eq. ( 15) would boil down to a constraint only on the expectations of Q agg 1 (Z 1 ) and P ⇥ 2:N (Z 1 ).

A.2 WAE OBJECTIVE:

Starting from the definition of the OT distance given in Eq. (3), and using the WAE approach with density written as Eq. ( 13), we find:

OT c (P D , P ⇥ ) = inf

EQUATION

 inf Given that the above does not depend on z >1 , the inf can be written over Q 1 (Z 1 |X) rather than the full Q 1 (Z 1 , Z 2 , . . . Z n |X). Replacing the constraint in the inf with the sufficient condition according to Eq. ( 6), which amounts to replacing Eq. ( 15) with Eq. ( 17), we obtain:

Q 1 (Z1,Z2,...Zn|X), R Z 1 p⇥(x|z1)q agg 1 (z1)dz1=

OT c (P D , P ⇥ )  inf

EQUATION

Eq. ( 7) is then obtained by relaxing constraint in Eq. ( 19); replacing the hard constraint by a soft constraint via a penalty term added to the objective, weighted by a 1 :

EQUATION

+ 1 D 1 ✓ Q agg 1 (Z 1 ), P ⇥ 2:N (Z 1 )

◆ where D 1 is any divergence function between distributions on Z 1 .

B EXPERIMENTS B.1 MNIST EXPERIMENTS:

We train a deep-hierarchical latent-variable model with N = 5 latent layers whose dimensions are d Z1 = 32, d Z2 = 16, d Z3 = 8, d Z4 = 4 and d Z5 = 2, respectively. We parametrise the generative and inference models as:

q i (z i |z i 1 ) = N z i ; µ q i (z i 1 ), ⌃ q i (z i 1 ) , i = 1, . . . , 5 p ✓i (z i 1 |z i ) = N z i 1 ; µ p i (z i ), ⌃ p i (z i ) , i = 2, . . . , 5

p ✓1 (x|z 1 ) = x f ✓1 (z 1 )

For both the encoder and decoder, the mean and diagonal covariance functions µ i , ⌃ i are fullyconnected networks with 2 same-size hidden layers (consider f ✓1 as µ p 1 ). For i = 1, 2, 3, 4, 5, the number of units is 2048, 1024, 512, 256, 128, respectively. For the regularisation hyperparameters, we use Q n i=1 i = n rec /d zn for n = 1, . . . , 4 (for each reconstruction term in the objective), and Q 5 i=1 i = match (for the final divergence term). We then perform a grid search over the 25 pairs ( rec , match ) 2 {0.01, 0.05, 0.1, 0.5, 1}⌦{10 4 , 10 3 , 10 2 , 10 1 , 1} and find the best result (smallest Eq. ( 10)) is obtained with ( rec , match ) = (0.05, 10 4 ).

We choose the squared Euclidean distance as the cost function: c n (z n , zn ) = kz n zn k 2 L 2 . The expectations in Eq. ( 10) are computed analytically whenever possible, and with Monte Carlo sampling otherwise.

We use batch normalisation (Ioffe & Szegedy, 2015) after each hidden fully-connected layer, followed by a ReLU activation (Glorot et al., 2011) . We train the models over 5, 000 epochs using Adam optimiser (Kingma & Ba, 2015) with default parameters and batch size of 128.

:



Figures and Tables:

Figure 1: (a) Generative model (blue lines represent generative model parameters). (b) Inference model used in STACKEDWAE. (c) Standard WAE; the generative model has only one latent with prior p(z 1 ) = R p(z 1 |z 2 )p(z 2 |z 3 )p(z 3 )dz 1 dz 2 dz 3 . (d) Inference model with skips connections and parameter sharing with the generative model, as in Sønderby et al. (2016).

deep-hierarchical-latent variable model with N = 5 latent layers on raw (non-binarised) MNIST (LeCun & Cortes, 2010). The latent layers have dimensions: d Z1 = 32, d Z2 = 16, d Z3 = 8, d Z4 = 4 and d Z5 = 2. We chose Gaussian distributions for both the generative and inference models, except for the bottom layer of the generative model, which we choose to be deterministic, as in Tolstikhin et al. (2018). The mean and covariance matrices used are parametrised by fully connected neural networks (see Appendix B.1 for details). We choose the squared Euclidean distance cost function: c n (z n , zn ) = kz n zn k 2 L 2 . Expectations in Eq. (10) are computed analytically where possible, and otherwise with Monte Carlo sampling. LEARNING A DEEP-LATENT HIERARCHY Our results are shown in Figure 2, with the training curves in Figure

Figure 2: 5-layer STACKEDWAE. (a) Training curves; each term in Eq. (10) is shown throughout training. (b) Model samples. (c) Z 5 latent-space interpolations.

Figure 5: 6-layer STACKEDWAE. (a) Model samples. (b) Points interpolations; the first and last columns are actual data points. (c) Visualisations of latent spaces Z i , as in Figure 3a. (d) Reconstructions for different encoding layers, as in Figure 3b.

✓1 (x|z 1 )q 1 (z 1 |x)p D (x)dz 1 dxdx

Peer review (iclr2020):
In this paper, a hierarchical extension to Wasserstein Autoencoders (WAE) is proposed, where the latent variables are stacked in a multi-layer structure. In the proposed model, the divergence function in WAE is viewed as a relaxed WS distance. Therefore, another layer of WAE can be stacked to minimise the WS distance. In this way, a hierarchical model can be built to learn hierarchical representations. I think the idea of viewing the divergence in WAE as a relaxed WS distance and then minimising it with another WAE structure is interesting, intuitive and straightforward. However, the advantages of the proposed model over WAE and VLAE (S.Zhao et.al 2017) are less obvious to me. It is a bit hard for me to tell whether the hierarchical latent variables help to improve quantitative results, generate better images, or learn intuitive hierarchical representations, which is the main reason that I go to mild rejection. For example, I would expect to see similar things as in VLAE, where the representations in different layers capture hierarchical structures or disentanglements. But in the proposed model, it seems to be hard to see the differences between the hierarchical representations such as in Figure 3(b). Also in the two-dimensional visualisation of Figure 3(a), it is a bit hard for me to intuitively understand what the representations really capture. From the graphical model point of view, the proposed model is a hierarchical Gaussian model and the inference (although with WAE) is in the flavour of Gibbs sampling, which propagates information layer-wisely from bottom up. Conventionally, a hierarchical Gaussian model is hard to work with many layers such as 5. Therefore, I may suggest improving in case of fewer layers.