SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Table 1: Test performance in negative log-likelihood  $(-LL)$  mean  $\pm 1$  standard deviation (lower the better) over 4 runs 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±0.05	82.50±0.00	82.60±0.10	82.30±0.10
(4,16)	82.47±0.45	82.33±0.13	82.52±0.02	83.02±0.20
Omniglot	(8,8)	107.10±0.10	106.45±0.10	106.55±0.30	106.34±0.01
(4,16)	108.12±0.16	107.15±0.08	107.35±0.40	108.30±0.50
Caltech101	(8,8)	116.83±1.57	114.94±0.45	114.00±0.42	113.54±0.40
(4,16)	116.30±1.11	114.55±1.18	113.02±0.34	112.53±0.42
Proof of this theorem is 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  $T_{c} = O(NBLM^{2})$  with  $N$  training examples and  $B$  epochs. Our method is linearly scalable as  $ST_{c}$ . With a proper thresholding, the number of active layers  $L$  is relatively small in each sample.
We analyze the behavior of our inference framework across various tasks. We study how AdaVAE 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-the-art performance on benchmark datasets. Finally, we demonstrate the framework's compatibility with different types of backbone networks and VAE variants.
5.1 Adaptive VAE Network Structures
AdaVAE enables us to perform joint inference on both encoding/decoding network structures and latent variables. To investigate how network structures evolve during training epochs, we set the truncation level  $T = 25$  on MLP backbone nets with tanh non-linearities. We analyze AdaVAE's behavior on  $28 \times 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 2 Left shows the evolution of the encoding/decoding network structures for one trial with  $S = 8$ . AdaVAE initializes multiple hidden
2Implementation details are in the Appendix. Codes are provided.
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 \leq L \leq 5$ , the performance of other regularization methods initially improves but then starts to decline, suggesting they suffer from overfitting even for shallow structures.
Table 2: The best performance of our method and the VAE regularization methods in Figure 3. We also demonstrate the compatibility of our framework.
MNIST	Omniglot	Caltech 101
-LL ↓	MI ↑	-LL ↓	MI ↑	-LL ↓	MI ↑
MIWAE [18]	86.11±0.01	9.13±0.01	110.61±0.10	8.98±0.00	116.19±0.08	7.64±0.01
MIWAE+DO	90.99±0.01	9.13±0.01	110.89±0.01	8.98±0.00	116.00±0.40	7.63±0.01
DVAE [10]	87.67±0.17	9.07±0.00	112.04±0.13	8.96±0.00	113.71±0.43	7.53±0.01
CR-VAE [12]	87.67±0.05	9.03±0.01	109.94±0.19	8.97±0.00	117.46±0.10	7.37±0.01
IBP-DGM [14]	92.24±0.75	-	124.01±1.62	-	135.23±0.54	-
BB-VAE [15]	91.55±0.69	-	124.47±0.58	-	135.18±0.83	-
Ours	82.30±0.10	9.13±0.01	106.34±0.01	8.98±0.00	113.54±0.40	7.67±0.02
Ours+DVAE	83.30±0.20	9.07±0.02	107.80±0.40	8.96±0.00	111.92±0.37	7.63±0.00
Ours+CR-VAE	85.20±0.20	9.03±0.01	105.60±0.05	8.97±0.00	108.93±1.40	7.51±0.02
layers with sparsely activated neurons, gradually converging to fewer fully activated layers. Figure 2 right top presents the medians of the number of active hidden layers in the encoding/decoding networks, as well as the percentage of activated neurons in the truncation changing over epochs. The encoding network structure stabilizes at 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 2 right bottom shows the generative performance assessed by negative log-likelihood  $(-LL)$  converge faster with an increased structure sample size  $S$ , which is consistent with Theorem 1.
5.2 Effect of Structure Sample Sizes
We assess our estimator using three benchmark datasets: MNIST [47], Omniglot [48], and Caltech101 Silhouettes [49]. In each minibatch, we set a budget of  $S \times M_s \times K_s = 64$  total latent variable samples for each datapoint. We examine four settings of the VAE structure sample size  $S = \{1,2,4,8\}$ , along with latent variable sample sizes  $(M,K) = \{(8,8),(4,16)\}$  as in [18]. The truncation level is  $T = 25$  with a maximum width  $O = 200$ . The distribution over the output from the decoding networks is factorized Bernoulli. Table 1 indicates larger values of  $S$  generally yield better performance. For MNIST, there is no statistically significant difference between  $S = 1$  and  $S = 8$ . 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-VAE [12], and BB-VAE [15], MIWAE with dropout (MIWAE+DO), along with vanilla MIWAE [18]. as in Figure 3. All methods share the same maximum width of  $O = 200$  and a latent variable dimensionality of 50. For
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.
Figure 5: Influence of the maximum number of neurons per layer  $O$  on our method. When  $O$  is small (e.g.,  $O \leq 100$ ), we tend to have shallower encoding/decoding networks. As  $O$  becomes reasonably large (e.g.,  $O \geq 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, we 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  $L = T = \{1,2,3,4,5,10,15,20,25\}$ , allowing us to compare their effectiveness in both shallow (i.e.,  $L = T \leq 5$ ) and deep (i.e.,  $L = T \geq 5$ ) VAE network structures.
Figure 3 shows that when  $L = T \leq 5$ , the VAE regularization methods can mitigate overfitting. However, even for the shallow cases as  $L = 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 \leq 2$ , it still outperforms other methods. For deep VAE structures with  $T / L = \{10,15,20,25\}$ , Figure 3 shows that our method's performance is minimally impacted by large truncations. This robustness indicates our method's ability to mitigate overfitting for large VAE structure settings. In contrast, as the depth  $L$  increases, the steadily increasing negative log-likelihood  $(-LL)$  for other regularization methods suggests their inability to 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.
Table 2 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 to the input or with CR-VAE by integrating the consistency regularization term into our estimator, we observe enhanced performance on Omniglot and Caltech101 datasets. Furthermore, in Figure 4, 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. A 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). Figure 5 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 \leq 100$ , we tend to have shallower
(b)
Figure 6: Our framework's performance on different VAE backbone networks. (a) and (b) show VAE and  $\beta$ -VAE [51] with convolutional layers on MNIST and Omniglot datasets. (c) and (d) show VGAEs [52] with graph convolutional layers on Cora and Citeseer datasets.
encoding/decoding networks with a lower percentage of activated neurons in the truncation. However, when  $O$  is reasonably large as  $O \geq 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 Backbone Networks
We demonstrate AdaVAE's efficacy by applying it to VAEs with different types of encoding/decoding backbone networks. To infer the number of convolutional layers in a convolutional VAE (cVAE) using the beta process, we readily mask the convolution channels in layer  $l$  with  $\mathbf{z}_{.l}$ . Figure 6 (a) and (b) show that by adapting network structures to the data we improve the overall performance of cVAE and  $\beta$ -cVAE [51]. The adaptive backbone networks effectively prevent overfitting for deep structure settings (i.e.,  $T / L \geq 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 of an edge between the nodes. We combine our framework with VGAE by elementwisely multiplying the GC channels (i.e., the feature vectors) of layer  $l$  with  $\mathbf{z}_{.l}$  in both its encoding/decoding networks. We compare the hybrid method's performance with graph autoencoders (GAEs) [52] and vanilla VGAEs on two benchmark graph-structured datasets: Cora [53] and CiteSeer [54]. The AUC scores of link prediction over varying numbers of GC layer settings are shown in Figure 6 (c) and (d). Our framework enables the VGAE to 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 / L = \{6,8,10\}$ ).
5.6 Application to VAE Variants
We assess the performance of our inference framework by leveraging it to adapt the network structures of VAE variants to data. Specifically, for  $\beta$ -VAE, we apply layer-wise binary masks  $\mathbf{z}_{l}$  to the convolutional channels and infer the layer numbers using the beta process. In the case of LaddervAE (LVAE) [4], we adjust its depth by applying layer-wise binary masks to its deterministic layers in the bottom-up dependency structures and add
Table 3: Performance comparison of VAE variants with and without our inference framework in terms of  $-LL$ , MI, and KL divergence.
Methods	-LL ↓	MI ↑	KL ↑
β-cVAE (β = 2) [51]	106.50±0.12	8.46±0.01	18.01±0.21
Ours + β-cVAE (β = 2)	102.45±0.07	8.47±0.01	20.65±0.02
LVAE [4]	136.50±1.50	8.43±0.02	20.05±0.21
Ours+LVAE	121.48±0.67	8.50±0.01	21.00±0.47
SkipVAE [34]	112.68±0.80	8.50±0.01	22.82±0.79
Ours+SkipVAE	108.00±0.04	8.50±0.01	28.26±0.10
NVAE [6]	98.83±0.17	8.51±0.01	34.97±0.13
Ours+NVAE	99.10±0.20	8.51±0.01	37.85±0.50
skip connections between the stochastic layers. For SkipVAE [34], we model its depth by employing layer-wise binary masks and skip connections in both its encoding/decoding networks. The expressive network structures of NVAE [6] consists of multiple blocks of convolutional layers. We apply our
framework to infer the number of blocks for a 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 3 on FashionMNIST [55] demonstrate that by inferring the encoding/decoding network structures we significantly improve the density estimation performance of the VAE variants. Our framework also boosts their ability to mitigate posterior collapse as indicated by MI and KL divergence.
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, AdaVAE exhibits promising applications across various types of VAE backbone networks and VAE variants, including those with hierarchical structures such as LVAE. Notably, AdaVAE 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.
This material is based upon work supported by the National Science Foundation under NSF Award No.2045804 and Award No.1850492. We are thankful to Kishan KC for helpful discussion. We acknowledge Research Computing at the Rochester Institute of Technology [57] for providing computational resources.
[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 on Machine Learning (ICML), pages 1278-1286. PMLR, 2014.
[3] Jyoti Aneja, Alex Schwing, Jan Kautz, and Arash Vahdat. A 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 Maaloe, Marco Fraccaro, Valentin Lievin, and Ole Winther. BIVA: A very deep hierarchy of latent variables for generative modeling. In Advances in Neural Information Processing Systems (NeurIPS), volume 32, 2019.
[6] Arash Vahdat and Jan Kautz. NVAE: A deep hierarchical variational autoencoder. In Advances in Neural Information Processing Systems (NeurIPS), volume 33, pages 19667-19679, 2020.
[7] Corinna Cortes, Xavier Gonzalvo, Vitaly Kuznetsov, Mehryar Mohri, and Scott Yang. AdaNet: Adaptive structural learning of artificial neural networks. In Proceedings of 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 (NeurIPS), volume 31, 2018.
[50] 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.
[51] Irina Higgins, Loic Mcthney, Arka Pal, Christopher P. Burgess, Xavier Glorot, Matthew M. Botvinick, Shakir Mohamed, and Alexander Lerchner. beta-VAE: Learning basic visual concepts with a 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 construction of internet portals with machine learning. 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: a novel image dataset for benchmarking machine learning algorithms. arXiv preprint arXiv:1708.07747, 2017.
[56] Kai Xu, Akash Srivastava, and Charles Sutton. Variational russian roulette for deep bayesian nonparametrics. In International Conference on Machine Learning (ICML), pages 6963-6972. PMLR, 2019.
[57] Research Computing at Rochester Institute of Technology. https://doi.org/10.34788/OS3G-QD15.
This section is in the supplemental material.