Text-guided small molecule generation via diffusion model

Summary

The de novo generation of molecules with targeted properties is crucial in biology, chemistry, and drug discovery. Current generative models are limited to using single property values as conditions, struggling with complex customizations described in detailed human language. To address this, we propose the text guidance instead, and introduce TextSMOG, a new text-guided small molecule generation approach via diffusion model, which integrates language and diffusion models for text-guided small molecule generation. This method uses textual conditions to guide molecule generation, enhancing both stability and diversity. Experimental results show TextSMOG’s proficiency in capturing and utilizing information from textual descriptions, making it a powerful tool for generating 3D molecular structures in response to complex textual customizations.

Graphical abstract

Highlights

• •Combines language models and diffusion models into text-guided 3D molecule generation
• •Multi-modal conversion module translates text conditions into reference geometries
• •Outperforms existing methods, yielding more stable and diverse subject molecules
• •Initializes understanding general text descriptions to explore molecular space

Chemistry; Molecular modelling; Artificial intelligence

Introduction

De novo molecule design, the process of generating molecules with specific, chemically viable structures for target properties, is a cornerstone in the fields of biology, chemistry, and drug discovery.^1,2,3,4 It not only allows for the creation of subject molecules but also provides insights into the relationship between molecular structure and function, enabling the prediction and manipulation of biological activity. Constrained by the immense diversity of chemical space, manually generating property-specific molecules remains a daunting challenge.⁵ However, the generation of molecules that precisely meet specific requirements, including the creation of tailor-made molecules, is a complex task due to the vastness of the chemical space and the intricate relationship between molecular structure and function. Overcoming this challenge is crucial for advancing our understanding of biological systems and for the development of new therapeutic agents. In recent years, machine and deep learning methods have initiated a paradigm shift in the molecule generation,^{6,7,8,9,10,11} which enable the direct design of 3D molecular geometric structures with the desired properties.^8,12,13 Notably, diffusion models,^14,15 specifically equivariant diffusion models,^16,17 have gradually entered the center of the stage with its outstanding performance. The core of this method is to introduce diffusion noise on molecular data, and then learn a reverse process in either unconditional or conditional manners to denoise this corruption, thereby crafting desired 3D molecular geometries. Meanwhile, some conditional inputs (e.g., polarizability $\alpha =100{\text{Bohr}}^3$) could be applied for constraining the model to generate more specific molecules types.

However, despite the promise of these methods, a significant proportion of molecules generated by diffusion models do not meet the practical needs of researchers. For instance, they may lack the desired biological activity, exhibit poor pharmacokinetic properties, or be synthetically infeasible. This would be due to the fact that, on one hand, searching for suitable molecules in drug design typically requires consideration of multiple properties of interest (e.g., simultaneously characterized by specific polarizability, orbital energy, properties like aromaticity, and distinct functional groups).^18,19,20 On the other hand, humans seem to struggle with conveying their needs precisely to the model. While a text segment such as “this molecule is an aromatic compound, with small HOMO-LUMO gaps and possessing at least one carboxyl group” can accurately describe human requirements and facilitate communication among humans, it is still challenging to directly convey this “thoughts” to the model. Therefore, we aspire to develop a method that allows for the interactive inverse design of 3D molecular structures through natural language. In other words, we aim to create a system where researchers can describe the properties they want in a molecule using natural language, and the system will generate a molecule that meets these requirements. This aspiration prompts us to explore text guidance in diffusion models, emphasizing the necessity for models adept at precise language understanding and molecule generation.

Toward this end, we propose TextSMOG, a new text-guided small molecule generation approach. The basic idea is to combine the capabilities of the advanced language models^{21,22,23,24,25,26} with high-fidelity diffusion models, enabling a sophisticated understanding of textual prompts and accurate translation into 3D molecular structures. TextSMOG accomplishes this through integrating textual information with a conversion module that conditions a pre-trained equivariant diffusion model (EDM)¹⁶ following the multi-modal fusion fashion.^9,27,28,29 Specifically, at each denoising step, TextSMOG first generates reference geometry, an intermediate conformation that encapsulates the textual condition signal, through a multi-modal conversion module. Equipped with language and molecular encoder-decoder, corresponding to the textual condition. Then the reference geometry guides the denoising of each atom within the pre-trained unconditional EDM, gradually modifying the molecular geometry to match the condition while maintaining chemical validity. By incorporating valuable language knowledge into the pre-trained diffusion model, TextSMOG enhances the generation of valid and stable 3D molecular conformations that align with a spectrum of diverse directives. This is achieved without the need for exhaustive training on each specific condition, demonstrating the model’s ability to generalize from the language input. This integration allows for the incorporation of valuable language knowledge in the high-fidelity pre-trained diffusion model, thereby enabling the conditional generation contingent upon a spectrum of diverse directives, while enhancing the generation of valid and stable 3D molecular conformations, without specific exhaustive training of the condition.

We applied TextSMOG to the standard quantum chemistry dataset QM9³⁰ and a real-world text-molecule dataset from PubChem.³¹ The experimental results show that TextSMOG accurately captures single or multiple desired properties from textual descriptions, thereby aligning the generated molecules with the desired structures. Notably, TextSMOG outperforms leading diffusion-based molecule generation baselines (e.g., EDM,¹⁶ EEGSDE¹⁷) in terms of both the stability and diversity of the generated molecules. This is evidenced by higher scores on metrics such as the Tanimoto similarity to the target structure, the synthetic accessibility of the generated molecules, and the diversity of the generated molecule set. Furthermore, when applied to real-world textual excerpts, TextSMOG demonstrates its generative capability under general textual conditions. These findings suggest that TextSMOG constitutes a versatile and efficient text-guided molecular diffusion framework. As an advanced intelligent agent, it can effectively comprehend the meaning of textual commands and accomplish generation tasks, thereby paving the way for a more in-depth exploration of the molecular space.

