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

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.