Walking on the Fiber: A Simple Geometric Approximation for Bayesian Neural Networks

TL;DR

The paper tackles the challenge of intractable posteriors in Bayesian neural networks by introducing MetricBNN, which first explores the local MAP neighborhood via a lightweight drift-plus-gradient sampling scheme and then learns a structured latent posterior to enable rapid, non-iterative sampling. This two-stage approach captures non-Gaussian, low-dimensional loss-landscape geometry that traditional Laplace approximations miss, providing a scalable alternative to refinement-based posteriors. Across toy tasks, UCI datasets, and high-dimensional image and OOD benchmarks, MetricBNN achieves competitive negative log-likelihood and calibration (ECE) while reducing computational overhead compared to Hessian-based methods, HMC, and deep ensembles. The work demonstrates the practical potential of coupling geometry-aware sampling with autoencoder-based latent representations for robust Bayesian uncertainty quantification in deep learning.

Abstract

Bayesian Neural Networks provide a principled framework for uncertainty quantification by modeling the posterior distribution of network parameters. However, exact posterior inference is computationally intractable, and widely used approximations like the Laplace method struggle with scalability and posterior accuracy in modern deep networks. In this work, we revisit sampling techniques for posterior exploration, proposing a simple variation tailored to efficiently sample from the posterior in over-parameterized networks by leveraging the low-dimensional structure of loss minima. Building on this, we introduce a model that learns a deformation of the parameter space, enabling rapid posterior sampling without requiring iterative methods. Empirical results demonstrate that our approach achieves competitive posterior approximations with improved scalability compared to recent refinement techniques. These contributions provide a practical alternative for Bayesian inference in deep learning.

Figures (12)

• Figure 1: We propose a novel sampling scheme to approximate the posterior of a trained network. When the loss landscape of a trained network over its parameters (blue regions) is low on a very low-dimensional curve (white region), a simple Hessian ($H^{-1}$) fails in capturing its distribution. Our sampling scheme is composed of two iterative steps: randomly perturb a given solution (yellow points) then refine to minimize the loss function over the given dataset (orange points). This proposed scheme allows to estimate posteriors of arbitrary shapes and it is well-suited when the space of solutions in a network is much lower dimensional than the parameter space.
• Figure 2: Posterior samples for Regression task. The blue points represent the dataset, the orange lines are samples from the estimated posteriors. Our proposed sampling method correctly captures the uncertainty in the data gap.
• Figure 3: Estimated posterior for the regression task. The blue points represent the dataset and in orange the mean and standard deviation of posteriors sampled from the estimated distributions. A naive SVD approximation of the solutions found with the sampling scheme fails in correctly representing the true posterior. Our proposed MetricBNN posterior correctly approximates it.
• Figure 4: Estimated posterior for the Banana classification task. The orange and yellow points represent the data of the two classes in the classification dataset. In blue is the value of the uncertainty of the posteriors sampled from the estimated distributions.
• Figure 5: Trade-off between network size and computational complexity on the regression task. Each marker shows NLL (lower is better) for a network with the indicated number of hidden layers. MetricBNN’s cost grows only mildly with depth.