PDCG: A Diffusion Model Guided by Pre-Training for Molecular Conformation Generation

Abstract

Background: While machine learning has advanced molecular conformation generation, existing models often suffer from limited generalization and inaccuracies, especially for complex molecular structures. These limitations hinder their reliability in downstream applications. Methods: We proposed a molecular conformation model combined with a molecular graph pre-training module and a diffusion model (PDCG). Feature embeddings are obtained from a pre-trained model and concatenated with the molecular graph information. Fusion features are used for generating conformations in the model. The model was trained and evaluated on the GEOM-QM9 and GEOM-Drugs datasets. Results: PDCG significantly outperforms existing baselines, which shows markedly superior results. Furthermore, in downstream molecular property prediction tasks, conformations generated by PDCG yield results comparable to those derived from DFT-optimized geometries. Conclusions: Our work provides a robust and generalizable model for accurate conformation generation. PDCG offers a reliable tool for downstream computational tasks, such as the virtual screening of functional materials and drug-like molecules.

1. Introduction

Molecular conformation generation plays a key role in many tasks, including molecular property prediction, docking, and virtual screening [1,2]. Traditional approaches, such as experimental determination and quantum chemical calculation, are expensive and time-consuming [3]. Classical force field methods offer greater computing speed but suffer from insufficient accuracy [4,5]. These limitations present significant challenges to high-throughput molecular screening and rational drug design, adversely affecting both the robustness and efficiency of these processes [6,7]. Machine learning (ML)-based algorithms for conformation generation are increasingly employed to address the compromise between computational cost and accuracy [7].

Machine learning models that predict atomic coordinates encounter intrinsic challenges, including rotational and translational invariance [8]. To address this issue and improve the stability, two principal strategies involving intermediate geometric prediction have been adopted. The first strategy predicts interatomic distances and subsequently reconstructs conformations using SE(3)-invariant networks [9]. Representative implementations include methods that predict the pairwise distance matrices and recover coordinates via distance geometry (DG) [10], as implemented in GraphDG [11] and CGCF [12]. Alternatively, more robust conformation generation can also be achieved by predicting the distance gradient in the models of ConfVAE [13] and ConfGF [14]. However, while distance-based predictions impose valuable geometric constraints, they often fail to adequately capture other critical structural features, such as the bond and torsion angles.

A second strategy has emerged to predict these local geometric attributes directly. Models like GeoMol [15] and Tora3D [16] predict local structures and torsion angles by integrating local and global molecular information. Reinforcement learning has also been employed to sample torsion angles for conformational exploration iteratively [17,18]. A fundamental limitation in both strategies is that errors in predicting intermediate geometric values propagate through the reconstruction process, often causing the final generated conformation to deviate from the actual structural distribution [19].

Directly outputting the molecular conformation from neural networks can provide a more accurate and efficient strategy for eliminating prediction errors [20]. Direct generation involves two methods, one of which is one-hot generation. The model can adaptively aggregate bond and atomic information to directly predict atomic coordinates, as DMCG [21] and Conf-GEM [22]. The conformational quality generated by these one-hot methods has reached the level of traditional methods, while it still exhibits potential for improvement in the conformational diversity of macromolecules. The other is the diffusion method, which gradually generates molecular conformations through continuous sampling. This method can directly predict coordinates without intermediate distances by iteratively processing coordinates and interatomic distances, and simultaneously achieve precise conformational modeling of rotational and translational invariance, such as COSMIC [23], GeoDiff [24], SDEGen [25], and EC-Conf [26]. The diffusion process can also be further restricted to the torsional angle space to improve performance [27]. The diffusion method demonstrates improved conformational diversity and structural validity compared to one-hot generation for macromolecules.

In recent years, to address the poor generalization ability of conformation generation models, pre-training methods have been introduced to learn general molecular representations and reduce reliance on precisely labeled data [28]. Pre-training could be achieved by several strategies such as Auto-Encoding model [29], Autoregressive Modeling [30], Masked Component Modeling [31], Context Prediction [32], Contrastive Learning [33], Replaced Components Detection [34], and DeNoising model [35]. Various pre-trained models [36,37,38,39] have been proposed to integrate chemical domain knowledge through pre-training strategies, and to enhance the generalizability of molecular embedding spaces. A combined physical–geometric constrained pre-trained model and generative models could improve generalization in conformation prediction. Wang et al. [40] proposed a novel framework that pre-trains a model on 3D molecular graphs and then fine-tunes it on molecular graphs without 3D structures. Beyond physics-aware methods, Alhamoud et al. [41] investigated the limitations of the GeoMol method and introduced pre-training to enhance molecular graph embeddings. These two methods emphasize the role of pre-training in connecting 2D and 3D representations.

