LDMol: Text-to-Molecule Diffusion Model with Structurally Informative Latent Space

We sincerely thank the reviewer for the constructive and valuable comments, and please refer to our responses below.

W1: Complex model pipeline of LDMol.

We would like to assure the reviewer that although LDMol comprises several modules, the relationship and information flow between them are very simple and intuitive. And since each module(AE, diffusion model, condition encoder) works independently, one can easily replace or improve each module. Given that StableDiffusion(which has an almost identical workflow) is also widely used and fine-tuned by many, we believe that LDMol is highly accessible to end users.

W2: Generalization ability on different types of data.

We would like to assure the reviewer that our model is not overfitted to any particular dataset. Rather it showed high performance in generation(Figure 4) and retrieval(Table 2) with short and less specific inputs, even in the relative out-of-domain distribution. Also, please refer to our text-to-molecule generation comparison in other test datasets in the next response.

W3: Request for a wider range of benchmarks.

Thank you for the suggestion of an additional benchmark. However, since text-to-molecule(T2M) generation is a recently emerging field, there is a lack of appropriate benchmarks other than CHEBI-20, and almost all the papers presented so far conducted their benchmark study on this. That said, we evaluated LDMol and several baselines' T2M generation performances on the PCdes test set, which is another prepared dataset with detailed captions. Note that LDMol was not trained with these data. We found that our model still outperforms the other baselines, and added this result to Appendix B.3(Table 6) in the revised manuscript.

W4: The model's possible reliance on large datasets.

In fact, the 320k molecule-text paired data that LDMol was trained with is relatively very small compared to the other baselines; molXPT[1] and biot5+[2] used 8M and 3.2M molecule-text combined data from web, on top of other millions of natural language data. On the other hand, we only utilized a relatively small but well-curated dataset and still achieved high performance.

Q1: The model behavior on ambiguous or contradictory inputs.

When we fed a completely ambiguous input, the model spits out a variety of different but incoherent molecules with high validity. When we fed completely contradictory inputs that the model cannot follow, the output validity dropped significantly and the outputs were chimeric between contradictory prompts, which we regard as appropriate behavior. Please refer to the provided examples in Appendix B.3(Figure 11) in the revised manuscript.

Q2: Possibility for other important chemical domains.

Thanks for the insightful question, and we believe that our approach to the latent diffusion model is applicable to other important chemical data domains like materials and proteins. Here, note that our main contribution is to show that the latent space should not just be learnable, but also informative about the data domain. Therefore, a special encoder objective function tailored to the target domain could be suggested, just like we trained our encoder with contrastive learning to introduce structure invariants.

Q3: Computational requirements for training / deploying LDMol.

In the training phase, our model did not use more than 24GB of VRAM, so it can be trained on NVIDIA Geforce 3090. Meanwhile, molXPT and bioT5+ were trained on NVIDIA A100 with 80GB of VRAM (probably due to their large amount of training data). For the inference stage, our model can be used with 5GB of VRAM, which is not so different from our baselines of 2GB(bioT5+) and 4GB(molT5).

Q4: Possible limitations on the generated molecules.

Since our autoencoder is only trained on SMILES below a certain token length, LDMol does not generate output for several hundreds of tokens such as polymers. However, since diffusion models have already shown great scalability in the image domain, we think that the token length can be increased and this will not be a fundamental limitation in the future.

---

[1] MolXPT: Wrapping Molecules with Text for Generative Pre-training, ACL 2023, arxiv:2305.10688

[2] BioT5+: Towards Generalized Biological Understanding with IUPAC Integration and Multi-task Tuning, ACL 2024, arxiv:2402.17810

Dear reviewer UnHu,

We’d like to summarize our revisions that addressed the questions and concerns you have raised. Specifically,

1. We ensured the proposed LDMol can be trained with a much smaller dataset and only needs comparable computational requirements compared to recent baselines.

2. We conducted extensive additional experiments for output behavior for ambiguous and contradictory inputs and outperforming baselines in another benchmark for LDMol's generalization ability.

3. We emphasized that LDMol is not so complex compared to other widely-used latent diffusion models, and has great potential in many important chemical fields as one of the first diffusion models outperforming AR models in chemical instance generation.

As the deadline for the reviewer-author discussion phase is less than 72 hours away, we’d like to gently remind you and ask if we’ve adequately addressed your questions and concerns. We’re happy to provide clarification on any remaining questions you may have.

Best regards, Authors

Dear Reviewer UnHu,

