Rao-Blackwell Gradient Estimators for Equivariant Denoising Diffusion
Authors: Vinh Tong
Dung Trung Hoang
Anji Liu
Guy Van den Broeck
Mathias Niepert
Venue: NeurIPS 2025 poster ()
Keywords: Equivariant Diffusion Models
Equivariant Flow Matching
Rao-Blackwell
variance reduction
Primary area: deep_learning
TL;DR: We reduce training variance in equivariant generative models using a low-variance gradient estimator, improving stability and performance across molecular, crystal, and protein generation tasks.

Abstract: In domains such as molecular and protein generation, physical systems exhibit inherent symmetries that are critical to model. Two main strategies have emerged for learning invariant distributions: designing equivariant network architectures and using data augmentation to approximate equivariance. While equivariant architectures preserve symmetry by design, they often involve greater complexity and pose optimization challenges. Data augmentation, on the other hand, offers flexibility but may fall short in fully capturing symmetries. Our framework enhances both approaches by reducing training variance and providing a provably lower-variance gradient estimator. We achieve this by interpreting data augmentation as a Monte Carlo estimator of the training gradient and applying Rao–Blackwellization. This leads to more stable optimization, faster convergence, and reduced variance, all while requiring only a single forward and backward pass per sample. We also present a practical implementation of this estimator—incorporating the loss and sampling procedure—through a method we call Orbit Diffusion. Theoretically, we guarantee that our loss admits equivariant minimizers. Empirically, Orbit Diffusion achieves state-of-the-art results on GEOM-QM9 for molecular conformation generation, improves crystal structure prediction, and advances text-guided crystal generation on the Perov-5 and MP-20 benchmarks. Additionally, it enhances protein designability in protein structure generation. Code is available at https://github.com/vinhsuhi/Orbit-Diffusion.git.

Rao-Blackwell Gradient Estimators for Equivariant Denoising Diffusion
Vinh Tong 1,2 Trung-Dung Hoang"; Anji Liu", Guy Van den Broeck Mathias Niepert University of Stuttgart, IIPPR-IS, SULA, *University of Bern, NNational University of Singapore vinh. tong@ki wwi-stuttgart.de
Abstract
In domains such molecular and protein generation, physical systems exhibit inherent symmetries that are critical to model. Two main strategies have emerged for learning invariant distributions: designing equivariant network architectures and using data augmentation to approximate equivariance. While equivariant architec- tures preserve symmetry by design, they often involve greater complexity and pose optimization challenges. Data augmentation, on the other hand, offers flexibility but may fall short fully capturing symmetries. Our framework enhances both approaches by reducing training variance and providing provably lower-variance gradient estimator. We achieve this by interpreting data augmentation as a Monte Carlo estimator of the training gradient and applying Bao-Blackwellizzttin leads to more stable optimization, faster convergence, and reduced variance, all while requiring only a single forward and backward pass per sample. We also present a practical implementation method we call Orbir Diffusion. Theoretically, we guarantee that our loss admits equivariant minimizers. Empirically, Orbit Diffusion achieves tatte-of-th---o art results GEOM GEOM-QM9 for molecular confor- mation generation, improves crystal structure prediction, and advances text-guided crystal generation the Perov-5 and MP-20 benchmarks. Additionally, it en- hances protein designability in protein structure generation. Code is available at https://giis httgggttt....... wwtpw//wwwww////eiiooo.o. wtw.://rrbit--iffuoo.o...o git
This
this sstimator-incorpraatin the loss and
sampling wrocedure-throuse
Introduction
Diffusion models have emerged as powerful methods for modeling complex distributions (Ho et al.
2020; Song et al. 2021a by Karras al. 2022); with applications in domains such as molecular and
protein generation (Zhang et al. 2023 Vignac et al. 2023] Anand & Achim 2022). Many physical systems, such as molecules or crystals, exhibit inherent symmetries. For example, molecule's physical properties remain unchanged under rotations in 3D space (Hoogeboom et al. [222]] Jing cct al.) 2022). Modeling such data requires learning distributions that are invariant under the action of group G. This setting is naturally captured by the notion of G-invariant distribution ((x0) that are
invariant under transformations from G (Chen et al. 2024 Köhler et al. 2020).
Two main strategies have emerged for learning G-invariant distributions: (1) designing equivariant network architectures, and (2) using data augmentation to approximate equivariance. Equivariant architectures, such as equivariant denoisers, ensure symmetry by construction (Hoogeboom et al.
2022, Klein et al. 2024 [gashov et al.) 2024), but are often less efficient due to increased architectural complexity and can pose optimization challenges (Brehmer al.) 2024 Abbe & Boix-Adsera (022). In contrast, data augmentation is flexible and widely used alternative that approximates equivariant training by sampling transformed versions of the input. While this approach scales easily and has
Equul contribution.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

become increasingly popular in the community, especially for large-scale models (Abramson et al. 2024b Geftner (025). its effectiveness capturing symmetries may vary depending on the group and application domain. In some settings, such as molecular dynamics and structural biology, explicit equivariance remains beneficial (Anand Achim] 2022 Batatia et al. 2022 Zaverkin et al.) 2024). In this work, we propose a framework that provably improves both approaches: we introduce a novel form of implicit data augmentation by computing the denoising target as a weighted average over group orbits, which reduces variance and improves the training of equivariant denoisers.
We revisit data augmentation from a principled perspective and interpret it as Monte Carlo estimator of the gradient of symmetrized loss. This loss is defined over fully symmetrized dataset, which yields an empirical distribution invariant under G (Chen et al.) (024)) While traditional data augmen- tation uses one or few samples from this augmented dataset, we instead apply --lakkellizzatioo to derive a new estimator with provably lower variance. Specifically, we decompose the gradient as an outer expectation over data and an inner conditional expectation over group actions. Replacing the noisy sample-based target with its conditional expectation yields a lower-variance gradient estimator, while preserving equivariance.
To translate our theoretical insights into practical method, we develop an efficient implementation of the proposed gradient estimator. This approach integrates the symmetrized loss and variance reduction into a modified training objective, without increasing the computational cost. Our implementation requires only a single forward and backward pass per sample, and it is compatible with both cquivariant and non-equivariant architectures. We refer this practical method as Orbit Diffusion. By implicitly applying
Blackwelllizaiinn through tailored loss formulation and sampling scheme, Orbit Diffusion enables stable optimization and improved generalization across a wide range of symmetry groups and tasks.
We provide theoretical guarantees that our symmetrized loss admits equivariant minimizers and that our gradient estimator has strictly smaller variance compared to existing methods. Empirically, we demonstrate strong performance across multiple domains. Our method achieves state-of-the- -frtte-art results on GEOM 1-M9 for molecular conformation generation, enhances crystal structure prediction, and improves text-guided crystal generation on the Perov-5 and MP-20 benchmarks. Moreover, our approach is compatible with non-equivariant denoisers; in particular, it improves the designability of protein structures generated by PROTEINA (Geffner et al. 2025).
2 Background
Groups. A groupp is a mathematical structure comprising a set G and a binary operation m : G G that combines two elements of G. A group action g E G defines how the group G acts on a set Q. such as a set of geometric objects We restrict our attention to locally compact isometry groups. Locally compact isometry groups encompass broad class of transformation groups that preserve distances and possess a well-behaved topological structure. Examples include the permutation group S,,, the orthogonal group (dd), and the special orthogonal group (O(d).
Invariance and Equivariance. A function f Ω + R is said to be G-invariant for all g E G and x E Ω, it satisfies f(g o x) = fxx). This means the function value does not change under the action of any group element. A function f Ω + Ω is said to be G-equivariant if for all x E 9, f commutes with any group action g E G: f(g x) g f(x).
Invariant and Equivariant Distributions. A probability distribution p(z) defined on a set Q said to be G-invariant under the action of a group G if the probability of any measurable subset A C Ω remains unchanged under the transformation induced by any group action g E G p(g o x E A) =
p(x E A). A conditional distribution p(y x). where E Ω and y E Q. is said to be G-equivariant for all g E G, the following condition holds: p(g y g x) = p(y æ)..
Group Symmetrization. Let SG be the symmetrization operator under group G, transforming any distribution p(x) into a C--mvariant-invariinn distribution, denoted as the G-symmetrized distribution:
(afp)(>) = / p(g x) duc(o), where Ng is the Haar measure G.
WWhen the group acts on a vector space V. we do not distinguish between the abstract group element E G and its linear representation p(g): G + GL(V). ). For simplicity, we write g o to denote the action (())
2

---

