Implicit Variational Inference for High-Dimensional Posteriors
Authors: Anshuk Uppal
Kristoffer Stensbo-Smidt
Wouter Boomsma
Jes Frellsen
Venue: NeurIPS 2023 spotlight (2023)
Keywords: Implicit models
Variational Inference
Bayesian Deep Learning
Bayesian Inference
Generative Modelling
TL;DR: High dimensional expressive variational approximations that perform better than existing methods.

Abstract: In variational inference, the benefits of Bayesian models rely on accurately capturing the true posterior distribution. We propose using neural samplers that specify implicit distributions, which are well-suited for approximating complex multimodal and correlated posteriors in high-dimensional spaces. Our approach introduces novel bounds for approximate inference using implicit distributions by locally linearising the neural sampler. This is distinct from existing methods that rely on additional discriminator networks and unstable adversarial objectives. Furthermore, we present a new sampler architecture that, for the first time, enables implicit distributions over tens of millions of latent variables, addressing computational concerns by using differentiable numerical approximations. We empirically show that our method is capable of recovering correlations across layers in large Bayesian neural networks, a property that is crucial for a network's performance but notoriously challenging to achieve. To the best of our knowledge, no other method has been shown to accomplish this task for such large models. Through experiments in downstream tasks, we demonstrate that our expressive posteriors outperform state-of-the-art uncertainty quantification methods, validating the effectiveness of our training algorithm and the quality of the learned implicit approximation.

