UniGEM: A Unified Approach to Generation and Property Prediction for Molecules

## Abstract

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.

## 1 Introduction

Artificial intelligence, particularly through deep learning technologies, is advancing various applications in drug discovery. This encompasses two major tasks: molecular property prediction (Zaidi et al., ; Feng et al., ; Ni et al., ; ) and molecule generation (Hoogeboom et al., ; Guan et al., ; Gao et al., ). The objective of molecular property prediction is to learn functions that accurately map molecular samples to their corresponding property labels, which can facilitate the virtual screening process. Meanwhile, molecule generation aims to estimate the underlying 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.

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 a deep understanding of molecular structures, enabling the creation of good representations during the generation process. Recent research findings provide support for our perspective. 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., ; Hudson et al., ; Mittal et al., ). In the molecular domain, studies have indicated that generative pre-training can enhance molecular property prediction tasks (Liu et al., ; Chen et al., ), particularly through data augmentation using diffusion models. However, these methods often require additional fine-tuning to achieve optimal predictive performances. Additionally, while predictive tasks can guide molecule generation (Bao et al., ; Gao et al., ), 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 a critical question: can we develop a unified model to enhance the performance of both tasks?

A straightforward approach to combining these two tasks is to use a multi-task learning framework, where a model is trained to be simultaneously supervised by both property prediction loss and generation loss. However, as shown in 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.

Figure: Figure 1: The two-phase generative process of UniGEM. We treat molecule generation as a two-phase problem: nucleation and growth, defining the separation time as nucleation time. Properties are only well-defined in the growth phase, so during training, property and atom type predictions are incorporated only in the growth phase.
Refer to caption: extracted/5925495/intro2.jpg

We attribute the suboptimal results to the inherent inconsistency between the generative and predictive tasks. In the generative task, the model operates throughout the entire diffusion process, transitioning molecules from a disordered state to a structured form. However, for predictive tasks, meaningful molecular properties can only be defined once the molecular structure has been fully established. Disordered molecules in earlier stages lack meaningful properties and can introduce errors. 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, ; Oord et al., ), a representation is 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 zero at larger time steps, indicating that representations at higher noise levels (i.e., disordered molecules) are 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 a novel two-phase generative approach to unify property prediction and generation, as shown in Figure 1. We divide the molecular 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, ), 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 is developed based on this scaffold. These phases are separated by a ‘nucleation time’. We then introduce a new diffusion model to formulate these two phases.
Before the nucleation time, the diffusion model progressively generates 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, our approach focuses solely on the diffusion of coordinates, treating atom type as a separate prediction task. This separation is justified by the observation that atom types can often be inferred based on the determined scaffold (Bruno et al., ; Vaitkus et al., ). Specifically, prior to nucleation, the diffusion process aims to reconstruct the coordinate. Following nucleation, we incorporate predictive losses for both atom type and properties into the unified learning framework.

Our experimental results demonstrate that UniGEM achieves superior performance in both the predictive 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., ). 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 molecular generation errors and producing more stable and accurate molecular structures.

To the best of our knowledge, we are the first to propose an effective unified model that significantly enhances of the performance of both molecular generation and property prediction tasks, demonstrating a 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.

## 2 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 diffusion approach, we adopt a notation scheme consistent with the E(3) Equivariant Diffusion Model (EDM) (Hoogeboom et al., ). We denote a 3D molecule with $M$ atoms as ${\bm{z}}=({\bm{x}},{\bm{h}})$, where ${\bm{x}}=({\bm{x}}_{1},\cdots,{\bm{x}}_{M})\in\mathbb{R}^{3M}$ represents the atomic positions, and ${\bm{h}}=({\bm{h}}_{1},\cdots,{\bm{h}}_{M})\in\{{\bm{e}}_{0},\cdots,{\bm{e}}_{
H}\}^{M}$ encodes the atom type information. Each ${\bm{h}}_{i}$ 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.

#### 2.1.1 Two-phase Generative Process

In contrast to traditional diffusion based model, UniGEM adopts diffusion process to generate only the atomic coordinates ${\bm{x}}$ rather than the whole molecule ${\bm{z}}$. We define the forward process for adding noise to the molecular coordinates as follows:

$$ $q({\bm{x}}_{0:T})=q({\bm{x}}_{0})\prod_{t=1}^{t_{n}}q({\bm{x}}_{t}|{\bm{x}}_{t -1})\prod_{t=t_{n}+1}^{T}q({\bm{x}}_{t}|{\bm{x}}_{t-1}),$ (1) $$

where $t_{n}$ denotes the nucleation time to indicate when the molecule forms its scaffold. We define the growth phase as $t\in[0,t_{n}]$ and the nucleation phase as $t\in[t_{n},T]$. During the growth phase, we will incorporate atom type and property information, as these attributes are clearly defined at this stage.
The conditional distribution is defined as $q({\bm{x}}_{t}|{\bm{x}}_{0})=\mathcal{N}_{x}({\bm{x}}_{t}|\alpha_{t}{\bm{x}}_{
0},\sigma^{2}_{t}{\bm{I}})$, $t=1,\cdots,T$, where
$\alpha_{t}$ and $\sigma_{t}$ represent the noise schedule following Hoogeboom et al. () and satisfy the condition $\alpha_{t}^{2}+\sigma_{t}^{2}=1$, with $\alpha_{t}$ decreases monotonically from 1 to 0. ${\bm{z}}_{0}=({\bm{x}}_{0},{\bm{h}}_{0})$ represent an unperturbed molecule in the dataset. $\mathcal{N}_{x}$ represents the Gaussian distribution in the zero center-of-mass (CoM) subspace satisfying $\sum_{i}^{M}{\bm{x}}_{i}=\textbf{0}$, to ensure translation equivariance.

In UniGEM, we train a unified model, denoted as ${\bm{\phi}}_{\theta}$, that is capable of both generative and predictive tasks. An overview is provided in Figure 1. Following previous molecular generative models (Hoogeboom et al., ; Wu et al., ; Xu et al., ), we adopt an EGNN architecture (Satorras et al., ) for our implementation.

For the coordinate generation task, we train the model to predict noise from the noisy inputs using a mean squared error (MSE) loss at each timestep:

$$ ${\mathcal{L}}_{t}^{(x)}=\mathbb{E}_{q({\bm{x}}_{0},{\bm{x}}_{n})}\|{\bm{\phi}} ^{(x)}_{\theta}({\bm{x}}_{t},t)-{\bm{\epsilon}}_{t}\|^{2},\quad{\bm{\epsilon}} _{t}=({\bm{x}}_{t}-\alpha_{t}{\bm{x}}_{0})/\sigma_{t},\quad t\in[1,T]$ (2) $$

where ${\bm{\epsilon}}_{t}$ is the standard Gaussian noise injected to ${\bm{x}}_{0}$ and ${\bm{\phi}}^{(x)}_{\theta}({\bm{x}}_{t},t)$ is the equivarant output of the backbone network. The denoising loss is equivalent to the score estimation loss $\simeq\mathbb{E}_{q({\bm{x}}_{0},{\bm{x}}_{n})}\|{\bm{\phi}}^{(x)}_{\theta}({
\bm{x}}_{t},t)-(-\sigma_{t}\nabla_{{\bm{x}}_{t}}\log q({\bm{x}}_{t}))\|^{2}$, where $\nabla_{{\bm{x}}_{t}}\log q({\bm{x}}_{t}))$ refers to the score function. This is derived through the equivalence between score matching and conditional score matching (Vincent, ), with a formal proof presented in Ni et al. (); Feng et al. (). The learned score function is crucial for reversing the noise addition process for coordinate generation.

For the atom type prediction task, we use an $L_{1}$ loss for regressing the one-hot vector of the atom types:

$$ ${\mathcal{L}}^{(h)}_{t}=\mathbb{E}_{q({\bm{x}}_{0},{\bm{x}}_{n})}|{\bm{\phi}}_ {\theta}^{(h)}({\bm{x}}_{t},t)-{\bm{h}}|,\quad t\in[1,t_{n}],$ (3) $$