Results

In this section, we present the architecture and the experimental results of our proposed TextSMOG model, showcasing its ability to generate molecules with desired properties.

Architecture

To evaluate our model, we employ the QM9 dataset,³⁰ which is a standard benchmark containing quantum properties and atom coordinates of over 130K molecules, each with up to 9 heavy atoms (C, N, O, F). For the purpose of training our model under the condition of textual descriptions, we have curated a subset of molecules from QM9 and associated them with real-world descriptions. These descriptions are sourced from PubChem,³¹ one of the most comprehensive databases for molecular descriptions, and are linked to the molecules in QM9 based on their unique SMILES.

PubChem aggregates extensive annotations from a diverse array of sources, such as ChEBI,³² LOTUS,³³ and T3DB.³⁴ Each of these sources offers an emphasis on the physical, chemical, or structural attributes of molecules. Additionally, we have employed a set of textual templates to generate corresponding descriptions based on the quantum properties of the molecules, thereby enriching the content of the dataset and supplementing textual context for those molecules lacking real-world descriptions. This process has enriched QM9 into a dataset of chemical molecule-textual description pairs. Our proposed TextSMOG model, illustrated in Figure 1, is built upon the pre-trained unconditional diffusion model EDM.¹⁶ It integrates the textual information into the conditional signal of diffusion models by employing a reference geometry that is updated at each step based on the textual prompt. The final molecular geometry is generated by gradually denoising an initial geometry, while noise is added at each step during the forward process until the molecular geometry is fully noise-corrupted.

Figure 1.

Architecture of our text-guided small molecule generation via diffusion model (TextSMOG)

The model starts with an initial geometry $\left({\mathcal{G}}_T\right)$ and gradually denoises it to generate the final molecular geometry. The reference geometry $\left({\mathbf{c}}_{\mathbf{P}}\right)$, updated at each step based on the textual prompt $\left(\mathbf{P}\right)$, is employed to integrate the textual information into the conditional signal of diffusion models. Flame denotes tunable modules, while snowflake indicates frozen modules.

Experiment on single quantum properties conditioning

Following EDM,¹⁶ we first evaluate our TextSMOG on the task of generating molecule conditioning on a single desired quantum property in QM9. Then we compare our TextSMOG with several baselines to demonstrate the effectiveness of our model on single quantum properties conditioning molecule generation.

Setup

We follow the same data preprocessing and partitions as in EDM,¹⁶ which results in 100K/18K/13K molecule samples for training/validation/test, respectively. In order to assess the quality of the conditional generated molecules w.r.t. to the desired properties, we use the property classifier network ${\varphi }_p$ introduced by.³⁵ Then for the impartiality, the training partition is further split into two non-overlapping halves ${\mathbb{D}}_a$ and ${\mathbb{D}}_b$ of 50K molecule samples each. The property classifier network ${\varphi }_p$ is trained on the first half ${\mathbb{D}}_a$, while our TextSMOG is trained on the second half ${\mathbb{D}}_b$. This ensures that there is no information leak and the property classifier network ${\varphi }_p$ is not biased toward the generated molecules from TextSMOG. Then ${\varphi }_p$ is evaluated on the generated molecule samples from TextSMOG as we introduce in the following.

Metrics

Following,¹⁶ we use the mean absolute error (MAE) between the properties of generated molecules and the ground truth as a metric to evaluate how the generated molecules align with the condition (see the supplemental information for details). We generate 10K molecule samples for the evaluation of ${\varphi }_p$, following the same protocol as in EDM. Additionally, we then measure novelty,³⁶ atom stability,¹⁶ and molecule stability¹⁶ to demonstrate the fundamental molecule generation capacity of the model (also see the supplemental information for details).

Baseline

We compare our TextSMOG with a direct baseline conditional EDM¹⁶ and a recent work EEGSDE which takes energy as guidance.¹⁷ We also compare two additional baselines “U-bound” and “#Atoms” introduced by.¹⁶ In the “U-bound” baseline, any relation between molecule and property is ignored, and the property classifier network ${\varphi }_p$ is evaluated on ${\mathbb{D}}_b$ with shuffled property labels. In the “#Atoms” baseline, the properties are predicted solely based on the number of atoms in the molecule. Furthermore, we report the error of ${\varphi }_p$ on ${\mathbb{D}}_b$ as a lower bound baseline “L-Bound”.

Results

We generate molecules with textual descriptions targeted to each one of the six properties in QM9, which are detailed in the supplemental information. As presented in Figure 2, our TextSMOG has a lower MAE than other baselines on five out of the six properties, suggesting that the molecules generated by TextSMOG align more closely with the desired properties than other baselines. The result underscores the proficiency of TextSMOG in exploiting textual data to guide the conditional de novo generation of molecules. Moreover, it highlights the superior congruence of the text-guided molecule generation via the diffusion model with the desired property, thus showing significant potential. Furthermore, as indicated in Figure 3, our proposed TextSMOG exhibits commendable performance in terms of novelty and stability. The text guidance we introduced has transformed the exploration of the model in the molecule generation space, generally enhancing the novelty of the generated molecules while maintaining their stability.

