Geometric Representation Condition Improves Equivariant Molecule Generation
Authors: Zian Li
Cai Zhou
Xiyuan Wang
Xingang Peng
Muhan Zhang
Venue: ICML 2025 spotlightposter (2025)
Keywords: molecule generation
equivariant generative models
representation
geometric deep learning
diffusion models
Primary area: deep_learning->generative_models_and_autoencoders
TL;DR: We propose a two-stage, model-agnostic generative approach that effectively leverages molecule representations to improve the generation quality of molecule generative models.

Abstract: Recent advances in molecular generative models have demonstrated great promise for accelerating scientific discovery, particularly in drug design. However, these models often struggle to generate high-quality molecules, especially in conditional scenarios where specific molecular properties must be satisfied. In this work, we introduce GeoRCG, a general framework to improve molecular generative models by integrating geometric representation conditions with provable theoretical guarantees. We decompose the generation process into two stages: first, generating an informative geometric representation; second, generating a molecule conditioned on the representation. Compared with single-stage generation, the easy-to-generate representation in the first stage guides the second stage generation toward a high-quality molecule in a goal-oriented way. Leveraging EDM and SemlaFlow as base generators, we observe significant quality improvements in unconditional molecule generation on the widely used QM9 and GEOM-DRUG datasets. More notably, in the challenging conditional molecular generation task, our framework achieves an average 50\% performance improvement over state-of-the-art approaches, highlighting the superiority of conditioning on semantically rich geometric representations. Furthermore, with such representation guidance, the number of diffusion steps can be reduced to as small as 100 while largely preserving the generation quality achieved with 1,000 steps, thereby significantly reducing the generation iterations needed.

Geometric Representation Condition Improves Equivariant Molecule Generation
Zian Li Cai Zhou *34 Xiyuan Wang 12
Xingang Peng Muhan Zhang
Abstract
1. Introduction
Recent advances in molecular generative models have demonstrated great promise for accelerating scientific discovery, particularly in drug design. However, these models often struggle gener- ate high-quality molecules, especially in condi- tional scenarios where specific molecular proper- ties must be satisfied. In this work, we introduce GeoRCG, a general framework to improve molec- ular generative models by integrating geometric representation conditions with provable theoret- ical guarantees. We decompose the generation process into two stages: first, generating infor- mative geometric representation; second, gener- ating a molecule conditioned on the representa- tion. Compared with single-stage generation, the easy-to-generate representation in the first stage guides the second stage generation toward a high- quality molecule in a goal-oriented way. Lever- aging EDM and SemlaFlow base generators, we observe significant quality improvements in unconditional molecule generation on the widely used QM9 and GEOM-DRUG datasets. More notably, in the challenging conditional molecular generation task, our framework achieves an av- erage 50% performance improvement over state- of-the-art approaches, highlighting the superiority of conditioning on semantically rich geometric representations. Furthermore, with such repre- sentation guidance, the number of diffusion steps can be reduced to as small as 100 while largely preserving the generation quality achieved with 1,000 steps, thereby significantly reducing the gen- eration iterations needed. Code is available
Recent years have seen rapid development generative modeling techniques for molecule generation (Garcia Sator- ras et al., 2021; Hoogehoom al., 2022; Luo & Ji, 2022; Wu et al., 2022; Xu et al., 2023; Le al., 2023; Morehead & Cheng, 2024), which have demonstrated great promise in ac- celerating scientific discoveries such as drug design (Graves et al., 2020). By representing molecules as point clouds of chemical elements embedded in Euclidean space (poten- tially with edges (Vignac et al., 2023; Irwin et al., 2024)) and employing equivariant models such as EGNN (Satorras et al., 2021) backbone architectures, these approaches ensure the 0(3)- (or SO(3)-) invariance of the modeled molecule probability and have shown significant progress in both unconditional and conditional molecule generation tasks.
Despite the advances, precisely modeling the molecular distribution ((M) still remains a challenge, with current models often falling short of satisfactory results. This is especially true in more practical scenarios where the goal is to capture the conditional distribution q(Mlcc for condi- tional generation, with e representing a desired property such as the HOMO-LUMO gap. In such cases, recent models still produce molecules with property errors signifi- cantly larger than the data lower bound (Hoogeboom et al., 2022; Xu et al., 2023). This challenge arises in part because molecules are naturally supported in lower-dimensional manifold (Mislow, 2012; De Bortoli, 2022; You et al., 2023), yet they are emhedded in a 3D space with much higher am- bient dimensions (N x (3 + d), where N is the number of atoms and the atom feature dimension). Consequently, directly learning these distributions without additional guid- ance or conditioning solely a single property can result in substantial errors (Song et al., 2021), often leading to unstable or undesirable molecular samples.
Equal contribution 'Inttitute for Artificial Intelligence, Peking University, Beijing, China "School Intelligence Seience and Technology, Peking University. Beijing, China 'Department of Electrical Engineering and Computer Science. Massachusetts In- stitule of Technology, Cambridge, MA, USA 'Department of Au- tomation, Tsinghua University, Beijing, China. Correspondence to: Muhan Zhang <muhan @ pku.edu.cn
In this work, we propose GeoRCG (Geometric- Representation- Conditioned Molecule Generation), general framework for improving the generation quality of molecular generative models by leveraging geometric representation conditions for both unconditional and conditional generation; see Figure for an overview of the framework. At a high level, rather than directly learning the extrinsic molecular distribution, we aim to
Proceedings of the 42nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Geometric Representation Condition Improves Equivariant Molecule Generation
a) Training (parallel)
b) Sampling (sequential)
Representation Generator
Representation Generator
T. P/4
Pre-trained
Fi (r)
Pa
rep. condition
Geometric
O
rep. condition
Encoder
Molecule Generator
Molecule
Generator
Po(M Ir)
Pa
M, P=00 Ir=x)
M = (t, E E x
Figure 1: Training and sampling procedure of GeoRCG for unconditional molecule generation. a) During training, each molecule M is mapped into representation r by pre-trained, frozen geometric encoder E. The representation distribution is then learned by a lightweight representation generator. The molecule generator is trained in a self-conditioned manner, generating a molecule M conditioned on its own representation E(M). b) During sampling, an informative representation first generated, which subsequently guides the molecule generator to produce high-quality molecules.
first transform it into a more compact and semantically meaningful representation distribution, with the help of a well-pretrained geometric encoder E such as Unimol (Zhou et al., 2023) and Frad (Feng et al., 2023). This distribution is much simpler because it does not exhibit any group symmetries, such as (3)) and S(N) groups which are present in extrinsic molecular distributions. As result, a lightweight representation generator (Li et al., 2023) can effectively capture this simple distribution. In the second stage, we employ a conditional molecular generator to achieve the ultimate objective: molecular generation. Unlike conventional approaches, our molecular generator is directly informed by the first-stage geometric representation, which encapsulates crucial molecular structure and property information. This guidance enables the generation of high-quality molecular structures with improved fidelity.
unconditional and conditional generation, offering a rigor- ous understanding of the improved performance.
To illustrate the effectiveness of our approach, we select one of the simplest and most classical equivariant gen- erative models, EDM (Hoogeboom et al., 2022), as the base molecular generator of GeoRCG, For better perfor- mance on the more challenging dataset GEOM-DRUG, we also apply GeoRCG onto the recent state-of-the-art (SOTA) model SemlaFlow (Irwin et al., 2024). Experimentallyy our method achieves the following significant improvements:
Substantially enhancing the quality (e.g., molecule sta- bility) of the generated molecules on the widely used QM9 and GEOM-DRUG datasets. On QM9, GeoRCG not only improves the performance of EDM by large margin, but also significantly surpasses several recent baselines with advanced performance (Wu et al., 2022; Xu et al., 2023; Morehead & Cheng, 2024; Song et al., 2024a). On GEOM-DRUG, GeoRCG also significantly improves EDM's s performance, and consistently en- hances SemlaFlow's 's already SOTA results.
Our approach is directly inspired by RCG (Li et al., 2023), which, however, focuses on image data with fixed sizes and positions and does not necessitate handling Euclidean and permutation symmetries factors that are markedly different in molecular data. Compared to recent work GraphRCG (Wang et al., 2024) which applies the RCG framework to 2D graph data, we explicitly handle 3D ge- ometry that is more complex due to the additional Eu- clidean symmetry. Moreover, we avoid the complicated step- wise bootstrapped training and sampling process proposed in Wang et al. (2024) that requires noise alignment, sequen- tial training, and simultaneous encoder training. Instead, we adopt a simple and intuitive framework that enables parallel training and leverages advanced pre-trained geometric en- coders containing valuable external knowledge (Zaidi et al., 2022; Feng et al., 2023), thus achieving competitive results without complex training procedures. Notably, while Li et al. (2023) primarily focuses on empirical evaluation, we also provide generic theoretical characterizations of the representationnconnittneeedd diffusion model class for both
More remarkably, in conditional molecule generation tasks, GeoRCG yields an average 50% improvement in performance (i.c., difference of generated molecule's property with conditions), while many contemporary models struggle to achieve even marginal gains.
By incorporating classifier-free guidance into the molccule generator (Li et al., 2023) and employing low- temperature sampling for representation generation (In- graham et al., 2023), GeoRCG demonstrates a flexible trade-off' between molecular quality and diversity on QM9 dataset without additional training, which is es- pecially advantageous in specific molecular generation tasks that prioritize quality over diversity.
With the assistance of the representation guidance, GeoRCG significantly reduces the number of diffusion steps required by approximately 10x, while preserving the quality of molecular generation.
2

---

Geometric Representation Condition Improves Equivariant Molecule Generation
2. Related Works
2023), and complex hybrid noise (Ni al., 2024; Feng et al., 2023; Jiao et al., 2024). Notably, Zaidi et al. (2022) showed that denoising equilibrium structures effectively corresponds to learning the underlying force field, thereby produeing molecular representations that are physically and chemically informative.
Molecular Generative Models. Early work has primarily focused on modeling molecules as 2D graphs (composed of atom types, connections, and edge types), utilizing 2D graph generative models to learn the graph distribution (Vi-
gnac et al., 2022; Jang et al., 2023; Le et al., 2023; Jo et al., 2023; Luo et al., 2023; Zhou et al., 2024). How- ever, since molecules inherently exist in 3D space where physical laws govern their behavior and spatial geometry provides critical information related to key properties, re- cent research has increasingly focused on leveraging 3D generative models to directly learn the geometric distribu- tion by modeling molecules point clouds of chemical elements. Notable early autoregressive models include G- SchNet (Gebauer et al., 2019) and G-SphereNet (Luo & Ji, 2022). More recently, diffusion models have demonstrated effectiveness in this domain, as evidenced by models like EDM (Hoogeboom et al., 2022) and subsequent advance- ments that enhance EDM with latent space (Xu et al., 2023), prior information (Wu et al., 2022) and more powerful back- bones (Morehead & Cheng, 2024). Furthermore, recent advances in flow methods (Lipman et al., 2022; Liu et al., 2022b) have inspired the development of geometric, equiv- ariant flow methods including EquiFM (Song et al., 2024b) and GOAT (Hong et al., 2024), which enable much faster molecule generation speed. Beyond these, there are also methods that jointly model 2D and 3D information (Vignac et al., 2023; You et al., 2023; Huang et al., 2024; Irwin et al., 2024) (also called 3D graph (You et al., 2023)), where rep- resentative methods include MiDi (Vignac et al., 2023) and SemlaFlow (Irwin et al., 2024) that jointly learn atom types, bond types, formal charges and coordinates.
Latent Generative Models. At high level, our framework can also be viewed as a latent generative model, where data distributions are learned in a latent space (our stage 1) and decoded back through some decoder (our stage 2). Most prior work in this domain either focuses on regular data forms (e.g., images) with fixed positions and sizes (Van Den Oord et al., 2017; Razavi et al., 2019; Dai & Wipf, 2019; Aneja et al., 2021; Rombach et al., 2022; Li et al., 2023), or on graph data without Euclidean symmetry and requires explicit modeling (Wang et al., 2024). Molecular data, however, presents unique challenges in both aspects. One of the key issues in this context is how to define the latent space-defining it as "latent coordinates and features" as in GeoLDM (Xu et al., 2023) still results in a geometri- cally structured and thus complex space, while defining it on representations as we do introduces the challenge of effec- tively "decoding" global, non-symmetric embedding back into geometric objects. LGD (Zhou et al., 2024) trains diffusion model on a unified Euclidean latent space obtained by jointly training a powerful encoder and a simple decoder, and performs both generation and prediction tasks focusing on 2D graphs. LDM-3DG (You et al., 2023) adopts repre- sentation latent space but employs a cascaded (2D+3D) auto- encoder (AF) framework, where the decoder is designed (or trained) to be deterministic, rendering poor performance on the 3D part as evidenced in our experiments. In contrast, we model the decoder as a powerful generative model, focusing solely on geometric learning while demonstrating superior effectivvness.
Pre-training for Molecular Encoders Learning mean- ingful molecular representations is crucial for downstream tasks like molecular property prediction (Fang et al., 2022). The strategy of pre-training on large-scale datasets followed by fine-tuning on smaller, task-specific datasets has been proven to significantly improve model performance in vi- sion and language domains (Kenton & Toutanova, 2019; Brown, 2020; Dosovitskiy, 2020). Building this suc- cess, recent studies have explored pre-training methods for molecular data, aiming to achieve similar performance im- provements (Zhou et al., 2023; Feng et al., 2023; Liu et al., 2022a; Fang et al., 2022; Jiao et al., 2024; Ni et al., 2024). Common pretext tasks involve masking and recovering atom types, bond bond lengths, or bond angles (Fang et al., 2022: Zhou et al., 2023). However, since molecules exist in continuous 3D space, more effective approach is introduced by adding carefully crafted noise into the molecular coordinates and training the model to denoise it. Examples of such noise types include isotropic Gaussian noise (Zaidi et al., 2022; Zhou et al., 2023), Riemann- Gaussian noise (Jiao et al.,
3. Methods
3.1. Preliminaries
In this work, we represent molecules point clouds of chem- ical elements in 3D space, denoted by M = (x, h). where x = (XI...... xx)) E R N3 represents the atomic coordi- nates of N atoms, and h = (hi.... hy) E T xd cap- tures the node features of dimension d, such as atomic num- bers and charges. This formulation follows the approach of Hoogehoom et al. (2022); Xu et al. (2023); Morehead & Cheng (2024) and is widely adopted in molecular represen- tation learning (Thomas et al., 2018; Li et al., 2024a; Zaidi et al., 2022), facilitating the integration of pre-trained molec- ular encoders (Zaidi et al., 2022; Feng et al., 2023). After generating point clouds of chemical elements, these methods infer bond types using lookup tables based on atom types and pairwise distances, or relying on advanced packages
3

---