where ${\bm{h}}$ is the one-hot vector of atom types corresponding to the coordinates ${\bm{x}}_{0}$ in the dataset, and ${\bm{\phi}}_{\theta}^{(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 ${\bm{z}}_{0}$ as $c$, and utilize an $L_{2}$ loss for regressing the property:

$$ ${\mathcal{L}}^{(c)}_{t}=\mathbb{E}_{q({\bm{x}}_{0},{\bm{x}}_{n})}||{\bm{\phi}} _{\theta}^{(c)}({\bm{x}}_{t},t)-c||^{2},\quad t\in[1,t_{n}],$ (4) $$

where ${\bm{\phi}}_{\theta}^{(c)}({\bm{x}}_{t},t)$ is the output from the backbone network with another two-layer MLP prediction head. For $t\in[t_{n},T]$, we set ${\mathcal{L}}^{(c)}_{t}$ and ${\mathcal{L}}^{(h)}_{t}$ as zero.

In this way, we can seamlessly integrate predictive and generative tasks using a single model, and avoid inconsistencies of the two tasks when the molecule structure is not fully formed. However, employing traditional diffusion strategies to 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.

#### 2.1.2 Training Strategies

In our framework, the growth phase occupies a small portion of the overall training process, with the best-performing configuration being $t_{n}/T=0.01$. If 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, we oversample the time steps in the growth phase. Thus, the final loss can be written as the following equation:

$$ ${\mathcal{L}}=\mathbb{E}_{t\sim\frac{1}{2}U(1,t_{n}]+\frac{1}{2}U(t_{n},T]} \left({\mathcal{L}}^{(x)}_{t}+{\mathcal{L}}^{(h)}_{t}+{\mathcal{L}}^{(c)}_{t}\right)$ (5) $$

However, we 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 a multi-branch network architecture. The network shares parameters in the shallow layers but splits into two branches in the deeper layers, with separate sets of parameters for each branch. These branches are activated in different phases of training: one branch is dedicated to 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.1.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 on the generated coordinates. Specifically, we express the reverse process as $p_{\theta}({\bm{x}}_{0:T})=p_{\theta}({\bm{x}}_{T})\prod_{t=1}^{T}p_{\theta}({
\bm{x}}_{t-1}|{\bm{x}}_{t})$, where the noise at each step is approximated by the trained neural network ${\bm{\phi}}_{\theta}^{(x)}({\bm{x}}_{t},t)$. The reverse transition at each time step is given by:

$$ $p_{\theta}({\bm{x}}_{t-1}|{\bm{x}}_{t})=\mathcal{N}_{X}\left({\bm{x}}_{t-1} \bigg{|}\frac{1}{\alpha_{t|s}}{\bm{x}}_{t}-\frac{\sigma^{2}_{t|{t-1}}}{\alpha_ {t|{t-1}}\sigma_{t}}{\bm{\phi}}_{\theta}^{(x)}({\bm{x}}_{t},t),(\frac{\sigma_{ t|t-1}\sigma_{t-1}}{\sigma_{t}})^{2}{\bm{I}}\right).$ (6) $$

Here $\alpha_{t|{t-1}}=\alpha_{t}/\alpha_{t-1}$ and $\sigma^{2}_{t|{t-1}}=\sigma^{2}_{t}-\alpha^{2}_{t|{t-1}}\sigma^{2}_{t-1}$.
The prior distribution is approximated by $p({\bm{x}}_{N})=\mathcal{N}_{X}(x_{N}|{\bm{0}},{\bm{I}})$.

For the molecular generation task, the atom types are predicted at the final time step, producing a complete molecule as ${\bm{z}}=({\bm{x}}_{0},{\bm{\phi}}_{\theta}^{(h)}({\bm{x}}_{0},0))$, as shown in Figure 4.
For property prediction, the inference is done with the time step fixed at zero during testing, i.e., using the output ${\bm{\phi}}_{\theta}^{(c)}({\bm{x}}_{0},0)$.
Notably, this approach incurs no additional computational overhead for either the generative or predictive tasks. The total inference time is the same as the baseline, ensuring that both tasks are performed efficiently without increasing the overall runtime.

### 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

In UniGEM, we incorporate property prediction during the diffusion process, specifically after the nucleation time. We analyze the reason why training with denoising diffusion loss can effectively learn molecular representations, as well as why property prediction is significantly more beneficial after nucleation time.

To facilitate this study, we analyze the problem from an Information Maximization (InfoMax) (Linsker, ; Oord et al., ) perspective, which has been shown to closely relate to the quality of representation learning. According to the InfoMax theory, effective latent representations $\zeta$ in the diffusion-based molecule generation approach are achieved by maximizing the mutual information (MI) between the original molecular coordinates ${\bm{x}}_{0}$ and $\zeta$, which are derived from the noisy coordinates ${\bm{x}}_{t}$ in the intermediate layers of the denoising network. We theoretically demonstrate that good representations can be learned during denoising diffusion training, but are only achieved at smaller time steps, indicating that the latent representations learned from later time steps are insufficient to support property prediction. These findings not only provide a foundation for using diffusion to learn molecular representations but also expose the inherent inconsistency between molecule generation and property prediction tasks. This explains the failure of the multi-task approach and emphasizes the need for a two-phase modeling process.

###### Theorem 2.1 .

The mutual information between ${\bm{x}}_{0}$ and $\zeta$ can be expressed as follows, with a subsequent lower bound:

$$ $\displaystyle I(\zeta,{\bm{x}}_{0})$ $\displaystyle=I({\bm{x}}_{0};{\bm{x}}_{t})-\mathbb{E}_{q({\bm{x}}_{t},\zeta)} \left[D_{\text{KL}}(q({\bm{x}}_{0}|{\bm{x}}_{t})\|q({\bm{x}}_{0}|\zeta))\right]$ (7) $\displaystyle\geq I({\bm{x}}_{0};{\bm{x}}_{t})-\mathbb{E}_{q({\bm{x}}_{t}, \zeta)}\left[D_{\text{KL}}(q({\bm{x}}_{0}|{\bm{x}}_{t})\|p({\bm{x}}_{0}|\zeta) )\right],$ $$

where $q(x_{0},x_{t})$ are data distribution defined by the forward process of diffusion, $q(\zeta|{\bm{x}}_{t})=\delta_{g_{\theta}({\bm{x}}_{t})}$ and $p({\bm{x}}_{0}|\zeta)$ represent the estimated representation and denoising distributions by the denoising network. In practice, our denoising network models $g_{\theta}({\bm{x}}_{t})$ and the mean of the denoising distribution $\mathbb{E}_{p}(\mathbf{x}_{0}|g_{\theta}(\mathbf{x}_{t})):=\frac{\mathbf{x}_{t
}-\sigma_{t}\varphi_{\theta}(\mathbf{x}_{t},t)}{\alpha_{t}}$.

We analyze the lower bound on the mutual information. As we show in appendix D.3, when $p({\bm{x}}_{0}|\zeta)$ and $q({\bm{x}}_{0}|{\bm{x}}_{t})$ follow Gaussian distributions with the same variance $\sigma$, minimizing the KL divergence is equivalent to minimizing the denoising diffusion loss. Thus, the second term $\mathbb{E}_{q({\bm{x}}_{t},\zeta)}\left[D_{\text{KL}}(q({\bm{x}}_{0}|{\bm{x}}_
{t})\|p({\bm{x}}_{0}|\zeta))\right]$ can be minimized during the diffusion training process. This suggests that effective representations can be implicitly learned throughout this diffusion process. However, according to the data processing inequality (Thomas & Cover, ), the first term $I({\bm{x}}_{0};{\bm{x}}_{t})$ decreases monotonically as $t$ increases, eventually approaching zero. As a result, $I(\zeta,{\bm{x}}_{0})$ becomes small for large values of $t$, indicating that it is increasingly challenging to learn meaningful representations from noisy molecules. This explains the failure of multi-task learning in such cases. In contrast, when $t$ is small, $I(\zeta,{\bm{x}}_{0})$ is larger, suggesting that the two tasks can align to facilitate learning robust representations during the molecular growth phase $t\leq t_{n}$.

###### Theorem 2.1 .

#### 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, we derive the upper bound of the generative error for two molecular generation approaches: UniGEM and traditional diffusion-based models, based on results from Chen et al. (). While comparing these upper bounds may not completely reflect real-world behavior, it provides a close approximation that yields valuable insights, enhancing our understanding of the relative performance of these methods.

###### Theorem 2.2 (Generative Error Analysis) .

With mild assumptions on the data distribution $q$ provided in appendix E.1, the molecular generative error measured by total variation between the UniGEM generated data distribution $p_{\theta}({\bm{z}}_{0})$ and ground-truth data distribution $q({\bm{z}}_{0})$ is bounded by the following terms.

$$ $\displaystyle\text{TV}(p_{\theta}({\bm{z}}_{0}),q({\bm{z}}_{0}))\lesssim$ $\displaystyle\underbrace{\sqrt{\text{KL}(q({\bm{x}}_{T})||p_{\theta}({\bm{x}}_ {T}))}e^{-\tilde{T}}}_{\parbox{113.81102pt}{\centering prior distribution error\@add@centering}}+\underbrace{(L_{x}\sqrt{d_{x}}l+L_{x}m_{x}^{2}l)\sqrt{ \tilde{T}}}_{\parbox{85.35826pt}{\centering discretization error \@add@centering}}$ (8) $\displaystyle+\underbrace{\sqrt{l\sum_{t=1}^{T}{\mathcal{L}}_{t}^{(x)}}}_{ \parbox{56.9055pt}{\centering coordinate score \\ estimation error\@add@centering}}+\frac{1}{2}\underbrace{\mathbb{E}_{q({\bm{x} }_{0})}{\mathcal{L}}^{(h)}({\bm{x}}_{0})}_{\parbox{56.9055pt}{\centering atom type \\ estimation error\@add@centering}}$ $$

where $m_{\cdot}^{2}=\mathbb{E}_{q(\cdot)}\|\cdot\|^{2}$ is the second moment of $q(\cdot)$. Traditional diffusion models generate both atomic coordinates and atom types simultaneously, as detailed in Appendix C. The total variation distance between the molecular distribution generated by the traditional diffusion model, denoted as $r_{\theta}({\bm{z}}_{0})$, and the ground-truth data distribution $q({\bm{z}}_{0})$ can be bounded by the following terms.

$$ $\displaystyle\text{TV}(r_{\theta}({\bm{z}}_{0}),q({\bm{z}}_{0}))\lesssim$ $\displaystyle\underbrace{\sqrt{\text{KL}(q({\bm{z}}_{T})||r_{\theta}({\bm{z}}_ {T}))}e^{-\tilde{T}}}_{\parbox{113.81102pt}{\centering prior distribution error\@add@centering}}+\underbrace{(L_{z}\sqrt{d_{z}}l+L_{z}m_{z}^{2}l)\sqrt{ \tilde{T}}}_{\parbox{85.35826pt}{\centering discretization error \@add@centering}}$ (9) $\displaystyle+\underbrace{\sqrt{l\sum_{t=1}^{T}{\mathcal{L}}_{t}^{(x|h)}}}_{ \parbox{56.9055pt}{\centering coordinate score \\ estimation error\@add@centering}}+\underbrace{\sqrt{l\sum_{t=1}^{T}{\mathcal{L }}_{t}^{(h|x)}}}_{\parbox{56.9055pt}{\centering atom type score\\ estimation error\@add@centering}}$ $$

Our theoretical analysis identifies four factors contributing to generation errors in both UniGEM and traditional diffusion-based models: prior distribution error, discretization error, coordinate score estimation error, and atom type score estimation (or prediction) error. 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 we supplement with experimental results showing that UniGEM outperforms traditional models in these aspects:

A critical issue with traditional diffusion-based models is their treatment of discrete atom types as continuous variables, which can lead to 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 a 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 predictions.

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 is smaller.

**Table 1: Comparisons of coordinate and atom type errors.**
|  | UniGEM | EDM |
| --- | --- | --- |
| Atom type | 1.3E-5(0.0001) | 0.0077(0.0012) |
| Coordinate | 0.2247(0.0408) | 0.2430(0.0434) |

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 a separate prediction task.

###### Theorem 2.2 (Generative Error Analysis) .

## 3 Experiments

### 3.1 Main Experiments

Datasets In our experiments, we utilize two widely used datasets. The QM9 dataset (Ruddigkeit et al., ; Ramakrishnan et al., ) 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 as LUMO (Lowest Unoccupied Molecular Orbital), HOMO (Highest Occupied Molecular Orbital), HOMO-LUMO gap, polarizability ($\alpha$). These properties are used as ground-truth labels for property prediction. In contrast, the GEOM-Drugs dataset (Axelrod & Gomez-Bombarelli, ) focuses on drug-like molecules, providing a 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., ; Satorras et al., ; Anderson et al., ), dividing the data into training (100K), validation (18K), and test (13K) sets. Similarly, for the GEOM-Drugs dataset, we follow the approach outlined in Hoogeboom et al. () to randomly split the dataset into training, validation, and test sets in an 8:1:1 ratio.

Implementation Details We adopt EGNN (Satorras et al., ) as our backbone and modify it into a multi-branch network. As illustrated in Figure 3, different branches handle the diffusion loss at different ranges of time steps, with shared layers preceding the branches. In the branch for $t\leq t_{n}$, 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 $t$ is $[t_{n},T]$ , and for the second half, the sampling range for $t$ is $[0,t_{n}]$, ensuring that each branch receives the same samples and is sufficiently trained. All losses are optimized in parallel, with each loss weighted equally at 1. We implement UniGEM with 9 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 8 layers. Optimization is performed using the Adam optimizer, with a batch size of 64 and a 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 $t_{n}$ to 10 and the total time steps $T$ to 1000, ensuring that the prediction of atom types and property occurs within the first 10 time steps. The atom type and molecular property predicted in the last time step are used as the final predictive result.

#### 3.1.1 Experimental Results for Molecule Generation

Baselines
UniGEM is implemented based on the codebase of a classic 3D molecular diffusion algorithm, EDM (Hoogeboom et al., ). Thus, our baseline comprises EDM and its variants, GDM (Hoogeboom et al., ), EDM-Bridge (Wu et al., ) and GeoLDM (Xu et al., ). GDM utilizes non-equivariant networks for training, EDM-Bridge enhances EDM through the technique of diffusion bridges. GeoLDM introduces an additional autoencoder to encode molecular structures into latent embeddings, where the diffusion process is conducted in the latent space. Additionally, we include G-Schnet (Gebauer et al., ), an autoregressive generation method, and Equivariant Normalizing Flows (E-NF) (Garcia Satorras et al., ) 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 distances between each pair of atoms to predict the bond type-whether it is single, double, triple, 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 are 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.

Results
The results in Table 2 show that UniGEM consistently outperforms all baselines across nearly all evaluation metrics for both QM9 and GEOM-Drugs. Notably, compared to EDM variants, UniGEM is significantly simpler, as it neither relies on prior knowledge nor requires additional training for an autoencoder, yet it achieves superior performance compared to EDM-Bridge and GeoLDM, highlighting UniGEM’s superiority.

**Table 2: Comparison of generation performance on metrics, including atom stability, molecule stability, validity, and validity*uniqueness. Higher values indicate better performance.**
|  | QM9 | 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 | 89.8 | 95.0 | 93.2 | 85.1 | 98.4 |

#### 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., ). We compare UniGEM against the EGNN model trained from scratch, referred to as EGNN without confusion, as well as several recent pre-training methods, that have demonstrated superiority in learning molecular representations, thereby improving the property prediction performance. These methods include GraphMVP (Liu et al., ), 3D Infomax (Stärk et al., ), GEM (Fang et al., ), and 3D-EMGP (Jiao et al., ). It is important to note that all baseline methods utilize the same backbone and data splits to ensure a fair comparison.

**Table 3: Performance (MAE, $\downarrow$) on 6 property prediction tasks in QM9. The best results are in bold.**
| Task<br>(Units) | $\alpha$<br>(bohr^3) | $\Delta\epsilon$<br>(meV) | $\epsilon_{\text{HOMO}}$<br>(meV) | $\epsilon_{\text{LUMO}}$<br>(meV) | $\mu$<br>(D) | $C_{v}$<br>($\frac{cal}{mol\cdot K}$) |
| --- | --- | --- | --- | --- | --- | --- |
| EGNN | .071 | 48 | 29 | 25 | .029 | .031 |
| GraphMVP | .070 | 46.9 | 28.5 | 26.3 | .031 | .033 |
| 3D Infomax | .075 | 48.8 | 29.8 | 25.7 | .034 | .033 |
| GEM | .081 | 52.1 | 33.8 | 27.7 | .034 | .035 |
| 3D-EMGP | .057 | 37.1 | 21.3 | 18.2 | .020 | .026 |
| UniGEM | .060 | 34.5 | 20.9 | 16.7 | .019 | .023 |

#### 3.1.3 Comparisons with General Unified Approaches

In this section, we compare UniGEM with two general approaches for unifying generation and property prediction tasks. Table 4 outlines these approaches: the first treats generation and property prediction as a multi-task joint training network, referred to as ‘Multi-task’. Additionally, we incorporate the property prediction task into a pre-trained generation network for continued training, specifically freezing the backbone to preserve the network’s generation capability, referred to as ‘Gen Pre-train’.

Table 4 presents our experimental results, indicating 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 simple integration negatively impacts the performance of both tasks. While ‘Gen Pre-train’ preserves the model’s generation capabilities by freezing the backbone, it also restricts the network’s expressiveness, which in turn limits its performance in property prediction.

**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 | Atom sta(%) | Mol sta(%) | Valid(%) | V*U(%) | $\epsilon_{\text{LUMO}}$ |
| EDM | 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 |

### 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 generation and property prediction tasks, comparing our 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 utilizing only atom type prediction loss (referred to as ATP) and the other using only molecular property prediction loss (referred to as MPP). The comparison results are presented in Table 5, with LUMO selected as the target property for performance evaluation.

As shown in Table 5, both ATP and MPP enhance generation performance compared to previous baseline models, with ATP demonstrating a greater improvement than MPP. The combination of these two losses functions yields the best results, as evidenced by UniGEM. These findings underscore 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 prediction facilitating the learning of better molecular representation, which ultimately benefits property prediction.

**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(%) | $\epsilon_{\text{LUMO}}$ |
| 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.2 Analysis of Training Strategies in UniGEM

To ensure sufficient training steps for property prediction, we oversample the time steps after the nucleation time. Additionally, we design a multi-branch network to mitigate the negative impact of oversampling on 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 6, where the first two rows represent the baselines, and the third row indicates that single-branch with normal time sampling(referred to as ‘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. 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 a multi-branch architecture, allowing different branches to handle different time steps. This adjustment eliminates the negative impact of time step oversampling, resulting in superior performance 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 molecular coordinates experience only minor adjustments. In practical applications, accurately determining the true nucleation time is challenging. Therefore, we conduct an empirical analysis to 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.

Figure: Figure 2: Comparison of performance across different nucleation time selections for both the generation and property prediction tasks.
Refer to caption: x1.png

**Table 6: 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(%) | Valid(%) | V*U(%) | $\epsilon_{\text{LUMO}}$ |
| 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 |

## 4 Conclusion

This paper introduces UniGEM, the first effective unified model that significantly improves the performance 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 molecular representations. The traditional inconsistencies of these tasks are overcome by a two-phase generative process with atom type and property prediction losses activated in the growth phase. UniGEM’s enhanced performance is supported by solid theoretical analysis and comprehensive experimental studies. We believe that the innovative two-phase generation process and its corresponding models offer a new paradigm that may inspire the development of more advanced molecule generation frameworks and benefit more specific applications of molecular generation.

## Appendix A Complimentary Figures

Figure: Figure 3: The training process of UniGEM.
Refer to caption: extracted/5925495/training2.jpg

Figure: Figure 4: The molecular generative process of UniGEM.
Refer to caption: extracted/5925495/generation.jpg

## Appendix B Related Work

### B.1 3D Molecule Generation

The technological development of generative AI has driven numerous research efforts in 3D molecular generation. Luo & Ji (); Gebauer et al. () use autoregressive methods to generate molecules, while Garcia Satorras et al. () employs normalizing flows (Chen et al., ) for molecule generation. Recently, many works have applied diffusion models (Ho et al., ; Sohl-Dickstein et al., ), which have achieved success in image generation, to 3D molecular geometric generation. Hoogeboom et al. () proposed an equivariant diffusion model (EDM) for 3D molecular generation. Building on this, EDM-Bridge (Wu et al., ) proposes adding prior bridges (Peluchetti, ) to improve EDM performance. GeoLDM (Xu et al., ) suggests using an autoencoder to encode molecules into a latent space, and then performing diffusion generation in the latent space. GeoBFN (Song et al., ) introduces a new generative technique, Bayesian Flow Networks (Graves et al., ), to 3D molecular generation, applying Bayesian inference to the parameter space of molecular distributions.

### B.2 3D Molecular Pre-training

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., ; Hou et al., ; Rong et al., ; Feng et al., ), contrastive learning (Stärk et al., ; Liu et al., ; ; Li et al., ), and 3D denoising. Among them, 3D denoising, due to its interpretability in learning force fields, has been proven to be a very effective pre-training strategy (Zaidi et al., ; Feng et al., ; Ni et al., ; ), especially for quantum chemical properties (Feng et al., ). Recently, some pre-training works have integrated diffusion as a pre-training task (Song et al., ) or incorporated diffusion into the pre-training process (Liu et al., ; Chen et al., ), demonstrating that useful molecular representations can be learned during the molecular generation process.

## Appendix C Introduction of Joint Diffusion for Molecules

A 3D molecule is represented by coordinates and atom types. Formally, we denote a molecule with $M$ atoms as ${\bm{z}}=({\bm{x}},{\bm{h}})$, where ${\bm{x}}=({\bm{x}}_{1},\cdots,{\bm{x}}_{M})\in\mathbb{R}^{3M}$ represents the atomic positions, and ${\bm{h}}=({\bm{h}}_{1},\cdots,{\bm{h}}_{M})\in\{{\bm{e}}_{0},\cdots,{\bm{e}}_{
H}\}^{M}$ encodes the atom type information. Each ${\bm{h}}_{i}$ 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 (Sohl-Dickstein et al., ; Ho et al., ) to 3D molecular data, employing joint diffusion to simultaneously generate molecular coordinates and atom types (Hoogeboom et al., ; Guan et al., ; Gao et al., ).
The diffusion model independently injects noise to the coordinates and atom types respectively via a forward process:

$$ $\quad q({\bm{z}}_{t}|{\bm{z}}_{0})=\mathcal{N}_{x}({\bm{x}}_{t}|\alpha_{t}{\bm {x}}_{0},\sigma^{2}_{t}{\bm{I}})\cdot\mathcal{N}({\bm{h}}_{t}|\alpha_{t}{\bm{h }}_{0},\sigma^{2}_{t}{\bm{I}}),\quad t=1,\cdots,T$ (10) $$

where $z_{0}$ is a molecule in the dataset, $\alpha_{t}$ and $\sigma_{t}$ represent the noise schedule and satisfy $\alpha^{2}_{t}+\sigma_{t}^{2}=1$ monotonically decrease from $1$ to 0 0. $\mathcal{N}_{x}$ represents the Gaussian distribution in the zero center-of-mass (CoM) subspace satisfying $\sum_{i}^{M}{\bm{x}}_{i}=\textbf{0}$. $\cdot$ refers to joint distribution.

To generate samples, the forward process is reversed using a Markov chain $r_{\theta}({\bm{z}}_{0:T})=r_{\theta}({\bm{z}}_{T})\prod_{t=1}^{T}r_{\theta}({
\bm{z}}_{t-1}|{\bm{z}}_{t})$ with a noise term approximated by a neural network ${\bm{\phi}}_{\theta}^{(h)}({\bm{z}}_{t},t)$:

$$ $r_{\theta}({\bm{z}}_{t-1}|{\bm{z}}_{t})=\mathcal{N}_{X}\left({\bm{x}}_{t-1}| \tilde{{\bm{\mu}}}\left({\bm{x}}_{t},{\bm{\phi}}_{\theta}^{(x)}({\bm{z}}_{t},t )\right),\tilde{\sigma}^{2}{\bm{I}}\right)\cdot\mathcal{N}\left({\bm{h}}_{t-1} |\tilde{{\bm{\mu}}}\left({\bm{h}}_{t},{\bm{\phi}}_{\theta}^{(h)}({\bm{z}}_{t}, t)\right),\tilde{\sigma}^{2}{\bm{I}}\right)\\$ (11) $$

where $\tilde{{\bm{\mu}}}\left(i_{t},{\bm{\phi}}^{(i)}_{\theta}({\bm{z}}_{t},t)\right
)=\frac{1}{\alpha_{t|s}}i_{t}-\frac{\sigma^{2}_{t|{t-1}}}{\alpha_{t|{t-1}}
\sigma_{t}}{\bm{\phi}}_{\theta}^{(i)}({\bm{z}}_{t},t)$, for $i\in\{{\bm{x}},{\bm{h}}\}$,
$\tilde{\sigma}=\sigma_{t|t-1}\sigma_{t-1}/\sigma_{t}$,
$\alpha_{t|{t-1}}=\alpha_{t}/\alpha_{t-1}$ and $\sigma^{2}_{t|{t-1}}=\sigma^{2}_{t}-\alpha^{2}_{t|{t-1}}\sigma^{2}_{t-1}$.

The prior distribution is approximated by $p({\bm{z}}_{N})=\mathcal{N}_{X}({\bm{x}}_{N}|{\bm{0}},{\bm{I}})\cdot\mathcal{N
}({\bm{h}}_{N}|{\bm{0}},{\bm{I}})$. The noise prediction network is trained using a mean squared error (MSE) loss $\min_{\theta}\mathbb{E}_{t\sim U(0,T]}{\mathcal{L}}_{t}$:

$$ ${\mathcal{L}}_{t}=\mathbb{E}_{q({\bm{z}}_{0},{\bm{z}}_{n})}\|{\bm{\phi}}_{ \theta}({\bm{z}}_{t},t)-{\bm{\epsilon}}_{t}\|^{2}\simeq\mathbb{E}_{q({\bm{z}}_ {0},{\bm{z}}_{n})}\|{\bm{\phi}}_{\theta}({\bm{z}}_{t},t)-(-\sigma_{t}\nabla_{{ \bm{z}}_{t}}\log q({\bm{z}}_{t}))\|^{2},$ (12) $$

where $\epsilon_{t}=({\bm{z}}_{t}-\alpha_{t}{\bm{z}}_{0})/\sigma_{t}$ is the standard Gaussian noise injected to ${\bm{z}}_{0}$ and $\nabla_{{\bm{z}}_{t}}\log q({\bm{z}}_{t}))$ refers to the score function. The equivalent loss is derived using the equivalence between score matching and conditional score matching (Vincent, ), with a formal proof provided in Ni et al. (); Feng et al. (). Thus ${\mathcal{L}}_{t}$ is termed denoising loss or equivalently score estimation loss. Loss in equation 12 can be further decomposed as atom type and coordinate noise prediction:
${\mathcal{L}}_{t}=\mathbb{E}_{q({\bm{z}}_{0},{\bm{z}}_{n})}\|{\bm{\phi}}^{(x)}
_{\theta}({\bm{z}}_{t},t)-{\bm{\epsilon}}^{(x)}_{t}\|^{2}+\|{\bm{\phi}}^{(h)}_
{\theta}({\bm{z}}_{t},t)-{\bm{\epsilon}}^{(h)}_{t}\|^{2}:={\mathcal{L}}^{(x|h)
}_{t}+{\mathcal{L}}^{(h|x)}_{t}$.

