UniGEM: A Unified Approach to Generation and Property Prediction for Molecules
Authors: Shikun Feng
Yuyan Ni
Lu yan
Zhi-Ming Ma
Wei-Ying Ma
Yanyan Lan
Venue: ICLR 2025 Poster (2025)
Keywords: Molecular generation
representation learning
Primary area: generative models

Abstract: Molecular generation and molecular property prediction are both crucial for drug discovery, but they are often developed independently. Inspired by recent studies, which demonstrate that diffusion model, a prominent generative approach, can learn meaningful data representations that enhance predictive tasks, we explore the potential for developing a unified generative model in the molecular domain that effectively addresses both molecular generation and property prediction tasks. However, the integration of these tasks is challenging due to inherent inconsistencies, making simple multi-task learning ineffective. To address this, we propose UniGEM, the first unified model to successfully integrate molecular generation and property prediction, delivering superior performance in both tasks. Our key innovation lies in a novel two-phase generative process, where predictive tasks are activated in the later stages, after the molecular scaffold is formed. We further enhance task balance through innovative training strategies. Rigorous theoretical analysis and comprehensive experiments demonstrate our significant improvements in both tasks. The principles behind UniGEM hold promise for broader applications, including natural language processing and computer vision.

Published as a conference paper at ICLR 2025
UNIGEM: A UNIFIED APPROACH TO GENERATION AND PROPERTY PREDICTION FOR MOLECULES
Shikun Feng 1 Yuyan Ni Yan Lu Zhi-Ming Ma², Wei-Ying Ma', Yanyan Lan 15
Institute for AI Industry Research (AIR), Tsinghua University
2 Academy of Mathematics and Systems Science, Chinese Academy of Sciences
3 UUniversity of Chinese Academy of Sciences
*Department of Computer Science and Technology, Tsinghua University
Beijing Academy of Artificial Intelligence
ABSTRACT
Molecular generation and molecular property prediction are both crucial for drug discovery, but they are often developed independently. Inspired by recent studies, which demonstrate that diffusion model, prominent generative approach, can learn meaningful data representations that enhance predictive tasks, we explore the potential for developing a unified generative model in the molecular domain that effectively addresses both molecular generation and property prediction tasks. However, the integration of these tasks is challenging due to inherent inconsisten- cies, making simple multi-task learning ineffective. To address this, we propose UniGEM, the first ciftusion-bbsedd unified model to successfully integrate molec- ular generation and property prediction, delivering superior performance in both tasks. Our key innovation lies in novel wwo-phase generative process, where pre- dictive tasks are activated in the later stages, after the molecular scaffold is formed. We further enhance task balance through innovative training strategies. Rigorous theoretical analysis and comprehensive experiments demonstrate our sigmificant improvements in both tasks. The principles behind UniGEM hold promise for broader applications, including natural language processing and computer vision!
INTRODUCTION
Artificial intelligence, particularly through deep learning technologies, is advancing various appli- cations in drug discovery. This encompasses two major tasks: molecular property prediction (Zaidi) let al. 2022 Feng et al. 2023al Ni et al. 2023 2024) and molecule generation (Hoogeboom et al.) Gao et al. 2024). The objective of molecular property prediction is to learn Tunctions that accurately map molecular samples to their corresponding property labels, which can facilitate the virtual screening process. Meanwhile, molecule generation aims to estimate the under- lying molecular data distribution, offering significant potential for automatic drug design. Although considerable research has been conducted in these areas, they have largely progressed independently. overlooking their intrinsic correlations.
2022 Guan et al. 2023
In our opinion, the essence of these two tasks lies in molecular representations. On the one hand, the effectiveness of various molecular pre-training methods demonstrates that molecular property prediction relies on robust molecular representations as a foundation. On the other hand, molecule generation requires deep understanding of molecular structures, enabling the creation of good representations during the generation process. Recent research findings provide support for our per- spective. For instance, researchers in the computer vision field have shown that diffusion models inherently possess the ability to learn effective image representations (Chen et al., 2024 Hudson al., 2024 MMttal et al., 2023). In the molecular domain, studies have indicated that generative pre- training enhance molecular property prediction tasks (Liu et al.) 2023) Chen et al. 2023a). par- ticularly through data augmentation using diffusion models. However, these methods often require
"Equal contribution. Work was done while Vuyan was a research intern at AIR, 1 Correspondence to 1 a mnynnynn@aa r. ts tet.tinnnhii inguug inggua.edu.co edu.cn. TTee code is released publicly at attttps: thub
http://:///gtthub.cm////e////eee kun /un//nnnei GHM

---

Published as a conference paper at ICLR 2025
Molecule Growth
Molecule Nucleation
nucleation time
property well-defined
property ill-defined
Xtt
xtn+
XT
Coorcinate Denoleing eq?
data= h, c)
equivariant network Φb
Activated if
Atom Type Prediction eq.3
Property Prediction eq.4
Figure I: The two-phase generative process of UniGEM. We treat molecule generation as a two- phase problem: nucleation and growth, defining the separation time nucleation time. Properties are only well-defined the growth phase, so during training, property and atom type predictions are incorporated only in the growth phase.
additional fine-tuning to achieve optimal predictive performances. Additionally, while predictive tasks can guide molecule generation (Bao et al.) 2022, Gao al. 2024), it remains unclear whether they can directly enhance generative performance. Therefore, existing studies have not sufficiently elucidated the relationship between generation and prediction tasks, raising critical question: can we develop unified model to enhance the performance of both tasks?
straightforward approach to combining these two tasks is to use a multi-task learning framework, where model is trained to be simultaneously supervised by both property prediction loss and gen- eration loss. However, as shown Table 4 our experiments indicate that this approach (third row) significantly degrades the performance of both the generation and property prediction tasks. Even if we freeze the weights of the generation model and add a separate head for the property prediction task to maintain the generation performance (fourth row), we observe that the property prediction performance does not improve compared to training from scratch.
We attribute the suboptimal results to the inherent inconsistency between the generative and pre- dictive tasks. In the generative task, the model operates throughout the entire diffusion process, transitioning molecules from disordered state structured form. However, for predictive tasks, meaningful molecular properties can only be defined once the molecular structure has been fully established. Disordered molecules earlier stages lack meaningful properties and can introduce er- rors. Consequently, merely adopting a simple combination approach results in an incorrect mapping between perturbed molecules and their properties, negatively impacting both molecule generation and property prediction. To validate the conjecture, we conduct a theoretical analysis of the mutual information between intermediate representations within the denoising network and the original data during diffusion training. According to the InfoMax principle (Linsker) 1988; Oord et al.) 2018), a representation better if the mutual information between data and the representation is larger. We theoretically demonstrate that the diffusion model implicitly maximizes the lower bound of this mutual information, suggesting the potential to unify generative and predictive tasks. However, the mutual information decreases monotonically throughout the diffusion process, approaching larger time steps, indicating that representations at higher noise levels (i.e., disordered molecules) ineffective. Thus, both intuition and theory suggest that generative and predictive tasks align only at smaller time steps when molecules remain more structured.
Based on these theoretical and empirical findings, we propose novel two-phase generative ap- proach to unify property prediction and generation, as shown in Figure I We divide the molecu- lar generation process into two phases: the 'molecule nucleation phase" and the 'molecule growth phase'. This division is inspired by the crystal formation process in physics (De Yoreo & Vekilov) 2003), allowing us to systematically capture the complexities of molecular assembly and enhance the efficiency of generation. In the molecule nucleation phase, the molecule forms its scaffold from a completely unstructured state, after which the complete molecule developed based on this scaf- fold. These phases are separated by 'nucleation time'. We then introduce a new diffusion model to formulate these two phases. Before the nucleation time, the diffusion model progressively gener- ates the molecular structure. After nucleation, it continues refining the structure while incorporating prediction losses into the diffusion process, optimizing both generation and prediction tasks. Unlike traditional generative models that typically perform joint diffusion of atom types and coordinates,
2

---

