Self-Reflective Variational Autoencoder
Authors: Ifigeneia Apostolopoulou
Elan Rosenfeld
Artur Dubrawski
Venue:  (2021)
Keywords: deep generative models
variational inference
approximate inference
variational auto encoder

Abstract: The Variational Autoencoder (VAE) is a powerful framework for learning probabilistic latent variable generative models. However, typical assumptions on the approximate posterior distributions can substantially restrict its capacity for inference and generative modeling. Variational inference based on neural autoregressive models respects the conditional dependencies of the exact posterior, but this flexibility comes at a cost: the resulting models are expensive to train in high-dimensional regimes and can be slow to produce samples. In this work, we introduce an orthogonal solution, which we call self-reflective inference. By redesigning the hierarchical structure of existing VAE architectures, self-reflection ensures that the stochastic flow preserves the factorization of the exact posterior, sequentially updating the latent codes in a manner consistent with the generative model. We empirically demonstrate the advantages of matching the variational posterior to the exact posterior---on binarized MNIST self-reflective inference achieves state-of-the-art performance without resorting to complex, computationally expensive components such as autoregressive layers. Moreover, we design a variational normalizing flow that employs the proposed architecture, yielding predictive benefits compared to its purely generative counterpart. Our proposed modification is quite general and it complements the existing literature; self-reflective inference can naturally leverage advances in distribution estimation and generative modeling  to improve the capacity of each layer in the hierarchy.

Under review as a conference paper at ICLR 2021
SELF-REFLECTIVE VARIATIONAL AUTOENCODER
Anonymous authors
Paper under double-blind review
ABSTRACT
The Variational Autoencoder (VAE) is a powerful framework for learning prob- abilistic latent variable generative models. However, typical assumptions on the approximate posterior distributions can substantially restrict its capacity for infer- ence and generative modeling. Variational inference based on neural autoregres- sive models respects the conditional dependencies of the exact posterior, but this flexibility comes at cost: the resulting models are expensive to train in high- dimensional regimes and can be slow to produce samples. In this work, we in- troduce an orthogonal solution, which we call self-reflective inference. By re- designing the hierarchical structure of existing VAE architectures, self-reflection ensures that the stochastic flow preserves the factorization of the exact posterior, sequentially updating the latent codes in a manner consistent with the generative model. We empirically demonstrate the advantages of matching the variational posterior to the exact posterior- binarized MNIST self-refiective inference achieves state-of-the-art performance without resorting to complex, computation- ally expensive components such as autoregressive layers. Moreover, we design a variational normalizing flow that employs the proposed architecture, yielding pre- dictive benefits compared its purely generative counterpart. Our proposed mod- ification is quite general and it complements the existing literature; self-reflective inference can naturally leverage advances in distribution estimation and generative modeling to improve the capacity of each layer in the hierarchy.
INTRODUCTION
The advent of deep learning has led to great strides in both supervised and unsupervised learning. One of the most popular recent frameworks for the latter is the Variational Autoencoder (VAE), in which a probabilistic encoder and generator are jointly trained via backpropagation to simultane- ously perform sampling and variational inference. Since the introduction of the VAE (Kingma & Welling 2014), or more generally, the development of techniques for low-variance stochastic back- propagation of Deep Latent Gaussian Models (DLGMs) (Rezende et al. 2014), research has rapidly progressed towards improving their generative modeling capacity and/or the quality of their varia- tional approximation. However, as deeper and more complex architectures are introduced, care must be taken ensure the correctness of various modeling assumptions, whether explicit or implicit. In particular, when working with hierarchical models it is easy to unintentionally introduce mismatches the generative and inference models, to the detriment of both. In this work, we demonstrate the ex- istence of such a modeling pitfall common to much of the recent literature on DLGMs. We discuss why this problem emerges, and we introduce simple-yet crucial-modifican architectures to address the issue.
to the existing
Vanilla VAE architectures make strong assumptions about the posterior citrrbutioonsspcifffeee is standard to assume that the posterior is approximately factorial. More recent research has inves- tigated the effect of such assumptions which govern the variational posterior (Wenzel et al., 2020) or prior (Wilson & Izmailov 2020) in the context of uncertainty estimation in Bayesian neural net- works. In many scenarios, these restrictions have been found to be problematic. A large body of recent work attempts to improve performance by building a more complex encoder and/or decoder with convolutional layers and more modern architectures (such as ResNets (He et al., (016)) (Sal- imans al. 2015) Gulrajani et al.) 2017) or by employing more complex posterior distributions constructed with autoregressive layers (Kinmma et al.] 2016:] Chen et al., 2017). Other work (Tom-)
czak & Welling 2018 Klushyn al. 2019a) focuses on refining the prior distribution of the latent

---