## Appendix D Representation Learning Analysis for Molecular Coordinate Generation

We analyze from the perspective of information maximization (InfoMax) (Linsker, ; Oord et al., ) the reasons why training with denoising diffusion loss can effectively learn molecular representations that enhance property prediction, as well as why property prediction is more effective after ready time.

### D.1 InfoMax Target

In accordance with the principles of information maximization (InfoMax), our objective is to select a representation $\zeta$ that maximizes the mutual information (MI) between the input data and its representation. In the context of UniGEM, this objective translates to maximizing the mutual information between the original molecular coordinates ${\bm{x}}_{0}$ and the learned latent representations $\zeta$. Here, $\zeta$ is derived from the intermediate layers of the denoising network, with the input consisting of the noisy coordinates ${\bm{x}}_{t}$ at time step $t$. This latent representation $\zeta$ is subsequently utilized in the denoising task. A graphical representation of this model is illustrated in Figure 5.

Figure: Figure 5: Graphical model of representation learning in the diffusion model.
Refer to caption: extracted/5925495/graphic.jpg

We define $q({\bm{x}}_{0},{\bm{x}}_{t})$ as the given forward distribution, while $q(\zeta|{\bm{x}}_{t})$ and $p({\bm{x}}_{0}|\zeta)$ represent the network-modeled latent representation distribution and the denoising predictive distribution, respectively. We denote $q({\bm{x}}_{0},{\bm{x}}_{t},\zeta_{t}):=q({\bm{x}}_{0},{\bm{x}}_{t})q(\zeta|{
\bm{x}}_{t})$, and $q({\bm{x}}_{0}|\zeta):=\frac{\int q({\bm{x}}_{0},{\bm{x}}_{t},\zeta_{t})d{\bm{
x}}_{t}}{\int q({\bm{x}}_{0},{\bm{x}}_{t},\zeta_{t})d{\bm{x}}_{t}d{\bm{x}}_{0}}$. The mutual information target to be maximized can be expressed as follows:

