Tokenizing 3D Molecule Structure with Quantized Spherical Coordinates

We sincerely thank the reviewer for their constructive feedback and valuable suggestions.

W1.1: Atom token ordering

We did not design a specific algorithm for determining the order of atom tokens; instead, we used the default canonicalization in RDKit, which employs a DFS-like algorithm. Additionally, it is important to note that our method is not dependent on any particular order—random orders can also be applied.

W1.2, 1.3: Paper typo

Thanks for pointing out our typos. We have fixed in our revised manuscript.

W1.4: Lack of justification of local atomic environment

Our goal is to develop a general-purpose molecular tokenizer, not solely focused on molecule generation. Understanding the 3D spatial neighborhood structure is crucial for capturing the local atomic environment, which plays a significant role in understanding tasks. However, the generation features themselves do not consider information about neighboring atoms. Therefore, we include the understanding features to incorporate this information. Removing these features simplifies the task and could improve performance on generation tasks. However, this comes at the cost of diminished capability for enhancing understanding tasks, undermining the tokenizer's broader applicability.

W2: Lack of comparison with Geo2Seq

Firstly, from a procedural standpoint (as per ICLR Reviewer Guidelines), Geo2Seq was uploaded to arXiv on August 19, and it has not been published in a peer-reviewed venue. Therefore, it is not mandatory to treat Geo2Seq as a baseline. Nevertheless, we have appropriately cited and discussed their work in the paper. Additionally, Geo2Seq has not provided open-source code. Based on the results reported in their paper, our method achieves comparable performance when utilizing GPTs as LM.

We can, however, further discuss the limitations of Geo2Seq:

1. Lack of Generalization in Global Spherical Coordinates: Geo2Seq's use of global spherical coordinates makes it difficult to generalize across scales.
2. Inefficiency of Number-Based Tokenization: Directly tokenizing numerical values results in several issues:
   1. Precision is bounded by the number of decimal places, and inconsistent decimal precision hinders generalization.
   2. Segmenting numbers into tokens often produces long sequences, especially for large numbers or high-precision decimals. This increases computational burden and complicates modeling.

In contrast, our approach employs discretization (e.g., VQVAE), which significantly reduces the length of numerical representations and computational overhead.

Moreover, our method is designed as a general-purpose tokenizer. We pretrain the vocabulary on the PCQM dataset, ensuring generalization across diverse datasets. Geo2Seq, however, relies on hash-based discretization trained only on specific datasets, limiting its generalizability. Furthermore, Geo2Seq does not serialize bond features, which are crucial for representing 2D graphs in real-world scenarios (Geo2Seq relies on RDKit to reconstruct chemical bonds). Finally, our method includes understanding features, which Geo2Seq lacks. While these design choices for generalization may slightly reduce performance on generation tasks, they expand the applicability of our method to a broader range of scenarios.

W3: Pre-split bin-based tokenization

Thanks for your suggestion. While pre-split bins can discretize continuous values, they cannot reconstruct the original continuous descriptors from the discretized bins. This makes it impossible to restore molecular structures. In contrast, VQ-VAE allows for the reconstruction of continuous descriptors from discrete tokens, preserving the structural information.

If we were to use the midpoint of each bin as the reconstructed continuous value, we can estimate the vocabulary size required to achieve reconstruction errors comparable to those of VQ-VAE. For simplicity, let us consider only the generation descriptors: In our method, VQ-VAE is trained jointly on understanding and generation descriptors. The reconstruction RMSD for generation descriptors is approximately 0.01Å (with a range of 0.9Å∼2.7Å), with angular errors of 0.06 (with ranges of [0,π] and [−π,π] for elevation and azimuthal angles, respectively). These can be considered upper bounds for the reconstruction errors achievable by VQ-VAE when trained only on generation descriptors.

To match this error using pre-split bins under the assumption of uniform data distribution: The bin interval for distances must be 0.04 Å, corresponding to 45 bins. The bin interval for angles must be 0.24, corresponding to 13 bins for elevation and 26 bins for azimuth. Taking the Cartesian product of these bins yields a vocabulary size of 45×13×26=15,210, which far exceeds the VQ-VAE vocabulary size of 256.

This illustrates that VQ-VAE adaptively learns an optimal vocabulary, achieving better compression performance while maintaining low reconstruction error. Such adaptability is crucial for effectively handling diverse molecular datasets and achieving compact yet accurate representations.

We perform additional experiments on pre-split-based discretization as described above, better with using VQ-VAE.

Dear reviewer 8GHQ,

