Geometry Informed Tokenization of Molecules for Language Model Generation

Dear Reviewer CPUN,

Thank you for your appreciation of our work and insightful comments! We have made efforts to thoroughly improve our work accordingly and provide responses for each concern here. Please also refer to our added experiments in this Link and our responses to other reviewers.

Some existing molecule 3D molecule tokenizing methods

• Thank you for the advice. We have included the discussions in the paper revision and also briefly discuss here.
• 3D-MolT5: Towards Unified 3D Molecule-Text Modeling with 3D Molecular Tokenization: 3D-MolT5 focuses on text-based molecule-related downstream tasks, including molecular property prediction, 3D molecule captioning, and text-based molecule generation. It does not consider 3D molecular generation tasks. thus we do not adopt it as a baseline method. 3D-MolT5 is designed to handle 3D-dependent tasks with text instructions using the T5 model. Its tokenization method is based on the Extended 3D Fingerprint (E3FP) algorithm, where the embeddings of the same atom in both 1D and 3D tokens are summed to form the final joint representation.
• Tokenizing 3D Molecule Structure with Quantized Spherical Coordinates: This is a concurrent work with our submission, submitted to arXiv in Dec 2024, thus we do not adopt it as a baseline method. This work uses SMILES and coordinates to build sequences and a VQ-VAE model to discretize the continuous coordinates. Compared to our method, this work uses VQ-VAE to learn structure tokens, which lacks guarantee of structural completeness.
• In summary, our method differs by (1) extending canonical labeling to encode 3D structural isomorphism, (2) enabling reversibility between sequences and 3D isomorphic structures, and (3) establishing theoretical guarantees of structural completeness and geometric invariance. We believe our formulation complements theirs, and we thank the Reviewer for encouraging this discussion.

W1: Scalability of Geo2Seq

• Thank you for the point. This is an interesting question we have been thinking as well. Indeed, discretization could impose limits on very large molecules due to increased sequence length and reduced resolution. However, on GEOM-DRUGS dataset, Geo2Seq still achieves competitive performances. Moreover, theoretically, if we use larger vocabulary size, the discretization would not be a limitation for scalability. Larger molecules only bring longer context length of upto 750, which can be handled given the capability of LLMs.
• To our understanding, the reason for the suboptimal performances on GEOM-DRUGS could be that the size of the GEOM-DRUGS dataset requires larger LMs than what we are using. The QM9 data size is ~100k and we are using ~90M parameter LMs, while GEOM ~7M and we are using ~100M parameter LMs. While this benefits efficiency, GEOM-DRUGS might perform optimally with larger model sizes. Due to the time limit of the rebuttal, we will explore further improving the performance on GEOM-DRUGS dataset with larger LLMs in the future.

Paper Representation

• Thank you for the suggestion! We have shortened the discussion in Sec.3 and moved essential experimental ablations including Table 4 to the main text.

Typos

• Thank you for pointing out! We have revised the paper and corrected the typos.

Q1: train from scratch or pre-trained?

• We train from scratch. We propose a molecular tokenization method different from the NLP tokenization used by GPT and Mamba. Pre-trained checkpoints need to be used together with the pre-training tokenizer, thus not applicable in our setting. In addition, those checkpoints include little molecular 3D structural knowledge, thus not suitable for our molecular tasks.

Q2: text-to-molecule generation tasks?

• Thank you for your insightful comments! We appreciate your suggestion to explore the interesting applicability of our Geo2Seq towards text-molecule tasks. Indeed, Geo2Seq be applied to text-to-molecule generation tasks. The main difference would be extending the tokenization with for text, which can be enabled with the BPE or SentencePiece tokenizer. On the other hand, text-to-molecule generation tasks poses a different setting with various baselines such as MolT5 and LDMol. We believe this opens a promising future direction and we will include the discussion in the paper revision. Within the time limit of the rebuttal, we focus on the field of 3D molecule generation in this work and leave this exploration for future work.

We sincerely thank you for your time! Hope we have addressed your concerns through practical efforts and shown the contributions and significance of our work. We look forward to your reply and further discussions, thanks!

Sincerely,

Authors

Dear Reviewer aNgc,

Thank you for your appreciation of our work and insightful comments! We address each point here and added experiments are in this Link.

Real number generalizability

We have conducted more studies on this point.

• LMs require tokenization which limits number resolution. To mitigate this, we study flexible number-tokens and show even with coarse discretization, our model maintains strong performance. Please see experiments in [Link] Table 1.
• We also analyze different discretization schemes. See experiments in Appx C Table 5 and [Link] Table 2.
• We can extend to additional datasets in future work.

