📊 statistics

Inversion-Free Natural Gradient Descent on Riemannian Manifolds

This paper proposes an inversion-free stochastic natural gradient method for Riemannian manifolds that maintains an online approximation of the inverse Fisher information matrix using transport operations to handle tangent space differences, achieving almost-sure convergence rates while enforcing geometric constraints and demonstrating superior performance in variational Bayes and normalizing flows compared to Euclidean counterparts.

Original authors: Dario Draca, Takuo Matsubara, Minh-Ngoc Tran

Published 2026-04-06

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Dario Draca, Takuo Matsubara, Minh-Ngoc Tran

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to find the lowest point in a vast, foggy valley to set up a camp. This is essentially what machine learning algorithms do: they try to minimize an "error" to find the best possible model.

In the old days, most algorithms assumed the ground was flat like a sheet of paper (a Euclidean space). They would just take a step downhill, check the slope, and take another step. This works fine on flat ground.

But in many real-world problems, the "ground" isn't flat. It's curved, like the surface of the Earth, or it has strict rules, like "you must stay on the surface of a sphere" or "you cannot step off the edge of a cliff." In math, these curved, rule-bound landscapes are called Riemannian Manifolds.

The Problem: The "Flat Earth" Mistake

If you try to walk down a curved mountain using flat-Earth logic, you might take a step that looks downhill on the map but actually sends you flying off a cliff or walking in circles.

Furthermore, standard methods often get stuck because they don't understand the shape of the data. Imagine trying to navigate a maze where some paths are wide and easy, while others are narrow and slippery. A standard walker takes the same size step everywhere. A smart walker (using Natural Gradient Descent) knows to take tiny, careful steps on the slippery paths and big, confident strides on the wide paths.

However, calculating these "smart steps" usually requires a heavy, slow calculation called matrix inversion. It's like trying to solve a complex puzzle every single time you take a step. For huge problems, this is too slow and computationally expensive.

The Solution: "Inversion-Free" Walking on Curves

This paper introduces a new, clever way to walk on these curved landscapes without needing to solve that heavy puzzle every time.

Here is the analogy:

The Map (The Manifold): Think of the parameter space as a curved surface, like the skin of a balloon. You can't just draw a straight line on it; you have to follow the curves (geodesics).
The Compass (The Fisher Information): This is a special compass that tells you not just which way is "down," but which way is "down" relative to how sensitive the data is. It accounts for the curvature.
The Heavy Puzzle (Matrix Inversion): Usually, to read this compass correctly, you have to do a massive calculation to "invert" the map. This is slow.
The Trick (Inversion-Free): The authors realized you don't need to solve the whole puzzle at once. Instead, you can learn the compass as you walk.
- Imagine you are walking and taking small notes about the terrain every few steps.
- Instead of recalculating the whole map, you just update your notes with the new information you just gathered.
- This update is fast (like adding a new line to a notebook) rather than slow (like redrawing the whole map).

The "Transport" Challenge

Here is the tricky part: Because the ground is curved, the "notes" you take at one spot (Tangent Space A) are written in a different "language" than the notes at the next spot (Tangent Space B).

Analogy: Imagine you are walking around the Earth. At the North Pole, "North" is one direction. At the Equator, "North" is a different direction relative to your feet. If you want to compare your notes from the Pole with your notes from the Equator, you have to rotate them so they match.
In the paper, this is called Vector Transport. The authors developed a way to rotate these "notes" (score vectors) so they can be combined correctly, even as you move across the curved surface.

Why This Matters

The authors proved that this method:

Converges: It actually finds the bottom of the valley (the optimal solution) reliably.
Is Fast: It avoids the slow "puzzle solving" (inversion) by using a quick, incremental update.
Handles Constraints: It naturally keeps you on the right path (e.g., ensuring a matrix stays "positive" or "orthogonal") without you having to manually force it.

Real-World Examples

They tested this on two specific types of "curved landscapes":

Gaussian Distributions (The Bures-Wasserstein Manifold): Used for approximating complex data distributions (like predicting stock prices or weather). They showed their method finds better answers faster than standard methods.
Stiefel Manifold (Orthogonal Matrices): Used in neural networks where weights must be perfectly perpendicular (like a set of perfectly aligned compass needles). Their method kept the needles aligned while learning, whereas standard methods often knocked them out of alignment.

The Bottom Line

This paper gives machine learning a pair of smart, adaptive hiking boots for curved terrain. Instead of struggling with heavy calculations or falling off cliffs, the algorithm learns the shape of the world as it goes, taking the perfect step size and direction every time, all while keeping the math light and fast.

1. Problem Statement

The paper addresses the optimization of objective functions $L(\theta)$ defined over a family of parametric probability distributions $Q = \{q_\theta : \theta \in \mathcal{M}\}$ , where the parameter space $\mathcal{M}$ is a Riemannian manifold rather than a Euclidean space.

Key Challenges:

Constraints: Many statistical models require parameters to satisfy strict constraints (e.g., covariance matrices must be symmetric positive definite (SPD), rotation matrices must be orthogonal). Standard Euclidean optimization often violates these constraints or requires complex projection steps.
Identifiability: In models with redundant parameterizations (e.g., low-rank matrix factorizations), the Fisher Information Matrix (FIM) can be singular in Euclidean space. Manifold formulations naturally eliminate non-identifiable directions.
Geometry: Objective functions that appear non-convex in Euclidean space may exhibit geodesic convexity on specific manifolds (e.g., the Bures-Wasserstein manifold for Gaussian distributions), offering better convergence properties.
Computational Bottleneck: The standard Natural Gradient Descent (NGD) requires computing and inverting the FIM ( $I_F^{-1}$ ). In high-dimensional settings, this inversion is $O(d^3)$ , making it computationally prohibitive. While "inversion-free" methods exist for Euclidean spaces (using Sherman-Morrison updates), extending them to Riemannian manifolds is non-trivial because score vectors at different iterates belong to different tangent spaces.