Figure 2.

Comparison of MAE for the generated molecules targeted to desired property

Statistics of baselines are from their original papers. The performance of EEGSDE varies depending on the scaling factor, and we report its best results. The numerical values are provided in the supplemental information.

Figure 3.

Comparison of novelty (Novel, %), atom stability (A. Stable,%), and molecule stability (M. Stable,%) on generated molecules targeted to the desired property

Statistics of baselines are from EEGSDE. The performance of EEGSDE varies depending on the scaling factor, and we report its best results.

Experiment on multiple quantum properties conditioning

The capacity to generate molecules, guided by multiple conditions, is a crucial aspect of the molecule generation model. When guided by textual descriptions, characterizing the condition with multiple desired properties is highly intuitive and flexible. Following the same setup and metrics in the previous section, we evaluate our TextSMOG on the task of generating molecules with multiple desired quantum properties in QM9. Then we compare TextSMOG with two baselines to showcase the effectiveness of our model in generating molecules conditioned on multiple quantum properties.

As shown in Table 1, our TextSMOG has a remarkably lower MAE than the other two baselines, thereby demonstrating the superiority of our model in generating molecules with multiple desired properties. This also further substantiates that, without necessitating additional targeted interventions, textual conditions can be utilized in our model to guide molecule generation that conforms to multiple desired properties.

Table 1.

Comparison of MAE on the generated molecules targeted to the multiple desired properties

Method	MAE1 $\downarrow$	MAE2 $\downarrow$
Condition	$C_v\left(\frac{\text{cal}}{\text{mol}}\mathrm{K}\right)$ and $\mu \left(\mathrm{D}\right)$
EDM	1.079 $\pm 0.007$	1.156 $\pm 0.011$
EEGSDE	0.981 $\pm 0.008$	0.912 $\pm 0.006$
TextSMOG	0.645$\pm 0.014$	0.836$\pm 0.017$
Condition	$\alpha \left({\text{Bohr}}^3\right)$ and $\mu \left(\mathrm{D}\right)$
EDM	2.76 $\pm 0.01$	1.158 $\pm 0.002$
EEGSDE	2.61 $\pm 0.01$	0.855 $\pm 0.007$
TextSMOG	2.27$\pm 0.01$	0.809$\pm 0.010$
Condition	$\Delta \epsilon \left(\text{meV}\right)$ and $\mu \left(\mathrm{D}\right)$
EDM	683 $\pm 1$	1.130 $\pm 0.007$
EEGSDE	563 $\pm 3$	0.866 $\pm 0.003$
TextSMOG	489$\pm 4$	0.843$\pm 0.009$

Statistics of baselines are from EEGSDE. Boldface indicates the best performance.

Additionally, as highlighted in Table 2, our proposed TextSMOG maintains superior performance in terms of novelty and stability, when generating molecules targeted at multiple desired properties. The results indicate that the flexible integration of multiple conditions through textual description does not compromise the stability of the generated molecules. Furthermore, this approach enhances novelty when compared to the baseline.

Table 2.

Comparison of novelty (Novel, %), atom stability (A. Stable, %), and molecule stability (M. Stable,%) on the generated molecules targeted to the multiple desired properties

Method	Novel $\uparrow$	A. Stable $\uparrow$	M. Stable $\uparrow$
Condition	$C_v\left(\frac{\text{cal}}{\text{mol}}\mathrm{K}\right)$ and $\mu \left(\mathrm{D}\right)$
EDM	85.31 $\pm 0.43$	98.00$\pm 0.07$	77.42$\pm 0.80$
EEGSDE	85.62 $\pm 0.86$	97.67 $\pm 0.08$	74.56 $\pm 0.54$
TextSMOG	85.79$\pm 0.66$	97.89 $\pm 0.10$	77.33 $\pm 0.72$
Condition	$\alpha \left({\text{Bohr}}^3\right)$ and $\mu \left(\mathrm{D}\right)$
EDM	85.06 $\pm 0.27$	97.96 $\pm 0.00$	75.95 $\pm 0.30$
EEGSDE	85.56 $\pm 0.56$	97.61 $\pm 0.04$	72.72 $\pm 0.27$
TextSMOG	85.64$\pm 0.64$	98.01$\pm 0.07$	75.97$\pm 0.44$
Condition	$\Delta \epsilon \left(\text{meV}\right)$ and $\mu \left(\mathrm{D}\right)$
EDM	85.18 $\pm 0.35$	98.00 $\pm 0.06$	77.96 $\pm 0.33$
EEGSDE	85.36 $\pm 0.03$	97.99 $\pm 0.06$	77.77 $\pm 0.26$
TextSMOG	85.44$\pm 0.41$	98.06$\pm 0.04$	78.03$\pm 0.29$

Statistics of baselines are from EEGSDE. Boldface indicates the best performance.

Generation on general textual descriptions

To further assess our model, we undertake additional training on a vast dataset of over 330K text-molecule pairs we gleaned from PubChem.³¹ Then, we generate molecules based on general textual descriptions to observe the capacity of our model to generate from generalized textual conditions.