Uniqueness of QM9

• For the uniqueness of QM9, we believe it is because conversion from real numbers to discrete tokens limits 3D search space, especially on small datasets like QM9. Evidence is that we achieves 99.77% uniqueness on GEOM-DRUGS. This reflects that richer database or vocabulary enlarge search space of 3D structures and enhance uniqueness.
• Moreover, following EDM, we emphasize validity/stability, thus setting temperature=0.7. This can be adjusted for validity-diversity trade-off. We can easily improve uniqueness with larger temperature (e.g., temp=1.0 enhances uniqueness from 81.7% to 86.5%).

Error Case

• We now include error case studies. Please see [Link] Table 7.
• We also conduct quantitative studies. Among 100 error cases, 61% stem from incorrect distance-angle combinations, 25% geometric inconsistencies, and 14% subsequences repetition. Errors are rarer if model well converges: when trained for 150 and 250 epochs, the model generates ~15% and <2% invalid samples, respectively.

Comparison with graph-based methods

• Previously we focus on comparing with SoTA methods. Here we compare with more graph-based methods, as in [Link] Table 3.
• Meanwhile, we clarify that our baselines also directly use GNN representations, where GNNs are backbones while other techniques enable generation. E-NF is based on E(n)-GNN, G-Schnet uses GNN SchNet, EDM and GEOLDM parametrized by EGNN, and GDM by non-equivariant MPNN.

LMs like llama/qwen

• We provide experiments extending Geo2Seq with LLaMA. Please see [Link] Table 4 and our Response to Reviewer 39Ws.

More Metrics

Thanks for the advice. We already include these metrics in Appx. D.

• Table 6, 7, 8, 9, 10, 11 in Appx. D report evaluations including novelty, structural diversity metrics, distribution distance metrics, etc. Property prediction accuracy is as Table 2 of the paper. Vast results show we can outperform other methods across various metrics.

Def 3.4 Validity

We clarify its soundness regarding theoretical alignment with graph theory and physical relevance in molecular modeling.

• Our definition builds upon colored graph isomorphism. Lemma 3.2 show attributed CL retains bijectivity and different-attributed molecules cannot be conflated under this formulation. This is consistent with graph matching literature. We extend to 3D graphs requiring valid SE(3) transformation mapping coordinate matrices. It is formalized in Def 3.4 and supported by Lemma 3.3 & Thm 3.5. Our definition corresponds to label- and coordinate-preserving isomorphism.
• To further substantiate, it aligns with physical indistinguishability in chemistry, where two conformers differing only by spatial orientation are considered identical. While our formulation is novel, it mirrors physical practices in geometric deep learning, which treat structures up to SE(3) equivalence.

Proof Completeness

• We clarify that our proofs do include sufficiency & necessity explicitly. Thm 3.5: lines 883-899 prove sufficiency, and lines 900-928 necessity. Cor 3.6: lines 967-987 sufficiency and 988-1021 necessity.

Pretraining

• We clarify that we do explore the impact of pretraining. See Appx D.5 and Table 13.

Limited Controlled Generation

• We focus on quantum properties because they are available conditions relating 3D information. Here we provide more experiments exploring generalizability across conditions.
• The first experiment studies generalization across multiple conditions (see [Link] Table 5). The second experiment tests generalization to unseen conditions (see [Link] Table 6). While generalization brings challenges, we can capture certain multi-condition knowledge with robustness.

Limited Limitations

• We now include a detailed limitations section. Key challenges include discretization loss, high-precision 3D geometry generalization, and solutions include tokenizer learning with continuous embeddings or vector quantized codes.

Figure Captions

• We have included more informative captions. For Figure 2, see [Link] Figure for updated caption.

Thank you for your time! Hope we have addressed your concerns with practical efforts and shown our work’s significance. We look forward to further discussions.

Sincerely,

Authors

Dear Reviewer hgAt,

Thank you for your appreciation of our work and insightful comments! We have made efforts to thoroughly improve our work accordingly and provide responses for each concern here. Please also refer to our added experiments in this Link and our responses to other reviewers.

Claims And Evidence