Geometric Representation Condition Improves Equivariant Molecule Generation
like OpenBabel (O' (O'Boyle et al., 201 Notably, approaches like MiDi (Vignac al., 2023) and SemlaFlow (Irwin al., 2024) additionally represent molecules with explicit bond types, enabling joint learning and generating of 2D and 3D information, which typically results in improved perfor- mance. We use to denote the underlying data distribution, such as molecule distributions q(M),, and p to denote the approximated distributions captured by parametrie models.
size of the representation's 's underlying molecule and the size of the molecule it guides to generate. Moreover, molecules with different sizes often have distinct modes in structures and properties (Hoogeboom et al., 2022), which is reflected in their geometric representations learned by modern pre- trained geometric encoders (Zhou et al., 2023; Feng et al., 2023), as shown in Figure 2. From the figures, is evident that by conditioning on N. the learning process for the rep- resentation generator becomes simpler and more effective, leading to the following loss function of our representation generator:
We denote the pre-trained geometric encoder E ( x (RN - Rd,, which embeds a molecule M with an arbitrary number of nodes N into a represen- tation vector r of fixed dimension d,. The geometric en- coder exhibits E(3)- (or SE(3)-) invariance, suggesting that F(M) = [xx, x, h) = F(xR" + t, h) for any E R and R E ((3) (or (3)3)) where ((3 is the set of orthogonal matrices (and S0(3) being the set of special orthogonal matrices).
Cup E ww. (/N [||r- [111-156h Solritt N))1
(I)
where D (E(M) N(M) representing atom number of . and Dimo de- noting the molecule dataset. Here, S&p is the MLP back- bone (Li et al., 2023), and r = Vayr + -T- oze is the noisy representation computed with the predefined schedule a₁ E (0, 1].
E Dine ). with
3.2. GeoRCG: Geometric- Semenetri-Reepreeeaaaioo-oooiooo Molecular Generation
Molecule Generator. Since the ultimate goal of our frame- work is to generate molecules from (MM), we decompose the molecular distribution as (MM) I ((M(r)q(r) dr to explicitly enable geometric-reeresennatiio conditions. Con- sequently, a geeometie-rppreeenttiinooooioiii ular generator pe (M]r is required. In principle, we can use many modern molecule generators (Hoogeboom et al., 2022; Xu et al., 2023; Morehead & Cheng, 2024; Irwin et al., 2024), as these models can all take additional conditions.
Geometric Representation Generator. To improve the quality of the generated molecules, we propose to first trans- form the geometrically structured molecular distribution ((M) into a non- genmeeremttic representation distribution (r)) using a well-pretrained geometric encoder F. that maps each molecule M its representation r. Learning the representa- tion distribution (r)) is considerably easier, since represen- tations do not exhibit any symmetry as in explicit molecular generative models (Hoogeboom et al., 2022). We thus lever- age a simple yet effective MLP-based diffussin architecture as proposed in (Li et al., 2023) for the representation gen- crator Pv (r). which follows DDIM schemes (Song et al., 2020a) for training and adopts predictor-corres frame- works for sampling (Song et al., 2020b).
molec-
To illustrate the effectiveness of our approach, we choose a relatively simple model EDM (Hoogeboom et al., 2022) as the base generator and primarily demonstrate our method with it. Furthermore, we showcase the generality of our ap- proach by adapting a recent flow-matching based SOTA model, SemlaFlow (Irwin et al., 2024), emphasizing its ability to consistently improve SOTA models' performance.
EDM is designed to ensure the ((3)-invariance, i.e., for any R E O()), PM(M) po h) = Pa pa(xR h). To accommodate EDM to representation conditions, we use the following training objective:
-25
-189
Eiol = E ((Mr) Din 101 (o
Figure 2: t-SNE visualizations of the representations pro- duced by Frad (Feng et al., 2023) for the QM9 dataset (left) and by Unimol (Zhou et al., 2023) for the GEOM-DRUG dataset (right). The representations exhibit clear clustering based on node count.
(2)
Puabl ) and sam- (ar) (h) from pling from No 1) entails drawing eo = (Eb") ..5
where Dmoo-rrpp ((M,
D
(O, I). adjusting ((²) by subtracting its geometric cen- ter to obtain c(*), and setting e (<(r),e)") This en- sures the zero center-of-mass property, as the distribution is defined on this subspace to ensure translation invari- ance (Hoogeboom et al., 2022). The noisy molecule
One additional design compared to previous practices (Li et al., 2023; Wang et al., 2024) is that we condition the rep- resentation generator on the molecule's 's node number N by default', This is crucial to ensuring consistency between the
and mathematical derivations for statement simplicity, as its inclu- sion does not affect the overall framework and conclusions.
'ee omit the condition N in our probability decompositions
4

---

Geometric Representation Condition Improves Equivariant Molecule Generation
(M) given by M, = of x, hj + oi ( schedules of (M) and oz (M), while the diffusion backbone fo. which is instantiated with EGNN (Satorras et al., 2021), is
e, with time-dependent
Define po which is the ending point of the reverse pro- cess starting from 4hp instead of Gaussian noise. Here, 4r is the T-th step in the forward process starting from 0 == A J. ( (r)) dr, where A is the normaliza- tion factor. Denote the k-dim isotropic Gaussian distribu- tion as x Then the following holds.
conditioned on r.
Combining the Two Generators Together: The represen- tation generator P.: (r) and the molecule generator pe (M]r) together model the molecular distribution (p) S (p dr, which approximates the data distri- bution (M) = ((M)r))(()) dr that we aim to capture. One notable advantage of the framework is that the decom- position enables parallel training of the two generators. The entire training and sampling procedure is summarized in Algorithm
(p)) 3 zy ((L(t))ll
convergonce of forward process
(3)
(Lx NLa 3)h + L, m
discretization crror
(4)
Theoretical Analysis of GeoRCG. There are several key properties of GeoRCG that facilitate high-quality molecule generation. First, GeoRCG preserves symmetry properties of the base molecule generator po(M):
6 conn VT
(5)
conditional score estimatioo error
TV (go 90)
(6)
Proposition 3.1. (Symmetry Preservation) Assume the original molecular generator po(M) is 0(3)- or SO(3)- invariant. Then, the two-stage generator PM. is also 0(3)- or -Oiia-invaiinnt
representation generation crror
Balancing Quality and Diversity of Molecule Genera- tion. In many scientific applications, researchers priori- tize generating higher-quality molecules over more diverse ones. To facilitate this, we introduce a feature that allows fine-grained control over the trade-off between diversity and quality in the sampling stage (thus without retraining). This is achieved by integrating two key techniques: low- temperature sampling (Ingraham et al., 2023) (controlled via the temperature T) for the representation generator, and classifier-fre guidance (Ho & Salimans, 2022; Zheng et al., 2023) (controlled via the coefficient w) for the molecule gen- erator. We provide more details about the two techniques in Appendix A. The combination of the two techniques enables flexible and explicit control, which we refer to as "Balancing Controllability" and demonstrate its effective- ness in Section 4.2.
Proof. This result follows directly from the defini- tion. Specifically, (M. = f po Po(M)r (r) dr = [ .h/r) (r( dr Px.o(xR h) for any R e 0(3) (or (3()). The second equality holds due to the symmetric property of (MM), which remains valid when additional non-symmetric conditions are applied.
diffusion models can achieve no higher overall total variation distance than tradi- tional diffusion models, and can arguably yield better results, as the representation encodes key data information that may further reduce estimation error. We present the rigorous bound in Theorem 3.2, and provide corresponding proof and detailed discussions in Appendix D. Remarkably, this is a generic theoretical characterization that applies to prior experimental work (Li et al., 2023). For models that account for equivariant symmetries such as EDM, we build upon results from (Feng et al., 2024; You et al., 2023) to establish finer-grained bounds, as detailed in Theorem D.14.
Moreover, representatiioncondiiioi
Handling Conditional Molecule Generation. The frame- work discussed thus far focuses on unconditional molecule generation, where no specific property e (e.g., HOMO en- ergy) is prespecified. However, for molecule generation, more practical and desired scenario is conditional (also called controllable) generation, where additional condi- tions c, such as the HOMO-LUMO gap energy, are intro- duced, and our objective shifts to generating molecules from the distribution g(Mcc). In GeoRCG, this condi- tional generation is naturally decomposed as (M]l) := I po (e(M)) (r|c) dr, , suggesting that we first generate a prroperty-maaningee
Theorem 3.2. Consider the random variable R (d+-3) ~ q(x), and assume that the second moment my of x is bounded as m := Eg(ts) I|| - =||]] < oc, where = Eg(a)(a [r). Further, assume that the score V la q(xx) is L. Li-siischitz,fr all t, and that the score estimation error in the second-stage diffusion bounded by 9.odd such that (r-p,(r), a (s)) (0 (zt,,,r) r)
E
In q( (PP1 < 2 "....cond Denote the step size as h == T/Na, where T is the total diffusion time and Na is the number of dis-
molecular representation r, which is then independently used to condition the second-stage molecule generation; see Figure 3 for an illustration. A key advantage of this modeling approach is that, when different properties (e.g., HOMO, LUMO, GAP energy) need to be captured, only the representation generator
cretization steps, and assume that h 3 L/L. Suppose that we sample x ~ po (2)r) from Gaussian noise, where r ~ (r),), and denote the final distribution of x as (x))
5

---

Geometric Representation Condition Improves Equivariant Molecule Generation
requires retraining under the new conditions. This retrain- ing is highly efficient due the lightweight nature of the representation generator. Notably, GeoRCG demonstrates outstanding conditional generation performance, as shown in Section 4.3. Moreover, we theoretically demonstrate that, under mild assumptions, the representation generator can provably estimate the conditional distribution and generate representations that lead to provable reward improvements toward the target, which subsequently benefits the second- stage generation. Further theoretical details are provided in Appendix D.2.
10k randomly sampled molecules, averaged over three runs, with standard errors reported in parentheses.
Instantiation of the Pre-trained Encoder. We employ Frad (Feng et al., 2023), which was pre-trained on the PCQM4My2 dataset (Nakata & Shimazaki, 2017) using a hybrid noise denoising objective, as the geometric en- coder for QM9 dataset. For GEOM-DRUG, we adopt Uni- mol (Zhou et al., 2023) architecture but perform our own pretraining using the dataset from (Zhou et al., 2023), with GEOM-DRUG included as an additional pretraining dataset. This is because GEOM-DRUG contains unique chemical elements not found in PCQM4My2 or other commonly used pretraining datasets such as ZINC or ChemBL (Li et al., 2021). We note that, when using Frad (Feng et al., 2023) as the encoder, GeoRCG also leads to significant improve- ments on the GEOM-DRUG dataset, although with slightly lower performance compared to Unimol; see Appendix C.
Rep. Generator
Malecule Generator
Rep. Generator
uncond.
Rep. Generator
Baselines, A direct comparison is made with our base molecule generators, EDM or SemlaFlow. For GeoRCG (EDM), we compare it against generative models that, like EDM, do not explicitly generate bonds but instead infer them based on bond lengths. Although this approach may be less effective in generating valid molecules, it is widely adopted and presents a greater challenge for generative models in learning 3D geometric cestributions-pprciisell where our geometric representation guidance offers the most significant improvement. These models include: (1) the non- equivariant counterparts of EDM and GeoLDM (Xu et al., 2023). specifically GDMCAUUGHHoogeooom et al., 2022) and GraphL.DM(-AUGGXX et al., 2023); (2) the autoregres- sive method G-SchNet (Gebauer et al., 2019); (3) advanced equivariant diffusion models such as GeoLDM (Xu et al., 2023), EDM-Bridge (Wu et al., 2022), and GCDM (More- head & Cheng, 2024); (4) fast equivariant flow-based meth- ods like E-NF (Garcia Satorras et al., 2021), EquiF (Song et al., 2024b), and GOAT (Hong et al., 2024); and (5) the recently introduced Bayesian-based method GeoBEN (Song et al., 2024a).
Figure 3: A single molecule generator can be employed for both unconditional and conditional molecule genera- tion with respect to various properties. For conditional generation, only the representation generator is re- trained on (molecule, property) pairs, allowing it to conditionally sample cooerrty-maaningee representations during the sam- pling stage.
4. Experiments
4.1. Experiment Setup
Datasets and Tasks. As a method for 3D molecule generation, we evaluate GeoRCG on the widely used datasets QM9 (Ramakrishnan et al., 2014) and GEOM- DRUG (Gebauer et al., 2019; 2022; Axelrod & Gomez- Bombarelli, 2022). We focus on two tasks: unconditional molecule generation, where the goal is to sample from q(M), and conditional (or controllable) molecule gener- ation, where property e is given, and we aim to sample from (MMc)).
For GeoRCG (Semla), we compare it with recent advanced 2D&3D methods that directly generate bonds to produce higher-quality samples similar to SemlaFlow, including MiDi (Vignac et al., 2023) and EQGAT-diff (Le et al., 2023).
We use "GeoRCG (EDM)" to denote the variant of GeoRCG that employs EDM as the base molecule generator, "GeoRCG (Semla)" to refer to its application built upon Sem- laFlow (Irwin et al., 2024), and "GeoRCG" when the context clear or for general purpose. To ensure fair comparisons, we follow the dataset split and configurations exactly as in Anderson et al. (2019); Hoogeboom et al. (2022); Xu et al. (2023). Without further clarification, we bold the highest scores and underline the second-highest one. Addi- tionally, to highlight the direct improvement over the base model, we display green numbers next to the score to indi- cate the improvement, and red numbers to denote a decrease. Without further elarification, results are calculated based on
Note that we intentionally separate the comparison between EDM-like 3D-only models and SemlaFlow-like 2D&3D models, focusing on the improvements brought by GeoRCG to the base model. This is because combining the com- parisons would be unfair, as 2D&3D models additionally learn bond information, which reduces the complexity of generating valid molecules (Morehead Cheng, 2024).
We provide further experiments, including ablation studies on the pre-trained encoder, in Appendix C.
6

---

Geometric Representation Condition Improves Equivariant Molecule Generation
Table I: Unconditional molecule generation on QM9 and GEOM-DRUG. The gray cells denotes the base molecule generator
employed in GeoRCG.
QM9
DRUG
Metries
Atom Sta (%)
Sta (%)) f
Valid (%)
Valid Unique (%)
Atom Sta (%)
Valid (%)
Methods
Data
95.2
97.7
97.7
86.5
99.9
G-Schnet
95.7
68.1
85.5
80.3
GDM
97
63.2
75
90.8
GDM-AUG
97.6
71.6
90.4
89.5
77.7
91.8
GraphLDM
97.2
70.5
83.6
82.7
76.2
97.2
GraphL
97.9
78.7
90.5
89.5
79.6
98
EDM
98.7
82
91.9
90.7
81.3
92.6
EDM-Bon Bridge
98.8
84.6
90.7
82.4
92.8
GeoLDM
98.9(0.11
89. 89.4(0.5)
93.8(0.4)
92.7(0.5)
84.4
99.3
GCDM
98.7(0.0)
85.7(0.4)
94.8(0.2)
93.3(0.0)
89
95.5
ENF
85
4,9
40.2
39.4
EquiFM
98.9(0.1)
88.3(0.3)
94.7(0.4)
93.5(0.3)
98.9
GOAT
98 4
84.1
90.9
89,99
81.8
96.0
GcoBFN
99.08(0.03)
90.87(0.1)
95.31(0.1)
92.96(0.1)
85.6
92.08
GcoRCG (EDM)
99.12(0.03) 0.43%
92.32(00.06) 12.598
96.52(0.2)
92.45(0.22
84.3(0.12)
98.5(0..12)
4.2. Unconditional Molecule Generation
Table Unconditional molecule generation on GEOM- DRUG for 2D&3D methods. Molecule stability and validity reported percentages, while energy and strain energy are expressed in kcal-mo mol- Results marked with were reproduced in our own experiments.
We first evaluate the quality of unconditionally generated molecules from GeoRCG, with the commonly adopted va- lidity and stability metrics for assessing molecules' qual- ity (Hoogeboom al., 2022). See Appendix for detailed descriptions of these metrics.
MIDI EQGAT-aif
99.8
91.6
77.8
We present the main results of GeoRCG on the QM9 and GEOM-DRUG datasets in Table 1 and Table 2. Below, we highlight the key findings: (i) Improvement over the base model: By leveraging geometric representations, GeoRCG significantly outperforms the base model, on both QM9 and GEOM-DRUG datasets. Notably, on QM9, GeoRCG (EDM) increases stable molecules from 82% to 93.9% and validity from 91.9% 97.4%. while also improving molecule uniqueness. (ii) Superior performance com- pared to advanced methods: GeoRCG (EDM) also sur- passes included advanced models on the QM9 dataset. On the GEOM-DRUG dataset, GeoRCG (EDM) outper- forms models such as EDM-Bridge and GOAT, and gets high score in validity. Although GeoRCG (EDM) falls short of achieving the best performance, we attribute this to the relatively limited capabilities of EDM. To address this, we replace EDM with the recent SOTA flow-matching based model, SemlaFlow (Irwin et al., 2024), as the base model on the GEOM-DRUG dataset, as shown in Table 2. As demonstrated. GeoRCG (Semla) consistently enhances SemlaFlow's SOTA performance across all metrics on the GEOM-DRUG dataset.
99.800.11
93.4(021))
94.60024)
1488009
140.200.7
SemlaFlow
(98/0.00)
97.400077
94.40.17)
95.72(1.24)
56.42(1.07)
GeRRGG (Semia)
99.800.0)
97.60.00)
95.30.13)
(86(1)33)
47.64(1110)
(see Appendix C for the extended figure that includes va- lidity). The results indicate a clear trend: increasing w and decreasing T improve validity and stability at the expense of uniqueness, allowing for -raiied flexible control over molecule generation. At its best, this approach achieves molecule stability of 93.9% and validity of 97.42%, ap- proaching the dataset's 's upper bound, with trade-off in lower validity&uniqueness of 86.82%.
Figure 4: Balance controllable generation on QM9 of GeoRCG (EDM). Increasing w and decreasing T enhances stability, with the cost of a reduction in uniqueness.
Molecule Stabdity
Chaetteer quideno contictee
56.36
91.53
91.38
36.76
36.00
65.04
55.45
05.45
91.72
9137
92.15
46.37
22:45
23.35
20.30
We proceed to investigate the "Balancing Controllablilit" " feature of GeoRCG introduced in Section 3.2. To this end, we conducted grid search by varying both w and T on
W0.00
22.23
Invecss Iareperature
invess Iamperature
QM9 dataset for GeoRCG (EDM), as depicted in Figure 4
In Appendix C, we present additional experiments on QM9,
7

---

Geometric Representation Condition Improves Equivariant Molecule Generation
Table 3: Conditional molecule generation on QM9. The metric used is the MSE between the target property value and the classifier predicted value. The gray cells denote the baseline molecule generator employed in our proposed approach. Models marked with indicate results obtained from our own experiments; these are provided only as a coarse reference due to potentially differing evaluation criteria, see Appendix B for details.
Properties
Methods
Aa
FHOMO
LLUMO
C2
QM9 (lower bound)
0.1
64
39
36
0.043
0.04
Random
9.01
1470
646
1457
1.616
6.857
N_atoms
3.86
866
426
813
1.053
1.971
EDM
2.76
356
1.111
1.101
GeoLDM
2.37
587
340
522
1.108
1.025
GCDM
1.97
602
344
479
0.844
0.689
EquiPM
2.41
591
337
530
1. 106
1 1.033
GOAT
2.74
605
350
534
1.01
0.883
LDM-3DG
12.29
1160
583
1093
1.42
5.74
GeoBFN
2.34
S77
328
516
0.998
0.949
GeoRCG (EDM)
0.86(0.01)
325.2(3.4)
50.35%
202.2(1.2)
43.20%
257.9(5.5) 55.84%
(.805(0.000) 27.54%
0.475(0.000)) 56.86%
demonstrating that GeoRCG enhances distribution-level geometric metries, such as BondAngleWI, which under- score GeoRCG's improved geometric learning capabilities.
molecule from the training dataset, allowing the molecule generator to simply recover its full conformation based on that representation. This could lead to small property er- rors but lack of novelty, To address this, we conducted a thorough evaluation of the generated molecules across each property. finding that the novelty (the proportion of new molecules not present in the training dataset) remains com- parable to other methods. Additionally, the conditionally generated molecules demonstrate much higher molecule sta- bility than EDM (Hoogeboom et al., 2022). Further details can be found Appendix C.
4.3. Conditional Molecule Generation
We now turn to more challenging task: generating molecules with a specific property value e from (MMcc.. We strictly follow the evaluation protocol outlined in (Hooge- boom al., 2022). Speicifically, QM9 is split into halves, and an EGNN classifier (Satorras et al., 2021) is trained on the first half for evaluating the generated molecules' property, while the generator is trained on the second half. We focus on six properties: polarizability (a), orbital energies (EhOMo, êTuMO), their gap (Ae), dipole mo- ment (μ), and heat capacity (C.).
Table 4: Unconditional molecule generation on QM9 with fewer diffusion steps. The blue cells indicate the highest value among methods with the same number of diffusion steps, while bold font emphasizes values that outperform all other methods across all diffusion steps.
The results are presented in Table 3. The first three base- lines, as introduced by EDM (Hoogeboom et al., 2022), represent the classifier's inherent bias as the lower bound for performance, the random evaluation result as the upper bound, and the dependency of properties on N. For more details, please refer Appendix
Metrics
Meihods
Stops
Atom Sta (%)
Mol Sta (%)
Valid (96)
Data
99
95.2
97.7
EquiFM GOAT
200
98 910.
88.300.)
94,700.4
198
EDM
50
97.000
66.4(0:2)
Bridge
97.3/0.1
69.2/0.2)
As shown, GeoRCG (EDM) nearly doubles the performance of the best existing models for most properties, with an erage 50% improvement over the best ones. This is task where many recent models struggle to make even modest improvements, as evidenced in the table. Notably, for differ- ent properties, we only re-train the representation generator, as demonstrated in Section 3.2, significantly saving training time. In Figure 5, we visualize the generated samples, which exhibit minimal property errors and display clear trend as the target values increase. Additional randomly generated molecules are provided in Appendix E.2.
GcoBEN
50
98.28/0.1 1)
(5..5
92.27(0.4)
GeeRCG
95.05(0.35)
EDM
100
97 30
69.8(0:2)
EDM Bridge
100
97. 97.900.11
72.3(0.2)
GooBFN
100
98.61/0.0 1)
87.21(0.3)
93.03/0.3)
GoORCG (EDM)
100
99.08/0.33)
91 (55855304)
96.49(0.27)
EDM
5X 5/0
81.2(0.1)
EDM Bridge
500
98 73,700.11 1)
83.7(0.1)
GooBFN
98.78/0.8)
$8.12/0.2)
93.35(0.2)
GaRRC (EDM)
5001
99 w9.0%/0.01 nosrs
91,8 (1.8900.24) USUP
96.57(0.12)
FDM
T0KI
MA.7
R2
EDM-Bridge
1000
98.8
84.6
GooBFN
1000
99 99.08/0.05)
90.87/0.2)
95.31(0.1)
GeoRCG (EDM)
1000
99.12/0.03) aoss
92.32(0.00) 13.95
96.52(0.2)x10z
4.4. Fewer-Step Generation
A potential concern is that for given property value c, P. (r(c) may produce a representation corresponding to a
With geometric representation condition, it is reasonable to expect that fewer discretization steps of the reverse diffu-
8

---

Geometric Representation Condition Improves Equivariant Molecule Generation
64.20
69.99 70.10
75.77 75.60
81.56 81.00
87.34 87.41
93.13 93.24
98.91 97.44
104.70
64.45
104.73
Figure 5: Conditionally generated molecules on property using GeoRCG (EDM). The black number indicates the condition value, the green number represents the oracle property value for the generated molecule conformer.
Table 5:
Unconditional molecule generation GEOM-
Table 6: Sampling time (in seconds) for 5k samples using SemlaFlow and GeoRCG (Semla) across different numbers of sampling steps, measured on single NVIDIA RTX 4090.
DRUG with fewer diffusion steps.
# Steps
Atom Sta (
Valid (
GeoBEN
GeoRCG (EDM)
GeoBEN
GeoRCG (EDM)
Steps
Method
Rep.
Mol. Time
CFG
50
75.11
81,44(0.100
91.66
95.70(0,70)
SemlaFlow
610
610
100
100
78.89
(3.02(0.06)
93.05
96.30(0,70)
GeoRCG (Semla)
97
770
500
81.39
84.03(0.37)
93.47
97.57(0.90)
SemlaFiow
310
310
1000
85.6
12)
92.08
98.5(0.12)
GeoRCG (Semla)
97
380
SemlaFiow
152
152
GeoRCG (Semla)
97
315
189
sion SDE (Song et al., 2021) would still yield competitive results. Therefore, we reduce the number of diffusion steps and evaluate the model's performance. The results are pre- sented in Table 4 and Table 5. We provide the fewer-step performance GeoRCG (Semla) on GEOM-DRUG dataset in Appendix C.
Limitations. We discuss two limitations of GeoRCG. First, as a representation-guddeee generative method, its generation quality may depend heavily on the quality of representa- tions. For instance, in the GEOM-DRUG dataset, an insuf- ficiently pre-trained encoder may produce less meaningful representations. As result, the benefits of low-temperature sampling and classifier-free guidance in enhancing gener- ation quality and controllability may be less pronounced. Future work could investigate more effective pre-training strategies beyond standard denoising or enhanced repre- sentation regularization techniques to mitigate this issue. Second, although the additional conditioning module intro- duces small overhead, the use of classifier-fis free guidance re- quires doubling the batch size, resulting in roughly twice the resource consumption (memory and computation), see Ta- ble 6. Nonetheless, for many cases such as GeoRCG (EDM) in QM9, performance gains are substantial even without employing classifier-free guidance. With ongoing advance- ments in hardwares and infrastructures, we expect the over- head introduced by parallelism to be further minimized.
As demonstrated, with the geometric representation condi- tion, GeoRCG consistently outperforms other approaches across almost all step numbers. Notably, Table 4, with approximately 100 steps, the performance our method nearly converges to the optimal performance observed with 1000 steps, which already surpasses all other methods across all step numbers. This demonstrates the strong poten- tial of GeoRCG to reduce the number of iterations required by sequential generative methods.
5. Conclusions and Limitations
Conclusions. In this work, we present GeoRCG, a simple yet effective framework to improve the generation quality of arbitrary molecule generators by incorporating geometric representation conditions. We use EDM (Hoogeboom et al., 2022) and SemlaFlow (Irwin et al., 2024) as base generators and demonstrate the effectiveness of our framework through extensive molecular generation experiments. conditional generation tasks, GeoRCG achieves a remarkable 50% per- formance boost compared to recent SOTA models. Addition- ally, the representation guidance enables sampling with 10x fewer diffusion steps while maintaining near-optimal perfor- mance. Beyond these empirical improvements, we provide theoretical characterizations of pppresentation-coniitton generative models, which address key gap in the existing empirical literature (Lii et al., 2023).
A Acknowledgement
This work is supported by the National Key R&D Program of China (2022ZD016030)) and National Natural Science Foundation of China (62276003).
Impact Statement
This paper presents work whose goal is to advance the field of molecule generation. There are many potential societal consequences of molecule generation improvement, such as accelerating drug discovery, and developing new material.

---

Geometric Representation Condition Improves Equivariant Molecule Generation
None of them feel need to be specifically highlighted here for potential risk.
Fang, X., Liu, L.,, Lei, J., He, D., Zhang, S., Zhou, J., Wang, F., Wu, H., and Wang. H. Geometry-enhanced molecular representation learning for property prediction. Nature Machine Intelligence, 4(2):127-134, 2022.
References
Abbasi-yadkori, Y., Pál, D., and Szepesvári, C. Im- proved algorithms for linear stochastic bandits. In Shawe-Taylor, J., Zemel, R., Bartlett, P., Pereira, F., and Weinherger, K. (eds.), Advances in Neural Information Processing Systems, volume 24. Curran Associates, Inc., neurips. 2011. URL htttps://proceeddngg http://rrooeddnnn...rrr.rppppppp oc hc/ppaperfilles//peeeeeeeeesseee fiaeeeerffLlee/apae////////// 21d5be1cc7f24556706686 pdf.
Feng, S., Ni, Y., Lan, Y., Ma, Z.-M., and Ma, W.-Y. Frac- tional denoising for 3d molecular pre-training. In Interna- tional Conference on Machine Learning, pp. 9938 9961. PMLR, 2023.
Feng, S., Ni. Y., Lu, Y., Ma, Z.-M,, Ma, W.-Y., and Lan, Y. Unigem: A unified approach to generation and property prediction for molecules. arXiv preprint arXiv:2410. 10516, 2024.
Anderson, B., Hy, T. S., and Kondor, R. Cormorant: Co- variant molecular neural networks. Advances in neural information processing systems, 32, 2019.
Garcia Satorras, V., Hoogeboom, E., Fuchs, F., Posner, I., and Welling, M. E (n) equivariant normalizing flows. Advances in Neural Information Processing Systems, 34: 4181-4192, 2021.
Aneja, J., Schwing, A., Kautz, J., and Vahdat, A. con- trastive learning approach for training variational autoen- coder priors. Advances in neural information processing systems,, 34:480-4999 2021.
Gebauer, N., Gastegger, M., and Schütt, K. Symmetry- adapted generation of 3d point sets for the targeted dis- covery molecules. Advances in neural information processing systems, 32, 2019.
R. Geom, energy- Axelrod, and conformations for property predic- annotated molecular tion and molecular generation. Scientific Data, 911):185, 2022.
Gebauer, N. W., Gastegger, M., Hessmann, S. S., Müller, K.-R., and Schütt, K. T. Inverse design of 3d molecular structures with conditional generative neural networks. Nature commmnications, 13(1):973, 2022.
Beaudry, N. and Renner, An intuitive proof of the data processing inequality. arXiv preprint arXiv: 1rXiv:1107.077 107.0740 2011.
Graves, J., Byerly. J., Priego, E., Makkapati, N., Parish, S. V., Medellin, B., and Berrondo, M. review of deep learning methods for antibodies. Antibodies, 9(2):12, 2020.
Brown, T. B. Language models are few-shot learners. arXiv preprint arXiv:2005.14100 14165, 2020.
Ho, J. and Salimans, T. Classifier-free diffusion guidance. arXiv preprint arXiv:2207.12598, 2022.
Chen, S., Chewi, S., Li, J., Li, Y., Salim, A., and Zhang, A. Sampling is as easy as learning the score: theory for diffusion models with minimal data assumptions. In The Eleventh International Conference on Learning Represen- tations, 2023. URL https:////........ /ttes://oppprree http/ppeeevvveeeeeeeeeeeeeeeeeee net. Dorum?idd-zLLVVgs2200
Hong, H., Lin, W., and Tan, K. C. Fast 3d molecule gen- eration via unified geometric optimal transport. arXiv preprint arXiv:2405.15252 2024.
Hoogeboom, E., Satorras, V. G., Vignac, C., and Welling, M. Equivariant diffusion for molecule generation in 3d. In International conference on machine learning, pp. 8867- 8887. PMLR, 2022.
Chen, X., Xie, S., and He, An empirical study of training self-supervised vision transformers. In Proceedings of the IEEE/CVF international conference computer vision, pp. 9640-9649, 2021.
Huang, H., Sun, L., Du, B., and Lv, W. Learning joint 2-d and 3-d graph diffusion models for complete molecule generation. IEEE Transactions on Neural Networks and Learning Systems, 2024.
Dai. B. and Wipf, D. Diagnosing and enhancing vae models. arXiv preprint arXiv: 1903.0578. 2019.
De Bortoli, V. Convergence of denoising diffusion mod- els under the manifold hypothesis. arXiv preprint Sttv:2208.0531444 2022.
Ingraham, J. B., Baranov, M., Costello, Z., Barber, K. W., Wang, W., Ismail, A., Frappier, V., Lord, D. M., Ng- Thow-Hing, C., Van Vlack, E. R., et al. Illuminating protein space with a programmable generative model. Nature, 12700189) 2023.
Dosovitskiy, A. An image is worth 16x16 words: Trans- formers for image recognition at scale. arXiv preprint arXiv:20/0.1 /9900/0//999,.....100 1929, 2020.
10

---

Geometric Representation Condition Improves Equivariant Molecule Generation
Irwin, R., Tibo, A., Janet, J.-P., and Olsson, S. Efficient 3d molecular generation with flow matching and scale opti- mal transport, arXiv preprint urXiv:2406.07266. 2024.
Liu, X., Gong, C., and Liu, Q Flow straight and fast: Learning generate and transfer data with rectified flow. arXiv preprint arXiv:2209.0300
2022b.
Jang, Y., Kim, D., and Ahn, S. Hierarchical graph generation with k² -trerss In ICML 2023 Workshop Structured Probabilistic Inference
Luo, T., Mo, Z., and Pan, S. J. Fast graph generation via spectral diffusion. IEEE Transactions on Patter Analysis and Machine Intelligence. 2023.
{ Generative Modeling, 2023.
Luo, Y. and Ji, S. An autoregressive flow model for 3d molecular geometry generation from scratch. In Inter- national conference on learning representations (ICLR), 2022.
Jaa R., Han, J., Huang, W., Rong, Y., and Liu, Y. Energy- motivated equivariant pretraining for 3d molecular graphs. In Proceedings of the AAAI Conference on Artificial In- telligence, volume 37, pp. 8096-8100 2023.
Mislow, K. Introduction to stereochemistry. Courier Corpo- ration, 2012.
Jiao, R., Kong, X., Yu, 7., Huang, W., and Liu, Y. Equivari- ant pretrained transformer for unified geometric learn- ing on multi-domain 3d molecules. arXiv preprint arXiv:2402. 12700....2.... 2024.
Morehead, A. and Cheng. J. Geometry-complete diffusion for 3d molecule generation and optimization. Communi- cations Chemistry, 7(1):150, 2024.
Jo, J., Kim, D., and Hwang, S. Graph generation with dif- fusion mixture. arXiv preprint Xiiv:22302.03996 2023.
Nakata, M. and Shimazaki, T. Pubchemqe project: a large- scale first-principlen electronic structure database for data- driven chemistry. Journal of chemical information and modeling, 57(6)) 2017.
Kenton, J. D. M.-W. C. and Toutanova, L. K. Bert: Pre- training of deep bidirectional transformers for language understanding. In Proceedings of naacL-HLT, volume I, pp. 2019.
Ni, Y., Feng, S., Hong, X., Sun, Y., Ma, W.-Y., Ma, Z.-M., Ye, Q., and Lan, Y. Pre-training with fractional denoising to enhance molecular property prediction. arXiv preprint arXiv:2407. /086, 2024.
Le. T., Cremer, J., Noé, F., Clevert, D.-A., and Schütt, K. Navigating the design space of equivariant diffusion- based generative models for de novo 3d molecule genera- tion. arXiv preprint arXiv:2309. 17296, 2023.
O' O'Boyle, N. M., Banck, M., James, C. A., Morley, C., Van- dermeersch, T., and Hutchison, G. R. Open babel: An open chemical toolbox. Journal of cheminformatics, 3: 1-14, 2011.
Li, P., Wang, J., Qiao, Y., Chen, H., Yu, Y., Yao, X., Gao, P., Xie, G., and Song. S. An effective self-supervised frame- work for learning expressive molecular global representa- tions to drug discovery. Briefings in Bioinformatics, 22 (6):bbab100 109, 2021.
Peebles, W. and Xie, S. Scalable diffusion models with transformers. In Proceedings of the IEEE/CVF interna- tional conference on computer vision, pp. 2023.
Li, T., Katabi, D., and He, K. Self-conditioned image gen- cration via generating representations. arXiv preprint arXiv:2312.03701, 2023.
Ramakrishnan, R., Dral, P. O., Rupp. M., and Von Lilienfeld, O. A. Quantum chemistry structures and properties of 134 kilo molecules. Scientific data, 1(1):17, 2014.
Li, Z., Wang, X., Huang, Y., and Zhang, M. Is distance matrix enough for geometric deep learning? Advances in Neural Information Processing Systems, 36, 2024a.
Razavi, A., Van den Oord, A., and Vinyals, O. Generating diverse high-fidelity images with vq- Advances neural information processing systems, 32, 2019.
Li, Z., Wang, X., Kang. S., and Zhang, M. On the complete- ness of invariant geometric deep learning models. arXiv preprint XxXx:2442.00436, 2024b.
Rombach, R., Blattmann, A., Lorenz, D., Esser, P., and Ommer, B. High-resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF con- ference on computer vision and pattern recognition, pp. 10684-10695, 2022.
Lipman, Y., Chen, R. T., Ben-Hamu, H., Nickel, M., and Le, M. Flow matching for generative modeling. arXiv preprint ar arXiv:2210.02744 2022.
Satorras, V. G., Hoogeboom, E., and Welling, M. E (n) equivariant graph neural networks. In International con- ference on machine learning, pp. 2021.
Liu, S., Guo, H., and Tang, J. Molecular geometry pretrain- ing with se (3)-invariant denoising distance matching. arXiv preprint arXiv:2206. /3602, 2022a.
PMLR,

---

Geometric Representation Condition Improves Equivariant Molecule Generation
Song, J., Meng, C., and Ermon, S. Denoising diffusion im- plicit models. arXiv preprint arXiv:2010.02500, 2020a. 2020a.
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, 35: 36533-36545, 2022.
Song, Y., Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Er- mon, S., and Poole, B. Score-based generative modeling through stochastic differential equations. arXiv preprint 133456 2020b.
Xu, M., Powers, A. S., Dror, R. O., Ermon, S., and Leskovec, J. Geometric latent diffusion models for 3d molecule gen- eration. In International Conference Machine Learn- ing. pp. 11392-33610.. PMLR, 2023.
Song. Y. Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Ermon, S., and Poole, B. Score-based generative mod- eling through stochastic differential equations. In In- ternational Conference on Learning Representations, 2021. URL httpp:://pen.revee.... forum? id=PxTIG12RRHS,
You, Y., Zhou, R., Park, Xu, H., Tian, C., Wang, Z., and Shen, Y. Latent 3d graph diffusion. In The Twelfth International Conference Learning Representations, 2023.
Song. Y., Gong, J., Qu Y., Zhou, H., Zheng. M., Liu, J., and Ma, W.-Y. Unified generative modeling of 3d molecules via bayesian flow networks. arXiv preprint arXiv:2403. 15441, 2024a.
Yuan, H., Huang. K., Ni. C.,, Chen, M., and Wang, M. Reward-directed conditional diffusion: Provable distri- bution estimation and reward improvement. In Thirty- seventh Conference Neural Information Processing Systems, 2023.
Song, Y., Gong, J., Xu, M., Cao, 7.,, Lan, Y., Ermon, S., Zhou, H., and Ma, W.-Y. Equivariant flow matching with hybrid probability transport for 3d molecule generation. Advances in Neural Information Processing Systems, 36, 2024b.
Zaidi, S., Schaarschmidt, M., Martens, J., Kim, H., Tch, Y. W., Sanchez-Gonzalee A., Battaglia, P., Pascanu, R., and Godwin, J. Pre-training via denoising for molecular property prediction. arXiv preprint A.Kv:2206600133: 2022.
Thölke, and De Fabritiis, G. Torchmd-net: equivariant transformers for neural network based molecular poten- tials. arXiv preprint arXiv:2202.0254/ 2022.
Zheng, Q., Le, M., Shaul, N., Lipman, Y., Grover, A., and Chen, R. T. Guided flows for generative modeling and decision making, arXiv preprint arXiv:2311.13443, 2023.
Thomas, N., Smidt, T., Kearnes, S., Yang, L., Li, L., Kohlhoff, K., and Riley, P. Tensor field networks: Rotation-and manslation-equivarian neural networks for 3d point clouds. arXiv preprint arXiv: 182fv.1802.08219, 2018.
Zhou, C., Wang, X., and Zhang. M. Unifying generation and prediction graphs with latent graph diffusion. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024.
Van Den Oord, A., Vinyals, O., et al. Neural discrete rep- resentation learning. Advances in neural information processing systems, 30, 2017.
Zhou, G., Gao, Z., Ding, Q., Zheng, H., Xu, H., Wei, Z., Zhang, L.., and Ke, G. Uni-mol: A universal 3d molecular representation learning framework. 2023.
Vignac, C. and Frossard, P. Top-n: Equivariant set and graph generation without exchangeability. arXiv preprint WxXiv:2/10.00000 2021.
Vignac, C., Krawezuk, I., Siraudin, A., Wang, B., Cevher, V., and Frossard, P. Digress: Discrete denoising diffusion for graph generation. arXiv preprint A1t1:22299.173:4. 2022.
Vignac, C., Osman, N., Toni, I.., and Frossard, P. Midi: Mixed graph and 3d denoising diffusion for molecule generation. In Joint European Conference on Machine Learning and Knowledge Discovery in Databases, pp. 560-576. Springer, 2023.
Wang, S., Tan, Z., Zhao, X., Chen, T., Liu, H., and Li, J. Graphreg: Self-conditioned graph genera- tion via bootstrapped representations. arXiv preprint arXiv:2403.0107/, 2024.
12

---

Geometric Representation Condition Improves Equivariant Molecule Generation
A. Algorithms
High-level Algorithm for Parallel Training and Sequential Sampling We provide the high-level training and sampling algorithm for GeoRCG in Algorithm 1.
Algorithm 1 Parallel Training and Sequential Sampling for GeoRCG
Input: Molecule dataset train C UN=1 (RNx3, x RRNxd), pre-trained geometric encoder E, initial representation
generator Peo (r), molecule generator Pa, (M]r).
Output: Trained representation generator Pr(r), molecule generator po(M)r), and molecule samples from Pp.0 (M).
Parallel Training:
Pre-process to obtain:
The representation dataset D
train {(E(M), (MM))-
train
The mol-rep dataset mmmoo-rop = M))MM E D
train
Train the representation generator Peo (r) with D using loss Crep in Equation (1).
Train the molecule generator po, (r) with Dimokrrp using loss Lomot in Equation (2), while applying the training techniques outlined below, including representation perturbation and representation loss. Sequential Sampling:
Sample a representation r= ~ (r)) with low-temperature sampling technique outlined below.
Sample a molecule M. ~ po (Mir.) conditionally, with classifier-fres guidance technique outlined below.
Return: Trained representation generator (r(r), molecule generator pa (M(M)) and generated molecule sample.
Training: Representation Perturbation Unlike typical conditional training scenarios, GeoRCG faces a unique challenge: the representations that condition the molecule generator during training may not always coincide with those generated by the representation generator during the sampling stage. This issue is particularly pronounced in molecular generation than image case (Li et al., 2023), where pre-trained encoders are typically not trained on that large datasets with advanced regularization techniques like MoCo v3 (Chen al., 2021). Consequently, the molecule generator is susceptible overfitting to the training representations, as evidenced by our preliminary experiments on QM9 molecule generation shown in Table 7.
Table 7: Quality of molecules generated by GeoRCG trained on the QM9 dataset without using the representation perturbation technique, comparing different representation sources. "Training Dataset" refers representations sampled from Deo while "Rep. Sampler" refers to representations generated by the trained representation generator (r).
Metrics
Mol Sta (%) (t)
Valid (%) (t)
Rep source
Training Dataset
93.20 (0.50)
97.07 (0.32)
Rep. Sampler
86.93 (0.50)
89.12 (0.21)
We find that a simple caaniique-petuurbinng the geometric representation during training the molecule generator with some Gaussian noise cep where (0, 7) and rep is relatively small variance--s particularly effective for solving this problem. Formally, after sampling a data point (E(M), from Demot-rop, we use (M, E(M) + rmp? for training, Ablation study in Appendix C show this simple method can effectively prevent overfitting and ensure that performance on novel representations matches those from the training dataset.
train
In practice, we set Grep to 0.3 for QM9 dataset and Trep to [0.3, 0.5) for GEOM-DRUG dataset.
Training: Representation Loss Training molecular generative models typically involves predicting a clean molecule from a noisy molecule input (in noise parameterization or vector field parameterization, an equivalent formulation exists for constructing clean molecule predictions). The loss is computed by minimizing the distance (e.g., MSE for coordinates) between the predicted clean molecule and the true clean molecule. To further strengthen supervision of the representation, we introduce an additional representation loss during training. This loss is defined as the MSE between the predicted clean molecule's representation and the actual clean molecule's representation.
In practice, for GeoRCG (EDM), we do not apply this technique, whereas for GeoRCG (Semla), we incorporate it with a relatively small coefficient.
13

---

Geometric Representation Condition Improves Equivariant Molecule Generation
Sampling: Cow-eempee Sampling We adopt the low-temperature sampling algorithm introduced by Chroma (In- graham et al., 2023) to the representation generator. However, we apply it an MLPPbased diffusion model rather than the equivariant diffusion model that processes geometric objects Chroma.
The objective of low-temperaturee sampling is to perturb the learned representation distribution Pp(r) by rescaling it with an inverse temperature factor, 7. where T is a tunable temperature parameter during sampling, and finally enables sampling
from Zyp where 2y is a normalization constant. The method proposed in Chroma (Ingraham et al., 2023) scales the score estimated at each diffusion time step using a fime-dependent factor A The approach is derived from and has theoretical guarantees for simplified toy distributions, and its performance on complex distributions, though lacking strict guarantees, has shown consistent results when combined with annealed Langevin sampling (Song et al., 2021). Here we briefly introduce it for conf-conta,, and recommend the readers to Ingraham et al. (2023) for detailed derivation and illustration.
Consider the vanilla reverse SDE used in DDPM sampling (VP formulation) (Song al., 2021):
dr 1 -- 3 V, log (:())t BBdww
where W is reverse-time Wiener process, (r(r) denotes the ground-truth representation distribution time i, and S, represents the time-dependent diffusion schedule. To incorporate low-temperature sampling, we utilize the following Hybrid Langevin Reverse-time SDE:
2 Ber -0x-
Aowi 2
dr =
Ax
+
Be V. log (+()d + BB(11 + b,dw,
(8)
where X₂ is time-dependent temperature parameter defined as a2+(1-af with Ao = 7 or satisfies 1, d'log The parameter 1/ controls the rate of Langevin equilibration per unit time, and as shown in Ingraham et al. (2023), it helps align more effectively with the reweighting objective in complex distributions. In our implementation, we employ the explicit annealed Langevin process (the corrector step from (Song et al., 2021)) achieve similar results.
In practice, for the unconditional QM9 and GEOM-DRUG generation results of GeoRCG (EDM) shown in Table 1, we set T = 1.0 for QM9 and T = 0.5 for GEOM-DRUG, For conditional generation in Table 3, we set T = 1.0. The effect of varying T on QM9 is detailed in Table 10. For GeoRCG (SemlaFlow) results in Table 2, we use T = 1.0.
Sampling: Classifier-Free Guidance We employ the classifier-fres guidance algorithm, as introduced in (Ho & Salimans, 2022), for our molecule generator. Specifically, we introduce a trainable "fake" representation, denoted as 1, which serves as the unconditional signal. During the training phase, / is initialized as learnable parameters, and with probability of Plake the true representation r is replaced by l. This ensures that the model is capable of generating molecules unconditionally, ie., (ML() approximates (.M). During sampling, the final score estimate produced by the molecule generator is adjusted using the formula (1 + ()fo wfo(Me, I), allowing flexible control over the strength of the representation guidance.
In practice, for unconditional generation in Table 1, we set w = 1.0 for QM9 and w = 0.0 for GEOM-DRUG, The impact of varying w on QM9 is shown in Table 10. For conditional generation in Table 3, we set w = 2.0. For GeoRCG (SemlaFlow), we use w -0.9 (note that -10 < < 0.0 indicates subtle representation guidance that exerts a small but still meaningful influence the model's behavior). Further tuning of these hyperparameters may yield improved performance.
B. Experiment Details
Metrics and Baseline Descriptions We adopt the evaluation metrics, guidelines and baselines commonly used in prior 3D molecular generative models to ensure a fair comparison (Hoogeboom et al., 2022).
In the unconditional setting, we assess the generated molecules using several key metrics:
Atom Stability: The proportion of atoms with correct valency.
Molecule Stability: The proportion of molecules where all atoms within the molecule are stable.
Validity: The proportion of molecules that can be converted into valid SMILES using RDKit.
14

---

Geometric Representation Condition Improves Equivariant Molecule Generation
Validity & Uniqueness: The proportion of unique molecules among the valid molecules.
Energy: The energy U(x) of conformation Lower energy values typically correspond to more stable and physically plausible conformations that are closer to what would be observed in nature. Following (Irwin et al., 2024), the energy is calculated using the MMFF94 force field within RDKit, a commonly used molecular modeling framework.
Strain: A measure of how distorted generated conformation x compared to its relaxed (optimized) state The relaxed conformation E is obtained by applying energy minimization using the MMF1994 force field, where the molecular structure is iteratively adjusted reduce its energy until reaching local minimum. Mathematicallly, the strain energy is defined as (x)) - U (T), where (x) is the energy of the generated conformation and U (*) is the energy of the relaxed conformation. Lower strain energy values imply that the generated conformations are closer to being physically accurate and require minimal correction during optimization.
Following the approach of Hoogehoom al. (2022); Vignac & Frossard (2021), do not report Novelty scores in the main text, since QM9 represents exhaustive enumeration of molecules satisfying a predefined set of constraints, therefore, "novel" molecule would often violate at least one of these constraints, which indicates that model fails to fully capture the properties of the dataset. For reference, we observe that the novelty of GeoRCG (EDM) on QM9 is approximately 42% when w = 0.0, T = 1.0, compared to about 65% for EDM.
When comparing with 2D & 3D models in Table 8, we evaluate two 3D metries introduced by MiDi (Vignac et al., 2023), which directly assess the geometry learning ability:
enggh The weighted I-Wasserstein distance between the bond-length distributions of the generated molecules and the training dataset, with weights corresponding to different bond types. Formally, it is defined as:
LondleengthsW = Σ
gd (v)))((Daaaal)) Daia(y)),
yEoond types
where g (y) the proportion bonds of type in the training set, Dais(y) is the generated distribution of bond lengths for bond type y, and Dais(y) is the corresponding distribution from the test set.
BondAAgleeW:: The weighted 1-Wasserstein distance between the atom-centered angle distributions of the generated molecules and the training dataset, with weights based on atom types. Formally, it is defined as
BondAnglesW1
Σ
xxxx (z)) angles (z), (angles(x)),
(10)
Eatom types
where xxxxx) (x)) denotes the proportion of atoms of type r in the training set, restricted atoms with two or more neighbors, and (zr) represents the distribution of geometric angles of the form Z(rk ri, : rt), where i is atom of type and k and j are neighbors of i.
In the conditional generation setting, as described in (Hoogeboom al., 2022), we evaluate our approach on the QM9 dataset across six properties: polarizability a, orbital energies Eномо, ELUMO, and their gap Ac, dipole moment p. and heat capacity C., The generative model is trained conditionally on the second half of the QM9 dataset, and an EGNN (Satorras et al., 2021) classifier, trained the first half, is employed to evaluate the MAE property error of the generated samples.
Three baselines are adopted Table 3:
QM9 (lower bound): The mean error of a classifier trained on the first half of the QM9 dataset and evaluated on the second half. This baseline represents the inherent bias/error of the classifier, setting a lower bound for model performance and reflecting the best possible performance model can achieve.
Random: The classifier's performance when evaluated on the second half of QM9 with randomly shuffled molecule property labels. This baseline provides an upper bound, representing the worst achievable performance.
-
N_atoms: The performance of a classifier trained exclusively on the number of atoms N and evaluated using only N as input. This baseline captures the intrinsic relationship between molecular properties and the number of atoms, which a generative model must surpass to demonstrate effectiveness.
15

---

Geometric Representation Condition Improves Equivariant Molecule Generation
Model Architectures, Hyperparameters and Training Details
Representation Generator. We use the same architecture for the representation generator as the MLP-based diffusion model proposed in Li et al. (2023). We use 18 blocks of residual MLP layers with 1536 hidden dimensions, 1000 diffusion steps, and a linear noise schedule for Bt. The representation generator is trained for 2000 epochs with a batch size of 128 for both the QM9 and GEOM-DRUG datasets. Training on QM9 takes approximately 2.5 days on single REOM-DRUG takes around 4 days single Nvidia A800. Training time can indeed be further reduced, the model shows minimal progress after approximately half of the reported time.
Nvidia 4090, while training
Molecule Generator. We adopt EDM (Hoogeboom et al., 2022) as the base molecule generator, using the same EGNN (Satorras et al., 2021) architecture, with the exception of the conditioning module. Specifically, we introduce a simple gated feedforward layer to incorporate the representation condition, inserting it between each EGNN block to enhance regularization and improve model expressiveness.
For the EGNN hyperparameters, we use 9 layers with 256 hidden dimensions for QM9 and 4 layers with 256 hidden dimensions for GEOM-DRUS DRUG. The number of diffusion steps is set to 1000 (except for cases in Table that generate molecules with fewer steps), and we employ the polynomial scheduler for at (MM), Notably, all model hyperparameters are identical to those in EDM for fair comparison.
During training, we use a batch size of 128 and 3000 epochs on QM9, and a batch size of 64 and 20 epochs on GEOM-DRUG. Training takes approximately 6 days on QM9 using a single Nvidia 4090, and around 10 days on GEOM-DRUG using two Nvidia A800 GPUs.
For GeoRCG (SemlaFlow), we select the letter-performing architecture between a gated feedforward layer and an AddaNN-Zer-llie module (Peebles & Xie, 2023) for conditioning on representations and place it between every block of Semla. During training, we use the batch cost 2048 and trains for 300 epochs.
Evaluation of LDM-3DG (You et al., 2023) We evaluate the performance of LDM-3DG (You et al., 2023) in Table and Table 3, an auto-Ennoder-based method that also leverages the compactness of the representation space to achieve good performance.
For the unconditional results in Table 8, we utilize the 3D conformations unconditionally generated by LDM-3DG (You et al., 2023) and compute the bond information using the look-up table method from EDM (Hoogeboom et al., 2022). Notably, although LDM-3DG predicts both the 2D molecular graph and the 3D conformation, we do not use the bond information it predicts for the following reasons:
1. For the calculation of 3D geometry statistics, we observe significant inconsistencies between the generated 2D graphs and 3D geometries (e.g.. valid molecules with bond lengths exceeding 100 m), leading to unreliable statistics (e.g.. BondLengthWi exploding to 3900).
2. For stability and validity metrics, which are fundamentally 2D and computed based on molecular graphs (atoms and bond types), using the generated 2D) graph would ignore the contribution of the 3D) module, preventing an evaluation of its 3D learning performance.
3. Most critically, their 2D module is explicitly designed to filter out invalid (ub-))moeccules during generation using the RDKit method. This means that if invalid molecules or sub-molecules are generated, they are regenerated. This explicit filtering deviates from our standard evaluation criteria and is unsuitable for a fair comparison.
For the conditional results in Table 3, we first note a potential issue with LDM-3DG (You et al., 2023): The model cannot explicitly specify the node number N during molecule generation, as it uses an auto-regressive 2D generator that automatically stops adding atoms/motifs when deemed sufficient. However, the evaluation in Table 3 requires specifying both N and property c, following the ground-truth distribution q(N, e) from the training dataset. To ensure fair evaluation, conditions feeding to LDM-3DG must also satisfy this distribution. As the authors claim the model can implicitly learn q (() and thus q(N]cc,, we first sample 10,000 values from (c) and feed them to LDM-3DG, expecting it to infer N from implicitly as argued, and thus matching the q( (N, c) conditions.
C. Additional Experiment Results
Comparison of GeoRCG (EDM) with 2D&3D Methods We compare GeoRCG (EDM) with recent 2D&3D methods such as MiDI (Vignac et al., 2023) and LDM-3DG (You et al., 2023). As discussed in Section 4.1, such comparison not
16

---

Geometric Representation Condition Improves Equivariant Molecule Generation
fair, since these models learn and generate both 2D bond structures and 3D geometries, which is beneficial for metrics like validity and stability. Considering that they rely on external chemistry toolkits like RDKit or OpenBabel (O' Boyle et al., 201 for bond determination at the input stage and continue to leverage this bond information throughout training and generation, we report GeoRCG (EDM) using the same external tools for accurate bond computation in the generated 3D conformations, rather than relying on simple look-up tables, to narrow the comparison gap (though not completely addressed since GeoRCG (EDM) still do not explicitly learn to generate bonds).
Furthermore, as GeoRCG (EDM) essentially captures 3D geometric distributions, we place more emphasis on 3D metries that directly evaluate 3D learning capabilities, including BondLengthW[ and BondAngleWI proposed by Vignac et al. (2023) and detailed in Appendix B.
The results in Table 8 demonstrate that GeoRCG not only significantly outperforms MiDi and LDM-3DG on 3D metrics, highlighting the advantages of using a pure 3D model for learning 3D structures, but also further enhances EDM's performance, which has already shown considerable promise in 3D learning.
Table 8: 3D geometry statistics and generated molecule quality on QM9 across different methods. Models marked with indicate results obtained from our own experiments; see Appendix B for the evaluation guidelines. The stability metrics for EDM are higher than in Table due to using the MiDi codebase for evaluation, which permits more valency for atoms.
Metries
Angles (
Bond Length (e-2 A
Mol Sta (%)
Atom Sta (%) f
Validity (5)
Uniqueness (%)
Methods
Data
~00.1
~0
98.7
99.8
98.9
99.9
MiDi (uniform)
0.67(0.02)
1.6(..)
(6.1(
99.7(0.0)
96.6(0.2)
97.6(0.1)
MiDi (adaptive)
0.62(0.02)
0.3(0.1)
97.5(0.1)
99.8(0.0)
97.9(0.1)
97.6(0.1)
LDM-3DG'
3,56
0.2
94.03
99.38
94.89
97.03
EDM
0.44
0.1
90.7
99.2
91.7
98.5
EDM + OBahel
0.44
0.1
97.9
99.8
99.0
98.5
GeoRCG (EDM)
0.21 (4) 52
0.04(00.) 601s
95.832 95.82(00166) 5.08
(9.59)0.02)
0,39%
96.54(0.27)
5.29%
95.74(0.1
(774(0.18)).84 2.8%
GeoRCG (EDM)+ + OBabel
0.20(0.04) 54.5%
0.07((.06) stra
(0.09) co 5014
99.888 0.08%s
99.0(0.04) 0.0%%
95.74(0.11) 2.8%
Balancing Controllability We present a more comprehensive figure that includes molecule stability, atom stability, validity, and
in Figure 6.
Figure 6: Balance controllable (unconditional) generation on QM9 dataset of GeoRCG (EDM). Increasing w and decreasing T enhances stability, with the cost of a reduction in uniqueness.
Molerale Stabitty
Atom Stabiitty
Validryy
Velidiiy&Gni Ceiigeeeenn
30
$0:46
11.55 5.5
66.30 16
67 P9
10
17.11 11
#5.70
19:25
00
24
$7.18
20,37
98.37
92.32
$3.06
03.20
00.20
$9.13
59.22
90.24
$3.25
56.52
17.11
07.04
07.20
90:45
$1.30
$0.53
Classter
10:04
95.00
29.11
93.61
92:23
91.70
61
Intwrru iamparrdtrr
Irowm betrparstioe
Intwree Iemparytur
luverse tersperutare mi
Fewer-step Sampling of GeoRCG (Semla) We present the performance of GeoRCG (Semla) across varying numbers of sampling steps in Table 9. The results demonstrate consistent improvements over SemlaFlow, highlighting the effectiveness of GeoRCG on this advanced method with reduced sampling steps.
Representation Encoders Geometric representations play pivotal role in GeoRCG. To evaluate the importance of representation quality, we conduct an ablation study comparing the quality of molecule samples generated by GeoRCG trained under different geometric encoder configurations.
Ablation Study:
We first assess the benefits provided by the pre-training stage. Specifically, we utilize the pre-trained encoder Frad (Feng et al., 2023), trained on the PCQM4Mv2 dataset (Nakata & Shimazaki, 2017) with a hybrid coordinates denoising task (Feng et al., 2023). This approach has been proven to equivalently learn force fields (Feng et al., 2023; Zaidi et al., 2022), and
17

---

Geometric Representation Condition Improves Equivariant Molecule Generation
Table 9: Comparison between SemlaFlow and GeoRCG (Semla) across varying numbers of sampling steps. Results of SemlaFlow is obtained by our experiments.
# Steps
Method
Energy 4
Strain +
Atom-Stab.
Mol.-Stab. t Validity
kcal mol
kcal-mol-
%
%
SemlaFlow
95.72(1.24)
56.42(1.07)
99.8(0.00)
97.4(0.07)
94.4(0.17)
GeoRCG (Semla)
88.6(1.03)
47.64(1.10)
99.8(0.00)
97.6(0.00)
95.3(0.13)
SemlaFlow
100.60(0.55)
60.31(0.13)
99.8(0.00)
11)
94.6(0.17)
GeoRCG (Semla)
91.60(1.03)
50.50(0.65)
99.8(0.00)
97.2(0.20)
95.3(0.22)
SemlaFlow
117.03(1.37)
76.99(1.18)
99.7(0.01)
95.4(0.17)
93.2(0.30)
20
GeoRCG (Semla)
(9.83(0.91)
62.88(0.56)
99.7(0.00)
95.5(0.13)
94.2(0.01)
therefore expected to produce informative representations that capture high-level molecular information. We train GeoRCG (EDM) using representations from a well-pretrained Frad and a Frad with randomly initialized weights. The molecule generation quality on QM9, as shown in Table 10, clearly underscores the critical role of pre-training on large datasets with advanced techniques in improving representation quality, ultimately enhancing GeoRCG's performance.
Table 10: Quality of molecules generated by GeoRCG (EDM) with different encoders trained on the QM9 dataset. "Random" indicates that the weights were initialized randomly without any pre-training.
Metrics
Atom Sta (%)
Mol Sta (%)
Valid (%)
Valid & Unique (%)
Encoder
Random Enc
98.55(0.01)
78.66(0.07)
94.68(0.09)
55.99(0.83)
Pretrained Enc
99.10(0.02)
92.15(0.22)
96.48(0.08)
92.45(0.21)
Next, we investigate the impact of different pre-trained encoders, which could vary in model structure and proxy tasks used for pre- pre-training. Specifically, we compare Unimol (Zhou et al., 2023), which employs message- message-passing neural network framework incorporating distance features (i.e., DisGNN in (Li et al., 2023; 2024b)) and primarily uses naive coordinates denoising, with Frad (Feng et al., 2023), which adopts TorchMD (Thölke De Fabritis, 2022) as the backbone and utilizes carefully designed hybrid-denoisee pre-trained the GEOM-DRUG dataset until convergence. We visualize the t-SNE of the representations generated for GEOM-DRUG. As shown in Figure 7, the t-SNE of the Unimol representations exhibits clearer clustering pattern based on node numbers compared to the Frad representations, which may suggest better representation learning. To further investigate, we utilize both encoders to train GeoRCG and subsequently evaluate the quality of molccule generation. The Frad-based GeoRCG achieves a Validity of 96.9(0..4) and Atom Stability of 84.4(0.27), while the Unimol-based GeoRCG achieves a Validity of 98.5 (0.12) and Atom Stability of 84.3(0.12): Although the Frad-based GeoRCG produces slightly higher atom stability, its high variance and significantly lower validity suggest inferior performance. These findings, along with our main results, offer insights into the types of representations more effective for guiding molecule generation, suggesting that sensitivity to molecule size may be critical factor.
tasks. Both Unimol (Zhou et al., 2023) and Frad (Feng et al., 2023) are
Table 11:
Quality of molecules generated by GeoRCG (EDM) trained on the QM9 dataset, with and without representation perturbation and representation condition dropout.
Metrics
Atom Sta (%)
Mol Sta (%)
Valid (%)
Valid & Unique (%)
Syyper-prrameerss
rep noise X rep noise cond. dropout /
cond. dropout X
98.53(0.08)
86.93(0.5)
93. 03.69(0.09)
89. 19.12(0.21)
98.62(0.08)
87.9(0.35)
94.64(0.18)
90.15(0.02)
rep noise v
cond. dropout x
99.05(0.01)
91.69(0.08)
96.48(0.11)
92.38(0.12)
rep noise cond. dropout v
99.10(0.02)
92.15(0.23)
96.48(0.08)
92.45(0.211
Ablation Study: Representation Perturbation As discussed in Appendix A, we investigate the effectiveness of the straightforward representation perturbation technique by introducing random noise to perturb the representations during training. Additionally, we apply extra dropout in the conditioning module of our molecule generator to mitigate overfitting on the representation conditions, Ablation experiments presented in Table 11 demonstrate the efficacy of these simple yet impactful methods.
18

---

Geometric Representation Condition Improves Equivariant Molecule Generation
100
150
60
10
75
40
Node Number (Log Scale)
Node Number (Log Scale)
80
20
25
60
40
-20
-25
-40
50
20
-60
-75
88
-60
-40
-20
20
40
-100
-50
50
100
Figure 7: t-SNE visualization of representations produced by the pre-trained encoders for the GEOM-DRUG dataset, colored by node number. The left plot corresponds to Unimol (Zhou et al., 2023), the right plot corresponds to Frad (Feng et al., 2023).
Table Supplementary evaluation of conditionally generated molecules from GeoRCG (EDM). The right side reports metrics for unconditionally generated molecules from other methods for reference. Note that conditional models (left) were trained on half of the QM9 dataset, while unconditional models (right) were trained on the full dataset, which may account for slight decreases in stability and validity metrics.
Ae
C.
BDM
GeoLDD
EquiFM
Atom Sta (%)
98 (3)04)4)
98.95(0.03)
98.93(0.02)
98.99(0.01)
98.90(0.02)
98.98(0.02)
98.7
98.9(0.1)
98.9(0.1)
Mol Sta (%)
90.64(0.36)
90.38(0.19) 90.46(0.24)
90.94(0.23)
90.02(0.22)
90.87(0.12)
82
89.4(0.5)
88.3(03))
Valid (%)
95.40(0..4)
95.46(0.19) 95.44(0.11)
95.74(0.16)
95.26(0.04)
95.62(0.08)
91.9
93.8(0.4)
94.7(0.4)
Valid & Unique (5)
90.47(0.44)
90.09(0.13)
90.06(0.11)
90.32(0).2)
89.93(0.24)
90.20(0.16 16)
90.7
92.7(0.5)
93.5(0.3)
Valid & Unique & Novelty (%)
50.03(0.26)
50.81(0.15)
50.90(0.33)
50.59(0.08)
51,70(0.79)
51.10(0.29)
65.7
58.1
$7.4
Quality of Conditionally Generated Molecules Detailed molecular metrics for conditionally generated molecules are provided in Table 12. For comparison, we also include the stability metrics of molecules conditionally generated by EDM, which highlight a notable improvement in stability with GeoRCG. Specifically, EDM's stability scores are: a (80.4%), As (81.73%), Eномо (82.81 %), fLUM (83.6%), f (83.3%), and C. (81.03%).
D. Theoretical Analysis
In this section, we provide rigorous theoretical analysis on pnvresentation-coniiinnn diffusion models. Our theory is not limited to molecule generation, and is the first theoretical breakthrough for the RCG framework (Li et al., 2023).
Our analysis is organized as follows. Appendix D.I, we analyze the generation bound of pngresentation-condiioonn diffusion models unconditional generation tasks by showing: (i) the representation can be well generated by the first-stage diffusion model with mild assumptions (Appendix D.1.1); (ii) the second-stage representation its.eentttiioncooniiiionneed diffusion model exhibits higher generalization error than traditional one-stage diffusion model, and arguably achieve lower leveraging the informative representations (Appendix 1.1.2). Then in Appendix D.2, we analyze conditional generation tasks as follows: (ii) under mild assumptions of representations and targets, we provide novel bound for score estimation eтror (Appendix D.2.1); (iv) generated representations have provable reward improvement towards the target, with the suboptimality composed offline regression error and diffusion distribution shift (Appendix D.2.2), thus would improve second stage of conditional generation (Appendix D.2.3).
Notations. In this section, we use SDE and score matching formulations of diffusion models to present our theoretical results, given their equivalence with the DDPM family (Song et al., 2021). We consider the random variable x E E x(++3) and use q(-) to denote the ground truth distributions, p(-) to denote the posterior distribution predicted by diffusson models. For instance, ((x) is the ground truth distribution of the underlying data x. while P. (r) is the predicted distribution latent representations. We use T denote the total time of diffusion models, and N₂ to represent discretization step number.
19

---

Geometric Representation Condition Improves Equivariant Molecule Generation
We consider SDE with continuous time [0, 7], as well its discretized DDPM which has N& diffusion steps with step size = T/Na. The forward process is denoted as (2+))=[0,TT) ~ Qt, and the reverse process is denoted (FF))E) Pt- If the reverse process is predicted by the score matching network, we use its parameters the subscript. Please note that there are two different initialization of the reverse process: the end of forward process qr and standard Gaussian noise yd. We use superscript fr to differentiate the former from the latter.
D.1. Unconditional Generation
D.1.1. PROVABLE GENERATION OF REPRESENTATIONS
Recall the two-stage generation process of representatooncconitioned generation: p(z,r) = po(2)r)p ) To quantitatively evaluate the generation process, we consider two stages separately. In this subsection, we first provide theoretical analysis on the provable generation of representations P. (r).
Assumption D.1. (Second moment bound of representations.)
m == Ey(r)[lrr -F113] oo
(1I)
where (()) is the ground truth distribution of the representations, and ;= Br(r)rr. Assumption D.2. (Lipschitz score of representations). For all t 2 0, the score v g(r))) is L, Lippschchi
where g(r)) is the distribution of noisy latent r, at diffusion step in the forward process.
Finally, the quality of diffusion models obviously depends on the expressivity of score network &o with prediction s Assumption D.3. (Score estimation error representations) For all t [].
Fecrelllss Vlng(rr)11) 2
(12)
p,ccre
These are similar assumptions to the ones in (Chen et al., 2023).
Proposition DA4. Suppose Assumption D. Assumption D.2, Assumption D.3 hold, and the step size h := T/Na satisfies 3 Ly Then the following holds,
TV(p)(r)) 3 v (xx(((r))lxx)))(((771 dz + M.y
Esp.score estimation
(13)
comergence of forward
discretization
Here a(-) 3 b(-) means there exists constant such that ax-) S Cb(.) always holds. This is direct conclusion from (Chen al., 2023). In typical DDPM implementation, we choose h = 1 and thus T = Na. Remarkably, Proposition D.4 indicates the benefit of generating the representation first: since d, < the generation quality (measured by the TV distance in Proposition D.4) of the low-dimensional representation easily outperform directly generating the high-dimensional data points z. The theorem also accounts for applying lightweight MLP as the denoising network while in the stage of generating the representation.
D.1.2. PROVABLE SECOND- DACOND-STAGE GENERATION
Tractable Training Loss. Now we analyze the generation quality of the second-stage diffusion model. Since we sample from p,(x, r), have representations as conditions even for unconditional generation tasks. To learn the score function conditioning the representations, consider the following loss for score matching,
[(0)= = S (1 Va, log
[(0)
(14)
However, since (x,lr) is intractable, use the following equivalent losses:
(()) /
T A(t)EE.o,r [E [1150 r, -z. log (2)22)1)))) [ool
(15)
20

---

Geometric Representation Condition Improves Equivariant Molecule Generation
Proposition D.5. (Tractable represeenation-comddittnnedd score matching loss.)
C(6) ! x(t)Eu..:llso(((n11 - V.. log (x)l)))))
(16)
1 X(t) Bo [11e, r, Vx, log
( (m)) di C
(17)
Proof.
The key is the following important property holds since the gradient taken w.r.t. 3e only:
Va log Qej (xelr) Vz, log yt. (J),
(18)
The remaining of the derivation parallels
to
traditional DDPM. We can replace Vx, log (2,,7) with
Vx, log (r,r r(zo):
Var log (t,, r) = Ezo.r
Va, log ( r r(or))
(19)
Thus,
V log ferr (x)l1v11)) E, F. Fo ( ir) Ex (eee)(ee) 11180 (Fs, r; 1) - Va, log der- (2) E, m (r) Re, ( (30,, Vz, log dr (2x [xo))1
E, (igh,e.gttt
(20)
(21)
(22)
which is equivalent to our tractable score matching loss.
Rigorous Error Bound for Second-Stage Generation. Utilizing Proposition D.5, analysis of the second-stage diffusion parallels to the first stage, except that the score network takes additional inputs
Assumption D.6. (Second moment bound of molecule features.)
m
(gtx)) 11| [] < xc
(23)
where ((z) is the ground truth distribution of the molecule features, and æ := Bg(z)
Assumption D.7. (Lipschitz score of second stage). For all 2 0, the score V In g(::) is L, Lipschittz
where ((x)) is the distribution noisy latent 2, at diffusion step in the forward process.
Finally, we make some assumptions of the score network estimation error
Assumption D.8. (Score estimation error second-stage diffusion). For all E [0,T],
(x (1so (ano)--nnnlllllaala (in < e score
(24)
This assumption contains the error brought by generating representations, i.e., the TV distance shown Proposition D.4. Later Theorem D. 12 we explicitly deal with the error brought by representation generation, which results in a more fine-grained error bound.
We now present key lemma which facilitates analysis and the proof of the central Theorem D.12.
Lemma D.9. Suppose Assumption D.6, Assumption D.7, Assumption D.8 hold, and the step size := T/Na satisfies 3 1/I. Suppose we sample po(xlr) (r) from Gaussian noise where (r( and denote the final distribution of x as (x). Then the following holds,
(v), 2 ((a (Iez V 3)h + M.y
60,score
(25)
convergence of forward process
discretization error
SCORe esimattion crror
21

---

Geometric Representation Condition Improves Equivariant Molecule Generation
Proof. Recall the notation that (a.p(a) J, Po)- (r) Po predicted by denoising networks 0, sp starting from Gaussian noise 7 (d4+3) Consider the reverse process 22 (mo) starting from fr instead of N(d+3))
TV(po.g(z;) S (V p + TV(2"" .40)
(26)
Using the convergence of the OU process in KL divergence (see (Chen al., 2023)), the following holds for the first term,
TV p & TV(yN(d+3) .4r) S (II( exp(-T)
(27)
The second term is caused by score estimation emror and discretization error, which can be bounded by
(V q0)² 2 z (L(gollp8") (E,, + L2 3)/ L
(28)
We start proving Equation (28) by proving
kb
E go
|100
(kh)
(Ixas kh, V Vlng( 3 (2 (() + 12 3)h h2
(29)
1k-
For E [(k: 1)h, kh),, we decompose
Eqprrp [|ls (KA) (IRns kh, r) V In (xx)))
(30)
go pa [1lsg (Kh)
((a)) (IRh, kh,r) VOka (##n) ) |1 + g. Vqui (xxa) - (x
(31)
+ E |ll (e((znn) [1
(32)
Deh ||
(IrZ)) )|2 |13] + B [|| - 1x
3 9,core
+ Egg
(33)
Qt
Note that we omit the term in expectation last two terms because they are independent Utilizing Lemma 16 from (Chen al., 2023), we bound
V In
gkh (#EA)
= (d 3)h + 12h231 + (L2 Vin
(34)
For the last term,
| (xxx)) | |1
||Vg+++ Vin (x;))l?
(35)
= V In d (xx) 2 10 12112
(36)
where the second term is absorbed into the third term of the decomposition Equation (30). Thus,
En ||10 (kh)
(IRn, kh, r) In (2x, ||
(37)
36 scoor
(d 3)/ L. J [2] 2 ||| ( B || - - 1x
(38)
36 3 8,scrr
- 12, (d + 3)
3) + m + L 3)1² + 1m + +Ndd+ 3)h)
(39)
S,ooor [d +3)++++++ L
(40)
Analogous to (Chen et al., 2023), using properties of Brownian motions and local martingales, we can apply Girsanov's theorem and complete the stochastic integration. Since qu is the end of the reverse SDE, by the lower semicontinuity the KI. divergence and the data-procent processnn inequality (Beaudry & Renner, 201 1), take the limit and obtain
KL(gollpoT) = (Ey, + L2 3)4 + L2,2 m))
(41)
We finally conclude with Pinsker's inequality 2TV² KI.).
22

---

Geometric Representation Condition Improves Equivariant Molecule Generation
This result holds for general represennation-conddioeeedd diffusion models, and our best knowledge are the first to provide theories for represeetatioonccodiioeeeddd generation, which is general generation framework suitable for various domains such as images (Li et al., 2023) and graphs.
Lemma D.9 quantitatively characterizes the bound on generalization error in representation-condiine suggests that the error of representation-conditonn This is because the first two components of the generalization error (i.e., the convergence of the forward process and the discretization error) of the
diffusion. It directly diffusion will be no higher than that of its one-stage counterpart.
ceapenttattoononitiinnnnn diffusion model align with those of traditional DDPM, provided that both are parameterized using the same diffusion processes. Furthermore, the third component (score estimation error) can be made identical if we simply set all representation--elevan parameters in 89 zero and disregard representation's 's impact. We therefore have the following conclusion,
Corollary D.10. Self error than one-stage diffusion model.
diffusson model can have the same or a lower generation distribution
We now give a more between r and a that enables our further qualitative analysis for the argubly better performance.
error bound analysis of representation-conddtinee diffusion, given the relationship
Assumption D.11. (representation-condtiiinee score estimation error of second-stage diffusion). For all & E [0,
V (Ial(( S
2 6.0.cond
(42)
The following main theorem is novel and precise since it (i) deals with the generation error of first-stage representations explicitly; (ii) takes advantages of the conditional distribution ((x]r) in the denoising network.
Theorem D.12. (Theorem 3.2 in the main text) Suppose Assumption D.6, Assumption D.7, Assumption D.11 hold, and the step size h : a/ satisfies h 3 1/ T.g. Suppose we sample x ~ po po(r)) from Gaussian noise where r ~ p.((r), and denote the final distribution of x as (x)) Define Po grie, which is the end point of the reverse process starting from gripp instead of Gaussian. Here qr is the T-th step in the forward process starting from 4o = A Sr q(200((0))) p) (1) (r)rr where the normalization factor. factor. Then the following holds,
TV(po.p(0), / (KT exp(-T)+ + (Ix V 3)h hIVT
(43)
convergence of forward pricess
discrctization error
Cp.0.cond VT
+
TV q0
(44)
conditional score estimation ervor
representation generation crror
Proof. The proof sketch parallels that of Lemma D.9, except that in the first step decompose the TV distance follows,
TV(po. ((x)) < TV(po. (Tile (or)) + TV grle (00pp) + TV 90))
(45)
We complete the proof analogously to the proof of Lemma D.9.
Remarkably, when 201e> i.e., Py fully recovers the ground truth marginal distribution of representations q(r), Theorem D.12 has the same format as Lemma D.9 but with cond (0,sooe-- This is because the former is the score estimation error based on explicit relationship between x and r while the latter learns implicitly. Thus, Theorem D.12 is a much tighter bound for representation condiiionnd generation. To the best of our knowledge, this is the first rigorous theoretical analysis on RCG (Li et al., 2023). We now provide some qualitative discussions on why representations can arguably lead to better generalization error.
Typically, representations are powerful (and sometimes even complete) as they encode key information about with potential additional knowledge via pretraining tasks (for example, coordinates denoising for molecules (Zaidi et al., 2022; Feng et al., 2023)). Therefore, it is reasonable to expect that score estimation conditioned on representations can be more accurate (i.e., f0,,c could be significantly smaller than when estimating the score without representation conditioning). If the representations are cempplete-werr same applies when r can be properly transformed back to x by a neural network. More generally, there are intermediate cases where r reflects partial information about x (e.g., a multiset of atoms and bonds), which would still aid in improving prediction.
special case would be = x-this would greatly assist in predicting the noise. The
23

---

Geometric Representation Condition Improves Equivariant Molecule Generation
Extension to Equivariant Diffusion Models. The previous conclusions are generic and can be applied general oppeseenation-coonitttnned generation. However, far we only consider traditional diffusion models without taking into account the permutation II and SE(3) transformation S invariancedequivvrianee of the diffusion model. We thus extend our theory specifically to equivariant diffusion models that operate on symmetry structures, which is the case for our experiments. Moreover, the previous results assume that both the diffusion process and the denoising model treat atom coordinates x and atom type features identically, while the fact is that (1) the noises of atom coordinates always have zero center of mass (CoM), thus actually lies in a subspace with degree of freedom 3(N - 1) as opposed to 3NN; (2) in the forward process, x; and h, are conditional independent give Xo and hg. and since the denoising network processes coordinates X, and h, differently, the score estimation error term can be further decomposed in Assumption 1.13.
Assumption D.13. (fine-grained represeettttion-coddiioon score estimation error of second-stage diffusion). For all t E [0,
(rr (xtnilll - =Eropo (r), (o (x,,t,r) Xx, ln (x,l))1 + (xn,tr) Vh, In (bbl)r))
(46)
(47)
2(E),),cond)) + ( (oon)))
(48)
where sg (-) and $g (.) refer to the predicted score of coordinates and atom type features by the score network, respectively,
Combining all the pieces, conclude the following result Theorem D.14 for equivariant diffusion models that generate 3D coordinates described above.
Theorem D.14. Suppose Assumption D.6, Assumption D.7, Assumption 1.13 hold, and the step size h == T/ satisfies h 3 L/ Suppose we sample x ~ (o(x))) from Gaussian noise with zero CoM in the coordinate subspace where r pr)()), and denote the final distribution of as (x)) Define Po 9719 which is the end point of the reverse process starting from drie instead of Gaussian. Here qripp is the T-th step in the forward process starting from gol := A S, g( )p (r)dr N(d++3)
where A is the normalization factor. Denote "lo-ood the N(d+ ((d+3 3) dimenssionll Gaussian but with zero center of mass in the N 3-dimensional subspace for coordinates. Denote 50-) as the distribution after acting permutation group and SE(3) transformation Qi on the data from p(.). Then the following holds,
TV (1). [(x)) a II, (po.))/((p... ((x))
(49)
v (T 1l sp 11
Nid+3) exp(- + (L2 N(d+ 3)h + L)yT
(50)
comvergence of forward process
discretizaioon ervor
(,o,.oo + c.ood VT VT
TV (golp: go)
(51)
conditional score estimation ervor
representation generation error
where ax.)) E [0,1],
Proof. The proof parallels the proof of Theorem 1.12 with three distinct parts.
1. Due to the existence of the permutation group and the SE(3) group Q, we need to consider the distribution pe.p (2). Note the definition appe,p, 11, 22) := TV (2))o (() and by data processing inequality (Beaudry & Renner, 2011), ),, (p II, 22) [). Specifically, when the denoising model po., is constructed invariant/equivariant permutations IT and SE(3) transformations Ω (which means the model treats all elements in one equivalent class the same), ap 11, ?) reaches the minimum; see (You et al., 2023) for further explanations.
2. Since the Gaussian noises for coordinates sampled from subspace with zero center of mean, the prior distribution NN (d+3)
y (d++3) in Theorem D.12 should be replaced with T0-CoM- the [ + 3) dimmenssoa Gaussian but with zero center of mass in the N x 3-dimensional subspace for coordinates.
[d+33 notable that the degree of freedom of 7o- CoM is actually (d + 3)- 3, and the remaining of the proof still holds (Feng et al., 2024).
3. As explained before, in a properly designed forward process, Xt and h, can conditional independent give xo and ho- Meanwhile, the denoising network processes coordinates x, and h, differently, the score estimation error term can be
24

---

Geometric Representation Condition Improves Equivariant Molecule Generation
further decomposed in Assumption D.13; see (Feng al., 2024) for more details. Therefore, the term 6,,8, VT in Theorem D.12 can be replaced by / /,,,0,,oon) + (.o.,,ood) which is further bounded by (4%),, p,o, + (ood)
In conclusion, we provide a detailed characterization of the generalization error of representation-condiit diffusion models in this subsection. estimation errors, are not constants; they are functions of the SDE total time I and the number of diffusion steps Na. Our conclusions also explain why second-stage diffusion steps Na is decreased for faster generation. This is because the score estimation error can remain small even when the number of diffusion steps is reduced, given the guidance from representations. As result, reducing Na causes a slower increase in cond (Na) compared to the score estimation error without representation conditioning, leading to representation
important to note that some parameters in our assumptions, such as Lipschitz scores and
innpeettttioncconitionnne generation methods remain competitive even when the number of
opprsennationnccondttonnne conditioned generative models' strong performance with fewer steps.
D.2. Conditional Generation
In this subsection, we aim to prove that conditional generation using our representaaiooncconiiioos generation have probable reward improvement. While we used e to denote conditions in the main text, we use the notation y instead for the targets or "reward" to keep coordinate with existing literature. Denote da := g(-ly a) as the ground truth conditional distribution, and Pa : pl-ly = a) the estimated distribution. Suppose the ground truth distribution satisfies y == (x, which can be decomposed
(x, r) g" (m)) h (1)
(52)
where we denote r = 2 when r ~ q(z), r = rii when r ~ q(r), and fz r) = (xx r) when x ~ g(z), ~ gr)).
To start with, we assume a linear relationship between and y. which reasonable thanks to the powerful pretrained model (which makes the representations helpful in predicting molecule properties and even complete) if some noises are allowed. In detail, the reward is f (2,r) w + and & ~ (0,,v)) In some cases, may further make Gaussian assumptions on r (WLOG, r ~ (1, Σ)) but is not necessary.
D.2.1. PARAMETRIC CONDITIONAL SCORE MATCHING ERROR
First, we give a detailed form of the representation score estimation error presented in Assumption D.I under the assumptions above.
Lemma D.15. For 8 2 O, with probability 1 - 5, the score estimation error 2 Expsccore is bounded by
T 1o tto ((r Iog (r(el)))) (,(r,,, = = (
('
N(S, 1 )dl log
(53)
n
where to the early stopping time of the SDE, n is the number of samples, the parametric function class of denoising network, and N(S, 2,) the log covering number S. S. When S is linearly parameterized, N(S, 1,) = OOd" log( ya )).
Proof. This is direct extension of Lemma C. from (Yuan al., 2023). Note that we consider the special case where the low-dimensional subspace is the original space (i.e., A = La, and d D d, in their paper), and our noised linear assumption between r and y is identical to their pseudo labeling setting (i.e., ÿ = TTr T + & where & ~ NO v²)). We only provide the proof sketch here.
When follows the Gaussian design, some algebra gives
a(t) h(t)
B, (a(t)r+ + h(t)) D² yw)
V, log (((r,))))
(54)
h(t)
where a(t) = exp(-t/2),A(t = exp(-t), B(t) = (a²(t) +
MCE
ww T + h(t))=1) We then bound the
DF
estimation error with PAC-learning concentration argument by using Dudley's entropy integral bound the Rademacher
25

---

Geometric Representation Condition Improves Equivariant Molecule Generation
complexity, and obtain
1 2 = ( to
N(S,
d log
1
(55)
n2
Further, the log covering number of under Gaussian design satisfies
n Mis, 1,) - de log
d to to
[
(56)
where 0 Amin is the smallest eigenvalue of 2, and typically the early stopping time O In (Yuan al., 2023) the authors assume v² = 1/d, which states that the variance v² of regression residuals reduces when we increase the representation dimensions, which is reasonable.
Lemma 1.15 provides detailed score estimation error given the linear assumption between and y. which serves as special case of G,ccore Substituting it into Proposition D.4, we can obtain quantitative result of representation generation error.
D.2.2. REWARD IMPROVEMENT VIA CONDITIONAL GENERATION
Next, we want obtain the reward guarantees of the generated samples give the condition y. Define the suboptimality of distribution P as
SubOpt(P:z") = y² P/"
(57)
where y" is the target reward value (condition) and f is the ground truth reward function. We use the notation Pa := (rlyy" = a), then we have the following result for SubOpt(fa: = a). which can also be viewed as a form of off- off-policy regret.
Lemma D.16. (Theorem 4.6 in (Yuan et al., 2023).)
SubOpt(P.:v" = a) S Erwga [tu - Tr| + Eyng. (r((r])] - Eropa [g (m11)]
+ +Erpp [h (r1))
(58)
82
Ea
Proof. Recall the notation qa = 9(11 = a), we have
E E [f']
(59)
E Tu 4u (r)] EErpp (r)] Er [f]]
(60)
En [f(r) Errngo [[((r) - (r) Erapa (x)] - Eynug [f*(r)]]
(61)
F T.g. [f()] Eyrng [[((r) g" [r(]]] Errgg [g" (m1)) Errp, [g (n Bynop (r1))
(62)
where w is the estimated w by Ridge regression, Er-g, (f(r)) = [a and r TI, (r) g" (r) when r Qa.
Here we give brief interpretation of the decomposition. & is the prediction and generalization error coming from regression, which is independent from the diffusion error. Both &g and &3 come from the diffusion process, where & reflects the disparity between P. and la on the subspace support while Eg measures the off-subspace error in P... The following results give concrete bounds for the terms in Lemma D.16.
Bounding Regression Error with Offline Bandits.
By estimating w with Ridge regression, we have
i (RTR+ R - (Rw" +n)
(63)
where R = (ri...., Tw) and r = (f....., En) where & ~ N.O, Define VA = R X E = 1V and
Eng, = Ryrng rr and take 1, we have
26

---

Geometric Representation Condition Improves Equivariant Molecule Generation
Proposition D.17. With high probability.
O11[0"11 + u² /d, d, log
& V Tr(E5'Inn Eg.)
(64)
Vn
Further when has Gaussian design r N(,,E))
a²
1 0
+ dr)
(65)
| Σ
when n (max Amatn
w d. }).
Proof. First we have
& 11g, 1r (w- 107--() w))l S Bp, Brlllrl -110* villus
(66)
where
E3, Irllv S V
T) ~ v (Ea.)
(67)
Hence we only need to prove
1|0- wllv 3 =0 v² v2 log
(68)
Using the closed form expression, we have
④ w-w- w = -V '
(69)
Thus,
I|l* ullvs S || allyy + \||/> I|
(70)
where Ily S VAllu"I|, and according (Abbasi-yadkorii et al., 2011).
| 7/ All = || vll(n ( O(v? ddr logn)
(71)
with high probability. We hence conclude the first part of proof.
Further, when has Gaussian design (N, E), according to Lemma C.6 of (Yuan et al., 2023), prove the results.
The key here is the conditional distribution follows the Gaussian below,
(r(r(((r)= = a) = -N(# Eni(" Ene + 12) v m), E - Eni(w
Ene + Tn
(72)
Thus
1/22 ΣΟ 1 (1100 - 22)2
a
(race(Ei>22in trace
(1
0
d+
(73)
v Ammin
|
Notice that || 2 ||w0 11x - 1110- 1|x. We have
|| - w"|lx o |110 /1110011
D² d, log(n)
(74)
Vn
where can prove || 2 3||ww|]
1|x when 22(
dr [|w-
Remarkably, this a more precise bound improving the results (Lemma C.5 and C.6) in (Yuan et al., 2023), where we make less assumptions on the relationship between y and r, explicitly taking ||w|| and 22 into account.
27

---

Geometric Representation Condition Improves Equivariant Molecule Generation
Bounding Distribution Shift in Diffusion. restricted under function class L as
We define the distribution shift between arbitrary distributions Pt and
Ex [x))
(2) == sup LEC
(75)
E Ex (x))
We have the follow lemma. Lemma D.18. Under the assumption that follows Gaussian design, then
TV( ya) -a)= [o( T(a(r, aa),
a), Aryi S)
crr
(76)
Amin
where E, is defined in Lemma D.15. We can bound En with:
Eg = =0 (TV((( vo0
(77)
where M(a) = O( o² bound
||w=|| + d). By plugging in e, Ö Ö
de ) when to = (dt/n) 10, Eg admits the best trade off with
lov
0VV T(g(r, = a)
x Aryi
& 6
(dt/n)
(78)
Amin
Proof. The proof directly follows from Lemma C.4 and Lemma C.7 in (Yuan et al., 2023). However, the conclusion is slightly different as we do not assume a low dimensional subspace of r.
One can also verify that when follows Gaussian design, (q((r, y = gryi (rij 5) Oaa v dz).
D.2.3. SECOND STAGE OF CONDITIONAL GENERATION
Now that we have proved that: (i) the first-stage diffusion model can estimate the score function with provable error bound (Appendix 1)2.1); and (ii) the generated representations have provable reward improvement (Appendix D.2.2). We continue to show that the ultimate generated samples also have distribution shift towards the desired target in the following contexts. Particularly, we want to answer the question: why utilizing the conditionally generated representations is enough for the second stage generation?
First, when we use the generated representations as the only condition of the second stage diffusion model, the generation process is identical the second stage of unconditional generation. Therefore, the results in Appendix D.1.2 can be directly applied to analyze the second stage generation of conditional generation, which states that the generation conditioning on representations has small TVV distance error compared with ground truth conditional distribution. Thus, when we have high-quality first-stage generation, the corresponding second stage generation would introduce almost no additional error, which implies provable reward improvements towards the desired target. In addition, the well-pretrained encode ensures that the correspondence between representations and data points is good, which makes it possible to rigorously construct the data points given the representations special case would be r the complete representation of x))
We then partially answer this question from the information theoretic perspectives. We use H(-) to denote the information entropy, and ((::-) denote the mutual information between two variables.
I(x;r) 2 [xx;)) On the one hand, r contains enough information to recover the targets y thanks to the results in Appendix D.2.2, thus we do not explicitly need y for the second stage. One the other hand, benefit from the pretraining task, the representations obviously contains more information in addition to y. This assumption is valid if w" in previous analysis is sparse (there are components in independent of y), i.e., H > Hyy) Therefore, generating conditioning is much easier than generating conditioning y (traditional one stage generation), the score estimation error of the former one would obviously be much smaller than the latter.
[xx v) > T(r;y). Recall Equation (52) which states the target property depends on both x and Hence, may contain additional information of obtained from pretrained tasks that is hard to (or cannot) be directly extracted from x - the complex pretrained model assists in extracting relevant information in our two-stage generation, while one-stage generation solely relies on the single denoising model to do so. Therefore, by leveraging representations with provable error bounds, we can better shift the distribution towards the target.
28

---

Geometric Representation Condition Improves Equivariant Molecule Generation
In summary, is ideal middle state connecting æ and y- it is easy to recover from y (Appendix D.2.1) and to recover from y (Appendix D.1.2), and vice versa. In comparison, it is somewhat more diffieult to directly recover from y or predict y from x. Consequently, may be better indicator of y compared with due to the aforementioned reasons.
Indeed, one-stage diffusion models generate directly from conditions need to optimize highly complex score v log p(rly) where x and are highly non-linearly mapped. As Theorem E.4 (Yuan et al., 2023) points out, the nonparametric SubOpt of æ generated by deep neural networks larger than our results Appendix D.2.2, which further validates the advantage first generating that can be well mapped from y.
E. Visualization
E.1. Representation Visualization
To illustrate how well (r) fits q(r), we sample from both ((r) (i.e., the representations produced by the pre-trained encoder for the QM9 and GEOM-DRUG datasets) and the trained representation generator (r). We then visualize them in Appendix E. with colors indicating whether the samples are from ((r) or (r()). We compute the Silhouette Score of the clustering results, scaled by 102 for clarity. A score close to zero suggests that the two clusters are difficult to distinguish, indicating a good fit between (r( and q(r). Similarly, we provide the visualization of conditionally generated representations in Figure 9
-0.0001
0.6503
Figure 8: I-SNE visualizations of representations unconditionally generated by the representation generator (T 1.0) vs. those produced by the rre-trained encoder the QM9 and GEOM-DRUG datasets. The Silhouette Score is scaled by 10 for clarity.
E.2. Visualization of Molecule Samples
In this section, we provide additional random molecule samples to offer deeper insights into the performance of GeoRCG. Figure 10 and Figure 11 show unconditional random samples generated by GeoRCG trained on the QM9 and GEOM-DRUC datasets, respectively. Figure 12 presents random samples conditioned the a property, along with their corresponding errors.
29

---

Geometric Representation Condition Improves Equivariant Molecule Generation
0.9062
0.0478
0.2397
-0.1558
0.2507
-0.5989
Figure 9: t-SNE visualization of representations conditionally generated by the representation generator vs. those produced by the pre-trained encoder on the QM9 dataset: (a) a, (b) Ac, (c) <HoMo. (d) fLUMO. (e) μt and (f) C.,. The Silhouette Score scaled by 102 for clarity.
30

---

Geometric Representation Condition Improves Equivariant Molecule Generation
Figure 10: Unconditional random samples from GeoRCG trained on QM9. The number of nodes is randomly sampled from the node distribution (NV).
31

---

Geometric Representation Condition Improves Equivariant Molecule Generation
Figure 11: Unconditional random samples from GeoRCG trained on GEOM-DRUG. The number of nodes is randomly sampled from the node distribution (NV.
32

---

Geometric Representation Condition Improves Equivariant Molecule Generation
74.27
76.51
68.79
76.89
71.89
68.15
78.27
75.44
73.84
76.0
68.04
76.77
71.43
66.37
77.71
77.01
85.02
76.89
91.32
78.9
71.5
80.32 32 80.47
72.93
87.54
85.57
74.24
91.46
79.08
64.27
71.93
88.19
86.44
68.95
88.33
54.34
55.86
76.36
81.19
71.36
86.72
67.94
88.88
54.32
55.5
77.35
82.95
69.35
79.86
72.46
80.81
56.03
74.47
60.17
62.81
76. 11.11
79.67
72.34
80.96
56.35
75.2
59.8
63.17
76.36
73.3 73.44
78.01 79.09
76.26 75.27
77.77 78.53
78.3 78.4
84.41 84.52
85.95 84.23
65.99 65.99
Figure 12: Conditional random samples from GeoRCG trained on QM9 dataset and a property. Black numbers indicate the specified property value condition, while green numbers represent the evaluated property value of the generated samples. The number of nodes and property value conditions are randomly sampled from the joint distribution ,(N,c).
33


Reviews:
[{"review_id": "l8xGGQCtf7", "reviewer": "Reviewer_eCAm", "summary": "The paper presents GeoRCG, a novel framework for 3D small molecule generation, applicable to both unconditional and conditional settings. The key innovation is a two-step generation process: first, generating a geometric representation, then sampling the 3D molecular structure conditioned on this representation. This approach aims to simplify the generative process, improve sample quality, and speed up molecule generation.\n\nKey benefits in downstream evaluation include:\n\n- Reduced computational cost in comparison to diffusion-based methods.\n- Comparative evaluation against state-of-the-art methods, showing significant improvements in conditional generation tasks.\n- Theoretical bounds for the proposed diffusion setup.\n\n## Update after rebuttal\n\nI have reviewed the additional tables and experiments provided by the authors, and they have satisfactorily address my concerns about fair comparisons and missing metrics. I suggest that authors revise their claims to provide more context in the final version, and incorporate the additional tables in the main paper/appendix. I have raised my rating to 4 (Accept).", "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": "Dear Reviewer eCAm,\n\nWe sincerely appreciate your thorough and insightful review! Below, we address each of your comments in detail. **We have provided additional tables in an anonymous link: https://anonymous.4open.science/r/rebuttal-8746. (Please refer to our response to Reviewer zxks for an *alternative link* in case the previous one encounters issues.)**\n\n---\n\n###  **1. Claim Clarifications and Conditional Experiments for SemlaFlow**\nWe acknowledge that some claims need additional context for clarity, and are committed to modifying them in the next revision. Here are brief clarifications:\n- In Claim 1, 2, we refer to “improvements” and \"faster\" **over base models**, where our GeoRCG consistently enhances the performance of base models.\n- Regarding Claim 3, we have conducted our own experiments for SemlaFlow's conditional generation for a more comprehensive comparison. Please refer to **Table 4 in the link.** The results demonstrate SemlaFlow’s superior performance in this setting; however, **GeoRCG consistently enhances SemlaFlow’s performance in conditional generation.**\n### **2. Comparing all the methods head-on.**\nWe clarify that we have split the experimental results for several important reason:\n- **Differences in dataset processing and evaluation:** The models in Table 3 were trained with different processing pipelines than those in Table 1 (e.g., Table 3 used the 5 lowest-energy conformations per structure for DRUG, while Table 1 used the 30 lowest). Additionally, dataset splits and evaluation criteria (e.g., allowable valences, molecule sanitization) differ significantly.\n- **Differences in model architecture, i.e., 2D&3D vs. 3D-only:** Models in Table 3, such as SemlaFlow, jointly learn both 3D conformations and 2D bonds, while 3D-only models like EDM focus solely on 3D conformations. **Direct comparison is unfair** since metrics like validity and stability depend on 2D graphs, making 2D&3D models generally perform better.\n- **Difference in evaluation focus:** Tables 4 and 5 report performance with fewer sampling steps, focusing on efficiency, while Table 1 shows best-case performance. However, **for a more comprehensive comparison of 3D-only models, we have reorganized the results into Table 5 in the link** respectfully for your review.\n### **3. Why we did not include energy and strain metrics in Table 1, 2, 4, 5?**\nWe acknowledge the absence of these metrics, and would like to provide the following clarifications:\n- Energy and strain metrics are introduced by SemlaFlow due to the **lack of effective 3D metrics for evaluating the 3D conformations generated by these 2D&3D models**, reasons as stated above. However, for the 3D-only methods presented in Tables 1, 2, 4, 5, the stability, validity, and property MAE metrics **already provide a reflection of the quality of the generated 3D conformations**, since the 2D bonds or properties used in these calculations are inferred from the generated 3D conformations.\n- For completeness, we evaluate our method along with selected baseline models (EDM, GeoLDM) on energy/strain metrics, please see **Table 3 in the link**. Notably, **GeoRCG continues to show consistent improvements under these metrics as well**.\n- We appreciate your suggestion regarding “Per Atom Metrics” and have added per atom metrics. Please see **Table 1, 3 in the link.**\n### **4. Clarifications regarding sampling time and steps of SemlaFlow experiments.** \nThank you for your valuable comment regarding this! In response, we have annotated the sampling steps and times in **Table 1 in the link**. We would like to highlight several key points:\n- GeoRCG consistently improves upon SemlaFlow **under the same sampling steps.**\n- The first stage (rep sampling) incurs significantly less computational cost than the second stage, even for SemlaFlow, which is already highly efficient in molecule generation.\n- GeoRCG (Semla) takes about twice the time of SemlaFlow with the same sampling steps, mainly due to classifier-free guidance, which doubles the batch size to enable guidance [1]. However, we argue that: \n    - Even with **half the sampling steps** (roughly the same time), GeoRCG **still outperforms base models** (Tables 4 and 5 in the main paper, and Table 1 in the link). \n    - Even **without CFG** (same steps → similar time), GeoRCG often surpasses its base model (Table 2 of the link). \n    - With advancements in GPUs, the overhead from CFG will become less significant due to parallel acceleration, while the sequential nature of generative models will continue to be a limiting factor.\n\nIn summary, GeoRCG holds its efficiency advantage for SemlaFlow, too. We will add a detailed discussion of SemlaFlow’s sampling time and include a more thorough analysis of the impact of CFG in the revision.\n    \n---\n\nOnce again, we sincerely appreciate your valuable comments and look forward to further discussion!\n\n[1] Ho, Jonathan, et al. \"Classifier-free diffusion guidance.\"", "claims_and_evidence": "The proposed claims needs some revision to provide more context:\n- Claim: The proposed method significantly improves molecule generation quality.\n    - These results hold only for the setup of conditional generation, the results are comparable to other SOTA models in unconditional setup. This should be clarified.\n\n- Claim: The proposed method achieves faster generation with fewer steps.\n   - Speed-up should only be claimed in comparison to other diffusion models, as some other methods (e.g., SemlaFlow) using flow matching use the same number of steps, and could arguably sample faster due to a single step sampling procedure instead of two. \n\n- Claim: The paper reports an average 31% improvement over SOTA baselines.\n     - The comparison with SemlaFlow is missing, and the results might not be as strong if SemlaFlow is also trained in conditional setup and using the same number of sampling steps. However, this is justifiable since SemlaFlow paper does not report conditional generation results.", "methods_and_evaluation_criteria": "- The overall proposed framework is reasonable for the problem. Representation generation prior to sampling is unexplored in molecular generation context, and the design of the paper is well motivated. \n- The main evaluation datasets are QM9 and Geom-Drugs as is the standard practice for papers in this area.\n- The evaluation primarily uses validity, stability, and uniqueness metrics.:\n     - Energy and strain metrics are missing from Table 2. \n     - It is recommended that energy and strain **per atom** metrics are used as is the standard practice.\n - Clarity in evaluation setup and comparisons:\n    - Table 3 and Table 4 should be integrated into Table 2 and all the methods should be compared together to provide accurate context for the improvements. The claimed speedups only hold when flow based methods such as SemlaFlow are also added to the table. \n- Comparison with SemlaFlow needs more depth, particularly regarding:\n    - Can SemlaFlow also be used for conditional generation ? \n    - It should be clarified if both the proposed method and SemlaFlow have similar computational costs.\n    -  Whether representation learning provides a real benefit over SemlaFlow", "theoretical_claims": "I did not carefully check the proofs, but the theoretical claims are mostly derivative from the previous works.", "experimental_designs_or_analyses": "The experimental design is overall justified, however there are missing details in the evaluation and claims as described in the previous section.", "relation_to_broader_scientific_literature": "The paper addresses the problem of conditional and unconditional 3D small molecule generation. Other works in this area primarily develop diffusion/flow based generative modeling for learning the molecular distributions and subsequently sampling. Contrary to previous works, this paper proposes a 2-step procedure, for first generating a representation and subsequently sampling the molecule.", "essential_references_not_discussed": "The paper discusses all the relevant essential references.", "other_strengths_and_weaknesses": "- Strengths:\n\n    - Novel approach: The two-step generation process is well motivated.\n    - Strong empirical results: Significant improvements in conditional generation.\n    - Efficient generation: The model reduces diffusion steps without quality loss.\n\n- Weaknesses:\n    - Missing clarity in evaluation: Too many tables, and all the methods are not compared head-on. \n    - Major claims require more context, as stated in the previous section.", "other_comments_or_suggestions": "The clarity of the paper can be improved. Several key modeling details are available in the Appendix, the readers would benefit from providing more information in the main paper. The evaluation and results are not well presented and several tables can be integrated to provide a clearer picture of head-on-head comparisons.\nThe figure needs improvement and does not provide a good picture in describing the method.", "questions_for_authors": "- Can SemlaFlow be used in conditional settings? \n- Could the authors comment on the sampling steps for SemlaFlow and the overall generation time in comparison ?\n- I would like to review all the methods presented in Tables 1,3, and 4 compared against each other, by clearly indicating the number of samplings steps for each and other important metrics (including stability and energy). \n\nI would consider raising my rating if the above concerns are satisfactorily 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": "IiilboeATX", "reviewer": "Reviewer_5wZu", "summary": "This paper proposes a new framework, GeoRCG (Geometric-Representation-Conditioned Molecule Generation), to enhance the performance of molecular generation models. The proposed approach decomposes molecular generation into a two-step process as follows:  \n\n- **Geometric Representation Generation**: A pre-trained geometric encoder (such as Uni-Mol or Frad) is used to encode molecular geometric information into a compact representation.  \n- **Molecule Generation Based on Geometric Representation**: The representation obtained in the first step is used as a condition for a diffusion model (such as EDM or SemlaFlow) to generate molecules.  \n\nIn diffusion model-based molecular generation, any feature vector can be used as a condition. The proposed method embeds molecules into a latent space using a molecular encoder that considers equivariance during training. A loss function is designed to enable sampling from this latent space for generation. As a result, the distribution in the latent space exhibits significantly better properties than the original molecular space, leading to improved molecule generation.  \n\nThe main contributions of this work are as follows:  \n\n- **Improved molecular generation accuracy**: Achieves significantly better performance than existing methods (such as EDM and SemlaFlow) on the QM9 and GEOM-DRUG datasets.  \n- **Enhanced conditional molecular generation**: Demonstrates a 31% improvement over state-of-the-art methods in generating molecules while considering properties such as the HOMO-LUMO gap, polarity, and heat capacity.  \n- **Reduced computational cost**: Maintains nearly the same generation quality even when reducing the number of diffusion steps from 1,000 to 100.", "questions": "", "limitations": "", "rating": 5, "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": "Dear Reviewer 5wZu,\n\nWe sincerely appreciate your expert and thoughtful review and are grateful for your recognition of our work.\n\nIn response to the comments from the other reviewers, we have provided additional experimental evidence at https://anonymous.4open.science/r/rebuttal-8746 (Alternative link in case the previous one encounters issues: https://docs.google.com/document/d/e/2PACX-1vQFD87aQep2Q11albXeuuTKUr9HrrLIALteNrVVsbguQ92c_8ArXr_H43J8xv0WrnuRxbzwAOgBYsav/pub), which we believe will enhance the clarity and depth of our paper.\n\nWe remain open to any further discussions or clarifications. Thank you once again for your valuable review!", "claims_and_evidence": "- **Demonstrated Performance Improvement**: The proposed method outperforms state-of-the-art (SOTA) approaches such as GeoLDM, EDM-Bridge, GOAT, and GeoBFN. Notably, significant improvements in molecular stability and validity are observed.  \n\n- **Successful Conditional Generation**: Previous methods struggled with achieving target properties accurately. GeoRCG reduces property error by 31%, enabling more precise conditional molecular generation.  \n\n- **Reduced Computational Cost**: By incorporating geometric representations, the number of diffusion steps is reduced by 90% while maintaining the same generation quality as existing methods.", "methods_and_evaluation_criteria": "The proposed framework can broadly be considered a **latent generative model**, as it learns the data distribution in a latent space (first stage in this study) and reconstructs it into its original form using a decoder (second stage).  \n\nOne key limitation of existing molecular generation models is that molecules inherently exist on a **low-dimensional manifold** (Mislow, 2012; De Bortoli, 2022; You et al., 2023), yet most approaches model them as distributions in a high-dimensional 3D space with \\( N \\times (3 + d) \\) dimensions.  \n\nThis work draws direct inspiration from **RCG (Li et al., 2023)**; however, RCG is designed for fixed-size, fixed-position image data and does not need to account for molecular-specific challenges such as Euclidean symmetry and permutation invariance. Similarly, **GraphRCG (Wang et al., 2024)** extends the RCG framework to **2D graph data**, whereas this study explicitly handles **3D geometric information with Euclidean symmetry**.  \n\nFurthermore, while **RCG (Li et al., 2023)** primarily focuses on empirical evaluations, this study generalizes the theoretical properties of **representation-conditioned diffusion models** for both **unconditional and conditional generation**, offering a more rigorous understanding of performance improvements.", "theoretical_claims": "I briefly checked the content of Theorem 3.2 but did not rigorously check its proof in the supplementary information.", "experimental_designs_or_analyses": "I checked Section 4 for \"Experiments\", where the QM9 and GEOM-DRUG datasets are used for a detailed comparison of the generated molecules in terms of stability, validity, and property error. Although using only two datasets might be considered a limitation, **QM9** and **GEOM-DRUG** are representative benchmarks for evaluating **3D molecular graphs**. More importantly, the experiments comprehensively assess the proposed method from multiple perspectives, including **unconditional and conditional generation**.", "relation_to_broader_scientific_literature": "While the approach of sampling from a latent space (as in RCG) is not new, to the best of my knowledge, it is novel in the field of **molecular generation**. Given its well-founded motivation, this method has the potential for **broad impact** in the field.", "essential_references_not_discussed": "None", "other_strengths_and_weaknesses": "None", "other_comments_or_suggestions": "None", "questions_for_authors": "None", "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": "gE9KRtbbFv", "reviewer": "Reviewer_zxks", "summary": "This paper introduces the GeoRCG framework which is a method for generating molecules by firstly generating a conditioning vector and then generating a molecule based on this condition. The authors show how the previously introduced RCG framework can be applied to molecule generation in 3D space by training a diffusion model to generate a vector which matches the molecular encoding provided by a pretrained geometric molecule embedding model. A 3D molecule generation model can then be trained in parallel to generate a molecule based given the embedding as conditioning vector. This representation makes it simpler to apply a pretrained 3D generative model to conditional generation tasks since only the conditioning vector generative model needs to be updated. For both diffusion and flow-matching molecule generation models, GeoRCG shows improvements in generative performance and shows strong performance on conditional generation tasks.", "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": "Dear Reviewer zxks:\n\nWe sincerely appreciate your expert and detailed review! Below, we address each of your comments in detail. **We have provided additional tables in an anonymous GitHub repository: https://anonymous.4open.science/r/rebuttal-8746. (Alternative link in case the previous one encounters issues: https://docs.google.com/document/d/e/2PACX-1vQFD87aQep2Q11albXeuuTKUr9HrrLIALteNrVVsbguQ92c_8ArXr_H43J8xv0WrnuRxbzwAOgBYsav/pub)**\n\n---\n\n### **1. Regarding the “very small improvements for unconditional generation”**\nThank you for the thoughtful comment on this! However, we would like to respectfully clarify that, when we refer to improvements, we are emphasizing the benefits of GeoRCG **over its base model, since we aim to improve any models’ performance.** In this context, **the improvement in unconditional settings is substantial**: GeoRCG achieves nearly a 13% improvement over EDM in unconditional generation for molecule stability, elevating it from one of the weakest models to the most powerful one in our comparison.\n\nTo further enhance generation quality, one can certainly apply GeoRCG to a more advanced model, such as SemlaFlow,  and expecting even better performance, as shown in **Table 3 in the paper** and **Table 1 in the anonymous link**. This approach consistently improves result quality and pushes new SOTA results. \n\nWe also respectfully invite you to refer to  **Tables 3 and 4 in the anonymous link** for our supplementary experiments in response to Reviewer eCAm, highlighting the benefits of GeoRCG in additional settings.    \n\n### **2. Is the incorporation of CFG and low-temperature sampling *fair*, and how does GeoRCG perform without them?**\nWe very appreciate this suggestion and have investigated the impact of **totally disabling these components** under the most basic configuration (CFG = 0.0, Temperature = 1.0) to evaluate GeoRCG’s performance on both QM9 and DRUG datasets. The results are presented in **Table 2 in the anonymous link**. Our experiments demonstrate that GeoRCG **consistently enhances the base model** and remains competitive with more advanced methods.\n\nHowever, we would like to respectfully emphasize that CFG and low-temperature sampling (**note that low-temperature is for rep sampling**) are **fundamental to our approach**, as they represent how strongly we enforce representation guidance in the model and how varied the representation should be. Note that it is not that other molecular generative methods “did not use these techniques,” but rather that they “**cannot use these techniques**”: These methods lack representation conditioning and, as a result, are unable to control the conditioning scale or temperature of the representations. This means that the performance boost provided by these techniques **does not imply unfairness or “not brought by GeoRCG”.** In fact, in the computer vision domain [2, 3], these techniques are considered negligible and directly influence the effectiveness of the proposed methods.\n    \n### **3. Sampling time for the first-stage generation and GeoRCG (Semla)**\nWe greatly appreciate your comment on this point. In **Table 1 of the anonymous link**, we provide measurements of the sampling steps and times for both stages.\nNotably, the **overhead** introduced by the first stage and by “adapting the SemlaFlow model to allow conditioning on this vector” **remains small**, even for a model as efficient as SemlaFlow: as evidenced in the “Rep. Sampling Time” column and “Molecule Sampling Time w/o CFG”. For more details, we respectfully refer you to our response to Reviewer eCAM under “Clarifications regarding sampling time and steps of SemlaFlow experiments.” \n    \n### **4. Additional pre-training dataset for UniMol [1]**\nIn our setting, UniMol’s pre-training dataset includes not only GEOM-Drugs but also additional public or purchasable datasets used in its own pre-training settings, as described in its original paper [1]. \n\n---\nOnce again, we sincerely appreciate your valuable comments and look forward to further discussion!\n\n[1] Zhou, Gengmo, et al. \"Uni-mol: A universal 3d molecular representation learning framework.\"\n\n[2] Li, Tianhong, et al. \"Return of unconditional generation: A self-supervised representation generation method.\"\n\n[3] Rombach, Robin, et al. \"High-resolution image synthesis with latent diffusion models.\"", "claims_and_evidence": "Yes, the authors claims are well supported by their extensive evaluations.", "methods_and_evaluation_criteria": "The evaluation has mostly been performed in a way that is fair and consistent with existing models, although GeoRCG uses a number of additional strategies to improve performance such as low temperature sampling and classifier-free guidance that many existing methods do not use. Discussed further below.", "theoretical_claims": "The theoretical results in the main appear to be correct although I did not check them thoroughly.", "experimental_designs_or_analyses": "The experiments seem to be designed correctly and consistently with existing methods.", "relation_to_broader_scientific_literature": "- The GeoRCG method takes a lot of inspiration from the existing RCG method, but the authors extend this to 3D equivariant generation models and introduce an error bound for the representation-conditioned generative model.\n\n- SemlaFlow focuses a lot of its evaluation on the efficiency of 3D molecule generation, but no evaluation was performed with GeoRCG to check how the new method impacts the inference time. How much overhead is there for generating a representation vector, and for adapting the SemlaFlow model to allow conditioning on this vector?\n\n- The improvement over existing methods for unconditional generation is very small, although the improvement over other methods in the conditional setting is more significant.", "essential_references_not_discussed": "NA", "other_strengths_and_weaknesses": "- In addition to conditioning on a representation, the authors introduce a number of additional tricks into the model, including low temperature sampling, classifier-free guidance and representation perturbation. It would be better to see how GeoRCG performs without these tricks and show a full ablation table with various combinations of training strategies and compare to existing methods. Figure 5 is a good start but it is important to see how these results compare against existing methods in an otherwise identical setup. This is especially important since one of the main comparisons in the paper is to the existing EDM model since the setup is otherwise the same, but it is difficult to tell how useful the RCG method is if GeoRCG uses other sampling tricks that EDM does not.", "other_comments_or_suggestions": "- Equation 6 should use a subscript on the first $q$ term.", "questions_for_authors": "- What dataset was the UniMol pretrained encoder trained on for the GEOM drugs model, in addition to GEOM drugs?", "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": "7GWVLLZ3IP", "reviewer": "Reviewer_Tpcv", "summary": "This paper proposes a general framework, namely GeoRCG, for 3D molecule generation. Basically, it factorizes a 3D molecule generation into two stages: the first is to generate a geometric representation and the second step is to generate 3D molecules conditioned on the geometric representation. Such factorization can make the task easier and the resulting model has been shown to achieve great performance on both unconditional and conditional generation.", "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": "Dear Reviewer Tpcv,\n\nWe sincerely appreciate your insightful and thorough review, and we are grateful for your recognition of our work. Below, we provide detailed responses to each of your comments.\n\n---\n\n### **1. Does additional pretraining data in the representation encoder introduce unfairness?**\n\nThank you for raising this important point! We believe this question is crucial and warrants further clarification. We argue that **additional pretraining data does not introduce unfairness** for the following reasons:\n+ It’s not that “other methods didn’t use any data for the pre-trained encoder,” but rather that they **cannot leverage valuable prior knowledge from additional data** due to the absence of an effective representation guidance mechanism. In contrast, GeoRCG successfully utilizes pre-trained encoders and advanced pre-training techniques for generation purposes. This approach is consistent with practices in the CV domain [1].\n+ While it may seem that we use additional pretraining data, generally speaking, **no customized pre-training** or **additional training beyond generative training** is necessary for GeoRCG’s development. We simply leverage **public** pre-trained models, like Frad, that have been trained using **general** pre-training methods! Although we did pre-train models for the DRUG dataset, this was due to the lack of public checkpoints containing rare elements (e.g., Bi in some drugs). However, we believe that as molecular pre-training progresses and the effectiveness of our work, more advanced pre-trained checkpoints will become publicly available, eliminating the need for pre-training ourselves, similar to the CV domain [1, 2].\n+ Even when seriously considering the effect of additional data in generative training, we emphasize that the extra training data **solely contributes to forming a meaningful and structured latent representation manifold in the encoder**. During GeoRCG’s training, we still focus exclusively on generative data (i.e., QM9 and DRUG), with the encoder simply translating molecules from this data into structured, meaningful representations.\n\n### **2. Yes, a universal encoder works for both benchmarks!**\n+ We use Frad for QM9 experiments and UniMol for GEOM-DRUG experiments simply **to achieve optimal performance.**\n+ When **using Frad for both benchmarks**, GeoRCG **consistently shows improved performance:** As shown in lines 929-943, Frad-based GeoRCG achieves an Atom Stability of 84.4 (compared to EDM’s 81.3) and Validity of 96.9 (compared to EDM’s 92.6). While these results are strong, UniMol-based GeoRCG achieves slightly higher Validity (98.5) with similar Atom Stability (84.3), which is why we selected it as our primary generator.\n+ Our analysis in Figure 6 and lines 942-943 further demonstrates that pre-trained encoders with more structured representation distributions tend to yield superior performance.\n\n---\n\nOnce again, we deeply appreciate your valuable feedback and look forward to continued discussions!\n\n[1] Li, Tianhong, Dina Katabi, and Kaiming He. “Return of unconditional generation: A self-supervised representation generation method.”\n\n[2] Li, Tianhong, et al. \"Mage: Masked generative encoder to unify representation learning and image synthesis.\"", "claims_and_evidence": "The main claim is that by the above factorization, the approach can achieve more effectiveness and faster generation. All these claims have been supported by experimental results.", "methods_and_evaluation_criteria": "The main method contributions are:\n\n* The novel idea of decomposing molecule generation into two substeps and each of these two steps make sense.\n* The two generators, including representation generator and molecule generator, can be trained in parallel.\n* The framework can work on both unconditional and conditional generation. While on different conditions, only the representation generator needs to be retrained. This sounds exciting.\n\nThe evaluation is performed on widely used benchmarks for molecule generations. The metrics looks valid to me.", "theoretical_claims": "I haven’t checked the details of the proof for Theorem 3.2 but the resulting takeaway sounds reasonable to me.", "experimental_designs_or_analyses": "I checked all the experimental designs, including unconditional and conditional generation, and they are valid to me.", "relation_to_broader_scientific_literature": "The proposed two-stage generation strategy are quite novel to the molecule generation field. It brings insights and important performance gains. Even though such two-stage idea has been explored in other domain such as image generation, I believe extending this to 3D molecule generation should be recognized.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "I have included the strengths before. One remaining concern is that since GeoRCG used additional data for the pre-trained encoder, the comparison to other existing molecule generation models might not be strictly fair as they didn't use any data for the pre-trained encoder.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "Why do you use different pre-trained encoders for QM9 and GEOM-DRUG benchmarking experiments? Does a universal encoder work for both benchmarks?", "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 submission presents an approach to decompose molecular generation into two stages. In the first stage a geometric representation is generated and in the second a molecule is generated conditioned on the geometric representation. The paper is well written and the experimental evaluation supports the claims of the paper. Reviewers collectively agreed on the value of this submission to the ICML community.
Decision: Accept (spotlight poster)