SymDiff: Equivariant Diffusion via Stochastic Symmetrisation
Authors: Leo Zhang
Kianoosh Ashouritaklimi
Yee Whye Teh
Rob Cornish
Venue: ICLR 2025 Poster (2025)
Keywords: Equivariance
Diffusion Models
Symmetrisation
Molecular Generation
Markov Categories
Primary area: generative models
TL;DR: We propose SymDiff, a novel method for constructing equivariant diffusion models using the recently introduced framework of stochastic symmetrisation.

Abstract: We propose SymDiff, a method for constructing equivariant diffusion models using the framework of stochastic symmetrisation. SymDiff resembles a learned data augmentation that is deployed at sampling time, and is lightweight, computationally efficient, and easy to implement on top of arbitrary off-the-shelf models. In contrast to previous work, SymDiff typically does not require any neural network components that are intrinsically equivariant, avoiding the need for complex parameterisations or the use of higher-order geometric features. Instead, our method can leverage highly scalable modern architectures as drop-in replacements for these more constrained alternatives. We show that this additional flexibility yields significant empirical benefit for E(3)-equivariant molecular generation. To the best of our knowledge, this is the first application of symmetrisation to generative modelling, suggesting its potential in this domain more generally.

Published as a conference paper at ICLR 2025
SYMDIFF: EQUIVARIANT DIFFUSION VIA STOCHASTIC SYMMETRISATION
Leo Zhang Kianoosh Ashouritaklimi Yee Whye Teh Rob Cornish Department of Statistics, University of Oxford
ABSTRACT
We propose SYMDIFF, method for constructing equivariant diffusion models using the framework of stochastic symmetrisation. SYMDIFF resembles a learned data augmentation that is deployed at sampling time, and is lightweight, computa- tionally efficient, and easy to implement on top of arbitrary off-the-shelf models. In contrast to previous work, SYMDUFF typically does not require any neural network components that are intrinsically equivariant, avoiding the need for com- plex parameterisations the use of higher-order geometric features. Instead, our method can leverage highly scalable modern architectures as drop-in replacements for these more constrained alternatives. We show that this additional flexibility yields significant empirical benefit for -33e-quuivarian molecular generation. To the best of our knowledge, this is the first application of symmetrisation generative modelling, suggesting its potential in this domain more generally.
INTRODUCTION
For geometrically structured data such as N-body systems of molecules or proteins, it is often of interest to obtain diffusion model that is equivariant with respect some group actions (Xu jet al. 2022; Hoogeboom et al. [022] Watson el al.] 2023 Yim et al. (023). By accounting for the large number of permutation and rotational symmetries within these systems, equivariant diffusion models aim to improve data efficiency (Batzner et al., 2022) and generalisation (Elesedy & Zaidi 2021 In previous work, equivariant diffusion models have been implemented using intrinsically equivariant neural networks, each of whose linear and nonlinear layers are constrained individually to be equivariant, so that the overall network also equivariant as a result (Bronstein et al. 2021).
There is growing literature on the practical limitations of intrinsically equivariant neural networks (Wang et al.) [024] Canez et al.) 2024 Abramson et al., 2024, Pertigkiozoglou et al.) 2024). It has been observed that these models can suffer from degraded training dynamics due to their imposition of architectural constraints, as well as increased computational cost and implementation complexity
[Dvval et al.] 2023a). To address these issues, there has been recent interest in symmetrisattin techniques for obtaining equivariance instead (Murphy et al. 20188] Puny et al.] [202 Duval et al.
2023b Kim al., 2023 Kaba et al.] 2023 Mondal al., [023]] Dym al. [024] Gelberg al.,
2024]. These approaches offer a mechanism for constructing equivariant neural networks using
subcomponents that are not equivariant, which previous work has shown leads better performing
models (Yarotsky) 2022] Kim et al.] 2023). However, the advantages of symmetrisation-bert equivariance have not yet been explored for generative modelling. This is possibly in part because previous work has solely focused on deterministic equivariance, rather than the more complex condition of stochastic equivariance that is required in the context of generative models.
In this paper, we introduce SYMDIFF, a novel methodology for obtaining equivariant diffusion mod- els through symmetrisation, rather than intrinsic equivariance. We build on the recent framework of stochastic symmetrisation developed by Cornish using the theory of Markov categories (Fritz, 2020). Unlike previous work on symmetrisation, which operates on deterministic functions, stochas- tic symmetrisation can be applied to Markov kernels directly in distribution space. We show that this leads naturally to a more flexible approach for constructing equivariant diffusion models than is pos- sible using intrinsic architectures. We apply this concretely to obtain an (3)))-vrrvari diffusion architecture for modelling N --bdy systems, where E(3) denotes the Euclidean group. For this task, formulate custom reverse process that is allowed to be -an-Gaasi,n,, for which we derive

---

