NExT-Mol: 3D Diffusion Meets 1D Language Modeling for 3D Molecule Generation

W5: Use of Questionable Metrics: The metric reported by JODO shows a 2.8% 3D molecule stability for the GEOM-Drugs dataset, rendering it practically meaningless. I strongly encourage avoiding the propagation of this metric in new papers. According to the code, the metric is based on a predefined bond length table, and a bond length is considered "good" if it is within 0.05 Å of the table value. However, the optimal distances between atoms are primarily defined by the energy landscape underlying the data—for GEOM-Drugs, it's GFN2-xTB—and depending on atom configurations, deviations in bond lengths can exceed 10%. While I'm uncertain about the validity of the 3D FCP metric, it's not entirely clear that it's completely off when compared with 3D molecule stability. Instead of attempting to report every possible metric to compare with other methods, focus on identifying the most important ones and emphasize the significance of the margins your model is achieving.

Q4: Analysis of Metrics Used: Could you analyze the metrics used in your study, especially focusing on 3D molecular stability and the FCP metric? Since the FCP metric relies on a neural network, is it capable of handling your data distribution, and is it up-to-date with recent developments?

Response: Thank you for the valuable suggestions and comments.

On the validity of our evaluation metrics. Our evaluation metrics follow the previous works [3,4,8,9] in 3D molecule generation. We discuss them below:

• On the 3D molecule stability metric. We agree that this metric might be problematic, and we generally agree on your analysis. Following your suggestion, we remove this metric in our revised main part of the submission, and move it to the Appendix B.5 as a backup. We find that other metrics, including FCD, atom stability, the distribution similarity of bond length, bond angle, and dihedral angle, can still showcase the improvement of our method.
• On the FCD metric. We agree that FCD is a more important metric for evaluating the quality of generated molecules.
  • Introduction to Fréchet ChemNet Distance (FCD). It is widely used for evaluating the quality of generated molecules. It measures the similarity between the distributions of molecular representations of generated molecules and the ground truth molecules [6]. It is calculated using activations of the penultimate layer of the deep neural network ChemNet trained to predict biological activities of drugs which captures properties of the compounds. ChemNet was trained to predict biological activities, therefore the learned molecule representations contain both chemical and biological information.
  • Highlighting FCD. Thank you for suggesting to highlighting our improvement of FCD. Following you suggestion, we have moved the presentation of FCD to the left most column in every comparison, to highlight its importance. On the GEOM-Drugs dataset, NEXT-Mol outperforms baselines by 5.39 FCD scores, a 26.5% relative improvement.
  • Alignment between FCD and our target data distribution. The ChemNet used in FCD is trained on subsets of ZINC, PubChem, and ChEMBL. However, their filtered datasets are not open-source, therefore making the direct comparison infeasible. Here we draw some empirical insights: GDB-17, a superset of QM9 [10], overlaps with 57% of PubChem’s 2017 release and 60% of ChEMBL’s 2017 release. Similarly, GEOM-DRUGS likely overlaps with ChEMBL and ZINC, given their shared focus on drug-like molecules.
• On the distribution similarity of bond length, bond angle, and dihedral angle. We measure the MMD distribution similarity for these geometry values. In general, these metrics can correctly reflect their description and should be reliable. In our understanding, they are included as surrogate metrics to evaluate the conformer quality, because the ground truth conformer is unavailable.

W6: Inconsistencies in Metric Definitions and Reporting: Providing detailed descriptions of the metrics used in the supplementary material is crucial due to inconsistencies across different papers. For example, MiDi reported the Wasserstein distance for bond angles and bond length distributions for all bonds and angles, whereas this paper (and at least JODO) uses MMD for the most frequent bonds, angles, and torsions. Additionally, the way MiDi computes atom stability and molecule stability differs from the JODO code. Performing kekulization at the beginning can alter the valencies of atoms, leading to different results (e.g., if you manually define H:O:H, where ":" is an aromatic bond, kekulization converts it to the valid water molecule H-O-H). Some models measure atom and molecule stability for raw data, while others use RDKit preprocessing before measuring stability, which can inflate results. It is essential to be consistent in comparisons. While this may not significantly change your results, it could artificially boost JODO's stability results, skewing the comparison. Overall, ensure that the same versions of the metrics are used across all comparisons and that they are well-defined in the supplementary material.

Q5 Consistency and Clarity of Metrics: To ensure consistency, could you provide detailed descriptions of the metrics used in your study and confirm that the same versions are applied across all comparisons? Additionally, please clarify any preprocessing steps that might affect the results.

Response: Thank you for the question.

• Distribution similarity measures. We use the MMD distance, following JODO, when reproducing MiDi and other baselines for consistent comparison.
• Regarding Kekulization.
  • 3D molecule. We do not include this step when computing 3D atom and molecule stability. The corresponding code can be found at the get_3D_edm_metric function in the evaluation/eval_functions.py file in our anonymous code. This setting is consistently used for other baselines.
  • 2D molecule. Kekulization is used here following JODO. The same evaluation code is used for all the baselines, therefore the comparison should be fair. Due to the time limitation in rebuttal, we are still working on a reproduction without Kekulization for 2D molecules. Thank you for the suggestion. We will publish the results on OpenReview if they become available before the discussion deadline.
• RDKit preprocessing.
  • We use canonicalized SMILES for both the generated molecules and the training dataset when computing novelty and uniqueness of molecules. This is to ensure the correctness of the reported performance. We consistently use this measure for all the baselines.
  • Other than canonicalization and Kekulization, no other RDkit preprocessing is used.

We have revised our Appendix D.2 to include these important details of evaluation metrics.

W7: Omission of EQGAT-Diffusion in Comparisons: EQGAT-Diffusion is a recent model—published well before the ICLR deadline—that should be included in the comparisons for 3D unconditional molecule generation to provide a more comprehensive evaluation.

Q6: Inclusion of EQGAT-Diffusion in Comparisons: EQGAT-Diffusion is a recent model relevant to 3D unconditional molecule generation. Could you include it in your comparisons to provide a more comprehensive evaluation?

Response: Thank you for pointing to this important related work and baseline model. We have included EQGAT-Diffusion in our comparisons for de novo 3D molecule generation to provide a more comprehensive evaluation. The results have been added to the revised manuscript in Table 2. We appreciate your feedback and have revised the manuscript accordingly.

W3: Mischaracterization of Computational Complexity: The paper states that structures were obtained using computationally intensive geometry optimization with DFT. However, the GEOM dataset was designed using the CREST software for conformer sampling, followed by geometry optimization with GFN2-xTB—a semi-empirical tight-binding method, not DFT. Moreover, compared to some deep learning models, this approach is not computationally intensive; on a reasonable workstation, geometry optimization takes about 0.5 seconds per average GEOM-Drugs structure.

W4: Missing Performance Metrics for Conformer Generation: In conformer generation, computational performance is extremely important. OpenEye Omega remains one of the most popular software tools for this task because of its speed. The paper lacks performance metrics related to speed and efficiency, which are necessary for a fair comparison with existing tools.

Q3: Inclusion of Performance Metrics: Considering that computational performance is crucial in conformer generation tasks, could you include speed and efficiency metrics for your model, particularly in comparison with tools like OpenEye Omega? I would also add OpenEye Omega conformer generation + consequent xTB geometry optimization time.

Response: We appreciate your attention to the computational complexity of our conformer generation process and the comparison with other conformer optimization toolkits. We would like to take a further discussion on these charming topics:

• The ground truth conformer sampling process (RDKit -> xTB -> CREST) in GEOM: For a molecule in GEOM-DRUGS, the initial conformers are generated from SMILES strings using RDKit, followed by geometry optimization with GFN2-xTB. The lowest energy conformer is then selected as input to CREST for further optimization.
• Computational efficiency of GEOM-Drugs conformer generation: Although the GFN2-xTB method is fast enough to be used in long metadynamics runs, the whole progress (RDKit -> xTB -> CREST) may take much more time. Particularly, the CREST program takes 89 core hours on Knights Landing Nodes or 8.2 core hours on Cascake Lake/Skay lake node [5] on average for each molecule in GEOM-DRUGS. Therefore, given a desktop of 8 core CPUs, it takes 1 -- 10 hours for conformer sampling, based on the CPU's performance and the complexity of the molecule.
• Time comparison with OpenEye Omega: We conducted a time comparison between our model and OpenEye Omega for conformer generation on the test set of the GEOM-Drugs dataset, which includes 1000 molecules. The results are shown in the Table below.

Table: Comparison of conformer generation time on the test set of the GEOM-Drugs dataset using various methods.

These experiments were performed on a platform with an 8-core Intel Xeon Processor@2.90GHz CPU and an NVIDIA A100 GPU and the time is measured in minutes and seconds. Please note that the Omega and GFN2-xTB are only run on the CPU only, while DMT and Mollama are run on the GPU. So the results may vary depending on the hardware.

We have revised our Appendix B.4 to incorporate these results.

W2: Overemphasis on 100% Validity: The paper focuses on achieving 100% validity, but in real-world applications, validity filtering is an extremely simple process due to how validity is defined. Consequently, there is no significant practical difference between achieving 90% validity and 100% validity.

Q2: Practical Advantages of 100% Validity: Given that validity filtering is straightforward in practice, what are the practical advantages of achieving 100% validity over, for example, 90% validity in real-world applications?

Response: Thank you for your feedback. We agree that validity filtering is an simple process. Here we clarify that achieving 100% validity offers critical advantages beyond increasing validity itself:

1. Improving validity could improve other 2D metrics, like SNN, Frag, and Scaf (see Table 2). These metrics measure the distributional similarity of 2D molecular structures of valid molecules. If a model still generate invalid molecules, it is likely the model does not capture the true target distribution, which contain only valid molecules. 100% validity helps the model learn from and sample from the valid molecular structures, which is essential for molecule generation tasks.
2. Improving validity could improve 3D geometry learning. The improved validity also leads to better learning of 3D molecular geometry, because it grounds 3D structure prediction on valid 2D structures. Other joint 2D and 3D prediction methods [3,4] can easily encounter invalid 2D structures when sampling 3D structures, therefore leads to worse 3D structure prediction. This is demonstrated by NEXT-Mol’s significant improvements in geometry similarity metrics (e.g., bond angle and bond length distributions) in Table 2.

These improvements beyond validity are crucial for real-world applications. We appreciate your feedback and have clarify the corresponding paragraph in the introduction and Appendix C.1 in the revised manuscript.

Table 2: Performances for de novo 3D molecule generation. We highlight the results of new 1D LM baselines in $\textcolor{blue}{\text{blue}}$. * denotes our reproduced baseline results using their source codes.

(a) Performances on the GEOM-DRUGS dataset. New baselines are evaluated for 2D-metric only.

(b) Performances on the QM9-2014 dataset. New baselines are evaluated for 2D-metric only.

As shown in Table 2(a) and 2(b) above, NEXT-Mol consistently outperforms the new baselines. Specifically, NEXT-Mol achieves relative FCD score improvement by 49% over MolGen and 62% over MolGPT on the GEOM-DRUGS datasets. These significant improvements are attributed to NEXT-Mol’s benefits from the scaling law: it is pretrained on a large dataset of 1.8 billion molecules, and features a large model size of 960M parameters. Additionally, NEXT-Mol will be released to support open-source research.

We have revised our submission to include these new results.

Thank you so much for the positive ratings and the extensive comments and suggestions about our submission. Here we provide a point-by-point response to address your concerns and incorporate your suggestions. For clarity, we have consolidated similar questions and weaknesses where appropriate, labeling them with corresponding identifiers (e.g., W1 for Weakness 1 and Q1 for Question 1) based on your original comments.

W1: Lack of Comparison with Other Molecular Language Models: Although the paper introduces a 1D molecule generation component (MoLama), a 3D conformer generation model (DMT), and a transfer learning technique, it primarily showcases the performance of the conformer generation part and the advantages of transfer learning. However, it lacks a comparison with other molecular language models concerning the quality of the generated SELFIES representations.

Response: Thank you for the suggestions. We agree that including the comparison with other molecular language models (LMs) will provide further empirical evidence to support our method. To address your concern, we include MolGPT [1] and MolGen [2] as new baselines for de novo 1D/2D molecule generation on the GEOM-DRUGS and QM9-2014 datasets. MolGPT is a decoder-only transformer trained on SMILES sequences, and MolGen is an encoder-decoder transformer trained on SELFIES sequences. For MolGen, we fine-tune its public checkpoint on the two datasets; for MolGPT, we first reproduce its pretrained model using its source code, and then fine-tune it on the two datasets. Note that, we do not directly compare with MolGen’s reported performances, because MolGen employs a different experimental setting to fill-in corrupted spans in SELFIES sequences from the test set, which is unusual. Considering these two baselines cannot generate 3D structures, we report only 2D-metric of them for the de novo generation task, and do not include them for conditional 3D molecule generation (Table 5), which requires 3D structures. The results are presented below:

Q1: Comparison of MoLama with Other Models: Could you provide a comparison between your 1D molecule generation model (MoLama) and other models like Equiformer on tasks beyond 3D molecule generation, such as molecule classification (or any suitable task where you can compare MoLama and let's say Equiformer)? Including such a comparison would highlight the standalone value of MoLama.

Response: Thank you for the suggestion. To evaluate MoLlama's capabilities beyond 1D molecule generation, we apply it to molecular property prediction tasks, highlighting the quality of its molecular representations. Following the setup in [11], we fine-tune MoLlama on four MoleculeNet datasets: FreeSolv, ESOL, Lipo, and QM7. We adopt the same experimental settings and dataset splits as [11], reporting mean performance and standard deviation over 10 random seeds. For each run, MoLlama is trained for 100 epochs, with test performance selected based on the validation dataset. We use a fixed learning rate of 1e-4 with the AdamW optimizer, and fine-tune MoLlama using LoRA (LoRA r=8 and $\alpha=32$) applied to all linear layers of the model. Following Section 3.3, we attach a single-layer bi-directional self-attention layer after MoLlama to improve its encoding ability. After that, we apply a linear layer on the mean embedding of all molecule tokens for property prediction.

Table 10: Molecule property regression results on four MoleculeNet datasets. Baseline results are from [11]. Lower $\downarrow$ is better.

Observation. As shown in Table 10, MoLlama outperforms baseline methods, achieving relative improvements of 6.5%, 4.7%, 3.5%, and 16.9% on the FreeSolv, ESOL, Lipo, and QM7 datasets, respectively. Notably, our baselines include LM-based, GNN-based, and pretrained GNN-based methods, and MoLlama's better performance demonstrates its advantages derived from the extensive pretraining.

We did not use choose Equiformer as a baseline because Equiformer is designed for 3D molecule property prediction, while MoLlama can only process 1D molecules without 3D. Therefore, we compare MoLlama with baselines that also process on 1D/2D molecules to ensure fairnes.

We have included these results in Appendix B.1 to strengthen our presentation. Thank you for your insightful suggestion.

Reference:

[1] MolGPT: Molecular Generation Using a Transformer-Decoder Model. In Journal of Chemical Information and Modeling 2021.

[2] Domain-Agnostic Molecular Generation with Chemical Feedback. In ICLR 2024.

[3] Learning Joint 2D & 3D Diffusion Models for Complete Molecule Generation. In TNNLS 2024

[4] MiDi: Mixed Graph and 3D Denoising Diffusion for Molecule Generation. In ECML 2023.

[5] GEOM, energy-annotated molecular conformations for property prediction and molecular generation. In Scientific Data 2022.

[6] Fréchet ChemNet Distance: A Metric for Generative Models forMolecules in Drug Discovery. In Journal of Chemical Information and Modeling 2018. [7] Molecular Sets (MOSES): A Benchmarking Platform for Molecular Generation Models. [8] Equivariant Diffusion for Molecule Generation in 3D. In ICML 2022.

[9] Geometric Latent Diffusion Models for 3D Molecule Generation. In ICML 2023.

[10] Enumeration of 166 Billion Organic Small Molecules in the Chemical Universe Database GDB-17. In Journal of Chemical Information and Modeling 2022.

[11] MolPROP: Molecular Property prediction with multimodal language and graph fusion. In Journal of Cheminformatics 2024.

W6: Additionally, the way MiDi computes atom stability and molecule stability differs from the JODO code. Performing kekulization at the beginning can alter the valencies of atoms, leading to different results (e.g., if you manually define H:O:H, where ":" is an aromatic bond, kekulization converts it to the valid water molecule H-O-H). Some models measure atom and molecule stability for raw data, while others use RDKit preprocessing before measuring stability, which can inflate results. It is essential to be consistent in comparisons. While this may not significantly change your results, it could artificially boost JODO's stability results, skewing the comparison.

Updated Response: Thank you for the suggestion. Following your feedback, we removed the kekulization step from the evaluation script. Additionally, we modified the RDKit MolFromSmiles function to disable the sanitization step when loading SMILES into RDKit molecule objects. The molecule sanitization is a preprocessing and includes kekulization, disabling it improves the fairness when comparing 1D-based methods and 2D-based methods, which do not use MolFromSmiles.

Observation. The evaluation results showed minimal changes for de novo 3D molecule generation. Only MolGPT exhibited notable differences, which is presented in the Table below:

Table: Comparing MolGPT’s performance difference when using the new evaluation script.

• We can observe that, on the GEOM-DRUGS dataset, MolGPT’s atom stability increased from 0.957 to 0.979, and molecule stability increased from 0.957 to 0.977.
• This change occurs because the sanitization step in RDKit imposes stricter constraints when loading molecules, leading to more SMILES-to-RDKit transformation failures. Disabling sanitization allows more molecules generated by MolGPT to be loaded successfully, improving its stability metrics.
• Other models were minimally affected because they already generate nearly 100% stable molecules. For these models, the removal of sanitization and kekulization steps does not alter results significantly.

Why Removing Kekulization Has Little Effect? We believe this is because most models already learn to distinguish aromatic bonds from other bonds effectively during training. As a result, their performance remains consistent regardless of the kekulization step.

We have revised the manuscript to include this latest evaluation results. Other than MolGPT, MiDi’s performance also has a slight difference. The newly edited part are highlighted in $\text{\color{red}red}$. Thank you again for your valuable suggestions.

W4: Lack of Discussion on Model Limitations and Future Extensions: Although the model shows improvements, potential challenges such as memory overhead in larger molecules or limitations in certain chemical property predictions are not thoroughly discussed.

Q3: Are there specific molecular properties or types of molecules where NEXT-Mol struggles to perform as well? Identifying any limitations or edge cases where the model’s performance drops would clarify its practical scope.

Response: Thank you for the valuable suggestion. We identify the following limitations of the existing NEXT-Mol model:

• Limited Exploration on Diffusion Guidance. Our DMT model utilizes i.i.d. sampling, without exploring advanced sampling methods like classifier guidance [12] and particle guidance [13]. However, particle guidance [13] demonstrates that a well-tuned guidance method can improve the conformer prediction by 10% precision. This is because the 3D molecular conformational space is large, and a guidance method with appropriate chemical inductive bias can improve the sampling efficiency. We leave this exploration as a future work.
• Computational Cost when Incorporating MoLlama for 3D Conformer Prediction. Incorporating MoLlama, a large LM with 960M parameters, increases training time. For example, training DMT-B alone (55M parameters) takes 52 seconds per epoch on an A100 GPU, while DMT-B with MoLlama takes 210 seconds. We mitigated this problem by using a pretrained DMT-B, instead of training it from scratch, to reduce the training epochs when incorporating MoLlama. Yet, we will need improvement when transferring 1D representations from a large LM.
• Quadratic Memory Complexity of DMT’s Pair Representation. This pair representation incurs an additional $O(N^2)$ GPU memory cost than the standard transformer, compared to the standard transformer’s $O(N)$ memory complexity when using FlashAttention, where N is the node number of molecular graphs. While we encountered no memory issues on the GEOM-DRUGS dataset (molecules with hundreds of nodes), this could be a bottleneck for molecules with thousands of nodes. Potential solutions include smaller batch sizes and model parallelism.

Thank you for the valuable suggestion, we have incorporated this new discussion into our limitation section (Appendix A).

Reference:

[1] STRATEGIES FOR PRE-TRAINING GRAPH NEURAL NETWORKS. In ICLR 2020

[2] PRE-TRAINING MOLECULAR GRAPH REPRESENTATION WITH 3D GEOMETRY. In ICLR 2022.

[3] GraphMAE: Self-Supervised Masked Graph Autoencoders. In KDD 2022.

[4] Llama 2: Open Foundation and Fine-Tuned Chat Models. In Arxiv 2023.

[5] BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding. In NAACL 2019.

[6] Improved denoising diffusion probabilistic models. In ICLR 2021.

[7] Denoising diffusion probabilistic models. In NeurIPS 2020.

[8] Equivariant Diffusion for Molecule Generation in 3D. In ICML 2022.

[9] Swallowing the Bitter Pill: Simplified Scalable Conformer Generation. In ICML 2024.

[10] Do Transformers Really Perform Bad for Graph Representation? In NeurIPS 2021.

[11] Uni-Mol: A Universal 3D Molecular Representation Learning Framework. In ICLR 2023.

[12] Diffusion Models Beat GANs on Image Synthesis. In NeurIPS 2021.

[13] PARTICLE GUIDANCE: NON-I.I.D. DIVERSE SAMPLING WITH DIFFUSION MODELS. In ICLR 2024.

W3: Limited Exploration of Alternative Architectures: The use of RMHA and the specific structure of DMT are well-motivated but not directly compared to alternative architectures. A comparative study could clarify if these design choices are optimal for all molecular generation tasks.

Q4: Could you elaborate on why RMHA was chosen over other potential attention mechanisms? A comparison or justification of this design choice could strengthen the architectural motivation.

Response: Thank you for the suggestion. We agree that including a comparison with alternative architectures would strengthen the presentation. Unfortunately, due to the time limitation in rebuttal and the significant efforts required for re-implementation and hyperparameter tuning, we regret to note that having this comparison is difficult. Below, we answer the other question regarding the elaboration of RMHA’s advantages over other attention mechanisms.

We first revisit the original multi-head self-attention (MHA) and elaborate on how RMHA enhances MHA to incorporate pair information of graphs. Further, we present the limitations of three representative graph attention methods to justify our design choice.

Revisit MHA. Given a input sequence of node representations $\mathbf{H}\in \mathbb{R}^{N\times d}$, the original MHA can be written as follows:

\ $ [\mathbf{Q};\mathbf{K};\mathbf{V}]&=[\mathbf{W}\_q;\mathbf{W}\_k;\mathbf{W}\_v]\mathbf{H}^\top, \quad (1)\\\\ a\_{ij} &= \text{softmax}\_j\left(\frac{\mathbf{Q}\_i\mathbf{K}\_j^\top}{\sqrt{d}}\right), \quad (2) \\\\ \mathbf{O}\_i & = \sum\_{j=1}^N{a\_{ij}}\mathbf{V}\_j, \quad (3) \ $

where $\mathbf{O}_i\in \mathbb{R}^d$ is the output for the i-th input node. However, MHA cannot directly incorporate node pair’s information, such as bond existence or type, limiting its suitability for graph-based tasks.

Revisit RMHA. RMHA addresses the limitation by taking a pair representation $\mathbf{E}\in \mathbb{R}^{N\times N\times d}$ as input, alongside the node representations $\mathbf{H}\in \mathbb{R}^{N\times d}$. $\mathbf{E}$ includes the edges of the fully connected structure of the corresponding graph. Specifically, $\mathbf{E}_{ij}\in \mathbb{R}^d$ encodes pair information like, whether $i$ and $j$ is connected by a bond, bond type (if applicable), and the Euclidean distance. RMHA modifies the MHA as follows:

\ $ [\mathbf{Q}; \mathbf{K}; \mathbf{V}] & = [\mathbf{W}\_{q}; \mathbf{W}\_{k}; \mathbf{W}\_{v}] \mathbf{H}^\top, \quad (4)\\\\ [\mathbf{Q}^E;\mathbf{V}^E] & = \tanh([\mathbf{W}\_{eq};\mathbf{W}\_{ev}] \mathbf{E}^{\top}), \quad (5)\\\\ a\_{i,j}&=\text{softmax}\_{j}(\frac{(\mathbf{Q}^E\_{i,j}\odot \mathbf{Q}\_i){\mathbf{K}}\_j^{\top}}{\sqrt{d}}), \quad (6)\\\\ \mathbf{O}\_i &= \sum\_{j=1}^{N} a\_{i,j}(\mathbf{V}^E\_{i,j}\odot \mathbf{V}\_j). \quad (7) \ $

Here, the pair representation $\mathbf{E}$ modifies the node-level queries ($\mathbf{Q}_i$) and values ($\mathbf{V}_j$) through element-wise multiplication ($\odot$), enabling RMHA to fully incorporate pairwise information:

• Pair information affects attention scores via $(\mathbf{Q}^E_{ij} \odot \mathbf{Q}_i) \mathbf{K}_j^\top$.
• Pair information influences aggregated attention values via $\mathbf{V}^E_{ij} \odot \mathbf{V}_j$.

By incorporating the rich pair information directly into the attention mechanism, RMHA significantly enhances the modeling of molecular graphs, addressing the limitations of MHA.

Limitations of Previous Works.

• Limitations of MCF [9]. MCF is the previous state-of-the-art method for 3D conformer prediction.
  • MCF uses the standard MHA, and incorporates graph structures by including graph eigenvectors as node features. However, graph eigenvectors lack bond type information, which can help discriminate bond lengths.
  • MCF’s implementation uses a fixed subset of all the eigenvectors to handle variable-length molecules. This causes a lossy representation of graph structures, and might demonstrate insufficient when encountering large molecules.
• Limitations of Graphormer [10] and UniMol [11]. Both methods enhance the original MHA by modifying attention scores to include pair representations. However:
  • Their pair representations are constrained to scalars, significantly limiting their representation capacity of the multi-dimensional pair features.
  • Their pair representations cannot influence the aggregated values in attention (see Equation 3 and Equation 7 above).
  • In contrast, RMHA uses vector-based pairwise representations, which are more expressive and influence both attention scores and aggregated values, enabling richer modeling of molecular graphs.

Thank you very much for the suggestion. We have revised Section 3.2 to incorporate the insights from this elaboration.

• DMT-B vs DMT-L.
  • Performance: DMT-L delivers better performance (4.7 % improvement on COV-P Median) than DMT-B, as shown in Table 3 in the manuscript.
  • Inference Time: As shown by the Table above, DMT-L takes about 3 times as long as DMT-B for conformer generation on the GEOM-DRUGS dataset.
• DMT vs DMT + MoLlama.
  • Performance: Adding MoLlama improves 3D conformer predictions by ~1% in COV-R for both DMT-B and DMT-L on GEOM-DRUGS (see Table 6 in the manuscript).
  • Inference time: Incorporating MoLlama increases inference time. For example, DMT-B’s inference time increases from under 200 minutes to 467 minutes.
• Conclusion. The choice between DMT-B and DMT-L, as well as the inclusion of MoLlama, depends on task-specific requirements, particularly the trade-off between computational efficiency and 3D conformer prediction accuracy
• Following your suggestion, we have revised Appendix B.3 to include a comparison of the inference time of our model and some representative baselines.

Different noise schedules at inference time. We test DMT-B’s robustness to different noise schedulers at inference, using two representative options: the linear [7] and polynomial [8] schedulers. The original noise scheduler, based on the cosine function, follows [6]. Due to time constraints during the rebuttal, we use the existing DMT-B checkpoint without retraining the model with these new schedulers, so the results are suboptimal.

Table 11: Performance of DMT-B on the GEOM-DRUGS dataset with different noise schedulers at inference.

As shown in Table 11, the polynomial scheduler achieves performance close to the cosine scheduler, likely because their curve shapes are similar. However, the linear scheduler results in a significant performance drop, suggesting that retraining DMT-B with the linear scheduler is necessary to achieve better results.

The Influence of Batch Size to 3D Conformer Prediction. We evaluate the performance of DMT-B with different batch sizes. The original batch size of 256 was chosen to maximize GPU utilization. To assess the impact of batch size, we tested two variations: (1) reducing the batch size to 128, and (2) increasing it to 512 using gradient accumulation.

Table 12: DMT-B’s 3D conformer prediction performances on the GEOM-DRUGS dataset when using different batch sizes.

As shown in Table 11, the performance with a 512 batch size is slightly worse than the original model. This is likely due to underfitting caused by fewer training steps. We keep the number of training epochs the same as the original experiment (256 batch size, epoch 3000), therefore the larger batch size results in fewer gradient updates, leading to reduced model performance. Other than this observation, using the 128 batch size does not lead to a significant difference than the original model.

Thank you for your suggestion. We have revised Appendix B.2 of our submission to include these experiments, and they have significantly improved our presentation.

Thank you very much for the positive comments and valuable feedbacks on our submission. Here we provide a point-by-point response to address your concerns and incorporate your suggestions. For clarity, we have consolidated similar questions and weaknesses where appropriate, labeling them with corresponding identifiers (e.g., W1 for Weakness 1 and Q1 for Question 1) based on your original comments.

W1: Limited Theoretical Insight: The paper lacks a theoretical explanation of why combining 1D and 3D modeling via transfer learning improves performance. Further theoretical analysis could provide deeper insights into the effectiveness of this architecture and potential limitations.

Q1: Could you provide a theoretical justification for transfer learning between 1D molecular sequences and 3D conformers? An in-depth explanation could clarify why this cross-modal transfer is effective.

Response: Thank you for the suggestion. We agree that a theoretical analysis for this 1D to 3D transfer learning can improve the rationality and presentation of our submission. While a comprehensive theoretical analysis is beyond the scope of this work, we offer an explanation combining theoretical insights and empirical observations.

The Rationale behind transfer learning between 1D molecule sequences and 3D conformers. The final goal of this transfer learning is to leverage the billion-scale 1D/2D molecule dataset to improve the 3D conformer prediction performance, which is constrained by limited 3D data. For clarity, we decompose the rationale into the following chain of arguments:

1. 3D conformers are theoretically governed by 2D molecular graphs under quantum mechanics (QM). 3D molecular properties and structures are fundamentally rooted in QM. Using (approximated) QM-based methods, like DFT, we can accurately predict 3D conformers from 2D molecular graphs, though at a high computational cost. This establishes the critical role of 2D representations in determining 3D structures.
2. 3D conformer prediction relies on high quality 2D molecule representations. Deep learning models predict 3D conformers from 2D graphs, and their performance is heavily influenced by the quality of 2D molecular representations. Transfer learning can enhance 2D molecular representations, as demonstrated by prior works [1,2,3].
3. 1D molecular representations can be converted to 2D molecular representations, and contribute to 3D prediction. 1D molecule sequences encode the same information as 2D molecular graphs, and the 1D to 2D transformation can be achieved by a deterministic toolkit, like RDkit. Leveraging RDkit and our proposed cross-modal projector, which is detailed in Section 3.3 and Appendix C.3, we can transform 1D molecular representations to 2D molecular representations, and therefore contribute to the 3D prediction. We have demonstrated this improvement in Table 4, where using the pretrained 1D representations improves 3D conformer prediction.
4. 1D pretraining scales more effectively than 2D. Given the billion-scale 1D/2D molecule dataset, we mostly prioritize scalability when selecting the pretraining method. After the literature review, we find that 1D LM-based pretraining methods, like Llama [4] and BERT [5], are extensively demonstrated for scalability and effectiveness. Therefore, we opt to 1D pretraining instead of 2D pretraining.

We have included this discussion in Appendix C.3 for enhanced presentation.

W2: Absence of Ablation Studies on Model Size and Hyperparameters: While the paper shows promising results with two model sizes, a more detailed examination of how model size or key hyperparameters (e.g., noise schedule, batch size) impact performance would provide more guidance on model tuning.

Q2: What are the potential computational trade-offs for using larger models (DMT-L) in terms of scalability and inference speed? Including a computational analysis of DMT-B versus DMT-L could highlight the scalability limits.

Response: Thank you for the suggestion. In response, we discuss the influence of model size, noise schedule, and batch size on the 3D conformer prediction task.

Computational Trade-offs: DMT-B vs. DMT-L and MoLlama Integration. We compare the performance and inference time of DMT-B, DMT-L, and their enhanced versions with MoLlama (DMT + MoLlama) using 1000 molecules from the GEOM-DRUGS test set.

Table: Comparison of conformer generation time on the test set of the GEOM-DRUGS dataset using various methods.

Thank you very much for your thoughtful feedbacks and detailed review comments. Below, we provide point-by-point responses to address your concerns.

W1: The comparison results of 3D molecular generation in the paper are unfair because the authors completely ignore the influence of 3D structures on molecular representation during the generation process. The results of molecular generation should be compared with generation models based on 1D molecular representations.

Response: Thank you for the thoughtful feedback. In response, we (1) add new baselines using 1D language models (LMs) for 1D/2D molecule generation, (2) clarify that our comparisons are fair since we focus on baselines that can generate 3D molecules, which previous 1D LMs cannot, and (3) clarify that our model considers 3D structures through the 3D conformer generation process.

Table 2: Performances for de novo 3D molecule generation. We highlight the results of new 1D LM baselines in $\textcolor{blue}{\text{blue}}$. * denotes our reproduced baseline results using their source codes.

(a) Performances on the GEOM-DRUGS dataset. New baselines are evaluated for 2D-metric only.

(b) Performances on the QM9-2014 dataset. New baselines are evaluated for 2D-metric only.

New baselines of 1D molecule LMs. To address your concern, we include MolGPT [1] and MolGen [2] as new baselines for de novo 1D/2D molecule generation on the GEOM-DRUGS and QM9-2014 datasets. MolGPT is a decoder-only transformer trained on SMILES sequences, and MolGen is an encoder-decoder transformer trained on SELFIES sequences. For MolGen, we fine-tune its public checkpoint on the two datasets; for MolGPT, we first reproduce its pretrained model using its source code, and then fine-tune it on the two datasets. Note that, we do not directly compare with MolGen’s reported performances, because MolGen employs a different experimental setting to fill in corrupted spans in SELFIES sequences from the test set, which is unusual. We report only 2D-metric for these two baselines because they cannot generate 3D molecules.

New comparison between TD w/ PD and DMT.

• Reproduced results for TD w/ PD. Following your feedback, we reproduce the TD w/ PG baseline (denoted as TD w/ PG*) and include the results in Table 3(a). The reproduced results differ from the originally reported values. We include both, graying out the original results as we could not replicate them. We assure you that we used their source code and public checkpoint. For transparency, we upload the generated 3D conformers and the evaluation log into our anonymous share link OSF, under the “onedrive/TD with PG Results” folder.
• New comparison to TD w/ PD*. Using the reproduced results, we can observe that DMT-B achieves comparable COV-P to TD w/ PD*, while outperforming it by 11% in COV-R. Additionally, DML-L outperforms TD w/ PD* across all metrics, demonstrating our model’s scalability. We hope these new results highlight the effectiveness of our DMT model, and address your concerns.
• Particle Guidance (TD w/ PD) [3] is an orthogonal contribution to our method. Particle guidance is a diffusion guidance method like the classifier guidance [4]: it modifies torsional diffusion’s original i.i.d. sampling process. Our DMT model uses i.i.d sampling, and focuses on improving neural architectures and transfer learning from 1D data. Therefore, our contributions are orthogonal.
• Typo in original submission. Thank you for noticing the typo: we have incorrectly bolded the best COV-P performance in Table 3(a) in the original submission, which is now fixed.

Thank you for your thoughtful observations. We have updated Table 3(a) and Section 4.3 in our submission to address your suggestions, which we believe have significantly improved the presentation.

Q1: I think the authors should provide new evidence of the advantages of this method as a 3D model for molecular generation.

Response: Thank you for the suggestion. Following your suggestions, we have provided the following new evidence to support NEXT-Mol’s advantage for 3D molecule generation.

• New results for 1D LM. In our previous response to weakness 1 (W1), we include new baselines of MolGPT [1] and MolGen [2], and demonstrate NEXT-Mol outperforms these 1D LM-based molecule generation methods.
• New comparison with TD w/ PG*. In our previous response to weakness 2 (W2), we include a new comparison between DMT and our reproduced TD w/ PG*. We show that DMT-B outperforms TD w/ PG* by 11% in COV-R, and achieves comparable COV-P scores. Additionally, DMT-L outperforms TD w/ PG* across all metrics.

All these results are incorporated in our revised submission, which contributes to improving the empirical evidence of our methodology.

Reference:

[1] MolGPT: Molecular Generation Using a Transformer-Decoder Model. In Journal of Chemical Information and Modeling 2021.

[2] Domain-Agnostic Molecular Generation with Chemical Feedback. In ICLR 2024.

[3] PARTICLE GUIDANCE: NON-I.I.D. DIVERSE SAMPLING WITH DIFFUSION MODELS. In ICLR 2024.

[4] Diffusion Models Beat GANs on Image Synthesis. In NeurIPS 2021.

As shown in Table 2(a) and 2(b) above, NEXT-Mol consistently outperforms the new baselines. Specifically, NEXT-Mol achieves relative FCD score improvement by 49% over MolGen and 62% over MolGPT on the GEOM-DRUGS datasets. These significant improvements are attributed to NEXT-Mol’s benefits from the scaling law: it is pretrained on a large dataset of 1.8 billion molecules, and features a large model size of 960M parameters. Additionally, NEXT-Mol will be released to support open-source research.

We have revised our submission to include these new results.

Our existing benchmark is fair because it focuses on baselines capable of generating 3D molecules. We initially excluded 1D LMs as baselines since they cannot generate 3D molecules. However, following your recommendation, we have now included them to enhance the presentation. Additionally, in the de novo molecule generation task, we report both 2D metrics, which evaluate the generated 2D structures, and 3D metrics, which evaluate the generated 3D structures. This ensures fairness by comparing 2D generation methods to NEXT-Mol’s 2D outputs and 3D generation methods to NEXT-Mol’s 3D outputs.

Our model considers 3D structures during 3D generation through the 3D conformer generation process. As elaborated in the introduction, we explore a two-step solution for 3D molecule generation: initially generating a 1D molecule (a subset of a 3D molecule) using an LM and subsequently predicting its 3D conformer using a diffusion model. Importantly, the generation of a 1D molecule is not the final step in 3D molecule generation. This is because a molecule can have many possible 3D conformers, therefore accurately predicting the desired 3D conformer is crucial for completing the 3D molecule generation.

Thank you for your feedback. We have revised the introduction accordingly to improve the clarity.

W2: The improvement in molecular conformation generation results mentioned in the paper is actually very small. In Table 3(A), DMT-B on GEOM-Drugs only improves by 1.4% compared to MCF-B, COV-R. Compared to Par. Guid, COV-P decreases by 3.7%, but the model parameters increase by 30 times. The significance of this method is limited.

Response: Thank you for the feedback. We address your concerns by (1) highlighting that DMT achieves better performance with fewer parameters than MCF, and (2) revising our comparison with particle guidance (TD w/ PG) [3] to include reproduced baseline performance and fix a typo. The updated results for 3D conformer prediction on the GEOM-DRUGS dataset are shown in Table 3(a) below.

Table 3(a): 3D conformer prediction results. * denotes reproduction using their codes. -R←Recall and -P←Precision. GEOM-DRUGS dataset. $\text{\color{blue}\small TD w/ PG denotes torsional diffusion with particle guidance.}$ We omit some baselines in the original submission for brevity.

DMT outperforms MCF with fewer parameters. The comparison between DMT and MCF should consider both our better performances and our smaller model sizes. Specifically, DMT-B outperforms MCF-B by 1.4% in COV-R while using only 85% of its parameters, and DMT-L outperforms MCF-L by 1.1% in COV-R with just 61% of its parameters. This demonstrates that DMT achieves better performances with more efficient parameter utilization than MCF, which is a significant contribution to the field.

Dear Reviewer boXt,

Thank you for your insightful feedbacks on our submission, especially for advising us to 1) adding more comparisons to 1D molecular LMs, and 2) demonstrating the improvements compared to previous 3D conformer prediction methods. These valuable suggestions have improved the significance and clarity of our work. We hope that these improvements will be taken into consideration.