$$ $\displaystyle I(\zeta,{\bm{x}}_{0})$ $\displaystyle=\mathbb{E}_{q({\bm{x}}_{0},{\bm{x}}_{t},\zeta)}\left[\log\frac{q ({\bm{x}}_{0}|\zeta)}{q({\bm{x}}_{0})}\right]$ (13) $\displaystyle=H({\bm{x}}_{0})-\mathbb{E}_{q({\bm{x}}_{0},{\bm{x}}_{t},\zeta)} \left[\log\frac{q({\bm{x}}_{0}|{\bm{x}}_{t})}{q({\bm{x}}_{0}|\zeta)}\right]+ \mathbb{E}_{q({\bm{x}}_{0},{\bm{x}}_{t})}\left[\log q({\bm{x}}_{0}|{\bm{x}}_{t })\right]$ $\displaystyle=H({\bm{x}}_{0})-H({\bm{x}}_{0}|{\bm{x}}_{t})-\mathbb{E}_{q({\bm{ x}}_{0},{\bm{x}}_{t},\zeta)}\left[\log\frac{q({\bm{x}}_{0}|{\bm{x}}_{t})}{q({ \bm{x}}_{0}|\zeta)}\right]$ $\displaystyle=I({\bm{x}}_{0};{\bm{x}}_{t})-\mathbb{E}_{q({\bm{x}}_{t},\zeta)} \left[D_{\text{KL}}(q({\bm{x}}_{0}|{\bm{x}}_{t})\|q({\bm{x}}_{0}|\zeta))\right]$ $$

