SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Broadly, works in implicit VI have proposed novel ways of estimating the ratio of the variational approximation to the prior (regularisation-term), usually referred to as density-ratio estimation (also referred to as the prior-contrastive formulation by Huszar, 2017). To the best of our knowledge, we are the first to propose and test an entropy approximation for an implicit density used as a variational approximation. In particular, Shi et al. (2018), Tran et al. (2017), and Pawlowski et al. (2017) have successfully demonstrated implicit variational inference using hypernetworks and align well with our goals. Tran et al. (2017) opt for training a discriminator network to maximally distinguish two distributions given only i.i.d. samples from each. This approach, though general, adds to the computational requirements and becomes more challenging in high dimensions (Sugiyama et al., 2012). To mitigate the overhead of training the discriminator for each update of the ELBO, many works limit the discriminator training to a single or few iterations. Furthermore, this approach entails an adversarial objective that is notoriously unstable (Mescheder et al., 2017b). Pawlowski et al. (2017) treat all weights as independent and find that a single discriminator network is inaccurate at estimating log ratios compared to the analytical form of Bayes by backprop (Blundell et al., 2015), and opt to use a kernel method that matches the analytical form more closely.
Shi et al. (2018) propose a novel way of estimating the ratio of the two densities using kernel regression in the space of parameters which obviates the need for a min-max objective but we expect to be inaccurate in estimating high-dimensional density ratios especially given a limited number of samples from both the densities as well as the RBF kernel. Pradier et al. (2018) are also motivated by the possibility of compressing the posterior in a lower-dimensional space and consider the parameters of the generator to be stochastic. They use an inference network with the generator/decoder and hence require empirical latent samples to train which doubles the training steps and limits the scalability. SIVI (Yin et al., 2018) and D-SIVI (Molchanov et al., 2019) use Monte Carlo (MC) averaging to approximate the variational approximation hence precluding an approximation over an intractable entropy as their resultant ELBOs only contains the explicit conditional  $q_{\gamma}(\theta \mid z)$ . Their novelties lie in using a general semi-implicit formulation to model all the weights of the network but as noted earlier this MC estimator is biased and has high variance. Yin et al., 2018 propose upper and lower bounds on the ELBO to train with this variational approximation and Molchanov et al. (2019) extend the approach by using an implicit prior.
Recently, several works have considered alternatives to high-dimensional inference by heuristically limiting randomness, e.g., Izmailov et al. (2020), Daxberger et al. (2021a,b), and Sharma et al. (2023). Our work is somewhat orthogonal to these, as we present a general method for performing high-dimensional inference. Indeed, while our focus here has been on inference in full networks, our bounds can readily be applied to subspaces or subnetworks too. A notable difference from the works building on the Laplace approximation, however, is that these works linearise the model during prediction and inference, whereas we do not. We linearise the neural sampler,  $g_{\gamma}(z)$ , but only when it is used to estimate the intractable entropy and its gradients, see Eqs. (10) to (12). Thus, the resulting posterior approximation is, still a highly flexible (i.e., non-linearised) implicit distribution.
We now test the proposed posterior approximation (Section 4) for its generalisation and calibration capabilities in downstream tasks on different established datasets. We use Bayesian Neural Networks (BNNs) for our analysis due to the immense number of global latent variables (i.e., the weights and biases) that are required by modern BNN architectures to perform well with larger datasets, validating our method on previously unreachable and untested scales for implicit VI.
We evaluate multiple posterior approximation schemes for solving various downstream tasks, with the final performance of the BNN serving as a metric for assessing the quality of each method's approximation of the true posterior. Ideally, we wish to understand whether an expressive posterior approximation helps uncertainty quantification in downstream tasks. As this is a challenging question to address directly, we compare LIVI against less flexible baselines to determine whether capturing posterior complexity leads to improved predictive performance and uncertainty estimates.
Model definition Briefly, BNNs are standard neural networks with a prior on their parameters, which, post-training, results in a posterior distribution. In a notation similar to that of Section 2, BNNs can be presented as  $p(\pmb{\theta} \mid (\pmb{x}, \pmb{y})) \propto p(y \mid f_{\pmb{\theta}}(\pmb{x})) p(\pmb{\theta})$ , where  $\pmb{\theta}$  represents the parameters of a
neural network denoted by  $f$ . The likelihood,  $p(y \mid f_{\theta})$ , depends on the given task. In this section, we test proxies to the true posterior  $p(\pmb{\theta} \mid (\pmb{x}, y))$  for how well they 1) model the dataset, and 2) model regions away from the dataset.
For regression tasks, we use a Gaussian likelihood and learn the homoscedastic variance,  $\eta^2$ , using type-II maximum likelihood,  $p(y|f_{\theta}(\boldsymbol{x})) = \mathcal{N}(y|\mu = f_{\theta}(\boldsymbol{x}),\eta^2)$ . For classification tasks, we use a categorical likelihood function for the labels with logits produced by  $f_{\theta}$ .
Generator architecture To efficiently generate millions of correlated parameters, we introduce a novel architecture for the generator  $g_{\gamma}(z)$ . We sample a low-dimensional noise matrix, which is gradually scaled up through the generator network using matrix multiplications (Shi et al., 2018). To generate all the parameters of the BNN, we gradually scale up this noise through the generator. The output from the first matrix multiplication layer is split and fed into a series of disconnected matrix multiplication networks that produce parameters for the individual layers of the BNN. This equates to pruning dense connections deeper in the network. This architecture is designed such that it can capture both within-layer correlations in the BNN (through the individual subnetworks), and across-layer correlations (through the first matrix multiplication), whilst having a manageable number of parameters. The architecture is described in detail in Appendix D and has been used for experiments in Sections 6.5 and 6.6.
Method comparison For LIVI, the generator or hypernetwork,  $g_{\gamma}(z)$ , represents the implicit distribution over network parameters for the BNN, which we sample from. We explicitly mention which of the ELBOs, Eqs. (15) and (16), we use in every experiment. We compare our method with other scalable uncertainty quantification methods: Adversarial Variational Bayes (AVB, Mescheder et al., 2017a), last-layer Laplace approximation (LLLA, Daxberger et al., 2021a), deep ensembles (DE, Lakshminarayanan et al., 2017), mean-field VI (MFVI, Saul et al., 1996; Osawa et al., 2019), as well as a simple MAP estimate. We chose DEs over other methods, such as SWAG (Maddox et al., 2019), as Daxberger et al. (2021a) found this to be a stronger baseline. Where possible, we compare against Kernel Implicit Variational Inference (KIVI, Shi et al., 2018); these results are taken from their paper. The only other implicit VI method is AVB, which, in all experiments, uses the exact same generator architecture as LIVI. Where feasible, we compare with HMC as it is the gold-standard posterior approximation technique.
Overview of experiments The first two sections focus on regression; a toy example in Section 6.2 and tests on UCI regression datasets in Section 6.3. Next, we move to classification experiments, focusing on out-of-distribution (OOD) benchmarks. Section 6.4 shows results on MNIST, then Section 6.5 considers images from CIFAR10, allowing us to perform OOD testing of the posterior approximations in millions of dimensions. In Section 6.6, we test on CIFAR100, evaluating our approach on even larger dimensional BNN posteriors<sup>1</sup>. We discuss the metrics used in Appendix E.1.
6.2 Toy data
In Fig. 1, we compare posterior predictions of our method (LIVI) against the standard baselines for posterior inference on a simple sinusoidal toy dataset. We qualitatively assess the epistemic uncertainty, or the variance in the function space, in regions away from training data. The BNN contains 105 parameters and is thus small enough that we can use  $\mathcal{L}'$ , Eq. (15), as the objective function for LIVI, (for the remaining details of the architectures and the training, see Appendix E.3). We observe that both AVB and MFVI dramatically underestimate uncertainties, whereas LIVI is much closer to Hamiltonian Monte Carlo (HMC). We further observe that LIVI's learned posterior approximation has multiple modes and heavy tails, see Appendix F.4.
6.3 UCI datasets
Here, we follow the setups in Lakshminarayanan et al. (2017) and Shi et al. (2018) and choose a single hidden-layered MLP as the BNN for all datasets. Due to the different dimensionalities of the datasets, the number of BNN parameters vary between 501 and 751. Still, at this scale, it is feasible to work with whole Jacobian matrices, which allows us to test  $\mathcal{L}'$ , Eq. (15), as the objective function and
Supporting code is available here - https://github.com/UppalAnshuk/LIVI_neurips23
Table 1: UCI regression datasets. We report RMSE  $(\downarrow)$  and log-likelihoods  $(\uparrow)$  on the test set and average across three different seeds for each model to quantify the variance in the results.
Method	Boston	Concrete	Energy	Kin8nm	Naval
Test RMSE	LIVI (L')	2.32 ± 0.07	4.24 ± 0.17	0.41 ± 0.27	0.03 ± 0.00	0.00 ± 0.00
LIVI (L'')	2.40 ± 0.09	4.62 ± 0.13	0.44 ± 0.11	0.08 ± 0.01	0.00 ± 0.01
HMC	2.26 ± 0.00	4.27 ± 0.00	0.38 ± 0.00	0.04 ± 0.00	0.00 ± 0.00
DE	3.28 ± 1.00	6.03 ± 0.58	2.09 ± 0.29	0.09 ± 0.00	0.00 ± 0.00
KIVI	2.80 ± 0.17	4.70 ± 0.12	0.47 ± 0.02	0.08 ± 0.00	0.00 ± 0.00
Test LL	LIVI (L')	-2.16 ± 0.05	-2.79 ± 0.11	-1.17 ± 0.13	1.24 ± 0.04	6.74 ± 0.04
LIVI (L'')	-2.40 ± 0.09	-2.99 ± 0.13	-1.37 ± 0.11	1.15 ± 0.01	5.84 ± 0.06
HMC	-2.20 ± 0.00	-2.67 ± 0.00	-1.14 ± 0.00	1.27 ± 0.00	7.79 ± 0.00
DE	-2.41 ± 0.25	-3.06 ± 0.18	-1.31 ± 0.22	1.28 ± 0.02	5.93 ± 0.05
KIVI	-2.53 ± 0.10	-3.05 ± 0.04	-1.30 ± 0.01	1.16 ± 0.01	5.50 ± 0.12
Table 2: OOD Test M-C1: We report out-of-distribution performance for models trained on MNIST (and evaluated on FMNIST, KMNIST and EMNIST) or CIFAR10 (and evaluated on SVHN, LSUN and CIFAR100). The metrics have been averaged across the OOD datasets.
Method	Confidence ↓	AUROC ↑
MNIST	CIFAR10	MNIST	CIFAR10
MAP	72.10 ± 0.36	81.40 ± 0.16	96.32 ± 0.22	86.73 ± 0.64
MFVI	69.23 ± 0.24	74.71 ± 0.23	96.53 ± 0.16	87.50 ± 0.25
LLLA	67.40 ± 0.19	53.6 ± 0.3	96.67 ± 0.27	89.03 ± 0.51
DE	63.14 ± 0.11	67.17 ± 0.21	97.52 ± 0.08	89.61 ± 0.11
LIVI	55.03 ± 0.13	43.47 ± 0.28	97.91 ± 0.27	91.83 ± 0.41
AVB	70.68 ± 0.45	NA	95.5 ± 0.4	NA
to analyse the performance difference with  $\mathcal{L}''$ , Eq. (18). Further details of the generator architectures and baseline methods are presented in Appendix E.4.
Our results are summarised in Table 1. Both LIVI bounds,  $\mathcal{L}'$  and  $\mathcal{L}''$ , perform similarly, giving empirical justification for the lower bound in Eq. (17). Furthermore, both bounds are comparable to HMC. We train our objective with far fewer samples compared to KIVI and outperform it for even small hypernetwork architectures.
6.4 MNIST Dataset
We use three different out-of-distribution tests to examine the estimated uncertainty in downstream tasks. We use the LeNet architecture for the Bayesian neural network with a total of 44000 parameters, necessitating the use of our second objective  $\mathcal{L}''$ , Eq. (18). We train on MNIST, achieving an indistribution test accuracy of  $\sim 99.1\%$  for all methods except MFVI, which achieves  $98.6\%$ .
OD Test M-C1 In this experiment, we consider three OOD datasets, FMNIST, KMNIST, EMNIST and Table 2 reports the averaged results. We outperform all the methods we compare against in both categories. This signifies that the implicit approximation makes the BNN less over-confident than competing methods when predicting over datasets that the BNN has not encountered in training.
OD Test M2 We compare our method on another OOD benchmark by rotating the MNIST image by different angles (Daxberger et al., 2021a). For this benchmark, we plot the negative log-likelihoods and expected calibration errors for all rotation angles in Fig. 2, testing all approaches at varying extents of drift. We conclude that we perform the best in these two categories as well.
OD Test M3 As the last benchmark on MNIST, we opt to plot the empirical CDF of predictive entropies across OOD images (Lakshminarayanan et al., 2017; Louizos et al., 2017). Under covariate shift, a well-calibrated model should predict a uniform distribution for never-before-seen data. So, with a model trained on MNIST we calculate predictive entropies on OOD datasets, plotting each method's fraction of high-entropy predictions in Fig. 3. In other words, the smaller the area between the curve and the dashed line, the better the method. We find that LIVI consistently produces
Figure 2: OOD Test M2: Rotated MNIST benchmark. LIVI performs as well as, or better, than competing methods.
Figure 3: OOD Test M3: Empirical CDF plot. On the  $x$ -axis is the entropy of a single prediction and empirical CDF on the  $y$ -axis. With this plot, we count the number of high entropy (uniform distribution) predictions that each method makes on OOD data. As we count over the whole dataset, we would like to encounter more high entropy predictions closer to the value 2.30 on the  $x$ -axis.
higher-entropy predictions than competitors, showing that it is less overconfident for any OOD image.
We elaborate on our experimental set-ups in Appendix E.5, keeping the critical training parameters the same across methods and also reporting our computational requirements for these experiments.
6.5 CIFAR10
To test our approach at a larger scale, we consider the CIFAR10 image classification dataset. Here, for the BNN we use the same WideResNet architecture as Daxberger et al. (2021a) with 16 layers and a widen factor of 4 (Zagoruyko et al., 2016). This network architecture contains over 2.7 million trainable parameters, meaning that we again use our second objective  $\mathcal{L}''$ , Eq. (18). We generate all parameters with our novel hypernetwork architecture, described in Appendix D, with training details in Appendix E.6. All the methods reached very similar accuracies of  $\sim 91.5\%$  on the in-distribution test set when trained for 200 epochs, except for MFVI, which reached  $\sim 88.2\%$ .
OOD Test M-C1 We report averaged AUROC and confidence to test if the models assign wrong labels with high probability, see Table 2. The set of OOD datasets contains SVHN, LSUN, and CIFAR100. LIVI again outperforms competitors by being less overconfident about OOD predictions.
OD Test C2 We test the models' performances when presented with corrupted images from CIFAR-10-C (Hendrycks et al., 2019). The negative log-likelihoods (NLL) and expected calibration errors (ECE) for different corruption intensities are plotted in Fig. 4. LIVI retains a low NLL across all corruption intensities, showing that its predictive uncertainty is better calibrated than competitors. Similarly, LIVI achieves consistently low ECEs, showing that its accuracy and confidence are well-aligned (i.e., decrease at a similar rate), even as the corruption intensity increases.
6.6 CIFAR100
Finally, we test our approach on CIFAR100 using a WideResNet(28,10) as the BNN. Not only is this dataset more complex to model, the BNN contains roughly 36.5 million parameters, demonstrating that our approach can scale to large networks. The generator used for LIVI has only about twice
Figure 4: OOD Test C2: Corrupted CIFAR10 benchmark. OOD performance for methods trained on CIFAR10 and making predictions for CIFAR-10-C images corrupted with Gaussian blur (Hendrycks et al., 2019). LIVI performs as well or better than competitors.
Table 3: CIFAR100 results. Posterior testing with a WideResNet(28,10) architecture on the CIFAR100 dataset. The numbers reported here are on the test set, with those for LIVI being averaged across three seeds.
Method	Accuracy (%) ↑	NLL ↓	ECE ↓
MAP†	79.8	0.875	0.086
MFVI†	77.8	0.944	0.097
LIVI	78.7 ± 0.4	0.774 ± 0.080	0.036 ± 0.004
Ensemble (n = 4)†	82.7	0.666	0.021
$\dagger$  Results from Nado et al. (2021).
as many (variational) parameters as the BNN latent variables – the same amount of variational parameters required for a mean-field Gaussian posterior.
We compare our results with those presented by Nado et al. (2021), who use the same BNN architecture, see Table 3. This architecture is also used by Sharma et al. (2023), who test performance on this dataset in various sub-stochastic settings and report on the same metrics, with which our results can therefore also be compared. The results indicate that LIVI's learned posterior approaches the performance of an ensemble with a combined 146 million parameters, nearly twice the roughly 75 million parameters required by LIVI. Moreover, LIVI's posterior approximation appears better than a simpler approximation like MFVI, even though this uses the same amount of parameters.
In this paper, we present a novel method to scale implicit variational inference to latent variable models with millions of dimensions by circumventing the need for a discriminator network and an adversarial loss. We find that modelling the posterior with a highly flexible approximation indeed does have benefits in the OOD and data drift scenarios, which may be due to the approximate posterior capturing richer information like correlations and multiple modes. Our method performs better than deep ensembles in a number of OOD benchmarks. Unlike conventional probabilistic methods, we do not fall short on in-distribution accuracy. We empirically evaluated our bound on Bayesian neural networks with millions of latent variables. Our method can also be applied for inference in other latent variable models, such as VAEs (Kingma et al., 2014; Rezende et al., 2014; Mescheder et al., 2017a). Future work includes scaling our methods to models with hundreds of millions of parameters. For this, one could consider independent hypernetworks for each target model layer, thus reducing computational requirements at the cost of losing correlations across the layers. One could also consider efficient ways of putting priors over deep networks to curb dimensionality, such as the works by Karaletsos et al. (2018) and Trinh et al. (2020).
The authors thank Søren Hauberg and Pierre-Alexandre Mattei for providing their insight multiple times during the development of this project. We acknowledge EuroHPC Joint Undertaking for
awarding us access to Karolina at IT4Innovations, Czech Republic. AU, WB, and JF acknowledge funding from the Novo Nordisk Foundation through the Centre for Basic Machine Learning Research in Life Science (MLLS, NNF20OC0062606). KSS acknowledges funding from the Novo Nordisk Foundation (NNF20OC0065611). Furthermore, JF was in part funded by the Novo Nordisk Foundation (NNF20OC0065611), the Innovation Fund Denmark (0175-00014B) and Independent Research Fund Denmark (9131-00082B).
Pradier, Melanie F., Weiwei Pan, Jiayu Yao, Soumya Ghosh, and Finale Doshi-velez (2018). Projected BNNs: Avoiding weight-space pathologies by learning latent representations of neural network weights.
Rezende, Danilo Jimenez, Shakir Mohamed, and Daan Wierstra (2014). "Stochastic backpropagation and approximate inference in deep generative models". In: International Conference on Machine Learning.
Saul, Lawrence K., Tommi Jaakkola, and Michael I. Jordan (1996). "Mean Field Theory for Sigmoid Belief Networks". In: J. Artif. Int. Res. 4.1, pp. 61-76.
Sharma, Mrinank, Sebastian Farquhar, Eric Nalisnick, and Tom Rainforth (2023). "Do Bayesian Neural Networks Need To Be Fully Stochastic?" In: Proceedings of The 26th International Conference on Artificial Intelligence and Statistics. Ed. by Francisco Ruiz, Jennifer Dy, and Jan-Willem van de Meent. Vol. 206. Proceedings of Machine Learning Research. PMLR, pp. 7694-7722.
Shi, Jiaxin, Shengyang Sun, and Jun Zhu (2018). "Kernel Implicit Variational Inference". In: International Conference on Learning Representations.
Sugiyama, Masashi, Taiji Suzuki, and Takafumi Kanamori (2012). Density Ratio Estimation in Machine Learning. Cambridge University Press.
Tipping, Michael E and Christopher M Bishop (1999). "Probabilistic principal component analysis". In: Journal of the Royal Statistical Society: Series B (Statistical Methodology) 61.3, pp. 611-622.
Titsias, Michalis K. and Francisco Ruiz (2019). "Unbiased Implicit Variational Inference". In: Proceedings of the Twenty-Second International Conference on Artificial Intelligence and Statistics. Vol. 89. Proceedings of Machine Learning Research, pp. 167-176.
Tran, Dustin, Rajesh Ranganath, and David Blei (2017). "Hierarchical implicit models and likelihood-free variational inference". In: Advances in Neural Information Processing Systems 30.
Trinh, Trung, Samuel Kaski, and Markus Heinonen (2020). "Scalable Bayesian neural networks by layer-wise input augmentation". In: arXiv preprint arXiv:2010.13498.
Wilson, Andrew G and Pavel Izmailov (2020). "Bayesian Deep Learning and a Probabilistic Perspective of Generalization". In: Advances in Neural Information Processing Systems. Vol. 33, pp. 4697-4708.
Yin, Mingzhang and Mingyuan Zhou (2018). "Semi-Implicit Variational Inference". In: Proceedings of the 35th International Conference on Machine Learning. Vol. 80. Proceedings of Machine Learning Research, pp. 5660-5669.
Zagoruyko, Sergey and Nikos Komodakis (2016). "Wide Residual Networks". In: Proceedings of the British Machine Vision Conference (BMVC).
A Derivation of the entropy approximation
In this section, we show that the approximation of the differential entropy in Eq. (13) can be further approximated by Eq. (14). Starting from Eq. (13), the entropy of a Gaussian DLVM can be approximated by  $\tilde{H}[\tilde{q}(\theta)]$ , that is
$$
H \left[ q _ {\boldsymbol {\gamma}} (\boldsymbol {\theta}) \right] = - \mathbb {E} _ {\boldsymbol {z} \sim q (\boldsymbol {z})} \mathbb {E} _ {\boldsymbol {\theta} \sim q _ {\boldsymbol {\gamma}} (\boldsymbol {\theta} \mid \boldsymbol {z})} [ \log q _ {\boldsymbol {\gamma}} (\boldsymbol {\theta}) ] \approx \underbrace {- \mathbb {E} _ {\boldsymbol {z} \sim q (\boldsymbol {z})} \mathbb {E} _ {\boldsymbol {\theta} \sim q _ {\boldsymbol {\gamma}} (\boldsymbol {\theta} \mid \boldsymbol {z})} [ \log \tilde {q} _ {\boldsymbol {z} ^ {\prime} = \boldsymbol {z}} (\boldsymbol {\theta}) ]} _ {=: \hat {H} [ \tilde {q} (\boldsymbol {\theta}) ]}, \tag {A.1}
$$
where we use the subscript  $z' = z$  to specify that  $\tilde{q}$  is a linearised distribution that is linearised at  $z$  and distinct for each sample from the base distribution  $q(z)$ . Using the definition of  $\tilde{q}$  from Eqs. (11) and (12), we have
$$
\log \tilde {q} _ {\boldsymbol {z} ^ {\prime}} (\boldsymbol {\theta}) = - \frac {m}{2} \log 2 \pi - \frac {1}{2} \log \det  C (\boldsymbol {z} ^ {\prime}) - \underbrace {\frac {1}{2} (\boldsymbol {\theta} - \boldsymbol {\mu} (\boldsymbol {z} ^ {\prime})) ^ {\intercal} C (\boldsymbol {z} ^ {\prime}) ^ {- 1} (\boldsymbol {\theta} - \boldsymbol {\mu} (\boldsymbol {z} ^ {\prime}))} _ {=: h (\boldsymbol {\theta}, \boldsymbol {z} ^ {\prime})}, \tag {A.2}
$$
where, c.f. Eq. (12),