We are writing to kindly remind you regarding our responses to your comments on our manuscript. We comprehensively addressed your comments and questions with rich additional experiments, and emphasized our contribution with its broad applicability to other domains.

As the deadline for finalizing the reviews is approaching within 48 hours, we would greatly appreciate it if you could timely consider our revision and inform us of your decisions. Please let us know if there are any further clarifications needed.

Best regards, Authors

We sincerely thank the reviewer for the insightful and constructive comments. Please refer to our responses below.

W1: Discussion about achieving structural consistency in the latent space.

Thanks for the important comments that highlight our innovation. We found that naive VAE shows structural inconsistency because knowing how to reconstruct input doesn't mean a complete understanding of the data domain. Even in LLMs with Autoregressive modeling, it is known that even if the model has successfully learned the grammar of the natural language, the model doesn't understand its meaning and often hallucinates nonsensical output[1]. Similarly, VAE is prone to only reconstruct the SMILES strings while not knowing what atoms and bonds the input SMILES imply. We added this discussion in the revised manuscript(line 239~243).

W2: Hyperparameter for the encoder contrastive learning.

Thanks for the constructive comments. It is known that lower temperature (tau) in contrastive learning penalizes hard negatives more strongly, making the learned feature more sensitive to fine-grain details[2]. We considered this as an appropriate property for our latent space and used a small tau of 0.07. Using too big tau reduced the model's ability to distinguish different molecules and made the training loss converge to a much higher value, but the values around 0.07 work well. This result can be found in Appendix A.2(Figure 7) in the revised manuscript.

W3: Justification and possible alternatives for the linear compression layer.

Thanks for giving an opportunity to clarify. We determined the feature dimension empirically by monitoring reconstruction accuracy. Please refer to the table below, which we included in Appendix A.2(Table 5) in the revised manuscript.

When we compressed it into 32, the reconstruction rate dropped to 94.8%. When we compressed it into 128, the reconstruction rate was slightly higher, but the training time for the subsequent diffusion model increased quadratically. Considering that the failed 2% for the current model were mostly very long molecules, we decided $d_z$ to 64.

Meanwhile, an alternative compression method could be a transformer layer such as Perceiver-Resampler[3][4]. However, we had to keep our compression layer as simple as possible, because adding another complex layer makes the latent space deviate from the former informative and well-regulated learnable space. In fact, when we replaced our linear layer with the transformer module, the performance was significantly decreased as shown below, as the output space became less learnable.

And considering that our linear layer has much fewer parameters(66K) than the transformer module(22M), we concluded that the linear compression was sufficient for our work.

W4: Incorporating additional 3d descriptors for performance improvement.

Thank you for your valuable suggestion, and we acknowledge that we could improve our model performance by injecting more external information into the constructed latent space. Unfortunately, we have not yet found a trainable 3d descriptor database that provides structural information beyond what SMILES already provides. Also, there should be another contribution to combine that information with the current latent space as a structural invariant(e.g. features from coordinate value should be SE(3)-invariant, etc.), so we found that this is beyond the scope of our current manuscript. We emphasize our current work as the first successful text-to-molecule diffusion models outperforming AR baselines, and hope you understand that we leave this approach as future work.

W5: Discussion about several related works.

Thank you for the suggestion, and we have added the works for molecule generation using bidirectional RNNs [5] and LLMs [6] into the Related Works section of the manuscript(line 146~154).

---

[1] A survey on hallucination in large language models: Principles, taxonomy, challenges, and open questions, arxiv:2311.05232

[2] Understanding the Behaviour of Contrastive Loss, CVPR 2021, arxiv:2012.09740

[3] Flamingo: a Visual Language Model for Few-Shot Learning, NeurIPS 2022, arxiv:2204.14198

[4] BLIP-2: Bootstrapping Language-Image Pre-training with Frozen Image Encoders and Large Language Models, PMLR 2023, arxiv:2301.12597

[5] Bidirectional molecule generation with recurrent neural networks. Journal of chemical information and modeling 2020

[6] DrugLLM: Open Large Language Model for Few-shot Molecule Generation, arxiv:2405.06690.

Dear reviewer CrH4,

We’d like to summarize our revisions that addressed the questions and concerns you have raised. Specifically,

1. We included theoretical support and experimental demonstration of the temperature(tau) selection for the encoder training.

2. We explained the choice of our linear compression dimension size according to the reconstruction rate and discussed possible alternatives with various experiments.

3. We followed your suggestions and included more related works and extensive discussions about achieving structurally consistent latent space.