This expression reveals that in order to learn a good representation $\zeta$, it is essential to maximize the term $I({\bm{x}}_{0};{\bm{x}}_{t})$ while minimizing the Kullback-Leibler (KL) divergence term.

- •
The first term, $I({\bm{x}}_{0};{\bm{x}}_{t})$, quantifies the information retained between the original and noisy molecular coordinates and is monotonically decreasing with respect to time $t$. Given that the KL divergence is nonnegative, $I({\bm{x}}_{0};{\bm{x}}_{t})$ serves as the upper bound for $I(\zeta,{\bm{x}}_{0})$, a result known as the data processing inequality. Therefore, at later stages of the generative process, such as when $t\leq t_{r}$, $I(\zeta,{\bm{x}}_{0})$ can become larger, particularly when the second KL divergence term is minimized.
- •
The second term, however, is challenging to minimize since $q({\bm{x}}_{0}|\zeta)$ is intractable. To address this, we employ a variational approach to derive a tractable upper bound for the KL divergence term. Furthermore, we will demonstrate how this KL divergence term can be minimized within the denoising task, thereby validating the representation learning capability of the denoising process.

### D.2 Variational Upper Bound on the KL Divergence Term

Since $\mathbb{E}_{q({\bm{x}}_{t},\zeta)}\left[D_{\text{KL}}(q({\bm{x}}_{0}|\zeta)\|p
({\bm{x}}_{0}|\zeta))\right]\geq 0$, the following inequality holds:

$$ $\int q({\bm{x}}_{0}|\zeta)\log q({\bm{x}}_{0}|\zeta)\,d{\bm{x}}_{0}\geq\int q( {\bm{x}}_{0}|\zeta)\log p({\bm{x}}_{0}|\zeta)\,d{\bm{x}}_{0}$ (14) $$

Thus, we obtain the following expression:

$$ $\mathbb{E}_{{\bm{x}}_{0},{\bm{x}}_{t},\zeta}\left[\log\frac{q({\bm{x}}_{0}|{ \bm{x}}_{t})}{q({\bm{x}}_{0}|\zeta)}\right]\leq\mathbb{E}_{{\bm{x}}_{0},{\bm{x }}_{t},\zeta}\log\frac{q({\bm{x}}_{0}|{\bm{x}}_{t})}{p({\bm{x}}_{0}|\zeta)}= \mathbb{E}_{q({\bm{x}}_{t},\zeta)}\left[D_{\text{KL}}(q({\bm{x}}_{0}|{\bm{x}}_ {t})\|p({\bm{x}}_{0}|\zeta))\right]$ (15) $$

Consequently, by minimizing this variational upper bound in the right hand side of equation 15, we can effectively reduce the KL divergence term in equation 13.

### D.3 Relation to Denoising Loss

In this section, we discuss the relationship between
$\mathbb{E}_{q({\bm{x}}_{t},\zeta)}\left[D_{\text{KL}}(q({\bm{x}}_{0}|{\bm{x}}_
{t})\|p({\bm{x}}_{0}|\zeta))\right]$ and the denoising loss.
Given that $q({\bm{x}}_{t}|{\bm{x}}_{0})={\mathcal{N}}_{x}(\alpha_{t}{\bm{x}}_{0},\sigma_{
t}^{2}{\bm{I}})$, we can apply Tweedie’s formula (Efron, ; Luo, ) to obtain $\mathbb{E}_{q}[{\bm{x}}_{0}|{\bm{x}}_{t}]=\frac{1}{\alpha_{t}}({\bm{x}}_{t}+
\sigma_{t}^{2}\nabla\log q({\bm{x}}_{t}))$. Please note that while the distribution is Gaussian in the zero center-of-mass (CoM) subspace, this does not alter the form of the Tweedie’s formula. A detailed derivation is provided below:

$$ $\displaystyle\mathbb{E}_{q}[{\bm{x}}_{0}|{\bm{x}}_{t}]=$ $\displaystyle\int{\bm{x}}_{0}q({\bm{x}}_{0}|{\bm{x}}_{t})d{\bm{x}}_{0}=\int{ \bm{x}}_{0}\frac{q({\bm{x}}_{t}|{\bm{x}}_{0})q({\bm{x}}_{0})}{q({\bm{x}}_{t})} d{\bm{x}}_{0}$ (16) $\displaystyle=$ $\displaystyle\frac{\int{\bm{x}}_{0}\frac{1}{(2\pi\sigma_{t}^{2})^{3(M-1)/2}} \exp(-\frac{1}{2\sigma_{t}^{2}}||{\bm{x}}_{t}-\alpha_{t}{\bm{x}}_{0}||^{2})q({ \bm{x}}_{0})d{\bm{x}}_{0}}{q({\bm{x}}_{t})}$ $\displaystyle=$ $\displaystyle\frac{\int[(\alpha_{t}{\bm{x}}_{0}-{\bm{x}}_{t})+{\bm{x}}_{t}] \frac{1}{(2\pi\sigma_{t}^{2})^{3(M-1)/2}}\exp(-\frac{1}{2\sigma_{t}^{2}}||{\bm {x}}_{t}-\alpha_{t}{\bm{x}}_{0}||^{2})q({\bm{x}}_{0})d{\bm{x}}_{0}}{\alpha_{t} q({\bm{x}}_{t})}$ $\displaystyle=$ $\displaystyle\frac{\int\sigma_{t}^{2}\nabla_{{\bm{x}}_{t}}q({\bm{x}}_{t}|{\bm{ x}}_{0})q({\bm{x}}_{0})d{\bm{x}}_{0}}{\alpha_{t}q({\bm{x}}_{t})}+\frac{\int{ \bm{x}}_{t}q({\bm{x}}_{0}|{\bm{x}}_{t})d{\bm{x}}_{0}}{\alpha_{t}}$ $\displaystyle=$ $\displaystyle\frac{\sigma_{t}^{2}\nabla_{{\bm{x}}_{t}}\int q({\bm{x}}_{t}|{\bm {x}}_{0})q({\bm{x}}_{0})d{\bm{x}}_{0}}{\alpha_{t}q({\bm{x}}_{t})}+\frac{{\bm{x }}_{t}}{\alpha_{t}}=\frac{\sigma_{t}^{2}\nabla_{{\bm{x}}_{t}}\log q({\bm{x}}_{ t})+{\bm{x}}_{t}}{\alpha_{t}}$ $$

