Information-Geometric Optimal Control for Diffusion Models: Unified Framework via Fisher-Rao Geodesics

Kaustubh S. Bukkapatnam, Laksh Patel

Keywords: Diffusion Models, Stochastic Optimal Control, Information Geometry

TL;DR: We train diffusion models via information-geometric optimal control in Fisher–Rao space, yielding provably faster convergence and significantly better few-step sampling without changing architectures.

Abstract: We introduce a unified framework connecting stochastic optimal control, information geometry, and manifold learning for diffusion models through the Fisher-Rao metric on probability space. By formulating diffusion training as a control problem respecting information-geometric structure, we derive geometry-aware paths with provable improvements. Our contributions: (1) a rigorous optimal control formulation establishing Hamilton-Jacobi-Bellman equations in infinite dimensions with existence guarantees; (2) dimension-independent convergence rates $\kappa = \Omega(m)$ versus standard $O(m/d)$, with explicit Wasserstein bounds; (3) a practical algorithm requiring only $O(d)$ overhead per iteration. Information-geometric control yields statistical geodesics that reduce discretization error, enabling better few-step sampling. Experiments on CIFAR-10 show consistent improvements, with FID gains most pronounced at low function evaluation counts. Our framework unifies flow matching, Schrödinger bridges, and standard diffusion as special cases under different metric choices.

Submission Number: 122