As the deadline for the reviewer-author discussion phase is approaching less than 72 hours, we would like to extend a warm reminder and ask whether we have addressed your questions and concerns adequately. We would be happy to clear up any additional questions.

Best regards, Authors

Dear Reviewer CrH4,

We are writing to kindly remind you regarding our responses to your comments on our manuscript. We addressed your concerns on our model with thorough experimental results, and revised our manuscript with additional discussions to support and emphasize our contribution.

As the deadline for finalizing the reviews is approaching within 48 hours, we would greatly appreciate it if you could timely consider our revision and inform us of your decisions. Please let us know if there are any further clarifications needed.

Best regards, Authors

We sincerely thank the reviewer for thorough reading and insightful comments, and please refer to our responses below.

W1: Convoluted pipeline with AR transformer.

We like to assure the reviewer that no additional complex models are added compared to the mentioned text-to-image(T2I) LDM. Both comprise an autoencoder(AE) and a conditional diffusion model, where the decoder part of our AE is an autoregressive(AR) transformer. This AR decoder doesn't model the relationship between text and molecules at all, and only recovers SMILES strings from given latent. We employed an AR transformer as our decoder because SMILES itself is sequential data with a varied length. Since our decoder is only used once at the end of sampling without a special AR sampling technique(e.g. beam search), we would like to assure the reviewer that the required time for LDMol was not so long compared to the other baselines.

W2: Modification on top of the existing models.

We agree that LDMol indeed has a similar pipeline to T2I latent diffusion models. However, our biggest contribution is preparing a latent space that embeds the useful and consistent chemical information, improving the diffusion model performance. This allows LDMol to be one of the first diffusion models outperforming AR models in sequence generation. Note that without our contrastive encoder training, the model trained on the naive latent space did not surpass the previous baselines(Table 3). We believe that the proposed SMILES-enumeration contrastive learning is a domain-specific and novel approach to encoding molecular information.

Q1: The necessity of the separate linear compression layer.

The reason we didn't create an encoder with the compressed dimension from the beginning is that we trained our encoder with contrastive learning to extract as informative and structure-invariant features as possible. When we directly applied contrastive loss with 64-dimensioned features, the trained encoder failed to learn informative features, and the corresponding decoder couldn't reconstruct the input molecule. Please refer to the table below, which we also included in Appendix A.2(Table 5) in the revised manuscript.

Q2: The usage of a separate AR transformer besides of the trained SMILES decoder.

We want to emphasize that we only trained a single SMILES decoder, and didn't build a "separate" autoregressive transformer; the trained SMILES decoder is an autoregressive transformer. You can consider "SMILES encoder + linear compression layer" as the total encoder part of our AE, and "autoregressive transformer" as the decoder part. Therefore, the suggested “feeding the DiT output directly back into the SMILES decoder (via cross-attention)” is what we’re doing currently.

Q3: The speed and runtime of the proposed model.

For the inference time, despite the usage of the latent decoder and Classifier-Free Guidance (CFG) that doubles the diffusion model iteration, LDMol showed comparable inference speed to other AR transformer models; as shown in the table below, LDMol was slower than bioT5+ but faster than the model of molT5$_{large}$. Furthermore, various studies and techniques have been proposed to reduce the diffusion model inference speed [1][2], so we believe that future research can further address the issue of LDMol’s inference time.

About the required time in the reverse diffusion and AR decoding step, we observed that about 70% of the runtime is used for the reverse diffusion and 30% for the AR decoding steps.

---

[1] SDXL: Improving Latent Diffusion Models for High-Resolution Image Synthesis, ICLR 2024, arxiv:2307.01952

[2] Consistency models, ICML 2023, arxiv:2303.01469

Dear reviewer aYeM,

We’d like to summarize our revisions that addressed questions and concerns you have raised. Specifically,

1. We emphasized our biggest contribution which is constructing a novel latent space to embed the useful and consistent chemical information, and demonstrating this can improve the performance of the diffusion model on it.

2. We rationalized our AR decoder as a replacement of the conventional latent decoders for length-varied SMILES, and justified the introduction of the linear compression layer with an experiment.

3. We showed that LDMol can operate with a comparable inference speed with current baseline models, even with the usage of the latent decoder and CFG.

As the deadline for the reviewer-author discussion phase is less than 72 hours away, we’d like to gently remind you and ask if we’ve adequately addressed your questions and concerns. We’re happy to provide clarification on any remaining questions you may have.

Best regards, Authors

Dear Reviewer aYeM,

