Riemannian diffusion models achieve dimension sharp convergence rates
Sharp Convergence and Sampling Trade-offs for Riemannian Diffusion under Nonnegative Ricci Curvature
Machine Learning
Summary
Diffusion models help generate complex data, and extending them from flat spaces to curved spaces (Riemannian manifolds) is useful but challenging. The authors found a way to better understand how many calculations a Riemannian diffusion model needs to closely match the desired data distribution. Their framework improves efficiency, matching the best-known rates in flat spaces under certain curvature conditions. This means the process can be done more precisely and possibly faster by taking multiple small steps in the curved space.
What this means in practice
- •For machine learning engineers: Design Riemannian diffusion models with guaranteed convergence rates competitive with Euclidean models under nonnegative curvature assumptions.
- •For statistical physicists: Implement sampling schemes on curved spaces with sharper complexity bounds relying on geodesic random walks for improved simulations.
A theory result. No direct application yet.
Authors
Yuhao Liu, Longbo Huang
Abstract
Diffusion models have emerged as state-of-the-art generative models, with recent extensions from Euclidean spaces to Riemannian manifolds. However, existing convergence guarantees for Riemannian diffusion models typically require $\tilde{O}(\mathrm{poly}(d,T)/ε^2)$ score evaluations, with potentially unfavorable dependence on the dimension. In this work, we develop a general framework that separates score discretization from Brownian-motion simulation and allows multiple geodesic random-walk steps per score evaluation. Under nonnegative Ricci curvature assumption and an exact Brownian-motion simulation oracle, we show that $\tilde{O}(d/ε^2)$ score evaluations suffice to achieve an $ε^2$ KL divergence from the target distribution, matching the existing convergence rate of Euclidean diffusion models. We further show that $\tilde{O}(d^4T/ε^2)$ geodesic random-walk steps suffice to approximate the required drifted Brownian motion to $ε$ total variation error. Combining these results yields a sampling scheme with $\tilde{O}(d/ε^2)$ score evaluations and $\tilde{O}(d^4T/ε^2)$ geodesic random-walk steps, motivating multiple random-walk steps between consecutive score evaluations. Our results provide a sharper characterization of the convergence and sampling complexity of Riemannian diffusion models.