Tokenizing 3D Molecule Structure with Quantized Spherical Coordinates

Authors: Kaiyuan Gao, Yusong Wang, Haoxiang Guan, Zun Wang, Qizhi Pei, John E. Hopcroft, Kun He, Lijun Wu
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
Keywords: molecule structure tokenization, generation, spherical coordinates, VQ-VAE

Abstract:
The application of language models (LMs) to molecular structure generation using line notations such as SMILES and SELFIES has been well-established in the field of cheminformatics. However, extending these models to generate 3D molecular structures presents significant challenges. Two primary obstacles emerge: (1) the difficulty in designing a 3D line notation that ensures SE(3)-invariant atomic coordinates, and (2) the non-trivial task of tokenizing continuous coordinates for use in LMs, which inherently require discrete inputs.
To address these challenges, we propose Mol-StrucTok, a novel method for tokenizing 3D molecular structures. Our approach comprises two key innovations: (1) We design a line notation for 3D molecules by extracting local atomic coordinates in a spherical coordinate system. This notation builds upon existing 2D line notations and remains agnostic to their specific forms, ensuring compatibility with various molecular representation schemes. (2) We employ a Vector Quantized Variational Autoencoder (VQ-VAE) to tokenize these coordinates, treating them as generation descriptors. To further enhance the representation, we incorporate neighborhood bond lengths and bond angles as understanding descriptors. Leveraging this tokenization framework, we train a GPT-2 style model for 3D molecular generation tasks. Results demonstrate strong performance with significantly faster generation speeds and competitive chemical stability compared to previous methods. 
Further, by integrating our learned discrete representations into Graphormer model for property prediction on QM9 dataset, Mol-StrucTok reveals consistent improvements across various molecular properties, underscoring the versatility and robustness of our approach.

Decision: reject
Reviewer scores: 5, 5, 8, 5

Reviews and discussion:
<thread>
creation date:
1730592894769

thread_type:
official_review

<thread_object>
creator: reviewer_8ghq

<content>
summary:
This work proposes a novel language model based 3D molecule generation framework. Novel spherical coordinate based 3D line notation method combined with VQ-VAE model is adopted to tokenize 3D molecules, and GPT-2 models are applied to this tokenization framework. Benchmark experiments are conducted to show the promising performance of the proposed method.

strengths:
Originality:  
This work makes a good originality contribution of  proposing a novel 3D molecule tokenization method that is applicable for language models.

Quality:  
This work presents sufficient theoretic analysis and proofs for the proposed 3D tokenization method, and benchmark experiments are conducted to demonstrate the promising performance.

Clarity:  
The writing of this paper is fairly good. The paper gives sufficient description to help readers understand the key pipeline and steps of 3D molecule generation.

Significance:  
The proposed method is useful in leveraging the power of large language models to the generation and modeling of complicated 3D molecule structures.

weaknesses:
(1) Some important methods details are not clearly clarified or presented. Authors are encouraged to add explanations to the following questions: 
- How the order of atom tokens in equation (2) is decided? By BFS or DFS traverse algorithm on 2D topology graphs?
- The $x_i$ in line 266 is actually $g_i$ in equation (7)?
- Why the sign of torsion angle $sign(\varphi_i)$ appears in the generation descriptor of equation (7) but not in the combined descriptor of equation (9)?
- Why local atomic environment (line 266-267) is needed? Can authors discuss the negative impacts of not including it and conduct ablation studies if possible?

(2) In benchmark experiments, an important baseline Geo2Seq [1] (it also proposes a spherical coordinate based 3D molecule tokenization framework) is not presented and compared. Particularly, Geo2Seq has higher molecule stability than authors' method, so more discussion about differences and advantages over Geo2Seq is needed.

(3) VQ-VAE based vector quantization is applied to atom descriptor in equation (9) to form the final generation target. A natural question is what is the advantages of VQ-VAE over simple discretization (i.e., map floating numbers to pre-split bins, e.g., 0-30, 30-60, 60-90, ... for angles) ? Authors are encouraged to discuss different possible quantization methods and conduct ablation experiments to justify the advantages of VQ-VAE.

[1] Li, Xiner, et al. "Geometry Informed Tokenization of Molecules for Language Model Generation." arXiv preprint arXiv:2408.10120 (2024).

questions:
See Weaknesses part.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730693390816

thread_type:
official_review

<thread_object>
creator: reviewer_8laj

