Inversion-Free Natural Gradient Descent on Riemannian Manifolds

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.

• The paper presents an inversion-free natural gradient descent framework that avoids costly Fisher Information Matrix inversion on Riemannian manifolds.
• It employs a rank-1 update with vector transport to efficiently maintain the inverse-FIM estimate, reducing per-iteration complexity from cubic to quadratic.
• Empirical experiments in variational inference and normalizing flows validate improved convergence speeds and enhanced optimization under manifold constraints.

Inversion-Free Natural Gradient Descent on Riemannian Manifolds: A Technical Summary

Introduction and Motivation

This paper develops an inversion-free stochastic natural gradient descent (NGD) framework applicable to probability distributions with parameters constrained to Riemannian manifolds. Classical NGD exploits the Fisher-Rao metric to yield parameterization-invariant optimization, but typically assumes Euclidean parameter spaces. Many applications, however, require or benefit from manifold-valued parameterizations (e.g., positive definite matrices, orthogonal matrices, simplex constraints). These structures enforce constraints or regularities, resolve identifiability, and in some cases render non-convex objectives geodesically convex.

The core innovation here is an intrinsic Riemannian NGD algorithm that eliminates repeated explicit matrix inversion of the Fisher Information Matrix (FIM) by maintaining and updating a low-rank recursion for the inverse FIM using score-vector outer products, while handling transport between tangent spaces via vector transport operations. The method is analyzed within the context of stochastic optimization (as in variational inference) and is not reliant on embedding in vector spaces, nor does it depend on Kronecker or low-rank approximations in the ambient space.

Methodology

Riemannian Fisher Information Metric

For statistical models with parameters $\theta$ on a smooth manifold $\mathcal{M}$ and distributions $\{q_\theta\}$, the Fisher information is an intrinsic, nonnegative-definite symmetric bilinear form on $T_\theta \mathcal{M}$:

$$F_\theta[u, v] = \mathbb{E}_{q_\theta}[D\ell_Y(\theta)[u] D\ell_Y(\theta)[v]], \quad u, v \in T_\theta \mathcal{M},$$

where $\ell_Y(\theta)$ denotes the log-likelihood and $D\ell_Y(\theta)$ is the differential. If $\mathcal{M}$ is equipped with a Riemannian metric, the FIM can be represented as a self-adjoint linear operator (the Fisher information operator). Critically, in non-Euclidean settings the FIM may reside in different tangent spaces as the iterate moves along the manifold.

Inversion-Free Natural Gradient Updates

Rather than computing and inverting $I_F(\theta)$ at each step—an $O(d^3)$ operation—the algorithm maintains a running estimate of the inverse FIM using a rank-1 (Sherman-Morrison) update:

$\mathcal{M}$0

where $\mathcal{M}$1 is a sampled score vector at $\mathcal{M}$2. This maintains quadratic rather than cubic per-iteration complexity. On a manifold, transporting $\mathcal{M}$3 into $\mathcal{M}$4 is necessary, accomplished via differentiated retraction-based vector transport (and, optionally, its adjoint for self-adjointness preservation).

A limited-memory variant preserves only the most recent $\mathcal{M}$5 rank-1 updates, reducing memory/storage to $\mathcal{M}$6. All operations, including transport and retraction, are intrinsic to the manifold and do not require identification with an ambient Euclidean space.

Iterative Scheme and Convergence

At each iteration, the following steps are performed:

1. Transport the inverse-FIM estimate to the new tangent space.
2. Update the inverse-FIM estimate with new score vector(s).
3. Estimate the (possibly stochastic) gradient.
4. Precondition the step with the running inverse-FIM estimate and update via a manifold retraction.

The authors establish almost sure convergence rates. For step size schedules $\mathcal{M}$7 with $\mathcal{M}$8, and under mild regularity conditions typical in stochastic approximation, they show that iterates achieve $\mathcal{M}$9 convergence in squared Riemannian distance to the minimizer. The approximation to the Fisher operator is shown to converge at rate $\{q_\theta\}$0 (for any $\{q_\theta\}$1).