Visual observations, as depicted in Figure 4, illuminate the impressive aptitude of our TextSMOG in aligning molecule structures with the desired property within the textual descriptions. For instance, when the textual description includes affirmatively mentioned terms such as “simple chain structure”, “at least one carboxyl group”, and “soluble in water”, the generated molecules consistently exhibit chain structures with at least one carboxyl group, and characteristics indicative of water solubility.

Figure 4.

Generated molecules targeted to text description excerpts

Moreover, when the textual description includes “polycyclic heteroarene” and specifies the solubility and heat capacity of the molecule, TextSMOG generates a variety of polycyclic aromatic hydrocarbon molecules. The ubiquitously present amino and nitro groups attest to a certain degree of solubility of the molecules. Referring to structurally similar molecules, their expected specific heat capacity is also relatively low.

Lastly, when the text description explicitly demands multiple nitrogen atoms and a low energy gap, the molecules generated by TextSMOG not only possess the required polycyclic structure and multiple nitrogen atoms, but also the rings on the same plane denote the low-energy structures of these molecules that are difficult to excite.

The remarkable alignment between the conditions and the generated molecule stands as a testament to the exceptional generative capabilities of TextSMOG. The result demonstrates that TextSMOG is equipped to deeply explore the chemical molecular space in a text-guided manner, thereby generating prospective molecules for subsequent applications. This capability could potentially expedite drug design and the discovery of materials.

The results highlight TextSMOG’s versatility in generating a wide variety of molecular structures, from simple chain structures to complex polycyclic compounds, under the guidance of general text descriptions. This underscores the model’s potential to perform well even when the conditions deviate significantly from the distribution of the training set.

Discussion

The translational impacts of TextSMOG are particularly significant for the field of drug discovery and materials science. By enabling the generation of molecular structures directly from textual descriptions, TextSMOG can streamline the early stages of drug design where rapid prototyping and iterative testing are crucial. This approach can facilitate the discovery of subject drug candidates by allowing researchers to quickly generate and evaluate provided molecules based on specific desired properties mentioned in literature or derived from expert knowledge.

Furthermore, TextSMOG can aid in the development of materials with tailored properties by generating molecules that meet specific criteria. This capability is valuable in industries such as polymers, nanomaterials, and catalysts, where precise molecular structures can significantly influence material performance.

In drug discovery, TextSMOG’s ability to generate molecules that align with complex textual prompts can accelerate the identification of compounds with potential therapeutic effects. This is especially relevant for targeting diseases with well-characterized biochemical pathways, where detailed descriptions of molecular interactions and desired properties are available. By generating candidate molecules that meet these criteria, TextSMOG can help narrow down the pool of potential drugs, reducing the time and cost associated with experimental validation.

Additionally, TextSMOG’s flexibility in handling diverse textual inputs can facilitate interdisciplinary research, where insights from different fields can be integrated into the molecule generation process. For instance, combining insights from biology, chemistry, and pharmacology can lead to more informed and effective drug design strategies.

Despite the complexity of translating textual prompts into accurate molecular structures, we have successfully integrated advanced language models with high-fidelity diffusion models in TextSMOG, a text-guided diffusion approach for 3D molecule generation. Our experiments on the QM9 and PubChem datasets demonstrate the superior performance of TextSMOG over leading baselines, affirming its efficacy in capturing desired properties from textual descriptions and generating corresponding valid molecules.

Limitations of the study

The integration of textual information with the denoising process of a pre-trained EDM allows TextSMOG to generate valid and stable molecular conformations that closely align with diverse textual directives. This initial success paves the way for significant advancements in the exploration of chemical space and the development of compounds. Nevertheless, our findings are not without limitations.

Our work was constrained by the scarcity of high-quality data linking real-world 3D molecules to their corresponding textual descriptions. This limitation impacted our ability to fully train the model on a diverse set of text-3D molecule pairs, potentially affecting the accuracy of the generated molecules in generating molecules that accurately align with complex textual descriptions. Moreover, the relative slowness of the sampling process due to the iterative nature of the total diffusion steps can pose a challenge in scenarios requiring rapid molecule generation, such as high-throughput drug discovery or material design.

In addition to these limitations, the current design of TextSMOG necessitates that the properties to condition on must be known upfront during the training phase. This might not always be feasible in practical settings, where specific properties linked to a particular drug discovery target may only become available later on, and often with very limited sample data. The generalization of TextSMOG to more complex and real-world scenarios also needs further exploration.

Looking ahead, we are optimistic about the potential of text-guided 3D molecule generation to revolutionize drug discovery and related fields. Future work will focus on overcoming these challenges by expanding and enhancing the quality of datasets linking textual descriptions to molecular structures, improving the efficiency of the sampling process, and making TextSMOG more adaptable to real-world applications. Addressing these limitations will not only enhance the performance of TextSMOG but also contribute significantly to the advancement of text-guided molecule generation technology.

Resource availability

Lead contact

Further information and requests for resources and reagents should be directed to and will be fulfilled by the lead contact, Wenjie Du (duwenjie@mail.ustc.edu.cn).

Materials availability

This study did not generate new unique reagents.

Data and code availability

• •The datasets generated during this study are available at HuggingFace and are publicly available as of the date of publication. The url is listed in the key resources table.
• •All original code has been deposited at GitHub and is publicly available as of the date of publication. The url is listed in the key resources table.
• •Any additional information required to reanalyze the data reported in this paper is available from the lead contact upon request.

Acknowledgments

This research is supported by the National Natural Science Foundation of China (92270114). This research was also supported by the advanced computing resources provided by the Supercomputing Center of the USTC.