• Thanks for your valuable comments. Beyond standing as a valid lossless tokenization, Geo2Seq has SE(3)-invariance and adopt spherical design, which can significantly benefit in achieving stronger empirical generation. We have conducted comprehensive ablation studies here. We explain in detail below.
• In Ablation on 3D representation of Appx. C, we explore using simpler or different coordinate encodings to represent 3D molecular structures. We compare the spherical coordinates in Geo2Seq with directly using the 3D Cartesian coordinates from xyz data files. We also study whether normalizing the xyz coordinates is effective by subtracting the xyz coordinates with the mass-center coordinates of each molecule. Additionally, we compare with using the SE(3)-invariant Cartesian coordinates that are projected to the equivariant frame proposed in Section 3.2. We also explore adopting to manage distances in a more local scheme, which reduces the scale of the distances.
• Results in Table 4 of Appx. C demonstrate that LLMs achieve the best performance on spherical coordinates, showing the superiority of invariant spherical coordinates over invariant Cartesian coordinates. We believe this is due to that the numerical values of distances and angles of spherical coordinates lie in a smaller region than coordinates, which reduces outliers and makes it easier for LLMs to capture their correlation. From these empirical results, we can analyze that the representation of azimuth and polar angles has brought sufficient advantage for LM learning over Cartesian coordinates, thus spherical representations with both distance schemes are showing promising performances.
• In Sec 3.2, we discuss the advantage of spherical coordinates. Compared to Cartesian coordinates, spherical coordinate values are bounded in a smaller region, namely, a range of $[0,\pi]$/$[0,2\pi]$. Given the same decimal place constraints, spherical coordinates require a smaller vocabulary size, and given the same vocabulary size, spherical coordinates present less information loss. This makes spherical coordinates advantageous in discretized representations and thus easier to be modeled by LMs.

Experimental Designs Or Analyses

Thanks for your comment. This is an important point and we have conducted various experiments and analyses including measuring geometry consistency and novelty to verify the capabilities of the method.

• Table 6 and 7 reports further random/controllable generation results including novelty metrics. Results show that our method achieves reasonably high novelty scores, which demonstrates that our method is not simply memorizing or reproducing token patterns.
• Also, experiments indicate that our generated molecules satisfy the internal geometry consistencies well. In addition to our main validity/stability results, we provided more evaluation results on various chemical metrics in Appx. D, Table 6, 7, 8, 9, 10, and 11, including diversity metrics, distribution distance metrics, bond lengths and angles, reasonable internal energy, steric hindrance, etc. Vast results show that we can outperform existing methods across metrics regarding various chemical constraints and geometry consistencies.
• In addition, in Appendix F.2, we provided UMAP visualizations of learned (atom type, distance, and angle) token embeddings, which indicates that the model has successfully learned structure information in 3D space. Figure 8 shows similar angle tokens (e.g., 1.41° and 1.42°) are placed next to each other, overall structure of all angles is a loop, and π-out-of-phase angles (3.14°, -3.14°, and 0°) are placed near each other. For atom type tokens, the model appears to capture the structure of the periodic table.
• We believe the current results are sufficient to prove the correctness and capabilities of our design. Geo2Seq with LMs can model 3D molecular structure distribution and capture the underlying chemical rules.

We sincerely thank you for your time! Hope we have addressed your concerns through practical efforts and shown the contributions and significance of our work. We look forward to your reply and further discussions, thanks!

Sincerely,

Authors

Dear Reviewer 39Ws,

Thank you for your appreciation of our work and insightful comments! We have made efforts to thoroughly improve our work accordingly and provide responses for each concern here. Please also refer to our added experiments in this Link and our responses to other reviewers.

Geo2Seq with larger or recent language models such as LLaMa

• Thanks for your valuable comments. We would like to provide more experimental results extending Geo2Seq to LLaMA. Due to the limited time of the rebuttal, here we provide the results of Geo2Seq with LLaMA for the controllable generation task, which have a smaller training data size of 50K. We will update the results of Geo2Seq with larger LLaMA/Qwen models as well as for other tasks and hyperparameter settings in the paper revision later. In scale with our data size, we use the LLaMA implementation from HuggingFace with 768 hidden size, 8 hidden layers, and 8 attention heads. The setting is the same as GPT used in the paper. We train the model for 200 epochs.

• As shown above, Geo2Seq with LLaMA achieves significantly better results over the baselines and performs similarly as Geo2Seq with GPT for most properties, without careful tuning of hyperparameters. This is expected given the functional architectures of LLaMA and GPT are very similar. We have included the results and discussions in the paper revision.

We sincerely thank you for your time! Hope we have addressed your concerns through practical efforts and shown the contributions and significance of our work. We look forward to your reply and further discussions, thanks!

Sincerely,

Authors