Title: VIKING: Deep variational inference with stochastic projections
Venue: NeurIPS 2025
Accepted: yes
Keywords: Bayesian deep learning

overparametrization

reparametrization invariance

linear algebra

matrix-free methods
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.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\nQuality:\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\n\nQUESTIONS:\n-\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\n\nLIMITATIONS:\n-\tThe authors do not explain clearly how much slower their method is.\n-\tStep-by-step explanation of the method could be clearer.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 5, "CLARITY": 3, "IMPACT": 4, "ORIGINALITY": 4, "DATE": "2025-06-26"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS AND WEAKNESSES:\n### **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.\n\nQUESTIONS:\n- 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?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 4, "ORIGINALITY": 3, "DATE": "2025-06-30"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\n1. 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.\n\nQUESTIONS:\nMiani 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.\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 2, "IMPACT": 1, "ORIGINALITY": 1, "DATE": "2025-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\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.\n\nQUESTIONS:\nI 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.\n\nLIMITATIONS:\nn/a", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\n**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.\n\nQUESTIONS:\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- 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.\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-07-02"}