<content>
summary:
The authors express 3D structure using spherical coordinates, then tokenize them with a VQ-VAE. The tokens are then used for training a GPT-2 model for 3D molecular generation tasks. They also use the discrete representations in a Graphormer model for property prediction on QM9, and observe consistent improvements across various molecular properties.

strengths:
1. The authors conduct an extensive set of experiments. They measure validity+uniqueness of generated molecules with different bond assignment methods, perform PoseBusters tests, evaluate quantum mechanical properties, and measure MAE for QM9 property prediction. They achieve state-of-the-art results in most experiments.
2. They also perform additional analysis regarding the inference speed of their method and the effect of the generation temperature on balancing quality and diversity.

weaknesses:
1. This is a hand-crafted tokenization scheme and should be compared to other tokenizers (e.g. BPE-based tokenizers), not just diffusion models and MPNN-based methods.
2. It may also be helpful to compare with structures expressed in other coordinate systems. I'd imagine that without SE(3) invariance there would be a wider range of possible tokenized sequences, making it harder for the GPT-2 model to learn.

questions:
1. To the best of my knowledge, PoseBusters has an energy ratio test, which tends to be very difficult to satisfy. Have you tried on that test? I think that test would be more meaningful since most methods already perform very well on the tests you demonstrated in Table 2.

2. What's the motivation for a purely MLP-based quantized auto-encoder architecture? Most other autoencoder-based tokenizers use some form of GNN or sequence model. What are the tradeoffs between element-wise discretization and discretization with message propagation?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730698984991

thread_type:
official_review

<thread_object>
creator: reviewer_7hba

<content>
summary:
The authors introduce a novel method for tokenizing 3D molecular structures. The authors design a 2D line notation for 3D molecules by extracting the local atomic spherical system, then employ VQ-VAE to tokenize upon this 2D line notation to tokenize the spherical coordinates. Results show the competitive performance compared to several SOTA methods.

strengths:
The combination of spherical line notation with vector quantization enables language models to process complex 3D data, which is challenging to discretize. This approach stands out from traditional graph-based or continuous-coordinate models by providing a discrete representation for language models without losing SE(3)-invariant information. Particularly, the augmented tokens incorporate both generation and understanding descriptors, including local spherical coordinates, bond lengths, and angles, allowing a better capture of molecular topology and spatial arrangements.

weaknesses:
### Major
The authors should clarify the rationale behind selecting exactly four neighbors for the atomic descriptor and explicitly address how the descriptor $\mathbf{z}_i$ is defined for atoms with fewer than four neighbors. This is essential, as molecules with varying coordination environments will likely have different numbers of neighbors, impacting the generality of the descriptor across datasets.

### Minor
1. The paper’s notations are somewhat inconsistent and could benefit from simplification and unification. For instance, symbols like $\mathcal M$ are used only within specific sections, such as Section 3.1, and do not appear elsewhere. Additionally, the notation in Appendix C is different from that in Section 3.1. A more streamlined notation would improve readability and coherence across sections.
2. Line 177, the phrase should be rephrased as "we tokenized the molecular graph into a sequence of atom tokens $\mathbb A$ and non-atom tokens $\mathbb B$".
3. In Line 248 and Line 268, references are made to "Section X" and "Figure X," which seem placeholders. These should be updated with specific section and figure numbers
4. Line 239, "Inspired by previous works" is ambiguous. The authors should specify which studies or methods they are referring to here.

questions:
See weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730796612405

thread_type:
official_review

<thread_object>
creator: reviewer_gmze

<content>
summary:
This paper focuses on the task of 3D molecule generation with language models (LMs). The main step is to convert 3D molecules to sequences of discrete tokens. The authors first use SMILES (or SELFIES) to convert a molecule's atom and bond information to a sequence. They then use spherical coordinates and VQ-VAE to convert 3D information to discrete tokens. After obtaining these 3D molecule sequences, they train a GPT-style model for molecule generation by next-token-prediction. The model is evaluated on the random generation and conditional generation tasks.

strengths:
1. This paper is well-written and easy to follow, with clear and informative tables and figures. 
2. The proposed method performs well, especially on the conditional generation task.
3. The ablation study is thorough and provides useful insights.

