Quotient-Space Diffusion Models
Authors: Yixian Xu
Yusong Wang
Shengjie Luo
Kaiyuan Gao
Tianyu He
Di He
Chang Liu
Venue: ICLR 2026 Oral (2026)
Keywords: Diffusion Models
Generative Modeling
Geometric Deep Learning
Molecular Structure Generation
Primary area: generative models
TL;DR: We propose a principled way to leverage group symmetry of the target distribution by defining a diffusion model on the quotient space, which achieves both easier learning and correct sampling for the first time.

Abstract: Diffusion-based generative models have reformed generative AI, and have enabled new capabilities in the science domain, for example, generating 3D structures of molecules. Due to the intrinsic problem structure of certain tasks, there is often a symmetry in the system, which identifies objects that can be converted by a group action as equivalent, hence the target distribution is essentially defined on the quotient space with respect to the group. In this work, we establish a formal framework for diffusion modeling on a general quotient space, and apply it to molecular structure generation which follows the special Euclidean group SE(3) symmetry. The framework reduces the necessity of learning the component corresponding to the group action, hence simplifies learning difficulty over conventional group-equivariant diffusion models, and the sampler guarantees recovering the target distribution, while heuristic alignment strategies lack proper samplers. The arguments are empirically validated on structure generation for small molecules and proteins, indicating that the principled quotient-space diffusion model provides a new  framework that outperforms previous symmetry treatments.

Published as a conference paper at ICLR 2026
QUOTIENT-SPACE DIFFUSION MODELS
Yixian Xu';
Yusong Wang2.6,, Shengjie Luo', Kaiyuan Gao", Tianyu He',
Di Hel! Chang Liu
State Key Laboratory of General Artificial Intelligence, Peking University, Beijing, China
2 State Key Laboratory of Human--Mcching Hybrid Augmented Intelligence, Institute of Artificial Intelligence and Robotics, Xi'an an Jiaotong University
3 Huazhong University of Science and Technology
4 Microsoft Research Asia
5 Zhongguameun Academy
ABSTRACT
Diffusion- generative models have reformed generative AI, and have enabled new capabilities in the science domain, for example, generating 3D structures molecules. Due the intrinsic problem structure of certain tasks, there is often symmetry in the system, which identifies objects that can be converted by group action as equivalent, hence the target distribution is essentially defined on the quo- tient space with respect to the group. In this work, we establish a formal frame- work for diffusion modeling on general quotient space, and apply it to molecular structure generation which follows the special Euclidean group SE(3) symmetry. The framework reduces the necessity of learning the component corresponding to the group action, hence simplifies learning difficulty over conventional group- equivariant diffusion models, and the sampler guarantees recovering the target distribution, while hcuristic alignment strategies lack proper samplers. The argu- ments are empirically validated on structure generation for small molecules and proteins, indicating that the principled quotient-space diffusion model provides new framework that outperforms previous symmetry treatments.
INTRODUCTION
Diffusion models have emerged as the dominant approach for modeling distributions in high- dimensional spaces. Building on their success real-world domains such as images (Ho et al., 2020; Song et al., 2021), audios (Kong et al., 2021; Evans et al., 2024), and videos (Ho et al., 2022; Li al., 2023), diffusion models are now increasingly adopted in scientific applications, ranging from fluid field solving (Bastek et al., 2025), electronic structure prediction (Kim al., 2025), molecular structure generation (Xu et al., 2022; Abramson et al., 2024; Hassan et al., 2024; Geffner al., 2025), and thermodynamic ensemble modeling (Zheng et al., 2024; Lewis et al., 2025).
Compared with general tasks, scientific applications often exhibit inherent symmetry structures, wherein objects that can be related through specific transformations Consider molecular structure generation as a representative example. A molecular structure can be represented vector 3R by concatenating the 3D coordinates of its N atoms. However, because the choice of coordinate system is arbitrary, vectors in R that differ only by a global 3D translation or rotation of all atoms correspond the same underlying structure. Mathematically, such transformations typically form a Lie group for example. the special Euclidean group SE(3) in the case of molecular structures, which formally characterizes the symmetry.
regarded equivalent.
The common treatment is putting the target distribution in the original space but assigning the same probability to equivalent objects, resulting in distribution that invariant under group action. This can be implemented by augmenting training data by applying randomly chosen group actions (Abramson et al., 2024), or using a group equivariant model (Xu et al., 2022; Hoogeboom et al., 2022b), which guarantees invariance if the starting prior distribution is invariant (Köhler et al., 2020). Nevertheless, we shall show that this treatment still has room to improve, as the neural network model, which is intended for updating the sample in each diffusion simulation step, still needs to
"quual contribution. Correspondence Di He <dihe@pku.edee edu.. edu..cnu cn>, Chang Liu <1 uuchang@bza edu. en>.

---

