Hierarchical VAE with a Diffusion-based VampPrior
Authors: Anna Kuzina
Jakub M. Tomczak
Venue: Accepted by TMLR ()

Abstract: Deep hierarchical variational autoencoders (VAEs) are powerful latent variable generative models. 
In this paper, we introduce Hierarchical VAE with Diffusion-based Variational Mixture of the Posterior Prior (VampPrior).
We apply amortization to scale the VampPrior to models with many stochastic layers. 
The proposed approach allows us to achieve better performance compared to the original VampPrior work and other deep hierarchical VAEs, while using fewer parameters. 
We empirically validate our method on standard benchmark datasets (MNIST, OMNIGLOT, CIFAR10) and demonstrate improved training stability and latent space utilization.

Published in Transactions on Machine Learning Research (11/2024)
Hierarchical VAE with Diffusion-based VampPrior
Anna Kuzina
u.kuzina@uu.tt
Department of Computer Science,
Vrije Universiteit Amsterdam, Netherlands.
Jakub M. Tomczak
j.m.omczakBtue..c
Deportment Mathematics ond Computer Science, Eindhoven Universiby of Technology, Netherlands.
Reviewed on OpenReview: htips: // openreview. net/ mett forum? id- NUkEo27Toa
Abstract
Deep hierarchical variational autoencoders (VAEs) are powerful latent variable generative models. In this paper, we introduce Hierarchical VAE with Diffusion-based Variational Mixture of the Posterior Prior (VampPrior). We apply amortization to scale the VampPrior to models with many stochastic layers. The proposed approach allows us to achieve better performance compared to the original VampPrior work and other deep hierarchical VAEs, while using fewer parameters. We empirically validate our method on standard benchmark datascts (MNIST, OMNIGLOT, CIFAR10) and demonstrate improved training stability and latent space utilization.
1 Introduction
Latent variable models (LVMs, Ms) parameterized with neural networks constitute a large group deep gener-
ative modeling [omczak 2022 One class of LVMs, Variational Autoencoders (VAEs) (Kingma & Welling,
2014 Rezende et al. 2014), utilize amortized variational inference efficiently learn distributions
vari-
data modalities, c.g., images (Kingma & Welling 2014), audio {Van Den Oord et al. 2017 , molecules et al. 2018). The expressive power of VAEs can be improved by introducing hierarchy of latent variables. The resulting hierarchical VAEs, such as ResNET VAEs (Kingmaa et al. 2016), BIVA (Maalec et al. 2019), very deop VAE (VDVAE) (Child 2021), or NVAE (Vahdat & Kautz, 2020 achiove state-of-lhe-art performance on images in terms of the negalive log- lig-likelihood (NLL). However, hierarchical VAFs are known have training instabilities (Vahdat Kautz 2020). To mitigate these issues, various tricks were proposed, such as gradient skipping (Child (2021). spectral normalization (Vahdat & Kautz 2020), or softmax parameterization of variances (Hazami al. 2022). In this work, propose different approach that focuses on two aspects of hierarchical VA (i) the structure of latent variables, and (ii) the form of the prior for the given structure. We introduce several changes to the architecture of parameteriza- tions (i.e. neural networks) and the model itself. As result, we can train a powerful hierarchical VAE with gradient- gradient-based methods and ELBO as the objective without any hacks.
In the VAE litcrature, it a well known fact that the choice of the prior plays an important role in the resulting VAE performance (Chen et al., 2017 Tomezak (022).. For example, VampPrior (Tomezak & Welling 2018), form of the prior approximating the aggregated posterior, has been shown to consistently outperform VAES with standard prior and a mixture prior. In this work, we extend the VampPrior to deep hierarchical VAEs in an efficient manner. We propose utilizing non-trainable linear transformation like Discrete Cosine Transform (DCT) to obtain pseudoinputs. Together with our architecture improvements, we can achieve state-of-the-ar performance, high utilization of the latent space, and stable training of deep hicrarchical VAEs.
The contributions of the paper the following:

---

Published in Transactions on Machine Learning Research (11/2024)
We propose new VampPrior-like approximation of the optimal prior (i.e., the aggregated posterior) which can efficiently scale to deep hierarchical VA
We propose & latent aggregation component that consistently improves the utilization of the latent space of the VAE.
Our proposed hicrarchical VAE with the new class of priors outperforms other hicrarchical VAE models (without additional data augmentations).
2 Background
2.1
Hierarchical Variational Autoencoders
Let us consider random variables x E X (e.g. = R). We observe N x's sampled from the empirical distribution r(x). We assume that each x has L corresponding latent variables Zi:L = (Z1...., ,g)) where z R: and M, the dimensionality of each variable. We aim to find a latent variable generative model with unknown parameters 0, po(x, ZI.L. = 2o (ZI.L.)
In general, optimizing latent-variable models with nonlinear stochastic dependencies is non-trivial. A possible solution is approximate inference, c.g.. in the form of variational inference (Jordan et al., 1999 with a family of variational posteriors over latent variables [1o(Z).1|x) This idea is exploited in Variational Auto-Encoders (VAFs) (Kingma & Welling 2014: Rezende et al. 2014), in which variational posteriors are referred to as encoders. As a result, we optimize a tractable objective function called the Evidence Lower BOund (ELBO) over the parameters of the variational posterior, Φ, and the generative part, 8, that is:
Er(x) [x)] 2 (x)))) Bystenn In (e((zz)))))) - DKL [go (l)
(1)
where r(x) is the empirical data distribution. Further, avoid clutter, we will use Ex H instead of (E)())) H. meaning that the expectation is taken with respect to the empirical distribution.
2.2 VampPrior
The latent variable prior (or marginal) plays a crucial role in the VAE performance, which motivates the usage of data-dependent priors. Note that the KT--divergence term in the ELBO (see Eq. 1 can be expressed as follows [Hoffman & Johnson 2016):
DKL. la. (z)(2)))lp((2) DRL. [[o(2)]lp(z)) + Ex Dia. (as())x)la(())
(2)
Therefore, the optimal prior that maximizes the ELBO has the following form:
(z) = Ex [g0(2]x)).
(3)
The main problem with the optimal prior in Eq. g is the summation over all N training datapoints. Since could be very large (e.g., tens hundreds of thousands), using such prior is infeasible due to potentially very high memory demands. As an alternative approach, Tomczak & Welling (2018) proposed VampPrior, new class of priors that approximate the optimal prior in the following manner:
p" (z) (x)x) R (z(u))
(4)
where is pseudoinpul, i.e., variable mimicking real data, r(u) = E us) is the distribution of u in the form of the mixture of Dirac's deltas, and {ux 1K are learnable parameters (we will refer to them as pscudoinputs as well). K is hyperparameter and is assumed to be smaller than the size of the training dataset, K < N. Pseudoinputs are randomly initialized before training and are learned along with model parameters by optimizing the ELBO objective using a gradient-based method.
In follow-up work, Egorov et al. (2021) suggested using a scparate objective for pscudoinputs (a greedy boosting approach) and demonstrated the superior performance of such formulation in the continual learning setting. Here, we will present different approximation to the optimal prior instead.
2

---