Recall that $q(\mathbf{\zeta}|\mathbf{x}_{t})$ and $p(\mathbf{x}_{0}|\mathbf{\zeta})$ represent the network-modeled distributions. In practice, our neural network does not involve random variables; rather, we define it to parametrize the expectation of $p(\mathbf{x}_{0}|\mathbf{\zeta})$ and $\mathbf{\zeta}$ is a deterministic function of $\mathbf{x}_{t}$, denoted as $\mathbf{\zeta}=g_{\theta}(\mathbf{x}_{t})$. Furthermore, $\mathbb{E}_{p}(\mathbf{x}_{0}|\mathbf{\zeta})=\mathbb{E}_{p}(\mathbf{x}_{0}|g_
{\theta}(\mathbf{x}_{t}))$ can be expressed as $\frac{\mathbf{x}_{t}-\sigma_{t}\varphi_{\theta}(\mathbf{x}_{t},t)}{\alpha_{t}}$. Thus, the denoising diffusion loss can be rewritten as

$$ $\displaystyle{\mathcal{L}}_{t}$ $\displaystyle=\mathbb{E}_{q({\bm{x}}_{0},{\bm{x}}_{n})}\|{\bm{\phi}}^{(x)}_{ \theta}({\bm{x}}_{t},t)-(-\sigma_{t}\nabla_{{\bm{x}}_{t}}\log q({\bm{x}}_{t})) \|^{2}$ (17) $\displaystyle=\mathbb{E}_{q({\bm{x}}_{0},{\bm{x}}_{n})}\|{\bm{\phi}}^{(x)}_{ \theta}({\bm{x}}_{t},t)-\frac{{\bm{x}}_{t}-\alpha_{t}\mathbb{E}_{q}[{\bm{x}}_{ 0}|{\bm{x}}_{t}]}{\sigma_{t}}\|^{2}$ $\displaystyle=\mathbb{E}_{q({\bm{x}}_{0},{\bm{x}}_{n})}\frac{\alpha_{t}^{2}}{ \sigma_{t}^{2}}\|\mathbb{E}_{q}[{\bm{x}}_{0}|{\bm{x}}_{t}]-({\bm{x}}_{t}- \sigma_{t}{\bm{\phi}}^{(x)}_{\theta}({\bm{x}}_{t},t))/\alpha_{t}\|^{2}$ $\displaystyle:=\mathbb{E}_{q({\bm{x}}_{0},{\bm{x}}_{n})}\frac{\alpha_{t}^{2}}{ \sigma_{t}^{2}}\|\mathbb{E}_{q}[{\bm{x}}_{0}|{\bm{x}}_{t}]-\mathbb{E}_{p}[{\bm {x}}_{0}|g_{\theta}({\bm{x}}_{t})]\|^{2}$ $$

In this formulation, the denoising loss is defined as the $L_{2}$ distance between the expectations $\mathbb{E}_{q}[{\bm{x}}_{0}|{\bm{x}}_{t}]$ and $\mathbb{E}_{p}[{\bm{x}}_{0}|{\bm{x}}_{t}]$. It is important to note that, in general, the KL divergence cannot be upper-bounded by the $L_{2}$ distance between expectations for arbitrary distributions. However, in the case of Gaussian distributions with identical variance, these two quantities can be equivalent. Specifically, when both $p({\bm{x}}_{0}|\zeta)$ and $q({\bm{x}}_{0}|{\bm{x}}_{t})$ follow Gaussian distributions with the same variance $\sigma$, the following relationship holds:

$$ $\displaystyle\mathbb{E}_{q({\bm{x}}_{t},\zeta)}\left[D_{\text{KL}}(q({\bm{x}}_ {0}|{\bm{x}}_{t})\|p({\bm{x}}_{0}|\zeta))\right]=\mathbb{E}_{q({\bm{x}}_{t}, \zeta)}\left[\frac{1}{2\sigma^{2}}\|\mathbb{E}_{q}[{\bm{x}}_{0}|{\bm{x}}_{t}]- \mathbb{E}_{p}[{\bm{x}}_{0}|\zeta]\|^{2}\right]$ (18) $\displaystyle=\frac{1}{2\sigma^{2}}\mathbb{E}_{q({\bm{x}}_{0},{\bm{x}}_{n})}\| \mathbb{E}_{q}[{\bm{x}}_{0}|{\bm{x}}_{t}]-\mathbb{E}_{p}[{\bm{x}}_{0}|g_{ \theta}({\bm{x}}_{t})]\|^{2}=\frac{\sigma_{t}^{2}}{2\sigma^{2}\alpha_{t}^{2}}{ \mathcal{L}}_{t},$ $$

The distribution $q({\bm{x}}_{0}|{\bm{x}}_{t})$ is typically approximated as Gaussian in empirical Bayes methods, and a similar Gaussian approximation is employed during the reverse process in diffusion models. We assume that $p({\bm{x}}_{0}|\zeta)$ also follows a Gaussian distribution with the same variance as $q({\bm{x}}_{0}|{\bm{x}}_{t})$.

Thus, minimizing the denoising loss corresponds to minimizing the KL divergence term, which contributes to optimizing the lower bound of the InfoMax objective in equation 13.

### D.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 in prior work (Zaidi et al., ; Feng et al., ; Ni et al., ; Arts et al., ). This process can be interpreted as approximating atomic forces, with the accuracy of these forces being highly sensitive to the noise scale applied during training (Feng et al., ; Ni et al., ). Specifically, the noise scale must strike a balance—neither too large nor too trivial—in order to yield pre-trained representations that are 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.

## Appendix E Generative Error Analysis of UniGEM and Joint Diffusion Model

### E.1 General Diffusion Error Bound

We follow (Chen et al., ) to derive the error bound for diffusion model and make the following mild assumptions on the data distribution $q(y_{0})$. To distinguish it from the notation of molecules, we use $y$ to represent general data generated by the diffusion model.

Assumption 1 (Lipschitz score). For all $t\geq 0$, the score $\nabla\ln q(y_{t})$ is $L$-Lipschitz.

Assumption 2 (second moment bound). For some $\eta>0$, $\mathbb{E}_{q(y_{0})}[\|Y_{0}\|^{2+\eta}]$ is finite. Denote $m^{2}_{2}=\mathbb{E}_{q(y_{0})}[\|Y_{0}\|^{2}]$ for the second moment of $q(y_{0})$.

$\tilde{T}$ refers to the timestep of the simple Ornstein-Uhlenbeck (OU) process as defined by Chen, described by the equation

$$ $d\bar{Y}_{t}=-\bar{Y}_{t}\,dt+\sqrt{2}\,dB_{t},\quad t\in[0,\tilde{T}],$ $$

while our forward process is given by

$$ $d\bar{Y}_{t}=-{g(t)}^{2}\bar{Y}_{t}\,dt+\sqrt{2}\,g(t)\,dB_{t},\quad t\in[0,T],$ $$

where g(t) is determined by our noise schedule $\alpha_{t}$. The timesteps of the two processes differ by a time reparameterization. In EDM, $T$ also refers to the discretization steps of the diffusion process.

###### Theorem E.1 ( Chen et al. ( 2023b ) , modified) .

Suppose that Assumptions 1 and 2 hold. Suppose that the step size $l:=\tilde{T}/T$ satisfies $l\lesssim 1/L$, where $L\geq 1$. Then, it holds that

$$ $TV(p_{\theta}({\bm{y}}_{0}),q({\bm{y}}_{0}))\lesssim\underbrace{\sqrt{KL(q({ \bm{y}}_{T})||p_{\theta}({\bm{y}}_{T}))}\exp(-\tilde{T})}_{\text{prior distribution error}}+\underbrace{(L_{y}\sqrt{d_{y}}l+L_{y}m_{y}^{2}l)\sqrt{ \tilde{T}}}_{\text{discretization error}}+\underbrace{\sqrt{l\sum_{t=1}^{T}{ \mathcal{L}}_{t}}}_{\text{score estimation error}},$ (19) $$

where $f_{1}\lesssim f_{2}$ denotes that there is a universal constant $C$ such that $f_{1}\leq f_{2}$, $d_{y}$ is the dimensionality of data $y$.

In Theorem E.1, we modified the last term compared to original paper. The original formulation is $\sqrt{\max_{t}\{{\mathcal{L}}_{t}\}\tilde{T}}$, but we adopted a tighter expression used in the proof of Theorem 10 in Chen et al. (), considering the average error across all timesteps: $\sqrt{\frac{\tilde{T}}{T}\sum_{t=1}^{T}{\mathcal{L}}_{t}}$.
This adjustment facilitates a more precise comparison of the error differences between the UniGEM method and the joint generation approach.

###### Theorem E.1 ( Chen et al. ( 2023b ) , modified) .