Author contributions

Conceptualization, Y.L. and J.F.; methodology, Y.L. and S.L.; investigation, Z.L. and S.L.; writing—original draft, Y.L. and J.F.; writing—review and editing, J.W., A.Z., and W.D.; funding acquisition, X.W.; resources, J.W., A.Z., and W.D.; supervision, X.W.

Declaration of interests

The authors declare no competing interests.

STAR★Methods

Key resources table

REAGENT or RESOURCE	SOURCE	IDENTIFIER
Deposited data
Processed PubChem Dataset	This paper, PubChem	https://pubchem.ncbi.nlm.nih.gov/
QM9 Dataset	Quantum Machine	http://quantum-machine.org/datasets/
Software and algorithms
TextSMOG Model	This paper	https://github.com/lyc0930/TEDMol

Method details

In this section, we elaborate on the proposed text-guided small molecule generation approach via diffusion model (TextSMOG), as illustrated in Figure 1. It integrates the textual information (i.e., text guidance) into the conditional signal of diffusion models by employing the reference geometry that is described in the first subsection following. Subsequently, we introduce an efficient learning approach that incorporates both the encoded conditional signal and pre-trained unconditional signal in the reverse process, to generate molecules that are not only structurally stable and chemically valid but also align well with the specified conditions, as presented in the second subsection.

Notation and background

We begin with a background of diffusion-based 3D molecule generation, introducing the fundamental concepts of the diffusion model and delving into equivariant diffusion models. See the comprehensive literature review on these topics in the Section Related works in supplemental information. In accordance with prior studies,^12,16,17 we use the variable $\mathcal{G}=\left(\mathbf{x},\mathbf{h}\right)$ to represent the 3D molecular geometry. Here $\mathbf{x}=\left(x_1,\dots ,x_M\right)\in {\mathbb{R}}^{M\times 3}$ signifies the atom coordinates, while $\mathbf{h}=\left(h_1,\dots ,h_M\right)\in {\mathbb{R}}^{M\times k}$ denotes the atom features. These features encompass atom types and atom charges, characterizing the atomic properties within the molecular structure.

Diffusion model

The diffusion model^14,15 emerges as a leading generative model, having achieved great success in various domains.^{59,60,61,62,63,64} Typically, it is formulated as two Markov chains: a forward process (aka. noising process) that gradually injects noise into the data, and a reverse process (aka. denoising process) that learns to recover the original data. Such a reverse process endows the diffusion model with enhanced capabilities for effective data generation and recovery.

Forward Process

Given the real 3D molecular geometry ${\mathcal{G}}_0$, the forward process yields a sequence of intermediate variables ${\mathcal{G}}_1,\dots ,{\mathcal{G}}_T$ using the transition kernel $q\left({\mathcal{G}}_t|{\mathcal{G}}_{t-1}\right)$ in alignment with a variance schedule ${\beta }_1,{\beta }_2,\dots ,{\beta }_T\in \left(0,1\right)$. Formally, it is expressed as:

$$q\left({\mathcal{G}}_t|{\mathcal{G}}_{t-1}\right)=\mathcal{N}\left({\mathcal{G}}_t|\sqrt{1-{\beta }_t}{\mathcal{G}}_{t-1},{\beta }_t{\mathbf{I}}_n\right)\text{,}$$ (Equation 1)

where $\mathcal{N}\left(·|·,·\right)$ is a Gaussian distribution and ${\mathbf{I}}_n$ is the identity matrix. This defines the joint distribution of ${\mathcal{G}}_1,\dots ,{\mathcal{G}}_T$ conditioned on ${\mathcal{G}}_0$ using the chain rule of the Markov process:

$$q\left({\mathcal{G}}_1,\dots ,{\mathcal{G}}_T|{\mathcal{G}}_0\right)=\prod _{t=1}^{T}q\left({\mathcal{G}}_t|{\mathcal{G}}_{t-1}\right)\text{.}$$ (Equation 2)

Let ${\alpha }_t=1-{\beta }_t$ and ${\overline{\alpha }}_t:=\prod _{s=1}^t{\alpha }_s$. The sampling of ${\mathcal{G}}_t$ at time step $t$ is in a closed form:

$$q\left({\mathcal{G}}_t|{\mathcal{G}}_0\right)=\mathcal{N}\left({\mathcal{G}}_t|\sqrt{{\overline{\alpha }}_t}{\mathcal{G}}_0,\left(1-{\overline{\alpha }}_t\right){\mathbf{I}}_n\right)\text{.}$$ (Equation 3)

Accordingly, the forward process posteriors, when conditioned on ${\mathcal{G}}_0$, are tractable as:

$$q\left({\mathcal{G}}_{t-1}|{\mathcal{G}}_t,{\mathcal{G}}_0\right)=\mathcal{N}\left({\mathcal{G}}_{t-1}|\tilde{\mu }\left({\mathcal{G}}_t,{\mathcal{G}}_0\right),{\tilde{\beta }}_t{\mathbf{I}}_n\right)\text{,}$$ (Equation 4)

where

$$\tilde{\mu }\left({\mathcal{G}}_t,{\mathcal{G}}_0\right)=\frac{\sqrt{{\overline{\alpha }}_{t-1}}{\beta }_t}{1-{\overline{\alpha }}_t}{\mathcal{G}}_0+\frac{\sqrt{{\alpha }_t}\left(1-{\overline{\alpha }}_t\right)}{1-{\overline{\alpha }}_t}{\mathcal{G}}_t,{\tilde{\beta }}_t=\frac{1-{\overline{\alpha }}_{t-1}}{1-{\overline{\alpha }}_t}{\beta }_t\text{.}$$ (Equation 5)