Under review as a conference paper at ICLR 2021
codes. Taking a different approach, hierarchical VAEs (Rezende al. 2014 Gulrajani al. 2017 Sønderby et al. 2016 Maalde et al. 2019;; Klushyn et al. 2019b) leverage increasingly deep and in- terdependent layers of latent variables, similar to how subsequent layers in discriminative network are helieved to learn more and more abstract representations. These architectures exhibit superior generative and reconstructive capabilities since they allow for modeling of much richer latent spaces. While the benefits of incorporating hierarchical latent variables is clear, all existing architectures suf- fer from a modeling mismatch which results in sub-optimal performance: the variational posterior does not respect the factorization of the exact posterior distribution of the generative model.
In earlier works on hierarchical VAEs (Rezende al.) 2014), inference proceeds bottom-up, counter to the top-down generative process. To better match the order of dependence of latent variables to that of the generative model, later works (Sønderby et al. [016 BBachman (016) split inference into two stages: first deterministic bottom-up pass which does necessary precomputation for evidence encoding, followed by stochastic top-down pass which incorporates the hierarchical latents to form closer variational approximation to the exact posterior. Crucially, while these newer architectures ensure that the order of the latent variables mirrors that of the generative model, the overall varia- tional posterior does not match because of the strong restrictions on the variational distributions of each layer.
Contributions. In this work, we propose to restructure common hierarchical VAE architectures with a series of bijective layers which enable communication between the inference and generative networks, refining the latent representations. Concretely, our contributions are as follows:
We motivate and introduce a straightforward rearrangement of the stochastic flow of the model which addresses the aforementioned modeling mismatch. This modification substantially com- pensates for the observed performance gap between models with only simple layers and those with complex autoregressive networks (Kingma et al. 2016 Chen al. 2017
We formally prove that this refinement results in hierarchical VAE whose variational posterior respects the precise factorization of the exact posterior. To the best of our knowledge, this is the first deep architecture to do without resorting computationally expensive autoregressive components or making strong assumptions (e.g., diagonal Gaussian) on the distributions of each layer (Sønderby et al. 2016) assumptions that lead to degraded performance.
We experimentally demonstrate the benefits of the improved representation capacity of this model, which stems from the corrected factorial form of the posterior. We achieve ofate-of-the-ar perfo- mance on MNIST among models without autoregressive layers, and our model performs par with recent, fully autoregressive models such as Kingma et al. (2016). Due to the simplicity of architecture, we achieve these results for fraction of the computational cost in both training and inference.
We design a hierarchical variational normalizing flow that deploys the suggested architecture in order to recursively update the base distribution and the conditional bijective transformations. This architecture significantly improves upon the predictive performance and data complexity of a Masked Autoregressive Flow (MAF) PPapamakarios et al. 2017) CIFAR-10.
Finally, it should be noted that our contribution is quite general and can naturally leverage recent advances in variational inference and deep autoencoders (Chen et al. 2017]] Kingma et al. 2016 Tomezak & Welling 2018, Burda et al., 2016) Dai & Wipf, 2019; van den Oord et al.) 2016a
Rezende & Viola 2018) as well as architectural improvements to density estimation ( (Gulrajani et al.
Dinh et al. 2017 Kingma & Dhariwal 2018,) Durkan et al., 2019 van den Oord et al. 201
Gregor et al. 2015). We suspect that combining our model with other state-of oaee--of-hh--nee art methods could
further improve the attained performance, which we leave to future work.
VARIATIONAL AUTONENCODERS
A Variational Autoencoder (VAE) Kingmma & Welling 2014 2019) is generative model which capable of generating samples æ E R from a distribution of interest p(=) by utilizing latent variables z coming from prior distribution p(z). To perform inference, the marginal likelihood
2

---

Under review as a conference paper at ICLR 2021
should be computed which involves integrating out the latent variables:
p(a) = p(x,, z) dz.
In general, this integration will be intractable and lower bound on the marginal likelihood is max- imized instead. This is done by introducing an approximate posterior distribution q(= x) and applying Jensen's 's inequality:
p(x) = log / pla, 2) dz dz= /
g(= x
/ g(> x)) log
p(æ x)p(z)
pla, 2) dz 2
dz
g(= x
<(= x))
- >logp(x) 2 // (og= z)] D (9(z x)) p(z)) & [(x; 0, Φ).
(2)
where 0, Φ parameterize p(a;, 2; 0) and g( x:; Φ) respectively. For ease of notation, we may omit 0, Φ in the derivations. This objective is called the Evidence Lower BOund (ELBO) and can be optimized efficiently for continuous 2 via stochastic gradient descent (Kingma & WWelling 2014 Rezende al. 2014).
REFLECIVEE VARIATIONAL INFERENCE
With this background, abilistic model which ensures that the variational posterior matches the factorization of the exact posterior induced by its generative model. We refer to this architecture the Self-Reflective Varia- tional Autoencoder (SeRe-Var VAE). We expound upon its components in the following subsections.
now ready introduce our main contribution: the first deep prob-
3.1 GENERATIVE MODEL
Figure displays the overall stochastic flow of the generative network. A detailed illustration of our model is provided in Figure S3
Our generative model consists of a hierarchy of L stochastic layers, as in Rezende et al. (2014). However, in this work, the data
(I1, 22. 1 ... . FL) E RP is partitioned into L blocks, with each layer generating only 2; E D: with D variables e E N; are first sampled from simple prior distribution (prior layer) and subsequently transformed to latent variables 21 E RM by bijective function fi R + RR
D. At each layer 1, Ni-dimensional latent
To distinguish between the two sets of latent variables in our model, throughout this paper we refer the latent codes. For example, for affine transformation
to 6 the base latent variables and fi the latent codes are given by 2 f((ct) = Q + (diag(d) ) + uuu?) x C,, with ct ui, dj E N and d; > 0 to ensure bijectivity. The latent codes 2 are subsequently passed the stochastic layer responsible for generating the observed data æ; (data layer).
Moreover, the layers in the hierarchy connected in three ways: i) the prior layer I can access the latent codes 2;-1 defining conditional distribution p(e: =1-1) ii) 21-1 is fed to the next bijection fi defining conditional transformation z; fi(E; z1-1) iii) the data layer receives the data block
- generated by the previous data layer defining conditional distribution p(æ: 21-1, 21-1). Intuitively, this choice justified because the latent codes 2 of layer I, conditioned on 21-1, will be successively refined based on how well _T- reconstructed æ;-1, yielding progressively more meaningful latent representations.
the following subsections, we describe these steps in detail. The joint distribution of the base latent variables (EI, ......, 6r.) and the observed data of the generative model is:
p(x, c) = p(e1) p(11/21) x II p(e 2111) x p(an 21-1, 21-))
(3)
1=2
3.2 INFERENCE MODEL
The inference network is identical to the generative network shown in Figure I except that the prior layers replaced by posterior layers, that are additionally conditioned the observed data ær, for
3

