L EARNING D EEP -L ATENT H IERARCHIES BY S TACKING W ASSERSTEIN A UTOENCODERS
Anonymous authors
Paper under double-blind review
A BSTRACT
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.
I NTRODUCTION
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 efﬁcient 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.,
; 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.
, 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 S TACKED WAE.
Our contributions are two-fold: ﬁrst, we introduce S TACKED WAE, a novel objective function based on OT, designed speciﬁcally for generative modelling with latent hierarchies, and show that it is able to fully leverage its hierarchy of latents.
Second, we show that S TACKED WAE performs signiﬁcantly 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).
S TACKED WAE
S TACKED WAEs 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 deﬁning the probabilistic models considered in this work, then introduce
OT, and ﬁnally discuss how to train probabilistic models with deep-latent hierarchies using OT, the method that we refer to as S TACKED WAE.
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. G ENERATIVE 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:
[Equation]
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 S TACKED WAE 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 S TACKED WAE.
This simple modelling setup is famously difﬁcult to train, as is discussed extensively in the VAE framework (see for example Burda et al.
Sønderby et al.
; Zhao et al.
The difﬁculty in training such models comes from the Markovian latent structure [Figure] [Equation] [Equation]
[Equation] [Equation] [Equation] [Equation]
[Equation]
[Equation]
of the generative model; in particular, the difﬁculty 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 Zp ✓ 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 ﬁnd 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 difﬁculty 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 sufﬁciently deep-latent hierarchies (discussed in Maaløe et al.
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 ﬂow 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.
Still, the standard WAE encounters the same difﬁculties 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. W ASSERSTEIN A UTOENCODERS
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 deﬁned by:
[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:
[Equation]
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 sufﬁcient condition for satisfying the second marginal constraint is (see Appendix A.1):
[Equation]
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):
[Equation]
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 ﬁrst 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 insufﬁcient.
2. S TACKING WAE S 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:
[Equation]
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 ⇥ .
Speciﬁcally, we choose D 1 in Eq.
(7) to be the relaxed Wasserstein distance c W c 1 , which following the same arguments as before, requires the introduction of a new variational distribution Q  2 ( Z 2 | Z 1 ) :
[Equation]
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 [Equation] Z [Equation]
=
R [Equation] Q  2 ( Z 2 | z 1 ) q agg 1 ( z 1 ) , as in Eq. 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.
This process can be repeated by using c W c 2 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 ﬁnal layer of the hierarchical-latent-variable model gives the S TACKED WAE objective:
[Equation]
where we denote ( Z 0 , Z 0 , z 0 )
=
( X , X, x ) and we deﬁne the empty product[-(cid:81) 
-]
[-i =
 λ i def-][Equation] = 1 .
Each p ✓ n is the n th layer of the generative model given in Eq.
The q  n ’s are the inference models introduced each time a WAE is "stacked", which combine to make the overall S TACKED WAE
Markovian inference model given in Eq.
(2), and the aggregated posterior distributions are deﬁned as
[Equation]
and
[Equation]
Note that the S TACKED WAE objective function in Eq.
(10) still requires the speciﬁcation 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.
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 justiﬁcation for the minimisation of a Wasserstein distance in the prior space, we found that it did not result in signiﬁcant improvement and comes at an extra efﬁciency 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.
E XPERIMENTS
We now show how the S TACKED WAE 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 S TACKED WAE approach performs signiﬁcantly better than the standard WAE when training deep-hierarchical-latent models.
3. MNIST E XPERIMENTAL SETUP
We trained a deep-hierarchical-latent variable model with N = 5 latent layers on raw (non-binarised)
MNIST (LeCun & Cortes, 2010).
The latent layers have dimensions:
d Z 1 = 32 , d Z 2 = 16 , d Z 3 = 8 , d Z 4 = 4 and d Z 5 = 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.
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 , ˜ z n ) = k z n  ˜ z n k 2 L 2 .
Expectations in Eq.
(10) are computed analytically where possible, and otherwise with Monte Carlo sampling.
L EARNING A DEEP - LATENT HIERARCHY
Our results are shown in Figure 2, with the training curves in Figure 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.
S TACKED WAE generates compelling samples and learns a smooth manifold.
Note that the choice d Z 5 = 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 S TACKED WAE 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 S TACKED WAE 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 speciﬁcally, 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
[Figure]
[Figure] [Figure]
towards ﬁtting 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, ﬁltering 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.
A BLATION STUDY : S TACKED WAE VERSUS STANDARD WAE
In this section, we compare S TACKED WAE with the original WAE framework for training generative models with deep-hierarchical latents.
In particular, we train a WAE using the objective deﬁned 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 S TACKED WAE.
The quality of the generated samples from the WAE is poor in comparison with that of the S TACKED WAE.
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 ﬁrst latent layer.
This behaviour is similar to that of the Markov HVAE as described in Zhao et al.
They show that, for Markov HVAEs to learn interpretable latents, one needs additional constraints beyond simply maximising the likelihood.
3. S TREET V IEW H OUSE N UMBERS
We now show that S TACKED WAE 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 S TACKED WAE 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 , ˜ z n ) = k z n  ˜ z n 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 ⇥ 2 , 4 ⇥ 4 ⇥ 1 and 2 ⇥ 2 ⇥ 2 .
These make for relatively high-dimensional latent spaces for the ﬁrst 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.
[Equation]
We ﬁnd 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
[Figure]
latent-dimension choices (see Table 6b), the ﬁrst 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 ﬁrst 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 (ﬁrst 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-ﬁdelity 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.
C ONCLUSION
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 signiﬁcantly more effective at learning smooth hierarchical latents than the standard WAE.
A S TACKED WAE DERIVATION DETAILS A. M ARGINAL CONSTRAINT
The space of couplings  2 P ( P D , P ⇥ ) that deﬁnes the OT distance in Eq.
(3) is constrained according to Eq.
WAE assumes a joint density of the form given in Eq.
(5), which automatically satisﬁes the p D marginal constraint, but requires the further sufﬁcient condition given in Eq.
(6) in order to satisfy the p ⇥ marginal constraint.
To see that Eq.
(6) is indeed a sufﬁcient condition for the p ⇥ marginal constraint, note that from the Markovian assumption of the generative model (see
Eq.
(1)), we can write  as
[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 ﬁrst variable, 8 x 2 X , [Equation]
[Equation]
Thus, to satisfy Eq.
(4), we need:
[Equation]
where the deﬁnition of the generative model from Eq.
(1) was used and we introduced: 8 z [Equation] Z 1 ,
[Equation]
To satisfy Eq.
(15), one obvious sufﬁcient condition on the aggregated posterior distribution Q agg 1 ( [Equation] 1 ) deﬁned in Eq.
(8) is Eq.
(6), namely that
[Equation]
which is what we sought out to show.
However, Eq.
(17) is a sufﬁcient condition, not a necessary one: indeed Eq.
(17) must only hold under
[Equation] So for example, if P
[Equation]
P [Equation] 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. WAE OBJECTIVE
Starting from the deﬁnition of the OT distance given in Eq.
(3), and using the WAE approach with density  written as Eq.
(13), we ﬁnd:
[Equation]
Given that the above does not depend on z > [Equation] 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 sufﬁcient condition according to Eq.
[Equation] which amounts to replacing Eq.
(15) with Eq.
[Equation] we obtain:
[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]
where D 1 is any divergence function between distributions on Z 1 .
B E XPERIMENTS B. MNIST EXPERIMENTS
We train a deep-hierarchical latent-variable model with N = 5 latent layers whose dimensions are d Z 1 = 32 , d Z 2 = 16 , d Z 3 = 8 , d Z 4 = 4 and d Z 5 = 2 , respectively.
We parametrise the generative and inference models as:
[Equation]
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
[Equation]
=1  i =  n rec /d
z n for n = 1 , . .
, 4 (for each reconstruction term in the objective), and Q 5 i [-=
-]=1  i =
match (for the ﬁnal 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 ﬁnd 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 , ˜ z n ) = k z n  ˜ z n 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 .
[Figure]
Layer i
Filters dim.
Resampling
Output dim.
Layer5 ⇥ 5 ⇥down / up 16 ⇥ 16 ⇥
[-Layer
 ×-]Layer3 ⇥ 3 ⇥down / up 8 ⇥ 8 ⇥
[-Layer
-][Figure]
Layer3 ⇥ 3 ⇥ [Figure] down /[-up 4 × 4 × Layer
-]
[Figure] 4 ⇥ 4 ⇥ [Figure]
[-Layer
-]Layer
[Figure]
⇥ [Figure] None[-4 × 4 × Layer
-]
[Figure] None [Figure]
B. SVHN EXPERIMENT SVHN SETUP
We train a 6-layers S TACKED WAE on the SVHN dataset using both the training dataset (73, digits) and the additional training dataset (531,131 digits).
The mean and covariance functions of the inference networks and generative models are parametrised by 3-layer ResNet-style (Kaiming et al., 2015) neural networks.
A M -layer residual network is composed of M  1 convolutional blocks followed by a resampling convolution, and a residual connection.
The outputs of the two are added and a last operation (either fully connected or convolution layer) is applied on the result.
A convolutional block is composed of a convolution layer followed by batch normalisation (Ioffe & Szegedy, 2015) and a ReLU non-linearity (Glorot et al., 2011).
We also use batch normalisation and ReLU after the sum of the convolutional blocks output and the residual connection.
See Figure 6a for an example of a 3-layers residual network with a last convolution operation.
When resampling, we use a convolution layer with stride 2 in both the skip connection and the resampling covolution for the inference networks and a deconvolution layer with stride 2 in both the skip connection and the resampling convolution in the generative models.
If no resampling is performed, then the resampling convolution is a simple convolution layer with stride 1 and the skip connection performs the identity operation.
The latent dimensions are then given by the dimensions and the number of features in the last convolutional operation.
More speciﬁcally, we use M =convolutional blocks with the dimensions of the ﬁlters speciﬁed in Table 6b.
Networks in layers 1 , 2 ,have 96 convolution ﬁlters while whose in layers 4 , 5 , 6 have 128 ﬁlters, each ﬁlters having the same size within each residual network, and doubling (in the inference networks) or divide by 2 (in the generative models) their number in the resampling convolution if any resampling is performed.
The latent layers 1 , 2 , 3 and 6 have a stride of 2 and we choose the number of features to be 1 , 2 , 4 , 2 , 1 ,for the latent layers i = 1 . . .
6 (see Table 6b for the full details).
For the regularisation hyperparameters, we use[-(cid:81) pi =
-]
[Equation]
for p = 1 , . . .
, 5 (each reconstruction term in the objective), and[-(cid:81) 6 i =
-]
Q [Equation]
match (the ﬁnal divergence term).
The choice for the reconstruction weights is motivated by the fact that the effective regularisation hyperparameters scale exponentially.
Thus, to avoid the collapse (or blowing up for  rec > 1 ) of the corresponding reconstruction terms, we choose the weights to scale like O [Equation]
.
We found that [Equation]
10  1 , 10  worked well with our experimental setup.
We train the models over 1000 epochs using Adam optimiser (Kingma & Ba, 2015) with default parameters and batch size of 100 .