SlowGuess/NeurIPS_2023_OCR huggingface dataset:
LIVI We use our novel generator architecture for this experiment with a noise dimension of  $130 \times 130$ . This architecture and its configuration file for matrix multiplication layers are present in the code. We train with a batch size of 256 and with 2 samples per minibatch. We use Eq. (17) else as the objective for training. Our generator architecture produces all the parameters of the BNN, implicitly learning the correlation between and across layers. This hypernetwork contains about 7 million parameters, in other words, these are the number of variational parameters that tune the implicit distribution.
DE We train five WideResNet (Zagoruyko et al., 2016), each 16 layers deep with different initialisations using the same optimiser hyperparameters i.e. with NAdam optimiser with weight decay of  $1 \times 10^{-4}$ , a learning rate of  $7.8 \times 10^{-4}$  and with cosine annealing of the learning rate. We average predictions from these 5 networks to represent deep ensembles in our experiments.
MAP We train one WideResNet with the same weight decay and optimiser hyperparameters as above and use that for representing the MAP estimate.
LLLA For the last-layer Laplace approximation, we use a network trained similarly to a MAP network and then run post hoc Laplace on it, inferring a full Hessian matrix for all of the last-layer weights. Again this is done to strike a middle ground between inferring a very faithful posterior and keeping the computational requirements in check.
AVB Even though AVB is theoretically very close to our variational approximation, the size of the discriminator to distinguish 2.7 million dimensional samples from the prior and the implicit posterior to estimate the regularisation term in Eq. (2) poses a prohibitively large overhead. Due to this, we could not train AVB for this set of experiments.
Training Interestingly, our method trained using Eq. (16) takes less time to train than a deep ensemble consisting of five networks with the same training parameters and provides better uncertainty quantification than a non-Bayesian ensemble of neural networks. Keeping everything else the same, each network of the ensemble takes 34 minutes to train on CIFAR10 using an A5000(24 gb vram) for a total of 2.83 GPU hours whilst our method/architecture requires 2.30 GPU hours to train on the same dataset using the same GPU. We use Eq. (17) to train our implicit approximation and consume 4.1 GB of memory with a batch size of 256 and training with 2 elbo samples per minibatch. A vanilla WideResNet (MAP) consumes 1.4 GB to train with the same batch size.
E.7 Details of Fig. 1
In Fig. 1, we compare both the objective functions presented in this work for training with implicit variational approximations to different methods for uncertainty quantification for neural networks. All models were trained for five thousand iterations and had to learn observation noise present in the toy sinusoidal dataset. For all methods, we train with a homoscedastic assumption and obtain an observation noise parameter using type-II maximum likelihood. We deliberately removed a part of the data to see if the models tested were able to find in-between uncertainties. All methods were given the same sized 2 hidden layered MLP with 7 and 10 units respectively. We trained 5 MAP networks with different seeds for Deep Ensembles and average their predictions to make the plot. The variance of the predictions was then used for the epistemic uncertainty in blue. We also train the model with an observation noise parameter. For MFVI, we annealed the KL divergence between the prior and the posterior with a weighting term to help it to converge. The weighting term reaches one as the training nears completion. For HMC we sample 10000 samples using the library Pyro. We also tried to make multiplicative normalizing flows converge for this dataset, but with even 20K parameters and training for ten thousand iterations with a very small learning rate did not help. We even tried KL down weighting to reduce the effect of the prior in the initial iterations but that did not work either.
After training, we also plot a KDE-plot of the samples from the generator in Appendix F.4. As a sanity check, we do confirm that the generator is capable of representing non-trivial distributions as we can spot heavy tails and multiple modes in these plots. In this experiment, we qualitatively compare the epistemic uncertainty resulting from the respective inference schemes in the data-scarce regions.
E.8 Computation Graph
Here we provide some details about how the combination of the joint generator-BNN model works. The Bayesian neural network classes for all types of architectures feed-forward, convolutions, etc.) require a generator in the init function. As such, the generator networks reside inside the BNN and reparametrize it with a simple sample_parameters function. The most important part of this kind of implementation was the layers themselves. PyTorch provides different kinds of mutable layer implementations in nn-module but these layers do not expose their state i.e. their parameters in a manner that allows changes on the fly during training. We reimplemented the layers allowing such resampling to occur with the generator. In the init function of the BNN, we generate one set of parameters with the generator, package it in a dict that has the weight sample as well as a index to know the number of weights used by a previous layer. This counter index is updated by each layer in their init and sample_parameters function. As such, only the parameters of the generator are trainable, the parameters of the BNN are switchable and relay gradients to the generator via the likelihood or the entropy term.
Table F.1: UCI regression datasets. We report RMSE  $(\downarrow)$  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
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
MNF	3.31 ± 0.10	5.82 ± 0.04	1.04 ± 0.01	0.08 ± 0.01	0.01 ± 0.00
Table F.2: UCI regression datasets. We report log-likelihood  $(\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
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
MNF	-2.66 ± 0.08	-3.24 ± 0.09	-1.34 ± 0.07	1.10 ± 0.01	5.01 ± 0.00
In this section, we report additional results. In particular, we consider multiplicative normalising flows (MNF, Louizos et al., 2017), as well as two versions of full Laplace, 1) a non-linearised full network Laplace using a low-rank approximation to the Hessian (FLAP), and 2) a linearised full network Laplace using a Kronecker factored structure of the Generalised Gauss-Newton approximation to the Hessian (FLLA).
While we show results for linearised Laplace, as this method often achieves state-of-the-art performance, it must be emphasised that, for LIVI, we do not linearise the posterior; we linearise the neural sampler  $g_{\gamma}(z)$ , but only when it is used within the estimation the intractable entropy and its gradients, see Eqs. (10) to (12). The resulting posterior approximation is, therefore, still a highly flexible (that is, non-linearised) implicit distribution. Linearised Laplace and LIVI are, therefore, fundamentally different approaches. Whereas linearised Laplace linearises the BNN and the predictive function, we linearise the neural sampler/hyper-network within the variational distribution  $q_{\gamma}(\theta \mid z)$ , thus obtaining the approximation  $\tilde{q}_{z'}(\theta \mid z)$ , but only when it is needed to approximate the entropy. This is also the reason why we do not linearise our BNN for predictions, which is common practice with linearised Laplace.
F.1 Additional results for UCI datasets
In Tables F.1 and F.2, we report additional results for experiments on the UCI datasets. Overall, LIVI  $(\mathcal{L}^{\prime},$  Eq. (15)) performs competitively, either outperforming or being on par with HMC, which is considered the gold standard. The multiplicative normalising flow (MNF) is difficult to get to converge, and both the training time and memory consumption are worse than for LIVI (see Appendix F.3).
We do not show results for full Laplace here, as we repeatedly ran into errors thrown by the second-order backends (BackPack and Asdl) of the laplace library (Daxberger et al., 2021a) when used regression tasks.
F.2 Additional results for OOD experiments
We show further results for the OOD experiment Test M-C1 in Table F.3, as well as Test M2 and M3 in Figs. F.1 and F.2, respectively. Note that we were unable to scale all the methods listed in Table F.3 to work on CIFAR10. In the table and both figures, the non-linearised version of full Laplace (FLAP) underfits due to the quite crude approximation to the posterior over the full network and generates
Figure F.1: Additional results for OOD Test M2: Rotated MNIST benchmark.
Table F.3: Further experiments for OOD Test M-C1. FLLA: full linearised Laplace; FLAP: full non-linearised Laplace; MNF: multiplicative normalising flows.
Method	Confidence ↓	AUROC ↑
MNIST	CIFAR10	MNIST	CIFAR10
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
MAP	72.10 ± 0.36	81.40 ± 0.16	96.32 ± 0.22	86.73 ± 0.64
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
MFVI	69.23 ± 0.24	74.71 ± 0.23	96.53 ± 0.16	87.50 ± 0.25
FLLA	52.63 ± 0.17	NA	97.57 ± 0.19	NA
FLAP	22.40 ± 0.32	NA	59.55 ± 0.23	NA
MNF	61.77 ± 0.26	NA	96.37 ± 0.15	NA
poor samples. The linearised version (FLLA) works slightly better than LIVI in terms of ECE and NLL on the rotated-MNIST benchmark (Fig. F.1) and in the OOD entropy-CDF test (Fig. F.2). As noted before in Section 5 and ??, our approach is fundamentally different from linearised Laplace approaches as we do not linearise the model during prediction and inference and hence these methods are not directly comparable.
F.3 Runtime analysis
We report training times until convergence as well as memory consumption for LIVI and four competitors on MNIST in Table F.4. The methods were all run on the same single-GPU (NVIDIA RTX A4000) machine. Compared to other expressive VI methods like AVB and MNF, LIVI is faster to train and consumes less memory. In particular, the hypernetworks used for LIVI in all experiments only contain about twice as many parameters as the networks they are modelling the posterior over the same amount of parameters required for mean-field VI. Interestingly, even though the memory requirements for LIVI are higher than that of DE, LIVI trains in a shorter amount of time.
Figure F.2: Additional results for OOD Test M3: Empirical CDF plot.
Table F.4: Training times (until convergence) and memory consumptions on MNIST.
AVB	MFVI	LIVI	DE	MNF
Training time (minutes)	64.1	10.2	12.3	15.6	18.7
Memory (GBs)	2.63	0.81	0.93	0.60	1.77
Kernel Density Estimate(single weight)
Figure F.3: Multimodal densities over weights using KDE
F.4 KDE plot
Figure F.3 shows a KDE plot of weights randomly chosen from samples obtained from a trained generator used for the toy dataset, see section 6.2.
F.5 Testing the entropy approximation
Because our proposal relies on approximating the entropy of an intractable variational approximation through linearisation, we test our approximation in isolation against an importance-sampled estimate of the entropy of a VAE's generative model.
$$
\begin{array}{l} - \mathbb {E} _ {q _ {\gamma} (\theta)} [ \log q _ {\gamma} (\theta) ] = - \mathbb {E} _ {q (z)} \mathbb {E} _ {q _ {\gamma} (\theta | z)} \left[ \log \int_ {z} q _ {\gamma} (\theta | z) p (z) d z \right] (F.1) \\ = - \mathbb {E} _ {q (z)} \mathbb {E} _ {q _ {\gamma} (\theta | z)} \left[ \log \mathbb {E} _ {\tilde {q} (z | \theta)} \left(\frac {q _ {\gamma} (\theta | z) q (z)}{\tilde {q} (z | \theta)}\right) \right] (F.2) \\ \end{array}
$$
For Fig. F.4, we estimate the entropy of a VAE decoder with a very small output variance using importance sampling as the ground truth. We choose a small output variance as this emulates our generator Section 3 which is, by construction very similar to a Gaussian VAE. Noteworthy here, the implicit variational approximation presented Section 3 does not involve an encoder architecture that parameterises the reverse distribution  $\tilde{q}(z|\theta)$  and is only part of this experiment. For this analysis, the VAE was trained on HMC samples from a small BNN trained on toy sinusoidal data similar to Fig. 1. As such this VAE's decoder learns the generative distribution over parameters of a BNN very similar to the generator that parameterises the implicit variational approximation presented in this paper. The architecture of this VAE is also the same as the generator architectures used for experiments. The results here depict that even though the entropy estimated by the importance sampled ground truth and our bound in Eq. (17) are numerically different, our approximation Eq. (15) and the smallest singular value lower bound do retain and follow the characteristics of the ground truth making them viable estimators of this quantity.
Figure F.4: Comparisons of different entropy approximations for decoder/generator through iterations. Refer to the paper for Eqs. (15) and (17) for the different entropy estimators we propose. The importance sampled estimate is presented here.
Scaling our method to models with hundreds of millions of parameters is an important line of future research. For this, one could consider independent hypernetworks for each layer of the BNN, 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).
We also leave for the future applications of this implicit approximation to sequential learning problems like continual learning(see section 4.4 of Daxberger et al., 2021a) and also the several possible extensions of Eq. (18). We use a very simple yet potent individual singular value to approximate a log det term but expect that this term can be better approximated minimising information loss from the Jacobian matrices without compromising computational efficiency.