Published as a conference paper at ICLR 2025
approach focuses solety the diffusion of coordinates, treating atom type separate predic- tion task. This separation is justified by the observation that atom types can often be inferred based on the determined scaffold (Bruno et al.) [011] Vaitkus et al. 2023). Specifically, prior to nucleation, the diffusion process aims to reconstruct the coordinate. Following nucleation, we incorporate pre- dictive losses for both atom type and properties into the unified learning framework.
Our experimental results demonstrate that UniGEM achieves superior performance in both the pre- dictive and generative tasks. Specifically, in molecule generation, UniGEM produces significantly more stable molecules, with molecular stability improving by approximately 10% over the diffusion- based molecular generation model EDM (Hoogeboom et al. 2022). In terms of property prediction, our method significantly outperforms the baseline trained from scratch, even achieving accuracy comparable to some pre-training methods without introducing any additional pre-training steps and large-scale unlabeled molecular data.
To explain the improved generation performance of our unified model, we derive upper bounds on the molecular generation error for both UniGEM and traditional diffusion models. Our analysis reveals that UniGEM significantly reduces the dimensionality of diffusion data and simplifies the treatment of the discrete atom types by framing it as a prediction task, leading to lower total molec- ular generation errors and producing more stable and accurate molecular structures.
To the best of our knowledge, are the first to propose effective diffusion- based unified model that significantly enhances of the performance of both molecular generation and property prediction tasks, demonstrating novel approach in this domain. Our treatment of molecular generation as comprising two phases, molecular nucleation and molecular growth, provides a novel paradigm that may inspire the development of more advanced molecule generation models. Furthermore, the analysis and design of how to unify generative and predictive models may also impact other AI applications, like natural language processing and computer vision.
METHOD
2.1 UNIGEM FRAMEWORK
UniGEM is a two-phase diffusion based generative approach for property prediction and molecular generation. For the convenience of comparing with the traditional joint joint diffusion approach, we adopt notation scheme consistent with the E(3) Equivariant Diffusion Model (EDM) (Hoogeboom al. 2022). We denote a 3D molecule with M atoms as z (x. h)) where = (#1.'.. .MM E 3add represents the atomic positions, and h = (hi,... has) E {eo. ... e ] encodes the atom type information. Each h; is a one-hot vector of dimension H, corresponding to the type of the atom i, where is the total number of distinct atom types in the dataset.
2.1.1 TWO-PHASE GENERATIVE PROCESS
In contrast to traditional diffusion based model, UniGEM adopts diffusion process to generate only the atomic coordinates æ rather than the whole molecule z. We define the forward process for adding noise to the molecular coordinates as follows:
fm
g( = g( II (1)
II t +1
g(æ,|<x--1.
where i, denotes the nucleation time to indicate when the molecule forms its scaffold. We define the growth phase E [0, fm]] and the nucleation phase E [4n.7]. During the growth phase, we will incorporate atom type and property information, as these attributes are clearly defined at this stage. The conditional distribution defined as q (z.) N2 (z, aa I), t 1,... ,T, where a
and represent the noise schedule following Hoogeboom et al. (2022) and satisfy the condition of o = 1, with or decreases monotonically from 1 to O. 20 (#o, ha) represent an unperturbed molecule in the dataset. N₂ represents the Gaussian distribution in the zero center-of-mass (CoM)
subspace satisfying Σ . æ; 0, to ensure translation equivariance.
In UniGEM, we train a unified model, denoted as do. that is capable of both generative and pre- dictive tasks. An overview is provided in Figure T Following previous molecular generative mod-
3

---

Published as a conference paper at ICLR 2025
els (Hoogeboom al. 2022 Wu et al. 2022 Xu et al. 2023), adopt EGNN architecture Satorras et al. 2021) for our implementation.
For the coordinate generation task, train the model to predict noise from the noisy inputs using mean squared error (MSE) loss at each timestep:
(r))
(2) (zz,
E, = (x; a,20)/or. ( t [1.7]
(2)
where E is the standard Gaussian noise injected to zo and o6 (x), () is the equivarant out- put of the backbone network. The denoising loss is equivalent to the score estimation loss (-o,Vas Ve, log q(x)))1 where Vx, log (a:)) refers to the score func- tion. This is derived through the equivalence between score matching and conditional score match- ing (Vincent (011b), with a formal proof presented in Ni et al. (2024); Feng et al. (2023a). The learned score function is crucial for reversing the noise addition process for coordinate generation.
E By(rs,,,,, .r., ) (*) (xx,t) (x,,
For the atom type prediction task, use an Lı loss for regressing the one-hot vector of the atom types:
(1) (xt, t) - hl, t e (1) c (h) = Ba =, |
where h is the one-hot vector of atom types corresponding to the coordinates æp in the dataset, and 2b (h) is the invariant output from the backbone network with an additional two-layer MLP prediction head.
In the property prediction task, we denote the property of the molecule c, and utilize an I,2 loss for regressing the property: ( 26 (c) ( c| E [1,4n),
(c)
(e) (ax, t) is the output from the backbone network with another two-layer MLP prediction where 2
head. For t E [t,, TJ, we set c( and c(Ah) as zero.
In this way, avoid inconsistencies of the two tasks when the molecule structure is not fully formed. However, employing traditional diffusion strategies train this unified model would fail. To overcome this challenge, we have developed a series of training strategies aimed at optimizing the balance between the two phases and maximizing the benefits of integration.
seamlessly integrate predictive and generative tasks using single model, and
2.1.2 TRAINING STRATEGIES
In our framework, the growth phase occupies small portion of the overall training process, with the best rest-performing configuration being tn = 0.01. we follow the standard diffusion training procedure and sample time steps uniformly. the number of iterations for the predictive task only takes 1% of the total training process which will significantly degrade the model's performance on this task. Therefore, to ensure sufficient training for the predictive task, oversample the time steps in the growth phase. Thus, the final loss can be written as the following equation:
(c!*))
c(()) + C(O))
E -E-30000 (00...
(5)
However, observed that oversampling can lead to imbalanced training across the full range of time steps, which adversely affects the quality of the generative process. To address this issue, we propose multi-branch network architecture. The network shares parameters in the shallow layers but splits into two branches the deeper layers, with separate sets of parameters for each branch. These branches are activated in different phases of training: one branch is dedicated the nucleation phase, while the other handles the growth phase, as shown in Figure
3 This design ensures that the predictive and generative tasks are trained effectively without negatively influencing each other.
2.3.3 INFERENCE OF UNIGEM
In UniGEM, molecular generation is achieved by reversing the forward diffusion process through a Markov chain to reconstruct atomic coordinates, followed by predicting atom types based the generated coordinates. Specifically, express the reverse process pa (For))
4

---

Published as a conference paper at ICLR 2025
po(xr)) IT= Po(2s-1 where the noise at each step is approximated by the trained neural net-
work 68 (ax,, t). The reverse transition at each time step is given by:
T d/--1 ( (ax, (), ( )21
po (()))) (10x)=1a NNo
3t-1 D2
(6)
0t|[-1
01
Tt
Here 0:/-11 = 01 a1/00- and of oic-- of a 1 The prior distribution is approximated by
p(æN) = (xx
For the molecular generation task, the atom types predicted the final time step, producing complete molecule z (æoo Φ (zx. 0)), as shown in Figure For property prediction, the inference is done with the time step fixed at zero during testing, i.e., using the output (x0, 0). Notably, this approach incurs no additional computational overhead for either the generative or predictive tasks. The total inference time is the same the baseline, ensuring that both tasks performed efficiently without increasing the overall runtime.
(æo,
2.2 THEORETICAL ANALYSIS
In this section, we present our theoretical results addressing two questions: first, why diffusion is capable of learning representations while the direct multi-task approach fails; and second, why the proposed UniGEM effectively enhances performance in the challenging generation task.
2.2.1 INCONSISTENCY BETWEEN MOLECULE GENERATION AND PROPERTY PREDICTION
Our study an Information Maximization (InfoMax) perspective (Linsker 1988 Oord et al., 2018), which serves as a criterion of the representation quality. According to the InfoMax theory, effective latent representations Çe in the diffusion-
molecule generation approach are achieved by maximizing the mutual information (MI) between the original molecular coordinates vo and Sr- Se are derived from the noisy coordinates æ; in the intermediate layers of the denoising network.
Theorem 2.1. The mutual information between To and St can be expressed follows, with subsequent lower bound:
I(Gr, Ro) = I(Roi x..) - (g( (.) [KLL
2 [(xo: æx)) - (x) (g [ppp))))p([[[[))))
where q( are data distribution defined by the forward process of diffusiio, q(Sa =()alar))80u(e) go and p(Zol(t) represent the estimated representation and denoising distributions by the denoising network. In practice, our denoising network models (xx)) and the mean of the denoising distribu- tion (xolgog (x:)) xx_0()0(X(()))
ax
As we show in appendix when p p(aolCt) and g( follow Gaussian distributions with the same variance minimizing the KL divergence is equivalent to minimizing the denoising diffu- sion loss. Thus, the second term (. [DxL (g(xo xx)) (ee)lllpppoo//)))) can be minimized during the diffusion training process. This suggests that effective representations be implicitly learned throughout the diffusion training.
However, according to the data processing inequality (Thmas & Cover 2006). the first term I [wo: x,) decreases monotonically as increases, eventually approaching zero. As arssult I(Ge: zn) becomes small for large values of i, indicating that it is increasingly challenging learn meaningful representations from noisy molecules. This explains the failure multi-task learning in such cases. In contrast, when is small, [C,, æo) larger, suggesting that the two tasks can align to facilitate learning robust representations, which emphasizes the need for a two-phase modeling approach.
2.2.2 THEORETICAL ANALYSIS ON GENERATION ERROR
To further investigate why UniGEM improves performance on the generation task, we conduct a theoretical analysis of the generation error. Specifically, derive the upper bound of the generative emrr for two molecular generation approaches: UniGEM and traditional diffusion- -aased models, based results from Chen al. (2023b). Complementary details provided in appendix G
5

---

Published as a conference paper at ICLR 2025
Theorem 2.2 (Generative Error Analysis). With mild assumptions on the data distribution q pro- vided in appendix G.I the total variation error between the UniGEM generated data distribution Po(zo) and ground-trice truth data distribution ((zo) is bounded by the following terms.
T
TV(po(zo).
KL(g(aa)) po(ar))e
(Le L.M.IVV+++++
Bigiao (() (mr) (8)
prior prteriitribttion crnor
discreniation CTRNT
atom type extmmstion crear
conrdinatt vour extimation rro
For Iraditional diffusion models generating coordinates and atom types simaltaneously with distri- bution re (z)) the error bound becomes:
TV(ra(zo). S /Klig(zo]lra(zr))
(L.vdd L.m.IVV
(9)
distrilation
discrctioatinn
score tipe estionaiin
estimatioo ervor
Here for y E {a, 2), my is the second moment of q(yo): Ly is the Lipschitz constant of In ((3)),
Vt 2 0. T is the terminal timestep; T is the discretization steps. We assumed the step size I : T/T 2 1y and Ly 11
Our theoretical analysis identifies four factors contributing to generation enrors in both UniGEM and traditional diffusion-basee models. The prior distribution and discretization errors both decrease with lower data dimensionality. Since UniGEM only generates atom coordinates, unlike traditional models that jointly generate both coordinates and atom types, UniGEM deals with significantly smaller data dimensionality, leading to reduced errors. As for the latter two errors, it's harder to compare through theory alone, so supplement with experimental results QM9 datasets in Table I showing both of them are smaller in UniGEM.
The reduced coordinate and atom type errors in UniGEM be intuitively explained. A critical issue with traditional biffusion-based models is their treatment of discrete atom types as continuous variables. which can lead instability during generation. For example, in our experiments, we observe that even at the later stage of generation, the predicted atom type (determined by the highest probability) oscillates between two categories, indicating suboptimal learned distribution. Since the coordinates in the next step depend on both the current coordinate and atom type, this oscillation introduces additional errors in coordinate estimation. In contrast, UniGEM avoids this issue, as it only generates the coordinates, potentially leading to more accurate and stable coordinate predictions.
Table 1: Comparisons of coordinate and atom type errors.
UniGEM
EDM
Atom type Coordinate
1.3E-5(0.000))
0,0077(0.00 (01)))))
0.2247(0.0408) 0.2430(0.0434)
For the atom type error, UniGEM approaches this as a prediction task, which is significantly simpler than the generation approach used in traditional methods, especially after the nucleation phase when molecule scaffolds have formed. Therefore, it is not surprising that the atom type error in UniGEM smaller.
Regarding the above points, we can conclude that it is both reasonable and necessary to decouple the discrete atom types from the joint generation process and handle them through separate prediction task.
3 EXPERIMENTS
3.1 MAIN EXPERIMENTS
Datasets our experiments, we utilize two widely used datasets. The QM9 dataset (Ruddigkeil al. 2012 Ramakrishnan et al., 2014) serves both molecular generation and property prediction tasks, comprising approximately 134,000 small organic molecules with up to nine heavy atoms, each annotated with quantum chemical properties such LUMO (Lowest Unoccupied Molecular Orbital), HOMO (Highest Occupied Molecular Orbital), HOMO-LUMO gap, polarizability (a)). These properties are used ground-truth labels for property prediction. In contrast, the GEOM-

---

Published as a conference paper at ICLR 2025
Table 2: Comparison of generation performance on metrics, including atom stability, molecule sta- bility, validity, and validity uniqueness. Higher values indicate better performance.
QM9
GEOM GEOM-Drugs
#Metrics
Atom sta
Mol sta %)
Valid(%)
V*U(%)
Atom sta( Valid(%)
E-NF
85.0
4.9
40.2
39.4
G-Schnet
95.7
68.1
85.5
80.3
EDM
98.7
82.0
91.9
90.7
81.3
92.6
GDM
97.6
71.6
90.4
89.5
77.7
91.8
EDM-Bridge
98.8
84,6
92.0
90.7
82,4
92.8
GeoLDM
98.9
89.4
93.8
92.7
84.4
99.3
UniGEM 99.0 -0.3%
89.8 +7.8%
95.0
+3.1%
93.2
+2.5%
85.1 43.8%
48.4445.8%
Drugs dataset (Axelrod & GomezzBombarellii 2022) focuses on drug-like molecules, providing large-scale collection of molecular conformers. It includes around 430,000 molecules, with sizes ranging up to 181 atoms and an average of 44.4 atoms per molecule.
The splitting strategies for both benchmarks follow the previous practices. For the QM9 dataset, we adopt the same split as in prior methods (Hoogeboom et al., 2022 Satorras et al.) 2021 Anderson ct al. 2019), dividing the data into training (100K), validation (18K), and test (13K) sets Simi- larly, for the GFOM-Drugs dataset, we follow the approach outlined in Hoogeboom et al. (0022) to randomly split the dataset into training, validation, and test in 8:1:1 ratio.
Implementation Details We adopt EGNN (Satorras et al. 2021 1) as our backbone and modify it into multi-branch network. As illustrated in Figure 3 different branches handle the diffusion loss different ranges of time steps, with shared layers preceding the branches. In the branch for S ln we add predictive losses, including molecular property prediction and atom type prediction. During training. we apply uneven sampling for each batch. For the first half of the batch, the sampling range for is [t,, I'. and for the second half, the sampling range for t [0, t,), ensuring that each branch receives the same samples and is sufficiently trained. All losses are optimized in parallel, with each loss weighted equally 1. We implement UniGEM with layers, each consisting of 256 dimensions for the hidden layers. The shared layer count is set to 1, while each separate branch, corresponding to different ranges of time steps, has layers. Optimization is performed using the Adam optimizer, with batch size of 64 and learning rate of 10-4. We train UniGEM on QM9 for 2,000 epochs, while training on the GEOM-Drugs dataset is limited to 13 epochs due to the larger scale of its training data. We set the nucleation time ln to 10 and the total time steps T to 1000, ensuring that the prediction atom types and property occurs within the first 10 time steps. The atom type and molecular property predicted in the last time step used as the final predictive result.
3.1.1 EXPERIMENTAL RESULTS MOLECULE GENERATION
Baselines UniGEM is implemented based the codehase of classic 3D molecular diffusion al- gorithm, EDM (Hoogeboom et al. 2022). Thus, our baseline comprises EDM and its variants, GDM (Hoogeboom et al.J 20022), EDM-Bridge (Wu al.) (002)) and GeoLDM (Xu et al. (023). GDM utilizes non-equivariant networks for training, EDM-Bridge enhances EDM through the tech- nique of diffusion bridges. GeoLDM introduces an additional autoencoder encode molecular structures into latent embeddings, where the diffusion process is conducted in the latent space. Ad- ditionally, we include G-Schnet (Gebauer et al. 2019). an autoregressive generation method, and Equivariant Normalizing Flows (E-NF) (Garcia Satorras et al. 2021) in our comparisons.
Metrics Following the approach of these baselines, we sample 10,000 molecules and evaluate atom stability, molecule stability, validity, and uniqueness of valid samples. Specifically, we utilize the single, double, triple,
distances between each pair atoms to predict the bond type-whether or non-existent. Atom stability is calculated as the ratio of atoms exhibiting correct valency, while molecule stability reflects the fraction of generated molecules in which each atom maintains stability. Validity and uniqueness assessed using RDKit by converting the 3D molecular structures into SMILES format, with uniqueness determined by calculating the ratio of unique generated molecules among all valid samples after removing duplicates.

---

Published as a conference paper at ICLR 2025
Table 3: Performance (MAE, 4) on 6 property prediction tasks in QM9. The best results are in bold.
Task
a
Ae (meV)
éHOMO (meV)
FLUMO (me
(Units)
(bohr")
(D)
cat
muLK K
EGNN
0.071
48
29
25
0.029
0.031
GraphMVP
0.070
46.9
28.5
26.3
0,031
0,033
3D Infomax
0.075
48.8
29.8
25.7
0.034
0.033
GEM
0.081
52.1
33.8
27.7
0,034
0,035
3D-EMGP
0.057
37.1
21.3
18.2
0.020
0.026
UniGEM
0.060-15.5%
34.5
28.1%
20.9
-27.9%
-33.%%
0.019 -34.5%
0.023 -25.8%
Results The results in Table [ show that UniGEM consistently outperforms all baselines across nearly all evaluation metrics for both QM9 and GEOM-Drugs, with the gray cells indicating the molecule generator used in our approach. Notably, compared EDM variants, UniGEM is sig- nificantly simpler, as it neither relies on prior knowledge nor requires additional training for an autoencoder, yet it achieves superior performance compared to EDM DDMdBridde and GeoLDM, high- lighting UniGEM's superiority, A visualization of the generation process is presented in Figure [ To demonstrate the flexibility of UniGEM in adapting
various generation algorithms, we apply UniGEM to Bayesian Flow Networks (BFN) Graves et al. (2023); Song et al. (2024a) and achieve SOTA results on the QM9 dataset, as detailed in Appendix B.1
Furthermore, we test UniGEM on conditional generation task in Appendix B.2, by using pre-trained property prediction module as guidance during sampling. This approach eliminates the need for retraining conditional generation model and avoids additional guidance networks.
3.1.2 EXPERIMENTAL RESULTS FOR MOLECULAR PROPERTY PREDICTION
Baselines The property prediction module of UniGEM is built on the widely used equivariant neural network EGNN (Satorras et al. 202I). We compare UniGEM against the EGNN model trained from scratch, referred to as EGNN without confusion, as well several recent pre-training methods, that have demonstrated superiority in learning molecular representations, thereby improving the property prediction performance. These methods include GraphMVP (Liu al. 2021). 3D Infomax (Stärk) ct al. 2022), GEM (Fang et al. 2022). and 3D-EMGP (Jiao et al. (023). is important to note that all baseline methods utilize the
backbone and data splits to
a fair comparison.
Results We evaluated UniGEM on six properties of the QM9 dataset, using the Mean Absolute Error (MAE) on the test set the evaluation metric. As shown in Table 3 UniGEM significantly outperforms EGNN trained from scratch, with the results highlighted in the grey cell, demonstrating the effectiveness of unified modeling. Surprisingly, UniGEM also surpasses most advanced pre- training methods (italicized), even though they utilize additional large-scale pre-training datasets. This highlights the strength of its unified approach, which effectively leverages the capabilities of molecular representation learning during the generation process, without requiring additional data and pre-traine pre-training steps. The relationship with denoising pre-training is discussed in Appendix B.3
3.1.3 COMPARISONS WITH GENERAL UNIFIED APPROACHESS
In this section, we compare UniGEM with two general approaches for unifying generation and property prediction tasks. Table g outlines these approaches: the first treats generation and prop- erty prediction as a multi-task joint training network, referred to 'Multi-task'. Additionally, we incorporate the property prediction task into a pre-trained generation network for continued train- ing, specifically freezing the backbone to preserve the network's generation capability, referred to as 'Gen Pre-train'.
Results presented in Table 4 demonstrate that both simple approaches fail to achieve better results than the baseline. The performance of the multi-task approach aligns with our previous analysis: there is an inconsistency between the generation process and the property prediction task, and sim- ple integration negatively impacts the performance of both tasks. While 'Gen Pre-train' preserves the model's 's generation capabilities by freezing the backbone, it also restricts the network's 's expres- siveness, which in turn limits its performance in property prediction.
8

---

Published as a conference paper at ICLR 2025
Table 4: Comparison of UniGEM with other general unified approaches on both generation and property prediction tasks. The best results are in bold.
Generation
Prop Pred
#Metrics EDM
Atom sta(%)
Mol sta(%)
Valid(%)
V*U(%)
ÉLUMO
98.7
82.0
91.9
90.7
EGNN
25
Multi-task
98.0
76.0
89.9
88.8
51.8
Gen Pre-train
98.7
82.0
91.9
90.7
51.2
UniGEM
99.0
89.8
95.0
93.2
16.7
Table 5: Ablation study about the impact of different predictive loss functions on the generation and prediction performance.
Generation
Prop Pred
#Metrics
Atom sta(%)
Mol sta(%) (%)
Valid(%)
V*U(%)
GLUMO
EDM
98.7
82.0
91.9
90.7
EGNN
25
MPP
98.7
85.7
94.2
92.6
20.6
ATP
98.7
89.3
94.6
92.8
MPP+ATP
99.0
89.8
95.0
93.2
16.7
3.2 ABLATION STUDIES
We conduct three ablation studies to assess the impact of predictive loss, time sampling strategy, and nucleation time setting using the QM9 dataset. These studies include validation on both gener- ation and property prediction tasks, comparing results to the generation baseline EDM and the property prediction baseline EGNN trained from scratch.
3.2.1 IMPACT OF PREDICTIVE Loss ON PERFORMANCE
In UniGEM, we incorporate two types of predictive loss: atom type prediction loss and molecular property prediction loss. To validate their impact, we compare our model against two baselines: one utillizigg only atom type prediction loss (referred to as ATP) and the other using only molecular property prediction
(referred to as MPP). The comparison results are presented Table with LUMO selected as the target property for performance evaluation.
Results are shown in Table both ATP and MPP enhance generation performance compared to pre- vious baseline models, with ATP demonstrating greater improvement than MPP. The combination of these two losses functions yields the best results, as evidenced by UniGEM. These findings un- derscore the importance of both predictive losses, particularly highlighting the critical role of atom type prediction.
In molecular property prediction, MPP enhances the original performance, and incorporating the atom type prediction loss further improves results. This improvement may stem from atom type pre- diction facilitating the learning of better molecular representation, which ultimately benefits property prediction.
3.2.2 ANALYSIS OF TRAINING STRATEGIES IN UNIGEM
To ensure sufficient training steps for property prediction, oversample the time steps after the nucleation time. Additionally, we design multi-branch network to mitigate the negative impact of oversampling generation performance. We conduct ablation studies to demonstrate the necessity of oversampling and branch splitting strategies in ensuring superior performance for both generation and prediction tasks.
Results are shown in Table where the first two rows represent the baselines, and the third row indi- cates that single-brane branch with normal time sampling(referred to 'Normal Sampling (S)') achieves higher generation performance than the baseline but exhibits degraded prediction performance. This suggests that the proportion of structured time is too small for sufficient property prediction training.

---

Published as a conference paper at ICLR 2025
Table Performance of different training strategies on the generation and prediction tasks, with the best results highlighted in bold.
Generation
Prop Pred
#Metrics
Atom sta(%)
Mol sta %)
(a
V*U(%)
fLLM
EDM
98.7
82.0
91.9
90.7
EGNN
25
Normal Sampling(S)
98.9
89.7
94.8
92.9
40.5
Oversampling(S)
98.3
80.1
90.1
88.8
19.2
UniGEM
99.0
89.8
95.0
93.2
16.7
In the fourth row, adding oversampling to the molecular growth phase(referred to as 'Oversampling (S)") improves property performance compared to the baseline, but lowers generation performance due to the effects of unbalanced time step sampling. Finally, in UniGEM (last row), we change the backbone to multi-branch architecture, allowing different branches to handle different time steps. This adjustment eliminates the negative impact of time step oversampling, resulting in superior per- formance on both tasks.
3.2.3 EFFECT OF NUCLEATION TIME
Nucleation time is defined as the moment when the molecular scaffold forms, after which the molec- ular coordinates experience only minor adjustments. In practical applications, accurately determin- ing the nucleation time challenging. Therefore, we conduct an empirical analysis assess its impact on model performance. We set the total training step to T=1000, and compare models with nucleation times of 1, 10, and 100.
Results are outlined Figure Across generation and property prediction criteria, nucleation time of 10 achieves optimal performance. Setting the nucleation time too large may incorporate time steps prior to the complete formation of the molecular scaffold, leading to suboptimal outcomes. Conversely, using a nucleation time that is too small also degrades performance, potentially due insufficient input noise, which is crucial for learning a more robust type and property mapping. From the perspective of force learning interpretation of the denoising task, discussed in Appendix [:4 a moderate level of noise is advantageous for capturing molecular
24
stability
2o mevia
90
88
88
87
86
20
40
80
100
Time step
Figure 2:
Comparison of performance across different nucleation time selections for both the generation and property pre- diction tasks.
representations, thus enhancing the predictive tasks.
4 CONCLUSION
This paper introduces UniGEM, the first effective unified model that significantly improves the per- formance of both molecular generation and property prediction tasks. The underlying philosophy is that these traditionally separate tasks are highly correlated due to their reliance on effective molec- ular representations. The traditional inconsistencies of these tasks are overcome by two-phase generative process with atom type and property prediction losses activated in the growth phase. UniGEM's enhanced performance supported by solid theoretical analysis and comprehensive ex- perimental studies. We believe that the innovative two-phase generation process and its correspond- ing models offer a new paradigm that may inspire the development of more advanced molecule generation frameworks and benefit more specific applications of molecular generation.
ACKNOWLEDGEMENTS
This work is supported by Beijing Academy of Artificial Intelligence (BAAI).
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Brandon Anderson, Truong Son Hy, and Risi Kondor. Cormorant: Covariant molecular neural networks. Advances in neural information processing systems, 32, 2019.
Marloes Arts, Victor Garcia Satorras, Chin-Wei Huang, Daniel Zügner, Marco Federici, Cecilia Clementi, Frank Noé, Robert Pinsler, and Rianne den Berg. Two for one: Diffusion
models and force fields for coarse-grained molecular dynamics. Journal of Chemical Theory
and Computation, 19(18)6111-6155) 2023, doi: 1..02//acs.cctc.3000//00
URL https:
111.11011111/000 jotc. 3c00702
PMID: 37688551.
Simon Axelrod and Rafael Gomez-Bombarely Geom, energy-annotated molecular conformations for property prediction and molecular generation. Scientific Data, 9(1): 2022.
Fan Bao, Min Zhao, Zhongkai Hao, Peiyao Li, Chongxuan Li, and Jun Zhu. Equivariant energy- guided sde for inverse molecular design. In The eleventh international conference on learning representations, 2022.
Jan J Bruno, Gregory P Shields, and Robin Taylor. Deducing chemical structure from crystallograph- ically determined atomic coordinates. Acta Crystallographica Section B: Structural Science, 67 (4):333-349, 2011.
Andrew Campbell, Jason Yim, Regina Barzilay, Tom Rainforth, and Tommi Jaakkola. Gen- erative flows discrete state-spaces: Enabling multimodal flows with applications pro- tein co-design. In Forty-first International Conference on Machine Learning, 2024. URL
https://opeerevieet net forum?id==<QWSSr00B
Changyou Chen, Chunyuan Li, Liqun Chen, Wenlin Wang, Yunchen Pu, and Lawrence Carin Duke. Continuous-time flows for efficient inference and density estimation. In International Conference on Machine Learning, pp. 923-833. PMLR, 2018.
Jiujiu Chen, Xuecang Zhang, Zhi-Ming Ma, Shengchao Liu, et al. Molecule joint auto-encoding: Trajectory pretraining with 2d and 3d diffusion. Advances in Neural Information Processing Systems, 36:55077-55096, 2023a.
Sitan Chen, Sinho heerr,,Jerry Li, Yuanzhi Li, Adil Salim, and Anru Zhang. Sampling is as easy as learning the score: theory for diffuusion models with minimal data assumptions. In The Eleventh International Conference on Learning Representations, 2023b. URL pttttps://openrevves
net//orum?? 2d
Xinlei Chen, Zhuang Liu, Saining Xie, and Kaiming He. Deconstructing denoising diffusion models for self-supervised learning. arXiv preprint arXiv:2401. 2024.
James J De Yoreo and Peter G Vekilov. Principles of crystal nucleation and growth. Reviews in mineralogy and geochemistry, 54(1):57-93, 2003.
Bradley Efron. Tweedie's formula and selection bias. Journal of the American Statistical Associa- tion, 106(496):602-16111 2011.
Xiaomin Fang, Lihang Liu, Jieqiong Lei, Donglong He, Shanzhuo Zhang, Jingbo Zhou, Fan Wang, Hua Wu, and Haifeng Wang. Gehmetry-enhanced molecular representation learning for property 2022.
prediction. Nature Machine Intelligence,
Shikun Feng, Yuyan Ni, Yanyan Lan, Zhi-Ming Ma, and Wei-Ying Ma. Fractional denoising for 3d molecular pre-training. In International Conference on Muchine Learning, pp. 9938-9961. PMLR, 2023a.
Shikun Feng, Lixin Yang. Weiying Ma, and Yanyan Lan. Unimap: Universal smiles-graph repre- sentation learning. arXiv preprint arXiv:2310. 14216, 2023b.
Shikun Feng, Yuyan Ni, Minghao Li, Yanwen Huang, Zhi-Ming Ma. Wei-Ying Ma, and Yanyan Lan. Unicor: A unified contrastive learning approach for milti-view molecular representation learning. arXiv preprint arXiv:2405. 1034433... 2024.

---

Published as a conference paper at ICLR 2025
Bowen Gao, Minsi Ren, Yuyan Ni. Yanwen Huang, Bo Qiang, Zhi-Ming Ma, Wei-Ying Ma, and Yanyan Lan. Rethinking specificity in SBDD: Leveraging delta score and energy-guicee -rrgy-guided dif- fusion. In Forty-first International Conference Machine Learning, 2024. URL https: //openreview. ew. net. dd=
d= NASSSWNN11XX999 2XA9r
Victor Garcia Satorras, Emiel Hoogeboom, Fabian Fuchs, Ingmar Posner, and Max Welling, E (n) equivariant normalizing flows. Advances in Neural Information Processing Systems, 34:4181- 4192, 2021.
Niklas Gebauer, Michael Gastegger, and Kristof Schütt. Symmetry-adapted generation of 3d point sets for the targeted discovery of molecules. Advances in neural information processing systems, 32, 2019.
Alex Graves, Rupesh Kumar Srivastava, Timothy Atkinson, and Faustino Gomez. Bayesian flow networks. arXiv preprint arXiv:2338.07003,, 2023.
Jiaqi Guan, Wesley Wei Qian, Xingang Peng. Yufeng Su, Jian Peng, and Jianzhu Ma. 3d equiv- ariant diffuusion for larget-aware molecule generation and affinity prediction. arXiv preprint arXiv:2303.03543, 2023.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840-688 1000068111200 2020.
Emiel Hoogeboom, Victor Garcia Satorras, Clément Vignac, and Max Welling. Equivariant diffu- sion for molecule generation in 3d. In International conference on machine learning, pp. 8867- 8887. PMLR, 2022.
Zhenyu Hou, Xiao Liu, Yukuo Cen, Yuxiao Dong, Hongxia Yang, Chunjie Wang, and Jie Tang. Graphmae: Self-supervised masked graph autoencoders. In Proceedings of the 28th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, pp. 594-604, 2022.
W Hu, Liu, Gomes, Zitnik, P Liang, V Pande, and Leskovec. Strategies for pre-training graph neural networks. In International Conference on Learning Representations (ICLR), 2020.
Drew A Hudson, Daniel Zoran, Mateusz Malinowski, Andrew K Lampinen, Andrew Jaegle, James I. McClelland, Loic Matthey, Felix Hill, and Alexander Lerchner. Soda: Bottleneck diffusion mod- els for representation learning. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 23115-23127, 2024.
Ross Irwin, Alessandro Tibo, Jon Paul Janet, and Simon Olsson. Efficient 3d molecular generation with flow matching and scale optimal transport. In ICML 2024 Al for Science Workshop, 2024. URL httt:ppss///ppeeeiiiiiie
net / forumor i d=CxA dkkqu
Rui Jiao, Jiaqi Han, Wenbing Huang, Yu Rong, and Yang Liu. Energy-motivated equivariant pre- training for 3d molecular graphs. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 37, pp. -104, 2023.
Tuan Le, Julian Cremer, Frank Noé, Djork-Aréé Clevert, and Kristof Schütt, Navigating the design space of equivariant diffusion- llffussio--basdd generative models for de novo 3d molecule generation. 2024.
Shuangli Li, Jingbo Zhou, Tong Xu, Dejing Dou, and Hui Xiong. Geomgel: Geometric graph contrastive learning for molecular property prediction. In Proceedings of the AAAI conference on artificial intelligence, volume 36, pp. 4541 1-449, 2022.
Haowei Lin, Shanda Li, Haotian Ye, Yiming Yang, Stefano Ermon, Yitao Liang, and Jianzhu Ma. TFG-flow: Training-free free guidance in multimodal generative flow. In The Thirteenth International Conference on Learning Representations, 2025. UKUbtttps://ppenreiiiee net nwt//orummttmm forum?
idd=K55i 7t
Ralph Linsker. Self-organization in perceptual network. Computer, 21(3):105-117, 1988.
Shengchao Liu, Hanchen Wang, Weiyang Liu, Joan Lasenby, Hongyu Guo, and Jian Tang. Pre- training molecular graph representation with 3D geometry. In International Conference on Learn ing Representations, 2021.
12

---

Published as a conference paper at ICLR 2025
Shengchao Liu, Weitao Du, Zhi-Ming Ma, Hongyu Guo, and Jian Tang. A group symmetric stochas- tic differential equation model for molecule multi-mode modal pretraining. In International Conference on Machine Learning. pp. 21497-21526. PMLR, 2023.
Calvin Luo. Understanding diffusion models: A unified perspective.
arXiv preprint
arXiv:2208. 11970, 2022.
Youzhi Luo and Shuiwang Ji. An autoregressive flow model for 3d molecular geometry generation from scratch. In International conference on learning representations (ICLR), 2022.
Sarthak Mittal, Korbinian Abstreiter, Stefan Bauer, Bernhard Schölkopf, and Arash Mehrjou. Dif- fusion based representation learning. In International Conference on Machine Learning, pp- 24963-24982. PMLR, 2023.
Yuyan Ni, Shikun Feng. Wei-Ying Ma, Zhi-Ming Ma, and Yanyan Lan. Sliced denoising: A physics- informed molecular pre-training method. arXiv preprint arXiv:2311.02//11 2023.
Yuyan Ni, Shikun Feng, Xin Hong, Yuancheng Sun, Wei- Wei-Ying Ma. Zhi-Ming Ma, Qiwei Ye, and Yanyan Lan. Pre-training with fractional denoising to enhance molecular property prediction. Nature Machine Intelligence. pp. 1-10, 2024.
Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predic- tive coding, arXiv preprint
2018.
Stefano Peluchetti. Non-denoising forward-ting -emerd-timeeaae diffusionss. arXiv preprint arXiv:2312. 2023.
Raghunathan Ramakrishnan, Pavlo O Dral, Matthias Rupp, and O Anatole Von Lilienfeld. Quantum chemistry structures and properties of 134 kilo molecules. Scientific data, 1(1):177, 2014.
Yu Rong. Yatao Bian, Tingyang Xu, Weiyang Xie, Ying Wei, Wenbing Huang, and Junzhou Huang. Self-supervised graph transformer on lcage-scale molecular data. Advances Neural Information Processing Systems, 33:12559- 2020.
Lars Ruddigkeit, Ruud Van Deursen, Lorenz C Blum, and Jean-Louis Reymond. Enumeration of 166 billion organic small molecules in the chemical universe database gdb-17. 17. Journal of chemical information and modeling, 52(1 11:2864-28777 2012.
Victor Garcia Satorras, Emiel Hoogeboom, and Max Welling. E (n) equivariant graph neural net- works. In International conference on machine learning, pp. 22222.33. PMLR, 2021.
Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan. and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamies. In International conference on machine learn- ing, pp. 2256-2265. PMLR, 2015.
Yuxuan Song, Jingjing Gong, Hao Zhou, Mingyue Zheng, Jingjing Liu, and Wei-Ying Ma. Uni- fied generative modeling of 3d molecules with bayesian flow networks. In The Twelfth Interna- tional Conference Learning Representations, 2024a. URL (tLh:tts:://e/ htttps::/ppenreviee..nt net/ Forum22d--UVtmmnzeR
Zixing Song, Ziqiao Meng, and Irwin King. A diffusion- pre-training framework for crystal property prediction. In Proceedings of the AAAI Conference Artificial Intelligence, volume 38, pp. 3999-90001 2024b.
Hannes Stärk, Dominique Beaini, Gabriele Corso, Prudencio Tossou, Christian Dallago, Stephan Günnemann, and Pietro Liò. 3d infomax improves gnns for molecular property prediction. In International Conference on Machine Learning, pp. 20479-20502, PMLR, 2022.
Joy A. Thomas and Thomas M. Cover. Elements of Information Theory. Iitey-Intersccence, 2nd edition, 2006. Chapter 2, Section 2.8, p. 32: Entropy, Relative Entropy, and Mutual Information: Data Processing Inequality.
Antanas Vaitkus, Andrius Merkys, Thomas Sander, Miguel Quirós, Paul A Thiessen, Evan E Bolton, and Saulius Gražulis. A workflow for deriving chemical entities from crystallographic data and its application to the crystallography open database. Journal of Cheminformatics, 15(1) 1:23, 2023.
13

---

Published as a conference paper at ICLR 2025
Clement Vignac, Nagham Osman, Laura Toni, and Pascal Frossard. Midi: Mixed graph and 3d denoising diffusion for molecule generation. In Joint European Conference on Machine Learning and Knowledge Discovery in Databases, pp. 560-576. Springer, 2023.
Pascal Vincent. A connection between score matching and denoising autoencoders, Neural Compu- tation, 201 la.
Pascal Vincent. A connection between score matching and denoising autoencoders. Neural compu- tation, 23(7):166111677, 201
Lemeng Wu, Chengyue Gong, Xingchao Liu, Mao Ye, and Qiang Liu. Diffusion-based molecule generation with informative prior bridges. Advances in Neural Information Processing Systems, 33336533-36543 2022.
Can Xu, Haosen Wang, Weigang Wang, Pengfei Zheng, and Hongyang Chen. Geometric-facilitated denoising diffusion model for 3d molecule generation. Proceedings of the AAAI Conference on Artificial Intelligence, 38(1):33--344 2024. doi: 10.1609/aaii.v38i v3871.27787. URL https:
//doi. org. 10. 1609/ aaai v3811. 277877
Minkai Xu, Alexander $ Powers, Ron O Dror, Stefano Ermon, and Jure Leskovec. Geometric latent diffusion models for 3d molecule generation. In International Conference Machine Learning, pp. 38592-38610. PMLR, 2023.
Yuning You, Ruida Zhou, Jiwoong Park, Haotian Xu, Chao Tian, Zhangyang Wang, and Yang Shen. Latent 3d graph diffusion. In The Twelfth International Conference on Learning Representations, 2023.
Sheheryar Zaidi, Michael Schaarschmidt, James Martens, Hyunjik Kim, Yee Whye Teh, Alvaro Sanchez-Gonzaler, Peter Battaglia, Razvan Pascanu, and Jonathan Godwin. Pre-training via de- noising for molecular property prediction. arXiv preprint 2,X1:22206.00333, 2022.
14

---

Published as a conference paper at ICLR 2025
Appendix
Table of Contents
Complimentary Figures
16
B Supplementary Experiments
16
B.1 Applying UniGEM to Advanced Generative Methods
16
Conditional Generation
17
B.3 Relationship with Denoising tre-training
C Hyperparameter Settings
18
D Related Work
19
D.I 3D Molecule Generation
19
D.2 3D Molecular Pre-training
20
Introduction of Joint Diftusion for Molecules
F
Representation Learning Analysis for Molecular Coordinate Generation
21
F.T
InfoMax Target)
21
F2
Variational Upper Bound on the KL Divergence Term
22
F33
Relation to Denoising Loss
22
F.4
The Benefits of Diffusion for Property Prediction: A Force Learning Perspective
23
Generative Error Analysis of UniGEM and Joint Diffusion Model
23
General Diffusion Error Hound
23
2
Error Bound for UniGEM and Joint Diffusion
24
G.3
Generative Eror Comparison
26
15

---

Published as a conference paper at ICLR 2025
A
COMPLIMENTARY FIGURES
data= h, c) 28
40
Coordinete
Branch #2
Denolaing
(o)1.) f)
Loss
(h)
Atom Type
Shared layers
p
Prediction Loss
Property
Prediction Loss
Branch #1
Figure 3: The training process of UniGEM.
Molecule Nuclestion
Molecule
tn
Layers
Proparty Prediction head
layers phoso
type
Figure 4: The molecular generative process of UniGEM
Molecular Gonaration Process of UniGEM
Acdiborab
100
200
300
400
500
600
700
300
900
950
Figure 5: A visualization of the generation process of UniGEM.
B SUPPLEMENTARY EXPERIMENTS
B. APPLYING UNIGEM TO ADVANCED GENERATIVE METHODS
As UniGEM is a flexible framework that can adapt to various generation algorithms, we expect to produce better results when adapted to more advanced generation algorithms. To illustrate this, we conducted additional experiments with UniGEM using the current SOTA generation algorithm, BFN Graves et al., 2023 Song et al. 2024a). Specifically, replaced EDM's coordinate gener- ation schedule with BFN's coordinate generation algorithm. The results, presented the table 7 highlight two key findings:
16

---

Published as a conference paper at ICLR 2025
Table 7: Molecular Generation Performance on Q.9.
Mol sta(%) Atom sta( %)
Valid(%)
V=U(%)
E-NF
4.9
85.0
40.2
39.4
G-Schnet
68.
95.7
85.5
80.3
GDM
71.6
97.6
90.4
89.5
EDM-Bridge
84.6
98.8
92.0
90.7
EquiFM
88.3
98.9
94.7
93.5
GFMDiff Xu et al. 2024)
87.7
98.9
96.3
95.1
GeoLDM
89.4
98.9
93.8
92.7
EDM
82.0
98.7
91.9
90.7
UniGEM (w/ HDM)
89.8 +7.8%
99.0 +0.3%
95.0
93.2 +2.5%
GeoBEN
90.9
99.1
95.3
93.0
UniGEM(w/ BFN)
93.7 +2.8%
99.3 +0.2%
97.3 +2.0%
93.0+0.0K% +0.0%%
Firstly, UniGEM can still achieve a significant improvement on top of a stronger generation algo- rithm, showeasing its robustness and adaptability across diverse frameworks.
Secondly, by leveraging diffusion for continuous coordinate generation and classification for discrete atom types, UniGEM effectively addresses the simultaneous generation of discrete and continuous variables. Its decoupled and straightforward approach outperforms GeoBFN (Song et al.) 2024a), which specifically tailors its generative algorithm for discrete atom types and discretized atom charge variables.
B.2 CONDITIONAL GENERATION
To assess the effectiveness of the proposed method, we evaluate it on a property-conditione gener- ation task using the QM9 dataset. This task provides more rigorous test of the model's ability to learn representations related to specific conditions while also capturing the underlying data distribu- tion-objectives that align closely with the goals the proposed method.
Unlike traditional methods, UniGEM integrates property prediction directly into its framework, en- abling the of its pre-trained property prediction module guide molecular coordinate updates. This approach bypasses the need to retrain a conditional generation model and eliminates the requirement for additional guidance network, both property prediction and generation share the same network.
(4) (az, t)
The update step for the molecular coordinates given by 21-1 ocjt-1 Dz
Giit-101 10 %
XVL (c)
oep 10-- where [c) is the property prediction loss of UniGEM.
Table 8: Performance of conditional generation on QM9 properties.
GLUMO (eV)
Gооо (eV)
Ac
a (bohr3)
C₁
(eV)
(D)
cal mot:
Conditional EDM(Average MAEI))
0.606
0.356
0.665
1.111
2.76
1.101
Guided UniGEM (Average MAE!)
0.592
0.233
0.511
0.805
2.22
0.873
Guided UniGEM (Mol sta(% 1))
81.5
86.3
80.4
83.2
83.0
80.9
Guided UniGEM (stable MAE.)
0.596
0.285
0.499
0.764
2.15
0.877
As shown in Table UniGEM achieves performance better than that of the conditional EDM trained explicitly with property conditions as input. This demonstrates the flexibility and capability of UniGEM in property-conditioned molecular generation.
To improve the evaluation of conditional property prediction tasks, we introduce two new metrics: Mol Stable and Stable MAE, addressing limitations in traditional evaluation methods. Existing
17

---

Published as a conference paper at ICLR 2025
approaches often rely an auxiliary property classifier (e.g., EGNN) to generate pseudo-labels labels for the generated molecules, comparing these labels with the target conditions to assess generation quality. However, these classifiers only utilize atom types and coordinates, meaning that even invalid molecules can yield seemingly accurate property predictions. To simultaneously evaluate how well the generated molecules satisfy the target conditions and ensure their chemical validity, propose the following metrics:
Mol Stable: The proportion of molecules with correct valence, serving as the most reliable indicator of unconditional molecular generation quality.
The mean absolute error (MAE) of properties, calculated only for stable molecules. By focusing chemically valid molecules, this metric mitigates the risk of models hacking classifiers (e.g.. generating invalid molecules with misleading pseudo- labels).
Stable MAE:
Average MAE: The conventional property MAE across all generated molecules, including unstable ones, which does not account for molecular stability.
The results in Table 8 demonstrate that UniGEM have the ability generate both mundtiinn-aligndd and chemically valid molecules.
B.3 RELATIONSHIP WITH DENOISING TRE-TRAINING
UniGEM leverages denoising diffusion generative models to learn molecular representations, draw- ing parallels to denoising pre-training approaches (Zaidi et al. 2022 Feng et al., 2023a] Ni et al.) 2023 To better understand UniGEM, explore its framework from the perspective of pre- training
Unlike conventional denoising pre-training methods, UniGEM introduces several distinctions. incorporates a broader range of noise scales during the molecular growth phase and integrates an atom prediction task. However, does not employ the specialized noise distributions proposed by Feng et al. (2023a); Ni et al. (2023) and utilizes a different model backbone. Moreover, UniGEM jointly trains supervised property prediction with "unsupervised" generation tasks, without relying on additional unsupervised datasets, diverging from the typical -en-train-andfinetun paradigm.
An interesting question arises when comparing these two "pre- "pre-training" approaches: does denoising with multiple noise scales, in conjunction with atom type prediction, lead to better representations? To investigate this, we compare UniGEM with Frad, a denoising framework specifically designed for denoising pre-training. For fair comparison, we aligned the network architecture by pre-training Frad with EGNN. The results, presented in Table 9 show that UniGEM outperforms Frad, despite the latter's focus on pre-training with specialized noise distributions. This outcome suggests that UniGEM may offer a promising approach to representation learning.
Table 9: Performance the QM9 property prediction tasks.
# pretrain data
ELUMO
CHоМ
Ae
Avg.
Uni-MOL (Uni-MOL)
19M
127.1
Frad (TorchMD-NET)
3.4M (PCQMv2)
13.7
15.3
27.8
18.9
Frad (EGNN)
3.4M (PCQMv2)
22.1
21.1
36.8
26.7
UniGEM (EGNN)
0.1M (QM9)
16.7
20.9
34.5
24.0
C HYPERPARAMETER SETTINGS
The hyperparameter settings are detailed in Table 10 Our model is based on EDM (Hoogeboom et al., 2022), with two additional tunable hyperparameters introduced by training algorithm: nucleation time and shared layers. Notably, these hyperparameters are guided by prior knowledge and do not require extensive tuning, highlighting the practical convenience of our approach.
18

---

Published as a conference paper at ICLR 2025
Table 10: Network and training hyperparameters.
Network Hyperparameters
Value
Embedding size
256
Layer number
for QM9, 4 for Seom-Drugs
Shared layers
Training Hyperparameters
Value
Batch size
64 for QM9, 32 for Geom-Drugs
Train epoch
2000 for QM9, 13 for Geom-Drugs
Learning rate
1.00 x 10
Optimizer
Adam 1000
T (sample times)
Nucleation time
10
Oversampling ratio Loss weight
0.5 for each branch
for each loss term
The nucleation time, ins has clear physical interpretation: it represents the moment when the molecular scaffold begins form during the denoising process. Thus the selection of in is not an arbitrary hyperparameter adjustment but is guided by the underlying physical principles. This parameter can be reasonably estimated from the noise schedule of the diffusion model and is largely independent of dataset size, model architecture, or optimization algorithms. Observing the noise process in EDM's coordinate denoising suggests that f,, E [0, 100] is an appropriate range where the molecular structure remains discernible.
Further, prior studies (Zaidi et al. 2022 Feng et al. 2023a) indicate that adding noise with a standard deviation of E [0.02, 0.05] preserves molecular semantics. Using EDM's noise schedule o(),, this corresponds to t / x 0.01 0.03, which translates to t, = 10 ~ 30. Our experiments further confirm that the choice of is robust variations in other parameters (Figure Therefore, the value t can be almost determined before training.
We conducted experiments with three different choices of shared layers (k- 1, and 5) and found minimal impact on property prediction, with k and k: = 3 outperforming i = 5 in generation.
For QM9, we trained for 2000 epochs with batch size of 64, which took approximately 7 days a single A100 GPU. For GEOM-drugs, we trained for 13 epochs with a batch size of 22x4=128, taking 6.5 days four A100 GPUs.
D RELATED WORK
D.1 3D MOLECULE GENERATION
The technological development of generative Al has driven numerous research efforts in 3D molec- ular generation. Luo & Jil (2022); Gebauer et al. (2019) use autoregressive methods to generate molecules, while Garcia Satorras et al. (2021) employs normalizing flows (Chen et al. 2018) for molecule generation. Recently, many works have applied diffusion models (Ho et al. 2020 Sohl- Dickstein et al., 2015). which have achieved success in image generation, to 3D molecular geo- metric generation. Hoogeboom et al. (2022) proposed an equivariant diffusion model (EDM) for 3D molecular generation. Building on this, EDM-Bridge (Wu et al.) 2022) proposes adding prior bridges (Peluchetti] (2023) to improve EDM performance. GeoL.M (Xu al. 2023) suggests using an autoencoder encode molecules into latent space, and then performing diffusion generation in the latent space. GcoBFN (Song et al. 2024a) introduces a new generative technique, Bayesian Flow Networks (Graves et al. 2023). 3D molecular generation, applying Bayesian inference to the parameter space of molecular distributions.
Generating atom types in diffusion- based molecular generation remains significant challenge due to their discrete nature. EDM represents atom types as one-hot vectors and generates them through diffusion model. In contrast, other approaches specifically design methods for discrete data gen-
19

---

Published as a conference paper at ICLR 2025
eration. The BFN algorithm introduces tailored generation strategies respectively for discrete and discretized data, inspiring GeoBFN to treat atom types as atomic numbers and employ the BFN algo- rithm for discreted data, which incorporates binning technique map continuous probabilities to discrete values. Additionally, Discrete Flow Models (Campbell et al., 2024) have been proposed for discrete data generation using Continuous Time Markov Chains and have been applied molecular generation in Lin et al. (2025). These methods generate atom types simultancously with atomic co- ordinates, inevitably requiring specialized algorithms for discrete data generation. UniGEM, on the other hand, circumvents this challenge by generating only atomic coordinates and subsequently pre- dicting atom types based on the generated structures. This approach effectively leverages molecular properties to bypass the difficulties associated with generating discrete data.
The methods following EDM adopt evaluation framework and dataset settings. However, re- cently, alternative approaches to 3DD molecular generation have emerged with different setups (Vi-) gnac al. 2023; Le al. 2024| Irwin et al., 2024). In terms of input, these methods incorporate bond information in addition to atomic types and coordinates for representing 3D molecules. Re- garding evaluation, they use bond inference strategy distinct from EDM's rules. EDM defines bonds strictly based on interatomic distances, implicitly enforcing constraints bond lengths and steric hindrance. In contrast, these newer methods rely on infeered bonds without such con- straints, significantly influencing stability and validity metrics /Vignac al., (023). Consequently, evaluating these methods requires additional energy-related metries and geometric distribution con- straints to enable a fair comparison with EDM-based methods.
D.2 3D MOLECULAR TRE-TNAINING
Molecular pre-training methods aim to learn molecular representations that are useful for various downstream tasks, especially property prediction. Common molecular pre-training methods include graph masking (Hu et al. 2020, Hou et al. 2022 Rong et al. 2020 Feng et al.) (023b). contrastive learning (Stäkk et al., 2022] Liu et al. 2021 2023; Li et al, 2022). and 3D denoising. Among them, 3D denoising, due to its interpretability in learning force fields, has been proven to be very effective pre-training strategy (Zidii et al.] [022] Feng et al.) 2023a [i et al.] [023] (024)) especially for quantum chemical properties (Feng et al. 2024). Recently, some pre-training works have integrated diffusion as a pre-training task (Song et al. 2024b) or incorporated diffusion into the pre-training process (Liu al. 2023) Chen et al.) 2023a), demonstrating that useful molecular
representations can be learned during the molecular generation process.
E INTRODUCTION OF JOINT DIFFUSION FOR MOLECULES
A 3D molecule is represented by coordinates and atom types, Formally, we denote molecule with M atoms as (x, h)). where (x1,.... xM) R3M represents the atomic positions, and h = (h1, ,,r) E {eo.... .... E AMd encodes the atom type information. Each h; is a one-hot vector of dimension H, corresponding to the type of the atom i, where H is the total number of distinct atom types in the dataset.
Recently, several works have applied diffusion models to 3D molecular data, employing joint diffusion to simultaneously generate molecular coordinates and atom types (Hoogeboom et al. 2022] Guan et al., 2023: Gao et al. 2024). The diffusion model independently injects noise to the coordinates and atom types respectively via a forward process:
al., 2015 Ho et al.) 2020)
((x =0) = Na. (wx (ax20, of -N (helaahho of I),
t= 1,
(10)
where 2p is a molecule in the dataset, a, and o, represent the noise schedule and satisfy of + of monotonically decrease from I to 0. Nr represents the Gaussian distribution in the zero center-of-
mass (CoM) subspace satisfying ΣM M x; 0. refers to joint distribution.
To generate samples, the forward process is reversed using a Markov chain re( = ro(=)) IT ro r0(z111/z)) z1) with noise term approximated by neural network tfth (z,, I):
10(2)-1/2)) = No (z.-1 j xz, 16 (2),1)
,21) 21) - (he- (Be. (000000(000 1)) ,02,)))))))))))) (11)
20