2. Methodology

The authors propose a Stochastic Inversion-Free Natural Gradient Descent algorithm tailored for Riemannian manifolds. The core idea is to maintain an online, low-rank approximation of the inverse FIM without ever performing a full matrix inversion.

Key Algorithmic Components:

Intrinsic Formulation: The method treats the FIM as a bilinear form on the tangent space $T_\theta \mathcal{M}$ . It utilizes a baseline Riemannian metric $\langle \cdot, \cdot \rangle_\theta$ to define score vectors (gradients of the log-likelihood).
Online Approximation of Inverse FIM:
- Instead of inverting the FIM, the algorithm maintains a running estimate $H_s$ of the FIM (or its inverse directly) using a Sherman-Morrison rank-1 update.
- At each iteration $s$ , a score vector $\phi_{s+1}$ is sampled. The inverse approximation is updated as:
  $H_{s+1}^{-1} = H_{s+1/2}^{-1} - \frac{H_{s+1/2}^{-1} \phi_{s+1} \phi_{s+1}^\top H_{s+1/2}^{-1}}{1 + \langle \phi_{s+1}, H_{s+1/2}^{-1} \phi_{s+1} \rangle}$
- This update has quadratic complexity $O(d^2)$ , significantly cheaper than cubic inversion.
Vector Transport:
- Since tangent spaces $T_{\theta(s)}\mathcal{M}$ and $T_{\theta(s+1)}\mathcal{M}$ differ, the inverse FIM approximation cannot be directly added to the new tangent space.
- The algorithm employs a vector transport operation $T$ (associated with a retraction $R$ ) to move the approximation $H_s^{-1}$ from the previous tangent space to the current one:
  $H_{s+1/2}^{-1} = T_{s, s-1}^* \circ H_s^{-1} \circ T_{s, s-1}$
- This ensures the approximation remains self-adjoint and valid in the current geometry.
Limited-Memory Variant:
- To handle high dimensions where storing a $d \times d$ matrix is infeasible, the authors propose a sliding window approach.
- Instead of storing the full matrix, they store the $K$ most recent score vectors and update coefficients. This reduces storage complexity to sub-quadratic ($O(dK)$).

3. Key Contributions

Purely Intrinsic Approach: Unlike previous works (e.g., Hu et al., 2024) that rely on embedding manifolds in a higher-dimensional Euclidean space, this method is purely intrinsic. It does not require an ambient space, making it applicable to general manifolds (SPD, Stiefel, Grassmann).
Inversion-Free Stochastic Optimization: It extends the Euclidean inversion-free NGD (Godichon-Baggioni et al., 2024) to Riemannian manifolds, handling the geometric complexities of vector transport and curvature.
Theoretical Convergence Rates:
- Proves almost-sure convergence of the iterates to a stationary point.
- Establishes convergence rates for the squared distance to the minimizer: $O(\frac{\log s}{s^\alpha})$ when the step size exponent $\alpha > 2/3$ .
- Proves the asymptotic consistency of the approximate FIM, accounting for errors introduced by vector transport (holonomy).
Limited-Memory Extension: Introduces a memory-efficient variant suitable for high-dimensional problems, reducing storage from $O(d^2)$ to $O(dK)$.

4. Experimental Results

The authors validate the method on two distinct manifold settings:

Gaussian Variational Inference (Bures-Wasserstein Manifold):
- Task: Bayesian logistic regression using Gaussian variational approximations.
- Manifold: The space of Gaussian distributions equipped with the 2-Wasserstein metric ( $BW(\mathbb{R}^d)$ ).
- Findings: The Riemannian Inversion-Free NGD (BW-NGD Approx) converged faster and achieved lower Negative ELBO (NELBO) than its Euclidean counterpart (Euc-NGD Approx). Crucially, the Riemannian method remained stable with larger step sizes, whereas the Euclidean approximation diverged on larger datasets (e.g., Adult, Mushrooms). This highlights the benefit of geodesic convexity in the BW space.
Normalizing Flows (Stiefel Manifold):
- Task: Variational inference for Bayesian neural networks using Sylvester normalizing flows, where weight matrices are constrained to be orthogonal (Stiefel manifold).
- Manifold: Stiefel manifold $St(d, d)$.
- Findings: The Riemannian Natural Gradient method significantly outperformed both standard Riemannian Stochastic Gradient Descent (RSGD) and an "Extrinsic" NGD method (which projects Euclidean updates onto the manifold). The intrinsic approach achieved the lowest NELBO across multiple datasets (German Credit, Census, etc.).

5. Significance and Impact

Bridging Geometry and Efficiency: The paper successfully bridges the gap between the theoretical benefits of Riemannian optimization (constraints, convexity, identifiability) and the practical need for scalable, inversion-free algorithms.
Robustness: By leveraging the intrinsic geometry, the method avoids the numerical instability and divergence often seen when applying Euclidean approximations to constrained problems.
Scalability: The limited-memory variant makes natural gradient methods feasible for high-dimensional statistical models where storing the full Fisher matrix is impossible.
Future Directions: The authors suggest extending this framework to infinite-dimensional settings, such as the Wasserstein space of probability measures, potentially benefiting particle-based inference methods like Stein Variational Gradient Descent (SVGD).

In summary, this work provides a rigorous, efficient, and geometrically sound framework for optimizing complex probabilistic models on manifolds, overcoming the computational bottlenecks of traditional natural gradient methods while preserving their statistical advantages.