### E.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 2.2-1.
With Assumptions 1 and 2, the total variation between the generated data distribution and ground-truth data distribution is bounded by:

$$ $\displaystyle TV(r_{\theta}({\bm{z}}_{0}),q({\bm{z}}_{0}))\lesssim$ $\displaystyle\underbrace{\sqrt{KL(q({\bm{z}}_{T})||r_{\theta}({\bm{z}}_{T}))} \exp(-\tilde{T})}_{\text{prior distribution error}}+\underbrace{(L_{z}\sqrt{d_ {z}}l+L_{z}m_{z}^{2}l)\sqrt{\tilde{T}}}_{\text{discretization error}}$ (20) $\displaystyle+\underbrace{\sqrt{l\sum_{t=1}^{T}{\mathcal{L}}_{t}^{(x|h)}}}_{ \text{coordinate score estimation error}}+\underbrace{\sqrt{l\sum_{t=1}^{T}{ \mathcal{L}}_{t}^{(h|x)}}}_{\text{atom type score estimation error}},$ $$

where $q({\bm{z}}_{T})=\mathcal{N}_{xh}(\alpha_{t}{\bm{z}}_{0},{\bm{I}})$, with $\alpha_{t}$ being a small value close to zero, $\mathcal{N}_{xh}$ represents the joint distribution of the corresponding subspace Gaussian of ${\bm{x}}$ and the Gaussian of ${\bm{h}}$, and $m_{z}^{2}=\mathbb{E}_{q({\bm{z}})}\|{\bm{z}}\|^{2}$ is the second moment of $q({\bm{z}})$.

###### 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 molecular coordinates must 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({\bm{x}})={\mathcal{N}}_{x}({\bm{x}}|{\bm{0}},{\bm{I}})=\frac{1}{(2\pi)^{3(M
-1)/2}}\exp(-\frac{1}{2}||{\bm{x}}||^{2})=(2\pi)^{3/2}{\mathcal{N}}({\bm{0}},{
\bm{I}})$, where ${\bm{x}}\in\mathbb{R}^{3M}$ and has zero CoM. In this case, the bound still holds, with further details available in You et al. ().

Secondly, the score estimation error can be further decomposed. We begin by decomposing the joint score function as follows:

$$ $\displaystyle\nabla_{{\bm{z}}_{t}}\log q({\bm{z}}_{t})=(\nabla_{{\bm{x}}_{t}} \log q({\bm{z}}_{t}),\nabla_{{\bm{h}}_{t}}\log q({\bm{z}}_{t}))=(\nabla_{{\bm{ x}}_{t}}\log q({\bm{x}}_{t}|{\bm{h}}_{t}),\nabla_{{\bm{h}}_{t}}\log q({\bm{h}} _{t}|{\bm{x}}_{t}))$ (21) $$

In the forward process $x_{t}$ and $h_{t}$ are conditional independent given $x_{0}$ and $h_{0}$. Therefore, we can express the conditional probability as $q({\bm{x}}_{t}|{\bm{h}}_{t})$$=$$\int q({\bm{x}}_{t}|{\bm{h}}_{t},{\bm{x}}_{0},{\bm{h}}_{0})q({\bm{x}}_{0},{\bm
{h}}_{0})d{\bm{x}}_{0}d{\bm{h}}_{0}$$=\int q({\bm{x}}_{t}|{\bm{x}}_{0},{\bm{h}}_{0})q({\bm{x}}_{0},{\bm{h}}_{0})d{
\bm{x}}_{0}d{\bm{h}}_{0}$$=\int q({\bm{x}}_{t}|{\bm{x}}_{0})(\int q({\bm{x}}_{0},{\bm{h}}_{0})d{\bm{h}}_
{0})d{\bm{x}}_{0}$$=q({\bm{x}}_{t})$. Thus, we obtain

$$ $\displaystyle\nabla_{{\bm{z}}_{t}}\log q({\bm{z}}_{t})=(\nabla_{{\bm{x}}_{t}} \log q({\bm{x}}_{t}),\nabla_{{\bm{h}}_{t}}\log q({\bm{h}}_{t}))$ (22) $$

Consequently, the score estimation error can be decomposed into the following components:

$$ $\displaystyle{\mathcal{L}}_{t}:=$ $\displaystyle\mathbb{E}_{q({\bm{z}}_{0},{\bm{z}}_{n})}\|{\bm{\phi}}_{\theta}({ \bm{z}}_{t},t)-(-\sigma_{t}\nabla_{{\bm{z}}_{t}}\log q({\bm{z}}_{t}))\|^{2}$ (23) $\displaystyle=$ $\displaystyle\mathbb{E}_{q({\bm{z}}_{0},{\bm{z}}_{n})}(\|{\bm{\phi}}_{\theta}^ {(x)}({\bm{z}}_{t},t)-(-\sigma_{t}\nabla_{{\bm{x}}_{t}}\log q({\bm{x}}_{t}))\| ^{2}+\|{\bm{\phi}}_{\theta}^{(h)}({\bm{z}}_{t},t)-(-\sigma_{t}\nabla_{{\bm{h}} _{t}}\log q({\bm{h}}_{t}))\|^{2})$ $\displaystyle:=$ $\displaystyle{\mathcal{L}}_{t}^{(x|h)}+{\mathcal{L}}_{t}^{(h|x)}$ $$

As a result, we can now establish the proof of Theorem 2.2.

$$ $\displaystyle TV(r_{\theta}({\bm{z}}_{0}),q({\bm{z}}_{0}))$ (24) $\displaystyle\lesssim$ $\displaystyle{\sqrt{KL(q({\bm{z}}_{T})||r_{\theta}({\bm{z}}_{T}))}\exp(-\tilde {T})}+(L_{z}\sqrt{d_{z}}l+L_{z}m_{z}^{2}l)\sqrt{\tilde{T}}+\sqrt{l\sum_{t=1}^{ T}({\mathcal{L}}_{t}^{(h|x)}+{\mathcal{L}}_{t}^{(x|h)}))}$ $\displaystyle\lesssim$ $\displaystyle\sqrt{KL(q({\bm{z}}_{T})||r_{\theta}({\bm{z}}_{T}))}\exp(-\tilde{ T})+(L_{z}\sqrt{d_{z}}l+L_{z}m_{z}^{2}l)\sqrt{\tilde{T}}+\sqrt{l\sum_{t=1}^{T} {\mathcal{L}}_{t}^{(h|x)}}+\sqrt{l\sum_{t=1}^{T}{\mathcal{L}}_{t}^{(x|h)}}$ $$

∎

Theorem 2.2-2.
Denote the reverse diffusion process of ${\bm{x}}$ is $p_{\theta}({\bm{x}}_{0:T})=p_{\theta}({\bm{x}}_{T})\prod_{t=1}^{T}p_{\theta}({
\bm{x}}_{t-1}|{\bm{x}}_{t})$. The atom types are predicted by a prediction network $\phi_{\theta}^{(h)}({\bm{x}}_{0}):=\phi_{\theta}^{(h)}({\bm{x}}_{0},t=0)=(p_{
\theta}({\bm{h}}_{0,1}|{\bm{x}}_{0}),\cdots,p_{\theta}({\bm{h}}_{0,M}|{\bm{x}}
_{0}))$. With Assumptions 1 and 2, the total variation between the generated data distribution and ground-truth data distribution is bounded by:

$$ $\displaystyle TV(p_{\theta}({\bm{z}}_{0}),q({\bm{z}}_{0}))\lesssim$ $\displaystyle\underbrace{\sqrt{KL(q({\bm{x}}_{T})||p_{\theta}({\bm{x}}_{T}))} \exp(-\tilde{T})}_{\text{prior distribution error}}+\underbrace{(L_{x}\sqrt{d_ {x}}l+Lm_{x}^{2}l)\sqrt{\tilde{T}}}_{\text{discretization error}}$ (25) $\displaystyle+\underbrace{\sqrt{l\sum_{t=1}^{T}{\mathcal{L}}_{t}^{(x)}}}_{ \text{coordinate score estimation error}}+\underbrace{\frac{1}{2}\mathbb{E}_{q ({\bm{x}}_{0})}{\mathcal{L}}^{(h)}({\bm{x}}_{0})}_{\text{atom type estimation error}},$ $$

where ${\mathcal{L}}^{(h)}({\bm{x}}_{0})=|\phi_{\theta}^{(h)}({\bm{x}}_{0})-f^{(h)}({
\bm{x}}_{0})|$ is the atom type prediction error, with $f^{(h)}$ providing the ground truth atom type $i$ corresponding to ${\bm{x}}_{0}$ in the dataset.

###### Proof.

