VIKING: Deep variational inference with stochastic projections
Authors: Samuel G. Fadel
Hrittik Roy
Nicholas Krämer
Yevgen Zainchkovskyy
Stas Syrota
Alejandro Valverde Mahou
Carl Henrik Ek
Søren Hauberg
Venue: NeurIPS 2025 poster ()
Keywords: Bayesian deep learning
overparametrization
reparametrization invariance
linear algebra
matrix-free methods
Primary area: probabilistic_methods
TL;DR: We develop a highly scalable reparametrization invariant variational training scheme for deep learning.

Abstract: Variational mean field approximations tend to struggle with contemporary overparametrized deep neural networks. Where a Bayesian treatment is usually associated with high-quality predictions and uncertainties, the practical reality has been the opposite, with unstable training, poor predictive power, and subpar calibration. Building upon recent work on reparametrizations of neural networks, we propose a simple variational family that considers two independent linear subspaces of the parameter space. These represent functional changes inside and outside the support of training data. This allows us to build a fully-correlated approximate posterior reflecting the overparametrization that tunes easy-to-interpret hyperparameters. We develop scalable numerical routines that maximize the associated evidence lower bound (ELBO) and sample from the approximate posterior. Empirically, we observe state-of-the-art performance across tasks, models, and datasets compared to a wide array of baseline methods. Our results show that approximate Bayesian inference applied to deep neural networks is far from a lost cause when constructing inference mechanisms that reflect the geometry of reparametrizations.

VIKING Deep variational inference with stochastic projections
Samuel G. Fadel, Hrittik Roy, Nicholas Krämer, Yevgen Zainchkovskyy, Stas Syrota, Alejandro Valverde Mahou Technical University of Denmark {samma, hroy, pekra,yeze, stasy, avama)/ddu........aaaaaaaaaa dk
Carl Henrik Ek University of Cambridge che290cam. ac uk
Søren Hauberg
Technical University of Denmark sohau@dtu.de
Abstract
Variational mean field approximations tend to struggle with contemporary over- parametrized deep neural networks. Where Bayesian treatment is usually associ- ated with high-quality predictions and uncertainties, the practical reality has been the opposite, with unstable training, poor predictive power, and subpar calibration. Building upon recent work on reparametrizations of neural networks, we propose a simple variational family that considers two independent linear subspaces of the parameter space. These represent functional changes inside and outside the support of training data. This allows to build fully-correlated approximate posterior reflecting the overparametrization that tunes easy-to-interpee hyperparameters. We develop scalable numerical routines that maximize the associated evidence lower bound (ELBO) and sample from the approximate posterior. Empirically, we observe state-of-the-art performance across tasks, models, and datasets compared to a wide array of baseline methods. Our results show that approximate Bayesian inference applied deep neural networks is far from lost cause when constructing inference mechanisms that reflect the geometry of reparametrizations.
The troubles of the overparametrized Bayesian
Bayesian inference provides an attractive framework for learning from data, marginalizing a data- dependent likelihood with data-inden iata-indeeendent prior leads to combined probability measure that be used for learning. For most model classes, the marginalization is intractable, and we resort to approximate methods of integration. While this has shown itself to be very successful approach, overparametrized models have subtle characteristics that lead to significant challenges in practice.
Specifically, overparametrizationn lead to many-to-one to-one relationship where several parameter con- figurations describe the same function. Dinh et al. (2017) exemplify that the neural network f(z) wi ReLU(w> æ) can be reparametrized as f(x) wi / ReLU w2 x) for any 7 0, implying that the weights (wi, w2) and (w1/>, 7w-2) yield identical functions. Thus, despite having two trainable parameters, the neural network only has degree of freedom. Making matters worse, the discrepancy between the number of parameters and degrees of freedom generally grows with the be performed analytically, the resulting measure will reflect
model size. If the marginalization this structure; however, if we need to resort to approximate methods, we need to carefully design our approximation that the correct geometry of the problem be parametrized. If not, the resulting object will not be valid probability measure over the space of functions and application of this
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

object will lead to pathological solutions where the same function has different densities dependent on its parametrization (Watanabe, 2009).
In recent work, Roy et al. (2024) showed that the parameters describing the same function are continuously connected sets in the weight space. This has damaging effects traditional posterior approximations like Laplace approximations (MacKay, 1992) or mean-field approximations (Blun- dell al., 2015) cannot accurately reflect the reparametrizations,, which lead the disappointing performance of approximate Bayesian inference in these types of models.
In this paper, we study variational inference applied to overparametriiee deep neural networks. We propose an approach that takes the intractable posterior and factorizes it into two orthogonal parts, a data-dependent and a data-inded iatt-inndeeedeen measure. These can, intuitively, be seen as modeling uncertainties on the training data in one subspace and elsewhere in the other. Importantly, this leads to an easy-to-interpret structure analogous the likelihood and the prior that lead to the posterior in the first place. Despite having an interpretable and simple decomposition, our approach is fundamentally fully correlated and captures dependencies between all model parameters, addressing the pathological behavior of previous methods. This introduces a computational challenge, which we address through stochastic extension of the alternating projections algorithm that was recently proposed for post hoc posterior approximation (Miani et al., 2025). The result is a scalable, architecture-agnosii approach to Bayesian training of deep neural networks that achieves sfte-of-the-as art performance on wide array of contemporary models and datasets. The both code for reproducing our experiments' and an casy-to-use library? are available online.
2 Background and related work
The aspiration of Bayesian deep learning is to combine the predictive advantages of deep neural networks with the theoretical justifications of the Bayesian framework, hoping achieve the benefit of both. While the individual layers in a neural network are simple mathematical objects, the compositional structure of the layers renders posterior inference analytically intractable for any interesting model. There has been a range of different approaches to recover posterior estimates of the neural network weights (Hernandez-Lobaal and Adams, 2015; Hron et al., 2018; Li et al., 2015; MacKay, 1995), but often the benefits have failed to materialise, some even restricting the Bayesian treatment to a subset of parameters (Zhao al., 2024). Rather, it is common to experience either over- or underfitting (Daxberger et al., 2021; Kristiadi et al., 2022; Wenzel et al., 2020; Zhang et al., 2024), and most Bayesian training schemes are reported to be brittle (Warburg et al., 2023). Success in achieving the promises of a Bayesian treatment of deep neural networks has been so elusive that it raises the question if there is something inherent in their structure that renders numb to a Bayesian approach (Sharma et al., 2023).
Singular learning theory (Watanabe. 2009) argues that the troubling structure comes from over- parametrizations, where a single function can be described in more than one way. While mathemati- cally elegant, the algebraic geometric approach has been difficult to turn into computational practice. Instead, Roy et al. (2024) framed the problem in terms of differential geometry by studying how the geometry of the function space embeds itself into the parameter space. From this, the authors were able to characterize the manifold of reparametrizations in a manner that lends itself to computation. Specificallyy, the manifold can be characterized locally by the kernel of the Fisher-Rao metric over the network parameters 0 (Amari, 2016),
a log po (ylx) ô Olo llogpo(ylx)
Fe Ey~po(yix)
RDDD
(1)
00
00
Here log Po x) is the log-likelihood of the data parametrized by 0 E CRD. The kernel of this matrix infinitesimally gives directions in weight space where log Po (ylx) is constant for all observations in the support of the data. This provides a characterization of why Bayesian learning has not succeeded, showing that for deep learning harvest the benefit of a Bayesian treatment, the geometrical structure needs to be reflected in the approximate posterior. For detailed discussion, we refer to (Kunstner al., 2019; Miani et al., 2025).
'httpa://github httpa////giccuub..oom//ge/e////s com//eugne//vkking peppe--xxpering
hattpa://github wewettuu..oomfffaee www com/fadel/viking
TThooughou the paper, kernel refers the null space of matrix, while image its orthogonal complement.

---

