Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Efficient Diffusion Models for Symmetric Manifolds

Oren Mangoubi, Neil He, Nisheeth K. Vishnoi

TL;DR

This work tackles the high computational cost of diffusion models on non-Euclidean symmetric manifolds by introducing a projection-based forward diffusion with a spatially varying covariance that bypasses the heat-kernel. By applying Itô's lemma, the authors derive a tractable training objective requiring only $O(1)$ gradient evaluations per step and enabling per-iteration arithmetic that scales near-linearly in the ambient dimension, e.g., $O(d^{1.19})$ for $ ext{SO}(n)$ and $ ext{U}(n)$, or $O(d)$ for the torus and sphere. They prove an average-case Lipschitz condition under manifold symmetries to obtain sampling guarantees with polynomial-in-$d$ accuracy and runtime, using an optimal transport-based coupling rather than Girsanov's transformation. Empirically, the method delivers substantial speedups and improved sample quality on synthetic datasets over $ ext{SO}(n)$, $ ext{U}(n)$, torus, and sphere, narrowing the gap to Euclidean diffusion models in both efficiency and performance.

Abstract

We introduce a framework for designing efficient diffusion models for $d$-dimensional symmetric-space Riemannian manifolds, including the torus, sphere, special orthogonal group and unitary group. Existing manifold diffusion models often depend on heat kernels, which lack closed-form expressions and require either $d$ gradient evaluations or exponential-in-$d$ arithmetic operations per training step. We introduce a new diffusion model for symmetric manifolds with a spatially-varying covariance, allowing us to leverage a projection of Euclidean Brownian motion to bypass heat kernel computations. Our training algorithm minimizes a novel efficient objective derived via Ito's Lemma, allowing each step to run in $O(1)$ gradient evaluations and nearly-linear-in-$d$ ($O(d^{1.19})$) arithmetic operations, reducing the gap between diffusions on symmetric manifolds and Euclidean space. Manifold symmetries ensure the diffusion satisfies an "average-case" Lipschitz condition, enabling accurate and efficient sample generation. Empirically, our model outperforms prior methods in training speed and improves sample quality on synthetic datasets on the torus, special orthogonal group, and unitary group.

Efficient Diffusion Models for Symmetric Manifolds

TL;DR

This work tackles the high computational cost of diffusion models on non-Euclidean symmetric manifolds by introducing a projection-based forward diffusion with a spatially varying covariance that bypasses the heat-kernel. By applying Itô's lemma, the authors derive a tractable training objective requiring only gradient evaluations per step and enabling per-iteration arithmetic that scales near-linearly in the ambient dimension, e.g., for and , or for the torus and sphere. They prove an average-case Lipschitz condition under manifold symmetries to obtain sampling guarantees with polynomial-in- accuracy and runtime, using an optimal transport-based coupling rather than Girsanov's transformation. Empirically, the method delivers substantial speedups and improved sample quality on synthetic datasets over , , torus, and sphere, narrowing the gap to Euclidean diffusion models in both efficiency and performance.

Abstract

We introduce a framework for designing efficient diffusion models for -dimensional symmetric-space Riemannian manifolds, including the torus, sphere, special orthogonal group and unitary group. Existing manifold diffusion models often depend on heat kernels, which lack closed-form expressions and require either gradient evaluations or exponential-in- arithmetic operations per training step. We introduce a new diffusion model for symmetric manifolds with a spatially-varying covariance, allowing us to leverage a projection of Euclidean Brownian motion to bypass heat kernel computations. Our training algorithm minimizes a novel efficient objective derived via Ito's Lemma, allowing each step to run in gradient evaluations and nearly-linear-in- () arithmetic operations, reducing the gap between diffusions on symmetric manifolds and Euclidean space. Manifold symmetries ensure the diffusion satisfies an "average-case" Lipschitz condition, enabling accurate and efficient sample generation. Empirically, our model outperforms prior methods in training speed and improves sample quality on synthetic datasets on the torus, special orthogonal group, and unitary group.

Paper Structure

This paper contains 43 sections, 9 theorems, 111 equations, 5 figures, 6 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Results
  3. Problem setup and projection framework
  4. Training algorithm and runtime analysis
  5. Training.
  6. Sampling algorithm and theoretical guarantees
  7. Sampling procedure.
  8. Symmetry and forward diffusion structure.
  9. Average-case Lipschitzness.
  10. Verifying the assumption on common manifolds.
  11. Theoretical guarantees.
  12. Comparison with prior work.
  13. Extension beyond symmetric manifolds.
  14. Derivation of training and sampling algorithm
  15. Overview of proof of sampling guarantees (Theorem \ref{['thm_sampling_Manifold_best']})
  16. ...and 28 more sections

Key Result

Theorem 2.2

Let $\varepsilon >0$, and suppose that $\varphi : \mathbb{R}^d \rightarrow \mathcal{M}$ satisfies Assumption assumption_Lipschitz for some $L_1,L_2 \leq \mathrm{poly(d)}$ and $\alpha \leq \varepsilon$, and $\psi(\mathcal{M})$ is bounded by a ball of radius $\mathrm{poly}(d)$. Suppose that $\hat{f} Moreover, Algorithm alg_sampling_manifold takes iterations. Here, each iteration requires one eval

Figures (5)

  • Figure 1: C2ST scores when training on datasets of quantum evolution operators on $\mathrm{U}(n)$ (top). Lower scores indicate better-quality generated samples (range is $[0.5,1]$). For $n \geq 9$, our model has the best C2ST scores. Generated samples are plotted for $n=15$ (bottom); axes are two matrix entries.
  • Figure 2: Points generated by different models when training on a dataset sampled from a wrapped Gaussian target distribution on the torus of different dimensions $d \in \{2,10,50,100,1000\}$.
  • Figure 3: Points generated by different models trained on a Gaussian mixture dataset on $\mathrm{SO}(n)$ for different values of $n$.
  • Figure 4: Points generated on $\mathrm{U}(n)$ for different values of $n$, when training on datasets comprising time-evolution operators of quantum harmonic oscillators with random potentials. For $n = 9$ and $n = 15$, we observe that our model generates samples resembling the data distribution, while the Euclidean, RSGM, and TDM models generate lower-quality samples.
  • Figure 5: A probability density $\mu$ with one mode (blue) on the torus. The map $\psi$, which maps points in the $d$-dimensional torus $\mathbb{T}_d$ to Euclidean space $\mathbb{R}^d$, may break up the single mode on the torus into up to $2^d$ separated modes in $\mathbb{R}^d$. This can make the task of learning the pushforward of the target distribution on $\mathbb{R}^{d}$ much more challenging than the task of learning the original target distribution on the torus, as the distribution in $\mathbb{R}^{d}$ may have exponentially-in-$d$ more modes.

Theorems & Definitions (14)

  • Theorem 2.2: Accuracy and runtime of sampling algorithm
  • Corollary 2.3
  • Lemma 3.1: Itô's Lemma
  • Lemma 6.2
  • proof
  • Lemma 6.3: Gronwall-like inequality for SDEs on a manifold of non-negative curvature
  • proof : Proof of Lemma \ref{['Gronwall']}
  • Lemma 6.4: Average-case Lipschitzness
  • proof
  • Proposition 6.5: Proposition 20 of chen2023sampling
  • ...and 4 more