Walking on the Fiber: A Simple Geometric Approximation for Bayesian Neural Networks
Authors: Alfredo Reichlin
Miguel Vasco
Danica Kragic
Venue: Accepted by TMLR ()

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.

Published in Transactions on Machine Learning Research (07 (07/2025)
Walking on the Fiber: A Simple Geometric Approximation for Bayesian Neural Networks
Alfredo Reichlin KTH Royel Institute of Technology
alfrei@kth.ss
Miguel Vasco
miguelsv@kth..s
KTH Royal Institute of Technology
Danica Kragic
dani@kth.cc
KTH Royal Institute of Technology
Reviewed on OpenReview: https: // openreview. net/ forum? id= NsuPyyrrr
Abstract
Bayesian Neural Networks provide 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, revisit sampling techniques for posterior exploration, proposing a simple variation tailored to efficiently sample from the posterior in pee-paaraeterized networks by leveraging the low-dimensional structure of loss minima. Building on this, we introduce model that learns a deformation 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 practical alternative for Bayesian inference in deep learning.
1
Introduction
Bayesian Neural Networks (BNNs) have emerged as a priucipled framework for uncertainty quantification deep learning by treating model parameters as random variables and inferring their posterior distribution. This Bayesian perspective is crucial for tasks requiring reliable decision-makee
under uncertainty, including medical diagnosis (Leibig et al. 2017). autonomous driving (Feng et al. 2018), and weather forecasting [Cofino et al. 2002). Despite their appeal, the high dimensionality and non-linearity of modern neural network parameter spaces render exact posterior inference intractable, requiring the development of efficient approximation techniques.
A widely used method make a deterministic, already-trained neural network Bayesian is the Taplace approximation, which estimates the postcrior distribution of the paramcters as a Gaussian centered the Maximum Posteriori (MAP) estimate (MacKay) 1992. Daxberger et al., 2021a). This approach is particularly convenient due to its simplicity and post-hoc applicability. However, the Laplace approximation faces significant challenges in modern deep networks. The computation and inversion of the Hessian scale poorly with network size (Martens 2010), and its reliance on local curvature limits its ability to capture the complex, non-linear geometry of the posterior in over-parameterized models Fort & Jastrzebski 2019 Bergamin et al. 2024).
In contrast, sampling techniques, though less frequently used in recent years, can be surprisingly effective in exploring the posterior when the parameter space of loss minima resides much lower-dimensional

---

Published in Transactions on Machine Learning Research (07 (07/2025)
Q2
θ1
01
22
H-1
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 given solution (yellow points) then refine to minimize the loss function over the given dalaset (orange points). This proposed scheme allows to estimate posteriors of arbitrary shapes and it is well-suited when the space of solutions in network is much lower dimensional than the paramcter space.
manifold than the full paramcter space (Garipov et al. 2018). By leveraging this property, sampling mothods can efficiently explore the posterior with fewer samples, especially in werrpparameterized settings. However, relying on sampling for posterior estimation can still be computationally expensive if every new inference requires a fresh sampling process. These limitations highlight the need for method that combines the efficiency of sampling-basee based exploration with more scalable and flexible posterior representation.
In this work, we propose variation of sampling method to explore the geometry of the loss landscape in the parameter space of neural networks. Our approach perturbs parameters around the MAP estimate using a set of drift directions and refines them with gradient updates, effectively maintaining computational efficiency as the network size increases. Using the data collected from our sampling process, we introduce a novel objective function learn structured latent reprosentation of the posterior by deforming the original parameter space. This learned representation enables the direct generation of new parameter samples without relying computationally expensive iterative sampling restrictive variational inference approximations. We refer to these two components collectively as MetricBNN. Empirically, we demonstrate the efficacy of our method in estimating uncertainty from trained models in both regression and classification tasks. In particular, our approach produces better-calibrated uncertainty estimates than Gaussian approximations, especially on real-world datasets and high-dimensional tasks.
2 Related Work
Bayesian Neural Networks. BNNs provide principled approach to uncertainty quantification by treating the network parameters 0 random variables and modeling their posterior distribution p(0 DD) However, exact inference is intractable due to the high dimensionality of parameter spaces in modern neural networks. Various methods have been proposed to approximate the posterior. Variational Inference (VI) approximates
the posterior by optimizing a surrogate distribution, often in the form of a mean-field Gaussian (Graves) 2011 Blundell et al. 2015). Extensions to hierarchical and amortized variational methods have improved scalability but often struggle to capture complex postcrior structures (Zhang et al. 2018). The Laplace Approximation
(LA) defines the posterior locally around the MAP estimate using Gaussian distribution (MacKay 1998
Daxberger et al. 2021a). Despite its simplicity and efficiency, it is limited by its reliance local curvature
2

---