Reverse Process

To recover the original molecular geometry ${\mathcal{G}}_0$, the diffusion model starts by generating a standard Gaussian noise ${\mathcal{G}}_T\sim \mathcal{N}\left(\mathbf{O},{\mathbf{I}}_n\right)$, then progressively eliminates noise through a reverse Markov chain. This is characterized by a learnable transition kernel $p_{\theta }\left({\mathcal{G}}_{t-1}|{\mathcal{G}}_t\right)$ at each reverse step $t$, defined as:

$$p_{\theta }\left({\mathcal{G}}_{t-1}|{\mathcal{G}}_t\right)=\mathcal{N}\left({\mathcal{G}}_{t-1}|{\mu }_{\theta }\left({\mathcal{G}}_t,t\right),{\mathrm{\Sigma }}_{\theta }\left({\mathcal{G}}_t,t\right)\right)\text{,}$$ (Equation 6)

where the variance ${\mathrm{\Sigma }}_{\theta }\left({\mathcal{G}}_t,t\right)={\tilde{\beta }}_t{\mathbf{I}}_n$ and the mean ${\mu }_{\theta }\left({\mathcal{G}}_t,t\right)$ is parameterized by deep neural networks with parameters $\theta$:

$${\mu }_{\theta }\left({\mathcal{G}}_t,t\right)={\tilde{\mu }}_t\left({\mathcal{G}}_t,\frac{1}{\sqrt{{\overline{\alpha }}_t}}\left({\mathcal{G}}_t-\sqrt{1-{\overline{\alpha }}_t}{ϵ}_{\theta }\left({\mathcal{G}}_t,t\right)\right)\right)=\frac{1}{\sqrt{{\alpha }_t}}\left({\mathcal{G}}_t-\frac{1-{\alpha }_t}{\sqrt{1-{\overline{\alpha }}_t}}{ϵ}_{\theta }\left({\mathcal{G}}_t,t\right)\right)\text{,}$$ (Equation 7)

where ${ϵ}_{\theta }$ is a noise prediction function to approximate the noise $ϵ$ from ${\mathcal{G}}_t$.

With the reverse Markov chain, we can iteratively sample from the learnable transition kernel $p_{\theta }\left({\mathcal{G}}_{t-1}|{\mathcal{G}}_t\right)$ until $t=1$ to estimate the molecular geometry ${\mathcal{G}}_0$.

Equivariant diffusion models

The molecular geometry $\mathcal{G}=\left(\mathbf{x},\mathbf{h}\right)$ is inherently symmetric in 3D space — that is, translating or rotating a molecule does not change its underlying structure or features. Previous studies^65,66,67 underscore the significance of leveraging these invariances in molecular representation learning for enhanced generalization. However, the transformation of these higher-order representations usually requires computationally expensive approximations or coefficients.^16,35 In contrast, equivariant diffusion models^16,68,69 provide a more efficient approach to ensure both rotational and translational invariance. The approach rests on the assumption that, with the model distribution $p\left(\mathcal{G}\right)=p\left(\mathbf{x},\mathbf{h}\right)$ remaining invariant to the Euclidean group E(3), identical molecules, despite being in different orientations, will correspond to the same distribution. Based on this assumption, translational invariance is achieved by predicting only the deviations in coordinate with a zero center of mass, i.e., $\sum _{i=1}^M{x}_i=0$. On the other hand, rotational invariance is accomplished by making the noise prediction network ${ϵ}_{\theta }(·)$ equivariant to orthogonal transformations.^16,35 Specifically, given an orthogonal matrix $\mathbf{R}$ representing a coordinate rotation or reflection, the conformation output $a^{\mathbf{x}}$ from the network ${ϵ}_{\theta }\left(\mathcal{G}\right)={ϵ}_{\theta }\left(\mathbf{x},\mathbf{h}\right)=\left(a^{\mathbf{x}},a^{\mathbf{h}}\right)$ is equivariant to $\mathbf{R}$, if the following condition holds for all orthogonal matrices $\mathbf{R}$:

$${ϵ}_{\theta }\left(\mathbf{\text{Rx}},\mathbf{h}\right)=\left(\mathbf{R}{a}^{\mathbf{x}},a^{\mathbf{h}}\right)\text{.}$$ (Equation 8)

A model exhibiting rotational and translational equivariance means a neural network $p_{\theta }\left(\mathcal{G}\right)$ can avoid learning orientations and translations of molecules from scratch.^16,35 In this paper, we parameterize the noise prediction network ${ϵ}_{\theta }$ using an E(n) equivariant graph neural network as introduced by,³⁵ which is a type of Graph Neural Network⁷⁰ that satisfies the above equivariance constraint to E(3).

Equivariant diffusion model for molecule generation

Diffusion models, formulated as two Markov chains—a forward process that gradually injects noise into the data and a reverse process that learns to recover the original data—have been successfully applied to various domains, including molecule generation. This process is particularly effective in the context of molecule generation, where the forward process adds noise to the molecular geometry at each step until it is fully noise-corrupted. The reverse process then gradually denoises the initial geometry ${\mathcal{G}}_T$ to generate the final molecular geometry ${\mathcal{G}}_0$.

