Inversion-Free Natural Gradient Descent on Riemannian Manifolds
| Authors | Dario Draca et al. |
| Year | 2026 |
| Field | Statistics / ML |
| arXiv | 2604.02969 |
| Download | |
| Categories | stat.ML, cs.LG, stat.CO, stat.ME |
Abstract
The natural gradient method is widely used in statistical optimization, but its standard formulation assumes a Euclidean parameter space. This paper proposes an inversion-free stochastic natural gradient method for probability distributions whose parameters lie on a Riemannian manifold. The manifold setting offers several advantages: one can implicitly enforce parameter constraints such as positive definiteness and orthogonality, ensure parameters are identifiable, or guarantee regularity properties of the objective like geodesic convexity. Building on an intrinsic formulation of the Fisher information matrix (FIM) on a manifold, our method maintains an online approximation of the inverse FIM, which is efficiently updated at quadratic cost using score vectors sampled at successive iterates. In the Riemannian setting, these score vectors belong to different tangent spaces and must be combined using transport operations. We prove almost-sure convergence rates of O(\log{s}/s^α) for the squared distance to the minimizer when the step size exponent α>2/3. We also establish almost-sure rates for the approximate FIM, which now accumulates transport-based errors. A limited-memory variant of the algorithm with sub-quadratic storage complexity is proposed. Finally, we demonstrate the effectiveness of our method relative to its Euclidean counterparts on variational Bayes with Gaussian approximations and normalizing flows.
Engineering Breakdown
Plain English
This paper solves a key limitation of natural gradient descent: it only works efficiently in Euclidean spaces, but many statistical models live on constrained surfaces called Riemannian manifolds (like positive-definite matrices or orthogonal groups). The authors propose an inversion-free stochastic natural gradient method that works directly on these manifolds by building an online approximation of the Fisher information matrix inverse, updated at quadratic cost using score vectors from each iteration. This approach enforces constraints automatically (positive definiteness, orthogonality) without penalty terms and guarantees nice properties like geodesic convexity, making optimization both mathematically cleaner and computationally more efficient than Euclidean methods with constraint handling.
Core Technical Contribution
The core novelty is an intrinsic formulation of the Fisher information matrix that operates directly on Riemannian manifolds rather than embedding the problem in Euclidean space. Instead of explicitly inverting the FIM at each step (which is expensive and numerically unstable), the method maintains a running low-rank approximation of the inverse FIM updated via rank-one corrections using score vectors—similar to quasi-Newton methods but adapted to the manifold geometry. This sidesteps the computational bottleneck of matrix inversion while respecting the manifold structure, ensuring that iterates stay feasible and exploit the geometry to accelerate convergence. The theoretical guarantee that parameters stay on the manifold while being identifiable and satisfying regularity conditions is a significant advantage over Euclidean approaches that require post-hoc projection or regularization.
How It Works
The method starts with a probability model whose parameters live on a Riemannian manifold (e.g., covariance matrices, rotation matrices). At each iteration, the algorithm computes a score vector—the gradient of the log-likelihood with respect to the manifold coordinates. The inverse Fisher information matrix approximation is updated using a rank-one or low-rank correction formula (similar to BFGS or SR1 in quasi-Newton optimization) that costs O(d²) per step rather than O(d³) for full matrix inversion. The natural gradient step is then computed by multiplying this approximate inverse by the score vector, moving along a geodesic (shortest path on the manifold) rather than a straight line. The key insight is that by maintaining the approximation online and updating it incrementally, the method avoids both the numerical instability of explicit inversion and the geometric mismatch of projecting Euclidean steps back onto the manifold. The stochastic variant uses batches of score vectors to make the approximation robust to sampling noise.
Production Impact
For engineers building probabilistic models with constrained parameters—Bayesian neural networks, covariance estimation, Gaussian processes on manifolds, or equivariant models—this method eliminates the engineering burden of manually enforcing constraints or tuning penalty parameters. Instead of training a model in Euclidean space and then projecting results (covariance matrices to positive-definite, rotation matrices to orthogonal), you train directly on the constraint surface, which is both faster and gives better final models because the optimizer never wastes steps on infeasible regions. The quadratic-cost update rule makes it practical at scale: for a d-dimensional parameter space, each update costs O(d²) instead of O(d³), which is material for problems with 100–1000 parameters. In a production pipeline, this reduces both per-iteration compute time and wall-clock convergence time, and it eliminates post-processing steps to enforce constraints, lowering operational complexity. The main trade-off is implementation complexity: you need to implement manifold operations (exponential map, parallel transport, geodesics) for your specific parameter space, which is non-trivial for custom constraints but straightforward for standard manifolds (positive-definite matrices, orthogonal groups).
Limitations and When Not to Use This
The method assumes you can compute score vectors (gradients of log-likelihood) for your model, which is not always cheap or differentiable; for black-box likelihoods or discrete latent variables, this breaks down. The approach also requires choosing and maintaining a low-rank approximation structure for the inverse FIM—the paper doesn't fully characterize how to set rank or detect when the approximation is stale, which could matter in highly non-stationary settings or early in training when the curvature changes rapidly. The manifold geometry must be well-understood: the method needs the exponential map, geodesic distance, and parallel transport operators for the parameter space, so it doesn't easily extend to implicitly-defined or high-dimensional manifolds without explicit structure. Finally, there is no analysis of how the method scales when the manifold intrinsic dimension is much smaller than the ambient dimension (e.g., learning on low-rank matrix manifolds), and the stochastic convergence rate compared to Euclidean natural gradient in practical settings is not empirically validated in the abstract.
Research Context
This work extends natural gradient descent, a classical technique in information geometry pioneered by Amari, into the manifold-constrained setting. It builds on quasi-Newton methods (BFGS, SR1) which use low-rank approximations to avoid explicit Hessian inversion, and applies that strategy to the Fisher information matrix on Riemannian manifolds. The paper fits into a recent wave of research on Riemannian optimization (following work by Udrişte, Absil, Boumal and others on manifold gradient descent) but with a focus on the stochastic, second-order case. The inversion-free approximation idea is particularly relevant given the renewed interest in natural gradient methods for neural networks and variational inference, where the Fisher matrix is too large to invert directly; this paper shows how to do it efficiently on structured parameter spaces.
:::tip Subscribe Get weekly breakdowns of papers like this in AI Letters - the newsletter for engineers building production AI systems. :::