Riemannian Laplace Approximation with the Fisher Metric

This paper addresses the limitations of the Riemannian Laplace Approximation using the Fisher metric, which produces biased and overly narrow estimates, by introducing two corrected variants that ensure asymptotic exactness and demonstrate practical improvements across various experiments.

The Gist

Imagine you are trying to draw a map of a mysterious, foggy island based on a few scattered clues. Your goal is to create a "best guess" map that shows where the treasure (the most likely data points) is hidden and how spread out the rest of the island might be.

In the world of machine learning and statistics, this "map" is called a posterior distribution. It's often a complex, wiggly, high-dimensional shape that is very hard to calculate exactly.

The Old Way: The "Flat Map" (Laplace Approximation)

For decades, the standard tool for this job has been the Laplace Approximation. Think of this as a cartographer who looks at the highest peak of the island (the most likely spot) and says, "Okay, I'll just draw a perfect, round circle around this peak."

• The Good: It's incredibly fast. You just find the peak and draw a circle.
• The Bad: Real islands aren't perfect circles. They have valleys, ridges, and weird curves. If the island is actually shaped like a banana or a funnel, a simple circle is a terrible map. It misses the shape entirely.

The New Idea: The "Curved Map" (Riemannian Geometry)

Recently, researchers tried to fix this by using Riemannian Geometry. Imagine instead of drawing on a flat piece of paper, you are drawing on a flexible, stretchy rubber sheet that bends and twists to fit the shape of the island perfectly.

This method (introduced by Bergamin et al. in 2023) takes your simple circle and stretches it along the "geodesics" (the shortest paths on that curved sheet) to match the island's shape.

• The Promise: It should give a much better map.
• The Problem: The specific rubber sheet they used (called the Monge metric) was a bit defective. It was like using a rubber sheet that was too stiff in the wrong places. It ended up drawing maps that were too small (underestimating the uncertainty) and biased (pointing in the wrong direction), even when you had a lot of data. It was like trying to fit a square peg in a round hole, but the peg kept shrinking.

The Solution: The "Fisher Metric" (The Perfect Rubber Sheet)

The authors of this paper, Hanlin Yu and colleagues, realized the problem wasn't the idea of using a curved sheet, but the type of sheet they were using. They proposed a new, better rubber sheet based on something called the Fisher Information Matrix (FIM).

Here is the analogy:

• The Old Sheet (Monge): Was based on how steep the hill was right where you were standing. It reacted too aggressively to small bumps, causing the map to shrink and warp incorrectly.
• The New Sheet (Fisher): Is based on the average sensitivity of the entire island. It's a "statistically natural" way to measure distance.

Why is the Fisher Metric better?

1. It's Honest: If the island is actually a perfect circle (a Gaussian distribution), the Fisher Metric map is exactly a circle. It doesn't shrink or warp it.
2. It Handles Weird Shapes: If the island is a "banana" shape or a "funnel" shape, the Fisher Metric stretches the rubber sheet perfectly to hug those curves.
3. It's Efficient: Surprisingly, even though it's more complex mathematically, it often requires fewer computational steps to draw the map than the old, flawed method.

The "Hausdorff" Twist

There was one other small issue. When you stretch a rubber sheet, the "center" of your map might shift slightly depending on how you look at it. The authors also introduced a way to find the "true center" (called the Hausdorff MAP) that doesn't get confused by how you rotate or stretch your coordinate system. This ensures the map stays centered on the treasure, no matter how you look at the island.

The Results: A Better Map for Everyone

The team tested their new method (called RLA-F) against the old methods on several tricky problems:

• The Banana Distribution: A classic test where the data is shaped like a curved banana. The old methods drew a tiny, straight line; the new method drew the perfect banana curve.
• Neural Networks: They tested this on AI models (Neural Networks). The new method produced predictions that were almost indistinguishable from the "gold standard" (which takes hours to compute), but did it in seconds.

The Takeaway

Think of this paper as upgrading the cartographer's toolkit.

• Old Tool: A ruler and a compass (draws perfect circles, fails on curves).
• Previous Upgrade: A stretchy rubber sheet, but it was the wrong kind of rubber (it shrank and warped).
• This Paper's Upgrade: A smart, self-adjusting rubber sheet (Fisher Metric) that knows exactly how to stretch to fit the terrain, whether it's a simple hill or a complex, winding canyon.

It's a way to get the speed of a simple guess with the accuracy of a complex, detailed map, making it a powerful new tool for AI and data science.

Technical Summary

Here is a detailed technical summary of the paper "Riemannian Laplace Approximation with the Fisher Metric" by Yu et al.

1. Problem Statement

The Laplace Approximation (LA) is a computationally efficient method for Bayesian inference that approximates a target posterior distribution with a Gaussian distribution centered at the Maximum A Posteriori (MAP) estimate. While LA is asymptotically exact for infinite data (due to the Bernstein-von Mises theorem), it often provides a poor approximation for finite-data posteriors or complex, non-Gaussian targets because the family of Gaussian distributions is too restrictive.

Recent work by Bergamin et al. (2023) attempted to generalize LA using Riemannian geometry. Their approach, RLA-B, transforms samples from a Gaussian distribution using geodesic paths induced by a specific Riemannian metric (the "Monge metric," based on the outer product of gradients). However, the authors of this paper identify two critical flaws in RLA-B:

1. Bias: Even for simple Gaussian targets, RLA-B produces biased approximations that underestimate uncertainty.
2. Asymptotic Inexactness: Unlike standard LA, RLA-B does not converge to the exact posterior even in the limit of infinite data.

