Efficient Diffusion Models for Symmetric Manifolds
Authors: Oren Mangoubi
Neil He
Nisheeth K. Vishnoi
Venue: ICML 2025 poster (2025)
Keywords: Diffusion Models
Symmetric Manifolds
Random Matrix Theory
Score-Based Generative Models
Efficient Sampling
Primary area: theory->probabilistic_methods

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
Oren Mangoubi'
Neil He? Nisheeth K. Vishnoi
Abstract
manifold M, such as robotics (Feiten al., 2013), drug discovery (Cheng et al., 2021), and quantum physics (Cran- mer et al., 2023), where configurations are often represented on symmetric-space manifolds like the torus, sphere, spe- cial orthogonal group SO(n), or unitary group U(n) where d RJ n²" A common approach enforces manifold constraints by mapping samples from Euclidean space IR M, but this often degrades sample quality due to distortions intro- duced by the mapping (see Appendix D for details).
We introduce 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 evalua- tions or exponential -n--ti arithmetic operations per training step. We introduce new diffusion model for symmetric manifolds with spatially-varying covariance, allowing us to leverage a projection of Euclidean Brownian motion to bypass heat ker- nel computations. Our training algorithm mini- mizes a novel efficient objective derived via Ito's Lemma, allowing each step run in (1) gradi- ent evaluations and nearly-linear-ine in (O (9) arithmetic operations, reducing the gap between diffusions on symmetric manifolds and Euclidean space. Manifold symmetries ensure the diffusion satisfies an "average-case enabling accurate and efficient sample generation. Empirically, our model outperforms prior meth- ods in training speed and improves sample quality on synthetic datasets on the torus, special orthog- onal group, and unitary group.
To address this, several works have developed diffusion models constrained to non- -unlliddd Riemannian manifolds (De Bortoli et al., 2022; Huang et al., 2022; Lou et al., 2024; Zhu et al., 2025; Yim et al., 2023). However, a significant gap remains between the runtime and sampling guarantees of Euclidean and manifold-basee diffusion models. For in- stance, while Euclidean models have a per-iteration runtime of (dd) arithmetic operations and O(1) evaluations of the model's gradient, objectives of manifold diffusion models often require exponential-in-d arithmetic operations, or eval- uating Riemannian divergence operators which require (d) gradient evaluations. Reducing this gap, particularly for symmetric manifolds, remains an open challenge.
Lipschitz condition,
To understand the technical difficulty, first consider the Eu- clidean case. A diffusion model consists of two components: a forward process that adds noise over time T >> 0 until the data is nearly Gaussian, and a reverse process that starts from a Gaussian sample and gradually removes the noise to generate samples approximating the original distribution T. A discrete-time Gaussian latent variable model is used approximate the reverse diffusion. In the manifold case, the forward process corresponds to standard Brownian motion on the manifold, and the reverse diffusion is its time- riverraversal. However, Gaussians are not generally defined on manifolds. To address this, previous works move to continuous time, where infinitesimal updates converge to a Gaussian on the tangent space. The reverse diffusion is then governed by a stochastic differential equation (SDE) involving the mani- fold's 's heat kernel. The heat kernel P= (-16) represents the density of Brownian motion at time т, initialized at point b. Training the reverse diffusion model thus requires mini- mizing an objective function dependent on the heat kernel.
1. Introduction
Recently, denoising diffusion- based methods have achieved significant success in generating synthetic data, including highly realistic images and videos (OpenAI, 2023). Given a dataset D sampled from an unknown probability distribution ,, a diffusion generative model aims to learn a distribution that approximates π and generates new samples from v. While most diffuuion models operate in Euclidean space RRd (Ho et al., 2020; Rombach et al., 2022), several applications require data constrained to d-dimensional non-Fuclidean
Worcester Polytechnic Institute, USA. 'Yale University, USA. Correspondence to: Oren Mangoubi <omangoubi @ gmail.com>. Nishecth K. Vishnoi <nisheeth. vishnoi @ gmail.comm>..
Proceedings of the 42nd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).
Even in the Euclidean case, the training objective is non- convex, and there are no polynomial-in -in-ommiilii-dimmsiioo runtime

---