Beyond these updates above, we have also conducted new experiments on molecule property prediction, detailed below, to provide further evidence of the advantages of our approach.

If our responses have adequately addressed your concerns, we would be grateful if you could reconsider the rating of our paper. Of course, we are more than happy to engage in further discussions if you have additional questions or comments.

Molecule Property Prediction Results. To evaluate MoLlama's capabilities beyond 1D molecule generation, we apply it to molecular property prediction tasks, highlighting the quality of its molecular representations. Following the setup in [1], we fine-tune MoLlama on four MoleculeNet datasets: FreeSolv, ESOL, Lipo, and QM7. We adopt the same experimental settings and dataset splits as [1], reporting mean performance and standard deviation over 10 random seeds. For each run, MoLlama is trained for 100 epochs, with test performance selected based on the validation dataset. We use a fixed learning rate of 1e-4 with the AdamW optimizer, and fine-tune MoLlama using LoRA (LoRA r=8 and $\alpha=32$) applied to all linear layers of the model. Following Section 3.3, we attach a single-layer bi-directional self-attention layer after MoLlama to improve its encoding ability. After that, we apply a linear layer on the mean embedding of all molecule tokens for property prediction.

Table 10: Molecule property regression results on four MoleculeNet datasets. Baseline results are from [1]. Lower $\downarrow$ is better.

