Scalable Non-Equivariant 3D Molecule Generation via Rotational Alignment
Authors: Yuhui Ding
Thomas Hofmann
Venue: ICML 2025 poster (2025)
Keywords: Non-equivariant diffusion
3D molecule generation
Primary area: deep_learning->generative_models_and_autoencoders

Abstract: Equivariant diffusion models have achieved impressive performance in 3D molecule generation. These models incorporate Euclidean symmetries of 3D molecules by utilizing an SE(3)-equivariant denoising network. However, specialized equivariant architectures limit the scalability and efficiency of diffusion models. In this paper, we propose an approach that relaxes such equivariance constraints. Specifically, our approach learns a sample-dependent SO(3) transformation for each molecule to construct an aligned latent space. A non-equivariant diffusion model is then trained over the aligned representations. Experimental results demonstrate that our approach performs significantly better than previously reported non-equivariant models. It yields sample quality comparable to state-of-the-art equivariant diffusion models and offers improved training and sampling efficiency. Our code is available at: https://github.com/skeletondyh/RADM

Scalable Non-Equivariant 3D Molecule Generation via Rotational Alignment
Yuhui Ding' Thomas Hofmann
Abstract
language (von Rütte al., 2025), molecules (Hoogeboom et al., 2022), proteins (Yim al., 2023).
Equivariant diffusion models have achieved im- pressive performance in 3D molecule generation. These models incorporate Euclidean symmetries of 3D molecules by utilizing an SEuua-equi denoising network. However, specialized equiv- ariant architectures limit the scalability and effi- ciency of diffusion models. In this paper, we pro- pose an approach that relaxes such equivariance constraints. Specifically, our approach learns a sample-dependent SO(3) transformation for each molecule to construct aligned latent space. A non-equivariant diffusion model is then trained over the aligned representations. Experimental results demonstrate that our approach performs significantly better than previously reported non- cquivariant models. It yields sample quality com- parable to state-of-the-art equivariant diffusion models and offers improved training and sam- pling efficiency. Our code is available at https: / //git hubb cwm/ wttp://keeeetoo/aaaaaatt
In this paper, we focus on the task of 3D molecule genera- tion, which has been an important and active topic in drug discovery. Specifically, our goal is to generate the 3D atomic coordinates and atom types molecule from scratch. Dif- ferent from data with grid-like structures (e.g.. images and text sequences), 3D molecules pose unique challenges to generative modeling due to the Euclidean symmetry group of R i.e., SE(3). Particularly, the chemical properties of a molecule remain unchanged no matter how the molecule is translated or rotated in the three-dimensional space. To in- corporate the Euclidean symmetries into the model, Hooge- boom et (2022) proposed an equivariant diffusion model EDM that aims learn an (0(3)-invariant distribution, which requires the score function of the diffusion model to be (e(3)--quivariant (Köhler et al., 2020; Garcia Satorras et al., 2021). To satisfy such a constraint, EDM utilized an EGNN (Satorras et al., 2021) to parameterize the denoising network. The impressive performance of EDM on the 3D molecule generation task has inspired a series of follow-up works (Wu al., 2022; Xu al., 2023; Vignac al., 2023b) that build equivariance into diffusion models. Consequently, equivariant diffusion models have become the dominant approach for tasks related to 3D molecules.
1. Introduction
Diffusion-based generative models (Ho et al., 2020; Song et al., 2021) have achieved rapid and significant progress in recent years. These models feature a forward process and a reverse process. The forward process is a fixed Markov process that gradually adds noise to each data sample until the original information is totally corrupted. The reverse process is parameterized by denoising neural network that is trained to reverse the forward process step by step and recover the original data from noise. Once trained, new samples can be generated by performing the reverse process starting from non-in informative noise. Motivated by the suc- cess of diffusion models in vision tasks (Dhariwal & Nichol, 2021; Rombach et al., 2022), researchers have sought extend them to broader areas beyond vision, such as natural
While equivariant diffusion models account for the symme- tries of 3D molecules, an important question arises naturally: Is equivariance necessary for an effective molecular gener- ative model? We argue that it is not, as the probability of molecule is determined by the total probability of all its possible 3D positions, regardless of whether each position has an equal probability. Furthermore, equivariant architec- tures have drawbacks. Compared with their non-equivariant counterparts, they have more complex parametrization and lack standardized implementations. Additionally, equivari- ant architectures tend to be less efficient and less scalable than non-equivariant ones. In contrast to the vision and text domains which have been unified by transformers (Vaswani et al., 2017; Dosovitskiy et al., 2021), irregular data from scientific domains, with equivariant models being the main- stream choice, have benefited less from recent advances in architectural optimization (Dao et al., 2022; Peebles Xie, 2023).
Depaarmeent Computer Science, ETH Zurich. Correspon- dence to: Yuhui Ding <yuhui.ding ehh.ce.h..ch...
Proceedings of the 42nd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Scalable Non-Equivariant 3D Molecule Generation via Rotational Alignment
Driven by the interest in further investigating the capacity of non-equivariant diffusion models and motivated by the pos- sibility of connecting advances in different domains, in this paper we propose an approach to improving non-equivariant diffusion models for 3D molecule generation. Our inspi- ration comes from 3D vision models that typically rely on clean and well-aligned datasets such as ShapeNet (Chang et al., 2015). We hypothesize that aligning the representa- tions of 3D molecules could mitigate the variation in 3D coordinates caused by arbitrary and unknown Euclidean transformations, thereby making the data distribution more learnable for non-equivariant generative model.
2.1. Problem Definition
We are interested in generating 3D molecules from scratch. We consider each molecule collection of points in the three dimensional space. Specificallyy, a molecule that has N atoms is represented as z = (#1, T2, Zxn ) E Nx3, which correspond to the atomic coordinates in RR, and h = (h1, hg..... ,N) E IR RNxd, which correspond to the features of atoms (c.g., types and charges). The atomic co- ordinates z are affected by translations and rotations in IR, while the atom features h are invariant to such transforma- tions. Let R E R3x3 be an orthogonal rotation matrix and t E R the transformed coordinates can be written as':
However, aligning 3D molecules is inherently challenging due to the absence of supervision signals. To address this challenge, we propose to construct aligned representations in unsupervised manner with autoencoder. Autoen- coders have been widely used to reduce input dimensions and improve efficiency for latent diffusion models (Rom- bach et al., 2022). In our approach, the autoencoder shapes the latent space that is more suitable for non-equivariant dif- fusion models. More specifically, we use separate graph neural network to generate a sample-dependent SO(3) trans- formation (in the form of rotation matrix) rotate each molecule, and then train non-equivariant autoencoder to reconstruct the rotated molecules. The network that gener- ates the rotation is trained jointly with the autoencoder to minimize the reconstruction error, so that the model learns to arrange molecules in a way that facilitates the learning of a non-equivariant autoencoder. The aligned latent space then allows us train non-equivariant diffusion model for which non-equivariant architecture (such as vanilla graph neural network or transformer (Peebles & Xie, 2023)) can be employed as the denoising network, offering improved efficiency and better scalability. We refer to our approach as RADM (Rotationally Aligned Diffusion Model). We empir- ically test our approach on molecule generation benchmarks, and experimental results demonstrate that our approach sig- nificantly outperforms previously reported non-equivariant models and achieves sample quality comparable to state- of-the-art equivariant diffusion models. As expected, our non-equivariant diffusion model exhibits better scalability and improves the sampling efficiency significantly.
Rx +t= = (Ræ, taz + tt..... REN + +t). (I)
2.2. Diffusion Models
The forward process of the diffusion model starts with a generic data point æ and adds increasing levels of Gaussian noise to it. We use 2n to denote the noisy version of æ for I: T:
q = ((Zz
where 0g controls how much information of the original data point æ is retained and og defines the level of added noise. In this paper we follow previous works (Ho et al., 2020; Kingma et al., 2021) and let af + of = 1. or is close to 1 at f = 0 and then monotonically decreases to or that is close to O, which means the original data point is gradually corrupted until it becomes almost pure noise. The forward process is defined to be Markovian and the joint distribution of all noisy variables can be written as:
T (=2, 21, :1, (1,,,,)...,,,,,,,) = (==0(=) II (ol Z1=) (3)
Then the process (2) can be equivalently expressed using transition distribution q(z))z)) for 0 < $ <<tt.
{(21(z)) = (1 03.0 T2
(4)
where 0zjs = az/a, and o2 = of - o2. Given the distributions above and using Bayes' rule, we can derive the true posterior distribution of the forward transition (4), conditioned on z::
2. Preliminaries
We review the necessary background in this section before we present our proposed approach. In Section 2.1, we for- mally describe the definition of the problem. In Section 2.2, we introduce the basics of diffusion models. To simplify the expressions, we use the notion of signal-to-noise ratio following Kingma et al. (2021). In Section 2.3, we summa- rize how existing works incorporate equivariance into the diffusion model, upon which we build our own approach.
g(2, (z).2z..)) x)) = =N(zzlMaiizz.. (.
(5)
where Moje(zz, a)) and of can be analytically computed:
a of
o2 o
03
Ms]((zt, x))
(6)
of
of
'Here R..: + t denotes the group action of (R, t) on xr.
2

---

Scalable Non-Equivariant 3D Molecule Generation via Rotational Alignment
The reverse process of diffusion models is also Markovian. It aims to generate all intermediate variables 2n and x in the reverse direction:
2.3. Equivariance
As mentioned in Section 2.1, the atomic coordinates æ of molecule can be arbitrarily translated and rotated in the three-dimensional space without affecting its chemical prop- erties, which poses challenges to generative modeling. The translation transformation can be handled trivially by pro- jecting the coordinates of the molecule into the (N -N- x 3 dimensional linear subspace where >. H = 0. The projec- tion can be made by subtracting the center of gravity from each x..- As proven by Garcia Satorras et al. (2021) and Xu et al. (2022), sampling from normal distribution in the (N -) x 3 dimensional subspace can be done by first sam- pling from a normal distribution in the N x 3 dimensional space and then subtracting the center of gravity. The diffu- sion model in the subspace can then be derived in the same way as in Section 2.2, except that the center of gravity needs to be subtracted from the part of the predicted noise that corresponds to atomic coordinates, to restrict intermediate variables to the subspace.
T p(20, 21,...,, 2T, 2)) = pl 2T ( 20) II ((=)) (7)
where p(zr) is chosen to be a standard normal distribution. Since the exact reverse transition p(za |z+) (0 S t - T) is unknown, it is parameterized by neural network that is trained to approximate the true posterior conditioned on x (5):
= (2, (z((Ze))) ())
(8)
In (8) the network Φ is expected to predict the original data point x from its noisy version 2,. Ho al. (2020) found that predicting the noise added to yielded better performance than predicting æ itself. We use the same noise prediction parametrization in our model, and a (ze, t) in (8) is further rewritten as:
26 (Z,,))) 1)
Zt Tt
(z(,))))
To deal with rotations, most existing works (Hoogeboom et al., 2022; Xu et al., 2023) add equivariance constraints to the diffusion model. Specifically, they model a distribution that invariant to rotation R:
Using similar techniques VAEs, the evidence lower bound (for a single data point æ)) can be derived for the diffusion model as:
p(œ) = p(Rx)
(13)
Köhler et al. (2020) showed that applying an equivariant invertible transformation to an invariant distribution would result in another invariant distribution. Formally, a function f is SO(3) equivariant if it satisfies (Ræ) = Rf(()) for any R SO(3). Xu et al. (2022) further showed that for Markov chain with an invariant initial distribution, if the Markov transition kernel is equivariant:
- logp(x;) - EL)OO()) Cpriee + Lo Σ Le (10)
where Lpeiee DKL. (g( (((2r/=)]ll (=)r))))))((TT) denotes the KL di-
vergence between the final distribution q(zr (x) generated by the forward process and the prior distribution p(ær),
= |- logp(æ> =0)] represents a reconstruction loss, and L, E, (ax) Dxr (g) |2,, x>) (1]1p(z1111 (2,))- Let SNR(t) be the signal-to-noise ratio (Kingma et al., 2021) panded (using the noise parametrization) as:
p( (x) (Rax' [Rx),
(14)
defined as SNR(t) = / /02, the term Le can be further ex-
then the marginal distribution any time step is invariant. In the context of diffusion models, the starting distribution of the generative process p(zr) is isotropic Gaussian that naturally satisfies (13). To ensure that the final distribution defined by the entire generative process is invariant, the tran- sition density function p(z. zz) (8) needs to be equivariant, which requires the noise prediction network & (9) be an equivariant neural network. EGNN (Satorras et al., 2021) has been the most popular architecture for ø. One EGNN layer that takes x', h' as inputs and outputs x [+1 h4+1 is defined as:
1 (SN((1- 1) (NR())
-1)
Eo [1] (11)
C. === =FF--NOO
1
where es denotes eø <a + ore, t), the predicted noise that is output by the denoising network Φ.
To obtain the training objective, we follow the practice of DDPM (Ho et al., 2020) and discard the weighting in (I)).
Note that Liirr is constant irrelevant of optimization and p(æ|20) can be parameterized by separate Gaussian dis- tribution to make Lo have a similar form to (11), the final training objective is:
mij = (h n2y diy. aij);
n th (hi, Σ E:jm
((x)= =
[||c
c0ll"]
(12)
541
(15)
x x
x
1+1
x +
oz (h/ n/y. dry azz);
where is uniformly sampled between and T.
d' +
3

---

Scalable Non-Equivariant 3D Molecule Generation via Rotational Alignment
where d, = |||| 21l12 ll denotes the Euclidean distance, dtj optional edge attribute and ē. = Aiar (m.j) reweighs mes- sages coming from different atoms. Learnable parameters are composed of de, Ub, Φr. Pint, which are all implemented as ML.ss
Algorithm 1 Training algorithm for the autoencoder.
Inputs: atomic coordinates æ,, atom features h
Learnable parameters: rotation network Rs, encoder
En. decoder De
while not converged do
R t Rx h)
3. Method
p2: 14 + E, h)
Subtract center of gravity from μa.
While equivariant diffusion models have become the main- stream choice for 3D molecule generation, the equivariance constraints (13) and (14) may not be necessary for a good 3D molecule generator. In this section, we propose a method to relax such constraints and make non-equivariant diffusion models work comparably well. We are inspired by 3D vi- sion models that typically benefit from well-aligned datasets such as ShapeNet (Chang et al., 2015) and thus want to con- struct aligned representations for molecules. In Section 3.1, we introduce how we learn alignment with an autoencoder. In Section 3.2, we describe the non-equivariant diffusion models that run in the aligned latent space.
(Ez, 6n) (0 I)
Subtract center of gravity from Ex
Zy. 24, f- 4 + dc
Calculate the reconstruction loss [(0, n, v) (20)
Update 8, 1, vb end while
let M be the output of a vanilla GNN (i.e., the rotation net- work). Let u = [at, ht] denote the concatenation of atomic coordinates and features, GNN layer gives:
3.1. Aligned Latent Space
mij = (, u²)), uttt = Iuluu' Σ Enym))) (17) f
In this section, we aim to address the following two key questions: (1) How should we represent an alignment op- eration? and (2) How can we learn an effective alignment for each molecule? Following the practice introduced in Section 2.3, we deal with translations by subtracting the center of gravity for each molecule and restricting the diffu- sion model to the subspace. Then, the variation in atomic coordinates can only be caused by rotation R E SO(3). Therefore, aligning a molecule can be thought of as rotating it into a particular orientation. There are multiple ways to represent a rotation (e.g., Euler angles, exponential coor- dinates), but not all of them are good for gradient-band based learning (Brégier, 2021; Geist et al., 2024). Based on recent analysis (Geist et al., 2024), we choose to parameterize rotation with an arbitrary matrix M E R333 which is then projected to SO(3) using the singular value decomposition (SVD):
where Φe and 6.. are MLPs and ěgg is parameterized simi- larly to (15). Comparing (17) (15), the only difference is that the vanilla GNN layer discards the equivariant update of atomic coordinates and is thus non-equivariant, but it is much simpler. We take the average of all atom representa- tions output by the last layer and feed it into a 2-layer head to obtain M.
With a properly parameterized rotation representation, the next challenge is how learn good sample-dependent rota- tions. This is non-trivial since there is direct supervision signal. To overcome this, we propose to learn rotations in an unsupervised manner through an autoencoder. The in- tuition is that well-aligned molecules should facilitate the learning of a non-equivariant autoencoder. We note that existing work GeoLDM (Xu al., 2023) introduced a latent diffusion model for 3D molecule generation. However, both the encoder and the decoder used by GeoLDM are equivari- ant (parametrized by EGNNs (15)). Therefore, the relative orientations between different molecules cannot be adjusted in the latent space, and their latent diffusion model is also an equivariant one. Different from them, our approach can add more flexibility to the latent space, making more suitable for non-equivariant diffusion models.
R = SVD+ (M)- = Udiag (1, 1, dev V (16)
where U, VE E 2333 are obtained by SVD: M- = UEV
and det (UVT) ensures that det (R)= =
The above rotation representation is good for gradient-based optimization in the sense that SVD* (M) is smooth where det(M) * 0 (Levinson et al., 2020). Brégier (2021) pro- vided an efficient toolbox2 that contains this method and is well compatible with PyTorch's 's automatic differentiation. Furthermore, there is no restriction on the input M, mak- ing suitable for building on top of neural network. To generate sample-dependent rotation for each molecule, we
Formally, use 0 to denote the parameters of the GNN that generates rotation representations and use R R (zx h) to denote the rotation matrix generated for the molecule (x. h). Then the molecule after applying R, is denoted as (R9æ, h). Let E, denote the encoder parameterized by n and Dy denote the decoder parameterized by t/. then the
Phttps://gtt thub. com naver/rmaa
4

---

Scalable Non-Equivariant 3D Molecule Generation via Rotational Alignment
encoding and decoding processes are described as:
choice is transformer model since the attention mecha- nism is naturally permutation--qqiivarinn, We experiment with the diffusion transformer (DiT) proposed by Peebles & Xie (2023). The DiT essentially follows the standard transformer architecture (Vaswani et al., 2017) and injects conditional information (e.g., time steps) through adaptive normalization parameters output by a 2-layer MLP that takes as imput the conditional information. For a more detailed illustration of the DiT architecture, we refer readers to the original paper (Peebles & Xie, 2023), To adapt DiT to our setting, we remove its patch embedding module and covari- ance prediction head, and add attention masks to deal with varying atom numbers of different molecules. Since the latent representation こn (18) already encodes the positional information of atoms, we don't apply positional encoding.
(0, Zh h.) = N (E, (Rox, h), II) (Rga, h| Z,z 2h.) = Pv (Rox, h (Zz, 2n))
(18)
(19)
where 2nr E ERNX3 2n E Nxdd denote the latent represen- tations for x and h respectively. Note that both the input to the encoder & and the target of the reconstruction are the rotated molecule (Rææ, h). In (18) we parameterize the joint distribution of Z2; 2/, as an isotropic Gaussian with a fixed variance o², which allows (2,, Z4 2,, h)) be de- composed into the product of two Gaussians: 90.n (Z (xx h) and qo.n (z (e, h). As mentioned in Section 2.3, we restrict the distribution (02,n2z x.. h) to the subspace where the cen- ter of gravity is zero. The reconstruction distribution (19) also admits a factorization: Pa. (Rex) Rg P (f) R 2.).
The autoencoder and the latent diffusion model are trained separately. We first train the autoencoder following Algo- rithm 1 and then train the latent diffusion model with the loss function (12), over the latent representations that are sampled according to (18) from the fixed encoder. To gen- erate a molecule, we first sample a molecule size N from the categorical distribution p(N) of molecule sizes on the training set, and then run the reverse process of the latent diffusion model with N fixed. The output of the final step of the reverse process is decoded back to the original space.
with Pe (Raz 22, =n.) being a Gaussian in the zero-center- of oraggravty subspace and Py (h)=,, zn) being a categorical distribution if h is in the form of one-hot encoding for atom types. The reconstruction loss [(0, 77, 1/ is defined as:
C- = - - (e,,,, [log p (Rox, h 12|z,, =x]]
(20)
We write down the entire training algorithm for our autoen- coder in Algorithm 1. Following Xu et al. (2023), we adopt an early stopping training strategy as the regularization for the encoder.
4. Experiments
Regarding the specific architectural choices for E, and Das we use the same encoder architecture as GeoLDM (Xu et al., 2023) for the purpose of ablation study, which is a 1-layer EGNN (15). As for the decoder, we use a non-equivariant GNN (17) with the same number of layers as the GeoLDM decoder. The non-equivariance of the decoder is necessary to make the autoencoder sensitive to rotational transforma- tions, thus enabling the learning of aligned representations..
In this section, we present experimental results that empiri- cally validate the effectiveness of our proposed approach. In Section 4.1, we introduce the experimental setup, including the datasets, baselines and implementation details. In Sec- tion 4.2 and 4.3, we present the main results on molecule generation benchmarks. In Section 4.4, we show the results of ablation studies. In Section 4.5, we demonstrate the effi- ciency and scalability of our non-equivariant model. Finally in Section 4.6, we show the conditional generation results.
3.2. Non-Equivariant Latent Diffusion Model
With the aligned latent space, we hypothesize that non- cquivariant diffusion models can learn the data distribu- tion more easily. The noise prediction network of a non- equivariant diffusion model has more flexible architectural design space, since permutation-squivaaiinn only inductive bias of the architecture. A straightforward option is -quivaaivar GNN (17) that simply concate- nates atomic coordinates and features as the new atom fea- tures and discards the equivariant update of EGNN. The time step is treated as a scalar feature and is appended to the feature vector of every atom. Another promising
4.1. Experimental Setup
Datasets We first evaluate our approach using the QM9 dataset (Ramakrishnan et al., 2014) which is standard molecule generation benchmark widely used by related works. QM9 contains 130K molecules with up to 9 heavy atoms (29 atoms including hydrogens). Each molecule has 3D coordinates, atom types (H, C, N, O, F) and (integer- valued) charges as atom features. We split the dataset in the same way as Hoogeboom et al. (2022), with 100K, 18K, 13K samples for the train, validation and test partitions re- spectively. Next we evaluate our model on the larger GEOM- Drugs dataset (Axelrod & dataset contains 430K molecules with up to 181 atoms and 44.4 atoms on average. Following the setup of Hoogeboom et al. (2022), for each molecule we select the 30 lowest
is now the
'Let é = (x)))) denote the reconstructed coordinates. The reconstruction loss reduces to L2 loss - - x x=|3. If E and D were both equivariant, the 1.2 loss would be invariant to any rotation of (since rotation matrices are orthogonal). making the rotation network unable to receive any supervision signal.
2022). This
5

---

Scalable Non-Equivariant 3D Molecule Generation via Rotational Alignment
Table I: Results of atom stability, molecule stability, validity and validity x uniqueness. Higher numbers are better. We use distinct colors to indicate the best equivariant baseline, the best non-equivariant baseline and our proposed approach. Baseline results are copied from respective papers.
QM
GEOM-Drugs
Atom Sta (%)
Mol Sta (%)
Valid (%)
Valid & Unique (%)
Atom Sta (%)
Valid (%)
Data
99.0
95.2
97.7
97.7
86.5
99.9
ENF
85.0
4.9
40.2
39.4
G-SchNet
95.7
68.1
85.5
80.3
EDM
98.7
82.0
91.9
90.7
81.3
92.6
EDM-bridge
98.8
84.6
92.0
90.7
82.4
92.8
GeoLDM
98.9
89.4
93.8
92.7
84.4
99.3
GDM
97.0
63.2
75.0
90.8
GDM-aug
97.6
71.6
90.4
89.5
77.7
91.8
GraphLDMM
97.2
70.5
83.6
82.7
76.2
97.2
GraphLDM-aug
97.9
78.7
90.5
89.5
79.6
98.0
RADM
98.2++0.1
83.4:10.2
92.
90.6-10.3
83.8
98.9
RADM
98.5:10.0
87.3:10.2
94.1-10.1
91.7-10.1
85.0
99.3
energy conformations, For both datasets, model learns to generate the 3D coordinates and atom types of complete molecules with explicit hydrogens.
4.2. QM9
We train our model to generate molecules unconditionally. Following Hoogeboom et al. (2022), we use the distance between each pair of atoms and the atom types to decide the bond type (single, double, triple, or none). We don't ' use any chemical software (e.g., Open Babel) to edit the gen- erated molecules. We report atom stability (the proportion of atoms with the correct valence) and molecule stability (the proportion of molecules for which all atoms are sta- ble) to represent sample quality. We also report validity and uniqueness of the generated molecules (as measured by RDKit).
Baselines We compare our model with several representa- tive equivariant generative models for 3D molecules, includ- ing Equivariant Normalizing Flows (ENF) (Garcia Satorras et al., 2021), G-SchNet (Gebauer et al., 2019), EDM (Hooge- boom et al., 2022), EDM-bridge (Wu et al., 2022) and Ge- oLDM (Xu et al., 2023). Among them, EDM and GeoLDM are state-oftthe-a equivariant diffusion models. Besides, we compare with EDM's 's non- con-equivariant counterpart GDM, which uses non-equivariant GNN (17) the denoising net- work, and GDM-aug. which randomly rotates each molecule during training as data augmentation. Similarly, we com- pare with GeoLDM's non-equivariant versions GraphLDM and GraphL.DM-aug.
On QM9, we train autoencoder for 200 epochs using batch size of 64. Then we evaluate two variants of the non-equivariant diffusion model based on the same trained autoencoder: RADMDNTSS and RADMOTT.. RADM uses the small version DiT (defined in the DiT paper (Pee- bles & Xie, 2023)) as the noise prediction network, while RADMpIT- uses the base version. Compared with DiT-S, DiT-B has twice the hidden size and twice the number of heads. We adopt a batch size of 256 used in the DiT paper, and train both RADM and RADMpTT for around 5500 epochs.
Implementation Details Our implementation is based on the open-source code base of EDM* and GeoLDM to keep a fair comparison. We use the same hidden dimension and number of layers for the autoencoder as GeoLDM, and the number of layers of the rotation network is 2 on both datasets. As for the diffusion model, we use the same noise schedule and number of time steps as EDM/GeoLDD. To implement DiT the noise prediction network, we follow the official code b and make necessary modifications explained in Section 3.2. We train the autoencoder (and the rotation network) using the Adam optimizer with a learning rate of 1 x 10-1 and cosine annealing schedule. The latent diffusion model is also trained using Adam with a learning rate of x 10-4
The results are shown in Table 1. We report the average performance and standard deviation across three runs, each sampling 10000 molecules. For case of illustration, we cate- gorize baselines into equivariant models and non-equivariant models, and use distinct colors indicate the best model of each group. As we perform much better than ENF and G-SchNet, and equiv- ariant baselines significantly outperform non-equivariant baselines. Notably, both RADMDITTS and RADMDTT im- prove the performance drastically compared with previous
from the table, diffusion models
*htpps: / thub. com ccm//hooeeboom//e.. for aeformmelecte difuusion.formmmeeculees melecules
'https://githul com/Mo fom/Minnn nkai Xu wt/ GeoLDD
6

---

On GEOM-Drugs, the autoencoder is trained for 4 epochs using a batch size of 32. Then we train RADMOITS and RADMETT for around 55 epochs with a batch size of 256. Scalable Kon-Equivariant 3D Molecule Generation via Rotational Alignment
Figure 1: Molecules generated by RADMorT- QM9 (the three the left) and GEOM-Drugs (the three on the right).
non-equivariant models. RADMDITSS outperforms EDM molecule stability and outperforms both EDM and EDM- bridge validity. We find that scaling from DIT-S to DiT-B consistently boosts performance on all metrics. RADM greatly reduces the gap with the best equivariant model. It performs slightly worse than GeoLDM on molecule sta- bility, but achieves the best validity. We note that most syper-parametery used in our experiments simply follow the equivariant baselines (GeoLDM and EDM), and more tuning may further improve the performance of our non- equivariant model. We visualize the learned rotations in Appendix A.
Section 4.2. We refer this variant RADMGNN. The non- equivariant baselines GraphLDM and GraphLDM-art used the same GNN architecture the noise prediction network, but were trained in a latent space without learned alignment. We also experiment with a pre-processing method that ap- plies PCA to atomic coordinates and calculates new coordi- nates by treating the principal components as the new basis, The same diffusion model then trained on molecules with coordinates, which we name RADMANN (PCA). We train the diffusion models for 3000 epochs using batch size of 64. The results are shown in Table 2. We can see that RADM (PCA) outperforms GraphI.DM and GraphLDM-aus which supports the benefit of alignment. However, PCA ignores the dependencies between atomic coordinates and atom types that should be generated to- gether, and the signs of the principal components are biguous. The performance of RADMONN (PCA) is worse than RADMANN. which further supports our approach that learns alignment. Interestingly, we also find that RADMGNN performs better than RADMDITS. This is not very surpris- ing since the GNN (17) also has a global receptive field and RADMGNN is trained with more gradient updates (the batch size is smaller).
4.3. GEOM-Drugs
Following previous works, we report the atom stability and validity in Table 1. The performance is averaged over three runs, each sampling 10000 molecules (the standard devia- tion is negligible after rounding). From the table we observe that RADM exhibits very strong performance on this much larger dataset with more complex molecules. RADMDTT-S is already highly competitive compared with the best equivari- ant model GeoLDM. After scaling to DiT-B, RADMDITE outperforms GeoLDM on atom stability while obtaining the same validity.
Table 3: Comparison of model size, average training time per epoch and sampling efficiency.
#Params
Training
Sampling
Table 2: Results (as percentages) of the ablation study QM9 with the same diffusion backbone.
EDM GeoLDM
5.3M
107. 00.9s
55s / 100 samples
11.4M
118.2s
49s / 100 samples
RADMGNN
8.7M
79.6s
24s / 100 samples
Atom Sta
Mol Sta
Valid
RADM
37.6M
38.8s
7s / 100 samples
RADMpIT-B
134M
68.9s
19s / 100 samples
GraphLDM GraphL.DM-au
97.2
70.5
83.6
97.9
78.7
90.5
RADMGNN (PCA)
98.1
81.9
91.7
4.5. Efficiency Comparison
RADMGNN
98.6
85.8
93.8
In this section, we compare the training and sampling effi- ciency of EDM. GeoLDM and RADM. We list the param- eter counts, average training time per epoch and sampling speed in Table 3. All the numbers are measured RTX 4090 GPU. We add the number of parameters of the autoencoder but exclude its training time, since training the autoencoder takes only a small portion the time required to train the latent diffusion model. The sampling speed is measured by the average time used to generate 100 samples
4.4. Ablation Study
single
To validate the effectiveness of the alignment itself, we con- duct an ablation study using the same architecture for the dif- fusion model. Specifically, we use basic non aoneeviivaiant GNN (17) as the noise prediction network and train the dif- fusion model based on the same trained autoencoder as in

---

Scalable Non-Equivariant 3D Molecule Generation via Rotational Alignment
as one mini-batch, and we use 1000 sampling steps for all models. From the table we can see that in general our non- equívariant model RADM is significantly more efficient than EDM and GeoLDM. While RADMDrTS and RADMo have larger model sizes, due to the parallelizable transformer architecture, they remain highly efficient and offer further improved sampling speed. These results demonstrate that our equivariinn model has better scalability and great potential for further improvement.
ing objective via parameterizing a noise prediction model. Song al. (2021) unified DDPM and score matching with Langevin dynamics (Song & Ermon, 2019) under general framework defined through continuous stochastic differen- tial equations. Concurrently, Kingma et al. (2021) derived a continuous-time variational lower bound for diffusion mod- els by considering infinitely deep VAEs and simplified the expression in terms of the signal-to-noise ratio. Latent diffu- sion models (Rombach et al., 2022) reduced computational requirements by training on latent representations output by autoencoders. In addition to reducing the input dimension for diffusion models, autoencoders allow a more flexible latent space which could benefit the training of diffusion models. In contrast to most previous works that used a con- volutional U-net as the backbone for the denoising network, Peebles & Xie (2023) proposed to replace the U-net with transformer and obtained better performance while being more compute-efficient.
Table 4: Mean Absolute Error for molecular property predic- tion by a pretrained predictor (lower is better). The baseline performance is borrowed from EDM and GeoLDM.
Property
As
EHOMO mcV
ELUMO meV
C
Unit
Bohr²
VeV
D
cal K
QM9
0.10
64
39
36
0.043
0,040
Random
9.01
1470
645
1457
1.616
6.857
EDM
2.76
655
356
584
1.111
1.101
GcoLDM
2.37
587
340
522
1.108
1.025
5.2.
Molecule Generation
RADMOTE
1.98
458
290
383
0.814
0.869
Molecule generation has always been a central topic in drug discovery. The relevant literature be categorized accord- ing to target tasks. 2D graph generative models (Jin et al., 2020; Jo et al., 2022; Vignac et al., 2023a) only consider node types and edge types, and permutation equivariance is the symmetry of interest. Among them. Jin et al. (2020) trained a VAE to extend a molecular graph in an autoregres- sive way and used structural motifs as the building blocks. Jo et al. (2022) directly applied continuous-time diffusion models (Song et al., 2021) to node features and adjacency matrices, and Vignac et al. (2023a) utilized discrete dif- fusion models (Austin et al., 2021). Compared with the 2D structure, 3D coordinates of a molecule contain richer information about its biological activities, so the genera- tion of 3D molecules has attracted increasing interest. To deal with the Fuclidean transformations (i.e., translations and rotations) that can affect the 3D coordinates arbitrarily, EDM (Hoogeboom et al., 2022) used an EGNN (Satorras et al., 2021) as the noise prediction network of a diffusion model to ensure an invariant target probability distribution. Wu et al. (2022) enhanced EDM by designing physics in- formed prior bridges. GeoLDM (Xu et al., 2023) utilized latent diffusion models for 3D molecule generation. Both the autoencoder and the diffusion model of GeoLDM are still equivariant. Different from them, we demonstrate that equivariance is not necessary for a good molecule diffusion model. Concurrently, Zhang et al. (2024) and Joshi et al. (2025) also consider DiT as the backbone of the molecule diffusion model and learn stochastic equivariance through data augmentation. Our paper provides a different perspec- tive that improves nou-vvuivariant diffusion models even without data augmentation. Another line of related work is conformation generation (Xu et al., 2022; Wang et al., 2024) where the 3D conformation of a molecule is predicted
4.6. Conditional Generation
In this section, evaluate the ability of our model to gener- ate molecules with desired properties. To assess the quality of the generated molecules regarding the target property, we follow the protocol of EDM (Hoogeboom et al., 2022). The QM9 training set is equally divided into two parts. An EGNN is trained on the first half as a property predictor, while the generative model is trained on the other half. The property predictor is then evaluated on the samples drawn from the generative model. We report the MAE by the pre- trained predictor as the evaluation metric in Table 4. The numbers for "QM9" are obtained by evaluating the prop- erty predictor directly on the second half of the training set, which can be considered as lower bounds for the model's performance, while "Random" represents a worst-case base- line that evaluates the property predictor using randomly shuffled property labels. We test RADMpTT on this task. Specifically, our conditional model uses the same pretrained autoencoder as in Section 4.2, but trains the noise prediction network of the latent diffusion model conditioned on the target properties which are appended to each atom's 's latent representation. From Table 4, we observe that RADMDIT consistently outperforms both EDM and GeoLDM, further proving the effectiveness of our proposed approach.
5. Related Work
5.1. Diffuuion Models
Diffusion models were first invented in the context of ther- modynamics (Sohl-Dickstein et al., 2015). DDPM (Ho et al., 2020) established a connection between diffusion models and denoising score matching and derived a simple train-
8

---

Scalable Non-Equivariant 3D Molecule Generation via Rotational Alignment
from its 2D structure. Notably, Wang al. (2024) observed that a non-equivariant model performed better than equiv- ariant baselines. However, in their setting the 2D structure is given during both training and sampling. There are other works (Vignac et al., 2023b; Le et al., 2024) that model 2D and 3D information jointly. These models have access to the 2D structures during training, which different from our setting.
Su, H., al. Shapenet: An information-ricc 3d model repository. arXiv preprint arxiv: 2015.
Dao, T., Fu, D., Ermon, S., Rudra, A., and Ré, C. Flashat- tention: Fast and memory-efficient exact attention with io-awareness. In Advances in Neural Information Pro- cessing Systems, 2022.
Dhariwal, P. and Nichol, A. Diffusion models beat gans on image synthesis. In Advances in neural information processing systems, 2021.
5.3. Learned Canonicalization
Our work is also related to learned canonicalization (Kaba et al., 2023; Dym et al., 2024). In our approach, we don't force the aligned form of a molecule be exactly rotation- invariant, and more importantly, we learn the alignment in an unsupervised manner. A concurrent work (Sareen al., 2025) merges an equivariant canonicalization network with a non-equivariant denoising network of a diffusion model and therefore inherits the inefficiency of training equivariant diffusion models. In contrast, our method learns alignment using a lightweight non-equivariant autoencoder and then trains fully non-equivariant diffusion models in the aligned latent space.
Dosovitskiy, A., Beyer, L.,, Kolesnikov, A., Weissenbomn D., Zhai, X., Unterthiner, T., Dehghani, M., Minderer, M., Heigold, G., Gelly, S., Uszkoreit, J., and Houlsby, N. An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference on Learning Representations, 2021.
Dym, N., Lawrence, H., and Siegel, J. W. Equivariant frames and the impossibility of continuous canonicaliza- tion. In International Conference on Machine Learning, 2024.
Gareia Satorras, V., Hoogeboom, E.,, Fuchs, F., Posner, I., and Welling, M. E (n) equivariant normalizing flows. In Advances in Neural Information Processing Systems, 2021.
6. Conclusion
In this paper, we introduce a novel approach that signifi- cantly boosts non-equivariant diffusion models on the 3D molecule generation task, through constructing a rotation- ally aligned latent space. Our approach performs compa- rably to state-of-the -hrt equivariant diffusion models, with favorable scalability and improved efficiency.
Gebauer, N., Gastegger, M., and Schütt,, K. Symmetry- adapted generation of 3d point sets for the targeted dis- covery of molecules. In Advances in neural information processing systems, 2019.
Impact Statement
Geist, A. R., Frey, J., Zhobro, M., Levina, A., and Martius, G. Learning with 3d rotations, hitchhiker's guide to SO(3). In Forty-first International Conference on Ma- chine Learning, 2024.
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none of which feel must be specifically highlighted here.
Ho, J., Jain, A., and Abbeel, Denoising diffusion proba- bilistic models. In Advances in neural information pro- cessing systems, 2020.
References
Hoogeboom, E., Satorras, V. G., Vignac, C., and Welling, M. Equivariant diffusion for molecule generation in 3d. In International conference on machine learning, 2022.
Austin, J., Johnson, D. D., Ho, J., Tarlow, D., and Van Den Berg, R. Structured denoising diffusion models in discrete state-spaces. In Advances in Neural Information Processing Systems, 2021.
Jin, W., Barzilay, R., and Jaakkola, T. Hierarchical gen- eration of molecular graphs using structural motifs. In International conference on machine learning. 2020.
Axelrod, S. and annotated molecular conformations for property predic- tion and molecular generation. Scientific Data, 2022.
R. Geom, energy-
Jo, J., Lee, S., and Hwang, J. Score-based generative modeling of graphs via the system of stochastic differen- tial equations. In International conference on machine learning, pp. 10362-10383, PMLR, 2022.
Brégier, R. Deep regression on manifolds: a 3d rotation case study. In 2021 International Conference on 3D Vision (3DV), 2021.
Chang, A. X., Funkhouser, T., Guibas, L., Hanrahan, P., Huang, Q. Li, Z., Savarese, S., Savva, M., Song, S.,
Joshi, C. K., Fu, X., Liao, Y.-L., Gharakhanyan, V., Miller, K., Sriram, A., and Ulissi, Z. W. All-atom diffusion
9

---

Scalable Non-Equivariant 3D Molecule Generation via Rotational Alignment
transformers: Unified generative modelling of molecules and materials. arXiv preprint arXiv:2503.0m 03366, 2025.
Song, Y. and Ermon, S. Generative modeling by estimating gradients of the data distribution. In Advances in neural information processing systems, 2019.
Kaba, S.-O., Mondal, A. K. Zhang, Y., Bengio, Y., and Ravanbakhsh, S. Equivariance with learned canonicaliza- tion functions. In International Conference Machine Learning, 2023.
Song, Y., Sohl-Dickstein, J., Kingma, P., Kumar, A., Er- mon, S., and Poole, B. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations, 2021.
Kingma, D., Salimans, T., Poole, B., and Ho, J. Variational diffusion models. In Advances in neural information processing systems, 2021.
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, u., and Polosukhin, I. Atten- tion is all you need. In Advances in Neural Information Processing Systems, 2017.
Köhler, J., Klein, L.., and Noé, F. Equivariant flows: exact likelihood generative learning for symmetric densities. In International conference on machine learning, 2020.
Vignac, C., Krawezuk, I., Siraudin, A., Wang, B., Cevher, V., and Frossard, Digress: Discrete denoising diffu- sion for graph generation. In The Eleventh International Conference on Learning Representations, 2023a.
Langley, P. Crafting papers machine learning. In Langley, P. (ed.), Proceedings of the 17th International Conference on Machine Learning (ICML 2000), pp. 1207-1216, Stan- ford, CA, 2000. Morgan Kaufmann.
Vignac, C., Osman, N., Toni, L., and Frossard, Midi: Mixed graph and 3d denoising diffusion for molecule generation. In Joint European Conference on Machine Learning and Knowledge Discovery in Databases, 2023b.
Le, T., Cremer, J., Noe, F., Clevert, D.-A., and Schütt,, K. T. Navigating the design space of equivariant diffusion- based generative models for de novo 3d molecule genera- tion. In The Twelfth International Conference on Learn- ing Representations, 2024.
von Rütte, D., Fluri, J., Ding, Y., Orvieto, A., Schölkopf, B., and Hofmann, Generalized interpolating discrete diffusion. arXiv preprint 2ttt:::25330044488. 2025.
Levinson, J., Esteves, C., Chen, K., Snavely, N., Kanazawa, A., Rostamizadeh, A., and Makadia, A. An analysis of svd for deep rotation estimation. In Advances in Neural Information Processing Systemx, 2020.
Wang, Y., Elhag, A. A., Jaitly, N., Susskind, J. M., and Bautista, M. Á. Swallowing the bitter pill: Simplified scalable conformer generation. In Forty-first Interna- tional Conference on Machine Learning, 2024.
Peebles, W. and Xie, S. Scalable diffusion models with transformers. In Proceedings of the IEEE/CVF Interna- fional Conference on Computer Vision, 2023.
Wu, L., Gong, C., Liu, X., Ye, M., and Liu, Q. Diffusion- based molecule generation with informative prior bridges. Advances in Neural Information Processing Systems, 2022.
Ramakrishnan, R., Dral, P. O., Rupp, M., and Von Lilienfeld, O. A. Quantum chemistry structures and properties of 134 kilo molecules. Scientific data, 2014.
Xu, M., Yu, L., Song, Y., Shi, C., Ermon, S., and Tang, J. Geodiff: A geometric diffusion model for molecular conformation generation. International Conference on Learning Representations, 2022.
Rombach, R., Blattmann, A., Lorenz, D., Esser, P., and Ommer, B. High-resolution image synthesis with la- tent diffusion models. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, 2022.
Xu, M., Powers, A. S., Dror, R. O., Ermon, S., and Leskovec, J. Geometric latent diffusion models for 3d molecule gen- eration. International Conference on Machine Learn- ing. 2023.
Sareen, K., Levy, D., Mondal, A. K., Kaba, S.-O., Akhound- Sadegh, T., and Ravanbakhsh, S. Symmetry-awe aware gener- ative modeling through learned canonicalization. arXiv preprint ar arXiv:250J.07772 2025.
Yim. J., Trippe, B. L., De Bortoli, V., Mathicu, E., Doucet, A., Barzilay, R., and Jaakkola, T. Se (3) diffusion model with application to protein backbone generation, Inter- national Conference on Machine Learning. 2023.
Satorras, V. G., Hoogeboom, E., and Welling, M. E (n) equivariant graph neural networks. In International con- ference on machine learning, 2021.
Zhang, L., Ashouritaklimi, K., Teh, Y. W., and Cornish, R. Symdiff: Equivariant diffusion via stochastic symmetri- sation. arXiv preprint arXiv:2410.0622 06262, 2024.
Sohl-Dickstein, J., Weiss, E., Maheswaranathan, N., and Ganguli, S. Deep unsupervised learning using nonequi- librium thermodynamics. In International conference on machine learning, 2015.
10

---

Scalable Non-Equivariant 3D Molecule Generation via Rotational Alignment
Figure Molecule samples from the training set.
Figure 3: Molecule samples after applying the learned rotations.
A. Visualization
We visualize the learned rotations in Figure and Figure 3. We indeed found that after alignment, molecules tended to arrange common structural semantics (e.g., rings) in similar orientations.


Reviews:
[{"review_id": "5clhBkeZRd", "reviewer": "Reviewer_t7B8", "summary": "The paper suggests learning a sample-dependent rotational transformation during molecule generation. This approach aligns the molecules to specific directions, eliminating the necessity of employing equivariant models. The proposed method demonstrates promising performance and efficiency on benchmark datasets.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank Reviewer t7B8 for taking the time to review our paper and provide valuable feedback. We respond to the concerns as follows.\n\n**Relation to canonicalization**: We thank the reviewer for pointing this out! We admit that we overlooked the literature on canonicalization, and we will add them to the related work. Furthermore, we want to highlight some differences between our paper and existing works on canonicalization:\n1. Existing canonicalization works mostly consider the supervised learning setting, while our paper studies a generative modeling problem.\n2. The main computational cost of equivariant baselines is the training of diffusion model (e.g., ~3000 epochs on QM9). Our approach uses an autoencoder to learn canonicalization in a few epochs before training the diffusion model. The additional overhead from learning canonicalization is marginal compared with the great speedup of training diffusion models (Table 3).\n3. While theoretical works have analyzed the sample complexity gain of equivariance under the supervised learning setting, their conclusions may not be trivially extendable to generative modeling. Empirically, recent works (e.g., notable AlphaFold 3) also show the good performance of non-equivariant generative models. We believe this problem is worth further investigation, and our work is one step in this direction.\n4. Regarding the smoothness of canonicalization, the rotation estimation by Eq (16) is smooth except when det(M)=0, which doesn’t notably impede training in practice [1]. Besides, we conjecture that the discontinuity issue of canonicalization is less of a concern in our setting, because we don't test canonicalization on unseen data and only use it to reshape the training data of which we have full control.\n\n**Ablations on architectural choices**: Our experiments already contain ablations. Specifically, our baseline GeoLDM uses equivariant networks for both encoder and decoder. We use the same equivariant encoder as GeoLDM to keep the architectural factors affecting the generation of latent representations the same, except that our model applies the learned rotations. Next, to learn rotations, our decoder has to be non-equivariant; otherwise, the reconstruction loss for atom coordinates (L2 loss) would be invariant to rotations and prevent the rotation network from learning anything. Note that although we have an equivariant encoder, it is only used to generate latent representations, and the probability distribution is learned by a non-equivariant diffusion model. This doesn't violate the claim of the paper.\n\nFor the ablation results, we kindly refer the reviewer to the performance of GraphLDM and ALDM$\\_{\\text{GNN}}$ in Table 1. GraphLDM is a non-equivariant baseline from the GeoLDM paper that replaces the original equivariant denoising network of the diffusion model with a non-equivariant network (Eq 17). ALDM$\\_{\\text{GNN}}$ also uses Eq (17) as the denoising network for diffusion and the same hyper-parameters (noise schedule, diffusion steps, etc.). As a result, the significant performance gain of ALDM$\\_{\\text{GNN}}$ comes from the rotations applied to the latent space, which validates the effectiveness of our proposal. We also tried to use DiT for GraphLDM, but the gap was still significant.\n\n**Use of a canonicalization algorithm (PCA)**: We considered PCA in our preliminary experiments. There are two issues with PCA. First, the directions (or signs) of the principal axes are ambiguous. For prediction tasks, frame averaging can address this ambiguity, but it's not applicable to generative modeling. Second, applying PCA to atom coordinates ignores the dependencies between coordinates and atom types, which should be generated together. To further address the reviewer’s concern, we report here the results based on PCA:\nPCA$\\_{\\text{GNN}}$ achieves 82.8% molecule stability on QM9 and 81.1% molecule stability on DRUGS. Its performance is better than the best non-equivariant baseline GraphLDM but still significantly worse than ALDM$\\_{\\text{GNN}}$ (87.4% on QM9 and 83.0% on DRUGS)\n\n**Use of an equivariant canonicalization network**: We agree that using an equivariant canonicalization network (the rotation network in the paper) makes the canonical form of a molecule rotation invariant. However, to learn the specific canonical forms for different molecules that facilitate the reconstruction of the training set, the canonicalization network still needs to be trained jointly with the VAE (and probably less efficient). We'll try it in the updated version. Since our paper mainly aims to improve non-equivariant diffusion models, we believe the existing results are sufficient to support our claim.\n\n[1] https://arxiv.org/abs/2006.14616\n\nWe hope our response can provide a clearer explanation of the key points of our paper and kindly request Reviewer t7B8 to consider revising their score if they think our response can address their concerns. We are happy to answer any further questions.", "claims_and_evidence": "The claims presented in this paper are substantiated by references and detailed experiments.", "methods_and_evaluation_criteria": "First and foremost, I would like to emphasize that the authors’ concept of “aligning” molecules is known as *canonicalization* [1], and the resulting molecules are referred to as “canonical forms”. The authors inadvertently overlooked an extensive body of research on canonicalization. Notably, their method of learning a rotational matrix using a network is a weaker version of *learned canonicalization* [2,3], where the authors do not strictly enforce rotational equivariance on the canonicalization network. Since the authors’ approach is essentially canonicalization, it inherits the drawbacks of canonicalization methods. For instance, it imposes additional computational burden on the network to learn canonicalization and is unlikely to be more sample-efficient than equivariant networks. Furthermore, it lacks smoothness, meaning that a slight perturbation of the input molecule can result in a significant change in its canonical form.\n\nThe authors propose employing a non-equivariant network to learn the rotation matrix, thereby the learned “canonical form” $\\mathbf{R}\\_\\theta\\mathbf{x}$ is not guaranteed to be invariant to rotations. In contrast, the authors could consider using an equivariant network for this purpose as well, which would guarantee that $\\mathbf{R}_\\theta\\mathbf{x}$ remains strictly invariant. However, using an equivariant network may be computationally inefficient, so the authors should also explore the use of a deterministic canonicalization algorithm. For instance, the simplest approach to “align” a molecule in 3D space would be applying Principal Component Analysis (PCA). This method involves identifying the top 3 directions with the highest variance in atom positions and aligning them with the standard basis (the $x,y,z$ axes). This method is even more efficient than the authors’ approach.\n\nFurthermore, the authors assert that equivariant models are not essential for molecule generation since we only care about the overall probability of all possible positions. This argument holds true. However, in Section 4.1, it appears that the authors are still using an equivariant network for the encoder. The authors should provide justification for this choice through ablation experiments. Additionally, if we employ an equivariant canonicalization network or a deterministic canonicalization algorithm, the use of an equivariant encoder becomes unnecessary, as the aligned molecules will be guaranteed to be invariant.\n\n---\n[1] Ma, G., Wang, Y., Lim, D., Jegelka, S., & Wang, Y. A Canonicalization Perspective on Invariant and Equivariant Learning. In The Thirty-eighth Annual Conference on Neural Information Processing Systems.\n\n[2] Kaba, S. O., Mondal, A. K., Zhang, Y., Bengio, Y., & Ravanbakhsh, S. (2023, July). Equivariance with learned canonicalization functions. In International Conference on Machine Learning (pp. 15546-15566). PMLR.\n\n[3] Sareen, K., Levy, D., Mondal, A. K., Kaba, S. O., Akhound-Sadegh, T., & Ravanbakhsh, S. (2025). Symmetry-Aware Generative Modeling through Learned Canonicalization. arXiv preprint arXiv:2501.07773.\n\n---\n\n**Update: I had previously overlooked that [3] is a concurrent work released on January 15 of this year. In light of this, I have updated my score accordingly.**\n\n**I would like to clarify my earlier rebuttal comment. When I stated that the paper \"lacks empirical evidence demonstrating that canonicalization reduces the variance of latent representations across different input orientations,\" I meant that the authors should explicitly measure and report the variance of the learned latent representations, and compare this to appropriate baselines. This would provide stronger support than the current anecdotal examples where certain molecular ring structures appear similarly aligned. Such qualitative examples are insufficient to substantiate the claim that performance gains arise from improved alignment in latent space due to canonicalization.**\n\n**I appreciate the additional ablation studies provided in the rebuttal and follow-up comment, and I recommend that the authors include these results in the paper, as they are essential for its completeness. Additionally, prior work on canonicalization should be acknowledged and discussed in the camera-ready version, which may require substantial revisions to the current writing.**", "theoretical_claims": "There are no theoretical claims in this paper.", "experimental_designs_or_analyses": "The authors provided sufficient details about their experiment settings, and the performance and efficiency of their method on real-world tasks seem promising. However, as mentioned, this paper lacks several ablation experiments.\n\n1. Comparing the use of a non-equivariant encoder or an equivariant decoder would help justify the authors’ architectural decisions.\n\n2. Comparing their method with an equivariant canonicalization network or a canonicalization algorithm which ensure the rotational invariance of the aligned molecules.\n\n3. Additionally, the authors claim that the aligned latent space reduces variations caused by Euclidean symmetries, but this assertion lacks empirical evidence. The authors could validate this claim through experiments, demonstrating a smaller variance in the latent space compared to non-equivariant baselines.", "relation_to_broader_scientific_literature": "The proposed method enhances the performance and efficiency of 3D molecule generation, benefiting the broader AI for science community.", "essential_references_not_discussed": "The paper lacks references of the research on canonicalization. This should encompass, but is not limited to, the following:\n- Ma, G., Wang, Y., Lim, D., Jegelka, S., & Wang, Y. A Canonicalization Perspective on Invariant and Equivariant Learning. In The Thirty-eighth Annual Conference on Neural Information Processing Systems.\n- Tahmasebi, B., & Jegelka, S. Generalization Bounds for Canonicalization: A Comparative Study with Group Averaging. In The Thirteenth International Conference on Learning Representations.\n- Dym, N., Lawrence, H., & Siegel, J. W. (2024, July). Equivariant Frames and the Impossibility of Continuous Canonicalization. In International Conference on Machine Learning (pp. 12228-12267). PMLR.\n- Ma, G., Wang, Y., & Wang, Y. (2023). Laplacian canonization: A Minimalist Approach to Sign and Basis Invariant Spectral Embedding. Advances in Neural Information Processing Systems, 36, 11296-11337.\n- Kaba, S. O., Mondal, A. K., Zhang, Y., Bengio, Y., & Ravanbakhsh, S. (2023, July). Equivariance with learned canonicalization functions. In International Conference on Machine Learning (pp. 15546-15566). PMLR.", "other_strengths_and_weaknesses": "The strengths and weaknesses are adequately discussed in the preceding points.", "other_comments_or_suggestions": "It would be more readable to bold the best performance in Tables 1 and 2.", "questions_for_authors": "I have no additional questions.", "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": "qc9nkys9Zz", "reviewer": "Reviewer_uMCt", "summary": "The paper explores a non-equivariant alternative to protein generation. It uses an explicit rotation network to rotate zero-centered molecule coordinates and builds a latent space on top of the rotated coordinates. The latent space is learned via VAE objective and encodes aligned features. This allows one to learn a non-equivariant diffusion model on the latent space, as motivated by those used in aligned 3D point cloud literature. The method achieves comparable performance as equivariant counterparts for generation while achieving faster sampling speed.\n\n## Update after rebuttal\n\nThe authors have adequately addressed my concerns regarding the non-equivariant decoder and scaling experiments. I increased score accordingly.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank Reviewer uMCt for taking the time to review our paper and provide valuable feedback. We respond to the concerns as follows:\n\n**Scalability of non-equivariant diffusion models**: In fact, we did experiments to show the scalability of our model. We kindly refer the reviewer to Table 3 on page 8, where we show the model sizes. We tested non-equivariant denoising networks of three different sizes for the diffusion model: GNN is from Eq (17), DiT-S is the small version model from the diffusion transformer paper, and DiT-B is the base version with a larger hidden dimension (4 times \\#params as DiT-S and more than 10 times \\#params as GeoLDM). Our largest model achieves the best performance on the DRUGS dataset and still maintains higher training/sampling efficiency than existing equivariant diffusion models.\n\n**Why is non-equivariant decoder necessary to learn alignment?**: Let $x$ denote the original atom coordinates and $x' = \\mathcal{D}(\\mathcal{E}(x))$ denote the reconstructed coordinates. The reconstruction error reduces to L2 loss $\\Vert x' - x \\Vert^2$ for continuous coordinates. Now suppose $x$ is rotated by $R$, i.e., $Rx$. If both the encoder and decoder are equivariant, the reconstruction would become $Rx'$, and the L2 loss between $Rx$ and $Rx'$ would be the same as $\\Vert x' - x\\Vert^2$ (because $R$ is orthogonal), making the rotation network unable to receive any supervision signals. Therefore, encoder and decoder should not both be equivariant. \n\nWe chose to make the decoder non-equivariant and keep the encoder equivariant because we wanted to compare directly with the baseline GeoLDM. GeoLDM uses equivariant networks for both encoder and decoder. We use the same equivariant encoder as GeoLDM to keep the architectural factors affecting the generation of latent representations the same, except that our model applies the learned rotations.\n\n**Ablations on the architecture and latent space**. Our experiments already contain ablations. We kindly refer the reviewer to the performance of GraphLDM and ALDM$\\_{\\text{GNN}}$ in Table 1. GraphLDM is a non-equivariant diffusion model baseline from the GeoLDM paper that replaces the original equivariant denoising network of the diffusion model with a non-equivariant network (Eq 17). ALDM$\\_{\\text{GNN}}$ also uses Eq (17) as the denoising network for diffusion and the same hyper-parameters (noise schedule, diffusion steps, etc.). Therefore, the difference between ALDM$\\_{\\text{GNN}}$ and GraphLDM lies in that the latent space of ALDM$\\_{\\text{GNN}}$ is aligned by rotations. The experimental results show a significant performance gain of ALDM$\\_{\\text{GNN}}$ compared with GraphLDM, which validates that the aligned representations improve the learning of non-equivariant diffusion models. We also tried to use DiT for GraphLDM, but the gap was still significant.\n\nAdditionally, we randomly sample a mini-batch of molecules from the training set and compare their original and rotated positions. The visualization is contained in the following anonymous link:\nhttps://docs.google.com/document/d/e/2PACX-1vQeBYa4zljEoIM_h-4cvlOH-unSQFJhDgwVuq0FxtISIXZJ-T9T8q53OoWOKtoMR0d4O1SA7Vp12xq3/pub\nWe indeed found that common structural semantics (e.g., rings) are easier to recognize after rotation. Since our encoder is equivariant, these rotations are applied to the latent representations in the same way and thus make them more aligned.\n\n**Novelty of the proposed method**.\nWhile the evaluation of novelty depends on one's own perspective, we want to emphasize that our paper is not simply “a change of architecture”. In fact, in this paper, we switch from equivariant models to non-equivariant models and propose an approach to close the gap between them. As we show in the experiments (explained above), the improvement brought by our approach is significant even if we use the same non-equivariant architecture (i.e., a basic GNN) for the diffusion model.\n\nBesides, our model not only improves the sampling speed, but also significantly saves training cost. On QM9, training the baseline equivariant diffusion model (e.g., GeoLDM) to convergence takes nearly 3000 epochs, which needs 4 days to finish on a single 4090 GPU (in Table 3 on page 8 we show the average training time per epoch). As a comparison, the training of our largest model ALDM$\\_{\\text{DiT-B}}$ takes 3 days on a 4090 GPU. We use a larger batch size and thus more epochs for ALDM$\\_{\\text{DiT-B}}$, but the total training cost is still reduced significantly. We believe this provides more support for our model.\n\nWe hope our response can provide a clearer explanation of the key points of our work and kindly request Reviewer uMCt to consider revising their score if they think our response can address their concerns. We are happy to answer any further questions.", "claims_and_evidence": "- The authors also claim that the non-equivariant diffusion model allows more \"scalability\" while no experiments are done to validate that this model \"scales\".\n\n- The authors claim that non-equivariant decoder is necessary to learn the aligned latent representation. Why is this the case? The paper does not further elucidate the intuition and does not design experiments to validate the claim. \n\n- The design of the encoder makes it such that the latent representation is supposedly aligned, which is the core of the method. However, the authors do not conduct experiments on investigating whether the learned latent spaces are actually aligned. Some simple PCA or other analysis can be done to investigate it further.", "methods_and_evaluation_criteria": "The paper mostly follows previous papers in terms of datasets and in terms of evaluating generation quality.", "theoretical_claims": "The method does not contain theories.", "experimental_designs_or_analyses": "Pros:\n- The paper validates that the non-equivariant diffusion is faster due to parallelizable transformer network by showing sampling wall-clock time. It similarly shows by experiments that the generation performance is comparable.\n\n\nCons:\n- One simple baseline the paper is missing is simply to train fully non-equivariant autoencoder and non-equivariant latent diffusion with the same architecture. This is to validate the claim that the supposedly aligned latent feature and the newly introduced rotation network is actually useful. it is not clear whether the retained performance comes from the aligned latent space or the new transformer itself.\n- Following the above point, it is also not clear whether the latent space is actually aligned or could it be optimized away by some non-equivariant encoders.", "relation_to_broader_scientific_literature": "The paper aims to advance core algorithms of drug generation and increase efficiency of generation while retaining quality.", "essential_references_not_discussed": "Related works are properly discussed.", "other_strengths_and_weaknesses": "Strength:\n- The writing is very clear and easy to follow.\n\nWeaknesses:\n- The proposal is incremental and novelty is limited.\n- The results show marginal improvement in terms of quality and some improvement in sampling speed only due to a change of architecture. The results do not strongly justify the reason for adopting this method.\n- The paper claims aligned feature helps, but does not actually conduct experiments to check that the latent space is aligned and the improvements are due to this design rather than other factors such as changing to a different transformer architecture.", "other_comments_or_suggestions": "As mentioned above.", "questions_for_authors": "I have listed all my concerns in previous sections and wish the authors to address. I can consider increasing score if they are properly addressed.", "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": "PtYOOksOsN", "reviewer": "Reviewer_Biv1", "summary": "This paper proposes learning a roto-aligned latent space for 3d molecule generation. The alignment is achieved through learning sample-wise rotation with auto-encoder. More specifically, ALDM adopts an equivariant encoder and a non-equivariant decoder, largely alleviating the constraints in architecture design, making Transformers potential for large-scale pretraining of small molecules. Experiments on QM9 and GEOM-DRUG shows ALDM performs better than non-equivariant baseline, and on par with SOTA equivariant models.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank Reviewer Biv1 for taking the time to review our paper and provide valuable feedback. We are glad to see that they found our approach interesting and our results good enough to support our claims. We respond to the concerns as follows:\n\n**Lack of stronger baselines**: We thank the reviewer for pointing this out! We will add more recent baselines (e.g., [1]) in an updated version of our paper. The performance of our non-equivariant models is still highly competitive compared with stronger equivariant baselines, and as the reviewer said, we believe the existing experimental results are sufficient to validate the effectiveness of our approach.\n\n[1] https://arxiv.org/abs/2403.15441\n\n**Lack of analysis of learned alignment and latent representation**. \nWe provide more analysis below. First, for a single molecule, its canonical form can be obtained (e.g., by applying PCA to atom coordinates). However, the challenge is that such canonicalization may not have consistent effects for different molecules. For example, the directions (or signs) of the principal components obtained by PCA are ambiguous. Moreover, applying PCA to atom coordinates ignores the dependencies between coordinates and atom types, which should be generated together. Therefore, we believe it makes more sense to learn canonical forms across the training set by maximizing the overall reconstruction quality.\nAs an ablation, we report the results based on PCA:\nPCA$\\_{\\text{GNN}}$ achieves 82.8% molecule stability on QM9 and 81.1% molecule stability on DRUGS. Its performance is better than the best non-equivariant baseline GraphLDM but still significantly worse than ALDM$\\_{\\text{GNN}}$ (87.4% on QM9 and 83.0% on DRUGS).\n\nTo analyze the rotations learned by the autoencoder, we randomly sample a mini-batch of molecules from the training set and compare their raw positions and the positions after rotations. We provide some visualization in the following anonymous link:\n\nhttps://docs.google.com/document/d/e/2PACX-1vQeBYa4zljEoIM_h-4cvlOH-unSQFJhDgwVuq0FxtISIXZJ-T9T8q53OoWOKtoMR0d4O1SA7Vp12xq3/pub\n\nWe indeed found that common structural semantics (e.g., rings) of some molecules are easier to recognize after the alignment.\nSince our encoder is equivariant, such rotations are applied to the latent representations in the same way and thus make them more applied\n\nWe hope our response can address the reviewer's concerns. We are happy to answer any further questions.", "claims_and_evidence": "Equivariance constraints may not be necessary for molecule generation.\nSuccessfully supported.", "methods_and_evaluation_criteria": "Yes", "theoretical_claims": "N/A", "experimental_designs_or_analyses": "1. Unconditional generation\nAblation of decoder choice is interesting, and it shows DiT's advantage over simple EGNN. Overall comparison is decent, ALDM achieves competitive results with GeoLDM. \n2. Conditional generation\nThe results look normal. Even though property-conditioned generation is not so useful.", "relation_to_broader_scientific_literature": "N/A", "essential_references_not_discussed": "There are some stronger baselines, yet the reviewer believes the results are enough to support main claims.", "other_strengths_and_weaknesses": "Strengths\n1. The idea is simple, but elegant and effective, making non-equivariant Transformer more applicable to large scale pretraining of molecules. The scaling law and emergent ability of Transformer have already been proved in other domains.\n2. Generally good results, making the claim convincing.\nWeakness\n1. Lack of stronger baselines.\n2. Lack of analysis of learned alignment (the sample-wise rotation) and latent representation. The reviewer is quite curious about: 1. Does the learned rotation have some physical or chemical insights? Is there a natural canonical orientation for 3d molecules? 2. Does the VAE really learn chemical/structural semantics in latent space? More analysis or visualization would be insightful.", "other_comments_or_suggestions": "No", "questions_for_authors": "See weakness.", "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": "IY4rnqT0mK", "reviewer": "Reviewer_bfzT", "summary": "This paper proposes Aligned Latent Diffusion Model (ALDM) to improve 3D molecule generation by relaxing SE(3) equivariance constraints in diffusion models. Instead of enforcing equivariance, ALDM learns a sample-dependent SO(3) transformation using an autoencoder to align molecular representations, enabling the use of non-equivariant diffusion models like vanilla GNNs and transformers. This approach significantly improves scalability and efficiency while maintaining state-of-the-art sample quality, outperforming previous non-equivariant models and matching equivariant baselines in molecule stability and validity.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank Reviewer bfzT for taking the time to review our paper and provide valuable feedback. We respond to the concerns as follows.\n\n**Not SOTA on QM9**: Our paper mainly aims to improve non-equivariant diffusion models. On QM9, compared with the best non-equivariant baseline GraphLDM-aug, all three variants of our model improve the performance drastically, which validates the main claims of our paper. In fact, our best model ALDM$\\_{\\text{DiT-B}}$ achieves the highest Validity (94.2%) on QM9 across all (equivariant and non-equivariant) baselines and is only negligibly worse than GeoLDM on other metrics. We kindly request the reviewer to also consider the efficiency improvement shown in Table 3. Overall, our model achieves the best performance-efficiency trade-off on QM9. Additionally, on GEOM-DRUGS, which is much larger and more challenging than QM9, our model achieves the best performance.\n\n**Difference between Eq. 17 and regular GNN**: Eq. 17 is indeed a regular GNN. We use it as a basic non-equivariant architecture (for the decoder and the noise prediction network of the diffusion model). Unlike EGNN, it simply concatenates atom coordinates and types and is thus non-equivariant. In Table 1, both ALDM$\\_{\\text{GNN}}$ and the baseline GraphLDM (from the GeoLDM paper) use Eq.17 as the noise prediction network, except that learned rotations are applied to the latent space of ALDM$\\_{\\text{GNN}}$. The performance gain validates that the aligned latent space can improve the training of non-equivariant diffusion models.\n\n**Alignment Quality Uncertainty**: The lack of supervision is an inherent challenge for generative modeling. Our approach manages to learn rotations in such an unsupervised setting, which we believe is part of our contribution rather than a weakness of our model. Furthermore, the theme of our paper is to show that “alignment” is better than “no alignment” for non-equivariant diffusion models, which is already supported by our experiments. We agree with the reviewer that investigating the “optimality” of the aligned representations would be interesting, but we think it is currently beyond the scope of this paper. We’d love to explore this in future work.\n\n**Limited Theoretical Justification on Performance Trade-offs**: In theory [1], diffusion models are able to learn the dataset distribution as long as the score function is learned well, without imposing any constraints (e.g., equivariance) on the target distribution. Therefore, non-equivariant diffusion models, in principle, have the full potential to learn the dataset distribution. On the contrary, previous equivariant diffusion models took equivariance for granted and didn’t justify why equivariance is necessary.\nAdditionally, a few recent works (e.g., [2] (AlphaFold 3), [3]) also demonstrate the superior performance of non-equivariant diffusion models, though the specific problems they address are different from ours.\n\n**Code and detailed experimental settings**: We already detailed the experimental setup in the main body of the paper (as mentioned by Reviewer t7B8: “The authors provided sufficient details about their experiment settings”). Specifically, we describe the dataset information in Section 5.1. We follow exactly the dataset preprocessing and splits of baselines EDM and GeoLDM and point to their public repositories. In Section 5.2, we introduce hyper-parameters (learning rate, number of layers, hidden dimension, batch size, etc.) we use before we present the experimental results. We will organize these information in a better format (e.g., using tables) in the next version. Our code will be released after the paper is accepted.\n\n[1] https://arxiv.org/abs/2011.13456 \\\n[2] https://www.nature.com/articles/s41586-024-07487-w \\\n[3] https://proceedings.mlr.press/v235/wang24q.html\n\nWe kindly request Reviewer bfzT to consider raising their score if they think our response can address the concerns. We are more than happy to answer any further questions!", "claims_and_evidence": "The claims made in the submission are supported by clear and convincing evidence", "methods_and_evaluation_criteria": "The proposed method is evaluated on appropriate datasets with other baselines.", "theoretical_claims": "Algorithms in this paper is correct.", "experimental_designs_or_analyses": "Yes, The proposed method is evaluated on appropriate datasets with other baselines.", "relation_to_broader_scientific_literature": "NA", "essential_references_not_discussed": "NA", "other_strengths_and_weaknesses": "Strengths:\n1. Improved Efficiency & Scalability – By removing SE(3) equivariance constraints, ALDM enables the use of simpler, more scalable architectures (e.g., GNNs, transformers), reducing computational costs and improving sampling speed.  \n2. Flexible Model Design – Unlike equivariant models that require specialized architectures (e.g., EGNNs), ALDM allows for standard non-equivariant GNNs and transformers, making implementation and optimization more accessible.  \n3. Well-Motivated Theoretical Insight – The paper challenges the necessity of equivariance in 3D molecule generation, providing a principled alignment-based approach that removes redundancy while maintaining symmetry-aware representations.  \n\nWeakness & Questions:  \n1.\tThe performance compromises. ALDM does not achieve the SOTA performance in each column on QM9 in Table 1.  \n2.\tWhat’s the difference between Eq. 17 and regular message-passing GNNs? EGNN mainly convert positions to distance to achieve the equivariance. Without this design, Eq. 17 becomes a regular GNN.   \n3.\tAlignment Quality Uncertainty – The autoencoder-based alignment is learned in an unsupervised manner, making it unclear whether the learned SO(3) transformations always provide an optimal latent representation for downstream diffusion modeling.  \n4.\tLimited Theoretical Justification on Performance Trade-offs – While the paper argues against the necessity of equivariance, it does not provide a rigorous theoretical analysis explaining why non-equivariant models can fully match or surpass equivariant ones under all conditions.  \n5.\tNo code available. No appendix showing more information. E.g. the dataset information, the detailed experimental settings, the optimal hyperparameters….", "other_comments_or_suggestions": "NA", "questions_for_authors": "See Weaknesses & Questions", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper presents an approach to unlock scalability in 3D molecule generative models by relaxing the equivariant constraints in previous architectures. The proposed approach first learns a sample dependent SO(3) transformation and then a generative model is trained in this aligned space. The paper is well written and the experimental results support the claims made in the paper. Reviewers have pointed out the simplicity of the approach and the elegant formulation of the problems as strengths. The AC agrees that this is paper is a valuable contribution for ICML.
Decision: Accept (poster)