Published as a conference paper at ICLR 2025
tractable optimisation objective. Our model is stochastically -quuivaaraat overall without need- ing any intrinsically B(3)-equivarian) neural networks as subcomponents. We also sketch how to extend SYMDIFF score and flow-based generative models (Song et al. 2020 Lipman et al. 2022).
To validate our framework, we implemented SYMDUFF for de novo molecular generation, and evaluated it
drop-in replacement for the F(3)-equivvriant diffusion of [Hooeeboom al. (2022), which relies on intrinsically equivariant neural networks. In contrast, our model is able to leverage highly scalable off-the-shelf architectures such Diffusion Transformers (Peebles & Xie 2023) for all of its subcomponents. We demonstrate this leads to significantly improved empirical performance for both the QM9 and GEOM-Druge datasets.
2
BACKGROUND
We provide here an overview of the underlying theory behind equivariant diffusion modelling. This theory most conveniently developed in terms of Markov kernels, whose definition we recall first.
2.1 EQUIVARIANT MARKOV KERNELS
Markov kernels At a high level, a Markov kernel k: : X' + y may be thought of as conditional distribution or stochastic map that, when given an imput x E X, produces random output y with distribution k(dy) x)) For example, given a function f X' x 3 - y' and random element of & there is Markov kernel : X + y for which k:dy x) is the distribution of f(x, n) case, every deterministic function f X' y may be thought of as Markov kernel X' Y also. When k(dy's x) has a density (or likelihood), we will denote this by kyy[x), although we note that can still reason about Markov kernels even when they do not admit likelihood in this sense.
As a special
Stochastic equivariance Let G be a group acting on spaces X' and V. We will denote this using "dot" notation, that the action on X' is function (g. x) + g x. Recall that function X is then equivariant
f(g x) g. f(x
for all X' and g G.
Here f is purely deterministic, and so this concept must be generalised in order to encompass models whose outputs are stochastic (Bloem-Reddy & Teh 2020). To this end, [ornish (2024) uses notion defined for Markov kernels: a Markov kernel k : X + Y' is stochastically equivariant (or simply equivariant) if
k(dylg x) g k(dylx) for all x E X and g E G,
(2)
where the right- hand side denotes the distribution of g when y ~ k(dylxx), or in other words the pushforward of k:(dy)x (x) under g. When i obtained from f and 1 as above, equation 2 holds iff
f(g x, g f(x. n)
for all x E X' and g E G.
where _ denotes equality in distribution If v is constant, this says is deterministically equivariant in the usual sense. Likewise, when l: has a conditional density, say k(y|))) xx) equation holld if k(g y x) k(ylx)
for all E X, y V. and g 9,
provided the action of G on y has unit Jacobian (Cornish) 2024 Proposition 3.1 as will be the case for all the actions we consider. This latter condition recovers the usual formulation of stochastic equivariance considered in the diffusion literature by e.g. Xu et al. (002)): Hoogeboom et al. 2002))
Stochastic invariance When the action of G on y is trivial, the condition in equation ferred to as invariance. This same idea carries over to Markov kernels also: we say that k above is stochastically invariant (or simply invariant) if k(dylg x) k(dylx) for all g G and X. Importantly, this differs from another natural notion of invariance that also arises in the stochas- tic context: we will say that a distribution p(dy) on y is distributionally invariant if it holds that g p(dy)= p(dy) for all E G. Given density p(y), this holds equivalently if
re-
plg y) = p(y) for all g E G.
again assuming the action on y has unit Jacobian.
For "nice" choices of Y, the converse also holds by noise outsourcing (Kallenberg 2002 Lemma 3.22). Note this more general than the almost sure condition from equation (17) of (llee-RReda & Teh 2020
2

---

Published as a conference paper at ICLR 2025
2.2 EQUIVARIANT DIFFUSION MODELS
Diffusion models construct a generative model po(zo) of an un- known data distribution Pdata (Z)) on space by learning to reverse an iterative forward noising process 2₁- Following the notation of [Knggma et al. (222)) the distribution of 2, is defined by ((z,) (z) [iz 2o I for some noise schedule olt, oe >> such that the signal-to-nosee ration SNR(t) = of oo is strictly monotonically decreasing. The joint distribution of the forward and
Denoising diffusion models
reverse processes respectively then have the forms:
T
q(Zo:T) = ((2)) II gg(i|2z-1 1=1
po(2o:T) p(2T) II po (o 2e (2x)
with ((z)) = Pdau (20), and q(z, (2-1) = M(2 1, -(111 with constants defined as adjt-1 == a1/ac-1 and - of 1 We take p(zz) to also be Gaussian. The reverse process then trained to maximise the ELBO of log po(o), which can be obtained as follows Sohl-DDickstein et al. 2015):
(oo)) 2 E Eg(z (20) llo [ DrL (rr(eerrzo))ll(err)
Σ Be(a (z) [KLL 20) (4)
This objective can be efficiently optimised when the reverse process is Gaussian. This is due to the fact that the posterior distributions y(2r-1/Z4: 2o) N(2, Mug (Zt, 2o). oo (() are Gaussian, and that the KL divergence between Gaussians can be expressed in closed form. To match these posteriors, the reverse process is then typically defined in terms of a neural network μ4 : Z Z as
po (z (2)) ((z pm))) (2)) (1)).
Invariant and equivariant diffusion It often desirable for pe (Zo) to be distributionally invariant
with respect to the action of group G. Intuitively, this says that the density po(2o) is constant the orbits of this action. For the model in equation 3 distributional invariance follows if po 2r) is itself distributionally invariant, and if each pe (2 (.) is stochastically G-equivariant (Xu et al. 2022
All previous work we are aware of has approached this by obtaining a deterministically equivariant μe. which then implies that the reverse process stochastically equivariant (Le al. 2023).
2.3 N-BODY SYSTEMS AND ((3)-EQUIVARIANT DIFFUSION
Equivariant diffusion models are often applied to model N-body systems such as molecules and proteins (Xu et al. 2022: Hoogeboom et al., 2022 Yim et al. 2023). This is motivated by the large number of symmetries present in these system. For example, intuitively speaking, neither the coordinate system nor the ordering of bodies in the system should matter for sampling. We describe the standard components of such models below.
N-body data In 3-dimensions, the state of an N N-body system can be encoded as a pair z = [x, h] E TR (3+4d) where x = (x(1)...., x()) R describes a set of N points 3D space, Each (6), For example, Hoogeboom et al. (2022) encodes molecules in this way, where each (i) denotes the location of an atom, and !t) some corresponding
and h = (h ...., h()) E Nxd describes a set of N feature vectors of dimension feature vector (() is associated with the point
properties such as atom type (represented as continuous quantities).
Center of mass free space Intuitively speaking, for many applications, the location of an N-body system in space should not matter. For this reason, instead of defining a diffusion on the full space of N-body systems directly. previous work (Garcia Satorras et al. [2011a] Xu et al.) 2022] Hoogeboom et al. 2022) has set Z = U, where U is the "center of mass (CoM) free" linear subspace of R (3++d consisting of [x, h] such that x == 1 Σ x = 0. In this way, samples from their model are always guaranteed to be centered the origin.
3

---

Published as a conference paper at ICLR 2025
CoM-free diffusions To construct their forward and reverse processes to now live entirely l instead of (3++)) Xu et al. (2022) defines the projected Gaussian distribution Nul I)),, for Ja and o² O, as the distribution of z obtained via the following process:
<~ (0,,)
z = ju+ + a (roju(E),
where proju centers its input in (3++33+)) the origin, i.e. proju ((x, hl) == (x-xxxxxx---.. x). h]. By construction, the projected Gaussian distribution is then supported on the linear subspace U. Xu et al. (202)) shows that this distribution has a density with Gaussian form NZy μu, I N ju, o'Ij defined for all living in the subspace U. This allows defining forward and reverse processes (<<t-1 (24) and po(2t-11 1 (2x) with exactly the same form in Section 2.2 before, but with Ner used everywhere in place of N. Since these processes are still Gaussian (albeit on linear subspace), the KL terms in equation 4 remain tractable. which allows optimising the ELBO in the usual way. This approach does require 1ig U > U now to be constrained to produce outputs in the subspace 4. Prior work has achieved this simply by taking μs to be neural network with projy as its final layer.
Invariance and equivariance Intuitively, the ordering of the N points and the orientation of the overall system in 3D space should not matter. To formalise this, let SN denote the symmetric group and O(3) denote the group of orthogonal 3 x 3 matrices. Their product Si x ((3) acts on N-body systems by reordering and orthogonally transforming points follows:
of permutations of the integers { N..,
(6) where a E SN and R E O3).. Previous work (Xu et al. 2022 Hoogeboom et al. 2022) has then chosen the model po(z)) to be distributionally invariant this action. They enforce this via the
(o, R) [x, h] = [Rx(o(1))
(o(1)) (x.,,,,(((()))) (N))
approach described in Section 2.2, by ensuring that μ4 is deterministically (Sx x which implies that the reverse process is stochastically equivariant also.
Following standard terminology in the literature (Garcia Satorras et al. 2021b Xu et al. 2022 Hoogeboom al. 2022), we refer to (SN + diffusion defindd on the CoM-free space LI as an E(3 eqquvariantt diffusion.
3 EQUIVARIANT DIFFUSION VIA STOCHASTIC SYMMETRISATION
In this section, we introduce SYMDIFF and apply it to the problem of obtaining equiivaraa diffusion models for N-body systems. We also discuss extensions score and flow-based generative models (Song et al. 2020 Lipman et al. 2022) in Appendix E
3.1 STOCHASTIC SYMMETRISATION -
Recently, Cornish 2024) gave general theory of neural network symmetrisation the framework of Markov categories (Fritz 2020). encompassing earlier approaches to symmetrisation based on
averaging or canonicalisation (Murphy al. 2018 Puny et al.] 202T] Kaba et al., [023 Kim al.]
2023). This theory applies flexibly and compositionally to general groups and actions, including
in the non-compact case, and extends to provide a methodology for symmetrising Markov kernels, which had not previously been considered.
In this work, we will make use of a special case of Example 6.3 Cornish (0224), which we state now before providing intuition. We denote by H x G the direct product of groups G and H. Recall that an action of H. x G on a space X' induces actions of both H and G on X' also. For example. H acts via h = (h, xx where eg is the identity element of G. We will also say that Markov kernel 7 X + 9 is an equivariant base case if it is both H-invariant and G-equivariant, where G acts on the output space of by left multiplication, i.e. g' g = g g. We then have the following result (see Appendix A. for a proof).
acts on X' and y. and 7 X G is an equivariant base case. Then every Markov kernel k : X + Y that is equivariant with respect to the induced action of H gives rise to a Markov kernel sym, X + y that is equivariant with respect to H. x G, where
Theorem 1. Suppose H
sym, (k)(dd)) xx) may be sampled from as follows:
g <dg' x)
y k(dylg 1g- x)
return 9 y-
4

---

Published as a conference paper at ICLR 2025
Intuitively, this result allows us start with Markov kernel that is equivariant only with respect to H. and then "upgrade" it to become equivariant with respect to both H. and G. As a special case, if H is the trivial group, then every Markov kernel k X + y is H-equivariant. Moreover, in this case H. x 9 = 9, and so Theorem 1 gives a procedure for obtaining G-equivariant Markov kernels from ones that are completely unconstrained. However, as our N-body example will illustrate, it can often be useful to symmetrise a Markov kernel that is already "partially equivariant", which motivates keeping H general here.
Beyond the existence of an equivariant base case, Theorem is completely generic and requires no assumptions on the groups and actions involved. As explained in Section 5.1 of Cfroiishh (2024), this is also the only natural procedure that can be defined in this way without further assumptions.
Recursive symmetrisation The symmetrisation procedure defined by Theorem | requires 7 al- ready to satisfy two equivariance constraints. In effect, this pushes back the problem of obtaining (H -)-equvvariant Markov kernels to the choice of % which mirrors the situation in the deter- ministic setting also (Puny et al. 2021: Kim et al. 2023). Whenever G is compact, Example 6.3 of Comish (2024) gives a suitable choice as (dg)) x) == X(dg). where À denotes the Haar measure on G (Kallenberg 1997). Other choices could also be made here on case-by-case basis, such as using intrinsically equivariant neural networks desired. To obtain greater modelling flexibility, Cornish
(024) also proposes recursive approach to obtaining 7. Specifically, the idea is to set
(7) where "oo 71 X - G are Markov kernels, and "ru is an equivariant base case (e.g. the Haar measure), but where now Yi is only required be H-invariant, and may behave arbitrarily with respect to G. We note this recursive approach exploits the stochastic nature of the procedure in Theorem and would not be possible using deterministic symmetrisation methods here instead.
Y := sym (71)
3.2 SYMDIFF: SYMMETRISED DIFFUSION
We propose to use stochastic symmetrisation to obtain a diffusion process as in Section 2.2 whose reverse kernels are stochastically equivariant. Specifically, suppose some product group H. x g acts on our state space Z. For each timestep t E (1....,7) T, we will choose some H-equivariant Markov kernel Ka Z - Z that admits a conditional density kg (4z-1/z)). Similarly, we will choose some Markov kernel "io Z + & that satisfies the conditions of Theorem when X' = Z. With these components, we will define our equivariant reverse process to be
pa (e(2t-1 == sym.,, (Z, (Z)))
which is guaranteed to be (H. x -)-equivariant by Theorem [ This defines conditional density,
not just a Markov kernel, as a consequence of the next result. For the proof, see Appendix A.2 Proposition 1. Assume the same setup as Theorem ! and for each fixed 9 E G, let k(dylg. x) be the distribution of the following generative process:
y ~ k(dylg-1 -1 return y. If k(dy))) has a density k(y)x)), then k(dylg g. has density k(ylg, (g. x), and (ym,,(4))dylbz has sym,(k)) (y') = Byd(ggb>) (olps))kk(lly [k(ylg.x))
as a density. the action on y has unit Jacobian, then we may write k(ylg. = k(g -ygg xx)
Training objective We would like learn the parameters 0 using the ELBO from equation 4
However, in general, we do not have access to the densities Pa(2,-1|2, from equation 8 in closed form, since this requires computing the expectation as in Proposition [ As such, we cannot compute the ELBO directly. However, since log is concave, Jensen's 's inequality allows us to bound
log po (Z (2:) 2 (dy((do))))) hoo
ko )-
Since the ELBO in equation 4 depends linearly on log
this allows us to also bound
log (o())) 2 E yzzz (odg)) [log Ko (2olg.222) -Do
E
(oo D 1,, Zu l 12. 2.)).
(9)
5

---

Published as a conference paper at ICLR 2025
We take this new bound as our objective used to train our SYMDIFF model. By a similar argument as in Remark 7.1 of [ornish (2024), if the model is sufficiently expressive, then optimising this new bound is equivalent to optimising the original ELBO (see Appendix D).
where the right-hand side is tractable lower bound to the original ELBO in equation 4
Comparison with deterministic symmetrisation An alternative approach to equation g would be to obtain μs in equation | via deterministic symmetrisation. However, many deterministic techniques (Puny et al. [202]] Kim et al., (222)) involve a Monte Carlo averaging step that requires multiple passes through the model instead, and then are only approximate, which would introduce
sampling bias here. In contrast, to sample from our method requires only single pass through na and ke, and involves no bias. The canonicalisation method of Kaba et al. (023) also avoids this averaging step, but instead suffers from pathologies associated with its analogue of the equivariant base case No. which Dym et al.] (222)) show must become discontinuous certain inputs. Addi- tionally, canonicalisation requires an intrinsically G-equivariant architecture for its analogue of yo- In contrast, whenever G compact, our stochastic approach allows "o to be obtained using the Haar measure in a way that does not suffer these pathologies, and may be implemented without any intrinsically G-equivariant neural network components at all, as we show concretely next.
3.3 SYMDIFF FOR N-BODY SYSTEMS
We now apply SYMDIFF in the setting of N -N-body systems considered in Section 2.3 Specifically take 2 = U, H Sv. and This means that we start with (o(2,-) | in equation that is already equivariant with respect reorderings of the N bodies, and then symmetrise this obtain (SN x ((3))-equivariant reverse kernel overall. We choose to symmetrise in this way because highly scalable Sx-equivariant kernels based on Transformer architectures can be readily constructed for this purpose (Vaswani
0(3), and consider the action 2 defined equation 6
et al. 2017 Lee et al. 2019 Peebles & Xie (023). whereas intrinsically O(3) equivvrrra neural networks have not shown the same degree of scalability to-date (Abramson et al. 2024).
Choice of unsymmetrised kernels It remains now choose (0, (z)). We now do so in way that equation 9 will resemble the standard diffusion objective in Ho et al. (2020), allowing for the scalable training of SYMDIFF. Specifically, we take
ko (o = 1: (1))),
(10)
where μs 2
1 is arbitrary S Sw-equivariant neural network, which in turn means ho is stochastically SN Sy-equivariant. We highlight that μ4 is otherwise unconstrained and can process x, jointyyy. In contrast, previous work using intrinsically (S,, -(3))-euivvarian components Satorras et al. 2021: Thölke & De Fabritiis
2022] Hua et al. 2024) has required complex parameterisations for μs that handle the and h inputs separately.
Form of KL terms We now show how our model yields a closed-form expression for the L, terms in equation Standard arguments show that each g(2,-1 |z, zo) (projected) Gaussian. We claim that our model also gives (projected) Gaussian
2o), og is
Ag (2z-1/ z;) = (u R Ho 2)),
(11)
Indeed, by the definition of this kernel in Proposition I and the definition of Ny in Section 2.3 and
since R Rd for R 0(3), for 2 Ka (2 R, 2,) have
d R. (R 2)) + +0, = R no po (RT 2:) + (.)R.
2,
where Ner (0, I). Since R c for R 0(3), equation 11 now follows. The same argument
as Hoogeboom et al. (2022) now yields the closed-form expression
1
|| Zu) - R Ro(" 2:) ):
L. Ez (ar, (o) m.)
(12)
(t)
We leave the dependence on implicit in notation throughout.

---

Published as a conference paper at ICLR 2025
We obtain unbiased gradients of this quantity whenever (d( is reparametrissbbe (inggaa 2013). In other words, we should define We(dR z to the distribution of (p (p E), where Pp deterministic neural network, and & is some noise variable whose distribution does not depend on 0. For a discussion of how we can handle the L1 term in equation in our framework, refer to Appendix C.T
c-parameterrsation
When po is taken have the e-form (Ho al.
2020 Kingma al.
2021
- Eg (zr),
ng (Zz) Ohgle-1
Ze
(13)
Odjt-10t
for some neural network Z + Z, the same argument given by Ho et al. (020 now allows us rewrite equation 12 in way that resembles the standard diffusion objective:
(t)l|= R co( =)]l"]
L. (G,
(14)
where 2, 0g ore, and w(t) = (1 SNR(L 1)/SNR(()). Recall we require μıg to be Sx-equivariant, which straightforwardly follows here whenever eg is. In practice, we set w((() during training as is commonly done in the diffusion literature (Kingma & Gao 2024h.
Recursive choice of 7 To apply Theorem D we require an equivariant base case that is Sy-
invariant and 0(3) equiivaraat To obtain this, we apply the recursive procedure from equation 7 where "/o Z - O(3) is obtained using the Haar measure on O(3) as described there, and (dR z
is defined as the distribution of fo(2,, 7), where is sampled from some noise distribution (dg) and no(:;n) Z 0(3) is Sx-invariant neural network for each fixed value of v). Both the Haar measure and v do not depend on 0, and so the overall reparametrisable by construction. We emphasise that fe is not required be ((3)-cquivariant in any sense, thus allowing for highly flexible choices such as Set Transformers (Lee et al. (019). At sampling time, we use the procedure from Section 5 of Mezzadri] (2006) sample from the Haar measure on 0(3), which is negligible overhead compared with the cost of evaluating fo.
obtained in this way is always
Algorithm 1 SYMDIFF training step
1: Sample Zg Paam(Z)), t Unif([1...., T)) and (0,
2: Ze +- 0gZo + are
3: Sample Rg from the Haar measure on O(3) and r) v(d1)
4: R + Ro fo (R] n)
5: Take gradient descent step with Velu(t)]le- R co cR RTT z=)l
Algorithm SYMDIFF sampling process
1: Sample ZT ~ I0,
2: for T, .....22 do
3: Sample Ro from the Haar measure on 0(3), n ~ v(dn), and NN((0,1)
4: + Ro fo Zz. T))
5: 2 a.11 2,
aie-1 10c
R CR(RR 2.) + (t)()
6: return 2g (2 (2))
See Appendix
C.1 for example of this output kernel
3.4 DATA AUGMENTATION SPECIAL CASE OF SYMDIFF
Data augmentation is popular method for incorporating "soft" inductive biases within neural networks. Consider again our setup from the previous section, but now using the unsymmetrised kernels Kp (2 (z,) from equation 10 in place of pu(zz-1/2 (2)) in backwards process. Suppose this model is trained using the standard diffusion objective, applying a uniform random orthogonal transformation to the input of eg before each forward pass. This is equivalent to optimising the

---

Published as a conference paper at ICLR 2025
following objective:
1
"[)t]le- -c zo oic)1
(15)
where À is the Haar measure on (3).. (More general choices of À could also be considered.) We then have the following result, proven in Appendix A.3 A..3
Proposition 2. When 7o (z( X(dIt) for all zt E Z, our SYMDIFF objective in equation 14 recovers the data augmentation objective exactly, so that =
As result, SYMDIFF may be understood strict generalisation of data augmentation in which a more flexible augmentation process To is learned during training. Moreover, our " is then also deployed sampling time in way that guarantees stochastic equivariance. In contrast, data augmentation is usually only applied during training, and the model deployed at sampling time then becomes only approximately equivariant.
EXPERIMENTS
We evaluated our ((3)-equivarian SYMDIFF model as a drop-in replacement for the E(3- equivariant diffusion (EDM) of [Hoogeboom et al. (2022) both the QM9 and GEOM-Drugs datasets for molecular generation. We implemented this within the official codebase of Hoogeboom
al. (2022 substituting our symmetrised reverse process for their one. We made minimal other changes to their code and experimental setup otherwise, and performed minimal tuning of our architecture. We found SYMDIFF led to significantly improved performance
both tasks. Our results were also on par or better than GcoLDM (Xu et al. 2023). END (Cornet al, 2024), and MUDiff (Hua et al. 2024). which bake in more sophisticated inductive biases related to molecular generation than EDM does. Our models were also more computationally efficient than these baselines, which all rely on intrinsically equivariant subcomponents. We give an overview of our setup now and provide full details in Appendix C Our code is available at: https:///áthub. wttptuuu..om/l/ooonnoog SaaraannMM//S/DiDiF//
4.1 QM9
Dataset QM9 (Ramakrishnan et al. 2014) is a common benchmark dataset used for evaluating molecular generation. It consists of molecular properties and atom coordinates for 130k small molecules with up to 9 heavy atoms and total 29 atoms including hydrogen. For our experi- ments, we trained our SYMDIFF method to generate molecules with 3D coordinates, atom types, and atom charges where we explicitly modeled hydrogen atoms. We used the same train-val-ter val-test split of 100K-8K-1KK as in Andersson et al. (2019).
Our model We took our core model (SymDiff) to be sym (k.) from equation 8 above, with kn and "o given as specified in Section 3.3 We chose the "backbone" neural network eg to be a Diffussinn Transformer (DiT) (Peebles & Xie 2023). which is Sw-cquivariant by construction. Likewise, we chose the "backbone" neural network fo the component " to be a DIT. We made this component S,,-invariant using Set Transformer (Lee et al. 2019) approach, thereby achieving the requirements described in Section 3.3 Our <g had 29M parameters, matching the smallest model considered by Peebles & Xie (2023), while our fo had 2.2M parameters. In this way, our Yo was much smaller than our kg, following a similar approach taken by earlier work on deterministic model had 31.2M parameters
symmetrisation (Kim al. 20023] Kaba al. 2023). Overall, in total. To test its scalability. we also trained a larger version of our method with a backbone ka having 115.6M parameters (SymDiff*), which matched the DiT-B model from Peebles & Xie 2023). We trained all our models for 4350 epochs to match the same number of gradient steps as Hoogeboom et al. (2022). For further details about our architecture, see Appendix [
Metrics To measure the quality of generated molecules, we follow standard practice (Hoogeboom ct al. 2022 Garcia Satorras et al. 2021a) and report atom stability, molecular stability, validity
8

---

Published as a conference paper at ICLR 2025
and uniqueness. We exclude results for the novelty metric for the same reasons discussed in Vignac & Frossard (202I) and refer the reader to these works for a more extensive discussion of these metries. For all of our metrics,
used 10,000 samples and report the mean and standard deviation over three evaluation runs. To demonstrate the efficiency of our approach, we also report the number of seconds per epoch, time taken to generate sample and vRAM.
Baselines As a baseline, we trained the 5.3M parameter EDM model using the original exper- imental setup of [Hoogeboom et al.] (2022). We also trained unsymmetrised reverse process with a 29M parameter DiT backbone (DiT), as well as the same model using data augmentation described in Section 3.4 (DiT-Aug)- Additionally, we trained SYMDIFF model with the same backbone ke as our SymDiff model above, but whose " was obtained using the Haar measure on 0()), rather than learning this component (SymDiff-H).
Results From Table [ we see that our SYMDIFF models comfortably outperformed EDM on all metrics, bar uniqueness. Additionally, our model was also competitive with the more recent, sophisticated baselines from the literature, outperforming all of them on validity. We attribute the improved performance of our method to the extra architectural flexibility provided by our approach to symmetrisation. Our largest model, SymDiff", outperformed all our baselines on atom stability and validity, and is within variance for molecular stability. We conjecture that similar performance improvements could be achieved by using our SYMDIFF approach reverse kernels in more sophisticated methods. Table [ also shows that DiT-Aug performed notably better than the DiT model on all metrics, highlighting its strength as a baseline. Despite this, our SymDiff model outperformed both SymDiff-H H and DiT-Aug on all metrics apart from uniqueness. This shows that approach has benefits that extend beyond merely performing data augmentation.
drop-in replacement for the
Computational efficiency Importantly, as Table2 our method was also more computation- ally efficient than the alternative methods considered, both in terms of seconds/epoch, sampling time, and vRAM. This is not surprising since these alternative models rely on intrinsically equiv- ariant graph neural networks that use message passing during training and inference, which is com- putationally very costly. In contrast, our symmetrisation approach allows us to use computationally efficient DiT components that parallelise and scale much more effectively (Fei al. 2024).
Table I: Test NLL, atom stability, molecular stability, validity and uniqueness on QM9 for 10,000 samples and 3 evaluation runs. We omit the results for NLL where not available.
Method
NLL Atm. stability (
Mol. stability (%)
Val. (%6
Uniq. (%) f
GeoL.MM
10.10
89.40 20.50
93.80
92.70 20.50
MUDiff
135,5002.10
98.80=
89.90=1.10
95.30=1.60
99.10=0.50
END
98.90 10.00
89.10 10.10
94,80 10.10
92.60 #:020
EDM
-100.70=1.50
98.70=n.10
82.00
01.90=0.50
90.70=0.00
SymDiff
-133.79.1.33
98.92
89.65
96.3620.27
07.66::0.22
SymDiff
-129.35 =1.07
98.74an.0a
87.49 20.20
95.75 #0.10
97.89 #0.26
SymDiff-H
-126.53=0.00
98.57=n.07
05.51
95.22#0.18
97.98
DiT-Aug
-126.81=1.60
98. 08.644.0.0a
85.85
15.10=0.17
97.98=0.008
DIT
-127.78-249
98.23-a0.04
81.03
94.71
97.98
Data
99.00
95.20
97.8
100
Table 2: Seconds per epoch, sampling time and vRAM for SymDiff and our baselines on QM9. Results for END are omitted as their code was not publicly available.
Method
Parameters
See. Sec./epoch (6)
Sampling time (s)
vRAM (GB)
GeoL.DM
11.M
210.93
0.26
27
MuDiff
9.7M
230.87
0.89
END
9.4M
EDM
5.4M
88.80
0.27
14
SymDiff
117.8 8M
53.40
0.21
16
SymDiff
31.2M
27.20
0.09
Ablations As an ablation study, we also tested the effect of making SYMDIFF smaller, and EDM larger. For SymDiff, we trained two models of 23.5M (SymDiff-) and 13.5M (SymDiff-~)

---

Published as a conference paper at ICLR 2025
parameters respectively. For EDM, trained two additional models with 9.5M (EDM*) and 12.4M (EDM**) parameters respectively. For full details see Appendix C.2.1 From Table that even smaller SymDiff models remained competitive. In particular, SymDiff gave comparable molecular stability as the second largest HDM model, EDM while being approximately times faster in terms of seconds/epoch.
Table 3: NLL, molecular stability, seconds per epoch, sampling time and vRAM for different sizes of SymDiff and EDM on QM9. For additional performance metrics see Appendix
Method
NLL
Mol. stability (%)
Sec./epoch (8)
Sampling time (8)
vRAM (GB)
EDM++
-19.112
85.68
160.60
0.56
23
EDM
-110.97=1.42
84.6349.16
192.60
0.46
23
EDM
-110.70=1.50
82.00 a:
88.80
0.27
14
SymDiff
-129.35=1.07
27.20
0.09
SymDiff
-125 02
83.51
24.87
0.08
SymDin
-110.68 29.55
71.25 19.50
20.60
0.07
4.2 GEOM-DRUGS
Dataset, model and training GEOM-Drugs (Axelrod & Gomez Bombarell 2022) is a larger and more complicated dataset than QM9, containing 430,000 molecules with up 181 atoms. We processed the dataset in the same way as Hoogeboom et al. (2022), where we again model hydrogen explicitly. We used the SymDiff model from earlier, which we trained for 55 epochs to match the same number of gradient steps as Hoogeboom et al. (2022).
Metrics and baselines We report the same metrics for QM9 but exclude molecular stability and uniqueness for the same reasons discussed Hoogeboom et al. (2022). We compared our method to the EDM model used by HHogeboom et al. (2022) for GEOM-Drugs, well as the other baseline architectures reported for QM9.
Results From Table we again see that our approach comfortably outperformed its EDM counterpart. It is also again competitive with the more sophisticated baselines, whose reported results we restate here. Like with QM9, our SymDiff models were significantly less costly in terms of compute time and memory usage compared with EDM (see Appendix CC33 In fact, when tried to run the EDM model resulted in out-of-memory errors on our NVIDIA H100 80GB GPU Hoogeboom al. (2022) avoid this by training EDM on x NVIDIA RTX A6000 48GB GPUs.)
Table 4: Test NLL,, atom stability and validity on GEOM- -rugs for 10,000 samples and 3 evaluation GeoLDM and EDM ran their results for just evaluation run. We omit the results for NL.L and validity where not available.
Method
NLL J Atm, stability (5)
Val. (%) 1
GeoLDM
84.4
99.3
END
87.8
92.920.3
BDM
137.1
81.3
SymDiff
-301.21
86.160
99.27
Data
86.50
99.9
5 CONCLUSION
We have introduced SYMDIFF: a lightweight, and scalable framework for constructing equivariant diffusion models based on stochastic symmetrisation. We applied this approach to E(3)- equivariance for --body data, obtaining overall model that is stochastically equivariant but that does not rely on any intrinsically equivariant neural network subcomponents. Our approach leads significantly greater modelling flexibility, which allows leveraging powerful off-the-shelf architectures such as Transformers (Vaswani et al.] (017). We showed empirically that this leads overall to improved performance several relevant benchmarks.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
The authors are grateful Tom Rainforth, Emile Mathieu, Saifuddin Syed and Ahmed Elhag for helpful discussions. LZ and KA are supported by the EPSRC CDT in Modern Statistics and Statis- tical Machine Learning (EP/S023151/1).
REFERENCES
Josh Abramson, Jonas Adler, Jack Dunger, Richard Evans, Tim Green, Alexander Pritzel, Olaf Ronneberger, Lindsay Willmore, Andrew J Ballard, Joshua Bambrick, et al. Accurate structure prediction of biomolecular interactions with alphafold 3. Nature, pp. 1-3, 2024.
Brandon Anderson, Truong Son Hy, and Risi Kondor. Cormorant: Covariant molecular neural
networks. Advances in neural information processing systems, 32, 2019.
Brian DO Anderson. Reverse-time diffusion equation models. Stochastic Processes and their Ap- plications, 12(3):313-326, 1982.
Simon Axelrod and Rafael for property prediction and molecular generation. Scientific Data, 9(1): 2022.
Geom, energy-annotated molecular conformations
Simon Batzner, Albert Musaelian, Lixin Sun, Mario Geiger, Jonathan P Mailoa, Mordechai Ko- rnbluhh, Nicola Molinari, Tess Smidt, and Boris Kozinsky. E (3)-equivrriant graph neural networks for data-efficient and accurate interatomic potentials. Nature communications, 13(1): 2453, 2022.
Benjamin Bloem-Reddy and Yee Whye Teh. Probabilistic symmetries and invariant neural networks. Journal of Machine Learning Research, 21(90):1-61, 2020.
Michael Bronstein, Joan Bruna, Taco Cohen, and Petar Veličković. Geometric deep learning:
Grids, groups, graphs, geodesics, and gauges. arXiv preprint arXiv:2104. 13478, 2021.
Diego Canez, Nesta Midavaine, and Thijs Stessen. Effect of equivariance training dynamics.
https: //gramm///mmmmm http:///ggaambbboooooooo www.boll..ppp...oooooooo........ htbl:hh22//bbl/222
rel axed-equi variance/
2024. Accessed: 2024-09-30.
Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David Duvenaud. Neural ordinary differential equations. Advances in neural information processing systems, 31, 2018.
François RJ Cornet, Grigory Bartosh, Mikkel N Schmidt, and Christian A Naesseth. Equivariant neural diffusion for molecule generation. In JCML 2024 Al for Science Workshop, 2024.
Rob Cornish. Stochastic Neural Network Symmetrisation in Markov Categories. arXiv preprint wrXiv:2406.118/4 2024.
Alexandre Duval, Simon Mathis, Chaitanya K Joshi, Victor Schmidt, Santiago Miret, Fragkiskos D Malliaros, Taco Cohen, Pietro Lio, Yoshua Bengio, and Michael Bronstein. A hitchhiker's guide to geometric gnns for 3d atomic systems. arXiv preprint arXiv:2312.05111, 2023a.
Alexandre Agm Duval, Victor Schmidt, Alex Hernández-Garcan Santiago Miret, Fragkiskos D Malliaros, Yoshua Bengio, and David Rolnick. Faenet: Frame averaging equivariant gnn for materials modeling. In International Conference on Machine Learming, pp. 9013-9033, PMLR, 2023b.
Nadav Dym, Hannah Lawrence, and Jonathan W Siegel. Equivariant frames and the impossibility of continuous canonicalization. arXiv preprint 160.777777 2024.
Bryn Elesedy and Sheheryar Zaidi. Provably strict generalisation benefit for equivariant models. In International conference on machine learning, pp. 2959-2966 PMLR, 2021.
Zhengcong Fei, Mingyuan Fan, Changqian Yu, Debang Li, and Junshi Huang. Scaling diffusion transformers to 16 billion parameters. arXiv preprint arXiv:2407. 1633, 2024.

---

Published as a conference paper at ICLR 2025
Tobias Fritz. A synthetic approach markov kernels, conditional independence and theorems sufficient statistics. Advances in Mathematics, 370:107239, 2020.
Victor Garcia Satorras, Emiel Hoogeboom, Fabian Fuchs, Ingmar Posner, and Max Welling, E(n) equivariant normalizing flows. Advances in Neural Information Processing Systems, 34:4181-
2021a.
Victor Garcia Satorras, Emiel Hoogeboom, Fabian Fuchs, Ingmar Posner, and Max Welling. E(n) equivariant normalizing flows. In M. Ranzato, A. Beygelzimer, Y. Dauphin. P.S. Liang, and Wortman Vaughan (eds.), Advances in Neural In- formation Processing Systems, volume 34, Pp. 4181-4192. Curran Associates, Inc.,
b02 URL
neurips wttcc///appa cc/paper. fi stllspppperr/222 file file/210 /eee2115666 0d (00 aaa.eaa-aaeaaaaaaaa--aa-eeeeee
Yoav Gelberg, Tycho FA der Ouderaa, Mark der Wilk, and Yarin Gal. Variational inference failures under model symmetries: Permutation invariant posteriors for bayesian neural networks. In ICML 2024 Workshop on chometry-groundee Representation Learning and Generative Mod- eling, 2024.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840-6851, 2020.
Emiel Hoogeboom, Victor Garcia Satorras, Clément Vignac, and Max Welling. Equivariant diffu- sion for molecule generation in 3d. In International conference on machine learning, pp. 8867- 8887. PMLR, 2022.
Chenqing Hua, Sitao Luan, Minkai Xu, Zhitao Ying. Jie Fu, Stefano Ermon, and Doina Precup. Mudiff: Unified diffusion for complete molecule generation. In Learning on Graphs Conference, pp. 33-1. PMLR, 2024.
G6kou-Oumar Kaba, Arnab Kumar Mondal, Yan Zhang, Yoshua Bengio, and Siamak Ravanbakhsh. Equivariance with learned canonicalization functions. In International Conference Machine Learning, pp. 15546-15566. PMLR, 2023.
Olav Kallenberg. Foundations of modern probability, volume 2. Springer, 1997.
Olav Kallenberg. Foundations of Modern Probubility. Springer, 2 edition, 2002.
Jinwoo Kim, Dat Nguyen, Ayhan Suleymanzade, Hyeokjun An, and Seunghoon Hong. Learning probabilistic symmetrization for architecture agnostic equivariance. Advances in Neural Informa- tion Processing Systems, 36: 18882-1861222 2023.
Diederik Kingma and Ruiqi Gao. Understanding diffusion objectives as the elbo with simple data augmentation. Advances in Neural Information Processing Systems, 36, 2024.
Diederik Kingma, Tim Salimans, Ben Poole, and Jonathan Ho. Variational diffusion models. Ad- vances neural information processing systems, 34:21696-21707, 2021.
Diederik P Kingma. Auto-cncoding variational bayes. arXiv preprint arXiv:1312.6114 2013.
Leon Klein, Andreas Krämer, and Frank Noé. Equivariant flow matching. Advances in Neural Information Processing Systems, 36, 2024.
Tuan Le, Julian Cremer, Frank Noé, Djork-Amé Clevert, and Kristof Schütt. Navigating the design space of equivariant diffusion- based generative models for novo 3d molecule generation. arXiv preprint arXiv:2309. 2tt.::2999.7766... 17296, 2023.
Juho Lee, Yoonho Lee, Jungtaek Kim, Adam Kosiorek, Seungjin Choi, and Yee Whye Teh. Set transformer: A framework for iitention-basee cermutation-invarian neural networks. In Interna- tional conference on machine learning, pp. 3744-3753. PMLR, 2019.
Yaron Lipman, Ricky TQ Chen, Heli Ben- Gen-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. arXiv preprint 2rXiv:2210.02747, 2022.
12

---

Published as a conference paper at ICLR 2025
Iyy Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Confer- ence on Learning Representations.
Shengjie Luo, Tianlang Chen, Yixian Xu, Shuxin Zheng, Tie-Yan Liu, Liwei Wang, and Di He. One transformer can understand both 2d & 3d molecular data. In The Eleventh International Conference on Learning Representations, 2022.
Francesco Mezzadri. How to generate random matrices from the classical compact groups. arXiv preprint math-ph/0609050, 2006.
Arnab Kumar Mondal, Siba Smarak Panigrahi, Oumar Kaba, Sai Rajeswar Mudumba, and Siamak Ravanbakhsh. Equivariant adaptation of large pretrained models. Advances in Neural Information Processing Systems, 36:50293-50399, 2023.
Ryan L Murphy, Balasubramaniam Srinivasan, Vinayak Rao, and Bruno Ribeiro. Janossy pool- ing: Learning deep permutation-invarian functions for variable-size inputs. arXiv preprint /tt::///1.0000 2018.
William Peebles and Saining Xie. Scalable diffusion models with transformers. In Proceedings of the JEEE/CVF International Conference on Computer Vision, pp. 4195-4206 2023.
Stefanos Pertigkiozoglou, Evangelos Chatzipantazis, Shubhendu Trivedi, and Kostas Daniilidss. Im- proving equivariant model training via constraint relaxation. arXiv preprint arXiv:2408. 13242, 2024.
Gabriel Peyré, Marco Cuturi, et al. Computational optimal transport: With applications to data science. Foundations and Trends@ Machine Learning, 115-66)355-607, 2019.
Omri Puny, Matan Atzmon, Heli Ben-Hamu, Ishan Misra, Aditya Grover, Edward J Smith, and Yaron Lipman. Frame averaging for invariant and equivariant network design. arXiv preprint 10.00:21 2021.
Raghunathan Ramakrishnan, Pavlo O Dral, Matthias Rupp, and O Anatole Von Lilienfeld. Quantum chemistry structures and properties of 134 kilo molecules. Scientific data, 1(1):177,2) 2014.
Victor Garcia Satorras, Emiel Hoogeboom, and Max Welling. E (n) equivariant graph neural net- works. In International conference on machine learning, pp. 9323 93322993.222 PMLR, 2021.
Selinger. A Survey of Graphical Languages for Monoidal Categories, pp. 289-355. Springer doi: 10. 111100779888---22-1882-1-299 URL
Berlin Heidelberg, 2010. ISBN 9783642128219.
http://dee doi www.orgg/1.11 10..000...09990 3--2------------21-------
Noam Shazeer. Glu variants improve transformer. arXiv preprint arXiv:2002.05202, 2020.
Jascha -ock-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamies. In International conference machine learn- ing, pp. 2256-2265. PMLR, 2015.
Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. basee generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.13456. 2020.
Philipp Thölke and Gianni De Fabritiis. Torchmd-net: equivariant transformers for neural network based molecular potentials. arXiv preprins arXiv:2202.02544/ 2022.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in Neural Informa- tion Processing Systems, 2017.
Clement Vignac and Pascal Frossard. Top-n: Equivariant set and graph generation without ex- changeability, arXiv preprint arXiv:2110.020%6 2021.
Pascal Vincent. connection between score matching and denoising autoencoders. Neural compu- fation, 23(7):1661-1674 2011.
13

---

Published as a conference paper at ICLR 2025
Yuyang Wang, Ahmed AA Elhag, Navdeep Jaitly, Joshua M Susskind, and Miguel Ángel Bautista. Swallowing the bitter pill: Simplified scalable conformer generation. Forty-firr International Conference Machine Learning, 2024.
Joseph L Watson, David Juergens, Nathaniel R Bennett, Brian L Trippe, Jason Yim, Helen E Eise- nach, Woody Ahern, Andrew Borst, Robert J Ragotte, Lukas F Milles, et al. De novo design of protein structure and function with rfdiffusion. Nature, 620(7976):1089-11
2023.
Minkai Xu, Lantao Yu, Yang Song. Chence Shi, Stefano Ermon, and Jian Tang. Geodiff: geo- metric diffusion model for molecular conformation generation. arXiv preprint arXiv:2203.02923, 2022.
Minkai Xu, Alexander S Powers, Ron O Dror, Stefano Ermon, and Jure Leskovec. Geometric latent diffusion models for 3d molecule generation. In International Conference Machine Learning, pp. 38592-38610. PMLR, 2023.
Dmitry Yarotsky. Universal approximations of invariant maps by neural networks. Constructive Approximation, 55(1):407-4744 2022.
Jason Yim, Brian I. Trippe, Valentin Bortoli, Emile Mathieu, Arnaud Doucet, Regina Barzilay, and Tommi Jaakkola. Sc(3) diffusson model with application to protein backbone generation. arXiv preprint arX1v:230..02277,200 2023.
14

---

Published as a conference paper at ICLR 2025
A PROOFS
A.1 PROOF OF THEOREM 1
Proof. Theorem is special case of Example 6.3 of Cornnihh (2024), whose notation and setup will import freely here. Note that Example 6.3 makes of string diagrams (Selinger) 2010), an introduction to which can be found in Section Cornish (2024). Intuitively, a string diagram represents a (possibly stochastic) computational processes that should be read up the page, with the inputs applied at the bottom, and outputs produced at the top.
To proceed, we first let the action p be trivial, that the semidirect product N x H becomes simply the direct product N x H by Remark 3.29 of Cornish] (2024). Now consider the following string diagram:
Y
Oy, H
k:
X
X
H
(16)
H
X
The upshot of Example 6.3 Cornish (2024) is that equation 16 is always equivariant with respect to the (N x H Haaactiinss x and ay whenever the following conditions all hold:
: X + H is equivariant to the H-actions ox, and *
Y X > H is invariant to the N-action ax,x;
k X + Y is equivariant with respect the N-actions ax,N and ay,N-
Here, in Example 6.3 of Cornish 2024):
We have decomposed the (N action axx H-action ax. followed by N-
action x using Remark 3.31 Cornish] 2024). In other words:
X
X
OX.N
QXX
OX.HH
N
X
N
H
X
We have decomposed ay similarly;
We denote by * and (-) the multiplication and inversion operations of the group H.
15

---

Published as a conference paper at ICLR 2025
To obtain our Theorem [ now simply instantiate things appropriately in the Markov category == Stoch, so that all the boxes appearing in equation [6 become Markov kernels. Concretely, we take X X, Y ), N H. and H G. A procedure for sampling from the symmetrised
Markov kernel equation 16 may then be read off as follows:
g g~g(gg(x))
y k(dy'g x)
return g Y
exactly as in the statement of this result.
A.2 PROOF OF PROPOSITION
Proof. For simplicity, we assume k(dy)s (x) has a density k x) with respect to the Lebesgue mea- sure j on y' R". To derive the density of k(dylg, where g E g is some fixed group element, we note that for ~ k(dy Ig. x), we have y = I Yo where Yo ~ k x). Hence, by the change-of-variables formula,
conclude that the density of k-y [g, x) exists and has the form:
0 1 y) dy
k(ygg x) k(g .ygg
Hence, that when the action of has unit Jacobian, the density of x(ylg. k(g ylg x)..
Further, suppose have ((dg (x) and y k(y g, x) i... sym, It the case that for an arbitrary (Borel) measurable set in V, we have
sym, (x) E A(x) 1 k(y Alg, (dggxx).
Since have shown above that k(dylg. x) has density, express this
sk),(()(y E Alx) 1. ( kly x) u(dx) 7(dglx)
(x(dg/)) [k(y x)] p(dy).
where we use Fubini's theorem for the second line as all quantities are negattgative. Hence, we can conclude that the density of sym, (k) exists and has the form sym, (k) = Eyy(gg xy)lg, O
A.3 PROOF OF PROPOSITION 2
Proof. The standard diffusion objective with data augmentation distributed according to the Haar measure À is given by
C
E (o.
wx)) E0 zo + t,)e)
(17)
w(t) c- co(R (- (a:z0 + R, =])
(18)
(G
2 u()) é - Eo(R (a.zo + a()
(19)
(e ((R))
w(t)]lc- R co (a,20 + a.d)
(20)
2
where we use the fact that e' Rd e is distributed according to (0, as R, are independent in the expectation, and that the action of E O(3) preserves the L2 norm. To conclude, we note that it is standard result that if R X, the inverse R is also distributed according to the Haar measure X. Hence, we see that L caug coincides with
H =u(Dill== R =o( z.)l?
C. =
(21)
where 0,20 Otc.
16

---

Published as a conference paper at ICLR 2025
B MODEL ARCHITECTURE
Below, outline the architectures used for 69 and 70- Both components rely on Diffusion Transformers (DiTs) (Peebles & Xie 2023) using the official Py PyTorch implementation at https: We also state the hyperparameters that we kept
//githab. com/ waw.bookreeearcchceeoooo. fixed for both our QM9 and GEog-Drags experiments. Any hyperparameters that differed between the datasets are discussed in their respective sections later in the Appendix.
We emphasise that our architecture choices were not extensively tuned as the main purpose of our experiments was to show that we can use generic architectures for equivariant diffusion models. We arrived at the below architecture through small adjustments from experimenting with DiT models in the context of molecular generation, which we stuck with for our final experiments.
B. ARCHITECTURE OF EB
As we need Ep to be Sx-equivariant architecture, we parametrise this in terms of DiT model which consists of Rlayers (intermediate) layers, Racad attention heads, hidden size Raize and a final output layer. We then project the outputs via projy to ensure the outputs lie in U. In addition, for the MLP layers, we use SwiGLU activations (Shazeer 2020) instead the standard GELU, where the ratio of the hidden size of the SwiGLU to Nsize and we do not the default Fourier embeddings for the inputs we pass our inputs directly into the model. We also use the default time embeddings. We refer to this model setup as DiT.
We also use the Gaussian positional embeddings from Luo et al. 2022) as additional features that concatenate the inputs of Di T. To compute this from x, we let
(-1
(x(0) x(0)
#
exp
V V
where & = 1...., K the number of basis kernels we use and n²* tt E are learnable parameters. Rx11 We then compute our positional embeddings by
We define (1.5) ((.s))"".... diK (kJ))
V. + Σラ= WD where Wn E B.K Xnem is learnable matrix, and concatenate these to form our embeddings = ,...,, x E Nxnas. We note that is (33)-invrriant.
Finally, provide pseudo-code for single pass through cg in Algorithm 3 where note that W, a learnable linear layer.
Algorithm 3 Computation of
Inputs: z = [x. h] where x E R Nx3 h E R NNx, te R
1: Compute V Nemb from x
2: z < xx. , where W, E (3+dd)xxn
3: z < (t, [z, I
4: return z
B.2 ARCHITECTURE OF
We construct ro following the recursive setup in Section 3.3 where we take the noise distribution // to be (0, I) on N N mmaie. We provide pseudo-code for a single pass through our fo Algorithm where we note that Wa. W1, W2 are learnable linear layers, we use the same embedding parameters for as before and D DiTwit hout FruutFinnat inalLayer the same as a DiT with Mayer (intermediate) layers, Whead attention heads, hidden size msize but where we do not apply the final layer.
B.3 HYPERPARAMETERS FOR AND ég
For both QM9 and GEOMsDrugs, we fixed the following hyperparameters.. For eg. we set K R3 Asize and Nemb Msize 7z.- For 700 Mmoise 3 and m/
2 maize-
17

---

Published as a conference paper at ICLR 2025
Algorithm 4 Computation of
Inputs: z = xx. h] where x E RNX33 h E RN R Ni ; R
1: Compute V R.V namb from
2: 2 < x. 1) L W where We E R (3+ M Mhank)
3: z < witWichouufFiallayerrtt..
4: 2 <- I"x where E N is vector of ones
Ensures Sy-invariance
"N"
5: z t- GELU(zW1)W2 where W1
R Wine W2 RR (3xx3)
6: R <- QRDecooposs i (z))
7: return R
C EXPERIMENTAL DETAILS
C.1 FIRST LIKELIHOOD TERM L1
We have a presented a framework for parametrising and optimising po (z (t) for t > 1 in Section 3.3 where pe is obtained via stochastic symmetrisation. This corresponds to the L; terms where > from our objective equation / However, we note that standard diffusion models usually choose different parametrisation for po [2) this corresponds to the final generation step. Depending on the modelling task, this requires a different approach compared to the other reverse kernels.
For example, in Hoogeboom et al. 2022). Pz (Z)) is defined the product of densities pg (z) bpg (1) (1) where po implements quantisation step converting the continuous latent 21 discrete values ho, while Pa you is still Gaussian distribution generating continuous geometric features Xy from In particular, we have that
cont Po
Xx
(21)
(xo: 01/a1c0 (21) (141),/////) aa
where Eg Z + output of eg.
is some Sy- -nviriati neural network and eg denotes the x component of the
We note that our proposed methodology can still account for this case by defining the symmetrised kernel by
Po 2o 21) sym..o (kio) (2o (2)),
ka (2ol2)) pe (xn 21 )1
disc (holz))
We can follow the same discussion Section
3.3 to conclude that
Po 1 = E (ddR(ar) [ko (zol R, 21)],
(z 21) = kepont (x R, 1h disc (hol R, 21).
where k.gon 20l R, 21) Nix:: x1 /1/an R en RF
( 21), /a and I (hol R, 21)
pdice (ho R 21). This allows us to decompose L1 into the form in equation and to tractably
optimise this objective since we have access to the density kig (zo|AR,, 21).
C.2
QM9 DETAILS
C.2.1 MODEL HYPERPARAMETER
For all of our experiments, we retain the diffusion hyperparameters as EDM (Hoogeboom et al., 2022) i.c. we use the same noise schedule, discretisation steps etc.
SymDiff Jsbles the hyperparameters for the cy backbone of the ky component used in the SYMDIFF models for QM9, The remaining hyperparameters were kept the same as in Appendix
B.3
For the fo backbone of the component, we set mssze = 128, Mlayers = 8, Mheads 4 for all models bar SymDiff* For SymDiff", we set msize = 216, Mayyers = 10, Whhaas =
EDM
Table shows the hyperparameters for the EDM models that used for QM9 exper- The remaining model hyperparameters were kept the same as those Hoogeboom et al.
iments.
2022
18

---

Published as a conference paper at ICLR 2025
Table 5: Choice of Nsizes Players Thhads for the eg of the SYMDIFF models used for Q.9.
Model
# Parameters
M/aize
Mayyrr
Yllccads
SymDiff
115.6M
768
12
12
SymDiff
29M
384
12
SymDitf
21.3M
360
10
SymDiff
11.3M
294
Table 6: Choices of the hyperparameters ny (# features per layer), ny (number of layers) for the EDM models used for QM9.
Model
Parameters
7
7M
EDM
12.4M
332
12
EDM
9.5M
256
16
EDM
5.3M
256
C.2.2 OPTIMISATION
For the optimisation of SYMDIFF models, we followed Peebles & Xie (2023) and used A [oshchilov & Hutter) with a batch size of 256. We chose learning rate of 2 x 10 and weight decay of 10 12 for our 31.2M parameter model by scarching over a small grid of values for each. To match the same number of steps in HHogghhoo al. (2022). trained model for 4350 epochs.
We applied the optimization hyperparameters from our 31.2M model to all other SYMDIFF models bar SymDiff", where we used learning rate of 10-4 For the EDM models, we followed the default hyperparameters from Hoogeboom al. (2022). In our augmentation experiments, first tuned the learning rate and weight decay for the DiT model, keeping all other optimization hyperparameters unchanged. These tuned values were then applied DiT-Aug.
C.3 FEOM-DRUGE DETAILS
For all of our experiments, we retain the diffusion hyperparameters as EDM (Hoogeboom et al. 2022 i.e. we use the same noise schedule, discretisation steps etc.
Like with QM9, report the seconds per epoch, sampling time (s) and vRAM for the models used in Table We exclude END as their code is not publicly available. We omit the results for EDM and GcoLDM were unable to run their code
NVIDIA H100 80GB GPU.
Table 7: Seconds per epochs, sampling time, and vRAM for different models on GEOM-Drugs.
Method
Parameters
Sec//epoch
Sampling time (s)
vRAM (GB)
GeoLM
5.5M
EDM
2.4M
SymDiff
31.2M
4336.82
0.39
63
For the SymDiff model, we used the same hyperparameters as for QM9 except for the learning where we used 10-4 as found this to result in a lower validation loss.
C.4 PRETRAIN-FNEEUNIIG
To further explore the flexibility of our approach, we experimented with using it in the pretrain- finetune framework, similar to Mondal et al. (2023). Using QM9, we took the trained DiT model from Table and substituted it as the for our SymDiff model, while keeping the same architecture and hyperparameters for fg. We tested two setups: finetuning both eœ and fo (DiT-FT) and freezing eg while tuning only fe (Die-FF-Freeze). The same training procedure and optimization hyperpa- rameters were used, except we now trained our models for only 800 epochs and used a larger grid for learning rate and weight decay tuning. Specifically, searched first for the optimal learning
19

---

Published as a conference paper at ICLR 2025
rate in [10-3,
3,00-,8,0 8 We found the optimal learning rate and weight decay to be 10-3 and x 10 10-12
10 1004 11=1]] and for the optimal weight decay in [0, 10 12 10 10-12]. 12].
Table 8: Test NLL, atom stability, molecular stability, validity and uniqueness QM9 for 10,000 samples and 3 evaluation runs.
Method
N
Atm. stability (%)
Mol. stability (%) f
Val. (%)
Uniq. (%)
SymDiff
-12
98.74
87.49.0.25
95.75
97.89
DiT-FT
-111.66+1.22
98.43=0.03
83.27 2:039
94.19
98. 18.17::0.20
DiT-FT-Freeze
-43.2943.78
95.68=0.02
55.02=0.38
90.48 20.24
99.06=00113
DiT
-127.78++2.99
98.23
81.033+0.25
94.7 40.31
97.98
From T we observe that finetuning both ey and So improves performance the DiT model, even with our minimal tuning. However, finetuning only fo leads to worse results, indicating that end-to-end training or finetuning the whole model necessary. This underscores the flexibility of our approach and its potential for easy and efficient symmetrisation of pretrained DiT models with an unconstrained fo.
D DISCUSSION ABOUT THE SYMDIFF OBJECTIVE
We explain here why the SYMDIFF objective in equation is reasonable use as a surrogate for the true ELBO in equation 4 The underlying idea is analogous to Remark 7.1 Cornish (2024). First, is straightforward check that our SYMDIFF objective recovers the ELBO exactly if either of the following two conditions are met:
is deterministie, i.e. (d(gg)z)) is Dirac distribution for every E 2; or is G-equivariant.
(For our model in Section 3.3 the latter holds if the function en 2 Z deterministically -euuiuuuir)))) It follows that the result of optimising our SYMDIFF objective will achieve least high an ELBO as the best performing for which either of these two conditions are met. Accordingly, if our model is powerful enough to express (or approximate) rich family of deterministic "o and G-equivariant kg. then it is reasonable surrogate objective. More generally, our model also has the ability to interpolate between these two conditions, allowing for potentially better overall optima than could achieved either case individually.
expect good performance from our
E EXTENSION TO SCORE MATCHING AND FLOW MATCHING
In this section, we discuss how to extend stochastic symmetrisation to score and flow-based gener- ative models to give an analogue of SYMDIFF to these paradigms. For clarity of presentation, we consider all models to be defined for N-bodd systems living in the full space 2 = N Nx3 where we wish to obtain SN ((3)-equivariant model i.e. we do not consider non-geometric features or translation invariance. Although, we note that the below discussion can be extended to such settings in the natural way as presented for diffusion models above.
E. SCORE MATCHING
Score-based generative models (SGMs) (Song et al. 2020) the continuous-time analogue of diffusion models. SGMs consider the forward noising process X; ~ p: for E [o, 7] defined by the following stochastic differential equation (SDE) with the initial condition Xo Pdau:
dx; = fxx, t) + g(t)dw,
(22)
for some choice of functions f [0, + and g [0, T] - R. and where w is a standard Weiner process.
20

---

Published as a conference paper at ICLR 2025
The corresponding backward process is shown in Anderson 1982 take the form:
dx; (x, t) (()) Vx x log (x dt + g(t)dw,
(23)
where is standard Weiner process and time runs backwards from T to 0. Hence, given samples Xy ~ Pr and access to the score of the marginal distributions Vx (x we can obtain samples from Pdam by simulating the backward process in equation 23
we can represent the sampling scheme of SGM in terms of the Markov chain (x(xx) II. p(xt,- Xt,), where the time-points t are uniformly spaced in [0, [ i.e. ti = i where At = T/n and the reverse transition kernels given by:
By considering the Euler-Maruyama discretisation of equation
(x.)) ((4)) A()). In what follows, we additionally assume that f(. t) is SN x 0(3) eqqiivriat for all t E [0, This is true for common choices of which take f to be linear in X
p(x (xx,) N (x,,1ixx At
Al (f(xx: (2) g(4))
Stochastic symmetrisation In order to learn an approximation to the transition kernels via stochastic symmetrisation, we can parametrise the reverse transition kernels, in a similar fashion as for diffusion models, by
sym 1xxe),
where we take ko [x (x,) = 1xxx ((:)))11)) We define ss by the following parametrisation
(x..)) xx,, + At A - (xx,)))
where we take so : Z + 2 to be a S Sx-equivariant neural network which aims to learn an ap- proximation to the true
Vx log Pt (x,). This ensures that ky is Sy-invariant. Additionally, assume that "/ 2 + +() is some choice of a Sx-invariant and (33)-qquivariant Markov kernel. Hence, conclude that pe Z + E is SN x -33)-equivariant Markov kernel by Theorem We can also guarantee that admits a density by Proposition
Training To learn 0 for pa (xt, (x,)), a natural objective is minimise the KL divergence be- tween the true reverse kernels and our parametrised reverse kernels
((0)= = - (o(t))
(())) = -Ep,, (xx,) (p(x,) (x.,)llooxxx
where Å is some time weighting function. We note that we run into the same issue as with S SIFFIFF in that we do not have access to pe px(x,__| xx,) an expectation. However, (,()) is linear function of - log po (x (xx.) function, we can apply Jensen's 's inequality again to provide the following upper bound our original objective
in clomed-form, since this is expressed in terms of
and log is convex
C'(0) = Σ ((). (0 = Epp, (x F=1
(ka R, X ))]
which now to train 0. To further simplify ()),,, note that ko(x,, x,,) N(x; .; R. po xx,,); g(.)) with similar derivation before. This allows to evaluate the KL divergences in closed form since kg are defined in terms of Gaussians. We can show that this gives
so(R" x1,) - Vx log (x² [L (o(t)) R-
() Ep,, (xx.)
(x.))
(24)
where we use the fact that f(,() is SN x -quueqquiaariant, To express equation 24 in tractable form (as we do not have access to the true score), we can apply the standard technique of employing the score matching identity (Vincent 201T) to give
00) (x)())) (pp(x))p(x4, (xxx,/ [Xa), )
22 (0(4) (1 R RR( xx.) Vx log Pt, (x,, [x))
+ C.
Samiilar to discussion diffusion models, leave the time dependency implicit in here.
21

---

Published as a conference paper at ICLR 2025
where p(x,, Xx) denotes the conditional distribution of xt, given Xy under the forward noising process p. and C2 is some constant. In practice, the choice of forward noising SDE in equation 22 is made ensure that we have access to p(x,, (xp) in closed-form and that the distribution is easy to sample from.
By making the choice that (t( (g(4)) for some suitable time weighting function X, we can show as At + O that our objective (0) will converge to (modulo some constant)
(U))),)) (x)) (dR(x.)
[()] R so (R x.) Vx p(x xx))|
(25)
where (0, I denotes the uniform distribution on [0, T). We see that our final objective in equa- tion 25 now resembles the standard score matching objective.
E.2 FLOW MATCHING
Continuous normalising flows (CNFs) (Chen et al. 2018) construct a generative model of data X1 ~
q = Púata by the pushforward of an ordinary differential equation (ODE) taking the form
d _ 0x(x) ue (x)),
(x(x) x,
(26)
where Ut 3 x [0, is the vector field function defining the ODE, and Φ. 2 x [0, [0 Z denotes the flow implicitly defined by solutions the above ODE. By letting po be some simple prior distribution, the above ODE defines a generative model X, ~ Pt by the pushforward of po through the flow Φ
pt [o
(27)
If dt is chosen in such a way that P1 q Pdsta we can then generate samples from Pdata
by sampling some Xo Po. then solving the ODE in equation [6 with this initial condition
Furthermore, as in previous work (Klein et al. (024). we assume that un SN -quiiiarri
By considering the Euler discretisation of equation 26 we can represent the generation process by the Markov chain Po (xx) II p(dx,, _x) where the time-pointe t; uniformly spaced in [0, T]
- i.c. t, = iAt where At T/n and the transition kernels are given by the Markov kernels
p(dx,, 1x__,) = 5(x,_, + ut,,, (tt),
(28)
where 8(-) denotes the Dirac measure at some point.
To learn the transition kernels induced by the vector field u, via stochastic symmetrisation, parametrise our transition kernels by
Stochastic symmetrisation
po (dxx, (x._,) sym., (ke) (xx, (x),,.)
(29)
where we take ka (cxx,, (x)_,) = 5(xt,_, + vi__, (x,_,, Att) ) in which vf 2 2 is some Sy- equivariant neural network which aims to learn an approximation to the true vector field u,. We further assume 79 + 0(3) some Sy-invariant and (33)-qquivariant Markov kernel. We can again conclude that pa Z + Z is a (SN x ((3)) Markov kernel by Theorem
Training A natural objective to learn po is to minimise the 2- 2-sssersttein distance W2 (Peyré al. 2019) between p(dx,, (x,,,) and po (dxe, (xt_-)) since these are defined in terms of Dirac measures. We can write our objective as
((0) Σ (,()). C,(0)= = ((0))=pp__
(x)_)) (x [M ( [)]]
where do is some time weighting function, and the 2-Wasserstein distance W/ is defined as W m2) inf, ||x y dx(x,y) where is taken over the space of possible couplings
'The use of time here reverse the convention used in the diffusion literature.
22

---

Published as a conference paper at ICLR 2025
between the measures T1, T2. We note that p(dx,, X ,) is Dirac, there only exists single cou- pling between the two kernels given by the product of the Markov kernels. This allows us to evaluate Li+1 as
(0) P.,, (x., (dR)))))
Ara R v/. (R. xx,) Me, (xx:)
(30)
To express equation 30 in tractable form (as do not have access to ue). we take ut to be constructed by the same setup used Flow Matching (Lipman et al. (022). This framework allows us express equation i0
in in the now tractable form
(0) gg(x)) (x) [
R vi (R. xx,) ue, (x,, [x1)]]] + C++1.
(31)
by the use of Conditional Flow Matching objective, where C++1 some constant. Here (x,))(x)) is family of conditional distributions where (x
Po equals our prior dis- tribution and p1 (x1 (x1) ss 8(x1). ). and for which u, (x(x((x)) is vector field generating Pe Pe(x; x1) by an ODE of the form in equation 27 These constructed to be easy to sample from and evaluate. The true vector field u4, which provides generative model of q Pdatas is then defined by some expectation of the conditional vector fields u4 (xx(xx))
(x,,, Xx) and g(x1).
Hence, by taking (t) 0, our objective L(0) will converge to (modulo some constant)
TAr for some suitable time weighting function X, we can show as At ->
A(t)
[A(t]] (R xx)) (x, x3 (x/xxx) ||]
(32)
We see that our final objective in equation 32 now resembles the standard flow matching objective.
23


Reviews:
[{"review_id": "kdR75WnsRK", "reviewer": "Reviewer_u6Xc", "summary": "This paper proposes a new way to build equivariant diffusion models. Whereas traditional methods rely either on simple data augmentation strategies or neural network architectures that are directly constructed to be equivariant, *SymDiff* uses stochastic symmetrization to learn equivariant diffusion models. Stochastic symmetrization is a recently introduced framework by Cornish (2024). On a high level, stochastic symmetrization can intuitively be seen as an advanced data augmentation method, where the augmentation is sampled from an additional, learnable Markov kernel and where the main denoiser network does not need to be equivariant itself. Moreover, the augmentation is also applied during sampling. The paper in detail introduces the approach and rigorously shows that it is principled. SymDiff is then validated on molecule generation, where the permutation equivariance is directly built into the employed networks, whereas rotation equivariance is enforced through stochastic symmetrization. The paper compares to various baselines, showing favourable results, and runs many appropriate and insightful ablation experiments.", "questions": "I have put almost all my questions in the ''Weaknesses'' section above already, but I have one very minor additional question. In 4.2, the authors write that, given the size of GEOM-Drugs, they train a smaller model than on QM9. But the GEOM-Drugs dataset is larger and more complex than QM9. Wouldn't one want to train a larger, more expressive model in that case, and not a smaller one?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "**Clarity**: The paper is overall well written and the necessary background is introduced in an appropriate manner, such that one can follow the overall paper. However, the stochastic symmetrization framework is introduced in a fairly abstract manner and I believe more intuitions could be discussed and examples given, which would likely help the broad ICLR audience (see questions below).\n\n**Originality**: Building equivariant diffusion models with stochastic symmetrization is an overall novel and original idea. Methodologically, the approach directly builds on the stochastic symmetrization framework from Cornish (2024). Hence, the novelty here lies in the application of the framework to diffusion models, specifically those for small molecule generation, a popular and important benchmark task.\n\n**Experimental Results**: Overall, the paper is able to show favourable results compared to various baselines, thereby successfully validating the method, and I liked the many ablation experiments that were run.\n\n**Quality and Significance**: Overall, this paper is of high quality and tackles an important problem. Enforcing equivariances and invariances in generative models can be crucial in applications with symmetries, as commonly encountered in scientific applications. However, it is difficult to directly build high-performance and scalable equivariant network architectures. Although quite related to data augmentation methods, stochastic symmetrization and SymDiff proposes a new way to address the problem, and I believe this paper has the potential to result in follow-up work. Hence, I believe this is sufficiently significant for ICLR.", "weaknesses": "I also see some weaknesses.\n\n- *More discussion of intuitions and more experimental analyses*: I think it would be great if the authors discussed some intuitions in more detail and analyzed their models in more detail. Specifically:\n  - In practice, training SymDiff boils down to the training procedure in Algorithm 1. $\\gamma_\\theta$ is trained as part of the ELBO objective, using reparametrization. What sort of $\\gamma_\\theta$ and $f_\\theta$ functions are we intuitively expected to learn in different situations? How does the model intuitively exploit its flexibility to produce more complex $R$ distributions compared to when only using the plain Haar measure? Can examples and intuitions be given?\n  - Moreover, is there anything preventing the model from simply learning $R = R_0 f_\\theta(...) = I$, this is, no ''augmentation'' at all? In other words, could $f_\\theta$ learn to ''undo'' the $R_0$ drawn from the Haar measure so as to remove any rotations? In that case, the training would reduce to standard training without the rotation equivariance, which may be beneficial from the objective's perspective considering that the network itself isn't equivariant? But this defeats the purpose of the method. Such potential edge cases should be discussed.\n  - To this end, can the authors maybe show and analyze what sort of $\\gamma_\\theta$ distributions SymDiff ends up learning in practice?\n  - The approach uses $\\gamma_\\theta$ (or $f_\\theta$) also during sampling. What is the intuition for that? And what would happen if we didn't apply $\\gamma_\\theta$ during sampling? This would be an interesting discussion or ablation. Moreover, in the case without learnable $\\gamma_\\theta$, i.e. when using only the Haar measure, I would expect that the model would still perform fine even if no $R$ is applied during sampling.\n\n- *Relation to data augmentation*: Related to the last point above, overall, it seems that SymDiff can be seen as a sophisticated data augmentation strategy (as also pointed out by the authors). However, the authors carefully distinguish the case between training with data augmentation (DiT-Aug) and SymDiff with Haar measure sampling only (SymDiff-H), pointing out that the only difference in that case is that SymDiff-H also applies augmentations during sampling, if I understood things correctly. However, in all experiments the performance of DiT-Aug and SymDiff-H is extremely close together. This makes me believe the application of the augmentation during sampling isn't necessary for SymDiff-H (also see point above), and that SymDiff-H really just corresponds to standard data augmentation. Related to that, it is not sufficiently well discussed, explained and analyzed how exactly this situation *qualitatively* changes when $\\gamma_\\theta$ is learnt (see point above).\n\n- *Experimental results*: The performance gains over the EDM baseline are meaningfully large, but at the same time compared to more sophisticated baseline models SymDiff does not perform much better. Further, the authors make the point that the comparison to EDM is fair. However, the very different network sizes still make a fair comparison somewhat difficult. How would performance between EDM and SymDiff compare for similar-sized neural networks?\n\n**Conclusion:** SymDiff is an interesting and novel method, tackling a significant problem, and I believe it is of interest to the ICLR audience. However, there are various fundamental questions about the method and the experiments (see above), that I hope the authors can answer and address. For now, I am leaning towards suggesting acceptance.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Pw2jnVSnWx", "reviewer": "Reviewer_RgSJ", "summary": "The paper “SymDiff: Equivariant Diffusion via Stochastic Symmetrisation” introduces SYMDIFF, a novel method for constructing equivariant diffusion models without the need for complex, intrinsically equivariant neural networks. Traditional approaches often rely on specialized architectures with complex parameterizations and higher-order geometric features to achieve equivariance, which can be computationally intensive and challenging to implement.\n\nSYMDIFF leverages the framework of stochastic symmetrisation, effectively functioning as a learned data augmentation technique applied during the sampling phase. This approach is lightweight, computationally efficient, and can be easily implemented on top of arbitrary off-the-shelf models. By not requiring the neural network components to be intrinsically equivariant, SYMDIFF allows the use of highly scalable modern architectures, providing greater flexibility and performance.\n\nThe authors demonstrate the effectiveness and scalability of SYMDIFF through empirical experiments on E(3)-equivariant molecular generation tasks using the well-known QM9 and GEOM-DRUGS datasets. The results show significant improvements over traditional methods, highlighting the benefits of SYMDIFF in generating molecular structures that are equivariant under the E(3) group (which includes translations, rotations, and reflections in three-dimensional space).\n\nThis work represents the first application of stochastic symmetrisation to generative modeling, suggesting that SYMDIFF has the potential for broader applications in the domain of equivariant machine learning models.", "questions": "- Do you have insights into how many training epochs are necessary for SYMDIFF to effectively learn equivariance in the molecular generation task?\n- How does the convergence speed of SYMDIFF compare to that of intrinsically equivariant models in learning equivariant properties?\n- In the context of molecular generation, do you consider the equivariant component of the task to be a minor or major challenge compared to learning the underlying chemical laws necessary for accurate atom placement?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "This paper introduces SYMDIFF, a simple yet powerful method for constructing equivariant diffusion models without relying on complex intrinsically equivariant neural network architectures. The strengths of the paper can be assessed across the following dimensions:\n\nOriginality:\nThe paper presents a novel application of stochastic symmetrisation to generative modeling, specifically in the context of diffusion models. This approach is original because it departs from traditional methods that require specialized equivariant architectures with complex parameterizations and higher-order geometric features. By leveraging stochastic symmetrisation as a learned data augmentation technique applied during sampling, the paper creatively combines existing ideas to overcome limitations of prior models. This opens up new possibilities for using standard neural network architectures in equivariant settings.\n\nQuality:\nThe paper demonstrates high quality in both theoretical and empirical aspects. It provides a solid theoretical foundation based on Markov kernels, offering a clear understanding of how SYMDIFF achieves equivariance without specialized architectures. The empirical evaluations are thorough, showcasing the method’s effectiveness on E(3)-equivariant molecular generation tasks using the well-known QM9 and GEOM-DRUGS datasets. The results indicate significant performance improvements and scalability advantages over traditional equivariant models. The experiments are well-designed, and the conclusions are well-supported by the data.\n\nClarity:\nThe paper is written with clarity and precision, making complex concepts accessible to the reader. It systematically introduces the theoretical background before delving into the methodological details of SYMDIFF. The explanations are clear, and algorithmic descriptions helps in understanding the implementation. The empirical results are presented in an organized manner, with insightful discussions that highlight the implications of the findings. Overall, the paper effectively communicates its contributions and significance.\n\nSignificance:\nThe significance of the paper lies in its potential to impact the field of equivariant machine learning models substantially. By eliminating the need for specialized equivariant architectures, SYMDIFF lowers the barrier to entry for researchers and practitioners interested in building equivariant models. It enables the use of arbitrary off-the-shelf neural network architectures, which are often more scalable and easier to implement. This flexibility can accelerate the development and deployment of equivariant models in various applications, including molecular generation, computer vision, and physics-based simulations. Moreover, being the first to apply stochastic symmetrisation to generative modeling, the paper paves the way for future research in this promising direction.", "weaknesses": "While the paper introduces SYMDIFF as a flexible and efficient method for constructing equivariant diffusion models, it lacks specific guidelines to help practitioners decide when to use SYMDIFF over traditional intrinsically equivariant models. This absence of decision-making support makes it challenging to understand the circumstances under which SYMDIFF is more advantageous.\n\nTo address this, I suggest the authors include a comparative analysis between SYMDIFF and traditional methods. This could be presented as a table or flowchart highlighting key factors such as:\n\n- Dataset Characteristics: Situations where SYMDIFF performs better or worse.\n- Model Complexity: Differences in architecture and scalability.\n- Computational Resources: Comparisons of memory usage, training time, and inference speed.\n- Performance Metrics: Quantitative comparisons on accuracy and sample quality.\n- Ease of Implementation: Practical considerations like code complexity.\n\nIncluding such a comparison would provide clear guidance for practitioners, making the paper more actionable and useful.\n\nAdditionally, while the paper claims that SYMDIFF is computationally efficient, it doesn’t provide empirical evidence to support this. I recommend the authors include:\n\n- Empirical Benchmarks: Runtime analysis and resource consumption statistics compared to intrinsically equivariant models.\n- Convergence Plots: Performance versus number of training epochs for both SYMDIFF and baseline methods.\n- Ablation Studies: Impact of stochastic symmetrisation on computational overhead and model performance.\n\nBy incorporating these elements, the paper would strengthen its claims and offer actionable insights, enabling readers to make informed decisions about adopting SYMDIFF.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "3hZWo3Owwd", "reviewer": "Reviewer_GD6U", "summary": "This paper proposes SymDiff, a novel method for constructing equivariant diffusion models via stochastic symmetrisation. It addresses limitations of prior work on geometric data by allowing more flexibility than using intrinsically equivariant neural networks. Experiments on QM9 and GEOM-Drugs datasets for molecular generation show SymDiff outperforms EDM and is competitive with other baselines, with better scalability. SymDiff can be seen as a generalisation of data augmentation, with a learned and deployed augmentation process ensuring stochastic equivariance.", "questions": "- Why SymDiff require larger model size to match the performance of EDM / GeoLDM etc.", "limitations": "", "rating": 5, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Originality: The paper proposes SymDiff, a novel method using stochastic symmetrisation for equivariant diffusion models, extending symmetrisation to generative modelling and differing from prior approaches.\n\n- Clarity: The paper is well-structured, clearly explaining concepts with examples and presenting algorithms, aiding understanding and implementation.\n\n- Quality: The paper presents a comprehensive framework with solid derivations and conducts rigorous experiments on multiple datasets.", "weaknesses": "- Performance gap: from Tab2 and 3, seems the model is not parameter efficient compared with baseline models. When reducing the model size, the performance will have a larger drop than the baseline model.\n\n- Unfair comparison: I didn't fully understand for SymDiff+++, ++, and +, why you cannot choose the same number of parameters of EDM. It seems your model SymDiff+ is already larger than EDM+++, but with worse performance.\n\n- Missing comparison: I think the author also needs to incorporate GeoLDM with different sizes into Tab 2 and 3, which potentially can work better since it acts on low-dimensional space.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)