VIKING: Deep variational inference with stochastic projections

Thank you for the positive assessment! We're glad you like the paper. We will reply to your questions below.

How many CG iterations were used in each optimization step?

We used ten CG iterations per optimization step.

How does the training wall-clock time compare with IVON and MAP?

We 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:

As 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.

Suggestion: visualize your approximate posterior in a toy case. It’s quite unusual and it would be good to get more intuition for it.

We 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.

Suggestion: please add pseudocode

We 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.

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.

Your 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.

Suggestion: please do not call the method KIVI.

We'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").

Again, thanks for your review. We highly appreciate your positive comments and assessment of our work!

We 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.

Maybe I'm being silly, but I know in 2D this is trivial (basically and ellipse). In 3D however it's a special kind of ellipse where one axis is "free" and the other two are tied to have identical width. In higher dimensions I don't really have a great intuition for what this looks like at all, i.e. a geometric intuition for the nature of this constraint. Not that it matter too much.

I maintain my score as strong accept. Thank you for answering my questions!

Thank you for the valuable review! In the following, we answer the questions individually.

The paper does not provide any runtime...

We 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:

As 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.

Can $\sigma^2_{\text{ker}}$ and $\sigma^2_{\text{im}}$ be interpreted...

While 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.

In Table 1, KIVI achieves higher accuracy on Imagenette ...

That'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.

IVON has been demonstrated on LLaMA-2 with 7B parameters...

Good 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.

Again, thanks for your valuable review! We hope that all concerns have been resolved adequately.

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.

The paper lacks sufficient discussion of prior work on Gaussian variational inference for Bayesian neural networks.

Thank you for pointing out relevant work we overlooked in our discussion. A reviewed version of the manuscript will appropriately mention those papers.

The author is not distinguishing between the expected value of the Fisher information matrix and its empirical estimate over the training samples.

Thank 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. $\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]$ corresponds 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.

To improve clarity, I recommend adding a subscript to the matrix...

We 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.

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).

Thank you for that comment. Please 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).

The last point relates to the answer to your question. To elaborate on that:

As 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).

Again, thanks for the thorough review! We hope that all concerns have been addressed, and we're looking forward to discussing.

Thank you for that thoughtful assessment. In the following, we will first reply to the weaknesses before answering the specific questions.

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.

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.

Proposed changes:

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.

We 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:

As 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.

I believe that Appendix A.1 yields a slightly sub-optimal gradient estimate, which should probably be discussed in the main text.

You 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.

The missing comparison to me is against ... [SLANG/MCMC]

Thank 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:

• 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.
• 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.

[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.

[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).

Next, we reply to the direct questions:

I believe that the first use of conjugate gradients...

Yes, 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.

I would generically think that your approach boils down to training a predictive mean ...

That 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.

Eq 2 – why do the authors not instead optimize ...

Thank 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?

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..

Thanks 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.

Eq 16: P\textsubscript{t} is not really defined previously...

We 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.

How accurate (in terms of forwards error) are the CG solves

We 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:

As 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).

Figure 5: why are the data points different for IVON and KIVI?

The 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.

Tables 1 and 2: I'm generally a bit surprised at the high ECE values ...

Yes, 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.

Writing comment

Fixed, thanks!

Again, thank you for your review. We hope that all concerns have been resolved; if not, please let us know!

Thank you for the detailed response, a few more comments and questions:

KIVI takes from 2-5 times the training time of IVON,

What is the relationship of IVON training time to something like a similar stochastic gradient step? From a quick skim, it seems like IVON should also be much slower.

methods that are generally considered state-of-the-art for BNNs

I would argue that SGMCMC methods are at least near state of the art for BNNs as they scale well and preserve pretty high accuracy, especially with multiple chains. This is true even if they are not precisely reflective of the "true posterior". See for example the discussions in Izmailov et al, '21, Pituk et al, '25 amongst others... Thus, I think it is a pretty clear baseline, especially as it lies on a different position on the accuracy / computational cost frontier than your method.

$\sigma_k UU^\top + \sigma_\text{im}^2(I - U U^\top) = \sigma_\text{im}^2 I + (\sigma_k - \sigma_\text{im}) UU^\top)$ are PD

I think im not totally sure as towhy the sum of the two matrices is going to stay PD. The parameterization I pointed out stays PD by construction. The concern is primarily if $(\sigma_k - \sigma_\text{im}) < 0$ then I'm not totally sure what happens.

forwards error

Is this relative error or absolute error?

it's unclear whether it improves scalability

Generally, this is a good point; thanks for the clarification. I do think you can do similar tricks for CG and SGD here - an adjoint style approach, so differentiation isn't slower.

related work of ntk + Jacobian solves

I brought these up because they point out some interesting properties of the variational posterior that you learned as well. Thus, I think it's not that they should be cited, but a bit of discussion is good. For example, the ntk literature suggests that weights don't really move that much for wide enough networks - in your variational posterior does the jacobian change during training from scratch at all?

I would simply like to add that in my opinion a comparison with IVON and Adam is good enough, and that an SGMCMC comparison, although nice to have, would not be completely necessary to recommend acceptance. I believe that they have taken one method on the "frontier" of VI (IVON) and shown competitive performance with a very different approach while staying in the VI family. That is not enough to claim superiority over all BNN inference methods, but I believe it is enough to claim significance and be worthy of conference publication.

my opinion a comparison with IVON and Adam is good enough, and that an SGMCMC comparison, although nice to have, would not be completely necessary to recommend acceptance.

I don't know how I feel about this. Reading the IVON paper (https://arxiv.org/pdf/2402.17641) i also don't see a SGMCMC comparison (which feels deeply strange to me)... It seems like ensembles are at least more highly performant (e.g. in Table 2), and maybe that's just the better baseline. 5x slower per iteraion is approximately the cost of a standard training ensemble of size 5 for example.

We thank the reviewer for the thoughtful comments. Below, we will reply to each point individually.

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?

Thanks 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.

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?

We 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.

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.

We propose adding the following paragraph to highlight the linear dependency on every hyperparameter in our method and clarifying the total computational complexity.

Proposed changes:

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.

We 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:

As 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.

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?

The 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!).

Notably, 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.

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.

Thank 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.

Thanks again for the thorough review! We hope that all concerns have been addressed adequately, and look forward to the discussion.