Published as a conference paper at ICLR 2026
Table 1: Comparison among different training strategies in presence of symmetry group. Learning difficulty is measured by whether the need to predict in the equivalent degrees of freedom (DoFs), induced by the group actions, is removed, and (if not) whether the variance on the equivalent DoFs is removed. Sampling compatibility means whether there is a sampler that exactly reproduces the target distribution. The denoising form of diffusion model Dg is used express the loss functions, where Ax( (Fq.. (11)) represents aligning x towards y. and A denotes treating as constant (i.e., nopp-raddinnt, The conclusions hold using either an equivariant architecture or a general architecture with data augmentation. See Sec. 3,4 for details.
Optimal
Reduction of learning difliculty
Sampling Removal of variance computibility
Training strategy for D,
solution of Dg
Removal of equivalent Dol's on equivalent Dol's
Conventional loss E)), - x
[[xx|x]]]
GeoDiff' aligmment loss (x),,( () - As. (xi)ll
AF3 alignment loss E Da(x,, 1) AD. o(xi))
g BlAx, (x1)
X
for arbitrary y C
quotient-space diffusion loss El| (D,( t)- x!]l"
E (x1))xg) + v.
for arbitrary v (er(()),,
learn a specific movement within the equivalent class (e.g.. rotating a molecular structure), which is unnecessary any such movement does not update the intrinsic system state (e.g., the shape of a molecular structure) hence is acceptable. In hope to remove this redundancy, there are a few heuristic treatments using alignment, i.e., adjusting the prediction target within its equivalent class according to a reference to remove these equivalent degrees of freedom (Xu et al., 2022; Abramson al., 2024). But we find that the corresponding sampling process becomes incompatible with such training strategies, even with heuristic fix attempts (Wohlwend et al., 2025).
In this work, we develop a principled approach building diffusion model considering the intrinsic symmetry of the system. In particular, we leverage the concept of quotient space, in which a set of equivalent objects (equivalent class) are treated as one element. It is the formal mathematical construction that reflects the intrinsic variability of the system. We first derive the diffusion process on a general quotient space based on the correspondence between the Wiener processes on the two spaces. Considering that the quotient space is generally not Euclidean. hence it is hard to directly carry out a simulation on it, we further leverage the mathematical construction of horizontal lift to induce a diffusion process back in the original space that can easily implement the quotient-space diffusion process. The resulting process effectively amounts to projecting the update vector in the original diffusion process onto the subspace that does not induce movement within the equivalent class (e.g., rotation). We show that this process guarantees producing the correct target distribution, meanwhile reduces learning difficulty by removing the necessity to learn specific movement within an equivalent class. A visualization example in the 2-dimensional plane with SO(2) symmetry is shown in Fig. 1. In this example, the lifted process only has radial movements (Fig. 1(Left) as the quotient space R²/SO(2) is isomorphic to the half real line and recovers the correct target distribution as conventional equivariant diffusion models (Fig. 1(Middle, Right)). A conceptual comparison with existing methods is shown in Table I. The quotient-space diffusion admits either an equivariant model or general model with data augmentation.
As a representative application, we deduce the specific training and sampling algorithms in the RAN /SE(3) scenario for molecular structure generation, which relaxes the model from learning a translation and rotation movement, while the sampling process keeps the structure with constant position and orientation. We study the empirical performance of quotient-space diffusion models on small molecule structure generation and protein backbone design tasks. The results show that our methods can consistently improve the generation performance in these applications over con- ventional equivariant diffusion models and using alignment strategies. Our method achieves 9%- 23% relative improvements of ET-Flow (Hassan et al., 2024) on GEOM-QM9 and GEOM-DRUGS datasets, surpassing previous heuristic alignment methods. For the protein structure generation task, our method surpasses the state-of-the-ar Proteína model (Geffner et al., 2025) with the same param- eter scale (60M) in large margin and also outperforms the much larger model (200M) on most key distributional metrics.
2

---

Published as a conference paper at ICLR 2026
Trajectory Comparison
Conventionel Equivarient Diffusion
Quotient-space Diffusion
Quster space titeee
Goraanann Fakarort Ditacin.
Generated Sumples
Generated Samples
Oustier-pue quics
Goesanionaal Eokkarert Ditaiion
Ouzten suece Binaes
10
A motivative illustration highlighting the behavior of the quotient-space diffusion model against the conventional equivariant diffusion model for modeling a distribution on R² (as M) with S0(2) (as 9) symmetry, whose density is represented by the gray scale. (Left) SDE sampling trajectories by the two diffusion models. The same color indicates the same starting point (the round dot). The quotient-space diffusion model moves each sample only along the ray from the origin, which can be understood as only traversing the quotient space R (2)SO(2), i.e., traversing over engiineeenered concentric circles, without moving within an equivalent class, i.e., an origin- centered circle. The conventional equivariant diffusion model moves each sample over the whole R22 space, requiring subtler simulation treatment. (Middle) Samples generated by the conventional equivariant diffusion model. (Right) Samples generated by the quotient-space diffusion model, which also recovers the data distribution as guaranteed. Moreover, the quotient-space diffusion simplifies learning difficulty: the neural network does not need to learn anything in the output subspace that responsible for intra saira-eqqivalenn--hee movement (Eq. (10)).
Figure 1:
2 BACKGROUND
2.1 DIFFUSION-BASEL GENERATIVE MODELS ON EUCLIDEAN SPACE
The main idea of diffusion models is to construct step-by-step transformation from a simple prior distribution to a complex target distribution. In this paper, we follow the Stochastic Interpolant framework (Albergo et al., 2023), which unifies diffusion models and flow matching models (Lip- man et al., 2023; Liu et al., 2023). Let Prarget (x) be the target distribution. The following linear interpolation is constructed:
0; + 8:x1 the (xx,xi) Pjolnt: E~ IO. t [0.
(1)
where Pjoint is a pre-defined joint distribution of x1) with marginals Xn- ~ Pprior and X ~ Pharget- The coefficients at, 8,, satisfy the boundary conditions on 1, 8o = 0, to 0, and on = 0. 3, = 1, = O. Under these conditions, the following ordinary differential equation (ODE) can transform Pprior to Ptarget (Albergo et al., 2023, Cor. 2.18):
dx, =v==v(x; t) dt, where v(x,,) t) == E[axxo + 8(x1 + + xx]-
(2)
The velocity vector field ((x,,) is typically trained with the objective: L()) Eptcy p (xx))
- (a/xo + 8/x1 + +2) where the prime denotes the time derivative, and p(t) and w()) control the sampling distribution and weighting over time. There is also stochastic process for sample generation, given by
dx, = (v t) + ns(,,, t)) dt + / 22 dwe, where t(x,,t)= == Vx, lo (x:)
(3)
is called the score function, and π 2 is non-negative smooth function (Albergo et al., 2023,
Cor. 20) In the special case where Pprior AO I) (the one-sided stochastic interpolant (Albergo
al., 2023, Def. 3.4)), contributions of xg and can be combined as x, â.e + f3x1,, where de
Va + + and the score function can be expressed by the velocity field: t) 8/x B.vxx
convenient variant to formulate the learning task is to define the vo(x,,t) model with neural network (x(x,,)) t) which reformulates the objective:
vo t) = d/x; (dL,3 -,,
(4)
de
[(0) (p() (04,3 =,1
Boxx (o --x- x11
(5)
3

---

Published as a conference paper at ICLR 2026
where p(x1, X;) is derived from Eq. (1) by integrating out X, and This objective conveys the intuition of recovering the clean-date sample x; from a noisy sample xt, hence Do(x,, t) is called denoising model and suits prevalent architectures. We adopt this form of diffusion model below.
2.2 FROM EUCLIDEAN SPACE TO QUOTIENT MANIFOLD
Tasks in scientific domains often involve inherent symmetry, where objects related by certain trans- formations are considered equivalent. A formal and inclusive description of symmetry in a system requires both the geometry of the configuration space and the algebraic structure of the transforma- tions, which leads to the concepts of manifolds and Lie groups.
Manifold and Lie groups. A (smooth) manifold a geometric object that generalizes the Euclidean space to allow spatial heterogeneity. Typically, manifold is endowed with a Riemannian metric, i.e., an inner product in each tangent space, which leads to common concepts like curve length, dis- lance, measure, gradient, Laplacian, and Wiener process the manifold (Appx. B. Symmetries are formally represented by transformations that connect equivalent (i.e., symmetric) objects, which constitute group. A continuously parameterized -arrretuusyy-araameteiizeedd group that is also a manifold is called Lie group.
We consider the general case where the configuration space of the system is an M M-dimensional Riemannian manifold .M. The symmetry of the system is represented by G-dimensional Lie group 9 acting on M. A distribution p on M is said G-invariant if p(g. x) = p(x), Vg E g, x E M. This invariance implies that all equivalent points {g- x E g}. collectively called an equivalent class, are assigned with the same probability.
Quotient space. The symmetry group defines an equivalent relation in M, i.e., X and X2 are equivalent, if there exists a group action E G such that X1 = X2, which is indeed an equivalent relation due to properties of group. The quotient space Q M/g treats equivalent objects under the action of G as one element, hence reflects the intrinsic variability of the system. There is a natural mapping called the projection connecting the two spaces: I T(x) == {g x g E g). Under appropriate conditions, the quotient space is a smooth manifold with dimension M G (Appx. C). However, defining a diffusion process on this space is non-trivial, necessitating the extension of "velocity" and Wiener process from Euclidean space to the manifold.
Tangent vector. On a manifold M, the velocity of a process at a certain point x is represented as tangent vector x, intuitively representing an infinitesimal movement. All tangent vectors constitute a linear space T&M called the tangent space at x. Since a manifold is typically curved, tangent spaces at different points are regarded different linear spaces, but with transformation on the manifold, e.g., a group action T.g x + g x, an associated mapping (I.g) (xg))x== : Tx.M - x between the tangent spaces be defined by linking infinitesimal movements around and around x by Ly (Appx. B). With this construction, we can define a G-equivariant vector field on M if is unchanged under the group action: (Lx)+x(Vx) = Vg.x- The projection mapping π naturally induces a projection of tangent vectors onto the quotient space by Txx T + Tx(x) Q.
Wiener process on manifold. In Euclidean space, the Wiener process is generated by the Lapla- cian operator A. The Saplace-Beltramy operator, defined from a Riemannian metric. serves as a counterpart on manifold, and defines the Wiener process to the manifold. Under a symmetry group 9, we require a meaningful stochastic process on the manifold M as G-invariant, meaning that its marginal distribution is G-invariant at any time step. See Appx. B for details.
3 METHODS
As the quotient space represents the "essential states" of a system with symmetry, a principled dif- fusion model for the system expected to be built on it. In this section, we unroll the development of the quotient-space diffusion model by deriving the projected diffusion process onto the quotient space, then lift back into the total space (i.e.. the original space) for convenient implementation. We then derive the specialization in the R3N /SF(33 case for molecular structure generation, fol- lowed by training and sampling algorithms. We highlight the merit of the quotient-space diffusion in reducing training difficulty and sampler soundness with comparative analysis with existing treatments considering symmetry.
4

---

Published as a conference paper at ICLR 2026
3.1 DIFFUSION PROCESS ON GENERAL QUOTIENT SPACE
If the diffusion process in M is G-invariant, the distribution at any time step can be viewed as distribution in the quotient space 2, then we can view the process as a stochastic process in Q. By leveraging the projection mapping
.M- + Q, we can map diffusion process {xe (xe,r)) in M (Eq. (3)) onto the quotient space as {y: == he 7.. This is stochastic process on Q. but
its expression as a diffusion process on Q using specifiers defining the diffusion process of x; is desired. The following theorem gives an explicit answer.
Theorem 1. Assume {x, (Xitee 0,21 is diffusion process on M. specified by the following SDE:
dx, = bx dt + dwt,
Ppriors
(6)
where by is a G-equivariant time-dependent vector field on M, We is the Wiener process on M that also G-invariant, and Perior is G-invariant distribution. Then the projected process {y: I (xx) (xx))e onto the quotient space Q == M/G is the solution to the following SDE:
dy: = (x.bb)(y))) ((==<(ra - o h(y))) di + at dww, Yo ~ T
2
where T. by is the projected vector field of bt onto Q induced by n, h(y) is the mean curvature vector field of Q reflecting the geometry of Q W. is the Wiener process Q, and i #prrio is the pushed-forware distribution of Pptior fi.e., Yu = π (Xo) where Xy ~ Ppeice).
See Appx. D.I for formal definitions of the concepts and the proof. Thm. I shows that the projected process is indeed a diffusion process on @. which consists of the projected vector field and corresponding Wiener diffusion process, and perhaps unexpectedly, an additional vec- tor field reflecting the curvature of Q. As the quotient space squeezes an equivalent class as one point, a pro- cess viewed on the quotient space should accommodate for the change of the volume of the equivalent class along the movement. This additional vector the gradient (i.e., the change rates in all movement directions) of the vol- ume of the equivalent class.
92 x
M
91 x
av
" (v)
π
n
horizontal lift
Q = M 5
TT,
z xxx)
Although the diffusion process on the quotient space is defined,
convenient simulate in the quotient space directly due to the non-trivial geometric structure of Q. Nevertheless, the quotient-space diffusion enables us a principled view to reduce the unnecessary movement within equivalent classes. A key observation from Thm.
Figure 2: Illustration of the relation be- tween the total space M and the quo- tient space Q and the correspondence of tangent vectors among them.
is that if b, v + b2 where Vx Ker x.x {v E T&M Tvx(v)= 0}. Vx E M, then the corresponding SDE in Eq. (6) has the same projection in the quotient space. This implies that the components in Ker n.x not really necessary.
For better characterization of the necessary component, we focus on the tangent space of M at The tangent space T&M is a linear space with the same dimensionality as M. Define the vertical space Vx == Ker T.x (G-dimensional) corresponding to the infinitesimal action of the group G. Since M has an inner product (because M is a Riemannian manifold), we can define the horizontal space Hx == (Ker + as the orthogonal complement of Vx. Then any tangent vector in T&M has an orthonormal decomposition v v vH, where v y is the vertical and horizontal component respectively; see Fig. 2 for visualization. Thus H is the necessary part of the vector field v.
Thanks to the quotient structure, we can leverage correspondence between the diffusion process on M and Q. For a diffusion process yr, there exists diffusion process 8, M such that (X.) = y: and & only has horizontal movement, which is called the horizontal lift of y, (see Appx. 1.2 for formal definitions). The horizontal lift of yt is given explicitly in the following thcorem.
Theorem 2. The horizontal lift of Eq. (7) has the following explicit expression:
(xx, ((x
o ((xx)) dt + o dw., 3o Pprior:
dž,
(8)
2
5

---

Published as a conference paper at ICLR 2026
where Px(v) := v is the horizontal projection in the tangent space of M. h is the horizontal lift of the mean curvature vector h in Eq. (7), We is the horizontal lift of the Wiener process on Q.
See Appx. D.2 for the proof. Comparing the expression between Eq. (6) and Eq. (8), we can observe that the lifted process is not simply given by adding a horizontal projection Px on each term of the the curvature of the quotient space arises. This term
SDE, and an additional term depending arises in Eq. (7) and remains after the horizontal lift. Intuitively, this term corrects the possible side effects by projection so that the resulting diffusion process still yields the desired target distribution. The horizontal projection Px and the mean curvature vector field can be calculated in specific cases, so Eq. (8) has explicit form when is specified.
As mentioned, Eq. (8) only has horizontal movements, in other words, it does not have any move- ment in the equivalent class. This process reduces unnecessary movement and helps to reduce sampling trajectory length. From this viewpoint, previous methods do not reduce these unnecessary movements, although they have the equivalent diffusion process in the quotient space. The formal results are summarized in the following corollary. See Appx. D.2 for proof.
Corollary 3. X1 (defined by Eq. (8)) has the same distribution on Q with x1 (defined by Eq. (6)). When o2 = O. Vxa E M, Eq. (8) has shorter trajectory length than Eq. (6).
3.2 SPECIAL CASE: THE SHAPE SPACE
The abstract results in the previous section give the direction for practical implementation. In this subsection, we focus the space of 3-dimensional coordinates of N points under the symmetry defined by the special Euclidean group SE(3), composed of the 3-dimensional translation group and the 3-dimensional rotation (3)3). By definition, element of this space is structured as x == (x(1), ..... ()). where x(() R and the SE(3) group acts on x by translating and rotating each x²) Since the translation group is not compact, there does not exist translational invariant distribution. We (as well as many others (Yim et al., 2023; Lin et al., 2024)) hence represent the quotient space w.r.t this group by considering the sitterr-f-maas(CoMMe
subspace M := {x R3N 1 =N x(6 0), and consider the SO(3) action on it. The resulting quotient space Q := (MM/0(0()) as the concrete construction for R /SE(3), is a smooth manifold under certain conditions (Appx. C.1). Each element in Q represents N N-point configurations that are equivalent under altogether translation and rotation, therefore, Q is regarded as the "shape space" reflecting the intrinsically different states of the N points. Now we can develop the correspondence between the diffusion process in M (Eq. (6)) and the its horizontal lift from the quotient space projection (Eq. (8)) The results are summarized in the following theorem.
Theorem 4. Assume X₃ is diffusion process in the CoM subspace M C R given by the following SDE: dx,- = (xx)) dt + or dw,, Xy ~ Pprior where by is a field. vt E [0, Pprior is the G-invariant prior distribution, w, is the standard Wiener process on COM. The horizontal lift of the process x xxx) is:
vector
of hxx:) dt + o2 OtPs, dw,, 2
dx: = PP (b.(k))
Xy ~ Ppriors
where the Px is the horizontal projection operator at x and h(xx) is the horizontal lift of the mean curvature vector. The explicit expressions of P and h are shown as follows:
N
Px(v) --
Σ
x x v
x x. Vv E TxM, and
N
N
= -
where J: Σ 1xx
Σ
x 0 (T)
E R3x3
1=1
See Appx. D.3 for proof. From the results of Thm. 4, we can deduce that π (xx)) has the same marginal distribution with rxx;) in Eq. (9) (Cor. 3). If we consider the generation process in Eq. (2) or Eq. (3) as Xt, we can construct the corresponding horizontal process &; that can generated the same target distribution the quotient space. Motivated by this fact, we can improve the training
and inference method of diffusion based generative models by leveraging the quotient structure.
Technically, guarantee proper structures, M needs to exclude an negligible subset; see Appx. C.I.

---

Published as a conference paper at ICLR 2026
3.3 PRACTICAL IMPLEMENTATIONS
Previous results describe how we can construct a diffusion process in the quotient space using the coordinates the total space. If we have a diffusion process on the total space, we can construct the horizontal lift of its projection process, which has no vertical velocity along its trajectory and the two processes are the same on quotient space, This fact implies that the vertical components of the original diffusion process are not dispensable and enables us to design a more efficient training and sampling algorithm of the diffusion model based on the quotient structure. In practice, we often set the total space the Euclidean space. Next, show the training and sampling methods for the special case Pprior = (0, I), and the general case is shown in Appx. E.
Training objective. The diffusion model on the total space M is trained by the objective Eq. (5). Since the vertical components of the velocity are not strictly needed, we propose to supervise the model only on the horizontal components and allow arbitrary vertical output of the model. We lever- age the horizontal projection operator Px (Thm. 4) and construct the horizontal training objective:
(())= ;= (p(e) ((t)) xx) Px, (Do(x,, (xxx)) f) - 11?
(10)
We can see that Dg + V has the same loss value with Dg, where v is an arbitrary vertical vector.
ODE sampler. After the training stage, Px, (Dg(x,, ()) is an approximation of the ground truth denoiser in the horizontal subspace. For the ODE sampler, we simulate the horizontal lift of the projected ODE, which is given by dx, Px, (x t)dt, where vo(x,.t) is given by Eq. (4). In practice, the ODE process is approximated by numerical solvers.
SDE sampler. For the stochastic sampler, the we need to simulate the horizontal lift of the projected original SDE Eq. (3). According to Thm. and Thm. 4, the lifted process is given by
dx,= Px (vo(xe, () + (050((4. f)) dt + 7m: + 22mm Px,dw,
where tx),, = x x.-A, (x...) and we introduce the hyperparameter y for protein generation
57 following Geffner al. (2025). The details are summarized in Algorithm and 3.
3.4 ANALYSIS ON EXISTING TREATMENTS FOR SYMMETRY
In this section, we make detailed analysis on existing methods that handle symmetry, and verify the conclusions in Table 1. In contrast to our quotient- space diffusion, we find that they either have not fully leveraged the symmetry to reduce model-learning difficulty, or do not have a proper sampler.
Conventional equivariant diffusion models and data augmentation. A common treatment is by assigning equal probability to equivalent objects, resulting in an invariant target distribution p(x1). This can be implemented by augmenting data samples by applying randomly chosen group actions, mimicking sampling from the invariant distribution, or using an invariant prior distribution and an equivariant architecture securing Dø(g x.t) g Dx t). The training strategy is the same as modeling a general distribution in the original space following Eq. (5), and the standard samplers by Eqs. (2, 3) remain valid. For each value of xx, this objective asks the model to minimize the average of [x;, t) x1 terms where x1 come from p(x (x). so the optimal solution is the conditional expectation [[[x||x)
Fig. 3 shows an example and reveals characteristics of the training strategy. The example considers generating the structure of a diatomic molecule, where the target distribution p(x1) concentrates on single structure up a uniform random orientation (Left). For given x,, samples of p(x) (x,) are x' structures posed in orientations distributed around the orientation of x, (Middle). Indeed, an X1 sample more closely oriented with x, would have a higher probability to produce the given x, in the diffusion process, so there is specific orientation correspondence between the learning target E|x||x, |xx] and xs. So the model is still asked to learn a correspondence in the equivalent degrees of freedom (DoFs) (i.e., rotation of the output), in contrast to the quotient-space case in Eq. (10) where the model is unconstrained in the vertical space (i.e., tangent space of the rotation group)- Moreover, the X, samples are not all posed in the orientation of x, because x" in other orientations can also generate this x, through the diffusion process. So the model learns the correspondence in the equivalent DoFs from samples with variance, leading to another aspect of learning difficulty.
GeoDiff alignment. To reduce the learning difficulty, some heuristic treatments are proposed based on alignment. The first representative alignment used in GeoDiff (Xu al., 2022) uses the following
7

---

Published as a conference paper at ICLR 2026
Conventional target B[x, [x,
Ag,
GeoDiff target
p(x.)
X
X1 (x]
for p(x,(x))) E(,4s,(x
Figure 3: Illustration of denoising learning target using conventional training and using GeoDiff alignment. (Left) The example considers the structure distribution p(x1) of diatomic molecule, which concentrates on a single structure x' up to uniform random orientation. (Middle) Given an X sample, the corresponding x samples distribute with a variance, and their expectation ||x| is the conventional learning target, which is not equivalent to x* (the bond is shorter). (Right) Given an x sample, all the X1 samples after alignment coincide with x" posed in the orientation of Xg, which is also the learning target of GeoDiff E[AAx, (x1)
training loss: xp(x1x)) x [Dx, t) - Ax. (x1) [I where the alignment operation is defined as:
Ay = argminl ((g
d(x', y);
(11)
where d(..-) -) is the distance metric on M. With an illustration in Fig. (Right,, the learning task can be understood as that for given value of xt, the model output needs to fit Ax, (x1) samples, which are all posed in the orientation of x,, and they all coincide with the x* structure in the orientation of X,. This supervises the model to the target B[AAx, (x1)[xx from samples with no variance in the equivalent DoFs (i.e., rotation of the output), hence reduces certain learning difficulty, Nevertheless, this target still requires the model to learn a specific mapping in the equivalent DoFss, hence does not enjoy the learning advantage in the quotient-space case that relaxes the learning in the DoFs.
caveat of this alignment approach is that a proper sampler needs to be developed, as the con- ventional samplers still require a model targeting E Exx, [x,]. which is different from B,Ax, (x1 [xx]] Fig. illustrates this difference: E(x, [ averages diversely oriented x* structures, resulting in different shape than x* (the bond is shorter), while [[[Ax, (x1) x,) is just x" in the orientation of x4-
AF3 alignment. Another alignment approach, which is used in Alphafold 3 (AF3) (Abramson et al., 2024), aligns the X₁ samples towards the model output: Ep(xx Do t) - (D,(x. (x1)]² where 0 is treated constant in optimization. This loss function allows the model output to differ by an arbitrary group action (e.g., rotation), hence removes the need to learn specific target in the equivalent DoFs. Indeed, for an arbitrary group action 9x.,t, a new denoising model 9x,,t (x(x,,)) achieves the same loss since 119x,,t D9(x,,t) - Agx.. (x) )l2 = 19x,, (x t)
(x.,()
Ix.. ADz(x..) 1| [|Dx(x;, L) (x1]] where the last equality holds since the group preserves metric (Appx. C). Up to this DoF, the learning target is the same as GeoDiff's E[Ax, (xx )x,], since all the x1 samples are averaged after aligned with the same reference.
In the sampling process, the arbitrariness in the equivalent Dol's (e.g., orientation) of the learned model D,(x., t) leads to an arbitrariness" in the vector field v*(x,, 1) through Eq. (4). Hence there is no guarantee of recovering the target distribution using conventional samplers. This problem is also noted by Boltz-1 1 (Wohlwend et al., 2025). which proposes to align the prediction Dø (x,, t) towards x, in the sampling process. As the AF3 target is the same as GeoDiff's up to an arbitrary rotation, this amounts to using the GeoDiff model for sampling, which still cannot guarantee producing the target distribution as concluded above. These discussions are summarized in Table
4 EXPERIMENTS
In this section, study the empirical performance of our quotient-space diffusion model. We carefully conduct several experiments covering different types data, scales and scenarios. To evaluate our quotient space diffusion model framework for real-world applications, we focus on the molecule structure generation protein backbone design tasks, in which we consider the diffusion models on 3 /SE(3) (Sec. 3.2). The details of all experiments shown in Appx. G.
4.1 STRUCTURE GENERATION FOR SMALL MOLECULES
Datasets. First, we evaluate our framework on the molecule structure generation task. In this sce- nario, our goal is to generate the 3D coordinates of molecule given the graph structure of the
This is not even an arbitrary group action (e.g., rotation) since does not vary together with the arbitrari- ness of (x f)-
8

---

Published as a conference paper at ICLR 2026
Table 2: The effect of the quotient-space diffusion scheme for molecular structure generation on the GEOM-QM9 and the GEOM-DRUGS datasets using the ET-Flow(SO(3)) and ET-Flow(O(3) architectures. We use the same sampling steps of 50 NFEs for fair comparison. Best results are marked in bold. Best results for the same architecture underlined.
Recall
Precision
Datasets
Methods
Coverage
AMR
Coverage
AMR 1
mcan
median
mcan
modian
mean
modian
mcan
median
CGCF
69.47
96.15
0.425
0.374
38.20
33.33
0.711
0.695
GeoDiff
76.50
100.00
0.297
0.229
50.00
33.50
1.524
0,510
GEOM QM9
GeuMol
91.50
100.00
0.225
0.193
87.60
100.00
0.270
0.241
(Positive samples
Torsional Diff.
92.80
100.00
0.178
0.147
92.70
100,00
0.221
0.195
within 0.5 A RMSD.)
MCF
95.0
100.00
0.103
0.044
93.7
100.00
0.119
0.055
ET-FHow(SO33)
95.98
100.00
0.076
0.030
92.10
100,00
0.110
0.047
+ Geodiff aligament
95.71
100.00
0.085
0,040
95.20
100,00
0.098
0.050
+ AF3 alignment
92.67
100.00
0.131
0.070
84.38
100,00
0.205
0.146
Quoticnt-spp spacce diffusion
96,40
100.00
0.069
0,024
93.30
100,00
0.096
0,036
GeoDiff
42.10
37.80
0.835
0.809
24.90
14.50
1.136
1.090
GeoMol
44.60
41.40
0.875
0.834
43.00
36.40
0.928
0.841
Torsional Diff.
72.70
80.00
0.582
0.565
55.20
56.90
0.778
0.729
MCF S (13M)
79.4
87.5
0.512
0,492
57.4
57.6
0.761
0.715
MCF- (62M)
84.0
91.5
0.427
0.402
64.0
66.2
0.667
0. 0.605
GEOM-DRUGS
MCF (242M)
84.7
92.2
0.390
0.247
66.8
71.3
0.618
0.530
(Pusitive samples
are within
ET-Flow (8.3M)
79.53
84.57
0.452
0.419
74.38
81.04
0.541
0.470
0.75 Á RMSD.)
+ Quotient-space spacc diffusion
reprodluction
78.94
79.86
84.24
85.71
0.489
0.459
0.472
0.433
66.24
72.70
79.63
70.42
0.651
0.565
0.595
0.501
FT- (3) (9. (M)
78.18
83,33
0.480
0.459
67.27
71.15
0.637
0.567
+ reproductton
74.91
80.90
0.541
0.515
60.33
62.71
0.724
0.665
+ Geodiff aligament
75.11
80.74
0.545
0.526
59.58
60:48
0.734
0.678
+ AF3 alignment
71.66
76.09
0.572
0.570
52.21
50.00
0.828
0.793
+ Quoticnt-space diffusion
78.50
84.20
0.477
0.455
67.35
71.42
0.635
0.563
molecule. We conduct the experiments on the GEOM datasets (Axelrod & Gomez-Bombarelli, 2022), which provides structure ensembles generated by metadynamies in CREST (Pracht et al., 2024) and we focus on the 1EOM-QM9 and GEOM-DRUGS datasets. Following the data pro- cessing and splits from Hassan et al. (2024), we use the random splits with rrain/validation/test of 243473/30433/100 for GEOM-DRUGS and (06586/11333/10000 for GEOM-QM9, In addition, data with disconnect molecule graph are removed for GEOM-DRUGS.
Setting. We primarily follow the setting in Hassan et al. (2024). We use an equivariant graph transformer architecture from ET-Flow (Hassan et al., 2024) and set the Gaussian distribution as prior distribution on GEOM-QM9 and use the harmonic prior for GEOM-DRUGS (Volk et al., 2023). We fix the architecture as (T-F)ow(SO(3)) for experiments on GEOM-QM9, and use the ET-Flow(O(3)). (T-F)) architecture on the GEOM-DRUGS dataset. Following Jing et al. (2022); Xu et al. (2022), we report the RMSD-based metrics, e.g. Coverage and Average Minimum RMSD (AMR) between the generated and ground truth structure ensembles.
Results. The results are presented in Table 2 for the GEOM-QM9 and GEOM-DRUGS datasets. As shown, our proposed quotient-space diffusion framework consistently outperforms prior methods and alignment techniques in terms of generation quality on both datasets. Our framework reduces learning difficulty by removing redundant components, enabling us to further improve the perfor- mance of the ET-Flow framework' on both datasets. On the GEOM-QM9 dataset, our quotient-space diffusion model framework surpasses strong basclines such MCF (Wang et al., 2023) and the ET- Flow framework with other heuristic alignment methods among most of the RMSD-based metrics. On the GEOM-DRUGS dataset, our framework not only significantly surpasses the ET-Flow base- line with heuristic alignment methods, since these methods are incompatible with training, but also achieves competitive performance against the larger MCF-L (242M) model (Wang et al., 2023) on the Precision metries.
4.2 PROTEIN BACKBONE DESIGN
Setting. To demonstrate the advantage of our quotient-space diffusion model for larger and more relevant molecules, we perform comparative analysis on the task of protein structure generation against the otate-of-the-art Proteína model (Geffner et al., 2025). We select their most efficient
'e reproduce the results using the released contigurations: nttps://githut con/ shenoyni khi 1/ETF1FFT low. Due to changes in the data processing pipeline, our reproduced results do exactly match those reported in the original paper.

---

Published as a conference paper at ICLR 2026
Table The effect of the quotient-space space diffusion scheme for protein structure generation using the Proteína model. Best results are marked in bold.
Settings
Methods
Designability (9))
FPSD vs.
ISD vs.
PDB.
AFDB
(CIA/T)
PDB.
AFDB.
FrameDiff
65.4
194.2
258. 1
2.46/5.78/2
335
1.04
1.42
FoldFlow (basc)
96.6
601.5
566.2
1.06/1.79/9.7 72
3.18
3.10
FoldFlow (stoe.)
97.0
543.6
520.4
1.21/2.09/1 1.59
3.69
2.71
FoldFlow (OT)
97.2
451,4
414.1
1.35/3.101
2.90
2.32
Representative
FrameFlow
88.6
129.9
159.9
2.52/5.88/22.00 00
0.68
0.91
References
ESM3
22.0
933.9
855.4
3.19/6.71/1..1
1.53
0.98
Chroma
74.8
189.0
184.1
2.344,95/88..1
1.00
1.08
RFDiffusion
94.4
253.7
252.4
02//5.06/19.83
1.21
1.13
Proteus
94.2
225.7
226.2
46/16.22
1.41
1.37
Gonic2
95.2
350,0
313.8
1.5//366/11.0 1.65
2.21
1.70
Proteina Mg 0.35
96.0
386.5
378.2 2
1.77/499//11 78
2.17
1.73
Quodient-spuce diffiusion
97.6
274.7
2.24/6.69/20.99
1.55
SDE Sampling
Proteina M 0.45
92.2
332.9
320.4
1.835 01/20,22
1.93
1.49
Quotient-space diffusion
92.6
244.5
246.3
2.24/6.68/23.47
1.43
1.28
Proteina Mitall 7
306.2
290.8
92/21.15
1.36
Quotient-space diffusion
90.2
228.0
228.7 7
2. 2.25/6 59/25. 24
1.32
1.17
Proteina Mss
19.6
85.4
21.4
2.51/5 2.51/5.65/27.38
0.59
0.09
ODE Sampling
Proteina M
13.8
83.2
21.9
2.45/5.630 13/31,76
0.58
0.12
AF3 alignment
3.8
229.0
82.4
1.35
0.36
Quotient-space difflssion
15.6
69.9
17.6
2.57/6.40/32.14
0.41
0.11
variant Manall 60M parameter transformer trained the Foldseek AFDB clusters (Drş) that forgoes triangle layers and pair representation updates, as a strong and relevant baseline. We train the quotient-space diffusion model from scratch using the identical architecture on the identical dataset. For evaluation, both our model and the officially released Proteína checkpoint are sampled using 400 steps with self-conditioning. We explore the baaggnabilit--dversse trade-off by testing a range of noise scales, y E {0.35, 0.45, 0.5]*. . To faithfully evaluate the distributional metrics proposed by Geffner et al. (2025), we utilize ODE sampling.
Results. The results Table 3 highlight the superiority of our quotient space framework, which, unlike alignment-based strategies (adapted from AF3 and Boltz-1), provides a theoretical guarantee for sampling the correct target distribution. The alignment- bution, with performance metrics falling short of even data-augmented, semi aemi-equvvrriant baselines. We attribute this failure to a fundamental incompatibility between their samplers and the learned models. Furthermore, our formulation effectively reduces learning difficulty by removing the need to learn a specific target in redundant spatial transformations, enabling the model to capture key structural features more efficiently than standard semi eunivrquivaiiant baselines. This advantage of ef- ficiency leads to significant results: our 60M parameter model not only surpasses its direct baseline across both SDE at all noise scales and ODE sampling setting, but also outperforms the much larger 200M Mrs model on most key distributional metrics. This provides compelling evidence that a quo- tient space framework ensuring both sampling fidelity and learning efficiency is key to advancing generative protein models.
methods fail to recover this distri-
CONCLUSION
In this work, we formally construct a framework for building diffusion models on the quotient space over a group, in hope for a principled approach to handle symmetry in a generative task. We ex- plicitly give the expression of the diffusion process the quotient space, then also construct corresponding diffusion process in the original space for easier implementation. The resulting train- ing algorithm reduces learning difficulty by removing the need predict the tangent vector in the direction along group action, and the resulting sampling process guarantees producing the target distribution while removes the unnecessary movement in the group-action direction. We instantiate the method in the case of R-N /SE(3) for molecular structure generation. Empirical results on struc- ture sampling for small molecules from the GEOM-QM9 and GEOM-DRUGS datasets and protein backbone generation demonstrate the better generation quality and design success rate over exist- ing conventional equivariant diffusion models and alignment-bue based approaches given equal or fewer training epochs, demonstrating the practical advantages from this principled framework to handling symmetry in diffusion models.
'Due to a known bug in a previous version of Foldseek (Daras et al.. 2025, Appendix B), our comparative analysis in the main text is focused solely on the designability. More comprehensive metries evaluating self-sampled structures are provided in Table 5.
10

---

Published as a conference paper at ICLR 2026
ACKNOWLEDGMENTS
This work is supported by Zhongguancun Academy (Grant No. C20250506). DH is supported by National Science Foundation of China (NSFC62376007), National Science Foundation of China (under Key Project No. 92570203), Beijing Natural Science Foundation (Z250001) and Beijing Major Science and Technology Project under Contract no. 120100000442004
6 ETHICS STATEMENT
This work adheres to the ICLR Code of Ethics. Our study does not involve human subjects, per- sonal data, sensitive demographic information. All experiments conducted publicly avail- able benchmark datasets, which are widely used in the machine learning community. No new data collection or human/animal experimentation was performed.
7 REPRODUCIBILITY STATEMENT
To facilitate the reproducibility of our research, we provide comprehensive details throughout the paper and its supplementary materials. We begin by establishing the necessary foundational knowl- edge in Sec. 2.1 and Appx. B. For all theoretical claims and proofs presented the main text, we offer detailed step-by-step derivations in Appx. D. Our experiments datasets, training procedures, and evaluation protocols are carefully described in Sec. 4 and Appx. G. Upon acceptance of this paper, commit to making full codebase and all model checkpoints publicly available to ensure that the community can fully reproduce our results.
thoroughly documented; the
THE USE OF LARGE LANGUAGE MODELS (LLMs)
In the preparation of this manuscript, LLMs were employed as writing assistant to refine the language and improve the grammar. Furthermore, we utilized LLMs assist in verifying our math- ematical formulas for notational consistency. Following this process, all textual and mathematical content was meticulously reviewed, revised, and validated by the authors, who assume full respon- sibility for the final work presented.
REFERENCES
Josh Abramson, Jonas Adler, Jack Dunger, Richard Evans, Tim Green, Alexander Pritzel, Olaf Ronneberger, Lindsay Willmore, Andrew J Ballard, Joshua Bambrick, et al. Accurate structure prediction of biomolecular interactions with AlphaFold 3. Nature, pp. 1-3, 2024.
Michael Albergo, Nicholas M Boffi, and Eric Vanden-Eijnden. Stochastic interpolants: A unifying framework for flows and diffusions. arXiv preprint arXiv:2303.08799, 2023.
Jacob Austin, Daniel D Johnson, Jonathan Ho, Daniel Tarlow, and Rianne Van Den Berg. Structured denoising diffusion models in discrete state-spaces. Advances in neural information processing systems, 14:119811-1993. 2021.
Simon Axelrod and Rafael Gomez-Bombare Bombarelli. GEOM, energy-annotated molecular conformations for property prediction and molecular generation. Scientific Data, 9(1):185, 2022.
Jan-Hendrik Bastek, WaiChing Sun, and Dennis Kochmann. Physics-informed diffusion models. In The Thirteenth International Conference on Learning Representations, 2025. URL https: / /pw.opveeee..eeeeoee wwte/foorm//teteeee forum? wttprr/mmm.dtttteemmim..
Fabrice Baudoin, Nizar Demni, and Jing Wang. Stochastic areas, horizontal Brownian motions, and hypoelliptic heat kernels. EMS Press, 2024.
Isaac Chavel. Riemannian geometry: modern introduction. Number 108. Cambridge university press, 1995.
Ricky TQ Chen and Yaron Lipman. Flow matching on general geometries. arXiv:2302. 03660, 2023.
arXiv preprint
11

---

Published as a conference paper at ICLR 2026
François Cornet, Federico Bergamin, Arghya Bhowmik, Juan Maria Garcia Lastra, Jes Frellsen, and Mikkel N Schmidt. Kinetic langevin diffusion for crystalline materials generation. arXiv preprint arXiv:2507.03602. 2025.
Giannis Daras, Jeffrey Duyang-Zhang, Krithika Ravishankar, William Daspit, Costis Daskalakis, Qiang Liu, Adam Klivans, and Daniel Diaz. Ambient proteins: Training diffusion models on low quality structures. bioRxiv, pp. 2025-07, 2025.
Valentin De Bortoli, Emile Mathieu, Michael Hutchinson, James Thornton, Yee Whye Teh, and Arnaud Doucet. Riemannian score-based generative modelling. Advances in neural information processing systems, 35:2406-2422, 2022.
Zach Evans, Cj Carr, Josiah Taylor, Scott H. Hawley, and Jordi Pons. Fast timing- conditondn latent audio diffusion. In Proceedings of the 41st International Conference on Machine Learn- ing, volume 235 of Proceedings of Machine Learning Research, pp. 12652-12665, 2024. URL
2015///vaa222 htmm.
Octavian Ganea, Lagnajit Pattanaik, Connor Coley, Regina Barzilay, Klavs Jensen, William Green, and Tommi Jaakkola. Geomol: Torsional geometric generation of molecular 3d conformer en- sembles. Advances in Neural Information Processing Systems, 34:13757-13769, 2021.
Tomas Geffner, Kieran Didi, Zuobai Zhang, Danny Reidenbach, Zhonglin Cao, Jason Yim, Mario Geiger, Christian Dallago, Emine Kucukbenli, Arash Vahdat, and Karsten Kreis. Proteina: Scal- ing flow-based protein structure generative models. In The Thirteenth International Confer- ence on Learning Representations, 2025. URL https: https///pppeeeeeeeeeeoom net.eeewwweeeeeeeeeeeeeeeeeeeeee www.nneneeet//t..ee forum?idt TVQLu34bdw
Majdi Hassan, Nikhil Shenoy, Jungyoon Lee, Hannes Stärk, Stephan Thaler, and Dominique Beaini. ET-Flow: Equivariant flow-matching for molecular conformer generation. Advances in Neural Information Processing Systems, 37: 17:128798-128824, 2024.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In Advances in Neural Information Processing Systems, volume 33, pp. 6840-6851, 2020.
Jonathan Ho, William Chan, Chitwan Saharia, Jay Whang, Ruiqi Gao, Alexey Gritsenko, Diederik P Kingma, Ben Poole, Mohammad Norouzi, David J Fleet, et al. Imagen video: High definition video generation with diffusion models. arXiv preprint
2022.
Emiel Hoogeboom, Victor Garcia Satorras, Clément Vignac, and Max Welling. Equivariant diffu- sion for molecule generation in 3d. In International conference machine learning, pp. 8867- 8887. PMLR, 2022a.
Emiel Hoogeboom, Victor Gareia Satorras, Clément Vignac, and Max Welling. Equivariant diffu- sion for molecule generation in 3D. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato (eds.), Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pp. 8867- 8887. PMLR, 17-23 Jul 2022b.
Elton P Hsu. Stochastic analysis manifolds. Number 38. American Mathematical Soc., 2002.
Chenqing Hua, Sitao Luan, Minkai Xu, Zhitao Ying, Jie Fu, Stefano Ermon, and Doina Precup. Mudiff: Unified diffusion for complete molecule generation. In Learning on Graphs Conference, pp. 33-1, PMLR, 2024.
Chin-War Huang, Milad Aghajohari, Joey Bose, Prakash Panangaden, and Aaron C Courville. Rie- mannian diffusion models. Advances in Neural Information Processing Systems, 35:2750-2761, 2022.
Bowen Jing, Gabriele Corso, Jeffrey Chang, Regina Barzilay, and Tommi Jaakkola. Torsional dif- fusion for molecular conformer generation. Advances in Neural Information Processing Systems, 35:24240-2444 -4253,, 2022.
12

---

Published as a conference paper at ICLR 2026
Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion- based generative models. Advances in neural information processing systems, 35:26565-2666 26577, 2022.
Seongsu Kim, Nayoung Kim, Dongwoo Kim, and Sungsoo Ahn. High-order equivariant flow match- ing for density functional theory Hamiltonian prediction. arXiv preprint arXiv:2505. 2025.
Jonas Köhler, Leon Klein, and Frank Noé. Equivariant flows: exact likelihood generative learning for symmetric densities. In International conference machine learning, pp. 5361-5370. ap.5361-5370.PMMLR, 2020.
Zhifeng Kong, Wei Ping, Jiaji Huang, Kexin Zhao, and Bryan C. Catanzaro. DiffWave: A versatile diffusion model for audio synthesis. In International Conference on Learning Representations (ICLR), 2021. URL hitps: wpwpr/eeeeeee..ee////e/e. wwre///orumiiiii a xFKBYmZ5J.
John M Lee. Smooth manifolds. In Introduction to smooth manifolds, pp. 1-29. Springer, 2003.
John M Lee. Introduction to Riemannian munifolds, volume 2. Springer, 2018.
Sarah Lewis, Tim Hempel, José Jiménez-Luna, Michael Gastegger, Yu Xie, Andrew Y. K. Foong, Victor García Satorras, Osama Abdin, Bastiaan S. Veeling, Iryna Zaporozhets. Yaoyi Chen, Soo- jung Yang, Adam E. Foster, Arne Schneuing, Jigyasa Nigam, Federico Barbero, Vincent Stimper, Andrew Campbell, Jason Yim, Marten Lienen, Yu Shi, Shuxin Zheng, Hannes Schulz, Usman Munir, Roberto Sordillo, Ryota Tomioka, Cecilia Clementi, and Frank Noé. Scalable emulation of protein equilibrium ensembles with generative deep learning. Science, 38966661) 2025. doi: 10. http://iennee.vv88110 URL https::///// /www///w.o www.scciiiccc..................e org htt//do///at/ /otpg//oo///////e///// /126/sciese
adv9817.
Xin Li, Wenqing Chu, Ye Wu, Weihang Yuan, Fanglong Liu, Qi Zhang, Fu Li, Haocheng Feng, Errui Ding, and Jingdong Wang. VideoGen: A reference-guide latent diffusion approach for high definition text-to-video generation. arXiv preprint arXiv:2309.00394 2023. URL https: // //arxiv... wwwvvvorg/aae abs 2209 00398.
Peijia Lin, Pin Chen, Rui Jiao, Qing Mo. Cen Jianhuan, Wenbing Huang, Yang Liu, Dan Huang, and Yutong Lu. Equivariant diffusion for crystal structure prediction. In Forty-first International Conference on Machine Learning, 2024.
Yaron Lipman, Ricky T. Q. Chen, Heli Ben-Hamu, Maximilian Nickel, and Matthew Le. Flow matching for generative modeling. In The Eleventh International Conference on Learning Repre- sentations, 2023. URL URL https:/ /tl:://ooeae opw.revieeeeeeeee///ee/eeeeeeeem n forum? httet//ooommmidpmmoooooo Pd-PPyMDDD.TTT.
Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. In The Eleventh International Conference Learning Repre- sentations, 2023. URL /ntttss://ppereevieee net/ pegpprrreieee..eetfffoooomool/l. /t
Philipp Pracht, Stefan Grimme, Christoph Bannwarth, Fabian Bohle, Sebastian Ehlert, Gereon Feld- mann, Johannes Gorges, Marcel Müller, Tim Neudecker, Christoph Plett, et al. Crest-a program for the exploration of low-energy molecular chemical space. The Journal of Chemical Physics, 160(11), 2024.
Arne Schneuing, Charles Harris, Yuanqi Du, Kieran Didi, Arian Jamasb, Ilia Igashov, Weitao Du, Carla Gomes, Tom I. Blundell, Pietro Lio, et al. Structure-based drug design with equivariant diffusion models. Nature Computational Science, 4(12):899-909, 2024.
Yang Song, Jascha Sohl-Dickstein, Diederik Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In Interna- tional Conference on Learning Representattinn, 2021.
Anton Thalmaier. Stochastic riemannian geometry. 2023.
Jos Torge, Charles Harris, Simon V Mathis, and Pietro Lio. Diffhopp: A graph diffusion model for novel drug design via scaffold hopping, arXiv preprint arXiv:2008.07416 2023.
13

---

Published as a conference paper at ICLR 2026
Amanda A Volk, Robert W Epps, Daniel T Yonemoto, Benjamin S Masters, Felix N Castellano, Kristofer G Reyes, and Milad Abolhasani. AlphaPlow: autonomous discovery and optimization of sulti-step chemistry using self-driven fluidic lab guided by reinforcement learning. Nature Communications, 14(1):140) 2023.
Yuyang Wang, Ahmed A Elhag, Navdeep Jaitly, Joshua M Susskind, and Miguel Angel Bautista. Swallowing the bitter pill: Simplified scalable conformer generation. arXiv preprint arXiv:231..7952., 2023.
Jeremy Wohlwend, Gabriele Corso, Saro Passaro, Noah Getz, Mateo Reveiz, Ken Leidal, Wojtek Swiderski, Liam Atkinson, Tally Portnoi, Itamar Chinn, et Boltz-1 democratizing biomolecular interaction modeling. BioRxiv, pp. 2024-11, 2025.
Lemeng Wu, Chengyue Gong, Xingchao Liu, Mao Ye, and Qiang Liu. Diffusion-based molecule generation with informative prior bridges. Advances in neural information processing systems, 35:36533-36545, 2022.
Minkai Xu, Lantao Yu, Yang Song. Chence Shi, Stefano Ermon, and Jian Tang. GeoDiff: geo- metric diffusion model for molecular conformation generation. In International Conference on Learning Representations, 2022.
Minkai Xu, Alexander S Powers, Ron O Dror, Stefano Ermon, and Jure Leskovee. Geometric latent diffusion models for 3d molecule generation. In International Conference Machine Learning, pp. 38592-38611 PMLR, 2023.
Jason Yim, Brian L Trippe. Valentin De Bortoli, Emile Mathicu, Arnaud Doucet, Regina Barzilay, and Tommi Jaakkola. SE(3) diffusion model with application to protein backbone generation. In International Conference on Machine Learning, pp. 40001-40039, 2023.
Shuxin Zheng, Jiyan He, Chang Liu, Yu Shi, Ziheng Lu, Weitao Feng, Fusong Ju, Jiaxi Wang, Jianwei Zhu, Yaosen Min, He Zhang, Shidi Tang, Hongxia Hao, Peiran Jin, Chi Chen, Frank Noé, Haiguang Liu, and Tie- Liu. Predicting equilibrium distributions for molecular sys- tems with deep learning. Nature Machine Intelligence, 2024. ISSN 2522-5839. doi: 10.1038/
-42256-024-00837-4
Yuchen Zhu, Tianrong Chen, Lingkai Kong, Evangelos Theodorou, and Molei Tao. Triv- ialized momentum facilitates diffusion generative modeling on lie groups arXiv preprint 1633 2024.
14

---

Published as a conference paper at ICLR 2026
APPENDIX
The organization of the appendix are as follows. In Appx. A, we briefly discuss the related work relevant to our research. In Appx. B, we review some background knowledge of Riemannian ge- ometry and stochastic calculus on the manifold. In Appx. C. we give the details of the Riemannian structures of the quotient space. In Appx. D, we give all the proofs of the theorems in the main text. In Appx. E, we show our methods for the general case. In Appx. F, we give some additional results and discussions. Finally, the details of the experiments given in Appx. G.
A RELATED WORK
Diffusion models on Riemannian manifolds. As the quotient has the Riemannian manifold struc- ture, several previous works construct the diffusion model on the Riemannian manifolds. De Bortoli et al. (2022) constructs diffusion models using different overlapping local coordinate systems of the manifold and requires geodesic random walk to simulate the forward process. Huang et al. (2022); Chen & Lipman (2023) construct diffusion models in an embedding space which allows a global representation but requires explicit geodesic formula of the manifold. Zhu et al. (2024) constructs the reverse of kinetic Langevin dynamics on a Lie group to perform generative modeling. Such an approach is not designed for and not readily applicable to the quotient space, which has different geometric structure from the Lie group. In our quotient space case, the specialty with a quotient structure enables us construct diffusion models using the coordinate systems of the total space without relying on an embedding of the quotient in the total space (unnecessarily an embedding space), which is more practical implement yet still general.
Geometric diffusion models. To ensure physical symmetry in the generation process, a main- stream strategy integrates fundamental physical constraints, such as SE(3) equivariance, directly into the diffusion model's architecture. This approach, pioneered by models like EDM (Hoogeboom et al., 2022a), typically employs an EGNN to operate directly on atomic coordinates, using tech- niques like zero center of mass adjustments guarantee translational invariance. This foundational concept was subsequently extended in several directions. For instance, the approach was adapted for Diffusion Bridges models like EDM-Bridge (Wu et al., 2022) and for diffusion in a latent space in models like GeoLDM (Xu et al., 2023). These equivariant diffusion techniques have been success- fully applied across range of molecular tasks. For structure generation, models like GeoDiff (Xu et al., 2022) predict 3D structures from molecular graphs. In molecular optimization, methods such as DiftHopp (Torge et al., 2023) refine existing molecules to enhance desired properties. For de novo design, a key advancement has been to combine discrete diffusion models (D3PM) (Austin et al., 2021) for 2D topology with continuous equivariant diffusion for 3D geometry, enabling joint generation as seen in models like DiftSBDD (Schneuing et al., 2024) and MUDiff (Hua al., 2024). A similar problem has also been considered in crystalline structure generation, where the intrinsic periodic translation invariance is an intrinsic symmetry. Lin al. (2024) highlighted the intrinsic periodic translation symmetry that has been omitted for a long time in the field of periodic crystalline structure generation. The work designed modified diffusion process that induces transition ker- nel that is invariant under periodic translation, leading to learning target for the score model that is invariant under periodic translation. Cornet et al. (2025) proposes a novel method that generalizes the Trivialized Diffusion Model framework for fractional coordinates to model the intrinsic periodic translation symmetry using flat coordinates. The proposed method considers the process with the velocity restricted to the CoM-free linear subspace. They have achieved the removal of variance on equivalent DoFs, but still asks the neural network model to learn to predict a specific target in the equivalent DoFs.
Learning with alignment To reduce learning difficulty, some heuristic treatments (learning with alignment) have been proposed hope to reduce the DoFs corresponding to the symmetry group action. The alignment strategy used in GeoDiff (Xu et al., 2022) aligns the target structure to the noisy input by finding an optimal rigid transformation that minimizes the distance between them. Another approach, proposed in AlphaFold 3 (AF3) (Abramson et al., 2024), aligns the target sam- ples to the model output structure. As discussed in the main text, these two alignment-bass training frameworks lack definite guarantee for recovering the correct target distribution, and is incompati- ble with the sampling process. Boltz-1 (Wohlwend et al., 2025), an open-source replication of AF3, noticed this issue and proposed modification in sampling to align the denoised structure to the
15

---

Published as a conference paper at ICLR 2026
structure in the current generation step before updating. Nevertheless, discussed in Sec. 3.4, this, together with the training protocol of AF3, amounts to the operation of GeoDiff, still questioning the sampling process.
B BACKGROUND IN RIEMANNIAN GEOMETRY AND STOCHASTIC CALCULUS
B. RIEMANNIAN GEOMETRY
In this section, we review some background on differential geometry and Riemannian geometry. For systematic treatment of the subject, please refer to standard textbooks Lee (2003; 201 8).
First, we give the formal definition of the smooth manifold. manifold general topological space that locally has Euclidean structure.
Definition 5. An M M-dimensional topological manifold is a topological space M such that:
MM is locally Euclidean, i.e. locally homeomorphic to RM,. Formally, Vr E M there exists an open neighborhood U M that homeomorphic some open set V M. We call the homeomorphism @ l - V c RRM coordinate system chart.
M is Hausdorff topological space.
.M has countable basis for its topology.
smooth manifold topological manifold with an additional smooth structure, which is defined as follows.
Definition 6. A smooth structure on a M M-dimensional topological space M is collection of coor- dinate systems E = ((((0), ola)): E A} which satisfies the following properties:
The collection & covers M: U.EA u(a) M;
For any a, 3 E A, the transition function gfa) 2ه00) is smooth map:
E is maximal collection, i.e. if (21 ②) is coordinate system such that for all A that the maps @ o p(a)-1 and sí()) o Φ-1 are smooth, then (lI, 6) E E.
The pair (M, 6) is called smooth manifold of dimension M.
there is a coordinate system (4, Φ) around a point x E M, then in this neighborhood of a, the manifold admits coordinate chart x² (x) d'(x) and manifold point in the neighborhood can be expressed as a vector x(x) = (x (x),
(x))
With the smooth structure, between smooth manifolds.
define smooth function the manifold and smooth mapping
Definition 7. Let M, be smooth manifolds with dimensions M, N respectively.
A function f M + R is called a smooth function if its vectorized form f a r-1
o (14) + R smooth on (2²(() RM for all smooth coordinate systems (U, Φ) of M. Denote all the smooth functions on M as (M))
A map F M - W'r is called a smooth map if its vectorized form to Foo-1 1 o (() (VV) is smooth for all smooth coordinate systems (4, Φ) of M and (V,$)
smooth map F: + N which is invertible and whose inverse smooth is called a diffeomor- phism. In this case say that M and N are diffcomorphic manifolds.
To define movement on smooth manifold M, we need to define tangent vectors on the manifold.
'On an abstract manifold, point is an abstract object and may not be a vector by itself, so we do not use a boldface notation. A vector representation as the coordinates is available after choosing (local) coordinate system.
16

---

Published as a conference paper at ICLR 2026
Definition 8. Let M be smooth manifold, and M is point, and U is neighborhood of it. A linear map v (2)) + R is called a derivative at x if satisfies
v( - (()) g(x)v())
The set of all the derivatives of C (2) in x, denoted by T,,, is vector space called the tangent space to M at x. An element of T.M is called a tangent vector at I.
Definition 9. Let N,N be smooth manifolds and F .M- + V be a smooth map. Let x' E M and V S N be neighborhood of F(x). Then induces a push-forward map over the tangent spaces, F.z: T.M- - TF(a) is defined as:
F.z
= = F), VSEC (V)) V
When a coordinate system (4, around and (V, around F(z) are chosen, the coordinate expression for F.x is just the Jacobian matrix of its vectorized form tb o F o Φ i.c., V( o F & (z)) So F., is also called the differential of F and also admits the notation d.F.
The tangent bundle TM a smooth manifold M is the union of the tangent spaces of each points, i.e. TM := Le T.M. Similar to the total derivative of the smooth map in Euclidean space, the differential of a smooth map between smooth manifolds is a linear map between tangent spaces.
A vector field on smooth manifold M is correspondence that associates to each point M a vector V. E M.M. The vector field is smooth the mapping v : M - TM is smooth. Denote all the smooth vector fields M by 2 (JM). With the definition of vector field, the solution of ordinary differential equation (ODE) on the manifold. The idea is similar to the definition in Euclidean space, the solution of the ODE is curve whose velocity each point the same as the vector field.
define
Definition 10. Let be smooth vector field on the smooth manifold M. An integral curve of v differentiable curve 7 : [0, T] + MM whose velocity at each point equal the value of v at that point:
(1)) Vy(1))
1T E [0,
Let TV M be the dual space of M,M, which called the cotangent space of M at æ. The cotangent bundle T M is the union of the cotangent space of each points, i.e. T M == UpeM M.
Definition 11. A I-form on smooth manifold M is a correspondence that associates to each point M covector +2 E T M. The 1-form is smooth if the mapping O M T M is smooth.
define continuous group with good properties. Definition 12. A Lie group smooth manifold G that also group with the property that the multiplication map G x G + G, (g, h) + h and the inversion map G + G,g + g are both smooth.
With the definition of smooth manifold,
Define the left multiplication mapping group element and as an action on a group element. A vector field v on G is said to be left-invariant it's invariant under all left multiplications, i.e. (Lg)=gg (Vg') = Vggg'
(h) = gh, which is introduced to differentiate g as Lie-
Definition 13. A Lie algebra a real vector space g endowed with a map called the bracket [: g x g + that satisfies the following properties for all X, Y, Z E g:
.
Bilinearity: Va, b E R.
JaX + +b,, Z]
21, 2, 68 a. bbZ,Y]:
Antisymmetry: [X,Y] = X|;
Jacobi Identity: [X, YY, 2], [] [Y, [2, [X,
The Lie algebra of all smooth left-invariant vector fields on Lie group G is called the Lie algebra of G. which has the same dimension with G.
Example 14. The Lie algebra of the group S3(3), denoted by so(3), is given by all the 3- dimensional antisymmetric matrices so(3) (A R3x3 A AT 0).
17

---

Published as a conference paper at ICLR 2026
Smooth manifold is topological structure. If we want define the "length of the velocity" and distance between two points on the manifold, a metric on the tangent space is required. Such a metric endows the metric with an additional geometry structure. The formal definitions are follows. Definition 15. A Riemannian metric on a smooth manifold a correspondence which associates to each point p of M an inner product -. that varies smoothly on M. In other words, for any two smooth vector fields u, v, (u, v) is a smooth function on M. smooth manifold with a given Riemannian metric is called a Riemannian manifold.
To define the "difference" between tangent space different points, need introduce a concept called affine connection.
Definition 16. An affine connection V Riemannian manifold is mapping
V L (M) 2 ((M) + ((M)
which is denoted by (u, v) which satisfies the following properties:
Vuv linear over C (M) u: V ((0 C (M) and n(), .(()) E (M))
70 f(1) V V V;
v linear over I in v: () (2)) E R and v(t), v E T (M).
V (a a (2) = ( vuv(1) + a v vvv(
satisfies the following product rule: Vf E (M),
Vu(/v) = /V.v (uf)v.
A connection is called the Levi-Civita connection if satisfies the following additional properties:
V is compatible with metric: V (v((), v(2) = (V.v v(2)) + (v(a), Vurle)):
V is torsion-free: V.v- V.u u(v(-)) v(u())).
The Levi-Civita connection is the connection with nice properties. Its existence and uniqueness is fundamental result of Riemannian geometry.
Theorem 17. (Fundamental Theorem of Riemannian Geometry (Lee, 2018, Thm. 5.10)) Assume (M,(.) is a Riemannian manifold. Then there exists a unique Levi-Civita connection.
As the end of this subsection, we introduce the Laplace- Laplace-Beltrami operator on the manifold, which is used define the Wiener process on the manifold.
Definition 18. Let V be the Levi-Civita connection on M. The Hessian of f e C (M) is defined by
(u, v) == v(u(f) (V) Vu, E E (M).
The Aaplace-Beltrami operator A is defined as the trace of Hessian. In other words, Af ΣM Hess(e,, e,) where (ei...., eM} is orthonormal basis for T.M.
B.2 STOCHASTIC CALCULUS
MANIFOLD
With the Riemannian structure defined in the previous section, stochastic differential equations (SDE) and diffusion processes on the manifold. For a systematic treatment of the subject, please refer standard textbooks Hsu (2002); Thalmaier (2023). First, recall the definition of SDE and diffusion process in Euclidean space.
consider the definition of
Definition 19. The infinitesimal generator of a stochastic process (x;); for a function d(x) is
E [[((X(+s) = x
Li((x) lim a-+0
where @ is suitably regular function. For an Itô process defined as the solution to the SDE dx, f(x,, t) dt + E(x,, t) dw,, the generator is
C, Σ (x +
Σ [xxx))xxx1) 0,o,
18

---

Published as a conference paper at ICLR 2026
On the other hand, the diffusion process can also be defined by its generator.
Definition 20. A D-dimensional stochastic process X with continuous sample path defined probability space (Ω, . IP) is called a diffusion process generated by smooth second-order elliptic operator C if the following hold: vf E C (x²)) the process
Mf = fxx)) - f(xo) - 1 (xf(x.) ds
a ---martingale.
To generalize the definition of SDE Riemannian manifold M, we need to define the second- order differential operator the manifold. Let M be M-dimensionall Riemannian manifold. A second-order partial differential operator on M of the form
R
L v
(0)
(k)3
where (M) (MM),
for some R E N+, The square of vector field is understood the decomposition of derivatives:
v(k) (f) v( (v(* Vf E (M))
The vector fields can be generalized to the time-dependent case. Now can extend the definition of diffusion process on a Riemannian manifold.
(FF) Definition 21. (Thalmaier, 2023, Def. 1.1.3) Let ($) (c2) ) be a probability space equipped with increasing sequence of sub-o-algebra F, $ An adapted continuous process x, taking values in M, is called L--diffusion for all test functions ʃ E C2 (M), the process
N/C f(x)) [(xo) 1 (L,f(xx.) ds,
t t 0,
is a martingale, i.e. EIN/ N/ $. 0, L.
For special case, we can define the Wiener process on the Riemannian manifold M.
Definition 22, A Wiener process w, on M is diffusion process with generator 1 A, where A is the Laplace-Beltran operator M, i.e. W; is a continuous stochastic process on M such that for any f (M),
[. A/(w,) ds,
[(x)) -
local martingale up valid time period.
For stochastic differential geometry, the Stratonovitch integral is more convenient than the Itô In- tegral. The Stratonovitch differential effectively subsumes the deterministic second-order effect of the Wiener process from the quadratic variation into the drift term, so that satisfies the ordinary chain rule of calculus. This property enables a clear correspondence between the diffusion process under diffeomorphism between two Riemannian manifolds. Next, we give the definition of the Stratonovitch integral on the Euclidean space and its generalization to Riemannian manifolds.
Definition 23. For continuous real-valued semimartingales and y, let dy dy |d(x, y] be the Stratonovitch differential. Here dy is the usual Itô differential, and d yx == dxdy is the quadratic co-variation of x and y. The integral
l x -
(xdy + faby++dl8s _dlxxyy1
called the Stratonovitch integral of x with respect
Proposition 24. Ito-Stratonovitch formula (Thalmaier, 2023, Prop. 1.2.10). Let x be a continuous -V--uue semimartingale and f E C (BR). Then (x (x), odx).
The Itô- Sid-Stratonovitch formula shows the advantage of the Stratonovich differential: it satisfies the usual chain rule of classical calculus. So at least formally, classical differential calculus can be applied in calculations involving Stratonovich differentials.
19

---

Published as a conference paper at ICLR 2026
Proposition 25. (Thalmaier, 2023, Prop. 1.2.11) Solutions to the Stratonovitch SDE dx: = b(x,,t) t) dt + [xx,,t) dw.
(12)
define L-diffusions for the operator
D
L = v 00)
v (2) where v (0)) =,, ((()
Σ E.
2
k=
From this result, we can see that Eq. (12) describes the same diffusion process as the following Itō SDE:
D
dx, (b(x,, t) + r " (v((t) dt + E(x,, f) dw..
Σ
2 k=1
(0, where v (v(K)) - ( o,
Now can generalize the definition of SDE to the Riemannian manifold case. An SDE manifold
M can be defined by vector fields (0), v ...., ((M) on M. Let w be the RMM-valued Wiener process and Xo be an -M-vllued random variable serving as the initial value of the solution. The equation is symbolically written as
dx; v (x), ()
Σ
(k)) (xt, dw
(13)
Definition 26. An M-valued semimartingale x, defined up to proper stopping time is solution to the SDE Eq. (13) up to t if for all f 6 (M),
f(xi) f(xo) L.' ([100 (1) 8) +
Σ
(1)) oow? O
k=1
Proposition 27. (Thalmaier, 2023, Cor. 1.2.19) to the SDE Eq. (13). Then for all C 1000000MM (M),
Let L= v(0) v00) Eh v(k) and x, be the solution
Let
N/ = f(x) f(xo) S'" (xx,) t 0,
is martingale. In other words, the solution of SDE Eq. (13) is L diffusion the operator = (0) ED (R) v()2
C CONSTRUCTION OF QUOTIENT SPACE
In this section, describe rigorous construction of the quotient space and endow it with man- ifold structure. Please refer to the standard textbooks Lee (2018) for the systematic treatments. Assume that the total space M is Riemannian manifold and is compact Lie group. First give the formal definition of the group action.
Definition 28. Let G be group and M is Riemannian manifold. A left action of on M is map x M + M. (g.x) + g xx satisfying 91 (92 xx) = (9192) x and e = x, Vg1, 92 E 9 E M. An action is smooth if its defining map x M + M is smooth. We also reload the notation T. (x) == x for distinguishing g from its action on the manifold.
In the case where the Lie group acts on a Riemannian manifold, to draw meaningful conclusions, we would expect some compatibility between group action and the Riemannian metric, which the concept of an isometric action. Moreover, to ensure the topological structure of the quotient space so as to define useful constructions on the quotient space, concepts free action and proper action are introduced.
Definition 29. (1) A smooth action is said to be an isometric action if the map Ly M M,X g x is an isometry for any E G, i.e.,
(u, v)x = (u) (T.g)xx(V))
(14)
20

---

Published as a conference paper at ICLR 2026
(2) A smooth action is said to be free if for any x E M, g indicates g (3) A smooth action is said to be proper if the map g x M + M x M. (g.x)
H (g x, is a proper map,
meaning that the preimage of every compact set is compact.
For the properness, there is a convenient characterization.
Proposition 30. (Lee, 2018, Prop. C.15) Assume 9 is a Lie group acting smoothly on the smooth manifold M. The action is proper if and only if the following condition is satisfid: if (pt) is a sequence in M and (gi) is a sequence in such that both (pi) and (gi . ps) converge, then subsequence of (g:) converges. Particularly, every smooth action by a compact Lie group on smooth manifold is proper.
The group action typically represents a symmetry in the sense that points that can be transformed to each other by a group action are regarded as symmetrie, i.e., they are equivalent. Therefore, we can define an equivalence relation on M as (1) x(2) if 2g x() 9 x(2), The equivalence class with representative x is defined as the set of all points that are equivalent to x. The quotient space Q = M/G (as a set) defined under this equivalence relation, which consists of equivalence classes under the relation ~. The original space M is referred to as the total space. There is a natural mapping called projection that connects the total space and the quotient space, which maps any x E M to the equivalent class it represents. In this case where the equivalence is defined by Lie group, the projection mapping can be written as:
π : + 2, T(x) == {g xx-x 1 g).
Due this expression, the equivalent class in such a case is the orbit of the Lie group 9 at x. Therefore, it can be understood that an equivalent class is a "representation" (literal meaning; not the mathematical concept) of the Lie group, hence can also adopt manifold structures of under the mentioned "good" conditions. Also, the (M, 2, π) structure forms fiber bundle, in which context the equivalent class is also called fiber at x (x)., and this special fiber bundle induced from a Lie group action is called principal G-bundle.
Moreover, under certain conditions, the quotient space inherits the Riemannian structure of the total space M through the projection mapping.
Theorem 31. (Lee, 2018, Cor. 2.29) Let M be a Riemannian manifold, and G be a Lie group acting smoothly, freely, properly, and isometrically on M. Then the quotient space MiG has unique smooth manifold structure and Riemannian metric such that π is a Riemannian submersion.
We will assume the conditions, i.e., G be a Lie group acting smoothly, freely, properly, and iso- metrically on M, in the following development. Given that Q is a smooth manifold, the projection mapping induces linear mapping т. between the tangent spaces of the two manifolds. It introduces more structures in the total space M. In each tangent space M, we can define a subspace of it, called the vertical space, by the kernel of n.:
1/2 Ker Txx-
By this definition, tangent vectors in the vertical space can be understood that it does not move x in a way that alters the projection onto Q by n, so the movement stays within the equivalent class. The vertical space can then be understood as the tangent space of the equivalent class. As mentioned above, in this case where the quotient space is induced from the Lie group G. the equivalent class is a "representation" of the Lie group, hence the vertical space is "mirror" of the tangent space of the Lie group, which is in turn isomorphic to the Lie algebra g of the Lie group G.
To complete the whole tangent space, concept of horizontal space Hx is expected. In general, the horizontal space Hx is a linear subspace of T that makes up M by direct sum with Vx:
TxM Vx + Hx.
Under this direct-sum construction, any tangent vector v E T&M can then be uniquely decomposed into the vertical and horizontal components, v v +v Correspondingly. a vector field on M is called a vertical/horizontal vector field if it takes a vertical/horizontal tangent vector at every point. Every smooth vector field v on M can be expressed uniquely in the form v = v + where both the vertical and horizontal vector fields are smooth (Lee, 2018, Prop. 2.25). For future reference, we assign a convenient notation to the horizontal projection within M itself:
Px(v) := v Vv E T.M
21

---

Published as a conference paper at ICLR 2026
Nevertheless, the horizontal space Mx subspace that makes up the tangent space M by the direct sum with Vx is not unique. Therefore, a smooth correspondence from x to such an Hx is an independent structure, referred to as the "connection" in the fiber-bundle context. In the current specific case where M endows a Riemannian structure, we can uniquely define the horizontal space the orthogonal complement under the inner product in the tangent space:
Hx = Vx (Te.
(15)
which gives canonical "connection".
As would be expected, in contrast to vertical tangent vectors, a horizontal tangent vector represents a movement through different equivalent classes, corresponding Q. Therefore, we can construct the concept of horizontal lift which establishes a correspondence from vector field
movement the quotient space
to horizontal vector field on M.
Definition 32. Given vector field on Q, vector field u on .M is called horizontal lift of u, if ú is a horizontal vector field, i.., Ux E Hxx for all x E M, and is n-related to u by Txx (ūx) (x))).
Proposition 33. (Lee, 2018, Prop. 2.25) Given smooth connection Hx and assuming M - Q is a smooth submersion, every smooth vector field on Q always has a unique smooth horizontal lift to M.
If the connection is induced from the Riemannian structure of M by Eq. (15) and if the group action is isometric, then a nice compatibility can be derived. For a quotient-space tangent vector E Q some y 2, consider two ways construct tangent vector some point r (y) in the equívalent class. The first way is directly by the horizontal lift, which gives üx. which the unique horizontal tangent vector such that (ūx) = u. The other way is to first horizontal lift u another point x' x (y) in the equivalent class, then push it forward to the tangent space at x by a transformation that maps x' to x. Since both points lie in the same equivalent class and the Lie group acts the manifold freely, there exists unique group action g E G such that x = g x' = (x)),, ), so the resulting tangent vector is (Lg)axe ). Noting that (Lg)=x) preserves the metric between Tx.M and T (see Eq. (14)). we know that it also preserves the horizontal spaces, i.e., (Lg)ax' (Hx) = Hx so (Lg)ax (ūx)) E Hx. Moreover, T.x (x ((Lg)xx (ūx)) (л o Lg)=xx (ūx)) )= = Txx (ūx²) ) = also projects to the quotient-space tangent vector (noting that o Lg x for any g E G. and recalling the definition of horizontal lift), by the uniqueness of the horizontal tangent vector that projects to u, we have ux (Lg) (ūx)) or
equivalently,
x (Lg)=x(ūx). Vx I
(xEE-"y)),EE (y),ge
(16)
The unique existence of the correspondence from Tx(x) Q back to T by horizontal lift (Prop. 33) allows us to introduce Riemannian structure on from that on M. For any y and u()), (2) Ty Q, define:
(1))) u (2) ( 1
(17)
for any x E -y). This is well-defined since the right hand side is independent the choice of x due the porizontal-iitt-pus--foore (
compatibility (Eq. (16) and isometry (Eq. (14)):
(i()) (2) xg = ((Lg)xx((t (ū("), (Ly),x (u )) (ū u M
Subsequent constructions on the quotient space can be induced from this Riemannian structure. Due to its compatibility with the original Riemannian manifold M, these constructions have direct connections to their counterparts on M. Particularly, the Levi-Civita connections on Q and M follow the relation below.
Proposition 34. (Lee, 2018, Exercise. 5.6) Let VM and 2o denote the Levi- Civi-Civita connections on M, Q, respectively, where ve is constructed from the Riemannian metric induced from that of M by Eq. (17). Then for any vector fields u1)), (2)) on Q, denoting their horizontal lifts to M as (1)) .u (2), we have:
V uCn
(V. ((2))
a() ū()ū(
where Qe uti (2)) denotes the horizontal lift of the vector field De ni (2) Q
22

---

Published as a conference paper at ICLR 2026
C.I THE SHAPE SPACE RAN /SE(3)
For concrete and practically highly concerned example, we consider the shape space R3N /SE(3). In this example, each RR element is structured as:
(N) E RRY, with each x R
x
which represents the 3-dimensional coordinates of N points in IR (point cloud). The SE(3) group is composed of the 3-dimensional translation group and the 3-dimensional rotation group SO(3). Since the translation group is not compact, there does not exist a probability distribution that is translation invariant. We (as well as many others (Yim et al., 2023; Lin al., 2024)) hence represent the quotient space w.r.t this group by suppressing these equivalent DoF's by choosing a canonical translational position by anchoring the center of mass (CoM) of the point cloud at the origin, and consider the resulting CoM-free subspace .M := {x E R3N
1 Σi= x(() = 0} "nd consider the S3(3) action on it. Since the constraint is linear, this space .M is a linear subspace of IR and it is naturally a Riemannian manifold with the standard inner product of R3N An element of the SO(3) group is given by a x 3 rotation matrix for which we reload the notation g. The natural action of g on x is defined as x = (gx(1), gx(2), ())( Le the rotation is applied on each point of the system.
N
Unfortunately, S0(3) does not act freely (see Def. 29) on M in some degenerate cases, egg. all the points lie on straight line. So define the subset D C MM that SO(3) does not have free action on it; i.e., for any X E 'D there exists a nontrivial action g * E SO(3) such that g x x, indicating that D contains points that have a higher symmetry beyond SO(3). For converging sequence } in D which converges to x E .M as i - 0o, there exists a sequence {g in G such that 9i X; X,. Since the group SO(3) is compact and the group action is continuous, {9:); has convergent subsequence that converges to E G. which satisfies g x x. Hence E DD and therefore, D is closed. Subsequently, M := =M \ D is still a smooth manifold. As D is measure-zero in M (since any y E SO(3) is non-singular, the equation g = x reduces degrees of freedom of x), it is unlikely for a real simulation in M to hit the set D, making negligible difference algorithmically.
By removing the degenerate set D, S3(3) can now act freely and smoothly on M. Moreover, since the SO(3) action is isometric in the Euclidean space and M inherits the same metric, SO(3) also acts isometrically on M. Since SO(3) is a compact group, by Prop. 30, the action is also proper. Now that the action is smooth, free, proper, and isometric, by Thm. 31, the quotient space Q == M/SO(3) is a Riemannian manifold and the projection æ : . + Q is a Riemannian submersion. Since the each element in this quotient space Q an equivalent class containing equivalent point- cloud configurations, we refer to this space Q as the "shape space".
By the projection mapping ग MM- Q, the vertical space Vx Ker T.x can already be defined, which reflects the infinitesimal movements in M by group actions, which amounts to movements within the equivalent class π (xx (Appx. B). Since M is Riemannian manifold (tangent space inner product is inherited from the standard Euclidean inner product), we can define the horizontal space Hxx : (Vx) -TxM as the orthogonal complement of V& in M. Since V/ and Hz recover T&M by direct sum, any tangent vector E T&M can thus be uniquely decomposed as the addition of vertical component and horizontal component.
On this concrete example, the vertical and horizontal spaces can be expressed explicitly. Since the vertical space induced from group action, which acts freely, so this space is isomorphic to the tangent space of the Lie group, i.e., the Lie algebra. For G = SO(3), the Lie algebra so(3) is the set of antisymmetric 3 x matrices. So the vertical space is given by:
Vx= = {(Ax!), Ax
Ax A E so(3))
Using the 3-dimensional representation a (al a²) matrix can be represented as A-
E R for 53)) any antisymmetric 3 x Following this representation, Ax axx((),
aa a² al
a
2 a
'Here we choose a simple form of CoM by treating atoms equally weighted to avoid unnecessary notation complexity. In fact, any choice to determine one point in R3 from the N points suffices the reduction of the translation DoFs (as long as proper permutational invariance is guaranteed).
23

---

Published as a conference paper at ICLR 2026
where "x" denotes the usual cross product on R so the vertical space can also be written as:
Vx {(a x(), a x x(2), a x) a R
The horizontal space, which is the orthogonal complement of the vertical space, is given by
N
Hxx fv= (v(1), v()) E RAN
Σ ( 0.
0 0).
since the v()) vectors should keep the CoM fixed, and the orthogonal complement, they should v(9
also satisfy E (a x()) (Ei= x(6) v()) 0 for any R
Given the construction of the horizontal space (i.e., a "connection"), the horizontal lift for quotient- space tangent vector (and vector field) can be derived. As (O(3) acts smoothly, freely, properly, and isometrically on M, a Riemannian structure can be induced for Q from that of M, which is inherited from the standard Euclidean metric.
D PROOFS
D.I PROOF OF TIM.
Theorem Assume {xchetOoT a diffusion process on M, specified by the following SDE:
dx, = b, (x,) dt + o, dw,, Xo ~ Pprior
(6')
where b, is a Geeqiivariant time-dependent vector field on M. w; is the Wiener process on M that is G-invariant, and Pprior is a G-invariant distribution. Then the projected process {y, := π (es)) onto the quotient space Q == M/G is the solution to the following SDE:
of h(x.)) di dw Yo #prior:
dy: (T.b))
(7')
where: (1) T. is the fushed-forward vector field of b, induced by x. i.e. (Tr1))(y))
Tax. by (x,), which is the same for any X₃ E π (y:) due to the G-equivariance of by; (2) h ==
Txx. (D=M-G+1 e, for any x2 E ar-1(y)) is the mean curvature vector yt. where {e M an orthonormal basis of xx, M such that Vz, span ) = (3) w, is the Wiener process on Q; and (4) #prior is the furwerdforward distribution of Pprior, Le., its samples can be produced by Yo = T (Xp) where Xg ~ Pprior-
Proof. As is diffusion process on M given by the the SDE dx, be(x,) di + or dw,, by Prop. 27, x₁ is Li-diffusion and the generator is
C A 2 AM
Let {e be an orthonormal basis of M such that Hx. (ean(e:) M. Vx, span (e Then by the Riemannian submersion construction of M Q (see Appx. C), {ē, == _ M-G is an orthonormal basis of T. Let V and 0 be the Levi-Civita connections on M, Q respectively, where ve is induced from the Riemannian met- ric inherited from M. Using the local expression of the Laplace-Belle Beltrami operator (Def. 18), the generator is given by
o² A b. L, = b. +
of
(e.)).))) V.M e,,
2
t=1
M
M
b,
00
MMM
61 C;
+
1
C
24

---

Published as a conference paper at ICLR 2026
Then the process is the solution to the following Stratonovitch SDE
M
dx, - v(0) (xe: f)dt
v (xw:
M
where v = b, 2
MMa ei
and v (1) o,e; for i= i= 1, M.
By Def. 26, for all f E C (M),
f(xx)= 1 f(xo) + 1/' (vO s)ds Σ s)
Let f (2), then f f== then
f(x(xx)) f(T((x)) f(x(x))) //' (fo s)ds + Σ
of (xx, 5) dww:
= i(x((x))) + 1/" (I. s)ds Σ (T. (() odw:
M
by Def. 9. Since f is arbitrary, by Def. 26, yt π (xx) is the solution
M
dy: T. n. (y,, t)dt + Σ n (0) (yc, dw
We first need check that the projected vector field is well defined. In fact, we only need to check that x. b is well defined. Since b is g-equivariant, then for any 9 E G. b..g).((x)) = be(g x). = (I o Lb, = n_ (x)).. where we have used the chain rule in the second-last step, and the last step holds since equivalent class (x and g x are in the same equivalent class). By a notational equivalence that T. (x)) = n, (T) we know that т, (y) is the same on the equivalent class regardless of the choice of x in π (y), which implies that the projected vector field т, b, is well defined.
Then m, (b: x)) = T.
Lg and π projects to the same
Next, we calculate the expression of the projected vector field. Since Hx span(ei, eM-c), Vx= = span
eM }. we have
(é,, ifi & M- G,
Txx
0, if M G,,,
SO (v(+)) = o, for i = 1.... M- G and T (v())) 0 for > M- - G + 1. For the drift term, using Prop. 34, we have
2 M
T. (0) =
(VMe) e.)
To
2
i=1
2 M G
M
=T,br(y)
Tx (..M e,)
01
Σ
To
(V.Me) e,)
2
i=1
:=M-G+1
M G
M
=m.b(ly) 2
of
V ej
ot
To (V.M)) es)
2
-=1
M+M
=,.bi(y) of G
0 2
_ é;
h(y).
6,
25

---

Published as a conference paper at ICLR 2026
So the generator of the process yı is
É. - R.b;
o
V e;
oi
0i
2 M
Σ
M G
M-G
T,
L
h
+
0
Σ
0
2
2
6
i=1
o h
T, by
+
A
2
2
Then we can conclude that the projected process y1 == π (x,, is the solution to the following SDE
o/ di
dy: (T..b.
h(y))) dt + og dw.,
2
where b, is the push-forward vector field, h(y:) is the mean curvature vector Yr and Wi is the standard Wiener process on the quotient space Q.
D.2 PROOF OF THM. 2
In Def. 32, we define the horizontal lift a vector field that generates a deterministic flow. fact, for a stochastic process Q, we define the horizontal lift for it similarly. First, need to define the stochastic line integral, which the integration of a one-form along the trajectory of stochastic process.
Definition 35. (Hsu, 2002, Prop. 2.4.2) Let + a 1-form (Def. 11) on M and x2 the solution to the equation
dx, = (x,, + Σ (0
tx
Then
D S 0x, 0' ds = (' ((((0)) ds ((v(l)) dw
==1
Definition 36. (Baudoin et al., 2024, Def. 3.1.9) A semimartingale (xt)r on M is called horizontal
for every 1-form on M whose kernel contains the horizontal space H. one has f Ox, ds O, such that yo is point of Q. Then for given
for all t > 0. Let (ye); be semimartingale starting point Xo E π (yo). there exists a unique horizontal semimartingale x; on M such that x starts from Xo and π (xx)) y: for all > 0. This process (Xr); called the horizontal lift of (y:). from Xg.
Theorem 2'. The horizontal lift Eq. (7) has the following explicit expression:
dx; (b. - oi h(x))) dt ordw,, X Pprior,
(8')
where Px(v) == H is the horizontal projection in the tangent space of M, h is the horizontal lift of the mean curvature vector, and w, is the horizontal lift of the Wiener process on Q.
Proof. We only need to check the definition of the horizontal lift (Def. 36). Again, assume (M is a local orthonormal basis of M and Hx = span(e ... .MMG), V/ = span
C.)). Then by the Riemannian submersion construction of क M -> Q (see Appx. C), (é, == म. e; im 1- is a local orthonormal basis of Q. Let V and 7 be the Levi-Civita connection on M, Q, respectively, where 22 is induced from the Riemannian metric inherited from M.
26

---

Published as a conference paper at ICLR 2026
Now caleulate the generator of the SDE in Eq. (8'):
oi 2
of
Pb:
h
P(e:))
PV
(18)
b)
o
h
+
01
M-G
Σ i=1
(e² - (Veyes
2
2
Its projection under n, given by
G-G
Le = n. b.
oi h
+
e
M - (V
2
2
i=
which is the generator of Eq. (7). So we have π (x,) = Yt. where y, is defined in Eq. (7).
Let ④ be I-form on .M whose kernel contains the horizontal space everywhere. From Eq. (18), X, is the following SDE
dx: = (x), t)dt
(1) (x,, t) dw
where (0)
[2/ o7 h
ot
M-G
V. e,) v(t) oje
2 i=1
Then the line integral
L/ Ox, ds / M
Σ ((((0) dw, = 0,
since v(0 E H. @(v(t)) 0. So we can conclude that X, is the horizontal lift of y..
Corollary 3'. X1 (defined by Eq. (8)) has the same distribution on Q with x1 (defined by Eq. (6)). When or 0, VXo M. Eq. (8) has shorter trajectory length than Eq. (6):
1/ (Px. (b, PR. (z.)) dt & L by(x))),((())))) dt.
Proof. By definition of horizontal lift, T(x;) y; T(x)), Vt [0, 1]. then x1 (defined by Eq. (8)) has the same distribution on Q with X1 (defined by Eq. (6)). Since (Xx) = T(x)), then
X, = 9: 9t E G. Then by the G-equivariant property of b, we have
[ (bx(),), bx(xx)))
(be(g Xx), belg. xx))) M
((Lg.),x, by (xx), (Lgg) by (x:))) dt
(x.),,), br
((((,))),,,,)**)))>)))))" [ ((b,(&)) (,, H)) + (b,) ()(x,)))))) dt
> (b,()))).bb((()))))) Mdi di
/ (Px, (bx())), Px, (b(((())))
D.3 PROOF OF THM.
For the calculation of the mean curvature vector, we can embed the equivalent class क (y) into the total space where y Q. Thus, define the embedding कx M
27

---

The mean curvature vector has a nice geometric relation to the Published as a conference paper at ICLR 2026
by $ (g) x. For (y) the horizontal lift of mean curvature vector is defined by h(xx = EE" Ci= e,,,,) where {er is a local orthonormal basis of T&M and Vx - -4t , .ar).
volume of the equivalent class that helps us to calculate it.
Definition 37. Let + - M be an immersion. A smooth variation of Φ is a smooth mapping F: : (-c,c) + M satisfying:
For any DE E (-C, & F(;t) an immersion;
Pu =(;00) = :;
Proposition 38. (First variation of volume (Chavel, 1995, Exercise. III,14) The mean curvature vector h(x) satisfies the following formula:
Vol(g, = xo1(9,1)=-1
(h, dVool(G
dt
where v (2))
In local orthonormal frame (e:) Gf the volume of G is defined by
Vol(G, I. /det(G.) dg' A... A dg
where G = (P, + e.) M dg' is the dual form of e ie dg' = 1 if = and dā;();) = i + j.
Theorem 4'. Assume x, is diffusion process in the CoM subspace M c R&N, given by the following SDE:
dx, = (x,) dt + o, dw,,
where (x,) is CoM. The horizontal lift of the process 7xx,) is given by the following SDE:
vector field 44 E [0, T). w, is the standard Wiener process on
dt + ar PxPxx, o - (i(x))
dx,-
Px. (x,().
(9')
where the Px is the horizontal projection operator at x and h(x) is the horizontal lift of mean curvature vector. The explicit expressions of P and h are shown as follows:
2v ( v(i) X
N
N
(0)x)= = - -I (0 where J= Σ
x
21 10
10
Σ 1=1
0 T
E Raa
Proof. Let (e,,..., em} } be an orthonormal basis for TxM, which ordered such that Hx span(e),... ... eM eM and Vx span(eM-G+1, J. Then by the Riemannian submersion construction of M + (see Appx. C). fc, of Q. Let V.M and 2 be the Levi-Civita connection on M, Q, respectively, where 42 is induced from the induced Riemannian structure from M on Q. As shown in Appx. D.2, the horizontal lift of Eq. (8) has the generator
a local orthonormal basis
M-G
L.
b
o h
+
0
Σ
e
Me,
2
2
i=1
By Prop. 34, M-G
o,, = then
o2
G
C,
1,
h
+
Σ
e - (V e,)
2
2
i=1
28

---

Published as a conference paper at ICLR 2026
Since M is Euclidean space, Veye, e, = Σド= e, where of is the j-th component of e; and 05- = =000x. Since (x) = (x), then the generator becomes
o7 h1(x)) (x)
of
M-G
4 =
+
a.)
2
2
i=1
of [((x))
M-
M
(2ab((x)
oi
Σ e((0,e5)0 + c.000a.. + e0(0,)))0.
j.k=1
(Radb(())
of
((((x))
1
2
M-G
M
0,000
2
2
1=1
j,k=1
M
= -(Pab((x)
of
h(x)
+
0
(Px)
0,
2
J,k=1
T is projection operator. Then Le is the generator of where we use Px - e
dR; = Ps. (b:( - oं n(x))) dt + or PrPxx,
2
For the explicit calculation, recall that in this case, the tangent space Tx. of M at x has the following decomposition:
The vertical tangent space Vx: 1/2 = {11 x x (1), x (2),
1x xxx!N N
E R
|| E R*).
The horizontal space Hx, which is the orthogonal complement of the vertical space:
N
Hxx = {v- (v ... v(N) E
RNN
x ( x (t) = =o.
1=1
The horizontal projection mapping is defined by P2(v) = v =V -V V, E Tx,, and we can
find an explicit form of it. By definition, Σ (1) x ((1) = 0. then
N
Σ x( (
Assume v²= (1 x x() 1x x
=
1 1x xxx ) then
N
x v((
Σ x1) v())
N
(( x (1xx1)
19,
x((00) )1 11 -1((x))))))))
Σ
1x 0 Σ 5x00x0 (0)())) 1.
1=1
where we use the identity x0 x (1 x x()
= (x(1) x(()) (x(0) 1)x((1) Denote
N J
Σ xxx0 II Σ x x
29

---

Published as a conference paper at ICLR 2026
And we have = J- J- i= x(1) v()), and
vv = (1 x x())
I x x22 ... x x (0))
N
= =J
Σ
x (1)
x v0) xx.
Then
Pvv=v==v-> Σx x E
For the calculations of the mean curvature vector h, use Prop. 38. As G = SO(3). its local frame (the norm of each vector v2) is given by the following matrices:
0 0
0
0 0
0 0
e
0
1
e2
0
03
0
1 0
-1
0
0
Then the Gram matrix G is defined by G = (e,x, eyx). By direct calculations, we have G J. Then by Prop. 38,
k(x) v log det G.
Using Jacobi's formula in matrix calculus, d log det G = tr(JJ' dJ). Then by
N
N
J =
xx
Σ
T
(2x)) ( I- (0)x(0) + -x
0x
where 6y E R a one-hot vector at j. Then
tr J-1 OJ
-2tr(J- )x//) 26 J xx2
2tr((J
0x (1)
Then we have
10 (x)
Vlog det G -(tr(J J-
2
E TRAINING AND SAMPLING METHOD IN GENERAL CASE
Training Objective The diffusion model on the total space M is trained by the denoising score matching objective. Since the vertical components of the velocity are not strictly needed, we pro- pose to supervise the model only on the horizontal components and allow arbitrary vertical output of the model. Recall that the horizontal projection operator Px projects vector to its horizontal component, i.e. Px(v)= vH Thus the improved training objective is given by
((0)= = En(o) w(t (7 Px, (ve(x,,t) (afxx + 8'1 + 112
ODE Sampler After the training stage, Px. (v, t)) is an approximation of the ground truth vector field in the horizontal subspace. For the deterministic sampler, need to simulate the horizontal lift of the projected ODE, which is given by
dx; df
Px, v(x,
In practice, the ODE process is approximated by numerical solvers, e.g. the Euler method and Lunge-Kutta methods.
SDE Sampler For the stochastic sampler, we need to simulate the horizontal lift of the projected original SDE in Eq. (3). According to Thm. and Thm. 4, the lifted process is given by
dx, Px. (vo(xe, t) + geso (x4, t)) dt mmh((x))d / Px. dw.,
30

---

Published as a conference paper at ICLR 2026
where introduce the hyperparameter y for protein generation following Geffner al. (2025). The training and sampling algorithm summarized in Algorithm 2 and
Algorithm Training for Pprior (0,1))
Algorithm 2 Training for general Pprior
repeat
repeat
2: (Xo, Xi) ~ Pycer
(Xn,X1) ~ Pjoim- C N)
3: ~~ Pt
3: [
4:
deXo + 3,x1
4: Xt axxo + Bix1 + n
5: Take gradient descent step
5:
0,xo 8ix; + T/E
V. w(() ||xx (Da(xe,) () x1
Take gradient descent step
6: until converged
Va w(() |r (vo v.)ll
7:
until converged
Algorithm 3 Sampling
1: Nu Petior
2: for i == to K
1 do
At, f++1 -t. t,
ODE sampling then
5: Xt..1 Xt, + Px., vo(xe, Ai)
6:
end if
7:
if SDE sampling then
d. Pxx, (va(X,,, (.) + n1, (0, (:)) + 19. h(x::)
9: ~)
10:
Xx; + d:A4; + 2 At, Px,
11: end if
12: end for
F ADDITIONAL EXPERIMENTAL RESULTS
F.1 EFFICIENCY AND COMPLEXITY ANALYSIS
Complexity analysis. In this subsection, we give a detailed discussion the computational cost of our method. As mentioned in Thm. 4, we need to compute the inversion of the matrix I and the cross product for the horizomtal projection operator Px and the mean curvature vector h(x). For the calculation of I notice that Z is always 3 x 3 matrix, so construction cost of I is only linear O(N), where N is the number of atoms (linear ONN) cost for constructing I, and constant O(1) cost for inversion). The cross product is conducted atom-wise moom-,iee,, its computational cost is also linear O(N). So we can conclude that the overall computational complexity is O(N) for both Px and h(x).
We would like mention that the alignment operation adopted in the heuristic alignment-based diffusion strategies also has the same complexity. To see this, for aligning x E 3 towards y E algorithm constructs the optimal rotation matrix as (H H + where H : yx E 3x3x3 requires linear (NN) cost. In practice, the O(N) computational cost is negligible compared to the cost of gradient back-propagation through the neural network. A comparison of practical training times shown the following table.
R3xN the Kabsch-Umeyann
Methods
Original diffusion
GeoDiff alignment
Af3 align- Quotient- space diffusion
ment
training speed (iters/s) 4.19
4.07
4.08
4.10
All the results are tested on a single Nvidia A100 GPU, From the results, we can see that the addi- tional computational cost brought by the alignment and projection negligible.
31

---

Published as a conference paper at ICLR 2026
train_loss_epoch
0.385
0.755
0.365
0.55
0.45
0385
40k
80k
120k
160k
200k
Figure 4: Training loss vs. training epochs. We find that our training is stable in practice.
Numerical stability. In quotient-space diffusion model framework, need to caleulate the matrix inversion of J. which may have numerical issues for near-collinear systems of points. In practice, add eI term before conducting matrix inversion, that is, we calculate (eI -)-1 in practice, where I is the 3 x 3 identity matrix. This treatment is widely adopted in algorithms facing similar situations, e.g., the practical implementation of the Mabsch-Umeyaan alignment. Our typical choice of is le-8, and we found that the training process is stable under this setting. We have shown the training curve of the model on the protein backbone generation task in Fig. 4, which indicates no numerical issues arise during the training process.
algorithm for
F.2 THE IMPLEMENTATION OF G-EQUIVARIANT VECTOR FIELD
In Thm. require that the vector field is (3)-eqquuvariin In practice, this can be implemented by using a 50(3) -qquivrauva network architecture or applying data augmentation. In this subsec- tion. we justify that both of these choices are valid, such that the diffusion model can generate -3nvaainn distribution.
Diffusion model with data augmentation. The optimal solution of the Euclidean diffusion model is given by (x,) E[x [x,] (Song et al., 2021; Karras al., 2022). When the data distribution augmented by random rotation, the data distribution becomes SO(3)-invrriant, Thus, the opti- mal diffusion model recover the -inva-att data distribution. When the transition den- sity p(x/)s x1) is -qquivvaan,,,,,,,, i.e. p(x{(x1) = p(g x11g x1), Vig E SO(3), the optimal To this, let g SO(3) be an arbitrary rotation matrix. Since (g x,) = R[x1 lg xz). by the Bayes formula,
network is SO(3) eqqiivaraan.
Exppa (x; [xiplg x xx]x
E[x1 Ig xx]
E P (x1) [p(g x, [x]]]
E p (x.) x x 1g- x]]
Ex (xx) pp Xx 1x1
9 pa [g
x1 X2
(x 11 x1)]
B (9 xx
1 xx1]
g E[x]
where use the equivariance property of the transition density to get the second equality and the invariance property of Prarget to get the third equality. Thus, we can conclude that optimal solution under these conditions is validation that a well-trained neural network becomes nearly equivariant even if architecture is not equivariant.
Geffner et al. (2025) also gives empirical
Equivariant architecture. When the model is required to be 50(3)-equivarian) the optimal so- lution of the diffusion model not E[x, [x.). To figure the optimal solution, we consider the training loss at time t. The loss function at is given by
L (0) = E Do(x,, t) x1
= f a d2Nxx xx) (x,, 1)1 ||x 2(Do(xe, 1). xi))
32

---

Published as a conference paper at ICLR 2026
The optimal solution satisfies
(x,,,
Do is SDG) equivariant (()
argmin
The training loss can be simplified using the equivariant constraint:
(x / daN x1
f dda xt p(x1,
(|Dx(x))] + 1x1112= 2(Doo(x).
-J JRNN /S0(3)
dr. / dg
d p(x1,9 r.) (I)Degg r.)l ||x1||- 2
SSO(3)
Since Do is (30(())equivvrrai)) Do rz) (r())) then we have
(0) = JaNN /50(3)
dr /sore) dg
ddNxx p(x1 g ry) (I + ||x1|22 22g (ro x1))
Define p(r,)= SSo(3) dg sd xI p(x1,9 r.), and p(x1, rz) = p(x(.gT))) p(re) Then we have
(.) = / J (SO(3)
dr p(rr)) 2(DD -((oo)rr))) / 30(3)
dg dan p(x1,9 rrig x1)
JaaN
+ |
dr, SO(3) RSN R3NNSSO(3) (3O(3)
dg danns p(xi, (oxx,,rr))le |x1
So we can conclude that
Dg(r,,()= = / (8((3) dg f daNxx p(xi,g re) g- X1,
Dig'' r.) rr)= [ J dg f a p(x1, g rz) g g x1, Vg SO(3).
(8o(3)
Notice that
Dā(r.) JSorn dg f p(x1,g re) g- -1
Iso(3) dg ddNxx, x1 plg x w) rr. 1.g xI)xi
(5o(3) dg ʃ daNx, plg xi (plg Fe xx)
(SO(3) dg s ddNxx p(g x1 (r) x1)x1
SSor3) dg f ddNx1 plg x1)p(r; xx))
which is equivalent to the case Pdnta- SSo(3) dg p(g x1), iee using the augmentation by random
SO(3) rotation.
F.3 TRAINING AND SAMPLING ACCELERATION
In this subsection, we study the training and sampling convergence speed of different methods. For the training convergence speed comparison, we plot the generation performance measured by the precision AMR median metric with respect to the training epochs for previous heuristic alignment methods and our quotient-space diffusion model in Fig. 5(Left). We only focus on the first 100 epochs for all the methods. These models are trained with the same architecture ET-Flow (SO(3)) and training configurations on the achieves a similar convergence speed to the AF3 heuristic method, because both methods reduce the learning difficulty of the model, shown in Table This theoretical benefit leads to faster conver- gence than the GeoDiff alignment method. We also notice that the AF3 alignment method starts to get worse generation performance after training epochs. This happens due to the incompatibility between the training loss and the generation performance metric, as the AF3 method is originally designed for the protein structure prediction task, which is not evaluated by distributional metrics.
dataset. The results indicate that our method
For the sampling convergence speed comparison, we plot the generation performance measured by the precision AMR median metric with respect to the number of function evaluations (NFE) for the sampling process in Fig. 5(ight)). For all these methods trained the GEOM-DRUGS dataset, we use the Flow Matching ODE sampler (Lipman et al., 2023) with Euler discretization. From
33

---

Published as a conference paper at ICLR 2026
Training ng Epoch Generation Performance
Sampling Steps vs. Generation Performance
alignment
0.00
GeoDiff alignment
Precision/MR median (A44
Quol ent-space Diffusion Model
Precision/ARR mecian (A)
0.75
6.82-
0.80
0.70
0.70
Crigin Diffusion
0.00
AF3 Aligmment
aligniment
6.72
Quationt-spe spcco Diffusion Mocal
0.50
20
30
40
600
50
300
10
20
Epoch
NFE
Figure 5: Training and sampling convergence speed comparison on GEOM-DRUGS (Left) The relationship between training epochs and generation performance measured by the precision AMR median metric. (Right) The relationship between the number of function evaluations (NFE) for sampling and generation performance measured by the precision AMR median metric.
the results, we can observe that models trained with different strategies exhibit similar convergence trends (performance gradually degrades NFE decreases), our quotient-space diffusion framework consistently outperforms all baselines across every NFE setting.
F.4 QUOTIENT SPACE BEYOND IR /SE(3)
Our framework can generalize quotient spaces generated by symmetry groups beyond the special Euclidean group SE(3). Possible examples include the U(1) symmetry in quantum wavefunctions, the SU(2) symmetry in particle physics, and the SO(3) symmetry in higher (> 3) representation spaces for tasks including the mean-field electron Hamiltonian matrix prediction. In this work, we focus on the SE(3) case for its significant relevance scientific research (Abramson et al., 2024). Applications of our framework on the mentioned more diverse systems above are left as future work.
G EXPERIMENTS
G.1 MOLECULAR STRUCTURE GENERATION
This appendix summarizes experimental setup, which strictly follows that of ET-Flow (Hassan et al., 2024). We detail the datasets, model architecture, training, sampling, and evaluation. For a more comprehensive discussion of each component, refer the reader the appendices of their original paper.
Dataset. First, we evaluate framework on the molecule structure generation task. In this sce- nario, our goal is to generate the 3D coordinates of molecule given the graph structure of the molecule. We conduct the experiments on the GEOM datasets (Axelrod & nomezz-ommarreli, 2022), which provide structure ensembles generated by metadynamics in CREST (Pracht et al., 2024), and we focus on the GEOM RMM and GEOM datasets. Following the data pro- cessing and splits from (Hassan et al., 2024), we use the random splits with train/valida vati//alidation//esst of 243473/30433/100 for GEOM-DRUGS and 106586/13325/1000 for GEOM-QM9. In addition, data with disconnected molecule graphs are removed for GEOM-DRUGS (Hassan et al., 2024). Our reproduction is based on the modified data vata-processing pipeline following the released configs thus different from the results reported in the original paper.
Settings. We primarily follow the setting in (Hassan et al., 2024). We set the Gaussian distribution as the prior distribution on GEOM-QMS and use the harmonic prior for GEOM-DRUGS (Volk et al., 2023). Following (Jing et al., 2022: Xu et al., 2022). report the RMSD-based metrics, e.g. Coverage and Average Minimum RMSD (AMR) between generated and ground truth structure ensembles. We parameterize vo by using equivariant graph transformer architectures from ET- Flow (Hassan et al., 2024). including the O(3) and SO(3) equivariant variants, which also serves as a verification that framework is compatible with different backbone models. For training,
34

---

Published as a conference paper at ICLR 2026
AdamW the optimizer, and set the hyper-parameter to 1e-8 and (31, 52) to (0.9,0.999). We the dynamic gradient clipping as (Hassan et al., 2024; Hoogeboom et al., 2022b). The peak learning rate is set to 5e-4 for GEOM-DRUGS and 7c-4 for GEOM-QM9. The batch size is set to 48 for HEOM--RUGG and 128 for GEOM-QM9. The weight decay set to 1e-8. The model trained for 1000 epochs for both datasets. The noise scale is Euler solver for sampling. All models are trained on 8 NVIDIA A100 GPUs.
0.1. We also 50 time steps with the
Baselines. Following (Hassan al., 2024), we choose strong baselines trained on GEOM-DRUGS and GEOM-QM9 for challenging comparison. We report the performance of GeoMol (Ganea al.. 2021), GeoDiff (Xu al., 2022). Torsional Diffusion (Jing al., 2022), and MCF (Wang et al., 2023).
G.2 PROTEIN
This appendix summarizes experimental setup, which strictly follows that of Proteína (Geffner et al., 2025). We detail the datasets, model architecture, training, sampling, and evaluation. For a more comprehensive discussion of each component, refer the reader to the appendices of their original paper.
G.2.1 DATASET
For training, we utilize the Foldseek AFDB clusters (Dys) dataset as curated and described in the Proteína. This dataset is high-quality, non-redundant subset of the AlphaFold Database (AFDB), containing 588,318 cluster-representatiee protein structures with lengths between 32 and 256 residues. The dataset is annotated with hierarchical CATH labels, which leveraged dur- ing training. Our data processing and handling strictly follow the pipeline detailed in Appendix M of (Geffner et al., 2025).
G.2.2 MODEL. ARCHITECTURE, AND TRAINING
Our model architecture is the same the efficient, non-equivariant transformer proposed by (Geffner et al., 2025). Specifically, we adopt the variant that forgoes the use of computationally expensive triangle update layers. The model is trained using the conditional flow matching (CFM) objective. Key aspects of the training protocol from Proteina are preserved, including their novel Beta-Unifomi mixture for the time- sampligg distribution p(t), the use of self-conditioning and data augmentation with random rotations. All model and training hyperparameters, such as embedding dimensions, number of layers, attention heads, and optimizer settings, are kept consistent with hy- perparameters saved in their released checkpoint Mamalll The hyperparameters for the Masmall model are detailed in Table 4, in comparison with the larger models from the original Proteína paper.
G.2.3 SAMPLING
To facilitate a direct comparison with the publicly available Proteína checkpoints, we trained our model with identical hierarchical fold class conditioning mechanism. However, to ensure fair assessment of foundational generative capabilities, all experiments reported in our main text were performed in strictly unconditional setting. We applied the same sampling protocol across all models, using 400 sampling steps and enabling self-conditioning which consistently improved per- formance. No other guidance techniques, such autoguidance, tie ODE sampling to assess distributional fidelity and SDE sampling to explore the designability- diversity trade-off. We adapt the SDE formulation and its Euler Maruya-a numerical scheme, de- tailed in Appendix I of (Geffner et al., 2025), for our quotient space framework, while retaining all other configurations, such as the sampling scheduler and g(t), from the original paper.
utilized. We determinis-
G.2.4 EVALUATION
We evaluate our models rigorously adheres the metrics established and validated in the Proteína paper. We assess model performance using the standard suite of metrics in protein design:
Designability, Quantified by the self-consistency RMSD (scRMSD) protocol, using Pro- teinMPNN for inverse folding and ESMFold for structure prediction, with a success thresh- old of scRMSD less than 2A.
35

---

Published as a conference paper at ICLR 2026
Table 4: Hyperparameters for Proteína model.
Hyperparameter
MES
Mon-tri FS
M
Proteína Architecture
sequence repr dim # registers
768
768
512
10
10
sequence cond dim sinusoidal enc dim idx. sinusoidal enc dim fold emb dim
512
512
128
256
256
196
128
128
196
256
256
196
pair repr dim
512
512
196
seq separation dim
128
128
128
pair distances dim (xr))
pair distances dim (T(z))
128
128
128
pair distances min (A)
pair distances max (A)
30
30
30
# attention heads
12
12
# tranformer layers
15
15
12
triangle layers
# trainable parameters
200M
200M
60M
Proteína Training
# steps
200K
360K
150K
batch size per GPU
4
10
5
GPUs # grad. acc. steps
128
96
16
Diversity, Measured in two ways: by the average pairwise TM-score among designable samples, and by the number of distinct structural clusters identified by Foldseek at TM- score threshold of 0.5.
Novelty. Assessed by calculating the maximum TM-score of each designable sample against reference structures in the PDB and AFDB databases.
We also adopt the novel probabilistic metrics introduced by (Geffner et al., 2025), to measure how well our model captures the true distribution of protein structures:
FPSD. Measured the distributional similarity between generated and reference structures in the feature space of a pre-trained fold class predictor.
IS. Evaluated both the quality and diversity of samples based on the confidence and entropy of fold class predictions.
IJSD. Quantified the similarity between the categorical fold class distributions of generated and reference sets.
is noteworthy that we have omitted the Diversity and Novelty metrics from our main text avoid comparisons with potentially inaccurate results in the literature. This decision is based on a bug recently identified in the alntmscore output of FoldSeek versions prior v10 (release 10-941cd33), which renders many previously reported TM-based metrics incorrect (also found in (Daras et al., 2025)). To provide a controlled and accurate benchmark, we conducted our own analysis using the FoldSeek v10 (release 10-941cd33) led33). We limited this re-evaluation to the released small Proteína model and our corresponding model trained in the quotient space. The full results of this comparison summarized in Table 5.
36

---

Published as a conference paper at ICLR 2026
Table Complete performance comparison of the released Proteína checkpoints against our version in the quotient space. Best results marked in bold.
Model
Designabiiity (%)
Divensity
Novelty
FPSD vs.
ES
ESD
Cluster?
PDBJ
AHDB]
AHDB]
SDE Sampling
JApr 0.35
96.0
0.44 (209)
0.50
0.86
0.91
386.5
378.2
97/7779
2.17
L.73
0.35 + ours
97.6
(197)
0.48
0.86
0.91
274.7
277.1
2.24/6.89/20 99
1.68
1.55
JApzal 0.45
92.2
(253)
0.49
0.84
0.90
332.9
320.1
18835 01/20.11222220 22
1.93
1.49
JMi 7 0.45 ouTS
92.6
(253)
0.90
244.5
246.3
2.2416 68/23.4
1.28
.M35 0.50
89.2
0.57 (255)
0.48
0.83
0.89
306.2
290.8
1866 92/21.15
1.81
1.36
.MA 7 OUISS
90.2
(231)
228.0
2,2516 59/25
1.17
ODE Sampling
0.90 (62)
21.93
Mgg ours
15.6
0.87 (68)
0.43
0.80
0.86
69.94
17.56
2.57/6.40/52144
0.41
0.11
37


Reviews:
[{"review_id": "0cb7xKXqB8", "reviewer": "Reviewer_Th24", "summary": "This paper establishes a formal mathematical framework for diffusion modeling on quotient spaces, with applications to molecular structure generation under SE(3) symmetry. The key idea is to construct diffusion processes directly on the quotient space M/G, then lift them back to the total space M for practical implementation using horizontal projections. The framework reduces learning difficulty by removing redundant components corresponding to group actions while guaranteeing correct sampling from the target distribution.", "questions": "1. How does the method perform on other symmetry groups beyond SE(3)?\n2. How does the computational cost of projection scale? Are there approximations or efficient implementations for large-scale problems?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "1. The paper provides a comprehensive mathematical framework grounded in Riemannian geometry and stochastic calculus. Theorems 1-4 give explicit characterizations of the projected diffusion process and its horizontal lift.\n2. The formulation of diffusion models via horizontal lifting on quotient spaces is innovative. It generalizes previous work on equivariant diffusion and provides a unifying geometric perspective that can be applied to various symmetry groups.\n3. Unlike heuristic alignment strategies, the proposed method guarantees correct sampling from the target distribution while reducing learning difficulty by removing redundant degrees of freedom.", "weaknesses": "1. The paper doesn't discuss the computational cost of computing the horizontal projection operator $P_x$ and mean curvature vector $h$ at each step. For the SE(3) case, these involve matrix inversions and cross products - what is the actual time comparison?\n2.  The heavy reliance on differential geometry, Riemannian manifolds, Lie groups, and stochastic calculus makes the paper inaccessible to much of the machine learning community. While mathematical rigor is valuable, the presentation could benefit from more intuitive explanations alongside formal derivations.\n3.  While consistent, the improvements over baselines are relatively small in some cases, which raises questions about practical significance.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "5lxhTNSJV1", "reviewer": "Reviewer_k95T", "summary": "The paper investigates the often overlooked issue of additional training difficulty due to symmetries in diffusion models. The paper proposes considering the projected diffusion dynamics to the quotient space of interest in order to remove the original symmetries and simplify the learning problem. The paper provides a mathematical analysis of this new diffusion process and provide an efficient approach to training and sampling within this quotient space. This approach is then validated empirically in conformation generation for small molecules and protein backbone generation demonstrating improved results.", "questions": "* Are there ever numerical instabilities from the projection operator or the horizontal lift of mean curvature due to matrix inverse for near collinear systems of points?", "limitations": "", "rating": 10, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "* I think the paper is very well done. It answers an important question (that I've also been considering) for diffusion models applied to common data modalities in scientific applications.\n\n* The mathematical derivation of quotient space diffusion is handled nicely and the resulting training and sampling algorithms are also nicely implemented in the ambient space to sidestep the complex form of the quotient space.\n\n* Very nice summary of previous approaches of handling the additional difficulty of learning symmetries in Section 3.4, as well as an explanation as to why there are issues with sampling time mismatches. This is a useful exposition as this has been less emphasised in the literature.\n\n* The empirical evaluation is well done with definite benefits reported for the relevant and practically important examples of conformation generation for small molecules and protein backbone generation.", "weaknesses": "* It is perhaps less clear how SMC/reweighing-based approaches for reward tilting could applied to this framework as I am not totally sure that transition probabilities computed from equation 9 represent the correct transition probabilities that you would want from the actual quotient space diffusion.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "CqC6CwGK6a", "reviewer": "Reviewer_mNjY", "summary": "This paper addresses the critical problem of intrinsic system symmetries (e.g., $SE(3)$ symmetry in molecular generation) that complicates generative modeling. Through a theoretically rigorous approach, it establishes a formal framework for diffusion models on a general quotient space $\\mathcal{M}/\\mathcal{G}$. The method involves projecting the standard SDE onto this quotient space and subsequently deriving a new \"horizontal lift\" SDE in the original space, which is constrained to movements that are purely horizontal to the group action. This design effectively reduces the learning space for the neural network, thereby lowering the learning difficulty. Critically, this new SDE guarantees consistency between the training objective and the inference process, ensuring sampling compatibility, which prior heuristic alignment methods lack. The paper provides a detailed and persuasive implementation, with exhaustive comparisons against these heuristic approaches, demonstrating excellent performance on key benchmarks: small molecule structure generation (on GEOM-QM9 and GEOM-DRUGS) and protein backbone design.", "questions": "1. **Guarantee of Recovering the Full G-Invariant Distribution**: The paper's primary focus is to ensure the SDE matches the quotient space distribution and maintains sampling consistency. This is a clear success. However, a question remains about learning the full G-invariant target distribution $p(x)$. For a finite dataset (which is not perfectly G-invariant), does the horizontal lift SDE guarantee that it learns the correct probabilities for all equivalent points $g \\cdot x$? For example, if the dataset contains only two samples $x_1, x_2$, the true target distribution should be non-zero and have the same probability on the entire orbit $ \\\\{g \\cdot x_1, g \\cdot x_2 | g \\in \\mathcal{G} \\\\} $. Does the proposed method recover this full manifold, or does it primarily learn the quotient-space projection well? A deeper discussion on this geometric probability coverage would be valuable.\n\n2. **Comparison to Other Lie Group Diffusion Methods**: Could the authors comment on how this quotient-space framework compares, in terms of theoretical advantages or disadvantages, to other recent work on diffusion for symmetric data? For instance, methods that operate directly on Lie groups (e.g., \"Trivialized momentum facilitates diffusion generative modeling on Lie groups,\" ICLR 2025, \"Flow matching on general geometries.\"ICLR 2024) also aim to simplify the learning process.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "1. **Significance and Theoretical Contribution**: This paper addresses a highly valuable and prevalent problem in generative modeling: the inconsistency between training objectives and sampling processes caused by intrinsic system symmetries. The work is exceptionally timely, as similar issues are being reported in various domains, such as periodic crystal generation [1, 2]. The authors are encouraged to cite similar works to further strengthen significance. The paper provides a powerful and unifying theoretical lens to explain and, more importantly, solve these inconsistencies.\n\n2. **Elegant and Sound Framework**: The proposed framework is mathematically elegant and theoretically sound. The derivation—projecting the standard diffusion SDE onto the quotient space $\\mathcal{M}/\\mathcal{G}$ and then deriving its corresponding \"horizontal lift\" SDE back in the original space—is a principled and clean solution.\n\n3. **Excellent Analysis of Prior and Heuristic Methods**: A major strength is the paper's insightful analysis of prior work within its new framework. The discussion (and Table 1) that masterfully categorizes and explains the shortcomings of previous methods (like data augmentation or inference-only alignment) is excellent. The critique of heuristic alignment strategies (GeoDiff, AF3) is particularly brilliant, clearly articulating how they introduce high variance or suffer from sampling incompatibility.\n\n4. **Clarity and Intuition**: The paper is well-written. Figures 1 and 2, in particular, are exceptionally clear and provide excellent intuition for the complex mathematical concepts involved, making the core idea highly accessible.\n\n5. **Strong Empirical Validation**: The authors provide thorough ablation studies and experiments on challenging, standard benchmarks (GEOM-QM9, GEOM-DRUGS, and protein design). The empirical results largely demonstrate the superiority of this principled approach over existing methods.\n\n[1]\"Equivariant diffusion for crystal structure prediction.\" ICML,2024\n\n[2]\"Kinetic Langevin Diffusion for Crystalline Materials Generation.\"ICML,2025", "weaknesses": "1. **Missing Discussion on GEOM-DRUGS Results**: In Table 3, the performance on the GEOM-DRUGS dataset, while strong, appears to be below that of the MCF baseline. This is a noteworthy result, but the paper does not provide any discussion or analysis for this specific comparison. A brief discussion of this would strengthen the experimental section.\n\n2. **Missed Visualization Opportunity**: While Figures 1 and 2 are excellent for conceptual understanding, the paper would be significantly improved by a new figure that visualizes the dynamics of the proposed horizontal diffusion process itself. For example, a qualitative visualization of the distribution's evolution on a simple manifold (e.g., $\\mathcal{M}/SO(2)$) would be highly illustrative and provide a powerful complement to the static diagrams.\n\n3. **Minor Formatting Suggestion**: In Table 3, the results for \"GEOM-DRUGS\" are a key contribution. Consider using bold formatting or out-standing rows to make it more clearly.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "37XXmTyEe9", "reviewer": "Reviewer_srZ8", "summary": "This paper addresses the problem of learning diffusion models for distributions with specific group symmetries. The proposed framework learns the score function on the quotient space with respect to the group, avoiding the need to model components corresponding to group actions. They start by projecting the standard forward SDE onto the quotient space using a Riemannian Ito's lemma. They then observe that the vector field on the quotient space admits an orthonormal decomposition into horizontal and vertical components, where only the horizontal component is relevant as the vertical component corresponds to infinitesimal group actions. Given the non-trivial geometry of the quotient space, they re-define a process in the original space that retains only the horizontal movement, called horizontal lift. Experiments focus on molecular structure generation, which is characterized by the SE(3) symmetry. Results demonstrate that training with this approach outperforms two alternative heuristics for handling symmetries.", "questions": "- Just to clarify my understanding, could the first paragraph of Section 3.4 be related to [1], where the authors highlight a mismatch between using a translation-invariant score network and a target score that is not translation-invariant? In other words, given $x_t$ and $g\\cdot x_t$, which are equivalent, the network is not able to distinguish between them while the corresponding target scores differ, leading to high variance during training?\n- By avoiding the modeling of these unnecessary degrees of freedom using the proposed method, do you observe that this leads to better results when sampling with fewer steps compared to models trained using the other heuristics?\n\n- line 961 'up' instead of 'bp'\n\n[1] \"Equivariant Diffusion for Crystal Structure Prediction\" Lin et al, ICML 2024", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper tackles the problem of handling symmetries in a novel and principled manner in contrast to previous approaches that rely on heuristics to deal with the symmetry problem. It is the first work I have seen that attempts to constrain the diffusion process within the quotient space.The authors also provide intuitive explanations of the main differences between their proposed method and existing approaches, supported by Table 1 and Section 3.4. Additionally, they include a comprehensive background on Riemannian geometry and stochastic calculus in the appendix to help readers follow the proofs of the theoretical claims presented in the main text. Experimental results demonstrate that applying their framework to both molecular and protein data improves the performance of the model compared to the use of conventional heuristic methods for handling symmetries.", "weaknesses": "I believe that, in principle, the approach proposed in the paper can be applied to any diffusion model by keeping the backbone architecture and changing the training, if I am not wrong. Therefore, it would have been interesting to see whether the improvements observed for the selected model also hold across different architectures. In relation to this, the results table could have been more informative if, for each baseline, the specific heuristics used to handle symmetries were explicitly stated. \n\nAdditionally, while the theoretical advantages of the proposed model are clearly illustrated in Table 1, I would have expected these advantages to be examined empirically as well, rather than being evaluated only based on final performance metrics. For example, I would expect a similar convergence behaviour between the proposed approach and the AF3 heuristic, since they are both removing the unnecessary DOFs and also lowering the variance. With respect to GeoDiff,I would expect the proposed method to achieve faster convergence during training, as it removes the equivalent DOFs. \n\nFinally, the proposed framework rely on a horizontal projection operator that involves an inverse of a matrix that depends on the input dimensions, i.e. number of points in the cloud, or atom in the molecule. It would have been useful to include a discussion on whether this operation could become a bottleneck when scaling the method to larger molecules, and whether the matrix inversion might introduce numerical instability.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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:

Decision: Accept (Oral)