Efficient Diffusion Models for Symmetric Manifolds
guarantees for the overall training process. However, the closed-form expression of the Fuclidean heat kernel allows each training iteration to run in O(d) arithmetic operations with (1) gradient evaluations. For non -on-Euclidean mani- folds, the lack of a closed-form heat kernel is a major bottle- neck. On symmetric manifolds like orthogonal and unitary groups, it can only be computed via inefficient series ex- pansions requiring exponential-ing runtimes. Alternatively, training with an implicit score matching (ISM) objective re- quires evaluating a Riemannian divergence, incurring (d) gradient evaluations per iteration. Due to these challenges, approximations are often used, degrading sample quality. Moreover, on manifolds with nonzero curvature, such as orthogonal and unitary groups, standard Brownian motion cannot be obtained via any projection from RR As a result, prior works rely on numerical SDE or ODE solvers to sam- ple the forward diffusion at each evaluation of the training objective, introducing significant computational overhead.
condition holds for the manifolds of interest (Lemma B.4).
Our paper introduces several new ideas. For our training result: (i) We define a novel diffusion on M. Unlike pre- vious works, our diffusion incorporates a spatially varying covariance term to account for the manifold's nonzero cur- vature. As a result, our forward diffusion can be computed as a projection yp of Brownian motion in R onto M, which can be efficiently computed via singular value decompo- sition when M is (O U(n). This enables efficient sampling from our forward diffusion in a simulation-free manner-withowe SDE or ODE solvers-by directly sam- pling from a Gaussian in R2 and projecting onto M. (ii) We introduce a new training objective that bypasses the need to compute the manifold's heat kernel. By applying Itô's Lemma from stochastic calculus, we project the SDE for reverse diffusion in Euclidean space onto M. The drift term of the resulting SDE is an expectation of the Fuclidean heat kernel. Since the Euclidean kernel has a closed-form expression and the projection yp can be computed efficiently, we evaluate the objective in time Odd)), ). (ii) While our covariance term is a x d matrix, we show that its structure,
In addition to denoising diffusions, several other generative models on manifolds leverage probability flows, including Moser flows (Rozen et al., 2021) and Riemannian normaliz- ing flows (Mathieu & Nickel, 2020; Ben-Hamu et al., 2022). More recent approaches include flow matching (Chen al., 2024) and mixture models of Riemannian bridge processes (Jo & Hwang, 2024). These models often achieve sample quality comparable to denoising diffusion models on mani- folds but frequently face similar computational bottlenecks.
arising from manifold symmetries, allows it to be computed in time (d subblinear in its dl entries.
For the sampling result, we show that the reverse SDE on the manifold M is deterministically Lipschitz, provided the pro- jection map satisfies our average-case Lipschitz condition (Lemma B.4). Since the projection introduces a spatially varying covariance in the SDE on M. prior techniques based on Girsanov's theorem cannot be used to bound accuracy. To address this, we develop an optimal transport-basee ap- proach, leading to novel probabilistic coupling argument that establishes the desired accuracy and runtime bounds. This approach differs fundamentally from previous proofs in Euclidean space (Chen et al., 2023b;a; Cheng et al., 2022: Benton et al., 2024) and manifold-basse diffusion models (De Bortoli et al., 2022), which rely on Girsanov's theorem.
Our contributions. We study the problem of designing efficient diffusion models when M is a symmetric- space manifold, such as the torus Ta, sphere Sa, special orthogo- nal group SO(n), and unitary group U(n), where d x n², as well as direct products of these manifolds, such as the special Euclidean group (F)) = RD x (O()). We present a new training algorithm (Algorithm 1) for these manifolds, achieving per-iteration runtimes of (d) arithmetic oper- (99 for SO(n) 2.37 is the matrix multiplication
ations for Ta and Sy. and O(d")
and U(n), where exponent. Each iteration requires only O(1) gradient evalu- ations of model for the drift and covariance terms of the reverse process. This significantly improves on previous methods (see Table 1). For SO(n) and U(n), our approach reduces gradient evaluations by a factor of d and achieves an exponential-in -nn-nttiil-indiiiiiin improvement in arithmetic operations, bringing runtime closer the Euclidean case. We also pro- vide sampling algorithm (Algorithm 2) with guarantees accuracy and runtime. Given an e-minimizer of our training objective, the algorithm attains an é x poly(d) bound on total variation distance accuracy in poly(d) runtime (The- orem 2.2), improving on the sampling accuracy bounds of (De Bortoli et al., 2022), which are not polynomial in d. Theorem 2.2 holds for general manifolds satisfying an average-case Lipschitz condition (Assumption 2.1). Us- ing techniques from random matrix theory, we prove this
Empirically, our model trains significantly faster per iter- ation than previous manifold diffusion models on SO(n) and U(z), staying within factor of 3 of Euclidean diffu- sion models even in high dimensions (d > 1000) (Table 3). Moreover, our model improves the quality of generated samples compared to previous diffusion models, achiev- ing improved C2ST and likelihood scores and visual qual- ity when trained on various synthetic datasets on wrapped Gaussian (mixture) models and quantum evolution opera- tors constrained to the torus, SO(n), and U(n) (Table 2 and Figure 1). The magnitude of the improvements in runtime and sample quality increases with dimension.
Thus, our results reduce the gap in training runtime and sam- ple quality between diffusion models on symmetric mani- folds and Euclidean space, contributing towards the goal of developing efficient diffusion models on constrained spaces.
2

---

Efficient Diffusion Models for Symmetric Manifolds
2. Results
2.2. Training algorithm and runtime analysis
We begin by describing the geometric setup, projection framework, and key assumptions used in our training and sampling algorithms. Notation is summarized in Ap- pendix G, and relevant background on Riemannian geometry and manifold diffusions is provided in Appendix H.
We now describe our training procedure and its computa- tional benefits for symmetric manifolds.
Training, We give an algorithm (Algorithm 1) that mini- mizes nonconvex objective function via stochastic gradient descent. This algorithm outputs trained models I f(x () and g(z,t) for the drift and covariance terms of our reverse dif- fusion, and passes these trained models as inputs to our sample generation algorithm (Algorithm 2). We show that the time per iteration of Algorithm is dominated by the computation of the objective function gradient (Lines 12 and 14 in Algorithm 1), which requires calculating the gradient of the projection map Vip as well as the model gradients Vof and Veg. where and e are the model parameters of I and g. When M is one of the aforementioned symmetric manifolds. Vi can be computed at each iteration within error ō in On log( ()) O(d'/2 log( ()) arithmetic oper- ations in the case of the special orthogonal group SO(n) or
2.1. Problem setup and projection framework
For a manifold M, we are given a projection map yp PM Id + M from a Euclidean space Rd of di- mension d' = (()im(M)), and a restricted-inverse map i = MMM= : M - Rd such that ((1((z)) = x for all x E M. We sometimes abuse notation and refer to the manifold's dimension as d rather than "O(d)", as this does not change our runtime and accuracy guarantees beyond a small constant factor. Denote by T.M the tangent space of . at x. For our sampling algorithm (Algorithm 2), we assume access to the exponential map exp(2,, v) on M for any x E M and e E T.M. In the setting where M is a symmetric space, there are closed-form expressions which allow one to efficiently and accurately compute the expo- nential map. For instance, on (O or U(z). the geodesic is given by the matrix exponential and can be computed in Onn") = O(d = !nn arithmetic operations. We are also given a dataset D S M sampled from with sup- port on M. These projection maps are efficient to compute and will be used throughout our framework for both training and sampling on M.
unitary group U(n), using the singular value decomposition of an n x n matrix, or in O(d (og())) operations for the sphere or torus. See Section 3.1 and Appendix E for details.
This significantly improves the per-iteration runtime of train- ing diffusion models on symmetric manifolds (Table 1). For instance, it achieves exponential-in-4 savings in arithmetic operations compared to Riemannian Score-based Genera- tive Models (RSGM) (De Bortoli et al., 2022), as RSGM requires summing 2(24) terms in truncated heat kernel ex- pansions for manifolds like the torus, sphere, orthogonal, or unitary groups, while our approach avoids this complexity. If RSGM is instead trained with an implicit score matching objective (ISM), which includes Riemannian divergence term that requires (dd) gradient evaluations, our model achieves a factor of d improvement in the number of gradi- eni evaluations. Similarly, we get factor of d improvement in gradient evaluations over Trivialized Momentum Diffu- sion Models (TDM) (Zhu et al., 2025), which also rely on ISM objectives, on manifolds like the orthogonal or unitary group. It also improves upon Scaling Riemannian Diffusion (SCRD) (Lou et al., 2024), where heat kernel computations for orthogonal or unitary groups involve expansions with 22(2)) terms (SCRD does not provide an ISM objective).
We set و Rdd + Rd and vb : Rd - Rd as identity maps when M = TR. For the torus Ta, [p( = =[i] mod 27r maps points to their angles, and e is its inverse on [0, 2 For the sphere Sd, ((r) = x
|||| and ab embeds the unit sphere into Rd. For the unitary group U(n) (and special orthogonal group SO())), we first define a map & which takes each upper triangular matrix X e Caxn (or X e RTxx), computes the spectral decomposition A of X + X", and outputs (xX)= = U. The spectral decomposition is unique only up to multiplication of each eigenvector Uj by a root of unity gd where the phases (Ø), Φn) lie on the n n-dimensional torus T, (or, in the real case, a subset of the torus). Thus, we define the projection map Φ Caxx x R" >> M to be the concatenated map yp- = (6, sP. ) where PP,, is the map defined above for the torus. The restricted- inverse map a takes each matrix U E M, computes U UU where A = 1 diag(nn,
Additionally, while RSGM and Riemannian Diffusion Mod- els (RDM) (Huang et al., 2022) use deterministic heat kernel approximations, these are asymptotically biased with fixed error bounds. Stochastic approximations the implicit objective introduce dimensionn-eependent noise (Lou et al., 2024). Our method improves the accuracy dependence from polynomial to logarithmic in J. Unlike solvers for SDEs or ODEs, which require polynomial-in- & iterations, our for- ward diffusion adds a Gaussian vector and projects onto the manifold, achieving high accuracy with only logarithmic cost in J.
1,..., 1), scales the diagonal by 2. and outputs the upper triangular entries of the result. For all of the above maps, ((M) is contained in a ball of radius (d)y Our general results hold under this assumption on vb. For manifolds M = Mi x M2, which are direct products of manifolds Mi and M2, where one is given maps (0,, 1/1 for Mi and yp2, 1/2 for M2, one can use the
concatenated maps yp = (421, (22 and ve = (v), 122).
3

---

Efficient Diffusion Models for Symmetric Manifolds
Table 1: Arithmetic operations plus model gradient evaluations to compute objective function's 's gradient within any error at each iteration of training algorithm, on the unitary group U(n), special orthogonal group (O()), sphere, or torus, of dimension d = n² (number of grad. eval. depends on algorithm but not on manifold).
Average-case Lipschitzness.
Assumption 2.1 (Average-case Lipschitzness). Vt E [0, T
there exists 22 C Rd4, whose indicator function 1, (z) de-
pends only on A = A(r), for which P(Zz E S4 A t e
[0, 2 1-a. For every æ E Qr we have [|Vp(z)]l222 S L1 ar V=(z)lla-2 S L1, [V 1 Lg. and
Grad. eval.
Arithmetic Operations
Algorithm
S0(n) or U(a)
Sphere
Torus
I1 V I L2. Moreover, d 2 - ||1|
RSGM (heut ker.)
1
".+yyoydd..
same
same
RSGM (ISM)
d
polyid, 3)
same
sume
Here |1112-22 denotes the operator norm and at. x denotes the derivative of the parameterization of x = (U( with respect to U E .M. This assumption allows us to show that the projected reverse diffusion is well-posed and numeri- cally stable, enabling sample quality guarantees.
TDM (heat ker.)
1
logí 7)
TDM (ISM)
d
polyíd. 2)
poly(d, }
log(
RDM
d
polyid, +)
same
same
SCRD
2d
+ poly(d, 2)
(dly 1) doog( (
This paper 1
d3 ()
d d'log()) (
d (og(())
Verifying the assumption on common manifolds. We choose projection maps yp that satisfy Assumption 2.1 with small Lipschitz constants. For example, for Ta, (x()[( = x [i] mod2m, i E [d] is 1-Lipschitz on all R trivially sat- isfying the assumption. For the sphere, sp(x) क
2.3. Sampling algorithm and theoretical guarantees
Next, we present our sampling algorithm, which uses the trained models to simulate the reverse diffusion process.
_ is
Sampling procedure. Our training algorithm (Algorithm 1) outputs trained models f(x, t) and g(z,t) .) for the drift and covariance terms of our reverse diffusion. We then use these models to generate samples. First, we sample a point from the stationary distribution of the Censtein-Uhlenbed process Z, on R., which is Gaussian distributed. Next, we project this point 2 onto the manifold to obtain a point y = p(()) and solve the SDE dY, = fY,, t)dt + (Y,, B) given by our trained model for the reverse diffusion's 's drift and covariance over the time interval [o, 77. starting at the initial point y. To simulate this SDE, can use any off- the-shelf numerical SDE solver, which takes as input the trained model for I and g. and the exponential map on M. We give one such solver in Algorithm 2, and prove guarantees for the accuracy of the samples it generates, and its runtime, in Theorem 2.2. Our guarantees assume the trained models fxx 1) and g(x, i) we hand to this solver minimize our training objective within some error e >
2- Lippschitz outside ball of radius around the origin, where the forward diffusion remains with high probability - -(2--). For U(n) (or SO())), xx(), which com- putes the spectral decomposition AU of X + +x*, has derivatives with magnitude bounded by the inverse eigen- value gaps x While singularities occur at points with duplicate eigenvalues, random matrix theory shows that eigengaps are w.h.p. bounded below by poll(d)) ensur- ing & satisfies the average-case Lipschitz assumption. For the unitary group, we show that Assumption 2.1 holds for L = Odd'5 1) and L2 = O(dTa 3) (Lemma B.4). For the sphere, it holds for L1 = L2 = O(a- 1). For the torus, it holds for L1 = L2 = 1. These bounds, derived in Appendix E, imply the assumption holds with high probability under standard random matrix models.
Theoretical guarantees. We denote by (MM == (x)) x E MM C Rdd the pushforward of M wi.t. th.
Symmetry and forward diffusion structure. Our theo- retical guarantees hold when M satisfies symmetry prop- crty and yp satisfies an "average-case" Lipschitz condition (Assumption 2). This symmetry property requires that each point 2 E IRd can be parametrized as z = z(U, where U = (02)) E .M and A = A(2) E A for some A C (d) is another parameter. For instance, on the sphere, U = is the projection onto the sphere, and
Theorem 2.2 (Aceuracy and runtime of sampling algo- rithm). Let O, and suppose that yp Rd - M satisfies Assumption 2.1 for some L1, L2 poly(d) and (MM) is bounded by ball of radius poly (d). Suppose that f and Å are outputs of Algorithm 2, and that f and 9 minimize our training objective for the target distribu- tion a with objective function value < e. Then Algorithm 2. with inputs f and 9. outputs generated sample whose probability distribution v satisfies
E, and
| = ||=|| is the distance to the origin. For SO(n) or U(n), the parametrization comes from the spectral decomposition UAU* where U E .M and A is a diagonal matrix. On the torus, U- = x(x) is the projection onto the torus, and A E 2x 2., 1. 2 0, is the Censtein-Uhlenbeck process on IRd Xx == (22) is our forward diffusion process on M, and Y, := X its time-reversal (see Section 3.1). This structure ensures that the reverse diffusion inherits behl-be properties from the Fuclidean process.
||~- < O(E(d t L2) log ) = =(Expoly((d)
Moreover, Algorithum 2, takes
O(id'T, + d2) log ( poly(d) x log (1)
iterations, where each iteration requires one evaluation of ff and g, one evaluation of the exponential map on M, plus O(d) arithmetic operations.
4

---

Efficient Diffusion Models for Symmetric Manifolds
Plugging in our bounds on the average-case Lipschitz con- stants in the case of the torus, sphere, SO(n), and U(n) (Lemma B.4) into Theorem 2.2, we obtain the following guarantees for the accuracy and runtime of our sampling algorithm for these symmetric manifolds:
3. Algorithm derivation and proof highlights
3.1. Derivation of training & sampling algorithm
Given a standard Brownian motion B, in R. a u R + Rd and R : Rd - Rdxd, a stochastic process X₁ satisfies the SDE dX;- ((X
Corollary 2.3. Suppose that M is T.. Sa. SO(n), U(n) with - vd. Suppose that sp and ub are chosen as specified above for these manifolds. Suppose that f and 9 are outputs of Algorithm 2, and that f and & minimize our training objective for the target distribution a with objective function value < e. Then Algorithm 2, with inputs I and I, outputs a generated sample whose probability distribution v satisfies
with initial condition I E if X, = x + So (x + + R...)dB..
Lemma 3.1 (Itō's Lemma). Let vb 12d - IR* be a second- order differentiable function, and let (() E Rd be an ltô diffussin. Then for all 2 and all i E [M], we have -
+x
dXx +
V d
The transition kernel Pal- (y]zz) is the probability (density) that X takes the value at time t conditional on X taking the value x at time т. Given an initial distribution т, the prob- ability density at time t is pe(a) = SM ( (x|=) (2)((()))))) (z) For any diffusion X., E [0, T), its time-reversal Y, is the stochastic process such that Y Xx- for E [0, T7. Y, is also diffusion, governed by an SDE. In the special case where X, has identity covariance, dX,- = b(X,)dt + dB,, the reverse diffusion satisfies (Anderson, 1982)
-lav S $O(E x log (1)
for the torus and sphere, and
| - IlllV < O(E x d log (4)
for (O(n) and Un... Moreover, Algorithm 2, takes Od log (e) iterations for the torus and sphere, and Od log (1) iterations for (O))) and U(n)). Here each iteration requires one evaluation of f and g, one evalua- tion of the exponential map on M, plus ((d) arithmetic operations.
dY, = d( llog (Y))dt dB,.
(1)
One can also define diffusions on Riemannian manifolds, in which case dB, is the derivative of Brownian motion on the tangent space (see (Hsu, 2002)). Below we show the key steps in deriving our diffusion model, training algorithm (Algorithm 1), and sampling algorithm (Algorithm 2).
Comparison with prior work. Theorem 2.2 improves on the accuracy and runtime guarantees for sampling of (De Bortoli et al., 2022) when M is one of the aforemen- tioned symmetric manifolds, since their accuracy and run- time bounds for sampling are not polynomial in the dimen- sion d (for instance, the "constant" term C' C(M, d) in (De Bortoli et al., 2022) has an unspecified dependence on the manifold and its dimension). Finally, we note that (Lou et al., 2024; Huang et al., 2022) do not provide guarantees on the accuracy and runtime of their sampling algorithm, and that the runtime bounds for the sampling algorithm in (Zhu et al., 2025) are not polynomial in dimension. Improv- ing the dependency on dimension remains an open question for future work.
Forward diffusion. Let Ze be a diffusion on R2d initialized at go = (T)) We choose Z, to be the Ornstein-Uhlenber process, dz, 2 + dB,,, whose stationary distribu- tion is (0, Z, is easy to sample as it has closed Gaussian transition kernel 4t/- Let X, == (Z.). the pro- jection of Ze onto M. X, is our model's 's forward diffusion.
Reverse diffuusinn SDE. Let Y; := Xr- denote the time- reversal of X, Y, is diffusion on M, and its distribution at time equals the target distribution л. It follows the SDE:
dY, = fY,, t)dt + g" g",Y,, t)dB.
(2)
Extension beyond symmetric manifolds. While the- oretical guarantees focus on symmetric manifolds, the al- gorithm itself applies more broadly. In Appendix F, we deseribe how projection maps y and exponential maps can be constructed for certain non-smooth or non-symmetric spaces, such as convex polytopes. Although proving Lip- schitz properties in these settings is more subtle due to curvature boundary singularities, our framework may still apply empirically. Extending theoretical guarantees to such general manifolds remains a promising direction for future work.
for some functions f* 1) M x [0, - T.M and g" (x,t): : M T.M x T.M. Here dB, is the derivative of standard Brownian motion on M's tangent space. We write dB, dBC when E M is clear from context.
We cannot directly apply (1) to derive a tractable SDE for the reverse diffusson Y, on M, as the transition kernel p. r of the forward diffussion X, on M lacks a closed form expression. Instead, we first use (1) to obtain an SDE for the reverse dif- fusion of Z, E Rd, dH,, = (H,/2 + 2V log (T.)(H + d.B. We use Itô's Lemma to project this SDE onto M. giving an SDE for Y, (see Appendix B.
An overview of the proof is given in Appendix A; the full proof appears in Appendix B.
dy, = ((vp((V T anI V (3)
5

---

Efficient Diffusion Models for Symmetric Manifolds
Training algorithm's objective function. From (3), we show one can train a model f. 9 for f*, 9" by solving an optimization problem (Lemma B.2). Here, f.g E C(M x [o, .T. T.M) are continuous functions from M x [0, T] to the tangent space TaM and & ~ Unif[o, 1]. Ju denotes the Jacobian of pp, that is, (x(x)) : I2 + T is the linear operator which maps any v E Id to the derivative of p in the direction of v..
(in practice, f 9 are neural networks fe Sg. and 0, outputs of Algorithm 1). We initialize this SDE at (o(N(O,, Ta)), the pushforward of (0, Ia) onto M with respect to y.
di; f(y t)dt+ t)dB, Yo ~ (0>((NN(, Ia)). (7)
To generate samples, we numerically simulate the SDE (7) for Yr by discretizing it with a small time-step A 0:
Di+1 = exxp(0t í(i. 1)A + g(:s. 4) AC),
(8)
miny E, (Vp(Zr-1)) Zp- (1b)) + ((tr((>p(Zz=-))) f(v(Zr-)), 1) = (4)
--
E {O..... 7/a}, initialized at jo (ON Ia)).
Algorithm Training algorithm
ming E,E [1) -g(( 11) =
Input: A way to compute a "projection" : : Rd -> M, and its
(5)
gradient
Input: A way compute map M + Rd s.L. (()(x))) =
Sublinear computation of training objective. For mani- folds with non-zero curvature, such as the sphere, SO(n), and U(z), our forward and reverse diffusions differ from prior works and incorporate a ppatially-varyinn covariance term to account for curvature. This allows the forward diffusion to be computed as a projection yp of the Ornstein- Uhlenbeck process in _ = [ (or C"**) onto the manifold. For (O(n) or U(2). yp is computed by one sin- gular value decomposition U of the Gaussian matrix Zr- + 2 requiring tions, where (n²²) is the manifold dimension. This en- ables computation of the drift term's gradient (4) in (dd) arithmetic operations and one gradient evaluation of f.
Vx E M
Input: Dataset D = (ad..... abd) } C M. Hyperparam. T >
Input: Models fa .M x [0 TM and 93 M x [0, +
TM x T.M Ô E R ā E R denote trainable parameters Input: Initial parameters Os E E du E I
Input: Hyperparameters: Number of stochastic gradient descent iterations E N. Step size > 0, batch size b 1: Define, 6Q E R"
E R b, x0 E M.
f E [0, the objective function F(O, b, 2, f) II(vp(2)) T 1-v e-(2--)
=
(F--0/2 + ((r(V=(((E) - [ i) 17
(n²*) = ((d*) ) arithmetic opera-
Define v0 R2 2 e Rd e [0. T]. the objective
G(0:b, 2, à, == (e)))))(()))))))))))))))())())) (13 0
3: Set + 0o á e- do
= io I. do 5: Sample random batch S C [m]] of size b. Sample ~ TT)
4:
To train the reverse diffusion's SDE, we also need to model the covariance term (5), dxdd = =n² xn² matrix. To achieve a per-iteration runtime sublinear in the d² = n² matrix entries, we leverage the special structure of the covariance matrix, which arises from the manifold's symmetries. For example, the forward diffusion (t) E SO(n) (or (t) E
7: for j E S do
8: Sample & ~ N La)
9: Set 2₁ + ((r)ee ((T-t) + v1 -0--i
10: Set + p(2))
11: end for
12: Compute + 1 Eyes VoF(0 25, 25,1)
Un)) is governed by the following system of SDEs:
13: 0 -- 0 '/I'
14: Compute T 1o2 Ejes 2,
dt,(t)= (f)=)xx EjjmnEa (at)(t)d u5 (t) Bg(l))
u: (t))dd (6)
15:
@ x- -+Φ- 1p-1II 1/T
16: end for
where (ty(t) = E E(/((00-x(()))((2)) U()) and (1)( = (0/00).000A,(0))((((()))) U()) Vi (m). To train model for this covariance term with sublinear run- time, we exploit the symmetries of the underlying group. These symmetries ensure the covariance term in (6) is fully determined by n scalar terms (t) for i,j E [r] and the n X n matrix U. Thus, it suffices to train a model A(U, E Raxn for these n2 terms by minimizing the ob- jective []AA(U,t) All where A is the n matrix with entries A.j 1/(x,(t))>(()),, and (t( is the ith diagonal entry of A = A(l). A similar method applies to efficiently train the covariance term for the sphere (see Appendix E).
17:
output: Parameters 8 & for the models fo and go
Algorithm 2 Sampling algorithm
Input: way to compute the value of the exponential map
exp(x, v) on some manifold M, for any x E M, v E T.M
Input: A way to compute the "projection" map ş Rd + M
[0, T] - TM and g3 . [0, + Input: Models fa M x
TM x T.M â E TR" á E R" denote trainable parameters
Input: Parameters 0, o (from output of Algorithm 1),
Input: T> 0. N E N. A > such that T/a NZ.
1: Sample 20 ~ (0, Ia), and Set ýo <- ((20)
2: for 0........... T/a
do 3: Sample E ~ N(0, Ia).
Sampling algorithm. To (approximately) sample from I,, we approximate the drift and covariance terms of the reverse diffusion (2) via trained models f, å obtained by solving (4)
4: Set ő841 < exp(ir. f(0,, iAjA + 9l., (A)AE
5: end for
6: output $2/A
6

---

Efficient Diffusion Models for Symmetric Manifolds
3.2. Highlights of proof of sampling guarantees
real-valued random matrices. During these intervals, the dif- fusion moves at exp(d) velocity. Despite this, we show that step size of (/poly(d, }) suffices to simulate a random solution to the SDE with a distribution ô-close to the correct one. During these intervals, interactions between eigenvec- tors nearly separate into slow-moving cigenvectors and pairs of fast-moving eigenvectors, with a simple transition kernel from the invariant measure on S0( (see Appendix B.7).
Girsanov transformations, used in prior works bound accuracy, do not apply to our diffusion due to its spatially varying covariance. We adopt an optimal transport approach, selecting an optimal coupling between the "ideal" diffusion Y,, governed by the SDE dY, f (Y,,, ()dl+)" (Y,, Br, and the diffusion Y, where f", are replaced by our trained model f. 9. within an error e. We first construct a simple coupling between Y and Y, by setting the underlying B, in their SDEs equal. Applying comparison theorems for manifolds of non-negative curvature to the coupled SDEs, we prove generalization of Gronwall's inequality to SDEs on manifolds (Lemma B.3).
4. Empirical results
We provide proof-of-concept simulations to compare the quality and efficiency of our algorithms to key prior works.
Datasets. We evaluate the quality of samples generated by our model on synthetic datasets from the torus Ta. the spe- cial orthogonal group SO(n), and the unitary group U(n). The torus provides simple, zero-curvature geometry for initial validation, while S0(n) and U(n) test the model on more complex geometries. Datasets include unimodal wrapped Gaussians on Ta, multimodal Gaussian mixtures on SO(n), and time-evolution operators of a quantum oscil- lator with random potentials on U(n), for varying dimen- sions d n². We also separately analyze per-iteration runtime to study scaling across dimensions, requiring only single training step and limited computational resources.
W2(Y,, Yi) S (p, Yo) + +e)e",
(9)
where P is geodesic distance on M and W2 the Wasserstein distance. (9) holds if g" are e-Lipschitz on all of M.
Showing "average-case" Lipschitzness. op is not in general Lipschitz. E.g., on (1) and U(n). (ZZ) has singularities at points where the eigengaps of Z + Z" vanish. By using tools from random matrix theory, we instead show y satis- fies an "average-case" case" Lipschitzness, on a set S2 on which the diffusion Ze E Rdd remains w.h.p. (Lemma B.4). Next, we show that for f*. g" to be c-Lipschitz everywhere on M, it is sufficient for o to only satisfy average-case Lipschitz- ness. To do this, we express I (and g²) as an integral over the eigenvalues A of 7,= + Zi= = UAU*,
For the torus Ta. following several works (De Bortoli et al., 2022; Zhu et al., 2025), we train diffusion models on data sampled from wrapped Gaussians on tori of different di- mensions d E {2, 10, 50, 100, 1000), with mean and co- variance 0. (See Appendix C.1 for definition of wrapped Gaussian). For U(n), following (Zhu et al., 2025), we use dataset on Unn) of unitary matrices representing time- evolution operators eitlf of quantum oscillator. In our sim- ulations, we consider a wider range of n E {3. 5, 9, 12, 15} (corresponding to manifold dimensions d' = =[n-1) {3. 10, 36, 66, 105}) when evaluating sample quality, and n E {3, 5, 10, 30, 50) (d E {3, 45,45, 435, 1225}) when evalu- ating runtime. Here is time, H = 2m - V is a Hamil- tonian, and A is the Laplacian. V is a random potential ((x) = - || - 10 with angular momentum w sampled uniformly on 32.3] and xo ~ (0, 1). As A, V are infinite- dimensional operators, matrices in U(n) are obtained by retaining the (discretized) top-n eigenvectors of A, V. For SO(n), following (Zhu et al., 2025), use datasets sam- pled from a mixture of small number k of wrapped Gaus- sians on SO(n). We set A 2 and use the same values of n as on U(n).
(((U,)) < / |V T V log 4o (0))) ... 11 (Z)dA.
Due to the manifold's s symmetries, we observe 2e (and the entire integrand) depend only on A, not the eigenvectors U E Un). This allows us to show the integral "smooths out"' the singularities of sp, and that f (U, t) (and g²) are poly (d)- Lipschitz at every U E U(n) on the manifold (Lemma B.6).
Improved coupling to obtain poly(d) bounds. While we have shown our model's SDE is poly (iiiishhti M, after times 7 > 1 polytdy our Wasserstein bound (9) grows exponentially with d. To overcome this, we define new coupling between Y,, y which we "reset" after time intervals of length T 1/c by converting our Wasserstein bound into a total variation (TV) bound after each interval. Key to converting the bound is to show that w.h.p. the Jacobian everywhere in a ball of radius 3/poly(d) around our diffusion. By alternating between Wasserstein and TV bounds, we get error bounds which grow proportional to T/r = (oly(d) (Lemma B.7).
projection o has voly(d)-Lipschh))
Al Algorithm. We use Algorithm I to train our model, and Algorithm 2 to generate samples. Each algorithm takes as in- put a projection map y and restricted inverse v with choices for Ta. SO(n), and U(n) detailed in Section 2. For SO(n) and U(x), we use the projection map &p defined in Section 2,
Handling instability on (n).. These proof ideas extend casily to the torus and sphere. For (O),,, an additional chal- lenge arises: w.h.p., gaps between neighboring eigenvalues become exponentially small in over short time intervals due to weaker "electrical repulsion" between eigenvalues of
7

---

Efficient Diffusion Models for Symmetric Manifolds
as it suffices generate high-quality samples in our simu- lations. In both algorithms, the drift function f(:--) for the reverse diffusion is parameterized by neural network. De- tails on the network architecture, training iterations, batch size, and hardware are provided in Appendix C.2.!
in e-g. (Leach et al., 2022) to evaluate sample quality of diffusion models on SO(n) for 3. It measures the abil- ity of a classifier neural network to differentiate between generated samples and samples in test dataset. This met- ric allows one to evaluate sample quality in settings where computing the likelihood may be intractable, as may be the case for diffusion models on SO(n) and U(n) for larger n. We use the C2ST metric on (0) and U(n) instead of log-likelihood, as found that computing log-likelihood for n >> 3 poses additional computational challenges, while C2ST can be computed efficiently. See Appendix C.3 for definitions of log-likelihood and C2ST, and additional de- tails.
The function -(::). which models the covariance in our model's is reverse diffusion SDE, vanishes on Ta. On SO()), U1z), 9 is n² X n² matrix. This matrix has a special structure (see Section 3.1), allowing it to be parameterized by d = n² numbers. We may thus parametrize (xx, () by neural net with inputs z of dimension d, of dimension 1, and output dimension d. Network architecture is the same as for f(:.-).
Results for generated sample quality. For the torus, we compare our model, Euclidean diffusion model, and RSGM dataset sampled from wrapped Gaussian on the d-dimensional torus for different values of d. Our model has the lowest negative log- ligelikelihood (NLL) for d 2 10, and its NLL degrades with the dimension at a much slower rate than the Euclidean model and RSGM (Table 2; lower NLL indicates better-quality sample generation).
Benchmarks. We compare samples generated by our model to those from RSGM (De Bortoli et al., 2022), TDM (Zhu et al., 2025), and "vanilla" Euclidean diffusion model. RSGM and TDM are included as they have demonstrated improved sample quality and runtime over previous mani- fold generative models, such as Moser Flow (Rozen et al., 2021) on Ta and (3). TDM also outperforms RSGM and Flow Matching (Chen et al., 2024) on higher-dimensional torus datasets and shows better sample quality on SO(n) and Un)), where RSGM and RDM (Huang al., 2022) lack experiments for n > 3. All models are trained with the same iterations, batch size, and architecture (see Appendix C.2).
For U(n), we train our model, a Euclidean diffusion model, RSGM, and TDM on dataset on U(n) comprised of dis- cretized quantum oscillator time-evolution operators, for n E {3, 5.9, 12, 15). For n 2 9, our model achieves the lowest C2ST score; a lower C2ST score indicates higher- quality sample generation (Figure I, top). Visually, we observe our model's generated samples more closely re- semble the target distribution than benchmark models' for 2 9 (see Figure 1, bottom, for n = 15, Appendix C.4 for other n). Improvements on SO(n) were similar to those on U(n) (see Appendix C.4).
For the Euclidean diffusion model, samples are constrained to the manifold M by preprocessing via q and postprocess- ing via の. For M Ta. [xx)[i]] = =[i] mod 2n, and i is its inverse on [0, 2z. For SO(n) and U(n), (xX) computes the U from the singular value decomposition
X X= UEV" and v is the usual embedding.
RSGM and TDM are trained using their divergence-bassee ISM objective, which is prohibitively slow for large dimen- sions. To fully train these models and evaluate their sample quality, we follow their implementation and use a stochastic estimator for the ISM divergence. We do not compare to TDM on the torus, as their specialized heat kernel objective does not generalize beyond the torus or Fuclidean space.
Table 2: Negative log-likelihood (NLL) when training on a wrapped Gaussian dataset on the torus of different di- mensions d. Lower NLL indicates better-quality sample generation. Our model's NLL increases more slowly with d, and is lower for d 2 10, than Euclidean and RSGM models.
Method d 2
10
d= 50
d = 100
1000
We do not compare to SCRD (Lou et al., 2024), as their implementation is limited to efficient heat kernel expansion for (O(n) and U(n) with n = 3. For n > 3, the cost of their expansion grows exponentially with n, making it infeasible for higher dimensions.
Euclidean 0.61#.02
0.61:.022
0.66+t.00
8.18
8.7+.37
RSGM
0.42 003
0.494:05
0.624.06
1.45.22
2.51
Ours
0.49+.04
0.47+.03
0.50++05
0.52 .99
4.97
Results for runtime. We evaluate the per-iteration training runtime on U(n) on a wider range of n 3, 5, 10, 30, 50 (corresponding to manifold dimensions of d n(n=1) [3, 10, 45, 435, 1225)). Our model's per-iteration runtime remains within factor of 3 of the Euclidean model's for all n. Per-iteration runtimes of TDM and RSGM, with ISM objective, increase more rapidly with dimension and are, respectively, 45 and 57 times greater than the Euclidean model's for n = 50 (Table 3). Similar runtime improve- ments were observed on 50(n) (Table 6 in Appendix C.2).
Metrics. For the torus, as in (De Bortoli et al., 2022; Zhu et al., 2025)), we evaluate the quality of generated samples by computing their log- likelihoodd For SO(n) and U(n). we use the Classifier Two-Sample Test (C2ST) metric (Lopez- Paz.& & Oquab, 2017). The C2ST metric was previously used
Our code can be found at Diffusion- Jiffusion-Moxees-ffon Models-for-Se Sear--Symttrii--aaaaene
8

---

Efficient Diffusion Models for Symmetric Manifolds
Table 3: Per-iteration training runtime in seconds on Un).. The fastest the Fuclidean model is in gray for comparison. Our model's runtime remains within a factor of 3 of the Euclidean model's for all n. Runtimes of TDM, RSGM increase more rapidly with dimension d = (nn²) and are, respectively, 45 and 57 times greater than the Euclidean model's 's for n = 50 (d- = 1225).
Our approach improves training runtime and sample quality symmetric manifolds, significantly narrowing the gap be- tween manifold-based diffusion models and their Euclidean counterparts. Furthermore, the model naturally extends to conditional generation: given a conditioning variable y, one can feed y as an additional input to the learned drift and co- variance functions, mirroring conditional diffusion models in Euclidean settings.
diffusion model is in bold;
Method
R ==3
R m=
R=
12 m
0 50
Several open directions remain. One is to extend our frame- work to more general manifolds-suce as the manifold of positive semi-definite matrices, or other domains admitting a projection oracle satisfying suitable average-case smooth- ness properties (see Appendix F). Another direction is to handle distributions supported on union of manifolds with varying dimensions, such as the GEOM-DRUGS dataset (Jing et al., 2022), which lies union of tori.
Fuclideun
0.19 1.01
0.19 1.01
1.20
0.20 1.001
0.211000
RSGM
1.0 1.02
1.222.00
1.51..03
19.9811119
11.35 3.1
TDM
0.91 2.08
1.07 1.06
2.46
4.09
3.77 1.17
9.43 1.23
Ours
0.36 1.00
0.36 1.00
0.36 1.00
0.46 1.01
0.60 1.01
Summary. We find that, as predicted by our theoretical training runtime bounds in Table 1, the per-iteration training runtime of our model is significamtly faster, and grows more slowly with dimension, than previous manifold diffusion models on Unn) (similar improvements were observed for (O(n)). Our algorithm's 's runtime remains within small constant factor of the per-iteration runtime of the Euclidean diffusion model, at least for < 50 (corresponding to manifold dimension of d < 1225), nearly closing the gap with the per-iteration runtime of the Euclidean model.
Finally, while our method yields polynomial-in-4 bounds on sampling accuracy-improving upon prior works that lacked such guarantees- guahantess-tightennngg tightening this dependence re- mains an important challenge for future research.
Impact statement
Diffusion generative models that generate data constrained to symmetric manifolds can significantly enhance the soci- etal impact of machine learning applications. These include areas such as drug discovery and robotics, where molecular or robotic configurations can be represented as points on symmetric manifolds, as well as quantum mechanies ap- plications like quantum chemistry, materials science, and microelectronics. Moreover, our focus on improving the efficiency of generative models has clear environmental ben- efits. Reducing the computational costs associated with training and sampling minimizes energy consumption, help- ing to lower the carbon footprint of machine learning ap- plications. This is particularly relevant as the demand for large-scale generative models continues to grow. However, generative modeling also presents potential risks. For exam- ple, such models can be used to generate fake news videos or to train robots and drones for harmful purposes when misused by malicious actors. While our work is primarily theoretical, limiting its direct negative impact, it important to remain mindful of both the positive and negative impli- cations of applying generative modeling technologies in sensitive domains. Balancing these considerations is essen- tial to responsibly hamess the potential of generative models in advancing critical scientific and societal challenges while mitigating risks and environmental impacts.
Moreover, we find that (except very low dimensions) our model is capable of improving on the quality of samples generated by previous diffusion models, when trained on dif- ferent synthetic datasets on the torus, (O(n) and Un)). The magnitude of the improvement increases with dimension.
Method
n ==
n
nL
Euclidean
69 10.00
.75 1.04
.87 1.04
.97 1.02
1.00
RSGM
.79 1.04
.92 1.04
.97 1.03
1.00 1.02
1.00 1:01
TDM
.73 1:04
.91 1.02
.99 1. 1.02
1.00 1.01
1.00 1:01
Ours
.75 +.05
.80 +:04
.84 + +.04
.88 +.05
.90 :.04
Target
Ours
Euclidean
RSGM
TDM
Figure 1: C2ST scores when training on datasets of quantum evolution operators on Unn) (top). Lower scores indicate better-quality generated samples (range is [0.5, 1)). For n Z 9. our model has the best C2ST scores. Generated samples are plotted for n = 15 (bottom); axes are two matrix entries.
5. Conclusion and future work
We introduce new diffusion model with spatially varying covariance structure, enabling efficient training on sym- metric manifolds with non-zero curvature. By leveraging manifold symmetries, we ensure the reverse diffusion satis- fies an "average-case" Lipschitz condition, which underpins both the accuracy and efficiency of our sampling algorithm.
Acknowledgments
OM was supported in part by Google Research Scholar award. NV was supported in part by NSF CCF-2112665.

---

Efficient Diffusion Models for Symmetric Manifolds
References
Dyson, F. J. A Brownian-motion model for the eigenvalues of random matrix. Journal of Mathematical Physics, (6):1191-119)) 1962.
Anderson, B. D. Reverse-time diffusion equation models. Stochastic Processes and their Applications, 12(3):313- 326, 1982.
Feiten, W., Lang, M., and Hirche, S. Rigid motion estima- tion using mixtures of projected Gaussians. In Proceed- ings of the 16th International Conference on Information Fusion, pp. 1465-1472. IEEE, 2013.
Anderson, G. W., Guionnet, A., and Zeitouni. O. An intro- duction to random matrices. Number 118, Cambridge university press, 2010.
Bakry, D., Gentil, I., Ledoux, M., al. Analysis and geome- try of Markov diffusion operators, volume 103. Springer, 2014.
Gronwall, T. H. Note on the derivatives with respect to a parameter of the solutions of system of differential equations. Annals of Mathematics, pp. 292-296, 1919.
Ben-Hamu, H., Cohen, S., Bose, J., Amos, B., Nickel, M.,
Ho, J., Jain, A., and Abbeel, P. Denoising diffusion proba- bilistic models. Advances in neural information process- ing systems, 33:6840-6833 2020.
Grover, A., Chen, R. T., and Lipman, Y. Matching nor-
malizing flows and probability paths on manifolds. In International Conference on Machine Learning, pp. 1749- 1763. PMLR, 2022.
Hsu, E. P. Stochastic analysis on manifolds. Number 38.
American Mathematical Soc., 2002.
Benton, J., De Bortoli, V., Doucet, A., and Deligiannidis, G. Nearly d-linear convergence bounds for diffusion models via stochastic localization. In The Twelfth International Conference on Learning Representations, 2024.
Huang, C.-W,, Aghajohari, M., Bose, J., Panangaden, P., and Courville, A. C. Riemannian diffusion models. Ad- vances Neural Information Processing Systems, 35: 2750-2761, 2022.
Chen, H., Lee, H., and Lu, J. Improved analysis of score- based generative modeling: User-friendly bounds under minimal smoothness assumptions. In International Con- ference on Machine Learning, pp. 4735-4763. PMLR, 2023a.
Jing, B., Corso, G., Chang, J., Barzilay, R., and Jaakkola, T. Torsional diffusion for molecular conformer generation. Advances neural information processing systems, 35: 24240-24253, 2022.
Chen, T., FAIR, M., and Lipman, Y. Flow matching on general geometries. In ICLR, 2024.
Jo, and Hwang, S. J. Generative modeling on manifolds through mixture of Riemannian diffusion processes, In International Conference on Machine Learning, 2024.
Chen, S., Chewi, S., Li, J., Li, Y., Salim, A., and Zhang, A. Sampling as easy as learning the score: theory for diffusion models with minimal data assumptions. In The Eleventh International Conference Learning Represen- tations, 2023b. URL /tsttpse//penrrviiee.. net./ forum? wiaiumddddyyyyMggss..
Leach, A., Schmon. M., Degiacomi, M. T. and Willcocks, G. Denoising diffusion probabilistic models on SO(3) for rotational alignment. In ICLR 2022 Workshop Geometrical and Topological Representation Leurning, 2022.
Cheng, X.,
Zhang, J., and Sra, S. Theory and algorithms for diffusion processes on Riemannian manifolds. arXiv preprint arXiv:2204. 13665, 2022.
Lopez-Paz, D. and Oquab, M. Revisiting classifier two- sample tests. In ICLR, 2017.
Cheng, Y., Gong, Y., Liu, Y., Song, B., and Zou, Q. Molec- ular design in drug discovery: a comprehensive review of deep generative models. Briefings in bioinformatics, 22 (6):bbab344, 2021.
Lou, A., Xu, M., Farris, A., and Ermon, S. Scaling Rieman- nian diffusion models. Advances in Neural Information Processing Systems, 36, 2024.
Cranmer, K. Kanwar, G., Racanière, S., Rezende, D. J., and Shanahan, P. E. Advances in machine- teaeeeeeeaannnbbaeee based sampling motivated by lattice quantum chromodynamics. Nature Reviews Physics, 5(9):526-535, 2023.
Mangoubi, O. and Vishnoi, N. K. Private covariance approx- imation and eigenvalue gap bounds for complex Gaussian perturbations. In The Thirty Sixth Annual Conference on Learning Theory, pp. 1522-1587. PMLR, 2023.
De Bortoli, V., Mathieu, E., Hutchinson, M., Thortoon, J., Teh, Y. W., and Doucet, A. Riemannian score-basee based generative modelling. Advances in Neural Information Processing Systems, 35:2406-2477 2022.
Mangoubi, O. and Vishnoi, N. K. Private low-rank approxi- mation for covariance matrices, Dyson Brownian motion, and eigenvaluc- gap bounds for Gaussian perturbations. Journal of the ACM, 72(2):1 188, 2025.
10

---

Efficient Diffusion Models for Symmetric Manifolds
Mathieu, E. and Nickel, M. Riemannian continuous normal- izing flows. In Advances in Neural Information Process- ing Systems, 2020.
Nash, The imbedding problem for Riemannian manifolds.
Annals of mathematics, 63(1):20-63, 1956.
OpenAl. Video generation models as world simula- tors, 2023. URL httpp:://ppeeaa..ccoo com /ndex/
vi me a- wwwoordddsmmmllaaae sis-ssmmmmmssss...
Petersen, P. Riemannian geometry, volume 171. Springer, 2006.
Rauch, H. E. A contribution to differential geometry in the large. Annals of Mathematics, 54(1):38-55, 1951.
Rombach, R., Blattmann, A., Lorenz, D., Esser, P., and Ommer, High-resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF con- ference computer vision and pattern recognition, pp. 10684-10695, 2022.
Rozen, N., Grover, A., Nickel, M., and Lipman, Y. Moser flow: Divergence-based generative modeling on mani- folds. Advances in Neural Information Processing Sys- tems, 34:17669-1766 2021.
Rudelson, M. and Vershynin, R. Hanson Wright inequality and sub-Gaussian concentration. Electronic Communica- tions in Probability, 8(82):1-9, 2013.
Vishnoi, N. K. Geodesic convex optimization: Differen- tiation on manifolds, geodesics, and convexity. CoRR, Abs/1806.00333, 2018. URL http://arxiv.org
06373.
Watson, J. L., Juergens, D., Bennett, N. R., Trippe, B. L., Yim, J., Eisenach, E., Ahern, W., Borst, A. J., Ragotte, R. J.. Milles, L. F., al. De novo design protein struc- ture and function with rfdiffusion. Nature, 620(7976): 1089-1100, 2023.
Yim, J., Trippe, B. L., De Bortoli, V., Mathieu, E., Doucet, A., Barzilay, R., and Jaakkola, Se (3) diffusion model with application to protein backbone generation. In In- ternational Conference Machine Learning, pp. 40001- 40039. PMLR, 2023.
Zhu, Y.. Chen, T., Kong, L., Theodorou, A., and Tao, M. Trivialized momentum facilitates diffusion generative modeling on Lie groups. In International Conference Learning Representations, 2025.
2025.

---

Efficient Diffusion Models for Symmetric Manifolds
A. Proof outline of Theorem 2.2
In the following, for any random variable X, we denote its probability distribution by Cx. As already mentioned, previous works use Girsanov's 's theorem to bound the accuracy of diffusion methods. However, Girsanov transformations do not exist for our diffusion as it has non-constant covariance term which varies with the position x. Thus, we depart from previous works and instead use an optimal transport approach based carefully chosen optimal coupling between the "ideal diffusion" Y, and the algorithm's 's process iit Specifically, denoting by μr the distribution of Yz and by Vi the distribution of Y, the goal is to bound the Wasserstein optimal transport distance W2 Pt) == inf (r,,, 8,y, [0p(( (Y,, Y2)l where K(H,2) 2/ is the collection of all couplings of the distributions μ and v, and p is the geodesic distance M. Towards this end, we would like to find coupling x which (approximately) minimizes ((ys-pe, ( (Y,, Y2) at any given time t.
As first attempt, we consider the simple coupling where we couple the "ideal" reverse diffusion Y,,
dy (Y,, (dd + g (Y,, ()B,,,
(10)
and the reverse diffusion y; given by our trained model f.g.
dÝ; = ()dtt + g(Y, t)dB,
(11)
To couple these two diffusions, we set their Brownian motion terms dB, to be equal to each other at every time t. In a similar manner, we can also couple Y and the discrete-time algorithm ý by setting the Gaussian term & in the stochastic finite difference equation (8) to be equal to & V A A(i+1) dB,dl for every i.
Step 1: Bounding the Wasserstein distance for everywhere-Lipschitz SDEs. To bound the Wasserstein distance W2(Y,, ii) S WY, Y,) + W2(Y,, (t), we first prove generalization of Gronwall's inequality to Stochastic differential equations on manifolds (Lemmaa B.3). Gronwall's 's inequality (Gronwall, 1919) says that if R [0, T + R satisfies the differential inequality a R(t) < R(t) for all & > 0, where the coefficient 3()): [0, 2] + R may also be a function of t, then the solution to this differential inequality satisfies R(t)= < R(0)e a(s)ds
Towards this end, we first couple Y1 and Y, by setting their Brownian motion terms dB, equal to each other and then derive an SDE for the squared geodesic distance p² (Y,, Y.) using Itô's lemma. Taking the expectation of this SDE gives an ODE for EBpp (X,, X.)l.
(X, X.)
(M(X.,,) (((x,t)) =E Vpp X.)
(Zl.. 9) (iP()x)))) (g((X.tt) ((x,,) 0) dt.
+ +5 ((x,,t)
(12)
To bound each term on the r.h.s., we first observe that, roughly speaking, due to the non-negative curvature of the manifold, by the Rauch comparison theorem (Rauch, 1951), each derivative on the r.h.s. is no larger than in the Euclidean case M Rd where p²(,,,
1| - X1 13. Hence, we have that
< <2|| (M(X..0)
V (X)))
Xill -(X f(X,, S 21 Xall(e|& - Xell +e),
as long as we can show that f* is Lipschitz for some e > 0 (see Step 2 below). Bounding the covariance term in similar manner, and applying Gronwall's lemma to the differential inequality, we get that
W2 Yi) < BEP YJ) < (2.. Yo) +
(13)
Step 2: Showing that our diffusion satisties an "average-case" Lipschitz condition. To apply (13), we must first show that the drift and diffusion terms f and g" are Lipschitz on M. Towards this end, we would ideally like to apply bounds on the derivatives of the projection map yp Rd - M which defines our diffusion Y,. Unfortunately, in general, y may not be differentiable at every point. This is the case for the sphere, where the map p(2) |=|l has singularity at = = 0. This
=
issue also arises in the case of the unitary group and orthogonal group, since the derivative of the spectral decomposition (=)) = U AU has singularities at any matrix 2 which has an eigenvalue gap X - -
12

---

Efficient Diffusion Models for Symmetric Manifolds
To tackle this challenge, we show that, for the aforementioned symmetric manifolds, the forward diffusion Z, in Rd remains in some set S2 C TR with high probability 1 a, on which the map yp( ((() has derivatives bounded by poly(d) (Assumption 2.1 and Lemma B.4). We then show how to "remove" the rare outcomes of our diffusion that do not fall inside Ωz As our forward diffusion X, (and thus the reverse diffusion Y;- = Xr-)) remains at every t inside Qg with probability 2 1 - removing these "had" outcomes only adds a cost of a to the total variation error.
Showing that y has poly(d) derivatives w.h.p. (showing that Assumption 2.1 holds). We first consider the sphere, which is the simplest case (aside from the trivial case of the torus, where the derivatives of ip are all ((1) at every point). In the case when data is on the sphere, which we embed
unit sphere in Rd, one can easily observe that e.g. l| [[V((=)] - ((1) for any outside ball of radius 2 Q11) centered at the origin. As the volume of a ball of radius r = is pm one can use standard Gaussian concentration inequalities to show that the Brownian motion X₃ will remain outside this ball for time T with probability roughly O (TTT)
We next show that the Lipschitz property holds for the unitary group U(n). We first recall results from random matrix theory, which allow us to bound the eigenvalue gaps of a matrix with Gaussian entries. Specifically, these results say that roughly speaking, if Xo is any matrix and X, Xo + B(t), where (t)) a symmetric matrix with i.i.d. N (0. t) entries undergoing Brownian motion, 2010; Mangoubi & Vishnoi, 2023; 2025)).
has that the eigenvalues 71(t) ... 2 Yhe (t) of Xt satisfy for all 0 (see e.g. (Anderson et al.,
poly(n)Vi) 1)
P
n++(()) - (tt))= S
1
s $0(ss)
(14)
rE(to.T]
1 Thus, if we define Ωe to be the set of outcomes of such that (t+ - (tt) S a² poly(n)vV
we have that P(X, E Si Vt e
[to, TT) 2 1 11--a a.
Our high-probability bound on Q2 allows us to show that oo satisfies a Lipschitz property at "most" points 24. However, if we wish to apply (13), we need to show that the drift term f* and the covariance term 9 in our diffusion satisfy a Lipschitz property at every point in Rd. Towards this end, first make a small modification to the objective function which allows us to exclude outcomes (X, (e of the forward diffusion such that X, E S2 for some t E [o. 7). Specifically, we multiply the objective function (4) by the indicator function (z). As determining whether point 2 E Ω4 requires only checking the eigenvalue gaps (when M is the unitary or orthogonal group), computing 1l,,, (z) can be done efficiently using the singular value decomposition.
Bounding the Lipschitz constant of J and g". Recall that (when, e.g., M one of the aforementioned symmetric manifolds) we may decompose any z Rd we will show that, as In, (=(U, A)) depends only on A, multiplying our objective function by Ln, does not make f and g* discontinuous (and thus does not prevent them from being Lipschitz). This is because f and g" are given by conditional expectations conditioned on U, and can thus be decomposed as integrals over A. Towards this end we express f as an integral over the parameter A,
z(U, A) where U E M. Note that La, (2) is not a continuous function of 2. However,
T S [Vo(=(T,A)) V VVll Vloggr-:o(s(U.A))) + 1 " (oroo/tt)))aa-((((C.a.. (e- (A)dA,
f" = CU JAEA
gr-
where cy is a normalizing constant. Differentiating with respect to U,
d au (U,()) = EQU.AMrr du (V>p(((U,A)) Vo log (ogT
d
1 + 2
V = U
(15)
where "... "...ludeeees three other similar terms. To bound the terms on the r.h.s. of (15), we apply Assumption 2.1 which says that the operator norms of V.p, 42y, a Vep and d V2 are all bounded above by poly(d) whenever 2 E Q, To bound au the term Vu log qr (U, A)) we note that v log (= A)) is the drift term of the reverse diffusion in Euclidean space. This term was previously shown to be d Lipschiit for all t 2 2()) when the support of the data distribution in Rd lies in a ball of radius C (see, e.g., Proposition 20 of (Chen et al., 2023b)). Thus, plugging in the above bounds into (15) we
13

---

Efficient Diffusion Models for Symmetric Manifolds
have that |l 11 (0,0)22)) poly(d). A similar calculation shows that implies that (U, 1) and 9d (U, () are poly (oly(dd-Lipscchtz every U E M.
poly(d). This immediately
Step 3: Improving the coupling to obtain polynomial-time bounds. Now that we have shown that f and g" are poly(d)- Lipsehitz, we can apply (13) to bound the Wasserstein distance: W2(Y, Y+++) S (D² Yi) + + Vr 2 0, where c S poly(d).
Moreover, with slight abuse of notation, we may define U4+7 to be a continuous time interpolation of the discrete process y. Applying (13) to this process, we get that, roughly, W2 e-r1 (1rr) S (0² (úc, Y)) + e + A) for T 2 0. Thus, we get bound on the Wasserstein error,
W2(Y(+r,0u+) S W2 Y(+r) (..++..00+)) S (p +E)++ A)CTT
+
T T 0.
(16)
Unfortunately, after times 7 > poly(aj: this bound grows exponentially with the dimension To overcome this
challenge, we define a new coupling between Y2 and Y; which we "reset" after time intervals of length r = 1 by converting our Wasserstein bound into a total variation bound after each time interval. Towards this end, we use the fact that if at any time the total variation distance satisfies Lv, L9, ||tv S a, then there exists coupling such that Y, Y, with probability at least the next time interval of without incurring exponential growth in time. Repeating this process I times, we get that [|Cr, - lirll < axx x T. where the TV error grows only linearly with T.
In other words, w.p. 1 - a, we have p(ic+Zi Y(+r) 0, and we can apply inequality (16) over
Converting Wasserstein bounds on the manifold to TV bounds. To complete the proof, we still need to show how to convert the Wasserstein bound into TV bound (Lemma B.7). Towards this end, we begin by showing that the transition kernel Perr+Alett (. (Het7) of the reverse diffusion He in TRd is close to a Gaussian in KI, distance:
DN + (H (ld)) ll l PETT+AIHT (-
(11((+))) S aT T
One can do this via Girsanov's 's theorem, since, unlike the diffusion Y, on the manifold, the reverse diffusion in Euclidean space H; does have constant diffusion term (see e.g. Theorem 9 of (Chen et al., 2023b)).
Next, we use the fact that with probability at least 1 - f the map in a ball of radius poly(d) about the point He+r has c-Lipschitz Jacobian where poly(d), and that the inverse of the exponential map exp(-) has O(1)-Lipschitz Jacobian, to show that the transition kernel pt of Y (HH)) satisfies
(KI PPPTT+ (- (V(++))) s (1 Dc) d aT
T S 2
if we choose Å < ( where v = -xpxxx, (N(Y++ + Á
( + Ia)).
Next, we plug in our Wasserstein bound W (Y 9--) & 0(s) into the formula for the KL divergence between two Gaussians to bound |LY,+T+A Cegrreallre lIv.. Specifically, noting that Corerssbe = exp (N(Ut + Af(i+t: t + T), (0rtit + Ia)),, we have that
(KI, f
= Tr (V "(())) ( + +T)
det 2+2 (Y +
-d + log det (2 (T) + w (Ag (Vo+(t)))) -w.. w.
where w : Ye+r Be+r + A(f* +++) f(itr, t + T)). Since with probability 2 af we have g (Y(++r)) 2
(dly plugging in the error bounds || (Yi+,, 1) - f(Ye+T; t)ll < e and l|g" (Vetr, 1) g(Yetr, 1)ll S 8 and the c- c-Lipschitz bounds on f* and g", where c = poly(d), (Assumption 2.1), we get that DKL(M L S ( Thus, by Pinsker's inequality, we have
|| - C llrv - 11Cr, - lrv
S + (ILLl O(ec).
(DLLluillPHtTARH (-
(V+tr))
(17)
14

---

Efficient Diffusion Models for Symmetric Manifolds
Step 4: Bounding the accuracy. Recall that d is the distribution of the forward diffusion Z₁ in Euclidean space after
time c, which is an Ornstein- Uhlenbeck process. Standard mixing bounds for the Omntein-Uhlenbeed process imply that, lgg N (0, TV I "(Ce=t) for all & > 0 (see e.g. (Bakry al., 2014)), where C' I poly(d) is the diameter of the support of (T)) Thus, it is sufficient to choose T = log( C) to ensure yr, - ~|/lv 1lqr NO, La)l S OEE).
As (17) holds for all t E rN, the distribution v = Lir of our sampling algorithm's output satisfies, since T = 1
= -0 - log (G)) - Ö( x poly(d))
π u||TV= |,yy, - <|ITV + 1|Cr, U1/TV 0o
E Ec
Step 5: Bounding the runtime. Since our accuracy bound requires T = log(dC/s), and requires a time-step size of A ed S poly(d) the number of iterations is bounded by 1" cdT S O (poly(d) x log (dC/=))
Step 6: Extension of sampling guarantees special orthogonal group. Similar techniques be used in the case of the special orthogonal group. However, in the case of the special orthogonal group encounter the additional challenge that, with high probability (111, the gaps between neighboring eigenvalues (++ - (t) may become exponentially small in d, over very short time intervals of length O ((22), Over these intervals, our diffusion moves at exp(d) velocity. Despite this, we show that a 1 1/poly(d, (d. step size is sufficient to simulate random solution to its SDE with probability distribution 5-close to the correct distribution. Specifically, from the matrix calculus formula for one can show that the SDE for the eigenvectors of the forward diffusion satisfy (these evolution equations, discovered by Dyson, are referred to as Dyson Brownian motion (Dyson, 1962))
(t)(t) = ddB((t) + Σ dt 1+1 (t(t) (t))
(18)
(t) Jit (t) (t) Ut(t)
dt jyti - (t))) ( Vi E [n].
dt;(t) =
Σ
(19)
2
Roughly speaking, this implies that only the interactions in (19) between eigenvectors with neighboring eigenvalues which fall below O 1,10) are significant, while interactions between eigenvectors with larger eigenvalue gaps are negligible over these short time intervals. Thus, one analyze the evolution of the eigenvectors over these short time intervals collection of separable two-body problems consisting of interactions between pair(s) of eigenvectors with a closed-form transition kernel given by the invariant measure on SO(2). For a detailed sketch of how one can extend our proof to the case of the special orthogonal group, see Appendix B.7.
B. Full proof of Theorem 2.2
In the following, we denote by p(z, y) the geodesic distance between x, y E M, and by (v)+y(v) the parallel transport of vector E Ta from y. For convenience, we denote sp:(-) := (-) Recall that we have assumed that (MM is contained in ball of radius C poly We will prove our results under the more general assumption (Assumption B. म, C)), which is satisfied whenever (MM) S C.
Assumption B.1 (Bounded Support (1b) म. C)). The pushforward of (TT) of x with respect to the map 1o M - RRd has support on a ball of radius C' centered at O.
B. Correctness of the training objective functions
Lemma B.2. I and g" are solutions to the following optimization prrpiiea
(V - ((C--)
min (d)(
c-(T-1)-
S( [(v(ZTr 12
+ マ vv (32((ZTT))) ((ap))))_(((((()_))
Zo Zo= -600)
[[(] -
Zo -6
min GEC(RA Rdxd) Ent Bows
where J.p denotes the Jacobian of sp
15

---

Efficient Diffusion Models for Symmetric Manifolds
Proof. Step 1: Obtaining an expression for the reverse diffusion SDE in R We cannot in general directly apply (1) to obtain a tractable expression for the SDE of the reverse diffusion Y, in M, since we do not have tractable formula for the transition kernel Pr of the forward diffusion X₁ on M. Instead, we will first obtain an SDE for the reverse diffusion of Ze in R", and then "project" this SDE onto M. Let H₃ := Zr-c denote the time-reversed diffusion of Z,. H2 is diffusion in R From (1), we have that the SDE for the reverse diffusion He on Rd is given by the following formula:
dH, 1(2 (111, + 2 log 1o )
dt + dW,,
(20)
where W3 is a standard Brownian motion on IRd Equation (20) can be re-written as
dH₁ = (2 H. 20 (11H,) V log (T- di + dW,
(21)
The r.h.s. of (21) is tractable since we have a tractable expression for the transition kernel @r-ejo (it is just a time re-scaling of Gaussian kernel, the transition kernel of Brownian motion).
Step 2: Obtaining an expression for the reverse diffusion SDE in M. Note that there exists a coupling between Z
and H, such that H, = - and that Y, = XT- for all & E [o, T). Thus, under this choice of coupling, we have that Y, = x- = +z) = (HH)) for all & E [o, Tj. In the special case when there is only one datapoint zp. the SDE for the reverse diffusion Y, on MM can be obtained by applying Itô's lemma (Lemma 3.1) D) to Y, = (H1)):
1 d = V.,, T dH, + 2 (dH (H:)dH
Vi E [d]
(22)
In the following, simplify notation, drop the "t" index from the notation (ps and dYj Unfortunately, the r.h.. of (22) is not (deterministic) function of Y, o0( H)), since is not an invertible map. To solve this problem, we can take the conditional expectation of (22) with respect to Y; = (HH)):
dy = E E(dY Yj] = EldY, (H.)) = E (H)) dH, + { (VH))((
(23)
The drift term on the r.h.s. of (23) is deterministic function of Y,. Denote this function by f :M x [0,7] + TM for any input x E . and output in the tangent space T.M of M at x.
Moreover, by (1), the covariance term on the r.h.s. of (23) must be the same as the covariance term for the forward diffusion Y, on M. This covariance term can be obtained from the covariance term dW, on R4, via Itô's lemma, implying that the covariance term is ((H, [[] dW, (p(H), i [d]. Using the notation for the Jacobian, this covariance term can be written more concisely as (H,, dW, (Hl)p(H))) Thus, the diffusion term is also deterministic function g" of Y, p (H,), where g" (Yi) is symmetric d x d matrix,
(H.) = g" (Y,,
(24)
where We is a standard Brownian motion on M.
Since dW, is the derivative of a standard Brownian motion in IR and dW, is the derivative of a standard Brownian motion on the tangent space of M, we have that
(25)
Thus, (23) can be expressed as:
I dY, = (H. 1
+
4 ((H)) = gf (a,,
(26)
In the more general setting when there is more than one datapoint, (26) generalizes to:
dY: E E Vo(H)) dH, (V) (HH)) Hr =
(27)
(' g" (Y.
(28)
16

---

Efficient Diffusion Models for Symmetric Manifolds
Since Y (H(H),, bring (Y,,,, l)dl and g (Y,, ()W, inside the conditional expectation:
Vip(H)) dH, + 1 (VH)))((VVHH)))(l)) -M ((H.) Hr ]]] 9- "V,, (Yi,))a.
E E
We can rewrite this as
1 Emilectta) (Pp(H)) dH, + (dH)) (V))) -f (",, Hr. HT
Bx Be(me))
g" (Ye, t.dW,
This simplifies to
ExTnT (H,(H)) aH, (())))) (",, ()dd Hr = -b = (Y,, Daw.
(29)
where the expectation is taken over the outcomes H₁. Plugging (21) into (29), and separating the drift and the diffusion terms on both sides of the equation (and noting that the higher-order differentials (dt)² and dWidt vanish), get that the drift terms satisfy
BB (H)((H))) (H. 2V
2V log (og dt (aw,) (V - (V,, t)dt H7 b) =0. (30)
Noting that (d) dl and dW.[[]dW.b]] 0 for all + j. we get
Bdwx (Vi((H.)) (H, + 2V log H b)) dt 5 (r()))((H))) (Y(,, Hy =0.
(31)
Dividing both sides by dt, we get expression for the drift term
1 E (v)(H.)) (H, 2V loggorr--ii ( - b = 0.
(32)
Finally, from (25), we have that diffusion term g satisfies
E E Jo(H)) (H. - (o (V.. ((H)) (00
Hy = =b = 0.
(33)
Step 3: Training the drift term. From (32), we have that the function f is the solution to the following optimization problem:
min Ee in (V)() ( (- 2og 1o + - SYY, -(Y.
.1l HT 6) (34)
where the inner expectation is taken over b and over the outcomes of H, time conditioned on Hr b (Note that Y, (H,, is deterministic function of H,).
Now, Hr b} has the same probability distribution as Zr = b} (and that Y2 = b} has the same probability distribution as el b)). Thus, we can re-write (34) as
min Un.. f
(V)( (Zr + 2V log qr (ZT
1 tr( tr( (((r)(())))))))))))))))))))))))) (
(X (XT-1:1 1) / Zo =6-
b] 0]
(35)
2
17

---

Efficient Diffusion Models for Symmetric Manifolds
Step 4: Training the diffusion term. From (33) we have that g" is the solution to the following optimization problem:
min Un) (H. - (o((,, 1))
H7 HT
where || is the Frobenius norm. Since H = b) has the same probability distribution as ZT Zr-oliZo b} (and
(//{ Hr = b) has the same probability distribution as Xr- 2o = b)), we can re-write (34) as
min Unif( (a. 1j) [| Jo (o(XT-,,t)) Zo b
9
B.2. Proof of Lemma B.3
In the proof of Theorem 2.2, we will the following lemma.
Lemma B.3 (Gronwall-like inequality for SDEs on a manifold of non-negative curvature). Suppose that M is a Riemannian manifold with non-negative curvature, and let p(x. y) denote the geodesic distance between any x. y E M. Suppose also that X, and X, are two diffusions on M such that
dX, = t) + ()
and
dE - - + (Xt. ()W,. where b is (itss-iis and C2(t)-Lipschitzz every time [0, Moreover assume that
[|b(z,t) - b(x, t)ll SE
and
((x,,))-(x; t)l? S
for all x E E, E [) Then there exists a coupling between X, and X, such that, for all t 2 0,
5Ea
(X,,
Xo)]
inf (E(0..] 2C1(.) + 3322())+2)
Proof of Lemma B.3. We first couple X₁ and X, by setting their underlying Brownian motion terms dW, to be equal to each other. Next, we compute the distance (x,,,, X₂) using Itô's Lemma. Letting h(r, y) = px y). by Itô's Lemma we have that
dp²((x,(1)
di(
(b(x,tt) dt (((8.:1)
X,)
1 Tr
(((xett) (((x,)))
2) (v2h(Ee, Xi)l
(x(Xe, 6(Xe,
3) dt
(x()))t (((x.t)) 0)
W. (w,)
x))XX,,,))
Therefore,
(b(Xett) (6(221)) ]
(x,, Xi) = E
Vh(X,, XT)
dt
9) (V'h(Xe.xx)) XiIl 2 t)),
(X),
(o(X,,t)
I]] dt
E
Tr
t((),t)
+0.
(36)
18

---

Efficient Diffusion Models for Symmetric Manifolds
Now, since M has non-negative curvature, by the Rauch comparison theorem we have
V,) (b(Xe.t))
< 2p(Xe; Xe) [16(X,; 1) I xx- x, ())Il
6(Xx,,)
s 20(X, x,) (16(X,;:) - x (6(X,, f) [Ib(x,,) b(x,,
s 2p(X,, x x (Ci(t)p(Xi (X,, X,) + +)-
(37)
where the last inequality holds since b Ci (ft))-ipsshhit. Moreover, since M has non-negative curvature, by the Rauch comparison theorem, also have that
(X,,() t) 9) (2)((1, (o()),)) (o(X)it)
9)
2 Th
(((,,tt) I (0 |
(11(()1)) (x,-x 1))
+ 0(x -o 1) Fr) )²
(8x.)) -Ix-=2, 1))
+ (xx,,
S p Xx) + +
(38)
Plugging (37) and (38) into (36), we have
Elp² x,) -
x,) (C.
3=2 v4 1+20 0.
(39)
dt
Hence,
d
Blp BBp (X Xi)]
I =BC11//0..,,)),//,3---0000001110 x,)] + +52
dt
= 2BC11)/(KKKK)++,(kk.,)))+0088883 (xx, X,) + +52
(2] X.]
Let E [0, be some number, and define R(t) == (X, X,)] + inf s((orl] 3c1 for all / E [0, r]. Then we have,
d R(I) (2C) + + ((),())
V20
(40)
dt
Thus, plugging (40) into Gronwall's lemma, we have, for all 2 0.
R(t) R0))66(2200(0+++3300)))+
5
+ inf sE(o.r] + +)2))))
e (000)
Thus,
5g2
Elp²(X,, X.)] + inf
sE(0.7] 20 (s) + 3C 22
5c2
(Xo: Xu)] + inf
c (010)))33)))))+
eE(0,]]
(s) 3S)
Hence, for all t 0,
5ea
E(, X.) <
Elp²(Xu, Xo)] + inf sE(U,7] 2C1 3S) +2
eld (00+
Plugging in 7 the above equation, we have, for all 2 0,
5=2
(X,, X.)l < [Blo"(Xo. Xo)] + inf .E(0.
20 (s)
3C2 +2
19

---

Efficient Diffusion Models for Symmetric Manifolds
B.3. Proof that average-case Lipschitzess holds on symmetric manifolds of interest (Lemma B.4)
Lemma B.4 (Average-case Lipschitzness). For the unitary group, Assumption 2.1 holds with L1 OOd'.s 1/3) and T/2 O(PTa ). For the sphere, it holds for In = L = O ((). For the torus, it holds for I,1 = 1.2 = 1.
Proof. For the torus, the map (xx) has Vip(ar) = La at every æ E IRd, which implies that Assumption 2.1 is satisfied for L1 L2 =1.
Sphere. In the case of the sphere, which embed via the map /b unit sphere in Rdd, can easily observe that e.g. [[v((=)] I O(1) for any outside ball of radius 2 (11) centered at the origin. As the volume of ball of radius T= is pd times the volume of the unit ball, one can use standard Gaussian concentration inequalities to show that the Omstein-UUllenbeck process Zz which is a Gaussian process, will remain outside this ball for time T with probability at
1
least 1 - -
Moreover, by standard Gaussian concentration inequalities (Rudelson & Vershynin, 2013), we have that 1 < [ with probability least 2a for all & [0, This motivates defining the set Q {a E Rd (4 TT < ||=|| < log {)), as we then have
P(Z. VIE
Since ||=|| (4 ) for any Ω and any t e [), must have that
31
[V<(<(,,A))l2-2
3
=h,
aT
$1
d Vp(2(U,A))
3
4
L1,
du
aT
2-2
1
22 (2((U,,))
8-2 3
4
L2
aT
d V)(((U, A))
1
3
4
T2
du
aT
2-22
d (z(U, A))
S |||||
do
2-+2
Unitary group. We next show that the Lipschitz property holds for the unitary group U(n). Similar techniques can be used for the case of the special orthogonal group, and we omit those details. We first recall results from random matrix theory, which allow us to bound the eigenvalue gaps of matrix with Gaussian entries. Specifically, these results say that, roughly speaking, if Zo E C= is any matrix and Z, = Zo + B(t), where Btt) is matrix with (complex) iid N (0, t) entries undergoing Brownian motion, one has that the eigenvalues (1()) 2 ... 2 (t())) of Z + Z satisfy (see e.g. (Anderson et al., 2010; Mangoubi & Vishnoi, 2023; 2025))
P 4Ello.27
inf (N+z((t) (t)) S poly(d) $ 5(sz)
Vs 2 0.
(41)
Thus, if we define 24 to be the set of outcomes of Z, such that w41(t) - (t(t) S a²
poly(n)ve we have that P(Z,
Ωr Vt E [to. TT) 2 1 1-a a.
From the matrix calculus formulas for Vv(UTAU)) d au VipUUT AU), VP AU), and av (UT AU), we have that,
20

---

Efficient Diffusion Models for Symmetric Manifolds
for all =(U,A) UAUT E El2
[VPp)(U,A))]22-2
E 14+1 A
S d Vta d L1
| du
d
V<(((U,AA))
$ ||||2-2
1
A++1 X;
2-+2
1 ())
x C+ Vog
Aitl - di S Vta = Li,
d
Σ
S &ta $ = L2
2-12
(A++1 - 1,)
i=1
[o()(U;A\)
$ ||All2-2
1
dt
(1i+1 - 2;
2-4
1=1
C+ Tod
(1)
()++1 2,)2 S &daa =1,,
c (=(UAA))
d
S ||||2--2
2-2
since Ait1 A a vavi
1
for all i E [d] and || 2VTdlog( whenever =(U, S
B.4. Proof of Lipschitzness of f and g" on all of M (Lemma B.6)
Recall that we denote by qr) (y/zz) the transition kernel of the Omstein-Uhlenbeen process Ze at any 2, y E Rd, and by
q: (xx) JM (x) the distribution of Zr any time O. We will use the following Proposition of (Chen et al., 2023b):
Proposition B.S (Proposition 20 of (Chen et al., 2023b)). Suppose that ()))) has support on a ball of radius C > O. For any a >> O, define the "early stopping time" lo min( g 2)). Then the drift term V log 4-()) of the reverse diffasion SDE in Euclidean space is O( az dC2 (Cin(C, vd)²)) Liipchizz at every time > to Moreover, W2(Ato: T) S a.
Recall that we denote by (u) the parallel transport vector from E to y. Lemma B.6. Suppose that Assumption (,p, L1, L2, a)) and Assumption B. B.1 π, C') both hold. Then every t E [to, TT. tx, - I - (f (x. 1)) s C x p(x,8),
Vx,y E M
(42)
and
(1g*(u.)) - (o 1) s C x p(x,1)
Vr,y M
(43)
where = (C + VTdlog( 22) L3 x L + (C + [ ))² La x L2 and to == min(g. o)), and Lg
O a d (min(C, V))
Proof. Recall that (when, e.g., M is one of the aforementioned symmetric manifolds) we may decompose any E Rd as = z(U, A) where U e M.
We have the following expression for (U,
/
(U,,)) cc
AEA
x (
21

---

Efficient Diffusion Models for Symmetric Manifolds
where cy = QT ((- (A) a normalizing constant. Then
d d MV,,) cu au S (VD(((U,AA)) (lgggr-(<((UAA)) + (
(
d (a x Ixed [TookcC.AN)]Tboggg-KUUU..) 2 +
(44)
For the first term on the r.h.s. of (44) we have,
Shea (VoI(((U,AA)
[oggrr + (=
CU x du NEA
x (
= - x f d
(Vo(s(t.A)) logor ( ((
NEA
aU
(uT-(E(U.A)LL(AAN
+ +ou x / (Vo(=(U.A)) (logo)(( tr ((= (J))
JEA
d du 9T ())d
cu /
( d NEA du (Vote(t.))))
+ V (((UA
x - (2(,
+ cu / (i(l:(uu))))Toooooo--( +
NEEA
x (10 (2(.))()))))(()))) (VVo(((U.A)))V (x
d
-E(UAA)NGT= dU
1 ( ) La(A)/V= v7
( (VVoacU:))) + (
+ E.(.
x (cllggn-((0..)))1((())>) = vl.
For the second term on the r.h.s. of (44) we have,
du cu c f SE ((T()))))))
d
d
2. f SeEA du
à AEA
=Q JEA
Vo (
cu x EE(U,A)~gr--
Vo Vo =U] U]
22

---

Efficient Diffusion Models for Symmetric Manifolds
and hence,
d (v JAEA
/
(vp(((ltAA)) Vooger-(((UUUU)) + ] ( (
x (
-E.(U.A)-r log (o
Be(UAA)er- [ (176(((7,A)))VVoBBr) Vlgg
+(U(V>((=((.A))) La(a) )aCAs) - u..
Thus
d (o
au (U, = -EQUAA
au
1 {(
+ ( ((((A((=(((A)))
+ (x
x Vu (
+ (E((A.))~rr- Vu (u.bggr-((((AA)))o)))))) =
x Ba(U,A)rrr ((Vo:al(t))))VVlQQ------
(45)
tr ( Ia(A) v 0
(46)
Moreover, by standard Gaussian concentration inequalities and Assumption B.1, without loss of generality we have that ||=(U, A S C + (ogl for all A) E Ωz. From Proposition B.5 we have that V PoggPrr(((=(U,, is Lg-Lipschitz where L3 = 0 d ))) and hence that
V logpr-qo(=((,
I du
(z(U, A))l2--22 X Poooppr- A))l2222
< || du (z(U,A))ll2-2 X La x ||)(U,
< Ls x A))(
$ La x (C+ ( ?
(47)
where the third inequality holds by Assumption 2.1, and the last inequality holds since |I=(U,A)|l= SC vddlgg()) + for
all zU, A) Q.. Thus, plugging Assumption 2.1 and (47) into (45), we have that
d a
(
<C + VIC+VTdlo(->))
x Ia - (C + VTdlog( ) x Lag x Lz
(48)
a
2-
Replacing 1* with 9* in the above calculation, we also get that
d (U,
(C VTdlogg 1)² I x Ln + (C + VTdlog( 1 ))2 x Log Tz
(49)
2-2
23

---

Efficient Diffusion Models for Symmetric Manifolds
Thus, (48) and (49) imply that
p(a,))
Ve,y E M
(50)
and
(y. -
2ly
Yx E M,
(51)
where C = (C++ vTdlog( x L3 41 + (C vTdlog())) x L3 x Lg
B.5. Wasserstein to TV conversion on the manifold (Lemma B.7)
Lemma B.7 (Wasserstein to TV conversion on the manifold). There is a number e S poly(d) such that for every E to, T] and any S 6 we have
|| L Sier+allTy
11Cr, CiellT
s DKL(UI || (-(Y=r)) + / Be O(e)).
(52)
Proof of Lemma B.7. Now that we have shown that f* and g" are poly (r))(dd-Ligg Lipschitz (by Lemmas B.4 and B.6), we can apply Lemma B.3 to bound the Wasserstein distance: W2 W2 Yo+r) S (p (Y,, Y,) + t 2 0, where I poly(d).
Moreover, with slight abuse of notation, we may define Wt+r to be a continuous interpolation of the discrete process ý. Applying (13) to this process, we get that, roughly, (Y(+++r ) s (p² Y) + for r 2 0. Thus, we get a bound on the Wasserstein error.
W2(YeT:: (l+r) S W2(FiT M(+)) + W2(Per.. (c+r) S (p,, Y2) + + D)e
T 720 0
(53)
Unfortunately, after times 7 >
poly(aj: this bound grows exponentially with the dimension
To overcome this challenge, we define a new coupling between Y, and Y which we "reset" after time intervals of length t by converting our Wasserstein bound into total variation bound after each time interval. Towards this end, we use the fact that if at any time the total variation distance satisfies Ce, C llv I a, then there exists coupling such that Y = Y4 with probability at least - a. In other words, w.p. 2 - a, we have p(ic+r, Y(+r) = O, and we can apply inequality (53) over the next time interval
without incurring an exponential growth in time. Repeating this process Z times, we get that |[ - ll S a x I. where the TV error grows only linearly with T.
Converting Wasserstein bounds on the manifold to TV bounds. To complete the proof, we still need to show how to convert the Wasserstein bound into a TV bound. Towards this end, we begin by showing that the transition kernel P (HH4)) of the reverse diffusion H, IRd close to a Gaussian in KL distance over short time steps Å:
(N + AVpr (H++7) ALd) ( (1+++) T
One can do this using Girsanov's 's theorem, since, unlike the diffusion Yú the manifold, the reverse diffusion in Euclidean space H; does have constant diffusion term (see e.g. Theorem 9 of (Chen et al., 2023b)).
Next, we use the fact that with probability at least 1 a the map in a ball of radius poly(a] about the point Herr has c-Lipschitz Jacobian where poly(d), and that the inverse of the exponential map exp(-) has O(1)-Lipschitz Jacobian, to show that the transition kernel Pt of Y; (dlt,) satisfies
DRL(U PEAT+D/e+r (.+++))) S (1
Ae)a aT T
T S
if we choose Å < (()),, where == expy, (N + (Y (Y + r)))))
Next, we plug in our Wasserstein bound W( 9--) & (()) into the formula for the KL divergence between two Gaussians to bound |LY Legrrta ||Tv. Specifically, noting that Lo = expa.rr (N(U++ + Áf +
24

---

Efficient Diffusion Models for Symmetric Manifolds
T), Agg(0rrr Td))),
have that
(1r +bor+)) Tr (Y +1) g +7)
det g (Y T) + w (X)r.)t))))
d + log det (Dttri +)
where = Y+++ B AV+"(Y+r; ( f(u+r,, 7)). Since with probability 2 1 aj we have 9 (Y++) t
poly(d), plugging in the error bounds IV f(Y t)ll and ||9 (Yi+.: (Y(,,,,,)))) t)lle I and the Lipschitz bounds on I* and 9" due to Lemmas B.4 and B.6, where = poly(d), we get that (KI, bigerts)) S O Thus, by Pinsker's inequality, we have
[[yYet+* --eer++alllr - 115r, fiillTV
S DKillM | (-1Ver)) (-DY+)))))))))))))
(DRa(uillf ((ec).
(54)
B.6. Completing the proof of Theorem 2.2
Bounding the accuracy. Recall that 4 is the distribution of the forward diffusion Z. in Euclidean space after time t, which is an Ornstein-Ullenbeck Uhlenbeck process. Standard mixing bounds for Masttin-Uhlennbcck process imply that
1/ - N(0, lallTV S ((CeTt)
for all (see c.g. (Bakry et al., 2014)). Thus, is sufficient to choose T log L
to ensure that 0e
Ly, - " lTv || - N(0, [a) S O(e).
As Lemma B.7 holds for all (4o 7). the distribution Lgr of our sampling algorithm's output satisfies
|| //|Tv ||Cr, - ~||lV
+ [|Cy, w/|Tv s O(s) 0 } O( poly(d)).
Bounding the runtime of the sampling algorithm. Since our accuracy bound requires T = log( (d)) and requires a time-step size A poly(d]. the number of iterations is bounded by
T s 0 poly(d)) log
d EC.) )
A
B.7. Proof sketch for extension of sampling guarantees to special orthogonal group
Similar techniques to those used in the case of the complex unitary group can be used to bound the accuracy and runtime of our sampling algorithm in the case of the real special orthogonal group. However, in the case the special orthogonal group we encounter the additional challenge that, due weaker "electrical repulsion" between eigenvalues of real-valued random matrices, with high probability 12(1) the gaps between neighboring eigenvalues 7-- (t) 7e(t) may become exponentially small in d, over very short time intervals of length 0 (())). To overcome this challenge, first note that one show that the gaps between non-neighboring eigenvalues do satisfy polynomial lower bound at every time w.h.p. (see e.g.
(Anderson et al., 2010; Mangoubi & Vishnoi, 2023; 2025)):
(1)2(1)m00) - S Tivi 1) s (8(...))
P
(55)
1E(Co,27
Moreover, one can also show that, except over at most O(n 1.5) "bad" time intervals [Ty, Ty + Ay}. each of length e.g. O (1)) the gaps between all neighboring eigenvalues are least n.0-
25

---

Efficient Diffusion Models for Symmetric Manifolds
From the matrix calculus formula for y one can show that the SDE for the eigenvectors of the forward diffusion satisfy (these evolution equations, discovered by Dyson, are referred to Dyson Brownian motion (Dyson, 1962))
(lx(1) = dBa(()) + Σ
dt
(56)
ni(l) 71))"
573
dB;,(l)
dt
du:(1) E J#L (1(1) - (t)
Uy(()
(7i(t) (t))2 ui(l) (1))a JØ
Vi E [n].
(57)
From (55), one can see that over the "bad" time intervals [as, bi]. each eigenvalue (1))) has at most one neighboring cigenvalue, say M++((), (t), with small gap (t) /+1(t) $ O v w.h.p. Roughly speaking, this implies that only the interactions in (57) between eigenvectors with neighboring eigenvalues that fall below O O(,,to) ) are significant, while interactions between eigenvectors with larger eigenvalue gaps are negligible over these short time intervals. Thus, one can analyze the evolution of the eigenvectors (41) over these short time intervals as a collection of separable two-body problems consisting of interactions between pair(s) of eigenvectors.
More precisely, using (56), one can show that over the bad time intervals [Ty, Ty + Ay}, any eigenvalue gap which falls 1 w.h.p. This is because eigenvalues
below y, also remains below n2 over a time sub-interval of length at least 52( (naya) (t)) and M+ (t) in (56) repel with an "electrical force" proportional (x())-)4))()) which implies that the eigenvalues gaps expand at rate proportional v (the stochastic term (t) also leads to the same vt expansion rate). Thus, using the evolution equations (57), one can show that, over the short bad intervals, the distribution of [u,(T, + Aj), u441 (Ty + (TT)) is polytd] -close in Wasserstein distance to the invariant (Haar) measure with respect to the action of SO(2) on (T)) (441(TT)) This is because (by the Itō isometry) the time-averaged variance in (t( over this time interval proportional to 1 A, ,,7 1T,+ (x,())-5(0))- dt fnto (vT)= dt log( . - log( = (10 8) log(n) = (1log(n)). But
the diameter of the 2-dimensional manifold (which is isomorphic to SO(2)) on which u; (t) and u,+1 (t) (approximately) lie is O()). Thus, after the time interval [T), Tj + Aj], the two neighboring eigenvectors [u,(T] + Aj). ui+1 (Tj + (T;) are (approximately) distributed according to the Haar measure with respect to the action of SO(2) on [u,(T)), (1)4((T))
) the transition kernel of both the Thus, one can show that as long as one uses a numerical step size A S O( (oooy(a))
continuous-time reverse diffusion Yú and the numerical simulation å over the time interval [Ty, + A,] will be very close
(within Wasserstein distance ,.o to the Haar measure on the action of SO(2) on (M,(T)) (T,)) As the Lipschitz
property does hold outside the bad intervals [T), Ty + Ay], the remainder of the proof follows in the same way as for the case of U(z).
C. Additional simulation details
C.1. Datasets
Given a d-dimensional Riemannian manifold M, number of mixture components & e N. points mi...., m E M and covariance matrices C,,..., : C& E Rdxdd we say that a random variable is distributed according to a wrapped Gaussian distribution with means M,,, ... my and covariances Ci, Cs (with equal weights on each component) its distribution equal that of random variable X sampled as follows:
1. Sample an index i at random from (1,...,k).
2. Sample Z ~ (0, C.)
3. Set X = expm. (Z).
where (x)...) denotes the exponential map at any point & E M.
Datasets on the Torus Ta. The synthetic dataset is sampled from a wrnppe Gaussian distribution, with mean the origin, (0..... . 0)T and covariance matrix 0.2Ig- A total of 30,000 points were sampled as the training dataset, and 10,000 were sampled as a test dataset to compute the log-likelihood of the generative model outputs.
Datasets on the special orthogonal group (O(n). The dataset is constructed by first picking 2 random means m1, m2 e SO(n) sampled from the uniform measure on the orthogonal group 50(n) (i.e., the invariant measure with respect to actions
26

---

Efficient Diffusion Models for Symmetric Manifolds
by the orthogonal group). We then sample 40,000 matrices from the wrapped Gaussian mixture distribution on (O(n) with means m my and covariance 0.2 30,000 of these matrices are used for training, and the remaining 10,000 matrices comprise the test dataset used to evaluate the C2ST score of the generative model outputs.
Datasets on the unitary group U(n). We use dataset on (n) of unitary matrices representing time-evolution operators eat of quantum oscillator. Here is time,
2m A V is a Hamiltonian, and A is the Laplacian. V is a random
potential (x) = || zoll² with angular momentum sampled uniformly on 3] and zo ~ AO 1). As A, V are dinint-admensional operators, matrices in Unn are obtained by retaining the (discretized) top-n eigenvectors of V.
C.2. Neural Network architecture, Training Hyperparameters, and hardware
Torus. In the case of the torus, the neural network architecture consists of 4-layer MLP with hidden dimension of k, with sin activation function. We set k = 512 for d < 1000 and k = 2048 for d = 1000.
The models were trained with batch size of 512, with an appropriate variance scheduler. For each model, we trained the neural networks for 50K iterations when d < 1000, and for 100K iterations when 1000.
Special Orthogonal Group. In the case of the special orthogonal group, the neural network architecture consists of 4-layer MIP with hidden dimension k 512, with a sin activation function.
For each model, the neural networks were trained for 100K iterations, with batch size of 512, and an appropriate variance scheduler.
Unitary Group. Following TDM (Zhu et al., 2025); when training each model on the unitary group, we use a more complicated neural network than on the torus and special orthogonal group, to accommodate the more complicated quantum evolution operator datasets used in our simulations on the unitary group. For both the drift and diffusion terms, let (X,, t,) be inputs into the neural network, where X; U(n) and t; E [O,T] is time. The output of the neural network is then given by
X: MMP (X.), Ambain(4))
where MLP is a 2-layer multi-layer perceptron of dimension D, MLPs is k skip-connected MLP layers of dimension D, Emb is a sinusoid embedding of dimension D, and NG denotes group normalization. In our simulations, we set h: = 8 and D = 512. For the drift term f(.-)) in each of the models, the final output is given by Xarin = projr,, (X), where projTx, U(n) is the projection onto the tangent space at X,. For the diffusion term o(.,)) in our model, the neural network outputs vector of dimension d.
For each model, the neural networks were trained for 80K iterations, with a batch size of 512, and an appropriate variance scheduler.
Hardware. Simulations evaluating sample quality on the Torus were run on an Apple MI chip with 10 cores. Simulations on the special orthogonal group and unitary group training runtime were run on a single RTX 3070 as well.
on single RTX 3070. All simulations evaluating per-iteration
C.3. Evaluation metrics
In this section, we define the metrics used in our simulations.
Log-likelihood metric. Let D {a) F2, ...,..... Fnd} } be synthetic dataset arising from target distribution with density function g. We train the generative model A on the datasct D. Next, generate points y1" ...., vn which are outputs of the trained model A. Since the points y Vd are generated independently. the likelihood of generating the points y v given the target distribution g is given by
II 9
(58)
The (average) log-likelihood of the generated points ,..... va with respect to the target density g is therefore
log (3.
(59)
27

---

Efficient Diffusion Models for Symmetric Manifolds
C2ST metric. Suppose we have two distributions P.Q. and sampled points Sp,Sq where Sp P.So ~ Q. We denote the number of sampled points by := |Sr| = | One motivation behind the Classifier Two-Sample Test (C2ST) metric (Lopez-Paz & Oquab, 2017) is to perform hypothesis test, where one wishes to decide whether reject or accept the null hypothesis P Q. By reporting the test statistic from this hypothesis test, the C2ST metric can also be used to evaluate the quality of samples generated by a generative model, we do in our simulations.
To compute the C2ST metric, construct a dataset D where
D 0|) U {1v, 1)}== (2):
and where 22 E Sp and y/ E Sq. Partition D randomly into training and test datasets Dtr and Dte where mie == 1eel denotes the number of points in the test dataset. Suppose f D + [O, 1] is a binary classifier trained on Der where f(z)) P(4 1/2;), then the value for the C2ST metric is computed as
E 1
{[1
((z))
> 2
=4
(60)
The (s..))DDe
where 1 is the indicator function. The null distribution is approximately N/ (2' amme): Then if P Q we would have f + 0.5. Whether to reject the null hypothesis can be done by performing p-value analysis using 60 and the null hypothesis. For the sake of comparison between model performance, we report the value of instead.
More specifically, the statistic & can take values greater than or smaller than 0.5, we report the value
(61)
2
2
where i is computed using 60.
C.4. Additional results
In this section, we provide additional empirical results.
Visual results the torus Ta. In 2 show points generated on the torus by the Euclidean diffusion model, the RSGM, and our model when trained on dataset sampled from a wrapped Gaussian distribution. For the 2D torus, we plot the result as a 2D scatter plot. For dinees-dd tori, for any sampled point x E Tr, the plot shows the first two angle coordinates (ıro; 11) of each generated point on the torus. We observe that, for dimensions d 2 100, our model appears to generate points that visually resemble those of the target distribution more closely than the points generated by the Euclidean diffusion model or the RSGM model.
C2ST score and visual results on the special orthogonal group (O(n). We train our model, a Euclidean diffusion model, RSGM, and TDM on a dataset sampled from a mixture of two wrapped Gaussian distributions on SO(n) for n E {3. 5,9, 12, 15). For n 2 9, our model achieves the lowest C2ST score; a lower C2ST score indicates higher-quality sample generation (Table 4).
The visual results for our model, the Euclidean model, the RSGM, and TDM are shown in 3. Here we plot the first and second entries of the first of each generated matrix in SO(n) Note that the target distribution is bimodal, yet the two are only plotting two of the matrix coordinates.
modes appear visually single mode
Additional visual results on the unitary group U(n). Visual sample generation results on U(n) were shown in Figure of Section 4 for n = 15. Figure 4, we give visual results for additional values of n.
Runtime on the torus Tr. 5 gives the per-iteration runtime of the Euclidean model, our model, TDM, and RSGM on the torus, for dimensions d E {2, 10, 50, 100, 10003. We observe that for each of these dimensions, our method has similar runtime to the Fuclidean model, whereas RSGM is roughly times slower than the Fuclidean model.
Runtime on the special orthogonal group (O)... ). 6 shows the per-iteration runtime of the Fuclidean model, our model, RSGM and TDM, on the special orthogonal group (1))), for n E (3,5, 10, 30, 50). We observe that our model's 's per- iteration training runtime remains within a factor 1.3 of the Euclidean model for all n. However, the per-iteration training runtimes of TDM and RSGM increase more rapidly with dimension and are, respectively, 51 and 66 times greater than the Fuclidean model for 50.
28

---

Efficient Diffusion Models for Symmetric Manifolds
5.0
5.0
5.0
5.0
2
2.5
2.5
2.5
2.5
0.0
0.0
0.0
0.0
0
5
0
5
0
5
5
10
5.0
5.0
5.0
II
2.5
2.5
2.5
0.0
0.0
0.0
0
5
0
5
50
5.0
5.0
5.0
II
2.5
2.5
2.5
0
0.0
0.0
0.0
0
5
0
5
100
5.0
5.0
5.0
5.0
II
2.5
2.5
2.5
2.5
d
0.0
0.0
0.0
0.0
0
5
5
5
1000
5.0
5.0
5.0
2.5
2.5
2.5
II
0.0
0.0
0.0
0
5
Target
Ours
Euclidean
RSGM
Figure Points generated by different models when training on a dataset sampled from wrapped Gaussian target distribution the torus of different dimensions d {2, 10, 50, 100, 1000}.
Table C2ST scores when training on wrapped Gaussian mixture dataset on SO(n). Lower scores indicate better-quality sample generation (range is [0.5, 1]. and 0.5 is optimal). For
9, our model achieves the best C2ST scores.
Method = 3
5
12
15
Euclidean
.51. +.01
.62. +.02
.64. 0.02
.72. 4.02
RSGM
1..1.01
.57. 4.02
.74. +.01
1.02
.90. +.02
TDM
.52. 1.01
.5.
.699 4.03
.73. #.02
.79. 4.03
Ours
.55. +.01
.6. +.02
.60. +.03
61.
.67. +.03
29

---

Efficient Diffusion Models for Symmetric Manifolds
d =
1
1
1
1
0
0
0
0
-1
-1
-1
-1
0
1
n
1
0
0
-1
-1
0
1
1
0
1
-1
O
1
0
1
-1
0
1
n 12
0
0
1
-1
1
0
1
1
0
1
0
1
0
1
1
0
1
1
1
1
15
0
0
0
-1
-1
1
c
1
0
1
-1
1
O
1
1
0
1
Target
Ours
Euclidean
RSGM
TDM
Figure 3: Points generated by different models trained on Gaussian mixture dataset on () for different values of n.
n
1
055
0.5
0.5
0.5
0.5
-05
-0.5
0.5
-0.5
0.5
Target
Ours
Euclidean
RSGM
TDM
Figure 4: Points generated on 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.
30

---

Efficient Diffusion Models for Symmetric Manifolds
Table 5: Per-iteration training runtime in seconds on the Torus. The manifold-constrrained diffusion model with the fastest
runtime is in bold; the Euclidean model is in gray for comparison. For each dimension d, our model achieves a similar runtime the Fuclidean model, whereas RSGM is roughly times slower.
Method
d =
d = 10
d =0=5 50
1
d = 1000
Euclidean ::00
0.17
016 +.00
0.17
0.18 +.01
RSGM
1.36
1.35
1.39
1.36
1.42
Ours 0.15 ..0
0.15 000 0.16 .01
0.15 0001
0.16 1.01
Table 6: Per-iteration training runtime in seconds on SO(n). The manifold-constraineed diffusion model with the fastest
runtime is in bold; the Euclidean model is in gray for comparison. Our model's 's runtime remains within factor of 1.3 of the Fuclidean model for all n. Runtimes of TDM and RSGM increase more rapidly with dimension and are 51 and 66 times greater than the Euclidean model for n2 50.
Method
11 3 3
n = 5
n1 = 10
==
n = 50
Euclidean 0.13
0.01
0.00
0.12
0.15
RSGM
0.73
0.96
1.18
2.12
9.89
TDM
0.62
0.78 +.05
1.67 #.04
2.85
7.63 1.12
Ours
0.13 +.0
0.13 .01
0.14
0.14
0.20 +.01
D. Challenges encountered when applying Euclidean diffusion for generating points constrained to non-Euclidean symmetric manifolds
The following examples illustrate why using Euclidean diffusion models to enforce symmetric manifold constraints may be insufficient.
Example 1. Consider the problem of generating points from a distribution j on the d-dimensional torus Tú = S x x S given a dataset sampled from μ. A naive approach is to map the dataset D from the torus to Euclidean space via the map v,, which maps each point on the torus to angles in [0, 2x) C RRd. One can then train a Fuclidean diffusion model on the dataset (D).
However, the map vb can greatly distort the geometry of μ. To why, let be unimodal distribution on Ta with mode cenetered near (0, ... 0). The pushforward of under consists of distribution with 2d modes, each near the 2 corners of the d-cube [0, 2n) (see Figure 5). Thus, Euclidean diffusion model needs to learn a multimodal distribution, which may be much harder than learning a unimodal distribution.
Example 2. Another example is the problem of generating samples from a distribution on the manifold SO(3) of rotation matrices. There is a natural map vi from SO(3) to IR3 which maps any M E SO(3) to its three Euler angles (a, b, c) E [-,, T] x I- 2. 3] x [-,, C R.. However, &p has singularity at b 5. which may make it harder to learn distributions with region of high probability density passing through this singularity, i may separate this region into multiple disconnected regions.
Additionally, it has been observed empirically that applying Euclidean diffusion models to generate Euler angles in 1 leads to samples of lower quality than those generated by diffusion models on the manifold SO(3); ; see e.g. (Leach et al., 2022), and (Watson et al., 2023).
E. Illustration of our framework for Euclidean space, torus, special orthogonal group, and unitary group
1. Euclidean space R the Euclidean case, our algorithm (with the above choice of yo, 1)) recovers the algorithms of diffusion models on d from prior works (e.g., (Ho al., 2020; Rombach al., 2022)). The forward diffusion
31

---

Efficient Diffusion Models for Symmetric Manifolds
v
T d
10, 2T)² S R
Figure 5: A probability density with one mode (blue) on the torus. The map v. which maps points in the d-dimensional torus Ta to Euclidean space R2, may break up the single mode on the torus into up to 2d separated modes in Rd This can make the task of learning the pushforward of the target distribution on R much more challenging than the task of learning the original target distribution on the torus, as the distribution Rd may have exponentially-ine more modes.
is the
process with SDE dz, = - Z + dB, initialized at the target distribution и, where B, is the standard Brownian motion. The training objective for the drift term fzz t) of the reverse diffusion is given by II(:T 2-be
e -- ((r-c) f(2 e)I where a point sampled from the dataset and & is a point sampled from Zr- = b) which Gaussian distributed as (be ((b--)(T-)) (1-e-(F-o)) (see Section 3.1). The number of arithmetic operations compute the training objective is therefore the same as for previous diffusion models in Euclidean space.
2. Torus Ta. For the torus, the forward and reverse diffusion of our model are the same as the models used in previous diffusion models on the torus (De Bortoli al., 2022) (Lou et al., 2024). The Forward diffusion is given by the SDE dX,- - X + dB, on the torus, initialized at the target distribution
The only difference is in the training objective function. To obtain our objective function, observe that X₃ is the projection X, = <(()) of the Ornstein-Uhlenbee, diffusion on Rd via our choice of projection map yp for the torus.
The drift term f for the reverse diffusion can be trained by minimizing the objective function || T B-p(()e 2 f(p()),t) where è ~ [be (TT-t, (er--t,,--es
0-1
Ia). Our objective function can be computed in ((d) arithmetic operations, improving by an exponential factor the per-iteration training runtime of (De Bortoli et al., 2022) which relies on an inefficient expansion of the heat kernel which requires an exponential-in-4 number of arithmetic operations to compute, and matching the per-iteration training runtime of (Lou et al., 2024) who derive a more efficient expansion for the heat kernel in the special case of the torus.
3. Sphere S.-1-
Forward diffusion. We first choose the projection map P Rdd - Sa-1 to be (p(x) = for Sa-1. and
bb : S + Rd to be the usual embedding of the unit sphere into Rd. We define our forward diffusion to be the projection X, = (z() of the Fuclidean-space Ornstein-UUlembee initialized at the pushforward (T) of the target distribution π onto Rd. Since the Ornstein-Ullenbeel distribution Z. is a Gaussian process, each sample from our forward diffusion can be computed by drawing a single sample from a Gaussian distribution and computing the projection map y once.
diffusion Z, onto the manifold M, where Z, is
The forward and reverse diffusion of our model on the sphere are different than those of prior diffusion mod- els the sphere. The evolution of our forward diffusion X, on the sphere is governed by the SDE dX, a(X,, [x,dtt + dB.) initialized at the target distribution n, where the coefficient a(t) is given by the con- ditional expectation a)X,,)) :- E
Tz. T (v()) = -x). Our forward (and reverse) diffusion has a (time-varying and) spatially-varying covariance term a (X,,, t)dB; not present in prior models (De Bortoli et al., 2022) (Lou et al., 2024). This covariance term, which accounts for the curvature of the sphere, allows our forward diffusion to be computed as a projection of Euclidean Brownian motion onto the sphere despite the sphere's non-zero curvature.
Training the model. The SDE for the reverse diffusion of our model has both a drift and a covariance term. To train a model ff for the drift term, we first sample a point from the dataset D at a random time E [o, T), and point from the Ornstein- Uhlenbeck diffusion Z, initialized at (b), which is Gaussian distributed. Next, we project this sample é to obtain a sample (p()) from our forward diffusion X, on the manifold. Finally, we plug in the point (o()), and
32

---

Efficient Diffusion Models for Symmetric Manifolds
the datapoint b into the training objective function for the drift term f. which is given by the closed-form expression
--p(b)e [a]| (I - [[a]a y c- - f()()), The model for the drift term f is trained by minimizing the 1
1
expectation of this objective function over random samples of b D and 2 Zz To learn the SDE of the reverse diffusion, we must also train model for the spatially-varyinn covariance term, which is given by d x d covariance matrix. Learning& a dense matrix model for this covariance term would require at least dl arithmetic operations. However, as result of the symmetries of the sphere, the covariance matrix has additional structure: it is a multiple a(X,, t) of the x d identity matrix. Thus, to learn this covariance term, it is sufficient to train a model a(X,, t) for a(X,,t), t). This can be accomplished by minimizing the objective function (á((()), ) Evaluating our objective functions for the drift term and covariance terms can thus be accomplished via a single evaluation of the projection map (x(x)= ||=]|
TET
z which requires (d 3)) arithmetic operations to compute within accuracy 6 > O when
generating the input to our training objective function, which is sublinear in the dimension d2 of the covariance term.
In contrast, the forward diffusion used in prior diffusion models on the sphere (De Bortoli et al., 2022) (Lou et al., 2024) cannot be computed as the projection of Euclidean Brownian motion and must instead be computed by solving an SDE (or probability flow ODE) on the sphere. This requires number of arithmetic operations, which is higher-order polynomial in the dimension d and in the desired accuracy } or ODE solver used). As their training objective function requires samples from the forward diffusion as input, the cost of computing their objective function is therefore at least a higher-order polynomial in d and & (for (De Bortoli et al., 2022) it is exponential in d, since their training objective relies on an inefficient expansion for the heat kernel which takes 2d arithmetic operations to compute).
à (the order of the polynomial depends on the specific SDE
Sample generation. Once the models fxx t) and g( t) for the drift and covariance terms of our reverse diffusion are trained, we use these models to generate samples. First, we sample point from the stationary distribution of the process Z, on RR which Gaussian distributed. Next, we project this point 2 onto the manifold to obtain point y <pz),, and solve the SDE dY; fY,, t)dt + g(Y,, t)dB, given by our trained model for the reverse diffusion's s drift and covariance over the time interval [o, Ti. starting at the initial point y. To simulate this SDE we can use any off-the-shelf numerical SDE solver. The point yr computed by the numerical solver at time T is the output of our sample generation algorithm.
4. Special orthogonal group (O(n) and unitary group U(n).
For the special orthogonal group (O()) and unitary group U(n), the forward and reverse diffusion of our model are also different from those of previous works, as our model's diffusions have a spatially-varying covariance term to account for the non-zero curvature of these manifolds. As result of this covariance term, forward diffusion computed as a projection yp of the (U(n)). This projection can be computed via a single evaluation of the singular value decomposition of a n x n matrix, which requires at most (nn)) = O(d*) ) arithmetic operations, where w 2 2.37 is the matrix multiplication exponent and d = n?2 is the manifold dimension.
be process in Rd = _ (or C= onto the manifold SO(n)
our The forward diffusion (t) E SO(n) (or (t) E U(z)) of our model is given by the system of stochastic differential equations
du:(1)= =
Σ
Ong oug (()dB
1 Σ
(e)()),
(62)
2
jE(m],iA
jEim],jzi
where (t) = E - 1o(2,) = utt) and B,y(t) = | (x(Z)) u(t) for every i,j E [m].
A model for the drift term ʃ for the reverse diffusion can be trained by minimizing the objective function || DD t))l where R is the matrix with i'th column R. - to- -
U(N
(b)),, for each 1 X for each E [n]. Here,
E [n]. and D is the diagonal matrix with i'th diagonal entry D3s = Eje(n) = be - + (1- where G is a Gaussian random matrix with i.i.d. N(0, 1) entries and UA denotes the spectral decomposition of 5*
To learn the SDE of the reverse diffusion, we must also train a model for the covariance term, which is given by a d x d = n² x n² covariance matrix. To train a model for this covariance term with runtime sublinear in the number of matrix entries 4 we observe that as a result of the symmetries of the orthogonal (or unitary) group, the covariance term in (6) is fully determined by the n² sealar terms asy (t) for i. E [n) and the x n matrix U. Thus, to learn the
33

---

Efficient Diffusion Models for Symmetric Manifolds
covariance term, it is sufficient to train model A(U, Ix*" for these n² terms, which can be done by minimizing the objective function ||A(UU, i) All?, where A is the matrix with ji,t)t entry Aij X; A for i,j [n]. and A; denotes the i'th diagonal entry of A.
The training objective function for both the drift and covariance term can thus be computed via a singular value decomposition of an nxn matrix (and matrix multiplications of nxn matrices), which requires at most (nn)) = (dd*)
arithmetic operations, where w 2.37 is the matrix multiplication exponent and d = n2² is the manifold dimension.
In contrast, the training objectives in prior works, including (De Bortoli et al., 2022) (Lou et al., 2024), require an exponential in dimension number of arithmetic operations compute, as they rely on the heat kernel of the manifold, which lacks an efficient closed-form expression. Instead, their training algorithm requires computing an expansion for the heat kernel of these manifolds, which is given as sum of terms over the d-dimensional lattice, and one requires computing roughly 24 of these terms to compute the heat kernel within an accuracy of (1).
F. Generalization to non-symmetric manifolds
While our theoretical framework and guarantees developed for symmetric Riemannian manifolds, it is natural to ask whether the approach can extend to more general geometries. In this section, we outline the minimal set of geometric and analytic conditions required for our guarantees- on runtime, simulation accuracy, and Lipschitz continuity- to continue holding on non-symmetric manifolds. We also illustrate, via a concrete example, how two of these conditions can be satisfied even on non-smon non-smooth, non- non-Riemannian domains, and discuss the challenges that arise in the absence of continuous symmetries.
Our guarantees rely on the following three key properties:
1. Exponential map oracle. An oracle for computing the exponential map on the manifold M..
2. Projection map oracle. A projection map :p: IR + .M, where d = (d)mm)M)) along with efficient computation of its Jacobian Jo (x) and the trace of its Hessian tr (=²)))))) both of which appear in our training objective.
3. Lipschitz SDE on M. The projection Y, sHH,) of the time-reversed Euclidean Brownian motion H, must satisfy a stochastic differential equation on M whose drift and diffusion coefficients are L-Lipschitz everywhere on M, with L growing at most polynomially in d. This is essential for the reverse diffusion process to be simulated accurately and efficiently.
Conditions (I) and (2) may hold even when M lacks the high degree of symmetry assumed in our main results. For example, suppose M is the boundary of a compact convex polytope K C Rd. which contains ball of radius r > 0 centered at point p. Although such polytope is not smooth manifold due to singularities at vertices and edges, its boundary is composed of piecewise flat (d - 1)-dimensional faces. Geodesics restricted to a single face are linear and computable efficiently, satisfying the spirit of property (1).
For property (2), one can define a projection y IRd -M that maps any x E Rdd the point where the ray emanating from p and passing through intersects the boundary M. This projection can be computed efficiently, e.g., via binary search or ray-casting techniques.
However, property (3) is significantly harder to satisfy in such domains. The drift of the reverse SDE projected onto M exhibits discontinuities at the vertices and diver-dimessionll faces of the polytope. In our analysis, we crucially rely on the continuous symmetries of the manifold to "smooth out" such irregularities and to prove average-case Lipschitz continuity (see the discussion on "average-case" Lipschitzness on page 6).
Even among smooth Riemannian manifolds, generalizing beyond symmetric spaces remains nontrivial. Examples include surfaces of revolution with varying curvature (e.g., tori with non-uniform cross-sections) or higher-genul manifolds such as a double torus. These lack the homogeneous structure exploited in our proofs and pose new challenges for both analysis and algorithm design. Extending our framework to such settings represents a promising direction for future research.
G. Notation
1. Tangent space T.M. Given smooth manifold M and point & E M, the tangent space at x denoted by T.M.
34

---

Efficient Diffusion Models for Symmetric Manifolds
2. Riemannian manifold. A Riemannian manifold is smooth manifold M equipped with Riemannian metric g, which assigns each point M positive definite inner product Oz TzM TzM IR.
3. Exponential map exp(x, u). Given I E M and v E T,M, there exists a unique geodesic y such that -(0) = a and (0) = v. The exponential map is defined as exp(æ, v) := Y(1), i.c., the point reached by traveling along the geodesic for unit time.
4. Parallel transport (v). For y E M and v E TaM, (v) denotes the parallel transport of along the (unique) distance-minn minamiiinnn geodesic from z to y. This transport yields vector in TyM.
5. Geodesic distance The geodesic distance p(z, v) between points I, E -M is the length of the shortest path (geodesic) connecting them on the manifold.
6. Jacobian J.. Let M + N be differentiable map between Riemannian manifolds. The Jacobian (or differential) at x E M is the linear map J T.M - TTtay/, defined by the directional derivative of ạp at x. In coordinates, (0.,), = 2z where 0z, denotes the ith basis vector of T. In the special case M = IR. and Wn = =RR", Ju is the matrix whose (i, entry is 3.,//xz..
7. Indicator fuunction Lz(z). Given set A C X, the indicator function 1A : - {0, 1 is defined by:
{ - otherwise.
1 i E A,
La(x)=
8. Total variation distance. Given two probability measures μ and on a measurable space X, the total variation distance between them is defined as
|| 1/|TV sup (m(A) v(A),,
ACX
where the supremum is taken over all measurable subsets A X.
9. KL divergence. For probability measures μ and measurable space X, with μ < (i.e., μ is absolutely continuous with respect v), the Kullback-Leibler (KL) divergence from to y is defined as
dp DKL(H (1) = / log
du (z)
du(x),
where de denotes the Radon Nikodyn derivative of μ with respect to v.
de
10. Pinsker's inequality. For any two probability measures μ and v,
|| -u-vllrv S (2DKL(H
This inequality provides an upper bound on the total variation distance in terms of the KL divergence.
11. Wasserstein distance (H, Let μ and v be probability measures on a metric space (M, p). and let E N. The k- Wasserstein distance is defined as:
WR(m,v) = inf rEQ(u..)
(xx
where ((μ,,) v) denotes the set of all couplings of and v.
12. Operator norm ||Al|2-2- Given multilinear map V12 x ... Vv + W between normed vector spaces, the operator norm is:
| 1x)12
| ==
sup (EM EV
| |
13. Partial derivative d In parameterizations of the form x = <U,, A), we write ab respect to U E M. For example, if .M- = SO(n) and zU, A) = UAUT, then this derivative corresponds to projecting UA + AUT onto the tangent space of (0(m).
d (U, A) for the derivative with
35

---

Efficient Diffusion Models for Symmetric Manifolds
H. Primer on Riemannian geometry and diffusions on manifolds
Let M be a topological space equipped with an open cover {oa and a corresponding collection homeomorphisms 6. Un + Rd. Each pair (U,, Φa) is called chart, and the collection {(U,, @a)) is referred to as an atlas for the manifold. For an optimization-orienee
overview of smooth manifolds, geodesics, and differentiability, see (Vishnoi, 2018).
We say that MM is smooth manifold if the transition maps Φg 0x² are -moooh functions on their domain for all overlapping chart pairs ot, B E A.
A real-valued function f .M- - IR is differentiable at a point x E M it is differentiable in some chart (Uas Qa) containing x. Similarly, a curve y [0,1] > .M is differentiable if du (y((t) is a differentiable curve in R for all such that y(t) E Ua.
The derivative of a differentiable curve passing through E M defines a tangent vector at x. The collection of all tangent vectors is the tangent space M, which is isomorphic R
A Riemannian manifold is pair (M.g) g) consisting of smooth manifold M and smooth function g. called the Riemannian metric, which assigns to each point E M positive-definit inner product 9x Ma x T&M + R.
By the fundamental theorem of Riemannian geometry (see, e.g., Theorem 2.2.2 in (Petersen, 2006)), there exists a unique torsion-frre affine connection on (M, g). known as the Levi-Civita connection, which enables isometric parallel transport between tangent spaces.
The Riemannian metric y induces length any differentiable curve 7 via:
length(7) = L.'' Vorcoot
The distance between two points 2, y E M defined as the infimum of the lengths of all curves joining them:
inf p(z,y) v) == (1)== length(r).
geodesic is a curve (t) such that parallel transport of the initial velocity 7((0) along 7 yields the velocity vector (t) at all times. Given any initial velocity v T,M, there exists a unique geodesic y with 70) = x and 7'(0) = The endpoint unit time defines the exponential map:
exp(2, v) = (1).
Given smooth map y - + V between Riemannian manifolds, the Jacobian (or differential) x E M is linear map Jop TaM Tp(a)/V defined as the directional derivative of yo at In local coordinates, may write:
J (O..) Ps
where 0z, is the ith coordinate basis vector in TM and yj is the jth component function of s.. In the special case M Rm and V = RR, the Jacobian becomes the matrix whose (i,j)th entry is de ozy"
Riemannian manifolds possess notion of curvature that is intrinsic to the manifold and independent of any ambient embedding. This curvature is described by the Riemannian curvature tensor, multilinear map that encodes how the manifold bends locally.
At any point x E M, the Riemannian curvature tensor is a multilinear map
Rs TaMM x TAM x T.M + T,M,
which assigns to each pair of tangent vectors u v E T.M linear operator
R(u, v) T.M + T.M.
This operator acts on third tangent vector E T.M according the formula:
R(u, v v)w-
RU [u, W,
36

---

Efficient Diffusion Models for Symmetric Manifolds
where is the aformentioned Levi-Civita connection and [u, u] is the Lie bracket of vector fields and v. Intuitively, this expression measures the failure of second covariant derivatives commute, and hence captures the intrinsic curvature of the manifold.
Let M be Riemannian manifold and let M be point. For any two linearly independent tangent vectors v E Ta the sectional curvature (u, v) is defined as the Gaussian curvature of the 2-dimensional surface M obtained by exponentiating the plane spanned by u and v at æ. Formally, the sectional curvature of the plane II span o} S TaM given by:
(R(u, e)u, u) K(u, v) |,|||||0|2 u),
where is the Riemann curvature tensor, and (...) -) is the Riemannian metric on T.M and || its associated 2-norm.
vector field (t) along geodesic 7(t) on Riemannian manifold M is called Jacobi field it satisfies the second-order differential equation:
D dia
R(J(),
where de denotes the covariant derivative along % and the Riemann curvature tensor. Jacobi fields describe the infinitesimal variation of geodesics and are used to analyze how nearby geodesics converge diverge. The behavior of Jacobi fields encodes information about the curvature of the manifold.
A fundamental result relating curvature to the behavior of geodesics is the Rauch comparison theorem. It states that the rate at which geodesics deviate from one another depends on the sectional curvature of the manifold. Formally, let M1 and M2 be two Riemannian manifolds of the same dimension, and suppose that along corresponding geodesics the sectional curvatures satisfy K1 < K2. Then, for Jacobi fields J1 (t), J2(() orthogonal to the geodesics with the same initial length and vanishing initial derivative, we have:
J (1) 2 || for all
Intuitively, this means that geodesics spread apart more quickly spaces with lower curvature. In particular, manifolds with non-negative sectional curvature constrain the divergence of nearby geodesics, fact that we use our analysis of diffusion processes on symmetric manifolds.
Given two probability measures p, on M and an integer & E N. the k- d-ssssesteein distance between μ and is:
We(uv) == inf (E(x,r)~ (X,))) meb(a.))
where ((1,,)) v) is the set of all couplings (joint distributions) with marginals μ and
Diffusions manifolds can be defined analogously to Euclidean settings, by interpreting d. infinitesimal Brownian motion in the tangent space (see (Hsu, 2002)). In particular, ItO's Lemma extends to maps u) M + V between Riemannian manifolds via the Nash embedding theorem (Nash, 1956), which ensures that any d-dimensional Riemannian manifold can be isometrically embedded in 2d
37


Reviews:
[{"review_id": "UTp1ZxohKX", "reviewer": "Reviewer_LvwE", "summary": "This work proposes a new efficient algorithm to generate symmetric manifold data, which enjoys $O(1)$ gradient evaluation and nearly $d$ arithmetic operations (exactly $d$ for sphere and torus data) and significantly improves previous results. The main intuition is to take advantage of Ito Lemma and the projection operator to allow the use of closed-form score on the Euclidean space. They also use synthetic experiments to support their theoretical results.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your valuable comments and suggestions. We are glad that you appreciate that our method significantly improves previous results, the intuition behind the algorithm design, and the technique novelty. We answer your specific questions below.\n\n>  It would be better to add a notation part at the beginning of the appendix.\n\nThank you for the suggestion. We will add a notation section at the beginning of the appendix in the final version.\n\n>  It seems that this work does not introduce the definition of $J_\\varphi$\n\nWe appreciate the opportunity to clarify. The Jacobian $J_\\varphi$ is the matrix whose $(i,j)$-th entry is the partial derivative $\\frac{\\partial \\varphi_i(x)}{\\partial x_j}$. While this is briefly mentioned at the top of page 5 (end of the first paragraph in the right column), we agree that the definition can be made more explicit. We will revise the explanation on page 5 and also include a formal definition in the new notation section.", "claims_and_evidence": "Yes.", "methods_and_evaluation_criteria": "Yes.", "theoretical_claims": "I partly check the correctness of the proof (proof sketch) for theoretical claims.", "experimental_designs_or_analyses": "I have checked the soundness of the analysis and experimental designs.", "relation_to_broader_scientific_literature": "Previous works on the manifold data either have $O(d)$ gradient evaluation or exponential $O(d)$ arithmetic operations. In this work, they improve these two terms at the same time.", "essential_references_not_discussed": "No.", "other_strengths_and_weaknesses": "Strengths 1: The intuition behind the algorithm design is helpful, and the technique novelty is clear (Sec. 3.2.).\n\nWeakness 1: It would be better to add a notation part at the beginning of the appendix.", "other_comments_or_suggestions": "Comment 1: It seems that this work does not introduce the definition of $J_{\\varphi}$", "questions_for_authors": "Please see the weakness part.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "UYxZFMeD9z", "reviewer": "Reviewer_zeqM", "summary": "This paper introduces a new method for producing scalable diffusion models on Riemannian manifolds with certain symmetries. The method is constructed by placing a diffusion process on a Euclidean space that can be projected entirely onto the manifold, with a partial inverse.\n\nThe training speed of this algorithm is significantly faster than the current best methods for applicable manifolds. It significantly improves on the runtime growth rate with manifold dimension, for special cases to be the same rate as Euclidean diffusion models.\n\nThe authors contribute several pieces of new theory to prove results regarding the runtime and accuracy of the model and the training objective.\n\nExperimental validation proves out the theory by showing excellent results in high dimension scaling.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your valuable comments and suggestions. We are glad that you appreciate the significant runtime improvement, our contribution of several pieces of new theory, and excellent experimental results in high dimension scaling.  We answer your specific questions below.\n\n\n> One missing experiment would be on real data of some kind - for example the earth sciences dataset typically used for Riemanian diffusion models, or something like the the experiments on molecule generation in papers such as Torsional Diffusion [1]. I suggest this as usually real data is significantly more mode-concentrated than synthetic data and it is useful to see if methods can handle this. I do not see this as a barrier to publication, but practitioners may find it more convincing to use the method if such experiments are included.\n\nThank you for this helpful suggestion. We agree that evaluating our method on real-world datasets would enhance its practical relevance.\n\nWe considered the datasets you mentioned. Our theoretical results and experiments on synthetic datasets indicate that our method yields greater gains in sample quality and training runtime as the dimension of the manifold increases—for example, on tori with $d \\geq 10$, and on $\\mathrm{U}(n)$ and $\\mathrm{SO}(n)$ for $n \\geq 9$ (see Table 1 and Figure 1). As such, low-dimensional datasets (e.g., on $\\mathbb{S}^2$) are less likely to reveal the advantages of our approach.\n\nIn contrast, the GEOM-DRUGS dataset from the paper *Torsional Diffusion for Molecular Conformer Generation* [1] offers a promising setting. It consists of molecules whose torsion angles can be represented as points on tori of varying dimensions (average $d \\approx 8$, with some above $d = 30$), provided one first applies a preprocessing model [1] that infers torsion angles from 3D molecular structures. However, one cannot directly apply our framework to this dataset, as the torsion angles of different molecules lie on tori of different dimensions. Applying our model to this dataset would require extending our framework (and code) to handle a union of tori of varying dimensions, and adapting it to perform conditional generation based on molecular graphs.\n\nWe are actively exploring this direction and will include a discussion of such extensions and limitations in the final version.\n\n> Not essential, but I would have like to see a mention of [1], as they also develop a model for placing density on the torus with runtimes on the same order as Euclidean space. \n\nThank you for pointing us to this work. We will include a discussion of this work in Section 1 (Introduction) and Section 2 (Results) of the final version.\n\n\n> I found the presentation of sections 2 & 3 a little confusing... I would perhaps reorder these section to present the method first, followed by the results and then the proof highlights. \n\nThank you for this feedback. We will revise the paper so that the algorithmic details precede the presentation of results.\n\n\n> I was also unsure about the use of the word oracle throughout the paper to describe what to me are just known functions, such as the exponential map and the projection maps...\n\nYou are right — the term \"oracle\" is not necessary here, since the exponential and projection maps are known and computable in closed form for the manifolds considered. We will remove the term in the final version.", "claims_and_evidence": "I believe all the claims are well supported.", "methods_and_evaluation_criteria": "They do.", "theoretical_claims": "I checked some of the derivations, that of the loss, the training and sampling algorithm.\n\nThe results on run time and accuracy are beyond my expertise.", "experimental_designs_or_analyses": "The evaluation section is well conceived and demonstrates the advantage of the model well.\n\nOne missing experiment would be on real data of some kind - for example the earth sciences dataset typically used for Riemanian diffusion models, or something like the the experiments on molecule generation in papers such as Torsional Diffusion [1]. I suggest this as usually real data is significantly more mode-concentrated than synthetic data and it is useful to see if methods can handle this. I do not see this as a barrier to publication, but practitioners may find it more convincing to use the method if such experiments are included.\n\n[1] https://arxiv.org/abs/2206.01729", "relation_to_broader_scientific_literature": "The methods here relate to the body of work modeling densities on Riemannian manifolds.Typically these methods have struggled to scale to higher dimensions outside of the special case of the torus. This method significantly improves on prior work on the subset of manifolds it applies to.", "essential_references_not_discussed": "Not essential, but I would have like to see a mention of [1], as they also develop a model for placing density on the torus with runtimes on the same order as Euclidean space. This method is clearly different and applicable to other settings however.\n\n[1] https://arxiv.org/abs/2206.01729", "other_strengths_and_weaknesses": "I found the presentation of sections 2 & 3 a little confusing. I found it strange to present the results ahead of the derivation of the algorithm and found myself constantly referring forwards. I would perhaps reorder these section to present the method first, followed by the results and then the proof highlights. This is just a suggestion however.\n\nI was also unsure about the use of the word `oracle` throughout the paper to describe what to me are just known functions, such as the exponential map and the projection maps. Could the authors explain why it is needed to term these maps this way?", "other_comments_or_suggestions": "None", "questions_for_authors": "Please see the relevant sections of the review.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Y10pf54XgT", "reviewer": "Reviewer_q8R3", "summary": "To improve the efficiency and accuracy of diffusion model on manifold, this work first defined a novel diffusion process on a so-called symmetric manifold by applying the projection map with mild smoothness condition. For the reverse process, instead of considering the manifold's head kernel that has no closed form, they design a new training objective to obtain the drift and covariance term. Furthermore, they obtain the complexity bound of the Wasserstein distance between the generated distribution and target distribution by using the comparison theorem for Riemannian manifolds, which is bounded by $O(\\text{poly}(d))$.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your valuable comments and suggestions. We are glad you appreciate the novel approach to diffusion on manifolds, and our significant improvement in computational efficiency. We answer your specific questions below.\n\n>…motivation of the symmetry property of $\\mathcal{M}$?\n\nThank you for this thoughtful question. The symmetry property, together with the average-case Lipschitz property for the map $\\varphi$ (Assumption 2.1), allows us to show that the SDE for the projected reverse diffusion $Y_t=\\varphi(H_t)$ on $\\mathcal{M}$ satisfies a Lipschitz property at *every* point on $\\mathcal{M}$ (We prove this in Lemma C.6; see also the paragraph *Showing \"average-case\" Lipschitzness* in Section 3.2). This everywhere-Lipschitz property in turn allows us to bound the numerical error of our algorithm's SDE solver.\n\nRoughly, the average-case Lipschitz property (Assumption 2.1) says the projection map $\\varphi$ satisfies a Lipschitz condition which holds on a subset of “average-case” points $\\Omega_t\\subset\\mathbb{R}^d$ which contains the Euclidean-space forward diffusion $Z_t$ with high probability. Moreover, $\\Omega_t$ satisfies a symmetry property: the indicator function $1_{\\Omega_t}(x),$ which determines if a point $x\\in\\mathbb{R}^d$ is in $\\Omega_t$, is independent of the projection $U:=\\varphi(x)$ (e.g., for the sphere, $1_{\\Omega_t}(x)$ depends only on the radial magnitude $\\Lambda=||x||$ and not on the projection $U=\\frac{x}{||x||}$ onto the sphere).\n\nWe show Assumption 2.1 holds for symmetric manifolds studied in our paper (see the paragraph following Assumption 2.1, and Lemma C.4).\n\n>I don't understand $z=z(U,Λ)$…\n\nWe appreciate the opportunity to clarify. The dimension $d$ of the Euclidean space $\\mathbb{R}^d$ is within a small constant factor of the dimension of the manifold $\\mathcal{M}$ ($d=O(\\mathrm{dim}(\\mathcal{M})$). The dimension of the Euclidean space $\\mathbb{R}^{d-\\mathrm{dim}(\\mathcal{M})}$ that contains the $\\Lambda$'s, and the dimension of $\\mathcal{M}$, add up to $d$. In our paper we sometimes abuse notation and refer to the manifold's dimension as $d$ rather than \"$O(d)$\", as this does not change the runtime bounds beyond a small constant factor. We will clarify this.\n\nWhen $\\mathcal{M}=\\mathrm{SO}(n)$, each element of $\\mathrm{SO}(n)$ is an $n\\times n$ orthogonal matrix. The map $\\varphi:\\mathbb{R}^d→\\mathcal{M}$ takes as input an $n\\times n$ upper triangular matrix, and outputs an $n\\times n$ orthogonal matrix in $\\mathcal{M}=\\mathrm{SO}(n)$. As each upper triangular matrix has $\\frac{n(n+1)}{2}$ nonzero entries, the space of upper triangular matrices (in vector form) is $\\mathbb{R}^d$ with $d=\\frac{n(n+1)}{2}.$\n\nThe dimension of $\\mathcal{M}=\\mathrm{SO}(n)$ is $\\mathrm{dim}(\\mathcal{M})=\\frac{n(n+1)}{2}-n$ (as $n\\times n$ orthogonal matrices have $\\frac{n(n+1)}{2}-n$ degrees of freedom). As $\\Lambda$ are the diagonal entries of an $n\\times n$ diagonal matrix, $\\Lambda\\in\\mathbb{R}^n$. Thus, $\\Lambda\\in\\mathbb{R}^{d-\\mathrm{dim}(\\mathcal{M})}$, as $d-\\mathrm{dim}(\\mathcal{M})=\\frac{n(n+1)}{2}-(\\frac{n(n+1)}{2}-n)=n$.\n\n>how to check…symmetry property\n\nIn this paper, we assume the constraint manifold is known a priori and is one of several standard symmetric manifolds (e.g., sphere, torus, $\\mathrm{SO}(n)$, $\\mathrm{U}(n)$ or their direct products). This is typical in many applications—e.g., molecular data on tori or quantum evolution matrices in $\\mathrm{U}(n)$. In such applications, numerical methods are not required to verify the symmetry property.\n\n>…previous knowledge of the exponential map…\n\nIn the setting where the manifold is a symmetric-space, there are closed-form expressions which allow one to efficiently and accurately compute the exponential map. E.g., on $\\mathrm{SO}(n)$ or $\\mathrm{U}(n)$, this map is given by the matrix exponential.\n\n>In Theorem 2.2 (Line 171), $\\varphi:\\mathbb{R}^d→\\mathcal{M}$\n\nWe will fix this typo.\n\n>…(Line 173) $\\phi(\\mathcal{M})$?\n\nThe map $\\psi$, defined in Section 2, is the (restricted) inverse of $\\varphi,$ where $\\varphi(\\psi(x))=x$ for all $x\\in\\mathcal{M}$. $\\psi(\\mathcal{M}):=${$\\psi(x):x\\in\\mathcal{M}$}$\\subset\\mathbb{R}^d$ is the pushforward of $\\mathcal{M}$ w.r.t. $\\psi$.\n\n>…preliminary knowledge…appendix.\n\nWe will include a brief primer on Riemannian geometry and manifold-based diffusion in the appendix.\n\n>$‖\\cdot‖_{2→2}$ … Is it the operator norm\n\nYes, this is the operator norm (induced 2-norm) of a matrix. We will clarify this in the final version.\n\n>…what $\\frac{d}{dU}\\varphi(x)$ and $\\frac{d}{dU}x$ exactly mean.\n\nIn the decomposition $x=x(U,\\Lambda)$, we define the partial derivative $\\frac{d}{dU}x(U,\\Lambda)$ as the derivative of the parameterization with respect to $U\\in\\mathcal{M}$. For instance, when $\\mathcal{M}$ is the special orthogonal group $\\mathrm{SO}(n)$, we have $x(U,\\Lambda)=U\\Lambda U^\\top$, and the derivative corresponds to projecting $U\\Lambda+\\Lambda U^\\top$ onto the tangent space of $\\mathrm{SO}(n)$.", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "Yes", "theoretical_claims": "- (Line 179-192) What is the motivation of the symmetry property of  $\\mathcal{M}$? Here, I don't understand $z = z(U,\\Lambda)$ and $\\Lambda \\in \\mathcal{A}$ with $\\mathcal{A} \\subset \\mathbb{R}^{d - \\text{dim}\\mathcal{M}}$. For the example of $SO(n)$,d $z = U\\Lambda U^*$ but $\\Lambda \\in \\mathbb{R}^d$ not in $\\mathbb{R}^{d - \\frac{d(d-1)}{2}}$. Or you mean $\\mathcal{A} \\subset \\R^d$ is a $d - \\text{dim}\\mathbb{R}$ submanifold? Does this concept come from the concept of symmetric space?", "experimental_designs_or_analyses": "Yes", "relation_to_broader_scientific_literature": "This work provides a novel approach to consider the diffusion on manifolds. Unlike the previous works that requires the heat kernel on manifolds, which has no closed form so that it requires a lot of calculations, their forward process, reverse process, and training strategy are much more computationally efficient. So it has the potential to be applied to unknown data manifolds.", "essential_references_not_discussed": "No.", "other_strengths_and_weaknesses": "### Strength:\n- This work provides another approach to consider the diffusion on Riemannian manifold by training the drift and variance terms to avoid the problem of computing the manifold's heat kernel which has no closed form, so it improves the efficiency significantly.\n\n- The assumptions are not too strict, because it does not require the smoothness of the construction maps $\\phi \\colon \\mathbb{R}^d \\rightarrow \\mathcal{M}$, $\\psi \\colon \\mathcal{M} \\rightarrow \\mathbb{R}^d$. So it may be applied to considering the learnable $\\phi, \\psi$ for unknown manifold $\\mathcal{M}$.\n\n### Weakness:\n- In general, the problem is how to check a data manifold satisfying the so-called symmetry property. \n\n- The sampling algorithm needs the previous knowledge of the exponential map, which may prevent the algorithm from being applied to unknown data manifold. The problem are whether we can train an exponential map and how the error effects the final accuracy.", "other_comments_or_suggestions": "- In Theorem 2.2 (Line 171), $\\varphi \\colon \\mathbb{R}^d \\rightarrow \\mathcal{M}$, and (Line 173) $\\phi(\\mathcal{M})$?\n- I thought it is better to provide a brief preliminary knowledge of geometry and diffusion on manifold in the appendix part.", "questions_for_authors": "- About the Assumption 2.1, what is the norm $\\Vert\\cdot\\Vert_{2 \\rightarrow 2}$ meaning? Is it the operator norm of a matrix, because I thought $\\nabla \\phi(x) \\in \\mathbb{R}^{d \\times d}$ is a matrix when extending $\\phi \\colon \\mathbb{R}^d \\rightarrow \\mathcal{M}$ to $\\phi \\colon \\mathbb{R}^d \\rightarrow \\mathbb{R}^d$? Furthermore, I am not sure what $\\frac{d}{dU}\\varphi(x)$ and $\\frac{d}{dU}x$ exactly mean.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "qGW9MJmfxB", "reviewer": "Reviewer_xdLt", "summary": "The paper introduces a new framework for designing efficient diffusion models on symmetric manifolds, including torus, sphere, special orthogonal group and unitary group. The paper incorporates a spatially varying covariance structure that allows efficient training without computing the manifold heat kernel. In addition, the forward process involves a projected Euclidean Brownian motion, and thus avoids costly numerical solvers.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your valuable comments and suggestions.  We are glad that you appreciate that the motivation and findings of this work is significant, and the benefits of our methods in terms of efficiency and sample quality.  We answer your specific question below.\n\n> Could the authors comment on how to generalize the method to non-symmetric manifolds?\n\nThank you for this important question. Our theoretical guarantees for runtime and sampling accuracy rely on the following three properties:\n\n1. *Exponential map oracle:* An oracle for computing the exponential map on the manifold $\\mathcal{M}$.\n\n\n2. *Projection map oracle:* A projection map $\\varphi: \\mathbb{R}^d \\rightarrow \\mathcal{M}$, where $d = O(\\mathrm{dim}(\\mathcal{M}))$ , which can be computed efficiently, along with its Jacobian $J_\\varphi(x)$ and the trace of its Hessian $\\mathrm{tr}(\\nabla^2 \\varphi(x))$, as these appear in our training objective.\n\n\n3. *Lipschitz SDE on $\\mathcal{M}$:* The projection $Y_t = \\varphi(H_t)$ of the time-reversal $H_t$ of Euclidean Brownian motion should satisfy a stochastic differential equation (SDE) on $\\mathcal{M}$ with drift and covariance terms that are $L$-Lipschitz at every point on $\\mathcal{M}$, with $L$ growing at most polynomially in $d$. This ensures accurate and efficient simulation of the reverse diffusion process.\n\nConditions (1) and (2) can, in principle, be satisfied even when $\\mathcal{M}$ is not a symmetric manifold. For example, consider the setting where $\\mathcal{M}$ is the boundary of a compact convex polytope $K \\subseteq \\mathbb{R}^d$ as a structured non-symmetric domain. $K$ is assumed to contain a ball of some small radius $r>0$ centered at some point $p$. While not a Riemannian manifold due to non-smooth points at vertices and edges, the polytope boundary $\\mathcal{M}$ is composed of piecewise flat $(d-1)$-dimensional faces. Geodesics within faces are linear and can be computed efficiently, satisfying the spirit of (1).  For (2), one can define the projection $\\varphi: \\mathbb{R}^d \\rightarrow \\mathcal{M}$ to be the map which maps a point $x \\in \\mathbb{R}^d$ to the intersection of the ray $\\ell$ with the polytope boundary $\\mathcal{M}$, where $\\ell$ is the ray extending from the center $p$ and passing through $x$.  This projection can be computed efficiently, for example via binary search.\n\nHowever, condition (3) is harder to guarantee in such cases. The drift of the projected reverse SDE has discontinuities at the vertices (and lower-dimensional faces) of the polytope. In our paper, we rely on continuous symmetries of the manifold to \"smooth out\" such irregularities and prove the necessary Lipschitz properties (see the paragraph \"Showing 'average-case' Lipschitzness\" on page 6). \n\nEven among smooth Riemannian manifolds, generalizing beyond symmetric spaces is nontrivial. Examples include surfaces of revolution with non-uniform curvature (e.g., a torus with a varying cross-section) or higher-genus compact manifolds (e.g., a double torus), which lack the high degree of symmetry leveraged by our current analysis.\n\nExtending our framework to such settings is a compelling and challenging direction for future research. We will include this discussion in the final version of the paper.", "claims_and_evidence": "The proposed diffusion model based on manifold projection for symmetric manifolds is claimed to be efficient and there exists evidence both in terms of complexity analysis and runtime comparison in experiments. The method seems also improve sample quality on synthetic datasets compared to prior manifold-based diffusion models.", "methods_and_evaluation_criteria": "The proposed methods are based on properties of symmetric manifolds and appear to be correct. The evaluation is standard as in existing works.", "theoretical_claims": "Did not closely check the proofs.", "experimental_designs_or_analyses": "Experiments seem to verify the benefits of the methods in terms of efficiency and sample quality.", "relation_to_broader_scientific_literature": "Given the need for efficiency and scalability in generative models, the motivation and findings of this work is significant.", "essential_references_not_discussed": "NA", "other_strengths_and_weaknesses": "Because I am not an expert in the field of manifold diffusion models, my judgement on the paper may not be less confident. In general, I found the paper to be well-written and the method has merits even though the idea may seem natural (as projection and its inverse are used for efficient training and sampling).", "other_comments_or_suggestions": "NA", "questions_for_authors": "1. Could the authors comment on how to generalize the method to non-symmetric manifolds?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper introduces a new method for producing scalable diffusion models on Riemannian manifolds with certain symmetries. Previous works in this direction have O(d) complexity of gradient evaluation and around O(d) arithmetic operations. In this work, the authors improve these two terms at the same time.

This paper was reviewed by the four reviewers. The reviewers agreed that the the proposed diffusion model is efficient and there exists evidence both in terms of complexity analysis and runtime comparison in experiments. The method seems also improve sample quality on synthetic datasets compared to prior manifold-based diffusion models. 

Based on their 3 "accept"s and 1 “weak accept”, the decision is to recommend the paper for acceptance. 

The reviewers did raise some valuable concerns that should be addressed in the final camera-ready version of the paper. The authors are encouraged to make the necessary changes to the best of their ability.

In particular, 
- experiments on real world data, demonstrating properties of the method, should really help to promote the method among practitioneers (reviewer zeqM)
- it could be good first to present the algorithm, and then the results concerning its properties, but not vice versa (reviewer zeqM)
- discuss motivation for the symmetry property (reviewer q8R3)
- possibility to generalize to the case of non-symmetric manifolds (reviewer xdLt)

 We congratulate the authors on the acceptance of their paper!
Decision: Accept (poster)