VIKING: Deep variational inference with stochastic projections

AI summary

Overview

Research area: Bayesian deep learning / approximate inference (stat.ML), building on singular learning theory, information geometry, and the geometry of overparametrized neural networks.

Technical level: Advanced. The paper assumes familiarity with variational inference, Fisher–Rao metrics, null spaces, and matrix-free numerical linear algebra.

Scope (one sentence): The paper proposes a variational family for Bayesian neural networks whose covariance is built from the kernel (null space) and image of an empirical Fisher–Rao metric, and derives a scalable stochastic alternating-projection algorithm to fit it.

Authors are affiliated with the Technical University of Denmark and the University of Cambridge (arXiv:2510.23684v1, 27 Oct 2025, licensed CC BY-SA 4.0). Code is available at https://github.com/eugene/viking-paper-experiments and a library at https://github.com/fadel/viking.

What This Paper Is About

Neural networks are overparametrized: many different weight configurations compute the exact same function, so a single function corresponds to a continuous set of parameters. The authors argue this is why Bayesian treatments of deep networks have historically produced unstable training, poor predictive power, and subpar calibration — approximate posteriors such as Laplace or mean-field approximations cannot represent this redundant geometry.

The goal is a variational posterior that explicitly reflects this structure: one part of the covariance models uncertainty that leaves predictions on the training data unchanged, and another part models uncertainty everywhere else, while remaining computationally tractable for modern architectures.

Key Contributions

1. A geometry-aware Gaussian variational family. The posterior covariance is Σ = σ_ker² U Uᵀ + σ_im² (I − U Uᵀ), where U spans the kernel of an empirical Fisher–Rao / Gauss–Newton estimate. It is fully correlated across all parameters yet controlled by only two interpretable scalars, σ_ker and σ_im.

2. A closed-form, cheap ELBO. Because the covariance is a sum of scaled projection matrices, the KL term is evaluated in closed form using only the kernel dimension R, estimated with Hutchinson's trace estimator (reusing vectors already computed for the reconstruction term).

3. A stochastic extension of alternating projections (the "stochastic projections" in the title), which makes kernel projection compatible with mini-batched ELBO optimization by blending the previous projection with fresh noise through a hyperparameter γ ∈ [0, 1].

4. A scalable, architecture-agnostic implementation using conjugate gradients with full reorthogonalization, avoiding materialization of the D × D projection matrix or the N × N matrix J Jᵀ.

Main Findings

• Two scalars can be surprisingly competitive. Although the authors call calibrating a network through two scalars "overly simplistic," VIKING is reported to be state-of-the-art against a wide array of baselines, which they interpret as evidence that reflecting overparametrization explicitly matters more than covariance richness alone.

• Stochastic projections are essential. On a small CNN classifier trained on Fashion MNIST, the no-noise variant (γ = 1.0) reaches the highest training accuracy but yields a poorer approximate posterior and generalization, because its projected samples only change once per epoch. The paper also evaluates γ = 0.0 (only noise) and γ = 0.5.

• Full ELBO optimization beats post-hoc tuning. Tuning only σ_ker and σ_im on a fixed model is not "per se problematic" and yields a well-performing model, but optimizing the variational mean with the ELBO improves on it.

• Maximum-likelihood warmup helps, to a point. Warmup via maximum likelihood is an effective way to quickstart ELBO learning, but there is a sweet spot: full convergence on maximum likelihood can cause the ELBO optimization to get stuck, which the authors conjecture is due to sharp minima disrupting the expectation term.

• Classification accuracy. Averaged over three runs, VIKING reports 0.991 ± 0.001 on MNIST, 0.900 ± 0.001 on Fashion MNIST, 0.960 ± 0.001 on SVHN, 0.877 ± 0.004 on CIFAR-10, and 0.887 ± 0.003 on Imagenette — the highest accuracy in each of those tables. Baselines include MAP, the post-hoc method of Miani et al. (2025), IVON, SWAG, and a last-layer Laplace approximation. On MNIST, IVON obtains a better NLL (0.043 ± 0.002) than VIKING (0.055 ± 0.003).

• Calibration gains where overparametrization is most prominent. VIKING is described as "particularly effective" against baselines on SVHN and CIFAR-10, where the model has around 220k parameters. On CIFAR-10 its ECE is 0.041 ± 0.004, compared with 0.075 ± 0.012 (MAP), 0.086 ± 0.014 (IVON), 0.559 ± 0.006 (Miani et al.), 0.694 ± 0.018 (SWAG), and 0.704 ± 0.000 (last-layer Laplace).

• Scaling to a 21.7 million parameter model. VIKING is trained on ResNet34 on Imagenette and shows superior accuracy and NLL over IVON, but inferior calibration metrics (ECE 0.077 ± 0.001 vs 0.069 ± 0.011; MCE 0.612 ± 0.162 vs 0.464 ± 0.230).