We are writing to kindly remind you regarding our responses to your comments on our manuscript. We have assured to resolve your concerns about our model design choices, and emphasized our contribution to novel molecule latent space enabled a successful text-to-molecule diffusion model.

As the deadline for finalizing the reviews is approaching within 48 hours, we would greatly appreciate it if you could timely consider our revision and inform us of your decisions. Please let us know if there are any further clarifications needed.

Best regards, Authors

We sincerely thank the reviewer for thorough reading and insightful comments. Please refer to our replies to your concerns below.

W1: Text-molecule paired datasets are often limited to structural descriptors.

We agree that the prepared dataset had fewer descriptions for bonds and atoms. However, even in text-molecule paired data set there were still a number of descriptions for functional groups and motifs which enabled LDMol to achieve SOTA-level generation performance. Given that text-conditioned molecule generation is a recently emerging field and high-quality large-scale text-molecule pair data is only now being presented, we believe our diffusion model methodology of LDMol can be further advanced in future works by incorporating more structurally and biochemically enriched data.

W2: Effectiveness of LDMol with vague or less specific text conditions.

We would like to assure the Reviewer that many text descriptions in the PubchemSTM training dataset were very short and vague; 29% of the PubchemSTM captions have less than 50 characters, and 68% of the captions have less than 100 characters. We emphasize that LDMol could generate molecules with both highly detailed conditions (Table 1) and less specific prompts (Figure 4). LDMol also performed well in the retrieval downstream task which includes unseen and out-of-domain text, indicating that our model's comprehension ability was not overfitted to the training data.

Q1: Model performance & training efficiency for small and less-detailed datasets.

Most recent transformer-based baselines used millions of single-modal and multimodal data from various databases. For example, Biot5+[1] used 8M and 3.2M molecule-text combined data from web crawling on top of millions of natural language data. MolT5[2] used 100 million molecule data to model the SMILES data distribution. Compared to these works, the 320k text-molecule pairs we introduced are actually a very small amount of data, and yet LDMol outperformed all baselines.

Q2: Modifications for the model's stereoisomer sensitivity.

We first emphasize that our introduced hard-negative samples in the encoder training were sufficiently helpful for LDMol to distinguish subtle structure differences such as stereoisomers. More specifically, we found that our LDMol encoder distinguishes the given molecule's stereoisomer from the SMILES enumeration of itself in 87.7% of the cases, whereas 50.3% for the encoder trained without hard-negatives. This improved the exact match ratio for the text-to-molecule generation benchmark by more than 25%p(Table 3), enabling LDMol to outperform previously suggested models.

That said, the stereoisomer discrimination problem is definitely a hindrance to the model performance, when downstream tasks require very subtle structure differences. We expect that injecting explicit 3-dimensional descriptor information (e.g. coordinates) effectively into the latent space might further resolve this problem, which could be an interesting future work.

Q3: Computational efficiency of LDMol.

Please refer to the table below that compares the computational efficiency of LDMol and several baselines.

In terms of memory usage, our model can operate with less than 5GB of VRAM, which is comparable to other transformer models with 2GB for biot5+ and 4GB of molT5_large. The required time was also comparable to transformer-based models, even with the latent decoder and the Classifier-Free Guidance(CFG) that doubles the diffusion model usage. Also, many works have been published to reduce the inference time of diffusion models [3][4]. Therefore, as one of the first successful text-to-molecule diffusion models, we believe that the inference time can be further improved by future research.

---

[1] BioT5+: Towards Generalized Biological Understanding with IUPAC Integration and Multi-task Tuning, ACL 2024, arxiv:2402.17810

[2] Translation between molecules and natural language, EMNLP 2022, arxiv:2204.11817

[3] Accelerating parallel sampling of diffusion models, ICML 2024, arxiv:2402.09970

[4] Consistency models, ICML 2023, arxiv:2303.01469

Dear reviewer A8d4,

We’d like to summarize the addressed questions and concerns you have raised. Specifically,

1. We ensured that LDMol is trained with a much smaller amount of data including vague texts compared to previous works, and yet able to generate molecules that faithfully follow the given condition, from less-detailed texts to highly specific prompts.

2. We explained how our introduced stereoisomer hard-negatives improved the sensitivity of the constructed latent space and suggested possible improvements in further research.

3. We showed that LDMol can operate with efficient time and memory usage* which is comparable to baseline models.

As the deadline for the reviewer-author discussion phase is approaching less than 72 hours, we would like to extend a warm reminder and ask whether we have addressed your questions and concerns adequately. We would be happy to clear up any additional questions.

Best regards, Authors