Towards your last message, we have provided the following response several days ago, please refer here again: https://openreview.net/forum?id=UqrSyATn7F&noteId=SsmAmGbExG.

However, unluckily, though we provided explanations and responses to your questions, we haven't received your further comments and feedback about our message.

We do try our best to give you thoughtful explanations and we do look forward your engagement to review our response.

Sincerely look forward to your comment.

Best,

Authors.

We sincerely thank the reviewer for their time and effort in reviewing our manuscript.

W1: Unclear rationale and handling of fixed neighbor selection

We chose four neighbors for the atomic descriptor in the understanding task because using only three atoms is insufficient to capture the local environment effectively. The most common atom, carbon, typically has four neighbors, and increasing this number beyond four yields similar performance. We wanted to limit the feature dimensionality of the understanding descriptors to ensure it does not overshadow the generation descriptors during reconstruction.

While our method cannot handle molecules with fewer than four neighbors, such cases are exceedingly rare (approximately 0.03% in the QM9 dataset, assuming hydrogen is included). For these rare corner cases, we opted to simply remove them from our experiments.

W2: Inconsistent and unclear notations.

Thank you for your suggestion. We agree that consistency in notation is important. However, we believe that the use of M is essential for maintaining the clarity and flow of the discussion in Section 3.1, so we have decided to retain it. On the other hand, we agree that the notational differences in Appendix C could cause confusion. Since removing Appendix C does not impact the completeness of the paper, we have opted to remove that section to ensure overall coherence.

W3, W4: Paper typo

Thanks for your careful review. We have updated these notions in our revised manuscript.

W5: References to prior works

Thank you for your suggestion. In the revised manuscript, we have clarified this statement by providing a detailed description and appropriate citations of the previous methods that inspired our work.

Dear Reviewer 7HbA,

Thank you once again for your valuable feedback. With only two days remaining until the discussion period ends, we kindly request confirmation that our responses have addressed all of your concerns. We are open to any further discussions if needed.

We truly understand how time-consuming and challenging the process of reviewing papers can be. We sincerely appreciate the effort each reviewer puts into reading and providing suggestions, as it is through your insights that the quality of our work improves.

At the same time, as an author submitting to this conference, we fully recognize the immense dedication involved in preparing a manuscript and responding to each reviewer’s queries. The paper we submitted, along with the responses we have provided, reflects the hard work and commitment of our team. It is of great significance to us, and we simply hope for a fair and constructive evaluation.

Thank you for your understanding and time. We look forward to your feedback and wish you all the best.

Best regards,

Authors

We sincerely thank the reviewer for their positive and encouraging feedback.

W1: Alternative BPE-based tokenization methods.

Thank you for pointing this out. Tokenization is indeed an interesting perspective. Our paper primarily focuses on the serialization and discretization of 3D molecules, which is why we defaulted to using the SELFIES token-based tokenizer. During the rebuttal, we also implemented a BPE-based tokenizer. The results, as shown in the table below, indicate that the choice of tokenization method has only a minor impact on performance.

W2: Alternative coordinate systems.

You are correct that using other coordinate systems can make the task more challenging to learn. To demonstrate this, we have conducted experiments using Euclidean coordinates. The results show that the LM struggles to generate stable molecules under this representation. In contrast, the spherical coordinate system constrains the numerical range, making it easier for the model to learn effectively.

Q1: PoseBusters' energy ratio test

The Reasonable Internal Energy term is actually the metric you mentioned. We directly used the test script from Symphony. We're sorry that we didn't elaborate on it in the paper. We have updated the detailed description of the metric in our paper.

Q2: Unclear motivation of MLP-based autoencoder

The motivation for adopting a purely MLP-based quantized auto-encoder architecture is twofold:

1. The task itself is not particularly complex, and node-level modeling already achieves excellent RMSD performance.
2. We aim for the discretized structure vocabulary to exhibit generalizability across different atomic types in different molecules rather than being tailored to individual molecules. To further illustrate this, we have included additional heatmap visualizations in the Appendix.

We thank the reviewer for their time and detailed feedback on our manuscript.

W1: Lack of novelty due to similarities with existing methods

We think that your comment regarding novelty is somewhat biased, so we would like to reiterate our contribution: our primary innovation lies in serializing and discretizing 3D molecular structures to support language model-based modeling. We openly acknowledge that our work is inspired by FoldSeek and thoroughly discuss FoldSeek and related methodologies in the Introduction, Related Work, and Experiments sections of our paper. However, applying similar ideas in FoldSeek is far from trivial on molecules. To clarify the differences between our method and FoldSeek:

• Serialization of 3D Structure: While FoldSeek operates on protein structures, which are naturally sequential, molecules are graphs with no inherent order between atoms. Serializing the 3D structure of molecules is not trivial. Our approach requires constructing SE(3)-invariant descriptors, which involves using both 1D line notation and 2D graph representations to achieve robust local structural encoding. This is a key distinction from FoldSeek, FoldToken, and previous AR methods for 3D molecule generation, as discussed in Section 3.2.
• Distinct Features: We diverge from FoldSeek in our feature design. Our model architecture utilizes a similar MLP structure, but with fine-tuned hyperparameters and an additional head to predict the sign of the azimuthal angle, which significantly aids in the reconstruction of azimuthal angles We also normalize length features to the log space and angle features to the 0-1 range. This normalization ensures that the model learns descriptors with consistent scales, which also benefits the VQ-VAE learning process.
• Training Objective: FoldSeek uses a Gaussian negative log-likelihood loss to model the distributions of aligned amino acids. In contrast, our goal is to reconstruct the local structure of source atoms, so we use MSE loss, which better reflects the quality of the reconstruction.

These design choices collectively highlight the distinctiveness of our approach compared to FoldSeek and FoldToken, particularly in how we handle molecular structure serialization, feature processing, and training objectives.

W2: More benchmark results

Thank you for your suggestion. We have already started conducting additional experiments on GEOM-Drug dataset. However, due to limited computational resources, it may take some time to pre-processing and complete these experiments. We will include the results as soon as they are available.

W3: Lack of comparison with Geo2Seq and BindGPT

Firstly, from a procedural standpoint (as per ICLR Reviewer Guidelines), Geo2Seq was uploaded to arXiv on August 19, and neither Geo2Seq nor BindGPT has been published in a peer-reviewed venue. Therefore, it is not mandatory to treat Geo2Seq as a baseline. Nevertheless, we have appropriately cited and discussed their work in the paper. Additionally, Geo2Seq has not provided open-source code. Based on the results reported in their paper, our method achieves comparable performance when utilizing GPTs as LM.

W4: Further FoldSeek-like visualization

Our approach differs fundamentally from FoldSeek. FoldSeek is trained not for reconstruction but to learn aligned distributions, employing a Gaussian negative log-likelihood loss. In their case, a single code decodes into a distribution, which allows visualizations like Figures S2 and S3 in their work.

In contrast, our task focuses on reconstruction, where the generation descriptors are designed to restore the molecular structure. Each descriptor can only be mapped to a vector, meaning visualizing atoms grouped under the same code does not provide meaningful insights for our method.

Q1: Mistaken description of SE(3)-invariance

We truly appreciate you pointing this out. We acknowledge that it was our mistake. The issue has been rectified in the revised manuscript.

Q2: Model size

The model sizes used in our method are as follows:

• VQ-VAE: 300k parameters
• GPT-2: 124M parameters
• Graphormer: 50M parameters

We have updated this information in our revised manuscript.

Dear Reviewer gMZe,

Thank you for your thoughtful feedback on our manuscript. We have carefully addressed all of the concerns you raised in our responses, and we hope that the clarifications provided sufficiently resolve the points of uncertainty and potentially alter any negative impressions that may have arisen during the initial review process. Given that the discussion period is nearing its deadline, we kindly request that you review our responses at your earliest convenience.

We greatly appreciate your time and effort in reviewing our work and look forward to your feedback.

Best regards,

Authors

Dear Reviewer gMZe,

We appreciate your feedback for acknowledging the improvements we've made in addressing your concerns. Regarding the observation that our work appears similar to previous studies, we would like to emphasize that both works were developed concurrently. The similarity you noted only reinforces the idea that this is an actively explored topic. In our paper and rebuttals, we have provided detailed explanations of how our approach differs from those previously proposed. As the Area Chair has not provided explicit feedback, we understand that we must accept your judgment regarding this matter.

Once again, thank you for your time and consideration. Wishing you all the best.

Kind regards,

Authors

Dear AC and SAC,

We truly appreciate the reviewing process and the efforts paid by reviewers. With only one day remaining until the discussion period ends, we kindly request your help in helping with the author-review discussion. We have paid much efforts to add more experiments and tried our best to address the concerns raised by reviewers. Even as the experiments compared to the paper uploaded on Arxiv on Aug, which is not necessarily needed that avoids the policy. However, it is sad that the reviewers do not give much attention on our rebuttals and follow-up responses.

We do concern that we have received some biased reviews and we sincerely hope you could help to guide the discussion period and also give a fair evaluation on our paper, instead of only considering the reviewers' initial comments. We do look forward your kind help.

Appreciate a lot.

Authors.