Published in Transactions on Machine Learning Research (11/2024)
2.3 Ladder VAEs (a.k.a. dop-down VAEs)
We refer to models with many stochastic layers L as deep hierarchical VABs. They differ in the way the prior and variational posterior distributions are factorized and parameterized. Here, we follow the factorization proposed Ladder VAE (Sondorby et al. 2016) that considers the prior distribution over the latent variables factorized in an autoregressive manner:
L-1
( ZL) = (ZL)
II
po
(5)
1=1
Next, using the top-down inference model results in the following variational posterior (Sonderby et al. 2016):
(o(2),..., 21|x) go (zz II
go
(6)
This factorization has been previously used by successful deep hierarchical VAEs, among others, NVAE (Vah- dat & Kautz, 2020) and Very Deep VAE (VDVAE) (Child] 2021). It was shown empirically that such a formulation allows one to achieve state-of-the-art performance of the hierarchical VAEs on several image datasets.
3 Our model: Diffusion-based VampPrior VAE
In this work, introduce the Diffusion-based VampPrior VAE (DVP-VAE). It is deep hicrarchical VAE model that approximates the optimal prior distribution at all levels of the hierarchical VAE in an efficient way.
3.1 VampPrior for hierarchical VAE
Directly using the VampPrior approximation (see Eq. 4 for deep hierarchical VAE can be very computation- ally expensive since it requires evaluating the variational posterior of all latent variables for K pseudoinputs at each training iteration. Thus, (Tomzzak Welling 2018) proposed a modification in which only the top latent variable uses VampPrior, namely:
p" (21.1) = (x se
(7)
where r(u) = 1 E. ō(u u&) with learnable pseudoinputs {u In this approach, there are a few problems: (i) how to pick the best number of pseudoinputs K, (ii) how to train pseudoinputs, and (iii) how train the VampPrior in a scalable fashion. The last issue results from the first two problems and the fact that the dimensionality of pseudoinputs is the same the original data, i.e., dim(u) dim(x).
Here, propose different prior parameterization to overcome all these three problems. Our approach consists of three steps in which we approximate the VampPrior at all levels of the deep hicrarchical VAE. We propose to amorlize the distribution of pseudoinputs in VampPrior and use them to directly condition the prior distribution:
p' p'(2..)) (x10,((2zz a Ex..ralajoo(a:r
where we use r(u) = Ex(r(uxx]]. Using r(u[x), which is cheap to evaluate for any input, we avoid the expensive procedure of encoding pseudoinputs along with the inputs x. Amortizing the VampPrior solves the problem of picking K and helps with training pseudoinputs.
To define conditional distribution, we treat pseudoinputs as a result of some noisy non-trainable trans- formation of the input datapoints Let us consider a transformation f: xD x + P i.c., f(x) de, where is standard Gaussian random variable and a is the standard deviation. Note that by applying e.g., linear transformation, can lower the dimensionality of pseudoinputs, din(u) < din(x), resulting in better scalability. As a result, we get the following amortized distribution:
r(ulx) (x)
(9)
3

---

Published in Transactions on Machine Learning Research (11/2024)
Algorithm 1 fdet: Create DCTT-ased pseudoin-
Algorithm 2 Sam: Invert DCT-based pseudoinputs
puts
Input: upCT Rexdxdd S E Rexddx,
Input: RexDDD , Boxdxd ded R
UDCT uDCT S
UDCT (
UpCT = zero perro_pad_(e
pap_pppddupuup..p (UDCT, - d)
UpCT = Crop(UDeT: d)
ux = =iDCT(uDCT)
UDCT S
sper
Return: Ux E DxD
Return: upcr E Rexdxd
The crucial part then how to choose the transformation f. It. non-trivial choice since we require the following properties of f: (i) it should result in dim(u) < dim(x), (ii) u should be a reasonable representation of x, (iii) it should be easily computable (e-g.; fast. for scalability). We have two candidates for such transformations. First, can consider a downsampled version of an image. Second, propose to use discrete cosine transform. We will discuss this approach in the following subsection.
Moreover, the outlined amortized VampPrior in the previous two steps seems to be a good candidate for efficient and scalable training. However, it does not seem to be suitable for generating new data. Therefore, we propose including pseudoinputs as the final level in our model and use a marginal distribution f(u) that approximates r(u). Here, we propose to use a diffusion basdd model for f(u).
3.2 DCT-based pseudoinputs
The first important component in our approach is the form of the non-trainable transformation from the input to the pseudoinput space. We assume that for to be reasonable representation of x means that u should preserve general patterns (information) of x, but it does not necessarily contain any high-frequency details of x. To achieve this, we propose to use a discrete cosine transform (DCT) to convert the input into a frequency domain and then filter the high-frequency component.
DCT DCT (Ahmd et al. 1974 is a widely used transformation in signal processing for image, video, and audio data. For example, it is part of the JPEG standard (Pennebaker & Mitchell 1992). For instance, consider a signal as a 3-dimensional tensor x E ReXDXD, DCT is linear transformation that decomposes each channel x; on a basis consisting of cosine functions of different frequencies: uDCT, = Cx,CT, where
for all pairs (k- = 0, n): Ck = Vb. b. and for all pairs (k, n) such that k > O: Ck. v D cos (3 (n + 1) k).
Our transformation We use the DCT transform as the first step in the procedure of computing pseu- doinputs. Let us assume that each channel of x is D x D. We select the desired size of the context d < D and remove (crop) D d bottom rows and right-most columns for each channel in the frequency domain since they contain the highest-frequency information. Finally, we perform normalization using the matrix S that contains the maximal absolute value of each frequency. We calculate this matrix once (before training the model) using all the training data: S maxxe [DCT(x)). The complete procedure is described in Algorithm D and we denote it as fdct.
In the frequency domain, pscudoinput has smaller spatial dimension than its corresponding datapoint. This allows us to use a small prior model and lower the memory consumption. However, we not em- pirically that conditioning the amortized VampPrior (see Eq. S on the pseudoinput in the original domain makes training casier. Therefore, the pseudo-inverse is applied as part of the TopDown path. First, we start by multiplying the pseudoinput by the normalization matrix S. Afterward, we pad each channel with zeros account for the "lost' high frequencies. Lastly, we apply the inverse of the Discrete Cosine Trans- form (iDCT). We denote the procedure for converting a pseudoinput from the frequency domain to the data domain fi and describe it in Algorithm
We consider the most widely used type-II DCT.
4

---

Published in Transactions on Machine Learning Research (11/2024)
23
Z
Z3
Z
72
72
Z2
21
Z1
Z1
21
(a) Ladder VAE
(b) Ladder VAE with pseudoinouts
Figure 1: Graphical model of the TopDown hierarchical VAE with three latent variables (a) without pseu- doimputs and (b) with pseudoinputs. The inference model (left) and the generative model (right) share parameters in the TopDown path (blue). The dashed arrow represents a non-trainable transformation.
3.3 Our Ladder VAE and Training Objective
We use the TopDown architecture and extend this model with deterministic, non-trainable function create the pseudoinput. In our generative model, pscudoinputs are treated as another set of latent variables:
(x 21:1, u) px( 21.1 )po (Z1.z (u)(())
(10)
(0(ZL.//u) po (z lu) II pe (zizzzz1. u).
(11)
Then, we choose variational posteriors in which pseudoinput latent variables are conditionally independent with all other latent variables given the datapoint x, that is:
qo (ZI.. u|x) = go (Z).z x)r(u(x).
(12)
L-1
(x111 = (o(22 xx)
II
(2 M+111,));
(13)
1=1
r(lxx) -x(ul[(xx)oTT)
(14)
In the variational posteriors, use the amortization of ruu) outlined in Eq.
Let us consider Ladder VAE (a TopDown VAE) with three levels of latent variables. We depict the graphical model of this latent variable model in Figure la with the inference model on the left and the generative model on the right. Note that inference and generative models use shared parameters in the TopDown path, denoted by the blue arrows.
In Figure 1b show graphical model of our proposed model that additionally contains pseudoinputs The generation model (Figure Tb right) is conditioned the pseudoinput variable at each level. This formulation is similar to the iVAF model (Khemakhem et al. 2020) in which the auxiliary random variable u is used to enforce the identifiability of the latent variable model. However, unlike iVAE, we do not treat u as a separate obscrved variable. Instead, we use a non-trainable transformation to obtain during training and introduce the pseudoinput prior distribution P-y (u) with learnable parameters y sample u test time. This allows us get unconditional samples from the model.
Finally, the Evidence Lower BOund for model takes the following form:
L(x, db, 0, r) = Bys (x) In (x (x)z).)) (uix) DEL (xx(211/x)) (p) [ - DKL [r())))]lruu)) (15)
The ELBO for our model gives a straight.orward objective for training an approximate prior by minimizing the Kullback-Leibler between r(ulx) and the amortized VampPrior r(u) = Ex(r(u[x)). However, calculating the KL-term with r(u) is computationally demanding and, in fact, using r(u) as the prior does not help us solve the original problem using the VampPrior. Therefore, in the following section, we propose use an approximation (u) instead.
5

---

Published in Transactions on Machine Learning Research (11/2024)
3.4 Diffusion-barte VampPrior
Even though pseudoinpuls are assumed to be much simpler than the observed datapoint x (e.g-, in terms of their dimensionalityy), very flexible prior distribution (u) is required to ensure the high quality of the final samples. Since we cannot use the amortized VampPrior directly, following (Vahdat ct al. 2021 Wehenkel & Louppe, 2021), we propose to use a diffusion-basee generative model (Ho et al. 2020 as the prior and the approximation of r(u).
Diffusion models are flexible generative models and, in addition, can be seen latent variable generative models [Kingma et al. 2021 As result, we have access to the lower bound its log-likelihood function:
log (u) 2 (u, styolu) [nr(u)y))] DaL. [d(v)lu)](r(>))
(16)
E (0 DKL | (0
We provide more details on diffusion models and the derivation of the ELBO in Appendix A We refer to this prior as Diffusion-bosed VampPrior (DVP). This prior allows us to sample infinitely many pseudoinputs, unlike the original VampPrior that uses a fixed set of pseudoinputs.
Now, we can plug this lower bound into the objective (Eq. 15 and obtain the final objective of our Ladder VAE with pseudoinputs and the Diffusion-based VampPrior (dubbed DVP-VAE):
max Ex (ala (xpo(x|z) Er(u(s) DKL. [x) [[((ux))) (ul Lu (u. 0.0-7
(17)
Note that [[[[[[[]]] & log(2n (ee(Texpo')) o2) where learnable parameter (see Eq. 9 and P the dimen- sionality of the pscudoinput.
4 Model Details: Architecture and parameterization
The model architecture and parameterization crucial to the scalability of the model. In this section, we discuss the specific choices we made. The starting point for our architecture is the architecture proposed in VDVAE CChild 2021). However, there are certain differences. We schematically depict our architecture in Figure 2a We consider a hierarchical TopDown VAE with L stochastic layers, namely, latent variables Zi...., ZL We assume that each latent variable has the same number channels, but they differ in spatial dimensions: 2 E Rexhxxu We refer to different spatial dimensions of the latent space as scales.
4.1 Bottom-up
The bottom-up part corresponds to calculations of intermediary variables dependent on x. We follow the implementation of Child (2021) for it. We start from the bottom-up path depicted Figure 2a (left), which is fully deterministic and consists of several ResNet blocks (see Figure 2ch. The input processed by Nenc blocks at each scale, and the output of the last resnet block of each scale passed to the TopDown path in Figure 2a (right). Note that here Nenc is separate hyperparameter that does not depend the nunber of stochastic layers L.
4.2 TopDown
The TopDown path depicted in Figure 2a (right) computes the parameters of the variational posterior and the prior distribution starting from the top latent variable Zy-
The first step is the pseudoinput block shown in Figure 2d Using the deterministic function fac.. it creates the pseudoinput random variable from the input x (see Algorithm[1 I that used to train the Diffusion- based VampPrior (u). At test time, a pscudoinput is sampled using this unconditional prior. The pscudoinput sample then converted back the input domain (see Algorithm 2 and used condition prior distributions at all levels po (Z):1/u).

---

Published in Transactions on Machine Learning Research (11/2024)
dooc
dec henv
Nasc
Av. Pool
Gome Lxx.1
Upsample
ht +
Men
ddcc
Av. Pool
(b) Top down block
Upsample
N
cow IXIXXXX
fdet
SAL
r(u
?(u)
u
SA
dct
D
(c) Resmet block
(d) block
Pscudoinput
(a) Model architecture
Figure 2: A diagram of the DVP-VAE: TopDown hierarchical VAE with the diffusion- VampPrior. (a) A Bottom Up path (left) and a TopDown path (right). (b) A TopDown block that takes features from the block above h dec, encoder features h (only during training) and a pscudoinput u as inputs. (c) A single Resnet block. (d) A single pscudoinput block.
Next, the model has L TopDown blocks depicted in Figure 2b Each TopDown block takes deterministic features from the corresponding scale of the bottom-up path denoted as henes the output of the pseudoinput block u, and deterministie features from the TopDown block above hase inputs. Our implementation of this block is similar to the VDVAE architecture, but there are several differences that we summarize below:
Incorporatngg pseudoinputs
We concatenate the pseudoinput (properly reshaped using average pooling) with haee to compute the parameters of the prior distribution.
Variational posterior parameters
We assume that both haee and hene have the same number of channels, allowing us to sum them instead of concatenating. This reduces the total number of parameters and the memory consumption.
Additional ResNet Connections Our TopDown block has three ResNet blocks (Figure 2cp. In contrast to architecture, in VDVAE only the block that updates haec has a residual connection.
We did not observe any training instabilities and did not apply the gradient skipping used in Child 2021)
4.3 Latent Aggregation in Conditional Likelihood
The last important element of our architecture is the aggregation of lotends. Let us denote samples from either the variational posteriors (o(Z).) x) during training or the prior po(zi.z during generating new data as 21....,2, ēz. Furthermore, let hy be the output of the last TopDown block. These deterministic features computed a function of all samples 21,..., Z7. Therefore, it be used to calculate the final likelihood value. However, observe empirically that in such parameterization some layers of latent variables tend to be completely ignored by the model. Instead, we propose enforce a strong connection between the conditional likelihood and all latent variables by explicitly conditioning on all of the sampled latent variables, namely:
NNN 1 VL E, 2))
p=(x(Z1:L) = NN
(18)

---

Published in Transactions on Machine Learning Research (11/2024)
We refer to this as the latent aggregation. We show empirically in the experimental section that this leads to consistently high ratio of active units.
5 Related Work
Latent prior in VA The original VAE formulation uses the standard Gaussian distribution as a prior over the latent variables. This can be an overly simplistic choice, as the prior minimizing the Evidence Lower bound is given by the aggregated posterior (Hoffmann & Johnson 2016 Tomezak & Welling 2018).
Furthermore, using a unimodal prior with multimodal real-world data can lead to non-smooth encoder or meaningless distances in the latent space Bozkurt et al. (2019). More flexible prior distributions proposed in
the literature include the Gaussian Mixture Model (Jiang et al. 2016 Nalisnick et al. 2016 Tran et al. 2022
the autoregressive normalizing flow (Chen et. al.
2017), the auloregressive model [Gulrajani et al.
2016
Sadeghi et al. 2019), the rejection sampling distribution with the learned acceptance function (Bauer & Mnih 2019), the diffusion-base prior (Vahdat ot al. 2021 Wehonkel & Louppe 2021). The VampPrior ( (Tomczak
& Welling 2018 proposes using an approximation of the aggregated posterior as a prior distribution. The approximation constructed using learnable pseudoinputs to the encoder. This work can be seen as an efficient extension of the VampPrior to deep hierarchical VAE, which also utilizes a diffusiin-based prior over the pseudoinputs.
Auxiliary Variables in VAES Several works consider auxiliary variables u as a way to improve the flexi- bility of the variational posterior. Maalee et al. (2016) use auxiliary variables with one-level VAE to improve the variational approximation while keeping the generative model unchanged. [Salimans et al. (015 use Markov transition kernel for the same expressivity purpose. The authors treat intermediate MCMC samples as an auxiliary random variable and derive evidence lower bound of the extended model. [Ranaaath et al.) (2016) introduce hierarchical variational models. They increase the flexibility of the variational approxi- mation by imposing prior on its parameters. In this setting, it assumes that the latent variable z and the auxiliary variable are not conditionally independent and the variational posterior factorizes, for example, as follows:
(u, z/x) = go (u (z)
(19)
In this work, in contrast, we use auxiliary variable to increase the prior flexibility and use conditional independence assumption in the variational posterior:
go (u, (o(u)x)gox/(x).
(20)
Khemakhem et al. (2020) consider the non-identifiability problem of VAÉs. They propose use auxiliary observation u and use it to condition the prior distribution. This additional observation is similar to the pseudoinputs that we consider in our work. However, we define a way to construct from the input and learn a prior distribution to sample it during inference, while Khemakhem et al. (2020) require u to be observed both during training and at the inference time.
Similarly to our work, (Klushyn et al., 2019) consider hierarchical prior (o(z)u)p(u). However, they treat u rather as second layer of latent variables and learn a variational posterior in the form (o(u,z/x) (o(u (z)xx)..
Latent Variables Aggregation There are different ways in which the conditional likelihood (x 21:2) can be parameterized. In LadderVAE (Sonderby et al. 2016). where TopDown hierarchical VAE was origi- nally proposed, the following formulation is used:
(x(x (x(x)NN()))).
(21)
That is, the conditional likelihood depends directly on the bottom latent variable z only.
Later, NVAE (Vahdat & Kautz 2020 and VDVAE CChid 2021 deterministic path in the TopDown
architecture in the conditional likelihood, namely:
(x(x (x (N(h))))
(22)
8

---

Published in Transactions on Machine Learning Research (11/2024)
Table 1: The test performance: negative log likelihood MNIST and OMNIGLOT, and bits-per-dimension
(BPD) on CIFAR10.
Results with data augmentation. Standard deviation reported in parentheses.
Results averaged over random seeds.
MNIST
OMNIGLOT
CIFARIO
MODEL
L
log ((x) S
SIZE
L
BPD S 4
DVP-VAE (ours)
8
77.10*
70.00"(0.05)
89.07' (0.10)
20M
28
2.73
Attentive VAE Aposlolopoulou el al.
2022
77.63
89.50
119M
16
2.79
CR-NVAE (Sinaa & Dieng 2021
15
76.93
131M
30
2.51
VDVAE Child 2021)
39M
45
2.87
OU-VAE [Pevve & Gavves 2021
5 81.10
96.08
10M
3
3.39
NVAE Vahdat & Kantz 2020)
15
78.01
30
2.91
BIVA Maalpe et al. 2019
6
78.11
91.34
103M
15
3.08
VampPrior Tommzak & Welling 2018
2
78.45
89.76
LVAE (Sonderby et al. 2016
5
81.74
102.11
IAF-VAE(KIN) et al. 2016
79.10
12
3.11
Note that deterministic features depend on all the latent variables. However, we propose to use a more explicit, dependency on latent variables in Eq. 18 Our idea bears some similarities with Skip- (Dieng et al. 2019). Skip- VAE proposes to add latent variable to each layer the neural network parameterizing decoder of the VAE with a single stochastic layer. In this work, instead, we add all the latent variables together to parameterize conditional likelihood.
6 Experiments
6.1
Settings
We evaluate DVP-VAE dynamically binarized MNIST (LeCun] 1998) and OMNIGLOT 2015).
[Lake et al.
Furthermore, we conduct experiments natural images using the CIFAR10 dataset Alex 2009
We provide all the hyperparameters for training DVP-VAE in Appendix C.1 and in the code repository
6.2 Main Quantitative and Qualitative Results
We report all results in Table where compare the proposed approach with other hierarchical VAEs. We observe that DVP-VAE outperforms most of the VAF model while using fewer parameters than other models. For instance, on CIFARIO, our DVP-VAE requires 20M weights to beat Attentive VAE with about times more weights. Furthermore, because of the smaller model size, we were able to obtain all the results using a single GPU. We show the unconditional samples Figure 3 (see Appendix D for more samples). The top row of cach image shows samples from the ieffusiion-base VampPrior (i.c., pseudoinputs), while the second row shows corresponding samples from the VAE. We observe that, as expected, pseudoinputs define the general appearance of image, while a lot of details are added later by the TopDown decoder. This effect can be further observed in Figure 4 where we plot the reconstructions using different numbers of latent variables. In the first row, only pseudoinput corresponding to the original image is used (i.e., ~ r(u(x)) while the remaining latent variables are sampled from the prior with low temperature. Each row below uses more latent variables from the variational posterior grouped by the scales. Namely, the second row uses the pscudoinput above and all the x 4 latent variables from the variational postcrior, then the third row uses additionally 8 latent variables, and so
6.3 Ablation Studies
Training stability and convergence In all of our experiments, we did observe many training in-
stabilities. Unlike many contemporary VAEs, we did not use gradient skipping (Child 2021), spectral
hnttps///itthbb com/ Awwp////zziiaapp...amma
NVAE model and considerably improves its performance by applying data augmen- tations. We did not use any angmentations compare fairly with all other hierarchical VAFs and able get better NLL than NVAE and recent. hierarchical VAE models.
SCCN VAE Sinha Dieng 2021

---

Published in Transactions on Machine Learning Research (11/2024)
&
1
9
8
0
1
8
T
X
A
Figure 3: Unconditional samples from the Diffusion-based VampPrior (top) and corresponding samples from the DVP- VAE (bottom).
Figure Generalive reconstructions. The top row using a pscudoinput sampled from r(uxx) only.
Table 2:
Test performance for the model trained for 500 epochs (or approximately 200k training iterations) on CIFARIO.
Table 3: Training settings for the model trained on CIFAR-10 compared to two VDVAE implementations.
VDVAE
EFFICIENT- VDVAE
DVP-VAE
Child 2021
Hazami et al.
2022
(ours)
L
SIZE
BPDS 4
AU
L
45
47
45
20
24M
2.99
94%
SIZE
39M
57M
38M
28
32M
2.94
93%
OPTIMIER
ADAMW
ADAMAX
ADAMAX
LEARNING RATE
2e-4
le-3
1e-3
36
40M
2.89
98%
GRAD. SMOOTHING
Yes
44
48M
2.84
97%
GRAD. SKIP
400
800
TRAINING ITER.
1.1M
0.8M
0.4M
TEST BPD
2.87
2.87
2.86
normalization Vahdat & Kautz, 2020), or softmax parameterization variances (Hazami ct al. 2022). Wo
use the Adamax version of the Adamn optimizer (Kingma & Ba 2015) following Hazami el. al. (2022), as it demonstrated much better convergence for the model with a mixture of discretized logistic likelihood.
First, observe consistent performance improvement stochastic layers. In Table 2. report test performance and the percentage of active units (see Section 6.3 for details) for models of different stochastic depths trained on the CIFAR10 dataset. We train each model for 500 epochs, which corresponds to less than 200k training itcrations. Additionally, we report gradicnt norms and training and validation losses for all four models in Appendix
increase the model size and the number of
To demonstrate the advantage of the proposed architecture, we compare our model closest deep hier- archical VAE architecture: Very Deep VAE[Child 2021). For this experiment, we chose hyperparameters closest to Table 4 in Child 2021 (CIFAR-10). That is, our model has 45 stochastic layers and comparable number of trainable parameters. Furthermore, following Child 2021), we train this model with batch size of 32, & gradient clipping threshold of 200, and an EMA rate of 0.9998. However, in DVP-VAE, we were able to eliminate gradient skipping and gradient smoothing. We report the difference in key hyperparameters and test performance in Table We also add a comparison with Efficient- VDVAE (see Table in Hazami et al. (2022)) We obscrve that DVP-VAE achieves comparable performance within much fewer training itcrations than both VDVAE implementations.
Latent Aggregation Increases Latent Space Utilization Next, we test the claim that latent variable aggregation discussed in Sec. 4.3 improves latent space utilization. We use Active Units (AU) metric (Burda et al. 2015) which can be calculated for a given threshold 8 as follows:
Mi EL 1=11 -==
[Ar
8)
AU
(23)
Σ1 Mi
where A, = Var (x) B
By go (x) Bys (as (21 [2]].
(24)
10

---

Published in Transactions on Machine Learning Research (11/2024)
Table 4: Active Units for the DCT-VAE with and without latent aggregation.
plar
105
95
Tea N.LL
Dhatosr your
910
100
LATENT AGGR.
SIze
L
AU 1
Test
95-
MNIST
80
76
x
0.7M
8
33.2%
11
Pseudoinput size
Paeutoinspt size
0.7M
8
91.5%
(a) MNIST
(b) OMNIGLOT
OMNIGLOT
Figure 5: Ablation study of for the pseudoinputs type (DCT and Downsampled image), pscudoinputs prior (Diffusion model and Mixture of Gaussians) and pseudoinpuls size (ranging from 3 x to 11 x 11). Each configuration is trained with four different random seods.
x
1.3M
8
71.3%
1.3M
8
93. 93.4%
CIFARIO
19.5M
28
98%
Seel D
Seal
Seal2 2
Seal
Sexl
Seal
Seal 2
Seal
6
8
Table 5: Test NLL for the model with and
8
7
4
4
9
2
G
5
9
9
0
8
+
9
without pseudoinputs. Averaged over ran- dom soods, standard deviation parenthesis.
8
4
1
8
4
4
3
5
7
4
4
0
8
4
9
8
4
5
8
5
3
3 9
6
7
9
4
5
8
+
PSEUDOINPUTS SIzE
NLLL
8
7
9
4
7
5
7
0
8
4 4
MNIST
8
7
4
3
5
3 5
6
9
6
8
4
X
0.6M
78.85 (0.24)
0.7M
77.10 (0.05)
(a) Learnable pseudoinputs
(b) DCT pseudoinputs
OMNIGLOT
Figure Samples from the pseudoinputs prior f(u) (top row) and correspouding samples from the model po (x = ú) (other rows). Columns corresponds to models trained with different random seeds.
X
1.1M
89.52 (0.23)
1.3M
89.07 (0.10)
Here Mj is the dimensionality of the stochastic layer 1, [B] is the Iverson bracket, which equals 1 if B is true and 0 otherwise, and Var stands for variance. Following Burda et al. (2015), we use the threshold = 0.01. The higher the share of active units, the more efficient the model is in utilizing its latent space.
We report results in Table 4 and observe that the model with latent aggregation always attains more than 90% of active units. Furthermore, latent aggregation considerably improves the utilization of the latent space if we compare it with exactly the same model but with the conditional likelihood parameterized using deterministic feature from the TowDown path (see Eq. 22
Amortized VampPrior Improves BPD Further, we test how the proposed amortized VampPrior im- proves model performance as measured by the negative log-likelihood. We report results in Table and observe that DVP-VAE always has better NLI metric compared to the deep hierarchical VAE with the same architecture and number of stochastic layers. Due to additional diffusion- DVP-VAR has slightly more trainable parameters. However, because of the small spatial dimensionality of the pseudoinputs, we were able to keep the size of the two models comparable.
prior over pseudoinputs,
Pseudoinputs type, size and prior We conduct extensive ablation study regarding pseudoinputs. First, train VAE with two types of pseudoinputs: DCT and Downsampled images. Moreover, vary the spatial dimensions of the pscudoinputs between 3 x 3 and 11 x 11. We expect that smaller pscudoinputs size will be an easier task for the prior f(u), but will constitute a poorer approximation of an optimal prior. The larger pseudoinput size, on the other hand, results in a better optimal prior approximation since more information about the datapoint x preserved. However, bocomes harder for the prior achieve good results since we keep the prior model size fixed. In Figure we observe that the DCT- transformation performs consistently better across various sizes and datasets.
11

---

Published in Transactions on Machine Learning Research (11/2024)
Trainable Pseudoinputs In the VampPrior, the optimal prior is approximated using learnable pseudoin- puts. In this work, on the other hand, we propose to use fixed linear transformation instead. To further verify whether a fixed transformation like DCT is reasonable, we checked a learnable linear transformation. We present in Figure G that the learnable lincar transformation of the input cxhibits unstable behavior in terms of the quality of learned pseudoinput. The top row of Figure G (a) shows samples from the trained prior and the corresponding samples from the decoder. We observe that only one out of four models with learnable pseudoinputs was able to learn a visually meaningful representation of the data (seed 0), which also resulted in very high variance of the results (rows below). For other models (e.g., Seed and Seed 2), the same pseudoinput sample corresponds to completely different datapoints.
This lack of consistency motivates us to use a non-trainable transformation for obtaining pseudoinputs. In Figure 0 (b); we show the expected behavior of sampling semantically meaningful pseudoinputs that is consistent across random sceds.
Table 6: Wall-clock time, memory consumption and Scalability Here, we study how DVP-VAE scales sampling time for DVP and VampPrior.
as we increase model size and input size in compari- son with VampPrior. For this, we implement Vamp- Prior as proposed by TTonezak & Welling (2018), where the top latent variable is trained with the VampPrior and the other layers with the conditional Gaussian (see Eq. 7 We use the same architec- ture as in main experiments for MNIST (32 channels) and CIFARIO (see Table 7 Additionally, we train model with doubled number of channels on MNIST (64 channels).
TRAIN TIME
GPU
SIZE
SAMPLE TIME
K
MEM.
MNIST (32 CHANNELS)
DVP
N/A
458
3Gb
0.7M
2.2
500
458
4Gb
1.0M
0.2
Vamp
1000
60s
6Gb
1.4M
0.3
MNIST (64 CHANNELS)
In Table 6 we report the training time (second per epoch), GPU memory utilization and total number of trainable parameters. We observe that Vamp- Prior almost always utilizes significantly more mem- ory and requires longer training time. The differ- ence is less visible on a small model and small input size (MNIST, 32 channels). However, as we double number of channels, both training time and memory utilization for VampPrior grows much faster. As result, for bigger input size (CIFAR10 dataset) model with more than 750 pscudoinputs does not fit into single A100 GPU. For 500 pscudoinputs, memory utilization is already 2.5 times higher for VampPrior compared to DVP-VAE.
DVP
N/A
50s
4Gb
2.4M
2.3
500
65s
7Gb
2.7M
0.5
Vamp
1000
90s
10Gb
3.1M
0.8
CIFARI0
DVP
N/A
370s
12Gb
19.5M
4.0
500
430s
29Gb
20.5M
1.5
Vamp
750
500s
38Gb
21.3M
1.8
1000
OOM
>40G
22.1M
7 Limitations
One limitation of DVP-VAE is the long sampling time reported Table 6 (scc. per 1000 images). It is direct consequence of using diffusion-based prior. While diffusson is powerful generative model that allows us to achieve outstanding performance, it is known to be slow at sampling. In our experiments, use 50 diffusion steps to generate pseudoinput, however, there are efficient distillation techniques (Salimans & 2022 Geng et al. 2024) that can be applied to mitigate this issue and reduce the number of diffusion steps to just a single forward pass through the model. We leave this optimization for future work.
Ho
Furthermore, within DVP-VAE, we add an extra learnable block to the model, namely, the diffusion-based prior over pseudoinputs, as a result, additional modelling choices should be made. We provide an ablation study to show the effect of these choices on the performance of the model.
12

---

Published in Transactions on Machine Learning Research (11/2024)
8 Conclusion
In this work, we introduce DVP-VAE, new class of deep hicrarchical VAEs with the diffusion-based Vamp- Prior. We propose to use a VampPrior approximation which allows us to use it with hierarchical VAEs with little computational overhead. show that the proposed approach demonstrate competitive performance terms of the negative log-likel on three benchmark datasets with much fewer parameters and stochastic layers compared to the best performing contemporary hierarchical VAEs.
13

---

Published in Transactions on Machine Learning Research (11/2024)
References
Nasir Ahmed, T Natarajan, and Kamisetty R Rao. Discrete cosine transform. IEEE Transactions Computers, 100(1):90 93, 1974.
Krizhevsky Alex. Learning multiple layers of features from tiny images. https://www. serprr//ttaamiing-ftteeeeeseee TR. pdf, 2009.
toronto.
Tfigencia Apostolopoulon, Ian Char, Elan Rosenfeld, and Artur Dubrawski. Deep attentive variational infer- ence. In ICLR, 2022.
Matthias Bauer and Andriy Mnih. Resampled priors for variational autoencoders. In The 22nd International Conference Artificial Intelligence and Slalislics, pp. 66 75. PMLR, 2019.
Alican Bozkurt, Babak Esmaeili, Jean-Baptiste Tristan, Dana H Brooks, Jennifer Dy, and Jan- Willem van de Meent. Rate-regularization and generalization in vaes. arXiv preprint 1xXi:1111.00594,, 2019.
Yuri Burda, Roger Grosse, and Ruslan Salakhutdinov. Importance weighted autoencoders. arXiv, 2015.
Xi Chen, Diederik P Kingma, Tim Salimans, Yan Duan, Prafulla Dhariwal, John Schulman, Nya Sutskever, and Pieter Abbeel. Variational lossy autoencoder. ICLR, 2017.
Rewon Child. Very deep vaes generalize autoregressive models and can outperform them on images. In ICLR, 2021.
Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans image synthesis. NewrIPS, 2021.
Adji Dieng, Yoon Kim, Alexander M Rush, and David M Blei. Avoiding latent variable collapse with generative skip models. In AISTATS, 2019.
Evgenii Egorov, Anna Kuzina, and Evgeny Burnaev. Boovae: Boosting approach for continual learning of vac. Advances Neural Information Processing Systems, 34:17889-17901, 2021.
Zhengyang Geng, Ashwini Pokle, and J Zico Kolter. One-step diffusion distillation via deep equilibrium models. Advances in Neural Information Processing Systems, 36, 2024.
Rafael Kömez-Bombarelly Jennifer N Wei, David Duvenaud, José Miguel Hernández-Lobato, Benjamín Sánchez-Lengeling Dennis Sheberla, Jorge Aguilora-Iparrguuiree, Timothy D Hirzel, Ryan P Adams, and Alán Aspuru-Guzit Automatic chemical design using a data-driven continuous representation of molecules. ACS central science, 2018.
Ishaan Gulrajani, Kundan Kumar, Faruk Ahmed, Adrien Ali Taiga, Francesco Visin, David Vazquez, and Aaron Courville, Pixelvae: A latent variable model for natural images. arXie preprint arXiv:1611.0500 100011,, 2016.
Louay Hazami, Rayhane Mama, and Ragavan Thurairatnam. Efficient vdvac: Less is more. arXiv preprint arXiv:2203.1377 137511, 2022.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. NewrIPS, 2020.
Matthew D Hoffman and Matthew Johnson. Elbo surgery: yet another way to carve up the variational evidence lower bound. In Workshop in Advances in Approzimate Bayesion Inference, NIPS, volume 1, 2016.
Emiel Hoogeboom, Victor Garcia Satorras, Clément Vignac, and Max Welling, Equivariant diffusion for molecule generation in 3d. In International conference on machine learning, pp. 8867-8887. PMLR, 2022.
Chin-Wei Huang, Jae Hyun Lim, and Aaron C Courville. A variational perspective on diffusion-basee generative models and score matching. NeurIPS, 2021.
Zhuxi Jiang, Yin Zheng, Huachun Tan, Bangsheng Tang, and Hanning Zhou. Variational deep embedding: An unsupervised and generative approach to clustering. arXiv preprint ar Xiv: 1111161111188 2016.
14

---

Published in Transactions on Machine Learning Research (11/2024)
Michael I Jordan, Zoubin Ghahramani, Tommi S Jaakkola, and Lawrence K Saul. An introduction to variational methods for graphical models. Machine learning, 37(2)1183-238 1999.
llees Khemakhem, Diederik Kingma, Ricardo Monti, and Aapo Hyvarinen. Variational autoencoders and nonlinear iea: A unifying framework. In International Conference on Artificial Intelligence and Statistics, pp. 2207-2217. PMLR, 2020.
Diederik P Kingma and Jimmy Lei Ba. Adam: method for stochastic gradient descent. In ICLR: international conference on learning representations, pp. 1-15, 2015.
Diederik P. Kingma and Max Welling. Auto-cncoding variational bayes. In ICLR, 2014.
Diederik P Kingma, Tim Salimans, Ben Poole, and Jonathan Ho. Variational diffusion models. In NewrIPS, 2021.
Durk Kingma, Tim Salimans, Rafal Jozefowicz, Xi Chen, 1lya Sutskover, and Max Welling. Improved variational inference with inverse auloregressive flow. NewrIPS, 2016.
Alexej Klushyn, Nutan Chen, Richard Kurle, Botond Cseke, and Patrick van der Smagt. Learning hicrarchical priors vaes. Advances in newral information processing systems, 32, 2019.
Brenden M Lake, Ruslan Salakhutdinov, and Joshua B Tenenbaum. Human-level concept learning through probabilistic program induction. Science, 350(6266):1332 1338, 2015.
Yann LeCun. The mnist database of handwritten digits. http://yamm. lecun. /enelb/mmmmms//.////mm//// 1998.
Lars Maalde, Casper Kaae Sønderby, Søren Kaae Sønderby, and Ole Winther. Auxiliary deep generative models. In International conference on machine learning, pp. 1445 444--143. PMLR, 2016.
Lars Maaløe, Marco Fraccaro, Valentin Liévin, and Ole Winther. Biva: A very deep hierarchy of latent variables for generative modeling. NeurIPS, 2019.
Eric Nalisnick, Lars Hertel, and Padhraic Smyth. Approximate inference for deep latent gaussian mixtures. In NIPS Workshop on Bayesian Deep Learning, volume 2, pp. 131, 2016.
William B Pennebaker and Joan L Mitchell. JPEG: Still image data compression standard. Springer Science & Business Media, 1992.
Adeel Pervez and Efstratios Gavves. Spectral smoothing unveils phase transitions in hicrarchical variational autoencoders. ICML, 2021.
Rajesh Ranganath, Dustin Tran, and David Blei. Hierarchical variational models. In International conference on machine learning, pp. 324 333. PMLR, 2016.
Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and approximate inference in deep generative models. In ICML, 2014.
Hossein Sadeghi, Evgeny Andriyash, Walter Vinci, Lorenzo Buffoni, and Mohammad H Amin. Pixelvae +: Improved pixelvac with discrete prior. arxiv, 2019.
Tim Salimans and Jonathan Ho. Progressive distillation for fast sampling of diffusion models. International Conference on Learning Representations, 2022.
Tim Salimans, Diederik Kingma, and Max Welling. Markov chain monte carlo and variational inference: Bridging the gap. In International conference on machine learning, pp. 1218-12211 1226. PMLR, 2015.
Samarth Sinha and Adji Bousso Dieng. Consistency regularization for variational auto-encoders. NewrIPS, 2021.
Jascha Sohl-Dickster Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In ICML, 2015.
15

---

Published in Transactions on Machine Learning Research (11/2024)
Casper Kaae Sønderby, Tapani Raiko, Lars Maalće, Søren Kaac Sønderby, and Ole Winther, Ladder varia- tional autoencoders. NeurIPS, 2016.
Jakub Tomezak and Max Welling. Vae with a vampprior. In AISTATS, 2018.
Jakub M. Tomezak. Deep Generalive Modeling. Springer Cham, 2022.
Linh Tran, Maja Pantic, and Mare Peter Deisenroth. Cauchy schwarz regularized autoencoder. Journal of Machine Learning Research, 23(11)::1-37, 2022.
Belinda Tzen and Maxim Raginsky. Neural stochastic differential equations: Deep latent gaussian models in the diffusion limit. arx, 2019.
Arash Vahdat and Jan Kautz. Nvae: A deep hierarchical variational autoencoder. NewrIPS, 2020.
Arash Vahdat, Karsten Kreis, and Jan Kautz. Score-based generative modeling in latent space. NewrIPS, 2021.
Aaron Van Den Oord, Oriol Vinyals, et al. Neural discrete representation learning. NeurIPS, 2017.
Antoine Wehenkel and Gilles Louppe. Dilfusion priors in variational autoencoders. In INNPICML, 2021.
16

---

Published in Transactions on Machine Learning Research (11/2024)
A Diffusion probabilistic models
Diffusion Probabilistic Models iiffusion-based Deep Generative Models (Ho et al. 2020
Sohl- Bohl-Dickstein
ct al. 2015 constitute class of generative models that can be viewed as special case of the Hicrarchical
VAEs (Hunng et al. 2021 Kingma et al. 2021 Tomezak
2022 Tzen & Raginsky 2019 Here, we follow
the definition of the variational diffusion model (Kingma et al. 2021 ). We use the diffusion model as a prior over the pseudoinputs u.
Forward diffusion process
The forward diffusion process runs forward time and gradually adds noise to the input as follows:
g(yelu) aN (1 aI)
(25)
where ye auxiliary latent variables indexed by time t [0, 1], ar is chosen in such way that the signal-to-noise ratio decreases monotonically over time.
Since the conditionals in the forward diffusion can be seen as Gaussian linear models, we can analytically calculate the following distributions for > 8:
(yrlyy,, u) NN (y ut,
(26)
0t (1 a²)
a - oه
where ju(y;, u)
0g (1 a?)
+
(1 at) )
u.
(27)
(a² -a7) a1
8lt,8)
(28)
(1 oi
Backward diffusion process
Similarly, define generative model, also referred to the backward (or reverse) process, Markov chain with Gaussian transitions starting with r(y1) N(y, We discretize time uniformlly into T timestamps of length 1/T:
T
r(yo.....y1) = r(y1) II ry (y1
(29)
where ry (y:) 1y (x,,,))
The Likelihood term
Common practice is to define the likelihood term as being proportional to the first step of the forward process: rulyo)) xx (yolu). Since we assume that the pseudoinput random variable u is continuous, we get the Gaussian likelihood distribution:
(ulyo) (V(llyo/on, (
(30)
Note that the same likelihood term was used for continuous atom positions in the equivariant diffusion model (Hoogeboom et al. 2022
17

---

Published in Transactions on Machine Learning Research (11/2024)
Training objective
We can use
25 and 26 to define the variational lower bound as follows:
(u) 2 (u, 7) Bolyol u) [Inr(ulyo]]
Del [
(31)
-Lo
Σ
Baterrl DEL (w---r//aommmllll/
by
Here refer to Tog as the reconstruction loss, L the prior loss, and Fop as the diffusion loss with T steps.
18

---

Published in Transactions on Machine Learning Research (11/2024)
B Training Stability: depth
We report the f2-norm of the gradient for each itcration training for models of different stochastic depth in Figure 7 We observe very high gradient norms for the first few training iterations but no spikes at later stages of training- This happens because we did not initiate the KT-term to be zero at the beginning of the training and it tends to be large at the initialization step for very deep VAEs. However, we did not observe our model diverge since the gradient is clipped reasonable values (200 in these experiments) and after the first few gradient updates, the KT-term goes to reasonable values. Moreover, we plot. training and validation losses in Figure
L 20
L 28
L 36
L 41
400
200
100000
100000
100000
100000 Training Iteration
Training Iteration
Training Ireration
Training Iteration
Figure 7: The gradient norm at each training iteration. Models were trained stochastic depths L.
CIFAR10 with different
Train
Validation
5
5
20
20
28
28
4
I. 36
4
= 36
L 11
L 11
3
3
0
100 200 300 400 500
0
100
200 300 400 500 Epoch
Epoch
Figure 8: The ELBO per pixel on train (left) and validation (right) dataset. Models were trained on CIFAR10 with different stochastic depths L.
19

---

Published in Transactions on Machine Learning Research (11/2024)
C Model details
C.1 Hyperparameters
In Table we report all hyperparameter values that were used to train the DVP-VAE.
The pseudoinput prior We use the diffusion generative model as a prior over psendoinputs. As a backbone, we use the UNet implementation from (Dhrriwa & Nichol 2021), available on GitHul with the hyperparameters provided in Table
Table 7: Full list of hyperparameters.
MNIST
OMNIGLOT
CIFARIO
# Epochs
300
500
3000
Batch Size (per GPU)
250
250
128
# GPUs
1
1
Optimizer
Adamax
Adamax
Adamax
Optimization
Scheduler
Cosine
Cosine
Cosine
Starting LR
le-2
le-2
3e-3
End LR
lc-5
lc-4
lc-4
LR warmup (epochs)
2
5
Weight Decay
1e-6
1e-6
1e-6
EMA rate
0.999
0.999
0.999
Grad. Clipping
5
2 -10
150
log a clipping
-10
-10
L.
8
28
10 32
Latents
4 142,
x 14
8 x 162.
Latent Sizes
x 72.
4 x 72.
6 x 82
4 x 42.
Latent Width (channels)
1
3
Context Size
1 7
1 x
7 x
Nene blocks
3
3
4
Architecture
ResBlock C.n
32
80
128
ResBlock Chid
32 SiLU Bernoulli
40
96
Activation
SiLU
SiLU
Likelihood
Bernoulli
Discretized Logisite
mixture comp
10
# Diffusion Steps # Scales in UNet
50
50
50
Context Prior
1
1
# ResBlocks per Scale
2
3
Channels
16
16
32
8 schedule
lincar
lincar
lincar
log SNR min
-6
-6
-10
log SNR max
7
7
20

---

Published in Transactions on Machine Learning Research (11/2024)
D
Samples
In Figure 9
we present non-chery-picked unconditional samples from our DVP-VAE and in Figure 10 we present non-cherry-picked unconditional samples from the pseudoinput prior f(u).
9
0
7
1
5
4
1
5
d
2
0
7
(
6
3
8
5
3
8
J
2
2
0
1
9
9
8
9
6
0
2
5
4
e
4
9
0
4
9
4
2
G
[
x
y
8
4
4
5
$
Q
3
2
10
K
1
a
3
9
3
2
4
4
9
8
b
G
]
4
8
6
4
7
1
8
8
V
0
VC
I
4
1
O
0 0
3
4
0
7
4
2z
V
2
Q
4
8
6
9
8
T
9
5
7
8
3
3
0
7
9
3
3
yp
w
3
6
b
7
r
(a)
MNIST
(b)
OMNIGLOT
(c) Cifar10
Figure 9: Uncoditional samples.
7 7
4
4
3
7
6
3
7
3 3
1
8
2
9
6
6
0
3
6
5
5
8
/
7
4
4
0
9
7
3
7
5
5
0
6
0
(a) MNIST
(b) OMNIGLOT
(c) Cifar10
Figure 10: Uncoditional samples of pseudoinputs.
21


Reviews:
[{"review_id": "bZX1PWhqc2", "reviewer": "Reviewer_xFeL", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "While the importance of the prior in VAEs is well-known and VampPrior has been shown to work empirically, it is still difficult to scale to deeper hierarchies, which limits the potential expressivity of the VAE. The authors address this issue. The proposed VAE method achieves good empirical performance on benchmarks while using fewer parameters than other VAE models. The ablation studies are solid, although additional studies would be beneficial.\n\nSome weaknesses:\n+ The explanation of the method is hard to follow, and could be improved. To provide a concrete example, consider the following sentence in Section 3.3: “However, unlike iVAE, we do not treat u as an observed variable. Instead, we introduce the pseudo-input prior distribution pγ(u) with learnable parameters γ. “\nThis is confusing for the reader. Based on the equation (9), it looks like u is a noisy version of x, where mean is given by deterministic mapping f. There are no trainable hyperparameters in f, which is also hinted by the sentence “the first important component in our approach is the form of the non-trainable transformation from the input to the pseudoinput space.”, so for the reader is unclear what is γ? Further, in the equation (15), there is no dependence on γ as hinted by the ELBO arguments. The authors later explain γ, and how the pseudo-point distribution is approximated, but the damage is already done: the story is hard to follow.\n+ There are many separate design proposal on top of hierarchical VAEs. To understand impact of these c, ablation studies are very useful. (thanks of having some in the manuscript). It would be useful to understand which of these are the most important in explaining the empirical performance in Table 1.\n+ Please report a measure of variation in the experimental results (Table 1). This is crucial as there are quite few trials (experiments replicated over 4 seeds)", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The authors propose an extension of the VampPrior (Tomczak & Welling, 2018), which allows for using more latents in hierarchical VAE. Futhermore, there multiple proposed desing choices for hierarchical VAEs such as \n+ Discrete cosine transform from datapoints to the mean of the pseudopoints distribution\n+ Diffusion-based approximation for the pseudo-point distribution", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "+ Report stds, especially in Table 1\n+ Improve the writing\n+ More ablation studies would be useful, e.g. break down the impact of design choices (how f modelled, prior model, etc.) in Table 1", "broader_impact_concerns": "No", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "usjjrK7zDH", "reviewer": "Reviewer_gcra", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "## Strengths\n\n1. To the best of my knowledge, the introduced VampPrior-like class of priors is novel and the latent aggregation component really does seem to improve the utilization of the latent space of the VAE (as shown in Table 4).\n\n2. The results of the proposed method are quite good (judging from Table 1 and Figures 3 and 4) and I am inclined to believe that this model is currently one of the best deep hierarchical VAEs in existence. However, I am not a fan of the \"state-of-the-art\" claims made throughout the paper. Amongst VAEs, there is at least one model that is better on CIFAR-10 (as measured by BPD), namely CR-NVAE. According to Papers with Code [a], there are also at least four diffusion models with a better BPD score.\n\n## Weaknesses\n\n1. Although the work has some novel elements, it ultimately relies quite a bit on existing methods. Namely, it borrows heavily from the Very Deep VAE (VDVAE) paper (Child, 2021) and the VAE with a VampPrior paper (Tomczak & Welling, 2018). In essence, one can view this paper as presenting advantageous modifications of these previous methods. This is still useful, but limits the novelty of this paper.\n\n## Verdict\n\nThis paper presents some novel modifications to existing deep hierarchical deep VAEs. In particular, it introduces a VampPrior-like class of priors and introduces a latent aggregation component that significantly improves the utilization of the latent space of the deep VAE. The empirical results are also quite solid; amongst VAEs tested on MNIST, CIFAR-10, and Omniglot, this model is one of the best. That being said, the paper borrows heavily from Very Deep VAE (VDVAE) (Child, 2021) and VampPrior (Tomczak & Welling, 2018), which limits the novelty, somewhat. I believe the claims made in the paper are reasonably well-supported by the empirical results. I also believe this paper is of interest to at least some portion of TMLR's audience, specifically those who work on VAEs and fusions of VAEs with diffusion models. As a consequence, I recommend an accept rating for this paper.\n\n### References\n\n[a] https://paperswithcode.com/sota/image-generation-on-cifar-10?metric=bits%2Fdimension", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposes a new hierarchical VAE, based closely on the modification of existing VAE components. In particular, the authors propose a new \"Variational Mixture of Posteriors\" prior (VampPrior)-like approximation of the optimal prior (i.e. the aggregated posterior), which can scale to deep hierarchical VAEs. They also propose a latent aggregation component that improves the utilization of the latent space of the VAE. These changes form the basis of the proposed hierarchical VAE; on the considered benchmark datasets of MNIST, CIFAR-10, and Omniglot, the proposed method achieves the best results of any deep hierarchical VAE.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "The writing in the paper is for the most part reasonable and the idea is described clearly enough; the illustrative figures are quite helpful (e.g. Figures 1 and 2). That being said, I would like the authors to tone down the \"state-of-the-art\" claims a bit, given the nuances I mentioned above in the \"Strengths and Weaknesses\" section. I would still like to provide a non-exhaustive list of minor corrections below:\n\n- Abstract: \"using much fewer parameters\" -> \"using fewer parameters\"\n- Section 2.2: \"a greedy booting approach\" -> \"a greedy boosting approach\"\n- Section 4.1: Please use \\citet for the Child, 2021 reference, since you are referring to the citation directly in text (as opposed to doing so parenthetically). This error occurs a few times (e.g. in Section 6.3 with \"following (Hazami et al., 2022)\"); please be careful with in-text citations (\\citet) versus parenthetical citations (\\citep or just \\cite).\n- Section 4.2: \"At the test time\" -> \"At test time\"\n- Section 5: \"using unimodal prior\" -> \"using a unimodal prior\"\n- Section 5: \"Original VAE formulation\" -> \"The original VAE formulation\"\n- Section 5: \"In this work, in contrary, we use\" -> \"In this work, in contrast, we use\"\n- Section 7: \"with little computations overhead\" -> \"with little computational overhead\"", "broader_impact_concerns": "This work deals with the development of a hierarchical deep VAE and is fairly model-oriented. There are no ethical concerns beyond the usual ones associated with any work that deals with generative AI. I do not believe this paper requires a broader impact statement.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "xauY1FUVVb", "reviewer": "Reviewer_4t7J", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n- The paper is well-structured and clearly written. The authors provide a comprehensive overview of the related work, the methodology, and the experimental results.  \n- The diffusion-based VampPrior provides a way of amortization that boost performance and scalability.\n- The paper provides thorough empirical validation of the proposed method across multiple benchmark datasets.\n\nWeaknesses:\n- Comparing wall time of training and inference would be a plus to demonstrate the claim of improved scalability.\n- Lack of substantial discussion, especially limitations. \n- Figures could be polished.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper presents a novel approach to deep hierarchical variational autoencoders (VAEs) by amortizing VampPrior. The authors aim to enhance the performance and stability of VampPrior. They empirically validate their method on standard benchmark datasets, demonstrating improved training stability and latent space utilization.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "- Comparing wall time of training and inference\n- Discussion about gain and loss, limitations. (Maybe use smaller fonts in the tables for more space.)\n- Figure 2: Use only either av.pool or av.pooling for all occurrences. Use consistent font for latex and non-latex text. (Optional) use edge color instead of fill color for boxes.", "broader_impact_concerns": "No broader impact concerns.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper is in general solid piece of research, and the reviewers had clear consensus on both its strengths and weaknesses. The technical solution is new and sound and it was found to work well in practice. However, the work is somewhat incremental and building heavily on previous deep VAE papers, the writing could still be improved, and the impact is reduced by limited discussion. Nevertheless, the paper is clearly suitable for publication in TMLR.
Decision: Accept (poster)