Observation. As shown in Table 10, MoLlama outperforms baseline methods, achieving relative improvements of 6.5%, 4.7%, 3.5%, and 16.9% on the FreeSolv, ESOL, Lipo, and QM7 datasets, respectively. Notably, our baselines include LM-based, GNN-based, and pretrained GNN-based methods, and MoLlama's better performance demonstrates its advantages derived from the extensive pretraining.

Reference:

[1] MolPROP: Molecular Property prediction with multimodal language and graph fusion. In Journal of Cheminformatics 2024.

Settings. We implemented the Predictor-Corrector (PC) sampler [1] to further enhance DMT’s performance in 3D conformer prediction. Specifically, we used the original checkpoint without retraining and modified only the inference code. By tuning the snr hyperparameter, the PC sampler enables a trade-off between diversity (recall) and precision. We experimented with snr values of {0.2, 0.3, 0.4}, observing that higher snr values lead to greater precision but reduced diversity.

Observations. The PC sampler can improve DMT-B’s sampling precision with minimal impact on recall. For example, DMT-B with the PC sampler (snr=0.3) significantly outperforms our original vanilla i.i.d. sampler and the TD w/ PD* baseline across all metrics, achieving a 2.4% increase in COV-P and an 11.8% increase in COV-R. However, increasing snr to 0.4 further boosts precision but causes a noticeable decline in recall. These results highlight the utility of hyperparameter tuning and may guide future research.