---

Under review as a conference paper at ICLR 2021
Figure D-separation be- tween stochastic layers. By the Bayes ball rule, all paths from él to Ea pass either through æi or 22, which D- separate them. Therefore, E1 IL E3 22, æ.
the generation of the base latent variables et. Specifically, the variational encoder of the SeRe-VAr is defined as follows:
1.
Q(EI, EE, x) = g(El 2)) x [[g(E 21-1, æ).
(4)
1=2
The formal justification of this factorization is deferred to section 3.3 Compared to other hicrar-
chical architectures, in the proposed model the inference layers are conditioned on the output of the preceding bijective layer these components are shared between the generative and the inference network (see also Figure (2). This choice allows for complex transformations of the latent variables and is theoretically motivated by the following proposition.
Proposition 1 Let p(c) and q(c) be two N-dimensional probability densities. Let f : :RN + RRN be an invertible, smooth transformation of the random variable с such that z f(c), yielding distributions p2) and o' (=) of respectively. Then, D DKI [/))) Dxr(a(E) p(e)).
Proof: From the definition of the Kullback-Leibler divergence and the change of variables formula Rudin 2006 Bogachev
20077:
log [x)2) = Eq(e) log (f(())
log ((c)x det J,(c) 1-1
Ey(e) log g(c) p(E) (5)
(x((s)
Eg(e)
p((2)
p(([(c))
p(E)x det ((()
where Jy(c) is the Jacobian matrix of f evaluated at e. Proposition]iipplne that the inclusion of the bijectors can help increase the conditional likelihood p(æ =) in equation 2 without increasing the KL term. Moreover-though not pursed in this work-it motivates the construction of normalizing flows for variational inference with non-linear time determinant of the Jacobian matrix, since the analytical form of the transformed distribution no longer needed for the computation of the KL-divergence. In this work, we assume Gaussian diagonal base distributions. In order to account for the two conditioning streams, the evidence æ and the latent factors 21-1, we employ a residual parametrization as described section 3.4.2
3.3 EXACT BAYES PROPAGATION
In this section, we provide the formal justification for the choice of equation we prove that backpropagation of our model preserves the factorization of the true posterior, without resorting to complex graph inversion as in Webb et al. (2018). We use the following straightforward lemma:
Lemma Let f NN + R be an invertible transformation such that both f and f-1 are differ- entiable everywhere. Then for any E R. p(e|z) = p=))f(=)).
Proof: By Bayes's Theorem and the change variables formula (Rudin) 2006 Bogachev 2007). p(E)f(z)) (=)()))e) x p(E) p(z|<) x det. (z x p(E) = p(E[z), p({())
p(=) x det (z)
where Jz( is the Jacobian matrix of f evaluated at 2, which has non-zero determinant by assumption.
We now present our main theoretical result, which says that the factorization our model's varia- tional posterior exactly matches that of the generative distribution.
4

---

Under review as a conference paper at ICLR 2021
Proposition 2 The factorization of the variational posterior defined in equation 4 torization of the exact posterior distribution induced by the generative model equation
respects the fac-
Proof: Let p(E1, . EL x) be the posterior distribution induced by the generative model defined in equation 3 as illustrated in Figure Then, according to the probability product rule the posterior distribution can be expressed
E2..., EL x) p(El
| x))
II 1=2
p(e: E<T,)),
(6)
where Edd A {EI, ......, E We will apply the Bayes ball rule (Jordan 2003 to simplify equation 6
Consider an arbitrary layer the hierarchy, Because fi-1 is a bijector, by Lemma
have
p(a | cet, 2) p(a (-1, p(c 2ı-1, Cc(-1,
Now, note that 6, is D-separated from 61-1,..., since all paths from e to 6<t pass through the observed nodes 21-1 21, 22,..., 21-1 (see Figure I for an example). Therefore, have
p(e) 21-1; 1, .) p(Ei 1))
Since this applies to every layer, it follows that the exact posterior equation 6 can also be expressed
I I[ #)-1,
p(El 62. . EL x) p(E1 x))
exactly matching the factorization of the approximate posterior in equation
3.4
IMPLEMENTATION DETAILS
3.4.1
AMORTIZED LAYERS
We use amortized parametrization to construct the conditional probability densities involved in the derivations above. In particular, for a probability density p(e 12 we take the parametrization
0 as a function of z: 0 = (zz). For example, conditional Gaussian distribution is defined as
p(E 2) =N(((z)) (=)) with ((z) = (J((2), (=)). The computational graph of an amortized
Gaussian layer is shown in Figure S4 Similarly, for a conditional bijector f'e z; ( we take 3 as
a function of z: 3 (zz). For example, for the affine bijector defined in section 3.1 we consider
(zz) (c(z), d(z), uz)).
3.4.2 RESIDUAL DISTRIBUTIONAL LAYERS
All but the first data layer p(æ;
21-1, and posterior layer q(Er ž1-1, a:) receive two streams
of conditioning factors one latent and observed. We ensure that each factor incrementally fines the distribution by adopting a residual parametrization. Here we describe the residual Gaussian distribution when conditioned the two factors z, Its probability density is given by
q(E|2, x) = ((u(z))o
which can be interpreted as follows. The first distribution (VHu(z), ((())))((()))) is corrected by the resid- uals Sa (xx), (xx);;; here we see the dependence on the conditioning factor æ. If x does not pro- vide additional information on (formally, p(c|2,æ)) p(c(z)), the two corrections collapse to and 0 respectively-that is, inducing no change. The reader may refer to Figure [S7 where we qualitatively illustrate the effect of the residual distributional layer that improves the conditional likelihood provided by the first one. To reduce the number of parameters, we consider networks for (22), (2)) that are shared between the prior and the posterior, yielding prior of the form p(E|z) = A [(p(z), (x)). Finally, we found experimentally that enforcing da (x: & 1 helps opti- mization by ensuring that only reduce the variance of the prior.
5

---

Under review as a conference paper at ICLR 2021
3.5 GENERAL REMARKS
Following the above analysis, we make some observations about the hierarchy of shared bijective layers in the model:
In contrast to Rezende et al. (0144) (see Figure(S]] in our model i) the prior layers are not indepen- dent, but rather are conditioned on the previous layers in the hierarchy; and ii) the transformational layers are restricted to be bijective.
The proposed model also differs from other hierarchical architectures GGurajaai et al. 2017 Sonderby et al. 2016 Maaløe et al. 2019); in these models the layers of the prior are condi- tioned upon the previous prior layers and not upon bijective layers that are shared between the generative and inference model.
One additional key difference between our model and all previous work is the coupling between the data layers. Therefore, the decoder can perceived layer-wise instead of pixel-wise autore- gressive rendering the sampling much more efficient (O(L) instead of O(D)). In section 4 we provide empirical results demonstrating the benefits of these modeling choices.
By reducing the set of conditioning variables from E<i to zt-1 theoretically justified man- ner, the hierarchical bijective layers offer convenient way to precisely and efficiently factorize the variational distribution, alleviating the bottleneck present in high-dimensional autoregressive approaches.
The model, albeit hierarchical, is less prone to posterior collapse, since each layer is responsible for the generation of different portion of the data. Experimental support for this observation is provided in Figure S8 where we plot the KL divergence for each layer of the architecture investigated in section 4.1.2
EXPERIMENTAL STUDIES
4.1 DYNAMICALLY BINARIZED MNIST
We empirically evaluate the SeRe- VAE on dynamically binarized MNIST. As in Burda et al. (2016); Sønderby al. (0166; Kingma al. (2016), the binary-valued observations sampled after each epoch with the Bernoulli expectations being set equal to the real, normalized pixel values in the dataset which prevents overfitting,
4.1.1 PERFORMANCE OF THE MLP VAR
To demonstrate that model's improved performance is due the restructuring of the stochas- tie flow and not sophisticated layers, we use simple multilayer- similarly forgo importance weighting (Burda et al. 2016). We adopt 10-layer architecture, with N₁ = 10 latent variables per layer, for a total of 100 latent features being passed to the decoder after being transformed by an affine bijector as described in section 3.1 We partition the image into = 10 equally sized blocks (except for the last one) from left to right a raster fashion. Finally, we use independent deterministic encoders for the data preprocessing. The full details of our imple- mentation are delegated to the supplementary material. We again emphasize the overall simplicity of our architecture, choosing instead to focus on the benefits of the corrected posterior factorization. As shown in Table I our model (SeRe-Vil VAE) outperforms existing models of the same complexity
(MLP) components; we
such as the DLGM and Ladder VAE (LVAE), those of higher complexity such as Inverse Autoregres- sive Flow (IAF), and models trained with importance weighted samples (IW-LVAE), Note that the architecture of the DLGM is identical to that of SeRe- to ensure a fair comparison, the DLGM was given larger feature maps the encoders to compensate for the additional bijective layer inputs in the BeRe-VAE. Therefore, the performance benefits are solely aftributed to the inclusion of the latent codes in subsequent stochastic layers in the hierarchy. Our model outperforms the LVAE models, despite using a smaller latent dimensionality (128 vs. 100) and being trained with a single importance sample. Moreover, our model exhibits superior performance compared to the autore- gressive IAF; this discrepancy could stem from the 1-layer architecture or the fact that standard normal prior was used. This result indicates that a prior of equivalent expressive capacity commu- nicating with the bijective layer could yield additional improvement. Finally, in our experiments the

---

Under review as a conference paper at ICLR 2021
Model
Details
log p(æ) 2
Self-Reflective
10 layers/ / 10 variables each, diagonal Gaussian prior
-81.17
Importance Weighted Ladder
layers / 128 variables total, #IW samples=10
-81.74
Ladder
5 layers / 128 variables total
-81.84
Self-Reflective IAF
10 layers/ 10 variables each, Standard Normal Prior
-81.96
Inverse Autoregressive Flow
layer/ 100 variables, Standard Normal Prior
-83.04
Deep Latent Gaussian Model 10 layers/ 10 variables each, diagonal Gaussian prior
-84.53
Relaxed Bernoulli VAEs
30 latent variables, exact factorization
-90
Table 1: Dynamically binarized MNIST Performance for VAEs without ResNet layers. 1000 im- portance samples were used for the estimation of the marginal likelihood. For the Ladder VAE performance, refer Table1 in Sønderby al. (2016). The models trained with a single importance sample unless otherwise noted (IW=l).
10-layer IAF took nearly twice long train compared to the SeRe- VAE. Finally, the Relaxed Bernoulli VAE (Webb et al. 2018) respects the factorization of the true posterior but scales up to 30 latent variables while not supporting recurrent refinement across layers. The learning curves, the architectural details and the training hyperparameters are provided in the appendix.
4.1.2 PERFORMANCE OF THE RESNET SERE-VAE
To demonstrate the capacity of model when combined with complex layers, replaced the MLPs with ResNets as in Salimans et al. (2015) while preserving the same number of latent vari- ables. As shown in Table our model performs better than all recent models that do not expensive coupling or pixel-Tevel autoregressive layers, either in the encoder or in the decoder, and on par with models of higher complexity. Especially for BIVA, it should be mentioned that more, 168 vs 100 of our model, latent variables are used. The full architectural details are provided in the appendix.
Model
log p(x)
Models with autoregressive (AR) or coupling (C) components
VLAE (Chen et al. 2017 |017
79.03
Pixel RNN vaa den Oord et al.
2016b
79.20
RQ-NSF (Durkn
2019
-9.663
Pixel VAF (Guirajan et
-9.66
RQ-NSF (AR) [Durkan
2019
-79.71
IAF VAE (Kingma el al_ 2016
79.888
DRAW (Gregor et al. 2015
80.97
Pixel CNN (van den Oord et al. 2016a
81.30
Models without autoregressive coupling components
SeRe-VAF
-79.50
BIVA (Maaloe 2019
-80.47
Discrete VAL [Rolle 2017
-81.01
Table 2: Dynamically binarized MNIST performance for VAEs with sophisticated layers. 1000 importance samples were used for the estimation of the marginal likelihood. All performances listed here are taken from Maalçe et al. (2019) and Durkan et al. (2019). All models were trained with single importance sample.
7

---

Under review as a conference paper at ICLR 2021
4.2 CIFARIO NATURAL IMAGES
4.2.1 ABLATION STUDY
In this section, we study the effect of the different couplings between the layers of the architecture presented in Figure on CIFAR-10 images which have dimension (32, 32. We consider 16- layer architecture with each layer generating a 8,38 3) patch of the image when partitioned in spatial checkerboard pattern. We use (8,8, (,8, 2) latent spaces per layer. For the decoder, we use the mixture of discretized logistic distributions (Salimans et al. 2017). In particular, investigate three different architectures:
case 1: there are no couplings (no vertical edges) between the layers and each patch is independently generated from the others,
case 2: there couplings only between the decoders (æ;-1 2) edges),
case 3: there is feedback from the previous inference layer both in the observed space
(x:-1 - zą edges) and the latent space (z1-1 + C₁2 and 21-1 + 2 edges).
In all of the above cases, we consider joint bijective layers between the inference and generative network. One observation that we would like make and turned out to be critical, when we tested our architecture on more complex regimes such as CIFAR-10, 10, and in order to obtain significant predictive benefits from case compared case 2 was that we had to consider a lower bound for the variance in the prior layers. In other words, a deep probabilistic should self-reflect by obtain- ing information from the previous inference layers but without being overly confident its prior assumptions. This can also be mathematically corroborated by examining the KL--divergence in the VAE objective of equation 2 for the residual parametrization introduced in section 3.4.2
(xl(g(E 2, 2xx) || p(c 2)) (ōp((z) - (1 50 (z)))((2))
1 000o(x)
1
2 (x))
2z²(2)
2
2
(10)
As it can be seen, bounding (z)) from below prevents making the first term of the KL arbitrarily large. In these experiments, we take a unit lower bound.
architecture
128 epochs
256 epochs
512 epochs
case I (no vertical edges)
4.56
4.47
4.47
case 2 (coupled decoders)
4.47
4.34
4.28
case 3 (SeRe- (SeRe-VAE)
4.19
3.79
3.68
Table Studying the impact in bits/dim of the connectivity between layers on the test of CIFAR- 10 data for a different number of training epochs. The KI. was linearly annealed (S.II.A. from 0.2 for the first half of the training.
In Table 3 we observe that utilization of information from previous layers in the hierarchy both in the evidence space and in the latent space consistently improves inference. Moreover, the gap the performance becomes larger as more training epochs are dedicated.
The attained performance could be further improved:
without increasing the complexity of the network i) by re-distributing the latent variables allocated per-layer that critical patches of the image given more latent variables ii) further finetuning, especially of the lower bound of the scale in the prior iii) investigat- ing block-coordinate descent optimization algorithms (with the parameters of each layer defining each block).
by increasing the complexity of the network, in particular i) by deploying a deeper archi- tecture ii) by increasing the receptive field of each inference layer that it is coupled not only with the previous inference layers responsible for the generation of the immediately adjacent left/above patches iii) by employing recent deep VAE architectures for each one the layer in our proposed scheme iv) by using more expressive, such as IAF, flows for the joint bijective layers v) by using pitoe-vetoeeeressivee decoders.
8

---

Under review as a conference paper at ICLR 2021
Please note that none of the aforementioned suggestions introduces modeling redundancies (large latent spaces with many of their dimensions collapsing to their prior counterpart) or modeling mis- matches between the true and the variational posterior.
4.2.2 PERFORMANCE OF A SELF-REFLEC,IVE VARIATIONAL MASKED ACTOREGRESSIVE FLOW ON CIFAR-I0
In this section, we introduce a hierarchical latent variable normalizing flow: the first VAE with decoder consisting of normalizing flow transformations transfomattion---caliiigg improvements over its purely generative counterpart. Due to space constraints we refer the reader to the appendix for a review of normalizing flows, as well as the full technical details of our architecture. A high-level description is provided here. The latent variables are generated by the proposed network shown in Figure I Subsequently, the latent variables z are incorporated in the flow in two ways: i) conditioning the base distribution and ii) conditioning the bijective transformations. In the case of a Masked Autore- gressive Flow (MAF) (Papamakarios et al. 2017) or Inverse Autoregressive Flow (Kingma al.) 2016), the latter amounts to designing conditional MADE layers (Germain et al. 2015) that account for a mask offset so that the additional inputs are not masked out. The first amounts building an amortized Gaussian layer. We used a layer hierarchy of 40 latent variables each. We adopted a unit rank Gaussian base distribution in the decoder parameterized in Equation (9) in Rezende et al. (2014) and diagonal Gaussian prior and posterior layers. We used neural spline bijective layers with coupling transformations (Durkan et al. 2019), which boosted the performance com- pared to affine transformations. We refer to our source code and the supplementary material for the implementation details. In Table 4 we compare against generative MAF models with the same or larger width, with without training dataset augmentation with horizontal image flips and different number of MADES. Our variational model exhibits significant improvement over the baselines.
Model
Variational
# MADE layers
Width
Flipped Images
Test Loglikelihood
Seke-MAF
Yes
10 flows, 5 layers)
1024
No
2 3190 (ELBO)
MAF
No
10
1024
No
2670
MAF (5)
(Papamakarins et al. 2017
No
2048
Yes
2936
MAF (10) ('spamakanios et al. 2017
No
10
2048
Yes
3049
Table 4: Performance of different MAFs on CIFAR-10.
5 CONCLUSION AND DISCUSSION
In this paper, we presented self-reflective variational inference that suggests a structural modifica- tion for hierarchical VAEs (SeRe-VAB) and combines top-down inference with iterative feedback between the generative and inference network through shared bijective layers. This modification increases the representation capacity of existing VAEs, leading smaller latent spaces and vast computational benefits without compromising the generative capacity of the model. We further in- troduced hierarchical latent variable normalizing flows which utilize the proposed architecture to recurrently refine the base distribution and the bijectors from the latent codes the previous layer. For our experiments, we used uncoupled deterministic encoders; it would be interesting explore any predictive benefits of bottom-up deterministic pass of the inference network, especially for modeling natural images. The architecture could be further refined by adopting hierarchical stochas- tic layers. Finally. integration of pigel-regressive decoders and importance-weighted variations of the proposed scheme constitute directions for future research.
REFERENCES
Philip Bachman. An architecture for Deep, Hierarchical Generative Models. In Proceedings of the 30th International Conference on Neural Information Processing Systems, 2016.
Vladimir Bogachev. Measure theory, volume 1. Springer Science & Business Media, 2007.
Yuri Burda, Roger B. Grosse, and Ruslan Salakhutdinov. Importance Weighted Autoencoders. In
4th International Conference Learning Representations, ICLR, 2016.

---

Under review as a conference paper at ICLR 2021
Xi Chen, Diederik P. Kingma, Tim Salimans, Yan Duan, Prafulla Dhariwal, John Schulman, llya Sutskever, and Pieter Abbeel. Variational Lossy Autoencoder. In 5th International Conference on Learning Representations, ICLR, 2017.
Bin Dai and David P. Wipf. Diagnosing and Enhancing VAE models. In 7th International Confer- ence on Learning Representations, ICLR, 2019.
Laurent Dinh, Jascha Sohl-Dickstein, and Samy Bengio. Density estimation using Real NVP. In 5th International Conference on Learning Representations, ICLR, 2017.
Conor Durkan, Artur Bekasov, lain Murray, and George Papamakarios. Neural spline flows. In Advances in Neural Information Processing Systems 32, 2019.
Mathieu Germain, Karol Gregor,
lain Murray, and Hugo Larochelle. MADE: Masked Autoencoder for Distribution Estimation. In Proceedings of the 32nd International Conference on Machine Learning, JCML,, 2015.
Karol Gregor, Ivo Danihelka, Alex Graves, Danilo Jimenez Rezende, and Daan Wierstra. DRAW: A Recurrent Neural Network for Image Generation. In Proceedings of the 32nd International Conference on Machine Learning, ICML., 2015.
Ishaan Gulrajani, Kundan Kumar, Faruk Ahmed, Adrien Ali Taíga, Francesco Visin, David Vázquez, and Aaron C. Courville. PixeIVAE: Latent Variable Model for Natural Images. In 5th Interna- tional Conference on Learning Representations, ICLR, 2017.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep Residual Learning for Image Recognition. The IEEE Conference on Computer Vision and Puttern Recognition, CVPR, June 2016.
Michael Jordan. An introduction to probabilistic graphical models, 2003.
Diederik Kingma and Prafulla Dhariwal. Glow: Generative flow with invertible 1x1 convolutions. In Advances in Neural Information Processing Systems 31, 2018.
Diederik P. Kingma and Max Welling. Auto-Encoding Variational Bayes. In 2nd International Conference on Learning Representations, ICLR, 2014.
Diederik P. Kingma and Max Welling, An Introduction to Variational Autoencoders. Founda- fions and Trends in Machine Learning, 12(4):307-392, 2019. doi: 10. 10.1561/22000000556 URL
Diederik P Kingma, Tim Salimans, Rafal Jozefowicz, Xi Chen, Ilya Sutskever, and Max Welling. Improved Variational Inference with Inverse Autoregressive Flow. In Advances in Neural Infor- mation Processing Systems 29, 2016.
Alexej Klushyn, Nutan Chen, Richard Kurle, Botond Cseke, and Patrick van der Smagt. Learning Hierarchical Priors in VAES. In Advances in Neural Information Processing Systems 32, 2019a.
Alexej Klushyn, Nutan Chen, Richard Kurle, Botond Cseke, and Patrick van der Smagt. Learning Hierarchical Priors in VAEs. In Advances Neural Information Processing Systems 32, 2019b.
Lars Maaløe, Marco Fraccaro, Valentin Liévin, and Ole Winther. BIVA: A Very Deep Hierarchy of Latent Variables for Generative Modeling. In Advances in Neural Information Processing Systems 32, 2019.
George Papamakarios, Theo Pavlakou, and lain Murray. Masked Autoregressive Flow for Density Estimation. In Advances Neural Information Processing Systems 30, 2017.
Danilo Jimenez Rezende and Fabio Viola. Taming VAFs. arXiv preprint 1800v:1810.00597, 2018.
Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic Backpropagation and Approximate Inference in Deep Generative Models. In Proceedings of the 31st International Conference Machine Learning, ICMI., 2014.
10

---

Under review as a conference paper at ICLR 2021
Jason Tyler Rolfe.
Discrete Variational Autoencoders. In 5th International Conference Learning Representations, ICLR, 2017.
Walter Rudin. Real and complex analysis. Tata McGraw-hill education, 2006.
Tim Salimans, Diederik Kingma, and Max Welling. Markov Chain Monte Carlo and Variational Inference: Bridging the Gap. Proceedings of the 32nd International Conference on Machine Learning, ICML, 2015.
Tim Salimans, Andrej Karpathy, Xi Chen, and Diederik P Kingma. Pixelenn++: Improving the pixelenn with discretized logistic mixture likelihood and other modifications. arXiv preprint arXiv:/701.05517, 2017.
Casper Kaae Sønderby, Tapani Raiko, Lars Maaløe, Søren Kaae Sønderby, and Ole Winther. Ladder Variational Autoencoders. In Advances in Neural Information Processing Systems 29, 2016.
Jakub Tomezak and Max Welling. VAE with VampPrior. In Proceedings of the Twenty-Fint International Conference on Artificial Intelligence and Statistics, volume of Proceedings of Machine Learning Research. PMLR, 2018
Aäron van den Oord, Nal Kalchbrenner, Lasse Espeholt, Koray Kavukcuoglu, Oriol Vinyals, and Alex Graves. Conditional Image Generation with PixeICNN Decoders. In Advances in Neural Information Processing Systems 29, 2016a.
Aäron den Oord, Nal Kalchbrenner, and Koray Kavukcuoglu. Pixel Recurrent Neural Networks. In Proceedings of the 33nd International Conference on Machine Learning, ICML, 2016b.
Stefan Webb, Adam Golinski, Rob Zinkov, N Siddharth, Tom Rainforth, Yee Whye Teh, and Frank Wood. Faithful inversion generative models for effective amortized inference. Advances in Neural Information Processing Systems, pp. 3070-3080, 2018.
Florian Wenzel, Kevin Roth, Bastiaan S Veeling, Jakub Šwiatkowski, Linh Tran, Stephan Mandt, Jasper Snoek, Tim Salimans, Rodolphe Jenatton, and Sebastian Nowozin. How Good is the Bayes Posterior in Deep Neural Networks Really? arXiv preprint arXiv:2002.02400 2020.
Andrew Gordon Wilson and Pavel Izmailov. Bayesian Deep Learning and Probabilistic Perspective of Generalization, arXiv preprint arXiv:2002.08791, 2020.


Reviews:
[{"review_id": "M0acz10e6UK", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "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": "1. Summary\nThis paper proposes an augmentation of the Ladder VAE model (LVAE [1]) using 1. more flexible variational distributions using normalizing flows (denoted $f$) 2. an autoregressive component (because of the generative dependency $p( x_{l} | z_{l}, x_{l-1}) $ (*data layers*), although this is not directly stated as a contribution. \nIt is argued the proposed factorization of the inference network (top-down inference) matches the factorization of the true posterior, that follows the top-down design of the generative model. The authors recorded competitive performances on dynamically binarized MNIST. \n\n2. [a] Strong Points\n- the authors introduce a very flexible architecture for DGMs by combining recent improvements from the literature. \n\n2. [b] Weak Points\n- the paper overall lacks clarity: the idea could be explained in much simpler terms\n- the contributions are weak: augmenting an LVAE model using flows brings little novelty \n- the autoregressive structure introduced in the paper (which can be interpreted as an instance of state-space models [2]) is not well discussed and is contradictory with the claim that no *autoregressive layer* is used.\n- an ablation study is required to disentangle the effect of 1. the flow augmentation and 2. the autoregressive structure\n- aiming for a better inference network aims at tightening the variational bound, which should be measured empirically \n- the diagonal Gaussian assumption for the variational family is not necessarily so limiting [3, 4] (as assumed in the paper). \n- the results reported for CIFAR-10 are not competitive with sota models [3, 4] (SeRe-MAF) $\\approx$ MAF (10) = 4.31 (original paper) $\\gg$ 3.08 (BIVA [4]) > 2.91 (NVAE without flow [3]) .\n\n3. Recommendation\nUnfortunately, I recommend rejecting this paper.\n\n4. Recommendation Arguments \nSeRe-VAE combines multiple architecture improvements (top-down model (LVAE), flexible variational distributions using flows, autoregressive components) into a final model, which is tested using a single dataset (excluding cifar10 results which are not competitive) and without performing an ablation study. This prevents the community from understanding the effect of each of the architecture choices and does not guarantee that such an architecture could be effectively adapted to other contexts. \n\n5. Questions to the Author\n- Figure 1: why considering the prior independent? the generative model adopts a hierarchical structure $p(z_l  | z_{l+1})$\n\n6. Feedback \nYour work is an engineering prowess, I am saddened to recommend rejection. I think your work could greatly benefit from an ablation study and from defining the model in more minimal terms. \n\nA few comments:\n- The inference network is conditioned on $\\mathbf{x}$, not the entire $\\mathcal{D}$.\n- you can measure the variational bound using the identity: $\\operatorname{KL}(q(z | x) | p(z | x)) = log p(x) - \\mathcal{L}(x)$\n- please report CIFAR10 results in bit per dimension, as stated in the literature\n- please report results on the more widely accepted Statistically binarized MNIST first, use dynamic MNIST as a second option.\n\n[1] Sønderby, Casper Kaae, et al. \"Ladder variational autoencoders.\" Advances in neural information processing systems. 2016.\n[2] Fraccaro, Marco, et al. \"Sequential neural models with stochastic layers.\" Advances in neural information processing systems. 2016.\n[3] Vahdat, Arash, and Jan Kautz. \"Nvae: A deep hierarchical variational autoencoder.\" arXiv preprint arXiv:2007.03898 (2020).\n[4] Maaløe, Lars, et al. \"Biva: A very deep hierarchy of latent variables for generative modeling.\" Advances in neural information processing systems. 2019.", "extra_scores": {}, "extra_text": {}}, {"review_id": "K0RK1jsjY1D", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "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": "**Contributions & Significance**:\n- The proposed paper introduces an encoder-decoder architecture for hierarchical VAEs based on bijective transformations, which preserves the true factorization of the posterior in its variational approximation (which is formally shown). (High)\n- The resulting method shows improved performance on dynamically binarized MNIST compared to models of similar and higher complexity (autoregressive models). (Medium)\n- A variation/extension of the model using normalizing flow transformations for the decoder shows significant improvement compared to Masked Autoregressive Flow (MAF) on CIFAR-10. (Medium)\n\n\n**Pros**\n- The proposed hierarchical architecture only relies on the latent code of the previous layer for conditioning, instead of all preceding latent codes such as in autoregressive models.\n- The method shows improvement in terms of Log-Likelihood by rearranging the stochastic flow instead of relying on a computationally expensive architecture.\n- The general method seems widely applicable for VAE-like models and as such can be combined with other architectural improvements or potential future work.\n\n**Cons**\n- The experiments are a bit limited in terms of the used datasets (Dynamically binarized MNIST & CIFAR-10).\n- Sections 3.1.1-3.1.3 could be written more clearly. In particular I think the notation of the parameters is overly confusing (e.g. $\\alpha_l$ and $\\beta_l$ as functions). For 3.1.1 maybe just directly use $(\\mu_l, \\sigma_l)$?\n- It is unclear how the number of data partitions L, their structure & order are selected. Does that have an effect, how large is this effect and what would the general recommendations be? What about the extreme cases (L=D, L=1)?\n\n**Style**\nIn general the paper is well written and reasoned for. \n\n**Minor Comments**\n- The positioning of Figure 6 seems a bit awkward. \n- Typo Figure 2: Prior Layers: I assume $p(\\epsilon|z_2) $should be $p(\\epsilon|z_1)$\n\n**Summary**\nThe paper is well written, has a promising and sound approach and seems generally applicable. As such I vote for accepting it. However, it could still benefit from some improvement in terms of clarity and number of experiments (i.e. other data sets).\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "FUEyAzzhgcb", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 5, "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": "**GENERAL**\nThe paper proposes a new hierarchical architecture for VAEs. The main idea is to use invertible components that could be shared by the encoder and the decoder. However, the paper is very confusing in many parts, and the experimental studies are not too strong.\n\n**Strengths:**\nS1: The authors propose a new architecture for stochastic layers in VAEs.\n\nS2: The idea is to use bijective layers shared by the generative and the variational parts of the VAE.\n\n**Deficiencies:**\nD1: Please correct Figures 1-5. First, the text seems to be \"squashed\" that hinders readability. Moreover, some words blend with each other. In Figure 5, rotating the text makes it almost impossible to read. Additionally, I believe including colorful squares or rectangles does not help; on the contrary, it makes it harder to read.\n\nD2: The hierarchical model is presented through Figures (e.g., Figure 1) that is very hard to follow. I believe the authors have some interesting ideas, but it is completely overshadowed by rather unreadable and pretty unclear figures. It would be much better to simply express the model mathematically. At the moment, I do not follow what is the semantics of nodes, edges and colors.\n\nD3: The text is hard to follow as well. For instance, the authors write that some components are implemented by ResNet or MLP. However, when MLP is used, and when ResNet is used? These statements are very confusing. Moreover, the authors introduce multiple notation for the same quantities that differ by one symbol, e.g., p(\\epsilon_l | ...) and p(\\epsilon ; ...). What is the purpose of using \"|\" and \";\" interchangeably?\n\nD4: The section about Residual Data Layers is written in a confusing manner. The authors introduce a two-step procedure to calculate \\gamma. First, they compute an estimation of it. Then, they calculate a second estimation \\delta \\gamma_l. Then, they combine these two estimations to define \\c_l^{\\gamma}. Why do the authors mention estimations? Why are the two quantities summed? It is totally unclear and written in a clumsy fashion.\n\nD5: The statements in Section 3.4 are very hand-wavy. The authors claim, e.g., that \"The model, albeit hierarchical, is less prone to posterior collapse, since each layer is responsible for the generation of a different portion of the data\". How do we know that? It is not obvious that this is the case.\n\nD6: I appreciate the experiments provided by the authors. However, the provided comparison on CIFAR10 is hardly acceptable. First, currently using Normal distribution for the decoder is not widely used. Second, MAF is a great idea and an extremely interesting paper, however, nowadays it cannot be treated as a strong baseline.\n\n*AFTER REBUTTAL* I would like to thank the authors for their rebuttal. I increase my score to 5. I am still unsatisfied with the presentation of the idea, because it is still rather hard to follow.", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper proposes a variant of the hierarchical VAE architectures. All reviewers felt that the paper's clarity was lacking. While the authors made very significant improvements during the feedback phase, which were recognized by reviewers, the paper could use a revision that takes clarity into account from the ground up. I also think that the ablation studies should be expanded (if SOTA is not the goal, then science should be), e.g., compare to the setting in which q does not share the bijective layers with p.
Decision: Reject