At present, some generative models have limitations, including insufficient generalization and inaccurate conformation generation for macromolecules. To improve the model’s generalization and generate molecular conformations more efficiently and accurately, we propose PDCG, a model comprising a 2D molecular graph pre-training model and a molecular conformation diffusion generation model. Our model starts from the Simplified Molecular Input Row Input System (SMILES) [42] and obtains molecular features through the pre-trained model. The features are concatenated with the molecular graph information, and the fused features are then used to generate molecular conformations in the model. This model can improve the quality and diversity of the generated conformations, promote superior 3D conformation generation, and provide conformations for downstream tasks.

2. Materials and Methods

Datasets: The experiments utilized two publicly available datasets for molecular conformation generation, GEOM-QM9 and GEOM-Drugs [43]. We adopted the established data split from the work of Xu et al. for both datasets [24]. From this split, a training set was constructed by randomly sampling 40,000 molecules, each represented by its SMILES string and associated with 5 conformations. A validation set of 5000 molecules was similarly prepared. For testing, we extracted 200 distinct molecules from the remaining pool. This yielded a total of 22,408 test conformations for GEOM-QM9 and 14,324 for GEOM-Drugs. Additionally, the standard QM9 dataset was employed to evaluate the performance of downstream molecular property prediction tasks using the generated conformations. For each molecule, the conformation with the lowest MAT-P was selected for our model training and evaluation.

Network architecture: The overall architecture of the proposed PDCG framework is illustrated in Figure 1. It primarily consists of two components: a pre-trained molecular language model MoLR [38] and a diffusion-based generative model [24]. MoLR is a chemical reaction-aware molecular representation learning model based on the USPTO-479k dataset. It uses a two-layer GNN encoder to complete self-supervised pre-training through “reactant embedding = product embedding” and conservation constraints of contrastive loss, providing initialization parameters for the model, shortening the training cycle. We employed a publicly available MoLR model as a fixed feature extractor, while its parameters were kept frozen throughout our experiments. Given a molecular SMILES string as input, this model directly outputs a high-dimensional molecular representation, denoted as H_G.

The extracted representation H_G is first projected into a lower-dimensional latent space via a multi-layer perceptron (MLP), resulting in a refined feature vector h_G. Concurrently, the input SMILES is converted into a 2D molecular graph. A graph encoder processes this graph to generate a complementary feature vector E_G. Subsequently, h_G and E_G are concatenated to form a fused molecular representation G, which serves as the conditional input for the diffusion model. The fused feature G guides a denoising diffusion probabilistic model to generate plausible 3D molecular conformations. This model learns to iteratively reverse a predefined noising process, transforming an initial Gaussian distribution into a target conformational distribution conditioned on G.

During the diffusion process, starting from the given initial conformation C⁰, random noise conforming to the Gaussian distribution is gradually added. This process can be divided into step T. The initial conformation C⁰ will diffuse to the chaotic state C^t after step t. This forward process is a posterior distribution q(C^1:T|C⁰). Specifically, we define it as a Markov chain according to a fixed variance schedule β₁, …, β_t:

$$q\left(\left.C^{1:T}\right|C^0\right)=\prod _{t=1}^{T}q\left(\left.C^t\right|C^{t-1}\right),$$

(1)

where $q\left(\left.C^t\right|C^{t-1}\right)$ is defined in Equation (2),

$$q\left(\left.C^t\right|C^{t-1}\right)=N\left(C^t;\sqrt{1-{\beta }_t}{C}^{t-1},{\beta }_{t}I\right).$$

(2)

During the generation process, under the condition of concatenating the molecular graph feature G, the conformation C⁰ is learned to be restored from the white noise C^T. This process is expressed as a learnable conditional Markov chain:

$$p_{\theta }\left(\left.C^{0:T-1}\right|{G,C}^T\right)=\prod _{t=1}^{\overline{T}}{p}_{\theta }\left(\left.C^{t-1}\right|G,C^t\right),$$

(3)

$$p_{\theta }\left(\left.C^{t-1}\right|G,C^t\right)=N\left(C^{t-1};{\mu }_{\theta }\left(G,C^t,t\right),{\sigma }_t^{2}I\right),$$

(4)

where µ_θ is the parameterized neural network, σ_t is the defined variance, and the initial distribution p(CT) is set as a standard Gaussian distribution and iterated through the anti-Markov kernel p_θ.

