Efficient Diffusion Models for Symmetric Manifolds

TL;DR

This paper presents a new efficient diffusion model for symmetric-space Riemannian manifolds that reduces computational complexity and improves sample quality, leveraging a novel covariance approach and projection of Euclidean Brownian motion.

Contribution

Introduces a diffusion model with spatially-varying covariance for symmetric manifolds, avoiding heat kernel computations and enabling faster, more efficient training and sampling.

Findings

Outperforms prior methods in training speed

Achieves higher sample quality on synthetic datasets

Reduces computational complexity to nearly-linear in dimension

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.