kernel samples
a
15
image samples
b
0.220
IVON variances
C
10-
5-
5-
0.310
0.205
Posterier
Weight
100000
200000
Figure/ 1: Panels a and show isotropic Gaussian samples that have been projected onto the kernel and image of the empirical Fisher-Rao metric associated with neural network trained on the shown data. The kernel samples retain the predictions of the neural network, while the image samples do not. Panel c shows the learned weight-varianees of IVON trained on CIFAR-10. These variances are near-identical across weights, suggesting that an isotropic approximate posterior has been learned.
Miani et al. (2025) take steps towards integrating knowledge of overparametrization into a post-hoc approximate posterior. They note that the kernel of per-datum loss-Jacobian, which is closely related the Fisher-Rao metric (I), consists of functions with identical training loss. Assuming noise-free data, they restrict the approximate posterior to only have probability mass in this kernel, resulting in a posterior approximation from which all samples have the same loss (see Figure la,b).
Variational inference (Blei et al., 2017) has shown itself to be an efficient and robust approach to approximate Bayesian inference and has been applied in several different ways to deep neural networks. Existing variational approximations predominantly use Gaussian variational distributions p(][x) x (0 )), see, c.g., (Blundell et al., 2015; Shen et al., 2024), where O and Σ are the variational parameters learned. To ensure that the approximation is tractable, the covariance Σ is commonly given diagonal or block- diigonial structure, i.e., a mean field approximation (Daxberger al., 2021). Current state-of- ff-the-art in this direction appears to be IVON (Shen et al., 2024), which develops Newton-like optimization scheme for the evidence lower bound (ELBO) associated with the mean field approximation. In our experience, IVON tends to learn near-identical variances for all network weights (Figure lc), suggesting that the posterior approximation is inaccurate. Further, mean field approximations tend be brittle and sensitive to the choice of hyperparameters. One hypothesis for this behavior is that the independence assumptions of mean field cause instabilities overparametrized models, where parameters, by construction, are highly correlated. In this work, we want to build a full variational inference method that respects the geometry induced by the overparametrization characterized by Roy et al. (2024). To achieve this, we split the parameter space into the kernel and image of the Fisher-Raa metric, considering these independent.
The numerical computations associated with non-trivial covariance structures can be daunting. The first challenge is that for large neural networks, the covariance matrix cannot be stored in memory as its size grows quadratically with the number of network parameters. One remedy is to note that some matrices, e.g., the Fisher-Rao metric (1) and the generalized Gauss-Newton matrix (Daxberger et al., 2021), are functions of the training data. This implies that vector products of such matrices can be implemented without instantiation non-trivial covariances at the cost of a higher compute budget, see, e.g., (Miani et al., 2024, 2025; Roy et al., 2024) for examples.
so-called matrix-free approach. This allows for working with
3 VIKING: Variational Inference with Kernel- and Image-spaces of numerical Gauss-Newton matrices
We next build a variational training loop designed with overparametrization in mind. We start from the ever-present prior p(0) = A/000, a 'II over parameters R.P, with precision a E R. As motivated earlier, we choose a variational family ((0) that is linked to the kernel of an empirical estimate of the Fisher-Ras (FR) metric at Ô denoted ker(Fa) CR.. Specifically,
(0) (o Eg), Σn = U.UJ + oi ( U.
(2)
Here, Ug R. Dx R is R-dimensional basis of the kernel, such that U a projection matrix that maps parameter vectors in IR to the linear subspace ker ker(F,). In other words, for any v E R we have Ug U v (er(F)), wherein the model's loss is unchanged at the data used to estimate the FR.
3

---

Algorithm 1: The VIKING algorithm
Input:
Initial values of (0, oeer mm), dataset with B mini-batches
Output: Optimized values of (0, oee Om)) representing q(0) (Equation 2)
for each epoch do
for each Monte Carlo sample = 1 to S do
Take c(..0) ~ (0, I)
for each mini-batch = 1 to B do Compute e(s,t) from e(s.t=1) (Equation 14)
> Alternating projections
(s.0) + els.B)
ker
7
for each mini-batch = to do for each Monte Carlo sample
S do (s.t-1) (Equation 18)
Compute ker
(s,t) from kker
(x,t) c(..0) (s,()
10
Im
ker
g(s) +- Tker ker
(s,t)
o Sim
(s,t)
A sample from g(0)
Keer; Jim f- Gradient step on L(O, oker Tim) using { (Equation 3)
12
This proposed approximate posterior is perhaps the simplest possible choice that is adapted to overparametrized models: it has scalar parameter, of E R+, that influences the predictive uncertainty on the training data, and another scalar parameter of E R4 that reflects uncertainty elsewhere. Arguably, it is overly simplistic to calibrate neural network through two scalars, but will see that is surprisingly competitive, suggesting that explicitly reflecting overparametrization is important in Bayesian approximations.
3.1
The VIKING ELBO
To estimate (0), maximize the usual (variational) lower bound C of the likelihood p(y)0, x)
L(Ú, Oker: Tmm) = Ee [logp(y 0, x)] - (L))())lP((O))
where KT. denotes the Kullback- Leibler divergence. We call the result Variational Inference with Kernel- and Image-spaces of numerical Gauss-Newton matrices (VIKING). The overall algorithmic steps involved in method summarized in Algorithm 1. We refer to the two terms of the ELBO as the reconstruction term and the KL term, respectively, and discuss their evaluation what follows.
The reconstruction term. Assume access to an algorithm to approximate products between the projection matrix UgU á and any vector E RDD. For 1,..., S, we can then estimate of the reconstruction
(x)) ~ N(O,
E
R
D
(4)
(projection onto kernel space)
ker
(+) U g ((5)
RR.
D
(5)
(image space orthogonal)
im
(a)
(I U&
)e(
(00)
-6
(s)
RR
(6)
( θ + oker (+)
+ Fim 2
R
D
(7)
s
Eo~g [og p x))
=
Σ logp(y (())
E
R
(8)
s
(s) Here, Equation 5 requires careful computation, to be addressed later. Once cke is computed, obtaining a sample ((s) from the approximate posterior trivial.
The KL term. Since we have both a Gaussian prior and a Gaussian approximate posterior, we can evaluate the KL term closed form with
(L(g(0))P((O)) = (TTr(EE) - D + 0 D Dlg((a) a)) log det Σg))
(9)

---

Recall the construction of Σg from Equation 2. Since Eg is sum of scaled projection matrices, it 2 has R eigenvalues equal to ker and D R eigenvalues equal to oi-- We can, thus, evaluate
T o R o (D R).. and log det [g) = 2R (orr) + 2(D-R) (imm)
(10)
(11)
requiring only an estimate of the kernel dimension to evaluate the full KL term in closed form. To estimate R. we note that U& U is a projection matrix, i.e., has R cigenvalues that are while 0
the remaining are 0. Consequently,
S ~
S
R = T
((0) S
U
c(+)
c(s)
(e)
(12)
S
kert
==1
8=1
where the approximation Hutchinson's trace estimator (Hutchinson, 1989). Note that we may reuse keer (Equation 5) from the reconstruction term.
3.2 Kernel projections
At first glance, the ELBO terms (Equations and 9) seem straightforward. The practical difficulty lies in projecting onto ker(Fa) as per Equation 5. One noticeable challenge is U U having x D entries, where D is the number of model parameters. With many modern models having millions or billions of parameters, it is not possible to instantiate this matrix, and resort to matrix-fres algorithms.
We start by defining the matrix containing a stack of per-datum loss gradients,
V log (V
J
E ERD
(13)
V log Po (y
This matrix is chosen such that the kernel of J JJJ estimates the kernel of the Fisher-Rao metric, kcr(Fa) (Equation 1), is common practice in Bayesian deep learning (Miani al., 2025). For deeper discussion of the implications of this choice, we refer to Kunstner et al. (2019). Projecting onto the kernel of J Jg can then be written as a least-squares problem over J 8.
ker U (+ arg min |I
1 subject to Jou - 0).
(14)
Using Lagrange multipliers, we can rewrite the loss in Equation 14 as
(uu A) u 1 1 + AJau.
with minima (+) and x that satisfy
= (E(x)
ker X
(15)
J
more explicitly, the optimal solution is
(+) ( x X ===
g (e)
(16)
The numerical bottleneck of the projection is thus inverting the matrix JJJT We never invert this matrix explicitly, but instead solve the linear system JJJ = Jye with the conjugate gradients algorithm (Hestenes et al., 1952). The conjugate gradients algorithm is a matrix- free method, which means that we never materialize the system matrix, and instead, require only matrix-vector products and rector-matrix products, which are accessible with automatic differentiation (since the system matrix consists of stacks per-data gradients).
Importantly, conjugate- -radiee solves can be numerically brittle they rely on well-conditioned systems. This frequently leads to inaccurate projections, in the sense that ( and c are not orthogonal, even if they should. The solution to this problem to either precondition the linear system to improve the numerical robustness of the conjugate gradients method with reorthogonal- ization - we found the latter to perform well, thus, experiments use conjugate gradients with full reorthogonalization (Gratton et al., 2021); see also Maddox et al. (2022).
5

---

3.3 Stochastic alternating projections
The projection algorithm performs most of the computation that eventually gets us samples from the approximate posterior q()). Unfortunately. following Equation 16 to project onto the FR kernel requires working with J g. N x N matrix. Recently, Miani al. (2025) proposed a scalable mini-batched algorithm for projecting onto matrix kernels, based on von Neumann's alternating projections (1949). In our setting, this algorithm would iteratively project onto the kernel of batched FR estimates,
0 U E,
(17)
where e is the vector to be projected and T U is the kernel-projection of the tth batch. This process can be shown to converge onto the kernel of the complete-data FR (Miani et al., 2025).
Unfortunately, the alternating projections algorithm still requires or passes through the entire dataset to project a single vector. As we require a sample from the approximate posterior in every optimization step, alternating projections is inapplicable to mini-batched ELBO optimization.
The fundamental issue is that in each ELBO optimization step, we change the approximate posterior mean Ó, such that the FR and its kernel also change. A naive procedure would draw an initial sample (0)) ~ (O, I).. and repeatedly project this onto mini-batch FR kernels throughout ELBO optimization, hoping that the posterior mean changes sufficiently slowly to maintain the properties of kernel samples. Empirically, we have not found this to be the case.
With this in mind, we propose a stochastic extension of the alternating projections algorithm suitable for batche ELBO optimization. Let e(t) denote a sample in the FR kernel associated with the posterior mean at step & (e.g., one parameter update step). We then update the sample
(1) = u (1) (C) (V>elt-1)+ (I-7(
n (0,,
(18)
where the hyperparameter 7 E [o, 1] controls the noise of this stochastic alternating projection. For Y 1, we recover the above naive approach that disregards the change in kernel, while for y 0 we merely project onto the kernel of the current batched FR estimate. In-between values implement soft "sliding window" along the history of projections, where previous projection information is preserved and then eventually forgotten after a certain number of steps.
What is the effect of additional noise? Empirically, we find the stochastic extension essential to training models suc- cessfully. In Figure 2, we show the accuracy curves over training of a small CNN classifier on Fashion MNIST, with dashed horizontal lines indicating performance on test data.
0.90
Accuracy
0.85
Y
0.0
The experiment evaluates values of y: no noise (7 = 1.0), only noise (7 0.0), and an in- in-between choice (y 0.5), per Equation 18. The model without noise quickly achieves the highest training accuracy, but its dependency the pro- jected samples, which change only once per epoch, leads to a poorer approximate posterior and generalization.
0.80
0.5
1.0
0.75
0
20
40
60
80
Epoch
Figure 2: Performance of values.
3.4 Post-hoc variational inference or full model fitting?
Given the computational effort evaluate and run the ELBO, is important to question whether this could be avoided entirely. One possibility to start from well-performing model, and only use the ELBO to tune the variational param- eters oee and Jin. i.c., post-hoc variational inference with a fixed model. This can be efficiently implemented as we only need to project one set of samples, which can be reused throughout optimization. With the same experimental setting from before, Figure shows the behavior of such a training procedure. Interestingly, post-hoc tuning is not per se prob- lematic and indeed leads to a well-performing model. However, the figure shows that optimizing the variational mean with the ELBO improves upon the host-hoc approach.
Accuracy
0.4
Post-hoc sigmas
Full ELBO
0
5
10
15 Epoch
20
25
30
Figure 3: Post-hoc tuning oeer Fun-

---

IVON (Shen et al., 2024)
(ours)
5.0
5.0
-5.0
-5.0
5d devitten
1.0
deviatio
1.0
0.5
Figure 5: A toy regression example sinusoid curve with 10 data points. Top: The curves show the training points black, with the mean fit as red line and 100 posterior predictive samples as blue lines. Bottom: The standard deviation of the predictions over each point in the horizontal axis.
Warming up as shortcut. observed, we can still ob- tain a well-behaved model from the post-hoc tuning proce- dure. Naturally, this informally suggests another way to save computational effort, resources, and time: "warming up" the model, first with regular maximum likelihood estimation (or through pretrained model), and then continuing with ELBO optimization.
0.90
Accuracy
0.85
warmup
0.80
20 epochs
50 epochs
0.75
We test this hypothesis experimentally, comparing a model trained from scratch using our ELBO and one that starts with a model first fitted using maximum likelihood (Figure 4). Our findings are twofold: (i) warmup using maximum likelihood an effective way to quickstart the ELBO learning process with the ELBO; (ii) there is sweet spot where full convergence of a model on maximum likelihood can lead to the ELBO optimization getting stuck. We conjecture that this behavior is due to the sharpness of some minima, which negatively affects the evaluation of the expectation term, as it attempts to sample regions of the weight space around the current model weights, which could have drastically different performance, negatively affecting optimization.
Epoch
Figure 4: Warmup by pretraining.
4 Experiments
We evaluate VIKING against popular Bayesian deep learning methods using standard benchmarks for Bayesian deep learning. Aiming to maintain consistency and simplify comparisons, we adopt the loss-Jacobian (see Section 3.2) when using our proposed variational family. Further details on models, hyperparameters used, and training procedures are in Section C.
Toy (sinusoid) regression. We first illustrate posterior samples in ID regression in Figure 5. We train the same model with both IVON (Shen et al., 2024) and VIKING. While IVON obtains good model fit, its posterior samples are uninformative with respect to uncertainties. On the other hand, VIKING shows higher variance close the boundary of data and beyond, where the model should be less confident.
Image classification. Using standard benchmarks, we evaluate our method against the maximum posteriori (MAP) estimate, the (loss-projected) post-hoc method from Miani et al. (2025), IVON (Shen et al., 2024), SWAG (Maddox et al., 2019), and a last-layer Laplace approximation. On MNIST (LeCun et al., 2010) and Fashion MNIST (Xiao et al., 2017), we train a LeNet (LeCun et al., 1989) model, while on SVHN (Netzer et al., 2011) and CIFAR-I0 (Krizhevsky and Hinton, 2009) we use a small ResNet (He et al., 2016). The models are compared using accuracy, negative log-likelihood (NLL), expected calibration error (ECE), and maximum calibration error (MCE) (Naeini et al., 2015).
The results in Table show that our ELBO generalize better in most cases, which we hypothesize is due to the interaction between evaluations the expectation term which constantly "explores" the weight space around the current mode through the sampling mechanism, whereas the kernel and image variances that allow this exploration are kept in check by the KI. term. calibration metrics, method is particularly effective against the baselines the SVHN and CIFAR- datasets, where
7

---

EMNIST Table 1:
Experimental results over three runs (mean and standard deviation) on in-distribution test data. MAP a point (model) estimate. VIKING and the post-hoc method from Miani et al. (2025) use the loss-projected variant.
Accuracy"
Conf.1
NLL
ECE
MCE+
MAP
0,986
0.00:
0.996
0,000
0.070
0,05
0.247
0.011
0,861 + 0.045
MNIST
VIKING (ours)
0.991
0.001
0.992
0,000
0.055
0.000
0.096
0.004
0.690
1: 0.102
Miani et al. (2025)
0.949
0.000
0.813 1. 0.018
1.225 1. 0.099
0.666
1. 0.007
0,894 1 0.011
IVON
0.989
0.001
0.990 = 0.001
0.043
= 0.002
0.077
= 0.005
0.651
+ 0.042
SWAG
0.982
0.000
0.982
= 0.001
0.064
=: 0,006
0.788
#: 0.005
0,906
+ 0.013
Last Layer LA
0.975
0.002
0.977 0.002
0.090
1: 0.005
0.784
1. 0,007
0.887 1. 0:008
Fashion MNIST
MAP
0.883 0.00
0.942 = 0.003
0.410 = 0.010
0.153 + 0.008
0.590 + 0.141
VIKING (ours)
0.900 0.001
0.928 = 0.001
0.332
#: 0.003
0,075 +: 0,002
0,611 + 0.360
Miani et al. (2025)
0.871 0.006
0.744 1: 0.0M
1.529 0.571
0.617 0.025
0,901 1. 0.013
IVON
0.897 0.004
0.926
0.335
0.073
0.683
SWAG
0.898
0
0.327
0.725
0.907 :: 0.003
Last Layer LA
0.896 0.002
0.931
0.339
0.727
0,902 0.004
MAP
0.947
0.963
0.201
0.055
0.608
VIKING (ours)
0.960
0.964
0.177
0.028
0.308 +:
SVIN
Miani et al. (2025)
0,949
0,948
0.191
0.734 0.017
0,880 0.012
IVON
0,943 0.002
0.888 0.007
0.302 0,016
0.082 0,004
0.492 0.248
SWAG
0.947 0,004
0.897 + 0.007
0.217 +: 0.014
0.745
+: 0.007
0.874
4: 0.005
Last Layer T.A
0.946 0.001
0.943 = 0.006
0.197
+ 0.009
0.740 + 0.007
0.899 + 0.009
MAP
0.824
0.012
0.869
0,000
0.536
0,055
0.075
0,012
0.619 1 0.243
Imagenette CIFAR-10
VIKING (ours)
0.877 0.004
0.893
0,000
0.407
+: 0.010
0.041
:: 0,004
0.331 1: 0.094
Miani et al. (2025)
0.855
0.00
0.701
0.03
2.643
+: 0.206
0,559
:: 0,006
0.802 1: 0.005
IVON
0.835
0.017
0.763 = 0.805
0.817
# 0.075
0.086
+ 0.014
0.436 + 0.244
SWAG
0.865
0.029
0.914 1 0.035
0.445
0.063
0.694
1. 0,018
0.881
1. 0.005
Last Layer L.A
0.894
0.001
0.944
0.001
0.406
0.704
0.880
MAP
0.852
0.002
0.883
0.807
0.481
#: 0.009
0.084 +: 0.010
0.717
+: 0.082
VIKING (ours)
0.887
0.003
0.906
6.002
0.403
0.010
0.077
1. 0,001
0.612 1. 0.062
IVON
0.876
0.021
0.849 0.013
0.656 + 0.136
0.069 :: 0.011
0.464 ::
the overparametrization of the model (around 220k parameters) is more prominent. Additional results including deep ensembles and SGLD (Welling and Teh, 2011) are in Section A.
Does scale to larger models? To evaluate the scalability of our approach, we also train a ResNet34 (21.7 million parameters) on Imagenette (Howard, 2019) using VIKING and IVON. We predominantly aim to demonstrate that VIKING scales to contemporary models despite learning a fully correlated covariance. The results show a superior performance of our ELBO-trained model, but inferior calibration metrics compared to IVON (Table 1).
Out-of-distribution detection. Next, evaluate VIKING in out-of distiiiiibttionn (OOD) detection tasks. We use the maximum variance of the softmax probabilities across output dimensions as an OOD score for both IVON and VIKING. In addition to the datasets from before, and CIFAR-100 as OOD data. For ease of comparison, our experimental setup replicates that of Miani et al. (2025), we additionally use their reported numbers for the other baselines well.
Table summarizes the obtained results. At a glance, VIKING performs on par with the best models, sometimes overperforming the other baselines by big margin (e.g.. MNIST - FMNIST and KMNIST) and other times being a close second. Overall, this further confirms the empirical performance of our variational family and ELBO optimization in obtaining robust models.
Generative modelling- Finally, we demonstrate our approach applied to generative modelling. We train a variational autoencoder (VAE, Kingma and Welling (2013)) until convergence. From that starting point, we further refine the decoder's 's 1.6 million parameters using both IVON and VIKING. We draw 32 samples from each of the resulting approximate decoder posteriors, which we use to both reconstruct images and measure uncertainties.
Figure 6 displays the reconstructed samples and the variance across the posterior samples per recon- structed image. Overall, both posteriors qualitatively recover well-performing models. The methods discern themselves more prominently upon inspection of the variances across the reconstructions of
8

---

Table 2: Area under ROC (1) performance on out-of-distributtin detection. VIKING and the post-hoc method from Miani al. (2025) use the loss-projected variant.
In-dist. Out-of-dist.
MNIST KMNIST
FMNIST KMNIST
FMNIST
EMNIST
MNIST
EMNIST
MAP
0.928
0.934 0.004
0.890
0.728 0:04
0.791
-0014
0.659 #
VIKING (ours)
0.972 à: 0.03
0.965
0,000
0.913
0.06
0.898
0/02
0.926
0.002
0.830 1 0M200
Miani et al. (2025)
0.899 d: 0.011
0.856
0.002
0.893
0.006
0.914
0.835
0.907
0.026
0.928 1: 0.007
IVON
0.948 1: 0.006
0.945
0.007
0.918
0
0.801
0.000
0.867
2: 0.00
0.739 1: 000S
SWAG
0.917 0.04
0.861
0.000
0.916
0.685
0.014
0.655
0.155
0.751 0.022
Last Layer LA
0.793 +: 0215
0.759 0.106
0.782
01
0.768
0:00
0.699
0.020
0.824
1: 0.06
Tn-dist.
SVHN
CIFAR-10
Out-of-dist.
CIFAR-10
CIFAR-100
SVHN
CIFAR-100
MAP
0,946 = 0.006
0.941 #: 0.000
0,831 =: 0.019
0.778 =: 0.013
VIKING (ours)
0.947 = 0.02
0.943 :: 0.827
0.805
# 0.003
Miani et al. (2025)
0,966 0.009
0.960
0.863
0.800 0.013
IVON
0.822 0,014
0.826
0.648
0.714 0.012
SWAG
0.777 0.029
0,787 : 0.027
0,798 0.040
0.782 0.042
Last Layer LA
0,914 0.005
0.908
+: 0.025
0.811 0.024
0.801
= 0.026
Reconstructions
Variances across channels
VIKING (ours)
IVON
Figure 6: Reconstructed samples taken from the latent space prior of variational autoencoders trained with IVON and VIKING. Using 32 posterior samples, the posterior variances compare key semantic features, such as contours between foreground and background, and eyes.
the different posterior samples. Notably, IVON tends to capture variance across all features of the image, including the background. VIKING focuses the variance facial features and outline instead, showing a disregard for the less-relevant aspects of the data.
To assess the model uncertainties, we compute the median value of the standard deviations per pixel produced by each of the 16k generated samples and summarize their distribution in Figure Our in-distribution samples come from a standard Gaussian sample, whereas dusti---iitiiitttin samples come from a Gaussian with twice the variance. The uncertainties produced by VIKING show well-separated distributions clearly detecting OOD samples, while IVON does not.

---

IVON
VIKING
00D
OOD
300
In dist.
600
dist.
Density
Density
200
400
100
200
0.005
0.010
0.015
0.020
0.025
0.005
0,010
0.015
0.020
0.025
of pixel-wise deviations
Median pixel-wise deviations
Figure 7: Histograms of uncertainties on 16K in- and out-of-distribution latent samples according to IVON and VIKING. Uncertainties are computed median of pixel-wise standard deviations computed by decoding the latent sample using 32 posterior model samples.
Discussion
This paper introduces variational family ('VIKING') that explicitly takes model overparametrization into account. This is achieved through a decomposition of the Fisher-Ras metric into parameter subspaces that capture directions in which the per-training data loss remains unchanged and direc- tions in which this loss changes. Intuitively, uncertainty along these directions can be interpreted uncertainty over the training data and general model uncertainty. This is similar the conven- tional splits into aleatoric and epistemic uncertainty commonly found, e.g., in Gaussian processes (GPs, Rasmussen and Williams 2006). The parameters estimate, thus, come with degree of interpretability.
Scaling VIKING to large models is, however, not straightforward, the approximate posterior correlates all parameters, yet we cannot afford to instantiate the covariance matrix due to excessive memory costs. We have developed stochastic extension to von Neumann's (1949) alternating projections algorithm that scales to large models through a matrix-free implementation. The approach readily applies to large contemporary models including ResNet with 21.7 million parameters and generative model with 1.6 million parameters.
Empirically, our results show that VIKING tends to be consistently good as or better than current slate-of-the-arr methods. This lends credibility to the hypothesis that overparametrization should be explicitly taken into account when designing approximate posteriors.
The key limitation is the additional computational overhead incurred by the projection algorithm. Standard gradient descent requires a single backward pass for one gradient update, while VIKING requires the number of CG iterations backward passes per sample. In practice, we train with one Monte Carlo sample and only need a few CG iterations, so the additional cost is manageable. Using more sophisticated linear algebra and custom gradients, significantly more efficient and numerically accurate least squares solvers can be employed (Roy et al., 2025), overcoming some of the time- efficiency and numerical issues our approach.
Acknowledgments
The authors thank Sebastian Mair for his comments early drafts. This work was sup- ported by the Danish Data Science Academy, which is funded by the Novo Nordisk Foundation (NNF21SA0069429), by research grant 42062 from VILLUM FONDEN; by the Novo Nordisk Foun- dation through the Center for Basic Machine Learning Research in Life Science (NNF200C0062006); by the European Research Council (ERC) under the European Union's Horizon programme (grant agreement 101 125993); and by the Horizon-EU EIC Pathfinder Open program to i-RASE: Intelligent Radiation Sensor Readout System, with reference number: HE EIC i-RASE 131530050.
References
Shun-ichi Amari. Information Geometry and Its Applications. Springer Tokyo, I edition, 2016.
10

---

David M. Blei, Alp Kucukelbir, and Jon D. McAuliffe and. Variational inference: A review for statisticians. Journal of the American Statistical Association, 112(518):55-8877 2017. doi: 10.108001661159.2277.12388
Charles Blundell, Julien Cornebise, Koray Kavukcuoglu, and Daan Wierstra. Weight uncertainty in neural networks. In International Conference on Machine Learning, pages 1613-1622. PMLR, 2015.
Erik Daxberger, Agustinus Kristiadi, Alexander Immer, Runa Eschenhagen, Matthias Bauer, and Philipp Hennig. Laplace redux-effortless bayesian deep learning. Advances in Neural Information Processing Systems, 34:20089-20111 2021.
Laurent Dinh, Razvan Pascanu, Samy Bengio, and Yoshua Bengio. Sharp minima can generalize for deep nets, In International Conference on Machine Learning, pages 1019-1028. PMLR, 2017.
Serge Gratton, Ehouarn Simon, David Ptleey-Peloquin, and Philippe L Toint. Minimizing convex quadratics with variable precision conjugate gradients. Numerical Linear Algebra with Applications, 28(1):+2337, 2021.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE Conference on Computer Vision and Puttern Recognition, pages 770-778, 2016.
Jose Miguel Hermandez-Lobate and Ryan Adams. Probabilistic backpropagation for scalable learning of bayesian neural networks. In Francis Bach and David Blei, editors, Proceedings of the 32nd International Conference on Machine Learning, volume 37 of Proceedings of Machine Learning Research, pages 1861-1869, Lille, France, 07-09 Jul 2015. PMLR.
Magnus R Hestenes, Eduard Stiefel, al. Methods of conjugate gradients for solving linear systems. Journal of research of the National Bureau of Standards, 19(6):409-433) 1952.
Jeremy Howard. Imagenette: A smaller subset of 10 easily classified classes from ImageNet, March 2019. URL https: //tp:://gtttttttt
com wtt.l/cco/aaaaaeeeeeee wtt.///sttiiimaaeee iimaaeeett.mmee....em.mm
Jiri Hron, Alex Matthews, and Zoubin Ghahramani. Variational Bayesian dropout: pitfalls and fixes. In Jennifer Dy and Andreas Krause, editors, Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pages 2019-2028. PMLR, 10-15 Jul 2018. URL https: /tttts::/proceeddiiit
mlr. www.spressseeee https//vv/////o118//
http:/oronnna....ee.. html.
Michael F Hutchinson. A stochastic estimator of the trace of the influence matrix for Laplacian smoothing splines. Communications in 1989.
and Computation, 18(3):1059-100
Pavel Izmailov, Sharad Vikram, Matthew D Hoffman, and Andrew Gordon Gordon Wilson. What are bayesian neural network posteriors really like? International conference on machine learning. pages 4629-4640. PMLR, 2021.
Diederik P Kingma and Max Welling. Auto-encoding variational bayes. arXiv preprint arXiv:1312.61 14, 2013.
Agustinus Kristiadi, Matthias Hein, and Philipp Hennig. Being bit frequentist improves bayesian neural networks. In International Conference on Artificial Intelligence and Statistics, pages 529-545. PMLR, 2022.
Alex Krizhevsky and Geoffrey Hinton. Learning multiple layers of features from tiny images. Technical report, Department of Computer Science, University of Toronto, 2009.
Frederik Kunstner, Lukas Balles, and Philipp Hennig. Limitations of the empirical fisher approxima- tion. CoRR, abs/1 905. 10005.12558.. 12558, 2019. URL httt:http//aaaet .rg. hargg/bb//e/ /190. 12558.
Y. LeCun, B. Boser, J. S. Denker, D. Henderson, R. E. Howard, W. Hubbard, and L. D. Jackel. Backpropagation applied to handwritten zip code recognition. Neural Computation, 144:5541-551, 1989. doi: 10. 10.1166 1662neco..
11

---

Yann LeCun, Corinna Cortes, and CJ Burges. Mnist handwritten digit database. ATT Labs [Online]. Available: http://yamn.eecu www.lllcccccoo.o.. weccn.coo//xxb//miiss. 2, 2010.
Yingzhen Li, José Miguel Hernández-Lobato, and Richard E Turner. Stochastic expectation propa- gation. In C. Cortes, N. Lawrence, D. Lee, M. Sugiyama, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 28. Curran Associates, Inc., 2015.
David JC MacKay. A practical bayesian framework for backpropagation networks. Neural computa- lion, 4(3):488447) 472, 1992.
David JC MacKay. Probable networks and plausible predictions-a review of practical bayesian methods for supervised neural networks. Network: computation in neural systems, 6(3):469, 1995.
Wesley Maddox, Pavel Izmailov, Timur Garipov, Dmitry P Vetrov, and Andrew Gordon Wilson A simple baseline for bayesian uncertainty in deep learning. Advances in neural information processing systems, 32, 2019.
Wesley J Maddox. Andres Potapcynski, and Andrew Gordon Wilson. Low-precision arithmetic for fast gaussian processes. Uncertainty in Artificial Intelligence, pages 1306-1316. PMLR, 2022.
Marco Miani, Lorenzo Beretta, and Søren Hauberg. Sketched Lanczos uncertainty score: low- memory summary of the Fisher information. In A. Globerson, L. Mackey, D. Belgrave, A. Fan, U. Paquet, J. Tomczak, and C. Zhang, editors, Advances in Neural Information Processing Systems, volume 37, pages 23123-23155 Curran Associates, Inc., 2024.
Marco Miani, Hrittik Roy, and Søren Hauberg. Bayes without underfitting: Fully correlated deep learning posteriors via alternating projections. International Conference on Artificial Intelligence and Statistics, pages 2062-2070. PMLR, 2025.
Mahdi Pakdaman Naeini, Gregory Cooper, and Milos Hauskrecht. Obtaining well calibrated proba- bilities using bayesian binning. Proceedings of the AAAI Conference on Artificial Intelligence, 29 (1), 2015.
Yuval Netzer, Tao Wang, Adam Coates, Alessandro Bissacco, Bo Wu, and Andrew Y. Ng. Reading digits in natural images with unsupervised feature learning. In NIPS Workshop on Deep Learning and Unsupervised Feature Learning, 2011.
Augustus Odena, Vincent Dumoulin, and Chris Olah. tifacts. Distill, 2016. doi: 10.23915/distil..0000..
Deconvolution and checkerboard ar- URL http://distill pub/201//
deconv-checkerboart
Carl Edward Rasmussen and Christopher KI Williams. Gaussian processes for machine learning, volume 2. MIT press Cambridge, MA, 2006.
Hrittik Roy, Marco Miani, Carl Henrik Ek, Philipp Hennig, Marvin Pförtner, Lukas Tatzel, and Søren Hauberg. Reparameterizatioo invariance in approximate bayesian inference. In The Thirty- eighth Annual Conference on Neural Information Processing Systems, 2024. URL https://
openreview wttpw/iiiiieeeetttneeeeee
Hrittik Roy, Søren Hauberg, and Nicholas Krämer. Matrix-fre least squares solvers: Values, gradients, and what to do with them. arXiv preprint arXiv:2510. 19634, 2025.
Mrinank Sharma, Sebastian Farquhar, Eric Nalisnick, and Tom Rainforth. Do Bayesian Neural Networks Need To Be Fully Stochastic? In International Conference Artificial Intelligence and Statistics (AISTATS, notable paper award), 2023.
Yuesong Shen, Nico Daheim, Bai Cong, Peter Nickl, Gian Maria Marconi, Clement Bazan, Rio Yokota, Iryna Gurevych, Daniel Cremers, Mohammad Emtiyaz. Khan, et Variational learning effective for large deep networks. arXiv preprint ArXiv:2402.1777 17641, 2024.
John von Neumann. On rings of operators. reduction theory. Annals of Mathematics, 50(2):401-445, 1949.
12

---

Frederik Warburg, Marco Miani, Silas Brack, and Søren Hauberg. Bayesian metric learning for uncertainty quantification in image retrieval. In Neural Information Processing Systems (NeurIPS), 2023.
Sumio Watanabe. Algebraic geometry and statislical learning theory, volume 25. Cambridge university press, 2009.
Max Welling and Yee W Teh. Bayesian learning via stochastic gradient langevin dynamics. In Proceedings of the 28th international conference machine learning (ICML-II), pages 681-688,
2011.
Florian Wenzel, Kevin Roth, Bastiaan S Veeling, Jakub Swiętkowski, Linh Tran, Stephan Mandt, Jasper Snock, Tim Salimans, Rodolphe Jenatton, and Sebastian Nowozin. How good the bayes posterior in deep neural networks really? arXiv preprint arXiv:2002.02405, 2020.
Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-mnist: novel image dataset for benchmarking machine learning algorithms. CoRR, abs/1708.07744 2017.
Yijie Zhang, Yi-Shan Wu, Luis A. Ortega, and Andres R Masegosa. The cold posterior effect indicates underfitting, and cold posteriors represent fully bayesian method to mitigate it. Transactions Machine Learning Research, 2024. ISSN 2835-8856. URL https: http:///ppeeeeeeeeeeeeee wwwprrvviieenneeffrrre.e wwt/forunffr id=GZORXGxHIHT.
Zheng Zhao, Sebastian Mair, Thomas B Schön, and Jens Sjölund. On feynman-kac training of partial bayesian neural networks. In International Conference on Artificial Intelligence and Statistics, pages 3223-3231. PMLR, 2024.
13

---

Table 3:
Experimental results over three runs (mean and standard deviation) on in-distribution test data. MAP a point (model) estimate. VIKING and the post-hoc method from Miani et al. (2025) use the loss-projected variant. Ensemble is over five models.
Accuracy
Conf.1
NLL
ECE+
MCE+
MAP
0,986
0.00:
0.996
0,000
0.070
0,05
0.247 +: 0,011
0,861 + 0.045
VIKING (ours)
0.991
0.00
0.992
0,000
0.055
1: 0:000
0.096
1: 0.004
0.690
1: 0.102
MNIST
Miani et al. (2025)
0.949
0.000
0.813 1. 0.018
1.225 1 0.099
0.666
1. 0.007
0,894 1 0.011
IVON
0.989
0.00
0.990 = 0.001
0.043
=: 0.082
0.077
= 0.005
0.651
+ 0.042
SGLD
0.990
0.000
0.975
= 0.001
0.062
#: 0,002
0.067
#: 0,002
0.684
+ 0.054
SWAG
0.982
0.000
0.982
0.001
0.064
1: 0.006
0.788
1. 0,005
0.906
1. 0.013
Last Layer LA
0.975
0.002
0.977 0.002
0.090
0.784
0.887
Ansemble
0.992
0.989
0.078
0.116
0.675
MAP
0.883
0.002
0.942
6.003
0.410
1: 0.010
0.153 1. 0,008
0.611 1 0.590 1. 0.141
Fashion MNIST
VIKING (ours)
0.900 0.001
0.928
0.332
0.075
Miani et al. (2025)
0.871 0.06
0.744
0.031
1.529 0:371
0.617
0.901
+: 0.013
IVON
0.897
0.926
0.335
0.073 0,005
0.683 0.004
SGLD
0.899 0.000
0.899
0.001
0.316
0,083
0.026
0,004
0.734 1 0.019
SWAG
0.898 0.001
0
+ 0,006
0.327
+: 0.001
0.725 1: 0.008
0.907 1: 0.005
Last Layer T.A
0.896 = 0.00
0.931
= 0.006
0.339
+ 0.011
0.727 + 0.004
0,902 1 0.004
Ensemble
0.910
0.913
0.670
0,052
0.793
MAP
0.947 0,004
0.963 + 0.004
0.201
+: 0.014
0.055 +: 0.0L0
0.608 J: 0.228
VIKING (ours)
0.960 0.00
0.964 = 0.001
0.177
= 0.002
0.028 :: 0.002
0.308 1: 0.024
Miani et al. (2025)
0.949 0.00
0.948 = 0.005
0.191 #: 0.007
0.734 + 0.017
0.880 +: 0.012
SVHN
IVON
0.943 0.002
0.888
0.007
0.302 1. 0.016
0,082 1. 0,004
0.492 1. 0.248
SGLDD
0.698 0.102
0.378 0.027
1.935
0.070
0.321
0.605
SWAG
0.947 0.04
0.897 0.00
0.217 0.014
0.745
0.874 0.003
Last Layer LA
0,946
0.943
0.197
0.740
0.899
Ensemble
0.965
0.955
0.231
0.051
0.769
MAP
0.824 0012
0.869
0.536 0.055
0.075
0.619
VIKING (ours)
0.877
0.893
0.407 0.000
0.041
0.331 + 0,802 + 0.005
CIFAR- 10
Miani et al. (2025)
0.855 0.002
0.701 0.013
2.643 0.205
0.559 0.006
IVON
0.835 Q.017
0.763 o.u05
0.817
= 0,075
0,086 + 0.014
0.436 + 0.244
SGLD
0.495 0.033
0.345 1: 0.020
2,409 1: 0.138
0,153 1 0,009
0,375 1 0.073
SWAG
0.865
0.029
0.914 = 0.835
0.445
+ 0.063
0.694 + 0.018
0.881 +: 0.005
Last Layer LA
0.894 0.001
0.944 0.001
0.406
+ 0.005
0.704 = 0.000
0.880 :: 0.007
Ensemble
0.919
0.891
0.486
0.064
0.780
Imagenette
MAP 0.852
0.002
0.883 = 0.807
0.481
#: 0.009
0.084 +: 0.010
0,717 +: 0.062
VIKING (ours)
0.887
0.003
0.906
= 0.002
0.403
0.010
0.077
+: 0,001
0.612 4: 0.362
IVON
0.876
0.01
0.849
0.013
0.656 : 0.136
0.069 : 0.011
0.464 : 0.230
A Additional experimental results
In Table 3 we detail the same experiments from the main paper, but with two additional basclines. The first is a simple deep ensemble, where five copies of the model are trained independently and their predictions are averaged. Here, the best-performing MAP model hyperparameters from the existing experiments are used to train the models again using additional seeds to build the ensemble.
The second is Stochastic Gradient Langevin Dynamics (SGLD, Welling and Teh, 2011), which is implemented in JAX using code provided by Izmailov et al. (2021). The hyperparameters are opti- mized using a grid search with a similar budget we use for VIKING, using experimental evidence by Izmailov et al. (2021) as a guideline. When collecting posterior samples from SGL.D, we allowed for burn-in period of 50 epochs, then collected each posterior sample after additional 50 epochs are done, totalling 300 epochs, which approximately corresponds to 50 epochs of ELBO training with VIKING.
These baselines are included as additional evidence that, given similar compute budgets, VIKING outperforms classical approaches such as SGLD and deep ensembles in calibration metrics, while still achieving competitive performance. Notably, deep ensembles achieve the best predictive performance, but that is not reflected in the calibration metrics. SGLD, on the other hand, struggled in the harder datasets with bigger models (SVHN, CIFAR-10), showing that the compute budget available is probably not enough, needing either significantly larger number of training steps to converge or hyperparameter tuning. When does properly converge, such as on Fashion MNIST and MNIST, its performance is not far behind the state of the art.
14

---

B Additional methodological details
Here expand the rationale and discussion present in the main paper. Our aim is address potential questions, issues, and minor challenges pertaining the formulation and use of our varia- tional family. For notational simplicity, omit the dependence on 0 in the analysis below where convenient.
B.1 the rank differentiable?
We know that the derivative of the rank R with respect to 0 is either zero undefined, as R is integer- valued and bound be piecewise-constane Its importance optimizing the KI. (Kullhack-Leibler divergence) term is clearer by first assuming that 8R/a0 O. Then,
OKI. 0
(TTr(E) - D + a a|l@|* --D Doo((a) log det logdet
(ec(((t)) }
00
00
2
0
{1
1
(alo R o R3) a a0022 2R log( 12(D- R) log log(Uim)
00
08 2 [10113]
a
(19)
Under this assumption, the KL term reduces a more usual Eg regularization. In practice, since the estimated R is not an integer, this assumption does not become problematic. In our implementation, we stop the gradients of the computed to ensure this.
B.2 Understanding the KL term in isolation
Let us consider the kernel component of the posterior in isolation and the KI. term of the ELBO & (Equation 3) separately to interpret what signal might exist in that term and why that influence our hyperparameter choices.
Suppose our posterior covariance is given by E, E DR and our prior covariance is given by Σp RDxD Due to the restrictions to the subspace, we can write E, UAUT and Σ, a a-I We want to find the A that minimizes the KI. term. We can write the terms that depend only on Σ, as
=T= = Tr - log E Egl-
Note that
Σ -, = =0 (T- -I UAUT = =U(AA)UT,
implying the eigenvalues of Ep11 Ey are given by aA, yielding
Tr(U((A)U') = Σ adi, and log U(aAUaa)'' Σ log(a>.).
i=1
For all i 1,..., D, a minimizes S(ade) ad: log(ads) The covariance that
minimizes the KI. divergence is thus
3
(20)
Since our posterior is Eg = 22 UUT 2 (1- UUT, we optimize a and o jointly by always ensuring
oeer + 1/o2ur
C Experimental details
For reproduction purposes, we detail each experiment (in order of appearance in the main paper). Note that the code implementing VIKING and for reproducing the experiments publicly available'
phttps://gitaiii wwtp://ooe///veiiiiii.... paper- paper-expeeeeeeeeee experiments
15

---

We start with few general experimental and implementation details used across the experiments. All optimization steps were done using the Adam optimizer. In what follows, we address specific components of the optimization that are shared across all experiments, unless otherwise noted.
oer and oi- all experiments, the variational parameter oee tied to the prior precision a, such that a 1ol is always true (for the reasoning, see Equation 20, in Section B) and use log a as an optimization variable. Furthermore, we observe that the posterior samples used to evaluate the expectation term of the ELBO can be quite sensitive to larger values of Jim- especially since only the kernel part of the projection guaranteed to not deviate arbitrarily from the MAP predictions. For this reason, we always initialize this variational parameter such that log Jm = 2 to prevent it from having a large influence in the posterior samples in the beginning, resulting in more stable optimization.
Variational parameters
Practical optimization with the KL term. Experimentally, we find that the ELBO as formulated can be quite often dominated by the KL term during optimization. This is mostly exacerbated when the number of model parameters grows. For this reason, in all experiments, we follow common procedure and instead optimize
C(O, keer Fim) Begggloogpe 0, x)] (KL))()))p((O))
(21)
where 0 <<p < < 1 (fixed) controls the relative strength of the KL. term.
Reorthogonalized conjugate gradients (CG). As discussed in Section 3.2, we use a variant of conjugate gradients that improves numerical stability by reorthogonalizing the basis vectors every iteration. Aiming for a good trade off between numerical robustness and performance, we use 10 iterations in all CG solves.
Posterior sampling. For computing the expectation term of the ELBO during training (either while optimizing just sigmas or the full model as well), S= 1. Empirically, we observe little gain over using multiple samples and this can be compensated by training for longer, with the benefit of accelerating alternating projections and reducing memory usage significantly. When evaluating on validation and test sets, we draw 20 posterior samples.
single sample from posterior, that is,
C.1 Architectures and data preprocessing
Unless otherwise noted, all data is standardized, that is, every data point is first subtracted by the mean of training data and the result is divided by the standard deviation of training data, both computed across dimensions. Naturally, this is performed also for the validation and test sets. On MNIST and Fashion MNIST, we train a LeNet; on SVHN and CIFAR-10, 10, we train a small ResNet.
LeNet architecture. Composed of two 5 x 5 convolution layers with 6 and 16 channels, respectively, each followed by 2 x 2 max-pooling with stride 2, then three fully fully-connected layers with 128 and 80 units, with the last one being the number of classes units; all layers have hyperbolic tangent activation functions. The total number of parameters is 44,426.
ResNet architecture (small variant). Composed of (3, 3, 3) residual blocks of channel sizes 16, 32, and 64, respectively; all layers have ReLU activation functions. The total number of parameters is 272.378.
C.2 Ablation experiments (Section 3)
This section details the experiments illustrated by Figures 2 to 4 in the main paper.
Stochastic alternating projections. For each y 0.0, 0.5, and 1.0, we train the LeNet model specified above on the Fashion MNIST dataset. Table 4 details the other hyperparameters used.
Post-hoc variational inference of oeer and Fim- Since we assess the need for full model ELBO optimization versus only post-hoc tuning of the variational parameters, we need to start from reasonably-performing model. In both settings, we thus begin with training the model using maximum
16

---

Table 4: Experiments evaluat-
Table 5: Post-hoc oeer Olm eX- periments (Section 3.4).
Table 6: Warmup experiments (Section 3.4).
ing 7 (Section 3.3).
Hyperparam. Value
Hyperparam.
Value
Hyperparam.
Value
10-4
10-1
B
10-4
Batch size 32
?
0.5
?
0.5
Init. log 4
Batch size
32
Batch size
32
Train. epochs 50
Init. log a
4
Init. log at
4
Learn rate
10-3
Warmup epochs
20
Learn. rate (MLE)
10
Learn rate
10-1
# epochs (ELBO)
50
Learn. rate (ELBO)
10
likelihood for 20 epochs, with the same batch size later, but a learning rate of 10-5 instead to prevent an early overfitting, Following that, the post-hoc setting optimizes the variational parameters (oker and mm) for 20 epochs. In the other setting, in order to draw a better parallel and more clearly highlight the benefits of full model optimization using our ELBO, we tune only the sigmas for epochs. For those epochs, the performance is thus comparable. We then continue with 15 more epochs where the model is also trained, for a total of also 20 epochs. Table 5 shows the values of all other hyperparameters used.
Warming up as a shortcut. The only change across different runs is the number of epochs the model was trained with maximum likelihood at the start (warmup): 0, 20, or 50. Table 6 contains the values of the other hyperparameters.
C.3 Main experiments
The experiments in Section 4 require further details for full reproducibility, given below.
C.4 Toy (sinusoid) regression
In these experiments, we train a neural network with three fully-connected layers, two with 10 units and one with a single output unit. The training data consists of 10 samples from the sinusoid y = 5innn(10z) + 2, where consecutive points in specified range of x are drawn, but only the first and last five are used, introducing the visible gap. The VIKING experiment uses the model Jacobian (not loss Jacobian as every other experiment) and a linearized predictive" to highlight the properties of guaranteed zero variance in the predictions when projecting onto the kernel of the GGN of the model on training data, The model is trained for 2,000 epochs with a learning rate of 10-2 and adaptive gradient clipping of rate 10-1, using 100 posterior samples each step. The log a prior (log oue -1/2 log a. kept tied as described earlier) is initialized with zero. For IVON, we train the same model for 5,000 epochs with the same learning rate, using posterior samples at each training step. The ESS hyperparameter is set to 10 (number of data points) and HESS-INIT is set to 15. The plots in Figure 5 depict the predictions of 100 posterior samples.
[0.35, 0.65).
(O, s), and s [10-1 -3 1]. from which 20
C,5 Image classification
Grid search. On MNIST, Fashion MNIST, SVHN. and CIFAR-10, 10, both VIKING and IVON were subjected to a grid search. On VIKING, the grid search is over batch size (16 or 32), number of warmup epochs (Section 3.4; 20 or 50), learning rate (10-3 3 or 10 4), (10 10 or 10 10-6) 0), and (0.2, 0.5. or 0.8). For IVON, the grid search is over batch size (16, 32, or 128), number of epochs (50 or 100), and learning rate (10 or 10 10-4). Table 7 breaks down the best hyperparameters per method and dataset.
3
WWit 0 q() denoting posterior sample, instcad of making predictions as y (x)),, we lincarize the model around è as y fa(x) +o (x) 0).
17

---

Table 7: Hyperparameters broken down per dataset. "MC samples" indicates the number of posterior samples used for training.
MNIST
F-MNIST
SVHN
CIFAR-10
Batch size
32
16
32
32
Warmup epochs
50
20
50
20
Learning rate
10
10-4
10
10
B 10-6
10-3
10-6
10-6
VIKING
Y
0.8
0.2
0.8
0.5
MC samples
Warmup learning rate
10-3
Initial log or Initial log din
4.0
2.0
Ker Im tuning
epochs 50 epochs
ELBO optimization
Batch size
16
16
16
16
Epochs
100
100
100
100
IVON
Learning
10-3
10-3
10-3
10-3
MC samples
ESS
Number of data points
HESSINIT
1.0
Table 8: Hyperparameters for the ResNet-34 experiment Imagenette. "MC samples" indicates the number posterior samples used for training. In VIKING, epochs and learning rates are split between warmup and ELBO, respectively.
VIKING
IVON
Batch size
128
128
Epochs (warmup, ELBO)
100, 50
150
Learning rate
10 20
10-1
10-7
N/A
0.2
N/A
MC samples
1
5
Large models. In this experiment, investigate the applicability of VIKING when training ResNet-34, whose number of parameters is 21,797,66 and compare it against IVON, Due to a limited compute budget with larger models, we do not perform extensive grid search. Table summarizes all hyperparameters, organized per method.
C.6 Generative modelling
Model architecture. Encoder: with five x 4 convolutional layers, the first two with 128 channels, followed by the next three with 256 channels, with the first and last having stride of one and the others a stride of two; two fully-connected layers, one with 256 units and the second with 64 units (latent dimensions) for the VAE approximate posterior mean and 64 for the VAE approximate posterior variance. Decoder: two fully-connected layers, one with 256 units and one with 16384 units; five 4 4 resize convolutional layers (Odena al., 2016), with the same stride configuration as the encoder, but with 256, 128, 64, and 32 channels, respectively. All layers in the encoder and decoder have ELU activation function. Total number of parameters: 6.5 million.
Maximum likelihood training. We train for 500 epochs with a maximum learning rate 10-3 and batch size 256 on raw data without augmentation. Furthermore, we used the Adamax optimizer and scheduled 1000-step (updates) linear warmup the maximum learning rate and then exponentially decayed it.
18

---

VIKING finetuning, From the maximum likelihood checkpoint, we train for 10 VIKING epochs using 8 posterior samples, Y 0.5, B 10 batch size 8, learning rate 10 and adaptive gradient clipping at 0.1. The value of log is 7.0 and log dim is -5.0.
IVON finetuning. From the maximum likelihood checkpoint, we train for 20 IVON epochs using Monte Carlo samples, using a batch size of 256 and learning rate of 10
19

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
The abstract predominantly discuss the problem and our methodology all of which appropriate. In terms of performance claims, approach is generally par or better than baselines as reported in the experiments.
Justification:
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results be expected to generalize other settings.
is fine include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
The last section of the paper has an explicitly annotated section concerning
limitations.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and prool's
Question: For each theoretical result, does the paper provide the full of assumptions and complete (and correct) proof?
20

---

Answer:
[NA]
Justification:
The paper does not contain theorems.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proof's can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided not)?
Question:
Answer: [Yes]
Justification: To the best of our ability, provide model and training details (predominantly in the supplements). We will further release code for full reproduceability.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, descríbing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e-g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing verifying the results.
5. Open access to data and code
21

---

Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes]
Justification:
We use standard benchmark data that is openly accessible.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips. cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" an acceptable answer. Papers cannot be rejected simply for including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https:
/ /nwpp...c////////o https/cccpplcccgudiccoeooo http:p/bliiiuuuieeeeeeeeeeuttaia
/Cooeeuboeeeeeeeeee cetsSSuumissiiiiio wwwiissssnnoiiiiieieoeeooo onPolicy) for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing much information possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: The supplementary material provides these details. We will further release code for full transparency.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can provided either with the code, in appendix, or as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: All quantitative experiments provide error bars to the extend possible. Some experiments designed to convey intuitions; these do not provide error bars.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
22

---

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
The supplementary material provides these details; All experiments are carried out conventional hardware (x86 architecture, NVIDIA GPU's).
Justification:
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental rums well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: https://neurd: /ttp::///eeuisueeeeeee ppuublic/thiissuiddeee
Answer: [Yes]
Justification:
The paper develops new methodology.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No. they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [No]
Justification: The paper is about improving general Bayesian model fitting. The method we propose has the same positive (as well as negative potential) other machine learning technologies.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
23

---

Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), faireess considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper deseribe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: Not applicable.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: Datasets are cited.
Guidelines:
The answer NA means that the paper does use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website). the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms use in package should be provided. For popular datasets, paperswithcode. thcode. com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.

---

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer:
[No]
Justification:
No assets introduced.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: Not applicable.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Bthies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: Not applicable.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, not include any information that would break anonymity (if applicable), such the institution conducting the review.
25

---

16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology. scientific rigorousness, originality of the research, declaration is not required.
Answer: [NA]
Justification: Such topics are not in scope for this paper.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to L.LM policy (https: http://eeurissspee heaprrpps.c///ooeeeeenneee////// for what should or should not be described.
26


Reviews:
[{"review_id": "Wdqf8VNL2n", "reviewer": "Reviewer_boCo", "summary": "This paper is motivated by the insufficiency of existing solutions in approximate inference for overparameterized deep neural networks. Aiming to address the pathologies of such methods in model training, predictive performance and calibration, the authors propose a novel variational family that decomposes into two orthogonal subspaces in the parameter space, spanning the in- and out-of training data supports. After developing a copious iterative algorithm to compute the corresponding posterior approximation, they demonstrate empirically that the novel scheme can often achieve higher predictive power and uncertainty quantification compared to state-of-the-art baselines across toy and image datasets in regression, classification and generative modeling tasks.", "questions": "- Projection matrix $U$ seems to be pre-maturely introduced into the notation and downplayed over experiments. I would recommend explaining it further and seeking for further insights in it over the experiments. Can it leverage potential sparsity on a data- and model-specific basis in practice?\n- It would be nice to see the effects of the data geometry on the efficacy of the VI method. Is it crucially benefited by clean-cut manifolds, as possibly the ones arising in image data? How would UQ behave in blurred versions of the images?\n- Practitioners and experts would benefit from a more thorough discussion of the computational complexity of the method, and potentially a time-comparison against considered baselines across the experiments.\n- The new method does not consistently outperform baselines in predictive and calibration performance. How are these findings explained? What makes SVHN a better benchmark to demonstrate the advantages of the methods compared to MNIST and Fashion MNIST?\n- A discussion on limitations of extending the method towards language models would help the reader get a clearer understanding of the scope and generalizability of the work.", "limitations": "Yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 2, "strengths_and_weaknesses": "**Strengths**\n\nThe writing of the paper is mostly clear. It begins with an engaging introduction to the problem, that offers some useful information-geometric intuition underneath the emerging pathologies in uncertainty quantification via approximate inference methods in overparameterized models. This leads to a clear motivation of why the examined problem is significant and paves the way towards the development of the proposed new variational family. Despite appearing as a simple contribution that relies on a fundamental idea of orthogonal decomposition and a minimalistic number of learnable variational parameters, Kernel-Image VI offers clear geometric interpretability. Moreover, it entails an involved computational scheme and can be efficiently combined with warm-up MLE-based training stages offering uncertainty quantification as an add-on on scalable training methods. \n\n**Weaknesses**\n\nThe paper has several areas that could be clarified and expanded to better support its contributions$-$consult also the following section on Questions. Notably, the projection matrix $U$ is introduced early in the formulation but is insufficiently motivated and receives limited attention in the experimental analysis. Additionally, the impact of data geometry on the proposed variational inference method remains underexplored. It would be informative to assess how the method performs under varying levels of manifold clarity, such as comparing clean-cut image structures with their blurred counterparts. The computational complexity of the approach is also not addressed in detail; a comprehensive analysis, including empirical runtime comparisons with baseline methods, would be beneficial for practitioners. Furthermore, the proposed method does not consistently surpass existing baselines in terms of predictive accuracy or calibration, and the reasons for these findings are not adequately discussed. Lastly, the manuscript would benefit from a critical examination of the limitations involved in extending the methodology to language models, thereby providing a clearer perspective on its generalizability.", "was_revised": true, "final_justification": "After reading the authors' response, I am increasing my final score to \"Accept.\" I appreciate the technical contribution of the work on overparameterized Bayesian Neural Networks and believe it meets the conference's standards for evaluation, reproducibility, and ethical considerations. I am agreeable to the authors' proposed edits regarding the discussion of computational complexity, though I believe it would be appropriate to include them in the appendix.", "rebuttal": "We thank the reviewer for the thoughtful comments. Below, we will reply to each point individually.\n\n> Projection matrix \\(U\\) seems to be pre-maturely introduced into the notation and downplayed over experiments. I would recommend explaining it further and seeking for further insights in it over the experiments. Can it leverage potential sparsity on a data- and model-specific basis in practice?\n\n\nThanks for raising this point. (Please note that the projection matrix onto the kernel (null) space is $U U^\\top$.) What we are interested in, based on existing theoretical developments on model reparametrization, is to project parameter vectors onto the kernel space of the model GGN, as it provides then a clear space of parameters which follow certain properties, e.g., maintain model predictions on data points used to build the stacked Jacobian $\\mathbf{J}_{{\\theta}}$. Indeed, any solution (or approximation) to Eq. 13 would be particularly interesting, as it would follow those properties, but without instantiating the projection matrix, since it is not tractable. \n\n\n> It would be nice to see the effects of the data geometry on the efficacy of the VI method. Is it crucially benefited by clean-cut manifolds, as possibly the ones arising in image data? How would UQ behave in blurred versions of the images?\n\nWe do not see reasons for our method to be dependent on data geometry or a specific task as we are focusing on uncertainty in overparametrized models. Thus, the amount of uncertainty will not influence the ability of our method to detect it. In other words, if the nature of the data leads to a higher uncertainty, we certainly expect this to be reflected in the $\\sigma^2_{\\text{ker}}$ and $\\sigma^2_{\\text{im}}$ parameters. \n\n> Practitioners and experts would benefit from a more thorough discussion of the computational complexity of the method, and potentially a time-comparison against considered baselines across the experiments.\n\nWe propose adding the following paragraph to highlight the linear dependency on every hyperparameter in our method and clarifying the total computational complexity.\n\n> ### Proposed changes:\n> **Computational complexity.** Assume a Jacobian-vector-product and a vector-Jacobian product through the model function at a single data-point each cost $\\mathcal{O}(D)$. Then, for batch-size $B$, the least-squares solver needs $BN_{\\text{lstsq}}$ of these matrix-vector and vector-matrix products to implement the projections. For each of the $S$ Monte Carlo samples (Equations 8-12) samples, a standard-Gaussian sample ${\\epsilon}^{(s)}$ is turned into a kernel-projection sample ${\\epsilon}\\_{\\text{ker}}^{(s)}$ in complexity $\\mathcal{O}(B N\\_{\\text{lstsq}} D)$. Afterwards, all $S$ kernel-projection samples are linearly combined to estimate the reconstruction-term $\\mathbb{E}\\_{{\\theta} \\sim q}\\left[ \\log p(y | {\\theta}, x) \\right]$ and the KL-term $\\mathrm{KL}( q({\\theta}) \\| p({\\theta}) )$. Therefore, a single ELBO-evaluation costs $\\mathcal{O}(SBN\\_{\\text{lstsq}} D)$. Our experiments solve the least-squares problem using $N_{\\text{lstsq}}=10$ Jacobian-vector and vector-Jacobian products, varying batch-sizes, and $S=1$ Monte Carlo samples. See Appendix B for details.\n\nWe additionally summarize the running times (total execution time from model initialization until the end of the last training epoch, in minutes) below, comparing IVON and KIVI:\n\n| Dataset/Architecture  | KIVI           | IVON          |\n|-----------------------|----------------|---------------|\n| MNIST/LeNet           | 27.27 ± 3.96   | 9.97 ± 0.35   |\n| FMNIST/LeNet          | 46.30 ± 7.08   | 9.71 ± 0.06   |\n| SVHN/ResNet_small     | 304.16 ± 0.96  | 108.87 ± 1.90 |\n| CIFAR-10/ResNet_small | 177.08 ± 11.68 | 70.17 ± 0.88  |\n\n\nAs shown by the table, KIVI takes from 2-5 times the training time of IVON, where aspects such as alternating projections (and whether the stochastic variant is used) and the number of MLE steps used for warmup can affect the total execution time.\n\n> The new method does not consistently outperform baselines in predictive and calibration performance. How are these findings explained? What makes SVHN a better benchmark to demonstrate the advantages of the methods compared to MNIST and Fashion MNIST?\n\nThe takeaway from, e.g, Tables 1 and 2 is that KIVI outperforms methods like SWAG and Miani et al. consistently, and performs on par with IVON and last-layer Laplace across benchmarks, beating them on SVHN. We don't claim that SVHN is a better benchmark than MNIST or Fashion MNIST (if there are misleading statements in the paper, please let us know and we'll correct them!).\n\nNotably, we emphasize that on Accuracy, NLL, and MCE (two \"quality-of-fit\" and one calibration metric), our method consistently performs best or second-best, while this level of reliability is not observed in the baseline methods. We emphasize that consistency is one of the key challenges facing BNNs.\n\n> A discussion on limitations of extending the method towards language models would help the reader get a clearer understanding of the scope and generalizability of the work.\n\nThank you for raising this crucial point. In principle, our method applies to large language models. Since we use implicit solvers to compute projections, our method has the same space complexity as standard optimization. While the projections introduce an additional computational overhead, the computational complexity is linear in parameter size, batch size, and the number of Monte Carlo samples. So training is only slower by a constant factor. Hence, it is feasible to use our method whenever it is feasible to train a model. \n\nThanks again for the thorough review! We hope that all concerns have been addressed adequately, and look forward to the discussion."}, {"review_id": "p4Vyvy0rWk", "reviewer": "Reviewer_rdDS", "summary": "This paper proposes KIVI, kernel image variational inference, which is a full rank Gaussian posterior approximation with a covariance that is a full rank basis constructed around the pseudo inverse of the Jacobian. Experiments are performed on toy regression tasks as well as smaller image based deep nets.", "questions": "I believe that the first use of conjugate gradients to solve systems of the form $(J’J)^{-1}J’v$ in a matrix free fashion is from Maddox et al, ’21 (https://arxiv.org/pdf/2103.01439), which does yield an interesting question in terms of the relationship between the posterior approximation and the learnable local Laplace approximations (e.g. Immer et al, ’20 [https://arxiv.org/abs/2008.08400], Antoran et al, ’22 [https://arxiv.org/abs/2210.04994] beyond just the ones cited in your work). In general, there’s a pretty deep connection to the first order Taylor approximation / linearization of the neural network and thus things like the local Laplace approximation and the neural tangent kernel.\n\n I would generically think that your approach boils down to training a predictive mean through something like the linearization in Maddox et al, ’21 rather than simply basing the covariance estimate off of a pre-trained model like they did. I wonder if the authors could connect their work to these types of linearizations a bit more?\n\nEq 2 – why do the authors not instead optimize $\\Sigma = \\sigma_K^2 UU’ + \\sigma_n I$ as $\\sigma_K = (\\sigma_{ker}^2 - \\sigma_{im}^2)$ in your parameterization and we need to ensure that the covariance is psd so $\\sigma_K^2 >= 0$ by construction.\n\nL147 / Eq 13: for further scalability, why not use SGD (or a cheap second order) method to minimize that loss (possibly with the relaxed constraint that \\lambda is small enough)? This should only require Jacobian vector products instead of Jacobian matrix solves..\n\nEq 16: P_t is not really defined previously. What is the “kernel-projection of the t th 168 batch”? Is this the reorthogonalization matrix from the previous solve?\n\nHow accurate (in terms of forwards error) are the CG solves after only 10 steps?\n\nFigure 5: why are the data points different for IVON and KIVI?\n\nTables 1 and 2: I’m generally a bit surprised at the high ECE values as these are much higher than whats reported in Table 3 of Maddox et al, ’19 (the original SWAG paper) but on checking Table 2 of https://arxiv.org/pdf/2410.16901 I see that they are the same. Could the authors comment on the difference (I believe that this is architectural but would like confirmation)?\n\nWriting comments:\n\n-\tL160:  “- we found.. “ turn into a new sentence.", "limitations": "n/a", "rating": 3, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Strengths:\n\nThe method is well described including all of the computational tricks that are needed to make it work.\n\nWeaknesses:\n\nNo discussion of the computational expense, which I believe should be quite expensive due (at a minimum 3x backwards * the number of CG steps per backwards passes). Scaling in my mind is pretty important here as it prevents doing some interesting experiments here either using the probabilistic predictions somehow or in doing optimization from large pretrained models.\n\nI believe that Appendix A.1 yields a slightly sub-optimal gradient estimate, which should probably be discussed in the main text. \n\nThe missing comparison to me is against a low rank (say rank 10) plus diagional (SLANG I believe, https://arxiv.org/abs/1811.04504) variational Gaussian posterior that just directly optimizes U instead of trying to structure it heavily. Additionally, I would probably argue for a MCMC based comparison like SGLD / SGHMC as those tend to be relatively cheap to run as well.", "was_revised": true, "final_justification": "I thank the authors for responding to the comments during the rebuttal period. Overall, I still think that the paper is quite borderline and further discussion with prior work (as I pointed out in the original review) and a comparison with deep ensembles/sgmcmc (rather than just vi methdods) is needed for acceptance.", "rebuttal": "Thank you for that thoughtful assessment. In the following, we will first reply to the weaknesses before answering the specific questions.\n\n> No discussion of the computational expense, which I believe should be quite expensive due (at a minimum 3x backwards * the number of CG steps per backwards passes). Scaling in my mind is pretty important here as it prevents doing some interesting experiments here either using the probabilistic predictions somehow or in doing optimization from large pretrained models.\n\n Thanks for bringing this up. We agree that providing more explanation of the complexity would improve the presentation. We've now added the following paragraph to highlight the linear dependency on every hyperparameter in our method and to clarify the total computational complexity.\n\n> ### Proposed changes:\n> **Computational complexity.** Assume a Jacobian-vector-product and a vector-Jacobian product through the model function at a single data-point each cost $\\mathcal{O}(D)$. Then, for batch-size $B$, the least-squares solver needs $BN_{\\text{lstsq}}$ of these matrix-vector and vector-matrix products to implement the projections. For each of the $S$ Monte Carlo samples (Equations 8-12) samples, a standard-Gaussian sample ${\\epsilon}^{(s)}$ is turned into a kernel-projection sample ${\\epsilon}\\_{\\text{ker}}^{(s)}$ in complexity $\\mathcal{O}(B N\\_{\\text{lstsq}} D)$. Afterwards, all $S$ kernel-projection samples are linearly combined to estimate the reconstruction-term $\\mathbb{E}\\_{{\\theta} \\sim q}\\left[ \\log p(y | {\\theta}, x) \\right]$ and the KL-term $\\mathrm{KL}( q({\\theta}) \\| p({\\theta}) )$. Therefore, a single ELBO-evaluation costs $\\mathcal{O}(SBN\\_{\\text{lstsq}} D)$. Our experiments solve the least-squares problem using $N_{\\text{lstsq}}=10$ Jacobian-vector and vector-Jacobian products, varying batch-sizes, and $S=1$ Monte Carlo samples. See Appendix B for details.\n\nWe summarize the running times (total execution time from model initialization until the end of the last training epoch, in minutes) below, comparing IVON and KIVI:\n\n| Dataset/Architecture  | KIVI           | IVON          |\n|-----------------------|----------------|---------------|\n| MNIST/LeNet           | 27.27 ± 3.96   | 9.97 ± 0.35   |\n| FMNIST/LeNet          | 46.30 ± 7.08   | 9.71 ± 0.06   |\n| SVHN/ResNet_small     | 304.16 ± 0.96  | 108.87 ± 1.90 |\n| CIFAR-10/ResNet_small | 177.08 ± 11.68 | 70.17 ± 0.88  |\n\nAs the table shows, KIVI takes from 2-5 times the training time of IVON, where aspects such as alternating projections (and whether the stochastic variant is used) and the number of MLE steps used for warmup can affect the total execution time.\n\n> I believe that Appendix A.1 yields a slightly sub-optimal gradient estimate, which should probably be discussed in the main text.\n\nYou are correct that not differentiating through the rank estimation leads to an inexact gradient. We opted to stop the gradient because the rank itself is not differentiable (only its estimate is), and we expected instabilities. We will emphasize this inexactness. We have investigated whether this choice matters, and found that it made no difference whether we differentiate through the rank estimation or not.\n\n> The missing comparison to me is against ... [SLANG/MCMC]\n\nThank you for the suggestions. While our focus has been on methods that are generally considered state-of-the-art for BNNs, we agree that these are interesting baselines, with the following reservations:\n* The general experience is that 'regular' variational methods for BNNs are quite unstable, leading to either underfitting or overfitting depending on initialization and choice of hyperparameters. For mean field approaches, this experience has been the driving motivation for IVON and similar methods. In our experience, SLANG suffers from similar issues, and we worry that reporting such results would not be informative, as highly different results could easily be obtained by changing initialization/hyperparameters. We are open to including such results, but they would come with notable caveats.\n* While SGMCMC methods are cheap to run, they can be quite difficult to tune and may not converge to reasonable samples [1, 2]. For this reason, SGMCMC methods are not commonly used for BNNs. As above, we are open to including results obtained with SGMCMC, but similar caveats would apply.\n\n[1] Bardenet, Rémi, Arnaud Doucet, and Chris Holmes. \"On Markov chain Monte Carlo methods for tall data.\" Journal of Machine Learning Research 18.47 (2017): 1-43.\n\n[2] Brosse, Nicolas, Alain Durmus, and Eric Moulines. \"The promises and pitfalls of stochastic gradient Langevin dynamics.\" Advances in Neural Information Processing Systems 31 (2018).\n\n\n**Next, we reply to the direct questions:**\n\n> I believe that the first use of conjugate gradients... \n\nYes, these are indeed related. The submission cites linearized Laplace approximations in Section 2. We'll add the papers you list in this section as we agree they should be mentioned.\n\n> I would generically think that your approach boils down to training a predictive mean ...\n\nThat is correct, we train (without MLE warm-up) from scratch using the ELBO. Aside from the differences in the covariance structure from our choice of variational family, concretely, we optimize the reconstruction (expectation) part of the ELBO using linearization around the current posterior mean, with the posterior samples parametrizing the linearized model. Thus, the posterior mean (when optimized with the ELBO) is allowed to be influenced by the covariance and not simply optimized in isolation with UQ done post-hoc, since the image component would also allow predictions to deviate from the mean predictions under the training data.\n\n> Eq 2 – why do the authors not instead optimize ...\n\nThank you for that suggestion. We keep $\\sigma_\\text{ker}$ and $\\sigma_\\text{im}$ separate to give different weights to image- and kernel-spaces, respectively. In the current parametrisation, both $\\sigma_\\text{ker}^2 U U^\\top$ and $\\sigma_\\text{im}^2(I - U U^\\top)$ are PSD (as projection matrices), so their sum is PSD. Does that help?\n\n> L147 / Eq 13: for further scalability, why not use SGD (or a cheap second order) method to minimize that loss (possibly with the relaxed constraint that \\lambda is small enough)? This should only require Jacobian vector products instead of Jacobian matrix solves..\n\nThanks for that suggestion! We'd like to emphasise that Equation 13 has a closed-form solution, hence we use conjugate gradients, not SGD. One could attempt to minimise the loss with something like SGD, but it's unclear whether it improves scalability (solving the system with 10 CG steps vs minimising with 10 SGD steps has similar complexity), and we need to ensure that the solves are of high quality. When we lose orthogonality of samples, training becomes very unstable. Also, differentiating through full SGD runs tends to be rather slow.\n\n\n> Eq 16: P\\textsubscript{t} is not really defined previously...\n\nWe apologize for the confusion caused by the notation. We have updated the text in line 168 to explicitly mention that $P^{(t)} = {U}{U}^{\\top(t)}$, where $t$ denotes Eq. 13 restricted for the $t$-batch (w.r.t. $\\mathbf{J}_{{\\theta}}$). Following the suggestion of another reviewer, we also include $\\hat{\\theta}$ as a subscript of $U$ to make it explicit that it depends on the current posterior mean.\n\n> How accurate (in terms of forwards error) are the CG solves\n\nWe have normalized the forward error with respect to the norm of the r.h.s. (the $b$ in the $Ax = b$ sense), which we refer to as \"residuals\" and summarized the average residual across different runs for the last epoch:\n\n| Dataset/Architecture  | KIVI            |\n|-----------------------|-----------------|\n| MNIST/LeNet           | 0.0750 ± 0.0295 |\n| FMNIST/LeNet          | 0.0104 ± 0.0021 |\n| SVHN/ResNet_small     | 0.0096 ± 0.0003 |\n| CIFAR-10/ResNet_small | 0.0317 ± 0.0047 |\n\nAs evidenced above, the errors are quite acceptable, and the performance observed in the other metrics confirms this as well. Relatedly, Gaussian process software uses coarse conjugate-gradients approximations for covariance matrices, and successfully so (see, e.g., the GPyTorch documentation regarding conjugate gradients).\n\n\n> Figure 5: why are the data points different for IVON and KIVI?\n\nThe data points in Figure 5 are the same, but the y-axes are scaled differently. Thanks for pointing this out. We have updated Figure 5 so that both y-axes have the same scale.\n\n> Tables 1 and 2: I'm generally a bit surprised at the high ECE values ...\n\nYes, that is correct, the ECE numbers are not comparable because of a difference in architecture and hyperparameters. In our experiments, we use the recommended hyperparameters from Maddox et al, ’19. \n\n> Writing comment\n\nFixed, thanks!\n\nAgain, thank you for your review. We hope that all concerns have been resolved; if not, please let us know!"}, {"review_id": "FhBluMk3qs", "reviewer": "Reviewer_4Nga", "summary": "This paper introduces a novel approach for Gaussian variational inference in the context of over-parameterized Bayesian neural networks. Motivated by the observation that neural network outputs remain invariant within the kernel of the Fisher information matrix, and building on recent work by Miani et al. (2024b), the authors propose a new form of approximate Gaussian posterior distribution. The covariance matrix of this distribution is constructed as a combination of two components: $UU^\\top$ and $I-UU^\\top$.\nSamples from the range of the former (i.e., the kernel of the Fisher matrix) do not affect the network’s prediction on the training data, whereas samples from the latter do. To reduce the computational cost of estimating the Fisher information matrix, the authors adopt a stochastic projection technique from Miani et al. (2024b). The proposed method is evaluated on out-of-distribution (OOD) detection and variational autoencoder tasks, demonstrating its effectiveness.", "questions": "Miani et al. (2024b) propose restricting the approximate posterior to the kernel of the Fisher-Rao matrix, ensuring that samples do not alter the model's predictions on the training data. This makes intuitive sense: it maintains predictive accuracy on the training data while allowing the model to express uncertainty on out-of-distribution (OOD) inputs.\n\nYour method, in contrast, allows posterior mass in both the kernel $R\\left(U U^{\\top}\\right)$ and its complement, even though the latter influences the model's predictions on the training set. Could you clarify the motivation for this choice? How does introducing variability in directions that change the training predictions contribute to the overall performance or uncertainty calibration? I request an explanation that is more than simply showing me the experimental results.", "limitations": "Yes.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 1, "clarity": 2, "significance": 1, "originality": 1, "strengths_and_weaknesses": "1. The paper lacks sufficient discussion of prior work on Gaussian variational inference for Bayesian neural networks. While this is a broad area, some key references should be acknowledged, at least in the appendix. Here is an inexaustive list: the linearized Lapalce approximation by https://arxiv.org/abs/2008.08400, https://arxiv.org/abs/1906.01930; the functional space variational inference by https://arxiv.org/abs/2312.17199. I am not asking the authors to empriicaly compare with these papers, but at least these closely related work should be properly discussed. \n\n2. The author is not distinguishing between the expected value of the Fisher information matrix and its empirical estimate over the training samples. I think the author is using the latter in the experiment. Since the experiments likely use the empirical version, it is important to clarify whether the desirable structural properties (e.g., predictive invariance within the kernel) still hold under this approximation.\n\n3. To improve clarity, I recommend adding a subscript to the matrix $U$, such as $U_{\\hat{\\theta}}$, to emphasize that it is computed at a specific parameter point $\\hat{\\theta}$. This is particularly important in Section 3.3, where the dependence of the Fisher kernel on the parameter estimate is critical to the algorithm.\n\n4. While I generally avoid making novelty the primary critique, in this case it is a significant concern. Much of the methodology, including the design of the covariance matrix and the stochastic projection trick, appears to be directly adapted from Miani et al. (2024b). The key difference-that this work uses a linear combination of $U U^{\\top}$ and $ \\mathrm{Id}-U U^{\\top}$ rather than restricting the posterior entirely to the kernel-is incremental and may not constitute a substantial enough contribution to justify publication. Given the heavy reliance on prior work, I do not believe the paper introduces sufficient novel insights, and I therefore recommend rejection.", "was_revised": true, "final_justification": "I appreciate the author's response, which has addressed most of my concerns. I now increase my score to accept.", "rebuttal": "Thank you for the valuable feedback! In the following, we'll briefly reply to what the review lists under weaknesses, before answering the specific questions.\n\n> The paper lacks sufficient discussion of prior work on Gaussian variational inference for Bayesian neural networks. \n\nThank you for pointing out relevant work we overlooked in our discussion. A reviewed version of the manuscript will appropriately mention those papers.\n\n\n> The author is not distinguishing between the expected value of the Fisher information matrix and its empirical estimate over the training samples. \n\nThank you for raising this point. In this paper, we target the kernel and image spaces of the Generalized Gauss-Newton matrix. The Fisher information matrix i.e.\n$\\sum_n  \\mathbb{E}\\_{y \\sim p(y|x\\_n)}  [\\nabla_\\theta \\log p(y|x\\_n) \\nabla_\\theta \\log p(y|x_n)^T]$\ncorresponds exactly to the generalised Gauss-Newton matrix in our setting; see Proposition 1 by Kunstner (\"Limitations of the Empirical Fisher Approximation for Natural Gradient Descent\", 2019). Note that this is distinct from the empirical Fisher i.e. $\\sum_n  [\\nabla_\\theta \\log p(y_n|x_n) \\nabla_\\theta \\log p(y_n|x_n)^T]$. So, we rely on empirical estimates over the training data, as is standard in deep learning, because we don't know the distribution over inputs, but we don't rely on empirical estimates over the training data and the labels as used in the empirical Fisher. We retain all the desirable structural properties of the loss-projected posterior from Miani 2024b since we use the same matrix.\n\n> To improve clarity, I recommend adding a subscript to the matrix... \n\nWe agree, thank you for the feedback. The notation around $U$ and its dependence on the current parameter point estimate is now explicit in the manuscript.\n\n> While I generally avoid making novelty the primary critique, in this case it is a significant concern. Much of the methodology, including the design of the covariance matrix and the stochastic projection trick, appears to be directly adapted from Miani et al. (2024b).\n\nThank you for that comment.\nPlease note that our work has several crucial differences from Miani et al. (2024b). Miani et al. (2024b) is a post hoc method, which means they sample from the proposed approximate posterior after MAP training. Our work constructs a variational family based on the approximate posterior proposed. While we borrow techniques for efficiently sampling from the variational family, the full ELBO training is a significant innovation that clearly distinguishes our contribution from Miani et al. (2024b).\n\n**The last point relates to the answer to your question. To elaborate on that:**\n\nAs stated by the reviewer, the kernel directions in the parameter space are where the predictions are locally constant. In Miani et al. (2024b), this property is exploited to build posteriors with 0 in-distribution uncertainty but high out-of-distribution uncertainty, leading to good OOD detection. However, since we propose a method for variational training, it is important that we converge to a distribution that has a good predictive performance and sensible uncertainty quantification. This is why it is crucial to include the image component in the posterior because that is the subspace where the predictions (and the loss) actually change (locally).\n\nAgain, thanks for the thorough review! We hope that all concerns have been addressed, and we're looking forward to discussing."}, {"review_id": "Cy3aKkgSGq", "reviewer": "Reviewer_9p1F", "summary": "The paper proposes KIVI, a variational inference method that accounts for overparameterization in deep networks by decomposing the parameter space via the Fisher-Rao metric into kernel and image subspaces. A fully correlated Gaussian posterior with separate variances in each subspace captures structured uncertainty. Using matrix-free stochastic projections and conjugate gradients, KIVI scales to large models. Empirical results show strong performance and improved calibration over baselines like IVON and SWAG.", "questions": "- Can $\\sigma^2_{\\text{ker}}$ and $\\sigma^2_{\\text{im}}$ be interpreted as epistemic and aleatoric uncertainty in practice?\n- In Table 1, KIVI achieves higher accuracy on Imagenette compared to IVON, but with worse calibration metrics. Can you provide an explanation for this trade-off? Does this suggest limitations in the variational family's expressivity?\n- IVON has been demonstrated on LLaMA-2 with 7B parameters. Given KIVI's reliance on loss-Jacobian projections and conjugate gradient solves, is it currently feasible to apply your method at that scale?", "limitations": "yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 4, "originality": 3, "strengths_and_weaknesses": "### **Strengths**\n\n- The method is grounded in a principled decomposition of the parameter space using the Fisher-Rao metric\n- Uses a minimal parameterization while still achieving competitive performance\n- Scalable to large models through matrix-free methods\n- Empirically outperforms or matches state-of-the-art methods\n\n### **Weaknesses**\n\n- Lacks theoretical analysis or guarantees on convergence and approximation quality\n- The paper does not provide any runtime or convergence comparisons between KIVI and IVON. While KIVI introduces additional computational overhead due to conjugate gradient-based projections, there are no wall-clock benchmarks or convergence plots to quantify the practical trade-offs.", "was_revised": true, "final_justification": "**Summary of how my questions were addressed**\n\n1. **Runtime & convergence trade-offs vs. IVON**  \n   Authors provided a runtime table showing KIVI is **2–5× slower** than IVON, explained the main cost factors (loss-Jacobian projections, CG iterations, warmup), and clarified the practical trade-offs.\n\n2. **Interpretation of $\\sigma^2_{\\text{ker}}$ and $\\sigma^2_{\\text{im}}$**  \n   Kernel subspace contains only *epistemic* uncertainty, image subspace contains a *mix* of epistemic and aleatoric uncertainty. Useful as an intuition, but not a strict decomposition.\n\n3. **Feasibility at LLaMA-2 scale**  \n   Not yet tested at 7B parameters. In theory, scales to that size, but would be 2–5× more expensive than IVON due to projections and CG solves.\n\n**Conclusion:**  \nAll questions were answered with concrete explanations or limitations. I was satisfied with the responses and kept my “Accept” score.", "rebuttal": "Thank you for the valuable review! In the following, we answer the questions individually.\n\n> The paper does not provide any runtime...\n\nWe summarize the running times (total execution time from model initialization until the end of the last training epoch, in minutes) below, comparing IVON and KIVI:\n\n| Dataset/Architecture  | KIVI           | IVON          |\n|-----------------------|----------------|---------------|\n| MNIST/LeNet           | 27.27 ± 3.96   | 9.97 ± 0.35   |\n| FMNIST/LeNet          | 46.30 ± 7.08   | 9.71 ± 0.06   |\n| SVHN/ResNet_small     | 304.16 ± 0.96  | 108.87 ± 1.90 |\n| CIFAR-10/ResNet_small | 177.08 ± 11.68 | 70.17 ± 0.88  |\n\nAs shown by the table, KIVI takes from 2 to 5 times the training time of IVON, where aspects such as alternating projections (and whether the stochastic variant is used) and the number of MLE steps used for warmup can affect the total execution time.\n\n> Can $\\sigma^2_{\\text{ker}}$ and $\\sigma^2_{\\text{im}}$ be interpreted...\n\nWhile we initially were optimistic that this would be the case, we eventually avoided this interpretation. The kernel contains only epistemic uncertainty, while the image can contain both epistemic and aleatoric uncertainty. In practice, we observe that the image space does contain a mix of uncertainties, and, depending on model complexity, this changes in degree. So, while the interpretation is a useful starting intuition, it should be followed with some care.\n\n> In Table 1, KIVI achieves higher accuracy on Imagenette ...\n\nThat's a good question. The calibration metrics are relatively similar (given the standard deviations), so concluding limitations on the expressivity may be a slight overinterpretation of these results. For example, on CIFAR-10 and SVHN, KIVI's calibration is better than IVON's. \n\n> IVON has been demonstrated on LLaMA-2 with 7B parameters...\n\nGood point. We haven't explored that scale for KIVI yet. You're right in that the loss-Jacobian projections and conjugate gradient solves make KIVI more expensive than IVON. Even though KIVI scales well, it costs about 2-5 times the training time of IVON; see our complexity analysis in the reply to Reviewer rdDS, which we will also add to the paper.\n\nAgain, thanks for your valuable review! We hope that all concerns have been resolved adequately."}, {"review_id": "5FsMrXCrMu", "reviewer": "Reviewer_bAN4", "summary": "This paper proposes a new variational posterior for Bayesian neural networks that has two variance parameters. One parameter controls the scale of the Gaussian distribution in the non-data-affecting directions in parameter space (null-space of the empirical Fisher) and the other controls the scale of a Gaussian distribution in the orthogonal complement subspace.\n\nIn order to perform scalable projections into the null space of the empirical Fisher, the authors take the method of alternating projections proposed in Miani et al 2024 and apply it within a conjugate gradient update scheme. To accommodate this in a minibatch setting where the variational posterior and therefore the Fisher changes at every iteration, the authors add an additional approximation where the noise variable has a memory and is updated with a kind of momentum.\n\nThe authors perform experiments to validate their method accuracy, uncertainty estimation and scalability. They show an ability to capture in-between uncertainty, and better uncertainty estimation than many competitive BNN inference techniques in image classification with CNNs.", "questions": "-\tHow many CG iterations were used in each optimization step?\n-\tHow does the training wall-clock time compare with IVON and MAP?\n-\tSuggestion: visualize your approximate posterior in a toy case. It’s quite unusual and it would be good to get more intuition for it.\n-\tSuggestion: please add pseudocode. There are layers of approximation here: VI, stochastic optimization, CG, alternating projections, stochastic alternating projections. It is quite hard to keep track of the exact steps that need to occur to run the algorithm.\n-\tSuggestion: could you extend the method to have more than 2 variance parameters? Right now there’s only two kinds of subspace, the null space and its complement. Perhaps there could be ways to partition the space more finely and introduce independent variational parameters for more subspaces.\n-\tSuggestion: please do not call the method KIVI. There is a name clash with Kernel Implicit Variational Inference: https://arxiv.org/abs/1705.10119", "limitations": "-\tThe authors do not explain clearly how much slower their method is.\n-\tStep-by-step explanation of the method could be clearer.", "rating": 6, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 3, "significance": 4, "originality": 4, "strengths_and_weaknesses": "Quality:\n\nThe paper is very well-written and a joy to read. The authors show a great understanding of the context of the field. The experiments are thorough, answering most questions that the reader is interested in, and comparing with good baselines. The authors are clear and informative regarding the limitations of their work, which is refreshing, and avoid over-claiming.\n\nClarity:\n\nThe paper is of above-average clarity, and the major ideas are communicated well in spite of their complexity.\n\nSignificance:\n\nHaving worked in the BNN field for 7+ years, I would have to say this is one of the top 5 most exciting papers I’ve seen in the field. It is quite rare to have a paper propose not one, but several interesting ideas and string them together this well. Performing good BNN inference is extraordinarily difficult, and the authors are right to note that many in the field believe that performant BNN inference is a lost cause. This work which targets in a perceptive way the exact directions in weight space where approximate posteriors can spread without affecting the loss is significant for breathing new ideas into an old field. Although their method is not “SotA on everything” I think that is almost besides the point in this case. The ideas are fresh and have been demonstrated to be worth pursuing further.\n\nOriginality:\n\nThe authors note that their work is a continuation of Miani et al 2024 and bears many similarities, although I believe the extension to ELBO optimization during training is significant and non-trivial, and another step in this exciting research direction. I hope to see this avenue developed further in future work. \nNote: other methods have proposed subspace inference before, but not in a targeted way, and with a very different inference methodology: https://arxiv.org/abs/1907.07504", "was_revised": true, "final_justification": "After reading the rebuttals and the other reviews, I maintain my strongly positive view of this paper and greatly hope it is accepted.", "rebuttal": "Thank you for the positive assessment! We're glad you like the paper. We will reply to your questions below.\n\n> How many CG iterations were used in each optimization step?\n\nWe used ten CG iterations per optimization step.\n\n> How does the training wall-clock time compare with IVON and MAP?\n\nWe summarize the running times (total execution time from model initialization until the end of the last training epoch, in minutes) below, comparing IVON and KIVI:\n\n| Dataset/Architecture  | KIVI           | IVON          |\n|-----------------------|----------------|---------------|\n| MNIST/LeNet           | 27.27 ± 3.96   | 9.97 ± 0.35   |\n| FMNIST/LeNet          | 46.30 ± 7.08   | 9.71 ± 0.06   |\n| SVHN/ResNet_small     | 304.16 ± 0.96  | 108.87 ± 1.90 |\n| CIFAR-10/ResNet_small | 177.08 ± 11.68 | 70.17 ± 0.88  |\n\nAs shown by the table, KIVI takes from 2 to 5 times the training time of IVON, where aspects such as alternating projections (and whether the stochastic variant is used) and the number of MLE steps used for warmup can affect the total execution time.\n\n> Suggestion: visualize your approximate posterior in a toy case. It’s quite unusual and it would be good to get more intuition for it.\n\nWe have visualizations of the predictive posterior in Figures 1 and 5, which give some indication of the weight posterior. A plot of the weight posterior itself may be a bit tame, as it is merely a sum of two Gaussians projected onto two disjoint subspaces.\n\n> Suggestion: please add pseudocode \n\nWe will summarize all the steps in pseudocode in a revised version of the manuscript, together with computational complexity analysis (see the reply to Reviewer rdDS for that one). Thank you for the suggestion.\n\n\n> Suggestion: could you extend the method to have more than 2 variance parameters? Right now there’s only two kinds of subspace, the null space and its complement. Perhaps there could be ways to partition the space more finely and introduce independent variational parameters for more subspaces.\n\nYour suggestion makes a lot of sense, thank you! The variance in the kernel subspace is only informed by the prior (which is usually an isotropic Gaussian), and it ensures that prior variance does not cause the model to underfit. Here a single scalar makes sense, though one could consider a more elaborate prior (e.g., one scalar per layer). But in the image subspace, each basis direction corresponds to different \"speeds\" along which the loss changes. Here, it would make a lot of sense to consider a more elaborate model, e.g., a low-rank covariance matrix or similar.\n\n> Suggestion: please do not call the method KIVI.\n\nWe've also noticed the name clash with the other KIVI since the paper submission and to avoid confusion, we'll rename our method to VIKING (\"Variational inference with kernel- and image-subspaces of numerical Gauss--Newton matrices\").\n\n\nAgain, thanks for your review. We highly appreciate your positive comments and assessment of our work!"}]

Author rebuttal:

Meta-review: This paper introduces KIVI, a variational inference method for Bayesian neural networks that decomposes the posterior covariance into two subspaces based on the kernel and image of the network's Jacobian.

The final reviews (post-discussion) show strong support for acceptance, with all reviews but one recommending acceptance. The reviewers in favor of acceptance praise the paper's originality. Reviewer bAN4 calls it one of the most exciting papers in the field. The work is viewed as a well-executed and valuable contribution that brings fresh ideas to BNN research. The primary dissenting opinion from reviewer rdDS centers on the need for additional baselines and a better discussion of prior work. They argue that given the method's computational cost, a comparison against non-VI methods like SGMCMC or deep ensembles is needed. I agree with reviewer rdDS and expect the updated manuscript to address these issues.

Two concerns that have not been raised by reviewers are the timeliness of this work and the experimental execution. While I agree with the reviewers that the submission introduces interesting ideas, I believe the experimental results are far below what I would expect to see in a Bayesian deep learning paper in 2025. The paper only considers toy and computer vision experiments, and the results (in terms of accuracy) are **significantly** below what would be expected in the deep learning literature, even a few years ago. I suspect the low predictive accuracy is due to the choice of model architecture (ResNet34 instead of ResNet18 or ResNet50) and other implementation decisions that are orthogonal to the proposed method. However, without results that show that the proposed methods outperforms MAP baselines that are on par with strong results reported in the deep learning literature (e.g., a 95% test accuracy on CIFAR-10 with a ResNet18 trained from scratch), this work is unlikely to have impact beyond the BNNs community. I strongly recommend to the authors to include such results in the updated manuscript. For reference points in the Bayesian deep learning literature, see for example https://arxiv.org/abs/2311.15990 and https://arxiv.org/pdf/2309.09814.

Despite these concerns, I believe the paper has merit and would be of interest to the NeurIPS community. I recommend acceptance.
Decision: Accept (poster)