Equivariant Neural Diffusion for Molecule Generation
Authors: François R J Cornet
Grigory Bartosh
Mikkel N. Schmidt
Christian A. Naesseth
Venue: NeurIPS 2024 poster (2024)
Keywords: Diffusion Models
Equivariant Neural Networks
Molecule Generation
Primary area: diffusion_based_models

Abstract: We introduce Equivariant Neural Diffusion (END), a novel diffusion model for molecule generation in 3D that is equivariant to Euclidean transformations. Compared to current state-of-the-art equivariant diffusion models, the key innovation in END lies in its learnable forward process for enhanced generative modelling. Rather than pre-specified, the forward process is parameterized through a time- and data-dependent transformation that is equivariant to rigid transformations.  Through a series of experiments on standard molecule generation benchmarks, we demonstrate the competitive performance of END compared to several strong baselines for both unconditional and conditional generation.

Equivariant Neural Diffusion for Molecule Generation
François Cornet Technical University of Denmark fr jr dk
Grigory Bartosh University Amsterdam g. barrosshuua..n. nl
Mikkel N. Schmidt
Christian A. Naesseth University of Amsterdam caassseth@uva.. nl
Technical University of Denmark mnsc@dtu. dk
Abstract
We introduce Equivariant Neural Diffusion (END), a novel diffusion model for molecule generation in 3D that is equivariant to Euclidean transformations. Com- pared to current state-of-the-art equivariant diffusion models, the key innovation END lies in its learnable forward process for enhanced generative modelling. Rather than pre-specified, the forward process is parameterized through time- and data detendepn transformation that is equivariant rigid transformations. Through a series of experiments on standard molecule generation benchmarks, we demon- strate the competitive performance of END compared to several strong baselines for both unconditional and conditional generation.
Introduction
The discovery of novel chemical compounds with relevant properties is critical a number of scientific fields, such as drug discovery and materials design (Merchant et al., 2023). However, due to the large size and complex structure of the chemical space (Ruddigkeit et al., 2012), which combines continuous and discrete features, it is notably difficult to search. Additionally, ab-initio Quantum Mechanics (QM) methods for computing target properties are often computationally expensive, preventing brute-force enumeration. While some of these heavy computations can be amortized through learned surrogates, the need for innovative search methods remains, and generative models have recently emerged promising avenue (Anstine and Isayev, 2023). Such models can learn complex data distributions, that, in turn, can be sampled from obtain novel samples similar to the original data. Compared other data modalities such images text, molecules present additional challenges as they have to adhere to strict chemical rules, and obey the symmetries of the 3D space.
Currently, the most promising directions for molecule generation in 3D are either rego-segees models (Gebauer et al., 2019; 2022; Luo and Ji, 2022; Daigavane et al., 2024) building molecules one atom at time, or Diffusion Models (DMs) (Hoogeboom et al., 2022; Vignac et al., 2023; Le al., 2024) that learn to revert a corruption mechanism that transforms the data distribution into noise. As both approaches directly operate in 3D space, they can leverage architectures designed for machine learned force fields (Unke et al., 2021), that were carefully developed to encode the symmetries inherent to the data (Schütt et al., 2017; 2021; Batzner et al., 2022; Batatia et al., 2022).
The success of DMs has not been limited to molecule generation, and promising results have been demonstrated on variety of other data modalities (Yang et al., 2023). Nevertheless, most existing DMs pre-specify the forward process, forcing the reverse process to comply with it. recent line of work has sought overcome that limitation and improve generation by replacing the fixed forward process with a learnable one (Bartosh et al., 2023; Nielsen al., 2024; Bartosh et al., 2024).
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Contributions In this paper, we present Equivariant Neural Diffusion (END), a novel diffusion model for molecule generation in 3D that (I) is equivariant Euclidean transformations, and (2) features a learnable forward process. We demonstrate competitive unconditional molecule generation performance on the QM9 and GEOM-Drugs benchmarks. For conditional generation driven by composition and substructure constraints, our approach exhibits a substantial performance gain compared to existing equivariant diffusion models. Our set of experiments underscores the benefit of learnable forward process for improved unconditional and conditional molecule generation.
Background
We begin by establishing the necessary background for generative modeling of geometric graphs. We first introduce the data representation and its inherent symmetries. We then discuss Diffusion Models (DMs), and more specifically the Equivariant Diffusion Model (EDM) (Hoogeboom al., 2022). Finally, we present the Neural Flow Diffusion Models (NFDM) framework (Bartosh et al., 2024).
2.1 Equivariance
Molecules as geometric graphs E(3) We consider geometric graphs embedded in 3-dimensional Euclidean space that represent molecules. Formally, each atomistic system can be described by tuple x = (r, h)) where r = (Ti...., TM) E RMx3 form a collection of vectors in 3D representing the coordinates of the atoms, and h (h1, hAs) E RMxD are the associated scalar features (e.g. atomic types or charges). When dealing with molecules, we are particularly interested in E33,, the Euclidean group in 3 dimensions, generated by translations, rotations and reflections. Each group element E(3 can be represented as a combination of a translation vector t E IR3 and an orthogonal matrix R 0(3) encoding rotation or reflection. While scalar features h remain invariant, coordinates r transform under translation, rotation and reflection as Rr+t (Rri + t.., R + +t).
X + V said to be equivariant to the action of group G, or G-equivariant, if g- f(x) f(g-w) x) Vg E G. is said to be G-invariant, if f(x) f(g.z), Vy E G. In the case of a function f: (IR (RMx5 x RMMx) - (RMx3 x MxDD operating on geometric graphs,
Equivariant functions A function f:
the function is said to be -quivvvaivri if
Ry + t, (h) 5 (Rr V E 0(3) and sadte t RR.,
where (r) and (h) denote the output related to and h respectively. There exists large variety of graph neural network architectures designed to be equivariant to the Euclidean group (Schütt et al., 2017: 2021; Batzner et al., 2022: Batatia al., 2022).
Equivariant distributions A conditional distribution (())))) is equivariant to rotations and reflec- tions when p(y(z) = p(Ry) Rx), A R E O(3), while a distribution is said to be invariant when p(x) pRR)), v R E O(3). Regarding translation, it is not possible have a cansslaaioniivvriaan non-zero distribution, as would require that p(x) = p(æ + t). Vt. E IR3 æ E RRM x3 which would mean that p(=)) cannot integrate distribution can be constructed in the linear subspace R, where the centre of gravity is fixed to 0 (i.e.
(Gareia Satorras et al., 2021). However, a caanllatonninvariin
zero CoM subspace): R = fr E M3 EM = 0} (Xu et al., 2022). As R can be shown to be intrinsically equivalent to R 1M-1)) (Bao al., 2023), we will consider in what follows that is defined in IR -3311) for ease of notation.
2.2 Equivariant Diffusion Models
Diffusion Models (DMs) (Sohl-Dickstein et al., 2015; Ho et al., 2020) are generative models that learn distributions through a hierarchy of latent variables, corresponding to perturbed versions of the data at increasing noise scales. DMs consist of a forward and reverse (or generative) process. The Equivariant Diffusion Model (EDM) (Hoogeboom et al., 2022) is a particular instance of DM, where the learned marginal (x: is made invariant to the action of translations, rotations and reflections by construction. Intuitively, this means that the likelihood given molecule under the model does not depend on its orientation.

---

Forward process The forward process perturbs samples from the data distribution, a ~ q(z), over time through noise injection, resulting in a trajectory of latent variables (2)))e(o.1)- conditional on æ. The conditional distribution for (z: [e)o1] given æ, can be described by an initial distribution ( (x) and a Stochastic Differential Equation (SDE),
d(z (r) (10 - f(t) [x dt +g((t) df ( "
where the drift ft) and volatility gt)) are scalar functions of time, and w((r and w(t) are two independent standard Wiener processes defined (M-1))33 and xDD respectively. Specifically, EDM implements the Variance-
Variance-Preservingg Preserving SDE (VP-SDE) scheme (Song et al., 2020), with f(t)
- ()t) and g(t) v for fixed schedule (t).. Due the linearity of the drift term, the conditional marginal distribution is known in closed-form (Särkkä and Solin, 2019), and can be reconstructed as
x "1(r. hl) = q( (() (n) = (2 | -(z, <(22)" 2
(M) loyh, oi
where a, = exp (- 2 1o [(s) ds) and oe = 1-e (- 1 3(s) ds) ds). It evolves from a low varvaarancce Gaussian centered around the data (200/z)) R (2 1) to an uninformative prior distribution (that contains no information about the data distribution), i.e. unit Gaussian g(21 ar)) (2 (0 T)-
Reverse (generative) process Starting from the prior (z(r),=(M (M)] NN I)
.(2 (M) o. II),
samples from g(ær) can be generated by reversing the forward process. This can be done by following the reverse-time SDE (Anderson, 1982).
dz, (F()) [5(0 (0))) (1) [ =/r) goz)( V af log logg(zz))) di +g(() d(wo t (0)),
where (r() and tō (1) are independent standard Wiener processes defined (M and RMXDD, respectively, with time flowing backwards. DMs approximate the reverse process by learning approximation of the score function F,, log ((x)) parameterized by a neural network (2, t). With the learned score function sa(ze, ). sample zo po (zu) = ((20) q(x)) can be obtained by first
sampling from the prior (T) 2 ( ~ N(2, (r) I0, N(2; J0, I), and then simulating the reverse SDE,
dz, = (1(0)
g20( (2,, ( (z,,
(z.: 1) dt +g(() df (1)) (0)))
where the true score function has been replaced by its approximation sy s0(21, t). In EDM, the approx- imate score is parameterized through equivariant function: so(=,,t) t) [s8 (zz, t), (h) (ze, such that so [R,,"", 2; ( [R.s(])
(r) (zz, 80 (zz, t)] V 0(3). Practically, this is realized
(h)
through the specific parameterization.
so a, Zz, t) Zt
o?
where the data point predictor d is implemented by an equivariant neural network.
Optimization The data point predictor žs. or sa, is trained by optimizing the denoising score matching loss (Vincent, 2011).
LDSM (u(t)-((m.))) x(t) (z),)) Z.. log g(z, x ||2]
where X(t) is a positive weighting function, and u(t) a uniform distribution over the interval [o, 1].
2.3 Neural Flow Diffusion Models
Neural Flow Diffusion Models (NFDM) (Bartosh et al., 2024) are based on the observation that latent variables in DMs, i.e. 2,, conventionally inferred through pre-specified transformation as implied by the chosen type of SDE and the noise schedule. This potentially limits the flexibility of the latent space, and makes the learning of the reverse (generative) process challenging.
Forward process In contrast to conventional DMs, NFDM define the forward process implicitly through learnable transformation (e t, x) of injected noise time t, and data point æ. The latent
3

---

variables 2, are obtained by transforming noise samples E, conditional on data point and time step 2; = (e, t, x). If F& differentiable with respect to E and t, and invertible with respect to e. then, for fixed æ and e, samples from (2( can be obtained by solving the following conditional Ordinary Differential Equation (ODE) until time 4,
t. dz:- = fo(ze, f, x) dt with fz(z,, t,x) O (F,
(1)
ät
e=FF (zx,,,,))
with 20 (z) While F., and q(E) define the conditional marginal distribution qu (z) (z), we need a distribution over the trajectories (=)))_(0.1) NFDM obtain this through the introduction of conditional SDE starting from 2o and running forward in time. Given access to the ODE in Eq. (1)
and the score function V log (zel=)))) the conditional SDE sharing the same conditional marginal distribution qe (2)(()))
given by
dz, = f (Z t, x) dt (tp dw with f (zc, t, x)) = fo( (t(() 2
vx, log qu (zr(z)) (2)
where the score function of (p(2) V., log (z,lz) = Vz, [lo[log(( - log J with (z,, t,), and JF' oF (2,,)..)
Sx.
Reverse (generative) process A conditional reverse SDE that starts from g(z1). runs backward in time, and reverses the conditional forward SDE from Eq. (2) can be defined as
50(1)
(0,, t, x) + (t) dw with for (z,,t, x) = fz t, x) dz = f0 (zi,
v. log go (z.(z)).
(3)
2
As is unknown when generating samples, dz: = (ze t) dt (t) dıo, where (z.p = f (ze, t,, 2o t)),
rewrite Eq. (3) with the prediction of instead,
(4)
where 2y is function that predicts the data point æ. Provided that the reconstruction distribution ([Zolæ)) and prior distribution (=2)) are defined, this fully specifies the reverse (generative) process.
Optimization The forward and reverse processes can be optimized jointly by matching the drift terms of the true and approximate conditional reverse SDEs through the following objective,
DIDDM = (u(t).np(e..s)) (tg 115 (ze, x)) - fo.p 1) [1]]
(5)
which can be shown be equivalent to minimizing the Kullback-Leibler divergence between the posterior distributions resulting from the discretization of Eqs. (3) and (4) (Bartosh et al., 2024).
3 Equivariant Neural Diffusion
We now introduce Equivariant Neural Diffusion (END), that generalizes the Equivariant Diffusion Model (EDM) (Hoogeboom et al., 2022), by defining the forward process through learnable transformation. Our approach is a synthesis of NFDM introduced in Section 2.3, and leverages ideas of EDM outlined in Section 2.2 to maintain the desired invariance of the learned marginal distribution (zo). By providing an equivariant learnable transformation F, and an equivariant data point predictor ée. we show that it is possible to obtain generative model with the desired properties. Finally, we propose a simple yet flexible parameterization meeting these requirements.
3.1 Formulation
The key innovation in END lies in its forward process, which is also leveraged in the reverse (generative) process. The forward process is defined through a Iearnable time- and data-dependent transformation (o t, x), such that the latent transforms covariantly with the injected noise (i.e. a collection of random vectors) and the data point x.
F F. t, Rı) = RFF t, x) = R2,, VR E 0(3).
We then define áp as another learnable equivariant function, such that the predicted data point transforms covariantly with the latent variable 2t, i.e. & = R t). Finally, we choose the noise and prior distribution, i.c. p(E) and p(z1), to be invariant to the considered symmetry group.

---

Invariance of the learned distribution With the following choices: (1) p(z1) an invariant distri- bution, (2) F.p an equivariant function that satisfies (2 t, R.x) = F t, and (3) &y an equivariant function, we have that the learned marginal (zo) is invariant as desired. This can be shown by demonstrating that the reverse SDE equivariant. We start by noting that the reverse SDE in END is given by
dz, fo.p (Z1, t) dt (gp(() duū
As the Wiener process isotropic, this boils down to verifying that the drift term, fe.p(zz (2t, t) is equivariant, i.e. f (Rzz, f) (2,,t) t). As the drift is expressed as sum of two terms, we inspect each of them separately. The first term is
fz 1, 20) = 2F, (z(1).
Ōt
l==Pip (,
If Fip is equivariant, then so is its time-derivative (see Appendix A.3.1). The same holds for its inverse with respect to (see Appendix A.3.2), such that we have F (R,,, Rx) Re. We additionally have that do is equivariant, by definition. As the equivariance of F implies the equivariance of qp from the second term of the drift, observe that, for y: Reg, we have
R.F F-111 (2t, 1, a))
V., log Up (yrl de (v:, RV log (=((
A O(3).
In summary, in addition an invariant prior, an equivariant Fg and an equivariant dæ ensure the equivariance of the reverse process, and hence the invariance of the learned distribution. In Appendix A.3.3, we additionally show that the objective function is invariant, i.e. LEND(Ræ)) (E))( VR ((3).
We note that alternative formulations for the drift of the reverse process, fo.g- exist. Most notably, it can be learned directly through an equivariant function, without explicit dependence on F.p, while maintaining the desired invariance.
3.2 Parameterization
We now introduce simple parameterization of Fg that meets the requirements outlined above,
t,) = te (2, t) (, (x, t)s,
(6)
where, due to the geometric nature of æ. (z, t) (M is structured matrix where each block a 3 x 3 matrix, ensuring a cheap calculation of the inverse transformation and its Jacobian. This is equivalent
block-diagonal
diagonal parameterization in the case of scalar features.
Similarly to EDM, our parametrization of F&p leads conditional marginal (2 (a:) that is conditional Gaussian with (block-) diagonal covariance, with the notable difference that the mean and covariance are now data- and dependenn,, and learnable through Fep-
(2,((a)) = (N (x,,
where (x, = (Z t)U (x, t), such that Σ,p (x, t) is also block- -iiggnall As F is linear in E, both (hip and UT must be equivariant functions whose outputs transform covariantly with æ. in order to ensure the desired equivariance of Feg.
F (Re, l, Ræ) = Mip (Rx, () + U, (Ræ, I)RE = Ryp (æc, () + +R (2,
[)e] = RF x)).
We can then readily check that the resulting 9, is equivariant, as V R E 0()), we have that
4e (z.(2=))= = Al -N (z)lup(ar (), (x. 1)) = / (R2,, (ar 1). (R,, (x,,
= 40 (Rz, Ra).
We note that other, and more advanced, parametrizations are possible, e.g. based on normalizing flows with flow architecture similar to that of Klein et al. (2024).
Prior and Reconstruction While not strictly required, it be advantageous to parameterize the transformation FG such that the prior and reconstruction losses need not be computed. To do so, design (5, 2; t) such that the conditional distribution evolves from a fow-variance Gaussian centered around the data, i.c. g(z0 x) se (z0la 121) to uninformative prior distribution (that
5

---

contains no information about the data distribution), i.c. unit Gaussian ((21)z) 22 (2 I0.
Specifically, we parameterize Fs through the following functions,
(= t) (1 t) + 11 (1,
(8)
(x 1) = (2,,) + (1((111111111
(aa1(((--ppp)) (x,
(9)
which ensure that (i) in O, pl 0) x and Σ. (x,
8 (ii in = 1. (1
and (x, 1) = I; while being unconstrained for t <|0, 1[. We note that this is only one possible parametrization for F.p. and that, by adapting MMp and U.p. richer priors can easily be leveraged - e.g. harmonic prior given molecular graph or scale-dependent prior (Jing et al., 2023; Irwin et al., 2024).
Implementation In practice, Fe is implemented as a neural network with an architecture sim- ilar to that of the data point predictor (0 !). but with a specific readout layer that produces (ap 1), p (xt, 1), (2( (2r,()). The mean output Pip (ar, () is similar that of &h (2z, (). For (z (), as E.p((r 1) (2, U) p (ar () should rotate properly, T o2(2, t) is positive invariant sealar, while (1, t) constructed as matrix whose columns are vectors that transform covariantly with x.
To ease notation, we introduced all notations in the linear subspace R. however in practice we work in the ambient space, i.e. IR Mx3 We detail in Appendix 1.5. how working in ambient space is possible. The training and sampling procedures are detailed in Algorithms I and 2 in the appendix.
3.3 Conditional Model
While unconditional generation is required stepping stone, many practical applications require some form of controllability. As other generative models, DMs can model conditional distributions p (x)|c), where is a given condition. Different methods exist for sampling from the conditional distribution, e.g. via guidance (Bao et al., 2023; Jung et al., 2024) or twisting (Wu et al., 2024), but the simplest approach consists in training a conditional model on pairs (x, c). In such setting, F&p and ê, simply receive an extra input representing the conditional information, such that they respectively become (e, f,, x, e), and ie (2,, I,, c). It is important note that, compared conventional DMs, the forward process of END is now also condition-dependen..
Experiments
In this section, we demonstrate the benefits of END with comprehensive set of experiments. In Section 4.1, we first display the advantages of END for unconditional generation on standard benchmarks, namely QM9 (Ramakrishnan et al., 2014) and GEOM-DRUGS (Axelrod and Gomez- Bombarelli, 2022). Then, in Section 4.2, we perform conditional generation in 2 distinct settings on QM9, Additional experimental details are provided in Appendix A.6.
Datasets The QM9 dataset (Ramakrishnan et al., 2014) contains 134 thousand small- and medium- sized organic molecules with up heavy atoms, and up to 29 when counting hydrogen atoms. GEOM-DRUGS (Axelrod and Gomez-Bombarede 2022) contains 430 thousand medium- and large- sized drug-like molecules with 44 atoms average, and up to maximum 181 atoms. We use the same data setup as in previous work (Hoogeboom et al., 2022; Xu et al., 2022).
4.1 Unconditional Generation
Task We sample 10 000 molecules using the stochastic sampling procedure detailed in Algorithm 2. As END is trained in continuous
we vary the number of integration steps from 50 to 1000. We repeat each sampling for 3 seeds, and report averages along with standard deviations for each metric.
Evaluation metries We follow previous work (Hoogeboom et al., 2022; Xu et al., 2023), and first evaluate the chemical quality of the generated samples in terms of stability, validity, and uniqueness (in Tables 1, 6 and 9). On QM9, we additionally evaluate how well the model learns the atom and bond types distributions by measuring the total variation between the dataset's 'ss and generated distributions, as well as the overall quality of the generated structures via their strain energy, expressed as the energy difference between the generated structure and a relaxed version thereof (in Tables 2 we additionally compute connectivity, total variation for atom types,
and 6). On and strain energy (in Tables 2 and 9). Connectivity accounts for the fact that validity easily

---

Table 1: Stability and validity results on QM9 and GEOM-DRUCH obtained over 10000 samples, with mean/standard deviation across 3 sampling runs. END compares favorably to the baseline across all metrics on both datasets, while offering competitive performance for reduced number of sampling steps. It reaches a performance level similar to that of current SOTA methods. $ denotes results obtained by our own experiments.
QM9
GEOM-DRUGS
Stability (1)
Val. /Uniq (f)
Stahility (t)
Val. / Conn. (t)
[%}
M [8]
V[S]
VxU
[S]
V V
VxC[S)
Model
Steps
Data
99.0
95.2
97.7
97.7
86.5
99.0
99.0
EDM
(Hoogeboom et ul., 2022)
1000
98.7
82.0
91.9
90.7
81.3
92.6
EDM- MBMMBRIDGE (Wu et al., 2022)
1000
92.0
92.8
GEOLDM (Xu al., 2023)
1000
98.9
89.4
93.8
92.7
84.4
99.3
15.8
50
98. 38.3
85. 85.1+..
92.3
90.7+.a
75.1
91.7
GEOBEN
98.61.1
87. :7
93.0 3
91.51.3
93.1
(Song el al., 2024)
500
98.8
88.44.2
93.4
91.82.2
81.4
93.5
1000
+..
90.9+.2
95.3
93.0+.1
85.6
92.1
50
97.62.o
77.6
90.2
89.2
84.7..9
93. 63.6
16.6g
100
98.11.0
81.9
92.1
90.91.2
85.21
93.8
56.2
EDM
250
98.3+.0
84.3
93.2
91.74.3
85.44.0
94.2=.1
61.44.6
500
98.4 98.4+.0
85.2+ +..
93.5+-2
92.2+.3
85.4+.0
94.3--2
63.4+.1
1000
98. ...0
85.3
93.5
91.91.1
85.3
91.2..1
61.22.6
50
98.61.0
84.6
92.7
91.41
87.1
84.6
66.0
100
98.8
87.4 a
94.1
92.34.2
87.24.1
87.0
73.72.4
END
250
98.9+.1
88.8 +.6
94.7+ 2
92.6+.1
87.1+.1
88.6-.2
77.4+.4
500
98. 9..
88 19.88-. +..
94.8
92.8
87.0+.a
88, 88.82.9
78. 78.6
1000
1.0
89.1
91.8
92.61.2
87.01.0
89.22.3
79.4
be increased by generating several disconnected fragments (where only the largest counts towards validity), the total variation ensures that the model properly samples all atom types, whereas the strain energy evaluates the generated geometries. On both datasets, we also report additional drug-related metrics as per Qiang et al. (2023) (in Tables 7 and 10). More details about evaluation metrics are provided in Appendix 1.
Baselines We compare END to several relevant baselines from the literature: the original EDM (Hoogeboom et al., 2022); EDM-BRIDGE (Wu et al., 2022), an improved version of EDM that adds a physics-inspired force guidance in the reverse process; GEOLDM (Xu et al., 2023), an equivariant latent DM; and GEOBFN (Song al., 2024), geometric version of Bayesian Flow Networks (Graves et al., 2024). A detailed discussion about related work can be found in Section 5.
Ablations of END In addition to baselines from the literature, we compare different ablated versions of END. As the key component to our method is the learnable forward process, the logical ablation is whether to include a learnable forward (=END), or not (=EDM). To ensure a fair comparison and that any difference in performance does not stem from an increase learnable parameters, an architectural change or the training paradigm, we implement our own EDM (Hoogeboom et al., 2022), denoted EDM It features the exact same architecture as END, the same amount of learnable parameters (i.e. through a deeper iha), and is similarly trained in continuous-time. Additionally, we provide EDM* + y, similar to EDM* but with a learned SNR (Kingma et al., 2021) 1) for each data modality (i.e. atomic types and coordinates), and END (H,p only), an ablated version of END where only the mean learned whereas the standard deviation the conditional marginal is pre-specified and derived from the noise schedule of EDM. Table 5 provides overview of the compared models.
Results on QM9 Our main results on the QM9 dataset are summarized in the left part of Tables and 2, where END is shown to significantly outperform the baseline and reach a level of performance similar to current SOTA methods, GEOLDM (Xu et al., 2023) and GEOBEN (Song et al., 2024), across stability and validity metrics. The ablation study in Appendix A.2.1 clearly reveals the benefits of learnable forward. On the one hand, the two variants of END are shown to outperform (or be on par with) all baselines across all metrics in Table 6 (in particular, in terms of the more challenging
7

---

Table 2: Additional results on QM9 and GEOM-DRUGS obtained over 10000 samples, with mean/standard deviation across 3 sampling runs. END matches better the training distributions, and generates less strained structures than baselines. denotes results obtained by our own experiments.
QM9
GEOM-DRUGS
TV (
Strain En. (L)
TV (
Strain En. (
A [10
[10
AE [kcal/mol]
A 1
[kcalmoll]
Model
Steps
Data
7.7
15.8
GEoLDM (Xu et al., 2023)
1000
1.6
1.3
10.13
10.6*
133.5
50
1 1.6
1.72.5
16. 16.4
10.5
134.2
3.
1.41.3
13.51.1
8.0g
110.910.3
EDM
250
2.8
1.34.4
12.3
6.7
98.9
2.6+.2
1.3+.4
11.7+.1
6.4+.1
96.3+0.1
1000
2.5
1.42.4
11.3
6.2+.0
92.9
1.41.1
1.91.
13.12.1
5.9
86.31. 6
100
1.11.0
1.52.2
11.1
4.5
67.91.9
END
250
1.0+.0
0.5+.1
10.3+.1
3.5+.0
58.9+.1
500
0.94.0
0.7+.1
10.0+.1
3.3+.0
56.4+.8
1000
0.9
1.0
9.7
3.0
55.0
molecule stability), and he in better agreement with the data distribution in Table (except for QED which is captured perfectly by all methods). On the other hand, observe that, with few as 100 integration steps, END is capable of generating samples that are qualitatively better than those generated by the simpler baselines in 1000 steps. few illustrative QM9-like molecules generated by END are displayed in Fig. 1.
Results on GEOM-DRUGS Our main results the more realistic SEOM-DRUGE dataset are presented in the right part of Tables I and 2, where we observe that END competitive against the compared baselines in terms of atom stability, while being slightly subpar in terms of validity. However, when accounting for connectivity (via the VxC metric), we observe that END does outperform the baseline with increased sucess rate of around 20% average, well better than the SOTA method GROLDM (Xu et al., 2023) by large margin. The generated molecules also better follow the atom types distribution of the dataset per the lower total variation, and feature better geometries than those generated by competitors as implied by the lower strain energy. As for QM9, the ablation study provided in Appendix A.2.3 illustrates the clear improvement provided by learnable forward process compared to a fixed one. In this more challenging setting, only learning the mean function yields slight decrease in performance similar agreement with the dataset in terms of the drug-related metrics in Table 10, compared to the full model. Furthermore, while each sampling step is currently slower than EDM (Table 11), the improved sampling efficiency afforded by END (in terms of integration steps) allows practical time gains on this more complex dataset. As a concrete example, samples obtained with END with only 100 steps are qualitatively better those generated by EDM* in 1000 integration steps, amounting to 3x time cut. Examples GEOM-DRUGS-llE molecules generated by END are provided in Fig.
the metrics reported in Table 9, and
Figure 1: Representative samples generated by END on QM9 (top), and GEOM-DRUGS (bottom).

---

4.2 Conditional Generation
Dataset and Setup We perform our experiments on the QM9 dataset, on 2 different tasks: composition-conditee conditioned and uubstructure-conddiooon
Table 3: On composition-cconditonee generation, cEND offers nearly perfect composition controllability. Matching refers to the % of samples featuring the prompted composition.
generation. Both tasks allow for direct validation with ground-truth properties without re- quiring expensive QM calculations, or approximations with surrogate models. In each case, we train conditional diffu- sion model as described in Section 3.3, i.e. where Fsp and ie are provided with an extra input corresponding to the con- dition. Additional details are provided in Appendix A.6.4.
Matching [%6 (1)
Model
Steps
50
69.6.0.6
100
73.0++..6
generation The model tasked to generate compound with a predefined com- position, i.c. structural isomers of given formula. The condition is specified as vector e = (e1,... CD) E Z where ca denotes the number of atoms of type d that the sample should contain. To evaluate the model, we generate 10 samples per target formula, and compute the proportion of samples that match the provided composition. Our results are provided in Table 3, where we observe that CEND sig- nificantly outperforms the baseline, and offers nearly fully controllable composition generation. Additionally, reducing the number of sampling steps has a very limited impact on the controllability. Finally, perform ablation whose results are presented in Table 8, where fixing the standard de- viation is shown lead to small decrease in performance, with respect to the full model while remaining significantly better than the baseline with fixed forward.
Task 1: composition-condiiion
CEDM
250
14.1
500
76.2+0.66
1000
75.5.0.5
50
89 89.2+0.8
100
90.1
CEND
250
91.2+0.8
500
91.5.20.
1000
91.0g0.9
Table 4: On substructure generation, CEND shows competitive performance, surpassing EEGSDE that uses an additional property predic- tor. Results are borrowed from Bao et al. (2023).
Tanimoto Sim.
Model
Steps
CEDM EEGSDE
1000
.671 .04
1000
.7501 0033
generation We adopt setup similar to that of Bao et al. (2023), and train con- ditional END, where the condition is a molecular fingerprint encoding structural information about the molecule. A fin- gerprint is a binary vector с = (ci.... CF) E {0, 1)",, where Cy is set to 1 if substructure f is present in the molecule, or to 0 if not. Fingerprints are obtained using OPENBA- BEL (O' OoBle et al., 201 We evaluate the ability of the compared models to leverage the provided structural infor- mation, by sampling conditionally on fingerprints obtained from the test set. We then compute the Tanimoto similarity between the fingerprints yielded by generated molecules and the fingerprints provided as conditional inputs. We com- pare CEND to EEGSDE (Bao al., 2023), an improved version of EDM (Hoogeboom et al., 2022), that performs conditional generation by combining a conditional diffusion model and regressor guidance. Our results are presented in Table 4, along with a handful of samples in Fig. 2. where CEND is shown to offer better controllability than the com- pared baselines, as highlighted by the higher similarity.
Task 2: gubstructure-condiiione
50
.601 4.000
100
6410..oz
CEDM
250
6663 o02
500
6694
1000
673+ 0022
50
.783+.00
100
807 o00
CEND
250
8194
500
825+ o01
1000
.828 001
Figure 2: conditioned samples, where CEND matches the provided substructure bet- ter (in terms of compositions and local patterns).
Excerpt of substructure-
Target
<DDM
<END
Related Work
The main approaches to molecule generation in 3D are auto- regressive models (Gebauer et al., 2019; Simm et al., 2020; Gebauer et al., 2022; Luo and Ji, 2022; Daigavane et al., 2024), flow-based models (Garcia Satorras et al., 2021), and diffusion models (Hoogeboom et al., 2022; Igashov et al., 2024). A notable exception to the geometric graph representation of 3D molecules are voxels (Skalic et al., 2019; Ragoza et al., 2022; O Pinheiro

---

et al., 2024), from which the 3D graph is extracted via some post-processe processing procedure. Recently. several works have shown that leveraging 2D connectivity information can lead to improved results (Peng et al., 2023; Vignac et al., 2023; Le et al., 2024). While not incompatible with END, we perform our experiments without modeling that auxiliary information, and therefore do not compare to these approaches directly. Other frameworks have also been tailored to molecule generation, such as Flow Matching (Lipman et al., 2022: Song et al., 2023; Irwin et al., 2024) or Bayesian Flow Networks (Graves et al., 2024; Song et al., 2024), also showing promises for accelerated sampling.
In the realm of diffusion models for molecules, et al., 2023) extend upon a continuous- time formulation of EDM (Hoogeboom et al., 2022), as END also does. Based on the observation that there exists an infinity of processes mapping from prior to target distributions, prior knowledge, i.e. part of the drift term is physically-inspieed force term. END can be seen as a generalization of EDDM-BRIGEE where the forward drift term is now learned instead of pre- specified. Through experiments, we show that a learnable forward performs better than fixed one, even when the latter is physics-inspired. EEGSDE specifically targets conditional generation by combining (1) a conditional model score model, (2) a method similar to classifier -eefferr- (requiring the training of auxiliary model). In CEND, we instead only learn conditional model. Finally, GEOLDM (Xu et al., 2023) is a latent diffusion model that performs diffusion in the latent space of an equivariant Variational
(Wu et al., 2022) and EEGSDE (Bao
constructs one such process that incorporates some
(VAE), and it can be seen as particular case of END, where (E t, x) = (x) + ETG, with Ex(x) denoting the dime-independent encoder of the VAE.
6 Conclusion
In this work, we have presented Equivariant Neural Diffusion (END), a novel diffusion model that is equivariant to Euclidean transformations. The key innovation in END lies in the forward process that is specified by learnable data- and time-dependent transformation. Experimental results demonstrate the benefits of our method. In the unconditional setting, we show that END yields competitive generative performance across two different benchmarks. In the conditional setting, END offers improved controllability, when conditioning on composition and substructure. Finally, as by-product of the introduced learnable forward, we also find the sampling efficiency (in terms of integration steps) to be improved, while that property is not actively optimized for in the design of the parameterization nor in the training procedure.
Avenues for future work are numerous. In particular, further leveraging the flexible framework of NFDM (Bartosh et al., 2024) constrain the generative trajectories, e.g. to be straight and enable even faster sampling. modelling bond information, or extending the conditional setting to other types of conditioning information, e.g. other point cloud or target property, are all promising directions.
Limitations From a computational perspective, END is currently slower to train and sample from compared to baseline with fixed forward with identical number of learnable parameters. Even convergence is similarly fast, each training (resp. sampling) step incurs a relative z 2.5x (resp. 3x). However, END usually requires much fewer number of function evaluations to achieve comparable (or better) accuracy, and alternative (and more efficient) parameterizations of the reverse process exist. In particular, the drift fe.p could be learned without direct dependence on fu. thereby leading to an improved training time and, more importantly, a very limited overhead with respect to vanilla DMs for sampling. In terms of scalability, END suffers from limitations similar to concurrent approaches. It operates on fully- culle-co graphs preventing its scaling to very large graphs, and models categorical features through an arbitrary continuous relaxation potentially suboptimal and scaling linearly with the number of chemical elements in the modeled dataset. Encodings that scale more gracefully, such as that of Analog Bits (Chen et al., 2023) (logarithmic) or Geol.DM (Xu et al., 2023) (learned low-dimensional), are good candidates to better deal with discrete features within END. In terms of data, the presented findings are limited to organic molecules and the metrics, while widely used in the community, also have some evident limitations. To fully assess the practical interest of the generated molecules, thorough validation with QM simulations would be required.
Broader Impact Generative models for molecules have the potential to accelerate in-silico discovery and design of drugs or materials. This work proposes an instance of such model. As any generative model, it also comes with potential dangers, as this could be misused for designing e.g. chemicals with adversarial properties.
10

---

Acknowledgments and Disclosure of Funding
FC acknowledges financial support from the Independent Research Fund Denmark with project DELIGHT (Grant No. 0217-00326B).
References
Amil Merchant, Simon Batzner, Samuel S Schoenholz, Muratahan Aykol, Gowoon Cheon, and Ekin Dogus Cubuk. Scaling deep learning for materials discovery. Nature, 624(7990):80-855 2023.
Lars Ruddigkeit, Rund Van Deursen, Lorenz C Blum, and Jean-Louis Reymond. Enumeration of 166 billion organic small molecules in the chemical universe database gdb-17. Journal of chemical information and modeling, 52(1 2879644-777, 2012.
Dylan M Anstine and Olexandr Isayev. Generative models as an emerging paradigm in the chemical sciences. Journal of the American Chemical Society, 145(16):8736-855)) 2023.
Niklas Gebauer, Michael Gastegger, and Kristof Schüu,. Symmetry-adapted generation of 3d point sets for the targeted discovery of molecules, Advances in neural information processing systems, 32, 2019.
Niklas WA Gebauer, Michael Gastegger, Stefaan SP 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.
Youzhi Luo and Shuiwang Ji. An autoregressive flow model for 3d molecular geometry generation from scratch. In International Conference on Learning Representations, 2022. URL https: www.arrewwwwe..ettooemeeaaaae wwweet/fooummicom www.oooommddssceeea....
//openreview.net wet.eeewweee//eeeeeeeeeeee
Ameya Daigavane, Song Eun Kim, Mario Geiger, and Tess Smidt. Symphony: Symmetry-equivariant point-centered spherical harmonics for 3d molecule generation. In The Twelfihh International Conference on Learning Representations, 2024. URL https: ptttss:/oppa http://ppppeeeeeeeeeeeeeeeeeeeet net/ wwt.:orummteeee.ccoo ffrmmfrfiiiid MIEnYt1Gyv.
Emiel Hoogeboom, Victor Garcia Satorras, Clément Vignac, and Max Welling. Equivariant diffusion for molecule generation in 3d. In International conference on machine learming, pages 8867-8887 PMLR, 2022.
Clement Vignac, Nagham Osman, Laura Toni, and Pascal Frossard. Midi: Mixed graph and 3d denoising diffusion for molecule generation. In Joint European Conference on Machine Learning and Knowledge Discovery in Databases, pages 560-576. Springer, 2023.
Tuan Le, Julian Cremer, Frank Noe, Djork-Arné Clevert, and Kristof T Schütt. Navigating the design space of equivariant diffusion-based generative models for de novo 3d molecule generation. In The Twelfth International Conference on Learning Representations, 2024. URL https://
openreview. / d=kzGuiRXZrd
Oliver T Unke, Stefan Chmiela, Huziel E Sauceda, Michael Gastegger, Igor Poltavsky, Kristof T Schuett, Alexandre Tkatchenko, and Klaus-Robert Mueller. Machine learning force fields. Chemi- cal Reviews, (21(11))1142--0018 2021.
Kristof Schütt, Pieter-Jan Kindermans, Huziel Enoc Sauceda Felix, Stefan Chmicla, Alexandre Tkatchenko, and Klaus-Rober Müller. Schnet: continuous-fin convolutional neural network for modeling quantum interactions. Advances in neural information processing systems, 30, 2017.
Kristof Schütt, Oliver Unke, and Michael Gastegger. Equivariant message passing for the prediction of tensorial properties and molecular spectra. International Conference on Machine Learning, pages 9377-9388. PMLR, 2021.
Simon Batzner, Albert Musaelian, Lixin Sun, Mario Geiger, Jonathan P Mailoa, Mordechai Kombluth, Nicola Molinari, Tess Smidt, and Boris Kozinsky. E (3)-equivari graph neural networks for data-efficient and accurate interatomic potentials. Nature communications, 13(1):2453, 2022.
11

---

Ilyes Batatia, David P Kovacs, 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, 35:1 1423-11436, 2022.
Ling Yang, Zhilong Zhang, Yang Song. Shenda Hong, Runsheng Xu, Yue Zhao, Wentao Zhang, Bin Cui, and Ming-Hsuan Yang. Diffusion models: A comprehensive survey of methods and applications. ACM Computing Surveys, 56(4):1-39, 2023.
Grigory Bartosh, Dmitry Vetrov, and Christian A Naesseth. Neural diffusion models. arXiv preprint 0rX00:2310.08337, 2023.
Beatrix Miranda Ginn Nielsen, Anders Christensen, Andrea Dittadi, and Ole Winther. Diffenc: Variational diffusion with learned encoder. In The Twelfth International Conference on Learning Representations, 2024. URL
Grigory Bartosh, Dmitry Vetrov, and Christian Naesseth. Neural flow diffusion models: Learnable forward process for improved diffusion modelling. arXiv preprint arXiv:2404. 12940, 2024.
Victor Gareia Satorras, Emiel Hoogeboom, Fabian Fuchs, Ingmar Posner, and Max Welling. E. (n) equivariant normalizing flows. Advances in Neural Information Processing Systems, 34:4181-4192, 2021.
Minkai Xu, Lantao Yu, Yang Song, Chence Shi, Stefano Ermon, and Jian Tang. Geodiff: A geometric diffusion model for molecular conformation generation. International Conference on Learning Representations, 2022. URL https:/
net net/forun' PZcvxEK2vCCC
Fan Bao, Min Zhao, Zhongkai Hao, Peiyao Li, Chongxuan Li, and Jun Zhu. Equivariant energy- guided SDE for inverse molecular design. In The Eleventh International Conference on Learning Representations, 2023. URL htRL:tt::// https://opeereviiee
Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In Francis Bach and David Blei, editors, Pro- ceedings of the 32nd International Conference on Machine Learning, volume 37 of Proceedings of Machine Learning Research, pages 2256-2265, Lille, France, 07-09 Jul 2015. PMLR. URL html.
mlr. wtw...press/v3///// httpe///////ooo--icccteeeee
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840-68511 2020.
Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations, 2020.
Simo Särkkä and Arno Solin. Applied stochastic differential equations, volume 10. Cambridge
University Press, 2019.
Brian DO Anderson. Reverse-time diffusion equation models. Stochastic Processes and their Applications, 12(3):3133326, 1982.
Pascal Vincent. A connection between score matching and denoising autoencoders. Neural computa- tion, 23(7):1661-1674 2011.
Leon Klein, Andrew Foong, Tor Fjelde, Bruno Mlodozeniec, Mare Brockschmidt. Sebastian Nowozin, Frank Noé, and Ryota Tomioka. Timewarp: Transfcrable acceleration of molecular dynamics by learning time-coarsened dynamics. Advances in Neural Information Processing Systems, 36, 2024.
Bowen Jing, Ezra Erives, Peter Pao-Huang, Gabriele Corso, Bonnie Berger, and Tommi S Jaakkola. Figenfold: Generative protein structure prediction with diffusion models. In ICLR 2023-Machine Learning for Drug Discovery workshop. 2023.
Ross Irwin, Alessandro Tibo, Jon Paul Janet, and Simon Olsson. Efficient 3d molecular generation with flow matching and scale optimal transport. In ICMI. 2024 AI for Science Workshop, 2024. URL https://openteii hhtp:://///eeeeieee/ wttp:/eteeeooomm/eeeeeaaaaaaaea
12

---

Hojung Jung, Youngrok Park, Laura Schmid, Jachyeong Jo. Dongkyu Lee, Bongsang Kim, Se- Young Yun, and Jinwoo Shin. Conditional synthesis of 3d molecules with time correction sampler. The Thirty-eighth Annual Conference Neural Information Processing Systems, 2024.
Luhuan Wu, Brian Tripee, Christian Naesseth, David Blei, and John P Cunningham. Practical and asymptotically exact conditional sampling in diffusion models. Advances in Neural Information Processing Systems, 36, 2024.
Raghunathan Ramakrishnan, Pavlo O Dral, Matthias Rupp, and Anatole Von Lilienfeld. Quantum chemistry structures and properties of 134 kilo molecules. Scientific data, 1(1):177, 2014.
Simon Axelrod and Rafael for property prediction and molecular generation. Scientific Data, 9(1):185, 2022.
Geom, energy-annotatte molecular conformations
Minkai Xu, Alexander S Powers, Ron O Dror, Stefano Ermon, and Jure Leskovec. Geometric latent diffusion models for 3d molecule generation. In International Conference on Machine Learning, pages 38592-38610. PMLR, 2023.
Bo Qiang. Yuxuan Song, Minkai Xu, Jingjing Gong, Bowen Gao, Hao Zhou, Wei- Yei-Ying Ma, and Yanyan Lan. Coarse-to-fine: hierarchical diffusion model for molecule generation in 3d. In International Conference on Machine Learning, pages 28277-28299, PMLR, 2023.
Lemeng Wu, Chengyue Gong, Xingchao Liu, Mao Ye, and qiang liu. Diffusion-based molecule generation with informative prior bridges. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors. Advances in Neural Information Processing Systems, 2022. URL https://openreviie net/fooummidd
Yuxuan Song, Jingjing Gong, Hao Zhou, Mingyue Zheng, Jingjing Liu, and Wei- Ying Ma. Unified generative modeling of 3d molecules with bayesian flow networks. In The Twelfth International Conference on Learning Representations, 2024. URL https: http////ppeeeeeeee//eee/ee net/forumiid= NSVtmmzeRB.
Alex Graves, Rupesh Kumar Srivastava, Timothy Atkinson, and Faustino Gomez. Bayesian flow networks, 2024.
Diederik Kingma, Tim Salimans, Ben Poole, and Jonathan Ho. Variational diffusion models. Advances in neural information processing systems, 34:21696-21707, 2021.
Noel M O' Boyle, Michael Banck, Craig A James, Chris Morley, Tim Vandermeersch, and Geoffrey R Hutchison. Open babel: An open chemical toolbox. Journal of cheminformatics, 3:1-14, 2011.
Gregor Simm, Robert Pinsler, and José Miguel Hernández-Lobato. Reinforcement learning for molecular design guided by quantum mechanics. In International Conference on Machine Learning, pages 8959 8969. PMLR, 2020.
Ilia Igashov, Hannes Stärk, Clément Vignac, Ame Schneuing, Victor Gareia Satorras, Pascal Frossard, Max Welling, Michael Bronstein, and Bruno Correia. Equivariant 3d-conditional diffusion model for molecular linker design. Nature Machine Intelligence, pages 1-11, 2024.
Miha Skalic, José Jiménez, Davide Sabbadin, and Gianni De Fabritii.. Shape-based generative modeling for de novo drug design. Journal of chemical information and modeling, 59(3): 1214, 2019.
Matthew Ragoza, Tomohide Masuda, and David Ryan Koes. Generating 3d molecules conditional on
receptor binding sites with deep generative models. Chemical science, 13(9):2701-2713 2022.
Pedro O O Pinheiro, Joshua Rackers, Joseph Kleinhenz, Michael Maser, Omar Mahmood, Andrew Watkins, Stephen Ra, Vishnu Sresht, and Saeed Saremi. 3d molecule generation by denoising voxel grids. Advances in Neural Information Processing Systems, 36, 2024.
Xingang Peng, Jiaqi Guan, Qiang Liu, and Jianzhu Ma. Moldiff: Addressing the atom-bond inconsistency problem in 3d molecule diffusion generation. In International Conference on Machine Learning, pages 27611-1 1-27229. PMLR, 2023.
13

---

Yaron Lipman, Ricky TQ Chen, Heli Ben- Gem-Hamu, Maximilian Nickel, and Matthew Le. Flow matching for generative modeling. The Eleventh International Conference on Learning Representations. 2022.
Yuxuan Song, Jingjing Gong, Minkai Xu, Ziyao Cao, Yanyan Lan, Stefano Ermon, Hao Zhou, and Wei- Ying Ma. Equivariant flow matching with hybrid probability transport for 3d molecule generation. In A. Oh, T. Neumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine, editors, Advances in Neural Information Processing Systems, volume 36, pages 549-568. Curran As- sociates, Inc., 2023. URL https::/ neurips oc oc/paper fecaper_filee www.s//ppppe//////... #1220006667788113333666666771712
pdf.
Ting Chen, Ruixiang ZHANG, and Geoffrey Hinton. Analog bits: Generating discrete data using diffusion models with ceff-conditioning. In The Eleventh International Conference on Learning Representations, 2023.
Arthur Gretton, Karsten M Borgwardt, Malte J Rasch, Bernhard Schölkopf, and Alexander Smola. A kernel two-sample test. The Journal of Machine Learning Reseurch, 13(1):723-773, 2012.
Greg Landrum et al. Rdkit: software suite for cheminformatics, computational chemistry, and predictive modeling. Greg Landrum, 8:31, 2013.
David Weininger. Smiles, chemical language and information system. 1. introduction methodol- ogy and encoding rules. Journal of chemical information and computer sciences, 28(1):31-36, 1988.
Peter Ertl and Ansgar Schuffenhauer. Estimation of synthetic accessibility score of drug-like molecules based on molecular complexity and fragment contributions. Journal of cheminfor- matics, 1:1-11, 2009.
Tuan Le, Frank Noe, and Djork-Arn Clevert. Representation learning on biomolecular structures using equivariant graph attention. In The First Learning on Graphs Conference, 2022. URL https::/openreetee net htte/forumiidkkv4xooPPP
14

---

A Appendices
A.1 Compared models
In Table 5, we detail all the compared diffusion models in terms of their transformation Fep.
Table 5: Compared models.
Comment
a. exp (- B(s) ((s) ds) exp 2 So ((s) da)
EDM (Hoogeboom et al., 2022) / EDM
0, + 0.E
and similar EDM p( N (z))) 671)
GEoLDM (Xu et al., 2023)
ax (x)) + oie
EDM 7/e
ag (t)x + op(()e
learned Y/p with outputs (r h)
END (jip only)
(o( dte
de similar EDM
END
(f( () + Up (2, l)e
introduced in Eq. (6)
A.2
Additional results
A.2.1
Ablation unconditional QM9
Table 6: Main results of the ablation study on QM9. Metrics are obtained over 10000 samples, with mean/standard deviation across sampling runs. The two variants of END compare favorably to baselines across metrics, while offering competitive performance for reduced number of sampling steps.
Stability (1)
Validity Uniqueness (t)
TV
Striin En. (L)
A [%]
M [%)
V [56]
[8]
A [10 3]
B [ 91
AE [kcal/mol]
Model
Steps
Data
99.0
95.2
97.7
97.7
7.7
50
97.62.0
77.61..
90.21.2
89.22.2
4.6
1.74.s
16.42.2
100
18.1+.0
81.9+.4
02.1+.2
90.9+.2
3.54.1
1.4.3
13.5+.1
EDM
250
98.3+..
84.3
93.2 93.21.4
91.71.3
2.8
1.3
12.34.4
500
98.41.0
85.21.5
93.31.2
92.21.3
2.6
1314
11.71.1
1000
98.44.0
85.31.3
93.54.1
91.92.1
2.5
1.4..4
11.3..1
97.7+.o
77.4+.3
91.14.4
90.24.4
4.34.1
1.5.2
15.54.2
100
98.24.a
82.64.2
92.9
91.64.2
3.2
1.24.2
12.8
EDM" + 7p
250
98.52.0
85.34..
93.91.1
92.41.1
252
1.0..1
11.32.
500
98.51.1
86.1
94.11.2
92.51.
2.2
1.04.3
11.1
1000
98.54.0
86.1
94.1
92.4+.2
2.1
1.4.11
10.8
50
98.5+.a
83.9
96.2+.2
93.8+.3
144.1
1.94.
13.1+.1
100
98.7
87.0
95.31.2
93.61.2
1.1.0
1.5..2
11.1..
END (M only)
250
98.91.
89.01.2
95.81.2
93.81.2
1000
0.5
10.3
500
98.94.0
88.64.a
95.64.1
93.5
0.9
0.7
10.0
1000
98.94.0
89.24.3
95.64.1
93.5+.1
9.7-22
98.6
84.64.1
92.7
91.4.1
1.2.1
1.94.
12.1
100
98.8
87.41.2
94.11.0
92.31.2
1.30
1.81.3
10.61.2
END
250
98.9
88.81.5
94.71.2
92.61.1
1.2
0.8
9.6_
500
98.94.0
88.8
04.84.2
92.8+.2
1.2..1
0.84.5
9.5-.1
1000
98.9
89.1
94.8
92.6t.2
1.2
0.8+.5
9.3
15

---

Table 7: Additional results of the ablation study with metrics from HIERDIFF (Qiang et al., 2023) and MMD (Gretton et al., 2012: Daigavane et al., 2024) on QM9, The two variants of END shows better, on par, agreement with the true data distribution compared the baselines.
SA (t)
QED
logP (t)
MW
MMD (J) [10-1]
Model
Steps
Data
0.626
0.462
0.339
122.7
7.7
50
0.6094.001
0.45644000
0..09 .08
125.7=.0
1.91
100
0.613+.001
0.458+4.00
0.049
124.9_.1
1.67
EDM
250
0.6174.000
4.01
0.107 0103
124.4=.0
1.52+.02
500
0.618
0.462
0.124
124.2
1.50+.04
1000
0.619 0000
0.462 .001
0.135
124.2
1.51+.02
50
0.612 000
0.454
0.053 .003
125.3_.0
2.04
100
0.616+.0o0
0.459
0.0-99+.010
124.6=.u
1.664.u1
EDM* + %
250
0.620
D.4614.cmm
0.1244..11
124.1=.0
1.544.u2
500
0.620+.001
0.461+.000
0.1444:
124.0
1.50#.02
1000
0.6221. 000
0.462:
0.1621..00
123.9
1.45..03
50
0.606+.001
0.456++.001
0.096+.008
123.8
2.80#.06
100
0,615
0.458+.000
0.157
123.6_.1
1.97
END (Hip only)
250
0.622 002
0.460+.o01
0.1984.o06
123.3_.1
1.48
500
0.6264.001
0.462+.000
0.219
123.1
1.36..03
1000
0.021
0.46244.001
0.225
123.0
1.36
50
0.602+.001
0.456+.001
0.074+.012
123.7
1.91
100
0.6134.001
0.4594.000
0.125+.010
123.3_.1
1.632.02
END
250
0.620
0.461
0.1644.
123.1
1.44..or
500
0.622 o02
0.462
0.193+.00s
123.2
1.41+.01
1000
0.623 4.01
0.462+.ou1
0.198+.u14
123.0
1.37
16

---

A.2.2 Ablation on composition-conditooned Q.M9
Table 8: Ablation on cepsstiionnnonitionned generation. The two versions CEND display better controllability. Fixing the standard deviation leads small decrease performance, with respect to the full model while remaining significantly better than the baseline with fixed forward.
Matching [%] (1)
Model
Steps
50
69.6-0.6
100
73.0
CEDM
250
14.1
500
76.2
1000
75.5
50
75.7=0.4
100
79.9
CEND (Mg only)
250
-2.7
500
83.
1000
03.5
50
89.2 0.8
100
90.1
CEND
250
91.2-008
500
91.5-0.8
1000
91.0_0.9
A.2.3 Ablations on GEOM-DRUGE
Table 9: Main results of the ablation study GEOM-DRUGS. Metries obtained over 10000 samples, with mean/standard deviation across 3 sampling runs. Most notably, END generates more connected samples, and less strained structures. Fixing the standard deviation leads to slight decrease performance.
Stability (t)
Val. / Conn. (t)
TV (L) A [10-2]
Strain En. (L) AE [kcal/mol]
A [%]
V [%]
VxC
Model
Steps
Data
86.5
99.0
99.0
15.8
50
84.74..
93.64.2
46.6
10.5
134.2
100
85.2
93.8+.3
56.21.4
8.0
110.94.3
*DM*
250
85.4+.0
94.2+.1
61.44.6
6.7
98.92.3
500
85.4+.0
94.3+.2
63.4
6.4
95.3
1000
85.3
94.44.11
64.24.6
6.24.0
92.941.1
50
85.64.1
87.8
66.04.4
7.94.
105.6
100
85.81.1
89.9
73.71.4
6.1
85.5.0.5
END (Mip only)
250
85.71.1
91.21.2
77.4
5.0
74.5
500
85.81.1
91.61.1
78.61.3
4.8
72.3
1000
85.8+
91.8+.1
79.4+.4
4.6+.0
71.0+0.6
50
87.1+.1
84.6
68.64.4
5.9+.1
86.3
100
87.2+.1
87.0+.2
76.74.5
4.5+.1
67.9..9
END
250
87.1+.1
88.5+.2
80.7+.6
3.5+.0
58.9+.1
500
87.0+.0
88.84.3
81.74.4
3.3+.0
56.4
1000
87.0+.0
89.2+.3
82.54.3
3.0+.0
55.0+.5
17

---

Table 10: Additional results of the ablation study with metrics from HIERDIFF (Qiang et al., 2023) on GEOM-DRUGS, The two variants of END shows better agreement with the true data distribution, compared to the baseline with fixed forward.
SA (t)
QED (t)
logP (1)
MW
Model
Steps
Data
0.832
0.672
2.985
360.0
50
0.590+.00
0.480+.001
1.105+.012
354.240.9
100
0.614+.002
0.5344.001
1.482+.02 .022
352.340.4
EDM*
250
0.630+.001
0.559+.003
1.7164.00 .009
351.440.4
500
0.637+.00
0.570+.003
1.794
350.740.3
1000
0.641 .002
0.5744.004
1.831 1.311.001 .013
350.5+1.4
0.634 0001
0.526+.001
1.435+.00 0007
352.440.6
100
0.664
0.568+.001
1.792+.020
351.340.5
END (M4 only)
250
0.681+.000
0.591+.002
1.996+.014
351.140.6
500
0.687 1.000
0.596
2.059 1.015
350.6
1000
0.690
0.602
2.0931.019
351.6 351.610.4
50
0.621
0.487
0.939
352.0+1.4
100
0.660+.001
2.050
1.530+.010
351.641.9
END
250
0.685+.001
0.578
1.832+.009
351.411.
500
0.695+.002
0.5864.004
1.945+.012
352.141.6
1000
0.698
0.590
2.0094.009
352.010.8
18

---

A.3 Equivariance / invariance proofs
A.3.1 Time-derivative of -uuuivrarrtt function
Let f x [0.1] 3/ be function that is equivariant to actions of the group 0(3), such that R f(x, t) = f(R æ, t), VR E 0(3).
Proof sketch
We need to show that & (S(R- x,))) R (f(x VR O(3) and Vo X.
a ((RR,,,))
0
R f(x, I))
ōt
0
R 0t
(f(z,t) t))
where the last equality follows by linearity.
A.3.2 Inverse of an -uu)-equiv function
Let x + y be a function that (1) equivariant to the action of the group 0(3), and (2) admits an inverse ;-1:y : + X,, then f-1 is also equivariant to the action of O(3).
Proof sketch We need to show that f y) R- (-), VR E 0(3) and Vy E y.
Since is invertible, we have that to any y y corresponds a unique a E X, such that y = f(x). and that f (y) ff())( x. As f is equivariant to the action of 0(3), we have that R f(x)(x fRR YRE
f y) (R. f(x))
= f x)).
R.
R f-)
A.3.3 inva-inve of the objective function
In this section, we show that the objective function invariant under the action of 0(3):
LEND (Ræ)) (æn)(x), VR O(3), provided that Ep and &o are equivariant.
19

---

(1 f (Z,, Ra) (t(()
CEND( = (2
Vx, log 2o (z.
2
,, 20 1)) + g (t)
Vz, log 9pp (2. (2), 1))
2
g2(1
(x
f (RR 2,, t, R.x)
xx log qup (RR 1R: (x(
(1)
2
- foRRR ze, Co (RR Zz, t)) + g6
Vz, log 9p
(RR -1 2e L (RR z2,
(RR-122- 1))
2
(4)
(R))
R R 1 2,, x.)
log qo (R (2(x))
2g
2
(t)
-Rf &, 12c t):
V., log qu (R
zz 2 -ze, 1))1
2
g3(1)
(Ry, (Ry, Rax))
(t)
RR t, t,
2
RV y log (V(l))
(1(
(y, t, (y t)) +
RV v log 2o (yr (3 t) 2
2
1
(t)
(m(
|1 (y),t,
Vz. log (v)
(1)
2
(1())
folus, 2o (v).
V. v. log (vs +)]112]
2
(x)
The first equality is obtained by replacing x by R. in the definition of the objective function Eq. (5). The second is obtained by multiplying by R.R and ãe are equivariant. We then perform a change of variable y: R zt. As rotation does preserve distances, we obtain the last equality.
The third equality by leveraging that for 9
20

---

A.4 Algorithms
Algorithm Training algorithm of END
Require: q(x), Fi.. Ze for training iterations do
~ ((=), ~ u(t), E ~ p(E)
2, - (2 + tU(,,tt)
C (r) 1 (=(,tt fe (Z 1)
Gradient step on and end for
Algorithm 2 Stochastic sampling from END
Require: Fs, Leg integration steps T, empirical distribution of number of atoms (N)
At = T.
N P(N)
21 ~ p(21)
for 1 1,... T
do
1 ~ I),
21-At c 2t (Z () go(()vv A
end for p())
21

---

A.5 Details about Fip
A.5.1
Working in ambient space
In this section, we omit the invariant features, and consider that æ; _ e RRMx are all collections of vectors.
In opposition to the notations introduced Section 3, we would like F& and &g (which ultimately neural networks) operate in ambient space directly, i.e. on "; R with R {u M xd M + v; = 0}. instead R(M--1xx as initially presented. To do so, we use the results from Garcia Satorras et al. (2021) that showed that the Jacobian of the transformation F'p can be computed in ambient space, for 2; and & that live in the linear subspace R. provided that the transformation F3 (invertible with respect to leaves the center of mass of e unchanged.
Considering flat representations 2t, n e E M-d the parameterization of F.p introduced in Eqs. (6), (8) and (9) can be adapted as follows to achieve such property
Fe t, = (a, t) + t2,
(10)
Mo(a,t) t) (I 1
TT TTTTT/L )i, (2x
(11)
M
1
TT (x, 1 TT. M
Up(2, t) = (I
(12)
M
where T E R-Waxdd = [[., Ia. Ia....J,, and 1 TT corresponds to the linear operator computing the center of mass. In Eq. (11), the unconstrained mean output Pp(xz, t) simply projected onto the O-CoM subspace, thereby inducing no translation. In Eq. (12), the unconstrained output Ü (x, t) is first projected onto the 0-CoM subspace, before being translated back to the initial center of mass.
M
With this adapted formulation, we need be able to (1) compute the latent variable 2; given E, (2) evaluate the Jacobian of the transformation F,,, and (3) evaluate the inverse transformation F-
Computing 2, from E Given that Rt obtaining ,, E R from e simply amounts to (1) computing 2, Ti (xr, 1) + (2, l)e, and then (2) removing the center of mass from ε, the second term in Eq. (12) induces no translation as e E R.
Computing Jr,, and Fi what follows, we shorten the notation, and denote (z, t) by U and (z( t) by U. To leverage known identities, start by reorganizing Eq. (12),
-= 1 TT (I Ú).
(13)
M
The Jacobian of F. is given by the determinant of U, and the latter can be derived by leveraging the Matrix Determinant Lemma,
det U = det Ū det (Ia + T (IM -
(14)
M
det Ū det (La + T (Ū-1(01111 -lM))))
(15)
M
A
det Ū det 1
Σ (1))))1).
(16)
M m=1
det - V,
(17)
M
II det 1m det V.
(18)
where Óm denotes the m-th d-block in Ü, and + EM (Ūm)-1 is a d-matrix. Computing the Jacobian therefore amounts to compute the inverse of M d x d-matrices and the determinant of M + d d-matrices. In practice, 3 such that all computations can be performed in closed-form.
22

---

Regarding F we do not need to evaluate the inverse transformation itself, but instead evaluate given 2/,
= - (2) 1a t))
The inverse U-1 can be obtained via the Woodbury matrix identity,
U- (Ü TT
a
M
= Ú- ---TV-'T - llM.)).
-1)))
M
Ú (I C)
where = Σm= (Úm)-1 previously defined in Eq. (18), and C' TV (0-1- -11--MMd LM-a). Given the specific structure of C, the computation of & can simplified to
U-! 211
(I-
Ü [c--
where = = p(x, t)) and Ca. acts as translation operator. We that computing the inverse transformation requires to invert M + 1 d d-matrices, but as d = 3 in practice, all computations can be performed in closed-form.
A.5.2 Invariant features
For simplicity, we omitted in Section 3 and Appendix A.5.1 that molecules are described as tuples: = (r. h), as only transform under Euclidean transformations. For the invariant features h,
use the following parameterization
(x, = (1- f(1 t) (h (xx t),
(19)
(1) (x, = 61-10 (,, 1)((1-1)
(20)
(h) which ensures that Mip (x. 0) h and 0ip
(h) (x, 0) 8; whereas Hig
(h) (h) (xr, 1) 0 and de (z. 1)
II.
Implementation As described in the main text in Section 3.2, Fg implemented as neural network with an architecture similar to that of the data point predictor Ca(2,, f), but with a specific readout layer that produces the outputs related to r (Ja t), σ (z( f)). Additionally, (x, produces P (x, t) and o (h) (x, f) as invariant outputs.
Inverse transformation writes
The logarithm of the determinant of the inverse transformation log JJ
MxD
log |J = log |Jp|
Σ
(b), (x,
log
du (x, t)) - log det [V], (21)
invariant features
vectorial features
where defined as Eq. (18).
23

---

A.6
Experimental details
In addition the details provided in this section, we release a public code repository with our implementation of END.
A.6.1 Evaluation metries
In this section, we describe the metrics employed to evaluate the different models:
Stability: An atom is deemed stable if it has charge of O, whereas molecule is stable if all its atoms have 0 charge. We reuse the lookup table from Hoogeboom et al. (2022) to infer bond types from pairwise distances.
Validity and Connectivity: Validity corresponds to the percentage of samples that can be parsed and samitized by rdkit (Landrum et al., 2013), after inference the bonds using the lookup table mechanism (Hoogeboom et al., 2022). It should be noted that the metric does not penalize fragmented samples as long as each individual fragment appears valid. larger compounds such those in
This can be problematic when running evaluation GEOM-DRU as models tend to generated disconnected structures. To account for that, we also report Connectivity, which simply check that valid molecule is composed of exactly one fragment.
Uniqueness: Uniqueness expressed as the proportion of samples that are valid and have unique SMILES string (Weininger, 1988) among all the generated samples. On GEOM-DRUGS, we do not report Uniqueness as all generated samples appear unique (as per previous work).
The total variation is computed the MAE between the (discrete) marginal obtained on the training data and on the generated samples. For bond types on QM9, we compute the ground truth and generated distributions using the lookup table mechanism (Hoogeboom et al., 2022).
Total variation:
The strain energy is expressed as the difference in energy between generated structures and relaxed version thereof obtained as per rdkit's MMFF (Landrum et al., 2013). From the generated samples, we infer rdkit mol objects using OPENBABER (O' (O'Boyle et al., 2011). We only evaluate the strain energy of valid and connected samples.
Strain Energy:
SA, QED, log P and MW: SA denotes the "Synthetic Accessibility Score", which is rule-based scoring function that evaluates the complexity of synthesizing a structure by organic reactions (Ertl and Schuffenhauer, 2009). We normalized its values between and 1, with 0 being "difficult to synthesize" and "easy to synthesize", QED denotes "Quantitative Estimation of Drug-likeness" log P denotes the octanol-water partition coefficient. MW denotes the molecular weight.
We employ the rdkit's implementation of all metrics. To do so, we convert the generated samples to rdkit mol objects using OPENBABEL (O' Boyle et al., 2011). We then only evaluate the different metrics for valid and connected samples.
MMD: On QM9, we follow the procedure of Daigavane et al. (2024), and compute the MMD Gretton et al. (2012) between true and generated pairwise distances distributions for the 10 most common bonds in the dataset: ["C-H:1.0",, "C-C:1.0"" "C-O:1.0", "C-N:1.0", "H-N:1.0", "C-O:2.0", "C-N:1.5", "H-O:1.0", "C-C:1.5", "C-N:2.0").
A.6.2 Architecture
Our forward transformation F6 and data point predictor ft share common neural network architec- ture that we detail here. The architecture is similar to that EQCAT (Le et al., 2022), and updates a collection of invariant and equivariant features for each node in the graph. We choose that architecture because it allows for an easy construction of U, (x, t) by linear projection of the final equivariant layer.
We follow previous work (Hoogeboom al., 2022) and consider fully-connected graphs. We initially featurize pairwise distances through Gaussian Radial Basis functions, with dataset-specific cutoff taken large enough to ensure full connectivity. In opposition to Hoogehoom et al. (2022), we do not update positions in the message-passing passing phase, but instead obtain the positions prediction through
24

---

linear projection of the final equivariant hidden states. The predictions for the invariant features are obtained by reading out the final invariant hidden states.
Optimization For all model variants, we employ Adam with learning rate of 10-4, We perform gradient clipping (norm) with a value of 10 on QM9, and a value of on GEOM-DRUGS.
A.6.3 Unconditional generation
We reuse the data setup from previous work (Hoogeboom et al., 2022; Xu et al., 2023).
QM9 QM9, we use 10 layers of message passing for EDM*, while the variants of END feature 5 layers of message-passing in F and 5 layers in it. For all models, we use 256 invariant and 256 equivariant hidden features, along with an RBF expansion of dimension 64 with cutoff of 12Á for pairwise distances. This ensures that the compared models have the same number of learnable parameters, i.c. 9.4M each. We train all models for most 1000 epochs with a batch size of 64.
GEOM-Drugs On REEM--DRUGE we use 10 layers of message passing for EDM*, while the variants of END feature 5 layers of message- paassiggg in F.p and 5 in ft. The hidden size of the invariant and equivariant features is set to 192, along with an RBF expansion of dimension 64 with cutoff 30Å for pairwise distances (as to ensure full-connectivity). Each model features 25 5.4M learnable parameters. We train all models for 10 epochs with an effective batch size of 64.
A.6.4 Conditional generation
We use 10 layers of message passing for EDM* while the variants of END feature 5 layers of geessageppassing in F.p and 5 in St. The hidden size of the invariant and equivariant features is set to 192, , along with an RBF expansion of dimension 64 with a cutoff of 10Å for pairwise distances. We train all models for 1000 epochs with a batch size of 64.
After initial encoding, the conditional information is introduced at the end of each message passing step, and alters the sealar hidden states through a one-layer MIP, that shares the same dimension as the hidden scalar state.
generation The encoding of the condition follows that of Gebauer et al. (2022). Each atom type gets its own embedding (of dimension 64), weighted by the proportion represents in the provided formula. The weighted embeddings of all atom types are then concatenated and flattened, and the obtained vector (of dimension 64 number of atom types is processed through a 2-layer MLP with 64 hidden units.
Composition-condiin
The compositions used at sampling time are extracted from the validation and test sets. For each unique formula, the model gets to generate 10 samples. The reported matching rate refers to the % of generated samples featuring the prompted composition.
Substtuuctue-conddttonnd generation The encoding of the condition follows that of Bao et al. (2023). Fach molecule is first converted to OPENBABEL object (O'Boyle et al., 2011) (solely based on positions and atom types), from which a fingerprint is in turn calculated. The obtained 1024- dimensional fingerprint is simply processed by 2 layer MLP with hidden dimensions [512, 256). and final linear projection to 192, i.e. the hidden size of the invariant features.
The model is evaluated by computing the Tanimoto similarity between the fingerprints obtained from generated samples and the fingerprints provided as conditional inputs.
A.7 Compute resources
All experiments were run single GPU. The experiments on QM9 were run on NVIDIA SM3090 with 24 GB of memory. The experiments on GEOM-DRUCH were run on NVIDIA A100 with 40 GB of memory. Training took up to 7 days.
Sampling The current implementation of END leads to 3x increase relative to EDM (with comparable number of learnable parameters) per function evaluation when performing sampling.
25

---

However, END usually requires much fewer number of function evaluations to achieve comparable (or better) accuracy, and we note that alternative parameterizations the reverse process are possible. In particular, the drift of the reverse process fa could be learned without direct dependence on fo- thereby leading to very limited overhead with respect to vanilla diffusion models for sampling only one neural network forward network would then required. For reference, we report sampling times on QM9 (1024 samples) in Table 11, for varying numbers of integration steps.
Table 11: Average sampling time in seconds for 1024 samples on QM9, The current implementation of END leads to = 3x increase relative to EDM.
sampling time [s]
Model
Steps
50
30.4
100
60.6
EDM
250
149.7
500
297.7
1000
593.8
88.6
100
179.6
END
250
445.9
500
886.4
1000
1765.8
Training A training step on QM9 (common batch size of 64) takes on average & 0.37s (END) vs 0.16s (EDM), this corresponds to Rs 2.3x relative increase. We observe the same trend on GEOM-Drugh (with an effective batch size 64), a training step takes on average s 0.40s for END & 0.15s for EDM (corresponding to a 2. 7x relative increase). In summary, other things equal, END leads to x 2.5x increase relative to EDM per training step while we find it to converge with a similar number of epochs. As for sampling, a direct parametrization of fe.pp would enable faster training while still requiring the evaluation of f&p obtain the target reverse drift term in Eq. (5).
26

---

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 method section shows that the proposed model has the properties we claim it has, while the experimental section demonstrates its performance. We additionally highlight limitations Section
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 be expected to generalize other settings.
is fine include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: In Section 6, we discuss the limitations the proposed approach, and experimental validation.
Guidelines:
The answer NA means that the paper has 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 separate "Limitations" section in 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 on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or 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 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 in the paper. The authors should use their best 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.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full of assumptions and complete (and correct) proof?
27

---

Answer: [Yes]
Justification:
When appropriate, we provide a proof sketch in Appendix.
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 in the statement of any theorems.
The proof's 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 a 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 the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided not)?
Question:
Answer: [Yes]
Justification: The data configuration follows from previous work, while most of the experi- mental details are provided in Appendix.
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 to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, descríbing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, provide access to the model. In general. releasing code and data is often one good way to 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 a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require 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 or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe 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 verifying the results.
5. Open access to data and code
28

---

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?
Question:
Answer: [Yes]
Justification:
We provide public code repository with our implementation.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.. cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" an acceptable answer. Papers cannot be rejected simply for including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https:
/ /nwpp...c//////////// wwtps/ccppblccggciiceooooo h
/CodeSuboeeeeeeee issiiissionn wswinnollll onPolicy) 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. If 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) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: The data configuration follows from previous work, while most of the experi- mental details are provided in Appendix.
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 provided either with the code, in 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: [Yes]
Justification:
All sampling results are averages over multiple seeds and we include standard
errors.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
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 with given experimental conditions).
29

---

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).
should be clear whether the error bar is the standard deviation or 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 to show in tables figures symmetric error bars that would yield results that are out of range (c.g. negative emror rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables 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:
We provide compute resources details in Appendix A.7.
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, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as 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 't make it into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: http///eeerppp.eeeeeeeeee httpppbii//tthhccsuueeeeeeeeeeee
Answer: [Yes]
Justification:
The paper conforms with the NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: There is dedicated paragraph in Section
Guidelines:
The answer NA means that there is no societal impact of the work performed.
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.
30

---

The conference expects that many papers will be foundational research and not tied 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 an 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 a 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 being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility MI..
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:
The paper does not pose such risks.
Guidelines:
The answer NA means that the paper poses 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 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 is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Lieenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: We have cited the original papers that released the datasets.
Guidelines:
The answer NA means that the paper does
existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
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.
31

---

If this information is 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: No assets introduced in the paper.
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 was 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. Crowdsourcing 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 did not require crowdsourcing.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information in 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 the NeurIPS Code of Ethies, 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 an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
No study participants were involved.
Guidelines:
The answer NA means that the paper does involve crowdsourcing research with human subjects.
Depending on 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.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS 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.
32


Reviews:
[{"review_id": "uRDlsv2pgG", "reviewer": "Reviewer_hMME", "summary": "This paper proposes an extension of diffusion models dubbed Equivariant Neural Diffusion, which leverages  a learnable forward diffusion process to enhance flexibility. The entire framework has been constructed such that physical symmetry, i.e., equivariance/invariance, of the density is preserved. Experiments are performed on QM9 and DRUGS datasets in the task of molecule generation from scratch as well as controllable generation.", "questions": "Q1. How does the method perform compared with GeoBFN, especially with different number of sampling steps?\n\nQ2. Could the method be applied to systems with larger scales, e.g., proteins?", "limitations": "The authors have discussed the limitations, including scaling, limited application scope, etc.", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. The presentation is mostly clear and the method is easy to follow.\n2. The method has been demonstrated to perform favorably in the controllable generation setting, even when the number of sampling steps is limited to, e.g., 50.", "weaknesses": "1. The proposed approach seems to be an incremental combination of geometric diffusion models and neural flow diffusion models. The way of combining these two flavors incurs limited technical novelty. The core design mostly lies in constructing $F\\_\\varphi$ which is equivariant.\n\n2. The performance on QM9 and DRUGS is a bit marginal compared with the selected baselines.\n\n3. Missing important baselines, e.g., GeoBFN [1], which has shown strong performance on the same task, i.e., molecule generation. Moreover, GeoBFN can also achieve high performance with very few sampling steps, even only 20.\n\n[1] Song et al. Unified generative modeling of 3d molecules with bayesian flow networks. In ICLR'24.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their constructive feedback.  We address the reviewer’s concerns/questions here below:\n\n**Limited technical novelty**  \nAs mentioned in our general rebuttal, END is indeed a combination of existing ideas. We however want to emphasize that (1) our work is the first to seek improvement by adding a learnable forward process, orthogonally to most previous work that has mostly focused on the reverse process, (2) designing an appropriate transformation was not straightforward, it required some care as to preserve the desired invariance of the learned distribution. \n\n**Marginal performance gain**  \nWhile we agree with the reviewer that the performance gains are limited on QM9, we have shown that END works considerably better than the considered baselines on GEOM-Drugs -- as summarized in the general rebuttal. Most notably, END significantly improves connectivity and geometry quality. It also shows better agreement with the training data on (newly added) drug-related metrics (SAScore, QED and logP).\n\n**Missing baseline GeoBFN**  \nUnfortunately, since GeoBFN was published on arXiv \\~2 months prior to the submission deadline a thorough comparison to this concurrent work is unfeasible. However, we extracted the relevant numbers presented in that paper and a direct comparison with GeoBFN in Tables 1 and 2 will be part of the final version. We summarize the comparison below. \n  \nOn QM9, END generally performs on par with GeoBFN on the stability / validity metrics, across sampling steps. We could not evaluate the additional metrics as we do not have access to samples / pretrained models for GeoBFN. \n\nOn the more realistic GEOM-Drugs, END is shown to outperform GeoBFN in terms of atom stability (first row in the table below) – especially for small numbers of steps. As we can not check whether GeoBFN generates disconnected fragments, solely based on validity GeoBFN seems to perform better (second row in the table below).\n\n|        | 50   | 100  | 500  | 1000 |  \n|--------|------|------|------|------|  \n| END    | 87.1 | 87.2 | 87.0 | 87.0 |  \n| GeoBFN | 75.1 | 78.9 | 81.4 | 85.6 |  \n|        |      |      |      |      |  \n| END    | 84.6 | 87.0 | 88.5 | 88.8 |  \n| GeoBFN | 91.7 | 93.1 | 93.5 | 92.1 |\n\nIt would be interesting to evaluate the connectivity / strain energy of the samples produced by GeoBFN for a full comparison.\n\n**Application to large systems**  \nThere is no restriction to the systems the END could be applied to.   \nAs all concurrent approaches, scaling to large systems is currently limited by the full connectivity of the message passing scheme.    \nAs for other methods treating categorical features as a continuous relaxation, we also note that for systems with a large number of atom types, alternatives to one-hot encoding (linear scaling wrt cardinality of space) might be beneficial. Methods such as Analog Bits \\[1\\] (logarithmic scaling), or continuous lower-dimensional embedding such as that of GeoLDM \\[2\\] are interesting approaches to deal with discrete features within END.\n\n\\[1\\] Analog Bits: Generating Discrete Data using Diffusion Models with Self-Conditioning, ICLR 2023\\.  \n\\[2\\] Geometric Latent Diffusion Models for 3D Molecule Generation, ICML 2023", "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": "8FHFtaFxGl", "reviewer": "Reviewer_K86T", "summary": "The Equivariant Neural Diffusion (END) model is a novel approach for molecule generation in 3D that maintains equivariance to Euclidean transformations. Unlike traditional diffusion models that use a pre-specified forward process, END introduces a learnable forward process, parameterized through a time- and data-dependent transformation. This innovation allows the model to adapt better to the underlying data distribution. Experimental results demonstrate that END outperforms strong baselines on standard benchmarks for both unconditional and conditional generation tasks, particularly excelling in generating molecules with specific compositions and substructures. This flexibility in modeling complex molecular structures suggests significant potential for applications in drug discovery and materials design.", "questions": "- Why was GeoDiff not included as a baseline in your comparisons? \n- Could you provide more detailed ablation studies that isolate the contributions of the learnable forward process and other key components of END?", "limitations": "**Computational Complexity:** The learnable forward process in the END model increases computational complexity, resulting in longer training times and higher resource usage. The authors should discuss potential optimizations to reduce this overhead, such as more efficient algorithms or hybrid approaches. Comparing training times and resource requirements with simpler models would provide useful insights into the model's efficiency.\n\n**Scalability Issues:** END's architecture limits its scalability to larger datasets and more complex molecular structures like proteins. To enhance scalability, the authors should explore strategies like sparse representations or hierarchical approaches. These methods could enable the model to handle larger and more complex datasets effectively.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "**Originality:** END introduces a novel learnable forward process, diverging from the fixed processes in traditional diffusion models. This allows the model to better adapt to the underlying data distribution, especially in complex 3D molecular structures. It creatively combines elements from Neural Function Matching Diffusion Models (NFDM) and Equivariant Diffusion Models (EDM), enhancing the generative process while maintaining E(3) equivariance.\n\n**Quality:** The model demonstrates superior performance in generating 3D molecular structures, outperforming strong baselines in both unconditional and conditional settings. It excels in generating stable, valid, and unique molecules, particularly evident in its results on the GEOM-DRUGS dataset. Comprehensive experiments and ablation studies confirm the robustness and reliability of the model.\n\n**Clarity:** The paper provides a detailed and clear exposition of the methodology, including the formulation of the learnable forward process, parameterization, and evaluation metrics. The inclusion of algorithmic steps and extensive experimental details facilitates replicability for researchers.\n\n**Significance:** END represents a substantial advancement in generative modeling for 3D molecules, addressing limitations of prior models by improving sample quality and generation speed. Its ability to maintain equivariance while achieving high performance has significant implications for applications in drug discovery and materials design, potentially transforming these fields.", "weaknesses": "**Performance Consistency:** The performance of END is not consistently superior to existing baselines. Although it shows competitive results, there are instances where traditional models, like EDM and its variants, outperform END, particularly in metrics such as validity and uniqueness across different datasets.\n\n**Complexity and Scalability:** The added complexity of a learnable forward process, while innovative, increases the model’s training time and resource requirements. END requires more computational resources and longer training periods compared to simpler, fixed-process models like EDM . The model operates on fully-connected graphs, which limits its scalability to larger datasets and more complex molecular structures. This constraint can hinder its applicability in more demanding real-world scenarios.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their constructive feedback. We are happy that our work was positively received by the reviewer.   \nWe address the reviewer’s concerns/questions here below.\n\n**Performance Consistency**   \nRegarding validity, we first note that (1) cheminformatics software implicitly adds hydrogens if an atom has a valence smaller than expected, (2) validity is usually computed on the largest connected fragment (thereby ignoring the remaining atoms).   \nOn QM9, mol. stability gives a better picture, as it ensures that each atom exactly has the desired variance. Models including a learnable forward pass are shown to perform significantly better than baselines without. Improvements are also observed across other metrics (Table 1 in submission).   \nOn the more challenging GEOM-Drugs, we believe that evaluating the model in terms of (validity x connectivity) is more realistic. According to that metric, END is clearly better than the baseline and samples far more connected molecules (Table 2 in rebuttal PDF). \n\n**GeoDiff**  \nTo the best of our knowledge, GeoDiff \\[1\\] is a diffusion model for conformer generation, i.e. generation of coordinates given a molecular graph. Perhaps the reviewer had another model in mind?\n\n**Ablation**  \nThe key component in our method is the learnable forward process. Hence, the logical ablation is whether to include a learnable forward (=END), or not (=EDM). \n\nTo ensure that the difference in performance does not come from an increase in parameters, we made sure that all EDM baselines shared the same amount of parameters as the models featuring a learnable forward. We did so by having 10 rounds of message-passing in EDM, against 5 in the forward and 5 in the reverse for END – all other things equal.\n\nWe additionally provided another ablated version of END on QM9, where only the mean is learned, whereas the standard deviation of the conditional marginal is pre-specified and derived from the same noise schedule as EDM.   \nAs suggested by another reviewer **\\[jNQ2\\]**, we conducted the same experiment in the conditional setting, where fixing the standard deviation is shown to lead to a small decrease in performance.  \nWe are currently running the same ablation on GEOM-Drugs, and hope to be able to report the results before the end of the discussion period. \n\n**Complexity and Scalability**  \n**Complexity:**   \nIn terms of training, we did not observe that END required more epochs than the baseline to reach convergence. However, it is correct that each training step takes longer, and that the total training time is increased. All other things equal, we evaluated that increase to \\~2.5x relative to EDM. In terms of resources, END could be trained on a single GPU for both datasets.\n\nRegarding sampling, our current implementation incurs a \\~3x increase relative to EDM.   \nHowever, we emphasize that END can be competitive with fewer integration steps. As a concrete example, on GEOM-Drugs, END with only 100 steps yields better samples than EDM with 1000 integration steps, i.e. in \\~⅓ of the time.  \nFinally, alternative efficient parameterizations of the reverse process are possible.  In particular, the drift of the reverse process $\\\\hat{f}\\_{\\\\theta, \\\\varphi}(\\\\boldsymbol{z}\\_t, t)$ could be learned without direct dependence on $f\\_{\\\\varphi}$, thereby leading to a very limited overhead with respect to vanilla diffusion models for sampling.\n\n**Scalability:** As all concurrent approaches, scaling to large systems is currently limited by the full connectivity of the message passing scheme.  \nAn element specific to END, is that, due to the learnable forward, the prior is no longer forced to be $\\\\mathcal{N}(0, I)$ as in conventional diffusion models and can instead e.g. become size-specific. For large molecules, one could imagine scaling the prior with the number of atoms present in the system, i.e. $\\\\mathcal{N}(0, sI)$ where $s$ is a function of the number of atoms. Combined with a distance-based cutoff function to determine the neighborhood, this would allow for a more graceful scaling – as early steps of the denoising process would not result in fully-connected message passing as it is effectively the case with a simple $\\\\mathcal{N}(0, I)$ prior.  \n\nAnother possibility to scale to systems such as large coordination complexes, is to design priors that “encode shapes”, e.g. square planar geometry could be built as a center and 4 groups of atoms representing the coordinated ligands, and thereby enabling more local message passing schemes. \n\n\\[1\\] GeoDiff: a Geometric Diffusion Model for Molecular Conformation Generation, ICLR 2022", "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": "dKbJtMDqMU", "reviewer": "Reviewer_jNQ2", "summary": "Paper presents, END, a diffusion models for 3D molecule generation that\n    \n- is equivariant to euclidean transformations and \n\n- includes a learnable forward process.\n  \n Specifically, the forward process in the presented models, is defined as a learnable transformation, dependent on both time and data such that the resulting latent representation $z_t$ transforms covariantly with the injected noise. This is the main difference between the forward pass of Neural Flow Diffusion Models and the proposed method.", "questions": "- The only form of ablation i see for the proposed model is in Table 1 where two versions of END are provided. However even here, it seems the END with $\\mu_\\phi$ performs similar or sometimes better (according to the presented metrics) than the full END model. Is the same pattern observed for the conditional generation tasks?\n \n- Relating to the first question above, can a much more thorough ablation of the proposed method be provided in both testing scenarios to really ascertain the utility of the components in the proposed method?\n\n- Can the training times for the benchmarked methods be provided for comparison? While it is stated passingly that END requires more training time, can this be cast in contrast with the baselines by providing the actual numbers?", "limitations": "Limitations are adequately discussed by the authors.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Can be used for both conditional and conditional molecules generation.\n\n- Improves on existing equivariant diffusion models.\n\n- Experimentally, the proposed method shows improvements on both conditional and unconditional generation on the QM9 and GEOM-DRUGS datasets.", "weaknesses": "- Lack of a thorough ablation of the proposed method.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their constructive feedback. We are happy that our work was positively received by the reviewer.   \nWe address the reviewer’s concerns/questions here below:\n\n**Ablation**  \nThe key component in our method is the learnable forward process. Hence, the logical ablation is whether to include a learnable forward (=END), or not (=EDM). \n\nAs stated in our global rebuttal, we made sure that all baselines shared the same amount of parameters as the models featuring a learnable forward to ensure that the difference in performance did not come from an increase in parameters. We did so by having 10 rounds of message-passing in EDM, against 5 in the forward and 5 in the reverse for END – all other things equal.\n\nAs pointed out by the reviewer, the ablated version where only the mean is learned (the standard deviation of the conditional marginal is pre-specified and derived from the same noise schedule as EDM) is shown to perform on par with the full version. We believe this is due to the relative simplicity of the task. We are currently running the same ablation on the more challenging GEOM-Drugs, and hope to be able to report the results before the end of the discussion period. \n\nAs nicely suggested by the reviewer, we conducted the same ablation in the conditional setting (composition-conditioned generation). The corresponding results are provided in the table here below – we will also add it to the final version. The ablated model is presented in the last row.\n\n| Matching [%] (↑) | 50   | 100  | 250  | 500  | 1000 |\n|------------------|------|------|------|------|------|\n| cEDM             | 69.6 | 73.0 | 74.1 | 76.2 | 75.5 |\n| cEND             | 89.2 | 90.1 | 91.2 | 91.5 | 91.0 |\n| cEND (mean only) | 75.7 | 79.9 | 82.7 | 83.0 | 83.5 |\n\nRegarding additional ablations, disabling equivariance in $F_\\varphi$ could also be a possible option, previous work \\[1\\] has however clearly shown that an architecture obeying the relevant symmetries is a useful inductive bias. We would be happy to hear the additional ablations that the reviewer would like to see in the final version of the paper. \n\n**Timing**  \nIn terms of timing, a training step on QM9 (common batch size of 64\\) takes on average \\~0.37s (END) vs \\~0.16s (EDM), this corresponds to a \\~2.3x relative increase.   \nWe observe the same trend on GEOM-Drugs (with an effective batch size of 64), a training step takes on average  \\~0.40s for END vs. \\~0.15s for EDM (corresponding to a \\~2.7x relative increase).  \nWe do not observe that END requires a larger number of epochs to converge, compared to EDM. All models were trained for the same number of epochs / steps.\n\nWith our current implementation, sampling with END requires slightly less than 3x the time required by EDM on average. We report sampling times on QM9 (1024 samples) as an example, for varying numbers of integration steps.\n\n| time \\[s\\] (↑) | 50   | 100   | 250   | 500   | 1000   |  \n|--------------|------|-------|-------|-------|--------|  \n| EDM          | 30.4 | 60.6  | 149.7 | 297.7 | 593.7  |  \n| END          | 88.6 | 179.6 | 445.9 | 886.4 | 1765.8 |\n\nHowever, we want to stress that END can achieve comparable (or better) accuracy with less integration steps. A concrete example of this can be seen in Table 2 in the rebuttal PDF (GEOM-Drugs), where with only 100 steps END yields better samples (connectivity and strain energy) than EDM with 1000 integration steps, i.e. in \\~⅓ of the time.  \nAs mentioned in the global rebuttal, we note that alternative parameterizations of the reverse process are possible. In particular, the drift of the reverse process $\\\\hat{f}\\_{\\\\theta, \\\\varphi}$ could be learned without direct dependence on $f\\_{\\\\varphi}$, thereby leading to very limited overhead with respect to vanilla diffusion models for sampling.   \n\n\\[1\\] Equivariant Diffusion for Molecule Generation in 3D. ICML 2022\\.", "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": "sWYKW0OssF", "reviewer": "Reviewer_5dRq", "summary": "This paper presents Equivariant Neural Diffusion (END), a novel diffusion model for 3D molecule generation. The major novelty of END over previous molecule diffusion models lies in adopting a learnable forward process based on neural flow diffusion models. Experiments show that END can achieve good performance on benchmark datasets.", "questions": "See Weaknesses part.", "limitations": "Yes.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Successfully incorporate neural flow diffusion models into the equivariant 3D molecule generation framework and demonstrate its usefulness in experiments.\n- Generally good, clear and well-organized writing.", "weaknesses": "- The novelty contribution of the proposed END method is not high, as END is largely a combination of EDM [1] and neural flow diffusion models. Particularly, the paper does not give a clear discussion or analysis about why adopting learnable forward diffusion process is useful and beneficial to 3D molecule generation, or what molecular structures can be additionally captured by END through learnable forward diffusion process compared with previous diffusion models.\n- There already exist some SDE based 3D molecule generation methods like EDM-BRIDGE [2] and EEGSDE [3]. Authors are encouraged to highlight the key difference in the diffusion process between END and these methods.\n- Compared with GEOLDM, END does not show better performance (Table 1 and 2), which weakens the claim about the advantages of using learnable forward process. Since the main evaluation metrics proposed by EDM [1] in 3D molecule generation are saturating in recent literatures, authors are encouraged to adopt metrics proposed by HierDiff [4] to better evaluate the quality of generated 3D molecules.\n\n[1] Equivariant Diffusion for Molecule Generation in 3D. ICML 2022.  \n[2] Diffusion-based Molecule Generation with Informative Prior Bridges. NeurIPS 2022.  \n[3] Equivariant Energy-Guided SDE for Inverse Molecular Design. ICLR 2023.  \n[4] Coarse-to-Fine: a Hierarchical Diffusion Model for Molecule Generation in 3D. ICML 2023.\n\n-------------------Post Rebuttal---------------------\n\nI appreciate authors' efforts in addressing my concerns and questions in rebuttal. After reading over authors' rebuttal responses and pdf, I think all my concerns have been addressed so I increased my score. I hope authors will carefully add all rebuttal updates (discussion, analysis and experiment results) to the revised version of paper in the future.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their constructive feedback.   \nWe address the reviewer’s concerns/questions here below:\n\n**Limited technical novelty**  \nAs mentioned in our general rebuttal, END is indeed a combination of existing ideas. We however want to emphasize that (1) our work is the first to seek improvement by adding a learnable forward process, orthogonally to most previous work that has mostly focused on the reverse process, (2) designing an appropriate transformation was not straightforward, it required some care as to preserve the desired invariance of the learned distribution. \n\n**Relevance of learnable forward process**   \nMost improvements to diffusion models for molecules have focused on designing more expressive denoisers, better noise schedules, or embedding discrete features in a continuous latent space. In END, we seek improvement in an orthogonal direction. Adding a learnable forward allows the hierarchy of latent variables to be learnt, instead of fixed as in classical diffusion models, thereby offering greater flexibility and enhanced generative modeling. We note that END can be combined with previous developments for even better performance.\n\nThis enhanced generative modelling is greatly showcased on the more challenging and realistic GEOM-Drugs, where the learnable forward greatly improves connectivity (\\~20% across all sampling schemes), and leads to 3D structures of better quality as evidenced by the significantly lower strain energies. \n\nFinally, the addition of a learnable forward process comes with two practically relevant by-products: (1) END achieves better performance with fewer integration steps, and (2) the learnable forward yields an improved conditional generative process – as demonstrated on two conditional tasks where cEND even outperforms  by a significant margin a model leveraging an auxiliary model through classifier-guidance (EEGSDE) (Tables 3 and 4 in the submitted manuscript). \n\n**Previous work**  \nWe first note that both papers were already included as baselines / previous work in the submitted manuscript, but, as suggested by the reviewer, we will add a more thorough discussion to the final version. We summarize it here:\n\n* EEGSDE is a continuous-time formulation of EDM, where conditional generation is performed by combining (1) a conditional model score model, (2) a method similar to classifier-guidance (requiring the training of an auxiliary model). In cEND, we instead only learn a conditional model, akin to classifier-free guidance.   \n* Similar to END, EDM-BRIDGE builds on the observation that there exists an infinity of processes mapping from prior to target distributions. EDM-Bridge constructs one such process that incorporates some prior knowledge, i.e. part of the drift term is a physically-inspired force term. END can be seen as a generalization of EDM-Bridge, where the forward drift term is now learned instead of pre-specified. Through experiments, we show that a learnable forward performs better than a fixed one, even when the latter is physics-inspired.\n\n**Additional metrics**  \nWe thank the reviewer for their suggestion. We added a citation to HierDiff \\[1\\] and computed additional metrics taken from the paper, namely: SAScore, QED, logP and MW (see rebuttal PDF). To further assess the quality of the generated 3D structures, we also added a metric that measures the strain energy – expressed as an energy difference between the generated geometry and its relaxation (obtained through force-field geometry optimization). \n\nThese additional metrics were computed on the valid x connected samples generated by each method. We converted the generated samples to SDF using OpenBabel, and read them in using RDKit. For the SAScore, we normalized the values between 0 and 1, with 0 being “difficult to synthesize” and 1 “easy to synthesize”.  \nAll the additional metrics are provided in the rebuttal PDF. We summarize here the main results. \n\n* On QM9, END shows better agreement with the data distribution (except for QED which is captured perfectly by all methods). In particular, the reduction in strain energy demonstrates that END yields better geometries than the baselines (Table 1 in rebuttal PDF).  \n* On GEOM-Drugs, in addition to greatly improved connectivity, the SAScore, QED and logP are shown to be in better agreement with the data distribution (Table 3 in rebuttal PDF). The geometries are also shown to be of better quality as per the significantly reduced strain energy (Table 2 in rebuttal PDF).\n\n\n**Comparison to GeoLDM**  \nWe downloaded the GEOM-Drugs checkpoint available with the official implementation of GeoLDM, and evaluated the generated samples. We added one line with the obtained results in Table 2 in the rebuttal PDF, and summarize the findings here.\nWhile all samples generated by GeoLDM were effectively deemed “valid”, only around **46%** were connected, against nearly **83%** for END. The strain energy of the samples generated by END was also significantly lower, 55 kcal/mol vs 133 kcal/mol for GeoLDM, underscoring the advantages of a learnable forward process. \n\n\\[1\\] Coarse-to-Fine: a Hierarchical Diffusion Model for Molecule Generation in 3D. ICML 2023\\.", "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:
We thank the reviewers for their constructive feedback. 

We are pleased that the reviewers [**hMME,5dRq,K86T**] found the presentation of END clear, its construction original **[K86T]**, that it is a potentially significant contribution to the field **[K86T],** and that it demonstrates experimental benefits over existing models in unconditional and conditional settings [**hMME,K86T,jNQ2**].

In this global rebuttal, we address the main concerns of the reviewers. We also provide a detailed answer to each reviewer below.

## Performance and metrics

Three reviewers [**hMME,5dRq,K86T**] remarked that we did not demonstrate a clear improvement over competing methods. We respectfully disagree, and to highlight the strong performance improvement of END, we have included several additional metrics. And both the original metrics and the newly included ones show that END strongly outperforms competitors.

### Additional metrics

We added 5 metrics. As specifically suggested by reviewer **5dRq**, they are taken from HierDiff [1]:

* **SAS**: Synthetic accessibility score;  
* **QED**: Quantitative estimation of drug-likeness;  
* **logP**: Partition coefficient;  
* **MW**: Molecular weight.

We also provided a metric measuring the quality of the generated 3D geometries: **Strain Energy**, as the difference in energy when relaxing the generated molecules (using RDKit MMFF)

### QM9

On the simpler QM9 benchmark, in terms of validity, stability, and uniqueness, END reaches a performance level similar to that of current SOTA methods (GeoLDM and GeoBFN) but these metrics are essentially saturated.   

More importantly END displays better agreement with the dataset distribution compared with EDM: For example, TV on atom types go from 2.5 (EDM) to 1.2 (END) at 1000 sampling steps (Table 1 in submitted manuscript). Furthermore, on all additional metrics, END shows better agreement with the data distribution (except for QED which is captured perfectly by all methods) as shown in Table 1 in rebuttal PDF. In particular, the reduction in strain energy demonstrates that END yields better geometries than the baselines.

### GEOM-Drugs

On the more challenging and realistic GEOM-Drugs we show **large improvements** (Tables 2 and 3 in rebuttal PDF):

- (i) ~20% improvement over EDM in terms of “Validity x Connectivity”;  
- (ii) ~50% improvement on “TV on atom types”;  
- (iii) improved 3D geometries as per the significantly lower strain energy;  
- (iv) improved SAS, QED and logP.

As suggested by reviewer **5dRq**, we conducted a more extensive comparison with GeoLDM by sampling using the publicly available checkpoint.  We summarize here our findings (Table 2 in the rebuttal PDF):
* V x C improves from 45.8% (GeoLDM) to 82.5% (END)  
* TV on atom types improves from 10.6 (GeoLDM) to 3.0 (END)  
* Strain Energy improves from 133.5 (GeoLDM) to 55.0 (END).

We believe these results demonstrate the strong improvement on the prior SOTA. All metrics are collected in tables in the rebuttal PDF.

## Missing baselines / Previous work

**hMME** requested a comparison to the very recent GeoBFN [2] (posted on arXiv ~2 months prior to the submission). While already part of the literature review, we will add the comparison in Tables 1 and 2 in the final version.

**5dRq** requested a comparison to two previous works [3, 4]. While already included in the submission as baseline/related work, we will additionally provide a paragraph discussing the key differences in the final version (see answer to **5dRq** for summary).

## Ablations

**K86T** and **jNQ2** mentioned that more detailed ablation studies could make the paper even stronger.   
First, we want to stress that we made sure that all baselines shared the same amount of parameters as the models featuring a learnable forward, as to ensure that the difference does not simply come from an increase in parameters. 
Following reviewer **jNQ2**'s suggestion, we ran an additional ablation study, and added the results in our rebuttal to **jNQ2**. We are currently running an ablated version of END (mean-only) on GEOM-Drugs – we hope to be able to report the results during the discussion period.

## Timing

Two reviewers  [**K86T, jNQ2**] requested more information about the increased training/sampling time. 
In summary, other things equal, END leads to ~2.5x increase relative to EDM per training step and requires the same number of epochs to converge.
Regarding sampling, our current implementation leads to ~3x increase relative to EDM per function evaluation. However, END usually requires much fewer number of function evaluations to achieve comparable (or better) accuracy and we note that alternative parameterizations of the reverse process are possible. In particular, the drift of the reverse process $\\hat{f}\_{\\theta, \\varphi}$ could be learned without direct dependence on $f\_{\\varphi}$, thereby leading to very limited overhead with respect to vanilla diffusion models for sampling. We will add these important considerations to the final version of the manuscript.

## Technical novelty

Some reviewers [**5dRq, hMME**] pointed out some limited novelty.   
END is indeed a combination of existing ideas. We however want to emphasize that (1) our work is the first molecule generation diffusion model to seek improvement by adding a learnable forward process, orthogonally to most previous work that has mostly focused on the reverse process, (2) designing an appropriate transformation was not straightforward, it required some care as to preserve the desired invariance of the learned distribution. 

[1] Coarse-to-Fine: a Hierarchical Diffusion Model for Molecule Generation in 3D. ICML 2023.  
[2] Song et al. Unified generative modeling of 3d molecules with bayesian flow networks. In ICLR'24.  
[3] Diffusion-based Molecule Generation with Informative Prior Bridges. NeurIPS 2022.  
[4] Equivariant Energy-Guided SDE for Inverse Molecular Design. ICLR 2023.

Meta-review: The reviewers all agreed that the method performs well (esp. in challenging conditions), and most reviewers noted that the paper is well written. The method is a novel combination of existing ideas, and different reviewers appraised this differently, with some finding it somewhat incremental (e.g. hMME) while others (e.g. K86T) finding it a creative combination. Many minor concerns raised by the reviewers were effectively addressed during the discussion, and in the end only hMME remained skeptical. Their primary concern is a comparison to the recent work GeoBFN. The authors have now shown their method to be competitive, and are working on running comparisons with additional metrics as requested.

In my view the method is sufficiently novel, and already fairly well benchmarked in the initial submission. The authors have been very active improving the paper and addressing reviewer concerns during the discussion period, and overall this now looks like a solid contribution. So I recommend the paper for acceptance at NeurIPS.
Decision: Accept (poster)