---

Published as a conference paper at ICLR 2025
where ji (z. 20 ()) = atj--1 it
o
0 (z, for (a. h], ā = 02/--1011/00
Po
10jt-10:
Ocjt-1 04/ 1 and off-11 of 11-1 1-11
The prior distribution is approximated by p(zN) N
(# (h (hx(0, I). The noise prediction
network is trained using a mean squared error (MSE) loss ming Ee Cr:
L. = (z (z), Eell E 10 (2) (-o, s. log (=()))
(12)
where 6, (2, 0,z0)/o /o, the standard Gaussian noise injected 2g and Vz, logg(zz))) (z,)) refers to the score function. The equivalent loss is derived using the equivalence between score matching and conditional score matching (Vincent 201Tb), with a formal proof provided Ni et al. (2024);
Feng et al. (2023a). Thus Ci termed denoising loss equivalently score estimation loss. Loss
in equation 12 can be further decomposed as atom type and coordinate noise prediction: C.
Bg(xa.=.)) @1 (F (zr,() - - 112 + 110b (z., e 1/ ((=lh) 41 ((hlz)
F REPRESENTATION LEARNING ANALYSIS FOR MOLECULAR COORDINATE GENERATION
We analyze from the perspective of information maximization (InfoMax) (Linsker 1988 Oord et al. 2018) the reasons why training with denoising diffusion loss can effectively learn molecular repre- sentations that enhance property prediction, as well as why property prediction is more effective after ready time.
F.1 INFOMAX TARGET
In accordance with the principles of information maximization (InfoMax), our objective is to select a representation Çt that maximizes the mutual information (MI) between the input data and its repre- sentation. In the context of UniGEM, this objective translates to maximizing the mutual information between the original molecular coordinates Zg and the learned latent representations Ç.- Here, & derived from the intermediate layers of the denoising network, with the input consisting of the noisy coordinates HG at time step t. This latent representation G subsequently utilized in the denoising task. A graphical representation of this model is illustrated in Figure 6
Neural Network
Network
X
X
CE
Given Forward Distribution (xx,,x))
Extracting representations ((G])x;)
Variational distribution p(r.Ks)
Figure 6: Graphical model of representation learning in the diffusion model.
We define q( ax) as the given forward distribution, while g(Gr|ax)) and p(æol(e) represent the
network-modeled latent representation distribution and the denoising predictive distribution, respec-
I 19 dd)a, The
tively. We denote g( Dt, &e) = q(20, x (S) (x), and g(zolGe) := 7 Tt mi Ge)ee,
mutual information target to be maximized can be expressed as follows:
T(Ge, 2o) = ax Joog g(<ol(<)) g(xx))
log g(20 (xa)
-x(x)) By(moo
+ (mg.
[log (
g
(13)
log g [x:))
H((o)) H (xoolx)))
B ate Ge)
q( (z (+)
=1(x0: xt) - Eg(m).))) DKL (g(20) (x; (x)))l9(20llG))
This expression reveals that in order to learn a good representation Çe, it is essential to maximize the term 7( m.)) while minimizing the Kullhack-Leible (KI.) divergence term.
21