Experiment setup: The model was implemented in PyTorch and trained on a single NVIDIA RTX 4080 GPU. The Adam optimizer was employed with a learning rate of 0.001 and a batch size of 64. The model architecture featured a global encoder with 6 convolutional layers and a local encoder with 4 convolutional layers, both with a hidden dimension of 128. For the diffusion process, the noise schedule parameters were set to a minimum variance β₁ of 1 × 10⁻⁷ and a maximum variance level β_T of 0.002. A spatial cutoff radius of 10 Å was applied for local structure modeling. The computational environment utilized Python 3.11, PyTorch 2.4.0, PyTorch Geometric 2.6.0, and RDKit 2022.09.1 [44].

Evaluation metrics: To evaluate the quality and diversity of the conformations generated by the model, we used an evaluation metric based on root mean square deviation (RMSD). The coverage (COV) and matching (MAT) scores measure diversity and accuracy, respectively. The COV score represents the percentage of the reference conformation generated by the prediction set. The MAT score represents the minimum RMSD between the generated conformation and the reference conformation. Based on the measurement of recall, COV-R and MAT-R can be defined as

$$COV-R\left(S_g,S_r\right)={\displaystyle \frac{1}{\left|S_r\right|}}C\in S_r\left|RMSD\left(C,\widehat{C}\le \sigma ,\widehat{C}\in S_g\right)\right|$$

(5)

$$MAT-R\left(S_g,S_r\right)={\displaystyle \frac{1}{\left|S_r\right|}}{\sum }_{C\mathsf{ϵ}{S}_r}\underset{\widehat{C}\mathsf{ϵ}{S}_g}{\min }\left(RMSD\left(C,\widehat{C}\right)\right)$$

(6)

S_g is the set of generated conformations, and S_r is the set of reference conformations of a molecule. σ represents the defined threshold. $\widehat{C}$ and C refer to a generated conformation and a reference conformation, respectively. Based on the measurement of precision, COV-P and MAT-P can be defined as

$$COV-P\left(S_r,S_g\right)={\displaystyle \frac{1}{\left|S_g\right|}}C\in S_g\left|RMSD\left(C,\widehat{C}\le \sigma ,\widehat{C}\in S_r\right)\right|,$$

(7)

$$MAT-P\left(S_g,S_r\right)={\displaystyle \frac{1}{\left|S_g\right|}}{\sum }_{C\mathsf{ϵ}{S}_g}\underset{\widehat{C}\mathsf{ϵ}{S}_r}{\min }\left(RMSD\left(C,\widehat{C}\right)\right).$$

(8)

A higher COV score or a lower MAT score indicates the generation of a more realistic conformation. Recall indicators focus more on diversity, while precision indicators rely more on quality.

Validation via a downstream task. The practical utility of the conformations was assessed by their ability to predict molecular properties using SphereNet [45]. Predictions were made using three distinct sets of input conformations for comparison: quantum-chemically optimized structures, force field-optimized (RDKit) structures, and structures generated by PDCG. SphereNet was implemented with the same parameters as in its original work. The accuracy for each property was quantified by calculating the mean absolute error (MAE) between the model’s predictions and the reference values.

3. Results and Discussion

3.1. The Performance with Different Diffusion Steps

A key hyperparameter in the diffusion framework is the number of diffusion steps. We examined its influence by training models with different step counts on the GEOM-QM9 and GEOM-Drugs datasets. Each model was evaluated on an independent test set of 200 molecules, with the iteration step size of the generation phase the same as that of the training phase.

For the GEOM-QM9 dataset, we evaluated the performance of the diffusion model with maximum diffusion step sizes of 3000, 5000, 7000, and 10,000 in the forward process. The iteration steps are the same in the generation and training phases. As summarized in

Table 1. Results of the diffusion model with different diffusion steps on the GEOM-QM9 dataset ¹.

Diffusion Steps	Sampling Steps	COV–R (%) ↑2	MAT–R (Å) ↓3	COV–P (%) ↑	MAT–P (Å) ↓
3000	3000	0.877	0.354	0.477	0.551
5000	5000	0.919	0.194	0.528	0.445
7000	7000	0.762	0.288	0.493	0.518
10,000	10,000	0.677	0.393	0.431	0.515

¹ The value with best performance is in bold. ² ↑ means higher is better. ³ ↓ means lower is better.

, an increase from 3000 to 5000 steps improved all evaluation metrics, with the latter yielding the best overall performance. Beyond this optimum, further increases in the step size degrade performance. This decline may be attributed to the accumulation of prediction errors across sequential sampling steps, which drives the generated structure away from the actual data distribution. Consequently, 5000 steps were established as the optimal setting for subsequent experiments, which is the same as the diffusion model GeoDiff [24].

