Technical Summary: Large-scale Score-based Variational Posterior Inference for Bayesian Deep Neural Networks

Problem Statement

Bayesian Deep Neural Networks (BNNs) offer significant advantages over point-estimate deep learning, including robust uncertainty quantification, noise resistance, and overfitting mitigation. However, exact posterior inference in BNNs is generally intractable. Variational Inference (VI) is a dominant approximate method that frames inference as an optimization problem by approximating the target posterior $\pi(\theta)$ with a tractable distribution $q_\lambda(\theta)$.

The current dominant approach, ELBO-based VI (minimizing the reverse KL divergence), suffers from a well-known "mode collapsing" issue, where the approximation fails to capture all modes of the target density. While score-based VI methods (which minimize the difference between the target score $\nabla_\theta \log \pi(\theta)$ and the model score $\nabla_\theta \log q_\lambda(\theta)$) can theoretically address mode collapsing, existing score-based approaches are not scalable to large-scale BNNs. Specifically:

1. Computational Intractability: Methods like Gaussian Score Matching (GSM) and Batch-and-Match (BaM) rely on matrix inversions or multiplications that scale cubically or quadratically with state dimension, making them prohibitive for high-dimensional neural network parameters.
2. Incompatibility with Stochastic Gradients: Existing score-based methods typically require exact scores or cannot handle noisy, unbiased mini-batch scores derived from stochastic gradients, which are essential for training on large datasets.
3. Reparametrization Constraints: Many methods rely on reparametrized sampling, which can be computationally expensive or unstable in high dimensions.

Methodology

The authors propose a novel Proximal Stochastic-Gradient Score-based VI algorithm designed to scale to large-scale BNNs, including Vision Transformers (ViT) and ResNets.

Core Optimization Objective

The method iteratively updates the variational distribution $q_t(\theta)$ to $q_{t+1}(\theta)$ by solving the following optimization problem at each iteration $t$:

$$q_{t+1} = \arg \min_{q \in \mathcal{Q}} \mathbb{E}_{\theta \sim q_t(\theta)} \left[ \alpha_t \|\nabla_\theta \log q(\theta) - \nabla_\theta \log q_t(\theta)\|_2^2 + \|\nabla_\theta \log q(\theta) - s(\theta)\|_2^2 \right]$$

Where:

• $s(\theta) = \nabla_\theta \log \pi(\theta)$ is the target score.
• $\alpha_t$ is an adaptive proximal impact parameter (scheduled linearly from 0 to a constant $\le 1$).
• The first term is a proximal penalty enforcing the new iterate to remain close to the current iterate in terms of scores.
• The second term is the score-matching loss.

Key Technical Innovations

1. Handling Noisy Mini-batch Scores: Unlike previous score-based methods that require exact scores, this approach utilizes noisy, unbiased mini-batch scores $\hat{s}(\theta)$ obtained via stochastic gradients. The linearity of the score term in the loss function ensures that the gradient of the loss with respect to variational parameters $\lambda$ remains an unbiased estimate, even with noisy scores.
2. Avoidance of Reparametrization: The proximal term allows the expectation to be taken with respect to the fixed current distribution $q_t$ rather than the variable $q$. This circumvents the need for the reparametrization trick (used in ADVI) and avoids differentiating twice with respect to $\theta$ (required by traditional Fisher/Score methods), significantly reducing computational overhead and memory footprint.
3. Scalability: By using Euclidean norms instead of covariance-induced norms (which require matrix inversion) and supporting stochastic gradients, the method scales to networks with hundreds of millions of parameters (e.g., ViT-L-32 with ~305M parameters).

Theoretical Convergence

The authors provide a convergence theorem (Theorem 3.1) stating that under the assumption that the variational family is rich enough to represent smooth vector fields, the score of the variational distribution converges to the target score as iterations increase. While the assumption of representational richness is ideal, the authors note it holds exactly if both the true posterior and the variational family are Gaussian.

Key Contributions

1. Novel Algorithm: Introduction of a proximal stochastic-gradient score-matching VI algorithm that combines score-matching loss with a proximal penalty.
2. Scalability to Large-Scale BNNs: The method is the first score-based VI approach demonstrated to work effectively on large-scale deep networks (ResNets, ViTs) and large datasets, overcoming the computational bottlenecks of GSM and BaM.
3. Robustness to Noise: The algorithm effectively handles noisy, unbiased mini-batch scores, enabling the use of stochastic gradients without requiring full-batch score computations.
4. Elimination of Reparametrization: The method avoids the reparametrization trick, simplifying the optimization landscape and improving efficiency compared to ELBO-based ADVI.

Experimental Results

The authors evaluated the method on synthetic data, the PosteriorDB benchmark, and real-world large-scale tasks:

• Synthetic Data: On Gaussian and Gaussian mixture targets, the proposed method demonstrated faster convergence and better robustness to initial parameter choices (specifically regarding small covariance eigenvalues) compared to ADVI, GSM, and BaM. It successfully handled model mismatch and high-dimensional state spaces where GSM/BaM failed due to numerical issues.
• MNIST (Small-scale BNN): In low-data regimes, the method achieved comparable test error rates to ADVI but demonstrated superior uncertainty quantification, evidenced by lower Expected Calibration Error (ECE) and Negative Log-Likelihood (NLL).
• Large-Scale Visual Recognition (ResNet-101 & ViT-L-32): On datasets like Pets, Flowers, Aircraft, and DTD, the method matched the test prediction accuracy of ADVI while consistently outperforming it in uncertainty metrics (ECE and NLL). Competing methods (GSM, BaM, traditional score matching) failed due to memory constraints or numerical explosions.
• Time-Series Forecasting: Using the Koopa network architecture, the method outperformed ELBO-based VI and non-Bayesian baselines in MSE, MAE, and NLL across multiple datasets (ECL, ETTh2, Ex. Rate, ILI, Traffic, Weather).
• Normalizing Flows: The method was extended to Normalizing Flow variational densities, showing consistent improvements over ADVI on time-series tasks.
• Efficiency: Despite maintaining both previous and current iterates, the wall-clock time and GPU memory footprint were comparable to ADVI, proving its practical viability for large-scale training.

Significance and Claims

The paper claims that this work provides a scalable, flexible alternative to the dominant ELBO-based VI for large-scale Bayesian deep learning. By addressing the mode-collapping issue inherent in reverse KL objectives and overcoming the computational limitations of previous score-based methods, the proposed algorithm enables rigorous Bayesian inference on modern, high-capacity architectures like Vision Transformers.

The authors emphasize that their approach is empirically effective across diverse benchmarks, offering better uncertainty quantification without sacrificing predictive performance. They modestly note that while theoretical convergence relies on strong assumptions about the variational family's expressiveness, the empirical results support the method's utility in practice. The work positions score-based VI as a viable, scalable pathway for Bayesian deep learning, moving beyond the constraints of Gaussian-only or full-batch requirements of prior art.