Riemannian Natural-Gradient Flow

• Riemannian natural-gradient flow is a geometry-aware steepest-descent system operating on manifolds with a Riemannian metric to optimize energy or loss functionals.
• It unifies diverse methodologies including information geometry, neural network training on matrix manifolds, and optimal transport through metric-specific gradient structures.
• Discretization schemes and convergence analyses ensure global minimization and accelerated dynamics in structured high-dimensional optimization settings.

A Riemannian natural-gradient flow is a canonical steepest-descent dynamical system for an energy or loss functional defined on a (possibly curved) manifold equipped with a Riemannian metric. In place of the usual Euclidean gradient, the flow follows the steepest descent direction dictated by the metric tensor, yielding geometry-aware optimization dynamics. This framework underlies a vast generalization of classical gradient descent, encompassing natural-gradient descent in information geometry, neural-network training on matrix manifolds, Wasserstein and Gromov–Wasserstein flows in probability spaces, and more. The mathematical unification of these domains offers globally convergent training regimes, structure-preserving flows, and deep connections to geometric mechanics and partial differential equations.

1. Definition and Core Principles

Let $M$ be a smooth finite- or infinite-dimensional manifold with a Riemannian metric $g_p$ on each tangent space $T_pM$. For a given smooth objective $f : M \rightarrow \mathbb{R}$, the Riemannian (natural) gradient $\nabla_g f$ is defined by

$g_p\big(\nabla_g f(p), v\big) = df(p)[v]$

for all $v \in T_pM$. The Riemannian natural-gradient flow is the solution of the ODE: $\dot{X}(t) = - \nabla_g f( X(t) )$ which realizes steepest descent of $f$ with respect to the geometry imposed by $g$. The squared norm of the instantaneous change is $g_p$0.

This construction contrasts with the classical gradient flow, which uses the canonical Euclidean metric. The Riemannian metric allows encoding statistical, algebraic, or symmetry structure inherent to the manifold or problem domain.

2. Examples Across Mathematical Domains

A. Deep Linear Networks as Flows on Fixed-Rank Matrix Manifolds

The end-to-end weight of a deep linear network $g_p$1 can be viewed as a point on the rank-$g_p$2 matrix manifold: $g_p$3 whose tangent space at $g_p$4 is

$g_p$5

For the squared loss $g_p$6, the Riemannian metric is defined via a self-adjoint, positive-definite operator $g_p$7, chosen so that the pushforward of the Euclidean gradient flow in parameter space coincides with a natural-gradient flow with respect to $g_p$8. The resulting ODE for $g_p$9 is: $T_pM$0 Almost all initializations yield convergence to global minimizers, exploiting the strict-saddle structure of non-minima (Bah et al., 2019).

B. Information Geometry and Natural-Gradient Flow

For statistical models $T_pM$1, the Fisher–Rao metric $T_pM$2 gives a canonical Riemannian structure on parameter space. The natural-gradient flow: $T_pM$3 tracks steepest descent in the Fisher metric, which is optimal under the geometry of the statistical model. Such flows admit connections to geodesic Hamiltonians, Jacobi–Maupertuis time reparametrizations, and replicator equations in evolutionary dynamics. Explicitly, one can connect the (dissipative) gradient trajectory to reparametrized geodesic flows via $T_pM$4 (Wada et al., 2021).

C. Wasserstein and Gromov–Wasserstein Gradient Flows

On the infinite-dimensional $T_pM$5-Wasserstein manifold $T_pM$6 of probability measures, the Riemannian metric induced by optimal transport yields: $T_pM$7 Gradient flows for energies $T_pM$8 satisfy

$T_pM$9

Parametric statistical models inherit a finite-dimensional Wasserstein metric $f : M \rightarrow \mathbb{R}$0 via pullback, enabling Wasserstein natural-gradient descent with theoretical links to Newton steps at convergence (Chen et al., 2018).

Extending this, the Gromov–Wasserstein (IGW) geometry introduces a global mobility operator $f : M \rightarrow \mathbb{R}$1 on $f : M \rightarrow \mathbb{R}$2, so the intrinsic Riemannian gradient is

$f : M \rightarrow \mathbb{R}$3

implementing collective nonlocal structure in the induced gradient flows (Zhang et al., 2024).

3. Structure of the Riemannian Metric and Gradient

The essential machinery of Riemannian natural-gradient flow relies on the following:

• Metric tensor $f : M \rightarrow \mathbb{R}$4: A smooth field of inner products; can be the Fisher–Rao, $f : M \rightarrow \mathbb{R}$5-Wasserstein, neural-tangent kernel, or problem-specific constructions.
• Gradient identification: In local coordinates, the Riemannian gradient at $f : M \rightarrow \mathbb{R}$6 is $f : M \rightarrow \mathbb{R}$7. For matrix or density manifolds, this may require solving PDEs or (pseudo-)inverting differential operators.
• Tangent spaces and projections: In function space (e.g., end-to-end maps for neural nets), the tangent space at $f : M \rightarrow \mathbb{R}$8 may take nontrivial forms (e.g., $f : M \rightarrow \mathbb{R}$9). The Riemannian metric is then determined by structure-preserving criteria, such as invariance under group actions or factorization symmetries.

Explicit forms for the metric and gradient can be highly nontrivial, as in convolutional networks (NTK metric) (Achour et al., 8 Jul 2025), Hopfield networks (diagonal, activation-based metric) (Halder et al., 2019), or spinor flows (infinite-dimensional $\nabla_g f$0 metric) (Ammann et al., 2012).

4. Numerical Discretization and Algorithms

Discretizations of Riemannian natural-gradient flows yield a variety of optimization algorithms:

• Full discretization (forward Euler): $\nabla_g f$1, recovering Amari's natural-gradient descent (Gunasekar et al., 2020).
• Partial (mixed) Euler / Mirror Descent: Direct integration in the dual chart (when the metric is Hessian), as in mirror descent.
• Proximal and JKO schemes: Gradient flows with respect to Riemannian (or Wasserstein) metrics admit time-discrete Moreau–Yosida iterations (JKO steps), highly relevant for density evolution, imaging, and stochastic networks (Celledoni et al., 2018, Halder et al., 2019).
• Accelerated flows: Recent developments include high-resolution ODEs and accelerated Riemannian gradient flows, where dynamics incorporate inertial and Hessian-driven damping terms, yielding provably faster $\nabla_g f$2 convergence rates in geodesically convex settings (Li et al., 8 Apr 2025).

Practical implementation of these discretizations often necessitates nontrivial linear solves, projections, or approximations (e.g., Kronecker-factored curvature in deep learning).

5. Convergence Properties and Theoretical Guarantees

Riemannian gradient flows generally enjoy strong theoretical properties:

• Monotonic decay of the objective: Along solutions, $\nabla_g f$3 (Celledoni et al., 2018).
• Convergence to critical points: For analytic $\nabla_g f$4 and under mild conditions (e.g., full-rank data, strict-saddle property), the flow converges to critical points (Bah et al., 2019).
• Almost sure convergence to global minimizers: If all non-minimum critical points are strict saddles, the set of initializations converging to them is measure zero. For deep linear networks this yields almost sure global convergence (Bah et al., 2019).
• Accelerated dynamics: Under geodesic convexity, accelerated flows with suitable damping exhibit $\nabla_g f$5 rates (Li et al., 8 Apr 2025).

In statistical settings (e.g., information geometry), replicator equations and mirror descent are shown to coincide with natural-gradient flows under Legendre duality, demonstrating the broad algebraic unity of these approaches (Wada et al., 2021, Gunasekar et al., 2020).

6. Connections to Information Geometry and Optimization

A central unifying perspective is the view of Riemannian natural-gradient flows as geometry-aware steepest descent on metric spaces or manifolds structured by statistical inference, group invariance, or optimal transport:

• Information geometry: The Fisher–Rao metric gives optimal local distinguishability of distributions; following the corresponding natural-gradient flow is asymptotically optimal for maximum likelihood and related objectives (Wada et al., 2021).
• Optimal transport: Wasserstein metrics transfer ground manifold geometry to statistical models; associated natural-gradient methods respect transportation costs over space (Chen et al., 2018, Li et al., 2018).
• Quantum and group settings: Riemannian flows on Lie groups (e.g., $\nabla_g f$6 for quantum circuits) exploit group symmetry in algorithmic updates (Wiersema et al., 2022).
• Neural architectures: Function-space natural gradient with respect to induced metrics (e.g., NTK) is intrinsic in certain deep learning regimes (Achour et al., 8 Jul 2025), and can diverge from parameter-based (Euclidean) updates unless structural conditions are satisfied.

The induced flows have found broad applications in deep learning, Bayesian inference, geometric PDEs, imaging, quantum optimization, and manifold-valued statistics.

7. Variants, Generalizations, and Open Directions

Variants and extensions of Riemannian natural-gradient flow include:

• General Riemannian/metric tensors beyond classical settings: Encompassing metrics from differential geometry, data geometry, ground metrics in OT, or non-Hessian structures (Gunasekar et al., 2020, Achour et al., 8 Jul 2025).
• Infinite-dimensional and measure-valued flows: Otto calculus for transport, energy evolution on spaces of probability measures, and gradient flows for free energies (Zhang et al., 2024, Halder et al., 2019).
• Manifold optimization and variational PDEs: Discrete Riemannian gradient methods preserve monotonicity and global convergence under coarse discretizations, useful for imaging and inverse problems (Celledoni et al., 2018).
• Interplay with modern optimization techniques: Acceleration (damping), preconditioning (e.g., NTK, Kronecker-factored, Wasserstein), and hybrid schemes align with convergence and scalability demands (Li et al., 8 Apr 2025).
• Geometry-specific limitations: The faithful reduction from parameter space to function space may break without structural conditions (e.g., unique factorization in convolutional nets, balancedness in fully connected nets), leading to parameter-dependent flows (Achour et al., 8 Jul 2025).

A plausible implication is that further developments may exploit problem-dependent structure in the design of Riemannian metrics and natural-gradient flows, optimizing both theoretical properties and empirical performance across domains.

---

References:

• (Bah et al., 2019): Deep linear networks as Riemannian flows on fixed-rank manifolds
• (Wada et al., 2021): Natural-gradient flows, geodesic Hamiltonians, replicator equations in information geometry
• (Gunasekar et al., 2020): Discrete-time discretizations, mirror descent, and natural gradient
• (Ammann et al., 2012): Spinorial energy gradient flow on infinite-dimensional bundles
• (Wiersema et al., 2022): Quantum circuit optimization on $\nabla_g f$7 via bi-invariant Riemannian metrics
• (Chen et al., 2018, Li et al., 2018): Wasserstein natural gradient, discrete and continuous statistics
• (Zhang et al., 2024): Riemannian flows in Gromov–Wasserstein geometry
• (Li et al., 8 Apr 2025): Accelerated natural-gradient flows and convergence rates
• (Achour et al., 8 Jul 2025): Function-space Riemannian geometry induced by deep convolutional networks
• (Halder et al., 2019): Hopfield dynamics as Riemannian and Wasserstein natural gradients
• (Celledoni et al., 2018): Discrete Riemannian gradient methods and dissipative ODEs