However, molecular geometries are inherently symmetric in 3D space—translations or rotations do not change their underlying structure or features. To take advantage of these invariances for improved generalization, we employ an equivariant diffusion model (EDM). The EDM ensures both rotational and translational invariance by predicting only the deviations in coordinate with a zero center of mass and making the noise prediction network ${ϵ}_{\theta }(·)$ equivariant to orthogonal transformations. This allows the model distribution $p\left(\mathcal{G}\right)$ to remain invariant to the Euclidean group E(3), meaning identical molecules in different orientations correspond to the same distribution.

In this work, the integration of textual information into the conditional signal of the equivariant diffusion model is achieved by employing a reference geometry ${\mathbf{c}}_{\mathbf{P}}$ that is updated at each step based on the textual prompt $\mathbf{P}$.

Integrating textual prompts into 3D molecular reference geometry

To ensure high-fidelity 3D molecule generation, the reverse process of the diffusion model is typically guided by tailored conditional information representing desired properties like unique polarizability. We represent this conditional information as $c$, which allows us to formulate the conditional reverse process as:

$$p_{\theta }\left({\mathcal{G}}_{t-1}|{\mathcal{G}}_t,c\right)=\mathcal{N}\left({\mathcal{G}}_{t-1}|{\mu }_{\theta }\left({\mathcal{G}}_t,c,t\right),{\tilde{\beta }}_t{\mathbf{I}}_n\right)$$ (Equation 9)

Unlike previous approaches relying on limited value guidance (i.e., property values), in this work, we aim to steer the reverse process with text guidance (i.e., informative textual descriptions), which can convey a broader range of conditional requirements. Intuitively, utilizing textual descriptions to specify conditional generation criteria not only provides greater expressivity but also better aligns the resulting 3D molecules with diverse and complex expectations.

Practically, we first introduce a textual prompt P describing desired 3D molecule properties. A multi-modal conversion module $\mathrm{\Gamma }$, pre-trained on 300K text-molecule pairs from PubChem, is then employed. This module is comprised of a GIN molecular graph encoder^71,72 and a language encoder-decoder extended from BERT.^21,73 It converts P into a reference geometry ${\mathbf{c}}_{\mathbf{P}}$, extracting specific information from the target conditions and refining the textual condition signal:

$${\mathbf{c}}_{\mathbf{P}}=\mathrm{\Gamma }\left(\mathbf{P}\right)\text{.}$$ (Equation 10)

Nevertheless, we should emphasize that valid and stable 3D molecules can hardly be obtained directly from ${\mathbf{c}}_{\mathbf{P}}$. The chemical fidelity in 3D molecular space may not be guaranteed. In what follows, we describe how to utilize ${\mathbf{c}}_{\mathbf{P}}$ for conditioning a pre-trained diffusion model to generate molecules that align with the desired properties, meanwhile alleviating the exhaustive training from scratch.

Conditioning with the reference of text guidance

To leverage ${\mathbf{c}}_{\mathbf{P}}$ for text-guided conditional generation while preserving the validity and stability of the synthesized molecule, TextSMOG employs the iterative latent variable refinement (ILVR)⁷⁴ to condition a pre-trained unconditional diffusion model meanwhile maintaining inherent domain knowledge in the unconditional model.

With the pre-trained unconditional diffusion model EDM,¹⁶ we could perform a step-by-step reverse process. Formally, at step $t$, we can sample an unconditional proposal molecular geometry:

$${\tilde{\mathcal{G}}}_{t-1}\sim {\tilde{p}}_{\tilde{\theta }}\left({\tilde{\mathcal{G}}}_{t-1}|{\mathcal{G}}_t\right)\text{.}$$ (Equation 11)

where $\tilde{\theta }$ is the fixed parameters of the pre-trained unconditional diffusion model.¹⁶ Then, to incorporate the condition signal ${\mathbf{c}}_{\mathbf{P}}$ in the reverse process, we introduce a linear operation ${\varphi }_{\theta }(·)$. Therefore the conditional denoising for one step at step $t$ can be formulated as:

$${\mathcal{G}}_{t-1}={\varphi }_{\theta }\left({\mathbf{c}}_{\mathbf{P}}\right)+\left(\mathcal{I}-{\varphi }_{\theta }\right)\left({\tilde{\mathcal{G}}}_{t-1}\right)\text{,}$$ (Equation 12)

where $\mathcal{I}(·)$ is the identity operation and $\left(\mathcal{I}-{\varphi }_{\theta }\right)(·)$ is the residual operation w.r.t. ${\varphi }_{\theta }(·)$⁷⁵. Accordingly, the condition signal ${\mathbf{c}}_{\mathbf{P}}$ is projected into the reverse denoising process by ${\varphi }_{\theta }(·)$, thus ${\mathcal{G}}_{t-1}$ is obtained as the generated 3D molecular geometry conditioned on ${\mathbf{c}}_{\mathbf{P}}$. Conceptually, the proposal geometry from unconditional generation ${\tilde{\mathcal{G}}}_{t-1}$ tries to push the atoms into a chemically valid position, while the reference geometry ${\mathbf{c}}_{\mathbf{P}}$ pulls the atoms toward the structure targeted to the condition.

By matching latent variables following Equation 12, we enable text-guided conditional generation with the unconditional diffusion model. Accordingly, the one-step denoising distribution conditioned on textual guidance at each step $t$ can be reformulated as:

$${\mathcal{G}}_{t-1}\sim p_{\theta }\left({\mathcal{G}}_{t-1}|{\mathcal{G}}_t,{\mathbf{c}}_{\mathbf{P}}\right)\text{.}$$ (Equation 13)

Training objective

To guarantee the quality of the generated molecules, the key lies in optimizing the variational lower bound (ELBO) of negative log likelihood, which equals minimizing the Kullback-Leibler divergence between the joint distribution of the reverse Markov chain $p_{\theta }\left({\mathcal{G}}_0,{\mathcal{G}}_1,\dots ,{\mathcal{G}}_T\right)$ and the forward process $q\left({\mathcal{G}}_0,{\mathcal{G}}_1,\dots ,{\mathcal{G}}_T\right)$:

$$\mathbb{E}\left[-\log p_{\theta }\left({\mathcal{G}}_0|{\mathbf{c}}_{\mathbf{P}}\right)\right]\le -\log \sum _{t\ge 1}\underset{:={\mathcal{L}}_{t-1}}{\underbrace{D_{\text{KL}}\left(q\left({\mathcal{G}}_{t-1}|{\mathcal{G}}_t,{\mathcal{G}}_0\right)\Vert p_{\theta }\left({\mathcal{G}}_{t-1}|{\mathcal{G}}_t,{\mathbf{c}}_{\mathbf{P}}\right)\right)}}+C\text{,}$$ (Equation 14)

where $C$ is a constant independent of $\theta$.

Note that we set ${\mathcal{L}}_0=-\log p_{\theta }\left({\mathcal{G}}_0|{\mathcal{G}}_1\right)$ as a discrete decoder following.¹⁵ Further adopting the reparameterization from,¹⁵ ${\mathcal{L}}_{t-1}$ can be simplified to:

$${\mathcal{L}}_{t-1}={\mathbb{E}}_{\mathrm{P},{\mathcal{G}}_0,ϵ}\left[{\Vert ϵ-{ϵ}_{\theta }\left(\sqrt{{\overline{\alpha }}_t}{\mathcal{G}}_0+\sqrt{1-{\overline{\alpha }}_t}ϵ,t,{\mathbf{c}}_{\mathbf{P}}\right)\Vert }^2\right]\text{.}$$ (Equation 15)

Evaluation metrics

Mean absolute error (MAE).⁷⁶ is a measure of errors between paired observations. Given the property classifier network ${\varphi }_p$, and the set of generated molecules $\mathbb{G}$, the MAE is defined as:

$$\text{MAE}=\frac{1}{\left|\mathbb{G}\right|}\sum _{\mathcal{G}\in \mathbb{G}}\left|{\varphi }_p\left(\mathcal{G}\right)-c_{\mathcal{G}}\right|\text{,}$$ (Equation 16)

where $\mathcal{G}$ is the generated molecule, and of which $c_{\mathcal{G}}$ is the desired property.

Novelty.³⁶ is the proportion of generated molecules that do not appear in the training set. Specifically, let $\mathbb{G}$ be the set of generated molecules, the novelty in our experiment is calculated as:

$$\text{Novelty}=\frac{\left|\mathbb{G}\mathbb{\cap }{\mathbb{D}}_b\right|}{\left|\mathbb{G}\right|}\text{.}$$ (Equation 17)

Atom stability.¹⁶ is the proportion of the atoms in the generated molecules that have the right valency. Specifically, the atom stability in our experiment is calculated as:

$$\text{AtomStability}=\frac{{\sum }_{\mathcal{G}\in \mathbb{G}}\left|{\mathbb{A}}_{\mathcal{G},\text{stable}}\right|}{{\sum }_{\mathcal{G}\in \mathbb{G}}\left|{\mathbb{A}}_{\mathcal{G}}\right|}\text{,}$$ (Equation 18)

where ${\mathbb{A}}_{\mathcal{G}}$ is the set of atoms in the generated molecule $\mathcal{G}$, and ${\mathbb{A}}_{\mathcal{G},\text{stable}}$ is the set of atoms in ${\mathbb{A}}_{\mathcal{G}}$ that have the right valency.

Molecule stability.¹⁶ is the proportion of the generated molecules where all atoms are stable. Specifically, the molecule stability in our experiment is calculated as:

$$\text{MoleculeStability}=\frac{\left|{\mathbb{G}}_{\text{stable}}\right|}{\left|\mathbb{G}\right|}\text{,}$$ (Equation 19)

where ${\mathbb{G}}_{\text{stable}}$ is the set of generated molecules where all atoms have the right valency.

Quantification and statistical analysis

The quantum properties in QM9 dataset

We consider 6 main quantum properties in QM9:

• •$C_v$: Heat capacity at 298.15K.
• •$\mu$: Dipole moment.
• •$\alpha$: Polarizability, which represents the tendency of a molecule to acquire an electric dipole moment when subjected to an external electric field.
• •${\epsilon }_{\text{HOMO}}$: Highest occupied molecular orbital energy.
• •${\epsilon }_{\text{LUMO}}$: Lowest unoccupied molecular orbital energy.
• •${\Delta }_{\epsilon }$: The energy gap between HOMO and LUMO.

Published: September 19, 2024

Supporting citations

The following references appear in the Supplemental Information:^{37,38,39,40,41,42,43,44,45,46,47,48,49,50,51,52,53,54,55,56,57,58}.