Published in Transactions on Machine Learning Research (07 (07/2025)
and the Gaussian assumption. Sampling-basel methods such Stochastic Gradient Langevin Dynamics (SGLD) (Welling & TTeh 2011) and Hamiltonian Monte Carlo (HMC) (Neal 2011 generate posterior samples by simulating stochastic dynamics. While these methods are more flexible, their computational cost often makes them impractical for large-scale neural networks.
Improving Posterior Samples. To address these limitations, recent works have sought to refine and extend the Laplace approximation by introducing more expressive approximate posteriors. [Kristiadi et al. (022) combines Laplace approximation with normalizing flows for non-Gaussial posterior, refining the base Gaussian distribution. Similarly, Immer et al. (2021) refines Laplace approximation by leveraging Gaussian variational methods and Gaussian processos, improving lincarized Laplace posterior accuracy. Miller et al. (2017) iteratively builds a mixture model to improve the posterior approximation by adding components to the variational distribution. Eschenhagen et al. (2021) combines multiple pre-trained models to form a weighted sum of Laplace approximations, improving posterior flexibility. [Havasi et al. (202T) introduced auxiliary variables to locally refine mean-field variational approximations, achieving better fit in regions of interest. Bergamin al. (2024) extends the Laplace approximation by leveraging Riemannian geometry to manifold, improving accuracy for non-linear loss landscapes.
model the posterior distribution
In this paper, we re-evalnate the efficiency of sampling methods for posterior estimation, particularly in the context of over-parameterized neural networks. Through empirical results, we demonstrate that our proposed simple sampling framework can outperform modern posterior refinement techniques in efficiency and accuracy. Furthermore, our novel use of an autoencoder moves beyond the Gaussian formulation of the posterior, enabling a flexible, easy-to-sample representation that significantly improves performance. This approach allows simplicity, addressing key limitations of both classical and modern refinement methods. Differently from classic hypernetwork approaches, MetricBNN learns a structured latent representation of near-optimal parameter regions explicitly for posterior approximation. This enables us to generate parameter samples efficiently while preserving a geometry-informed posterior structure, directly targeting calibrated Bayesian uncertainty estimation rather than task-conditioned adaptation.
capture the non-linear geometry of the posterior while maintaining computational
3 Preliminaries
BNNs provide a principled framework for quantifying uncertainty in neural network predictions by treating the model parameters as random variables. This section introduces the relevant background on BNNs, discusses the Laplace approximation for posterior estimation, and highlights its limitations.
Consider an independent and identically distributed (i.i.d)) dataset DD ((x, W where x, E R.D and Vi RC. Let fo R + RC denote parametric function (c.g., a neural network) with parameters 0 E R The goal is to model the predictive distribution: p(y' ,²,D) r(s' x', 0)p(0 D) d0, where p(0 () is the posterior distribution of the parameters given the dataset. Bayesian inference provides a framework for estimating the posterior p(0 D) using Bayes' theorem: p(0 D) p(D(o)p()) where p(D 0) is the
p(D)
likelihood of the data, p(0) is the prior distribution over the parameters, and p(D) is the marginal likelihood or
evidence. In practice, computing the evidence p(D) is often intractable, making direct posterior computation challenging.
3.1 Laplace Approximation
The Taplace approximation a classic method for approximating the posterior distribution p(0 | (D) using a Gaussian centered at the MAP estimate. Let L(0) denote the regularized negative log- li-likelihood::
[()) = - loggp(D - a log)()),
(1)
where a is a regularization coefficient. and depends on the choice of prior for the parameters. Assuming this to be the Normal distribution, the regularization can be rewritten as an Euclidean norm.
One of the main advantages of the Laplace approximation that it allows to define a posterior distribution given an already fully trained network, post-hoc. When given fully trained network we assume access to
3

---

Published in Transactions on Machine Learning Research (07/2025)
LA GN
ENSEMBLES
HMC
MetrieBNN Sampling
Figure 2: Postcrior samples for Regression task. The blue points represent the datasct, the orange lines are samples from the estimated posteriors. Our proposed sampling method correctly captures the uncertainty in the data gap.
the of parameters, 0" that minimize the loss function of Equation This is the MAP solution. The exponential of this loss function is proportional to the posterior distribution p(0 D) p(D 0)p(O). Laplace approximation methods propose to approximate the posterior distribution with second-order Taylor expansion around the MAP of the loss function. This results in Gaussian approximation of the parameters with the mean being the MAP solution and the covariance being the inverse of the Hessian computed in the MAP. The posterior can then be defined as:
= -N ((""
(2)
where = (=)(0)) is the Hessian of the loss function at 0"
3.2 Limitations of the Laplace Approximation
While widely used, the Laplace approximation has several limitations:
1. Local Approximation: It captures only the local curvature of the loss landscape around 0", ignoring the broader structure of the postcrior, which can be highly non-Gaussian in high-dimensional spaces (Wilson & Izmailoy 2020
2. Scalability: Computing and inverting the Hessian is computationally expensive for large-scale neural networks, limiting its applicability to small models (Martens) 2010).
Dositive-Defiuitennss:: In over- peraparameeeize networks, the Hessian is often not positive definite, making difficult define a valid Gaussian approximation. Regularization or approximate mothods sometimes used to mitigate this issue, but these approaches can introduce biases (Daxberger et al. 2021a
The limitations of the Laplace approximation motivate the need for alternative methods that better capture the true posterior distribution. Specifically, the posterior for BNNs often lies on a complex, non-linear manifold in parameter space, which the Laplace approximation fails represent. This motivates our approach, which combines efficient exploration of the paramcter space with flexible latent representation to better approximate the posterior.
4 Method
To address the limitations of the Laplace approximation and improve posterior estimation in BNNs, we propose MetricBNN, a two-step framework. The first step involves locally exploring the parameter space around the MAP estimate using a simple sampling method. The second step learns a latent representation to construct flexible posterior distribution that captures the complex geometry of the parameter space.
4.1 Exploring Neighbor Solutions
Given trained neural network with parameters 6* that minimize the regularized loss function (Equation D) the parameter space of deep networks is known to exhibit high degree of redundancy due to overparameterization
4

---

Published in Transactions on Machine Learning Research (07 (025)
Sampling Distribution
MetricBNN Distribution
SVD Approximation
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.
and reparametrization invariances Fort & Jastrzebski (2019); Garipov (2018). These properties create connected, Nower-dimensions manifold of near-optimal solutions surrounding the MAP estimate 6", This phenomenon, often referred to the linear connectivily assumplion, suggests that different sets of parameters can achieve similar performance, even when interpolating between them (Garipov et al. 2018 Fort &
Jastrzebski 2019 Brea et al. 2019).
In particular, empirical studies on the loss landscapes neural networks have shown that minima are often connocted by low-loss paths, forming smooth and structured regions in the paramcter space (Garipov al., 2018). This implies that the posterior distribution is not confined to single mode but instead spans broader, non-linear manifold. Capturing this geometric complexity is crucial for accurately modeling the posterior. By exploring neighboring solutions around 6", we aim to gather representative samples that reflect the underlying structure of this manifold, providing a richer and more accurate approximation of the posterior distribution.
Estimating the posterior distribution of the parameters amounts to identifying the distribution of these solutions. To achieve this, we propose collecting set of such solutions using the following sampling technique (MetricBNN sampling):
1. Initialization: Start with N particles, each initialized at the MAP estimate 8*, i.e., 0,,0-
Random Drift Assignment: Assign each particle a random drift vector d,, sampled as a unit vector with uniform orientation in the parameter space.
Iterative Exploration: For T time steps, update each particle's position as:
(a) Drift: Perturb the particle by adding the corresponding random drift.
0 .-1-1 = 01, a dg.
(3)
(b) Refinement: After each drift update, refine the particle's position using M steps of gradient descent to ensure alignment with the loss landscape.
gm+1 0,,+1 1) VoL(
(4)
0 0M1
(5)
This procedure generates a set of N x T viable solutions near 0* These samples represent a local exploration of the posterior distribution of interest, capturing the diversity of solutions around the MAP estimate. Figure on the right) illustrates an example of these solutions for a two-dimensional regression problem.
While sampling methods have traditionally been considered computationally expensive for high-dimensional neural networks, leading to their relative neglect in favor of variational inference and other approximation

---

Published in Transactions on Machine Learning Research (07 (07/202) (025)
Figure 4: Estimated posterior for the Banana classification task. The orange and ycllow 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.
techniques MacKay 1998 Blundell et al. 2015 argue that these concerns warrant re-evaluation in light of the linear connectivity assumption. Specifically, the high redundancy and structured geometry of neural network parameter spaces suggest that meaningful posterior exploration can be achieved through efficient sampling on a lower-dimensional, connected manifold of solutions (Garipov et al. 2018 Fort & Jastrzebski 2019). This implies that the computational complexity of sampling methods may scale more favorably with the dimensionality of modern networks than previously assumed.
The empirical covariance matrix of the collected solutions provides improved local Gaussian approximation of the posterior distribution. However, due to the extreme non-linearity of the solution manifold, Gaussian approximation alone fails to accurately capture the true structure of the posterior distribution. Addressing this limitation requires moving beyond the Gaussian assumption.
4.2 Defining a Posterior
While the empirical covariance of the collected samples provides a basic Gaussian approximation, it fails to capture the non-linear geometry of the posterior, see Figure 3 Addressing this limitation requires moving beyond the Gaussian assumption. Given that the drift applied in the Drift step (3a is small enough, we can assume that every element in between two sets of parameters is a viable solution. The curve defined by all the samples can, however, be arbitrarily non-linear. To overcome this, we propose to learn a representation that maps the parameters into structured latent space where the collected trajectories of parameters are linearly distributed.
We define a latent representation through an autoencoder framework:
p: Z: An encoder that maps neural network parameters to a latent space Z C R&
: θ: A decoder that reconstructs parameters 0 from the latent space.
The training dataset for the autoencoder consists of the collected samples, i.c., Dg. The goal is to learn a latent space where the posterior distribution well-structured and linearly interpolable. We achieve this by optimizing the following loss function:
C C= + A-C_ + Cds
(6)
with:
C+ = Ep, (11100) -(0')1 11 T
27
(7)
C. = Eng [- log (1):(()) ((0"))),,
(8)
Ca ED, [[ly (v(()) 0113]
(9)
where and 6' are two sets of parameters of the same trajectory and successive time step while d" another randomly sampled set of parameters. Both X and X_ scalar values. The role of these terms is as follows:

---

Published in Transactions on Machine Learning Research (07 (07/202) /025)
20
10-
3
GN LA NIM
HMC
Figure 5: Trade-off between network size and computational complexity on the regression task. Each marker shows NLL (lower is better) for network with the indicated number of hidden layers. MetricBNN's cost. grows only mildly with depth.
L+: encourages local distances between successive samples are preserved.
L_: encourages these sets of parameters to stretch in the learned latent space by maximizing every pair-wise distance.
Ld: encourages reconstruction, mapping latent representations back to their original parameter space.
Using the trained autoencoder, define the posterior distribution in the latent space Z, MetricBNN posterior. Each trajectory of samples is treated independent, with uniform probability assigned to each trajectory. Along cach trajectory, we place a uniform probability between the MAP solution and the last sample of such trajectory, i.e., Og, Sampling from the posterior involves: Sample a trajectory index i uniformly, Sample a scaling factor ~ Uniform(0, 1). Compute the parameter sample:
0 x- (v((0)) + ce. (((0..)) ()()))))
(10)
This method captures the broader geometry of the posterior while maintaining computational efficiency. We provide a pseudo-code the whole method Appendix
5 Experiments
In this section, we evaluate the ability of our proposed method to approximate the posterior efficiently. We first assess the quality of our sampling technique in capturing model uncertainty using simple regression and classification tasks. We then analyze the computational complexity of our approach as the network size increases. Finally, we present quantitative results on real-world datasets, evaluating the negative log-likelihood (NLL) and the expected calibration error (ECE) of our posterior distribution.
We benchmark our method against Bayesian inference baselines that collectively span Docal-Gaussial srcchastic-saplingg and iitermnnistic-esembling paradigms.
Gauss Newton Laplace (LA GN) (Daxberger et al. 2021a) fits Gaussian around the MAP solution using the Gauss-Newtt curvature, offering lightweight post-hoc posterior of any trained network. For some experiments, we also include low-rank estimate of the Laplace approximation (LA LowRank) and a lincarized version (Lin LA).
Riemannian Laplace (Riem LA) (Bergamin et al. 2024) extends this idea by replacing the Euclidean metric with an appropriate Riemannian version, yielding a curvature-aware Gaussian that better respects the manifold structure of the parameter space. We also include the linearized version of this approach (Riem Lin LA).

---

Published in Transactions on Machine Learning Research (07 (07/2025)
For fully stochastic references we include Hamiltonian Monte Carlo (HMC) (Neal 2011 and its scalable stochastio-graddiet variant. (SGLDD (Welling & Teh, 2011), both of which provide asymptotically exact posterior samples at the cost of increased computation.
Complementing these, Deep Ensembles [Lakshminarayanan et al. 2017 (ENSEMBLES) approximate the posterior with a set of independently trained networks, while (SWAG) (Maddox et al. (019) captures the first two moments of SGD iterates form low-rank Gaussian that can be sampled cheaply.
Finally, (ESPRO) (Benton et al. 2021 exploits mode-connecting simplices to generate functionally diverse checkpoints without extra training epochs, yielding a fast ensembling baseline.
We conduct experiments across five different settings:
Simple Regression: We evaluate our method on the onc-dimensional regression problem introduced Snelson & Ghahramani (2005). The dataset consists of 200 points, with 50 held out to assess the modcl's ability to estimate uncertainly. We use fully connected neural network with three hidden layers of 32 units and ReLU activations.
Simple Classification: We test our method on the Banana datasel, a two-dimensional binary classification task with 5,300 points, of which 30% are reserved for testing. We use a fully connected network with two hidden layers of 6 units and Tanh activations.
Classification on UCI Datasets: To evaluate performance on real-world structured data, we experiment. with six classification datasets from the UCI repository (Markelle et al. 2023). We use fully connected network with two hidden layers of 32 units and ReLU activations.
Image Classification: For high-dimensional problems, we experiment with the MNIST (LeCun 1998), FashionMNIST (Xiao et al., 2017) and CIFAR10 (Krizhevsky) 2009) datasets. For MNIST and FMNIST we use a shallow convolutional neural network with two convolutional layers followed by three fully connected layers with Tanh activations. For CIFAR10 instead use ResNet18 architecture (He et al., 2016 and compute the posterior on the last two layers of the network.
OOD Classification: We test the ability of the model in classifying MNIST images with an increasing amount of rotations (Daxberger et al. 2021b) and CIFAR10 images with 15 different realistic noises (Hendrycks
& Dietterich 2019).
5.1 Scalability of the Proposed Sampling Method
Approximating the posterior distribution in neural networks is challenging due to the high-dimensionall and non-linear nature of the parameter space. The standard Laplace approximation has been shown to struggle in these settings (Ritter al. 2018 Lawrence 2001), as assumes a Gaussian posterior, which may not align well with the actual distribution.
Our proposed sampling method provides more flexible alternative, allowing for tighter and better-calibrated posterior that is independent of the loss landscape's ' curvature. Figure 2 compares our method with the Laplace approximation on the toy regression task, demonstrating a significantly improved uncertainty estimation.
Beyond accuracy, scalability is key factor Bayesian inference for deep networks. Many posterior approximation techniques rely on computing second-order derivatives, making them computationally expensive the network size grows. To assess the computational trade-off, analyze the inference time and NLL performance as the number of layers in the network increases. Figure 5 presents these results for our method, LA GN, Riem LA, ENSEMBLES, HMC. Results show that Hessian-based methods and HMC scale poorly with size. ENSEMBLES have a higher temporal cost as they require the training of multiple models. Moreover, they suffer from a very -iw-dimeniional network as they converge to the same solution and thus result in low variance solutions. Our sampling approach maintains a reasonable computational cost while preserving accuracy. Furthermore, our learned latent posterior model enables rapid parameter sampling, further reducing the overhead compared to iterative sampling.
8

---

Published in Transactions on Machine Learning Research (07/2025)
Table 1: Quantitative results on the UCI datasets, including NLL, lower is better, and ECE, lower better.
Model
Australian
Breast
Glass
lonosphere
Vehicle
Waveform
LA GN
0.71 + 0.06
0.73 + 0.07
2.28 + 0.28
0.72 + 0.09
0.8 + 0.16
0.88 + 0.09
LA LowRank
0.69 + 0.03
0.64 0.10
2.06 0.15
1.08 + 0.21
0.76 0.05
0.96 + 0.04
Riem LA
0.54 = 0.07
0.59 = 0.1
1.42 = 0.21
0.17 + 0.02
0.65 + 0.02
0.3 + 0.01
Lin LA
0.69 + 0.04
0.04
3.59 1.17
0.42 1 0.04
0.68 1 0.01
0.02
Riem Lin LA
0.74 + 0.05
2.32 1 0.74
15.92 1 0.05
13.5 1 0.66
0.65 0.01
0.38 1 0.02
NLL
ENSEMBLES
0.51 + 0.07
+ 0.10
0.90 0.08
0.25 + 0.05
0.63 0.00
0.30 0.01
HMC
1.30 + 0.31
1.52 + 0.61
1.49 0.40
0.80 + 0.30
0.65 +
0.43 0.03
SGLD
0.64 + 0.16
0.65 + 0.10
1.36 0.03
0.26 + 0.07
0.73 0.02
0.72 0.09
ESPRO
0.76 = 0.25
0.63 = 0.12
0.17
0.22 0.10
0.64 +.03
0.32 + 0.01
SWAG
0.70 = 0.04
0.59 + 0.06
1.56 = 0.08
0.25 + 0.08
0.68 + 0.00
0.30 + 0.01
MetricBNN
0.66 + 0.05
0.57 + 0.04
1.07 + 0.07
0.17 + 0.05
0.69 + 0.02
0.3 + 0.01
LA GN
13.61 4.34
18.53 =
12.25 + 3.56
28.03 = 7.23
21.36 9.26
21.68 = 6.40
LA LowRank
17.12 + 3.17
16.66 + 9.26
17.72 1 10.25
49.68 1 5.82
18.88 1 6.29
13.82 1 4.38
Riem LA
10.28
1 2.91
21.30 1: 7.77
5.57 0.78
5.80 1.10
5.62 0.70
Lin LA
10.62 - 1.71
18.45 4.57
15.32 1: 7.21
28.80 + 1.67
7.31 + 1.30
23.00 1 1.82
Riem Lin LA
8.47 + 2.71
20.90 + 2.58
14.10 + 2.14
6.84 + 2.31
10.75 + 1.88
2.51 + 0.72
BCE
ENSEMBLES
8.09 + 1.37
15.46 3.77
13.47 + 1.57
8.24 1.14
2.90 + 1.30
1.39 + 0.26
HMC
16.03 = 2.87
26.04 = 4.98
18.32 + 2.43
9.24 2.10
8.46 + 1.17
6.53 + 0.96
SGLD
13.38 3.36
15.12 = 6.35
18.40 + 3.09
8.12 2.58
12.12 + 3.03
8.70 + 4.52
ESPRO
11.62 - 3.33
15.41 + 3.37
14.03 1 2.91
6.47 1 1.72
8.29 + 1.62
2.51 1 0.96
SWAG
9.80 1 6.35
14.29
1 2.49
3.87
1 2.75
7.44 1. 2.12
1.18 1 0.98
2.39 1 0.59
MetricBNN
27.33 = 8.08
12.42 + 4.38
11.51 + 5.46
5.74 + 0.97
10.18 + 10.20
1.50 + 0.53
Table Quantitative results on image datasets, including NLL, lower better, and FCC., lower better.
Model
MNIST
FMNIST
CIFAR10
LA GN
2.42 0.03
2.24 1 0.08
2.29 1 0.09
LA LowRank
1.35 0.04
1.67 + 0.06
2.19 + 0.08
Riem LA
1.11 + 0.11
1.20 0.02
0.02
0.84 + 0.03
Lin LA
0.08
0.13
1.93
NLL
Riem Lin LA
0.62 + 0.09
0.87 + 0.07
0.83 + 0.06
ENSEMBLES
0.17 0.00
0.47 + 0.00
1.55
SCLD
1.60 1 0.20
1.98 + 0.26
0.78 1 0.06
ESPRO
0.15 0.02
0.51 1
0.84 1
SWAG
0.24 + 0.01
0.89 + 0.11
1.97 + 0.33
MetricBNN
0.12 0.01
0.58 + 0.09
0.70 + 0.04
LA GN
4.78 + 2.97
12.31 + 2.39
4.59 + 1.88
LA LowRank
50.58 +
18.70 3.07
9.74 5.82
Riem LA
26.78 + 7.83
33.91 + 8.23
Lin LA
29.43 1.32
19.63 1.18
35.46 12.67
ECE
Riem Lin LA
21.20 - 2.00
15.19 + 1.92
32.10 + 3.15
ENSEMBLES
10.52 0.13
11.33 0.29
8.94
SGLD
13.83 1 7.47
14.55 6.06
7.45 1.25
ESPRO
1.56 0.20
4.36 + 0.79
10.40 + 0.33
SWAG
13.47 + 1.07
25.25 + 4.95
3.47 2.76
MBNN
1.57 0.31
13.57 + 6.89
9.26 + 0.45
5.2 Quality of the Proposed Geometric Posterior
A potential concern is whether the posterior obtained via our sampling method could be approximated using naive covariance computation, similar to the Laplace approximation. However, as shown in Figure X this approach still does not yield accurate uncertainty estimates. By contrast, latent posterior representation

---

Published in Transactions on Machine Learning Research (07 (07/2025)
effectively captures the true posterior structure, demonstrating the benefits of learning structured parameter space.
5.3 Performance on UCI Regression and Image Classification
We further evaluate our approach to real-world structured data by testing it on standard UCI regression benchmarks. Table reports NLL and ECE results, showing that our method remains compelitive with the baselines. In Appendix we provide details on the computational cost as well demonstrating that our method offers a good balance between accuracy and computational cost.
To performance in high-dimensional settings, Irain convolutional neural networks MNIST and INST and CIFAR1O. Table 2
presents NLL and ECE results.
Gaussion Noise
Shot Notse
Impulse Note
54.0
Un A
Defecus Noise
Frested Glass Blur
Mution Blur
50L0
EMSTMALDS 05
Un
Snow
Noase
3440
15/10
62.0
Brightress
N67 128.0
SWAG
15/10
SULD
EMSTH0 05
0.0
2.5
1.0
3.5
2.9
2.0
0.5
3.5
1.3
2.0
1.0
0.5
2.0
2.3
2.5
Nor 22NN
5WAU
ESPR
62.0
6.0
0.5
1.5
2.0
Figure 6: Robustness to semantic and photometric shift CIFAR10. Each panel reports test NLL across each corruption.
5.4 OOD Classification
We evaluate all methods under two canonical OOD settings. For every corruption we report the log-likelihood.
All networks are first trained on clean data only. Posterior approximations are then post-hoc.
10

---

Published in Transactions on Machine Learning Research (07/ 2025)
Geometric shift: progressively rotating MNIST test images by 00-900 (Daxberger et al.) 2021b). Figure 7 shows that Met has competitive log- logllikel performances for ro- tations up to 30 degrees. Performances drop when the rotations become more severe as distinguishing wetween the numbers is not always possible.
Semantic + photometric shift: applying the fifteen cor- ruption families of CIFAR10 (Hendrycks & Dietterich 2019). Figure D shows competitive performances for MetricBNN for some of the perturbations.
Glation
SOLD
SWAG
LALIOWRANC
ESPRO
MARN
Our experiments highlight three key takeaways: Our sampling- based postcrior exploration provides a more flexible and well- calibrated uncertainly estimate, particularly as network size increases. The proposed latent posterior model enables efficient posterior sampling, avoiding the computational overhead of iterative methods. Our approach remains competitive in terms of NLL and F.C on real-world datasets while being significantly more computationally efficient.
ENGONIL
Figure 7: Robustness to geometric shift (Ro- tated MNIST). Test NLL and standard devi- ations of MNIST image classification versus rotation angle; lower is better.
6 Discussion
In this work, we revisited scalable Bayesian inference in deep learning, benchmarking our method, MetricBNN, against a broad range of baselines including Laplace approximations, variational methods, deep ensembles, and sampling- based techniques. Across structured and high-dimensional tasks, MetricBNN consistently achieves competitive or superior negative log-likelihood and calibration while maintaining efficiency as network size increases. We observe that while Laplace-based methods and their geometric extensions can perform well in certain settings, their performance is notably sensitive to the parameter count, network capacity, and the quality of the MAP solution used for centering the approximation. Differences in MAP training conditions and model size can therefore lead to the lower scores observed in our experiments compared those reported in other Laplace-focusee focused studies. Our sampling- based exploration combined with a learned latent posterior helps mitigate these sensitivities, capturing non-Gaussian, non-linear postcrior structures more robustly across diverse regimes.
Across our experiments, MetrieBNN demonstrates strong empirical performance on structured datasets, high- dimensional image classification, and out-of-distribution (OOD) robustness benchmarks. On structured UCI datasets, MetricBNN achieves negative log-likelihood and calibration scores competitive with or better than Laplace (T.A GN, Riem L.A) and variational baselines, while providing a more flexible posterior representation. On high-dimensional datasets like MNIST, FashionMNIST, and CIFARIO, MetricBNNN maintains low NLL and ECE scores, matching outperforming deep ensembles and methods such as SWAG, while avoiding the compulational overhead of training multiple models or computing Hessians at scale. In OOD settings, MetricBNN shows consistent robustness under geometric shifts (rotated MNIST) and photometric and semantic corruptions (CIFARIO), providing well-calibrated uncertainty estimates across perturbation levels.
Compared to fully sampling- based methods such as Hamiltonian Monte Carlo (HMC) and SGLDD, MetricBNN achieves comparable uncertainty estimates while being significantly more computationally efficient, making it practical for large-scale neural network applications. In contrast to deep ensembles, which require training multiple networks to achieve uncertainty calibration, MetricBNN provides comparable calibration with a single network, improving efficiency and scalability. Compared to variational inference methods, which often rely on restrictive mean-field approximations, MetricBNN captures the non-Gaussian structure of the posterior by leveraging the geometry of the parameter space through structured latent representations.
Finally, while hypernetworks typically generate weights conditioned on input or task information for adaptation, our approach fundamentally differs by using an autoencoder to learn a structured latent space of near-optimal parameter regions for scalable posterior sampling. This allows MetricBNN to efficiently generate parameter
11

---

Published in Transactions on Machine Learning Research (07/2025)
samples that reflect the geometry of the loss landscape without the use of expensive iterative techniques variational methods.
7 Conclusions
In this work, we proposed a simple variation of sampling-basee based techniques tailored to explore the posterior geometry of Bayesian Neural Networks efficiently, even in over-parancterized scttings. By leveraging the low-dimensionaa structure of loss minima, our method achieves competitive posterior approximations while maintaining scalability as network size increases. Additionally, we introduced a model that learns a deformation of the parameter space based on the collected samples, enabling rapid posterior sampling without requiring iterative methods. Our empirical results demonstrate that this approach improves posterior accuracy and computational efficiency compared to recent refinement techniques. These contributions provide a practical and flexible framework for Bayesian inference in deep learning, offering new directions for scalable uncertainty quantification in complex models.
Acknowledgments
This work was supported by the Swedish Research Council (2019-00748), the Knut and Alice Wallenberg Foundation, the European Research Council (ERC-BIRD--84807) and the European Horizon 2020 CANOPIES project.
References
Gregory Benton, Wesley Maddox, Sanae Lotfi, and Andrew Gordon Gordon Wilson. Loss surface simplexes for mode connecting volumes and fast ensembling. In International Conference on Machine Learning, pp- 769 779. PMLR, 2021.
Federico Bergamin, Pablo Moreno- Munno, Søren Hauberg, and Georgios Arvanitidis. Riemannian laplace approximations for bayesian neural networks. Advances in Neural Information Processing Systems, 36, 2024.
Charles Blundell, Julien Cornebise, Koray Kavukcuoglu, and Daan Wierstra. Weight uncertainty in neural network. In International conference on machine learning, pp. 1613-1622. PMLR, 2015.
Johanni Brea, Berfin Simsek, Bernd Illing, and Wulfram Gerstner. Weight-space symmetry in deep networks gives rise to permutation saddles, connected by equal-loss valleys across the loss landscape. arXiv preprint arXis:1907.0291 2019.
Antonio S Cofino, Rafael Cano, Carmen Sordo, and Jose M Gutierrez. Bayesian networks for probabilistic weather prediction. In 15th European Conference on Artificial Intelligence (ECAI). Citeseer, 2002.
Erik Daxberger, Agustinus Kristiadi, Alexander Immer, Runa Eschenhagen, Matthias Bauer, and Philipp Hennig. Laplace redux-efforteess bayesian deep learning. Advances in Neural Information Processing Systems, 34:20089 20103, 2021a.
Erik Daxberger, Fric Nalisnick, James U Allingham, Javier Antorán, and José Miguel Hernándew-LLobato. Bayesian deep learning via subnetwork inference. In International Conference on Machine Learning, pp- 2510-2521. PMLR, 2021b.
Runa Eschenhagen, Erik Daxberger, Philipp Hennig, and Agustinus Kristiadi. Mixtures of laplace approxi- mations for improved post-hoe uncertainty in deep learning. ar preprint arXiv:2111.03577, 2021.
Di Feng, Lars Rosenbaum, and Klaus Dietmayer. Towards safe autonomous driving: Capture uncertainty the deep neural network for lidar 3d vehicle detection. In 2018 21st international conference on intelligent transportation systems (ITSC), pp. 3266 3273. IEEE, 2018.
Stanislav Fort and Stanislaw Jastrzebski. Large scale structure of neural network loss landscapes. Advances in Neural Information Processing Systems, 32, 2019.
12

---

Published in Transactions on Machine Learning Research (07/2025)
Timur Garipov, Pavel Izmailov, Dmitrii Podoprikhin, Dmitry P Vetrov, and Andrew G Wilson. Loss surfaces, mode connectivity, and fast ensembling of dnns. Aduances in newral information processing systems, 31, 2018.
Alex Graves. Practical variational inference for neural networks. Advances in neural information processing systems, 24, 2011.
Marton Havasi, Jasper Snoek, Dustin Tran, Jonathan Gordon, and José Miguel Hernández-Lobato. Sampling the variational posterior with local refinement. Entropy, 23(11):1475, 2021.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference computer vision and pattern recognition, pp. 770-778, 2016.
Dan Hendrycks and Thomas Dietterich. Benchmarking neural network robustness to common corruptions and perturbations. arXiv preprint arXiv: 12261, 2019.
Alexander Immer, Maciej Korzepa, and Matthias Bauer. Improving predictions of bayesian neural nets via local linearization. In International conference artificial intelligence and stafistics, pp. 703 711. PMLR, 2021.
Agustinus Kristiadi, Runa Eschenhagen, and Philipp Hennig. Posterior refinement improves sample efficiency in bayesian neural networks. Advances in Neural Information Processing Systems, 35:30333 30346, 2022.
Alex Krizhevsky. Learning multiple layers of features from tiny images. Technical Report TR- 2009, University of Toronto, Toronto, Ontario, 2009. URL https: https://wwtt http///wwwc.oo csw.cc.. toronto..o...ooooo edu/ edu/-kriz/ Learning- features--eeeeeee TR.9-TR.ppr
Balaji Lakshminarayanan, Alexander Pritzel, and Charles Blundell Simple and scalable predictive uncertainty estimation using deep ensembles. Advances in neural information processing systems, 30, 2017.
Neil David Lawrence. Variational inference in probabilistic models. PhD thesis, Citeseer, 2001.
Yann LeCun. The mnist database of handwritten digits. hllp://yamm. lecun. /ee//////m/is//../////////// 1998.
Christian Leibig, Vaneeda Allken, Murat Seçkin Ayhan, Philipp Berens, and Siegfried Wahl. Leveraging uncertainty information from deep neural networks for disease detection. Scientific reports, 7(1):1-14, 2017.
David JC MacKay. A practical bayesian framework for backpropagation networks. Neural compulation, 4(3): 448 472, 1992.
David JC MacKay. Choice of basis for laplace approximation. Machine learning, 33:77-86, 1998.
Wesley J Maddox, Pavel Izmailov, Timur Garipov, Dmitry P Vetrov, and Andrew Gordon Wilson. A simple baseline for bayesian uncertainty in deep learning. Advances in newral information processing systems, 32, 2019.
K. Markelle, Rachel, and N Kolby. Uci dataset. the uci machine learning repository, 2023.
James Martens. Deep learning via hessian-free optimization. In Proceedings of the 27th International Conference on International Conference on Machine Learning, pp. 735-742 2010.
Andrew C Miller, Nicholas Foti, and Ryan P Adams. Variational boosting: Iteratively refining posterior approximations. In International Conference on Machine Learning, pp. 2420 2429. PMTR, 2017.
Radford M Neal. Mcme using ensembles of states for problems with fast and slow variables such as gaussian process regression. Xix preprint arXiv: arXiv:1101.03877 101.0387, 2011.
Hippolyt Ritter, Alcksandar Botev, and David Barber. A scalable laplace approximation for neural networks. In 6th international conference on learning representations, ICLR 2018-conference track proceedings, volume 6. International Conference Representation Learning, 2018.
13

---

Published in Transactions on Machine Learning Research (07/2025)
Edward Snelson and Zoubin Ghahramani. Sparse gaussian processes using psendo-inputs. Advances in neural informalion processing systems, 18, 2005.
Max Welling and Yee W Teh. Bayesian learning via stochastic gradient langevin dynamics. In Proceedings of the 28th international conference on machine learning (ICMI-11), pp. 681-688. Citeseer, 2011.
Andrew G Wilson and Pavel lamailov. Bayesian deep learning and a probabilistic perspective of generalization. Advances in neural information processing systems, 33:4697 4708, 2020.
Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-mnist: novel image dataset for benchmarking machine learning algorithms. ar preprint 1t:i::108.07774/ 2017.
Cheng Zhang, Judith Bütepage, Hedvig Kjellström, and Stephan Mandt. Advances in variational inference.
IEEE transactions pattern analysis and machine intelligence, 41(8):2008-200 2026, 2018.
14

---

Published in Transactions on Machine Learning Research (07 (07/2025)
A
Experimental Details
This section provides additional dctails regarding the experimental sctup, including the models, hyper- parameters, and computational settings.
A.1 Experimental Setup
We evaluate our method across four types of lasks: toy regression, toy classification, structured classification from the UCI repository, high-dimensional image classification and OOD classification.
Toy Regression. We consider the Snelson 1D regression dataset (Snclson & Ghahramani 2005 consisting of 200 points, with 50 held out for evaluating uncertainly estimation. We use a fully connected neural network with three hidden layers of 32 units and ReLU activations. The model is trained using the Adam optimizer with a learning rate of 1 x 10-3, batch size of 200, and L2 regularization of 0.01 for 50,000 epochs.
For the sampling procedure of our approach, sampling steps. At each step, we updated the parameters using M = 10 gradient refinement steps to improve posterior estimates. We initialized N = 10 particles, each following a perturbed trajectory in the parameter space, with a drift step scaled by an inner learning rate of = 0.001. For the autoencoder-based posterior model, we projected the sampled parameter trajectories into a structured latent space of k = 32 dimensions. The autoencoder was trained using batch size 1024 for 1000 epochs, optimizing the contrastive loss with weight coelficients X+ 1.0 and X_ = 1.0. The archilecture consists of an encoder with two fully connected hidden layers of 256 units and ReLU activations, and decoder with symmetric structure of two hidden layers with 256 units and ReLU activations. Training was performed using the Adam optimizer.
used drift-scaling factor of a 0.1 and performed T 100
Toy Classification. For classification, we use the Banana dalasel, a 2D binary classification problem with 5300 points, of which 30% are used for testing. The network consists of two hidden layers with 16 units and Tanh activations. The training procodure follows the regression sctup, except with smaller batch size (32) and 2500 training epochs. For our model, we use the same parameters as in the regression task.
Structured Data (UCI Datasets). To evaluate our method in structured classification settings, we test six datasets from the UCI reposilory (Markelle et al. 2023). using fully connected architecture with two hidden layers of 32 units and ReLU activations. The training follows the same schedule as before, with a batch size of 32 and 1000 training epochs. Our method estimated with the same parameters as in the classification experiment except for = 50 steps.
High-Dimensional Image Classification, For large-scale experiments, we train Bayesian neural networks MNIST (LeCun 1998). : FashionMNIST (Xiao et al.) 2017) and CIFAR10 (Krizhevsky 2009). For MNIST and FAMNIST we shallow convolutional neural network consisting of two convolutional layers, followed by three fully connected layers, with Tanh activations. For CIFAR10 we use standard ResNet18 architecture. Training follows the same setup as UCI experiments, but with reduced number of epochs (100). Our method follows the same parameters as in the UCI experiment.
OOD Classification. We use the same networks described above for both MNIST and CIFAR10 as well the same parameters for our method.
A.2 Scalability and Computational Complexity
To conduct the experiment on the computational complexity depicted in Figure 5 we used a fully connected network with hidden layers from 1 to 9. Each layer with 64 units and ReLu as activation function. For each architecture, we measure the inference time and negative log-likelihood (NIL.) performance. To compute the posterior of each model we use 100 samples of the parameters. All experiments were conducted on NVIDIA RTX3080 GPU.
15

---

Published in Transactions on Machine Learning Research (07 (07/2025)
B
Pseudocode
Algorithm Proposed Posterior Approximation Method
Require: Trained neural network fe, dataset D, drift scaling a, sampling steps T, gradient steps M, particles
N, inner learning rate 7), autoencoder latent dimension k, training epochs E
Ensure: Approximate posterior g(f)
1: Phase 1: Parameter Sampling
2: Initialize N particles at 0" (MAP estimate)
for each particle l...... N do
Sample random drift direction d2 + TeT 2~ 1(0,
5: for each step = I,...,, do
Apply drift: 0.. 0 + ad, for each gradient step 1...., M do
Refine using gradient descent:
0i. t 0z,t (oL(())))
10: end for
11: end for
12: end for
13: Store all sampled parameters = (0..e)
14: Phase Learning Latent Posterior Representation
15: Train autoencoder yp O Z using:
16: for epoch 1,..., E do
17: Sample minibatch of subsequent parameters 0, 0' ~ θ
18:
Compute latent embeddings: 2 = (00) and L' = (00))
19:
Compute negative parameters by reshuffling the batch: 2"
20:
Compute contrastive loss:
21:
C+ = (1|= 1/7))
22:
C. = logl0 I
23:
La (x) 011
24: Update autocncoder using L = X4 + + Adha
25: end for
26: Posterior Approximation and Sampling
27: Sample trajectory index j Categorical ((Zmax.s L
f=1)
28: Sample interpolation factor e U(0,
29: Compute latent posterior sample: 2MAP (zmaxx 2MAP)
30: Map to parameter space: 0= - (z)
31: return Posterior sample 0
16

---

Published in Transactions on Machine Learning Research (07 (07/2025)
C Time Complexity Results
We provide more detailed comparison of the computational complexity and the accuracy of our method and the baselines for the UCI experiment and the images experiment.
Table 3: Logarithmic time complexity in seconds and accuracy for the UCI datasets.
Model
Australian
Breast
Glass
lonosphere
Vehicle
Waveform
LA GN
1.88 + 0.98
1.54 + 0.8
1.56
0.67
1.88 + 0.72
3.05 + 1.95
2.68 0.48
LA LowRank
2.0 + 0.22
1.24 + 0.28
0.98 + 0.37
1.47 + 0.48
4.69 + 0.1
3.68 + 0.16
Riem LA
4.36 + 4.06
2.42 1.61
3.69 + 3.63
2.38 + 1.61
3.33 3.0
3.26 2.45
Log Time
Lin LA
1.86 + 1.0
1.52 + 0.4
1.55 + 0.48
1.89 + 0.65
3.05 = 1.95
2.68 = 0.31
Riem Lin LA
1.9 + 1.04
1.67 0.54
1.93 = 0.65
3.08 1.42
2.69 1.53
ENSEMBLES
2.5 0.02
2.14 0.03
2.01 + 0.02
2.2 + 0.02
3.76 0.01
3.32 1 0.02
HMC
3.3 0.2
0.54 0.21
1.73 0.09
0.36 1 0.24
1.43 0.16
SGLD
1.38 + 0.03
1.06 0.07
0.95 + 0.05
1.11 + 0.04
2
2.16 0.03
ESPRO
0.61 + 0.1
0.33 - 0.16
0.19 0.22
0.31 + 0.18
1.66 0.01
1.27 = 0.03
SWAG
0.95 + 0.01
0.62 0.07
0.51 + 0.08
0.68 + 0.08
2.16 - 0.02
1.72 0.01
MetricBNN
1.66 + 0.34
1.61 + 0.29
1.64 + 0.32
1.76 + 0.35
1.8 + 0.36
1.82 = 0.41
LA GN
0.49 + 0.12
0.57 + 0.21
0.17 + 0.11
0.36 + 0.09
0.45 + 0.04
0.57 + 0.14
LA LowRank
0.55 0.05
0.60 0.23
0.12 0.06
0.33 + 0.06
0.46 0.04
0.57 0.05
Riem LA
0.57 = 0.04
0.72 + 0.01
0.38 + 0.01
0.93 + 0.02
0.61 + 0.01
0.86 = 0.00
Lin LA
0.56 1 0.06
0.62 0.08
0.11 0.09
0.91 1 0.00
0.60 1 0.01
0.01
0.84 0.01
Accuracy
Riem Lin LA
0.75 + 0.01
0.68 1 0.03
0.48 1: 0.18
0.91 0.03
0.02 0.58 0.02
0.86 T 0.01
ENSEMBLES
0.82 0.02
0.71 0.04
0.04
0.04
1 0.02
0.72 1 0.00
0.83 0.00
SGLD
0.73 + 0.03
0.52 0.08
0.92 0.02
0.50 +
0.64 = 0.03
ESPRO
0.82 + 0.02
0.77 0.01
0.01
0.67 0.08
0.70 + 0.00
0.86 0.01
SWAG
0.57 = 0.10
0.71 0.04
0.34 # 0.05
0.02
0.58 - 0.01
0.86 0.01
MetricBNN
0.50 = 0.04
0.76 + 0.01
0.57 + 0.04
0.95 + 0.01
0.58 + 0.01
0.85 = 0.01
Table 4: Logarithmic time complexity in seconds and accuracy for the image datascts.
Model
MNIST
FMNIST
CIFAR10
LA GN
3.57 + 0.02
2.24 + 0.08
4.79 0.16
LA LowRank
2.86 + 0.11
1.67 + 0.06
3.83 0.02
Riem LA
5.22 = 3.16
5.01 4.09
5.31 0.08
Log Time
Lin LA
3.64 + 2.80
3.67 + 2.96
4.81 + 0.06
Riem Lin LA
3.90 + 3.18
3.91 +
5.50 0.09
ENSEMBLES
4.07 0.04
4.03 0.05
5.35 + 0.01
SGLD
2.97 + 0.05
2.97 1 0.05
4.17 0.08
ESPRO
2.12 0.05
2.10 + 0.05
3.93 + 0.10
SWAG
2.51 + 0.06
2.51 + 0.06
4.20 0.10
MetricBNN
3.61 = 0.05
3.67 + 0.09
4.22 0.10
LA GN
0.19 + 0.02
0.19 + 0.04
0.16 0.04
LA LowRank
0.77 + 0.04
0.45 1 0.03
0.26 + 0.06
Riem LA
0.71 + 0.04
0.38 0.03
0.45 0.05
Accuracy
Lin LA
0.85 + 0.03
0.67 + 0.04
0.52 - 0.04
Riem Lin LA
0.92 + 0.02
0.75 0.06
0.82 0.07
ENSEMBLES
0.98 + 0.01
0.87 + 0.02
0.86 + 0.02
SGLD
0.38 + 0.14
0.25 + 0.09
0.85 0.10
ESPRO
0.96 + 0.02
0.83 + 0.03
0.86 + 0.03
SWAG
0.97 + 0.02
0.77 + 0.02
0.16 0.03
MetricBNN
0.97 = 0.03
0.82 + 0.02
0.85 0.40
17

---

Published in Transactions on Machine Learning Research (07 (07/2025)
D Sensitivity of the Hyper- Parame-ars for MetricBNN
We additionally provide experiments on the sensitivity of the hyper-parameteers of MetricBNN. Figures B and show the difference in the sampling distribution of MetrieBNN when varying the parameters a, N, T. As expected the variance increases with an increase in the number and length of the trajectories as well as the norm the drift term. Figures I0 and [1 provide the negative loglikelihood for the same varying parameters plus the number of hidden dimensions of the autocncoder & for the posterior distribution.
-douadd in Parks Ned Me.10
aett bwild bert
#otwE Mell
Figure 8: Samples from MetricBNNN on the Toy regression problem with varying length of sampling (T" and number of particles (N).
18

---

Published in Transactions on Machine Learning Research (07/202) 2025)
Figure 9: Samples from MetricBNN on the Toy regression problem with varying X and A_ for training the autocncoder (Aa fixed to 1).
19

---

Published in Transactions on Machine Learning Research (07/2025)
alpha 0.01
N=5 T=40
N=5 T=80
N=5 T=100
5.44
15.41
2.81
2.61
2.66
1.54
3.12
3.96
0,80
17.55
32
2.51
7.66
4.75
*
32
2.08
3.16
1.85
32
1.18
3.11
0.60
3.19
2.80
3.21
3,20
2.31
0.62
5.56
15.0
10
10
5
10
N-10 T-40
N=10 T-80
T=100
12.5
512
3.32
2.56
3.05
512
2.91
7.44
3.40
512
1.45
3.06
0.53
0.80
5.54
1.06
4.32
1.43
3.80
3.27
10.0
5.66
3.23
1.81
2.94
2.12
1.33
1.59
1.41
1.20
7.5
20
M
M
M
N=100 T=40
N=100 T=80
N=100 T=100
5.00
4.76
19.26
6.85
3.70
3.13
5.63
2.13
1.25
1.66
32
6.46
13.59
2.56
32
1.25
2.85
1.99
32
1.48
4.47
2.95
6.03
8.32
6.59
0.61
3.42
1.56
1.46
1.69
10
20
10
30
Negative log- ligllikeldhood on the Toy regression experiment for MetricBNN with a = 0.01 and different values of N and T (different subplots) and M and (within each subplot).
Figure 10:
20

---

Published in Transactions on Machine Learning Research (07 (07/2025)
alpha 0.1
N-5 T-100
512
1.26
1.69
0.85
512
0.96
1.46
1.41
512
1.10
1.21
1.44
32
1.65
1.38
0.91
32
1.13
1.42
0.90
32
1.28
0.99
1.28
2.0
1.03
1.79
0.87
0.91
1.26
1.09
1.22
1.03
1.17
1.8
-0
-P
1
N=10 T=40
N=10 T=80
N=10 T=
1.6
1.11
1.4
1.16
1.54
1.44
1.09
1.17
1.10
0.98
1.05
1.02
1.14
1.37
1.03
0.98
1.04
0.98
1.08
1.07
1.11
1.2
10
10
19
10
M
M
M
N=100
N=100 T=80
N=100 T=100
-10
512
0.58
0.71
1.61
$12
0.76
0.87
0.94
512
0.88
0.81
0.92
0.8
32
0.66
0.82
1.33
32
1.01
0.94
0.98
32
0.91
0.91
0.97
0.80
0.85
1.26
0.90
1.07
1.05
0.98
0.93
1.01
-0
P
10
4
Figure 11: Negative log- likelihood on the Toy regression experiment for MetricBNN with a 0.1 and different values of and (different subplots) and M and (within each subplot).
21

---

Published in Transactions on Machine Learning Research (07/2025)
E Principal Components Analysis on Estimated Posterior
We additionally present the two principal components of the posterior samples of both MetricBNN and HMC Given enough steps, the presented
for the Banana dataset with varying amounts of steps in Figure 12 method covers a large part of the parameter space.
20
15.
5-
5-
g
2
-10
-10
-10
-10 -5
-10
PC1
PC1
PCL
IIMC
MetricONN
MAP
Figure
22


Reviews:
[{"review_id": "ntYN0heZcD", "reviewer": "Reviewer_rEb8", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths:**\n\nThe writing is overall pretty clear especially in the methodological description. Good job.\n\nThere are several understanding experiments towards the practical implications of this type of posterior approximation.\n\n**Weaknesses:**\n\nThere’s unfortunately a fair amount of missing work here, and a few incorrect claims:\n\n- Pg 2: “Sampling-based methods such as stochastic gradient langevin dynamics … [are computationally costly] and impractical for large-scale neural networks.” Yes, HMC is extremely computationally costly (see Wilson & Izmailov, 20 and Izmailov et al, ’21). However, stochastic gradient MCMC is not actually much more costly than standard training (each gradient step is the same cost as standard training), its just the inference costs can be much higher due to the sample collection required – see Welling & Teh, ’11, Zhang et al, ’20.\n\n- Figure 4: Qualitatively, the SVD approximation here displays more of the behavior of a neural network that we would expect – it rapidly becomes unconfident off the data manifold. By comparison, the metricBNN distribution stays very confident in most places, even more so than the “sampling distribution”.\n\n- The overall idea is very similar to randomize – then – optimize in the applied mathematics literature, see Bardsley et al, 14, just with direct optimization instead of sampling thereafter, and to Pearce et al ‘20’s direct sampling around the posterior, except again with direct optimization instead of an injected prior. The idea is also pretty well known in the RL literature for neural linear models typically. \n\n- The autoencoder piece is novel; however, some of the follow up works to Garipov et al, ’18 were able to successfully get low dimensional distributions in the weight space of the neural networks without the need for autoencoders. See for example, Benton et al, ’21 and Wortsman et al, ’21 – the former is more explicitly probabilistic; however, both have Bayesian themes.\n\n- The resemblance to Bayesian hypernetworks (Kruger et al, ’18) should also be discussed as that approach directly trains the autoencoder from scratch to generate weight samples of the NN.\n\nReferences:\n“cyclical stochastic gradient mcmc for Bayesian deep learning,” zhang et al, iclr 2020.\n\n“what are Bayesian neural network posteriors really like,” izmailov et al, icml 2021.\n\n“Bayesian deep learning and a probabilistic perspective on generalization,” Wilson & izmailov, neurips, 2020.\n\n“Bayesian learning via stochastic gradient langevin dynamics,” welling & teh, icml, 2011.\n\n“uncertainty in neural networks: approximately Bayesian ensembling,” pearce et al, aistats, 2020.\n\n“Randomize-Then-Optimize: A Method for Sampling from Posterior Distributions in Nonlinear Inverse Problems”, Bardsley et al, SIAM Journal on Scientific Computing, 2014.\n\n“loss surface simplexes for mode connecting volumes and fast ensembling”, Benton et al, icml, 2021.\n\n“Learning neural network subspaces,” Wortsman et al, 2021.\n\n“Simple and Scalable Predictive Uncertainty Estimation using Deep Ensembles,” Lakshminarayanan et al, 2016\n\n“Bayesian hypernetworks,” Kruger et al, 2018. https://arxiv.org/pdf/1710.04759", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposes a sample then optimize approach for Bayesian neural networks. Specifically, a neural network is trained to completion, then noise is added to corrupt the solution, after which a few gradient steps are taken. When an explicit distribution, rather than samples, are desired, the authors train an autoencoder on a low dimensional space as a sort of hypernetwork to generate new sets of weights.\nExperiments are performed on a toy classification task, MNIST classification, and UCI regression.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Please compare the NLL and log time to some sampling based approaches, e.g. SGMCMC or something analogous. Similarly, a comparison with the low dimensional approach in Benton et al, ’21 would be quite useful as it’s similarly inspired by the mode connectivity work of Garipov et al, ’18.\n\nEnsembles (Lakshimarayan et al, ’16) are also a very good, missing baseline in the two quantitative experiments. \n\nFigure 5: What parameter sizes are the networks here? \n\nIs this the time taken to compute the NLL? With how many samples?\n\nCould accuracy metrics be reported for Table 2? \n\nCould a larger experiment, even CIFAR 10 with conv nets, be performed to test the practical scalability of the approach?\n\nHow many samples are used for training the metric BNN autoencoder, and what is its typical computational cost?", "broader_impact_concerns": "n/a", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "Se0hZOWSkH", "reviewer": "Reviewer_hmoS", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "[Strengths]\n\nElegant use of geometry: By “walking on the fiber” instead of relying on Hessian curvature, the method sidesteps the positive-definiteness issues and cubic costs that plague Laplace in over-parameterized nets.\n\nFast sampling after training: Once the auto-encoder is fitted, drawing a posterior sample is reduced to one latent-space interpolation and one forward pass through the decoder, enabling near-instant Monte-Carlo predictions.\n\nVersatility and minimal intrusiveness: The algorithm is completely post-hoc—no change to the original optimizer or loss—and is demonstrated on fully-connected, convolutional and shallow networks, as well as both regression and classification problems.\n\nCompetitive empirical performance: MetricBNN matches or outperforms four Laplace-based baselines on NLL while keeping wall-time low, and its uncertainty bands on toy problems visually capture data gaps that Laplace approximation misses.\n\n[Weaknesses]\n\nLimited theoretical grounding: The paper assumes, but does not rigorously justify, that solution sets form one-dimensional fibers and that linear interpolation in latent space remains within a high-probability posterior region.\n\nModest experimental scale: All image experiments use small-footprint CNNs on MNIST-like data; there is no evidence the approach remains tractable for modern large language/vision models. Even adding an experiment on small transformer model would greatly enhance the generalizability of their conclusions.\n\nHyper-parameter sensitivity: Performance may hinge on drift scale α, number of particles N, sampling steps T and auto-encoder dimension k, yet only a single setting is reported and no robustness study appears.\n\nComparative breadth: Baselines omit state-of-the-art stochastic methods such as SGLD, SGHMC, SWAG and deep ensembles, so it is unclear how MetricBNN fares against the broader Bayesian-deep-learning toolkit.\n\nEvaluation metrics: The study relies almost exclusively on NLL; calibration (ECE), out-of-distribution detection, or posterior-predictive checks are not reported, leaving uncertainty quality only partially assessed.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper introduces MetricBNN, a two-stage, post-hoc Bayesian procedure for neural networks. Stage 1 draws a cloud of weight vectors by repeatedly adding small random “drift” perturbations to the maximum-a-posteriori (MAP) solution and then re-optimizing for a few gradient steps, exploiting the empirical observation that good solutions lie on a low-dimensional, connected manifold in weight space. Stage 2 trains an auto-encoder on those samples so that, in the learned latent space, entire trajectories become nearly linear; new posterior samples are produced by simply interpolating between the latent MAP point and the endpoint of a chosen trajectory and decoding back to weights. Across toy regression/classification tasks, six UCI datasets, MNIST and Fashion-MNIST, MetricBNN delivers negative-log-likelihoods on par with, or better than, Laplace and Riemannian-Laplace refinements while requiring markedly less Hessian-related computation and scaling almost flatly with network depth.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Can you provide a sensitivity analysis for α, N, T, M and latent dimension k, and clarify how many samples are really needed to train a useful auto-encoder?\n\nWhat is the role of 1/T in Eq. (7)?\n\nThe drift-and-refine procedure seems local; how do you ensure that distinct posterior modes (if any) are visited rather than merely thickening a single basin?\n\nWhat prevents the decoder from mapping interpolations to weights with poor training loss, especially as interpolation length grows? Have you measured the actual loss of decoded samples before using them at test time?\n\nCould you compare against SWAG or SGLD on at least one dataset to illustrate where MetricBNN sits on the accuracy–cost frontier?\n\nFigure 5 should include a legend instead of explaining what each color represents. There seems to be a difference in the trend where the proposed sampling method shows decreasing NLL for larger networks/time, while both Laplace approximation methods see an increase in NLL. Why does sampling on a low-dimensional manifold fix the overfitting?", "broader_impact_concerns": "None.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "eCDLLaAAEp", "reviewer": "Reviewer_457L", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "### Strengths\n\n* The paper is clearly written.\n* The proposed method is simple yet interesting.\n\n## Weaknesses\n\n* **Framing** – While the proposed method is similar to the Laplace approximation in that both methods can be applied post-hoc, I found the constant comparisons with the Laplace approximation to be unnecessary/distracting. Of course, it is normal to compare a new method with relevant baselines, but it is usually a good idea to (a) compare with a range of relevant methods, and (b) motivate the method on its own strengths rather than another methods weaknesses. Furthermore, many of the comparisons (in both the text and results) were made with a very specific instance of the Laplace approximations (e.g., not linearised, and/or without any common further scalability approximations such as last-layer or block-diagonal structure). Ultimately, this makes the comparisons somewhat meaningless (and misleading) since these are the methods that are used in practice). \n* **Experimental Evaluation** – Unfortunately, the experimental evaluation for this paper does not meet the par for publication in my opinion. There are several areas that I feel are lacking. (1) Comparison with relevant baselines such as deep ensembles and HMC are missing. (2) Evaluations on NNs of even modestly large sizes by todays standard. (3) Evaluations on standard uncertainty estimation benchmarks such as OOD detection and dataset shift. (4) Ablations for the different components/modeling choices in MetricBNN. (5) Qualitative experiments that help to build intuition for how MetricBNNs works (and where it might not work). I provide some concrete suggestions for improving these issues below.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposes 'MetricBNN' a novel, sampling based approach for approximating the posterior distribution of a BNN. Like the Laplace approximation for BNNs, MetricBNN is applied post-hoc to MAP estimated NNs. MetricBNN has two components. First, a series of sample solutions near the MAP solution are found by a sequence of random perturbation + gradient update steps. Then, an AutoEncoder is fit to the sample solutions in order to create a low-dimensional latent space that can easily be sampled from. On small scale regression and classification problems, MetricBNN performs favorably compared to several Laplace approximation baselines.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1. [critical] Ensure that all existing comparisons to the Laplace approximation are to variants of the Laplace approximation that are used in practice (such as those discussed in Daxberger at al. (2021). \n2. [critical] Add Deep Ensembles and HMC as baselines for all existing experimental results. Deep Ensembles are a standard baseline for Bayesian Deep Learning papers as they are simple to implement, commonly used in practice, and provide a common point of reference. HMC is a highly relevant baseline for this work since it represents an alternative sampling method that practitioners might consider using.\n3. [critical] Add the rotated MNIST, corrupted CIFAR, and OOD detection benchmarks from https://arxiv.org/pdf/1906.02530 and https://arxiv.org/pdf/2010.14689. These are standard benchmarks in the Bayesian Deep Learning community. \n4. [critical] Add results for networks at least the size for ResNet18. This is a very small network by todays standards. MetricBNN is likely to struggle to scale to networks of this size, but the authors cannot make claims that their method is scalable when only toy sized networks are used in their results. \n5. [critical] Add an ablation to show the importance of both of the components of MetricBNN. For example, what happens when you try sample directly in parameter space without the AutoEncoder component? What happens as you vary $N$ and $T$? What about the choices for $\\lambda_+$, $\\lambda_-$, $\\lambda_d$?  \n6. [strengthen] Add qualitative results showing how well (or not) the samples from MetricBNN represent the posterior of the BNN. For example, compare the MetricBNN samples to those from HMC when projected into the top 2 principle components of the parameter space. \n7. [critical] Add details about the size of the AutoEncoder network. Add ablations for the size of this network relative to the BNN. Add ablations for the size of the latent space. \n\nI understand that these are big changes, but I strongly believe that this is the standard that minimum standard that the Bayesian Deep Learning community must hold itself to. \n\nI also want to clarify that if these changes are made, I would be happy to accept this paper, regardless of the results of the experiments. I would not expect the method to be the best in all benchmarks and against all baselines. However, it is important for the paper to accurately characterize the method so that practitioners can make and informed choice about whether to use the method and so that future work in the field can be built on strong foundations. I think that the proposed method is interesting and could have several strengths, but unfortunately it isn't possible to know that with the current experimental evaluation.", "broader_impact_concerns": "None", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Decision: Accept (poster)