Why We Did Not Implement Particle Guidance? This experiment demonstrates that DMT’s performance can be further improved using an advanced sampler, similar to how Particle Guidance [2] enhances torsional diffusion. However, we did not use Particle Guidance because its original implementation is designed for diffusion on torsion angles. Adapting it to our method, which operates in the full Euclidean space, is non-trivial and requires significant theoretical and engineering effforts.

Further Discussion. We hope these new results underscore the effectiveness of our methods and provide valuable insights for future work. We look forward to continued discussion and collaboration.

Reference:

[1] SCORE-BASED GENERATIVE MODELING THROUGH STOCHASTIC DIFFERENTIAL EQUATIONS. In ICLR 2021

[2] PARTICLE GUIDANCE: NON-I.I.D. DIVERSE SAMPLING WITH DIFFUSION MODELS. In ICLR 2024.

Q1: What is the rationale for combining 1D and 3D generation sequentially, and what benefits does this approach offer?

Response: Unlike previous works, NEXT-Mol decomposes 3D molecule generation into two steps: first generates 1D molecular sequences using MolLama (a large language model), and then conduct by 3D conformer prediction using a 3D diffusion model. The main motivation is to leverage the following two advantages of 1D generation:

1. 1D generation guarantees 100% molecular validity: Achieving 100% validity offers critical advantages beyond increasing validity itself:
   • Improving validity could improve other 2D metrics, like SNN, Frag, and Scaf (see Table 2). These metrics measure the distributional similarity of 2D molecular structures of valid molecules. If a model still generate invalid molecules, it is likely the model does not capture the true target distribution, which contain only valid molecules. 100% validity helps the model learn from and sample from the valid molecular structures, which is essential for molecule generation tasks.
   • The improved validity also leads to better learning of 3D molecular geometry, because it grounds 3D structure prediction on valid 2D structures. Other joint 2D and 3D prediction methods [15,16] can easily encounter invalid 2D structures when sampling 3D structures, therefore leads to worse 3D structure prediction.
2. 1D generation can be effectively large-scale pretrained:
   • Our large-scale pretraining leverages 1.8B molecules of 90B SELFIES tokens. It enables better learning of 1D/2D molecular patterns, compared to other methods that are constrained to the GEOM-DRUGS and QM9-2014 datasets, which are limited in size.
   • The pretrained 1D LM can be used to improve 3D conformer prediction by transfer learning. Our proposed cross-modal projector effectively transfers MolLama's molecular knowledge to guide 3D prediction. Table 4 shows consistent improvements across multiple 3D conformer prediction metrics when incorporating MolLama representations

This sequential pipeline effectively combines the strengths of SELFIES language modeling and 3D diffusion for accurate molecular generation.

Reference:

[1] Equivariant Diffusion for Molecule Generation in 3D. In ICML 2022

[2] Learning Joint 2D & 3D Diffusion Models for Complete Molecule Generation. In TNNLS 2024.

[3] Equivariant Energy-Guided SDE for Inverse Molecular Design. In ICLR 2023

[4] NAVIGATING THE DESIGN SPACE OF EQUIVARIANT DIFFUSION-BASED GENERATIVE MODELS FOR DE NOVO 3D MOLECULE GENERATION. In ICLR 2024

[5] GuacaMol: Benchmarking Models for De Novo Molecular Design

[6] Recent progress on exciplex-emitting OLEDs. In Journal of Information Display 2019.

[7] STRATEGIES FOR PRE-TRAINING GRAPH NEURAL NETWORKS. In ICLR 2020