Implicit Variational Inference for High-Dimensional Posteriors
Anshuk Uppal
Kristoffer Stensho-Smidt
Technical University of Denmark ansupedtu.cm dk
Technical University of Denmark krssedtu.dk
Wouter Boomsma
Jes Frellsen
University of Copenhagen wb0di. ku. dk
Technical University of Denmark jefr@dtu.dk
Abstract
In variational inference, the benefits of Bayesian models rely on accurately cap- turing the true posterior distribution. We propose using neural samplers that specify implicit distributions, which are well-suited for approximating complex multimodal and correlated posteriors high-dimensional spaces. Our approach advances inference using implicit distributions by introducing novel approximate bounds by locally linearising the neural sampler. This is distinct from existing methods that rely additional discriminator networks and unstable adversarial objectives. Furthermore, we present a new sampler architecture that, for the first time, enables implicit distributions over tens of millions of latent variables, ad- dressing computational concerns by using differentiable numerical approximations. Our empirical analysis indicates method is capable of recovering correlations across layers in large Bayesian neural networks, a property that is crucial for a network's 's performance but notoriously challenging achieve, To the best of knowledge, no other method has been shown to accomplish this task for such large models. Through experiments on downstream tasks, we demonstrate that our expressive posteriors outperform state-of tfetthrartt uncertainty quantification methods, validating the effectiveness of our training algorithm and the quality of the learned implicit distribution.
Introduction
In the Bayesian approach to statistical inference, information about the data is captured by the posterior distribution over a model's latent variables. By marginalising (averaging) over the variables weighted by the posterior, Bayesian inference can provide excellent generalisation and model calibration. This is particularly compelling for complex models, such as modern deep models, where extreme overparametrisation makes overfitting inevitable unless precaution is taken (Wilson et al., 2020).
Exact Bayesian inference is, however, intractable for all but the simplest models. Practitioners, therefore, rely on approximate inference frameworks, which can be broadly categorised as either sampling-based approaches or variational inference (VI, Saul et al., 1996; Peterson, 1987), the latter being the preferred approach for high-dimensional problems. VI seeks a tractable approximation to the intractable posterior, limiting us simple approximations like the isotropic Gaussian, which are often far too crude to satisfactorily cover the posterior, either only narrowly covering a single posterior mode or placing excessive mass in low-probability regions of the posterior (Foong et al., 2019, 2020).
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Mean-field VI
AVB
LIVI
HMC
-2
-4
0
1
-1
0
1
-1
1
-1
0
1
Mean prediction
Epistemic uncertainty
Epistemic alcatoric
Figure 1: A Bayesian neural network trained on small regression dataset using four different approximation methods: mean-field variational inference (VI), Adversarial Variational Bayes (AVB, Mescheder et al., 2017a), our proposed model (LIVI), and Hamiltonian Monte Carlo (HMC). The dashed line represents the true underlying function from which the training data (black points) were sampled. For each model, we show the mean prediction together with the epistemic uncertainty (dark shaded region) and the combined epistemic and aleatoric uncertainties (light shaded region).
To achieve more expressive approximations, implicit variational inference (IVI) defines distributions implicitly by transforming samples from a simple distribution. This provides much greater flexibility the cost of not having tractable likelihood for the posterior approximation (Diggle al., 1984; Mohamed et al., 2016) and thus often requires unstable adversarial objectives based on density ratio estimation, which are not easy scale to high dimensions (Sugiyama et al., 2012; Huszár, 2017).
The different approaches to posterior approximation can imagined lying on spectrum illustrated in Fig. 1. On the left, the mean-field VI approximation is much too confident in regions without data. Sampling- hampling-based approaches, here represented by Hamiltonian Monte Carlo (HMC) on the right, converge to the true posterior given enough compute, which, however, is typically prohibitively large. Implicit VI approaches, here represented by our proposed approximation, LIVI (Section 4), and Adversarial Variational Bayes (AVB, Mescheder et al., 201 which uses discriminator network for density ratio estimation, hit a tractable middle-ground between the two extremes.
LIVI can produce highly flexible implicit densities approximating the posteriors of complex models through general (non-invertible) stochastic transformations. jectives through approximating a lower bound to the model evidence as well novel generator architecture, which results in vastly improved capacity to model expressive distributions, even in millions of dimensions. Specifically, contributions are:
present novel, stable training ob-
We derive a novel approximate lower bound for variational inference using an implicit variational approximation to avoid unstable adversarial objectives.
To keep computational overhead at large scale in check, we further lower bound the approxi- mate lower bound reducing memory requirements and run times significantly.
We devise a novel generator architecture capable of specifying highly correlated distributions in high dimensions whilst using a very manageable number of weights and biases, necessity for modelling high-dimensional distribution.
The paper is organised as follows: Section 2 introduces the background on VI for Bayesian models in general, and VI with implicit distributions (IVI) in particular. Section introduces our proposals to tackle well-known challenges in IVI. Integrating these proposals in Section 4 yields our novel bounds for implicit VI, and Section 5 places our paper the context of the published literature. In Section we experimentally evaluate our proposed method before summarising and concluding in Section 7.
2 Variational inference for Bayesian models
Consider the supervised learning setting, where we have a training set D = {(x,, y.) X {æ, are the covariates (inputs) and Y "y are the labels (outputs), which we wish to fit using a Bayesian model parametrised by E Imm. That is, assuming a prior p(0) over the
parameters and a likelihood p(D) ()), we wish to compute the posterior p(0) (D) = p(D) ()p(0) As
x(D)

---

computing this posterior is intractable for all but the simplest models, we turn to the framework of variational inference (VI) which finds the best matching tractable variational distribution, 9-y (0) parametrised by Y to the true posterior, p(0 (D). A sensible measure of their closeness is the (intractable) Kullback-Leibler (KI.) divergence between the two distributions defined as
DKL [g- (0) || p D)] (D) (0) [logr(0 D) log 97 (0)
(I)
((>)
Since the evidence, pDD), does not depend on y, minimising the above divergence is equivalent to maximising [(y). referred to as the evidence lower bound (ELBO), which then becomes the objective function in VI. It can be further expanded into a likelihood term and a regularisation term as
[(>)= (0) [log
DRI. (,, p(0)]
likelihood term
regularisation term
where the likelihood term encourages the variational approximation to model the data well, and the regularisation term tries to keep the posterior close to the prior.
2.1 Implicit variational inference
In implicit VI (IVI), the variational distribution is only implicitly defined through its generative process which allows sampling from this distribution, but neither its density nor the gradients of its density can be evaluated. One way to draw from an implicit distribution is to push low-dimensional noise vector through a neural network (Mohamed et al., 2016; Tran al., 2017; Shi al., 2018). These models have also been called amortised neural samplers (Kumar et al., 2019), with the generative process z (z) and g-((2) and the corresponding density
4-
((2) dz,
00, 00.
00m
where ((z) is fixed base distribution and 9y Rd R is non-linear, typically non-invertible mapping due to which this integral non-trivial to solve. The resulting intractability of the density precludes the usage of popular objectives in probabilistic modelling. The likelihood term from Eq. (2) and its gradients can be estimated using Monte Carlo and the reparameterization trick (Kingma et al., 2014). However, the regularisation term involves the entropy of 9y, Hy[47]:
DkL. (o, (0) [(0] = -Ea (0) log (0)
p00) (0) (0) [log (0) (0) [()] (4)
which following the intractability of the density Eq. (3) is not explicitly available. Moreover, critical to consider, 4y is not well-defined density in the parameter space, as its support lies on a low dimensional manifold and has measure zero. Consequently, the KL divergence in Eq. (1) is not always well-defined (Arjovsky et al., 2017) giving rise to another noteworthy issue with IVI. To tackle this issue, standard IVI approaches (Sugiyama al., 2012; Huszár, 2017) rely on density ratio estimators based on discriminator (in GANs) to estimate the regularisation term. Geng et al. (2021) give tractable and differentiable lower bound on this entropy.
A deep latent variable model and its entropy
In this section, we address the issues noted above, I: The implicit density lies low dimensional manifold, leading to ill-defined KL (Eq. (1)) and II: The entropy of the implicit density and its gradients are intractable.
3.1 Defining a proper KI. divergence
To make the KL Eq. (1) well-defined, we add Gaussian noise to the output of our neural sampler 9r which has the effect of making the implicit density continuous in the ambient space of 0, see Arjovsky al. (Lemma 1, 2017). Accordingly, we switch to the corresponding variational distribution Gaussian deep latent variable model (DLVM). This is special case of the semi-implicit distribution
3

---

(Yin et al., 2018; Titsias et al., 2019) and equivalent to the generative model in a Gaussian variational autoencoder (VAE, Kingma et al., 2014; Rezende et al., 2014):
(-( = ( 2) 2)9(z) dz = Bz-a(=) 14 (0) =)]. where,
(5)
( z)) - ( gy Im) 9y -R--"
(6)
O I and the latent variable = E R. We assume Gaussian base density q(z) = (x 0, a.). g-y is the neural sampler (or generator) and 0² E Rt is the fixed, homoscedastic variance of the output density, which we take to be small. Thus, each sample from the base density, z, results in a Gaussian output density, (-(0)) (0]).. And the marginalisation leads to a very flexible g, (0) which informally, is infinite mixture of Gaussians (example 2, Im et al., 2017). Generally, we do not have a closed form for 9- (0) and consequently H [gy], due to the non-linear function By within 9- 2) in Eq. (5), so we still left with challenge II, which we address next.
3.2 Approximating the intractable entropy
With the DLVM, we can expand the entropy of the variational approximation in Eq. (4) as,
H H[g((0) = - -x))(()) (er [ [log (0)]
In principle, a Monte Carlo estimator of Eq. (5) could be used to approximate the qr (0) term and by extension, this entropy, but this estimator is biased and has high variance. This variance can be reduced with directed sampling via an encoder, but to avoid training an additional network, we derive an encoder-free approximation in the following. To do this, we need to work around the non-linearity of g-y that causes the intractability.
Lincarisation of the generator First, to estimate (0) during training with lower variance than direct Monte Carlo estimate, we consider an analytical approximation of Eq. (7) obtained via a local linearisation of the generator/neural sampler around 2':
(8) where (=') E Rmxxd is the Jacobian of By evaluated in z', which we assume exists. We can approximate (2) by the linear function T T (z) when z is close to x' and obtain a Gaussian approximation of the output density, 9- z) R ( 2) = | (z), I and substitute
(z) = (2') + Jz(2') (z -==22,, x²),
this into Eq. (5) to approximate q- (0):
9y (0) = Ez -(=) 10 =]] ~ [ (0- (0) =)]
Et (N (x')) J (z (z x²),0²1
(10)
Now the expectation in Eq. (10) can be analytically solved by integrating over the latent variable (Tipping al., 1999) and is Gaussian, which gives the closed form of the approximation
(0) 2 NO My (z'), Cz())) = (0),
(11)
where
(²(2') - (x(=))==(()))-/ Jz 2z))
_,')=,(()') + +
(12)
Approximation of the differential entropy the DLVM,
We use the above result to approximate the entropy of
(()+()) - FZzg(( Eg (z) 1z) (logg - Fargg, x) [log (lo (0)),
(13)
where log (0) Eq. (11) evaluated the samples from the outer expectation. Importantly, we do the linearisation of q- (0) around the latent value that is used to sample each 0 in the expectation, which means that for small ---vulues, is close to z', and the linear approximation will be good. This amounts to using a point-wise Gaussian to approximate the entropy at z in the training objective, whilst sampling from the DLVM noted in Eq. (5). Hence, important to consider here, that the linearisation does not restrict the expressivity of the posterior that we finally obtain.
We could, in principle, evaluate the approximation in Eq. (13) numerically, but it involves the matrix inverse of C'x²',, which is computationally demanding. In Appendix A. we show that for small values of the output variance o², we can further approximate the differential entropy by
H (0)] ~ ) log det (J)())J((z) + Im1m)
m + m log 2T
(14)
2
2
4

---

4 Linearised Implicit Variational Inference
Finally, we arrive at our novel approximation to the ELBO for IVI using Eq. (5) as the variational distribution. Using the entropy approximation from Eq. (14), we obtain the approximate ELBO denoted the Linearised Implicit Variational Inference (LIVI) bound
C' (1) = E0 (D) + log p(0) + Esmo(e) log det (J)()) +
+
(15)
where c= m m log 2T. As the determinant of a matrix is the product of its singular values, the log determinant is the sum of the log of these singular values. If sa(z) ... 2 (2) are the non-zero singular values of the Jacobian Jy (z), we can write the log determinant term as
2 7
1 log det det(J, Im 2 Σ 2
1 d
+ +2) m d log 0²
(z)
(16)
2
To avoid calculating all the singular values of large Jacobian matrix, we can follow Geng al. (2021, Eq. 10) and lower-bound Eq. (16) as
1 d
log(s ( + + + m d log o² 2 log(sf(z) + +00++ 2 logo²
(17)
2
We defer the details of Eq. (16) and the algorithm (LOBPCG, Knyazev, 2001) for obtaining 81 to Appendix C. This allows us to extend our variational approximation to high-dimensional parameter spaces, and using Eq. (17) we obtain a lower bound on (T)) given by
(7) = -Egnq,(0) [log (0) + logp(0) + Exng(*)
log log(s{(z) + 02) + (18)
2
We denote L'(7) the approximate LIVI bound with whole Jacobians and ("( the LIVI bound with a differentiable lower bound on the determinant. We draw comparisons between how well these bounds estimate the entropy of DLVM model in Appendix F.5, and write the reparameterised version of both bounds in Appendix B. The relationships between the bounds and the evidence are
log ((D) 2 (()) C'(r) 2 (7)
(19)
Based on the compute budget, the two LIVI bounds offer an accuracy vs compute trade-off. both cases, entropy maximisation promotes the exploration of the latent space. It is also worthwhile to note that Eq. (15), which uses the entropy approximation in Eq. (13), is fundamentally different from other linearisation approaches in the literature (e.g., Immer et al., 2021), model. Instead, our approach linearises the generator parameterising (0 | only when needed to approximate the differential entropy, thus leaving the variational approximation unchanged.
do not linearise the Bayesian
Related Work
The usage of secondary model generate parameters of primary model first appeared in the form of hypernetworks (Ha et al., 2017). Our approach is probabilistic and is hence closer to Bayesian hypernetworks (Krueger al., 2017) that mandate invertibility of the generator. The formulation is exactly like normalising flows and thereby sidesteps the complexities of estimating the entropy term. Flow-based approximations require particular focus on the design of the variational approximation to curb the dimensionality of the flow and the Jacobian matrices. Louizos et al. (2017) use an expressive flow on the multiplicative factors of the weights in each layer and not on all weights jointly. Our method does not necessitate invertibility, making it more general. We also share parallels with works manifold learning (Brehmer al., 2020; Caterini et al., 2021) that they also use noninvertible mappings and have to employ general form of the standard change of variable formula log (=) (z)) in the approximating density. Kristiadi et al. (2022) infers an expressive distribution over subset of weights by using post hoc Laplace approximation as the base distribution for a normalising flow. This approach, however, inherits the usual scalability limitations of flows.
5

---

Broadly, works in implicit VI have proposed novel ways of estimating the ratio of the variational approximation to the prior (regularisation-term) usually referred as density- estimation (also referred to as the prior-contrastive formulation by Huszár, 2017). To the best of our knowledge, we are the first to propose and test an entropy approximation for an implicit density used as 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 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 single discriminator network is inaccurate estimating log ratios compared to the analytical form of Bayes by backprop (Blundell et al., 2015), and opt to 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 nigh-dimensional density ratios especially given limited number of samples from both the densities as well as the RBF kernel. Pradier al. (2018) are also motivated by the possibility of compressing the posterior in dower-dimensiee 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 ELBOSs only contains the explicit conditional 9- (-(( z). Their novelties lie in using 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.
space and consider the parameters
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 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- but only when is used to estimate the intractable entropy and its gradients, see Eqs. (10) to (12). Thus, the resulting posterior approximation is, still highly flexible (i.e., non-linearised) implicit distribution.
Experiments
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 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 improved predictive performance and uncertainty estimates.
6.1 Experimental setup
Model definition Briefly, BNNs are standard neural networks with prior on their parameters, which, post-training, results a posterior distribution. In a notation similar to that of Section 2, BNNs can be presented as p(0 (x, v)) x ply (p()), where 0 represents the parameters of

---

neural network denoted by f. The likelihood, p(y fe), depends on the given task. In this section, we test proxies the true posterior p(0) (z. y)) for how well they model the dataset, and 2) model regions away from the dataset.
For regression tasks, we Gaussian likelihood and learn the homoscedastic variance, nf² using type-ll maximum likelihood, p(y fx)))) (y
(z), v/²). For classification tasks, we use a categorical likelihood function for the labels with logits produced by fo.
Generator architecture To efficiently generate millions of correlated parameters, introduce a novel architecture for the generator g- (r(z). We sample a low-dimensional noise matrix, which is gradually scaled up through the generator network using matrix multiplications (Shi al., 201 8).. 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 series of disconnected matrix multiplication networks that produce parameters for the individual layers of the BNN. This equates 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 hypernetwork, (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), use in every experiment. We compare our method with other scalable uncertainty quantification methods: Adversarial Variational Bayes (AVB, Mescheder et al., 20)7a), 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 stronger baseline. Where possible, we compare against Kernel Implicit Variational Inference (KIVI, Shi et al., 2018); 1);; 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 LIVI. Where feasible, we compare with HMC as 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 nuttof-distribution (OOD) benchmarks. Section 6.4 shows results on MNIST, then Section 6.5 considers images from CIFAR1O. allowing us to perform OOD testing of the posterior approximations in millions of dimensions. In Section 6.6, we test on CIFAR 100, evaluating our approach on even larger dimensional BNN posteriors' We discuss the metrics used in Appendix E.
6.2 Toy data
In Fig. 1, we compare posterior predictives 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 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 MIP as the BNN for all datasets. Due 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 test L', Eq. (15), the objective function and
SSuportinng code is available here htttt:://gttthue www.:au/ooooooupaaaanaaaaaaa httpaaallaauut/LLLLnooooooo
7

---

Table I: UCI regression datasets. We report RMSE (1) and log-likelihoods (1) 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
(C))
2.32 0.07
4.24 0.17
0,41 0.27
0.03 0.00
0.00 = 0.00
LIVI
2.40
4.13
0.14
0.08 0.01
0.00 = 0.01
HMC
2.2000000
4.27 0.00
0.38 + 0.00
0.04 0.00
0.00
DE
3.28 1.00
6.55
2.09 + 0.29
0.09 = 0.00
0.00
KIVI
0.10
4.70
0.47
0.08 0.00
0.00
LIVI (C))
-.016
0.11
-.11
+.04
6.74
Test LL
LIVI (C")
0.09
0.13
-..3
0.01
5.84 0.06
HMC
0.00
0.00
0.00
1.27
7.79
DE
-2.41 0.25
0..0
1.28 + 0.02
5.93 0.05
KIVI
0.10
0.00
-1.30
1.16
5,50
Table 2: OOD Test M-C1:
We report out-of-distributinn performance for models trained on MNIST (and evaluated FMNIST, KMNIST and EMNIST) CIFARIO (and evaluated on SVHN, LSUN and CIFAR100). The metrics have been averaged across the OOD datasets.
Confidence 4
AUROC
Method
MNIST
CIFARIO
MNIST
CIFARIO
MAP
72.10 0.36
0.16
9.22
86.73
MFVI
69.23 + 0.24
0.223
96.53
87.50 0.25
LLLA
67.40
53.6 0.3
96.67.0.22
89.03 +
DE
6.11
6..1
9..52
89.61 0.11
LIVI
5.10
4.247
97.91
91.83 0.41
AVB
70.68 # 0.45
NA 95.5 0.4
NA
to analyse the performance difference with C", 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, L' and 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 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 44 000 parameters, necessitating the use of our second objective L" Eq. (18). We train on MNIST, achieving an in- distribution test accuracy of 99.1% % for all methods except MFVI, which achieves 98.6%.
OOD Test M-C1 In this experiment, we consider three OOD datasets, FMNIST, KMNIST, EMNIST and Table reports the averaged results. We outperform all the methods 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.
We compare our method on another OOD benchmark by rotating the MNIST image by different angles (Daxberger et al., 202 For this benchmark. plot the negative ligeliossss 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.
OOD Test M2
OOD 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, well-calibrated model should predict a uniform distribution for never-before-seen data. So, with model trained on MNIST we calculate predictive entropies on OOD datasets, plotting each method's 's fraction of high-entropy predictions Fig. In other words, the smaller the area between the curve and the dashed line, the better the method. We find that LIVI consistently produces

---

NLL on R-MNIST
ECE on R-MNIST
7.5
LIVI
DE
5.0
LLLA
0.4
MFVI
2.5
AVB
0.2
MAP
0.0
60
120
180
0.0
120
150
180
Rotation angle
Rotation angle
Figure 2: OOD Test M2: Rotated MNIST benchmark. LIVI performs well as, better, than competing methods.
CDF counts over whole datasiet
EMNIST
FMNIST
KMNIST
1.00
0.75
0.50
MFVI
Uniform
0.00
Entropy
Entropy
Entropy
Figure 3: OOD Test M3: Empirical CDF plot. On the z-axis the entropy of a single prediction and empirical CDF the y-axis. With this plot, count the number of high entropy (uniform distribution) predictions that each method makes on OOD data. As we count over the whole dataset, would like encounter high entropy predictions closer the value 2.30 on the æ-axxis
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 CIFARIO
To test approach at larger scale, consider the CIFAR10 image classification dataset. Here, for the BNN we use the same WedeResNet architecture as Daxberger et al. (2021a) with 16 layers and widen factor of 4 (Zagoruyko al., 2016). This network architecture contains over 2.7 million trainable parameters, meaning that we again use our second objective L" Eq. (18). We generate all parameters with our novel hypernetwork architecture, described in Appendix in Appendix E.6. All the methods reached very similar accuracies of ~ 91.5% on the in-distribution test when trained for 200 epochs, except for MFVI, which reached 88.2%.
with training details
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 of OOD datasets contains SVHN, LSUN, and CIFAR100. LIVI again outperforms competitors by being less overconfident about OOD predictions.
OOD Test C2 We test the models' performances when presented with corrupted images from CIFAR-10-6 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 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.c., decrease at a similar rate), even as the corruption intensity increases.
6.6 CIFARI00
Finally, we test our approach CIFAR100 using a WideResNee((8.100 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

---

NLL with CIFAR-1O-C
ECE with CIFAR-10-C
LIVI
0.5
DE
LLLA
MFVI
0.3
MAP
0.2
Corruption Intensity
Corruption Intensity
Figure 4:
OOD Test C2: Corrupted CIFAR10 benchmark. OOD performance for methods trained on CIFARIO and making predictions for CIFAR-IO1 -O-C images corrupted with Gaussian blur (Hendrycks et al., 2019). LIVI performs well or better than competitors.
Table 3: CIFAR100 results. Posterior testing with a WideResNee(28.I) architecture on the CI- FARI00 dataset. The numbers reported here are on the test set, with those for LIVI being averaged across three seeds.
Method
Accuracy (%) f
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
Results from Nado et al. (2021).
as many (variational) parameters as the BNN latent variables the same amount of variational parameters required for mean-field Gaussian posterior.
We compare our results with those presented by Nado et al. (2021), who use the same BNN architec- ture, Table This architecture is also used by Sharma al. (2023), who 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 combined 146 million parameters, nearly twice the roughly 75 million parameters required by LIVI. Moreover, LIVI's posterior approximation appears better than simpler approximation like MFVI, even though this uses the same amount of parameters.
7 Conclusion
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 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 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., 201 4; 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 the works by Karaletsos et al. (2018) and Trinh et al. (2020).
8 Acknowledgements
The authors thank Søren Hauberg and Plerre-Alexandre Mattei for providing their insight multiple times during the development of this project. We acknowledge EuroHPC Joint Undertaking for
10

---

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 (MILSS NNF200C0062606), KSS acknowledges funding from the Novo Nordisk Foundation (NNF200C006561) I). Furthermore, JF was in part funded by the Novo Nordisk Founda- tion (NNF200C0065611, the Innovation Fund Denmark (0175-00014B) and Independent Research Fund Denmark (9131-00082B).
References
Arjovsky, Martin and Leon Bottou (2017). "Towards Principled Methods for Training Generative Adversarial Networks". In: International Conference on Learning Representations.
Bingham, Eli, Jonathan Chen, Martin Jankowiak, Fritz. Obermeyer, Neeraj Pradhan, Theofanis Karaletsos, Rohit Singh, Paul Szerlip, Paul Horsfall, and Noah D. Goodman (2018). "Pyro: Deep Universal Probabilistic Programming". In: Journal of Machine Learning Research. Blundell, Charles, Julien Cornebise, Koray Kavukcuoglu, and Daan Wierstra (2015). "Weight Un- certainty in Neural Network". In: Proceedings of the 32nd International Conference on Machine Learning. Vol. 37. Proceedings of Machine Learning Research, pp. 1613-1622.
Brehmer, Johann and Kyle Crammer (2020). "Flows for simultancous manifold learning and density estimation". In: Advances in Neural Information Processing Systems. Vol. 33, pp. 442-453. Caterini, Anthony L.., Gabriel Loaiza- Ganeem Geoff Pleiss, and John P. Cunningham (2021). "Rect- angular Flows for Manifold Learning". In: Neural Information Processing Systems. Daxberger, Erik, Agustinus Kristiadi, Alexander Immer, Runa Eschenhagen, Matthias Bauer, and Philipp Hennig (2021a). "Laplace Redux- Effortless Bayesian Deep Learning". In: Advances in Neural Information Processing Systems.
Daxberger, Erik, Eric Nalisnick, James U Allingham, Javier Antoran, and Jose Miguel Hernandez- Lobato (2021b). "Bayesian Deep Learning via Subnetwork Inference". In: Proceedings of the 38th International Conference on Machine Learning. Vol. 139, Proceedings of Machine Learning Research. pp. 2510-2521.
Detlefsen, Nicki S., Alison Pouplin, Cilie W. Feldager, Cong Geng, Dimitris Kalatzis, Helene Hauschultz, Miguel González-Duque, Frederik Warburg, Marco Miani, and Søren Hauberg (2021). "StochMan". In: GitHub. Note: http:///ghht..com/uuiiiee... wtwaaaaiinneemmnnnggeeeeeeeeeeee
Diggle, Peter J. and Richard J. Gratton (1984). "Monte Carlo Methods of Inference for Implicit Statistical Models", In: Journal of the Royal Statistcal Society. Series B (Methodologicul) 46.2, pp. 193-227.
Foong, Andrew, David Burt, Yingzhen Li, and Richard Turner (2020). "On the Expressiveness of Approximate Inference in Bayesian Neural Networks". In: Advances in Neural Information
Processing Systems. Vol. 33, pp. 15897-15908.
Foong, Andrew, Yingzhen Li. José Miguel Hemándezz-obate and Richard E Turner (2019). "'In- Between' Uncertainty in Bayesian Neural Networks". In: arXiv preprint arXiv:1906. 11537. Geng, Cong, Jia Wang, Zhiyong Gao, Jes Frellsen, and Søren Hauberg (2021). "Bounds all around: training energy- models with bidirectional bounds". In: Advances in Neural Information Processing Systems. Vol. 34, pp. 19808 19821.
Ha, David, Andrew M. Dai, and Quoc V. Le (2017). "HyperNerworks"" In: International Conference on Learning Reprexentations.
Hendrycks, Dan and Thomas Dietterich (2019). "Benchmarking Neural Network Robustness to Common Corruptions and Perturbations". In: Proceedings of the International Conference on Learning Representations.
Huszár, Ferene (2017). "Variational inference using implicit distributions". In: arXiv preprint 172.0:1702.08235.
Im. Daniel Jiwoong, Sungjin Ahn, Roland Memisevic, and Yoshua Bengio (2017). "Denoising Criterion for Variational Auto Encoding Framework". In: Proceedings of the Thirty-Firss AAAI
Conference on Artificial Intelligence. 1A San Francisco, California, USA: AAAI Press,
pp. 2059-2065.
Immer, Alexander, Maciej Korzepa, and Matthias Bauer (2021). "Improving predictions of Bayesian neural nets via local linearization". In: Proceedings of The 24th International Conference on Artificial Intelligence and Statistics. Ed. by Arindam Banerjee and Kenji Fukumizu. Vol. 130. Proceedings of Machine Learning Research. PMLR, pp. 703-711.

---

Izmailov, Pavel, Wesley J Maddox, Polina Kirichenko, Timur Garipov, Dmitry Vetrov, and Andrew Gordon Wilson (2020). "Suussacce inference for Bayesian deep learning". In: Uncertainty in Artificial Intelligence, pp. 1169-1179. Karaletsos, Theofanis, Peter Dayan, and Zoubin Ghahramani (2018). "Probabilistic Meta- Representations Of Neural Networks". In: arXiv preprint 1.XXiv:810.00000 Kingma, Diederik P and Max Welling (2014). "Auto-encoding variational Bayes". In: International Conference on Learning Representations.
Knyazev, Andrew V. (2001). "Toward the Optimal Preconditioned Eigensolver: Locally Optimal Block Preconditioned Conjugate Gradient Method". In: SIAM J. Sci. Comput. 23, pp. 517-541. Krishnan, Ranganath, Pi Esposito, and Mahesh Subedar (2022). Bayesian- Torch: Bayesian neural network layers for uncertainty estimation. https / gi thub com hanp/IntelLabb //abbb///ayaae bayesian- torch.
Krishnan, Ranganath, Mahesh Subedar, and Omesh Tickoo (2020). "Specifying weight priors in bayesian deep neural networks with empirical bayes". In: Proceedings of the AAAI Conference on Artificial Intelligence. Vol. 34. 04, pp. 4477-4488
Kristiadi, Agustinus, Runa Eschenhagen, and Philipp Hennig (2022). "Posterior Refinement Improves Sample Efficiency in Bayesian Neural Networks". In: Advances in Neural Information Processing Systems. Ed. by S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh. Vol. 35. Curran Associates, Inc.. pp. 30333-30344
Krueger, David, Chin- Huang. Riashat Islam, Ryan Turner, Alexandre Lacoste, and Aaron Courville (2017). "Bayesian hypemetworks" In: arXiv preprint arXiv:1710.04755
Kumar, Rithesh, Sherjil Ozair, Anirudh Goyal, Aaron Courville, and Yoshua Bengio (2019). "Maxi- mum entropy generators for mergy-baseed models". In: arXiv preprint arXiv: /0t::/99000000 0985008588. Iakshminarayanan. Balaji, Alexander Pritzel, and Charles Blundell (2017). "Simple and Scalable Predictive Uncertainty Estimation using Deep Ensembles". In: Advances in Neural Information Processing Systems. Ed. by I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett. Vol. 30.
Louizos, Christos and Max Welling (2017). "Multiplicative normalizing flows for variational bayesian neural networks". " In: International Conference on Machine Learning, pp. MacKay, David J. C. (May 1992). "Bayesian Interpolation". In: Neural Computation 4.3, pp. 415- 447.
Maddox, Wesley J, Pavel Izmailov, Timur Garipov, Dmitry P Vetrov, and Andrew Gordon Wilson (2019). "A Simple Baseline for Bayesian Uncertainty in Deep Learning". In: Advances in Neural Information Processing Systems. Ed. by H. Wallach, H. Larochelle, A. Beygelzimer, F, d' Alché Buc, E. Fox, and R. Garnett. Vol. 32. Curran Associates, Inc. Mescheder, Lars, Sebastian Nowozin, and Andreas Geiger (2017a). "Adversarial Variational Bayes: Unifying Variational Autoencoders and Generative Adversarial Networks". In: International Conference on Machine Learning (ICML).
-(201b)). "The Numerics of GANs". In: Advances in Neural Information Processing Systems. Vol. 30.
Mohamed, Shakir and Balaji Lakshminarayanan (2016). "Learning in implicit generative models".
In: arXiv preprint arXiv: /010::661000488.
Molchanov, Dmitry, Valery Kharitonov, Artem Sobolev, and Dmitry Vetrov (2019). "Doubly Semi- Implicit Variational Inference". In: Proceedings of the Tieety-Secondl International Conference on Artificial Intelligence and Statistics. Vol. 89. Proceedings of Machine Learning Research.
Nado, Zachary, Neil Band, Mark Collier, Josip Djolonga, Michael W. Dusenberry, Sebastian Farquhar, Qixuan Feng, Angelos Filos, Marton Havasi, Rodolphe Jenatton, Ghassen Jerfel, Jeremiah Liu, Zelda Mariet, Jeremy Nixon, Shreyas Padhy, Jie Ren, Tim G. J. Rudner, Faris Sbahi, Yeming Wen, Florian Wenzel, Kevin Murphy, D. Sculley, Balaji Lakshminarayanan Jasper Snock, Yarin Gal, and Dustin Tran (2021). "Uncertainty Basclines: Benchmarks for Uncertainty & Robustness in Deep Learning". In: arXiv preprint arXiv: 0et::220000011..
Osawa, Kazuki, Siddharth Swaroop, Mohammad Emtiyaz E Khan, Anirudh Jain, Runa Eschenhagen, Richard E Turner, and Rio Yokota (2019). "Practical Deep Learning with Bayesian Principles". In: Advances in Neural Information Processing Systems. Vol. 32.
Pawlowski, Nick, Andrew Brock, Matthew CH Lee, Martin Rajchl, and Ben Glocker (2017). "Implicit weight uncertainty in neural networks"" In: arXiv preprint 1 Petersen, K. B. and M. S. Pedersen (2012). The Matrix Cookbook. Tech. rep. Version 20121115. Peterson, Carsten (1987). "A mean field theory learning algorithm for neural networks". In: Complex systems pp. 9019.
12

---

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: Interna- tional 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". Journal of the Royal Statistical Society: Series (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 I 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 Per- spective of Generalization" In: Advances in Neural Information Processing Systems. Vol. 33, pp. 4697-4708.
Yin, Mingzhang and Mingyuan Zhou (2018). "Semi-Implicit Implicit Variational Inference". In: Proceedings of the 35th International Conference on Machine Learning. Vol. 80. Proceedings Machine Learning Research, pp. 1660-5669
Zagoruyko, Sergey and Nikos Komodakis (2016). "Wide Residual Networks". In: Proceedings of the British Machine Vision Conference (BMVC).
13

---

Implicit Variational Inference for High-Dimensional Posteriors
Supplementary material
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 Gaussian DLVM can be approximated by H [[]( that is
H H[ (0) = -B) (0)=)) 4, (0)] = (Ez-g(=)B (0) (oo (0)),
(A.1)
where we use the subscript x' = z to specify that q is a linearised distribution that is linearised at and distinct for each sample from the base distribution q(z). Using the definition of from Eqs. (11) and (12), we have
1 10- =(2')) p(z1)),
log (0)
log 2x
log det C(z')
(A.2)
2
2
=:((,x'))
where, c.f. Eq. (12),
Mz(22) = ((=')==((z'))-- -z
C(z') Jg Jn(2')J,('') + I
(A.3)
Our approximation of the entropy thus
#[((0)] log 2T
2 E Rx det C(z)] Ez-g(()) E (0) ( 2)]
(A.4)
We rewrite the last term of Eq. (A.4) by using Eq. (380) from Petersen al. (2012) solve the inner expectation obtaining
1 (0 ((0--(()) - -m(2))
Ez Ez~g(z) Bo - Bx) (0 =)0(00 = Ez-g(=) (01=)
(A.5)
E Ex)g(z)
- (0r
((z)) ((
(A.6)
E [x)g(*)
(I
(r(( -1 Im)
(A.7)
where, in Eq. (A.7), we substituted ((z) by Jz)
Now, we rewrite the two terms inside the expectation of Eq. (A.7) and apply the limit o - 0, following our small noise assumption in Section For the first term inside the expectation, we get
(=)) (J,(2)2)) (J,(2)J](z) + o ImI (Jz(2)z)
(A.8)
= -J,(z)))
(z)= (2)0 Im) Jz(2)z
(A.9)
= 1
21 + Im) 15(2)0-211)
(A.10)
72
- 11. [[())
(A.11)
=z
Ia -2-02 J1
(A.12)
14

---

where we obtained Eq. (A.11) using Eq. (159) by Petersen et al. (2012). Now, taking the limit of Eq. (A.12) as + obtain
lim 2T (o
ZTz. =Tz
(A.13)
For the second term in Eq. (A.7), we use some basic identities for matrix traces and inverse. Namely, constant term inside trace be moved outside, and the trace of matrix is the of its eigenvalues (Petersen et al., 2012, Eq. 12). Furthermore, eigenvalues are reciprocated when the corresponding matrix is inverted (Petersen et al., 2012, Eq. 287). We have that
tr( O2 = tr (Jg(z)
+ + Im
(A.14)
d
d
= 0²
Σ
Σ
1
(A.15)
o
0²
2
0²
d+1
-=1
d+1
2
(m- d) +
(A.16)
i=1 02
In Eq. (A.15), we use the squared singular values of [z(2) as per singular value decomposition (Petersen et al., 2012, Eq. 292). As before, we consider the limit Eq. (A.16) as a -> and obtain
lim tr(C(z) 02 m d. o=0
(A.17)
For small values of the output variance a², we approximate the last term of A [[(]] Eq. (A.4) by
its limit + O. Using the limits from Eqs. (A.13) and (A.17), have
E (0 (0, (, =)] ~ lim [Z2 Bo (0) [h(0, =]]
(A.18)
o=0
E Ex)((t)
(2 + m - d]
(A.19)
E g(]
[zTz] + d
+ m d
m
(A.20)
2
2
2
Applying the approximation in Eq. (A.20) Eq. (A.4), gives the final approximation
H ~ m +
m log 2TT + B Bz- log det (=)( +
):
(A.21)
2 2
which also stated in Eq. (14). It is important note here that every Jacobian (w.r.t. input =) matrix for generator map By is function of y, the variational parameters, and hence needs to be differentiable. To obtain whole Jacobian matrices, we use the nnj wrapper from the StochMan library (Detlefsen et al., 2021), which gives Jacobians in forward pass rather than requiring for loop over backward passes.
B Reparameterisation of the LIVI bounds
We note that it is trivial to reparameterise the LIVI bounds as the Gaussian DLVM is straightforward to reparameterise. Using the base variables z, 17, the generative process for Gaussian DLVM, c.f. Appendix D, can be reparameterise as
= (x),,, 0 = 0' + n,
N(O,
(B.1)
(0, Im))
(B.2)
The reparameterised LIVI with whole Jacobians, corresponding Eq. (15), is
C'(7) = E. (z=g(())
logr(D) (x(22) + n) + logp(17(z)) + +n)
(B.3)
+ log det (Jz())],(2) + + ) tc.
15

---

and the reparameterised LIVI bound with a differentiable lower bound on the determinant, conr sponding to Eq. (18),
log (y(2) + + (ogp(g><)) + n) + d log(ss(z) + +0))
(7) Ex-g(t),
(B.4)
C
Singular values explained
Here we write out the steps that lead from the entropy approximation in Eq. (15) to the efficient entropy approximation we actually use Eq. (16). The first step is to note that the determinant of square matrix can be written as the product of its eigenvalues and then the log determinant by extension can be written as the sum of the log of these eigenvalues. Let 81 (z)..... 8m 5m(2) be the
singular values of the Jacobian J (2)T (including zeros), using Petersen et al. (Eq.. 286, 2012) we can write the log determinant term in Eq. (15) as
det J Jg) = II s (2)
(C.1)
Strictly speaking, the matrix J JT is not full rank and has only d nen-zero singular values, and thus this determinant will be equal to Ó. The additional noise added at the output of the generator resolves this issue as appears in the (z)) in Eq. (A.1).
dJt + Im2
= i=1
II
8f (x))) + 0²] o 2(m--))
(C.2)
log det JJ + + 0²Imm)
log (5(2)) + +02] + +2(m d) > 2m1
(C.3)
As we do not wish to calculate all the singular values to estimate Eq. (C.3) exactly, because this would be computationally expensive, we resort to using lower bound in Eq. (17). The terms in the summation in Eq. (C.3) are positive, so estimating this with any one singular value should give us lower bound.
Depending on what singular value we wish to compute, we can use the LOBPCG algorithm, which iteratively optimises the generalised Rayleigh quotient and converges in linear time, All details about the convergence results are available Knyazev (2001). We also employ an efficient backward mode to get vector- jacobaaa products to work with the generalised Rayleigh quotient, and thus we never need to store large Jacobian matrices during the calculation of this singular value. The Rayleigh quotient is given as follows:
JJJJJy
p(u)=
(C.4)
TTv
D Novel generator architecture
Here we describe the novel generator architecture. A central property is that we drop dense con- nections deeper into the network. This resulted in an architecture that had fully connected matrix multiplication layers at the start and smaller individual networks deeper into the architecture. In other words, to decrease the number of parameters closer to the output of the network, dense connections pruned away. The final output vector of parameters is generated via the concatenation of outputs from smaller networks, which can be thought of as producing weights for each layer in the BNN.
The architecture is pictorially represented in Fig. D.I. All computational blocks, as well as the input block in this network, are a matrices, as we follow the standard recipe the Matrix multiplication neural network (Shi et al., 2018). Starting from the left, we have a large noise vector represented in matrix format. We use the -5x65-dimensional input noise matrix for MNIST experiments and 130x 130 input noise matrix for the CIFAR10 10 experiments. The noise is sampled from the standard normal distribution. Then we use a matrix multiplication layer to produce a higher dimensional output that is split and passed on to series of disconnected individual matrix multiplication networks
16

---

output
output
noise
MM layer
output
split
output
Figure D.I: Our novel CMMNN architecture.
Table E.1: Dual Networks Setups. With this table we clarify the architectures we use for the primary (hypernetwork) and the secondary (BNN) network for training with all the different datasets used for experiments.
Toy Data
UCI Reg
MNIST
CIFARIO
CIFAR100
Hypernetwork
MLP
MLP
MMNN LeNet
CMMNN
CMMNN
BNN
MLP
MLP
W1deResNee(16.4)
WideResNet(28, 10)
that produce weights and biases for individual layers of the WideResNet BNN. This way the task of capturing correlations within layers is taken care of by the individual sub-networks towards the end and capturing correlations across layers is taken care of by the homogeneous layer at the start of the network.
E
Experiment Details
E.1
Metrics
For regression, we test generalisation by measuring negative log-likelihoods and mean squared error over the test set.
For the classification experiments, we measure the averaged OOD AUROC and confidence (maximum softmax value) (Daxberger et al., 2021a) on OOD datasets. The OOD AUROC and confidence are crucial metrics understand the behaviour of the BNN on unseen data. A low confidence score indicates that the model is very uncertain about a classification, while high AUROC score suggests that the model does not misclassify unseen data with high probability or confidence. Such misclassifications would increase the false positive rate, thereby lowering the AUROC. Therefore, these metrics allow us measure the model's generalisation performance on OOD data hinting at the quality of the posterior. We further consider the negative likelikel (NLL) and expected calibration error (ECE) for some OOD classification experiments. These two metrics ensure that the implicit approximation is not placing mass in undesirable regions of the posterior and is learning the right variance in the function space. the NLL is lower for OOD points, it signifies that the functional variance farther from the data increases, which is desirable. At the same time, low ECE indicates that the model's predictive confidence falls with its accuracy as per expectations.
E.2 Experimental setup
Details of the architectures used for the generator/hyyemeetork. and the Bayesian neural network for all of the experiments in Section 6, can be found Table E. For reference, our novel CMMNN architecture implemented CorreMMGeneratoo in our code.
E.3
Toy experiments
use 2D toy sinusoidal dataset for this example with 70 training data points. We intentionally leave gap in the middle to visualise the epistemic uncertainty the diversity of functions in this
17

---

region as this has been reported by Foong et al. (2019) to be tricky or unachievable for simpler variational approximations.
The BNN architecture is shared across all the methods and a two-hiddee-layered MIP with 7 and 10 units and ELU activations trained with losses unique each of the methods. We also assume homoscedastic noise in the data that is estimated by all the methods using an extra parameter learnt via type maximum likelihood. This is set in the data generating process at 0.3.
LIVI: Here, a generator network, which is one-hidden layered MLP with an 80-dimensional noise input and ELU activations, produces a 105-dimensional output vector which then used to reparameterise the layers (weights and biases) of a BNN for training.
For mean-field VI, we take advantage of the sereeian-tordee library (Krishnan et al., 2022). We construct the exact same architecture for the BNN using the reparameterisable layers provided by this library and we slowly warm up the KI. term reaching towards the end of training. This is done to help the method converge.
MFVI:
HMC: We sample 10000 samples after a warmup of 2000 samples using the same and neural network architecture. We use the NUTS sampler available in the Pyro (Bingham al., 2018) package.
We also compare our method qualitatively against the conventional style of training implicit distributions which uses a density ratio estimating discriminator network. Using the same architecture as above for the BNN, we use one-hidden layered discriminator network that we retrain within each training loop of the BNN to discriminate 20 samples from the prior and the posterior (generator) each. We run this inner training loop for 10 epochs. The algorithm here is exactly the same as Algorithm 1 in Mescheder et al. (2017a).
AVB:
E.4 Details of UCI experiments
We use an MLP BNN with 50 units one hidden layer. We report the RMSE and liglliill on held-out data for our method. We use generator architectures with far fewer parameters than Shi al. (2018) and do not assume independence across layers, i.e., we use one MLP to generate all the weights of the BNN. All of the generator architectures are one hidden layered MLP with a slightly varying number of units depending on the dataset.
We train our method with homoscedastic assumption, i.c., the variance in the dataset is assumed to be constant and we train an observation noise parameter using type-II maximum likelihood.
E.5 MNIST experiments
Of the wide-ranging approximate inference approaches we compare against, some focus on capturing multiple modes like deep ensembles, while others like MFVI capture information around a single mode sacrificing all correlations for computational feasibility. Similarly, the Laplace approximation, a popular post hoc method (MacKay, 1992), captures unimodal information but can be made much faster by considering a sub-stochastic Bayesian neural network (Daxberger et al., 2021a). Moreover, there is a range of approximations for estimating the hessian that again, trade-off correlations for computational feasibility. In all of these cases, a single MAP estimate remains the easiest baseline to implement and train but represents the minimum amount of information that can be obtained about the posterior surface.
For all methods, we use the same LeNet architecture and for (LLLA, DE MAP) we train with equal weight decay of 8 x 10 For MFVI, we use a standard normal prior.
LIVI For all our MNIST experiments we use an MMNN with 65 65 input noise, one hidden layer of size 250 x 250 and produce an output matrix of size 350 x 127, the elements of which are used as weights and biases the LeNet BNN, For all experiments, we trained without dataset augmentation and with maximum of to 2 samples (from 4- (0)) per minibatch. This network generates close to 44 000 parameters for the LeNet BNN while keeping its own parameters appreciably low compared to an MLP hypernetwork conserving computation required for the objective Eq. (18). The weights
18

---

and biases of the hypernetwork are the variational parameters that tune the output implicit distribution and here they number slightly over six hundred and ten thousand i.e. these are the total number of weights and biases in the hypernetwork.
The full Hessian for the last layer can be obtained very cheaply as the last layer Hessian is simply the Generalised Gauss Newton matrix since the network is linear in the last-layer weights (Immer et al., 2021). Hence, we strike balance between expressivity and computational cost by obtaining the most accurate post-hoc covariance structure over sub-stochastic network. We use pretrained LeNet with weight decay of 8e-4 for this method.
LLLA
To represent deep ensembles we use five LeNet networks trained with different seeds and average over their predictions and use one of these networks to represent the MAP estimate. Each network is trained with a weight decay of 8e-4.
DE
MFVI We use weights from a pre- preiteained LeNet for initialising the mean and retrain it for several more epochs using the ELBO and standard normal prior to obtain an isotropic Gaussian posterior. Again this is done using bayesian-tone torch library and the MOPED functionality (Krishnan et al., 2020).
AVB We again compare our method to the existing adversarial training routine for implicit distribu- tions which is represented by AVB (Algorithm 1, Mescheder et al., 2017 For this method, use an MLP discriminator with two hidden layers containing 2500 and 4500 neurons for classifying if a sample comes from the generator or the prior (regularisation term in Eq. (2)). We use 300 samples from the prior and the posterior (generator) for full batch training of the discriminator for 5 epochs within each minibatch update of the main training loop.
KIVI To compare our method against Shi et al. (2018) we also use MLP BNNs as per their experimental setup and outperform on in-distribution test accuracy reported in Shi al. (2018, Figure 2, left table). To test fairly against KIVI we use MLP BNN architectures only for this experiment as a LeNet would've been easily able achieve better accuracy the in in-distribution test set. We find that with the same BNN and hypernetwork architectures we get better test accuracies on MNIST
Training To train LIVI, we use Matrix Multiplccation Neural Network (MMNN, Shi et al., 2018) 8) as the hypernetwork architecture. For the LLLA,, we use a complete ("full", Daxberger et al., 2021a) Hessian for the weights and biases in the last layer obtaining all correlations over a smaller parameter set to ensure feasibility. For deep ensembles, we average over five networks trained with weight decay. For one-to-one comparison with AVB we use the exact same generator and training hyperparameters as LIVI. To train LIVI with the aforementioned generator architecture takes 12.3 GPU minutes on Titax X (Pacal 12gb vram) GPU and consumes 980 megabytes of memory when training with batch size of 256. On the other hand, to train a LeNet deep neural network with the same batch size on the same GPU it takes 3.07 GPU minutes and 700 megabytes of GPU memory.
E.6 CIFAR1O Experiments
LIVI We use our novel generator architecture for this experiment with a noise dimension of 130x130. 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) elbo 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 initiali- sations using the same optimiser hyperparameters, i.e. with NAdam optimiser with weight decay of 1 x 10-4, a learning rate of 7.8 x 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.
19

---

For the last-layer Laplace approximation, we use network trained similarly to 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 very faithful posterior and keeping the computational requirements in check.
LLLA
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 deep ensemble consisting of five networks with the same training parameters and provides better uncertainty quantification than a non- non-Bayesian ensemble of neural networks. Keeping everything else the same, each network of the ensemble takes 34 minutes to train on CIFARIÓ 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. I, 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-ll maximum likelihood. We deliberately removed a part of the data to see if the models tested were able to find in in-between uncertainties. All methods were given the same sized hidden layered MLP with 7 and 10 units respectively. We trained 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 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, 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 accitectures(feee-forweee forward, convolutions, etc.) require a generator in the init function. As such, the generator networks reside inside the BNN and reparametrise it with simple sample_ samplee_prrmmee www.mmterreesseeeee ters 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 e but these layers do not expose their state i.c. their parameters in manner that allows changes on the fly during training. We reimplemented the layers allowing such resampling to oceur 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 index to know the number of weights used by a previous layer. This counter index is updated by each layer in their init and are trainable, the parameters of the BNN are switchable and relay gradients to the generator via the likelihood or the entropy term.
paramsters function. As such, only the parameters of the generator
20

---

Table F.1: UCI regression datasets. We report RMSE (L) on the test set and average across three different seeds for each model quantify the variance the results.
Method
Boston
Concrete
Energy
Kin8nm
Naval
LIVI (C))
0.07
4.24 0.17
0.41 0.27
0,03 0,00
0.00 0.00
LIVI (C")
2.09
0.13
0.44 0.11
0.08 = 0.04
0.00
HMC
2.06
4.27
0.38 + 0.00
0.00 2 0.00
DE
3,28
0.58
2.09 = 0.29
0.09 = 0.00
0,00 + 0.00
KIVI
0.17
4.10
0.47
0.08 0.00
0.00 0.00
MNF
3.31
+.82
1.04 0.01
0.08 0.01
0.01 0.00
Table F:2: UCI regression datasets. We report log-likelihood (f) the test and average across three different seeds for each model to quantify the variance in the results.
Method
Boston
Concrete
Energy
Kin8nm
Naval
LIVI (C')
-2.16 0.05
-1.17 0.11
1.24 0.04
6.74 0.04
(IV)
0.09
0.13
-.11
1.15 0.01
5,84 + 0.06
HMC
0.00
0.00
-..14
1.27 0.00
7.79 0.00
DE
0.25
-3.06
-.22
1.28
5.93
KIVI
0.10
-3.05
-..3
1.16
5.50
MNF
-.08
-3.24
-1.34
1.10
5.01 + 0.00
F
Additional results
In this section, we report additional results. In particular, we consider multiplicative normalising flows (MNF, Louizos et al., 2017), as well two versions of full Laplace, non-linearised full network Laplace using low-rank approximation to the Hessian (FLAP), and 2) a linearised full network Laplace using Kronecker factored structure of the Generalised Mewus- approximation the Hessian (FLLA).
While show results for linearised Laplace, as this method often achieves state-of-the ahe- perfor- mance, it must be emphasised that, for LIVI, we do nor linearise the posterior; we linearise the neural sampler g-y (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 highly flexible (that is, non-linearised) implicit distribution. Linearised Laplace and LIVI are, therefore, fundamentally different approaches. Whereas lincarised Laplace linearises the BNN and the predictive function, we linearise the neural sampler/hyper-netwonl obtaining the approximation & (0 2), 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.
within the variational distribution 9-y (0 =). thus
F1 Additional results for UCI datasets
In Tables F.1 and F.2, we report additional results for experiments on the UCI datasets. Overall, LIVI (C', Eq. (15)) performs competitively, either outperforming or being on par with HMC, which 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 F33, 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 F3 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
21

---

NLL on R-MNIST
ECE on R-MNIST
7.5
UMI
0.6
5.0
LLLA
0.4
ECE
NU
AVB
2.5
MAP
0.2
FLLA
0.0
30
90
120
150
180
30
60
120
150
180
Rotation angle
Rotation angle*
Figure F.1: Additional results for OOD Test M2: Rotated MNIST benchmark.
Table F.3: Further experiments for OOD Test M-CI. FLLA: full linearised Laplace; FLAP: full
non- non-linearised Laplace; MNF: multiplicative normalising flows.
Confidence
AUROC1
Method
MNIST
CIFAR10
MNIST
CIFAR1O
LLLA
67.40 0.19
53.6 0.3
96.67 0.27
8..0
DE
0.11
0.21
9.052
39.61
MAP
0.31
0.16
9.22
86.73 0.64
LIVI
0.13
43.47 + 0.28
0.271
91.83 0.41
AVB
70.68 0.45
NA
95.5 0.4
NA
MFVI
0.22
74.71 + 0.23
0.15
87.50 0.25
FLLA
0.17
NA
97.57
NA
FLAP
0.32
59.55
NA
MNF
0.27
NA
96.37 0.15
NA
poor samples. The lincarised version (FLLA) works slightly better than LIVI in terms ECE and NLL the rotated-MNIST benchmark (Fig. F.1) and in the OOD) entropy-CDF test (Fig. R.2). As noted before in Section and ??, our approach fundamentally different from linearised Laplace approaches 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 MNIST in Table F.4. The methods were all run the same single-GUU (NVIDIA RTX A4000) machine. Compared to other expressive VI methods like AVB and MNF, LIVI faster 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 amount of parameters required for mean-field VI. Interestingly, requirements for LIVI are higher than that of DE, LIVI trains a shorter amount of time.
though the memory
CDF/ counts over whole dataset
EMNIST
FMNIST
KMNIST
LIVI
LLLA
MPVI
0.25
Uniform
FILAA
0.00
0.0
0.5
1.0
1.5
2.0
0.0
0.5
1.0
1.5
2.0
0.0
0.5
1.0
1.5
2.0
Entropy
Entropy
Entropy
Figure F.2: Additional results for OOD Test M3: Empirical CDF plot.
22

---

Table F4:: Training times (until convergence) and memory consumptions on MNIST.
AVB
MFVI
LIVI
DE
MNF
Training time (minutes)
64.1
10.2
12.3
18.7
Memory (GBs)
2.63
0.81
0.93
0.60
1.77
Kernel Density Estimate(single weight)
030 80 10
Figure F.3: Multimodal densities over weights using KDE
F.4
KDE plot
Figure F.3 shows KDE plot of weights randomly chosen from samples obtained from 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-samplee estimate of the entropy of VAE's generative model.
log S. 1 (01=)(((())
(0) (0) = -By(s) (01=))
(F.I)
[g (0(2))(((())
= BBg By
log Eq(=10)
(FR2
7(=10)
For Fig. F.4, we estimate the entropy of a VAE decoder with a very small output variance using importance sampling the ground truth. We choose small output variance this emulates generator Section which is, by construction very similar to Gaussian VAF. Noteworthy here, the implicit variational approximation presented Section does not involve an encoder architecture that parameterises the reverse distribution q(z(0) 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 Fig. 1. As such this VAE's decoder learns the generative distribution over parameters of a BNN very similar the generator that parameterises the implicit variational approximation presented 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.
23

---

Entropy Estimators
-100
-105
Entropy of the model
-115-
-120-
-130
-135
3.5 VAE training iterations in 1000s
1.5
Figure F4: Comparisons of different entropy approximations for decoder//gnerrtoo 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.
G Future Work
Scaling our method to models with hundreds of millions of parameters is an important line of future research. this, one could consider independent hypernetworks for each layer 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 al. (2018) and Trinh et al. (2020).
We also leave for the future applications of this implicit approximation 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 approximate log det term but expect that this term can be better approximated minimising information loss from the Jacobian matrices without compromising computational efficiency.
24


Reviews:
[{"review_id": "aaDAScNZcW", "reviewer": "Reviewer_MstJ", "summary": "## Post Rebuttal Update \n\nI have engaged with the authors for the rebuttal, and found their responses informative, prompting me to increase my score from a 7 to an 8.\n\n## Original Review\n\nThe paper addresses two issues that are common in implicit variational inference methods such as amortised neural samplers, namely - \n\n1. The implicit density often lies on low dimensional manifolds making the KL infinite/ill-defined\n2. The gradients and entropy of the implicit density are intractable\n\nThe paper addresses these issues by (i) adding Gaussian noise to the output of the neural sampler, making it continuous w.r.t $\\theta$, and (ii) linearising the neural sampler, resulting in a Bayesian linear model approximation to the non-linear model, giving closed-form solutions for the distributions. On different benchmarks the authors show that these methods work well on a range of small scale experiments compared to other state-of-the-art, such as last-layer Laplace and deep ensembles.", "questions": "I think there are some really interesting connections between the obtained linearised implicit variational model, and existing methods such as linearised Laplace and regular variational inference, and I would love for there to be a small discussion about these things potentially. For example, in linearised Laplace, you assume a fully Bayesian posterior over a NN parameterised by $\\theta$, then linearise the model around the MAP estimate, resulting in a tractable Gaussian approximation to the posterior. However, this posterior is more rich compared to a mean-field variational approximation, because the covariance of the posterior is given by $J_{\\theta}^T \\Lambda J_{\\theta} + \\sigma^2 I$, where $\\Lambda$ is the prior over weights, and $\\sigma^2$ is the noise variance. This form of the posterior looks very close to what is obtained by implicit variational inference, where the neural sampler can be considered similar to the NN model in linearised Laplace. I would be really interested in seeing discussions about these potential connections if possible. In fact, optimising the marginal likelihood for a linearised Laplace model is akin to doing ELBO with a Gaussian approximation. I would be really interested in seeing these parallels.\n2. Given that the authors are using the entirety of the Jacobian of their neural sampler, and comparing to only last-layer Laplace, where only the last layer is modelled probabilistically, I would be really interested in seeing if they can run experiments on full linearised Laplace, especially on the small UCI and MNIST datasets, where this should be tractable. In fact, there are methods that perform full probabilistic inference for linearised Laplace using samples from the posterior, such as in [1].\n3. I am also really interested in seeing what approximations to the Jacobian are best for estimating log det (J J^T), such as last-layer only, or through samples using Hutchinson's estimate, or by using more than one singular value and ablating through what the optimal number of singular values to consider is.\n\n[1] Antorán, Javier, et al. \"Sampling-based inference for large linear models, with application to linearised Laplace.\" arXiv preprint arXiv:2210.04994 (2022).", "limitations": "N/A.", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper is well motivated, has clear mathematical development, and solves common issues in implicit variational inference models. I like the use of linearisation to make the entropy and gradients tractable by getting closed-form solutions. The experiments are clear, and the baselines are well-considered. The results are also quite impressive in this domain, as it is often really difficult to beat deep ensembles, and the method seems to consistently perform really well.", "weaknesses": "In general I think the paper is well-written, however I have one major criticism - \n\n1. I'm not convinced that the current approximation to $\\log(J J^T)$ has been well-motivated or properly ablated. It would be nice to see some ablations on smaller-scale problems where calculating the full log determinant is tractable, and comparing it to the approximation made using just the highest singular value. Or doing a sweep over adding subsets of singular values vs the highest. Or even, plotting the eigenspectrum of $J J^T$ to show that it is true that the highest singular value is often much more dominant than the others.\n2. If the paper is using the full-rank Jacobian of the neural sampler in order to estimate the entropy, I think a fairer comparison to make would be against full Laplace, not last-layer Laplace, which should be possible for UCI datasets and MNIST at least, using the Laplace library the authors cite.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your careful evaluation of our paper. We are happy that you liked our proposed method and found the results impressive. We respond to your specific questions and concerns below.\n\n**W1: The lower bound on $\\log\\det(JJ^\\top)$ has not been motivated sufficiently nor properly ablated.**  \nThe motivation for the bound on $\\log\\det(J J^\\top)$ comes from Geng et al. (2021) [1]. To clarify a possible source of confusion, the bound uses the *smallest* eigenvalue, not the largest. Essentially, it simply states that the sum of all log singular values is larger than (or equal to) the smallest singular value times the number of dimensions. While one can view this as a trivial and potentially quite loose bound, the smallest singular value is efficiently computed using the LOBPCG algorithm, and efficiency is very important for the high-dimensional problems we are focusing on in this work.\n\nTo empirically assess the effect of the bound, we now include experiments to compare our two bounds, $\\mathcal{L}'$ and $\\mathcal{L}''$, which demonstrate the effect of the bound on the log determinant. The results, which can be found in tables 1 and 2 of the attached PDF, show that models trained with either $\\mathcal{L}'$ or $\\mathcal{L}''$ perform quite similarly both in terms of test log-likelihood and test RMSE, giving empirical justification for the lower bound on $\\log\\det(J J^\\top)$.\n\n[1] Geng et al., \"Bounds all around: training energy-based models with bidirectional bounds\", NeurIPS 2021.\n\n\n**W2: Full Laplace is a fairer comparison than last-layer Laplace.**  \nThis is a great suggestion. Full Laplace using a full-rank Hessian is not feasible, even on MNIST, due to the size of the BNN that we use in our experiments. However, we have added results for two versions of full Laplace to the paper, 1) a non-linearised full Laplace using a low-rank approximation to the Hessian, and 2) a linearised full Laplace using a KFAC factorisation of a Generalised Gauss-Newton approximation to the Hessian,  see table 3 and figures 1 and 2 in the attached PDF. The non-linearised version of full Laplace underfits due to the quite crude approximation to the posterior over the full network. The linearised version works slightly better in terms of ECE and NLL on the rotated-MNIST benchmark (figure 1) and in the OOD entropy-CDF test (figure 2). Note, however, that linearised Laplace is not directly comparable to our proposed model, see the discussions for Q1 and Q2 below.\n\n\n\n**Q1: It would be nice to see a discussion on the connections between LIVI and existing methods such as linearised Laplace and regular variational inference.**  \nWe would like to clarify that 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. (9) and (12)). The resulting posterior approximation is, therefore, still a highly flexible (non-linearised) implicit distribution. It is correct that the terms appearing in a Laplace posterior and our linearised approximation are similar, but it is hard to draw direct analogies between the two approximate Bayesian methods as they are fundamentally quite different. Whereas linearised Laplace linearises the BNN and the predictive function, we linearise the neural sampler/hyper-network within the variational distribution $q_\\gamma(\\theta)$, thus obtaining the approximation $\\tilde{q}_z(\\theta)$, but only when it is needed to approximate the entropy. This is also the reason for why we do not linearise our BNN for predictions, which is common practice with linearised Laplace. We acknowledge that the distinction was not presented clearly enough in the original manuscript and will update the paper accordingly.\n\n\n**Q2: Please run experiments on full linearised Laplace on the UCI datasets and MNIST.**  \nThank you for this suggestion. We have added results for full Laplace without linearisation (low-rank Hessian) and full Laplace with linearisation (KFAC factorisation of a GGN approximation) to the paper and the PDF attachment here, see table 3 and figures 1 and 2 for results on MNIST.  We tried our best to make the Laplace library work for UCI regression but could not due to errors thrown by the second-order backends used (BackPack and Asdl), we will keep working on this and report back soon.While these are both relevant baselines, we would like to emphasise that, because of the reasons noted above, our results cannot be compared one-to-one with linearised Laplace as that posterior measures the uncertainty of a generalised linear model approximation and not the actual BNN.\n\n\n**Q3: Are other approximations to the Jacobian are better for estimating $\\log \\det (JJ^\\top)$?**  \nThis is a very valid question, although one that is perhaps best answered in future work given the scope of it. We did not consider a last-layer-only approximation, as we have focused on the conventional Bayesian treatment where all latent variables (weights and biases) of the respective BNNs in all experiments have been modelled probabilistically. Basing the bound on more singular values than just the smallest one is a good idea, as it might give us a tighter bound. However, a motivating factor for using just the smallest singular value is that this can be found in linear time using the LOBPCG algorithm and, as we show in tables 1 and 2 in the attached PDF, the bound empirically works well. Using Hutchinson's estimator is a good idea too, however, empirically this estimator provides an upper bound, not a lower bound; please see Geng et al. (2021) who compare the smallest singular value estimate against Hutchinson's estimator for $\\log \\det(J^TJ)$."}, {"review_id": "1jaKaLOZ7J", "reviewer": "Reviewer_C434", "summary": "The paper studies the problem of approximating high-dimensional multi-modal posteriors through neural samplers specifying implicit distributions. While Bayesian methods promise a variety of benefits in terms of generalization and calibrated predictions, in practice they have seen limited success due to intractability of exact Bayesian approaches and the tradeoffs made in approximate Bayesian methods. Implicit Variational Inference provides an alternative to approximating exact Bayesian posteriors by maintaining distributions implicitly by transforming samples from simple distributions, allowing it to admit much richer distributions. Implicit Variational Inference typically requires some density ratio based adversarial objectives, which can fail on high-dimensional problems (e.g. parameters of neural networks). The authors note two major issues in existing approaches a) KL being ill-defined due to the implicit density lying on a low dimensional manifold b) intractability of the entropy of the implicit density and its gradients. To tackle these issues, the authors first introduce Gaussian noise to the output of the sampler making it a Gaussian DLVM resulting in a well-defined KL over the parameter space.  Next the authors approximate the generator with a local linearization resulting in a Gaussian approximation of the output density, and obtain easy to compute approximation of the differential entropy of the output density. This results in a novel approximation to the ELBO based on the entropy approximation, which is scalable to high dimensional parameters. The authors discuss two variants based on computing the whole Jacobian or using a differentiable lower bound on the the determinant, which trade-off compute and quality of the approximation. Finally, the authors evaluate the method on a variety of tasks including impressive results on fairly large BNNs (WideResNet). ", "questions": "* What do you think would be the critical challenges in applying the method to even larger modern networks, e.g. transformers? \n* What is the challenge in implementing the KIVI baseline? (since that is not included in all the experiments as the authors note)\n* The runtime of the method is not explicitly mentioned anywhere in the paper except for the remark relative to Deep Ensembles in the appendix. This is important information which should be included in the main paper. \n* Minor typos: L117 “in the following” -> “in this section” L123 “aGaussian” -> “a Gaussian”", "limitations": "Despite the impressive results in the paper, there still remains a gap between the models studied in the paper and the size of models considered in practice. ", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* The paper studies the important problem of approximating expressive Bayesian posteriors on high-dimensional spaces. Due to the general applicability and promising results, the work is significant and relevant to the community. \n* To the best of my knowledge the main contributions of the paper namely addressing the ill-defined KL, local linearization for the entropy and the LIVI bound on the ELBO, are all novel.\n* Introducing Gaussian noise to the sampler to induce a Gaussian DLVM is a neat and simple way of fixing the KL with minimal additional restrictions to the model\n* Similarly, the local linearization of the neural sampler is a nice idea to obtain a cheaper estimator for the entropy. \n* The experimental results are quite impressive - in particular the results on the WideResNet. I also appreciate the authors including the code with their submission to aid reproducibility of the results. \n* Overall the paper is well-written with a clear exposition of ideas and most relevant details covered. ", "weaknesses": "* The paper proposes a local linearization to make the entropy computation tractable. However, what is not clear to me is how the local-linearization affects the expressivity of the posterior and general performance. The experiments indicate that the effect is not large, but these are still “relatively” simple tasks so a thorough study of this would be useful \n* The empirical comparisons also do not consider alternative approximate Bayesian methods. \n* Some recent work [1] proposes a closely related approach which might be worth discussing in the paper\n\n[1] Posterior Refinement Improves Sample Efficiency in Bayesian Neural Networks. Kristiadi et al., NeurIPS 2022. \n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the insightful review. We very much appreciate that you find our work significant and relevant to the community and that the manuscript is well-written. We answer your specific concerns below.\n\n**W1: How does the local linearisation affect the posterior in terms of expressivity and performance?**  \nThis is a great question. 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. (9) and (12)). However, this could indeed lead to sub-optimality in how we train the implicit approximation and the quality of the resulting distribution. We hope to obtain a highly flexible implicit distribution, and as you point out, our results indicate that the obtained posterior approximation is at least as good as the one found by our competitors.\n\nTo investigate this further, we have added a comparison of our method with HMC on the UCI datasets, see tables 1 and 2 in the attached PDF. The results show that the model trained with $\\mathcal{L}'$ performs close to or as well as HMC, suggesting that the local linearisation does not harm the expressivity dramatically. If the reviewer has suggestions for other experiments assessing the expressivity of the posterior, we would be happy to hear them.\n\n\n**W2: The empirical comparisons do not consider alternative approximate Bayesian methods.**  \nWe agree that comparing our proposed method to other approximate Bayesian methods is crucial. In the paper, we consider approximate Bayesian methods like MFVI, Laplace and Adversarial Variational Bayes, which also falls into the category of implicit methods. We now also include experiments with two versions of full Laplace, see tables 1 and 2 in the attached PDF. Furthermore, we have added results for multiplicative normalising flows (MNF, Louizos and Welling, 2017, [1]) on the UCI datasets (tables 1 and 2) and MNIST (table 3). The MNF are difficult to get to converge, however. All these results have been added to our paper too. If you have specific baselines in mind that you think we are missing to make a comprehensive benchmark, we will happily take them into consideration.\n\n[1] Louizos and Welling, \"Multiplicative normalizing flows for variational Bayesian neural networks\", ICML 2017.\n\n**W3: Posterior refinement (Kristiadi et al., 2022) is closely related and should be discussed.**  \nThank you very much for pointing us to this work. It is indeed relevant and we will add it to the related works section of the paper. Posterior refinement works by using the Laplace approximation as a clever base distribution for a normalising flow, which is then optimised to model the posterior of BNN. This can work well given a sufficiently expressive flow, which, however, typically comes at a high computational cost, especially in high dimensions. This is why Kristiadi et al. (2022) focuses on last-layer approximations. In contrast, LIVI uses a neural sampler to implicitly represent the posterior of the BNN, which means we can model the full posterior in an expressive manner. While there are pros and cons of both methods, it is difficult to imagine posterior refinement being scaled to millions of dimensions as we do here.\n\n\n**Q1: What is needed to apply LIVI to large, modern networks?**  \nThe hypernetwork poses the biggest challenge, in our opinion, as this is the network that is supposed to efficiently parametrise the high-dimensional implicit posterior. We do think that there are still ways to go beyond what we have done here. One approach would be to consider independent hypernetworks for each layer of the BNN, which would reduce the size of the individual hypernetworks at the cost of losing correlations across layers of the BNN. This itself should help such a framework to increase its modelling capacity significantly. One could also consider efficient ways of putting priors over deep networks to curb dimensionality like [1] that proposes implicit BNNs that have priors over activations rather than weights and biases, thus reducing the dimensionality of latent variables, and [2] which also assumes priors over units in a neural network and models weights using these latent distributions.\n\n[1] Trinh et al., \"Scalable Bayesian neural networks by layer-wise input augmentation\", arXiv 2020, https://arxiv.org/abs/2010.13498  \n[2] Karaletsos et al., \"Probabilistic Meta-Representations Of Neural Networks\", UDL workshop 2018, https://arxiv.org/abs/1810.00555\n\n\n**Q2: Why was KIVI not implemented?**  \nThe official code for KIVI is written in a probabilistic programming subpackage developed by researcher at Tsinghua University and used by researchers there. The package is built with TensorFlow 1 and hence is not easy to port. We tried to implement our own version, but we were not confident in our implementation and decided to not include results from this method when they were not available in the original paper.\n\n**Q3: The runtime of the method should be included in the paper.**  \nWe agree that this is very a important metric by today's standards and we will include this in the main paper if space allows, otherwise the supplementary. The only reason the runtime was not reported in the submission is that, during the deadline rush, the experiments were performed on different GPU devices and hence we did not have one-to-one time comparisons for all the methods. We have now completed benchmark measurements on MNIST, see table 4 in the attached PDF, which shows the training time required for convergence as well as the memory consumption for LIVI and four other methods. When compared with other expressive VI methods like AVB and MNF, LIVI is faster to train and consumes less memory. \n\n**Q4: Minor typos.**  \nThank you very much for pointing these out. They have been corrected."}, {"review_id": "hEjFIFe6QM", "reviewer": "Reviewer_Q4Sr", "summary": "This work proposes one approach to implicit variational inference that avoids using adversarial training. This is achieved by first applying a small amount of Gaussian noise to the implicitly generated weights $\\theta$ to ensure that they are equipped with a valid distribution (enabling valid divergences), and then the now valid ELBO is approximated through linearizing the entropy regularization term. A further lower bound to this approximation is also proposed that allows for better scaling in high-dimensional settings, which is common for applying VI to neural models.  ", "questions": "I do not have any direct questions, see the weaknesses for my main comments to be addressed please.", "limitations": "The authors did discuss the limitations adequately. Negative societal impact is not directly applicable here in my opinion. ", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "I felt that the approach presented was well justified, easy to follow in terms of motivation and development, and enables much more complex settings for variational inference to be applied to (in a stable manner) without having to suffer from the typical limiting mean-field assumptions.\n\nThe experimental results were impressive as well, seemingly achieving much better uncertainty quantification than the other competing methods while still retaining quality predictive performance. ", "weaknesses": "The paper states in the contributions (line 49) that it \"derive[s] a novel lower bound for variational inference\"; however, I believe this is not quite the case. It is my understanding that the \"novel lower bound\" being mentioned here refers to $\\mathcal{L}'$ and the further lower bound $\\mathcal{L}''$ used for better scaling. This statement is in direct contradiction to equation 18 that states:\n$\\log p(\\mathcal{D}) \\geq \\mathcal{L}(\\gamma) \\approx \\mathcal{L}'(\\gamma) \\geq \\mathcal{L}''(\\gamma)$. Strictly speaking, $\\mathcal{L}'$ is not a lower bound, but rather an approximation to an actual lower bound $\\mathcal{L}$. To be clear, I do not think this is bad by any means, it just needs to be communicated clearly. \n\nAside from this, I think the experiments can be bolstered with a few additional comparisons. Namely:\n- The impact of using $\\mathcal{L}''$ over $\\mathcal{L}'$ when the latter is still eligible (i.e., when using smaller models such as the UCI dataset tasks).\n- How much performance (both in and out of distribution) differs between your proposed method and HMC (or other MCMC methods). This again would need to be done in low-dimensional settings, but I believe it should be feasible for at least a smaller neural network or linear model on some of the UCI dataset tasks.\n\nLastly, one of the contributions listed cited a novel generator architecture (line 53); however, the details of this seem to be relegated to the appendix. In future revisions, I believe if you are going to cite this as one of the major contributions then the details should at least be partially included in the main paper. I understand this probably wasn't done due to space limitations, but it is worth considering in my opinion.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for a careful review of our work. We are delighted that you find our exposition easy to follow. We address you individual questions and suggestions below.\n\n**W1: The contributions state a novel lower bound, however, it is actually an approximation to a lower bound.**  \nThank you for highlighting this. We agree that we should have been more clear that $\\mathcal{L}'$ is an approximation to a lower bound. We will clarify this in the paper.\n\n**W2: An experiment on the impact of using $\\mathcal{L}''$ over $\\mathcal{L}'$ would be good.**  \nThank you for suggesting this experiment - it would indeed be informative. We now include experiments to compare our two bounds, $\\mathcal{L}'$ and $\\mathcal{L}''$ on the UCI datasets, see tables 1 and 2 of the attached PDF. The results show that models trained with either $\\mathcal{L}'$ or $\\mathcal{L}''$ perform quite similarly both in terms of test log-likelihood and test RMSE, giving empirical justification for the lower bound on $\\log\\det(J J^\\top)$.\n\n**W3: An experiment on the performance difference between LIVI and HMC in low-dimensional settings would be good.**  \nThis is also a great suggestion. For a qualitative assessment, Figure 1 of the main paper shows such an experiment for a toy problem. Furthermore, we now include experiments comparing our two bounds, $\\mathcal{L}'$ and $\\mathcal{L}''$, and HMC on the UCI datasets in tables 1 and 2 of the attached PDF. The results show that the model trained with $\\mathcal{L}'$ performs close to or as well as HMC, suggesting that the local linearisation does not harm the expressivity dramatically. A model trained with $\\mathcal{L}''$ is not much worse.\n\n**W4: The novel generator architecture should be briefly discussed in the main paper.**  \nThank you for pointing this out, it is a very good point. We will add a short description of the architecture to the main paper.\n"}, {"review_id": "CqSsWJWcH6", "reviewer": "Reviewer_M1uk", "summary": "The paper proposes LIVI a variational-inference-based approach for Bayesian NNs. It hinges on implicit VI (IVI) which allows to easily sample from an arbitrary distribution, yet obtaining the gradients w.r.t its parameters or the density is hard. As the ELBO in VI decomposes to two terms, an expected likelihood term and a KL term, under this framing the KL term cannot be evaluated, and even may be ill-defined. To remedy that, the authors propose to use a Gaussian deep latent variable model (DLVM), i.e., the generator network in VAEs, as a replacement for the implicit density. And, to approximate the entropy term in the KL divergence the authors propose to linearize the output of the network. The authors also propose a lower-bound to the ELBO by approximating the log-determinant of the covariance using the minimal singular value, which can be calculated more efficiently compared to all of the values. The authors compare their method to several Bayesian methods on in-distribution and out-of-distribution tasks showing improved performance over baseline methods.", "questions": "* How did your method of approximating the entropy is compared to standard MC sampling as suggested in the beginning of section 3.2?\n* How does your method perform compared to baseline methods in terms of memory and run-time?", "limitations": "The authors did not address the limitations of their method.", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "* A novel approach for using hyper-networks to learn a Bayesian model using VI.\n* Good results on OOD tasks compared to baseline methods.\n* The paper is written clearly and addresses relevant related studies.\n* The results seem to be reproducible - exact experiential details were given along with the code.", "weaknesses": "**Method**\n* Although I think the idea is nice a major shortcoming of this method is that it is challenging to apply it to large networks because of the use of Hyper-networks. Even in the experimental section, the largest network has 2.7 million parameters which is considered small these days. How do the authors propose to handle that issue and scale it to networks of 1-2 orders of magnitude larger?\n* To estimate the entropy term the std of the noise is taken to be zero on only one summand (the bilinear function), but not on the logdet of the covariance. How do the authors justify that?\n\n**Experiments**\n* The compared methods are somewhat outdated. Many Bayesian models were published in recent years, and all of them show that they beat deep ensembles in one way or another. I believe that more recent baselines should have been evaluated.\n* I find it a bit odd that the authors didn't compare LIVI to methods with similar pipelines, such as flow-based methods [1, 2].\n* The method was evaluated on UCI benchmarks, MNIST and CIFAR-10. In my opinion, it is not enough to showcase the advantage of the method and more challenging datasets should be considered.\n\n\n[1] Krueger, D., Huang, C. W., Islam, R., Turner, R., Lacoste, A., & Courville, A. (2017). Bayesian hypernetworks. arXiv preprint arXiv:1710.04759.\n[2] Louizos, C., & Welling, M. (2017, July). Multiplicative normalizing flows for variational bayesian neural networks. In International Conference on Machine Learning (pp. 2218-2227). PMLR.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for a detailed review. We answer your concerns below.\n\n**W1: How could we make the method scale to tens or hundreds of millions of dimensions?**  \nWe absolutely agree that current state-of-the-art DNNs are much larger than the networks our approach can serve. However, this does not make our approach any less relevant to the Bayesian deep learning community, as research in this area does not solely focus on transformers. No prior work in the Bayesian literature has been able to scale expressive variational approximations to such a large number of latent variables.\n\nTo scale our method even further, one could consider independent hypernetworks for each layer of the BNN, which would reduce the size of hypernetwork at the cost of losing correlations across layers of the BNN. This should help increase the modelling capacity significantly. One could also consider efficient ways of putting priors over deep networks to curb dimensionality like [1] that proposes implicit BNNs that have priors over activations rather than weights and biases, and [2] which also assumes priors over units in a neural network and models weights using these latent distributions.\n\nWhile we consider scaling to even higher dimensions as future work, we agree that the topic is important for our paper and have added a discussion on the issue of scaling to the paper.\n\n[1] Trinh et al., \"Scalable Bayesian neural networks by layer-wise input augmentation\", arXiv 2020, https://arxiv.org/abs/2010.13498  \n[2] Karaletsos et al., \"Probabilistic Meta-Representations Of Neural Networks\", UDL workshop 2018, https://arxiv.org/abs/1810.00555\n\n**W2: Why is $\\sigma_m^2$ not taken to be zero in the logdet of the covariance (e.g., in Eq. (A.18))?**  \nIt is correct that both $E_{z\\sim q(z)} [\\log \\det C(z)]$ and $E_{z\\sim q(z)} E_{\\theta \\sim q_\\gamma(\\theta | z)} [ h(\\theta, z) ]$ in Eq. (A.13) depend on $\\sigma^2$. For smaller networks, $E_{z\\sim q(z)} [\\log \\det C(z)]$ can be computed analytically, but $E_{z\\sim q(z)} E_{\\theta \\sim q_\\gamma(\\theta | z)} [ h(\\theta, z) ]$ involves a matrix inverse, which requires a computationally expensive and unstable procedure. We, therefore, choose to approximate this term, which we do by letting $\\sigma^2 \\rightarrow 0$. We could have chosen other approximations, but this one makes sense, as $\\sigma^2$ is already defined to be small. We do not apply the same approximation to $E_{z\\sim q(z)} [\\log \\det C(z)]$ since we do not have to (as long as we consider small networks), and introducing an unnecessary approximation makes little sense. \n\nTo derive our second bound, $\\mathcal{L}''$, rather than approximating $\\mathbb{E}_{z\\sim q(z)} [\\log \\det C(z)]$, we introduce a lower bound, which was presented by Geng et al. (2021) [1]. From an optimisation perspective, a bound makes more sense than an approximation, and this bound is furthermore efficient to compute.\n\n [1] Geng et al., \"Bounds all around: training energy-based models with bidirectional bounds\". NeurIPS 2021.\n\n**W3: The baselines are somewhat outdated. Many recent Bayesian models beat deep ensembles.**  \nThank you for raising this concern. We do our best to make fair comparisons by including both recent and well-performing models. We are not aware of the models you refer to, but if you have specific ones in mind, we will happily try to test them and report back to you.\n\n**W4: There should be comparisons to flow-based methods.**  \nNFs are powerful models, but scaling them to millions of dimensions is incredibly difficult and computationally demanding. In fact, these challenges were some of the main motivators for our work. Regarding the multiplicative normalising flows (MNF), we have added results for these on the UCI datasets (tables 1 and 2) and MNIST (table 3), see the attached PDF. The flows are difficult to get to converge, however. Additionally, the training time and memory consumption are shown in table 4. Compared to MNF, LIVI is both faster to train, consumes much less memory, and performs better.\n\n**W5: There should have been more challenging datasets in the experiments.**  \nWhile we certainly agree that more challenging datasets exist, we chose these ones for a number of reasons. The datasets are all commonly used in the literature, which makes it easier to find competing methods and compare our work to others. The simpler datasets, such as rotated MNIST, are easy to understand intuitively and can therefore provide us with insights into the models' behaviours and their abilities to learn inductive biases. However, if you have specific datasets in mind, we will happily try to test LIVI on them and report back here. We should say, however, given compute limitations, for any dataset, the models should be able to run on a single GPU.\n\n**Q1: How does LIVI compare to standard MC sampling?**  \nWe compare the LIVI bounds to an importance sampling estimate in figure F.1 in the supplementary. Although the two bounds systematically underestimate the entropy (which makes sense as they are lower bounds), crucially, they largely follow the behaviour of the sampling estimate. This gives us some confidence that they are useful objectives. Note, however, that the experiment is only possible to do for very small models as importance sampling is challenging to get to work in more than a few hundred dimensions.\n\n**Q2: How does LIVI compare to other methods in terms of memory and runtime?**  \nWe agree that these are very important metrics and the only reason the runtime was not reported in the submission is that, during the deadline rush, the experiments were performed on different GPU devices and hence we did not have one-to-one time comparisons for all the methods. We have now completed benchmark measurements on MNIST, see table 4 in the attached PDF, which shows the training time required for convergence as well as the memory consumption. Compared to other expressive VI methods like AVB and MNF, LIVI is faster to train and consumes less memory."}]

Author rebuttal:
We thank all the reviewers for carefully reading our paper and providing constructive feedback. We appreciate that the reviews found that “the paper is well motivated, has clear mathematical development, and solves common issues” (MstJ) and that “experimental results were impressive as well” (Q4Sr). We have replied to each of you individually and have further attached a PDF with additional results that were requested. We look forward to further discussions with you.



Meta-review: This paper proposed a new implicit VI method for Bayesian neural network approximate posterior inference. The idea is to use a neural sampler for the q distribution. The main contribution of this paper is a new methodology that solves two issues of implicit VI: (1) using semi-implicit posterior to solve the ill-defined KL divergence issue, and (2) using linearisation to approximate the intractable entropy term for the q distribution. Experiments on UCI and MNIST scale data demonstrates improved results regarding comparisons with other implicit VI methods and a few other baselines.

Reviewers all agree that the proposed approach is highly novel and sound, they commend that the method provides significant advances in implicit VI literature. One reviewer raised the question regarding small scale experiments and mentioned that experiments in current Bayesian deep learning literature is at least ResNet-18 or -50 level. 

Overall implicit VI with neural sampler has yet to scale to ResNets with >10M parameters which is a downside of this paradigm of methods. But within this paradigm the proposed approach has made solid progress in my opinion, and the approach of using linearisation for implicit distribution entropy estimation is quiet interesting. Therefore I would recommend acceptance for this paper and encourage the authors to consider larger scale experiments in future work.
Decision: Accept (spotlight)