Tubular Riemannian Laplace Approximations for Bayesian Neural Networks

TL;DR

The paper tackles the problem of calibrating uncertainty in Bayesian neural networks when the loss landscape is highly anisotropic and exhibits symmetry-induced valleys. It introduces Tubular Riemannian Laplace (TRL), a geometric posterior that forms a probabilistic tube along a low-loss valley, with a tangent direction governed by prior uncertainty and transverse directions governed by data-driven curvature via a Fisher/Gauss--Newton metric. TRL constructs a discrete spine to approximate the valley and uses a low-rank transverse subspace together with a tubular map to generate samples, enabling scalable inference through implicit curvature estimates (Lanczos) and Hessian-vector products. Empirically, TRL achieves ensemble-like calibration on high-dimensional tasks (ResNet-18 on CIFAR-100/10) at a single-model cost, outperforming static Laplace baselines and competing favorably with SWAG, while requiring significantly less training than Deep Ensembles. The work highlights a practical path to reliable uncertainty quantification in deep networks by geometrically aligning posterior mass with loss-valley structure, and it points to future work on scaling to billion-parameter models and integrating function-space priors.

Abstract

Laplace approximations are among the simplest and most practical methods for approximate Bayesian inference in neural networks, yet their Euclidean formulation struggles with the highly anisotropic, curved loss surfaces and large symmetry groups that characterize modern deep models. Recent work has proposed Riemannian and geometric Gaussian approximations to adapt to this structure. Building on these ideas, we introduce the Tubular Riemannian Laplace (TRL) approximation. TRL explicitly models the posterior as a probabilistic tube that follows a low-loss valley induced by functional symmetries, using a Fisher/Gauss-Newton metric to separate prior-dominated tangential uncertainty from data-dominated transverse uncertainty. We interpret TRL as a scalable reparametrised Gaussian approximation that utilizes implicit curvature estimates to operate in high-dimensional parameter spaces. Our empirical evaluation on ResNet-18 (CIFAR-10 and CIFAR-100) demonstrates that TRL achieves excellent calibration, matching or exceeding the reliability of Deep Ensembles (in terms of ECE) while requiring only a fraction (1/5) of the training cost. TRL effectively bridges the gap between single-model efficiency and ensemble-grade reliability.

Figures (5)

• Figure 1: Geometry of TRL. (a) Euclidean Laplace fits a local ellipsoid, while TRL models the posterior as a tube along the low-loss valley. (b) We decompose curvature into invariance ($v_\parallel$) and data-sensitive ($v_\perp$) directions.
• Figure 2: TRL as a tubular pushforward. Left: samples from an isotropic Gaussian in latent coordinates $(z_\parallel,z_\perp)$. Right: the corresponding samples in parameter space under a tubular map $T$, forming a curved tube along a low-loss valley.
• Figure 3: Empirical Hessian Spectrum (ResNet-18 on CIFAR-10). Computed via Lanczos iteration on the MAP estimate. The spectrum exhibits a characteristic rapid decay ("spiked" structure). Our choice of $k_\perp = 20$ (red dashed line) captures the dominant stiff directions ($\lambda_i > 5$) before the spectrum flattens into the uninformative bulk ($\lambda_i < 5$), balancing geometric fidelity with computational efficiency.
• Figure 4: Noisy sine regression. Predictive mean (blue) and two–sigma bands (shaded) for Euclidean Laplace (left), linearised Laplace (middle) and TRL with Hessian–based valley and parallel transport (right). The dashed black curve is the ground–truth sine.
• Figure 5: Two–moons classification. Top row: predictive probability $p(y=1\mid x)$ for ELA (left), LLA (middle) and TRL (right). Bottom row: corresponding predictive entropy. Black dots are training points. TRL produces sharp, well–aligned decision regions and concentrates uncertainty along the class boundary and far from the data.