AdaVAE: Bayesian Structural Adaptation for Variational Autoencoders
Authors: Paribesh Regmi
Rui Li
Venue: NeurIPS 2023 poster (2023)
Keywords: nonparametric Bayes
variational autoencoders

Abstract: The neural network structures of generative models and their corresponding inference models paired in variational autoencoders (VAEs) play a critical role in the models' generative performance. However, powerful VAE network structures are hand-crafted and fixed prior to training, resulting in a one-size-fits-all approach that requires heavy computation to tune for given data. Moreover, existing VAE regularization methods largely overlook the importance of network structures and fail to prevent overfitting in deep VAE models with cascades of hidden layers. To address these issues, we propose a Bayesian inference framework that automatically adapts VAE network structures to data and prevent overfitting as they grow deeper. We model the number of hidden layers with a beta process to infer the most plausible encoding/decoding network depths warranted by data and perform layer-wise dropout regularization with a conjugate Bernoulli process. We develop a scalable estimator that performs joint inference on both VAE network structures and latent variables. Our experiments show that the inference framework effectively prevents overfitting in both shallow and deep VAE models, yielding state-of-the-art performance. We demonstrate that our framework is compatible with different types of VAE backbone networks and can be applied to various VAE variants, further improving their performance.

AdaVAE: Bayesian Structural Adaptation for Variational Autoencoders
Paribesh Regmi
Rui Li
Rochester Institute of Technology (pr8537, rxlics]@rit.ee edu
Abstract
The neural network structures of generative models and their corresponding infen- rence models paired in variational autoencoders (VAEs) play critical role in the models' generative performance. However, powerful VAE network structures are hand-crafted and fixed prior to training. resulting in a one-size-fits fits-all approach that requires heavy computation tune for given data. Moreover, existing VAE regularization methods largely overlook the importance of network structures and fail to prevent overfitting in deep VAE models with cascades of hidden layers. To address these issues, we propose Bayesian inference framework that automati- cally adapts VAE network structures data and prevent overfitting as they grow deeper. We model the number of hidden layers with beta process to infer the most plausible encoding/decoding network depths warranted by data and perform layer-wise dropout regularization with conjugate Bernoulli process. We develop a scalable estimator that performs joint inference on both VAE network structures and latent variables. Our experiments show that the inference framework effectively prevents overfitting in both shallow and deep VAE models, yielding sate-of-the--al performance. We demonstrate that framework is compatible with different types of VAE backbone networks and can be applied to various VAE variants, further improving their performance.
Introduction
The inference models and the generative models paired in variational autoencoders (VAEs) com- monly constructed with neural networks, i.c., encoding networks and decoding networks, respectively 12 Extensive research efforts show that well-designed mccodin/decooiin network structures for VAE models can constantly achieve state-of- sere--of-hh-eatt generative performance compared to other generative models 4 6). However, powerful VAE network structures hand-crafted and fixed prior to training. The issue with fixed network structures is that shallow ones limit VAE models' expressiveness, whereas overly deep networks are slow to use and prone to overfitting. Traditional model selection by training different VAE network structures for given data is difficult since finding optimal hyperparameters for each candidate structure a daunting task, and training large VAE structures requires significant computation. On the other hand, network structure adaptation methods for discriminative model settings [7] 9 cannot straightforwardly applied to address the unique challenge posed by VAE estimation along with the latent variables.
Although the network structures play a critical role in the performance of VAE models, they are largely overlooked by current VAE regularization methods. This renders their failure to prevent overfitting when the network structures grow deeper. Amortized inference regularization (AIR) proposes two approaches: injecting random noise to the VAE objective for inference or directly restricting the inference models set of smooth functions [1] 11] Another approach in VAE regularization
"Corresponding author
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

process
Beta process
GO
Bernoulle pracess
SEROOA TIrtOnII
Encoding netwark
Decoding network
(b)
Figure I: (a): Demonstration of the VAE network structure inference framework. Beta processes induce infinite number of hidden layers for encoding/decodineg networks, and its conjugate Bernoulli process prunes the neurons in each layer with layer-wise activation probability from the beta process. Filled circles indicate activated neurons, corresponding to a sample of 1 from the Bernoulli process, while empty circles correspond to deactivated neurons, corresponding to sample of O. (b) and (c): Two settings the stick- breaking constructions beta process. The sticks on top are random draws from the process, which act the layer-wise activation probabilities. A stick location on corresponds to a hidden-layer function f/ E F, and the height denotes its activation probability mi The bottom shows the conjugate Bernoulli processes to activate or deactivate neurons in each layer moolumn-wisely).
incorporates additional constraints to the VAE objective, which enforces a similarity between original input and its semantic preserving transformation in the latent representation [12 [3]. Nonparametric Bayesian inference such as Indian buffet process deep generative models (IBP-DGM) [1] and beta- Bernoulli process VAE (BB-VAB [S] address overfitting by inferring the dimensionality of VAE latent variables using Indian buffet process, specifically, marginalized beta- eeta-Bernoulli process prior. While these regularization methods are effective for shallow network structures, they fail to prevent overfitting in deep VAE models.
We propose a novel Bayesian inference framework that automatically adapt VAE network structures by inferring the most plausible encoding/decoding network depths based on the given data, as demonstrated in Figure [ achieve this, we employ beta processes [6 17 to model the number of hidden layers in the encoding/decoding networks, allowing for infinite depths. A conjugate Bernoulli process is utilized to prune the neurons in each layer based on layer-wise activation probabilities generated by the beta process. In addition, to enable efficient joint inference both the network structures and the latent variables, we extend multiply importance weighted autoencoder (MIWAE) [8] by introducing an additional sample size used for Monte Carlo estimation of the network structures to be tuned. Our theoretical and empirical analysis shows that this novel gradient estimation scheme leads a tight lower bound with high signal-to-noise ratio of parameter gradients.
In summary, our contributions are: i) We propose AdaVAE, novel VAE structural adaptation strategy based on Bayesian model selection to enhance model performance. ii) We introduce scalable estimator that facilitates joint inference on both encoding/decoding network structures and latent variables. iii) We conduct comprehensive analysis of VAE's regularization capabilities and demonstrate its ability to effectively mitigate overfitting in both shallow and deep VAE models and achieve state-of of af-thh-art performance. iv) We showcase the versatility of Ada VAE by demonstrating its compatibility with different types VAE backbone networks. It can also be readily applied to various VAE variants, thereby enhancing their performance.
2 Related Works
Variational autoencoders EAEs) have gained popularity as generative models across wide range of applications [19] 20 22
24 Extensive research efforts have focused on enhancing the 31 32). Notably, well-designed encoding /dccodnng neural network structures constantly yield state-of-the-art generative performance [4 6 33].
23
performance of VAEs 25
26
27
28
29
30
Ladder VAE (LVAE) employs shared top-down dependency structure in both the inference and generative models to facilitate information sharing between layers 4. SkipVAF employs skip connections that connect latent variables to each layer in the generative model, mitigating posterior

---