$$ $\displaystyle TV(p_{\theta}({\bm{z}}_{0}),q({\bm{z}}_{0}))$ $\displaystyle=\frac{1}{2}\int|\ p_{\theta}({\bm{h}}_{0},{\bm{x}}_{0})-q({\bm{h }}_{0},{\bm{x}}_{0})\ |\ d{\bm{x}}_{0}\ d{\bm{h}}_{0}$ (26) $\displaystyle=\frac{1}{2}\int|\ p_{\theta}({\bm{h}}_{0}|{\bm{x}}_{0})p_{\theta }({\bm{x}}_{0})-q({\bm{h}}_{0}|{\bm{x}}_{0})q({\bm{x}}_{0})\ |\ d{\bm{x}}_{0} \ d{\bm{h}}_{0}$ $\displaystyle\leq\frac{1}{2}\int|\ p_{\theta}({\bm{h}}_{0}|{\bm{x}}_{0})p_{ \theta}({\bm{x}}_{0})-p_{\theta}({\bm{h}}_{0}|{\bm{x}}_{0})q({\bm{x}}_{0})\ | \ d{\bm{x}}_{0}\ d{\bm{h}}_{0}$ $\displaystyle\quad+\frac{1}{2}\int|\ p_{\theta}({\bm{h}}_{0}|{\bm{x}}_{0})q({ \bm{x}}_{0})-\delta_{(f^{(h)}({\bm{x}}_{0}))}q({\bm{x}}_{0})\ |\ d{\bm{x}}_{0} \ d{\bm{h}}_{0}$ $\displaystyle=\frac{1}{2}\int(\int p_{\theta}({\bm{h}}_{0}|{\bm{x}}_{0})\ d{ \bm{h}}_{0})|\ p_{\theta}({\bm{x}}_{0})-q({\bm{x}}_{0})\ |\ d{\bm{x}}_{0}$ $\displaystyle\quad+\frac{1}{2}\mathbb{E}_{q({\bm{x}}_{0})}\sum_{j=1}^{M}\sum_{ i=1}^{H}|\ p_{\theta}({\bm{h}}_{0,j}={\bm{e}}_{i}|{\bm{x}}_{0})-\delta_{(f^{(h )}_{j}({\bm{x}}_{0})={\bm{e}}_{i})}\ |$ $\displaystyle=TV(p_{\theta}({\bm{x}}_{0}),q({\bm{x}}_{0}))+\frac{1}{2}\mathbb{ E}_{q({\bm{x}}_{0})}{\mathcal{L}}^{(h)}({\bm{x}}_{0})$ $$

Then by using Theorem E.1 to provide the upper bound of the diffusion error of $x$, we can obtain equation 25
∎

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 a well-defined mapping also applies to the organic molecules of interest. Consequently, there exists $f$ such that for all potential coordinates ${\bm{x}}_{0}$, $q({\bm{h}}_{0}|{\bm{x}}_{0})=\delta_{(f^{(h)}({\bm{x}}_{0}))}$, providing a robust foundation for utilizing predictive methods to determine atom types.

###### Proof.

###### Proof.

### E.3 Generative Error Comparison

We compare the error bounds for the two methods by analyzing each component in the upper bounds:

a. Prior distribution error and discretization error: These errors are influenced by the dimensionality of the generated data. Since UniGEM operates on a 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: $d_{x}<d_{z}$.
- 2.
Since the distribution of ${\bm{x}}$ and ${\bm{h}}$ are independent in the forward process and so is the prior distribution of the reverse generative process, i.e.
$q({\bm{z}}_{T})=$$q({\bm{x}}_{T})$$q({\bm{h}}_{T})$, $r_{\theta}({\bm{z}}_{T})=$$r_{\theta}({\bm{x}}_{T})$$r_{\theta}({\bm{h}}_{T})$, $p_{\theta}({\bm{x}}_{T})=$$r_{\theta}({\bm{x}}_{T})$, we have
$KL(q({\bm{z}}_{T})||r_{\theta}({\bm{z}}_{T}))$$=KL(q({\bm{x}}_{T})||p_{\theta}({\bm{x}}_{T}))$$+KL(q({\bm{h}}_{T})||r_{\theta}({\bm{h}}_{T}))$. Thus $KL(q({\bm{x}}_{T})||p_{\theta}({\bm{x}}_{T}))\leq KL(q({\bm{z}}_{T})||r_{
\theta}({\bm{z}}_{T}))$ and $KL(q({\bm{z}}_{T})||r_{\theta}({\bm{z}}_{T}))$ scales with $d_{z}$.
- 3.
The second moment term satisfies $m^{2}_{z}=\mathbb{E}_{q({\bm{z}}_{0})}||{\bm{z}}_{0}||^{2}=\mathbb{E}_{q({\bm{
z}}_{0})}(||{\bm{x}}_{0}||^{2}+||{\bm{h}}||^{2})=m^{2}_{x}+m^{2}_{h}$, i.e. $m^{2}_{x}\leq m^{2}_{z}$ and $m^{2}_{z}$ scales with $d_{z}$.
- 4.
Since the distribution of $x$ and $h$ are independent in the forward process, the score function can be expressed as $\nabla_{{\bm{z}}_{t}}\ln q({\bm{z}}_{t})=(\nabla_{{\bm{x}}_{t}}\ln q({\bm{x}}_
{t}),\nabla_{{\bm{h}}_{t}}\ln q({\bm{h}}_{t}))$. As a result, the Lipschitz constant for $x$ is bounded by the Lipschitz constant for the joint variable $z$, i.e. $L_{x}\leq L_{z}$.
- 5.
Other variables $l$, $T$, $\tilde{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$\sqrt{l\sum_{t=1}^{T}{\mathcal{L}}_{t}^{(h|x)}}$ and prediction loss $\frac{1}{2}\mathbb{E}_{q({\bm{x}}_{0})}{\mathcal{L}}^{(h)}({\bm{x}}_{0})$. For coordinate estimation, we evaluate the results of two denoising losses $\sqrt{l\sum_{t=1}^{T}{\mathcal{L}}_{t}^{(x|h)}}$ and $\sqrt{l\sum_{t=1}^{T}{\mathcal{L}}_{t}^{(x)}}$. It is important to note that the input of the network differs in the two denoising losses, with ${\mathcal{L}}_{t}^{(x|h)}$ incorporating noisy atom type as input:

$$ $\displaystyle{\mathcal{L}}_{t}^{(x|h)}=\mathbb{E}_{q({\bm{z}}_{0},{\bm{z}}_{n} )}(\|{\bm{\phi}}_{\theta}^{(x)}({\bm{z}}_{t},t)-(-\sigma_{t}\nabla_{{\bm{x}}_{ t}}\log q({\bm{x}}_{t}))\|^{2}$ (27) $\displaystyle{\mathcal{L}}_{t}^{(x)}=\mathbb{E}_{q({\bm{x}}_{0},{\bm{x}}_{n})} (\|{\bm{\phi}}_{\theta}^{(x)}({\bm{x}}_{t},t)-(-\sigma_{t}\nabla_{{\bm{x}}_{t} }\log q({\bm{x}}_{t}))\|^{2}$ (28) $$

The detailed experimental comparison and
analysis is presented in Section 2.2, showing UniGem achieves lower atom types and coordinates estimation error.

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
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.
soundness: 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; 

(2) The theoretical analysis is well-developed and can explain the tradeoff phenomenon properly.
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.
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;

(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;

(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;

(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.

---

confidence: 3
contribution: 3
initial_rating: 8
presentation: 3
questions: No further questions.
soundness: 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.

2. 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.

3. Experiments validate the efficacy of the proposed framework.
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.
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.

2. 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.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: 1. How much time does it take the overall training?

2. Can EGNN be replaced by recent 3D GNNs, e.g., [4], to further improve performance?

3. Can the authors provide visualizations of the denoising process for generated samples?

[4] Liu et al., Spherical Message Passing for 3D Molecular Graphs, ICLR 2022
soundness: 3
strengths: - The problems of interest, molecular generation and property prediction, are important for real-world applications, e.g., drug discovery.

- 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.

- The findings are interesting; simultaneously learning both predictive and generative objective improve each other's performance.
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.
weaknesses: - Insufficient baselines.

My 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.

---
- Marginal improvements.

Considering the omitted baselines [1,2,3], the improvements seem marginal for me. Furthermore, it seems that [3] largely outperforms this paper. 

---
- Complexity of the framework.

UniGEM 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.

---
Overall, 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.

[1] Song et al., Equivariant Flow Matching with Hybrid Probability Transport for 3D Molecule Generation, NeurIPS 2023\
[2] Xu et al., Geometric-Facilitated Denoising Diffusion Model for 3D Molecule Generation, AAAI 2024\
[3] Irwin et al., Efficient 3D Molecular Generation with Flow Matching and Scale Optimal Transport, arXiv

Meta-review: The paper proposes a novel pretraining approach for molecule property prediction and generation that aims to better align two main types of pretraining, namely prediction-based and generation-base.

Reviewers found the proposal of integration of generation and prediction based pretraining as diffusion with deep supervision novel and promising.

The rebuttal phase has cleared all significant issues Reviewers have raised. All reviewers voted ultimately for accepting the paper. Perhaps most importantly, the rebuttal phase has cleared that the proposed approach allows achieving state of the art results, as the original version of the manuscript omitted certain baselines.

A remaining concern is that improvements remain somewhat marginal in certain benchmarks. However, the novelty of the approach warrants optimism about the potential impact of the paper.

All in all, it is my pleasure to recommend acceptance of the paper. Please make sure to address reviewers’ remarks in the camera ready version.

Decision: Accept (Poster)