---

Published as a conference paper at ICLR 2025
The first term, [( æt), quantifies the information retained between the original and noisy molecular coordinates and is monotonically decreasing with respect to time Given that the KL divergence is nonnegative, [( 2;) serves as the upper bound for [[e, æo), result known as the data processing inequality. Therefore, at later stages of the generative process, such when 4, T(C,, mo) become larger, particularly when the second KL divergence term is minimized.
The second term, however, is challenging minimize since g(Zu|<e) intractable. To address this, we employ variational approach to derive tractable upper bound for the KI. divergence term. Furthermore, we will demonstrate how this KI. divergence term can be minimized within the denoising task, thereby validating the representation learning ca- pability of the denoising process.
F.2 VARIATIONAL UPPER BOUND ON THE KL DIVERGENCE TERM
Since Be(a..s)) [Dri (a 2 the following inequality holds:
/ (( log (() dxp 2 / q(<oole) log p( dæo
(14)
Thus, we obtain the following expression:
(< xxx) ((2oo(())
< (m log (1)) p(
Bg(mo.mm.51) log
Bg(me.)) [DKL [DK [c))]]
(15)
Consequently, by minimizing this variational upper bound in the right hand side of equation 15 can effectively reduce the KL. divergence term in equation 13
F.3 RELATION TO DENOISING Loss
In this section, we discuss the relationship between Eg( [DKL (q(2z) [=)))(((2()))) and the de-
noising loss. Given that ga (z0)) = =N (o I) we can apply Tweedie's formula (Efron 2011 Luo 2022) to obtain F, 22 ar) 1 (ar, + of log q(æ;)). Please note that while the distribution
as
Gaussian in the center-of-mass (CoM) subspace, this does not alter the form of Tweedie's formula. A detailed derivation is provided below:
q( (20g(()))(( dx
Eg Du [1)
q(x1)
I mo (2mo7))3(a 1322 exp( 20 ||x, an
(
(xx))
[[aax0 - xx) + xx] (2 exp(
1 || -
20
(16)
(x:))
f V ( (2)))
f x |
+
(1q(x1)
02
of V g(
2t
of x, log ((x)) a
at
12
Recall that g(G((x;) and p(x)l(e) represent the network-modeled distributions. In practice, our neural network does not involve random variables; rather, we define to parametrize the expectation of p(xo Ge) which means x (x..t) (x..t. Ep [xo[C]] and & a deterministic function x,, denoted as & = go (x;). Thus, the denoising diffusion loss can be rewritten as the Lg distance between two expectations.
22

---

Published as a conference paper at ICLR 2025
[) = Bylas en 11
(00) (2 () ( (-oz Va, log a x 1?
Be(as, 11 (xe,t) (- V Vx. log qf xz )|( +C
St 04 q 2 Fit
Eg(as. 11 (x),1)
112 +C
(17)
Bg m..)
|
2 [2
(x. 0 (z) (Ie, t))//all
= ...)
| Fo [x
Ey pol C
where C' constant independent of 0 (Vincent 201 Is important to note that, in general, the KL divergence cannot be upper-bounded by the L2 distance between expectations for arbitrary distribu- tions. However, the case of Gaussian distributions with identical variance, these two quantities can be equivalent. Specifically, when both p(zulç) (Sa) and (( (z;) follow Gaussian distributions with the same variance a, the following relationship holds:
22 | [x.] Ep 2 kcll?
(18)
Ea (a:)]11 of frz
2o Ro(20....)) 1 [.,]
20
The distribution q(0n axt) is typically approximated as Gaussian in empirical Bayes methods, and similar Gaussian approximation is employed during the reverse process in diffusion models. We assume that p(xol(e) (e) also follows Gaussian distribution with the same variance as g( (z)).
Thus, minimizing the denoising loss corresponds to minimizing the KL divergence term, which contributes to optimizing the lower bound the InfoMax objective in equation 13
F.4 THE BENEFITS OF DIFFUSION FOR PROPERTY PREDICTION: A FORCE LEARNING PERSPECTIVE
The denoising task, particularly at specific time steps during diffusion training, has been shown to effectively learn meaningful molecular representations. as demonstrated prior work (Zaidi al. 2022 Feng et al. 2023a: Ni et al. 2024 Arts et al.) 2023). This process can be interpreted as approximating atomic forces, with the accuracy of these forces being highly sensitive the noise scale applied during training (Feng et al. 2023a Ni et al.) 2024). Specifically, the noise scale must strike: a
too large nor too trivial in order to yield pre-trained representations that beneficial for molecular property prediction. Thus, the denoising training steps in the growth phase may capture valuable molecular features by implicitly learning the atomic forces.
G GENERATIVE ERROR ANALYSIS OF UNIGEM AND JOINT DIFFUSION MODEL
G.1 GENERAL DIFFUSION ERROR BOUND
We follow CCen et al. 2023b) to derive the error bound for diffusion model and make the following mild assumptions the data distribution q(yo). To distinguish it from the notation of molecules, we use y to represent general data generated by the diffusion model.
Assumption (Lipschitz score). For all > the score V g(ve)))) is Lypsscchhizz
Assumption 2 (second moment bound). For some 17 0, Eg(vo) (1l [| is finite. Denote m Eg |ol | for the second moment of q(yo).
T refers to the timestep of the simple Ornstein-Uhlenber (OU) process as defined by Chen, de- scribed by the equation
dv. --, V2dBi, LE [0, T).
23

---

Published as a conference paper at ICLR 2025
while our forward process is given by
dY, = - + (2g(t) dB,, EE [0,TT]
where g(t) is determined by our noise schedule at, The timesteps of the two processes differ by time reparameterivatioo. In EDM, 7 also refers to the discretization steps of the diffusion process.
Theorem G.I (Theorem2 in Chen et al. (2023b)). Suppose that Assumptions and hold. Suppose that the step size I == 7 /T satisfies S 1/L, where I. Then, holds that
(P( x LKL(g(yo)lplll (xp(-T)-+ (L (Lyvagl L.M.DVI +
C
prior distriburion
discretization evror
Score estimation evror
(19)
where fi S f2 denotes that there is a universal constant C such that S f2. dy is the dimension- ality of data y-
In Theorem
modified the last term compared to original paper. The original formulation max, [. but we adopted a tighter expression used the proof of Theorem 10 in Chen
G.1
et al. (023b), considering the average error across all timesteps: T Ei= C. This adjustment Tacilitates more precise comparison of the error differences between the UniGEM method and the joint generation approach.
G.2 ERROR BOUND FOR UNIGEM AND JOINT DIFFUSION
Based on the results above, we can derive the molecular generation error bound for UniGEM and Joint Diffusion.
Theorem 1.2.1. With Assumptions 1 and 2, the total variation between the generated data distribu- tion and ground-trun truth data distribution bounded by:
TV(ro(zu), ((x)) S ///s///n (KL(g(=))]loo( exp(- + (L.
d.l (L.Vd.l
L.m.lVv
prier distribution ervor
discrctization ervor
T
(20)
C
(alk)
+
(klz)
1=1
1=1
coordinate scere estimation envor
atom type score estimation ervor
where q(ZT) Nu (o I), with Oe being a small value close to zero, Nan represents the joint distribution of the corresponding subspace Gaussian of and the Gaussian of h, and m Eg(=o) [0] |2 is the second moment ((zo).
Proof. This result is largely a direct application of the previous theorem, with two key distinctions. Firstly, instead of operating in a standard Gaussian distribution in 3M dimensional space, the molec- ular coordinates have to be projected into the zero center of mass (CoM) subspace. This projection introduces a constant factor difference in the probability density. Specifically, the distribution of the projected standard Gaussian in the zero CoM subspace is given by p(æ) (2 I) (2m))(M-1y2 exp( 1 || (2n)) I). where x E R&M and has zero CoM. In this case, the bound still holds, with further details available in You et al. 2023))
Secondly, the score estimation error can be further decomposed. We begin by decomposing the joint score function as follows:
=x, log q(z)) = aVa, log g(zt), Vn, log g(z))) = (V., logg(x; [hs), Vn, log ghe/(())))) (21) In the forward process It and ht are conditional independent given xo and he. Therefore, we can express the conditional probability
g(æx//.))= g( he, wo, ho)g(wo, ho dæodhoo= J y( ho)g(wo, ho)ddnddh== ( ([) g( ho ho)dho) ddoo)dro== g(xx)). Thus, we obtain
2, log ((x)) = (Va, log g(xx), Vn, log ((h))
(22)
24

---

Published as a conference paper at ICLR 2025
Consequently, the score estimation error be decomposed into the following components:
Lt 2t) 1) - (-o, V., log ((1)))1
((00 (1 (0)) (z1, (-o, or V log (o:)(=;))11 + (1 (zz, 8) (-o, or V log (1+))
- Li (h)z)
(23)
As result, we can now establish the proof of Theorem 2.2 (r(o(z))).
Zv / + del + L.m,
c (b/z)) + c((lb))))
T
S/KL(((o))loo(r))eepppp (L (
(hlz)
c
(m)h)
1=1 [=1
(24)
L
Theorem 2.2 2.
Denote the reverse diffusion process of x po(o)))
(o(TT) II Po x x2) The atom types are predicted by a prediction network
00 (xr)) == 00 (wo. t 0) (pa (ho. (zo. pa (ho, (o)).. With Assumptions and
2, the total variation between the generated data distribution and ground-truth data distribution is bounded by:
TV (po((z)), $ V (KL(((=))lPp(<z)) exp((T) + (LoV + Lmm!)VT
prior distribution
discretization
(25)
Bg(we) ((h) (aro)
atom type estimation crror
coordinate estimation
where (xx) 10b (00l((an (mo) - (xxo) is the atom type prediction error, with f()) providing the ground truth atom type corresponding to Do in the dataset.
Proof.
1
/
TV(po {(20))
po (ha. g(ho. xo) dxc dhg
2
1
/
pe pe(holzu)po(zz) - q x (xx))
dxo dho
2
1
Po (ho po (hoollo))((g (2))) daro dhs
2
/
+
2
po [po - )100 day dhs
(26)
1
po Be(holzz (20) dho)l po(xo) q(20) dx
2
1
E Ey(o
po(ho.s e, es(200) - 8 -00("((a0)==))
2
(po (mo)) ((x)))
(og c(h)) (200)
Then by using Theorem G. to provide the upper bound of the diffusion error of x, we can obtain
equation
25
25

---

Published as a conference paper at ICLR 2025
Remark of the deterministic mapping f. In the datasets involved in this paper, the coordinates of different molecules are distinct. This ensures that the mapping f from coordinates to atom types is well-defined. We posit that this notion of well-defined mapping also applies to the ganic molecules of interest. Consequently, there exists f such that for all potential coordinates En. q( (ho (o) ((a)) justifying architectural choice of utilizing predictive methods to deter- mine atom types.
G.3 GENERATIVE ERROR COMPARISON
We compare the error bounds for the two methods by analyzing each component in the upper bounds:
Prior distribution error and discretization error: These errors are influenced by the dimension- ality the generated data. Since UniGEM operates on lower-dimensional space i.e. generating coordinates only, its error in these terms is smaller.
1. The dimensionality of generated data is smaller for UniGEM: da < <d2.
2. Since the prior distribution of x and h are standard Gaussian distributions the reverse generative process,
have re(zr) =ra(x))ro(tz)), po(xr) =ro(xr)-
Then KL(g( [[a(zr))= [r)(())))) K L[g((r, (r))))r = LL(g( [[)((T)))) +
(ggan) (hg (ro) (ro()))). Thus K KL(g(wo) (()o(())) S KI(g(Za)]]ro(2)) and
K (q(
scales with d=.
3. The second moment term satisfies m = B |l = ()) (1 + ||]13) m i.e. m S m and m² scales with d2.
m +
4. Since the distribution of x and are independent in the forward process, the score function can be expressed =,, In (2x) = (Va, , q(z)), Vn, In gh+)) As a result, the Lipschitz constant for æ is bounded by the Lipschitz constant for the joint variable z, i.e. La S Lz.
5. Other variables T, T are the same for the two methods.
b. Type and coordinate estimation error:
According to Theorems 2.2 we decompose the score estimation error into two components: atom type estimation error and coordinate estimation error, which are compared experimentally. For atom type estimation, we compare the accumulated denoising loss /E c((b)z)) and prediction loss )11g ((h) (xo)). For coordinate estimation, evaluate the results of two denoising losses
VI 11 cl (z]h) and VIEI c(*). It is important to note that the input of the network differs in the denoising losses, with (mnl)) incorporating noisy atom type input:
L (b)b) Bg(xo = (110bb (z,, () (-o, V m. log q (z:) 12
(27)
By ..,) (1 (x,, (-o, log (zx)))
(28)
The detailed experimental comparison and analysis is presented Section 2.2, showing UniGem achieves lower atom types and coordinates estimation error.
26


Reviews:
[{"review_id": "0zrarm0pON", "reviewer": "Reviewer_wNED", "summary": "This paper proposes a unified model for molecular generation and property predictions. Since diffusion models require diffusing the 3D molecular conformers to a collapsed state, the molecular structures in the early time steps do not contain any meaningful structural information. As a result, the joint training of generation and property prediction can be inconsistent in the early stage. The author develops a theoretical framework to explain this issue. To solve this issue, this paper proposes a novel training paradigm that separates the diffusion time steps into two stages, which are nucleation time and growth time. The joint training would only be conducted during the growth time. Experimental results demonstrate the effectiveness of the proposed method.", "questions": "How do you determine the nucleation time before training diffusion models? Currently, it seems that the nucleation time cannot be determined before training the diffusion models. This gonna be extremely costly to select the most suitable nucleation time.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "(1) The presentation of the major idea is very clear. This reviewer can understand the motivation fluently through reading the introduction part. This article is well-organized and easy to follow; \n\n(2) The theoretical analysis is well-developed and can explain the tradeoff phenomenon properly.", "weaknesses": "(1) It appears that the nucleation time cannot be determined prior to training the diffusion models. This could prove to be highly costly if there is a need to repeatedly retrain the diffusion models with varying hyperparameters, which may be impractical in real-world applications;\n\n(2) While the theoretical framework is well-developed, this reviewer believes that the theoretical analysis section occupies an excessive portion of the paper. Given that the phenomenon observed by the authors is quite intuitive, experimental observations or straightforward explanations could suffice;\n\n(3) Although the experimental results demonstrate the effectiveness of the proposed method, UniGem only marginally outperforms the baseline methods across most datasets and tasks. Furthermore, the molecular generation benchmark is not entirely convincing at this stage, as the evaluation metrics used in this benchmark fail to distinguish the capabilities of various state-of-the-art generation models. Also, in the molecular quantum property prediction benchmark tasks, UniGem only slightly surpasses the baseline models, and many quantum properties are absent from the reported benchmark results. Overall, the current empirical results do not sufficiently convince readers of the superiority of the proposed method;\n\n(4) The novelty is a little bit limited from the algorithmic contribution side. The major contribution is splitting the diffusion training into two different stages and only conducting the joint training in the later stage. This is somehow not that significant in more general application domains.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "giYKBrhSpi", "reviewer": "Reviewer_Q62J", "summary": "This paper proposes UniGEM, a unified method designed to effectively address both molecular generation and property prediction tasks. Experimental results and ablation studies demonstrate the effectiveness of the proposed approach.", "questions": "No further questions.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper presents an innovative and well-reasoned approach. Denoising has recently emerged as a dominant pre-training method for molecular representation learning, with strong physical interpretability. Combining the diffusion model with representation learning is a promising direction and good first trial.\n\n2. Decoupling atom type and coordinates during generation is an interesting approach. Together with theoretical insights, this approach shows promise in reducing generation error. This technique resembles some protein generation models, where pure chain geometry is initially learned before predicting amino acid sequences.\n\n3. Experiments validate the efficacy of the proposed framework.", "weaknesses": "1. The experimental setup lacks depth. To better demonstrate the effectiveness of the proposed method, I suggest testing on a property-conditioned generation task using the QM9 dataset, a standard benchmark for generative models. This would offer a more challenging evaluation, testing both the model’s ability to learn representations related to specific conditions and its capacity to capture the underlying data distribution—goals that align very well with the proposed method’s objectives.\n\n2. The comparisons lack recent SOTA methods. It would be valuable to compare UniGEM’s performance in prediction tasks against other pre-trained methods that also leverage denoising, such as UniMol or Frad (I understand that the model backbones differ, yet this comparison could at least shed some lights). This will test whether UniGEM’s approach (i.e., diffusion-based denoising, with larger noise levels and more fine-grained steps) offers advantages over conventional coarser denoising techniques in representation learning.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "aTxEXzm2eK", "reviewer": "Reviewer_5XEL", "summary": "This paper proposes UniGEM, a unified approach for molecular generation and molecular property prediction. The authors carefully designed both generative and predictive objectives in a single unified framework. Interestingly, these two objectives exhibit a synergistic relationship, enhancing each other's performance.", "questions": "1. How much time does it take the overall training?\n\n2. Can EGNN be replaced by recent 3D GNNs, e.g., [4], to further improve performance?\n\n3. Can the authors provide visualizations of the denoising process for generated samples?\n\n[4] Liu et al., Spherical Message Passing for 3D Molecular Graphs, ICLR 2022", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The problems of interest, molecular generation and property prediction, are important for real-world applications, e.g., drug discovery.\n\n- Proposed method seems reasonable; (1) applying the predictive and atom type loss near t=0, (2) changing atom type loss to classification loss, and (3) learning only coordinates when t is far from 0.\n\n- The findings are interesting; simultaneously learning both predictive and generative objective improve each other's performance.", "weaknesses": "- Insufficient baselines.\n\nMy major concern is the omitted baselines. For example, [1,2,3] also study 3D molecular generation and often outperform this method. I think such 3D molecular generation baselines should be in main tables for comparison.\n\n---\n- Marginal improvements.\n\nConsidering the omitted baselines [1,2,3], the improvements seem marginal for me. Furthermore, it seems that [3] largely outperforms this paper. \n\n---\n- Complexity of the framework.\n\nUniGEM introduces many design choices, e.g., t_n and oversampling ratio. I think this may limit the generalizability to other datasets, i.e., require extensive hyper-parameter tuning.\n\n---\nOverall, I think the concept of this work is reasonably novel. However, due to the insufficient comparison with baselines, I lean towards rejection at this moment.\n\n[1] Song et al., Equivariant Flow Matching with Hybrid Probability Transport for 3D Molecule Generation, NeurIPS 2023\\\n[2] Xu et al., Geometric-Facilitated Denoising Diffusion Model for 3D Molecule Generation, AAAI 2024\\\n[3] Irwin et al., Efficient 3D Molecular Generation with Flow Matching and Scale Optimal Transport, arXiv", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)