collapse [34]. BIVA consists of a skip- connected generative model and an inference model formed by a bidirectional stochastic inference path [5]. NVAE designs expressive neural network structures tailored for VA based on deep residual networks G Furthermore, advancements in VAE network
structure design have also empowered recent works on bybrid models based on VAES 35 36 37
Current VAE regularization methods often overlook the overfitting effect caused by deep network structures [12] I0 15 [3]. Amortized inference regularization (AIR) re-interprets the amortized
inference model as regularization for maximum likelihood training. AIR encourages the smoothness of the encoding network to restrict the model capacity, effectively mitigating overfitting and improving predictive performance [10]. Specifically. AIR proposes a denoising variational autoencoder (DVAB) that modifies the VAE objective using random perturbation training. An alternative AIR technique is weight- normalzerd inference VAE (WNI- VAE) that directly restricts the encoding networks to set of smooth functions and achieves comparable performance to DVAE. Both approaches exhibit a lower inference gap than standard VAEs [39]. Consistency regularization for VAE (CR- VAE)
tackle the inconsistency problem of the inference models by enforcing the latent representations of an image and its semantic-preserving transformation to be similar, thereby yielding robust latent representations accounting for data variation 112.
extends the regularization techniques employed in semi-supervised learning 440 4
IBP-DGM 14 and BB- ---VAEBL L
as nonparametric Bayesian inference applied to VAE regularization focus on inferring the dimensionality of latent variable via an Indian Buffet Process (IBP) prior. The IBP prior derived by marginalization over beta process, resulting in a binary vector that masks the VAE latent variables. Without accommodating network structures, latent variable regularization alone is inadequate for mitigating overfitting.
Efficient VAE Estimators
Maximum likelihood estimation of a vector of parameters 0 of deep generative model pe(x. h) with x denoting observed variables and h denoting latent variables is intractable general due to the marginalization over h. Amortized variational inference optimizes an evidence lower bound (ELBO) (x) on the log marginal likelihood log (x() = distribution 2s (h(x):
log po (x, h)dh by introducing a variational
log po(x)= 2 f go (o(2)) log g (b|x) dh po po (x - (h(x))l
pa [x h)
(x)
where á denotes variational parameters. For VAES, qa (h(x) denotes an inference model commonly constructed with a neural network. (x(x(h denotes a generative model that can also be parameterized by neural network. The ELBO is optimized with muddent--bseed methods via reparameterizzttoo trick using stochastic Monte Carlo estimators of Co. ..
IWAE obtains a tighter lower bound using K K-sample importance weighting estimate of the log marginal likelihood 42
K pel (x hy )o (ba)
Ca.o By (h/)x) log
Go (bg [x)
where h& 4 (h|xx. It shows that the bound gets tighter with increasing K. However, [88 presents theoretical and empirical evidence that increasing the importance weighted sample size K to tighten the bound degrades signal-to signal-to-noiss ratio (SNR) of parameter gradients estimates for the encoding network, and hurts the learning process. A new estimator (MIWAE) is thus introduced to address the issue of diminishing SNR, and its gradient estimate is:
1
M
1 K
K
po po (x|h )o (hm,,k
AM.K
B. log
33)
M
4o (bm, x
m=1
where hn, 2o (bm,,(()) [x). For a fixed budget M x K of total number of hidden variable samples, the number of samples M reduces the variance estimating the ELBO gradient. K is the importance sample size in IWAE.
3

---

4 VAE Structure Inference Framework
Traditional model selection cannot effectively adapt pre-determined VAE network structures to data without incurring significant computation overhead. We thus propose AdaVAE that enables joint inference on the structures of the encoding/decoding networks using stochastic processes 43 9 and latent variables, and optimizes VAE objective without requiring additional expensive computation.
4.1 Formulation of the Inference Model
Let the prior over the latent variables h be zero-mean isotropic multivariate Gaussian Po(h) = Nh; 0, I). We formulate the inference model by letting the variational distribution over h be multivariate Gaussian with diagonal covariance structure:
ls (h|x, Z) ((,
(4)
where the mean μ and the standard deviation a are outputs of the encoding neural network (x)) with the variational parameters Φb The binary matrix Z [Fot E {0, 1]] denotes the network
structural variable, as in Figure I().
Let fi denote the /-th hidden layer of (x) composed of neurons (i.e., non non-linear activation functions)
f(+). The encoding network (x) has the form:
fi = f(W(((1)) G 2. Oz. LE
(1.2....00) oo]
(5)
where (W, ROxO), and W, is the layer I's weight matrix. @ denotes element-wise multiplication of two vectors, so that we drop out the I-th layer's outputs by multiplying them elementwisely by the column vector z. of Z. Each random variable zet takes the value with T E [0, 1] indicating activation prohability of the I-th layer, as in Figures F (b). O is the maximum number of neurons a layer that is set be the same for all hidden layers. We have skip connections to propagate the outputs of the hidden layers to the output layer. Note that when this network structure used the decoding network for the generative model pe (x h), then h and x are swapped, and denotes the weights parameters. The output-layer can be readily replaced with a logistic function and the normal distribution with a Bernoulli distribution for binary data.
4.2 Beta Process Prior over Layer Number
A beta process B = En Tírr, where &G is unit point mass at ft, is completely random measure over countably infinite set of pairs (fc, mi) [16), where fę denotes hidden-layer function and क is its activation probability m [0, 1j. Its conjugate Bernoulli process can be defined as Zo BeP(B), where Z. E, ōnzōr, is at the same locations or, as B where 2ot are independent Bernoulli variables with T7 being the probability of oot 1. As in Eqn. (5). 2ot activates the o ''th neuron in layer l. Computationally, we employ the stick- breckib construction Bernoulli process
17 beta process and its conjugate
Sol Ber(
T
II Uy in ~ Beta(a, B)
where are sequentially drawn from beta distribution. The hyperparameters and B can be set to balance the network depth and width. Specifically, Figure demonstrates that if > a > the network structure prior favors shallower but wider network structures with the first few layer-wise activation probabilities being high. a 3 > 1, activation probabilities tend to be low over larger number of active layers, and the prior prefers a deeper but narrower network, in (c.
We thus define the prior over the encoding network structural variable Z
0
Pa.9 (Z. v) = Pa.9 Pa.s (v)p( II Beta (v,,, 3)
II 0=1
Ber(2
To enable asymmetric encoding/decoding network structures, we independently apply the prior to both networks. Further analysis symmetric constraints can be found in the Appendix.

---

4.3 Joint Inference on VAE Network Structures and Latent Variables
We first expand the overall marginal likelihood over the VAE structures Z
Po. (x) K L[q (at))== {be) Po.s v/x)] + {aa,).(0))
(8)
where the first RHS term denotes a Bellbackk-eiblee (KL) divergence the approximate variational distribution from the true posterior of the VAE network structural variables. We specify the variational distribution
T (Z, UZ (bO = II (e aa b.) II
ConBer(Eaa (T)
==1
where 74 = IT' Vj and {at, b are the variational parameters. T denotes a truncation level for the maximum number of hidden layers [7]. We also relax the constraint the discrete variables by reparameterizing the Bernoulli distribution into concrete Bernoulli distribution ConBer( (T)) [44 45). This allows us to efficiently backpropagate the parameter gradients the estimator while generating network structure samples.
The second RHS term in Eqn.
(8) denotes the ELBO to the overall marginal likelihood:
((ash(b))))))
a Z, ( (at) (ac) (00g + log Pa, (Z,
log (Z, Maa ) (br (
(10)
The term log Pa [x] (Z) in Eqn. (10) is the marginal likelihood over the latent variable h. which is an
extension of Eqn. T) in terms of conditioning on the structure variable Z. Thus, the ELBO the
marginal likelihood log po is:
log po(x|Z) 2 / 9s (b|x,, Z) log
Po (x, h/Z) dh
4s (h|x, Z)
= -Ey. [log Pa (x (h, Z)] (L|xx(h|x, (h)p))(h(2)
= Lx,o(x (2)
(11)
Lemma Let Q(h(x) = S (hlx, (), vdZddv the variational distribution of the latent variable h marginalizing over Z, then
Po h|Z)
Px h/Z)
E BE (hlx log
< (Q(()x) log Qhl|x)
(12)
4a h h/x. Z)
Lemma] Tindicates that the overall ELBO we derived on the left-hand side in Eqn. 122 bounds the lower bound to the marginal likelihood over h. The proof is in the appendix. In partieular, qZ. v) is essentially a non-explicit mixing distribution [46). allows the variational distribution Qh|xx) to take complex form, and results in informative latent representation.
We adopt Monte Carlo estimation of the expectations over both Eqn. (TO) and in Eqn.
to estimate the overall ELBO. In particular, extend the MIWAE estimator in Eqn. (3) and introduce three sample sizes to tune: the number of samples S used for Monte Carlo estimation the expectation over the VAE network structure variable Z, the number of samples M, used for Monte Carlo estimation of the gradient of the latent variable ELBO conditioned on the structure samples in Eqn. (T). and the number of importance samples K, used for estimation of the expectation over the latent variables h. We thus express our gradient estimate in the general form as
1
Mr
K.
(x(x hm,k, Z. po (h (Z,)
A..MKK S
Vo. log K.
(13)
M
4o (hm, |x. Z,)
m=1
where hm,k ~ qo (h|x, Z). When S estimator becomes equivalent the MIWAE objective
in Eqn. (3) Since our estimator generates latent variable samples conditioned on network structure samples, increasing will not impact the SNR of the gradient estimator in Eqn. B
Theorem Let Lg be the lower bound with S structure samples of Z ~ q (Z, v). then:
Σ po (x, h (z,)
Ls < Ls+1 $ logPa, (x),, Ls = Eg E (h/x.2) log
(14)
5
go (b|x, Z.)
5

---

Epoch: 1
Epoch: 10
No of activetted layors
Encoding Network
Decoding
Encoding Network
Decoding Network
Act AcToval mIn Iwturmntt
W
222
h
Actusted neumchh
Epochr
Epoch:
Epoch:
Encoding Network
Decoding Network
Encoding Network
Decading Network
140
T
T
h
100
7
9
100
130 140
Epochs
Figure 2: Left: Evolution of the encoding/decoding network structures visualized through layer-wise activation probabilities иi (top) and neuron activations Z (bottom) with reconstructed sample. Right: the top shows the median change of the number of active layers over training epochs and the percentage of activated neurons in the truncation. The bottom shows the convergence of the proposed estimator in terms of negative tog-likelihode (-LL) for different structure sample size S.
Table I: Test performance in negative log-likelihood (-LL) mean +1 standard deviation (lower the better) over 4 with random initialization. The overall best result on each dataset is bolded.
Dataset
(M.K)
S=1
S=2
S=4
S=8
MNIST
(8,8)
82.25
82.50.0.00
82.60...10
82.30.10.10
(4,16)
82.4710.4
82.33
82.52
83.02.10.20
Omniglot
(8,8)
107.1040.0
106.45+0.10
106.55+0.30
106.34:..01
(4,16)
108, 12:6.16
107.15
107.35
108.30+0.50
Caltech101
(8,8)
116.83+1.57
114.94:1.45
114.00
113.5410.40
(4,16)
116.30
114.55
113.02.10.4
112.53:0.42
Proof of this theorem in the appendix. The theorem shows the convergence of our estimator. Specifically, increasing S leads to a tighter lower bound for the overall marginal likelihood. Training an encoding/decoding network with depth L and width M, the time complexity is Te O(N with N training examples and B epochs. Our method is linearly scalable as ST.. With proper thresholding, the number of active layers I, is relatively small in each sample.
Experiments
We analyze the behavior of our inference framework across various tasks. We study how Ada VAE facilitates the evolution of encoding/decoding network structures for inferring the most plausible depth from the given data, while generating expressive latent representations. Next, we explore the impact of the structure sample size S on the convergence of the proposed estimator in Eqn. 13). Then we show that AdaVAE effectively mitigates overfitting in both shallow and deep network settings, leading to state-of-ther tf-the performance on benchmark datasets. Finally, we demonstrate the framework's s compatibility with different types of backbone networks and VAE variants. ?
5.1
Adaptive VAE Network Structures
Ada VAE enables us perform joint inference both encoding/decoding network structures and latent variables. investigate how network structures evolve during training epochs, we set truncation level T 25 MLP backbone nets with tanh non non-linearities. We analyze ada'sEEs behavior on x 28 binarized MNIST images [47). employing structure sample sizes S = (1, 2, 4, 8}. We run the experiments 3 times and averaging the outcomes. Figure Left 2 shows the evolution of the encoding/decoding network structures for one trial with $ = 8. Ada VAE initializes multiple hidden
IImplementation details are in the Appendix. Codes provided.

---

MNIST
Omniglot
Caltech100
VIWAE
MINAEIDO
180
HIMAE
HIMAE
350
MIAAE
180
BP-DGM
ESP DGN
100
DVAE
CRIAE
300
CRAME
IPURE
140
250
Curs
120
100
150
15
100
20 100
Layers(T/L)
Layers(T/L)
Layers(T/L)
Figure 3: The performance of VAE regularization methods changes with network depths. Our proposed method effectively prevents overfitting for both small and large truncations T, consistently achieving the best performance. In contrast, for 1 < L S 5, the performance of other regularization methods initially improves but then starts to decline, suggesting they suffer from overfitting even for shallow structures.
Table The best performance our method and the VAE regularization methods in Figure also demonstrate the compatibility of our framework.
We
MNIST
Omniglot
Caltech 101
-LL +
MI
-LL
MI
-LL +
MI 1
MIWAE [18]
86.11.000
9.13.0001
110.61
8.98
116.1910.08
7.64
MIWAE+DO
90.99+0.01
9.13=0.01
110.89+0.01
8.98:+0.00
116.00++.440
7.63
DVAE [1]
87.67+0.17
9.07
112.04+0.13
8.96=+0.00
113.71+0.43
7.53
CR-VAE [22
87.67+0005
9.03
109.94
8.97+0000
117.46+t0.0
7.37
IBP-DGM [14]
92.2440.75
1...01
135.23++054
BB-VAE [15
91.55+0.69
124.47+0.58
135.18+0.3
Ours
82.30=1010
9.13:1001
106.344..01
8.98
113.54::0.00
7.67 =0.02
Ours+DVAE
83.3040.20
9.07
107.80+..40
8.96.1000
111.92:1..7
7.63
Ous+CC-VVE
85.20
9.03.0.01
105.60
8.9710000
108.93:1.40
1..02
layers with sparsely activated neurons, gradually converging to fewer fully activated layers. Figure right top presents the medians the number of active hidden layers in encoding/decoding networks, as well as the percentage of activated neurons in the truncation changing over epochs. The encoding network structure stabilizes two active layers, while the decoding network settles at six active layers. The decoding network tends to have more active layers compared to the encoding network on Omniglot as well (see Appendix). Figure bottom shows the generative performance assessed by negative log-likelihood (- LL)) converge Taster with increased structure sample size S, which is consistent with Theorem
5.2 Effect of Structure Sample Sizes
We assess estimator using three benchmark datasets: MNIST [47]. Omniglot [48, and Caltech
Silhouettes [49] each minibatch, we set a budget of S x M, K, 64 total latent variable samples
for each datapoint. We examine four settings of the VAE structure sample size {1, 2, 4, 8). along with latent variable sample sizes (M, K) (88 8), (4, 16) as in (18). The truncation level is T 25 with maximum width O 200. The distribution
the output from the decoding networks is factorized Bernoulli. Table indicates larger values of S generally yield better performance. For MNIST, there no statistically significant difference between S and S Among the two structure sample sizes, the best importance sample configuration is (M, K) (8,8).
5.3 On Preventing Overfitting
To assess the performance of our method across varying truncation level T. we compare with
existing VAE regularization methods: Denoising VAE (DVAE) [10), CR-V VAE [2] and BB-VAE [5]. MIWAE with dropout (MIWAF+DO), along with vanilla MIWAE [8] as Figure 3 All methods share the same maximum width of O 200 and latent variable dimensionality of 50. For

---

BB-VAE
Ours
DVAE
DVAE
CR-VAE
Ours+CR-VAL
Figure 4: Visualization of the latent representation via t-SNE embeddings for the MNIST dataset. The embeddings are colored based on class labels. DVAE and CR- VAE combined with our framework result in better representations.
MNIST
Omniglot
Caltech101
15
Aclive Layers (Encader)
Activaten neuron (%)
20
Active Layers (Encoder)
Activated neuroni 9%)
(Encoder)
Activated ('euronn
Active Layers
Active Layers (Decodar))
20
Active Layers
Active Layers (Dacodar)
Active Layers
Active Layers (Decoder)
25
9-
220
15
10
Activated neunon(%)
Activated nouron(S))
Activated neurnn())
250
150 Width
100
150
200
250
Width
Width
Figure 5: Influence the maximum number of neurons per layer O on our method. When O is small (c.g.. O < 100), we tend to have shallower encoding/decoding networks. As O becomes reasonably large (e.g., O 2 100), it tends not to have significant influence on the depth. Meanwhile, the percentages of activated neurons in the truncation become stable.
IBP-DGM and BB-VAE VAE, set the maximum dimensionality of the latent variables to 100, following the setup of [15]. In Figure 3 we set depth L for the regularization methods and truncation level T for our method over the range
T {1. 3, 5, 10, 15, 20, 25), allowing us to compare their effectiveness in both shallow (i.c., L = T' S. 5) and deep (i.c., L = =T 2 5) VAE network structures.
Figure 3 shows that when L T However, even for the shallow cases as = T = (4, 5). the performance of these methods is affected by overfitting, as evidenced by the "U"'shaped performance curves, indicating the classical variance-bias trade-off. Despite our method exhibiting slight underfitting issues for L = T < 2, it still outperforms other methods. For deep VAE structures with T/L {10, 15, 20, 25). Figure shows that our method's performance is minimally impacted by large truncations. This robustness indicates our method's ability mitigate overfitting for large VAE structure settings. In contrast, as the depth L increases, the steadily increasing negative log-lik ligeliiel (-LL) for other regularization methods suggests their inability prevent overfitting in deep VAE models. Overall, AdaVAE consistently delivers superior performance for both shallow and deep VAE structures. An analysis of the computational time required by our framework and its comparison to the baselines is presented in the Appendix.
the VAE regularization methods
mitigate overfitting.
Table highlights our superiority over other methods in terms of density estimation and mutual information (MI) [50]. Additionally, when incorporating our framework with DVAE through the addition of Gaussian noise the input with CR-VAE by integrating the consistency regularization term into our estimator, we observe enhanced performance on Omniglot and Caltech101 datasets. Furthermore, in Figure we visualize the latent representations of the VAE methods. DVAE and CR- VAE when combined with our framework result in well-separated clusters, indicating that application of our framework allows VAEs to learn a meaningful latent representation. quantitative evaluation of the latent representations is presented in the Appendix by analyzing the performance of VAE methods on a downstream classification task.
5.4 Effects of the Maximum Width O
We further investigate the influence of the width O (i.e., the maximum number of neurons per layer). shows the evolution of the medians of the number of active hidden layers (i.e., the hidden layers with activated neurons) as O increases. When O < 100, we tend to have shallower
8

---

MNIST
Omniglot
Cora
Citeseer
230
CVAE
CVAE
250
B-eVAE
190
BeVAE
0.8
150
Ours+ CVVVE
160
Ours+G-cVAE
GAE
GAE
VSAE
VEAE
100
Ours VGAE
VIAEE
10
20
10
15
20
25
03
10
0.5
Layers(T/L)
Layers(T/L)
Layers(T/L)
Layers/TALI
(a)
(b)
(c)
(d)
Figure 6: Our framework's performance on different VAE backbone networks. (a) and (b) show VAE and 3-VAE [5]] with convolutional layers MNIST and Omniglot datasets. (c) and (d) show with graph convolutional layers on Cora and Citeseer datasets.
VGAFs 52
encoding/decoding networks with lower percentage of activated neurons in the truncation. However, when O is reasonably large as O 2 100, it has no influence on the network depth. In particular, the percentage of activated neurons in the whole truncation also remains relatively stable. This suggests that our method can automatically balance network depth and width to maintain the best performance.
5.5 Application to VAE Backhone Networks
We demonstrate ada VAE's efficacy by applying it to VAEs with different types of encoding/decoding backbone networks. To infer the number of convolutional layers in convolutional VAE (cVAE) using the beta process, we readily mask the convolution channels in layer / with 2.,. Figure (a) and (b) show that by adapting network structures the data we improve the overall performance of cVE and BAE [51]. The adaptive backbone networks effectively prevent overfitting for deep structure settings (i.e., T/L Z 15).
Variational graph autoencoder (VGAE) [52] encodes graph nodes to latent embeddings with graph convolutional (GC) layers and re-creates the input graph from the embeddings by predicting the existence an edge between the nodes. We combine our framework with VGAE by elementwisely multiplying the GC channels (i.e., the feature vectors) of layer with z.І in both its encoding/decoding networks. We compare the hybrid method's performance with graph autoencoders (GAFs) [22 and
vanilla VGAEs two benchmark graph-structured datasets: Cora (33 and Citeseer [54] The AUC
scores link prediction over varying numbers of GC layer settings are shown in Figure 6 (c) and (d).
Our framework enables the VGAE maintain its best performance for all the network depth settings
by automatically adapting its structures to the data, whereas the performance of GAE and vanilla VGAE drops with the increase of the layer numbers (e.g., T {6,8, 103).
5.6 Application to VAE Variants
We assess the performance of our inference framework by leveraging it adapt the net- work structures of VAE vari- ants to data. Specifically, for 3-VAE, we apply layer-wise bi- nary masks z.; to the convolu- tional channels and infer the layer numbers using the beta process. In the case of Ladder- VAE (LVAE) 4. we adjust its depth by applying layer- wise binary masks its determin- istic layers in the bottom-up dependency structures and add skip connections between the stochastic layers. For SKi 34 we model its depth by employing layer-wise binary masks and skip connections both encoding/decoding networks. The expressive network structures of NVAE [6] consists of multiple blocks of convolutional layers. We apply our
Table 3: Performance comparison of VAE variants with and without
our inference framework in terms of - L.L MI, and KI divergence.
Methods
-LL
MI
KL
8-CVAE (B = 2) [11]
106.50
8.46:1001
18.01
Ours + 3-E (8 2)
102.45++.07
8.47
20.65+
LVAE [4]
136.50+
8.43+000
20.05+00.1
Ours+LVAE
121.48-+0.77
8.50+10.01
21.00:10.47
SkipVAE [34]
112.68+0.00
8.50+4..1
22.82+0079
Ours+Skip VAE
108.00+0.04
8.50+9.01
28.26+0010
NVAE 6
98.83
8,51
34.9740.13
Ours+NVAE
99.10+0.0
8.51
37.85

---

framework to infer the number of blocks for light NVAE version without top-down dependency between the blocks. The detailed settings of these VAE variants and additional results can be found in Appendix. The results in Table on FashionMNIST 55 demonstrate that by inferring the encod- ing/decoding network structures we significantly improve the density estimation performance of the VAE variants. Our framework also boosts their ability to mitigate posterior collapse indicated by MI and KL divergence.
Conclusion
We present a Bayesian inference framework and a scalable estimator that automatically adapts VAE network structures to data. The experiments demonstrate its effectiveness in preventing overfitting for both shallow and deep structure settings. Moreover, Vda exhibits promising applications across various types of VAE backbone networks and VAE variants, including those with hierarchical structures such as LVAE. Notably, Ada enhances the generative performance of these models without requiring pre-determined network structures prior to training. Our future work entails relaxing the constraint of truncation levels by incorporating the Russian roulette method (56) and scaling up the inference for large images.
7 Acknowledgements
This material is based upon work supported by the National Science Foundation under NSF Award No. No.2045804 and Award No. 185099.. We are thankful to Kishan KC for helpful discussion. We acknowledge Research Computing at the Rochester Institute of Technology [57] for providing computational resources.
References
[1] Diederik P. Kingma and Max Welling. Auto-encoding variational bayes. In Proceedings of the International Conference on Learning Representations (ICLR), 2014.
[2] Danilo J. Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and approximate inference in deep generative models. In International Conference Machine Learning (ICML), pages 1278-1286. PMLR, 2014.
[3] Jyoti Aneja, Alex Schwing, Jan Kautz, and Arash Vahdat. contrastive learning approach for training variational autoencoder priors. In Advances in Neural Information Processing Systems, volume 34, pages 480-493, 2021.
[4] Casper Kaae S, Tapani Raiko, Lars Maalø e, Søren Kaae Sø nderby, and Ole Winther. Ladder variational autoencoders. In Advances in Neural Information Processing Systems (NeurIPS), volume 29, 2016.
[5] Lars Maalee, 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 (NeurlPS). volume 32, 2019.
[6] Arash Vahdat and Jan Kautz. NVAE: deep hierarchical variational autoencoder. Advances Neural Information Processing Systems (NeurIPS), volume 33, pages 19667-19677 2020.
[7] Corinna Cortes, Xavier Gonzalvo, Vitaly Kuznetsov, Mchryar Mohri, and Scott Yang. AdaNet: Adaptive structural learning of artificial neural networks. In Proceedings the 34th International Conference on Machine Learning, volume 70 of Proceedings of Machine Learning Research, pages 874-883. PMLR, 2017.
[8] Konstantinos Panousis, Sotirios Chatzis, and Sergios Theodoridis. Nonparametric Bayesian deep networks with local competition. In International Conference on Machine Learning, pages 4980 4988. PMLR. 2019.
[9] Kishan KC, Rui Li, and Mohammad Mahdi Gilany. Joint inference for neural network depth and dropout regularization. In Advances in Neural Information Processing Systems (NeurIPS), volume 34, 2021.
[10] Rui Shu, Hung H Bui, Shengjia Zhao, Mykel J Kochenderfer, and Stefano Ermon. Amortized inference regularization. In Advances in Neural Information Processing Systems (NewriPS), volume 31, 2018.
10

---

[11] Pascal Vincent. Hugo Larochelle, Yoshua Bengio, and Fierre-Antoinn Manzagol. Extracting and composing robust features with denoising autoencoders. In Proceedings of the 25th International Conference on Machine Learning (ICML), pages 1096-1103, 2008.
[12] Samarth Sinha and Adji Bousso Dieng. Consistency regularization for variational auto-encoders. In Advances in Neural Information Processing Systems (NeurlPSS) volume 34, 2021.
[13] Salah Rifai, Pascal Vincent, Xavier Muller, Xavier Glorot, and Yoshua Bengio. Contractive encoders::: Explicit invariance during feature extraction. In International Conference Machine Learning (ICML), 2011.
[14] Sotirios P Chatzis. Indian buffet process deep generative models for semi-supervised classification. In 2018 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pages 2456-2460. IEEE 2018.
[15] Rachit Singh, Jeffrey Ling, Finale Doshi Velez. Structured variational autoencoders for the beta- bernoulli process. In NIPS 2017 Workshop on Advances in Approximate Bayesian Inference, 2017.
[16] Romain Thibaux and Michael Jordan. Hicrarchical beta processes and the Indian buffet process. In Artificial Intelligence and Statistics (AISTATS), pages 564-571, 571. PMLR, 2007.
[17] John W Paislcy, Aimee Zaas, Christopher W Woods, Geoffrey S Ginsburg, and Lawrence Carin, bick-breaking construction of the beta process. In International Conference on Machine Learning (ICML), 2010.
[18] Tom Rainforth, Adam Kosiorek, Tuan Anh Le, Chris Maddison, Maximilian Igl. Frank Wood, and Yee Whye Teh. Tighter variational bounds are not necessarily better. In International Conference on Machine Learning (ICML), pages 4277-4285. PMLR, 2018.
[19] Irina Higgins, David Amos, David Pfau, Sebastien Racaniere, Loic Matthey, Danilo Rezende, and Alexan- der Lerchner. Towards definition of disentangled representations. arXiv preprint arXiv:1812.02230. 2018.
[20] Fabian Falck, Haoting Zhang, Matthew Willetts, George Nicholson, Christopher Yau, and Chris C Holmes. Multi-facet clustering variational autoencoders. In Advances in Neural Information Processing Systems (NeurlPS), volume 34, pages 8676-8690, 2021.
[21] Ricky Chen, Xuechen Li. Roger B Grosse, and David K Duvenaud. Isolating sources of disentanglement in variational autoencoders. In Advances in Neural Information Processing Systems (NeurIPS), volume 31, 2018.
[22] Abhishek Kumar and Ben Poole. On implicit regularization in 9-VAEs. In International Conference on Machine Learning (ICML), pages 5480-5490. PMLR, 2020.
[23] Dazhong Shen, Chuan Qin, Chao Wang, Hengshu Zhu, Enhong Chen, and Hui Xiong. Regularizing variational autoencoder with diversity and uncertainty awareness. arXiv preprint arXiv:2110.12330 2021.
[24] Vaibhav Saxena, Jimmy Ba, and Danijar Hafner. Clockwork variational autoencoders. In Advances in Neural Information Processing Systems (NeurIPS), volume 34, pages 29246-29257, 2021.
[25] Chris Cremer, Xuechen Li, and David Duvenaud. Inference suboptimality in variational autoencoders. In Proceedings of the 35th International Conference Machine Learning, ICML 2018, StocKholmsmässon. Stockholm, Sweden, July 10-15, 2018, volume 80 Proceedings of Machine Learning Research (PMLR). pages 1086-1094, PMLR, 2018
[26] Joe Marino, Yisong Yue, and Stephan Mandt. Iterative amortized inference. Proceedings of 35th International Conference on Machine Learning (ICML), volume 80 of Proceedings of Machine Learning Research, pages 3403-3412, PMLR, 10-15 Jul 2018.
[27] Rajesh Ranganath, Dustin Tran, and David Blei. Hierarchical variational models. In Proceedings of The 33rd International Conference on Machine Learning (ICML), volume 48 of Proceedings of Machine Learning Research (PMLR), pages 333-333 20-22 Jun 2016.
[28] Yingzhen Li and Richard Turner. Rényi divergence variational inference. In Advances in Neural Information Processing Systems (NeurIPS), volume 29, 2016.
[29] Arash Vahdat, Evgeny Andriyash, and William Macready. Dvaeft: Discrete variational autoencoders with relaxed boltzmann priors. In Advances in Neural Information Processing Systems (NeurIPS), volume 31, 2018.

---

[30] Ali Razavi, Aäron van den Oord, Ben Poole. and Oriol Vinyals. Preventing posterior collapse with delta-vaes. In 7th International Conference on Learning Reprexentations, (ICLR) 2019, New Orleans, LA, USA, May 6-9, 2019, 2019.
[31] James Lucas, George Tucker, Roger B Grosse, and Mohammad Norouzi. Don't blame the elbo! linear vae perspective on posterior collapse. Advances Neural Information Processing Systems (NeurlPS), volume 32, 2019.
[32] 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, volume 32, 2019.
[33] Jakob D. Havtorn, Jes Frellsen, Søren Hauberg, and Lars Maaløe. Hierarchical VAES know what they don't know. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning (ICML), volume 1339 of Proceedings of Machine Learning Research (PMLR), pages 4117-4128, 18- 18-24 Jul 2021.
[34] Adji B. Dieng, Yoon Kim, Alexander M. Rush, and David M. Blei. Avoiding latent variable collapse with generative skip models. In Proceedings of the Tieent-Seeoon International Conference on Artificial Intelligence and Statistics (AISTATS), Proceedings of Machine Learning Research, pages 2397-2421 2905,2405, 2019.
[35] Durk P Kingma, Tim Salimans, Rafal Jozefowicz, Xi Chen, Ilya Sutskever, and Max Welling. Improved variational inference with inverse autoregressive flow. In D. Lec, M. Sugiyama, U. Luxburg, I. Guyon, and R. Garnett, editors, Advances in Neural Information Processing Systems (NeurlPS), volume 29, 2016.
[36] Vaden Masrani, Tuan Anh Le, and Frank Wood. The thermodynamic variational objective. In Advances Neural Information Processing Systems (NeurIPS), volume 32, 2019.
[37] Ali Razavi, Aaron van den Oord, and Oriol Vinyals. Generating diverse high-fidelity images with vq-vae-2. In Advances in Neural Information Processing Systems, volume 32, 2019.
[38] Arash Vahdat, Karsten Kreis, and Jan Kautz, Score-based generative modeling in latent space, In Advances in Neural Information Processing Systems, volume 34, pages 1 11287-11300 2021.
[39] Chris Cremer, Xuechen Li, and David Duvenaud. Inference suboptimality in variational autoencoders. In International Conference on Machine Learning (ICML), pages (078-1086. PMLR, 2018.
[40] Philip Bachman, Ouais Alsharif, and Doina Precup. Learning with pseuda-ensembles. In Advances in Neural Information Processing Systems (NeurIPS), volume 27, 2014.
[41] Mehdi Sajjadi, Mehran Javanmardi, and Tolga Tasdizen. Regularization with stochastic transformations and perturbations for deep semi-supervised learning. In Advances in Neural Information Processing Systems (NeurIPS), volume 29, 2016.
[42] Yuri Burda, Roger B Grosse, and Ruslan Salakhutdinov. Importance weighted autoencoders. In Proceedings of the International Conference on Learning Representations (ICI.), 2016.
[43] Tamara Broderick, Michael I Jordan, and Jim Pitman. Beta processes, stick-breaking and power laws. Bayesion Analysis, 7(2):439-476, 2012.
[44] Chris J Maddison, Andriy Maih, and Yee Whye Teh. The concrete distribution: A continuous relaxation of discrete random variables. arXiv preprint
2016.
[45] Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterizztion with gumbel-softmax. arXiv preprint 1:X:::661..01144 1.00144144 144, 2016.
[46] Mingzhang Yin and Mingyuan Zhou Semi-in Semi-implicit variational inference. In International Conference on Machine Learning (ICML), pages 5660-5669. PMLR, 2018
[47] Yann LeCun, Leon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient- based learning applied to document recognition. In Proceedings of the IEEE, volume 86 (11), page 2278 2324, 2014.
[48] Brenden M. Lake, Russ R. Salakhutdinov, and Josh Tenenbaum. One-shot learning by inverting a compositional causal process. In Advances in Neural Information Processing Systems (NeurIPS), 2013.
[49] Benjamin Marlin, Kevin Swersky, Bo Chen, and Nando Freitas. Inductive principles for restricted boltzmann machine learning. In Yee Whye Teh and Mike Titterington, editors, Proceedings of the Thirteenth International Conference on Artificial Intelligence and Statistics (AISTATS), volume 9, pages 509-5 516. PMLR, 13-15 May 2010.
12

---

[50] Matthew D Hoffmman and Matthew J Johnson. ELBO surgery: Yet another way carve up the variational evidence lower bound. In Workshop in Advances in Approximate Bayesian Inference, NeurIPS, volume 1, 2016.
[51] Irina Higgins, Loic Matthey, Arka Pal. Christopher P. Burgess, Xavier Glorot, Matthew M. Botvinick. Shakir Mohamed, and Alexander Lerehner. beta- VAE: Learning basic visual concepts with constrained variational framework. In Proceedings of the International Conference on Learning Representations (ICLR), 2017.
[52] Thomas N Kipf and Max Welling. Variational graph auto-encoders. arXiv preprint arXiv:1611.07308, 2016.
[53] Andrew Kachites McCallum, Kamal Nigam, Jason Rennie, and Kristie Seymore. Automating the construc- tion of internet portals with machine Icarning. Information Retrieval, 3:127-163, 2000.
[54] C Lee Giles, Kurt D Bollacker, and Steve Lawrence. CiteSeer: An automatic citation indexing system. In Proceedings of the third ACM Conference on Digital Libraries, pages 89 98, 1998.
[55] Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-MNIST: novel image dataset for benchmarking machine learning algorithms. arXiv preprint UrXiv:1708.007777 2017.
[56] Kai Xu, Akash Srivastava, and Charles Sutton. Variational russian roulette for deep bayesian nonparamet- rics. In International Conference on Machine Learning (ICML), pages 6963-6972. PMLR, 2019.
[57] Research Computing at Rachester Institute of Technology. http:///ollog////0444//000000000
A
Appendix
This section is in the supplemental material.
13


Reviews:
[{"review_id": "f9k6BKirVw", "reviewer": "Reviewer_9dFM", "summary": "The paper proposes a novel VAE structural adaptation strategy called AdaVAE based on Bayesian model selection to enhance model performance. It introduces a scalable estimator that facilitates joint inference on both encoding/decoding network structures and latent variables. The paper conducts a comprehensive analysis of AdaVAE's regularization capabilities and demonstrates its ability to effectively mitigate overfitting in both shallow and deep VAE models and achieve state-of-the-art performance. The versatility of AdaVAE is showcased by demonstrating its compatibility with different types of VAE backbone networks. It can also be readily applied to various VAE variants, thereby enhancing their performance.\n\nThe main contributions are:\n- Proposes AdaVAE, a novel VAE structural adaptation strategy based on Bayesian model selection to enhance model performance.\n- Introduces a scalable estimator that facilitates joint inference on both encoding/decoding network structures and latent variables.\n- Conducts a comprehensive analysis of AdaVAE's regularization capabilities and demonstrates its ability to effectively mitigate overfitting in both shallow and deep VAE models and achieve state-of-the-art performance.\n- Showcases the versatility of AdaVAE by demonstrating its compatibility with different types of VAE backbone networks.\n- Can be readily applied to various VAE variants, thereby enhancing their performance.", "questions": "Please refer to the above sections for detailed discussion.\n", "limitations": "Yes, there is a section at the end of the paper discussing limitations and future research opportunities.\n", "rating": 7, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The authors introduce a novel Variational Autoencoder (VAE) structural adaptation strategy, dubbed AdaVAE, which employs Bayesian model selection as a mechanism to enhance model performance. This innovative approach pushes the boundaries of current practices in the field and sets a precedent for future explorations.\n\n2. The study further contributes by proposing a scalable estimator. This facilitates joint inference on not only the structures of the encoding and decoding networks but also the latent variables. This dual focus enhances the model's applicability and comprehensiveness, potentially opening new avenues in inferential methodologies.\n\n3.A thorough analysis is presented on AdaVAE's regularization capabilities, showcasing its efficacy in mitigating overfitting across both shallow and deep VAE models. This is a critical point, as it demonstrates the proposed method's capability to achieve state-of-the-art performance across varying levels of complexity.\n\n4. The versatility of AdaVAE is effectively demonstrated, as the authors show its compatibility with different types of VAE backbone networks. This level of adaptability reinforces the model's potential for broad application and usability within the field.\nThe authors execute a rigorous evaluation of the proposed method on three benchmark datasets: MNIST, Omniglot, and Caltech101 Silhouettes. This broad evaluation allows a comprehensive understanding of AdaVAE's performance and robustly positions it in relation to other state-of-the-art methods.\n\n5. The paper includes a qualitative evaluation of the latent representations, an element that strengthens the argument for the proposed method's compatibility with VAE regularization methods. This kind of analysis enhances the empirical rigor of the study and provides an additional perspective on the utility of AdaVAE.\n\n6. The manuscript offers a theoretical analysis of the proposed method and derives a tight lower bound with a high signal-to-noise ratio for parameter gradients. This theoretical grounding ensures the model is not only empirically valid but also theoretically sound.\n\n7.The manuscript is particularly well-executed in terms of its structure and style. The clear writing, logical organization, and the authors' ability to elucidate complex concepts makes it accessible and easy for readers to grasp the proposed method and its evaluation.\n\nIn essence, the paper is a substantial contribution to the field, demonstrating methodological innovation, theoretical robustness, and a thorough evaluation, all of which underscore the potential and value of AdaVAE within the realm of VAE structural adaptation.", "weaknesses": "This manuscript is laudable for its comprehensive coverage, including a lucid motivation, rigorous theoretical proofs, and a thorough comparative analysis with benchmark models across various datasets. It demonstrates a high level of academic rigor and makes a compelling case for the proposed model.\n\nHowever, I do have a minor suggestion to enhance the paper further. While the evaluation conducted on Caltech101, among other datasets, certainly contributes to the robustness of the results, I believe there could be value in extending this evaluation to include more real-life, challenging datasets. Such datasets, replete with their inherent complexities, could serve to more thoroughly test and validate the model. This, in turn, would potentially bring the proposed model closer in comparison to other state-of-the-art generative models. It would also ensure that the model's performance is tested in conditions that mirror real-world applications, reinforcing the practical relevance and applicability of the model.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We'd like to thank Reviewer 9dFM for your time and your constructive comments. They are valuable for our future research.\n\nWe agree with Reviewer 9dFM that extending the evaluation to include more real-life, challenging datasets will be valuable. It is part of our future work. In addition, we conducted some additional experiments to test our framework on CIFAR-10. We implemented a simplified NVAE by designing an encoding and decoding network with a series of NVAE blocks. We report the performance as below:\n\n| | | | \n|:-:|:-:|:-:|\n| **Methods**   | **Reconstruction loss** | **KL** |\n| NVAE      | **13313$\\pm$09**       |  150$\\pm$03 |\n| Ours+NVAE | 13600$\\pm$30    | **171$\\pm$00** |\n\nWe also provided reconstructed samples from the two methods in the attached pdf. \n\nWe will add more additional results on real-life datasets in the camera-ready version of our paper."}, {"review_id": "zCM7Nk4pT0", "reviewer": "Reviewer_uhJ9", "summary": "This paper introduces a Bayesian structural adaptation framework that automatically adapts VAE network structures (both encoder and decoder) to the current data.  By modeling the number of hidden layers as a beta process and performing layer-wise dropout regularization with the conjugate Bernoulli process, the proposed model develops a joint ELBO that can optimize all parameters (network structures and latent variables) via the SGD algorithm. Empirical studies are conducted on three visual datasets and two graphical datasets. They also provide extensive ablation studies to show the robustness.", "questions": "(1) I suggest that the authors add a training algorithm, which can help the reader understand the algorithm more easily.\n\n(2) From Fig(6) (a) (b), we find that AdaVAE have similar results with cVAE. Can the authors provides deeply analysis?\n\n(3) Given that the proposed AdaVAE aims to find optimal structure for the dataset at hand, it is advisable to include a comparison of network parameters in Table 2.", "limitations": "None", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "(1) The paper addresses one of the common issues in VAE communities. The developed AdaVAE attempts to prevent the overfitting issues by jointly learning the network structures and latent variables under the Bayesian framework. This motivation sounds good, and the novelty of the method is generally ok.\n\n(2) The paper is written clearly and the empirical results show the improvement of the proposed model.", "weaknesses": "(1) One of the main concerns comes from the experiment section. The proposed model is only tested on several simple network structures (the number of layer usually less than 25). Additional results on more expressive neural network structures (such as BIVA and NVAE) will improve the quality of this paper.\n\n(2) One of the goal of the proposed model is to prevent the overfitting by adapting the network structures under the Bayesian framework. The authors compared the proposed model with existing VAE regularization methods, and the results show the improvements. Unfortunately, other baselines (such as network structure search methods [1]) are not included in the experiments.\n\n[1] Corinna Cortes et.al. AdaNet: Adaptive Structural Learning of Artificial Neural Networks.\n\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We'd like to thank Reviewer jhJ9 for your constructive comments. They are valuable for our future research. \n\n**Application to expressive VAE variants**: BIVA is an extension of LVAE on which we conducted rigorous tests. We have additional results on BIVA/LVAE on the MNIST dataset as follows:\n| | | |\n|:-:|:-:|:-:|\n| **Methods** |    **-LL** |  **KL**  |  \n| BIVA/LVAE      | 116.07$\\pm$2.21     | **23.05$\\pm$0.15**   |  \n| Ours+BIVA/LVAE | **86.07$\\pm$0.07**      | 19.78$\\pm$0.01  |  \n\nWe show our framework's performance on NVAE on the cifar10 dataset as below by implementing a simplified NVAE  by designing an encoding and decoding network with a series of NVAE blocks.\n\n| | | |\n|:-:|:-:|:-:|\n| **Methods**   | **Reconstruction loss** | **KL** | \n| NVAE      | **13313$\\pm$09**         |   150$\\pm$03 |  \n| Ours+NVAE | 13600$\\pm$30        |   **171$\\pm$00** | \n\nWe also provided reconstructed samples from the two methods in the attached pdf. These additional results suggest our framework's flexibility on more expressive VAE structures. We will add them to the camera-ready version of our paper.\n\n**Difference with AdaNet**: There are several major differences setting our method and AdaNet apart: 1). AdaNet employs an iterative procedure to search for optimal network structures. Namely, at each iteration, it selects a set of candidate subnetworks and trains/re-trains the expanded alternative models. This is a time-consuming procedure. In contrast, our framework models the growth of VAE network structures as a stochastic process. Consequently, we can jointly infer the network structures and latent variable in a single training pass. 2). AdaNet incrementally constructs feedforward neural network structures in supervised learning settings. It remains uncertain whether it can be effectively extended to VAEs, considering its iterative searching process. Additionally, scalability could be a concern, as AdaNet solely examines network configurations with a maximum of three layers. 3). AdaNet doesn't demonstrate if their method can be applied to different backbone networks, such as CNN and GCN. We will clarify these differences in the camera-ready version of our paper.\n\n\n**Training algorithm**: The pseudocode of our training algorithm was included in the Supplementary Material/Appendix Section 4. We conducted additional experiments to analyze its convergence in training in Appendix Section 6. We have also included our codes in the Supplementary Material. \n\n**Similar results between AdaVAE and cVAE in Figure 6**: The similar results between cVAE and AdaVAE for a smaller number of layers in Figure 6 (a) and (b) are due to the convolutional layers' preventive effect to overfitting. In comparison, Figure 3 shows that fully connected networks are more sensitive to overfitting. However, for deeper network structures, Figure 6 shows that cVAE still suffers from overfitting. AdaVAE, in contrast, can successfully mitigate it. We will elaborate on the results in the camera-ready version.\n\n**Comparison of network parameters**: First, we reported detailed parameter settings of the baseline methods and ours in Supplementary Material/Appendix Section 5.3. Second, since the baselines activate the whole network structures, so the total number of parameters for encoding and decoding network are $2\\times O\\times O\n\\times L$, where $O$ denotes the maximum number of neurons per layer (i.e., width) and $L$ is the number of layers. For layers $L=25$ and width of $O=200$, they have $2M$ parameters. With the same size of truncation, AdaVAE only activates a part of it in general and fits the activated structures to data. Thus, the activated number of parameters (neuron activation percentage) for $T=L=25$ are as follows:\n\n| Methods |   MNIST      | Omniglot     | Caltech101      |\n|---------|--------------|--------------|-----------------|\n| Ours    | 0.32M(16%) | 0.36M(18%) | 0.21M (10.5%) |\n\nAs the table shows AdaVAE uses a smaller number of parameters to achieve state-of-the-art performance. We will include the results in the camera-ready version of our paper."}, {"review_id": "aAUKmKpKGU", "reviewer": "Reviewer_p4pj", "summary": "The paper proposes a method called AdaVAE that adapt the sizes (depth and width) of the inference and the generative networks in VAEs.\nIt extends the idea of beta-Bernoulli process VAE to use a beta process over the expected number of activated neurons over depth and Bernoulli processes per layer for the actual activations, followed by a dense Gaussian latent layer.\nThe paper develops an estimator based on MIWAE to learn AdaVAE. \nResults show that AdaVAE can prevents overfitting and lead to performance improvement than previous regularization methods for VAEs.", "questions": "In Figure 2, the last few layers of the encoder and the first few layers of the decoder seem to have no neuron activated.\nCan the author(s) clarify if it is the case and if so why the VAE still works in this case?\n\nDo we have any computation-wise comparison between the proposed methods and alternative/naive ones?\nIt is useful to address practical concerns of actually using the method.", "limitations": "The paper mentions its limitations on the fixed truncation level and some future work on removing it.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "### originality\n\nThe proposed AdaVAE for adapting VAE sizes is novel.\n\n### quality\n\nThe proposed method is technically sound.\n\n### clarity\n\nThe paper did a good presentation of the proposed methods with technical details, which is easy to follow.\n\n### significance\n\nThe proposed method is generally applicable to different VAE architecture and is potentially useful as a standard regularization technique.", "weaknesses": "### originality\n\nThe proposed estimator is a straightforward combination of previous works.\n\n### clarity\n\nThe paper is very related to [11,36,54] but only limited discussion is given.\nI encourage the author(s) to give some discussions between the AdaVAE and share intuitions of why those previous works (that only adapts the size of the latent layer) are not enough.\n\nSome of the figures with small legends are a bit hard to read.\n\nFigure 4 is purely qualitative. \nIt's hard to convince the readers that one method yields better representation than others.\nIt looks like the figure it trying to give an idea on how different clusters are separated.\nIf it is the case maybe some clustering metrics could be used to assess the quality quantitatively.\n\n### significance\n\nThe proposed method can be still computationally heavy to use, compared to a naive one without any structure adaptation.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We'd like to thank Reviewer p4pj for your constructive comment. They are valuable for our future research. \n\n**Relationship to [11,36, 54]**: Our experiment results in Figure 3 shows the methods as [11,36] only imposing an IBP prior to regularize the dimensionality of latent variables is not sufficient to mitigate overfitting caused by deep encoding/decoding networks. An intuition is large-scale encoding networks tend to learn redundant information or higher levels of feature abstraction from training data and embed it to the latent representation, which causes overfitting problems. Therefore, a more effective way to prevent it is by directly regularizing the network structures, rather than constraining the latent variables. [54] proposes a way to relax the truncation of stick breaking process. A part of our future research is to utilize the technique to improve our method.  We will add this analysis to the camera-ready version.\n\n**Small legends**: We will enlarge all the legends and axis labels in the camera-ready version.\n\n**Quantitative evaluations of Figure 4**: The quantitative evaluations of the VAE variants' performance in Figure 4 are reported in Table 2. The purpose of Figure 4 is only to give readers a qualitative intuition. Moreover, in order to further evaluate the quality of latent representations, we conducted additional experiments and reported downstream classification accuracy in Supplementary Material/Appendix \n Section 10.3, Table 8. This is a more effective way to assess the clustering.\n\n**Skipping non-activated layers**: We have skip connections as in Eqn. 5, so that we can propagate the last activated layer to the output layer by skipping the non-activated layers in between. We will detail this point in the camera-ready version.\n\n**Computation-wise comparison**: We have computation complexity analysis in Lines 179-182. We also compared the running times between our method and alternative ones and reported and analyzed the results in Supplementary Material/Appendix Section 8."}, {"review_id": "stDyTXEFOz", "reviewer": "Reviewer_kqwq", "summary": "The paper proposes a structure adaptation algorithm specifically tailored to the well-known Variational Autoencoders (VAEs). The main motivation lies in the fact that the structure of a generative model plays a significant part in the overall performance, something that is very under-explored in related VAE literature. To this end, the authors turn to the solid Bayesian inference framework and utilize a structural adaptation approach based on the Beta Process to infer the optimal network depth and the Bernoulli process to prune neurons in each hidden layer. The experimental evaluation focuses on the adaptation abilities of the proposed framework, how the structure sample size affects convergence and its over-fitting prevention properties. Finally the authors discuss how to integrate the proposed approach on VAE backbones and other VAE variants in general.", "questions": "Please see the Weaknesses section.", "limitations": "Please see Weaknesses, especially concerning complexity and stability.", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "This work focuses on a very interesting aspect of modern architectures, that is their structure. The paper is overall easy to follow. The notation is clear and almost everything is well defined.\nBy now there has been a lot of work in the community that aim to address this challenge through different views such as pruning approaches, component omission mechanisms, e.t.c.  Compared to most of existing approaches the above formulation allows for simultaneously inferring both the network depth and width of each layer in a principled way via Bayesian arguments.", "weaknesses": "1) As the authors note, the idea of using a Beta-Bernoulli pair in order to adapt the model capacity is not something new. There exist several works in both discriminative and generative models that aim to tackle this issue [1,2,3,4,5]. The authors however fail to appropriately introduce and discuss the differences with and advantages/drawbacks compared to these different methods leading to a single two sentence mention of one of the most similar approaches, i.e., BB-VAE. Instead the authors seem to focus more to the regularization aspect of other dissimilar methods. \n\n2) In this context, the work presented in [4] should be the main focus on comparison because it is essentially the same exact method, with minor adaptations for the VAE structure. Apart from the difference in estimation, i.e., Gumbel Softmax vs MIWAE, the differences are minor and down to some notation changes. The authors do cite the paper at some point  without any kind of discussion. \n\n3) Taking into consideration that the method is very similar to [4], a core difference is the different estimator. Did the authors investigate the impact of other estimation methods? How does it compare to Gubmel Softmax in terms of both complexity and convergence?\n\n4) The authors should briefly expand on the chosen metrics in the respective tables. The negative log likelihood is a well known measure, while the MI and KL divergence can have substantially different interpretations. The authors cite the work of [6], which dissects the ELBO into three different terms that probably correspond to the Tables 2,3. Further clarifications for the reader are essential. For example, why is the largest KL the better? The purpose of optimization is the minimization of the KL divergence and in [6] it is also noted that \"whenever it is large it indicates a very strong and potentially unwanted regularization effect from the prior\"\n\n5) The stability of the training process might be an issue that isn't addressed in the main text, especially in the stick-breaking construction. Did issues arise during training due to very small or very large values? and how did the authors address them?\n\n6) There is no discussion about the computational and memory complexity of the approach. It is apparent that the introduction of the additional latent parameters and the KL terms will significantly contribute to the overall footprint of the method, especially with multiple samples due to the MIWAE estimation. Some wall time measurements are necessary for both training and inference. \n\n7) There is not a clear definition of the prediction process of the framework. Do you draw multiple samples from the learned posteriors and sample the active ones? do the authors use a threshold in \"each sample\"? (line 182)\n\n8) The performance of the method on different VAR backbones is not clear. The authors mention the masking of convolutional channels in order to infer the number of convolutional layers, without any further expansion on the specific formulation and results. Does adaptation take place in specific features maps similar to [2] or the options are to drop the layer or not?\n\n9) Apart from the t-SNE visualization, some reconstruction comparisons could be useful. \n\nConsidering these points and the similarity of this work to [4], I believe that the novelty and overall contribution of this work is very limited for publication. \n\n[1] Sotirios P Chatzis. Indian buffet process deep generative models for semi-supervised classification. In 2018\nIEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)\n\n[2] Konstantinos Panousis, Sotirios Chatzis, and Sergios Theodoridis. Nonparametric bayesian deep networks with\nlocal competition. In International Conference on Machine Learning, pages 4980–4988. PMLR, 2019.\n\n[3] Rachit Singh, Jeffrey Ling, and Finale Doshi-Velez. Structured variational autoencoders for the beta-\nbernoulli process. In NIPS 2017 Workshop on Advances in Approximate Bayesian Inference, 2017.\n\n[4] Kishan KC, Rui Li, and Mohammad Mahdi Gilany. Joint inference for neural network depth and dropout\nregularization. In Advances in Neural Information Processing Systems (NeurIPS), volume 34, 2021.\n\n[5] Xu, W., Chen, R., Li, X., & Duvenaud, D. (2022). Infinitely Deep Bayesian Neural Networks with Stochastic Differential Equations. International Conference on Artificial Intelligence and Statistics.\n\n[6] Matthew D Hoffman and Matthew J Johnson. ELBO surgery: Yet another way to carve up the variational evidence lower bound. In Workshop in Advances in Approximate Bayesian Inference, NeurIPS, volume 1, 2016.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Novelty of our framework comparing with other beta-Bernoulli based methods**: Our work diverges fundamentally from [1,2,3] by treating the expansion of VAE network structures as a stochastic process within a comprehensive framework, in contrast to their exclusive focus on the regularization of VAE latent variable, overlooking network structures. Moreover, our experiments in Figure 3 distinctly demonstrate that adapting network architectures is a notably more effective strategy to boost VAEs' performance than the one-size-fits-all methods. In comparison to [4], which infers feedforward network depths in supervised learning scenarios, we have innovatively introduced a novel generative learning approach tailored to diverse VAE variants. The novelty contribution of our VAE structural adaptation framework is non-trivial, because 1) we developed a new estimator to jointly infer both net structures and latent variables. 2) We also enabled the asymmetric encoding/decoding network structures in the inference procedure. 3) Our extensive experiments show our framework and our estimator boosted different VAE backbone networks' and various VAE variants' performance and achieved state-of-the-art.\n\n**Gumbel-softmax vs MIWAE**: Gumbel-softmax is a reparameterization trick, relaxing categorical distributions to continuous ones. It is not a VAE estimator we can employ or compare with. The technique is not involved or relevant to our work.\n\n**Discussion of chosen metrics**: We have a detailed discussion on the evaluation metrics in Appendix Section 5.1.\n\n**Stability with training**: We do not have any stability issues with our training algorithm. We provided a theoretical analysis of the algorithm in Theorem 1 and its proof in Appendix Section 2. We analyzed its convergence in training in Appendix Section 6. It shows no stability issues. We have the pseudocode in the Appendix Section 4 and we've also included the codes in the Supplementary Material. \n\n**Computational and memory complexity**: We have computational complexity analysis in Line 179-182. We also provided detailed comparison of the running times of the VAE variant with/without our framework in the Supplementary Material/Appendix Section 8. \n\n**Prediction process**: We described the prediction process in the Supplementary Material/Appendix Line 23-24. For prediction, we compute the ELBO i.e the IWAE  estimator in Eqn. (2) using 5000 importance samples and 1 structure sample. We will move it to the camera-ready version of our paper.\n\n**Performance on VAE variants**: We detailed the implementation of combining our framework with VAE variants in Appendix Section 5.5. \n\n**Reconstruction evaluation**: We demonstrated reconstructions in Appendix Section 10.1 (Figures 7 and 8) and provided qualitative comparison and analysis."}]

Author rebuttal:
Attached here is the reconstructed samples of the cifar-10 dataset (related to the additional experimental results.)

Meta-review: This is definitely a nice and clear paper. The authors worked hard to convince the reviewers about the novelty compared to existing work and they succeeded. Acceptance is therefore recommended.
Decision: Accept (poster)