Inversion-Free Natural Gradient Descent on Riemannian Manifolds

TL;DR

This paper introduces an inversion-free stochastic natural gradient method on Riemannian manifolds, enabling efficient optimization of probability distributions with manifold-structured parameters.

Contribution

It develops an intrinsic, inversion-free natural gradient algorithm that maintains an online approximation of the inverse Fisher information matrix on Riemannian manifolds.

Findings

Proves convergence rates of the proposed method with respect to the squared distance to the minimizer.

Establishes convergence rates for the approximate Fisher information matrix with transport-based errors.

Demonstrates the method's effectiveness on variational Bayes and normalizing flows.

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^\alpha)$ for the squared distance to the minimizer when the step size exponent $\alpha >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.