Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Polynomial Convergence of Riemannian Diffusion Models

Xingyu Xu, Ziyi Zhang, Yorie Nakahira, Guannan Qu, Yuejie Chi

TL;DR

This work develops a discrete-time theory for diffusion models constrained to compact Riemannian manifolds, proving that polynomially small steps suffice to achieve small total-variation sampling error under $L_2$-accurate score estimates, without requiring smooth or positive data densities. The key contributions are a TV-type convergence bound and an iteration complexity that scales polynomially in the manifold dimension and curvature, circumventing the exponential-in-dimension guarantees of prior manifold diffusion results. The analysis rests on Li–Yau gradient bounds for the manifold heat kernel, a drift-localization scheme across tangent spaces, and Minakshisundaram–Pleijel parametrix methods to compare manifold and Euclidean heat flows, under mild curvature assumptions. Numerical experiments on spheres, tori, and warped Gaussian mixtures corroborate the theoretical results and illustrate practical stability. Overall, the paper sharpens the theoretical understanding of diffusion models on non-Euclidean spaces and opens avenues for sharper manifold-specific diffusion analyses.

Abstract

Diffusion models have demonstrated remarkable empirical success in the recent years and are considered one of the state-of-the-art generative models in modern AI. These models consist of a forward process, which gradually diffuses the data distribution to a noise distribution spanning the whole space, and a backward process, which inverts this transformation to recover the data distribution from noise. Most of the existing literature assumes that the underlying space is Euclidean. However, in many practical applications, the data are constrained to lie on a submanifold of Euclidean space. Addressing this setting, De Bortoli et al. (2022) introduced Riemannian diffusion models and proved that using an exponentially small step size yields a small sampling error in the Wasserstein distance, provided the data distribution is smooth and strictly positive, and the score estimate is $L_\infty$-accurate. In this paper, we greatly strengthen this theory by establishing that, under $L_2$-accurate score estimate, a {\em polynomially small stepsize} suffices to guarantee small sampling error in the total variation distance, without requiring smoothness or positivity of the data distribution. Our analysis only requires mild and standard curvature assumptions on the underlying manifold. The main ingredients in our analysis are Li-Yau estimate for the log-gradient of heat kernel, and Minakshisundaram-Pleijel parametrix expansion of the perturbed heat equation. Our approach opens the door to a sharper analysis of diffusion models on non-Euclidean spaces.

Polynomial Convergence of Riemannian Diffusion Models

TL;DR

This work develops a discrete-time theory for diffusion models constrained to compact Riemannian manifolds, proving that polynomially small steps suffice to achieve small total-variation sampling error under -accurate score estimates, without requiring smooth or positive data densities. The key contributions are a TV-type convergence bound and an iteration complexity that scales polynomially in the manifold dimension and curvature, circumventing the exponential-in-dimension guarantees of prior manifold diffusion results. The analysis rests on Li–Yau gradient bounds for the manifold heat kernel, a drift-localization scheme across tangent spaces, and Minakshisundaram–Pleijel parametrix methods to compare manifold and Euclidean heat flows, under mild curvature assumptions. Numerical experiments on spheres, tori, and warped Gaussian mixtures corroborate the theoretical results and illustrate practical stability. Overall, the paper sharpens the theoretical understanding of diffusion models on non-Euclidean spaces and opens avenues for sharper manifold-specific diffusion analyses.

Abstract

Diffusion models have demonstrated remarkable empirical success in the recent years and are considered one of the state-of-the-art generative models in modern AI. These models consist of a forward process, which gradually diffuses the data distribution to a noise distribution spanning the whole space, and a backward process, which inverts this transformation to recover the data distribution from noise. Most of the existing literature assumes that the underlying space is Euclidean. However, in many practical applications, the data are constrained to lie on a submanifold of Euclidean space. Addressing this setting, De Bortoli et al. (2022) introduced Riemannian diffusion models and proved that using an exponentially small step size yields a small sampling error in the Wasserstein distance, provided the data distribution is smooth and strictly positive, and the score estimate is -accurate. In this paper, we greatly strengthen this theory by establishing that, under -accurate score estimate, a {\em polynomially small stepsize} suffices to guarantee small sampling error in the total variation distance, without requiring smoothness or positivity of the data distribution. Our analysis only requires mild and standard curvature assumptions on the underlying manifold. The main ingredients in our analysis are Li-Yau estimate for the log-gradient of heat kernel, and Minakshisundaram-Pleijel parametrix expansion of the perturbed heat equation. Our approach opens the door to a sharper analysis of diffusion models on non-Euclidean spaces.
Paper Structure (51 sections, 25 theorems, 211 equations, 2 figures, 1 table, 1 algorithm)

This paper contains 51 sections, 25 theorems, 211 equations, 2 figures, 1 table, 1 algorithm.

Table of Contents

  1. Introduction
  2. Our contribution
  3. Techniques.
  4. Related works
  5. Non-asymptotic convergence for Euclidean diffusion models.
  6. Sampling on Riemannian manifold.
  7. Backgrounds
  8. Diffusion models on Euclidean space
  9. Forward SDE and Fokker–Planck.
  10. Score and reverse process.
  11. Discretization with approximate score.
  12. Geometry and notation
  13. Heat flow, Brownian motion, and diffusion on $\mathcal{M}$
  14. Main Result
  15. Iteration complexity.
  16. ...and 36 more sections

Key Result

Theorem 1

Assume Assumptions assumption:regularity and assumption:score hold. There exists some universal constant $C, C' >0$ such that the following holds. If $T \ge \frac{C}{\lambda_1} (d\log (Kd) + K + \log(\frac{N}{\varepsilon}))$, then the output $Y_0$ of Algorithm algorithm obeys where $h$ is the discretization step size, $\lambda_1 > 0$ is the mixing rate of the geometric Brownian Motion on $\mathca

Figures (2)

  • Figure 1: Reset probabilities on spheres and tori. In \ref{['fig:exit2']}, we examine the relationship between $h^{-1/2}$ and the log of the reset probability of \ref{['algorithm']} on both sphere $\mathbb S^2$ and torus $\mathbb T^2$ under the reset rules of \ref{['algorithm']}. In both cases, we see that the reset probability decays exponentially, confirming the conclusion in \ref{['eqn: sum BM exit prob']}. In \ref{['fig:td']}, we examine the same statistics on high‑dimensional tori, and we find increasing $d$ only shifts the curves to the right but leaves the exponential decay rate in $h^{-1/2}$ unchanged.
  • Figure 2: TV distance on $\mathbb T^d$ with a warped Gaussian‑mixture target. The total variation is estimated with a kernel density estimator.

Theorems & Definitions (43)

  • Theorem 1
  • Lemma 1
  • Lemma 2
  • Lemma 3: Pinsker's inequality, polyanskiy2025information
  • Lemma 4
  • Lemma 5: Metric distortion in normal coordinates
  • proof
  • Lemma 6
  • proof
  • Lemma 7: Geodesics are almost straight in small balls
  • ...and 33 more