Learning Deep-Latent Hierarchies by Stacking Wasserstein Autoencoders
Authors: Benoit Gaujac
Ilya Feige
David Barber
Venue:  (2020)
Keywords: Generative modelling
Optimal Transport

Abstract: Probabilistic models with hierarchical-latent-variable structures provide state-of-the-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.

Under review as a conference paper at ICLR 2020
LEARNING DEEP-LATENT HIERARCHIES BY STACK- ING WASSERSTEIN AUTOENCODERS
Anonymous authors Paper under double-blind review
ABSTRACT
Probabilistic models with hieraachiica--atent-vviideee structures provide state-of- the-art results amongst non-autoregressivve 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, won-likelihood-baae -asedd framework for training generative models with appealing theoretical properties, in principle allowing easier training convergence between distributions. In this work we propose 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 in-so-doing, it is more effective than the original Wasserstein Autoencoder with Maximum Mean Discrepancy divergence.
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-variabl models use only one stochastic latent layer. Auto-regressee models (Larochelle & Murray 2011; Van Den Oord et al., 2016b at Salimans et | 2017), for example, produce state-of-the-ar samples and likelihood scores. However, ruto-regressive models suffer from poor scalability to sigh--imensional 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.
lece-latent-variabee 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-of- the-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. (017). to simply
the poor expressiveness of the models used ZZha 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) (016). sharing generative model and inference network parameters (Sonderby et al. 2016),, well as bidirectional inference network (Maaløc et al., 2019). Maaloc et al. (2019) managed to train very deep-hierarchical. state-of-the-arr 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.
models achieving near
Optimal Transport (OT) (Villani 2008) is mathematical framework to compute distances between distributions. Recently, it has been used as training method for generative models (Genevay et al.)

---