weaknesses:
1. The proposed method is quite similar to existing methods, such as FoldSeek and FoldToken. Specifically, similar to the SE(3)-invariant spherical coordinates here, FoldSeek also uses distances and angles computed based on reference nodes as SE(3)-invariant representations. In addition, Furthermore, both methods employ VQ-VAE to learn discrete tokens. These overlapping components limit the novelty of this work.
2. About the datasets: the proposed method is only evaluated on QM9 dataset, which includes only small molecules. It will be better to evaluate on larger molecules, such as geom-drug dataset.
3. More LM-based baseline methods should be included and compared, such as Geo2Seq and BindGPT.
4. Further analysis and potential visualizations of the learned structural alphabet, for example, which types of local structures can be mapped to the same code? Refer to Figure S2 and S3 of FoldSeek.


Other questions:
1. line 52: SE(3) invariance is invariance under rotation and translation, but not reflection
2. What is the model size?

questions:
See weaknesses
</content>
</thread_object>

</thread>
<thread>
creation date:
1732547542922

thread_type:
official_comment

<thread_object>
creator: authors

<content>
comment:
Dear Reviewers,

Thanks for the time, effort, and expertise you have invested in reviewing our paper.

As the author-reviewer discussion period approaches its deadline, we would greatly appreciate it if you could prioritize providing your feedback at your earliest convenience.

Thank you again for your valuable suggestions and contributions. 


Sincerely,

Authors
</content>
</thread_object>

</thread>
<thread>
creation date:
1732764657817

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Official Comment Regarding Reviewer Violation of ICLR's Contemporaneous Work Policy

comment:
Dear Area Chairs, Senior Area Chairs, and Program Chairs,

We deeply appreciate the review process and value the constructive feedback we have received. However, we would like to bring to your attention a concern regarding Reviewer **8GHQ,** whose evaluation appears to contravene the ICLR policy on contemporaneous work comparisons.

According to the ICLR Reviewer Guide (https://iclr.cc/Conferences/2025/ReviewerGuide), authors are not expected to compare their work to papers published within four months before the submission deadline or to non-peer-reviewed works (e.g., papers on arXiv). Specifically, the policy states:

> We consider papers contemporaneous if they are published within the last four months. That means, since our full paper deadline is October 1, if a paper was published (i.e., at a peer-reviewed venue) on or after July 1, 2024, authors are not required to compare their own work to that paper.
> 

The reviewer **has explicitly based his/her rejection recommendation on a direct comparison with a contemporaneous work posted on arXiv (not a publishment) on August 19, 2024**, citing the absence of one additional dataset and slightly lower performance on a few metrics as reasons. However, even at the time of submission, the overall performance of our work was comparable to the contemporaneous study (listed below). Furthermore, our experiments on QM9 were carefully designed to be representative of the task, in alignment with prior ICLR papers such as **Symphony** (https://openreview.net/forum?id=MIEnYtlGyv), which also evaluated methods solely on QM9.

| Unconditional | Atom Sta | Mol Sta | Valid | Valid & Unique |
| --- | --- | --- | --- | --- |
| Geo2Seq | 98.3% | **90.3%** | 94.8% | 80.6% |
| Mol-StrucTok | **98.5%** | 88.3% | **98.0%** | **85.1%** |

| Conditional | alpha | gap | homo | lumo | mu | Cv |
| --- | --- | --- | --- | --- | --- | --- |
| Geo2Seq | 0.46 | 98 | **57** | 71 | **0.164** | 0.275 |
| Mol-StrucTok | **0.33** | **89** | 64 | **62** | 0.285 | **0.169** |

While we respect the value of the contemporaneous work and acknowledge its inspiring contributions, **we believe it is inappropriate for this to form the basis of a rejection recommendation given that it directly conflicts with ICLR’s stated policy, and also given the results of both papers are comparable.** Our submission included a detailed discussion highlighting the distinctions between our work and the referenced study, demonstrating that both contributions are meaningful to the field.

We kindly request your guidance and consideration on this matter, as the evaluation may have been unduly influenced by factors outside the scope of ICLR’s policies. We remain open to constructive feedback and are committed to addressing valid concerns to improve our work further.

Thank you for your time and attention to this matter.

Sincerely,

Authors
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper studies how to convert molecules into tokens for language model generation. The overall idea is quite interesting, but the reviewers raised major concerns on the technical and experimental aspects of the work, which are not resolved after rebuttals. I agree that this work does not necessarily need to compare with Geo2Seq due to its relative timing. However, the reviewers are concerned with method novelty and experimental advantages compared to other prior methods, like FoldSeek, FoldToken, VQ-VAE, etc. in terms of general methods, loss functions, and empirical improvements. Due to these major unresolved concerns, a reject is recommended.