Numerical Experiments

The proposed method is benchmarked against exact NGD and first-order methods on both Euclidean and Riemannian geometries using common variational inference problems and normalizing flows.

Gaussian Variational Inference on Bures-Wasserstein Manifold

For Bayesian logistic regression, the space of non-degenerate Gaussian posteriors is equipped with the Bures-Wasserstein (BW) metric, which renders KL divergence geodesically convex when the underlying model is log-concave. The authors demonstrate that using BW geometry yields improved optimization landscapes compared to the Euclidean parameterization. Figure 1 compares the negative ELBO (NELBO) across datasets for exact NGD, first-order, and inversion-free approximate NGD updates under both Euclidean and BW geometries.

Figure 1: NELBO versus iteration for the Bayesian logistic regression model. Top row: Euclidean covariance parameterization. Bottom row: Bures-Wasserstein manifold. Datasets from left to right: Ionosphere, Breast Cancer, Sonar.

The experiments indicate that (i) Fisher preconditioning consistently accelerates convergence, (ii) the inversion-free NGD closely matches the exact version after an initial burn-in phase, and (iii) Riemannian geometry (BW) improves convergence speed and final ELBO compared to the Euclidean setup, especially in higher-dimensional settings and on large datasets.

Figure 2: Comparison of approximate/exact natural gradient algorithms on larger datasets.

Normalizing Flows on the Stiefel Manifold

The method is applied to variational inference with normalizing flows where the parameters (specifically weight matrices in Sylvester flows) are constrained to have orthonormal columns, i.e., they live on the Stiefel manifold. The full Riemannian structure, including Cayley retraction and isometric vector transport, is exploited in the inversion-free NGD. This is compared to Euclidean (extrinsic) methods with projected updates and standard SGD.

Figure 3: Negative lower bound plots of the three VB methods.

On several UCI datasets, the Riemannian inversion-free NGD method converges faster and achieves strictly better ELBO minima than extrinsic counterparts or vanilla Riemannian SGD, highlighting the practical impact of correct Fisher-respectful preconditioning and intrinsic implementation.

Theoretical and Practical Implications

This work significantly extends the applicability of intrinsic NGD to non-Euclidean (not necessarily embedded) parameter spaces with rigorous stochastic approximation guarantees. The presented method:

• Handles arbitrary Riemannian manifolds, including cases with strong constraints or non-identifiability.
• Avoids problematic cubic-cost FIM inversion.
• Provably converges to stationary points under stochastic gradients, with explicit rates at nontrivial step-size regimes.
• Enables scalable NGD in high-dimensional manifold models via a memory-efficient (sliding-window) recursion.

This framework supports scalable, constraint-preserving, and geometry-aware variational inference and other statistical optimization on matrix and shape manifolds, optimal transport spaces, and beyond.

Future Directions

Potential avenues for extension include:

• Non-parametric settings: e.g., adapting the method to the infinite-dimensional Wasserstein space, with application to particle-based variational methods (e.g., SVGD).
• Further structure-exploiting transport operations for computational acceleration on specific manifolds.
• Local and global rates in strongly convex or structured non-convex regimes.
• Central limit theorems and bias/variance analysis for averaged iterates on manifolds.

Conclusion

This paper delivers an explicit, intrinsic, and inversion-free stochastic natural gradient methodology for Riemannian manifolds, with proven theoretical guarantees and superior empirical performance in challenging variational inference scenarios. By decoupling the method from both embedding and explicit inversion, and supporting scalable memory usage, it broadens the practical and theoretical reach of information-geometric optimization in manifold-constrained learning and inference.

---

References:

"Inversion-Free Natural Gradient Descent on Riemannian Manifolds" (2604.02969)