Under review as a conference paper at ICLR 2020
2018; Bousquet al. 2017). Tolstikhin 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 non- linellikekhood based framework with appealing topological properties (Arjovsky et al. 2017, Bousquet et al. 2017), in theory allowing for easier training convergence between distributions, Gaujac al. (2018) trained two-latent-layer generative model using WAE, showing promising results in the capacity of the WAE framework to leverage latent hierarchy in generative modelling.
Following these early successes, we propose to train deep-hierarchical-Hane 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 Soonderby et al. (2016) Maaloe et al. (2019), we believe that deep-latent hierarchy offers the potential improve generative models, if they could be trained properly. In order to leverage the deep-latent hierarchy, we derive novel objective function by stacking WAEs, introducing an inference network 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-Latee 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-laten varie models with a deep hicrarchy of latents. They can be minimalistically simple in their inference and generative models, but trained using OT in novel way. We start by defining the probabilistic models considered in this work, then introduce OT, and finally discuss how train probabilistic models with deep-latent hierarchies using OT, the method that we refer to as STACKED
Throughout this paper, 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)).
calligraphic letters (e.g. X) for sets, capital letters (e.g. X) for random
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 Pe the parametric model with N latent layers, where θ = {O1...., }N we have:
(0())= = / Zn...EN
po, (α|²1) po, (2)122) ... Pa (ZN-1 p(ZN) dzl ...1....dyy. daN
(1)
where the data is X E X' and the latent variables are Z, E 2n and we chose p(2N) N(xx:
The corresponding graphical model for N = 3 given Figure la
We will be using variational inference through the WAE framework of Tolstikhin et al. (0118) introducing variational distributions, go (21,..., n (r), 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, qu (2), (Zi,..., 2N (x) can have a Markovian latent hierarchy
when following the STACKEDWAE approach. That is,
w.l.o.g. with STACKEDWAE
(2 2N
4o (21 96. (22/21) JoN (2))
(2)
where each gox (z/ 11) 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 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): (Senderyy et al. (d.](20))) Zhao et al. (0077). The difficulty in training such models comes from the Markovian latent structure
2

---

Under review as a conference paper at ICLR 2020
Z3
7
Z3
93
2
22
Z2
Q2
Z
Z1
q1
91
(a) Our gen. model
(b) Our inf. model
(c) Inf. without stacking
(d) Alternative inf.
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 2==) = f p(z1 (2)))((()))))))))))))))) zs )p(p))((())) (2)=))2z d22d23. (d) Inference model with skips connections and parameter sharing with the generative model, as in Sønderby al. (2016).
of the generative model; in particular, the difficulty of learning structure in the deeper latents. This is because, to generate samples I pz(z), only the joint pe (r, 21) is needed as pe(x) = Jz, Pe, (x|z) )pe )dzj. Learning a smooth structure in each latent layer not strict requirement for learning good generative model, however, 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 bierarchies (discussed in Maalee et al. (0199). 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-latens 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 Pn and the model distribution Pe, for a given cost function c, is defined by:
f (x
OT Pe) = inf TEP(Po.Pe)
where P(Pp. Pe) is the space of all couplings of Pp and Pe; namely, the space of joint distributions on X x whose densities have marginals pD and pe:
P(Po, Pe) P (x,
{rl
pD()) and / (x,, (1x dx (
pe(
(4)
In the WAE framework introduced by Tolstikhin et al. (0188), the space of couplings is constrained to joint distributions I, with density 7, of the form:
(x, = / p 2N) qs, (21, ......,,, en |x) (D dz ... day
(5)
ZN...ZN
where Qs, (Zi,..,, ZN| Xx) plays the same role as the variational distribution in variational inference.
3

---

Under review as a conference paper at ICLR 2020
Marginalising over è in Eq. automatically gives pD(x), however the second marginal constraint in Eq. 4 (that over æ giving pe) is not guaranteed. Due to the Markovian structure the generative model, sufficient condition for satisfying the second marginal constraint is (see Appendix A.1)
J q (211)) (21)2)pD())d=
poa (21 ... pew (2N-11 (p(N )d22 daN def pe (21)
Jz...ZN
(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 ).2)
(1 Qの (Z1lz) (x( dx, Pe (2))
W.(PD.Po) =
inf Qo (z))))
X1 DI
/ c(z, σ) Po. 1 4s. (21 (z) (x) dx dē dzp
xxxx 22
where D2 is any divergence function and A relaxation parameter. Note that the inf is taken only over 9o, (21 (a:) instead of the full go, (Z1..... 2N( because the expression does not depend on 2>1-
While Eq. 7 in-principle tractable (e.g. for Gaussian 9á, (21 |z) and sample-based divergence function such as MMD), it only depends on the first latent Z1. Thus it will Iearn only a good approximation for pe (r, 21) = po, (2121 )p(z), rather than a full hierarchy with each latent living on smooth manifold. We show empirically in Section 3.1 that Eq. 7 is indeed insufficient.
2.3 STACKING WAES FOR DEEP-GENERATIVE-LATEEE VARIABLE MODELLING
In theory (e-g. X + oo), Eq. does not depend on the choice of divergence Dj. 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, 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:
Qs (Z1 PD(x) dr (1 def /
Q
from which we can only access samples, and an analytically-knon known distribution Pe (Z1). which is analgous to where we started with the Wasserstein distance between PD and the full Pe.
Specifically, choose Dj in Eq. 7 to be the relaxed Wasserstein distance W which following the same arguments as before, requires the introduction of a new variational distribution Q² (Z2 Z1):
Di (Zz), PezI (Z,)(Z)))
inf (2212)))))
\ D2 22
28R (Zz), Pe (Zz))
$ 2, 22
C1 (2 21) pe, (21 aoz (22)21) Q1 222 (z1) dæj dā dz2
where the notation in Eq. (6) has been re- re-used for the prior Pe,. (Zz), and similarly for the aggregated
posterior, Q (Zz) def Jz, Q. Z2 2 )4
88y (z1). as in Eq. (8. Just as before, Q (22) Zi) does not need to provide: a distribution over the 2>2 without loss of generality.
The divergence D that arose in Eq. 7 between two distributions over Zı is thus mapped the latent at the next level in the latent hierarchy, Zz, via Eq. (9). This process can be repeated by
using W., again for D2 in Eq. (9) to get an expression that maps to Za. requiring the introduction of
another variational distribution Q (Z ZZ). Repeating this process until we get to the final layer of the hicrarchical-Latee latent-variable model gives the STACKEDWAE objective:
N
W WTTACKEDWAE(P... =
inf Qo(21...N11a ZN|z)
II 1=1
DN (O MER (ZN), (NN),,) P(ZN))
(10)
/
Ca(Eno Ponsa (F (znlznzz) Gouss (Zn+1l 28 (zn) dz, dā d2n+t1
2nx2,, Zut1
4

---

Under review as a conference paper at ICLR 2020
where we denote (Zo; 20, =0) = (X, x,)) and we define the empty product I A def Each pe, is the 1hhh layer of the generative model given in Eq. | The qo., 's are the inference models introduced each time WAE is "stacked", which combine make the overall STACKEDWAE Markovian inference model given in Eq. ② and the aggregated posterior distributions are defined as
Qo def PD
and
(Zn) def [ Eu
Qon (Znlen-n-
REE (2n džn-1
(11)
Note that the STACKEDWAE objective function in Eq. (1O) still requires the specification of divergence at the highest latent layer Dx. which we simply take to 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 theoretical justification for the minimisation of a Wasserstein distance 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 en each layer; for simplicity take all cost functions to be the squared Euclidean distance in their respective spaces.
3 EXPERIMENTS
We show how the STACKEDWAE approach of Section 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-bierarchhicll-ltton models.
3.1 MNIST
EXPERIMENTAL SETUP
We trained heep-hhirarchhiclllltee variable model with N = 5 latent layers raw (non-binarised) MNIST (LeC'nn & Cortes 2010). 'he latent layers have dimensions: dz, = 32, dz, = 16, dz, = 8, dz, and dz, 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 for details).
We choose the squared Euclidean distance cost function: Cn, žn) = || E..1.22. Expectations 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 2a samples from the generative model in Figure 2b and latent space interpolations in Figure 2c where digits reconstructed from points in Z5 taken evenly in grid varying between 2 standard deviations from the origin. STACKEDWAE generates compelling samples and learns smooth manifold. Note that the choice dz, allows for easy visualisation of the learned manifold, rather than being the optimal dimension for capturing all of the variance in MNIST.
Figure 2 shows that STACKED 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 STACKEDWAF leverages the full hierarchy in its latents, with structured manifolds Icarnt at each stochastic layer.
An advantage of deep-latent hierarchies is their capacity to capture information at different levels, aug- menting 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 2; using the generative model p( 21 )p( (22)... p 1(22-11/ 12,). Row i in Figure [3b 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 it moves
5

---

Under review as a conference paper at ICLR 2020
-AW,
5
4
607
A.W.
2,W,
G
8
4
A.W; MnDy
0 6
5
7
6.5
9 19000
0
48
443
9 999
-22
/
7
0 2
999999551
5555
3
1
8
7
S 55553
5
9
3
9
0
2
S S 5 5
5 8
7
060
5
7
9
7
0 / 2
6
4
7
0
5
0
2
2
0
5
7
0
3
8
/
6
3
7
7
0
0
Epachs
3000
500
2
5
8
(a)
(b)
(c)
Figure 2: 5-layer STACKEDWAE. (a) Training curves; each term in Eq.
10
is shown throughout
training. (b) Model samples. (c) Zg latent-space interpolations.
Latent space 21
Latent space Z:
Latent space Z:
Latent space 24
Latent space 2ç
(a)
0
2
6
6
6
0
3
6
500000555
2
9
7
6
6
0
6
8
0
6
8
4
/
6
0
02992722
4
50050055200
0
4
8
4
6
0
4
50
3
2
0
6
8
6
9
92 2
6
50
3
7
0
6
(b)
Figure 3: 5-layer STACKEDWAE (a) Visualisations of latent spaces Z, Each colour corresponds to digit label. dz, 2 can be directly plotted; for higher dimensions we use UMAP (Mcinnes & Healy 2018) plot two dimensional representation.
(b) Reconstructions for different encoding layers. The bottom row is data; the 1h row from the bottom is generated using the latent codes 23 which are from the th encoding layer.
3 1
7
2
6
0
9 0 0
R
6
4
2
3 22523
1
4
4
0
R
0 9
7
4
8
27
4
7
9
{
1
2
#
2
6
4
9
/
5
2
3 t
131111
221
5
7
6 4
5
8 5
19999990
#9
f
5
6 4
8
/
5
/
3
54 1
4
8
O
8797
Y Y 4/339
2
4
4
5536
2 7
8
5900
4 3
6
/
4
6
22
6
0
9
5
0
0
0
9
7
5 8
0
0 1
4
/
5554
44 /
8
5
9
0
I
2 6
5
6
6
9 2 9 7 6
0
3 3 8
336
0
2
0
t
2
3
0
6
4
9 9
7 6
2
6
8
9 3
9000000
8
2
2
3
/
3
5
6
3
2
(a)
(b)
(c)
Figure 4: 1-layer implicit-prior WAE. (a) Reconstructions (within pairs of rows, data is above with the corresponding reconstructions below). (b) Model samples. (c) 2g latent-space interpolations.

---

Under review as a conference paper at ICLR 2020
towards fitting the encoding into the 2-dimensional, unit-normal prior distribution. The dimension- ality of the deeper-laten 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 3bb 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,
train WAE using the objective defined in Eq. 7 and an inference distribution as in Eq. (2) the corresponding graphical model is shown in Figure lc 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 Hc 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). 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 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 deep-latent hierarchy on more real-world datasets, in particular Street View House Numbers (SVHN) (Netzer et al.] (011).. 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 6n (Zn, 2n) = || - 2. 11 and compute the expectations Eq. 10 analyti- cally 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: 16x16x1, 8x8x2, 4x4x4, 4x4x2, 4x4x1 and 2x2x2. 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 than its intrinsic dimension. One solution proposed in Rubensstei al. (2018) is include in the objective function Eq. (10) regulariser that maintains a non-zero variance. As in Rubenstein et al. (2018), we include a varraareence penalty term as given in Eq. (12:
space
N
6pen
log E [m]
(12)
1=
T2
We find that X motivated by the fact that the bigger the latent dimension, the more likely it is that dz > dz, where dz and dz are the intrinsic dimension of the input and the latent dimension, respectively. Given our
(2++) for i = 1, 1,... 6 works well in our setting. This choice has been
7

---

Under review as a conference paper at ICLR 2020
4
8
8
8
2
2
2
0
2
1
1
2
2
3
B
[
2
2
3
3
2
6
0
3
2
2
6
6
6
6
6
3
13
3
3
3
3
3
3
(a)
(b)
Latent space
Latent space
Latent space
Latent space 24
Latent space E5
Latent space
(c)
3
186
6
6
4
5
9
7
0
100
0
3
66
6
111
7
6
3
1
0
5
8
3
8
177
8
0
3
12
0
3
8
3
8
9
8
3
152
3
8
3
9
659
8
3
12
18
3
9
659
8
8
9
(d)
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 2,, as in Figure 3a (d)
Reconstructions for different encoding layers, in Figure
latent-diimensio choices (see Table 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 Samples from the generative model are shown in Figure Sa (oggureBbwar 55]with the anchor data points in the first and last columns; the interpolations are smooth as expected, though perhaps due to the choice of now-dimensional deepest layers. the reconstructions (second and second-to-las columns) may fail to match perfectly the anchor data points (first and last columns). The encoded latent space is shown in Figure for each latent layer, where the structure of SVHN can 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 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, can see that the deepest latent layer may not be large enough to enable high-fidelit reconstructions. Our intention is to show that the hierarchy of latents 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.
Interpolations across the deepest-latent layer are shown in Figure
4 CONCLUSION
In this work we introduced a novel objective function for training generative models with deep hierar- chies 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.
8

---

Under review as a conference paper at ICLR 2020
REFERENCES
M. Arjovsky, S. Chintala, and L. Bottou. Wasserstein generative adversarial networks, In International Conference on Machine Learning. 2017.
Bachman. An architecture for deep, hierarchical generative models. In Advances in Neural Information Processing Systems, 2016.
O. Bousquet, S. Gelly, I. Tolstikhin, C. J. Simon-Gabriel, and B. Schölkopf. From optimal transport to generative modeling: the VEGAN cookbook. 2017.
Y. Burda, Grosse R., and R. Salakhutdinov. Importance weighted autoencoders. In International Conference on Learning Representations, 2015.
M. Cuturi. Sinkhorn distances: lightspeed computation of optimal transport. In Advances in Neural Information Processing Systems, 2013.
L. Dinh, J. Sohl-Dickstein, and S. Bengio. Density estimation using real NVP. In International Conference on Learning Representations, 2016.
B. Gaujac, I. Feige, and D. Barber. Gaussian mixture models with Wasserstein distance. arXiv:1806.04460 04465, 2018.
A. Genevay, G. Peyre, and M. Cuturi. Learning generative models with Sinkhorn divergences. In International Conference Artificial Intelligence and Statistics, 2018.
X. Glorot, A. Bordes, and Bengio. Deep sparse rectifier neural networks. International Conference Artificial Intelligence and Statistics, 2011.
A. Gretton, K. M. Borgwardt, J. Rasch, Schölkopf, and A. Smola. kernel two-sample test. In Journal of Machine Learning Research, 2012.
S. loffe and C. Szegedy. Batch normalization:
Accelerating deep network training by reducing internal covariate shift. In International Conference on Machine Learning, 2015.
H. Kaiming. Z. Xiangyu, R. Shaoqing, and S. Jian. Deep residual learning for image recognition. In IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2015.
D. P. Kingma and J. Ba. Adam: a method for stochastic optimization. In International Conference Learning Representations, 2015.
D. P. Kingma and P. Dhariwal. Glow: Generative flow with invertible Ixl convolutions. In Advances in Neural Information Processing Systems, 2018.
D. P. Kingma and M. Welling. Auto-encoding variational Bayes. In International Conference on Learning Representations, 2014.
D. Kingma, T. Salimans, R. Jozefowicz, X. Chen, I. Sutskever, and M. Welling. Improved varia- tional inference with inverse autoregressive flow. In Advances in Neural Information Processing Systems, 2016.
H. Larochelle and Murray. The neural autoregressive distribution estimator. In International Conference on Artificial Intelligence and Statistics, 2011.
A. B. L. Larsen, K. Sønderby S., H. Larochelle, and O. Winther. Autoencoding beyond pixels using a learned similarity metric. In International Conference on Machine Learning, 2016.
Y. LeCun and C. Cortes. MNIST handwritten digit database. 2010.
L. Maaløe, M. Fraccaro, V. Liévin, and O. Winther. BIVA: very deep hierarchy of latent variables for generative modeling. In Advances in neural information processing systems, 2019.
L. McInnes and J. Healy. UMAP: uniform manifold approximation and projection for dimension
reduction. arXiv: 03426, 2018.

---

Under review as a conference paper at ICLR 2020
Y. Netzer, T. Wang, A. Coates, A. Bissacco, B. Wu, and A. Y. Ng. Reading digits in natural images with unsupervised feature learning. In Advances in Neural Information Processing Systems Workshop Deep Learning and Unsupervised Feature Learning. 2011.
G. Patrini, M. Carioni, Forré, S. Bhargav, M. Welling, Van Den Berg, T. Genewein, and F. Nielsen. Sinkhorn autoencoders. arXiv: 1810.01118, 2018.
D. J. Rezende and S. Mohamed. Variational inference with normalizing flows. In International Conference on Machine Learning. 2015.
D. J. Rezende, S. Mohamed, and D. Wierstra. Stochastic backpropagation and approximate inference in deep generative models. In International Conference Machine Learning, 2014.
K. Rubenstein, B. Schölkopf, and I. Tolstikhin. On the latent space of Wasserstein auto-encoders. In Workshop track International Conference on Learning Representations, 2018.
T. Salimans, A. Karpathy, X. Chen, and D. P. Kingma. Pixelcnn++: Improving the PixeICNN with discretized logistic mixture likelihood and other modifications. In International Conference on Learning Representations, 2017.
C. K. Sønderby, T. Raiko, L. Maalóe, S. K. Sønderby, and O. Winther. Ladder variational autoencoders. In Advances in neural information processing systems, 2016.
I. Tolstikhin, O. Bousquet, S. Gelly, and B. Schoelkopf. Wasserstein auto-encoders. In International Conference Learning Representations, 2018.
A. Van Den Oord, N. Kalchbrenner, and Kavukcuoglu. Pixel recurrent neural networks. In International Conference on Machine Learning. 201
A. Van Den Oord, N. Kalchbrenner, O. Vinyals, L. Espeholt, A. Graves, and K. Kavukcuoglu. Conditional image generation with PixcICNN decoders. Advances in Neural Information Processing Systems, 2016b.
C. Villani. Optimal Transport: Old and New. Springer Berlin Heidelberg, 2008
Zha, Song, and Ermon. Towards deeper understanding of variational autoencoding models. 2017.
S. Zhao, J. Song, and S. Ermon. Learning hierarchical features from deep generative models. In International Conference on Machine Learning, 2017.
10

---

Under review as a conference paper at ICLR 2020
A STACKEDWAE DERIVATION DETAILS
A. MARGINAL CONSTRAINT
The space of couplings I' P(PD. Pe) that defines the OT distance in Eq. 3 is constrained according to Eq. 4 WAE assumes joint density of the form given in Eq. 3 which automatically satisfies the pD marginal constraint, but requires the further sufficient condition given in Eq. (6) in order to satisfy the pe marginal constraint. To see that Eq. ( is indeed a sufficient condition for the pe marginal constraint, note that from the Markovian assumption of the generative model (see Eq. (T), we can write as
(a, 2) E X X.
70, अ) / Pdd =1) qoi (EI...., AN (D( dz1 daN
ZI...N
(13)
pad (5|=1) do (z1la) (D(() dz
3,
The constraint in Eq. 4 on the marginal is trivially true as the integral over the second variable be brought inside the integral over 21, after which all the integrals simply integrate unity leaving PD-
The constraint on the second marginal obtained by integrating Eq. (3) over the first variable, Vr E X,
J 7(x, x) dx /x] [ pe, (T|2) Ma. (2))))p da
= / pa, (ā]zz) / Io (2) (x)
(14)
21
Thus, to satisfy Eq. 4 we need:
Vī E X
[ pa (5)())) / gor (z) dzl need
pe pe(
(15)
dsl 11
(21)
def / po, (T)z1) PezN (21) day
where the definition of the generative model from Eq. used and introduced:
Vzz E 2,
Pea. (21)
def
po, (2. ..pp (zN-1lZN) p(2x) dz2 daN
(16)
Zx...IN
To satisfy Eq. 15 one obvious sufficient condition on the aggregated posterior distribution Q260 defined in Eq. is Eq. (6). namely that
V21 E Zi.
MEB 11
(21) = P02.N (21)
(17)
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 Sz, Pe 21 dzz So for example, if Pa, (X|21) ) = Pe, (X),
then Eq. (T5)) would boil down to constraint only the expectations of QA (Z1) and Pea.y (71).

---

Under review as a conference paper at ICLR 2020
A.2 WAE OBJECTIVE
Starting from the definition of the OT distance given in Eq. (3. and using the WAE approach with density y written as Eq. (13) we find:
I Jxxx
Po)
inf (eP(Po.Ps)
c(x, (x,
(18)
Sexx c(x, 8) / po, (z)21)90, (2) (211z)) dz) dadž
<
inf Qu, 22......... X), pe 2) 1 (2))
[2, Pe (a))
(21 /2
Given that the above does not depend on 2>1. the inf can written over Qs (Z) rather than the full QΦ, (71, 72 Zn Zn|Xx 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:
Jexx clx,
/ po, (z)z1)g0, (2 (=)rr) dædō
OTe(PD, Pe) s
inf Qoy (Z))), 2
(19)
Eq. 7 is then obtained by relaxing constraint in Eq. 19 replacing the hard constraint by soft constraint via penalty term added to the objective, weighted by a Ay;
xxx Iex
/ c(2, (8121) do (2112) pD (z) dz dodō
Wc(PD.
inf (01()))))
(20)
+ + Di ( Pea.N (21)
where Dj is any divergence function between distributions Z1.
B EXPERIMENTS
B. MNIST EXPERIMENTS
We train a deep-hierarchical latent-variable model with 5 latent layers whose dimensions are d2, 32, dz, 16, dz, 8, dz, and d2, 2, respectively. We parametrise the generative and inference models as:
do, (zs 11) -N (2- (=-1)), i= 1.,,..,, pe. (2-11 (2:) -((2)1111((2)) (2))) i= 2.....,
(21)
pa, (z)21) 8(x fo, (21))
For both the encoder and decoder, the mean and diagonal covariance functions p4, E, are fully- connected networks with 2 same-size hidden layers (consider fo, as uf). For 2, 4, 5, the number of units is 2048, 1024, 512, 256, 128, respectively.
For the regularisation hyperparameters, we use I[ A = /d, for n = 1...., 4 (for each recon- struction term in the objective), and I A grid search over the 25 pairs (Arec Amatch,) E {0.01, 0.05, 0.1, 0.5, 1)@ 10-3 10 20-2.100 10-1,11111 1} and find the best result (smallest Eq. (1O)) is obtained with (Arec Amarch) = (0.05, 10 +).
Amatch (for the final divergence term). We then perform
We choose the squared Euclidean distance as the cost function: (x Z,) = Zn - 2n 1..1122.. The expectations in Eq. (10) are computed analytically whenever possible, and with Monte Carlo sampling otherwise.
We use batch normalisation (loffe & Szegedy 2015) after each hidden fully-connected layer. followed by ReLU activation (Glorot et al. (2011). We train the models over 0000 epochs using Adam optimiser (Kingma & Ba (015) with default parameters and batch size of 1 128.
12

---

Under review as a conference paper at ICLR 2020
Inputs
Layer
Filters dim.
Resampling
Ouiput dim.
Cony. Block
Layer
5x5x64
down / up
16x
Layer
3x3x66
down / up
8x8x2
Cony, Ressmpling
Residual Connection
Layer
3x3x64
down / up 4x4x4
Layer 4
3x3x 128
None
4x4x2
Layer
3x8x
None
4x4x1
BatehNorm--Hie ReLU
Layer 6
3x3x 128
down /
2x2x2
Couvolation
Outposs
(a)
Figure 6: (a) Residual network with 3 hidden convolutions. (b) Details of the architecture use in
Section 3.2
B.2 SVHN EXPERIMENT
SVHN SETUP
We train a 6-layers STACKEDWAE on the SVHN dataset using both the training dataset (73,257 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. 2013) neural networks.
A M-llyeer residual network is composed of M 1 convolutional blocks followed by resampling convolution, and residual connection. The outputs of the two are added and last operation (either fully connected or convolution layer) is applied on the result. A convolutional block is composed of convolution layer followed by batch normalisation (loffe & 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 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 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 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 specifically, we use M convolutional blocks with the dimensions of the filters specified in 6b) Networks in layers 1, 2, 3 have 96 convolution filters while whose in layers 4, 5, have 128 filters, each filters 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 performed. The latent layers 1, 2, and 6 have stride of 2 and we choose the number of features to be 1, 4, 1, for the latent layers 1.... 6 (see Table 6b for the full details).
For the regularisation hyperparameters, we use IP >; = A 1)/+1)/+1 for = 1,...,5 ,,5 (each reconstruction term in the objective), and 11 A; Amatch (the final 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 A 1) of the corresponding
reconstruction terms, we choose the weights to scale like ( We found that (Arec Ammtch)) =
(10-1,10 100 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.
13


Reviews:
[{"review_id": "rJgzs5lwtr", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 1, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The paper aims to develop a deep generative model, which -unlike VAEs or GANs- comprises a hierarchy of latent variables rather than a direct map from the stochastic latent manifold to the observation space. To this end, the paper builds a training objective based on nesting the Wasserstein distance between the data distribution and its estimation arbitrarily many times. The generated objective corresponds naturally to a deep hierarchical generative model.\n\nThe principled approach followed to achieve the objective is solid and elegant. It is also intuitive and matches nicely with some valid observations highlighted in the paper such as insufficiency of by-passing intermediate latent variables (sentence above the Sec 2.3 title).\n\nOne major weakness of the paper is that it lacks a sufficient argumentation about how it differentiates from earlier attempts to nest Wasserstein distances. For instance,\n\nY. Dukler et al., \"Wasserstein of Wasserstein Loss for Learning Generative Models\", ICML, 2019\n\nApart from the theoretical argumentation, the paper should also compare their solution to this prior work on a number of benchmarks.\n\nAnother major weakness is that the paper lacks a quantitative evaluation scheme for its success. The experiments section starts with the claim that the proposed method \"significantly\" improves on the WAE, which I fail to see on the plots. \n\nLastly, Having said that the proposed method is novel and elegant, it is still a straightforward extension of the existing and well-known Wasserstein Auto-Encoder (WAE) approach. It extends WAEs by repetitively applying the tricks proposed by this earlier work, putting aside some minor additional adjustments.\n\nMinor on style: The abstract does not give any single hint about the methodological novelty of the work.\n\n---\nPost-rebuttal: Thanks to authors for their effort for clarifications. Yet, I'm afraid the author response does not touch at all to any of the concerns I have raised. There are well-known ways to compare the success of generative models, FID being one of them as the authors point out. Another could be the test log-likelihood of a synthetic data set the true distribution of which can be predesigned. I understand the issues the authors raise about the difficulties in comparing generative models, but I kindly disagree with the attitude that there are no ways to compare, so we are obliged to live with qualitative comparisons. If a one-score comparison is not enough, the right way to go is to provide multiple scores. If direct metrics are not feasible, one should go for indirect ones, but should still provide outcomes a reader can reproduce.", "extra_scores": {"experience_assessment": "I have read many papers in this area.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I made a quick assessment of this paper.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I did not assess the derivations or theory."}, "extra_text": {}}, {"review_id": "BJx5f3iAKr", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "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.\n\nI 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.\n\nFor 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. \n\nFrom 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.", "extra_scores": {"experience_assessment": "I have read many papers in this area.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}, {"review_id": "BJllJxEH5B", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper presents a deep, latent variable model for unsupervised data modeling problems. The problem with such latent, deep generative models is that they are difficult to train reliably. In this paper, the authors provide an approach based on stacked Wasserstein autoencoders to train deep latent variable models. Experimental results are demonstrated on various image datasets and the latent codes are demonstrated to have an interpretable meaning. \nI like the inference techniques in the paper and like the ideas presented in this paper.", "extra_scores": {"experience_assessment": "I do not know much about this area.", "review_assessment:_thoroughness_in_paper_reading": "I made a quick assessment of this paper.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "N/A"}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper received 6, 3, 1. The main criticism is the lack of quantitative evaluation/comparison. The rebuttal did not convince the last reviewer who strongly argues for a comparison. The authors are encouraged to add additional results and resubmit to a future venue.
Decision: Reject