• Out-of-distribution detection. Using the maximum variance of softmax probabilities across output dimensions as the OOD score, VIKING performs on par with the best baselines, sometimes by a large margin — for example, MNIST → Fashion MNIST (AUROC 0.972 ± 0.003) and MNIST → KMNIST (0.965 ± 0.002) — and is a close second elsewhere. EMNIST and CIFAR-100 are used as additional OOD datasets, replicating the setup of Miani et al. (2025).

• Generative modelling. Refining the decoder of a 6.5 million parameter VAE with IVON and VIKING produces qualitatively similar reconstructions, but IVON's variance spans all image features including backgrounds, whereas VIKING focuses variance on facial features and the outline. Measuring the median per-pixel standard deviation over 16k generated samples, VIKING's uncertainty distributions clearly separate in-distribution samples (standard Gaussian) from out-of-distribution samples (Gaussian with twice the variance), while IVON's do not.

• Diagnosed weakness in IVON. The authors report that IVON tends to learn near-identical variances for all network weights (shown for CIFAR-10), suggesting an inaccurate, effectively isotropic posterior.

Methodology in Plain English

The starting assumption is a standard Gaussian prior over weights with precision α. The posterior approximation is also Gaussian but not diagonal — its covariance has two pieces with different scales: one along directions in weight space where the loss does not change for the training data (the kernel of the empirical Fisher–Rao metric, closely related to the Gauss–Newton matrix), and one in all remaining, orthogonal directions (the image space). Only the mean θ̂ and the two scales σ_ker and σ_im are learned.

Because sampling from this posterior requires projecting random noise onto the kernel, the core numerical problem is a constrained least-squares projection. The projection is expressed through a stack of per-datum loss gradients J, and solved without ever forming the large matrices involved: the system is solved with conjugate gradients, and full reorthogonalization is used because conjugate gradients can otherwise produce projections that are not properly orthogonal.

Projecting onto the full-data kernel would require passing over the entire dataset for every sample, which is incompatible with mini-batched ELBO optimization that needs a fresh sample each step. The authors therefore run alternating projections in a stochastic fashion: at each step a new sample is a projection of a weighted mixture of the previous projected sample and fresh Gaussian noise, with γ controlling how much history is retained. Samples are then fed through the network to estimate the ELBO's reconstruction term, gradients update the mean and the two scales, and the KL term is computed in closed form using only the kernel dimension.

Why This Matters

Impact on research. The paper challenges the widespread view that approximate Bayesian inference for deep networks is a "lost cause." By showing that a simple decomposition aligned with the geometry of reparametrizations closes much of the gap, it reframes the problem as one of respecting overparametrization rather than of finding ever-richer covariance structures. It also provides a practical, mini-batch-compatible bridge between the abstract differential-geometry characterization of Roy et al. (2024) and actual training loops.

Real-world applications (potential, not evaluated in this paper):

• Medical imaging, where calibrated confidence and OOD detection matter more than raw accuracy, and where the paper's MNIST/Fashion MNIST/SVHN/CIFAR-10 calibration and OOD results are the closest proxy.
• Autonomous perception and safety-critical robotics, where the Imagenette ResNet34 experiment suggests the method reaches modern-scale vision models.
• Scientific regression from sparse observations, illustrated by the sinusoid example with 10 data points, where VIKING assigns higher variance near and beyond the data boundary.
• Generative models with unreliable inputs, illustrated by the VAE experiment separating in- and out-of-distribution face samples.

Industry relevance. The method is architecture-agnostic and matrix-free, and the release of a general-purpose library (fadel/viking) alongside reproduction code lowers the barrier to adoption. The training-time benefit of maximum-likelihood warmup is an immediately usable practical recipe, and the post-hoc tuning mode offers a cheap path for teams that already have trained checkpoints.

Future Directions

• Calibration at scale. VIKING has superior accuracy but inferior ECE and MCE to IVON on ResNet34/Imagenette, so the source of that calibration gap and whether the two-scalar parameterization needs refinement remain open.

• Computational cost. Each ELBO step involves conjugate-gradient solves with full reorthogonalization, and the D × D projection matrix cannot be instantiated for very large models. Reducing this cost — via preconditioning or better projections — is a natural next step.

• Choosing the kernel estimator. The paper adopts the loss-Jacobian J as a stand-in for the Fisher–Rao metric, citing Kunstner et al. (2019); the consequences of this approximation and of alternatives are not settled.

• Theory for the stochastic projection. The behavior of γ is currently justified empirically (with γ = 1.0, 0.0, and 0.5 compared), so a theoretical account of the optimal sliding-window memory, and of the "sharp minima" hypothesis behind the warmup sweet spot, would strengthen the method.

Target Audience

Researchers and graduate students in Bayesian deep learning, probabilistic machine learning, and approximate inference who are interested in why standard variational and Laplace approximations underperform on overparametrized models. It is also relevant to statisticians working on singular learning theory and the geometry of neural network parameter spaces, and to practitioners who need calibrated uncertainty or OOD detection from deep models and are willing to invest extra training compute for it. Beginners will find the paper demanding; the numerical sections on matrix-free conjugate gradients and alternating projections require a solid linear algebra background.

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