For the GEOM-Drugs dataset, we investigated model performance across a range of maximum diffusion steps (6000, 8000, 10,000, and 12,000) in the forward process, with sampling steps held constant during generation. As shown in

Table 2. Results of diffusion model with different diffusion steps on the GEOM-Drugs dataset ¹.

Diffusion Steps	Sampling Steps	COV–R (%) ↑2	MAT–R (Å) ↓3	COV–P (%) ↑	MAT–P (Å) ↓
6000	6000	0.901	0.795	0.687	1.286
8000	8000	0.904	0.785	0.689	1.083
10,000	10,000	0.916	0.770	0.752	0.988
12,000	12,000	0.805	1.065	0.663	1.358

¹ The value with best performance is in bold. ² ↑ means higher is better. ³ ↓ means lower is better.

, optimal performance across all four evaluation metrics was achieved at 10,000 steps. This configuration yielded the best balance between recall and precision, indicating superior overall performance for this more complex molecular set. Consequently, 10,000 steps were established as the optimal forward process setting for training on GEOM-Drugs. This contrasts with the 5000-step setting used in the GeoDiff model [24], which was kept consistent across both the GEOM-QM9 and GEOM-Drugs datasets.

3.2. Comparison with Baseline Methods

We evaluated PDCG on a diverse set of seven baselines, encompassing one-hot encoding (GraphDG, ConfVAE), iterative generation (CGCF, Conf-GF, SDE-Gen, EC-Conf), and a pre-training-enhanced model (GeoMol + MRL). The results show the advantages of the proposed method across key metrics.

For the GEOM-QM9 dataset, the results are summarized in

Table 3. Comparison of our model with the baseline model on the GEOM-QM9 datasets ¹.

Models	COV–R (%) ↑2	MAT–R (Å) ↓3	COV–P (%) ↑	MAT–P (Å) ↓
GraphDG	0.733	0.425	0.439	0.581
ConfVAE	0.778	0.415	0.380	0.622
CGCF	0.781	0.422	0.662	0.661
CONFGF	0.885	0.267	0.522	0.464
SDEGen	0.815	0.357	0.484	0.566
EC-Conf	0.813	0.324	0.794	0.330
GeoMol + MRL	0.826	0.298	0.837	0.310
Diffusion	0.919	0.194	0.528	0.445
PDCG	0.934	0.188	0.752	0.369

¹ The value with best performance is in bold. ² ↑ means higher is better. ³ ↓ means lower is better.

. Among all baseline methods, the standard diffusion model achieved the good recall-oriented metrics, with a COV-R of 0.919 and an MAT-R of 0.194. Our model sets a new state-of-the-art in recall, with scores of 0.934 and 0.188, which exceed the best baseline (diffusion model). Conversely, on precision-oriented metrics, the GeoMol + MRL model delivered the strongest performance, with scores of 0.837 (COV-P) and 0.310 (MAT-P), and is the best among all models. Our model obtained values of 0.752 and 0.369 on these metrics, slightly weaker than GeoMol + MRL. This comparative analysis suggests that incorporating pre-trained representations, as in GeoMol + MRL, provides a significant advantage in improving generation precision. As shown in Figure S2, we benchmarked the generation time for a representative molecule. GeoMol + MLR exhibits a significant efficiency advantage, which is primarily attributed to the substantially fewer denoising steps required by GeoMol + MLR.

We further evaluated our model on the more challenging GEOM-Drugs dataset, which includes macromolecular structures. As summarized in

Table 4. Comparison of our model with the benchmark model on the GEOM-Drugs datasets ¹.

Model	COV–R (%) ↑2	MAT–R (Å) ↓3	COV–P (%) ↑	MAT–P (Å) ↓
GraphDG	0.083	1.972	0.021	2.434
ConfVAE	0.552	1.238	0.230	1.829
CGCF	0.540	1.249	0.217	1.857
CONFGF	0.622	1.163	0.234	1.722
SDEGen	0.673	1.126	0.323	1.679
EC-Conf	0.864	0.902	0.701	1.108
GeoMol + MRL	0.815	1.132	0.756	1.045
Diffusion	0.916	0.770	0.752	0.988
PDCG	0.929	0.712	0.766	0.945

¹ The value with best performance is in bold. ² ↑ means higher is better. ³ ↓ means lower is better.

. For recall metrics, it achieves a COV–R of 0.929 and an MAT–R of 0.712, which exceeds the best baseline diffusion model with scores of 0.916 and 0.770, respectively. On precision metrics, our model has a COV–P of 0.766 and an MAT–P of 0.945, which is the highest among all the compared methods. These results indicate a clear advantage in generating accurate macromolecular conformations. Therefore, the effective mapping from 2D graphs to precise 3D conformations validates our core design. The diffusion model ensures extensive generative coverage of the conformational space, while the pre-trained component injects essential prior knowledge to ensure geometric stability. This combined approach is particularly effective at generating plausible conformations for large, flexible drug-like molecules.