Diffusion Models and Equivariant Diffusion Models. Diffusion models (Ho et al. 2020) are class of generative models that construct complex data distributions by iteratively transforming simple noise distributions through learned denoising process. Formally, given data g( the forward process generates a sequence 24 over time E [0. 7) using stochastic differential equation (SDE) (Song et al. 2021b) or a discrete Markov chain (Ho et al. 2020), such as:
Ft 04 + ore, ~~ (0,
where and de are a time-dependent scaling factor and noise factor that determines the level of noise added at each step, respectively. The noise should be increasingly added to the sample so that at time / = T, ((xT) R (0,
The reverse process, parameterized by neural network (22, 1), approximates a clean sample æo given its noisy version The training objective typically involves minimizing reweighted form of the denoising loss (Ho et al. 2020 Song et al.) 2021b Karras al. (022):
(1 xo |?
where (t) is a time-dependent loss weight. For notational simplicity, we omit this term throughout the remainder of the paper. To sample from the diffusion model, we begin with noise vector 2T NO I) and iteratively apply the learned reverse process to transform into a data sample 2n using the trained model db. The reverse process can involve solving an ODE or SDE numerically (Song
et al.
2021b Karras et al. 2022 Lu et al. 2022; Tong et al. 2025
Equivariant diffusion models extend standard diffusion by enforcing equivariance of the neural network denoiser. Specifically, the denoiser Φe is said to be G-equivariant if it satisfies dog 0g(go t) = g @0( t) for all g G.
3 Method
(b)
Algorithm I: Orbit Diffusion with RB.
2o (mn)
I1 tx 2o
1: Sample a data point Zo ~ q (xo)
xe xx ~ q (ste 2o)
2: Sample noise level t
|| 12, t) - 20l
3: Generate noisy sample ~ ( mo)
4: for i 1 N do
Sample group element g (9) ~ (g)
(c)
(1) To p( I xx))
Compute orbit sample ap = g o
Compute w = oh (2a a (v
(
(x1 260
8: end for
9: Approximate B[xe [xe] with SNIS:
Σ
() f0
Σ to
target Φ. (zo xx)
N
proposal p(zo w2)
E[zu 21)
④ (mo)
18 [1t [re]
zo
10: Backpropagate gradients of the loss:
| 1)
iB(2 m4]]
p(a. t) Biuo xx)]
low weight
high weight Iow weight
Figure 1: Gradient estimation strategies for training (approximately) equivariant diffusion models: (a) Sampling from the symmetrized joint distribution to obtain 2g and 24-- (b) The standard data augmentation approach, which directly uses these samples for training. (c) The proposed method, leveraging self-normalizing importance sampling (SNIS) to estimate the inner conditional expectation. Both (b) and (c) require single neural function evaluation per gradient step, but (c) has lower variance than (b). The pseudo-code for the Rao-Bla B-acklllwwl estimator with SNIS is shown on the right.
Let G be a symmetry group, such as the group of Euclidean rotations. Our goal to learn a G- invariant data-generating distribution g( However, the observed distribution q (from which we obtain training samples) is generally not G-invariant due to dataset biases. For example, in molecular datasets, each molecule may be stored in canonical but arbitrary orientation, even though physically

---

all rotated versions are equally probable under As a result, training directly on á would lead to model that may not respect the underlying symmetry.
This issue is widely recognized in the literature, and two standard solutions are:
1. Equivariant model: Train with the diffusion loss in Equation 2 using G-equivariant network ((x,, f).
2. Data augmentation: Train non-equivariant with Equation random group actions to encourage approximate G-invariance.
augmenting data with
The second approach has been widely adopted several high- profile convevaiva models achieve strong results through augmentation (Abramson et al. [2024b] Geffner et al.) (022)). However, empirical evidence shows that augmentation offers no benefit for already equivariant models, result we formally prove in Appendix B.2
3.1 From Symmetrized Loss Variane Gradient Estimators
To unify these approaches, consider the symmetrized data distribution from Equation (1), and the forward noising kernel (x, zo) (e.g.. a Gaussian in denoising diffusion) with marginal of (zz). The symmetrized diffusion loss at time
()) Go of (-1zg) [10(,,, t) zo1??]
(3)
Both the equivariant model and the data augmentation approach can be viewed as implicitly minimiz- ing (Ø) (Chen et al. 2024). The true gradient of this loss is
VoL. Ex Bz (((00,,) zo)],
where the expectations are taken over the full symmetrized joint distribution.
In practice, do have access the exact expectations in Equation ( Instead, construct Monte Carlo gradient estimator by sampling To ~ 7 (or qe in the augmented/equivariant case), then
sampling Vy ~ q. ro), and using the single-sample estimate VoC. = 2(((z,, - zo). This estimator is unbiased, but, as in other diffusion training setups, it can exhibit high variance (Kingma
ct al. 2021] Xu et al.) 2023 Nichol Dhariwal 202T). Increasing the batch size reduces variance but requires more forward passes of the neural network é, raising computational cost.
In this work, we introduce class of Rao Blackwellized gradient estimators that provably reduce variance while remaining unbiased, and can be applied to both the equivariant and augmentation- based training strategies.
3.2 Rao-Blackwellized Gradient Estimator
Our key observation is that Φ(2,, t) does not depend on This allows us to move it outside the inner expectation Equation (4. yielding
E, En ((0(21))) E En (anl (xgls))
(5)
B[zo [12]]
Replacing zi with its conditional mean EFzo æı] yields a Rao- Blackwellized (RB) gradient estimator, which remains unbiased and has variance no greater than the iiginal--ttricte less unless dn I is deterministic, a situation that rarely occurs in generative modeling where zy is typically stochastic given Iy. This improvement requires additional neural network evaluations, only more accurate target. To avoid custom backward passes, we can minimize:
RRB (1) = Ennng [10((1e;1) - Elzoo xz m=a]][]]
(6)

---

Variance reduction guarantee
Theorem 1. Let V be the Monte Carlo gradient from Equation and (RB)) Equation
be from If [20 æi] can be computed exactly, then
@
Var Var(V(nB)) VVa((Vo)
with strict inequality unless 2o Dt is a Dirac delta.
The challenge now is estimating the loss target E[zp zy] accurately and efficiently. We address this next using self-norm importance sampling (SNIS).
3.3 Estimating the Conditional Expectation
central challenge in computing the gradient estimator is evaluating the conditional expectation E of (sola;) [.r)]. which is generally intractable:
- 2o qe (zo xx) dxo.
E [
This expectation can be approximated by drawing independent samples (1) and computing the sample mean. Unfortunately, we cannot sample efficiently and directly from q. (xo 2.). Using Bayes' rule:
To ee (ra xa)
(N)
of (zo 21) x qe (x: xo) (Iro),
(8)
where d (xx mn) is available in closed form, but (#o) is intractable due to integration over the group orbit of zo-
To address this, we use self-normalized importance sampling (SNIS) with proposal distribution p( xz) that shares the same intractable orbit-integral structure as q (æo). allowing cancellation of the problematic terms in the importance weights, The conditional expectation is approximated as:
N E To wi() Σ wli
(((
E (x [zo] 2
w( - | (0))
(9)
p xT))
(8)
The design of the proposal p(ro 1r2) aims ensure that the quotient 8 (z)) p xx) becomes tractable. This can be achieved by first sampling from the original dataset D using some user-defined p( x2),, where p is designed have non-zero probability for all elements in the dataset. Once a sample E D is drawn, we sample group element g from the group G' uniformly at random, and apply this group action the sample. This results in
sample z!!² g zo- This proposal distribution inherits the same orbit-integral structure as q (zo). causing the intractable terms in the ratio qe (xx (z00))p((()) (x)
)/p(z()) xx) to cancel. Specifically, with the Dirac delta function:
((xo) Sc ö(g g' o xo) dua(l') p(a! (2) Iz) p(go o 1o x.) p(ro xe) Jc 8(g xo g' o xo) duc(gg)) p(rro xx))
(3) ( vp)
q(xo)
(10)
Beside, since ((xo) = D/ for any su E D, this term can be omitted from the importance weight. Thus, the final importance weight simplifies to wi(() qe (z, (.) 100))/((a x2). Importantly, all components of the importance weights are tractable: qf (Ix, m0
(6)) is the forward diffusion process;
((rn) corresponds to the empirical data distribution; and p(2n xx)) is user-defined and tractable.
Moreover, SNIS estimators based on these importance weights are always consistent.
An important instance results from setting p(zo xx) g(20) where the proposal recovers the exact symmetrized distribution: p(zo xx) = q (xx) and the importance weight simplifies to (+()) = ( (z, x
5

---

3.4 Practical Implementation Diffusion (OrbDiff)
We present Orbir Diffusion (OrbDiff) as a practical variant of our estimator. Although using q (ar) as the proposal is theoretically valid, it is inefficient and cumbersome in practice. For small t, the conditional 7. (ate 20) is sharply concentrated around the z6 that generated wy, so uniformly sampled zo rarely yield useful gradients. Furthermore, sampling from the full support generally requires drawing points outside the current mínibatch, adding non-trivial implementation complexity.
To improve efficiency, we fix Zo to the example that produced F; and sample candidates only from its orbit Ozo {g zo 9 E G}. This biases the proposal toward points with high likelihood under q. (xt I 2o) for-for for instance, small rotations in (03))-qquivarial settings or local permutations in discrete symmetry groups. Since such points dominate the conditional distribution small noise levels, orbit sampling greatly improves sample efficiency by prioritizing candidates with non-trin non-trivial importance weights. At high noise levels, contributions from outside the orbit may increase, but their weights are typically small, and expanding the proposal has shown little benefit.
Formally, OrbDiff replaces the intractable conditional expectation E(zo are] in Equation 6 with the orbit-weighted target
6" (F., Z( /
(g (go zo) q (z) gg zu) dpa(g).
(11)
where (0)) = fG 0f (5zz gooo)) dpa(1) is the normalization constant. This yields the OrbDiff loss:
[ CorbDiff (6) = E Be (20) En, (F(m0) [1o(2,, t) (zu, v4, {)]
(12)
Figure ! provides an illustration and pseudo-code of the practical implementation. When is G- equivariant (see Appendix B..neddxBB.2aaoe and Chen et al. (2024)), ), the loss in Equation (12) and its gradient are unchanged if xg is drawn from the empirical distribution ((xr)) rather than 7 (xo)
Even though 6² ("0g It, t) * E(ro xr). for an equivariant forward process the gradient of Equa- tion (12) matches that of Equation (5), ensuring that OrbDiff yields an unbiased gradient estimate. The orbit-weighted target is also equivariant, providing training signal aligned with the model's inductive bias. We formally prove both properties in Appendix B.4
Unbiased gradient and equivariance of OrbDiff
Theorem 2. Let G be a locally compact isometry group acting on data space Q, and suppose the forward kernels (x, g E G. Then:
(g (z, zo) are G-invariant: (g zt g o xo) = qf (It zo) for all
1. The OrbDiff larget (xxo x4, 1) satisfies: (20, ho 11, = h o= (xo. Nt. for all h E G.
2. The gradient of the OrbDiff loss (2) equals that of the ideal loss (5) i.e. OrbDiff provides unbiased gradient estimator.
We also explore non-uniform sampling of group elements g for approximating the conditional expectation. For small noise, we sample near the identity action, expanding the neighborhood as noise increases. These distributions, (g). depend on the noise schedule. While not all groups support closed-form expressions for the density of individual group elements, they exist for the translation group (sampled from a Gaussian) or S0(3) (sampled from the von Mises-Fisher distribution).
Next, we divide by (4(g) to account for the group sampling distribution, resulting in the importance weights (0) qf (Ira ) )/v(()) We also ensure that the identity group element is included in the sampled set. This strategy is effective in practice and provides computational advantages.
4 Experimental Results
Our experiments evaluate the generality and robustness of Orbit Diffusion across diverse generative tasks. We begin with controlled synthetic setup using standard diffusion model (Section 4.1) and

---

consider various isometry groups, including reflections, rotations, translations, and graph automor- phisms. We then extend our method to Flow Matching (Section (4.2) and diffusion models with non-standard forward processes (Section 4.3 Finally, we apply Orbit Diffusion to a non-equivariant
denoiser, demonstrating effectiveness without architectural symmetry (Section 4.4)
4.1 Experiments on Synthetic Data
0.4
EquiNet
Variant
RMSD (x10 ")
W2 (x10 3)
Training
[Ana]
EquiNet
9.05
9.47
1.150 1,581
0.2
[OrbDiiT]
+ [Aug]
9.50 9.63
0.853 1.156
0.0
+ [OrbDiff]
0.37 0.09
0.004 0.001
[ter=0k
50k
100k
Table 1: Synthetic experiment results: RMSD to the closest target in { -,,1] 1] and W2 distance to ground-tree truth distribution.
Figure 2: Learning curves.
We construct a dataset with a single ID sample To where the equivariant group reflection: g x = =r. We train denoiser (EquiNet) of the form D3 ( t) = fo t) fe(-,,,1 t), which is equivariant by design, with fo being simple 3-layer MLP.
We evaluate three training variants of EquiNet: default, trained without Orbit Diffusion or data augmentation; + [Aug]. with data augmentation only; and [OrbDiff], with Orbit Diffusion only.
Each model is trained for 100k iterations. After training, we generate 100k samples and compute the root mean square deviation (RMSD) from the closest target in (-1, 1}. and the Wasserstein-2 (W2) distance to the target distribution. We report mean and standard deviation for each model in Table 1
Figure 2 shows the training loss curves for all three EquiNet variants. EquiNet and its augmented version exhibit similar fluctuations and magnitudes, indicating that data augmentation does not reduce variance. This is consistent with our theoretical result (Appendix [B2)): for equivariant denoiser, augmented and non-augmented losses are equivalent. In contrast, the OrbDiff variant shows a smoother and lower loss curve, confirming that Lao-Blackwellizatiin reduces gradient variance and stabilizes training. As shown in Table OrbDiff achieves roughly 25x lower RMSD and 200x lower W2 distance, outperforming both EquiNet variants.
4.2 Molecular Conformer Generation
Molecular Conformer Generation. Molecular Conformer Generation aims to generate plausible 3D structures from 2D molecular graphs. crucial for drug discovery and property prediction due to the role of 3D geometry (Liu al. 2023 Axelrod & Gömez- Bombarelli 2020).
We evaluate on the GEOM-QM9 dataset (xxerod & Bomezz-eB 2022 respecting two key symmetries: invariance under global 3D rotations and equivariance to graph automorphisms, which permute atom indices without altering molecular identity. We compare against strong baselines. including GEOMOL (Ganea al. 20211) Torsional Diffusion (Jing al., 20222, , MCF (Wang et al.) 2024b), and ETFLOW (Hassan et al. 2004). J. Our method, Orbit Diffusion, is integrated into ETFLOW, strong equivariant flow matching model that employs harmonic prior for bonded atom proximity.
We finctune ETFLow with OrbDiff using their public checkpoint.
During training, apply symmetry-aware sampling by uniformly sampling 50 automorphisms and 200 (3) rotations per molecule, including the identity. These are applied to both 2D graphs and 3D conformers. All other settings follow ETFLOW; see Appendix C.I for details.
We benchmark against three versions of ETFLOW: the results reported the original paper, their released checkpoint, and our own reproduced results using the provided code and configuration Despite extensive effort, we were unable to match their reported performance, so we report all results under the same evaluation protocol.

---

Table 2: Molecular conformer generation performance GEOM- QM9. Reported in the original paper. Obtained using the published checkpoint. We train the public implementation from scratch.
Models
Recall
Precision
Cover.1 (1) Cov e0.5 (f)
AMR (.)
Cov (?)
Cov 60.5 (t)
AMR (1)
Mean Median Mean Median Mean Median Moun Median Mcan Median Moan Median
GeoMor
28.4
0.0
91.1
100.0
0.224
0.194
20.7
0.0
85.8
100,0
0.271
0.243
Tursional
MCF
37.7
81.9
100.0
94,9
100.0
0.103
0.049
78.6
93.8
12.5
93.9
100.0
100.0
0.221
0.113
0.195
0.055
Figure 3: Molecular conform-
ETFLOW"
96.5
100.0
0.073
0.047
94.1
100.0
0.088
0.039
ers generated by ETFLOW
ETFLOW
79.5
0.066
(left), + [OrbDiff] (center),
+ [OrbDief]]
ETFLOW
85.4
81.4
100.0
100.0
96.3
94.4
100.0
100.9
0.074
0.092
0.027
0.039
80.2
74.6
93.9
85.5
91.9
89.1
100.0
100.0
0.145
0.113
0,042
0.044
and ground-truth (right).
Table shows our method consistently improves both precision and diversity. OrbDiff achieves the best recall scores, including a 4% improvement in mean Cov@0. and the lowest Recall AMR (mean and median). It also maintains competitive precision at 0.1 Á. While MCF performs better 0.5 Á precision, OrbDiff achieves the lowest AMR overall. Further experimental details and comparisons with more baselines are in Appendix C.T
4.3 Crystal Structure Prediction (CSP)
CSP involves recovering 3D atomic positions and lattice parameters from chemical composition. Due to periodicity, suffices to predict the structure within single unit cell, where coordinates lie in the fractional domain [0, 1)3xM, To handle this, DiffCSP uses a Wrapped Normal diffusion that respects periodic translation symmetry. We integrate Orbit Diffusion into DiffCSP, demonstrating that our approach extends beyond Gaussian diffusion models. We test two variants: OrbDiff U, which samples uniformly over the translation group, and OrbDiff WN, a time-dependent Wrapped Normal centered at zero, concentrating around at low noise and spreading out at high noise. Details of the OrbDiff WN proposal are in Appendix C.2.3
We evaluate our method on two CSP benchmarks: Perov-5 (Castelli et al. 2012a bb and MP-20 (Jiin
al. 2013). We use the three strongest baselines from the DifICSP paper (Jiao et 2023): P-cG- SchNet (Gebauer et al. 2022), CDVAE (Xie et al., 2022), and DiffCSP, all with publicly available implementations. We evaluate performance using two standard metries: Match Rate (the proportion of correctly matched structures in the test set) and RMSD (the average atomic deviation for matched samples, normalized by lattice volume). Full metric definitions and details are in Appendix C.2.4
We further consider a relevant task, introduced by TGDMat (DAS et al. 2025), where crystal structures are generated conditioned additional text descriptions of the desired structures. In this task, two types of descriptions are considered: long and short, with the latter being easier to obtain than the former. We follow the same evaluation framework as for Non-text-guided CSP.
Table 3: Text-guided CSP with TGDMat.
Table 4: Crystal Structure Prediction (CSP).
Method
Perov55
Method
Perov-5
MP-20
Match (1) RMSE (L) Match (f) RMSE ())
Match (?) RMSE (L) Match (1) RMSE (L)
TGDMat (S)
59.39
0.066
59.90
0.078
P-cG-SchNet CDVAE
0.418
15.39
0.376
[0rbbiff of WN]
[D=bDiff__ U]
63.51
65.57
0,054
0.062
56.50
61.29
0.085
0.072
45.31
0.114
33.90
0.105
TGDMat (1.)
95.17
0.013
61.91
0.081
DifCCSP
52.02
0.076
51.49
0.063
[DrbDitt_ U]
95. 88
0.012
65.94
0.069
[OrtDifr_U] U]
52.29
0.078
54.47
0.054
[DrbDiff_ __
95.98
0.012
[OrbDiff_ _WN]
52.39
0.069
55.70
0.053
From Tables 3 and one can see OrbDiff. consistently enhances the performance in all cases, with a notable increase from 59.39% to 65.57% for TGDmat (S) on Perov-5 and from 61.91 % to 66.74% for TGDMat (L) MP-20. At the same time, OrbDiff U outperforms the baselines in 5 out of 6 cases, showing consistent benefits. Improvements are also observed consistently in RMSE.

---

0:
MP-20
T3DMAT
Figure 4: Qualitative comparison of Crystal Structure Predictions by models, including DiffCSP, TGDMat (S) and TGDmat (L) with baselines, OrbDiff_ _U, and OrbDiff ___N against ground-truth samples on randomly selected samples from Perov-5 and MP-20 dataset.
4.4 Protein Structure Generation with PROTEINA
Table 5: Protein Structure Generation. Full comparisons are in the appendix. "+ [finetune]" denotes MMS finetuned with the original loss; "+ [OrbDiff]" uses OrbDiff for finetuning. Full table can be found in Appendix C.3.4 C.3.4
CER
Designability
Diversity
Novelty
Model
Fraction (*) scRMSD Cluster
TM-score ()) PDB
AFDB (L)
FoldFlow (OT)
97.2
0.37
0.41
0.71
0.75
Figure 5: OrbDiff (orange) gen-
MaM
99.0
0.72
0.30
0.39
0.81
0.84
erated structures versus reference
Mggai 93.8
1.04
0.62
0.36
0.69
0.76
[linetune]
1.00
0.54
0.74
structure (green).
[OrbD1ff]
95.6
0.93
0.52
0.37
0.74
0.83
We adopt PROTEINA (denoted M) (Geffner al.) 2025).
the state-of-the-art model for protein
structure generation. as a backbone. Although PROTEINA is non-equivariant transformer, performs well through extensive data augmentation. We finetune the 200M-parameds version of POOTEINA-MIRO using OrbDiff, applying Rao-Blackwelllizatio with uniform proposal distribution over 53)3), sampling 10,000 group elements.
For comparison, we also include the best-performing equivariant baseline, FoldFlow (OT) (Bose et al. (2024). We evaluate protein structures using three metries: Designability (the feasibility of synthesizing the generated structures), Diversity, and Novelty. Designability is the most critical, while Diversity and Novelty are based the designable samples, ensuring that the generated structures synthetically plausible, diverse, and novel. Evaluation metric details are in Appendix C.3.3
Table shows that Orbit Diffusion boosts designability of MM to 95.6% and lowers scRMSD to 0.93, outperforming naive finetuning while maintaining competitive diversity (Cluster: 0.52) and novelty (PDB: 0.74). The state-of-the-ae M21M (400M parameters) achieves higher designability (99.0%) and scRMSD (0.72) but at the cost of much lower diversity (Cluster: 0.30) and novelty (PDB: 0.81), showing a trade-off between validity and structural variety.
4.5 Benefits of OrbDiff: Efficiency, Stability, and Equivariance
Base Modal
Overhead
DUBCSP
[OrbDif U]
[OrbDet WN]
PROTEINA
[GrbDii]]
Match (%)
Grad Var
PROTEINA
99.9%
101
(xta
77.
LTDIff
97:22%
5.0
RquiNet
93.8%
20-2
80.8%
Overhead
20
40
80
100
Epodi=0
20
40
80
100
1-20
420
820
1-0.0
1.0
(a)
(b)
(c)
(d)
Figure 6: (a) OrbDiff introduces minimal computational overhead. (b) Match rate (f) during DiffCSP training on the Perov-5 dataset: OrbDiff accelerates convergence. (c) OrbDiff reduces gradient variance across noise levels. (d) OrbDiff improves the equivariance of PROTHINA.

---

Adding OrbDiff introduces minimal overhead (Figure (u) about 20% for smaller models such as DiffCSP, and only 0.1% for larger ones like PROTEINA (200M parameters). run, this corresponds to just 1. 1.4 additional minutes. In practice, the extra memory and computation for sampling and averaging over group elements are negligible compared to the overall training cost. For example, in our PROTEINA experiments- one of the largest protein diffusion models we use 10,000 SO(3) samples per training example, which adds only 440 MB per GPU to the total memory usage of ~54 GB.
24-hour training
To better understand how OrbDiff accelerates convergence, we compare DiffCSP and DiffCSP + [OrbDiff] across training checkpoints. As shown in Figure integrating OrbDiff consistently improves match rates throughout training, especially in the early stages.
We also compare the empirical gradient norm variance of DiffCSP (Perov-5) with OrbDiff. _ and OrbDiff WN across the full training set and various timesteps. As expected, OrbDiff. WN achieves significantly lower variance at small to intermediate noise levels by sampling locally around Fo. while OrbDiff_ performs better at high noise levels due its more global sampling. Both methods substantially reduce variance compared to DifICSP. Finally, to assess equivariance preservation, we compare PROTEINA [finetune] and PROTEINA [OrbDiff] using an equívariance test from (Geffner et al. 2025):
Error, E @f (zo). z4 z (m, (zo): g=-)) RMSD(g o 6(2,, t), o o x,, t)]].
(13)
As shown in Figure (ot, OrbDiff substantially reduces equivariance error compared naive finetun- ing, indicating improved geometric consistency in the model's 's denoising.
Related Work
Equivariant neural networks have been extensively studied for their ability to encode symmetry priors
in domains such as vision (Cohen & Welling 2016 Worrall et al., 2017). 3D geometry (Thomas
al. 2018 Deng et al. 2021), and molecular modeling (Le et al. 2022 Kozinsky et al. 2023). More
recently, these ideas have been incorporated into diffusion models to improve generative performance
in structured domains like proteins and molecules (Corso al., 2024 Hoogeboom et al.) [002]]
Igashov et al. 2024).. By designing the denoising network to be equivariant under symmetry group, these models better align with the underlying data distribution. A comprehensive discussion of related
work is provided in Appendix
Conclusion
Orbit Diffusion is a framework for training generative models under symmetry constraints by reducing gradient variance through Rao-Blackwellization. The approach unifies equivariant architectures and data augmentation strategies within single probabilistic formulation, providing provably lower-variance estimator while maintaining computational efficiency. Theoretically, we show that the proposed loss admits equivariant minimizers and connects existing score-based and diffusion formulations. Empirically, Orbit Diffusion demonstrates strong and consistent performance across a diverse set of generative tasks, including molecular, crystal structure, and protein structure modeling. By bridging the gap between symmetry-aver modeling and optimization stability, our method improves both the scalability and practical applicability of equivariant generative models in scientific domains.
7 Limitations
To estimate the conditional expectation in our gradient estimator, we employ orbit sampling to reduce variance, which leads to improved performance. While efficient, our sampling scheme inherently constrains the choice of proposal distributions and may limit modeling flexibility, especially for tasks requiring diverse or adaptive noise structures. Furthermore, the number of orbit samples directly affects both computational cost and estimation accuracy, suggesting that dynamically adjusting the sampling strategy during training could yield better efficiency trade-offs. Future work could explore adaptive or learned proposal mechanisms to enhance generalization and robustness across broader data regimes.
10

---

Acknowledgements
This work was supported in part by the Deutsche Forschungsgeemeinscaatt (DFG, German Re- search Foundation) under Germany's 's Excellence Strategy EXC 2075 -307740016 the DARPA ANSR program under award FA8750-23-2-0004 the DARPA CODORD program under award HRO01 12590089; the NSF grant IIS- 1954994641; ; the National University of Singapore Start-up Grant (Award No. SUG-25 RES250); and gifts from Adobe Research, Cisco Research, and Amazon. We also acknowledge the support of the Stuttgart Center for Simulation Science (SimTech). VT and MN thank the International Max Planck Research School for Intelligent Systems (IMPRS-IS) for support. This work further received partial support from the Diabetes Center Berne and from the Ministry of Science, Research and the Arts Baden-Württent Württemberg through the Artificial Intelligence Software Academy (AISA).

---

References
Emmanuel Abbe and Enric Boix-Adserà. On the non unnivevrrslllty of deep learning: Quantifying the cost of symmetry. In Advances in Neural Information Processing Systems 35 (NeurIPS 2022),
2022.
Josh Abramson, Jonas Adler, Jack Dunger, Richard Evans, Tim Green, Alexander Pritzel, Olaf Ronneberger, Lindsay Willmore, Andrew Ballard, Joshua Bambrick, et al. Accurate structure prediction of biomolecular interactions with alphafold 3. Nature, 2024a.
Josh Abramson, Jonas Adler, Jack Dunger, Richard Evans, Tim Green, Alexander Pritzel, Olaf Ronneberger. Lindsay Willmore. Andrew Ballard, Joshua Bambrick, et al. Accurate structure prediction of biomolecular interactions with alphafold 3. Nature, 630(8016):493-500,, 2024b.
Namrata Anand and Tudor Achim. Protein structure and sequence generation with equivariant denoising diffusion probabilistic models. CoRR, abs/2205. 15019, 2022.
Simon Axelrod and Rafael Gomnez-oombarelii Molecular machine learning with conformer ensembles. Machine Learning: Science and Technology, 4, 2020. URL https://api
semanticscholar org/ /Cooooossd::21111000000
Simon Axelrod and Rafael Gomez-Bombarede Geom, energy-annotatee molecular conformations for property prediction and molecular generation. Scientific Data, 9(1):185, 2022.
Ilyes Batatia, David P Kovaes, Gregor Simm, Christoph Ortner, and Gábor Csányi Mace: Higher order equivariant message passing neural networks for fast and accurate force fields. Advances in neural information processing systems, 2022.
Avishek Joey Bose, Tara Akhound-Sadegl Guillaume Huguet, Kilian Fatras, Jarrid Rector-Brooks. Cheng-Hao Liu, Andrei Cristian Nica, Maksym Korablyov, Michael Bronstein, and Alexander Tong. Se(3) -too)-sto flow matching for protein backbone generation. In Proceedings of the 12th International Conference Learning Representations (ICLR), 2024.
Johann Brehmer, Sönke Behrends, Pim de Haan, and Taco Cohen. Does equivariance matter at scale? arXiv preprint arXiv:2410.23172 2024.
Ivano E. Castelli, David D. Landis, Kristian S. Thygesen, Søren Dahl, Ib Chorkendorff, Thomas F. Jaramillo, and Karsten W. Jacobsen. New cubic perovskites for one- and two-photon water splitting using the computational materials repository. Energy & Environmental Science, 2012a.
Ivano E Castelli, Thomas Olsen, Soumendu Datta, David D Landis, Søren Dahl, Kristian S Thyygsen, and Karsten W Jacobsen. Computational screening of perovskite metal oxides for optimal solar light capture. Energy & Environmental Science, 5(2):5814
2012b.
Ziyu Chen, Markos A. Katsoulakis, and Benjamin J. Zhang. Equivariant score-based generative models provably learn distributions with symmetries efficiently. arXiv preprint 4rXiv:2410.01244, 2024.
Taco Cohen and Max Welling, Group equivariant convolutional networks. In Proceedings of the 33rd International Conference on Machine Learning (ICML), 2016.
Gabriele Corso, Arthur Deng, Nicholas Polizzi, Regina Barzilay, and Tommi Jaakkola. Deep confident steps new pockets: Strategies for docking generalization. In International Conference on Learning Representations (IC (ICLR), 2024.
KISHALAY DAS, Subhojyoti Khastagir, Pawan Goyal, Seung-Cheol Lee, Satadeep Bhattacharjee, and Niloy Ganguly. Periodic materials generation using text-guided joint diffusion model. In The Thirteenth International Conference on Learning Representations, 2025.
Justas Dauparas, Ivan Anishchenko, Nathaniel Bennett, Hua Bai, Robert J Ragotte, Lukas F Milles, Basile IM Wicky, Alexis Courbet, Rob J de Haas, Neville Bethel, Patrick JY Leung, Thomas Huddy, Samuel Pellock, Daniel Tischer, Felix Chan, Brian Koepnick, Huy Nguyen, Andrew Kang, Balamurugan Sankaran, Abhishek K Bera, Neil King, and David Baker. Robust deep lcarning-basee protein sequence design using proteinmpnn. Science, 2022.
12

---

Congyue Deng, Or Litany, Yueqi Duan, Adrien Poulenard, Andrea Tagliasacchi, and Leonidas J. Guibas. Vector neurons: general framework for eqqiivaiaa networks. In Proceedings of the 2021 IEEF/CVF International Conference on Computer Vision (ICCV), 2021.
Fabian Fuchs, Daniel E. Worrall, Volker Fischer, and Max Welling. Se(3))transformer 3d roto- translation equivariant attention networks. In Advances in Neural Information Processing Systems 33 (NeurIPS 2020), 2020.
Octavian Ganea, Lagnajit Pattanaik, Connor W. Coley, Regina Barzilay, Klavs F. Jensen, William H. Green Jr., and Tommi S. Jaakkola. Geomol: Torsional geometric generation of molecular 3d conformer ensembles. In Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems 2021, NeurlPS 2021, December 6-14, 2021, virtual, 2021.
Benoit Gaujac, rereemie Dona, Liviu Copoiu, Timothy Atkinson, Thomas Pierrot, and Thomas D. Barrett. Learning the language of protein structure. ArXiv, abs/2405.15840, 2024. URL https //api. seagpessamnniccccoll... arg/Corrpssd::2706311
Niklas W. A. Gebauer, Michael Gastegger, Stefaan S. P. Hessmann, Klaus-Robert Müller, and Kristof T. Schütt, Inverse design of 3d molecular structures with conditional generative neural networks, Nature Communications, 13(1);973, 2022. doi: 10. 0011088841446--02222282533 URL
https://doi. orgg 10. 1038/ #414/44146777 285266y
Tomas Geffner, Kieran Didi, Zuobai Zhang, Danny Reidenbach, Zhonglin Cao, Jason Yim, Mario Geiger, Christian Dallago, Emine Kucukbenli, Arash Vahdat, and Karsten Kreis. Proteina: Scaling flow-based protein structure generative models. In Proceedings of the International Conference on Learning Representations (ICLR), 2025.
Mathis Gerdes, Pim de Haan, Corrado Rainone, Roberto Bondesan, and Miranda C. N. Cheng. Learning lattice quantum field theories with equivariant continuous flows. SciPost Physics, 15 (6):238, 2023. doi: A68/SciPostPPys..5.... SSciossPhyss 15. 6388
UtL/httts::/ssciptttt org/1/ 21468/
Jiaqi Guan, Wesley Wei Qian, Xingang Peng, Yufeng Su, Jian Peng, and Jianzhu Ma. 3d equivariant diffusion for target-awan molecule generation and affinity prediction. In Proceedings of the 11th International Conference on Learning Representations (ICLR), 2023.
Majdi Hassan, Nikhil Shenoy, Jungyoon Lee, Hannes Stärk, Stephan Thaler, and Dominique Beaini. Et-flow: Equivariant flow-matching for molecular conformer generation. Advances Neural Information Processing Systems, 2024.
Thomas Hayes, Roshan Rao, Halil Akin, Nicholas J. Sofroniew, Deniz Oktay, Zeming Lin, Robert Verkuil, Vincent Q. Tran, Jonathan Deaton, Marius Wiggert, Rohil Badkundri, Irhum Shafkat, Jun Gong, Alexander Derry, Raul S. Molina, Neil Thomas, Yousuf A. Khan, Chetan Mishra, Carolyn Kim, Liam J. Bartie, Matthew Nemeth, Patrick D. Hsu, Tom Sereu, Salvatore Candido, and Alexander Rives. Simulating 500 million years of evolution with a language model. Science, 373(6557), 2025.
Lingshen He, Yuxuan Chen, Zhengyang Shen, Yiming Dong, Yisen Wang, and Zhouchen Lin. Efficient equivariant network. Advances in Neural Information Processing Systems 34 (NeurIPS 2021), 2021.
Jan Hermann. Zeno Schätzle, and Frank Noé. Deep neural network solution of the electronic schrödinger equation. Nature Chemistry, 12(10):891-897 2020.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 33:6840 35:6840-6851, 2020.
Emiel Hoogeboom, Victor Garcia Satorras, Clément Vignac, and Max Welling, Equivariant diffusion for molecule generation in 3d. In International conference on machine learning, pp. 8867-8887. PMLR, 2022.
13

---

Guillaume Huguet, James Vuckovic, Kilian Fatras, Eric Bribodeau-Laufet Pablo Lemos, Riashat Islam, Cheng-Hao Liu, Jarrid Rector-Brooks, Tara Achound-Sadepa Michael M. Bronstein. Alexander Tong, and Avishek Joey Bose. Sequence- aureene- se(3)-flow matching for conditional protein backbone generation. In Proceedings of the 37th Conference on Neural Information Processing Systems (NeurIPS 2024). 2024.
Ilia Igashov, Hannes Stärk, Clément Vignac, Ame Schneuing, Victor Garcia Satorras, Pascal Frossard, Max Welling, Michael Bronstein, and Bruno Correia. Equivariant 3d-conditional diffusion model for molecular linker design. Nature Machine Intelligence, pp. 1-11, 2024.
John B. Ingraham, Max Baranov, Zak Costello, Karl W. Barber, Wujie Wang, Ahmed Ismail, Vincent Erik R. Van Vlack, Shan Tie, Vincent Xue, Sarah C. Cowles, Alan Leung, João V. Rodrigues, Claudio L. Morales-Perez, Alex M. Ayoub, Robin Green, Katherine Puentes, Frank Oplinger, Nishant V. Panwar, Fritz Obermeyer, Adam Root, Andrew I., Beam, Frank J. Poelwijk, and Gevorg Grigoryan. Illuminating protein space with programmable generative model. Nature, 2023.
Frappier. Dana M. Lord, Christopher Ng-Thow-
A Jain, SP Ong, G Hautier, W Chen, WD Richards, Dacek, S Cholia, D Gunter, Skinner, G Ceder, et al. The materials project: materials genome approach to accelerating materials innovation, apl mater. (2013) 01102,, NA.
Anubhav Jain, Shyue Ping Ong, Geoffroy Hautier, Wei Chen, William Davidson Richards, Stephen Dacek, Shreyas Cholia, Dan Gunter, David Skinner, Gerbrand Ceder, et al. Commentary: The materials project: A materials genome approach to accelerating materials innovation. APL materials, 1(1), 2013.
Rui Jiao, Wenbing Huang, Peijia Lin, Jiaqi Han, Pin Chen, Yutong Lu, and Yang Liu. Crystal structure prediction by joint equivariant diffusion. In Thirty-seventh Conference on Neural Information Processing Systems, 2023, URL httpttts///:pereeeeeee http://pppreveee.....ee www.we..weeeooooniiieeeeeeeeeeee http://oouuu111ddDMd22222222
Bowen Jing, Gabriele Corso, Jeffrey Chang, Regina Barzilay, and Tommi Jaakkola. Torsional diffu- sion for molecular conformer generation. In Advances Neural Information Processing Systems 35 (NewrlPS 2022), pp. 24240-24253, 2022. URL httttt::://pppeeee nips cc/ wwtp:pccppppppppeeeeee sfllse oastr//aaeeeteeeeeeeeeeeeaaaaana wwwncenwwweennnooooooooooooooeee .tmml
Bowen Jing, Ezra Erives, Peter Pao-Huang, Gabriele Corso, Bonnie Berger, and Tommi Jaakkola. Higenfold: Generative protein structure prediction with diffusion models. Machine Learning for Drug Discovery Workshop at the International Conference on Learning Representations (ICLR), 2023. URL
net/pdf (id--ii
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-26577 2022.
Nayoung Kim, Seongsu Kim, Minsu Kim, Jinkyoo Park, and Sungsoo Ahn. Mofflow: Flow matching for structure prediction of orgal-or frameworks. In Proceedings of the Thirteenth International Conference on Learning Representations (ICLR 2025), 2025. URL https: http://oppeeeiieeee net/ Horum?id=ddT3abDDee
Diederik Kingma, Tim Salimans, Ben Poole, and Jonathan Ho. Variational diffusion models. Advances in neural information processing systems, 34:21696-21707,, 2021.
Leon Klein, Andreas Krämer, and Frank Noé. Equivariant flow matching. Advances in Neural Information Processing Systems, 36, 2024.
Jonas Köhler, Leon Klein, and Frank Noé. Equivariant flows: Exact likelihood generative learn- ing for symmetric densities. In Proceedings of the 37th International Conference on Machine Learning, ICMI. 2020, 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learning Research, pp. 5361-5370. PMLR, 2020. URL http: htt.htt:://proeteeee httptt:://proccednngg. mlr, mlr.press/ presss v v119/kohler20a.ttm
14

---

Boris Kozinsky, Albert Musaelian, Anders Johansson, and Simon L. Batzner. Scaling the leading accuracy of deep equivariant models to biomolecular simulations of realistic size. In Dorian Arnold, Rosa M. Badia, and Kathryn M. Mohror (eds.), Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis, SC 2023, Denver, CO, USA, November 12-17, 2023, pp. 2:1-2:12. ACM, 2023. doi: 10. 10.1145/358178443270044 URL https://doos https://doi.orr/100 10. 10111111/11887 3627041
Tuan Le, Frank Noé, and Djork-Amé Clevert. Representation learning on biomolecular structures using equivariant graph attention. In Bastian Rieck and Razvan Pascanu (eds.), Learning on Graphs Conference, LoG 2022, 9-12 December 2022, Virtual Event, volume 198 of Proceedings of Machine Learning Research, pp. 30. PMLR, 2022. URL WRRLhttts://pocceediii mlw..mees/////// press/
mlr
v198/ v198/1e222 2e///ee22a.... html
Tuan Le, Julian Cremer, Frank Noé, Djork-Amé Clevert, and Kristof Schütt. Navigating the design space of equivariant diffusion baffed generative models for de novo 3d molecule generation. In Proceedings of the Twelfth International Conference Learning Represenlations (ICLR 2024), 2024.
Yeqing Lin, Minji Lee, Zhao Zhang, and Mohammed AlQuraishi. Out of many, one: Designing and scaffolding proteins at the scale of the structural universe with genie arXiv preprint arXiv:2405. 15489, 2024. Uttpttts::/arrxx www.:orggaaa....ea abs/ 2405 15489
Zeming Lin, Halil Akin. Roshan Rao, Brian Hie, Zhongkai Zhu, Wenting Lu, Nikita Smetanin, Robert Verkuil, Ori Kabeli, Yaniv Shmueli, Allan dos Santos Costa, Maryam Fazel-Zarandi, Tom Sercu, Salvatore Candido, and Alexander Rives. Evolutionary-scale prediction of atomic-level protein structure with a language model. Science, 2023.
Shengehao Liu, Hongyu Guo, and Jian Tang. Molecular geometry pretraining with SE(3)-invariant denoising distance matching. In Proceedings of the Eleventh International Conference on Learning Representations (ICLR 2023), 2023. URL httttttss///::eteeeee
Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu, Dpm-solver: A fast ode solver for diffusion probabilistic model sampling in around 10 steps. Advances in Neural Information Processing Systems, 35:5775-5787, 2022.
Alexander Quinn Nichol and Prafulla Dhariwal.
Improved denoising diffusion probabilistic models. In International Conference on Machine Learning, pp. 8162-8171. PMLR, 2021.
Artem R Oganov and Colin W Glass. Crystal structure prediction using ab initio evolutionary techniques: Principles and applications. The Journal of chemical physics, 124(24), 2006.
Arne Schneuing, Charles Harris, Yuanqi Du, Kieran Didi, Arian Jamasb, Ilia Igashov, Weitao Du, Carla Gomes, Tom L. Blundell, Pietro Lio, et al. Structure-based drug design with equivariant diffusion models. Nature Computational Science, 4(12):899-909, 2024.
Xuecheng Shao, Jian Lv, Peng Liu, Sen Shao, P. Gao, Hanyu Liu, Yanchao Wang, and Yanming Ma. A symmetry orrientt divide-and-comquer method for crystal structure prediction. The Journal of chemical physics, 156 1:014105, 2021. URL https://ppp semanticscholar.
238198823
Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. In Pro- ceedings of the 9th International Conference Learning Representations (ICLR), 2021a. URL httpss//opeerrviess Sttp://rrr...eeeteoooaaaaaaaaaaa StigddrrCHtiirr
Yang Song, Jascha Sohl-Dickstein, Diederik Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations (ICLR), 2021b.
Nathaniel Thomas, Tess E. Smidt, Steven Kearnes, Lusann Yang, Li Li, Kai Kohlhoff, and Patrick Riley. Tensor field networks: Rotation- and translationeequivrriant neural networks for 3d point clouds. In Advances in Neural Information Processing Systems (NeurIPS) 31, 2018.
15

---

Alexander Tong, Kilian Fatras, Nikolay Malkin, Guillaume Huguet, Yanlei Zhang, Jarrid Rector- Brooks, Guy Wolf, and Yoshua Bengio. Improving and generalizing flow-based generative models with minibatch optimal transport. Transactions on Machine Learning Research, 2024.
Vinh Tong, Trung-Dung Hoang, Anji Liu, Guy Van den Broeck, and Mathias Niepert. Learning to discretize denoising diffusion odes. In Proceedings of the 1 13th International Conference on Learning Representations, 2025.
Clément Vignac, Igor Krawczuk, Antoine Siraudin, Bohan Wang, Volkan Cevher, and Pascal Frossard. Digress: Discrete denoising diffusion for graph generation. In The Eleventh International Confer- ence on Learning Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023, 2023.
Chentong Wang, Yannan Qu, Zhangzhi Peng, Yukai Wang, Hongli Zhu, Dachuan Chen, and Longxing Cao. Proteus: exploring protein structure generation for enhanced designability and efficiency. bioRxiv, pp. 2024-02, 2024a.
Yuyang Wang, Ahmed A. A. Elhag, Navdeep Jaitly, Joshua M. Susskind, and Miguel Ángel Bautista. Swallowing the bitter pill: Simplified scalable conformer generation. In Forty-first International Conference on Machine Learning, ICML 2024, Vienna, Austria, July 21-27, 2024, 2024b.
Joseph I.. Watson, David Juergens, Nathaniel R. Bennett, Brian L. Trippe, Jason Yim, Helen E. Eisenach, Woody Ahem, Andrew J. Borst, Robert J. Ragotte, Lukas Milles, Basile I. M. Wicky, Nikita Hanikel, Samuel J. Pellock, Alexis Courbet, William Sheffler, Jue Wang, Preetham Venkatesh, Isaac Sappington, Susana Vázquez Torres, Anna Lauko, Valentin De Bortoli, Emile Mathieu, Sergey Ovchinnikov, Regina Barzilay, Tommi S. Jaakkola, Frank DiMaio, Minkyung Baek, and David Baker. De novo design of protein structure and function with rfdiffusion. Nature, 620(7976):1088 1100, 2023.
Lai Wei, Sadman Sadeed Omee, Rongzhi Dong, Nihang Fu, Yuqi Song, E. Siriwardane, Meiling Xu, Chris Wolverton, and Jianjun Hu. Cspbench: a benchmark and critical evaluation of crystal structure prediction. In NA, 2024. URL hhttps///apppp semanticschon ar Crrg/Corus Id: 270869939
Daniel E. Worrall, Stephan J. Garbin, Daniyar Turmukhambetov, and Gabriel J. Brostow. Harmonic networks: Deep translation and rotation equivariance. In 2017 IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2017, Honolulu, HI, USA, July 21-26, 2017, pp.
IEEE Computer Society, 2017. doi: 10. 10.11099.P.....1..555.. URL https: //oii org/10
1109/CVPR 2017 758
Kevin E. Wu, Kevin Kaichuang Yang. Rianne van den Berg, James Zou, Alex X. Lu, and Ava P. Amini. Protein structure generation via folding diffusion. Nature Communications, 15, 2022. URL
https://api. wwmanticscoolarr wwwp:/iisccoooaoooo//rooooee orgo :ttg/Corpusd:::522000 2526685511 ott.:p/cc......oooooooooooaaaaaa
Tian Xie, Xiang Fu, Octavian Setavian-Eugen Ganea, Regina Barzilay, and Tommi S. Jaakkola. Crystal diffusion variational autoencoder for periodic material generation. Proceedings of the 10th International Conference Learning Representations (ICIR), 2022.
Minkai Xu, Lantao Yu, Yang Song, Chence Shi, Stefano Ermon, and Jian Tang. Geodiff: A geometric diffusion model for molecular conformation generation. In The Tenth International Conference on Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022, 2022a.
URL https://openrevier: net httpa//rrrrrr.eeeooooooooo//ooop
Minkai Xu, Lantao Yu, Yang Song, Chence Shi, Stefano Ermon, and Jian Tang. Geodiff: A geometric diffusion model for molecular conformation generation. In International Conference on Learning Representations, 2022b. URL htttp::/ httppp///pppaareeeieeeeeerrrrrra www..e//oooommddddccettttatoeeae
Yilun Xu, Shangyuan Tong, and Tommi Jaakkola. Stable target field for reduced variance score estimation in diffusion models. arXiv preprint arXiv:2302.0080 002700000000,, 2023.
Jason Yim, Andrew Campbell, Andrew Y. K. Foong, Michael Gastegger, José Jiménez-Luna, Sarah Lewis, Victor Garcia Satorras, Bastiaan S. Veeling, Regina Barzilay, Tommi Jaakkola, and Frank Noé. Fast protein backbone generation with se(3) flow matching. arXiv preprint arXiv:2310.05297 2023a.
16

---

Jason Yim, Brian L. Trippe. Valentin De Bortoli, Emile Mathieu, Arnaud Doucet, Regina Barzilay, and Tommi Jaakkola. SE(3) diffusion model with application to protein backbone generation. In Proceedings of the 40th International Conference on Machine Learning. 2023b.
Viktor Zaverkin, Francesco Alesiani, Takashi Maruyama, Federico Errica, Henrik Christiansen, Makoto Takamoto, Nicolas Weber, and Mathias Niepert. Higher-rank irreducible cartesian tensors for equivariant message passing. In Conference Neural Information Processing Systems, 2024.
Mengchun Zhang, Maryam Qamar, Taegoo Kang. Yuna Jung, Chenshuang Zhang, Sung-Ho Ho Bae, and Chaoning Zhang. A survey on graph diffusion models: Generative Al in science for molecule, protein and material. CoRR, abs/2304.01565. 2023. doi: 10.4550/ARXXV.230..0.... URL
httpp://doo ftsps:://doioor/1000 10. 48550/ arXiv 2304.01566
17

---

RAO-BLACKWELL GRADIENT ESTIMATORS FOR EQUIVARIANT DENOISING DIFFUSION
ADDITIONAL MATERIAL
Related Work
19
B
Theoretical Proofs
19
Proof of Theorem
19
B.2
Equivalence of Symmetrized and Equivariant Diffusion Losses
20
B.3
Symmetrized Forward Diffusion Distribution,
22
B
Proof of Theorem
23
B.5
Unconstrained Non-Symmetrizae Diffusion Minimizer is not guaranteed G-equivaraint
25
Experimental Details
26
Molecular Conformer Generation (MCG) with Flow Matching
26
Crystal Structure Prediction (CSP)
29
Protein Structure Generation (PSG) with Non -qn-iiriari Denoiser PROTEINA
33
18

---

A Related Work
Equivariant Neural Networks. Equivariant neural networks have attracted significant attention
for tasks involving structured data, such as computer vision (Cohen & Welling 2016
Worrall
ct al. 2017). 3D modeling (Thomas et al., 2018 Fuchs et al. 2020; Deng et al. (021). quantum
mechanics and quantum field theory (Gerdes et al. [023] Hermann et al. 2020), biomolecular Geffner
design (Le et al, 2022. Kozinsky et al., [023] et al. 2025). These networks exploit group
symmetries to ensure consistent outputs under transformations such rotations, translations, and permutations, which are commonly encountered in many scientific domains. By incorporating symmetric inductive biases, equivariant networks enhance model generalization and reduce data requirements by naturally encoding symmetry constraints. However, challenges remain, such as high computational complexity (He et al., 2021 and difficulties in effectively learning with stochastic gradient descent (SGD) (Abbe &
2022).
Equivariance and Diffusion Models. Diffusion models have become dominant class of generative models, known for their effectiveness in modeling complex data distributions (Song et al. 220 la
Ho al,, 2020 Karras al., 2022). They work by gradually adding noise to data in forward process and learning to reverse this corruption using denoising network Φ9. Incorporating symmetry, particularly equivariance, into these models has shown significant benefits in domains where data lies on geometric manifolds with known symmetries -ucch structural biology (Corso et al.) [024] Yim
2023b Schneuing et al. 2024; Igashov et al.] 2024), molecular modeling (Hoogeboom et al. 2022; Guan al.) 2023 Le et 20044, and material design (Xie et al. [202]] Gebauer et al.) 2022
al.
Jiao al. (023). In such settings, where symmetries are typically governed by the Euclidean group or its subgroups, enforcing equivariance in the generative process helps ensure physical plausibility and improves generalization. A common approach is to design the denoiser Φy to be equivariant under symmetry group G, which encourages the learned distribution to converge toward the correct
invariant target (Hoogeboom et al. 2022 Xu et al. 2022a Bose et al., 2024 Chen et al. 2024).
Nonetheless, recent advances have shown that models without explicit equivariant constraints can still achieve strong empirical performance, thanks to greater flexibility in architectural design and effective use of data augmentation. These approaches may implicitly capture symmetry through training strategies rather than architectural bias, as demonstrated by state-of-the-art models such as PROTEINA (Geffner et al. 2025) and AlphaFold 3 (Abramson et al. 2024a).
B Theoretical Proofs
B.1 Proof of Theorem
Variance Reduction of
estimator.
Theorem 1. Let & (olC((0) and ((RB) (CS [C(( denote the gradient estimators of Equa-
tion and Equation (5), respectively. Suppose we can compute E.zo~ (oolas) [ro]. Then
[ S Var (5olC)
Var
Moreover, the inequality is strict unless of (xo xx) is a Dirae delta, which is rarely the case in generative modeling where zo is typically stochastic given It.
Proof. By definition of the Rao- Sho-Blackwellized estimator, we have
(RB) [ (())))
(+) E (wal (weal
VoLCE (+) x
Thus, the Rao- Moo-Blackwell theorem implies that conditioning reduces variance:
Var
[ Volc (+) zz
Var (voIc(
with equality if and only if v Volca (o)] is almost surely function of a (i.c.. deterministic given 22).
19

---

To verify this explicitly, consider the variance of the Rao- Rao-Blackwellized estimator: estiatoo
Var
[[ (d))
= Ex,
E
V (o(C
m. Vol5
Eze
B
[olc
vo1c x1)) VolC)
(FolC)
S Ezz
B
[0]]
Var (VolC( (+))
where the inequality follows from Jensen's 's inequality (applied to the convex function f(2) = 2²).
Therefore, the Rao-Blacke nro-Blaackwelized estimator has variance less than or equal to the original estimator, and strictly less unless the conditional variance given 24 zero.
B.2 Equivalence of Symmetrized and Equivariant Diffusion Losses
Proposition 1. Let c. (1) denote the Symmetrized Diffusion Loss, defined
C 2 (0) Be ( (f [|o(or,)) () Be [m)] 2
and let L.2² (0) denote the corresponding loss defined the original (non-symmetricee) data distribution q(zo):
[|o(crs)) - EE of0 [201
- (6) Erl ((mr)) En - (ail=))
Suppose o(z,,t) t) is G-equivariant function, i.e.,
ø(g o zt, 1) = ((x1, EgEG, I S.
Then (6) and cf²) (Ø) are equivalent in the sense that they share the same minimizer and gradient with respect 6.
To prove this, we need the following lemma:
Symmetrized Forward Diffusion Distributions
Lemma 1. Let ((xa) an empirical distribution and ÿ (xx) its symmetrized counterpart
under symmetry group G, defined by (x(o) Salğ (xo) Suppose forward diffusion process acting on q(ιo) yields dime-dependent marginal distributions q (Zt). Let similar process act on @ (Ii) to generate of (x,). Then, for all t 0, the following holds:
(x,)) Salā.](xe)
The proof can be found in Appendix B.3
We also need the following lemma:
Lemma Let be a group of isometries acting on S, and suppose the distribution Ge (xt 2o) is equivariant under the action of G. Then for any g E G. the following identity
holds:
E ( g (
(20
20

---

Proof. We start by expanding the conditional expectation:
(zolgom,)) (xolon
B of (zo
[5o] =
so
aG (2o g ( dro
/
(x dro-
To q (g 2o) q (g (go x4)
n
Now, apply the change of variable sip
g o Fn- Since g is isometry, the Lebesgue measure is
invariant, i.e., dæg = dēp. Therefore:
( Fo) drp (g
(g
Ta)
8 (goms lgollggan To)
n
/
p dīø
go To 9: (at
5o)
0. (xx,
g of [zo]
where the second equality follows from the equivariance of of and the invariance go under the group action. This concludes the proof.
Now the prove the Proposition
Proof. We begin by simplifying c(":
(0) (en [1 Exor (aul [zo]
(x |o(onn t) Bf ( [mro] (0)
Similarly,
(0) (d) E Ex [10(ar1)) BBo (2al 1
We now rewrite c.20 in integral form:
(0) = 1 [o(ean)))
-) -zo (mm) |co|] || a (x.) day.
Using the decomposition of the measure over the group action, we change variables:
[ f(g o 2z) dpc(g) drc. [ f(x)) J
Thus,
c (.) J / p6g o xe, of goz, (oo. [zo]]
2 qi q (g o x4) dua(g) dr.
= -
( 1 g b(x B (mol goz, [
of
(z,) duc(g) dra
where we used the G-equivariance of o and the invariance of q. Since G acts isometrically, we apply:
|g d( !) - Exg (m loo [sol] = |o(r),)) - g Exg (mo ggr)) (200 ||
Next, applying Lemma we have:
g ((r,, En (Zolgon,) [no]] || = | () - Exo (no)
21

---

We conclude:
L (Ø) J / ((F2: B ( |oo| || (xx))) dua dz.
= =1 (((2))t) - Beon || (x))
Next, we similarly transform (d):
(Φ) Sa | o Ie, () of (solson))) ) ( 2x) duols) dr
= - ((0n))) - ( (xo leoll (So de an) duolo) dx.
Using the theoretical result from Appendix B.3 we have:
q (x+) SG (x())(())= [ (algom)daall
conclude:
C 22 (0)
Jag [0(e,1)) Bf (mo (z)l [roo]] (xx) dr. (0).
/
Hence, the two loss functions equivalent in the sense that they yield the same gradients and minimizers with respect to é.
B.3 Symmetrized Forward Diffusion Distribution.
Below is the formal lemma and the proof for the symmetrized forward diffusion distribution.
Symmetrized Forward Diffusion Distributions
Lemma 1. Let ((z)) be empirical distribution and g (ma)) its symmetrized counterpart under a symmetry group G, defined by (zp)) = Scld) Suppose a forward diffusion process acting on q(20) yields time-dependent marginal distributions qu (æt). Let similar process act on qc (zo) to generate di (Ix). Then, for all t 2 0, the following holds:
( Sol0)((x)).
Proof. The marginal distribution at time step t diffusion process is defined as
(c(xx) [ qe (#x (a))( dro,
where qı( zo) = (x; a, 1) is the Gaussian diffusion kernel, which is equivariant under isometry group transformations. Specifically, for any g E G. we have
(x( xn) q: g zn).
The symmetrized marginal distribution at time is defined as
Sa(ād((x)= / (g o (x))))))))(
f [// a(goz, xo))(2o) dro duc(9)
1 / gelg 22 x))ā(z0) (g)d
22
22

---

Next, we compute the marginal distribution at time t when the forward process is applied to the symmetrized data distribution:
[ (x (o)Sc((((z)) dæ
(xx)
/ (( mo) / i - 2o)dpc())dzo
[ qe(ex zo)alg 20 (n)duo(g)da
[ [ q0(2s zo)ā(g zo do)dd())
Applying a change of variable vn + g- ro. we get
1 (1) g o zo)ā(g و zol)d(g- ro)dpa(g)
(x1))
a
( / (e [ ( (since the Jacobian of is 1)
{ (0lgo Ie drodppp))))) (by kernel equivariance).
Thus, have
(<,,) Sclā)((x))
This completes the proof.
B.4
Proof of Theorem
2
Unbiased gradient and equivariance of OrbDiff
Theorem Let G be a locally compact isometry group acting on data space Q and suppose
the forward kernels (xx,, mо) are G-invariant: q. (g xo) = qe (I, zo) for all
g E G. Then:
1. The OrbDiff target o² (Io, (I, f) satisfies: (z ze, = o= (zo: Fz. t) for all h E G.
2. The gradient of the OrbDiff loss (12) equals that of the ideal loss (5). i.e. OrbDiff provides an unbiased gradient estimator.
We first prove that the OrbDiff target is equivariant:
Proof. We compute (xo. h x,, using the definition:
(r h It, Z(h o Fz, vo) / (go (g) (hor go dp)(g).
1
By the equivariance of of, we have:
( a g xo) qt (xt h
o o zo).
Letting o' 11 g, so o g', and using the left-invariance the Haar measure μaa, we get:
[ (h o g' au) qf (x. Ig' o xx) dpo(gg))
("(zu, h oz Z(h o zı, 20)
Factoring out from the integrand gives:
ho Z(ho Iz, 20) / (g' xo) qt (z: i xo) duc(o')
23

---

It remains to show that Z(h Iz, 20) = [zx, mo), where:
(1r, = 1 ( ( 1g
Using the same substitution:
[hoo,, mn) Jc di (horr gooo) dualo)) = = l (xx 11 1 dua(g) = 1 41 (xxx 1 goo) dpa(s)) = 2(xr).
Therefore,
o (0 ho It, 1) = h o² (2o, t),
Next, we prove that OrbDiff yields an unbiased gradient estimator:
Proof. Our goal is to estimate the following gradient:
V CG (d) = 2E, 2B (a.))) [6(2),) t) En (x (sol (
(14)
Our proposed Rao-Blackwell loss function has the same gradient as 29 (Ø), but with reduced variance due to the use of the conditional expectation =f (nolus)
[|oten.)) t) Bzo ( [Fo ||]
CR (6) Esg -00 (1()) Pe (x.lz))
(15)
However, computing Ez (molas) [so] can be computationally expensive. To address this, we introduce OrbDiff, which uses biased proposal distribution to approximate this expectation using only samples from the orbit of the that generated z4. This yields the alternative target:
/ (g 2o) q. (g zu 24) dpc(g).
8" ("o, Ii,
(16)
Z( xz)
where the normalization constant is
Z(ro, x1) / qee zo 22) dg)(g).
(17)
This matches Equation which can be verified via straightforward transformations.
Although 6² differs from B (ms .x.) [Zo], we show that replacing the Rao-Blackwell target in Eq. (15) with 6* yields a loss whose gradient still matches the original gradient. This relies on the assumption that the forward conditional distribution is equivariant under the group action, i.c.,
(g 24 | 2 (x æo). Vg E G,
(18)
which is natural condition satisfied in many generative models such as diffusion models flow matching with isotropic Gaussian priors.
Under this assumption, the OrbDiff loss is given by:
c OrrDDiff
(Ø)
Exo (wrp) Eg 4f (a (mo)
[|0( t) 00(xo. Ft, t)[??]
(19)
Ez (m:) E - o (mal
[10( x" Fy, 1]
(20)
Taking the gradient with respect to o, we obtain:
V C OrbDiff (ő) = 2E 2 (x,) Pao [6(x,, t) 0² (ao, xz, t)]
(21)
2E (oz)) ((z,, t) (f 10 (xo, I'e t)]
(22)
Thus, ensure that OrbDiff yields the correct gradient, it suffices to show:
Bx (no ( (æo, Zz, t)] IE E (x [20]-
(23)
24

---

We compute:
/ 0" (g' ro, Fe. 1) " xx) duc(g' dzo.
(f (ro, Fr, 0) =
(24)
Substituting the expression for o* and applying the change of variables 9 g with inva-invariant Haar measure fG, we have:
- -aal
Z(g' vo. 2x) Ic (g o xxl)) o o xol 24) duc(o) qf (g' xo ax) dua (g')
dro
(25)
(0) (g 1z) duc(o)
(g' o xo x dp)(((())) dæo
(26)
Z(ro, x4t)
[/ (e' o vo x.) duc(s') doo
(g di (go mn dp)(())))
(27)
(( 2e)
(g x0) tl (g o 2o x dp)(()
(Iro, dro
(28)
/ (g xo) (g x④) dua(g) dag
(29)
E 0f [zo]-
(30)
Therefore, despite o² not being equal to the conditional expectation at each 2t, the gradient induced by the OrbDiff loss matches the desired V CJ (6), while using only samples from the orbit of zio-
gradient. OrbDiff thus provides unbiased estimate of
B.5
Unconstrained Non-Symmetrized Diffusion Minimizer is not guaranteed G-equivaraint
Lemma 2. Let
(((0) - [Ilo(z,, t) =all
be the unconstrained con-Symmetrized Diffusion loss, where @(zu, xx) is the empirical joint distribution of clean and noisy data. Then the minimizer 6² (I),t) of Le is the conditional expectation:
8" ('xx,, Ezz (10l
However, this minimizer is not guaranteed to be equivariant under the action of symmetry group G-equivariant.
Proof. To find the minimizer of the diffusion loss, we first compute the stationary point. The loss is given by:
(b) = (e [1lo(1), () xo 13] = Exeenge (wx) (20 (mol [1lo(xe, !) 2oll
The gradient of the internal expectation is:
(colas) [1lo(x,, t) xoll 2 (0(2t, vul
= 2 (6(0), t) - (moo) [zu])
Setting the gradient to zero gives the minimizer:
(2t, t) = (col)))))))) / Zoqi(Ho xx)) dap-
Next, provide counterexample to show that o (xz, is not guaranteed to be equivariant.
25

---

Counterexample: Translation in 1D
1. Data: Two points 23 = 0 and z 1, with uniform empirical distribution ((x)) 0.5.
2. Group action: Translation by 1, i.e., x
3. Diffusion kernel: (e(2z) zo) = ( 0,30, of)
First, rewrite the minimizer:
6 (x(, = L x2))
oo (xxa m ((rb)
(x(e
T (xt (g)((()))
(x(e
Substituting zi 0 and z 1:
1
(0- (0-N(zr;0, -0.5 + 1NN ar, -05). 6" (xr, (+())))
0.5
N,5-N(z;; o7)
Nzz
((t)))
Mz 0, +N
Now, compute (g (g Iı, t) by applying g oz = 1
N(x, + 1:0;; of)
o" (g = N(x, + 110 + N(xx 1;a0, <I
Next, compute g o² (xx, l):
N(z of) goo" M(z1;0,0F) N(mga +1>1 1.
Thus,
6" 6g Iz, go (x,.tt since o² (gonr.)) 1 and g o² (x,,t) 1. Consequently,
3 Ft, 1) # g Ø² (xt, l)-
This completes the counterexample, showing that x² (I () is not necessarily equivariant.
C Experimental Details
C.1 Molecular Conformer Generation (MCG) with Flow Matching
Molecular conformer generation is fundamental task in computational chemistry and drug discovery, where the goal is to generate plausible 3D structures (conformers) correspond to 2D molecular graph. Molecular conformer generation is essential for drug discovery and molecular property prediction, as the 3D structure greatly influences chemical behavior and interactions (Liu et al. 2023 Axelrod & Gómez- Bombarrlli 2020).
To model this task effectiively, it is crucial to respect the underlying symmetries of molecular structures. Since conformers are invariant under global rotations and translations, one might consider the full Euclidean group SE(3). However, in practice, molecular structures are typically zero- centered, effectively removing the need to model translation invariance. As result, it suffices to consider equivariance under the rotation group SO(3). addition, the molecular graph may exhibit symmetry under automorphisms- automophismss-perrtatoo permutationss of atoms that preserve the graph structure- structure-mmkknn it important to account for graph isomorphism to avoid redundant representations and ensure physically meaningful predictions.
26

---

C.1.1 MCG - Dataset
We evaluate our method on the GEOM-QM9 dataset (Axelrod & Gomez-Bombarell 2022). widely used subset of GEOM containing molecules with an average of 11 atoms. We follow the same traiin/validation/tes split as in (Ganea et al.) 2021 Jing et al. 2022). consisting of 106,586 / 13,323 / 1,000 molecules, respectively.
C.1.2 MCG - Basselnnes
We compare against strong recent baselines with publicly available code, including GEODIFF (Xu et al.
2022b), GEOMOL (Ganea et al. (022), Torsional Diffusion (Jing et al. 2022), MCF (Wang
al. 2024b) and ETFLOW (Hassan et al. 2024). GEODIFF generates structures using roto-
translationally invariant diffusion process, starting from an invariant initial density and evolving through Markov kernel that preserves this invariance. GEOMOL predicts 3D structures by modeling torsion angles conditioned on a molecular graph, offering fast inference and good geometric validity, Torsional Diffusion generates conformers using diffusion process torsion space. MCF directly models 3D coordinates using a diffusion model without enforcing equivariance, relying instead on model scale to achieve strong performance. ETFLOW, the strongest of these baselines, is an equivariant flow matching model that uses harmonic prior to encourage spatial proximity of bonded atoms. We integrate Orbit Diffusion into ETFLOW to build on its strong geometric foundation.
C.1.3 MCG - ETFFOW with Orbit Diffusion
While Orbit Diffusion is framed within the context of diffusion models, it naturally extends to flow matching. We introduce this extension through the design of ETFLOW, which employs a harmonic prior and flexible coupling to the data distribution. Unlike diffusion models, which fix the prior and reverse-time coupling, flow matching allows arbitrary choices for both (Tong et al. 2024), offering greater flexibility in the generative process.
Assume a coupling ((0. 21) between the base distribution qu and the data distribution qı. For each pair (Io, I, define the linear interpolation:
4 4 x1) = (1- (1-t)x0 + tx E 1].
Note: contrast diffusion models (where corresponds to data and 1 to noise), flow matching treats go as the prior sample and
Q1 as the data sample.
ETFLow defines the conditional distribution:
(x Zo. z1) N (z, 4 21), t)),
with small a, inducing the following velocity field:
1 2t e, e e~
vy (x,,) = x1 - Tu +
2 2//1(1=
Given the sampling equation:
Tt (1 -1-18 t)xo + tz + t ooott11 t)c,
can express
== Et - (1 -1-18 f)x00 - tæı
0 0 -t)
Substituting this into vy (xx) yields:
1 21
Uy (xx) xı
3o +
(Ir (1 - ()250 -txr)
2at(1 1
-)
1 21
1 1- 21
20t ) zo + (1- - 2/
22t(1 t)
1
20 - -t) 0
= =h(t)x - g(t) + f())x/
27

---

The model is trained to match this target velocity using the loss:
(Ø) (x ((100))))) [1 t) u2 (1))))
with gradient:
V Vol(o) =1)) B an) [2(((2r, vx(ee))]
A seooBlaakwellizeed gradient be derived by conditioning 21: V Vol(ø) = Em. [2 (0(zz, 6) - (B Ma. [l(t)x; g(t) + f(t)=1])
Ex [2 (Ø(#,, 6) - h(()x;; + (o [mo] - [t)
We a single sample zo estimate E2/ xr] then use the same estimation technique as Orbit Diffusion to compute E[ æz]. enabling efficient Rao-Blackwellized gradient estimation.
C.1.4 MCG Training Protocol
Instead of training from scratch, we finetune ETFlow using their public checkpoint. During training, we explicitly incorporate both forms of symmetry relevant to molecular data: discrete graph auto- morphisms and continuous spatial rotations. For each molecule, from its automorphism group using the pynauty library, capturing permutation symmetries in the 2D molecular graph structure. Simultancously, we sample 200 elements from the rotation group SO(3) to account for the contimuous rotational symmetries of its 3D conformation. This symmetry-aware augmentation is applied consistently across the dataset to ensure that the model learns to respect and exploit both types of equivariances. All other training settings, including optimizer configurations and learning rate schedules, follow the defaults of ETFLOW.
uniformly sample 50 elements
C.1.5 MCG Evaluation Protocol
In the test set, for each molecule with ground-truth conformers, we generate K = 2L conformers and evaluate their quality using standard metrics.
Evaluation Metrics. As conformer C represents assignment of each atom in the molecular graph to a point in 3D space, it can be viewed as set of vectors in R To evaluate molecular conformer generation, previous works have employed two key metries: Average Minimum RMSD (AMR) and Coverage (COV) for both Precision (P) and Recall (R). Given a molecular graph, we generate twice as many conformers as those provided by CREST. Let:
{Ci) 1E be the set of grounth- -truth conformers provided by CREST.
{= be the of generated conformers, where K 2L.
* 5 be a predefined RMSD threshold for considering conformer match.
COV-P: Measures the proportion of generated conformers that closely match at least one ground-truth conformer.
1 I[ E [] 31 E [1, RMSD( C7) < 8H1 K
COV-P
AMR-P:
Computes the average of the minimum RMSD values between each generated conformer and its closest ground-truth conformer.
1. min RMSD(C CT) 1=1
AMR-P
Measures the proportion of ground-truth conformers that have at least one close-enough generated conformer.
COV-R:
COV-R = |(1 E [1, 3k E [1, KJ. RMSD CP) < <8H
AMR-R:
Computes the average of the minimum RMSD values between each ground-truth conformer and its closest generated conformer.
1. K AMR-R
min RMSD(CCs Ci)
1=1
28

---

Figure 7: Molecular conformers generated by ETFFLOW (left), + [OrbDiff] (center), and ground-truth (right).
C.1.6
MCG Full results
Table 6: Molecular conformer generation performance GEOM-QM9. Reported in the original paper.
Obtained using the published checkpoint. From our reimplementation trained from scratch.
Recall
Precision
Models
1. (1)
Cov@0.5 (f)
AMR (L)
Cov@0, (f) Cov Cov@0.5 (?)
AMR (L)
Mean Median Mean Median Mean Median Mean Median Mean Median Mean Median
GEODIFF
76.5
100.0
0.297 0.229
50.0
33.5
1.524
0.510
GEOMoL,
28.4
0.0
91.1
100.0
0.224
0.194
20.7
0.0
85.8
100.0
0.271
0.243
Torsional Diff.
37.7
25.0
88.4
100.0
0.178
0.147
27.6
12.5
84.5
100.0
0.221
0.195
MCF
81.9
100.0
94.9
100.0
0.103
0.049
78.6
93.8
93.9
100.0
0.113
0.055
ETFLOW"
96.5
100.0
0.073
0.047
94.1
100.0
0.098
0.039
ETFLow
79.5
100.0
93.8
100.0
0.096 0.037
74.4
83.3
88.7
100.0
0.142
0.066
ETFLOW!
81.4
100.0
94.4
100.0
0.092
0.039
74.6
85.5
89.1
100.0
0.145
0.064
[OrbDiff]
85.4
100.0
96.3
100.0
0.074
0.027
80.2
93.9
91.9
100.0
0.113
0.042
C.2 Crystal Structure Prediction (CSP)
Crystal Structure Prediction (CSP) is the process of identifying the most stable three-dimensional arrangement of atoms in crystalline solid, given only its chemical formula. This task lies at the heart computational materials science, as the resulting crystal structure dictates key physical thermodynamic stability, electronic conductivity, and chemical reactivity (Wei et al. 2024 Shao et al. 2021 Kim et al. 2025). Despite its importance, CSP remains formidable challenge due to the immense combinatoril search space of atomic positions and the presence of complex symmetry constraints that define equivalent configurations (Oganov & Glass 2006). Efficiently navigating this landscape to discover low-energy, physically plausible structures continues to central focus of the field.
and chemical
We also evaluate our method on related task introduced by TGDMat (AS et al. 2025 where crystal structures are generated based on textual descriptions of the desired materials. This setting reflects more user-centric interface for materials design, where scientists domain experts can
29

---

specify target properties or structural features in natural language. The task includes two levels of textual conditioning: long descriptions, which provide detailed structural and compositional information, and short descriptions, which are more concise and easier to obtain. Supporting text-to- structure generation enables more accessible and flexible workflows in materials discovery, especially in scenarios where precise structural data may be unavailable.
C.2.1 CSP Dataset
We evaluate our method on two used CSP benchmarks: Perov-5 (Castelli et al. 2012a b) and MP- 20 (Jain et al., 2013). These datasets encompass a broad range of inorganic crystal compositions. Perov-5 comprises 18,928 perovskite structures, which share a common structural motif but vary in elemental composition. In contrast, MP-20 includes 45,231 stable inorganic materials curated from the Materials Project database (Jain et al. NA). offering more diverse set of crystal systems and chemistries.
C.2.2 CSP Baselines
For comparison, use the three strongest baselines from the DiffCSP paper (Tiao et al. (023) P-cG-SchNet (Gebauer et al., 2022), CDVAE (Xie et al. 2022). and DifCCSP, all with publicly available implementations.
C.2.3 CSP - Di and TGDMat with Orbit Diffusion
DiffCSP (Jiao et al. 2023) is diffusion-based framework for crystal structure prediction that jointly models lattice parameters and atomic positions while respecting the fundamental symmetries of crystalline materials. Since TGDMat is built upon DifCCS, we focus on describing how Orbit Diffusion integrates with DiffCSP; the same integration applies directly to TGDMat.
DiffCSP formulates the task as joint diffusion process with two interconnected components: one for the lattice and one for the atomic coordinates. The lattice defines the shape and scale of the unit cell, while the coordinates specify atomic positions in fractional units relative to the lattice vectors. To capture the relevant symmetries, the lattice diffusion (3)) equuivarra (invariant under rotations and reflections), and the coordinate diffusion is both permutation equivariant and periodic translation equivariant, reflecting atomic indistinguishabiiity and lattice periodicity.
Since the lattice involves only a few parameters, it is a relatively simple subtask. We therefore concentrate the more challenging component: generating atomic coordinates. To ensure periodic translation equivariance, DiffCSP defines forward diffusion process based on the kernel g0(2z:g(())) (x; æo). modeled as a Wrapped Normal distribution periodic analogue of the Gaussian-nnsuring the diffusion respects the toroidal geometry of fractional coordinates. Specifically,
||
+ =113
( 2O)
Σ exp =E2d
(31)
20
which defines valid density the torus TVd One can verify that ( g00z) = o w5)) for any in the periodic translation group (Jiao et al. 2023). confirming its equivariance. Consequently,
all our theoretical results (e.g., Theorem and Theorem 2 hold when applying Orbit Diffusion to
DifICSP under the periodic translation group.
Orbit Diffusion with non-uniform group sampling, Rather than sampling uniformly from the group G, for the periodic translation group we propose sampling translation elements from Wrapped Normal distribution. We refer to this variant as DrbDiff WN. Formally,
||mg + =gll3
ve(g) &
Σ exp
(32)
(t)2
where my is the translation vector corresponding to group element g. and oy is the time-dependent bandwidth. In our experiments with both DiffCSP and TGDMat, we set (t) = 25, To sample from this distribution, we first draw e No 7) in R&3, and then compute my = ()()) mod 1. We sample 1000 such group elements per step to form the group approximation.
30

---

More training details. All models were trained on a single NVIDIA GeForce RTX 4090 GPU. TGDMat was trained for 1.500 epochs, while DifiCSP was trained for 500 epochs. On the MP20 dataset, each epoch took roughly 15 seconds, resulting in total training times of 6.25 hours for TGDMat and 2.08 hours for DifiCSP. On the Perov-5 dataset, each epoch took about seconds, corresponding to 2.08 hours (TGDMat) and 0.69 hours (DiffCSP) of training time.
C.2.4 Evaluation Protocol
To evaluate crystal structure prediction, we randomly generate one sample for each structure in the test set. We then calculate two metries: the Match Rate and the average Root Mean Square Distance (RMSD) across the test set. We repeat this procedure three times and report the median values of these metrics for more reliability.
Match rate: The Match Rate is defined as the proportion of predicted structures that successfully match the corresponding ground-truth structures in the test set. Specifically, it is calculated follows:
Number matched structure pairs
Match Rate-
Total number of test samples
Following previous works (Jiao et al. from the pymatgen library to determine structure matching. The matching process is based on the following criteria:
2023 Xie al. 2022 we use the StructureMatchhr class
Length Tolerance (Itol): 0.5 (fractional length tolerance).
Site Tolerance (stol): 0.3 (fraction of the average free length per atom)
Angle Tolerance (atol): 10 (in degrees)
The StructureMatcher algorithm aligns the lattice vectors of two structures. If the tolerance criteria are satisfied, the structures are considered matched.
RMSD:
For an alignment between lattices two structures, StructureMatcher continues to align atoms compute the average RMSD. The process is repeated for all possible lattices to find the smallest RMSD. Then the Average RMSD is computed as the average of the smallest RMSD of all matched structure pairs.
Nussttad
RMSD
RMSD(ggneraaed, ground-truth;)
Nmaatcedd
Here, Nmmathhed is the total number of matched structures. Unmatched structures are excluded from the calculation.
Ideally, we aim for a high Match Rate and a low RMSD. low Match Rate with low RMSD is not useful because unmatched samples are effectively treated as having very high RMSD. Thus, RMSD alone cannot fully capture prediction quality. We emphasize that Match Rate is more critical, especially during initial screening, where we prioritize valid structures over perfectly matched ones.
C.2.5 CSP Full Results and Visualizations
For the TGDMat models (TGDMat (S) and TGDMat (L)), we trained both from scratch, experiment- ing with and without the proposed loss functions: [OrbDiff __] and [OrbDiff WN]. Meanwhile, the DiffCSP model was trained using our proposed losses, while the baseline models (P-cG-SchNet, CDVAE, and DiffCSP) rely on the results reported in the original DiffCSP paper. Quantitative results summarized in Tables and g with qualitative comparisons shown in Figures and
31

---

Table Text-guided CSP with TGDMat.
Table 8: Crystal Structure Prediction (CSP).
Method
Perov-5
MP-20
Perov-5
MP-20
(1)
RMSE ({)
(?)
(J)
Methed
Match (1)
RMSE (1)
Match (*)
RMSE (1)
TGDM (S)
[0xtDiff__
[bradsft_Ki]
63.51
59. 39
65.37
0.062
0.066
56.50
59.90
0.085
0,078
PchGeS
COVAE
48.22
0.418
15,39
0.305
0.376
TGDMa (L)
[CrbDaff_Nh
[OrbDiff_ .U]
93.88
95.17
0,012
0,012
0.013
65.94
61.91
0.069
0.081
DIFCSP
[OrtDifE UJ
[OxaDifr_ND)
$2.39
52.29
52.02
0.076
0.069
0.078
54.47
55.70
51.49
0.053
0.051
0.063
0.
bob
8
DIECSP
[[Obbiirr_u UJ [OrbBit WN] TODMAT (S) [Oraabi U] [Orbbirr WN] TCDMAT (L) [Graabin Ui
[GrIDir WN] Ground-ruth
Figure 8: Qualitative comparison of Crystal Structure Predictions by models, including DifICSP, TGDMAT short and TGDMAT long with baselines, OrbDiff U, and OrbDiff_Me against ground- truth samples on randomly selected samples from MP-20 dataset.
0
DifCCS?
[OrbDatt U) [OrbDic WN] ISDMNT [OrbDift [Orbbutt WN] TGDMAT [OrbDitt U] [OrbDit WN) Ground truth
Figure 9: Qualitative comparison of Crystal Structure Predictions by 9 models, including DiffCSP. TGDMat (S) and TGDmat (L) with baselines, OrbDiff. U, and OrbDiff WN against ground-truth samples on randomly selected samples from Perov-5 dataset.
32

---

C.3 Protein Structure Generation (PSG) with Non-Equivariant Denoiser PROTEINA
Protein structure generation focuses on sampling physically valid 3D conformations of proteins by learning a probabilistic distribution over either atomistic or coarse-grained representations. This generative task plays a crucial role de novo protein design and has broad implications for under- standing protein folding and function (Jing et al., 2023 Wu et al. 2022 Watson et al. 2023 Huguet cct al. 2024; Jing et al. 2023). Unlike traditional structure prediction, which aims to infer a single or most likely conformation, generative models must capture the full structural manifol--accountinn accounting for the inherent 53(3) rotational symmetry of protein backbones, maintaining biochemical realism, and respecting physical constraints such bond lengths and steric clashes (Gaujac et al., 2024 Geffner et al. (025). Effectively modeling these aspects ensures that generated structures are not only diverse but also biologically and physically plausible.
C.3.1 PSG Baselines
We list the baselines used by the work of Geffner al. 2025) as follows:
FrameDiff
Yim et
2023b
FoldFlow Bose et al. 2024)
FrameFlow (Yim et al.
2023a
ESM3
(Hayes et al. 2025
Chroma Ingraham al. 2023
RFDiffusion (Watson al. 2023
Proteus (Wang et 2024a)
Genie2 (Lin al. 2024)
PROTEINA. PROTEINA is large-scale, flow-based model for generating protein backbones, built on scalable transformer architecture. Unlike equivariant models, it does not enforce equivariance, granting more architectural flexibility. This choice allows the use powerful transformer networks with hundreds of millions of parameters, enabling PROTEINA to effectively learn from large datasets. As a result, excels at modeling complex protein structures by balancing expressiveness and computational computationss scalability.
There three variants of PROTEINA:
(i) Mps. a 200M-parameter transformer with an additional 15M parameters in triangle layers, trained on the Foldseek AFDB clusters dataset Drs. which includes 555,318 structures of lengths 32-256; (ii) Mapsinn simplified version without triangle layers, also with 200M parameters and trained on the same dataset;
(iii) M21M. 400M parameter transformer with 15M triangle parameters, trained on a high-quality filtered AFDB subset D2IM. comprising approximately 21M structures. M21M represents the current state of the art in designability modeling.
C.3.2 PROTEINA with Orbit Diffusion
We apply Orbit Diffusson to the simplest variant of PROTEINA, namely M by fine-tuning the public checkpoint released by Geffner et al. (025),, since we do not have access to the extensive compute resources used for the original training. Our fine-tuning setup uses 4 A100 GPUs for 24 hours on the same dataset, whereas the original training employed 96 GPUs.
The key symmetry group in protein structure generation is S0(3), which represents 3D rotations. To exploit this symmetry using OrbDiff, we apply 10,000 uniformly sampled random rotations each training sample as part of our Rao-Blackwell estimator, enhancing both the efficiency and stability of the flow matching process.
The only change we introduce to the original model is replacing the conditional flow matching loss with our proposed flow matching objective (see Appendix C.I.3)
For a fair comparison, we also fine-tune M (+ [finetune]) using the original loss under the same computational budget, with batch size of 8 and 32 gradient accumulation steps.
33

---

C.3.3 PSG Evaluation Protocol
To evaluate the quality of our generated protein backbones, we rely three widely used metrics: designability, diversity, and novelty. Following the protocol established by Geffner et al. (025). we generate 500 samples total-100 for each length in (50, 100, 150, 200, 250] and compute all metrics on this dataset. Among these, designability is the most critical metric, as it directly reflects the biological feasibility of the generated structures and serves the foundation for the other two metrics.
Designability. Designability measures whether backbone structure can realistically be encoded by an amino acid sequence. For each generated backbone, we produce 8 candidate sequences using ProteinMPNN (Dauparas et al., (2022) with a sampling temperature of 0.1. These sequences are then folded using ESMFold (Lin et 2003), and the root mean square deviation (RMSD) is calculated between each predicted structure and the original backbone. A backbone deemed designable least one sequence folds with an RMSD below 2Å, where this minimum RMSD known as the
self-consistency RMSD (scRMSD).
Since diversity and novelty are computed only designable samples, accurate assessment of designability is essential for interpreting the other metrics meaningfully.
The designability score for model is reported the fraction of samples deemed designable. Additionally. we report the average scRMSD across all samples, allowing for more nuanced comparison between our model and existing baselines.
Diversity. We evaluate diversity among the designable samples using two approaches. First, we compute the average pairwise TM-score within each length group (50, 100, 150, 200, and 250) as measure of structural variation. Lower average TM-scores indicate greater diversity.
Second, calculate diversity (cluster) by grouping designable samples into clusters based TM-score threshold of 0.5. Each cluster contains samples with pairwise TM-scores above this threshold. Diversity (cluster) is then defined the ratio of the total number of clusters to the number of designable samples. A higher ratio reflects larger number of distinct structural groups relative to the sample size, signaling increased diversity.
Novelty. Novelty measures how structurally distinet the designable samples are compared to known protein structures. For each designable backbone, we compute the TM-score its closest match in two reference sets: the Protein Data Bank (PDB) and the Drs dataset used for training. The average these maximum TM-scores across all designable samples reported as the novelty score. Lower values indicate the model generates structures that are more novel relative to both established experimental data and the training distribution.
C.3.4 PSG Full results
Table Protein Structure Generation. Full comparison with baseline models; all baseline results are taken from (Geffner et al. 2025 "+ [finetune]" indicates M finetuned with the original loss,
while "+ [OrbDiff]" denotes hnntuning with OrbDiff.
Designability
Diversity
Novelty
Model
Fraction (t)
scRMSD (L)
Cluster (t)
TM-score ())
PDB (L)
AFDB (L)
FrameDiff
65.4
0.39
0.40
0.73
0.75
FoldFlow (base)
0.42
0.75
0.75
0.77
FoldFlow (stoc.)
97.0
0.61
0.38
0.62
0.68
FoldFlow (OT)
97.2
0.37
0.41
0.71
0.75
FrameFlow
88.6
0.59
0.34
0.79
0.80
ESM3
22.0
0.52
0.57
0.70
0.75
Chroma
78.8
0.42
0.43
0.77
0.76
RFDiffusion
94.4
0.46
0.34
0.79
0.80
Proteus
94.4
0.42
0.43
0.77
0.80
Genie2
95.2
0.59
0.38
0.63
0.69
M2u
99.0
0.72
0.30
0.39
0.81
0.84
Mgg
93.8
1.04
0.62
0.36
0.69
0.76
[finctune]
93.8
1.00
0.54
0.37
0.74
0.83
[OrbDiff]
95.6
0.93
0.52
0.37
0.74
0.83
34

---

The state-of- ofte-offthheart model M21M. with 400M parameters and trained on large, high-quality dataset, achieves the highest designability (99.0%) and lowest scRMSD (0.72), reflecting its strong recon- struction capability. However, this comes at the expense of diversity and novelty: it exhibits the lowest diversity score (Cluster: 0.30) and higher novelty metrics (PDB: 0.81, AFDB: 0.84), indicating reduced structural variety and generalization.
In comparison, our base model Mis already achieves competitive performance (designability: 93.8%, Cluster: 0.62, PDB: 0.69), and naive finetuning with data augmentation ("+ [finetune]]") fails improve designability or diversity. Notably, our method (*+ [OrbDiff]") improves designability to 95.6% and reduces scRMSD to 0.93, while preserving competitive diversity (Cluster: 0.52) and novelty (PDB: 0.74). This highlights that OrbDiff is an effective finetuning strategy enhance functional accuracy without fully sacrificing structural diversity-cchieving achieving a better trade-off than both naive finetuning and heavily overparameterized models.
35

---

NeurIPS Paper Checklist
1. Claims
Question:
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The abstract and introduction clearly state the paper's main claims and contri- butions, which align with the theoretical and experimental results presented.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
We have discussed about limitation in the Section 6
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
36

---

Justification:
We provided comprehensive proof of our theoretical results in the Appendix
section.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: We provide detailed descriptions the experimental setup, model architecture, hyperparameters, training procedure, and evaluation protocols.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, or other means that appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
37

---

Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [NA]
Justification: The paper reuses existing benchmark datasets and builds on top of previous open-source models.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips... cc/
for more details.
While encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https SrlodeSuumisssionPP..
for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing much information possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Question:
Answer: [Yes]
Justification: We specify relevant training and testing details, including data splits, hyper- parameters, and optimizer choices. These details align with those used in the reused code repositories for the bascline methods.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, appendix, or as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [No]
Justification: Not all experiments include error bars or statistical significance measures due to the high computational cost of training some models, limiting repeated runs.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emoor bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
38

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not show in tables figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification:
The paper details the compute resources used of all experiments in Appendix
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics https: peattps://nurrre newrips........ wwt.rrisscccpuul///eeee scc/public/Bthhccchiideellll
Answer: [Yes]
Justification:
The research fully complies with the NeurIPS Code of Ethics in all aspects.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- cration due to laws or regulations their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: This work is foundational research without direct socictal impact or applica- tions that raise ethical concerns.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
39

---

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
This paper does not release data or models with high risk for misuse.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: This work reuses existing assets with MIT license.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URI.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
40

---

If assets are released, the license, copyright information, and terms of use in the thcode. has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
package should be provided. For popular datasets, paperswi
latasets.paperssit
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: This paper does not introduce or release any new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing or any research with human subjects.
Guidelines:
The answer NA means that the paper doess not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country institution) were obtained?
Answer: [NA]
Justification: This research does not involve human subjects or crowdsoureing.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
41

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification: LLMs were not used in any core or non-standard part the method. Their use, if any, was limited to minor writing edits and did not impact the research.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to our LLM policy attt:e///oorrpss.cc/ooofrrrreeee for what should or should not be described.
42


Reviews:
[{"review_id": "girsqk4Lyy", "reviewer": "Reviewer_j92t", "summary": "The paper proposes Orbit Diffusion as a Rao-Blackwellized gradient estimator to train equivariant diffusion models. The core idea is to Rao-Blackwellize the targets, not the denoiser, by conditioning on the symmetrized data distribution. \n\nFor a practical computation of the conditional expectation, self-normalised importance sampling is used with a user-defined proposal.", "questions": "How sensitive are experiment results to the choice of biased or unbiased proposals?", "limitations": "See weaknesses", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 3, "significance": 4, "originality": 4, "strengths_and_weaknesses": "Strengths\n\nThe paper is written and presented well, with sufficient background on equivariance, group symmetrization and diffusion models to describe the proposed method.\n\nThe derivation of the Rao-Blackwellized gradient estimator is detailed. \n\nEmpirical results support the proposed method.\n\nThe approximation strategy for the conditional expectation of $x_0 \\sim \\hat{q}_t^{G}(x_0|x_t)$ is stated clearly as the main limitation and an avenue for future work.\n\nWeaknesses\n\nWhile the use of a biased proposal is well-motivated (in lines 170-171), experiment results should include an ablation study for using a biased proposal vs an unbiased proposal.", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your valuable feedback and for pointing out the areas of improvement in our manuscript.\n\n### **Weakness & Question #1: Ablation study of biased proposal and unbiased proposal**\n\nThank you for raising this insightful point. We agree that an ablation study would strengthen the paper by illustrating how the biased proposal distribution impacts final performance. To clarify, whenever we refer to the proposal or proposal distribution, we mean the proposal used in the self-normalized importance sampler (SNIS) for estimating the inner expectation $\\mathbb{E}[x_0 \\mid x_t]$. This estimate in OrbDiff is biased because the proposal’s support does not fully cover that of the target distribution $\\hat{q}_t^G(x_0 \\mid x_t)$. However, as we will show later, this bias cancels out when computing the gradient, making OrbDiff an **unbiased estimator of the training gradient.**\n\nWe evaluated our approach on the Crystal Structure Prediction task using the Perov-5 dataset (DiffCSP Perov-5), which involves modeling a conditional distribution: given a list of atom types, generate a valid 3D crystal structure. Since an atom list can correspond to multiple valid structures, the task is inherently generative. In the Perov-5 training set, 3,330 out of 8,026 atom lists map to two distinct structures; the rest map to one.\n\nFor an atom list with two valid structures, $x_0$ and $x_0^\\prime$. Estimating $\\mathbb{E}[x_0∣x_t]$ ideally requires sampling from both orbits. However, OrbDiff’s biased proposal samples only from $x_0$​'s orbit.\n\nWe conduct an ablation study with the following model variants:\n- DiffCSP_bp_1k: Use biased proposal distribution with 1000 samples from the orbit of $x_0$.\n- DiffCSP_up_1k: Use unbiased proposal distribution with 1000 samples from both orbits of $x_0$ and $x_0^\\prime$.\n- DiffCSP_up_2k: Use unbiased proposal distribution with 2000 samples from both orbits of $x_0$ and $x_0^\\prime$.\n\n| Setting                          | Match Rate (↑)     | RMSD (↓)           |\n|----------------------------------|------------------|-----------------|\n| DiffCSP_bp_1k       |   **52.29**   |  0.078   |\n| DiffCSP_up_1k   |   52.02   |  0.090  |\n| DiffCSP_up_2k   |   52.13   |  **0.072**  |\n\nFrom the table, we observe that with the same sampling budget of 1,000, the model trained using the **biased** estimate of  $\\mathbb{E}[x_0 \\mid x_t]$ outperforms the one using the **unbiased** estimate. When doubling the sampling budget to 2,000, the unbiased variant (DiffCSP_up_2k) achieves a lower RMSD and a Match Rate close to the biased variant (DiffCSP_bp_1k), indicating that both approaches are effective depending on the target metric. This aligns with our theoretical analysis: both biased and unbiased estimates lead to the same **unbiased gradient in expectation** (see the proof below).\n\n\n## **Additional Clarification [Relevant to all Reviewers]:**\n\nIn revisiting our analysis during the rebuttal—prompted particularly by Reviewer 2’s question on the potential bias of OrbDiff—we discovered an important theoretical property of our method that was not recognized during the initial writing of the paper: the **OrbDiff gradient estimator is unbiased**, even though the OrbDiff SNIS estimation of the inner expectation $\\mathbb{E}[x_0 | x_t]$ is biased. \n\nSince OrbDiff samples $x_0$ only from the orbit that gave rise to $x_t$, rather than from the full symmetrized distribution $\\hat{q}^G(x_0)$, the support of OrbDiff's proposal distribution does not cover that of the original distribution. As a result, this approximation introduces bias in estimating the target value $\\mathbb{E}[x_0 \\mid x_t]$. **Crucially, however, the resulting gradient estimator remains unbiased.**\n\nThis distinction is important: for training to be statistically sound, what ultimately matters is the unbiasedness of the gradient.\n\nTo summarize, and to have a better clarity, we have the following table:\n#### Distinguishing RB Estimator vs. OrbDiff\n| Aspect       | Rao-Blackwell Estimator | OrbDiff |\n|--------------|--------------------------|---------|\n| **Target** | Exact $\\mathbb{E}_{x_0 \\sim \\hat{q}_t^G(x_0 \\mid x_t)}[x_0]$ | Monte Carlo estimate using samples **only from the orbit** of the true $x_0$ |\n| **Sampling** | From full $\\hat{q}^G(x_0)$, includes all possible orbits | From the orbit of the training sample $x_0$ |\n| **Estimator** | Unbiased for both target and gradients | Biased estimator of the target but **unbiased gradients** |\n| **Practicality** | Inefficient and nontrivial to implement | Efficient and easy to implement |\n\nOur goal is to estimate the following gradient (Eq.(6) in the paper):\n\n$$\n\\nabla_\\phi \\mathcal{L}_t^{G}(\\phi) = 2\\mathbb{E} _{x_t \\sim \\hat{q}_t^G(x_t)} \\left[  \\phi(x_t, t) - \\mathbb{E} _{x_0 \\sim \\hat{q}_t^G(x_0 \\mid x_t)} [x_0]  \\right].\n$$\n\nOur proposed Rao-Blackwell loss function (Eq.(13) in the paper) has the same gradient as $\\mathcal{L} _t^G(\\phi)$, but with reduced variance due to the use of the expectation $\\mathbb{E} _{x_0 \\sim \\hat{q}_t^G(x_0 \\mid x_t)}[x_0]$ :\n\n$$\n\\mathcal{L}_t^{\\mathrm{RB}}(\\phi) = \\mathbb{E} _{x_0' \\sim \\hat{q}^G(x'_0)} \\mathbb{E} _{x_t \\sim \\hat{q} _t^G(x_t \\mid x_0')} \\left[\\lVert \\phi(x_t, t) - \\mathbb{E} _{x_0 \\sim \\hat{q}_t^G(x_0 \\mid x_t)}[x_0]   \\lVert^2 \\right].\n$$\n\nHowever, computing $\\mathbb{E}_{x_0 \\sim \\hat{q}_t^G(x_0 \\mid x_t)} [x_0]$ can be inefficient and nontrivial—for example, it may require sampling from data orbits not present in the current batch. To address this, we introduce OrbDiff, which uses a biased proposal distribution to approximate this expectation using only samples from the orbit of the $x_0$ that generated $x_t$. This yields the alternative target:\n\n$$\n\\phi^*(x_0, x_t, t) = \\frac{1}{Z(x_0, x_t)} \\int_{G} (g \\circ x_0) \\hat{q}_t^G(g \\circ x_0 \\mid x_t) \\mathrm{d}\\mu_G(g),\n$$\n\nwhere the normalization term is\n\n$$\nZ(x_0, x_t) = \\int_{G} \\hat{q}_t^G(g \\circ x_0 \\mid x_t) \\mathrm{d}\\mu_G(g).\n$$\n\nThis matches Eq.(14) in the paper, which can be verified via straightforward transformations.\n\nAlthough $\\phi^*$ differs from \n$\\mathbb{E} _{x _0 \\sim \\hat{q} _t^G(x _0 \\mid x _t)} [x _0]$, \n\nwe show that replacing the Rao-Blackwell target with $\\phi^*$ yields a loss whose gradient still matches the original gradient. This relies on the assumption that the forward conditional distribution is equivariant under the group action (such as Gaussian kernel), i.e.,\n\n$$\n\\hat{q}_t^G(g \\circ x_t \\mid g \\circ x_0) = \\hat{q}_t^G(x_t \\mid x_0), \\quad \\forall g \\in G,\n$$\n\nUnder this assumption, the OrbDiff loss is given by:\n\n$$\n\\mathcal{L}_t^{\\mathrm{OrbDiff}}(\\phi) = \\mathbb{E} _{x_0 \\sim \\hat{q}^G(x_0)} \\mathbb{E} _{x_t \\sim \\hat{q}_t^G(x_t \\mid x_0)} \\left[ \\lVert \\phi(x_t, t) - \\phi^*(x_0, x_t, t) \\lVert^2 \\right] \\\\\n$$\n\n$$\n\\qquad \\qquad = \\mathbb{E} _{x_t \\sim \\hat{q}_t^G(x_t)} \\mathbb{E} _{x_0 \\sim \\hat{q}_t^G(x_0 \\mid x_t)} \\left[ \\lVert \\phi(x_t, t) - \\phi^*(x_0, x_t, t) \\lVert^2 \\right].\n$$\n\nTaking the gradient with respect to $\\phi$, we obtain:\n\n$$\n\\nabla_\\phi \\mathcal{L}_t^{\\mathrm{OrbDiff}}(\\phi) = 2 \\mathbb{E} _{x_t \\sim \\hat{q}_t^G(x_t)} \\mathbb{E} _{x_0 \\sim \\hat{q}_t^G(x_0 \\mid x_t)} \\left[ \\phi(x_t, t) - \\phi^*(x_0, x_t, t) \\right] \n$$\n\n$$\n\\qquad \\qquad = 2 \\mathbb{E} _{x_t \\sim \\hat{q}_t^G(x_t)} \\left[ \\phi(x_t, t) - \\mathbb{E} _{x_0 \\sim \\hat{q}_t^G(x_0 \\mid x_t)} [\\phi^*(x_0, x_t, t)] \\right].\n$$\n\nThus, to ensure that OrbDiff yields the correct gradient, we will show $\\nabla _\\phi \\mathcal{L} _t^{G}(\\phi) = \\nabla _\\phi \\mathcal{L} _t^{\\mathrm{OrbDiff}}(\\phi)$. It is suffice to verify:\n\n$$\n\\mathbb{E} _{x_0 \\sim \\hat{q}_t^G(x_0 \\mid x_t)} [x_0] = \\mathbb{E} _{x_0 \\sim \\hat{q}_t^G(x_0 \\mid x_t)} [\\phi^*(x_0, x_t, t)].\n$$\n\nWe compute:\n\n$$\nRHS = \\int _\\Omega \\int _G \\phi ^*(g' \\circ x _0, x _t, t) \\hat{q} _t^G(g' \\circ x _0 \\mid x_t)  \\mathrm{d}\\mu _G(g')  \\mathrm{d}x_0.\n$$\n\nSubstituting the expression for $\\phi^*$ and applying the change of variables $g \\mapsto g \\cdot g'$ with left-invariant Haar measure $\\mu_G$, and since $Z(x_0, x_t)$ is invariant to group transformations of $x_0$, we have:\n\n$$\nRHS = \\int_\\Omega \\int_G \\left[ \\frac{1}{Z(g' \\circ x_0, x_t)} \\int_G (g \\circ [g' \\circ x_0]) \\hat{q}_t^G(g \\circ [g' \\circ x_0] \\mid x_t) \\mathrm{d}\\mu_G(g) \\right] \\hat{q}_t^G(g' \\circ x_0 \\mid x_t) \\mathrm{d}\\mu_G(g') \\mathrm{d}x_0\n$$\n\n$$\n= \\int_\\Omega \\int_G \\left[ \\frac{1}{Z(x_0, x_t)} \\int_G (g \\circ x_0) \\hat{q}_t^G(g \\circ x_0 \\mid x_t) \\mathrm{d}\\mu_G(g) \\right] \\hat{q}_t^G(g' \\circ x_0 \\mid x_t) \\mathrm{d}\\mu_G(g') \\mathrm{d}x_0\n$$\n\n$$\n= \\int_\\Omega \\left[ \\frac{1}{Z(x_0, x_t)} \\int_G (g \\circ x_0) \\hat{q}_t^G(g \\circ x_0 \\mid x_t) \\mathrm{d}\\mu_G(g) \\right] \\left[ \\int_G \\hat{q}_t^G(g' \\circ x_0 \\mid x_t) \\mathrm{d}\\mu_G(g') \\right] \\mathrm{d}x_0\n$$\n\n$$\n= \\int_\\Omega \\left[ \\frac{1}{Z(x_0, x_t)} \\int_G (g \\circ x_0) \\hat{q}_t^G(g \\circ x_0 \\mid x_t) \\mathrm{d}\\mu_G(g) \\right] Z(x_0, x_t) \\\\mathrm{d}x_0\n$$\n\n$$\n= \\int_\\Omega \\int_G (g \\circ x_0) \\hat{q}_t^G(g \\circ x_0 \\mid x_t) \\mathrm{d}\\mu_G(g) \\mathrm{d}x_0\n$$\n\n$$\n= \\mathbb{E}_{x_0 \\sim \\hat{q}_t^G(x_0 \\mid x_t)} [x_0].\n$$\n\nTherefore, despite $\\phi^*$ not being equal to $\\mathbb{E}[x_0 | x_t]$, the gradient induced by the OrbDiff loss matches the desired gradient. **OrbDiff thus provides an unbiased estimate of $\\nabla_\\phi \\mathcal{L}_t^{G}(\\phi)$**, while using only samples from the orbit of $x_0$."}, {"review_id": "4kl4AkM5Tw", "reviewer": "Reviewer_yVWb", "summary": "This paper introduces Orbit Diffusion, a training framework for denoising diffusion models that respects symmetry in data, such as molecular or protein structures with SE(3) invariance. By interpreting data augmentation as a Monte Carlo approximation of a symmetrized loss and applying Rao–Blackwellization, the authors derive a low-variance gradient estimator that remains compatible with both equivariant and non-equivariant architectures. The method constructs denoising targets by averaging over symmetry group orbits using importance sampling, leading to improved stability and convergence.", "questions": "1. Can Orbit Diffusion handle approximate or partial symmetries—i.e., when the symmetry group is only approximately valid in the data?\n2. How sensitive is Orbit Diffusion to the choice of group sampling distribution $\\nu_t(g)$, especially in high-dimensional or non-compact groups like SE(3)?\n3. How does the method handle symmetries that act only on a subset of the input (e.g., local symmetries, hierarchical groups)?\n4. What would it take to extend this method to learn symmetry when the group is not known a priori (e.g., data-driven symmetry discovery)?", "limitations": "Yes!", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "STRENGTHS\n1. Introduces a provably lower-variance gradient estimator, grounded in classical statistical theory.\n2. Provides a clear theoretical link between data augmentation and group symmetrization.\n3. Enforces equivariance at the loss level, enabling training of both equivariant and non-equivariant models.\n4. Proposes self-normalized importance sampling (SNIS) over group actions, which is tractable and sample-efficient.\n5. Adds negligible computational cost to large models.\n\nWEAKNESSES\n1. The paper qualitatively discusses the bias–variance trade-off, but does not offer rigorous bounds or quantitative guarantees on the bias introduced. Figure 11 shows small empirical bias, but only for a toy molecule with 8 conformers. No real-world downstream impact of this bias is analyzed.\n2. Sampling from large symmetry groups (e.g., permutations, SE(3), periodic translations) requires thousands of samples (e.g., 10,000 for SO(3) in PROTEINA), which may not be feasible in large-scale 3D scenes or protein assemblies. The paper provides no ablation on the trade-off between number of orbit samples and performance (e.g., how many samples are needed before diminishing returns). No hardware or memory scaling benchmarks are provided.\n3. For several benchmarks (e.g., GEOM-QM9, MP-20), the evaluation is missing standard deviation/error bars, especially for high-variance metrics like AMR or RMSD. In some cases (e.g., ETFLOW), the reported baseline results could not be reproduced, raising concerns about reproducibility and benchmarking fairness.", "was_revised": true, "final_justification": "The authors have thoroughly addressed all of my concerns, and I have adjusted my rating accordingly.", "rebuttal": "Thank you for your valuable feedback and for pointing out the areas of improvement in our manuscript.\n\n### **Weakness #1: More insight and guarantee about the bias**\nThank you for raising this point. Upon revisiting our analysis, we realized that although OrbDiff's SNIS estimation of $\\mathbb{E}[x_0 \\mid x_t]$ is biased, **the resulting gradient estimator is actually unbiased**. This happens because the bias cancels out when computing the outer expectation of the gradient. In the original paper, the term ``bias'' referred only to the approximation of $\\mathbb{E}[x_0 \\mid x_t]$, not to the gradient itself. This distinction is crucial, as training depends on the unbiasedness of the gradient.\n\nTo further clarify, Figure~11 is meant to illustrate how well the OrbDiff target $\\phi^*$ approximates $\\mathbb{E}[x_0 \\mid x_t]$ in a toy setting---not to assess bias in the gradient estimator. The quantity that truly matters for training is the gradient bias, which we now prove to be zero.\n\nDue to space limitations, we include the full proposition and proof in our **response to Reviewer 1 (j92t)**—please refer to that for details.\n\n### **Weakness #2: Scalability and Efficiency of Symmetry Group Sampling**\nThank you for the thoughtful question. We agree that scalability is an important concern for large symmetry groups. However, it's worth highlighting that the main cost in large-scale training typically lies in the forward, backward, and optimization steps of large neural networks. OrbDiff, by contrast, computes the group-averaged target entirely outside the network, making it extremely efficient.\n\nIn practice, the additional memory and computation required for sampling and averaging over group elements are negligible relative to the overall training cost. For instance, in our ProteinA experiments—one of the largest protein diffusion models—we use 10,000 SO(3) samples per training example. The added memory is just ~40MB per GPU, compared to ~54GB used during training.\n\nImportantly, we also measured the runtime overhead directly, as shown in Figure 5a. Even with 10,000 orbit samples, OrbDiff adds only 0.1% overhead to the total training time for ProteinA. To put this into perspective, in a 24-hour training run, this amounts to just 1.4 additional minutes—demonstrating that the method is highly efficient in practice.\n\nWe will update the paper to highlight not only the minimal training-time overhead but also the low memory footprint of OrbDiff, to better address concerns about its scalability.\n\nTo further assess downstream performance, we conducted an ablation study on the Text-guided Crystal Structure Prediction task using the Perov-5 dataset (TGDMat Perov-5), where the goal is to generate a crystal structure given a set of atoms and a natural language description of a plausible material. In this study, we sampled uniformly from the periodic translation group (OrbDiff\\_U) to demonstrate the effect of sampling size on performance.\n\n**Performance on the TGDMat Perov-5.** \n| # Samples | Match Rate (↑)    | RMSD (↓)        |\n|-----------|---------------|--------------|\n| 10        | 63.38  | 0.062 |\n| 100       | 64.28  | 0.055 |\n| 1000      | 65.57       | **0.054**         |\n| 5000      | **65.60**   | 0.056 |\n\nAs shown above, increasing the number of samples from 10 to 1,000 improves the match rate from 63.38 to 65.57 and reduces RMSD from 0.062 to 0.054. However, beyond 1,000 samples, performance has converged, with no significant improvements observed at 5,000 samples. Therefore, using 1,000 samples offers a good balance between performance and efficiency.\n\n### **Weakness #3: Reproducibility and Statistical Rigor in Benchmark Evaluation**\n\nThank you for highlighting this important point. While prior work (including ETFlow and others) typically omits such statistics, we agree with the reviewer that including them helps better assess result variability and statistical significance.\n\nIn response, during the rebuttal phase, we have managed to add the standard deviations for the QM9 benchmark. For each method, we perform evaluation three times and report the mean and standard deviation, shown in the following table:\n\n**Molecular conformer generation performance on GEOM-QM9.** \n\n\\* Reported in the original paper.  † Obtained using the published checkpoint. ‡ We train the public implementation from scratch. \n\n| Model | COV\\@0.1-R_mean (↑) | COV\\@0.1-R_median (↑) | AMR-R_mean (↓) | AMR-R_median (↓) | COV\\@0.1-P_mean (↑) | COV\\@0.1-P_median (↑) | AMR-P_mean (↓)  | AMR-P_median (↓)|\n|-------|----------------|------------------|------------------|--------------------|------------------|--------------------|--------------------|----------------------|\n| GeoMol† | 28.4±0.9 | 0.0±0.0 | 0.224±0.001 | 0.192±0.002 | 20.8±0.2 | 0.0±0.0 | 0.270±0.001 | 0.240±0.004 |\n| Torsional Diff.† | 37.0±0.6 | 23.7±1.6 | 0.178±0.000 | 0.146±0.001 | 27.3±0.4 | 12.0±0.9 | 0.220±0.001 | 0.193±0.002 |\n| MCF† | 80.9±0.9 | **100.0±0.0** | 0.102±0.002 | 0.050±0.001 | 78.7±0.1 | **94.4±0.6** | **0.112±0.003** | 0.052±0.002 |\n| ETFlow* | - | - | 0.073 | 0.047 | - | - | 0.098 | 0.039 |\n| ETFlow† | 80.2±0.8 | **100.0±0.0** | 0.097±0.004 | 0.035±0.002 | 74.4±0.4 | 83.3±0.0 | 0.143±0.003 | 0.070±0.004 |\n| ETFlow‡ | 81.9±0.6 | **100.0±0.0** | 0.089±0.006 | 0.036±0.002 | 74.8±0.2 | 84.5±1.1 |0.143±0.002 | 0.064±0.002 |\n| + [OrbDiff] | **85.3±0.3** | **100.0±0.0** | **0.073±0.001** | **0.027±0.001** | **80.0±0.2** | 93.2±0.6 | **0.112±0.000** | **0.043±0.002** |\n\nETFLow + [OrbDiff] consistently outperforms state of the art models on various metrics. Notably, we observe that the standard deviations for AMR, especially for OrbDiff, are relatively small, indicating the stability in evaluation. If the reviewer still feels strongly about including standard deviations, we are happy to accommodate that in the revised version of the paper.\n\nRegarding ETFLow - reproducibility:\nIn the current version of the paper (Table 2), we have already taken care to address the reproducibility of ETFLOW by reporting three separate results:\n(i) the numbers reported in the original ETFlow paper,\n(ii) results obtained by running the official ETFlow code from scratch, and\n(iii) results obtained using a model checkpoint shared with us directly by the authors.\n\n\n### **Question #1: Orbit Diffusion and approximate or partial symmetries.**\nOrbDiff can work with partial symmetries as long as there is an efficient method for sampling from the corresponding set of group actions. For instance, if we want to restrict equivariance up to some degree of rotation, or if we have only partial permutation equivariance. The former by simple restrictions of existing algorithms for sampling rotations, the latter with the product replacement algorithm.\n\nExtending the framework to explicitly model approximate or data-driven symmetries is a promising direction for future work. Such extensions could involve adaptive group selection or learning symmetry parameters jointly with the model, thereby relaxing the assumption of exact group structure. This is, however, beyond the scope of our work. \n\n### **Question #2: The sensitivity of OrbDiff with the choice of group sampling distribution.**\n\nWe thank the reviewer for the insightful question. As noted in lines 69--70 of the paper, we focus on locally compact isometry groups, which cover a broad range of practical applications. Extending the method to non-compact groups is certainly possible and presents an interesting direction for future work. In high-dimensional groups, it may be necessary to increase the number of group samples to ensure sufficient coverage of the group distribution.\n\nNevertheless, we note that standard diffusion training can be viewed as sampling only the identity element of the group. In contrast, OrbDiff samples multiple group elements, which we have shown reduces the variance of the gradient estimator without introducing bias. Given this, along with the efficiency of OrbDiff, we believe it is generally advisable to prefer OrbDiff over standard training whenever group symmetries are known and we can sample from the set of group actions.\n\n\n### **Question #3: Orbit Diffusion and local symmetries or hierarchical groups.**\n\nWe thank the reviewer for this interesting question. In principle, the framework could be extended to handle local symmetries by defining group actions on different regions and applying our SNIS-based estimator accordingly. Hierarchical symmetries could be approached via multi-scale or nested sampling strategies. We consider these valuable directions for future work.\n\n### **Question #4: Extending this method to learn symmetry when the group is not known priori.**\n\nWe thank the reviewer for raising the possibility of learning symmetries from data. While our method assumes known group actions, it is orthogonal and potentially complementary to symmetry discovery methods such as [1], [2] or [3]. These methods offer promising directions for automatic symmetry identification. Future work could explore coupling our loss-level equivariant framework with a learned group structure, possibly enabling generalization to approximate or local symmetries. \n\n[1] Generative adversarial symmetry discovery\n\n[2] AtlasD: Automatic Local Symmetry Discovery\n\n[3] Automatic symmetry discovery with lie algebra convolutional network"}, {"review_id": "CdoHjdFdZL", "reviewer": "Reviewer_LfDH", "summary": "This paper considers a novel variance reduction technique for the gradients of the squared loss in diffusion models with symmetry through Rao-Blackwell's Theorem. It offers practical guidance called Orbit Diffusion that can be applied to either equivariant or non-equivariant neural network architecture. Numerical simulations on various datasets verify the usefulness of the proposed method. In general, this work offers a practical bias-variance trade-off dedicated to the equivariant diffusion models.", "questions": "1. Line 94, $q(x_T) =\\mathcal{N}(0,I)$ should be $q(x_T) \\approx\\mathcal{N}(0,I)$ since you would not achieve stationarity in finite time. Also, the notation is confusing, I think it might be better to replace $q(x_0), q(x_T)$ by $q_0(x), q_T(x)$. \n\n2. What is $D$ in line 123? Has it been defined before?\n\n3. For Eq. (4) to be equivalent to Eq. (5), you may need to make certain assumptions on the noising kernel.\n\n4. What's the computational burden for the SNIS?\n\n5. What if you don't bias as you did in lines 170-171? How is the performance improvement?", "limitations": "Yes.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strength: The paper is written clearly, and the idea is simple but interesting. It offers practical contributions to the training of equivariant diffusion models.\n\nWeaknesses: How the variance reduction of the gradient improves performance is still underexplored. Orbit Diffusion includes bias.", "was_revised": true, "final_justification": "I thank the authors for the detailed response and the additional experiments, both of which have addressed my questions and concerns.", "rebuttal": "Thank you for your valuable feedback and for pointing out the areas of improvement in our manuscript.\n\n### **Weakness #1: How the variance reduction of the gradient improves performance is still underexplored.**\n\nTo study the practical benefit of variance reduction, we conducted extensive experiments across diverse tasks—molecular conformer generation, crystal structure prediction, and protein structure generation. In all these settings, we observe consistent improvements in model performance when using OrbDiff, demonstrating that lower-variance gradients translate into better downstream results. Furthermore, from diffusion research literature, high-variance gradient estimators can lead to unstable or inefficient optimization. Specifically, prior work has shown that reducing gradient variance during the training of diffusion models can lead to improved final performance [1, 2, 3].\n\n[1] Variational Diffusion Models\n\n[2] Stable Target Field for Reduced Variance Score Estimation in Diffusion Models\n\n[3] Improved Denoising Diffusion Probabilistic Models\n\n### **Weakness #2: Orbit Diffusion includes bias**\n\nThank you for highlighting this important point. We emphasize that the gradient estimator used to train the model **remains unbiased**, even though OrbDiff’s SNIS-based estimate of the inner expectation $\\mathbb{E}[x_0 \\mid x_t]$ is biased. This happens because the bias cancels out when computing the outer expectation of the gradient. This subtle but crucial distinction was not clearly conveyed in the original version of the paper. We appreciate the reviewer for prompting this clarification and will revise the text to clearly state that OrbDiff yields unbiased gradient estimates, despite relying on a biased proposal for efficiency.\n\nPlease find the **proof in the response to Reviewer 1’s comment.**\n\n### **Question #1: Notation clarity**\n\nThank you for the helpful feedback. We agree with the first point and will revise the text to write $q(x_T) \\approx \\mathcal{N}(0, I)$ to reflect the fact that exact stationarity may not be achieved in finite time.\n\nRegarding the second point, we acknowledge the potential for confusion, but we follow the standard notation in the diffusion literature (e.g.,[1]), where $q(x_t)$ denotes the distribution of $x_t$​ at time $t$. We will clarify this in a footnote or early in the notation section to help readers unfamiliar with this convention.\n\n[1] Denoising Diffusion Probabilistic Models\n\n### **Question #2: What is D in line 123?**\n\nWe thank the reviewer for the question. It is the dataset. We will add the definition of the notation to the manuscript. \n\n### **Question #3: The equivalent of equations (4) and (5)**\n\nThank you for the thoughtful comment. We realize that the equivalence between Eqs. (4) and (5) may not be immediately clear due to the notation. In Eq. (4), we write $x_0 \\sim \\hat{q}^G(x_0)$, whereas in Eq. (5), the distribution of $x_t$ is written as $\\hat{q}^G_t(x_t)$, which could give the impression that different distributions are being used.\n\nTo clarify: both equations are expressing expectations under the same joint distribution over $(x_0,x_t)$, where $x_0 \\sim \\hat{q}^G(x_0)$, and $x_t \\sim \\hat{q}^G_t(x_t \\mid x_0)$, defined via a group-equivariant noising process. Since this defines a well-formed joint distribution $\\hat{q}^G_t(x_0, x_t)$, we can apply the law of iterated expectations in either order, giving equivalent forms: $\\mathbb{E} _{x_0 \\sim \\hat{q}^G(x_0)}\\mathbb{E} _{x_t \\sim \\hat{q}^G_t(x_t \\mid x_0)}[⋅]  = \\mathbb{E} _{x_t \\sim \\hat{q}^G_t(x_t)}\\mathbb{E} _{x_0 \\sim \\hat{q}^G_t(x_0 \\mid x_t)}[⋅]$\n\nWe will update the text to clarify this point explicitly and refine the notation to avoid potential confusion between $\\hat{q}^G(x_0)$ and $\\hat{q}^G_0(x_0)$, which we used interchangeably.\n\n### **Question #4: Computational burden of SNIS**\n\nThank you for the thoughtful question. It's worth highlighting that the main cost during training typically lies in the forward, backward, and optimization steps of large neural networks. OrbDiff with SNIS, by contrast, computes the group-averaged target entirely outside the network, making it extremely efficient.\n\nIn practice, the additional memory and computation required for sampling and averaging over group elements are negligible relative to the overall training cost. For instance, in our ProteinA experiments—one of the largest protein diffusion models—we use 10,000 SO(3) samples per training example. The added memory is just ~40MB per GPU, compared to ~54GB used during training.\n\nImportantly, we also measured the runtime overhead directly, as shown in Figure 5a. Even with 10,000 orbit samples, OrbDiff adds only 0.1% overhead to the total training time for ProteinA. To put this into perspective, in a 24-hour training run, this amounts to just 1.4 additional minutes—demonstrating that the method is highly efficient in practice.\n\nWe will update the paper to highlight not only the minimal training-time overhead (reported in Figure 5a) but also the low memory footprint of OrbDiff, to better address concerns about its scalability.\n\n### **Question #5: What if we do not bias?**\n\nFirstly, thanks to the reviewer's question, we take a closer look at how bias is OrbDiff and we find an interesting finding that **although OrbDiff's SNIS estimation of $\\mathbb{E}[x_0 | x_t]$ is biased due to the choice of an efficient proposal distribution, the resulting gradient estimator is in fact unbiased**. This happens because the bias cancels out when computing the outer expectation of the gradient. \n\nWe provide a detailed proof of this result in **our response to Reviewer 1 (j92t).**\n\nGiven this, we believe that the performance difference between using a biased versus an unbiased proposal for estimating $\\mathbb{E}[x_0 | x_t]$ is minimal when the sampling budget is sufficiently large (since both of them produce unbiased gradients). However, employing an unbiased proposal is both inefficient and not straightforward. For instance, computing the target for the loss function of a minibatch under an unbiased proposal requires sampling $x_0$ from outside the minibatch, complicating implementation. Moreover, as analyzed in the paper, the additional information gained from sampling external orbits is often negligible. \n\nGiven that **OrbDiff yields an unbiased estimator of the gradient**, we believe it is still valuable to include an ablation comparing the performance of biased and unbiased proposals.\n\nWe evaluated our approach on the Crystal Structure Prediction task using the Perov-5 dataset (DiffCSP Perov-5), which involves modeling a conditional distribution: given a list of atom types, generate a valid 3D crystal structure. Since an atom list can correspond to multiple valid structures, the task is inherently generative. In the Perov-5 training set, 3,330 out of 8,026 atom lists map to two distinct structures; the rest map to one.\n\nFor an atom list with two valid structures, $x_0$ and $x_0^\\prime$. Estimating $\\mathbb{E}[x_0∣x_t]$ ideally requires sampling from both orbits. However, OrbDiff’s biased proposal samples only from $x_0$​'s orbit.\n\nWe conduct an ablation study with the following model variants:\n- DiffCSP_bp_1k: Use biased proposal distribution with 1000 samples from the orbit of $x_0$.\n- DiffCSP_up_1k: Use unbiased proposal distribution with 1000 samples from both orbits of $x_0$ and $x_0^\\prime$.\n- DiffCSP_up_2k: Use unbiased proposal distribution with 2000 samples from both orbits of $x_0$ and $x_0^\\prime$.\n\n| Setting                          | Match Rate (↑)     | RMSD (↓)           |\n|----------------------------------|------------------|-----------------|\n| DiffCSP_bp_1k       |   **52.29**   |  0.078   |\n| DiffCSP_up_1k   |   52.02   |  0.090  |\n| DiffCSP_up_2k   |   52.13   |  **0.072**  |\n\nFrom the table, we observe that with the same sampling budget of 1,000, the model trained using the **biased** estimate of  $\\mathbb{E}[x_0 \\mid x_t]$ outperforms the one using the **unbiased** estimate. When doubling the sampling budget to 2,000, the unbiased variant (DiffCSP_up_2k) achieves a lower RMSD and a Match Rate close to the biased variant (DiffCSP_bp_1k), indicating that both approaches are effective depending on the target metric. This aligns with our theoretical analysis: both biased and unbiased estimates lead to the same **unbiased gradient in expectation**."}, {"review_id": "8sAOIVXVlV", "reviewer": "Reviewer_veFd", "summary": "This paper introduces OrbDiff, a framework for training group-invariant generative diffusion models that leverages Rao–Blackwellized gradient estimators to reduce variance. The method can be used in both equivariant model designs and data augmentation schemes and provably improves the variance of standard gradient estimators. OrbDiff is tested on molecular conformation generation, crystal structure prediction, and protein design. Experiments suggest improvements in gradient variance, and convergence speed.", "questions": "1. ___Clarification on the Data Assumptions:___ The paper uses symmetrized distributions $\\hat{q}^G$ to define the loss under data augmentation. Is the assumption that the underlying data distribution $q$ is already $G$-invariant? If so, symmetrization doesn't do anything. If not, what justifies learning a $G$-invariant distribution in the first place? A more formal treatment of this modeling assumption would clarify the motivation and benefit the paper. E.g., I'm assuming that the \"true\" data distribution is $G$-invariant but the samples a biased in some way.\n\n2. ___Justification of the High Variance Claim:___ The paper states that the standard gradient estimator (Equations (4) and (5)) suffers from high variance. However, this estimator is just the typical diffusion gradient with augmented samples. The claim that it has higher variance than in the non-invariant case is plausible but not theoretically motivated explicitly. Could the authors provide either a mathematical argument or additional experiments to support this claim? This would significantly improve the argument of the of paper. \n\n3. ___Clarity of Theorem 1:___ The statement of Theorem 1 compares the variance of two gradient estimators, but Equations (5) and (7) define expectations, not estimators. To analyze variance, one needs to specify how the estimators are constructed—e.g., which quantities are sampled via Monte Carlo. The theorem is likely correct with a reasonable interpretation, but a clearer statement would benefit the paper.\n\n4. ___Notes:___ \n    - The method is referred to as “Rao–Blackwell Gradient Estimation” but the Rao–Blackwell theorem is never explicitly mentioned in the main body (only in the proof of Theorem 1 in the Appendix). It would be helpful to summarize the relevant argument in the main text rather than relegating it to the Appendix.\n    - Consider presenting the RB loss $\\mathcal{L}_t^\\mathrm{RB}(\\phi)$ (Equation (13)) earlier in the paper, as it is more interpretable and makes the goal of the method clearer. Starting from the loss makes it easier to understand what the method is doing and how it works in practice.", "limitations": "yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**Strengths:**\n- Conceptually simple method with provable variance reduction guarantees.\n- The method is broadly applicable in symmetry-aware generative modeling. It works with both equivariant architectures and data augmentations.\n- Extensive empirical evaluation across diverse domains.\n\n**Weaknesses:**\n- The motivation for focusing on the variance in the gradient estimator is not clearly formalized or empirically quantified outside the synthetic setup.\n- The connection to the Rao–Blackwell theorem is not clearly explained in the main text, despite being central to the method.", "was_revised": true, "final_justification": "The paper presents a clean, general method for reducing gradient variance in symmetry-aware diffusion models, with solid theory and strong results across tasks. While the main text underexplains some theoretical points, the rebuttal clarifies them well. My main concern is impact: the paper doesn't clearly show that high variance is (a) especially severe in equivariant models or (b) meaningfully harmful. Overall, this is a well-motivated contribution, but the paper would benefit from earlier clarity and stronger evidence that the addressed problem is significant.", "rebuttal": "Thank you for your valuable feedback and for pointing out the areas of improvement in our manuscript.\n\n\n### **Weakness #1: The motivation for focusing on the variance in the gradient estimator is not clearly formalized or empirically quantified outside the synthetic setup**\n\nWe appreciate the reviewer’s comment regarding the motivation behind our focus on the variance in the gradient estimator. Our primary motivation is based on the insight that high-variance gradient estimators can lead to unstable or inefficient optimization. Moreover, prior work has shown that reducing gradient variance during the training of diffusion models can lead to improved final performance [1, 2, 3].\n\nIn revisiting our analysis during the rebuttal—prompted particularly by Reviewer 2’s question on the potential bias of OrbDiff—we identified an important theoretical property that was not highlighted in the original manuscript: **the OrbDiff gradient estimator is unbiased**, even though it relies on a proposal distribution that lead to a biased estimation of $\\mathbb{E}[x_0 \\mid x_t]$. This finding further supports our motivation for using OrbDiff, as it enables variance reduction without introducing bias into the gradient estimate. \n\nAdditionally, OrbDiff incurs only minimal extra computational and memory overhead. For the full proof of the unbiasedness result, please refer to **our response to Reviewer 1(j92t)**. \n\n[1] Variational Diffusion Models\n\n[2] Stable Target Field for Reduced Variance Score Estimation in Diffusion Models\n\n[3] Improved Denoising Diffusion Probabilistic Models\n\n\n### **Weakness #2: The connection to the Rao–Blackwell theorem is not clearly explained in the main text, despite being central to the method**\n\nThank you for pointing out this concern. We will include the RB theorem in the main paper to show the connection. \n\n### **Question #1: Clarification on the Data Assumptions**\n\nWe appreciate the reviewer’s insightful question. To clarify, we do assume that the underlying data-generating distribution $q$ is **G**-invariant. However, the observed distribution $\\hat{q}$​ (from which we obtain training samples) is generally not **G**-invariant due to dataset biases. For example, in molecular datasets, each molecule may be saved in a canonical but arbitrary orientation, even though physically all rotated versions are equally probable under $q$.\n\nAs a result, training directly on $\\hat{q}​$ would lead to a model that may not respect the underlying symmetry. This is a widely recognized issue in the literature, and standard solutions include:\n- Using equivariant networks, which enforce symmetry by architectural design.\n- Using data augmentation, where samples are transformed by group actions during training.\n\nBoth approaches can be seen as implicitly training on the symmetrized distribution $\\hat{q}^G$ [4]. We formulate the loss under the symmetrized distribution to provide a unified framework for analyzing existing methods like data augmentation and equivariant architectures. This perspective clarifies their implicit inductive biases and reveals new opportunities for improvement—most notably, the construction of lower-variance gradient estimators, which directly motivates OrbDiff.\n\n[4] Equivariant score-based generative models provably learn distributions with symmetries efficiently\n\n### **Question #2: Justification of the High Variance Claim**\n\nThank you for pointing this out. We apologize for the confusion caused by our original statement. The claim in the paper that “this estimator, used in data augmentation, suffers from high variance” at line 126 in the paper is imprecise. \n\nWe meant to convey that gradient estimators in diffusion model training in general, whether with or without data augmentation, are known to suffer from high variance, as supported by prior work [1, 2, 3]. \n\nWe will revise the statement to “Gradient estimators used in diffusion model training—including those with data augmentation—are known to suffer from high variance.”\n\nWe appreciate the reviewer’s careful reading and will include additional citations and empirical evidence to clarify this point.\n\n[1] Variational Diffusion Models\n\n[2] Stable Target Field for Reduced Variance Score Estimation in Diffusion Models\n\n[3] Improved Denoising Diffusion Probabilistic Models\n\n### **Question #3: Clarity of Theorem 1**\n\nThank you for the helpful observation. You're right that Equations (5) and (7) define expectations, while variance applies to their finite-sample estimators. In practice, we use standard minibatch estimators during training, and Theorem 1 compares their variances through the expectations they approximate.\n\nWe will clarify this in the paper and explain that the result is motivated by the Rao-Blackwell theorem: replacing a random group sample with its conditional expectation leads to a lower-variance estimator. This justifies our analysis at the level of expectations and explains the benefit of the OrbDiff estimator in practice. \n\n### **Question #4: Suggestion for improving the paper**\n\nThank you for the suggestion. We agree and will move the key idea behind the Rao–Blackwell theorem into the main text to better motivate our method. We’ll also present the Rao-Blackwellized loss (Eq. 13) earlier in the paper, as it provides a clearer and more intuitive starting point for understanding OrbDiff in practice."}]

Author rebuttal:

Meta-review: All the four reviewers lean towards acceptance. That said, reviewers have raised several areas for improvement that the authors should address carefully in the final version, including adding clarifications, new analysis and enhancing the presentation and organization of the paper to minimize confusion. My recommendation is acceptance.
Decision: Accept (poster)