Fisher-Rao Gradient Flow: Geodesic Convexity and Functional Inequalities

1. Problem Statement

The paper addresses the convergence analysis of Fisher-Rao (FR) gradient flows for probability density functions. While the convergence of Wasserstein gradient flows (e.g., Langevin dynamics) is well-understood through functional inequalities like the Log-Sobolev inequality and geodesic convexity (displacement convexity), the FR setting presents unique challenges:

• Non-locality: FR gradient flows are non-local differential equations involving global expectations (e.g., $E_\rho[\cdot]$), unlike the local PDEs of Wasserstein flows.
• Failure of Standard Tools: Existing analyses for FR flows often rely on explicit solution formulas (specifically for KL divergence) or strong assumptions on the initial distribution.
• Target Independence: A key goal is to establish convergence rates that are uniform and independent of the target distribution $\rho^*$, contrasting with Wasserstein flows where rates depend heavily on the geometry of $\rho^*$ (e.g., log-concavity or spectral gaps).

The paper investigates whether standard tools like geodesic convexity and gradient dominance (Polyak-Lojasiewicz) hold for FR flows driven by general $f$-divergences, and if not, what alternative functional inequalities can guarantee convergence.

2. Methodology

The authors employ a rigorous geometric and analytic approach:

• Geometric Framework: They define the Fisher-Rao metric on the space of probability densities $\mathcal{P}$ as $g^{FR}_\rho(\sigma_1, \sigma_2) = \int \frac{\sigma_1 \sigma_2}{\rho} d\theta$. They derive the gradient flow equation for general $f$-divergences $D_f[\rho\|\rho^*]$.
• Hessian Analysis: They compute the Riemannian Hessian of the energy functional under the FR metric to test for geodesic convexity.
• Counterexample Construction: To disprove convexity or gradient dominance, they construct specific counterexamples using Gaussian mixtures and piecewise constant densities (mollified to be smooth).
• Reduction to Discrete Simplex: For proving necessary and sufficient conditions for gradient dominance, they reduce the continuous non-local inequality to a discrete problem on a probability simplex using measure-theoretic truncation and linear programming arguments.
• Dual Inequality Construction: Recognizing the failure of standard inequalities for KL divergence, they introduce a novel Dual Gradient Dominance Condition, analyzing the evolution of a different $f$-divergence ($\bar{f}$) along the flow of the original divergence ($f$).

3. Key Contributions and Results

A. Negative Results: Failure of Standard Properties for KL Divergence

The paper proves that for the widely used KL divergence ($f(x) = x \log x$):

1. No Geodesic Convexity: The KL divergence is not geodesically convex under the Fisher-Rao metric, even in a small neighborhood of the target $\rho^*$.
2. No Gradient Dominance: The standard gradient dominance condition (Polyak-Lojasiewicz inequality) fails for KL divergence.
   Significance: This explains why previous proofs of exponential convergence for FR-KL flows relied on explicit solution formulas rather than general geometric inequalities.

B. Characterization of General $f$-Divergences

The authors derive conditions on the function $f$ that determine the behavior of the flow:

1. Geodesic Convexity: An $f$-divergence is geodesically convex if $x f'(x)$ is concave. If additionally $f''(1) > 0$, it is strongly geodesically convex with a constant depending only on $f$.
2. Gradient Dominance: They establish a necessary and sufficient condition for gradient dominance:
   $$(f'(y) - f'(x))^2 \geq \alpha \left(\frac{1}{x} - \frac{1}{y}\right) \left(\frac{f(x)}{1-x} - \frac{f(y)}{1-y}\right)$$
   for all $0 < x < 1 < y$.
   • They provide a simpler sufficient condition: $x^2 f''(x) > \alpha_s$ for $0 < x \leq 1$.
   • Result: KL divergence fails this condition, but many other $f$-divergences (e.g., those with $f''(x) = x^p$ for $p \leq -2$) satisfy it.

C. The Novel Contribution: Dual Gradient Dominance

Since KL divergence fails standard gradient dominance, the authors introduce a Dual Gradient Dominance Condition:
$$\frac{d}{dt} D_{\bar{f}}[\rho_t \| \rho^*] \leq -\alpha_f \left( D_{\bar{f}}[\rho_t \| \rho^*] + D_f[\rho_t \| \rho^*] \right)$$

• Mechanism: Instead of bounding the decay of the energy $D_f$ directly, they bound the decay of a dual energy $D_{\bar{f}}$ (e.g., Reverse $\chi^2$ or Reverse KL) along the flow of $D_f$.
• Key Findings:
  • For KL divergence ($f(x)=x \log x$), choosing the dual as the Reverse KL ($\bar{f}(x) = x \log(1/x)$) yields a dual gradient dominance with rate $\alpha = 1$.
  • For general $f$-divergences with $f''(1) > 0$, choosing the dual as the Reverse $\chi^2$ divergence yields a positive rate $\alpha_f$ depending only on $f$.
• Implication: This proves that the sum of the forward and reverse divergences decays exponentially fast.

4. Significance and Impact

• Uniform Convergence Rates: The most significant theoretical contribution is that the convergence rates derived are independent of the target distribution $\rho^*$. Unlike Wasserstein flows where the rate depends on the log-Sobolev constant of $\rho^*$ (which can be arbitrarily small for multimodal or anisotropic targets), FR flows exhibit uniform exponential convergence for a broad class of $f$-divergences.
• New Analytical Tool: The "Dual Gradient Dominance" is a new functional inequality specific to the Fisher-Rao geometry, filling a gap where standard convexity arguments fail.
• Sampling Applications: These results provide a rigorous theoretical foundation for using Fisher-Rao gradient flows (and their discretizations like Natural Gradient Descent) in Bayesian inference and sampling, particularly for challenging target distributions where Wasserstein-based methods (like Langevin dynamics) may suffer from slow mixing.
• Generality: The framework extends beyond KL divergence to general $f$-divergences, offering a unified view of convergence for various statistical distances.

5. Conclusion

The paper systematically extends the success of functional inequality techniques from Wasserstein geometry to Fisher-Rao geometry. By proving that standard convexity fails for KL divergence but introducing a powerful dual gradient dominance framework, the authors establish that Fisher-Rao gradient flows achieve uniform exponential convergence for general target distributions. This resolves a critical theoretical gap and validates the use of FR-based algorithms for robust sampling and inference.