[8] PRE-TRAINING MOLECULAR GRAPH REPRESENTATION WITH 3D GEOMETRY. In ICLR 2022.

[9] GraphMAE: Self-Supervised Masked Graph Autoencoders. In KDD 2022.

[10] Llama 2: Open Foundation and Fine-Tuned Chat Models. In Arxiv 2023.

[11] BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding. In NAACL 2019.

[12] Swallowing the bitter pill: Simplified scalable conformer generation. In ICML 2024.

[13] E(n) Equivariant Graph Neural Networks. In ICML 2021. [14] Accurate structure prediction of biomolecular interactions with AlphaFold3. In Nature 2024.

[15] Learning Joint 2D & 3D Diffusion Models for Complete Molecule Generation. In TNNLS 2024.

[16] MiDi: Mixed Graph and 3D Denoising Diffusion for Molecule Generation. In ECML 2023.

W4: Limited ML Novelty: The model presents minimal innovation from an ML perspective, as it mainly combines existing components—LLaMA and diffusion models. This combination, particularly in transferring 1D information to 3D, offers limited novelty and benefit for the conformer generation part.

Response: We respectfully disagree and argue that our model introduces significant ML novelties. Below, we outline the key novel contributions of our method:

1. Novel Diffusion Molecular Transformer (DMT) Architecture:
   • Our DMT introduces a novel architecture that preserves complete 2D molecular graph information through Relational Multi-Head Attention (RMHA). RMHA extends standard self-attention by incorporating pairwise information that describes atomic interactions.
   • DMT achieves superior performance while using fewer parameters than previous approaches (55M vs. MCF's 64M). Table 2 demonstrates significant improvements: FCD score improved from 0.655 to 0.334 on GEOM-DRUGS, with substantial gains in bond length, angle, and torsion MMD metrics
2. Innovative 1D-to-3D Transfer:
   • Our cross-modal projector and training strategy effectively leverages MolLama's 1D representations to enhance DMT's 3D predictions. This approach uniquely addresses 3D data scarcity by utilizing abundant 1D molecular sequences. Table 4 shows consistent performance on the 3D conformer prediction task on both GEOM-DRUGS and GEOM-QM9 datasets after incorporating MolLama representations. Table 6 shows consistent improvement on the 3D molecule generation task.
3. Large 1D Molecule LM:
   • To our knowledge, our MoLlama, which has 960M parameters and being pretrained on 1.8B molecules, is the largest SELFIES-based autoregresive LM for 1D molecule generation. It presents a significant contribution to the community.
4. Versatile Multi-task Learning:
   • Our NEXT-Mol is a versatile multi-task learner and uniquely excels in all three tasks: de novo 3D generation, conformer prediction, and conditional 3D generation, as demonstrated in our comprehensive evaluation, while previous works typically specialize in only one of these tasks.

Other Reviewers’ Opinion: we quote other reviewers’ comments regarding the novelty of our approach.

• Reviewer ZnjA: “Innovative Approach: The combination of a 1D language model with a 3D diffusion model is an original solution to ensuring chemical validity while efficiently generating 3D conformers.”
• Reviewer  iSmq: “Novel Methodology: The integration of 1D molecule generation through SELFIES with conformer prediction addresses the limitations of current datasets and methods in 3D molecule generation.”

Summary: Our work combines novel architectural components, innovative 1D-to-3D transfer learning, and large-scale pretraining to tackle critical challenges in 3D molecular generation. The model’s versatility, scalability, and demonstrated improvements across diverse tasks highlight its significant ML contributions. We hope this clarifies the novelty of our approach.

W3: Not accurate statement on Rotation Augmentation: The statement, “Following AlphaFold3 (Abramson et al., 2024), we apply random rotation augmentation on 3D conformers to help DMT obtain equivariance to rotated inputs by learning. While (Wang et al., 2024) report decreased performance given random rotations, DMT benefits from it, potentially due to the improved neural architecture,” is unclear. The authors imply that rotation can be achieved without using equivariant networks (those maintaining symmetry under rotation). However, AlphaFold claimed that it is not necessary to have equivariant networks, so it is essential to clarify how DMT benefits from rotation augmentation and to distinguish it from learned equivariance.

Response: Thank you for your question regarding equivariance. In this work, we adopt random rotation augmentation on 3D conformers to help DMT learn equivariance. We would like to clarify the following points:

1. DMT uses learned equivariance instead of using built-in equivariant networks: Our DMT model obtains equivariance to rotations by applying random rotation augmentation: we apply the same random rotation on the input 3D coordinates and the prediction target 3D coordinates. Therefore our model learns to output 3D coordinates that change equivariantly with the input. We do not use build-in equivariant networks, such as EGNN [13], which can achieve equivariance without any training.
2. Our rotation augmentation is consistent with AlphaFold3: We clarify that AlphaFold3 can also learn equivariance to rotations because it uses random rotation augmentation (see Algorithm 19 in their Appendix). We clarify that, the claim in their paper “we find that no invariance or equivariance with respect to global rotations and translation of the molecule are required in the architecture”, means that they do not use built-in equivariant architecture, like EGNN [13].
3. Clarification on the benefits of rotation augmentation: As shown in Table 7 below, the performance of DMT-B is improved with random rotation augmentation, which proves the effectiveness of this technique.

Table 7: Ablating random rotation augmentation for 3D conformer prediction on GEOM-QM9.

Thank you for the feedback. We have revised our description about rotation augmentation in Section 3.2 to improve clarity.

W2: 1D-to-3D Transformation Claim: The claim that converting a 1D sequence to 3D adds value is questionable, as the graph already provides all necessary information.

Q2: Could you clarify how transferring a 1D representation contributes to the overall model performance?

Response: Thank you for the feedback. We agree that 1D sequences and 2D molecular graphs convey the same information. However, the key here is transfer learning: the 2D molecular graphs do not have the knowledge learned by the 1D molecule LM from a billion-scale dataset, but this knowledge can be encoded in 1D representations generated by the 1D LM. Our transfer learning is to use these knowledge learned from a billion-scale dataset to help 3D conformer prediction, which is constrained by limited data. Below, we elaborate the rationale behind this transfer learning.

Rationale behind transfer learning between 1D molecule sequences and 3D conformers. The final goal of this transfer learning is to leverage the billion-scale 1D/2D molecule dataset to improve the 3D conformer prediction performance, which is constrained by limited 3D data. For clarity, we decompose the rationale into the following chain of arguments:

1. 3D conformers are theoretically governed by 2D molecular graphs under quantum mechanics (QM). 3D molecular properties and structures are fundamentally rooted in QM. Using (approximated) QM-based methods, like DFT, we can accurately predict 3D conformers from 2D molecular graphs, though at high computational cost. This establishes the critical role of 2D representations in determining 3D structures.
2. 3D conformer prediction relies on high quality 2D molecule representations. Deep learning models predict 3D conformers from 2D graphs, and their performance is heavily influenced by the quality of 2D molecular representations. Transfer learning can enhance 2D molecular representations, as demonstrated by prior works [7,8,9].
3. 1D molecular representations can be converted to 2D molecular representations, and contribute to 3D prediction. 1D molecule sequences encode the same information as 2D molecular graphs, and the 1D to 2D transformation can be achieved by deterministic toolkit, like RDkit. Leveraging RDkit and our proposed cross-modal projector, which is detailed in Section 3.3 and Appendix C.3, we can transform 1D molecular representations to 2D molecular representations, and therefore contribute to the 3D prediction. We have demonstrated this improvement in Table 4, where using the pretrained 1D representations improve 3D conformer prediction.
4. 1D pretraining scales more effectively than 2D. Given the billion-scale 1D/2D molecule dataset, we mostly prioritize the scalability when selecting the pretraining method. After literature review, we find that 1D LM-based pretraining methods, like Llama [10] and BERT [11], are extensively demonstrated for scalability and effectiveness. Therefore, we opt to 1D pretraining instead of 2D pretraining.

Thank you for your valuable comments and insightful feedback on our submission. Below, we offer a detailed, point-by-point response to address your concerns and integrate your suggestions. For clarity, we have grouped similar questions and identified weaknesses where applicable, assigning corresponding labels (e.g., W1 for Weakness 1 and Q1 for Question 1) based on your original comments.

W1: Focus on Core Objective: In the title the paper is presented as a generative model for molecules. The model is benchmarked as a conformer generating model. Currently, benchmarks appear misaligned from the title; consider established benchmarks like CheckPose or DrugPose for 3D generation.

Response: Thank you for the feedback. For response, we (1) clarify that we have already employed 3D generation benchmarks of de novo 3D molecule generation and conditional 3D molecule generation, (2) explain the rationale of using 3D conformer prediction as an additional benchmark, and (3) justify that our benchmarks and contents align with our title.

• We have already included 3D molecule generation benchmarks. Our method are evaluated on two-well established benchmarks:

  • De novo 3D generation (Section 4.2): this task complements virtual screening by providing a candidate molecule space of manageable size [5], enabling the exploration only on the mostly relevant molecules.
  • Conditional 3D molecule generation (Section 4.4): This task targets at quantum chemical properties (e.g., HOMO and LUMO), which are useful for designing new material [6].

  Both are well-established benchmarks in 3D molecular generation used by many previous works [1,2,3,4] and demonstrate the practical value of our approach.

• Why not PoseCheck or DrugPose? 3D molecule generation spans many tasks, such as antibody generation, de novo 3D molecule generation, conditional 3D molecule generation targeting quantum chemical properties, and structure based molecule generation (e.g., PoseCheck and DrugPose). In this work, we focus on de novo 3D molecule generation and conditional 3D molecule generation targeting quantum chemical properties. The benchmarks of PoseCheck and DrugPos involve different evaluation metrics and different baselines, making them outside the scope of this work. They can be studied as future works, as we mentioned in the Conclusion section.

• Content alignment with the title. We are willing to revise the title to “NEXT-Mol: 3D Diffusion Meets 1D Language Modeling for 3D Molecule Generation and Conformer Prediction”, if you agree that this new title can better reflect our used benchmarks on both 3D molecule generation and 3D conformer prediction. Our existing title “NEXT-Mol: 3D Diffusion Meets 1D Language Modeling for 3D Molecule Generation” can accurately reflect our used benchmarks of de novo 3D molecule generation and conditional 3D molecule generation, similar to the following recent studies:

  • Equivariant Diffusion for Molecule Generation in 3D. In ICML 2022.
  • GeoLDM: Geometric Latent Diffusion Models for 3D Molecule Generation. In ICML 2023.
  • Learning Joint 2D & 3D Diffusion Models for Complete Molecule Generation. In TNNLS 2024.
  • MiDi: Mixed Graph and 3D Denoising Diffusion for Molecule Generation. In ECML 2023.
  • Navigating the Design Space of Equivariant Diffusion-Based Generative Models for De Novo 3D Molecule Generation. In ICLR 2024.

  All the works above use the same benchmarks (de novo 3D molecule generation or conditional 3D molecule generation) as ours, and share the keywords of “molecule generation” and “3D”.

• Rationale of using 3D conformer prediction as an additional benchmark. The performance of 3D conformer prediction is critical to our final de novo 3D molecule generation performance. This is because our method is a two-step solution for 3D molecule generation: initially generating a 1D molecule (a subset of a 3D molecule) using an LM and subsequently predicting its 3D conformer with a diffusion model. Thus, we include 3D conformer prediction as an additional benchmark to evaluate the effectiveness of our second step.

Dear Reviewer pyyj,

Thank you for your thoughtful feedbacks on our submission, especially for advising us to 1) clarifying the reasoning of choosing the existing 3D molecule generation benchmarks, 2) clarifying the content alignment with the title, and 3) clarifying the rationale behind 1D-to-3D transfer learning. These valuable suggestions have improved the clarity and quality of our work. We hope that these improvements will be taken into consideration.