3.3. RMSD of Conformation Generation

To evaluate performance, we analyzed the RMSD of generated 3D conformations in the GEOM-QM9 dataset and its correlation with molecular size and complexity. RMSD values exhibited a positive correlation with the number of heavy atoms (Figure 2a). The median RMSD of atomic positions increased from 0.2 Å for molecules with a single heavy atom to 2.1 Å for those with nine. This increase was accompanied by a broadening of the interquartile range (IQR) and whiskers, indicating greater conformational heterogeneity in larger molecules. In contrast, RMSD showed minimal dependence on the number of rings (Figure 2b). The median RMSD remained consistently between 1.5 and 2.0 Å across molecules containing zero to seven rings, with negligible IQR variation. This stability suggests that ring systems impose conformational constraints that limit structural deviation.

Bond distances exhibited a monotonic increase with the number of heavy atoms (Figure 2c), with the median rising from 0.07 Å for single-atom molecules to 0.10 Å for those with nine atoms. This was accompanied by a progressive expansion of the interquartile range (IQR) and overall range, reflecting greater bond length variability in larger frameworks. In contrast, bond distances decreased as the ring count increased (Figure 2d). The highest median bond distance (0.11 Å) was observed in acyclic molecules (zero rings), which steadily declined to 0.06 Å for molecules with seven rings, indicating a ring-induced bond-shortening effect.

Bond angles exhibited a non-linear dependence on heavy atom count (Figure 2e). The median angle peaked at 12° for monatomic systems (one atom), dropped to 3–5° for molecules with 2–5 atoms, and subsequently rose to 10° for those with nine atoms. This trend highlights the distinct angular geometries adopted by small versus large molecular assemblies. For ring-containing systems (Figure 2f), bond angles increased with ring count, rising from a median of 6° for acyclic molecules (zero rings) to 15° for molecules with five rings, after which they plateaued. The interquartile range (IQR) expanded markedly for molecules with three or more rings, which underscores the role of ring systems in inducing angular distortion.

Collectively, these results show that molecular size, measured by heavy atom count, is the primary driver of conformational flexibility and bond elongation. In contrast, ring systems play a secondary but counterbalancing role by constraining geometric variability and promoting shorter and more rigid bond metrics. These trends were consistent across all statistical descriptors.

3.4. Properties Prediction by SphereNet

To qualitatively assess the accuracy of generated conformations for downstream tasks, we evaluated molecular property predictions based on conformations derived from DFT and our proposed PDCG.

As shown in

Table 5. Comparison of property prediction by SphereNet used conformation DFT optimization and our model.

Property	Unit	DFT	PDCG
μ	D	0.0245	0.0330
α	$a_0^3$	0.0449	0.0510
ϵHOMO	eV	0.0228	0.0329
ϵLUMO	eV	0.0189	0.0216
ϵgap	eV	0.0313	0.0511
<R2>	$a_0^2$	0.2680	0.4367
zpve	eV	0.0011	0.0020
Cv	cal/(mol·K)	0.0215	0.0255
U0	eV	0.0063	0.0634
U	eV	0.0064	0.0450
H	eV	0.0063	0.0568
G	eV	0.0078	0.0447

, the mean absolute errors (MAEs) of properties predicted from PDCG-generated conformations are very close to those obtained from DFT reference conformations. Compared to RDKit, PDCG achieves lower MAEs across most properties, particularly for the energy-related labels (U₀, U, H, and G). These results indicate that the conformation generated by PDCG could improve the precision of downstream predictive tasks.

4. Conclusions

This work introduces PDCG, a diffusion-driven framework for generating 3D molecular conformations combined with a molecular graph pre-training module and a diffusion model. The pre-training module used 2D molecular graphs to extract useful embeddings from SMILES representations. These embeddings and graph features were proceeded in a diffusion module to yield 3D conformations. Comprehensive evaluations reveal that PDCG achieves statistically significant enhancements in both COV and MAT metrics on a benchmark dataset of QM9-Drugs. Downstream validation through molecular property prediction confirms that the conformational accuracy of PDCG is closed to that of DFT-computed structures. Consequently, PDCG presents a promising and efficient alternative to support key computational chemistry applications, such as structure-based virtual screening and the design of functional materials. Future developments of PDCG could add the conditional generation of conformations to discover functional molecules with the desired properties. The model improvement for accurate multi-conformer generation is also a priority for our future research.