The paper aims to correct these shortcomings by developing a Riemannian Laplace Approximation that retains computational efficiency while achieving asymptotic exactness and reducing bias.

2. Methodology

The core idea is to define a Riemannian manifold over the parameter space $\Theta$ and map samples from a Gaussian distribution on the tangent space to the manifold using the exponential map. The quality of the approximation depends heavily on the choice of the Riemannian metric tensor $G(\theta)$.

The paper proposes and analyzes three variants:

A. RLA-B (The Baseline)

• Metric: Uses the Monge metric: $G(\theta) = I + \nabla \ell_y(\theta)\nabla \ell_y(\theta)^\top$.
• Issue: The geodesics in this metric cause the norm of the velocity vector to decrease as the gradient increases. This results in samples being "pulled" closer to the MAP than they should be, causing the approximation to be too narrow (underestimating variance) and biased.

B. RLA-BLog (Corrected Baseline)

• Approach: Keeps the Monge metric but introduces a logarithmic map correction.
• Mechanism: It samples a point on a manifold $P$ (where the metric is induced by the Gaussian approximation) and maps it to the target manifold $M$ (induced by the posterior) using $\text{Exp}_M \circ \text{Log}_P$.
• Pros/Cons: This corrects the bias for Gaussian targets, making it asymptotically exact. However, computing the logarithmic map is a boundary value problem that is computationally expensive and numerically unstable compared to the exponential map.

C. RLA-F (The Proposed Solution)

• Metric: Uses the Fisher Information Matrix (FIM) as the Riemannian metric.
  $$G(\theta) = \mathbb{E}_{Y|\theta}[-\nabla^2 \log \pi(Y|\theta)] - \nabla^2 \log \pi(\theta)$$
  This combines the expected Fisher Information (likelihood) with the negative Hessian of the log-prior.
• Theoretical Justification:
  • Invariance: The FIM transforms as a $(0,2)$ tensor under reparameterization, making the method invariant to coordinate changes.
  • Hausdorff MAP: To maintain invariance, the authors propose using the Hausdorff MAP (maximizing the density with respect to the Riemannian volume form) rather than the standard Euclidean MAP.
  • Exactness: The authors prove that for targets that are diffeomorphisms of a Gaussian (e.g., exponential family distributions), RLA-F is exact even for finite data. In the limit of infinite data, the metric becomes constant, and the geodesics become straight lines, recovering the exactness of standard LA.
• Computation: The FIM for complex models (like Neural Networks) is computed via the pullback metric using Jacobians and automatic differentiation.

3. Key Contributions

1. Identification of Bias: The paper rigorously demonstrates that the previously proposed RLA-B is biased and underestimates uncertainty, failing to be asymptotically exact even for Gaussian targets.
2. Theoretical Correction: Two solutions are proposed:
   • RLA-BLog: A fix using logarithmic maps (theoretically sound but computationally heavy).
   • RLA-F: A superior approach replacing the metric with the Fisher Information Matrix.
3. New Theoretical Results:
   • Proved that RLA-F is exact for any target distribution that is a diffeomorphism of a Gaussian (Theorem 3).
   • Proved that RLA-F is asymptotically exact for exponential family models (Theorem 2).
   • Established the invariance of the method under reparameterization when using the Hausdorff MAP.
4. Scalable Implementation: Provided a method to compute the Fisher metric and Christoffel symbols for Neural Networks using pullback metrics and automatic differentiation, enabling application to deep learning models.

4. Experimental Results

The authors evaluated RLA-F against ELA (Euclidean LA), RLA-B, and RLA-BLog on several tasks:

• Banana Distribution (2D): A classic non-Gaussian target. RLA-B was biased and narrow. RLA-F (especially with Hausdorff MAP) provided the most accurate approximation, closely matching the ground truth (NUTS samples).
• Bayesian Logistic Regression: Tested on 5 datasets (Ripley, Pima, Heart, etc.) with both standardized and raw inputs.
  • RLA-F consistently outperformed all other methods in Wasserstein distance.
  • RLA-B and RLA-BLog required significantly more function evaluations (up to 1000x more for raw inputs) to integrate the ODEs due to the complex curvature of the Monge metric.
  • RLA-F was faster and more stable.
• Neural Network Regression:
  • RLA-F matched the predictive performance of NUTS (the gold standard) with significantly fewer function evaluations.
  • RLA-B and RLA-BLog often overestimated predictive variance or produced unstable samples.
  • Scalability: RLA-F remained competitive in speed even for networks with $>1000$ parameters, whereas RLA-B became computationally prohibitive.

5. Significance

This paper makes a significant contribution to approximate Bayesian inference by bridging the gap between computational efficiency and accuracy in Laplace approximations.

• Theoretical Rigor: It corrects a fundamental flaw in the recent literature regarding Riemannian Laplace approximations, establishing the conditions under which these methods are exact.
• Practical Utility: By advocating for the Fisher Metric, the authors provide a method that is not only theoretically sound (asymptotically exact and invariant) but also practically superior. It avoids the numerical instability and high computational cost of the logarithmic map correction while outperforming the original gradient-based metric.
• Deep Learning Applicability: The ability to compute the Fisher metric for Neural Networks allows for scalable, uncertainty-aware Bayesian deep learning without resorting to expensive MCMC sampling or training separate variational inference networks.

In conclusion, the paper argues that RLA-F is the preferred variant for small-to-medium problems due to its accuracy and stability, and a viable candidate for large-scale problems where scalable Fisher approximations can be employed.