If our response has resolved your concerns on our paper, we will greatly appreciate it if you could re-evaluate our paper for a higher rating. We are also willing and ready to engage in discussions, if you have any further questions.

Followup Question. Could you please provide concrete examples where transfer learning demonstrably improves 3D generation?

My concern is that the 1D generation aspect is not novel, as it has been explored previously and is associated with well-documented failure points. While the 3D generation shows some improvement based on evaluation metrics, I am not convinced that new information is being learned by the model. This is because the conformer generation step already leverages the graph representation as input, which already contains the required information for the prediction.

Could you clarify which specific failure modes of graph-based models are being addressed or mitigated by your approach?

Response: Thank you for the follow-up discussions. Here we elaborate on the advantages of transfer learning from 1D to 3D generation in two aspects:

Argument 1: Pretrained 1D representations improves 3D conformer prediction. We understand this to be the reviewer’s primary concern. Our ablation study in Table 5 demonstrates improvements in conformer prediction evaluation metrics, validating this claim. To further support our argument, we provide additional analysis and concrete examples that highlight the benefits of incorporating pretrained 1D representations.

Table 6: 3D conformer prediction performance on GEOM-DRUGS’s test subsets, split by scaffold frequency in the training set. 68 low-quality samples are filtered following [1]. We report mean values of AMR-R and AMR-P.

• MoLlama's 1D Representation Improves 3D Prediction for Unseen Scaffolds. A molecule scaffold is the core structure of a molecule. Scaffold split, which ensures no scaffold overlap for molecules between the training and testing sets, is widely used to evaluate a molecular model's generalization ability to unseen structures [2]. To evaluate this, we divide GEOM-DRUGS's test set into subsets based on the test molecule's scaffold frequency in the training set. The results are shown in Table 6, we can observe that:
  • DMT’s failure mode. DMT-B's performance drops significantly for test molecules with unseen scaffolds: AMR-R and AMR-P increase by 0.086 and 0.236, respectively, compared to molecules with scaffold frequency $\geq$ 1.
    • Why does this failure happen? This is because deep learning methods are typically trained to mimic the training data, and often struggle when faced with distinct data distributions than training [2]. For molecular scaffolds absent from GEOM-DRUGS’s training set, DMT lacks sufficient prior knowledge to make accurate predictions. Even though all the bonds and atoms information are already provided in the molecular graph, the ability (1) to represent the 2D molecular structure, and (2) to decode 3D structures from the 2D structures, can still constrain a 3D conformer prediction model’s performance.
  • Transfer learning with MoLlama can mitigate this failure. Incorporating MoLlama can mitigate this issue, reducing AMR-R and AMR-P by 0.028 and 0.030, respectively, for molecules with unseen scaffolds.
    • Reason for MoLlama’s Improvement. This improvement is attributed to MoLlama’s pretraining on a large molecular dataset containing 1.8 billion molecules, compared to the 230k molecules in the GEOM-DRUGS training set. This extensive pretraining exposes MoLlama to a far more diverse range of scaffolds, enabling it to generate molecular representations that generalize more effectively. In contrast, DMT is limited to the scaffolds present in the smaller GEOM-DRUGS training set. The broader scaffold knowledge provided by MoLlama’s pretraining allows for better performance on unseen scaffolds.
  • Insights for seen scaffolds. For test molecules with seen scaffolds, MoLlama’s improvement is minimal. This is because DMT can effectively learn the 3D structure prediction of these molecules from similar examples in the training data.
  • **\textcolor{red**{Visualization of concrete examples}.} We have updated Figure 4 in Section 4 and Figure 7 in Appendix B.7 to include visual examples where DMT+MoLlama outperforms DMT, focusing on cases with scaffolds absent from the training set. Notably, the primary improvements are observed in torsion angle predictions, even though our design does not explicitly target torsions. This may be because bond lengths and angles are easier to predict, while torsion angles pose greater challenges.