Generating molecular structures with the desired properties is a critical task with broad applications in drug discovery and materials design (Hajduk & Greer, 2007;Mandal et al., 2009;Pyzer-Knapp et al., 2015). There is a growing interest in generative models to automate this task (You et al., 2018;Jin et al., 2018;2020;Bjerrum & Threlfall, 2017). Kusner et al. (2017) proposed a variational autoencoder that uses parse trees for probabilistic context free grammars to encode discrete sequences that specify the Simplified Molecular-Input Lineentry System (SMILES) (Weininger, 1988) notation of molecules use the resulting parse trees to generate SMILES codes that describe molecular graphs. Edwards et al. (2021) formulated molecule retrieval from textual description as a cross-lingual retrieval task.
Recent advances in language models have led to innovative approaches to generate molecular structures that match human-friendly textual descriptions of their properties, substructures, and biochemical activity (Edwards et al., 2022;Zeng et al., 2022;Zhao et al., 2023b;Fang et al., 2023a). These approaches learn mappings from textual descriptions to molecular structures using single transformer-based cross-lingual language model for molecule generation. Specifically, with the given text description of a molecule, these approaches produce SMILES strings that encode the molecular structure.
Despite promising initial results, such approaches suffer from some important limitations: (i) SMILES encoding is degenerate in that the a given molecular structure can have multiple SMILE encodings. For example, the structure of ethanol can be specified using CCO, OCC or C(O)C. Although in principle a unique canonical SMILES representation can be produced for a given molecular structure, the procedure used to do so relies on several arbitrary choices. Not surprisingly, molecules with identical structure may map to different SMILES strings. This makes SMILES notation less than ideal for generating molecules from their textual descriptions (Jin et al., 2018). Important chemical properties corresponding to the substructures of molecules are more straightforward to express using molecular graphs rather than linear SMILES representations (Jin et al., 2018;Du et al., 2022); (ii) Using language models to generate diverse and high-quality molecular structures matching a description is challenging, because the process used to decode the learned representation into the corresponding molecular structure, e.g., beam search or greedy search, often produces structures that are often too similar to each other (See Figure 2). Hence, there is an urgent need for better approaches for generating diverse, novel molecular structures from a given textual description (Brown et al., 2019). Against this background, this paper aims to answer the following research question: How can we generate a diverse and novel collection of molecular graphs that match a given textual description? We introduce 3M-Diffusion, a novel multi-modal molecular diffusion method, for generating high-quality, diverse, novel molecular structures from a textual description.
Inspired by text-guided image generation using a diffusion model on the latent space (Rombach et al., 2022), 3M-Diffusion is trained to generate molecular structures from the latent distribution of a molecular graph autoencoder (See Figure 1). This enables the diffusion process to focus on the high-level semantics of encodings of molecular structures. We find that this approach is particularly well-suited for discrete molecular graph modality because it relegates the challenge of modeling a discrete and diverse distribution to the autoencoder and simplifies the diffusion process by restricting it to the continuous, latent feature space. Since graph data consists of both node attributes and graph structure, which differs significantly from the sequential text data, to align the latent representations of textual descriptions and molecular graphs, 3M-diffusion employs contrastive learning on a large data set of pairs of molecular structures and their textual descriptions to pretrain molecular graph encoders and LLM text encoders. This enables 3M-diffusion to generate high-quality molecular graphs that match the given textual descriptions.
The key contributions of this paper are as follows: (i) To the best of our knowledge, 3M-Diffusion offers the first multimodal diffusion approach to generating molecular structures from their textual descriptions, surpassing the limitations of stateof-the-art (SOTA) methods for this problem. (ii) 3M-Diffusion aligns the latent spaces of molecular graphs and textual descriptions to offer a text-molecule aligned latent diffusion model to generate higher-quality, diverse, and novel molecular structures that match the textual description provided; and (iii) Results of extensive experiments using four real-world text-based molecular graphs generation benchmarks show that 3M-Diffusion outperforms SOTA methods for both text-guided and unconditional molecular structure generation. In particular, 3M-Diffusion achieves 146.27% novelty, 130.04% diversity relative improvement over the SOTA method with maintained semantics in textual prompt on the PCDes dataset.
Existing molecular graph generation methods can be broadly grouped into two categories: text-based and graph-based models. Text-based models typically utilize the SMILES
Figure 2: Qualitative comparisons to the MolT5-large of generated molecules on ChEBI-20.
Compared with the SOTA method MolT5-large, our results are more diverse and novel with maintained semantics in textual prompt. More results are provided in Appendix D.4 (Weininger, 1988) or SELFIES (Krenn et al., 2022) linear strings to describe each molecule (Bjerrum & Threlfall, 2017;G ómez-Bombarelli et al., 2018;Kusner et al., 2017;Flam-Shepherd et al., 2022;Fang et al., 2023b;Grisoni, 2023). Because a given molecular graph can have several linear string representations and small changes to the string representation can result in large changes in molecular graph being described, such linear string encodings are far from ideal for learning generative models for producing molecular graph from textual descriptions. Inspired by the success of learning enhanced graph representations using graph based models (Kipf & Welling, 2016;You et al., 2020;Xiao et al.;2024), several authors have employed deep generative models, including graph variational autoencoders (VAEs) (You et al., 2018;Jin et al., 2018;Xiao et al., 2021;Jin et al., 2020;Kong et al., 2022;Liu et al., 2018;Diamant et al., 2023), normalizing flows (Madhawa et al., 2019;Zang & Wang, 2020;Luo et al., 2021), generative adversarial networks (Guarino et al., 2017;Maziarka et al., 2020), diffusion models (Niu et al., 2020;Vignac et al., 2022;Jo et al., 2022) and autoregressive models (You et al., 2018;Xiao & Wang, 2021;Popova et al., 2019;Goyal et al., 2020) to learn distributions of molecular graphs from large databases of known graphs. However, many of these studies focus on generating molecular graphs with specific low-level properties, such as logP (the octanol-water partition coefficient) as opposed to those that match high-level textual descriptions of a broader range of molecular properties (Edwards et al., 2022), e.g., water solubility, chemical activity, etc. However, practical applications, such as drug design, call for effective approaches to generating diverse and novel molecular graphs that match a given textual description of such high-level properties.
Advances in deep learning and language models have inspired a growing body of work on applications of such models to problems in the molecular sciences (Schwaller et al., 2019;Xiao et al., 2023;Toniato et al., 2021;Xiao et al., 2024;Zhao et al., 2023a;Seidl et al., 2023;Schwaller et al., 2021;Vaucher et al., 2020). Advances in text-guided generation of images, videos, and audios (Radford et al., 2021;Rombach et al., 2022;OpenAI, 2023;Yin et al., 2023), together with the flexibility offered by natural language descriptions of molecular graphs have inspired a growing body of work on text-guided generation of molecular graphs (Christofidellis et al., 2023;Edwards et al., 2022;Fang et al., 2023a;Zeng et al., 2022;Su et al., 2022). For example, MolT5 (Edwards et al., 2021;2022) and ChemT5 (Christofidellis et al., 2023) pre-train on a substantial volume of unlabeled language text and SMILES strings for tasks such as molecule captioning and text-based molecule generation (Christofidellis et al., 2023;Edwards et al., 2022). However, as noted above, the degenerate nature of linear SMILES string representation of molecular graphs makes it less ideal for generating molecules from their textual descriptions. To address this problem, we propose 3M-Diffusion, a novel method that aligns the latent representation of textual descriptions of molecules with that of the corresponding molecular structures, yielding a powerful approach to text-guided molecular structure generation.
We consider generative modeling of 2D molecular graphs. Each molecule is represented
x be the node attribute matrix, where x i is the D x -dimensional one-hot encoding feature vector of v i , such as atomic type and chirality type. Similarly, a tensor A ∈ R |V |×|V |×b groups the one-hot encoding e ij of each edge, where each entry is a distinct edge type (bond types for molecule), with the absence of an edge encoded explicitly using a designated edge type. Given the dataset D = {(G i , T i ) N i=1 } where molecule G has specific text descriptions T characterizing the features of the molecule, our goal is to learn conditional generation models p θ (G | T ) for generating novel, diverse, and valid molecules while also aligning with the desired features outlined in text descriptions T . For brevity, in what follows, we omit the input subscript i when it is clear from context.
Diffusion models (DMs) (Ho et al., 2020;Song et al., 2020) are latent variable models that represent the data z 0 as Markov chains z T • • • z 0 , where intermediate variables share the same dimension. DMs involve a forward diffusion process q (z 1:T | z 0 ) = ∏ T t=1 q (z t | z t-1 ) that systematically adds Gaussian noise to samples and a reverse process p θ (z 0:T ) = p (z T ) ∏ T t=1 p θ (z t-1 | z t ) that iteratively "denoises" samples from the Gaussian distribution to generate samples from the data distribution. A formal description of diffusion models is given in the Appendix A. To ensure training stability, we can use a simple regression objective with a denoising network ẑθ with input (x t , t):
where t is the time step, α t ∈ [0, 1] is the noise schedule and ϵ ∼ N (0, I) is Gaussian noise and λ t is a time-dependent weighting term. Intuitively, the denoising network is trained to denoise a noisy state, z t = √ α t z 0 + √ 1α t ϵ, aiming to reconstruct the clean data z 0 with Equation (1) that prioritizes specific times t. After training, we can draw samples with ẑθ by the iterative ancestral sampling (Ho et al., 2020):
where α t|t-1 represents α t /α t-1 and σ(t -
The sampling chain is initialized from Gaussian prior z T ∼ p (z T ) = N (z T ; 0, I).
3M-Diffusion (See Figure 1) aims to learn a probabilistic mapping from the latent space of textual descriptions to a latent space of molecular graphs for text-guided generation of molecular graphs. However, direct mapping between these two latent representations using a diffusion model faces a key hurdle, namely, large mismatch between the latent spaces of textual descriptions and of molecular graphs. To overcome this hurdle, 3M-Diffusion adopts a two-step approach. The first step uses contrastive learning to produce text-molecule aligned variational autoencoder that align the representation of molecular graphs with that of their textual descriptions, and a decoder that maps the latent representation back to the corresponding molecular graphs. The second step uses the text-aligned latent representation of molecular graphs to learn a conditional generative model, that maps textual descriptions to latent representations of molecular graphs. As we will see later, 3M-Diffusion can produce diverse novel molecular graphs that match the given textual descriptions.
The proposed text-molecule aligned variational autonecoder comprises of three parts: the molecular graph encoder E g , LLM encoder E t and molecular graph decoder D. To bridge the representation gap between molecular graphs and their textual descriptions, the text and molecular graph encoders are trained on a large-scale dataset of molecule-text pairs. These encoders use contrastive learning to construct well-aligned text-molecular structure encodings that align the latent representation of molecular structures by aligning it with with that of their textual descriptions. The details of this process are given below.
Molecular Graph Encoder aims to map each molecular graph G with the adjacency tensor A and node features matrix X into a lower-dimensional latent space Z. This transformation maps the discrete space of graph structures and node attributes into a continuous latent space. Specifically, we use Graph Isomorphism Network (GIN) (Hu et al., 2019) as encoder network on the input adjacency tensor A and node features X of graph G to derive the mean and standard deviation of the variational marginals:
where µ g and σ g represent the mean and standard deviations, respectively. Then, the latent graph representation z is sampled from a Gaussian distribution as z ∼ q (z | A, X) ≃ N µ g , σ g . This graph encoding process can be expressed as z = E g (G) signifying the outcome of the encoding operation in the input context of G.
LLM Encoder maps the textual descriptions of molecular structures into a latent space that provides input to a conditional generative model of molecular graphs. This encoded information is instrumental in generating molecules that match a given textual description.
To inject potentially useful scientific knowledge from the literature into E t we initialize it with encoder-based LLM, i.e., Sci-BERT (Beltagy et al., 2019), which is pretrained on the text of scientific publications.
Aligning the representation of molecular graphs with that of their textual descriptions is crucial for establishing a shared representation for effective text-guided molecular graph generation. Graph data encodes both node attributes and graph structure which differs significantly from the sequential and contextual nature of text data. We aim to align the molecular representation from the molecule graph encoder E g with text representation from the LLM encoder E t using contrastive learning. Consider a molecular graph G and its textual description T . To align these two encoders and ensure a cohesive latent space, we employ the following contrastive learning loss over a batch B ∈ D of large-scale dataset containing molecule-text pairs to pretrain encoders:
where
) and c j = E t (T j ) are representations of corresponding molecule graph and texts, respectively. cos(•, •) denotes cosine similarity and τ is the temperature hyperparameter. The encoders are pretrained on molecule-text (300K) pairs from PubChem Liu et al. (2023). As we will see, this process yields a well-aligned moleculetext space that can support the generation of diverse and novel molecular structures from their textual descriptions.
Molecular Graph Decoder learns to use latent representation z to generate the corresponding molecular structure. In this work, we employ the HierVAE (Jin et al., 2020) decoder. Note however that our approach can be combined with other graph decoders (Li et al., 2018;Kong et al., 2022;Grover et al., 2019). Specifically, the molecular graph decoder generates a molecular structure G using a graph representation z from the encoder. Given the molecular graph G, we minimize the negative Evidence Lower Bound (ELBO) (Jin et al., 2020) to train both the encoder and the decoder:
Here Ĝ = D(z), in which D denotes the decoder parameterized by HierVAE (Jin et al., 2020). KL[q(•)||p(•)] is the Kullback-Leibler (KL) divergence between q(•) and p(•). We follow Jin et al. (2020) and set the weight of KL loss α = 0.1. Here, p(z) = N (z | 0, I) is the isotropic multivariate Gaussian prior, where I is the identity matrix. We optimize this first expectation term concurrently using the reparameterization trick (Kingma & Welling, 2013).
We next introduce latent diffusion for conditional molecular graph generation, a method that leverages the aligned molecule-text latent space of encoder-decoder described above, enabling us to preserve the crucial molecular properties and aligning the textual description of G. Specifically, we learn a probabilistic mapping from the text to the text aligned molecular structure latent space, and thus generate molecular structures that match the conditional text inputs.
The denoising network ẑθ focuses on generating latent graph representation z conditioned on the text representation c from the LLM encoder E t . The objective to be optimized is:
where λ t are time-dependent weights and we use MLPs for the network ẑθ . Specifically, we concatenate c to graph latent representation z t for timestep t as the input of ẑθ . The denoising network is trained to denoise a noisy latent, z t = √ α t z + √ 1α t ϵ to the clean z. c is the text representation from the fixed encoder E t pretrained by Equation (4).
Additionally, to improve sample quality we take advantage of classifier-free guidance (Ho & Salimans, 2022), to jointly train an unconditional network, ẑθ (z t , t), and a conditional network, ẑθ (z t , t, c). Conditioning information is randomly dropped with a probability of p = 0.1 during training. When conditioning information is dropped, we replace the embedded source text with its embedding.
Training. We proceed to describe the training procedure for 3M-Diffusion. While the learning objectives for graph encoder-decoder and multi-modal latent diffusion have already been specified by Equations ( 5) and ( 6), it is unclear whether the two components should be trained sequentially. Previous research on latent diffusion models for image generation indicates that a two-stage training strategy often results in superior performance (Rombach et al., 2022). Hence, we adopt a similar two-stage strategy. The first stage consists of training an encoder-decoder pair using Equation (5), and the second stage consists of training a multi-modal molecule latent diffusion with the fixed encoder-decoder from the first stage.
Inference. During inference, we use a weight w in computing the prediction as follows (Ho & Salimans, 2022)
where zt represents the newly predicted results, incorporating both conditional and unconditional information. Setting w = 1.0 corresponds to the conditional diffusion model, while setting w > 1.0 enhances the impact of the conditioning information. A smaller value of w contributes to more diverse samples. The new sampling equation is: where the sampling chain is initialized from Gaussian prior z T ∼ p (z T ) = N (z T ; 0, I). The detailed training and inference algorithms of 3M-diffusion are provided in Appendix B.
Datasets. We use PubChem (Liu et al., 2023), ChEBI-20 (Edwards et al., 2021), PCDes (Zeng et al., 2022) and Momu (Su et al., 2022) datasets. Following the evaluation methodology adopted with datasets (ZINC250K (Irwin et al., 2012) and QM9 (Blum & Reymond, 2009; Rupp et al., 2012)) used in previous studies of molecular structure generation, we limit our training, validation, and test data to molecules with fewer than 30 atoms. Statistics of datasets are shown in Table 5 and detailed descriptions are given in the Appendix C.1.
Evaluation. We evaluate the model's performance on two tasks: text-guided and unconditional molecule generation. Following previous works on text-guided molecule generation, we compare our 3M-Diffusion the SOTA baselines (Edwards et al., 2022;Christofidellis et al., 2023;Fang et al., 2023a). Text-guided molecule generation measures the model's capacity to generate chemically valid, structurally diverse, novel molecules that match the textual description provided. As in previous work (Chen et al., 2021;Fu et al., 2021;Wang et al., 2022), we adopt the following metrics for this task: Similarity, Novelty, Diversity, Validity. We call a predicted molecular structure, specifically one that is chemically valid (g ′ ), qualified with the ground truth graph g if f (g, g ′ ) > 0.5, where f measures the MACCS (Durant et al., 2002) cosine similarity between the two. If the f (g, g ′ ) similarity between g ′ and g is smaller than a threshold (0.8 in this paper), we call this predicted molecule g novel. Specifically, (i) Similarity: the percentage of qualified molecules out of 20K molecules; (ii) Novelty: the percentage of novel molecules out of all qualified molecules; (iii) Diversity: the average pairwise distance 1f (•, •) between all qualified molecules. (iv) Validity: the percentage of generated molecules that are chemically valid out of all molecules.
We report the results for unconditional molecule generation tasks using the GuacaMol benchmarks (Brown et al., 2019) to assess the model's ability to learn a generative model of molecular structures that can be used to generate realistic, diverse molecules. Specifically, Uniqueness gauges the ratio of unique molecules among the generated ones. Novelty measures the models' capacity to generate molecules not present in the training set. KL Divergence assesses the proximity between the distributions of various physicochemical properties in the training set and the generated molecules. Fréchet ChemNet Distance (FCD) computes the proximity of two sets of molecules based on their hidden representations in ChemNet (Preuer et al., 2018). Each metric is normalized to a range of 0 to 1, with a higher value indicating better performance.
Baselines. For text guided molecule generation, we compare 3M-Diffusion with the following baselines: MolT5 (Edwards et al., 2022), ChemT5 (Christofidellis et al., 2023) and Mol-Instruction (Fang et al., 2023a). We consider seven representative baselines for unconditional molecule generation based on Variational Autoencoder, including CharRNN (Segler et al., 2018), VAE (Kingma & Welling, 2013), AAE (Makhzani et al., 2015), LatentGAN (Prykhodko et al., 2019) BwR (Diamant et al., 2023), HierVAE (Jin et al., 2020) and PS-VAE (Kong et al., 2022). The details of the baselines are given in Appendix C.1.
We randomly initialize the model parameters and use the Adam optimizer to train the model. HierVAE is used as our variational autoencoder for all of the following experiments. For baselines of unconditional generation that did not report results in these datasets, we reproduce the results using the official code made available by the authors.
To ensure a fair comparison across all methods for unconditional generation, we select the best hyperparameter configuration solely based on the loss on the training set. We evaluate the performance of baselines for text-guided molecule generation using the pretrained parameters made available by the authors of the respective studies. Table 1 and 2 report Similarity, Novelty, Diversity, and Validity for text-guided molecule generation on four datasets. Each metric is computed as an average over 5 structures generated from a given textual description.
We see that 3M-Diffusion is competitive with SOTA baselines with respect to Similarity and outperforms them substantially with respect to Diversity and Novelty. Specifically, 3M-Diffusion achieves 146.27% (Novelty of PCDes), 130.04% (Diversity of PCDes) relative improvement over the second-best SOTA baseline while matching the textual prompt. We conclude that while the transformer based models that are pretrained on a large dataset excel in producing molecular structures that match the textual descriptions provided, 3M-Diffusion outperforms them with respect to the diversity and novelty of the molecular structures without sacrificing the match between the textual descriptions provided as input and the molecular structures generated as output. Also, as Figure 3 shows, 3M-Diffusion enjoys a substantial advantage in terms of its speed over transformer-based methods (See Appendix D.1 for details).
Table 3 shows the results of unconditional molecule generation on PubChem and ChEBI-20. 3M-Diffusion consistently outperforms or matches the performance of SOTA baselines across all four metrics on both datasets, suggesting that it is able to generate realistic, diverse, and novel molecules without overfitting the training data. Furthermore, 3M-Diffusion consistently outperforms HierVAE when both methods use the same decoder. This underscores the clear advantage of 3M-Diffusion in generating a diverse and novel molecular graphs (See Figure 10 for 50 samples from the learned distribution).
We conduct ablation studies to assess the impact of different model designs, by joint training of the diffusion model and Variational Autoencoder ("without two-stage training") and training the model without aligning the representation space of graphs and text using L con in Equation (4) ( "without representation alignment"). Additionally, we perform experiments without both two-stage training and representation alignment ("without twostage training & representation alignment") in Table 4. We find that joint training the aligned encoder-decoder alongside Multi-modal Molecule Latent Diffusion achieves comparable or superior performance relative to 3M-Diffusion in terms of diversity and novelty, but at the expense of reduced quality of match between the molecular structures produced and the textual descriptions provided. We find that training the model without the aligned encoderdecoder results in substantial drop in performance relative to 3M-Diffusion, underscoring the importance of aligning the representation of textual descriptions of molecules with that of molecular structures. Ablation of both two-stage training and representation alignment components leads to a substantial drop in performance (Similarity of 10.24), rendering the resulting model unsuitable for real-world applications. The full 3M-Diffusion model (last row) achieves the best performance with respect to all of the metrics. This confirms the importance of two-step training and representation alignment in training a model that can generate diverse, novel set of molecular structures from their textual descriptions.
Property Conditional Molecule Generation. In traditional molecule generation models, the generation process is often conditioned on specific molecular properties such as logP (Kusner et al., 2017) and QED. logP assesses a molecule's solubility, while QED evaluates its drug likeness (Bickerton et al., 2012). Previous methods have explored ways to optimize such models further in order to generate molecules with desired properties. In contrast, 3M-Diffusion utilizes textual prompts for conditional molecule generation (See Figure 2).
In particular, we generated 10 samples for each prompt and subsequently select the top 5 samples based on the value of the desired property. For MolT5-large, we adjust the temperature of the transformer in an attempt to encourage diversity while ensuring that invalid samples are controlled during the sampling process. We applied the same methodology used in our previous experiments to maintain consistency across the comparative analysis, as shown in Figure 2. When compared to the MolT5-large model, we find that 3M-Diffusion excels in generating a diverse range of molecules that also score better with respect to the desired properties, as evidenced by higher logP values Figure 2. This underscores our 3M-Diffusion's ability to produce molecules that are not only novel and diverse but also match the textual description provided. This unique approach distinguishes 3M-Diffusion from prior methods for molecular generation, presenting a fresh and promising perspective on molecular design. Additional experimental details and insights regarding prompts related to solubility and drug likeness are provided in Appendix D.2.
In Figure 2, we introduce more specific prompts related to molecular categories, such as "The molecule is a member of the class of flavones." These prompts bias the model in favor of generating useful drug candidates. Specifically, we choose a reference molecule from the training dataset that matches the given textual prompt. Subsequently, we generate ten additional molecules and select the top five molecules that exhibit the highest structural similarity (calculated by Rascal algorithm (Raymond et al., 2002)) to the reference molecule. These selected molecules are shown in Figure 2.
Comparison of 3M-The molecule is a dihydroxy monocarboxylic acid anion that is the conjugate base of (3,4-dihydroxyphenyl)acetic acid, arising from deprotonation of the carboxy group. It has a role as a human metabolite. It derives from a phenylacetate. It is a conjugate base of a (3,4dihydroxyphenyl)acetic acid. The molecule is a dihydroxyanthraquinone that is anthracene-9,10dione substituted by hydroxy groups at positions 1 and 5. It derives from a hydride of an anthracene. Text The molecule is an aliphatic alcohol that is octane substituted by a hydroxy group at position 3. It has a role as a metabolite. It is an aliphatic alcohol and an octanol. The molecule is an octanoate ester resulting from the formal condensation of the carboxy group of octanoic acid (caprylic acid) with the hydroxy group of propan-1-ol. It has a role as a metabolite. The molecule is the alphaamino-acid anion formed by proton loss from the carboxy group of 5-hydroxylysine. It is a conjugate base of a 5hydroxylysine. The molecule is a member of the class of xanthones that is 9H-xanthen-9-one substituted by hydroxy groups at positions 1 and 8, a methoxy group at position 3 and a methyl group at position 6. It has been isolated from Microdiplodia species. It has a role as a metabolite. It is a member of xanthones, a polyphenol and an aromatic ether. Diffusion with MolT5-large shows that 3M-Diffusion generates more diverse molecules that exhibit a higher degree of structural similarity with the reference molecule while producing novel molecular structures. Additional examples of prompts are provided in Appendix D.3.
Case Study. In Figure 4, we choose a set of examples from the ChEBI-20 dataset to visually illustrate the structural similarity between the molecules generated by 3M-Diffusion and the ground truth structures corresponding to the given textual descriptions. Specifically, we show the largest substructure shared by the generated and ground truth molecules (using RDkit (Landrum, 2010)). 3M-Diffusion demonstrates the capability to generate molecules that include the relevant substructure, as illustrated in Figure 4. Simultaneously, it displays the flexibility to generate molecules with novel substructures, thereby striking a balance between generating structures that are similar to those represented in the training data and suggesting novel structures that match the textual description provided. We include additional results in Appendix D.4.
Generating molecules with desired properties is a critical task with broad applications in drug discovery and materials design. Such applications call for methods that produce diverse, and ideally novel, molecules with the desired properties. We have introduced 3M-Diffusion, a novel multi-modal molecular graph generation method, to address this challenge. 3M-Diffusion first encodes molecular graphs into a graph latent space aligned with text descriptions. It then reconstructs the molecular graph and atomic attributes based on the given text descriptions using the molecule decoder. 3M-Diffusion learns a probabilistic mapping from the text space to the latent molecular graph space using the latent diffusion model. The results of extensive experiments show that 3M-Diffusion can generate high-quality, novel and diverse molecular structures that semantically match the textual description provided. Ablation studies underscore the critical role of aligning the latent representation of text with that of molecular graphs in 3M-Diffusion.
We present a formal description of diffusion (Ho et al., 2020;Song et al., 2020). Diffusion models are latent variable models that represent the data z 0 as Markov chains z T • • • z 0 , where intermediate variables share the same dimension. DMs involve a forward diffusion process q (z 1:T | z 0 ) = ∏ T t=1 q (z t | z t-1 ) and a reverse process p θ (z 0:T ) = p (z T ) ∏ T t=1 p θ (z t-1 | z t ). The forward process starting from z 0 can be written as:
where the hyperparameter α 1:T determines the magnitude of noise added at each timestep t.
The values of α 1:T are selected to ensure that samples z T converge to standard Gaussians, i.e., q (z T ) ≈ N (0, I). This forward process q is usually predefined without trainable parameters.
The generation process of DMs is defined as learning a parameterized reverse denoising process, which incrementally denoises the noisy variables z T:1 to approximate clean data x 0 in the target data distribution:
where the initial distribution p(z T ) is defined as N (0, I). The means µ θ typically are neural networks such as U-Nets for images or Transformers for text, and the variances σ 2 (t -1, t) typically are also predefined. As latent variable models, the forward process q(z 1:T |z 0 ) can be seen as a fixed posterior. The reverse process p θ (z 0:T ) is trained to maximize the variational lower bound of the likelihood of the data:
To ensure training stability, a simple surrogate objective up to irrelevant constant terms is introduced (Song et al., 2020;Ho et al., 2020;Nichol & Dhariwal, 2021):
where z t is used to denote √ α t z 0 + √ 1α t ϵ for simplicity. Intuitively, we train a neural network ẑθ (z t , t) to approximate the original data given some noisy latent and the timestep through Equation ( 12), ẑθ (z t , t) ≈ z. With a trained denoising network, we define the mean and variance function in Equation ( 10) following the previous work (Ho et al., 2020) as:
where λ t is the time-dependent weight and α t|t-1 = α t /α t-1 . For generation, we employ the standard DDPM sampler, also referred to as the ancestral sampler (Ho et al., 2020). The process involves sampling initial noise z T ∼ N (0, I) and iteratively applying the update rule:
where ϵ ∼ N (0, I). We use T = 50 for the sampling timesteps.
Algorithm 1 Training Algorithm of 3M-Diffusion 1: Input: Pairs of Molecule Graph data G = (V, E ) and Textual Description T 2: Initial: Graph Encoder network E g , LLM Encoder network E t and Decoder network D, denoising network ẑθ 3: Pretraining Stage: Aligned Text and Graph Encoder Training 4: while E g and E t have not converged do 5:
optimizer(L con ) 7: end while 8: First Stage: Autoencoder Training 9: Initialize E g from Pretraining Stage 10: while E g and D have not converged do 11: z ← E g (G) 12: G ← D(z) ▶ Decoding 13:
14: E g , D ← optimizer(L elbo ) 15: end while 16: Second Stage: Multi-modal Molecule Latent Diffusion Training 17: Fix Graph Encoder E g from Stage 1 and LLM Encoder E t from Pretraining Stage 18: while z θ have not converged do 19: zt = wẑ θ (z t , t, c) + (1w)ẑ θ (z t , t), 6:
▶ Decoding 9: return x, h
The complete training algorithm is depicted in Algorithm 1. Firstly, the process begins by training the graph and LLM encoders using contrastive learning. Following that, the Variational Autoencoder is trained to facilitate the reconstruction of molecule graphs. Lastly, the Multi-modal Molecule Latent Diffusion Model is employed to train the denoising network for the sampling process. This structured approach forms the core of our training procedure. Moreover, we present the sampling method of 3M-Diffusion in Algorithm 2.
The statistics of datasets, including their the number of samples for training, validation and test sets are given in Table 5.
ChEBI-20 (Edwards et al., 2021): This dataset is generated by leveraging PubChem (Kim et al., 2016) and Chemical Entities of Biological Interest (ChEBI) (Hastings et al., 2016). It compiles ChEBI annotations of compounds extracted from PubChem, comprising moleculedescription pairs. To ensure less noise and more informative molecule descriptions, it includes samples with descriptions exceeding 20 words.
PubChem (Liu et al., 2023): This dataset comprises 324k molecule-text pairs gathered from the PubChem website. As the dataset contains numerous uninformative texts like "The molecule is a peptide," it curates a high-quality subset of 15k pairs with text longer than 19 words for downstream tasks. This refined subset is then randomly split into training, validation, and test sets. The remaining dataset, which is more noisy, is utilized for pretraining. To prevent data leakage, we exclude samples that appear in the testing set of other datasets from the pretraining dataset.
PCDes (Zeng et al., 2022): This dataset compiles molecule-text pairs sourced from Pub-Chem (Kim et al., 2016), encompassing names, SMILES notations, and accompanying paragraphs providing property descriptions for each molecule.
MoMu (Su et al., 2022): This dataset comprises paired molecule graph-text data, with the textual information for each molecule extracted from the SCI paper dataset (Lo et al., 2019).
For all datasets, inspired by previous datasets (ZINC250K (Irwin et al., 2012) andQM9 (Blum &Reymond, 2009;Rupp et al., 2012)) on molecule generation, we preserve molecules with fewer than 30 atoms, creating new training, validation, and test sets from their original counterparts.
We firstly introduce the baselines used for text-guided molecule generation in this section.
MolT5 (Edwards et al., 2022): MolT5 addresses the difficult problem of cross-domain generation by linking natural language and chemistry, tackling tasks such as text-conditional molecule generation and molecule captioning. It presents several T5-based LMs (Raffel et al., 2020) by training on the SMILES-to-text and text-to-SMILES translations.
ChemT5 (Christofidellis et al., 2023): ChemT5 introduces a versatile multi-task model for seamless translation between textual and chemical domains. This innovative approach leverages transfer learning in the chemical domain, with a particular emphasis on addressing cross-domain tasks that involve both chemistry and natural language.
Then, we will also include seven baselines for unconditional generation to verify the effectiveness of our model.
CharRNN (Segler et al., 2018): CharRNN models a distribution over the next token based on the previously generated tokens. We train this model by maximizing the log-likelihood of the training data, which is represented as SMILES strings.
VAE (Kingma & Welling, 2013): Variational Autoencoder (VAE) consists of two neural networks: an encoder and a decoder. The encoder infers a mapping from high-dimensional data to a lower-dimensional space, while the decoder maps this lower-dimensional representation back to the original high-dimensional space. It optimizes reconstruction loss and regularization term in a form of Kullback-Leibler divergence. We use SMILES as input and output representations as discussed in (Polykovskiy et al., 2020).
AAE (Makhzani et al., 2015): Adversarial Variational Autoencoder (AAE) replaces the Kullback-Leibler divergence in VAE with an adversarial objective. An auxiliary discriminator network is trained to distinguish between samples from a prior distribution and the model's latent codes. Similar to VAE, we use SMILES as input and output representations as shown in (Polykovskiy et al., 2020).
LatentGAN (Prykhodko et al., 2019): The Latent Vector Based Generative Adversarial Network (LatentGAN) integrates an autoencoder with a generative adversarial network. It first pretrains an autoencoder to map SMILES structures to latent vectors. Subsequently, a generative adversarial network is trained to generate latent vectors for the pre-trained decoder.
BwR (Diamant et al., 2023): BwR aims to decrease the time/space complexity and output space of graph generative models, demonstrating notable performance improvements in generating high-quality outputs on biological and chemical datasets.
HierVAE (Jin et al., 2020): HierVAE employs a hierarchical encoder-decoder architecture to generate molecular graphs, leveraging structural motifs as fundamental building blocks. (Kong et al., 2022): PS-VAE automatically identifies regularities within molecules, extracting them as principal subgraphs. Utilizing these extracted principal subgraphs, it proceeds to generate molecules in two distinct phases.
Note that we use the code for CharRNN, VAE, AAE, LatentGAN by the Molecular Sets (MOSES) Benchmark Polykovskiy et al. (2020).
We use official implementation publicly released by the authors of the baselines or implementation from Pytorch. Note that we use the decoder of HierVAE for our model. We run experiments on a machine with a NVIDIA A100 GPU with 80GB of GPU memory. In all experiments, we use the Adam optimizer (Kingma & Ba, 2014). In particular, for our 3M-diffusion. we set the dimension of latent representation as 24, the number of layers for the denosing networks (MLP) as 4, α as 0.1, T as 100 for training and 50 for sampling process and the probability p = 0.8 for text-guided molecule generation. For unconditional generation, we set T as 10 or 5 for the sampling process.
In this section, we compare the average inference times of 3M-Diffusion, MolT5, and ChemT5 by generating 1,000 samples, as shown in Figure 3. The average time to generate each molecule is calculated and presented. We observe that 3M-Diffusion achieves a speed 45 times faster than baseline methods, confirming the efficiency of our proposed model. Our proposed model not only generates molecules more quickly than previous approaches but also maintains a high level of semantic alignment with textual descriptions.
In this section, we provide more prompt examples and more experiment details for Property Conditional Molecule Generation. We generate 10 samples for each prompt and subsequently select the top 5 samples based on the desired property values, as shown in Figure 2, 5, 6 and 7. For instance, when using the prompt "This molecule is not like a drug," we
QED: 0.11 QED: 0.05 QED: 0.05 QED: 0.11 QED: 0.13 Prompt: This molecule is not like a drug. QED: 0.77 QED: 0.77 QED: 0.89 QED: 0.89 QED: 0.89 MolT5-large 3M-Diffusion Reference Prompt: The molecule is an anthocyanidin cation.
In this section, we choose 20 more examples from the ChEBI-20 dataset to visually illustrate the structural similarity between the molecules generated by 3M-Diffusion and the ground truth molecules corresponding to the given text descriptions. All results are shown in Figure 4. We also observe that our model is proficient at generating molecules that share significant common structural elements, while simultaneously obtaining diverse and innovative molecular structures.
As shown in Figure 10, our sampled molecules present rich variety and structural complexity. This demonstrates our model's ability to generate diverse, novel, and realistic molecules.
The molecule is a hydroxyquinolone that is 3-hydroxyuinolin-2(1H)-one which is substituted at position 4 by a phenyl groups. Isolated from the mycelium of several Penicillium species, it exhibits strong antibiotic activity against M. tuberculosis and also against B. subtilis, S. aureus and S. cerevisiae. It has a role as an antibacterial agent and a Penicillium metabolite.
The molecule is a tripeptide composed of Lphenylalanine, L-alanine and L-proline joined in sequence by peptide linkages. It has a role as a metabolite. It derives from a Lphenylalanine, a L-alanine and a L-proline.
The molecule is an imidazoquinoline that is 3Himidazo[4,5-f]quinoline substituted by a methyl group at position 3 and an amino group at position 2. It has a role as a carcinogenic agent.
The molecule is a L-serine derivative. It has a role as an Escherichia coli metabolite. It is a conjugate acid of a N-(2,3dihydroxybenzoyl)-L-serinate.
The molecule is a phenylacetaldehyde in which the 3 and 4 positions of the phenyl group are substituted by hydroxy groups. It has a role as a human metabolite, an Escherichia coli metabolite and a mouse metabolite. It is a member of catechols, an alpha-CH2-containing aldehyde and a member of phenylacetaldehydes.
The molecule is a quinolinemonocarboxylate that is the conjugate base of xanthurenic acid, obtained by deprotonation of the carboxy group. It has a role as an animal metabolite. It is a conjugate base of a xanthurenic acid.
The molecule is an aldoxime resulting from the formal condensation of the aldehyde moiety of indol-3-ylacetaldehyde with hydroxylamine. It is a member of indoles and an aldoxime.
The molecule is a quaternary ammonium ion obtained by methylation of the tertiary amino function of reticuline. It derives from a reticuline.
The molecule is an N(2)-acyl-L-ornithine where the acyl group is specified to be acetyl. It has a role as a human metabolite, an Escherichia coli metabolite, a Saccharomyces cerevisiae metabolite and a mouse metabolite. It is an acetyl-L-ornithine and a N2-acyl-L-ornithine. It is a tautomer of a N(2)acetyl-L-ornithine zwitterion.
The molecule is a trihydroxyflavone with the hydroxy groups at positions C-5, -7 and -8. It has a role as an antioxidant and a metabolite.
The molecule is a leucoanthocyanidin that is the 3-deoxy derivative of (2R,3S,4S)leucocyanidin. It is a pentahydroxyflavan and a leucoanthocyanidin. It derives from a (2R,3S,4S)-leucocyanidin.
The molecule is a tetrahydroxyflavone that is apigenin with an extra hydroxy group at position 8. It has a role as a metabolite. It derives from an apigenin.
The molecule is an aromatic L-alpha-amino acid zwitterion resulting from transfer of a proton from the carboxy to the amino group of 3-O-methyldopa. Major species at pH 7.3 It has a role as a human metabolite. It is a tautomer of a 3-O-methyldopa.
The molecule is a 2-hydroxy carboxylate that is the conjugate base of 3-(indol-3-yl)lactic acid. It has a role as a human metabolite. It derives from a lactate. It is a conjugate base of a 3-(indol-3-yl)lactic acid.
The molecule is a tetrachlorobiphenyl that is biphenyl in which both phenyl groups are substituted by chlorines at positions 3 and 5. It is a tetrachlorobiphenyl and a dichlorobenzene.
The molecule is an aromatic ether that is the derivative of 2-phenylethylamine with methoxy substituents at the 3-and 4-positions. It is an alkaloid isolated from the Cactaceae family. It has a role as a plant metabolite and an allergen. It is a phenylethylamine, an aromatic ether and an alkaloid. It derives from a hydride of a 2-phenylethylamine.