Geometry Informed Tokenization of Molecules for Language Model Generation
Authors: Xiner Li
Limei Wang
Youzhi Luo
Carl Edwards
Shurui Gui
Yuchao Lin
Heng Ji
Shuiwang Ji
Venue: ICML 2025 poster (2025)
Keywords: Language models
Primary area: deep_learning->generative_models_and_autoencoders

Abstract: We consider molecule generation in 3D space using language models (LMs), which requires discrete tokenization of 3D molecular geometries. Although tokenization of molecular graphs exists, that for 3D geometries is largely unexplored. Here, we attempt to bridge this gap by proposing a novel method which converts molecular geometries into SE(3)-invariant 1D discrete sequences. Our method consists of canonical labeling and invariant spherical representation steps, which together maintain geometric and atomic fidelity in a format conducive to LMs. Our experiments show that, when coupled with our proposed method, various LMs excel in molecular geometry generation, especially in controlled generation tasks. Our code has been released as part of the AIRS library (https://github.com/divelab/AIRS/).

Geometry Informed Tokenization of Molecules for Language Model Generation
Xiner Li Limei Wang"
11
Youzhi Luo't Carl Edwards? 2
Shurui Gui Yuchao Lin Heng Ji
Shuiwang Ji'
Abstract
lined data processing and powerful generation capabilities, have shown success across various domains, particularly in natural language processing (NLP). Recently, large lan- guage models (LLMs) (Zhao et al. 2023b) show extraordi- nary capabilities in learning complex patterns (Zhang et al., 2024) and generating meaningful outputs (Touvron et al. 2023 Achiam et al.) 2023] Chowdhery et al., 2023). De- spite their potential, the application of LLMs to the direct generation of 3D molecules is largely under-explored. This is primarily due to the fact that geometric graph structures of molecular data are fundamentally different from texts. However, 3D geometric information is crucial in molecular tasks, since different conformations of the same molecule topology have different properties, such as per-atom forces. This gap reveals unique challenge of how to make use of the powerful pattern recognition and generative capabilities of LLMs to handle complicated molecular graph structures, especially geometries. On the other hand, solutions to this challenge with model-level modifications cannot effectively leverage the rapidly developing power of LMs. These so- lutions require specific module designs, which needs to be done separately for each LM architecture and can be infea- sible for modern LMs released via APIs.
We consider molecule generation in 3D space using language models (LMs), which requires dis- crete tokenization of 3D molecular geometries. Although tokenization of molecular graphs ex- ists, that for 3D geometries is largely unexplored. Here, we attempt to bridge this gap by proposing the Geo2Seq, which converts molccular geome- tries into S (E33)invvariant 1D discrete sequences. Geo2Seq consists of canonical labeling and in- variant spherical representation steps, which to- gether maintain geometric and atomic fidelity in a format conducive to LMs. Our experiments show that, when coupled with Geo2Seq, various LMS excel in molecular geometry generation, es- pecially in controlled generation tasks. Our code has been released as part of the AIRS library httpps//gthhh..om/dvvvea///AAs
1. Introduction
The generation of novel molecules with desired properties is an important step in drug discovery. Specifically, the design of hires-dimensional (3D) molecular geometries is particu- larly important because 3D information plays a critical role in determining many molecular properties. Different gen- erative models have been used for 3D molecule generation. Early studies such as G-SchNet (Gebauer et al.. 2019) use autoregressive generative models to generate 3D molecules by sequentially placing atoms in 3D space. It was observed that these models often yield results with low chemical va- lidity. Recently, diffusion models (Hoogeboom et al. 2022 Xu et al. (0223) achieve better performance in 3D molecule generation tasks. However, they typically need thousands of diffusion steps, resulting in long generation time.
In this work, we bridge this gap by applying LMs to the task of 3D molecule generation. We employ a novel approach translating the intricate geometry of molecules into format that can be effectively processed by LMs. This is achieved by our proposed tokenization method Geo2Seq, which con- verts 3D molecular structures into (3(3)-invariant one- dimensional (ID) discrete sequences. The transformation is based canonical labeling, which allows dimension reduc- tion with no information loss outside graph isomorphism groups, and invariant spherical representations, which guar- antees SE) -inariincce under the cquivariant global frame. By doing so, we harness the advanced sequence-processing capabilities and efficiency of LMs while retaining essential geometric and atomic information. Note that since Geo2Seq operates solely input data, our method is agnostic to the subsequent LMs used. and can seamlessly adapt to any state-of- saate-of-the-arrd tttthhe-rtttt sequence model, maximizing LM capabili- ties while avoiding additional architecture design or redun- dant computations. When combined with powerful modern LLMs, Geo2Seq can achieve highly accurate modeling of 3D molecular structures. In addition, Geo2Seq can benefit
Language models (LMs) (Vaswani et al.
2017 Devlin et al. 2018 Brown et al. 2020 Gu et al. 202 with their stream-
Eqqua contribution 'Texas A&M University Uniiversity of Illinois Urbana- Urbana-Champa Champaign. Correspondence to: Shuiwang Ji <sji saji@tamu.edux.
Proceedings of the 42nd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Geometry Informed Tokenization of Molecules for Language Model Generation
conditional generation by including real-world chemical properties in sequences because modern L.MSs capable of capturing long-context correlations to comprehend global structure and information in sequences. Our experimental results demonstrate these advantages. We show that using different LMs with Geo2Seq can reliably produce valid and diverse 3D molecules and outperform the strong diffusion- based baselines by a large margin in conditional generation. These results validate the feasibility of using LMs for 3D molecule generation and highlight the potential to aid in the discovery of new molecules, paving the way for applications such as drug development and material science.
S -33 in 3D molecule generation, existing stud- ies have proposed various strategies. Early studies pro- pose to generate 3D atom positions by S - features, such as interatomic distances, angles and torsion angles. They construct 3D molecular structures through either atom-by-atom generation (Gebauer et al. 2019) Luo & Jil 2022) or generating full distance matrices (Hoffmann) & Noč 2019) in one shot. Recently, more and more studies have applied generative models generate 3D atom coordi- nate directly. These studies include E-NFs (Satorras et al.) 2021a) and EDM (Hoogehoom et al. 2022). which com- bine equivariant atom coordinate alignment process with equivariant EGNN (Satorras et al.] (2021b) model for 3D molecule generation. Following EDM, many other stud- ies have proposed to improve diffusion-basee 3D molecule generation frameworks by stochastic differential equation (SDE) based diffusion models (Wu et al.
2. Preliminaries and Related Work
2.1. 3D Molecule Generation
2022 Bao et al.,
In this work, we study the problem of generating 3D molecules from scratch. Note that this problem is different from the 3D molecular conformation generation problem
2023) or latent diffusion models (Xu et al. 2023a Besides,
some recent studies QQiang et al. 2023) have explored gen- erating 3D molecules through generating and connecting fragments first, then aligning atom coordinates with soft- ware like RDKit. We refer readers to et al. (2022) Zhang et al. 2023b) for a comprehensive review.
studied in the literature (Mansimov et al.
2019 Simm &
Hornandez-Lobato 2020; Gogineni et al.
2020
Xu et al.
2021a b Shi et al. |021] Ganea et al. 2021 Xu et al. 2022
Jing et al. 2022), where 3D molecular conformations are
generated from 2D molecular graphs. We represent 3D
While generating 3D molecules in the form of 3D point clouds have been well studied, few studies have tried ap- plying powerful language models to this problem. In this work, different from mainstream methods, we convert 3D point clouds to (33)-iivariian ID discrete sequences, and show that generating sequences by LMs achieves promising performance in the 3D molecule generation task.
molecule with n atoms in the form of a 3D point cloud (i.e., set of points with different positions in 3D Euclidean space)
as G = (z, R). Here, z = [= 2nD E Z" is the atom type vector where z is the atomic number (nuclear charge number) of the i-th atom, and R = [r1,... rr.] E IR 3x is the atom coordinate matrix, where r, is the 3D coordinate of the i-th atom. Note that 3D atom coordinates R are com- monly called 3D molecular conformations or geometries in chemistry. We aim to solve the following two generation tasks in this work:
2.2. Chemical Language Model
LMs have catalyzed significant advancements across a spec- trum of fields. Recently, LLMs have revolutionized the land- scape of NLP and beyond (Touvron et al. 2023 Achiam et al., 2023 Chowdhery et al., 2023). Drawing inspira- tion from NLP methodologies, chemical language models (CLMs) have emerged as a competent way for represent- ing molecules (Bran & Schwaller 20033 2023]Janakarajj Janakarajan et al.,
Random generation. Given a 3D molecule dataset G = (G, (m we aim to learn an unconditional generative model pe (-) on G so that the model can generate valid and diverse 3D molecules.
Controllable generation. Given a 3D molecule dataset g = {(GG, 8 where is a certain property value of Gj. we aim to learn a conditional generative model (-(ss) on G so that for a given s, the model can generate 3D molecules whose quantum property values
2023 Bajorath [024]] Zhang et al., (004)). Due to the su- periority LMs show in generation tasks, most CLMs designed generative models. Variants of LMs have been adapted for molecular science, producing variety of works.
CLMs learn the chemical vocabulary and syntax used to represent molecules, as well as the conditional probabilities of character occurrence at given positions of sequences de- pending on preceding characters. This vocabulary covers all characters from the adopted molecule representation. All in- puts including chemical structures and properties should be converted into sequence form and tokenized for compatibil- ity with language models. Commonly, SMILES (Weininger) 1988) is used for this sequential representation, although
major technical challenge of 3D molecule generation lies in maintaining invariant to SE(3) transformations, includ- ing rotation and translation. In other words, ideal models should assign the same probability to G = (z, R) and G" = (z, R') if R' = QR + 61", where 1 is an n- dimensional vector whose elements are all one, E RR3 is an arbitrary translation vector, and Q E 3333 is a ro- tation matrix satisfying QQT I, Q = 1. To achieve
2

---

Geometry Informed Tokenization of Molecules for Language Model Generation
other formats like SELFIES (Krenn et al. 2019), atom type strings, and custom strings with positional or property val- ues are also viable options. To learn representations, CLMs are usually pre-trained on extensive molecular sequences through self-suuervised learning. Subsequently, models are fine-tuned on more focused datasets with desired proper- ties, such as activity against a target protein. Generative CLMs generally adopt an autoregressive training approach of next token prediction, i.e., iteratively predicting each sub- sequent token in a sequence based on the preceding tokens. Traditional autoregressive models use the Transformer archi- tecture with causal self-attention (Brown et al.) (020) due to its superior efficacy, while other sequence models like recurrent neural networks (RNNs) and state space models (SSMs) (Gu et al. 2021; Özęelik al., 2024 2023) also show considerable functionality.
at. 2023 Liang et al.) 2023 Liu et 2023a; Frey et al.
2023 which requires additional design and training tech-
niques to mitigate alignment issues. Some works extend the architecture of CLM to include -D-geometric-model like modules in the attention block (Fuchs et al. 2020) (Shi et al.
2022 Liao & Smidt 2022] Thölke & De Fabritiis
2021
Luo et al. 2022b; Masters et al. 2022; Ünlü et al., 2023
Zhao al. 2023a capturing 3D information as positional
encodings with considerable computations and framework design. In contrast, Flam- Plam-Shepherd & Aspuru-Guzik (2023) make an initial attempt showing language models trained directly on contents of XYZ format chemical files can gen- erate molecules with three coordinates, implying pure LMs' potential to directly explore 3D chemical space. In this work, we propose an invariant 3D molecular sequencing algorithm, Geo2Seq, empower CLMs with structural completeness and geometric invariance, showing LMs' capabilities of un- derstanding molecules precisely in 3D space. We extend beyond the conventional Transformer architecture of CLMs and additionally employ SSMs as LM backbones. Further- more, Geo2Seq operates solely on the input data, which allows independence from model architecture and training techniques and provides reuse flexibility.
Given dataset of sequences, U = U2,,U,.... ,NN. where U2 is transformed from the representation, property conditions and/or descriptions of a molecule G, with ny nodes, let U.- = {u1, U22.. , and all tokens u, be- long to vocabulary V. An autoregressive CLM has pa- rameters 0 encoding distribution with conditional proba- bilities of each token given its predecessors, p(U,;; 0) IT; p(u,) uo uy-1; 0). The optimization process in- volves maximizing the probabilities of the entire dataset 0(U;0) = II P(U.; Each conditional distribution p(uy/un : 1 0) is a categorical distribution over the vo- cabulary size [V|: thus the loss for each term aligns with the standard cross-entropy loss. To generate new sequences, the model samples each token sequentially from these con- ditional distributions. To introduce randomness and control into generation, the sampling process is typically modu- lated with Top-K (k) and temperature (T) hyperparameters, enabling a balance between adherence and diversity.
3. Tokenization of 3D Molecules
A fundamental difference between LMS and other models is that LMs use discrete inputs, i.e., tokens. In this sec- tion, we introduce our tokenization method to map input 3D molecules with atomic coordinates to discrete token sequences appropriate for LM learning.
A main challenge in tokenization design is to develop bijec- tive mappings between 3D molecules and token sequences, Le., obtaining the same token sequence for the same input 3D molecule, while obtaining different sequences for dif- ferent inputs. In this section, we present our solutions to tackle this challenge. We first reorder the atoms in the input molecule to a canonical order (Section [.T). such that any two isomorphic graphs result in the same canonical form, and any non-isomorphic graphs yield different canonical forms. We then convert 3D Cartesian coordinates to (3) invariant spherical representations, including distances and angles (Section (.2). Combining them together, we ob- tain our geometry informed tokenization method Geo2Seq (Section (.3). We provide rigorous proof of all theorems supporting the bijective mapping relation in Appendix B
Most existing CLM works consider chemical structures as well as other modalities such as natural language cap-
tions (Bagal et al.
2021 Li et al.] 2023a b
Edwards al.
2022
Xie et al.
2023 Chen et al.
2023b
Tysinger et al.
2023
Xu et a 2023b Chen et al.
2023a
Pei et al.) 2023
Liu et al.] 20233b [Wang et al. 2023). while some focus
on pure text of chemical literature (Luo et al.] (002a) or
molecule strings (Haroon et al. 20233 Mao et al. 2023b
Blanchard et al. 2023 Mazuz et al. 2023 F Fang et al. 2023
Kyro et al., 2023 Izdebski al. 2023 Yoshikai et al. 2023
2023 Mao et al. (203a)). Notably, all these works solely consider 2D molecules for representation learning and downstream tasks, overlooking 3D geometric structures which is crucial in many molecular predictive and genera- tive tasks. For example, different conformations of the same 2D molecule have different potentials and per-atom forces. In order use pivotal 3D information, another line of work incorporate geometric models such as GNNs in parallel with the CLM (Xia et al. 2023 Zhang et al. 2023a
Wu et al.
3.1. Serialization via Canonical Ordering
As the first step in 3D molecule tokenization, we need to transform graph to a 1D sequential representation. We resort to canonical labeling as a solution for dimension reduction without information loss.
Cao
3

---

Geometry Informed Tokenization of Molecules for Language Model Generation
Scrialization
Sphorical Representation
Training & Generation
Canonical Labeling
Global Frame
Spherical Coundinate
3D Sequenee
I Langunge Models
$33 00300 2ut ouer UL daph
Heaedar edar dair
Tokenizaaion
Input Molecule
(-0 -0.77, 1.26)
Iaput IDs
groeeation
m 4.25 9 1365
Gonaitio:: Fiecos poppaty "aaee put 2:00 oor
acncratios
HONC C HI C
Figure 1: Overview of Geo2Seq. We use the canonical labeling order to arrange nodes in row, fill in the place of each node with vector [zi, d;, 0s, dı], and concatenate all elements into a sequence. Each node vector contains atom type and spherical coordinates. Notably, the spherical coordinates are SE(3) (3)-invariann.
Canonical labeling (CL), in the context of graph theory, is process to assign unique form to each graph in way that two graphs receive the same canonical form only if they are isomorphic (McKay et al. 198T). The canonical form is a re-indexed version of graph, which is unique for the whole isomorphism class of a graph. The new indexes naturally establish the order of nodes in the graph. The order, which refer to canonical labels, is not necessarily unique if the graph has symmetries and thus has automorphism group larger than 1. However, all canonical labels are strictly equivalent when used for serialization. The canonical label essentially re-assigns an index & to each node originally indexed with i in graph G. Since canonical labeling can pre- cisely distinguish -somooomo graphs, fully contains the structure information of graph G. Thus, by arranging nodes with attributes in the labeling order &1, f2. ..., we obtain a sequential representation of attributed graphs with all structural information preserved.
CL processes can also be extended to mode/edge-attributee graphs, leading us to the guarantee below.
Lemma 3.2. /Canonical Labeling for Colored Graph Iso- morphism] Let G1 = (V1, E1, A1) and G2 = (V2, E2, A2) be two finite, undirected graphs where V; denotes the set of vertices, E, denotes the set of edges, and A; denotes the node attributes of the graph foror i 2. Let L G - C be a function that maps a graph G G. the set of all finite, undirected graphs, to its canonical label L(G) C, the set of all possible canonical labels, as produced by the Nauty algorithm. Then the following equivalence holds:
L(G,) = L(G2) G1 _ G2
where G1 _ G2 denotes that G1 and G2 are isomorphic.
Lemma 3.2 indicates that the CL process is both complete (sufficient to distinguish non-isomorphic graphs) and sound (not distinguishing actually isomorphic graphs). Note that if LG) corresponds to multiple automorphic labels, we can randomly select one since they are all equivalent and pro- duce the same sequence later through Geo2Seq, as detailed in Appendix B However, this is very uncommon case for real-world 3D attributed graphs like molecules.
The Nauty algorithm (McKay & Piperno 2014), tailored for CL and computing graph automorphism groups, presents a rigorous formulation of CL. In this paper, we adopt the Nauty algorithm for CL caleulation, while all analyses and derivations apply to other rigorous algorithms. The bijective mapping between Cl--obtained sequential repre- sentation and graph can be be proved based on graph iso- morphism. First, due to the geometric need here, we extend to define the isomorphism problem for attributed graphs.
3.2. Invariant Spherical Representations
In this section, we describe how incorporate 3D structure information into our sequences. One main challenge here is to ensure the -iivaaiiie property described in Sec- tion 2.1 Specifically, given a 3D molecule, if it rotated or translated in the 3D space, its 3D representation should be unchanged. Another challenge to ensure no information loss (Liu et al. 2022 Wang et al. 2022). Specifically, given the 3D representation, we can recover the given 3D struc- ture. If two 3D structures cannot be matched via SE(3) transformation, the representations should be different. This property is important to the discriminative ability of models.
Definition 3.1. |Graph Isomorphism] Let G1 (Vi, F1,A1 A1) and G2 = (V2, E2, A2) be two graphs, where V; denotes the set of vertices, E, denotes the set of edges, and A, denotes the node attributes of G2 for i = 2, Let attr(v) denote the node attributes of vertex w. The graphs G2 and G2 are said be isomorphic, denoted as G2 G2. if there exists a bijection b : V1 + Va such that for every vertex E V1, attr(v) A attr(b(v)) E Ag, and for every pair of vertices u, E Vi,
We address these challenges by spherical representations, iee. using spherical coordinates to represent 3D structures.
U) E F. + (b)) E Fz.
4

---

Geometry Informed Tokenization of Molecules for Language Model Generation
6-777 L.6
(0:51, (.25,3413)
atom in the global frame F. As shown in Figure 2 if the input coordinates rotated by a matrix Q and translated by vector b, the transformed spherical coordinates remain the same, so the spherical coordinates are S -invariiia
Next, we demonstrate that there is no information loss in our method. We show that given our (3) spher- ical representations, we can recover the given 3D struc- tures. For each node 4,, we convert the spherical coordinate [de, Be,, De, to coordinate r'g, in 3D space as [de, sin (0c,) cos( de, sin(Or,) sin(cr,) de, cos 0r.]
Equivariant Frame
Spharical Coordinates
Equivariant Frame
Sphenical Coordinates
e - 128.6
1907,0- o0- (r)
(d = 4.23, - 136 00)
Note that our reconstructed coordinate T/, may not be ex- actly the same the original coordinate rg,- However, there exists (F33)-rranfformation g. such that g(r))) = re, for all i. Note that the same transformation is applied to all nodes. Formally, by applying the function f(-) to the 3D coordinate matrix R. we can demonstrate the following properties of spherical representations.
Figure 2: Illustrations of the equivariant frame and invariant spherical representations. If the molecule is rotated and translated by rotation matrix Q and a translation vector b, the atom coordinates change accordingly. But our spherical representations remain invariant since the frame equivari- ant the SE(3) trranfformation.
Lemma 3.3. Let G = (z, R) be 3D graph with node type vector z and node coordinate matrix R. Let be the equivariant global frame of graph G built based on the first three con-colinnea nodes in L(G). f() our function that maps 3D coordinate matrix R of G to spherical representu- tions S under the equivariant global frame F. Then for any 3D transformation g E (SE(3), we have ((R) {(g (I))))
Compared to Cartesian coordinates, spherical coordinate values are bounded in smaller region, namely, a range of [0, T or [0, 2r]. This makes spherical coordinates advanta- geous in discretized representations and thus easier to be modeled by LMs. Given the same decimal place constraints, they require smaller vocabulary size, and given the same vocabulary size, they lose less information. This is also supported by empirical results and analysis in Appendix C
Given spherical representations S = (R). there exist transformation E FF3), such that f (S) RR).
Lemma 3.3 indicates that our spherical representation is and we can reconstruct (a transformation of) the original coordinates. Therefore, our method can convert 3D structures into SE(33-invarian representations with no information loss. Proofs are in Appendix B
We propose to maintain (EE(3)invariinn invvariance while ensuring no information loss. Given a 3D molecule G with atom types z and atom coordinates R, we first build a global coordinate frame F- = (x, y, z) based on the input. Specifically, as shown in Figure the frame is built based on the first three non-collinear atoms in the canonical ordering [G)). Let L1, f2, and t be the indices of these three atoms. Then the global frame F = (x, y, z) is calculated as x normalize (rta T
3.3. Geo2Seq: Geometry Informed Tokenization
In this section, we describe the process and properties of our 3D tokenization method, Geo2Seq. Equipped with canoni- cal labeling that reduces graph structures to ID sequences with no information loss regarding graph isomorphism and SE(3) invaariatt spherical representations that ensure no 3D information loss, we develop Geo2Seq, a reversible trans- formation from 3D molecules to ID sequences. Figure !] shows an overview of Geo2Seq. Specifically, given a graph G with nodes, Geo2Seq concatenates the node vector [zi, di, Oc, (i.) of every node in G to ID sequence by its canonical order, -,,-.. Ln. To formulate the properties of Geo2Seq, we extend the concept of graph isomorphism in Definition B.T to 3D graphs.
y normalize ((rep Tτ, x x),
x X y.
Here normalize()) is the function to normalize a vector to unit length. Note that the global frame is equivariant to the rotation and translation of the input molecule, as shown in Figure[ann and Appendix B.2 After obtaining the global function f(-) to convert the coordinates
frame, we of each atom to spherical coordinates d, 0, ó under this frame. Specifically, for each node l with coordinate Tł,, the corresponding spherical coordinate is
de. ||r r
Definition 3.4. [3D Graph Isomorphism] Let Gi and G2 = (#2, R2) be two 3D graphs, where z; is the node type vector and R, is the node coordinate matrix of the molecule G,. Let Vú, denote the set of vertices, A, denote node attributes, and no edge exists. Two 3D graphs
02. = arccos ((ri: (T,) =z/de), Φe. atan2 ((re, - -y, (re. - -r()) xx).
(2)
The spherical coordinates show the relative position of each
5

---

Geometry Informed Tokenization of Molecules for Language Model Generation
G1 and G2 are 3D isomorphic, denoted 11 3DD G2. if there exists a bijection b: Vi - V2 such that G₁ = G2 given A, = [zi, Rij. and there exists 3D transformation E (3) such that Ci g( b(8)
Corollary 3.6 extends Theorem 3.5 guarantees for the practical use of Geo2Seq. If we allow a round-up error below |10 |10-"1/2 /2 for coordinates when distinguishing 3D iso- morphism, all properties still hold. This implies that the practical Geo2Seq implementation retains near-complete geometric information and invariance, with numerical preci- sion of S 110-
G2 ). If small error e is
Gz ) allowed such that |r;" G. -9( ((7)()) graphs e-constrained 3D isomorphic.
I e, we call the two 3D
Considering Lemma 3.2 we specify G E, with A = [z, R. and define the CL function for 3D molecules as Lm. which extends the equivalence of Lemma 3.2 to Lm with 3D isomorphism. We formulate Geo2Seq and our major theoretical derivations below.
With discreteness incorporated, we can collect finite vo- cabulary covering all accessible molecule samples to enable tokenization for LMs. Specifically, we use vocabularies of approximately 1K-16K tokens consisting of atom type to- kens 'C, N, ON,O ... and spherical coordinate tokens such as 1-9''', '1.570' or --0.0320* Specifically, the vocabulary size is approximately 1.8K for the QM9 dataset, and 16K for the Gcom-Drug dataset. Note that we consider chirality for atoms and use the special token suffixes to distinguish clockwise and counterclockwise chiral cen- ters, for example, 'C@' and 'C@ @+ The numerical tokens range from the smallest to the largest distance and angle values with restricted precision of 2 or 3 decimal places. Experimental results show the benefits in using this level of tokenization, as detailed in Appendix C
Theorem 3.5. [Bijective Mapping between 3D Graph and S Following Definition 3.4 let G1 = (z1, R
and G2 = = R2) be two 3D graphs. Let LG) be the canonical label for 3D graph G and f : R + S be the function that maps 3D coordinates to its spheri- cal representations. Given a graph G with nodes and X =| ".,]" E Raxm, where in E Z, we define (m) @ x = (oncat(ze....... of the node labeled by L, (G). and concat(-) concatenates elements as a sequence. We define
"@' and '@ @'
where Eg is the index
Ge) = Lm) 2 (z. ((R)) Lm & X,
4. 3D Molecule Generation
where x; = =Zi, tz, dil. Then Geo2Seq G + U is a
surjective function, and the following equivalence holds:
Training and Sampling, Now that we have defined a canon- ical and robust sequence representation for 3D molecules, we turn to the method of modeling such sequences, U. Here, we attempt to train a model M with parameters @ to capture the distribution of such sequences, PU)), in our dataset. As this is well-studied problem within language modeling, we opt to use two language models, GPT (Radford et al. (0188) and Mamba (Gu & Dao 2023), which have shown effective sequence modeling capabilities on a range of tasks. Both models are trained using a standard next-token prediction cross-entropy loss for all elements in the sequence:
Geo2Seq(Gi) = Gee2SSe((G) <> G1 Man G2,
where G1 3DD G2 denotes G1 and G2 are 3D isomorphic.
Theorem 3.5 establishes the following guarantees for Geo2Seq: (1) Given 3D molecule, can uniquely con- struct a 1D sequence using Geo2Seq. (2) If two molecules are 3D isomorphic, their sequence outputs from Geo2Seq are identical.
(3) Given a sequence output of Geo2Seq, we can uniquely reconstruct 3D molecule. (4) If two constructed sequences from Geo2Seq are identical, their corresponding molecules must be 3D isomorphic. This en-
able sequential tokenization of 3D molecules, preserving structural completeness and geometric invariance.
min 0 WEU
E (Mo(un, ... ..s) us+1)
Due to the necessity of discreteness in serialization and tokenization for LMs, in reality, numerical values need to be discretized before concatenation. In practice, we round up numerical values to certain decimal places. Thus Theo- rem 3.5 can be extended with constraints, as below.
To sample from trained model, we first select an initial atom token by sampling from the multinomial distribution of first-tokens in the training data (we note that in almost all cases this is 'H'). We then perform a standard autoregressive sampling procedure by iteratively sampling from the condi- tional distribution Po (14+1/u,, , .u.) until the stop token or max length is reached. We sample from this distribution using top-k sampling (Fan et al. 2018) and softmax tem- perature (Ackley et al. 1985 Ficler & Goldberg 2017). Unless otherwise noted, 7 0.7 and k = 80.
Corollary 3.6. |Constrained Bijective Mapping between 3D Graph and Sequence] Following the notations and defi- nitions of Theorem 3.5] let spherical coordinate values be rounded up to b decimal places. Then Geo2Seq: G - U is a surjective function, and the following equivalence holds:
Geo2Seg (G))) Ge))Seq(G2) + G1 23D- 11/2 G2,
Controllable Generation. For controllable generation, we follow Bagal et al. (202T) and use a conditioning token for the desired property. This token is created by projecting the
where G1 230-11001122 G2 denotes graphs G1 and G2 are (100=\1/2-connstriine 3D isomorphic.
6

---

Geometry Informed Tokenization of Molecules for Language Model Generation
desired properties through a trainable linear layer to create vector with the model's initial token embedding space. This property token is then used as the initial element in the molecular sequence. Training and sampling are performed as before with this new sequence formulation. Sampling begins with the desired property's 's token as input.
In addition, we report the percentage of valid molecules that can be successfully converted to SMILES strings by RDKit, and the percentage of valid and unique molecules that can be converted to unique SMILES strings.
Baselines. We compare GPT and Mamba models with sev- eral strong baseline methods. Specifically, we compare with an autoregressive generation method G-SchNet (Gebauer ct al.
5. Experimental Studies
2019) and an equivariant flow model based method
In this section, we evaluate the method of generating 3D molecules in the form of our proposed Geo2Seq represen- tations by LLMs. We show that in the random generation
E-NFs (Satorras et al. 202 We also compare with some recently proposed diffusion based methods, including EDM (Hoogeboom et al. 2022). GDM (the non-equivariant variant of EDM) and GDM-AUG (GDM trained with ran- dom rotation as data augmentation). Besides, we compare with EDM-Bridge (Wu et al., 2022) and GROLDM (Xu et al. 2023a), which are two latest 3D molecule generation methods improving EDM by SDE based diffusion models and latent diffusion models, respectively. To ensure that the comparison is fair, our methods and baseline methods use the same data split and evaluation metrics.
task (see Section 2 the performance of Geo2Seq with
GPT (Radford et al. 2018) or Mamba (Gu & Dao 2023)
models is better than or comparable with state-of-the-art 3D
point cloud based methods, including EDM (Hoogeboom
al. 2022) and GEOLDM (Xu al.) 2023a). In addi- tion, in the controllable generation task (see Section 2.1
we show that Geo2Seq with Mamba models outperform previous 3D point cloud based methods by a large margin.
We present the random generation results of dif- ferent methods on QM9 and GEOM-DRUGS datasets in Table [ Note that for GEOM-DRUGS dataset, all methods achieve nearly 0% molecule stability percentage and 100% uniqueness percentage. Thus, following previous studies, these two metrics are omitted. According to the results in Table D on QM9 dataset, generating 3D molecules in Geo2Seq representations with either GPT or Mamba models achieve better performance than all 3D point cloud based baseline methods in molecule stability and valid percent- age, and achieves atom stability percentages close to the upper bound (99%). This demonstrates that our method can model 3D molecular structure distribution and capture the underlying chemical rules more accurately. It is worth noticing that our method does not achieve very high unique- ness percentage, showing that it is not easy for our method to generate a large number of diverse molecules. We be- lieve this is due to that the conversion from real numbers to discrete tokens limits the search space of 3D molecular structures, especially on small dataset like QM9, while it is casier to generate more diverse molecules for 3D point cloud based methods as they directly generate real numbers. This is reflected by the fact that our method achieves nearly 100% uniqueness percentage on the large GEOM-DRUGS dataset. On GEOM-DRUGS dataset, both GPT and Mamba models achieve reasonably high atom stability and valid percentage. The performance of our method is comparable with strong diffusion based baseline methods, showing that LLMs have the potential to model very complicated drug molccular structures well. We will explore further improv- ing the performance on GEOM-DRUGS dataset with larger LLMs in the future.
Results.
5.1. Random Generation
Data. We adopt two datasets, QM9 (Ramakrishnan et al. 2014) and GEOM-DRUGS (Axelrod & Gomez-Bombarcl/
2022), to evaluate performances in the random generation task. The QM9 dataset collects over 130k 3D molecules with 3D structures calculated by density functional theory (DFT). Each molecule in QM9 has less than heavy atoms and its chemical elements all belong to H. C, O, F. Fol- lowing Anderson et al. (2019), we split the dataset into train, validation and test sets with 100k, 18k and 12k samples, sep- arately. The GEOM-DRUGS dataset consists of over 450k large molccules with 37 million DFT-calculated 3D struc- tures. Molecules in GEOM-DRUGS has up to 181 atoms and 44.2 atoms on average. We follow Hoogeboom et al. (2022) to select 30 3D structures with the lowest energies per molecule for model training.
Setup. On the QM9 dataset, we set the training batch size to 32, base learning rate to 0.0004, and train a 12-layer GPT model and a 26-layer Mamba model by AdamW (Loshchilov & Hutter 2019) optimizers. On the GEOM-DRUGS dataset, we set the training batch size to 32, base learning rate to 0.0004, and train a 14 14yle GPT model and 28-layer Mamba model by AdamW optimizers. See Appendix Doo more information about hyperparame- ters and other settings. When model training is completed, we randomly generate 10,000 molecules, and evaluate the performance on these molecules. Specifically, we first trans- form 3D molecular structures to 2D molecular graphs using the bond inference algorithm implemented in the official code of EDM. Then, we evaluate the performance by atom stability, which is the percentage of atoms with correct bond valencies, and molecule stability, which is the percentage of molecules whose all atoms have correct bond valencies.
See Appendix D.3 for additional experiments and met-
7

---

Geometry Informed Tokenization of Molecules for Language Model Generation
Table 1: Random generation performance QM9 and GEOM-DRUGE datasets with 3 runs. Larger numbers indicate better performance. bold and underline highlight the best and second best performance, respectively. For molecule stability and unique percentage are close to 0% and 100% for all methods so they are not presented.
dataset,
QM9
GEOM-DRUGS
Method
Atom Sta (%)
Mol Sta (%)
Valid (%)
Valid & Unique (%)
Atom Sta (%)
Valid (%)
Data
99.0
95.2
97.7
97.7
86.5
99.9
E-NFs
85.0
4.9
40.2
39.4
G-SchNet
95.7
68.1
85.5
80.3
GDM
97.0
63.2
75.0
90.8
GDM-AUG
97.6
71.6
90.4
89.5
77.7
91.8
EDM
98.7++0.1
82.0=0.4
91,9++0.5
90.7-0.6
81.3
92.6
EDM-Bridge
98.8:10.1
8..6=0033
92.0:10.1
90.7=0.1
82.4
92.8
GEOLDM
98,9:10.1
09.4:0055
93.8:10.4
92.7=0.5
84.4
99.3
Geo2Seq with GPT
98.3++0.1
90.3-0.1
94,8+002
80.6-0.4 0.4
82.6
87.4
Geo2Seq with Mumba
98.9
9..2:002
97.1+-0.2
81.7-0.4
82,5
96.1
Table 2: Controllable generation performance on QM9 datasets. Smaller numbers indicate better performance,
Property (Units)
(Bohr")
Ac (meV)
сhомO (meV)
élumo (meV)
μa (D)
C, ( (oo!
K)
Data
0.10
64
39
36
0.043
0.040
Random
9.01
1470
645
1457
1.616
6.857
Natcens
3.86
866
426
813
1.053
1.971
EDM
2.76
655
356
584
1.111
1.101
GEOLDM
2.37
587
340
522
1.108
1.025
Geo2Seq with Mamba
0.46
98
57
71
0.164
0.275
Geo2Seq with GPT
0.53
102
48
53
0.097
0.325
rics (Huang et al. 2023 Vignac et al.
2023 ), Appendix C
molecules. Note that we use the trained EGNN based prop- erty prediction models to calculate the property values of the generated molecules.
for ablation studies, Appendix D for complexity analysis, and Appendix F for visualization.
Baselines. We compare our models with two equivariant diffusion models, EDM (Hoogeboom et al. 2022) and GE- OLDM (Xu et al. 2023a). In addition, we use several baselines that are based on dataset molecules. One baseline (Data) directly taking the molecules from the QM9 dataset and use their property values as conditions. The MAE met- ric simply reflects the prediction error of the trained property prediction model, which can be considered as a lower bound. The second baseline (Random) is taking the molecules from the dataset but uses the randomly shuffled property values as conditions, and its MAE can be considered as an upper bound. The third baseline (Natoms) uses the molecules from the dataset but property values predicted from the num- ber of atoms conditions. Achieving better performance than this baseline shows that models can use conditional information beyond the number of atoms.
5.2. Controllable Generation
Data. In the controllable generation task, we train our mod- els molecules and their property labels in the QM9 (Ra- makkishhan et al., 2014) dataset. Specifically, we try taking a certain quantum property value as the conditional input to LLMs, and train LLMs to generate molecules with the conditioned quantum property values. Following Hooge boom et al. (022), we split the training dataset of QM9 to two subsets where each subset has 50k samples, and train conditional generation models and EGNN (Satorras let al. 2021b) based quantum property prediction models these two subsets, respectively. We conduct the controllable generation experiments on six quantum properties from QM9, including polarizability (a), HOMO energy (<HoMO), LUMO energy (elUMO), HOMO-LUMO gap (Ae), dipole moment (p) and heat capacity at 298.15K (C.).
Results. Controllable generation results of different meth- ods are summarized in Table [ As shown in the table, among all six properties, method outperforms the strong diffusion based baseline methods EDM and GEOLDM by large margin. Our method moves a significant step in pushing the performance of controllable generation task towards the lower bound, í.e., Data baseline. As we use the same training set as EDM and GEOLDM to train the conditional generation model, the good performance of our
Setup. For the controllable generation experiment, train 16-layer Mamba (Gu & Dao 2023) models with the same hyperparameters as the random generation experiments in Section 5.1 To evaluate the performance, we sample 10000 quantum property values, generate molecules conditioned on these property values by trained models, and compute the mean absolute difference (MAE) between the given property values and the property values of the generated
8

---

Geometry Informed Tokenization of Molecules for Language Model Generation
method shows that LLMs have more powerful capacity in incorporating conditional information into the 3D molecular structure generation process. We believe that the powerful context correlation capturing structures from LLMs, e.g., attention mechanism, play significant roles in achiev- ing the good control of 3D molecule generation by the con- ditioned property values. The huge success of LLMs in controllable molecule generation will motivate broader ap- plications of LLMs in goal-directed or constrained drug design. See Appendix for visualization of molecules gen- erated from some given polarizability values.
References
Achiam, J., Adler, S., Agarwal, S., Ahmad, L., Akkaya, I., Aleman, F. L.,, Almeida, D., Altenschmidt, J., Altman, S., Anadkat, S., et al. GPT-4 technical report. arXiv preprint SBXiv:2303.087744 2023.
Ackley, D. H., Hinton, G. E., and Sejnowski, T. A learning algorithm for boltzmann machines. Cognitive science, (1):147-169, 1985.
Anderson, B., Hy, T. S., and Kondor, R. Cormorant: Covariant molecular neural networks. In Wallach, H., Larochelle, H., Beygelzimer, A., d'Alché-Bue F., Fox. E., and Gamett, R. (eds.), Advances in Neural Information Processing Systems, vol- ume 32. Curran Associates, Inc., 2019. URL http://ppoeeddnngnnnuuuppccppppp hils/papee/211199 flB033770022277666cccc2999122223 Paper.pdf
6. Conclusion and Discussion
Geo2Seq showcases the potential of pure LMs in revolution- izing molecular design and drug discovery when geometric information is properly transformed. The framework has certain limitations, particularly in the generalization abil- ities across the continuous domain of real numbers. Due to the discrete nature of vocabularies, LMs rely on large pre-training corpus, fine-grained tokenization or emergent abilitics for better generalization, as a trade-off to high pre- cision and versatility. Future works points towards several directions, such as expanding on conditional tasks and ex- ploring advanced tokenization techniques.
Axelrod, S. and Gomez-Bombarelll, R. GEOM, energy-
annotated molecular conformations for property predic- tion and molecular generation. Scientific Data, 91):185, 2022.
Bagal, V., Aggarwal, R., Vinod, P., and Priyakumar, U. D. MolGPT: molecular generation using a transformer- decoder model. Journal of Chemical Information and Modeling, 62(9):2064 2066 2021.
Acknowledgments
This work was supported partially by National Science Foun- dation grant 100-2006861 and National Institutes of Health grant UÓI AG0701 12 (to S.J.), and by the Molecule Maker Lab Institute (to H.J.): an AI research institute program supported by NSF under award No. 2019897. The views and conclusions contained herein are those of the authors and should not be interpreted as necessarily representing the official policies, either expressed or implied, of the U.S. Govermment. The U.S. Government is authorized to re- produce and distribute reprints for governmental purposes notwithstanding any copyright annotation therein.
Bajorath, J. Chemical language models for molecular design. Molecular Informatics, 43(1):=202300028, 2024.
Banck, M., Morley, C. A., Vandermeersch, T., and Hutchi-
son, G. R. Open Babel: An open chemical toolbox, 2011. URL https://doi.rrg/ 10. 10..1.111111..11 010111777-2946---3 2946-3-33
Bao, F., Zhao, M., Hao, Z., Li, P., Li, C., and Zhu J. Equivariant energy-guided SDE for inverse molecular design. In The Eleventh International Conference on Learning Representations, 2023. URL https://oppeereiiww net/forum \OW
Impact Statement
Blanchard, A. E., Bhowmik, D., Fox, Z., Gounley, J., Glaser, J., Akpa, B. S., and Irle, S. Adaptive language model training for molecular design. Journal of Cheminformat- ics, 15(1):1-12, 2023.
This paper presents work whose goal is to advance the field of Deep Learning, particularly 3D molecule gener- ation. While this research primarily contributes to technical advancements in generative modeling, has potential impli- cations in domains such drug discovery. We acknowledge that generative models in science could be misused if not carefully applied. However, we emphasize the importance of responsible deployment and alignment with ethical guide- lines in generative AL. Overall, our contributions align with the broader goal of machine learning methodologies, and we do not foresee any immediate ethical concerns beyond those generally associated with generative models.
Bran, A. M. and Schwaller, P. Transformers and large language models for chemistry and drug discovery. arXiv preprint 2rX.v:2310.06083, 2023.
Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., et al. Language models are few-shot learners. Advances in neural information processing systems, 33: 1877-1901, 2020.
9

---

Geometry Informed Tokenization of Molecules for Language Model Generation
Buttenschoen, M., Morris, G. M., and Deane, C. M. Pose- Busters: Al-hased docking methods fail to generate phys- ically valid poses or generalise to novel sequences, 2023.
Ficter, and Goldberg, Y. Controlling linguistic style aspects in neural language generation. arXiv preprint arXiv:1707.02633, 2017.
Cao, H., Liu, Z., Lu, X., Yao, Y., and Li, Y. InstructMol: Multi-modal integration for building versatile and reli- able molecular assistant in drug discovery, arXiv preprint arXiv:2311. 2023.
Flam-Shepherd, D. and Aspuru-Guzik, A. Language models can generate molecules, materials, and protein binding sites directly in three dimensions as xyz, cif, and pdb files. arXiv preprint arXiv:2305.05700 2023.
Chen, Y., Wang, 2., Zeng, Li, Y., Li, P. Ye, X., and Sakurai, T. Molecular language models: Rnns or trans- former? Briefings in Functional Genomics, pp. elad012, 2023a.
Frey, N. C., Soklaski, R., Axelrod, S., Samsi, S., Gomez- Bombarelli, R.. Coley, C. W., and Gadepally, V. Neural scaling of deep chemical models. Nature Machine Intelli- gence, 5(11):1297-11305 2023.
Chen, Y., Xi, N.. Du, Y., Wang, H., Jianyu, C., Zhao, S., and Qin, B. From artificially real to real: Leveraging pseudo data from large language models for low low-resource molecule discovery. arXiv preprint arXiv:2309.05200 0550030,, 2023b.
Fuchs, F., Worrall, D., Fischer, V., and Welling, M. Se (r)-transfr:m:rs: 3d roto-translation equivariant attention networks. Advances in neural information processing systems, 33:1970- 1981, 2020.
Ganea, O.-E., Pattanaik, L., Coley, C. W., Barzilay, R.,
Chowdhery, A., Narang, S., Devlin, J., Bosma, M., Mishra, G., Roberts, A., Barham, P., Chung, H. W., Sutton, C., Gehrmann, S., et al. PaLM: Scaling language modeling with pathways. Journal of Machine Learning Research, 24(240):1-113, 2023.
Jensen, K., Green, W., and Jaakkola, T. S. GeoMol: Torsional geometric generation of molecular 3D con- former ensembles. In Beygelzimer, A., Dauphin, Y., Liang, P., and Vaughan, J. W. (eds.), Advances in Neu- rul Information Processing Systems, 2021. URL https //opemreiie..ee//foummdlallll hhnggiiiidiaaanannnNn hng9tuNj
Cuthill, E. and McKee, J. Reducing the bandwidth of sparse symmetric matrices. In Proceedings of the 1969 24th national conference, pp. 157-172, 1969.
Gebauer, N., Gastegger, M., and Schütt, K. Symmetry- adapted generation of 3D point sets for the targeted dis- covery of molecules. Advances in neural information processing systems, 32, 2019.
Daigavane, A., Kim, S., Geiger, M., and Smidt, T. Sym- phony: Symmetry-equivarre point-centered spherical harmonics for molecule generation. arXiv preprint arXiv:2311.16199 2023.
Gogineni, T., Xu, Z., Punzalan, E., Jiang, R., Kammeraad, J., Tewari, A., and Zimmerman, P. TorsionNet: a reinforce- ment learning approach to sequential conformer search. In Larochelle, H., Ranzato, M., Hadsell, R., Balcan, M. F., and Lin, H. (eds.), Advances in Neural Information Pro- cessing Systems, volume 33, pp. 20142 221133 Curran Associates, Inc., 2020.
Devlin, J., Chang, M.-W., Lee, K., and Toutanova, K. Bert: Pre-training of deep bidirectional transformers for lan- guage understanding, arXiv preprint arXiv:1810.0480) 2018.
Dijkstra, E. W. A note on two problems in connexion with graphs. In Edsger Wybe Dijkstra: His Life, Work, and Legacy, pp. 287-290. 2022.
Gu, A. and Dao, T. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.00752,022 2023.
Du, Y., Fu, T., Sun, J., and Liu, S. Molgensurvey: A sys- tematic survey machine learning models for molecule design. arXiv preprint arXiv:2203. 14500, 2022.
Gu, A., Goel, K., and Ré, C. Efficiently modeling long sequences with structured state spaces. arXiv preprint arXiv:2111.00390 2021.
Edwards, C., Lai, T., Ros, K., Honke, G., Cho, K., and Ji, H. Translation between molecules and natural language. arXiv preprint arXiv:2204.1 /817, 2022.
Haroon, S., Hafsath, C., and Jereesh, A. Generative pre- trained transformer (GPT) based model with relative at- tention for de novo drug design. Computational Biology and Chemistry, 106: 106:107911. 2023.,,0023....
Fan, A., Lewis, M., and Dauphin, Y. Hierarchical neural story generation. arXiv preprint arXiv:1805.044 0005004333,... 2018.
Hilbert, D. and Hilbert, D. Über die stetige abbildung einer linie auf ein flächenstück.. Dritter Band: Analy- sis- Grundlagen der Mathematik- Physik Verschiedenes: Nebst Einer Lebensgeschichte, pp. 1-2, 1935.
Fang, Y., Zhang, N., Chen, Z., Fan, X., and Chen, H.
Molecular language model as multi-task generator. arXiv preprint arxiv:2301.1 1/259, 2023.
10

---

Geometry Informed Tokenization of Molecules for Language Model Generation
Hoffmann, M. and Noé, F. Generating valid Euclidean dis- tance matrices. arXiv preprint /rXiv39900001::3 2019.
Li, J., Liu, Y., Fan, W. Wei, X.-Y,, Liu, H., Tang, J., and Li, Q. Empowering molecule discovery for molecule- caption translation with large language models: A chatgpt perspective. arXiv preprint arXiv:2306 200.0:23306.06615, 2023a.
Hoogeboom, E.,, Satorras, V. G., Vignac, C., and Welling, M. Equivariant diffusion for molecule generation in 3D.
In Proceedings of the 39th International Conference on
Li, Y., Gao, C., Song, X., Wang, X., Xu, Y., and Han, S. DrugGPT: GPT-based strategy for designing potential ligands targeting specific proteins. bioRxiv, pp. 2023-06, 2023b.
Machine Learning, volume 162 of Proceedings of Ma- chine Learning Research, pp. 8867 8887. PMLR, 2022. URL http://moeeeddiin.....sssss////p
Liang. Y., Zhang, R., Zhang, L., and Xie, P. DrugChat: towards enabling ChatGPT-like capabilities on drug molecule graphs. arXiv preprint arXiv:2309.0300 20w.2330903307,..0 2023.
Huang, H., Sun, L., Du, B., and Lv, W. Learning joint 2d & 3d diffusion models for complete molecule generation. arXiv preprint arXiv:2305.1234 2023.
Liao, Y.-L. and Smidt, T. Equiformer: Equivariant graph attention transformer for 3d atomistic graphs. arXiv preprint arXiv:2206. 2022.
Izdebski, A., Weglarz-Tomezak, E.,, Szezurek, F., and Tom- czak, J. M. De novo drug design with joint transformers. arXiv preprint arXiv:2310.02066, 2023.
Liu, S., Nie, W., Wang, C., Lu, J., Qiao, Z., Liu, L., Tang, J., Xiao, C.,, and Anandkumar, A. Multi-modal molecule structure-text model for text-based retrieval and editing, Nature Machine Intelligence, 5(12):1447-1457, 2023a.
Janakarajan, N., Erdmann, T., Swaminathan, S., Laino, T., and Born, J. Language models in molecular discovery. arXiv preprint arXiv:2309. 16235, 2023.
Liu, Y., Wang, L., Liu, M., Lin, Y., Zhang, X., Oztekin, B., and Ji, S. Spherical message passing for 3d molecular graphs. In International Conference on Learning Repre- sentations, 2022. URL https://openreiie.ntttooommm givsRXsO19r
Jing, B., Corso, G., Chang, J., Barzilay, R., and Jaakkola, T. S. Torsional diffusion for molecular conformer genera- tion. In Oh, A. H., Agarwal, A., Belgrave, D., and Cho, K. (eds.), Advances in Neural Information Processing Systems, 2022. URL https://ppeneeeee.ee//fommm//// w6ff2r62r_HI
Liu, Z., Zhang, W., Xia, Y., Wu, L., Xie, S., Qin, T., Zhang, M., and Liu, T.-Y. Molxpt: Wrapping molecules with text for generative pre-training, arXiv preprint arXiv:2305.10683 2023b.
Kim, Y. and Kim, W. Y. Universal Structure Conver- sion Method for Organic Molecules: From Atomic Con- nectivity to Thres--Dimnnioonal Geometry. Bulletin of the Korean Chemical Society, 36(7):179-11777, 2015. doi:
Loshchilov, and Hutter, Decoupled weight decay regu- larization. In International Conference on Learning Rep- resentations, 2019. Ut.///upps//ppeeeiieeee//efoooo id=Bkg6RiCyY7,
10334. URL https:
fontineiiiaag.iwie..ccco/oo////d 10334
Krenn, M., Häse, F., Nigam, A., Friederich, P., and Aspuru- Guzik, A. Selfies: a robust representation of semantically constrained graphs with example application in chem- istry. arXiv preprint 1rx.v:1905.13777 13741, 1(3), 2019.
Luo, R., Sun, L., Xia, Y., Qin, T., Zhang, S., Poon, H., and Liu, T.-Y. BioGPT: generative pre-trained transformer for biomedical text generation and mining. Briefings in Bioinformatics, 23 2022a.
Kyro, G. W., Morgunov, A., Brent, R. I., and Batista, V. S. Chemspaceal: An efficient active learning methodology applied protein-specific molecular generation. ArXiv, 2023.
Luo, S., Chen, T., Xu, Y., Zheng, S., Liu, T.-Y,, Wang, L., and He, D. One transformer can understand both 20 & 3d molecular data, arXiv preprint arXiv:2210.0176:5 2022b.
Luo, Y. and Ji, S. An autoregressive flow model for 3D molecular geometry generation from scratch. In Interna- tional Conference on Learning Representations, 2022.
Langley, P. Crafting papers on machine learning. In Langley, P. (,d.), Proceedings of the 17th International Conference on Machine Learning 2000), pp. 1207-1216, Stan- ford, CA, 2000. Morgan Kaufmann.
Mansimov, E.,, Mahmood, O., Kang, S., and Cho, K. Molec- ular geometry prediction using a deep generative graph neural network. Scientific reports, 1-13, 2019.
Lee, C. Y. An algorithm for path connections and its appli- cations. IRE transactions on electronic computers, (3): 346-365, 1961.
Mao, J., Wang, J., Zeb, A., Cho, K.-H., Jin, H., Kim, J., Lee, O., Wang, Y., and No, K. T. Deep molecular generative

---

Geometry Informed Tokenization of Molecules for Language Model Generation
model based on variant transformer for antiviral drug design. Available at SSRN 43458/ 1, 2023a.
Satorras, V. G., Hoogeboom, E., and Welling, M. E(n) equivariant graph neural networks. In Meila, M. and Zhang, T. (eds.), Proceedings of the 38th International Conference on Machine Learning, volume 139 of Pro- ceedings of Machine Learning Rexearch, pp. 9323-9332. PMLR, 18-24 Jul 202 lb. URL httpss//ppooceecdingg..s mir press/v 1 39/satorras22a...o la.hmml
Mao, J., Wang, J., Zeb, A., Cho, k.-H., Jin, H., Kim, J., Lee, O., Wang. Y., and No, K. T. Transformer-base molecular generative model for antiviral drug design. Journal of Chemical Information and Modeling, 2023b.
Masters, D., Dean, J., Klaser, K., Li, Z., Maddrell- Mander, S., Sanders, A., Helal, H., Beker, D., Ram- pášek, L., and Beaini, D. Gps++: An optimised hy- brid mpun/transformer for molecular property prediction. arXiv preprint arXiv:2212.0222: 2022.
Schaeffer, R., Miranda, B., and Koyejo, S. Are emergent abilities of large language models a mirage? Advances in Neural Information Processing Systems, 36, 2024.
Shi, C., Luo, S., Xu, M., and Tang, J. Learning gradient fields for molecular conformation generation. In Meila, M. and Zhang, T. (eds.), Proceedings of the 38th Inter- national Conference on Machine Learning, volume 139 Proceedings of Machine Learning Research, pp. 9558- 9568. PMLR, 18-24 Jul 2021.
Mazuz, E., Shtar, G., Shapira. B., and Rokach, L. Molecule generation using transformers and policy gradient rein- forcement learning. Scientific Reports, 13(1):8799, 2023.
McKay, B. D. and Piperno, A. Practical graph isomorphism, ii Journal of symbolic compulation, 60:94-112, 2014.
Shi, Y., Zheng, S., Ke, G., Shen, Y., You, J., He, J., Luo, S., Liu, C., He, D., and Liu, T.-Y. Benchmarking graphormer on large-scale molecular modeling datasets. arXiv preprint 2203::2200.0881111 2022.
McKay, B. D. et al. Practical graph isomorphism. 1981.
Özşelik, R., de Ruiter, S., and Grisoni, F. Structured state- space sequence models for de novo drug design. 2023.
Simm, G. and Mermandez-Lobate J. M. A generative model for molecular distance geometry. In III H. D. and Singh, A. (eds.), Proceedings of the 37th International Confer- ence on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pp. 8949-8958. PMLR, 13 18 Jul 2020.
Özęelik, R., de Ruiter, S., Criscuolo, E., and Grisoni, F. Chemical language modeling with structured state spaces. 2024.
Pei, Q., Zhang, W., Zhu, J., Wu, K., Gao, K., Wu, L., Xia, Y., and Yan, R. Biot5: Enriching cross-modal integration in biology with chemical knowledge and natural language associations. arXiv preprint arXiv:2330.00 0.0777, 2023.
Thölke, P. and De Fabritis, G. Equivariant transformers for neural network based molecular potentials. In Interna- tional Conference on Learning Representations, 2021.
Qiang. B., Song, Y., Xu, M., Gong, J., Gao, B., Zhou, H., Ma, W.-Y., and Lan, Y. Coarse-to-fine fine: a hierar- chical diffusion model for molecule generation in 3D. In Krause, A., Brunskill, E., Cho, K., Engelhardt, B., Sabato, S., and Scarlett, J. (eds.), Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Re- search, pp. 28277-28299, PMLR, 23-29 Jul 2023. URL
Thomas, H. et al. Introduction to algorithms, 2009.
Touvron, H., Lavril, T., Izacard, G., Martinet, X., Lachaux, M.-A., Lacroix, T., Rozière, B., Goyal, N., Hambro, E., Azhar, F., et al. Llama: Open and efficient foundation lan- guage models. arXiv preprint arXiv:2302.13377 2023.
Tysinger, E. P., Rai, B. K., and Sinitskiy, A. V. Can quickly learn to "translate" bioactive molecules with transformer models? Journal of Chemical Information and Modeling, 63(6):1334-1444 2023.
Radford, A., Narasimhan, K., Salimans. T., Sutskever, I., et al. Improving language understanding by generative pre-training. 2018.
Uhrin, M. Through the eyes of a descriptor: Construct- ing complete, invertible descriptions of atomic envi- ronments. Phys. Rev. B, 104:144110, Oct 2021. doi: 10. 10.1103/hhy.Rev..044.I 104. URL https://ink..ps org//oi//0..1033/hyysseo 104.1..42
Ramakrishnan, R., Dral, P. O., Rupp, M., and Von Lilienfeld, O. A. Quantum chemistry structures and properties of 134 kilo molecules. Scientific data, 1(1):1-7, 2014.
Ünlü, A., Çevrim, E., Sangün, A., Çelikbilek, H., Güvenilir, H. A., Koyaş, A., Kahraman, D. C., Rifaioglu, A., and Olgaç. A. Target specific de novo design of drug candi- date molecules with graph transformer based generative adversarial networks. arXiv preprint arXiv:2302. 07868, 2023.
Satorras, V. G., Hoogeboom, E., Fuchs, F. B., Posner, I., and Welling, M. E(n) equivariant normalizing flows. In Advances in Neural Information Processing Systems, 10. URL http:s//oppacvveewwweer/rrr////d Row VA
12

---

Geometry Informed Tokenization of Molecules for Language Model Generation
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L.,, Gomez, A. N., Kaiser, t.,, and Polosukhin, I. At- tention is all you need. Advances in neural information processing systems, 30, 2017.
Xu, M., Yu, L.., Song, Y., Shi, C., Ermon, S., and Tang. J. GeoDiff: A geometric diffusion model for molecular conformation generation. In International Conference on Learning Representations, 2022. URL https://pperreiiew net/forum PeroooumnYddPPeevvMMMovoo
Vignac, C., Osman, N., Toni, L., and Frossard, P. Midi: Mixed graph and 3d denoising diffusion for molecule generation. In Joint European Conference on Machine Learning and Knowledge Discovery in Databases, pp- 560-576. Springer, 2023.
Xu, M., Powers, A., Dror, R., Ermon, S., and Leskovec, J. Geometric latent diffusion models for 3D molecule generation. In Proceedings of the 39th International Con- ference on Machine Learning, Proceedings of Machine Learning Research. PMLR, 2023a.
Wang, I., Liu, Y., Lin, Y., Liu, H., and Ji, S. ComENet: Towards complete and efficient message passing for 3d molecular graphs. In Advances in Neural Information Processing Systems, 2022. URL WRLLbtts:///peeeeveew.edd forum Moumm?ddmm22MqeWSH)
Xu, Z., Luo, Y., Zhang, X.,, Xu, X.,, Xie, Y., Liu, M., Dicker- son, K., Deng, C., Nakata, M., and Ji, S. Molecule3d: A benchmark for predicting 3d geometries from molecular graphs. arXiv preprint arXiv:2110.0171, 2021c
Wang, Z., Wang, Z., Srinivasan, B., Ioannidis, V. N., Rang- wala, H., and Anubhai, R. Biobridge: Bridging biomed- ical foundation models via knowledge graph. arXiv preprint ar arXiv:2310.03320, 2023.
Xu, Z., Lei, X., Ma, M., and Pan, Y. Molecular gener- ation and optimization of molecular properties using a transformer model. Big Data Mining and Analytics, 7(1): 142-155, 2023b.
Weininger, D. SMILES, chemical language and informa- tion system. 1. introduction to methodology and encoding rules. Journal of chemical information and computer sci- 1988.
Yoshikai, Y., Mizuno, T., Nemoto, S., and Kusuhara, H. Difficulty in learning chirality for transformer fed with smiles. arXiv preprint arXiv:2303. 11593, 2023.
ences,
Zhang, Q., Ding, K., Lyv. T., Wang. X., Yin, Q.. Zhang, Y.,
Wu, J.-N., Wang, T., Chen, Y., Tang, L.-J,, Wu, H.-L., and Yu, R.-Q. Fragment-based t-smiles for de novo molecular generation, arXiv preprint arXiv:2301.01829 2023.
Yu, J., Wang, Y., Li, X., Xiang, Z., et al. Scientific large language models: A survey on biological & chemical domains. arXiv preprint arXiv:2401.14650 2024.
Wu, L.. Gong, C., Liu, X., Ye, M., and qiang liu. Diffusion- based molecule generation with informative prior bridges. In Oh, A. H., Agarwal, A., Belgrave, D., and Cho, K. (eds.), Advances in Neural Information Processing Systems, 2022. URL httpss//ppeneevie...e//oonnnn//u TJUNiZiTKE
Zhang, W., Wang. X., Nie. W., Eaton, J., Rees. B., and Gu, Q. MoleculeGPT: Instruction following large language models for molecular property prediction. In NeurlPS 2023 Workshop on New Frontiers of Al for Drug Discov- ery and Development, 2023a
Xia, J., Zhu, Y., Du, Y., Liu, Y., and Li, S. A systematic survey of chemical pre-trained models. IJCAL, 2023.
Zhang, X., Wang. I.., Helwig, J., Luo, Y., Fu, C., Xie, Y., Liu, M., Lin, Y., Xu, Z., Yan, K., et al. Artificial intelli- gence for science in quantum, atomistic, and continuum systems. arXiv preprint arXiv:2307.08423 2023b.
Xie, T. Wan, Y., Huang, W. Yin, Z., Liu, Y., Wang, S., Linghu, Q. Kit, C., Grazian, C., Zhang, W., al. DAR- WIN series: Domain specific large language models for natural science. arXiv preprint 230iv:2308. 13565, 2023.
Zhao, H., Liu, S., Ma, C., Xu, H., Fu, J., Deng, Z.-H., Kong, L., and Liu, Q. Gimlet: A unified graph-tex model for instruction- molecule zero-shot learning. bioRxiv, pp. 2023-05, 2023a.
Xu, M., Luo, S., Bengio, Y., Peng, J., and Tang, J. Learning neural generative dynamics for molecular conformation generation. In International Conference on Learning Rep- resentations, 2021a. URL http://s////oeeeeiii.eeeeeeeee/e netfforumm id=pAbmlqtheGR
Zhao, W. X., Zhou, K., Li, Tang, T., Wang, X., Hou, Y.. Min, Y., Zhang, B., Zhang, J., Dong, Z., et al. A survey of large language models. arXiv preprint arXiv: 18223, 2023b.
Xu, M., Wang, W., Luo, S., Shi, C., Bengio, Y., Gomez- Bombarelli, R., and Tang, An end-to-end framework for molecular conformation generation via bilevel program- ming. In Meila, M. and Zhang, T. (eds.). Proceedings of the 38th International Conference on Machine Learn- ing, volume 139 of Proceedings of Machine Learning Research, pp. 15537-11547 PMLR, 18-24 Jul 2021b.
13

---

Geometry Informed Tokenization of Molecules for Language Model Generation
A. Broader Impacts and Limitations
Our work demonstrates the significant potential of pure language models (LMs) in revolutionizing molecular design and drug discovery by effectively transforming geometric information. The challenge of molecule design is particularly daunting when scientific experiments are cost-prohibitive or impractical. In many real-world scenarios, data collection is confined to specific chemical domains, yet the ability generate molecules for broader tasks where experimental validation is difficult remains crucial. Traditional diffusion- based models fall short in terms of efficiency, scalability, and the ability to learn from extensive databases or transfer knowledge across different tasks. In contrast, LMs exhibit inherent advantages in these areas. We envision the development of efficient, large-scale models trained on vast chemical databases that can function across multiple datasets and molecular tasks. By introducing LMs into the 3D molecule generation field, we unlock substantial potential for broad scientific impact.
Our research adheres strictly to ethical guidelines, with involvement of human subjects or potential privacy and fairness issues. This work aims to advance the field of Machine Learning and AI for drug discovery, with no immediate societal consequences requiring specific attention. We foresee no potential for malicious or unintended usage beyond known chemical applications. However, we recognize that all technological advancements carry inherent risks, and we advocate for ongoing evaluation of the broader implications of our methodology in various contexts.
We admit certain limitations, including that rounding up numerical values to certain decimal places bring information loss and discretized numbers impair generalization abilities across the continuous domain of real numbers. However, this trade-off betweeen advantages brought by our model-agnostic framework. Due the discrete nature of vocabularies, LMs depend extensive pre-training corpora, fine-grain gineigrined tokenization, or emergent abilities for better generalization, balancing high precision and versatility. Geo2Seq operates solely on the input data, which allows independence from model architecture and training techniques and provides reuse flexibility. This also means that we can effortlessly apply Geo2Seq on the latest generative language models, making seamless use of their capabilities. Future work points towards expanding on conditional tasks and exploring advanced tokenization techniques to enhance the model's performance and applicability.
B. Proofs
B.1. Proof of Lemma 3.2
First, we define the isomorphism problem for attributed graphs follows.
Definition B.1. [Graph Isomorphism] Let G₁ = E1) and G2 = (V2, Ez, A2) be two graphs, where Vi denotes the set of vertices, E, denotes the set of edges, and A; denotes the node attributes of G, for i 1,2, Let alir(v) denote the node attributes of vertex v. The graphs G1 and G2 are said to be isomorphic, denoted G₁ G2, there exists a bijection b: V1 1/2 such that for every vertex E V1, altr(v) A altr(b(v)) A2, and for every pair of vertices u, U Vi, (u, E E e ((u).b(v)) E Eg.
Next we prove Lemma 3.2
Lemma (Colored Canonical Labeling for Graph Isomorphism). Let G1 (V1, E1, A1) and =2=(va Eg, A2) be two finite, undirected graphs where V; denotes the set of vertices, E, denotes the set of edges, and A; denotes the node attributes of the graph fo i = 1,2 2. Let L: 9 - L be function that maps graph G E G. the set of all finite, undirected graphs, to its canonical labeling (()) E L, the set of all possible canonical labelings, as produced by the Nauty algorithm. Then the following equivalence holds:
L(G1) = L(G2) + G G2
where G _ G2 denotes that the graphs G1 and G2 are isomorphic.
The Nauty algorithm, tailored for CL and computing graph automorphism groups, presents rigorous mathematical underpin- nings to guarantee the CL properties. Here we leave out the proof of Nauty algorithm's rigor for canonical labeling, which is detailed in the work of McKKa & Piperno (2014). The key is the refinement process ensuring that the partitioning of the graph's 's vertices done in such way that any two isomorphic graphs will end with the same partition structure.
B.2. Proof of Lemma
Lemma. Let G = (z, R) be a 3D graph with node type vector x and node coordinate matrix R. Let F be the equivariant global frame of graph built based on the first three non-collinear nodes L(G). f(-) is our function that maps 3D coor-
14

---

Geometry Informed Tokenization of Molecules for Language Model Generation
dinate matrix R to spherical representations S under the equivariant global frame F. Then for any SE(3) transformation have f R) (g((R) Given spherical representations [R), there exist S SF(3) transformation 9, such that (S) = g R).
Proof. Let e l2, and &p be the indices of the first three non-collinear atoms G. Then the global frame (3x, y z)
a normalize(Pra - r(,) y normalize ((Fér r&)) x x1)
x x y
For (3) transformation g, let R' g R) QR b. Then the global frame F' (ar, (a',y', z²) is x' normalize(ra, rε,) = normalize(g(rr) - g((r,)) = Qx
y normalize ((rig Tt,) x x2) = normalize ((g(rez) ((r,,)) x x2) Qy
z' x' x y (Qx) (Qy) Q=
Thus ===0== QF. Here normalize(-) is the function to normalize vector to the corresponding unit vector. Then Vi, the spherical representations f e is
de; = ||r, - -rre 12
0:: arccos ((re, T(1) z z/de;)
Φe, atan2 ((re. -TE)) y. (re, rre,) xx)
Similarly, the spherical representations f(R')e, is
de I 121 = 119 (mr:) g (rr,) ||2 = di:
6% arccos ((r', - -r(,) xz x /d22,) = arccos ((g (Te,) - (re,)) - x2/ = tc,
6,, atan? ((re: - rez) yy, (rq: -re) x') atan2 ((g (rr.) - (rt,)) )) y (g (Fx,) - (rt,) x') de.
Therefore, we show that ((R) = [(((R)) Next, we consider the function - For all i, the three terms in (-) are de, sin(Ot,) c0s(
de, sin(Oe,) sin(dr.)
de, cos dr.
Then have ry: (S) F re. Therefore, show that there exist SF((3 transformation g. such that (S)) = R.
B.3. Proof of Theorem 3.5
First we establish lemma and provide its proof.
Lemma B.2. Let G1 = R1) and G2 = (Ez, R2) be two 3D graphs, where 2, is the node type vector and R, the node coordinate matrix of the molecule G. i Let L(G) be the canonical label of graph G. We have G1 _ G2- Let E; and E, denote the indexes of the node labeled i correspondingly in L(Gi) and L(G2), respectively. Let F be the equivariant global frame of graph G built based the first three non-collinear atoms in (G).. Let f : + S be surjective function that maps 3D graph G E G to its spherical representations $ = f(G) E S under the equivariant global frame F. Then the following equivalence holds:
Vi V1, f(r)e,,, = f(GG)e + G1 3sD G2
+
where G1 EdD G2 denotes that the graphs G1 and G2 are 3D isomorphic.
Proof. Let L(G) be the canonical labeling of graph G. Let & and E denote the index the node labeled correspondingly in LG1) and (G2), respectively. We have
G aa G2 + {
G M G2 and
G. there exists 3D transformation SE(3) such that r ===g(r(.)
Specifically, g(re,) Qrc, + b. Here Q a rotation matrix, and is a translation vector.
15

---

Geometry Informed Tokenization of Molecules for Language Model Generation
Let E,, f2, and Eo be the indices of the first three non-collinear atoms in G1- Then the equivariant global frame Fi
(11, y1, 21)
I normalize - Te,) -rG,
yi normalize ((reg T4 21)
11 X 1/1
Here normalize(-) is the function to normalize a vector to the corresponding unit vector. Then Vi, the spherical representations f(Gi),, is
de. l/r re, ll2
Or. arccos ((Ti, - -TE,)= =1/de,)
Φe. atan2 ((ri, - ) yı, (re. - - .11
Similarly, for G2, let G, 22 and Eo be the indices of the first three non-collinear atoms. Then the equivariant global frame F22 = (I2, y/2 x2) is 2²2 normalize(re, re) normalive(g(rr)--
Qan
y/2 = normalize ((re, re;) x x²2) = normalize ((g(rep) - g(re,)) ) x2) Qy
22 = 22 x y/2 = (Qæ1)= x (Qy1) = Q=1
Then Vi, the spherical representations f(G2)e is
dy = ||re Tę ||2 = 11g (Té,) g (re, 112 = de,
Og: arecos ((ré. - -r(,)) =2/dzz = arccos (lg (rr;) ) - (ra)) 22/de) Oc.
Φg. atan? ((Te - re;) y2 (re: -Te)) x2) atan2 ((g (re.) - 9 (rt)) y2, (re,) - (r(.)) -e.
Therefore, we show that G Snn G2 < vi, E Vi, (Gi)) f(G2)e; holds.
Then we prove Theorem
Theorem (Bijective mapping between 3D graph isomorphism and sequence). Let G Ri) and G2 = (zz. Ro) be two 3D graphs, where zj is the node type vector and R2 is the node coordinate matrix of the molecule G3 for j 1,2. Let (G) be the canonical label for 3D graph and its spherical representations. Given graph G with nodes and X = [Ix,...,, xnll' Raxm where m E Z. we define (m, & X (oncal(ze.......1 where E is the node index of the node labeled i in Lm (G), and concat(.) concatenates elements as a sequence. Define
S be the function that maps a 3D graph G to
Geo2Seq (G.) L. (G) (z.))
Lm (G) & X,
where x; [21, dg. 0z). Then GeolSeq 9 LA is a surjective function, and the following equivalence holds:
Geo2Seq G1) = Geo2Seq(Gz) = G₁ Snn G2
where G1 SDD G2 denotes that the graphs G1 and G2 are 3D isomorphic.
Proof. First, we prove that Geo2Seq 9 + 1/ is a surjective function. Given the definition Geo2Seq(G.) L. (G) @ (z, f(G)) L.. (G) X X,
where x; = =[2,, d.. f/, Φi), we need to prove that all operations are deterministic. X and z are defined to be deterministic, and g + S is a function. Lm,(Cr (G,) outputs the automorphism group of Gý's canonical label. By definition, the automorphism group contain different labels of the strictly identical graph. Let l and E describe two different sets of labels of the same automorphism group with nodes; since the graphs are identical,
[ze, de,, Ge., S2,) [=x; de; 0e, de bori = 1,..., n.
Thus concat(#; ( Lt,.) = concat (mtg ... Ce, ), i.e., different labels of one automorphism group produce identical sequences with Geo2Seq. Therefore, Geo2Seq + is a well-defined function; given a 3D molecule, we can uniquely construct a 1D sequence from Geo2Seq.
Next prove Geo2Seq's surjectivity. Given any output sequence & E li of Geo2Seq, the sequence is in the format concat((z1, d1, 01, oi,,..,, [Zn, dn, Onz dn.)).
For the nodes in q, we denote with S = [[d1. 01, Φ1)..., dd. fn dn]]. Given the surjectivity of the spherical representation function f : G -> and the defined f 1 S - G. there must be a unique G(z, R) E 9 where = f(G). Therefore,
16

---

Geometry Informed Tokenization of Molecules for Language Model Generation
output sequence q U there exists
R( s.t. q = Geo2Seq(G),
i.e., Geo2Seq is surjective; given sequence output of Geo2Seq, we can uniquely reconstruct a 3D molecule.
Now we prove the equivalence Geo2Seq(G,) = Geo2Seg(GG) = G1 SDD G2, starting from right to left. Considering
Lemma 3.2 for molecule G = (z, R), we specify G = F, A) with A = [z, R] and define the CL function for 3DD molecule graphs as L. which extends the equivalence in Lemma 3.2 to L on molecules with 3D isomorphism.
G 3n G2, i.e., graphs G and G2 are 3D isomorphic, then from Lemma 3.2 we know the canonical forms Lm (G1)
L G2). Let graphs G1 and G2 have numbers of node n. Let & and e be the denotations of a corresponding pair of canonical labelings from L, and L. (G2), respectively. Since graphs G1 and G2 are 3D isomorphic, from Def 3.4 we
know Vi E V(GI), =e, नęद and from Lemma B.2 we know vi V(G f(G2)e Thus, we have
Geo2Seq(Gi) Lm (G) @ (ZG,)
concat A, dd (f:)), (126,, de, Be, de.l)
(4)
(Ize, a. 0ef, 0c;1)
(G2) 2 (z2. (G2)) Geo2Seq(G2).
Note that if (G1) and L. (G2) contain automorphism groups larger than 1, include all possible labelings, which will all produce the same sequence later through Geo2Seq, as we have shown in detail above. However, this is a very rare case for real-world 3D graphs like molecules. Therefore, we have shown that if two molecules are 3D isomorphic considering atoms, bonds, and coordinates, their sequences resulting from Geo2Seq must be identical.
Finally, we prove the equivalence from left to right. We provide proof by contradiction. Given that Geo2Seq(Gi) Geo2See(G2)) (G2), we assume that the graphs G, and G2 are not 3D isomorphic. We denote with G2 (z R,) and G2 = (z2, R2). If G and G2 are not even isomorphic for A; z,, then from Der there does not exist a node-to-node mapping from G1 to G2, where each node is identically attributed and connected. And from Lemma 3.2 canonical forms Lm (G1) * Lm(G),, (G2). Thus for Geo2Seq(Gi) = concat, (,) (, (1=,, Oc. &t,l);
know the
and
Geo2Seq (22) = concat,, (e.d there must be at least one pair of ze, ZØ where zt, ze. Therefore, Geo2Seq(G1) o Geo2Seq(G2), which is contradiction to the initial condition that Geo2Seq (Gi) Geo2Sey((G) and ends the proof.
de: pr: Φr,1),
If G1 and G2 are isomorphic for A, = z., we continue with the following analyses. Let l, and e be the denotations of a corresponding pair of canonical labelings from (m1))) and (G2), respectively. Let f g -> S be the surjective function mapping a 3D graph to its spherical representations. Since G₁ and G2 are not 3D isomorphic, from Lemma B.2
we know there exists at least one
ie V1, s.t(i)) * f(G2)e
otherwise, would have contradicting the above condition. Thus for Geo2Seq(Gi) = co 3, d, 01) (1ze, dr, Oc. &r.)):
Vi E Vi, (GI)), f(G2)e (((2)) = G1 SDD G2,
and
Geo2Seq(G2) = concat, sx, Ef (1ze de; 0g Del):
Gi and G2 are isomorphic, so
=i= I, ze; = Zg
at least one pair of spherical coordinates does not correspond, so there must be at least one pair of (de, 0c,, Φe.) and (dz, 0g, Φx)) where
01, (e.) * (de, 0x).
Thus, Geo2Seq(G1) * (eo2Seq(Gz)). which contradicts the initial condition that Geo2Seq (eo) There- fore, we have shown that if two constructed sequences from Geo2Seq are identical, their corresponding molecules must be 3D isomorphic considering atoms, bonds, and coordinates. This ends the proof.
B.4. Proof of Corollary 3.6
Corollary (Constrained bijective Mapping between 3D graph and sequence). Let G1 (21, R1) and G2 = (x2, R2) be two 3D graphs, where Zj is the node type vector and R2 is the node coordinate mairix of the molecule Gj for
17

---

Geometry Informed Tokenization of Molecules for Language Model Generation
j 1,2. Let (G) be the canonical labeling for 3D graph and 9 - S be the function that maps 3D graph G to its spherical representations. Given graph G with nodes and X = [2n1,..., Fnn] E Ruxm, where E Z, we define (G) & X concat( 1....... Dt,)), where & is the node index of the node labeled i in Lum Lm(.)), and concat(.) concatenates elements as a sequence. Define
Geo2Seq(G., (G) @ (2, (c)) L. (
where x; [ di, O2 Ab). Let the truncation of spherical coordinate values be after b decimal digits. Then Geo2Seq : 9 - U is surjective function,, and the following equivalence holds:
Geo2Seg = Geo2Seq(G.) = G1 3D
3D-10 1200b G2
where G1 ~ an و G2 denotes that the graphs G1 and G2 are 110 2
constrained 3D isomorphic.
Proof. First, we prove that Geo2Seq: G- U is a surjective function, which resembles the proof for Theorem 3.5 Given the definition
(G) @ (z. f(G)) = (G) @ X; where x, [2s, d, 0, di), we need to prove that all operations
Geo2Seq(G:) = L..(G)
deterministic. @ and # are defined to be deterministic, and f + S with truncation after certain decimal places is still a well-defined function. L. (G;) outputs the automorphism group of G, 's canonical label. By definition, the automorphism group contain different labels of the strictly identical graph. Let & and f describe two different sets of labels of the same automorphism group with n nodes; since the graphs are identical,
[zt, de, 0e,, Φc,] = [ze, 02;- Ge. ffori = l....,
Thus concat( De,....,
concat(ag; : Fe, ). i.e., different labels of one automorphism group produce identical sequences with Geo2Seq. Therefore, Geo2Seq 9 - V is still well-defined function; given 3D molecule, we can uniquely construct a ID sequence from Geo2Seq.
Next we prove Geo2Seq's surjectivity. Given any output sequence q E li of Geo2Seq, the sequence is in the format q = concat([21 d1, f1, Φ1), [zn, dn. Bn Qnl))
For the nodes in g. define Strwn function f G - S and the defined f S + G. there must be unique G(z, R) E G where Strun f(G). Therefore, A output sequence 9 E U there exists
..., ....,,dd (dz, Ony 2n]] Given the surjectivity of the spherical representation [di 01, 01),...,
G(z, R) E g s.t. q Geo2Seq(G).
ie., Geo2Seq is surjective; given a sequence output of Geo2Seq, we can uniquely reconstruct a 3D) molecule.
Now we prove the equivalence Ge.)Seq(G.) = Geo2Seq(G.) = G San G2, starting from right to left. When a number is truncated after b decimal places, according to the rounding principle, the maximum error caused is S _ Considering Lemma 3.2 for molecule G = (z, R), specify G (V, E, A) with A = [z, R] and define the CL function for 3D molecule graphs as Lm, which extends the equivalence in Lemma 3.2 to Lm. on molecules with 3D isomorphism. If
[10-"|
G 3D 11- G2, i.e. graphs G2 and G2 are
[10-1 constrained 3D isomorphic, then from Lemma 3.2 we know G2 and 2
G2 are still isomorphic for A, = 2,, and the canonical forms Lm (G1) = Lm (G2). Let graphs G1 and G2 have numbers of node n. Let l, and E, be the denotations of a corresponding pair of canonical labelings from L,n (G1) and L,n (G2), respectively. Since graphs G1 and G2 are 2
1 constrained 3D isomorphic, from 3.f we know Vi E V(G)), 24, = zy;
error range allowed for each numerical value. and from Lemma B.2 we know Vi E V(G)), [GI),, = f(G2) with |10~1 Thus, we still have (G, = Lm (G1) @ (Gi))
2
concat, 20,
(lze, de, 0c, 0e,1)
(5)
(lzg, de; Og dc(])
(G2) 2 (z2, f (G2) ) = Geo)Seq
Note that if Lm (G1) and L. (G2) contain automorphism groups larger than 1, we can include all possible labelings, which have shown in detail above. However, this is very
will all produce the same sequence later through Geo2Seq, rare case for real-world 3D graphs like molecules. Therefore, we have shown that if two molecules are 3D isomorphic [10-"| considering atoms, bonds, and coordinates within the round-up error range 2 their sequences resulting from Geo2Seq must be identical.
Finally, we prove the equivalence from left to right. We provide proof by contradiction. Given that Geo2Seq(G1) =
18

---

Geometry Informed Tokenization of Molecules for Language Model Generation
Geo2Seg(G2), we assume that the graphs G1 and G2 are not |10 -conssccntted 3D isomorphic. We denote with G1 (21, R1) and G2 = (Z2, R2). G1 and G2 are not even isomorphic for A, = Zi, then from Def[B.I, there does not exist a node-to-node mapping from G1 to G2, where each node identically attributed and connected. And from Lemma 3.2 we know the canonical forms Lm (G1) ) # Lm (G2). Thus for
Geo2Seg( (1) = concat ( (Ize, de Be.: 06.1).
and
Geo2Seq(G2)) = c (i (lE de; reg 0r;1),
there must be least one pair of 26,, Zg, where zg, + ze. Therefore, Geo2Seq (G1) # Geo2Seq(G2), which contradiction the initial condition that Geo2Seq (G1) Geo2Sen(G2) and ends the proof.
If G1 and G2 are isomorphic for A, -i, continue with the following analyses. Let & and f be the denotations of corresponding pair of canonical labelings from L. (G,) and Lm (G2), respectively. Let f g - S be the surjective function mapping a 3D graph to its spherical representations. Since G| and G2 are not [10 cconstrained 3D isomorphic, from Lemma B.2 we know there exists at least one
i E V1, s.t. ((i)) + M(G2)e)
with error range
10
allowed: otherwise, would have Vi E (Gi)e, - f(G2)e; = Gi an G2,
contradicting the above condition. Thus for Geo2Seg(G1) (,n (1ze, de, B2. 0c,)).
and
Geo2Seq(Go) concatt,,ca.,,,00,,.e(((...... (lee de, 8r 0r,1),
G1 and G2 are isomorphic, so
Vi = 2/3
least one pair of spherical coordinates does not correspond, so there must at least one pair of (de,, Ob D2,) and (de; De, de;) where
110-1
min(Ide, 02 dbl lde; Oeg ohe
Thus Geo2Seq G1) * Geo2Seq(G2), which contradicts the initial condition that Geo2Seq(Gi) Geo2Seq((2)) There- fore, we have shown that if two constructed sequences from Geo2Seq are identical, their corresponding molecules must be 3D isomorphic considering atoms, bonds, and coordinates within the round-up error range |10-"| This ends the proof.
C. Ablation Studies
Table 3: Random generation performance with different atom generation orders.
Order
Atom Sta (%)
Mol Sta (%)
Valid (%)
Valid & Unique (%)
Canonical-locality
97.39
86.77
92.97
84.71
Canonical-nonlocality
96,45
81.36
90.89
83.37
Canonical-SMILLSS
97.35
85.86
92.97
84.05
DFS Thomas et al.
2009
95,95
81.54
90.45
82,48
BFS
Lee. 1961
96.85
80.92
90.49
76.13
Dijkstra Dijkstrr.
2022
95.29
77.25
88.97
73.52
Mcee (Cuthill& McKee 1969
93.56
71.57
85.36
76.23
Hilbert-curve (Hilbert & Hilbert 1935 Random
90.11
64.99
80.40
67.83
64.87
20.14
43.16
38.44
To study the effects of atom order, 3D representations and tokenization of Geo2Seq on the generation performance of LLMs, we conduct series of ablation experiments. Among all ablation experiments, we train 8-layer GPT models on QM9 dataset for 250 epochs with the same hyperparameters as Section 5.1 and use the random generation metries in Section 5.1 to
19

---

Geometry Informed Tokenization of Molecules for Language Model Generation
Table 4: Random generation performance with different 3D representations.
3D representation
Atom Sta (%)
Mol Sta (%)
Valid (%)
Valid & Unique (%)
Original coordinates
91.1
58.1
75.6
55.1
Normalized coordinates
92.7
63.2
83.1
72.5
Invariant Cartesian coordinates
96.0
78.5
89.7
74.1
Inv-spherical coordinates
97.3
83.4
91.0
82.7
Inv-spherical coordinates-local distances
97.1
82.8
91.7
79.6
Table 5: Random generation performance with different tokenization.
Tokenization
Atom Sta (%)
Mol Sta (%)
Valid (%)
Valid & Unique (%)
Char-tokenization
90.5
43.7
71.5
71.0
BPE
85.3
55.3
74.4
57.6
Sub-tokenization
96.4
80.3
89.9
74.4
Comp-tokenization
97.0
82.2
91.0
75.5
compare the performance under different settings.
Ablation on atom order. First, we show that our proposed canonical order of atoms in Geo2Seq sequence representation significant for LLMs to achieve good 3D molecular structure modeling. Specifically, we conduct an extended study of ordering algorithms, comparing our Geo2Seq with alternative canonicalization strategies well established traversing baselines. As specified in Sec 3.1 theoretically, analyses and derivations apply to all rigorous CL algorithms. In the paper, we select Nauty Algorithm because its implementation has the best time efficiency among all existing CI. algorithms. We implemented Nauty Algorithm for 3D molecules, where multiple strategies can be applied for the partitioning of graph vertices (a step in Nauty). We compare canonicalization strategies with/without locality considered. Canonicalization with locality considered can lead to better results, due to the importance of neighboring atom interactions in molecular evaluations. Given the similar nature, canonical SMILES produces very similar ordering with "Nauty with locality", thus close in performances. The traversing baselines includes Breadth-First Search (BFS) (Lee) 1961). Depth-First Search (DFS) (Thomas al.) 2009. Dijkstra's algorithm (Dijkstra 2022), Cuthill-McKee algorithm (Cuthill & McKee 1969), and Hilbert curve (Hilbert & Hilberi 1935). We also compare with Random sequence representation where atoms are randomly ordered. All the other settings of sequence representations remain the same. As Table 3 shows, canonicalization with locality considered can lead to better results, due to the importance of neighboring atom interactions in molecular evaluations. In addition, we can clearly observe that well-designed canonical ordering as in Gco2Seq significantly outperforms basic traverse strategies and the random order, which validates the significance canonical order.
Advantage of Nauty Algorithm. Note that in the paper, implement Nauty Algorithm for 3D molecules because: (1) its implementation has the best time efficiency among all existing CL algorithms; (2) it is naturally rigorous. The widely used canonical SMILES is based the Morgan CL Algorithm, which is proven to be incomplete for isomorphism corner cases (such as two triangles versus one hexagon). While canonical SMILES solve corner cases by manual restrictions, Nauty Algorithm is elegantly rigorous. Still, we emphasize that all rigorous CL. algorithms are usable for our method, while our contribution lies achieving structural completeness and geometric invariance for LM learning of 3D molecules.
Ablation on 3D representation. Besides, we explore using different methods to represent 3D molecular structures. We compare the spherical coordinates in Geo2Seq with directly using the 3D Cartesian coordinates of atoms from QM9 xyz data files in sequences. 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 5E((3 (e0(3)-invvriatt 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. We compare with "local distances". where our "distances to the global frame" are replaced with "relative distances to the previous atom" (except for the first atom) while the angles remain the same, Results in Table demonstrate that LLMSs achieve the best performance on spherical coordinates.
20

---

Geometry Informed Tokenization of Molecules for Language Model Generation
We believe this is due to that the numerical values of distances and angles of spherical coordinates lie in smaller region than coordinates, which reduces outliers and makes it easier for LLMSs to capture their correlation. Furthermore, both our spherical coordinates and that replaced with local distances achieve comparable results, while outperforming Cartesian coordinates. 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 addition, the similar performances could be attributed to that molecular systems often exhibit localized spatial structures (e.g., compact subunits or functional groups), which naturally constrain distances for most small molecules.
Ad Advattag of invariant spherical representations. The above experiments show the superiority of invariant spherical coordinates over invariant Cartesian coordinates. While invariant Cartesian coordinates when our proposed equivariant frame is applied can also 33(3)-invariance spherical coordinates are advantageous in discretized representations. Compared to Cartesian coordinates, spherical coordinate values are bounded in a smaller region, namely, a range of [o, T] or [0, 2r]. 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 be modeled by LMs. Lemma 3.3 and its proof aim to guarantee the validity that our proposed invariant spherical representations possess SE(3) -inaarrrriia We consider it as a part of our theoretical contribution towards the derivation of Theorem 3.5
Ablation on tokenization. Finally, we explore other ways to tokenize real numbers in spherical coordinates. Instead of simply taking the complete real number as a token (Comp-tokeninn comp-tooenization), we try splitting it by the decimal point and treat every part as an individual token Subbtooknizztion). We also explore the common NLP tokenization method, including treating each character as a token (Char-tokenizatie methods in Table|5 Results show that our used Comp-tokenization leads to better performance. This shows that treating the complete real number
and Byte-Pair Encoding (BPE). We compare these tokenization
individual token enables LLMs to capture 3D molecular structures more effectively.
Overall, through series of ablation experiments, we show that canonical atom order, spherical coordinate representation and Comp-tokenization in Geo2Seq are all very useful in parsing 3D molecules to good sequence representations.
D. Experimental Details and Additional Results
D.I. Hyperparameters and Experimental Details
In the random generation experiment (Section 5.1 we apply two LMs, GPT (Radford et al. 2018) and Mamba (Gu & Dao 2023). to our proposed Geo2Seq representations. For GPT models, adopt the architecture of GPT-1, set the hidden dimension to 768, the number of attention head to 8, and the number of layers to 12 and 14 for QM9 and GEOM-DRUGS datasets, respectively. For Mamba models, we set the hidden dimension to 768 and the number of layers to 26 and 28 for QM9 and GEOM-DRUGS datasets, respectively. On QM9 dataset, we set the batch size to 32, base learning rate to 0.0004, the number of training epochs to 600 and 210 for GPT and Mamba models, respectively. On GEOM- -RUGS dataset, we set the batch size to 32, base learning rate to 0.0004, the number of training epochs to 20 and 25 for GPT and Mamba models, respectively. During model training, we use Adam W (Loshchilov & Hutter 2019) optimizer and follow the commonly used linear warm up and cosine decay scheduler to adjust learning rates. Specifically, the learning rate first linearly increases from zero to the base learning rate 0,0004 when handling the first 10% of total training tokens, then gradually decreases to 0.00004 by the cosine decay scheduler. Besides, the tokenization of real numbers uses the precision of two and three decimal places for QM9 and GEOM-DRUGS datasets, respectively. In the controllable generation experiment (Section (.2). we train 16-layer Mamba models for 200 epochs, and all the other hyperparameters and settings are the same as the random generation experiment. Based on data statistics, we set the context length to 512 for QM9 dataset and 744 for GEOM-DRUGS dataset throughout the experiments. All experiments on the QM9 dataset are conducted using a single NVIDIA A6000 GPU. Experiments on the GEOM-DRUGS dataset are deployed on 4 NVIDIA A100 GPUs.
D.2. Licenses
We strictly follow all licenses when using the public assets in this work. The QM9 dataset is under license CC-BY 4.0. The GEOM-DRUGS dataset is under license CCO 1.0. The code of EDM, GEOLDM, JODO, and MiDi is under MIT License.
21

---

Geometry Informed Tokenization of Molecules for Language Model Generation
D.3. Experiments additional baselines and metries
We extend our experiments with two more baselines, JODO (Huang et al. 2023) and MiDi Vignac al. 2023). which are diffusion models jointly generating 2D and 3D molecular information. We exclude them in experiments of the main paper, since the setting is not the same as ours. Our method follows works on 3D molecule generation without 2D information, such as bonds.
We extend the metrics of our evaluation for more comprehensive comparisons on random generation of QM9 dataset. We report the percentage of valid, unique and novel molecules, i.e., that are not present in the training set. We also report the percentage of complete molecules in which all atoms are connected. Following JODO (Huang et al. 2023). we also include 2D metries. Frechet ChemNet Distance (FCD) measures the distance between the test set and the generated set with the activation of the penultimate layer of ChemNet. Lower FCD values indicate more similarity between the two distributions. Similarity to the nearest neighbor (SNN) calculates an average Tanimoto similarity between the fingerprints of generated molecule and its closest molecule in the test set. Fragment similarity (Frag) compares the distributions of BRICS fragments in the generated and test sets, and Scaffold similarity (Scaf) compares the frequencies of Bemis-Murcks scaffolds between them. Additionally, we include alignment metrics. For RDKit generated bonds, we compute the Maximum Mean Discrepancy (MMD) distances of the bond length (Bond), bond angle (Angle), and dihedral angle (Dihedral) distributions, and report their mean MMD distances. To ensure fair comparison, we evaluate the metrics of all methods on the generated 3D structures, and use RDKit to convert 3D structures 2D graphs if needed. We use the same model and settings the main paper for Geo2Seq, and follow the released codes for the baselines' respective hyperparameter and settings. Table 6 reports the random generation results on QM9 dataset. According to the results, though our model is not designed to directly learn 2D information, the performance of our method is better than or comparable with baseline methods on all metrics including the 2D metrics, which demonstrates the effectiveness of our design.
Table 6: Additional random generation results on QM9 dataset.
Metric
EDM
GEOLDM
JODO
MiDi
Geo2Seq with Mamba
Atom Sta (%)
98.7
98.9
98.9
98.2
98.9
Mol Sta (%)
82.0
89.4
89.0
83.5
93.2
Valid (%)
91.9
93.8
94.9
95.2
97.1
Valid & Unique (%)
90.7
91.8
92.8
92.8
81.7
Valid & Unique & Novel (%)
83.0
83.1
85.2
85.5
71.2
Complete (%)
90.9
93.3
94.4
94.4
97.3
Bond Length MMD
0.18
0.12
0.27
1.09
0.08
Bond Angle MMD
0.04
0.04
0.05
0.05
0.04
Dihedral Angle MMD
0.003
0.003
0.0022
0.0033
0.0011
FCD
1.16
0.94
1.55
1.28
2.04
SNN
0.47
0.49
0.47
0.47
0.49
Frag
0.94
0.94
0.94
0.94
0.83
Scaf
0.29
0.33
0.25
0.26
0.38
Reporting the percentage of novel molecules is important in showing that language models can generate new molecules instead of merely memorizing the training dataset. Given our improvements on controllable generation is significant, we explore whether the generated molecules are different from the molecules the training set. Thus we also extend the metric on controllable generation experiments. We use the same model and setting as the main paper. Table 7 presents the novelty results of controllable generation compared with EDM and JODO. Results show that our method achieves reasonably high novelty scores, which demonstrates that our method is not simply memorizing training data.
In addition, following Hoogebooo et al. 2022). we compare negative log-likelihood (NLL) performance on the random generation of QM9 dataset for Geo2Seq and baseline models that reports this metric. For this experiment, we use the same model and setting as the main paper. From Table g we can see the performance of our method is better than or comparable with all baseline methods, evidencing the validity of our model.
For more comprehensive comparisons, we also extend include the metrics of Symphony (Daigavane et
2023) in our
22

---

Geometry Informed Tokenization of Molecules for Language Model Generation
Table 7: Additional controllable generation results for the percentage of valid, unique, and novel molecules on QM9 dataset.
Method
Cx
CHоM
ÉLUMO
C₁
EDM
87.0%
84.1%
79.8%
84.7%
73.0%
68.0%
JODO
86.5%
87.3%
86.7%
86.2%
86.8%
85.6%
Geo2Seq with Mamba
82.8%
82.8%
83.6%
83.0%
83.3%
83.6%
Table 8: Additional Negative Log Likelihood (NLL) comparisons of random generation on QM9 dataset.
Method
NLL
E-NF
-59.7
GDM
-94.7
EDM
-110.7
GEOLDM
-335.0
Geo2Seq with Mamba
-242.0
evaluation. As shown Table 10 we compare the performances of baseline methods and Geo2Seq with Mamba on Symphony metrics. Multiple algorithms exist for bond order assignment: xyz2mo1 (Kim & Kim 2015). OpenBabel (Banck et al., 201T) and a simple lookup table based on empirical pairwise distances in organic compounds (Hoogeboom) al. 2022). We perform the comparison between these algorithms for evaluating machine-learning generated 3D structures. each of these algorithms infer the bonds and create molecule from generated 3D molecular structure. A molecule is valid if the algorithm could successfully assign bond order with no net resulting charge. We also measure the uniqueness to see how many repetitions were present in the set of SMILES strings of valid generated molecules. Buttenschoen et al. (2023) showed that the predicted 3D structures from machine-learned protein-ligand docking models tend to be highly unphysical. Table utjiliizss the PustBussers framework to perform the following samity checks to count how many of the predicted 3D structures are reasonable. The valid molecules from all models tend to be quite reasonable. Next, we evaluate models on how well they capture bonding patterns and the geometry of local environments found in the training molecules Table [1 We utilize the bispectrum (Uhrin) 2021) rotationally invariant descriptor of the geometry of local environments. Given a local environment with central atom u, all of the neighbors of projected according to the inferred bonds onto the unit sphere 52. Then, the signal f is computed as a sum of Dirac delta distributions along the direction of each neighbor. The bispectrum B(J) of f is then defined as B(J) EXTR AEEEAATSAAAAAE Thus, I captures the distribution of atoms around u, and the bispectrum B(J)) captures the geometry of this distribution. The bispectrum varies smoothly when f is varied and is guaranteed to be rotationally invariant. We follow Symphony and compute the bispectrum of local environments with atleast 2 neighboring atoms, and exclude the pseudoscalars the bispectra. For comparing discrete distributions, we use the symmetric Lensen-Shannol divergence (JSD) as Hoogeboom et al. (2022). Given the true distribution Q and the predicted distribution P. the SaasennShaanon divergence between r. (Q M) + DDz z. (P || M) where Dxr. is the Kullback- divergence and M Q+P is the mean distribution. For continuous distributions, estimating the Gensen-Shannen divergence from samples tricky without further assumptions on the distributions. We follow Symphony and use the MMD scores to compare samples from continuous distributions. Overall, the performance of our method is better than comparable with baseline methods across the metrics, showing the effectiveness of our 3D molecule generation.
In Table
& & f).
them is defined as: (ss(Q (Q P) 1 Drz
D.4. Generation Efficiency Analysis
We compare the generation efficiency of our method and the diffusion- based methods using single NVIDIA A100 GPU and batch size of 32. The results in Table 6 show that our method is much faster than diffusion-based methods, indicating the great efficiency of our method. Though we have take more memory compared to diffusion--aate methods, our time efficiency is much better than diffusion-basee to measure generation efficiency. In particular, Geo2Seq with Mamba is more than 100 times faster than diffusion- based methods, indicating the high throughput of our method, significant advantage in practical applications where speed
methods. Throughput, or samples per second, is one of the most important metrics
23

---

Geometry Informed Tokenization of Molecules for Language Model Generation
Table 9: Additional validity and uniqueness percentages of molecules following Symphony.
Metric 1
Symphony
EDM
G-SchNet
G-SphereNet
Geo2Seq
Validity via xyz2mo1
83.50
86.74
74.97
26.92
95.42
Validity via OpenBabel
74.69
77.75
61.83
9.86
83.84
Validity via Lookup Table
68.11
90.77
80.13
16.36
97.55
Uniqueness via xy22mo1 Uniqueness via OpenBabel Uniqueness via Lookup Table
97.98
99.16
96.73
21.69
98.88
99.61
99.95
98.71
7.51
99.91
97.68
98.64
93.20
23.29
98.95
Table 10: Percentage of valid molecules passing each PoseBusters test following Symphony.
Test f
Symphony
EDM
G-SchNet
G-SphereNet
Geo2Seq
All Atoms Connected
99.92
99.88
99.87
100.00
100.00
Reasonable Bond Angles
99.56
99.98
99.88
97.59
99.90
Reasonable Bond Lengths
98.72
100.00
99.93
72.99
100.00
Aromatic Ring Flatness
100.00
100.00
99.95
99.85
99.98
Double Bond Flatness
99.07
98.58
97.96
95.99
99.45
Reasonable Internal Energy
95.65
94.88
95.04
36.07
96.10
No Internal Steric Clash
98.16
99.79
99.57
98.07
99.33
Table 11: Additional comparison statistics of generated molecules the training set for QM9 dataset following Symphony.
MMD of Bond Lengths +
Symphony
EDM
G-SchNet
G-SphereNet
Geo2Seq
C-H: 1.0
0.0739
0.0653
0.3817
0.1334
0.0488
C-C: 1.0
0.3254
0.0956
0.2530
1.0503
0.0705
C-O: 1.0
0.2571
0.0757
0,5315
0.6082
0.0712
C-N: 1.0
0.3086
0.1755
0,2999
0.4279
0.1056
N-H: 1.0
0.1032
0.1137
0.5968
0.1660
0.0965
C-O: 2.0
0.3033
0.0668
0.2628
2.0812
0.0667
O-N: 1.5
0.3707
0.1736
0,5828
0.4949
0.1570
O-H: 1.0
0.2872
0.1545
0.7899
0.1307
0.0990
C-C: 1.5
0.4142
0.1749
0.2051
0.8574
0.0832
C-N: 2.0
0.5938
0.3237
0.4194
2.1197
0.2676
MMD of Bispectra +
Symphony
EDM
G-SchNet
G-SphereNet
Geo2Seg
C: C2.H2
0.2165
0.1003
0.4333
0.6210
0.0955
C: C1HH3
0.2668
0.0025
0.0640
1.2004
0.0011
C: C3.H1
0.1111
0.2254
0.2045
1.1209
0.0867
C: C2,
0.1500
0.2059
0. 1732
0.8361
0.1058
C: C1,H2,01
0.3300
0.1082
0.0954
1.6772
0.0802
O: CI.HI
0.0282
0.0056
0.0487
0,0030
0.0022
C: C2.HINNI
0.1481
0.1521
0.1967
1.3461
0.1111
C: C2.H1
0.2525
0.0468
0.1788
0.2403
0.0851
C: C1.22I11
0.3631
0.2728
0.1610
0.9171
0.1285
N: C2,H1
0.0953
0.2339
0.2105
0.6141
0.1081
Jensen-Shannon Divergence J.
Symphony
EDM
G-SchNet
G-SphereNet
Geo2Seq
Atom Type Counts
0.0003
0.0002
0.0011
0.0026
0.0002
Local Environment Counts
0.0039
0.0057
0.0150
0.1016
0.0035
24

---

Geometry Informed Tokenization of Molecules for Language Model Generation
erucial.
Table 12: Generation efficiency comparison between diffusiin-basee methods and our LM-based method.
QM9
DRUG
Method
Parameters
Memory
Sample/second
Parameters
Memory
Sample/second
EDM
5.3M
1.5GB
1.4
2,4M
7.4GB
0.1
GeoLDM
11.4M
1.5GB
1.4
5.5M
8.4GB
0.1
Geo2Seq with GPT
87.7M
2.4GB
8.3
105.4M
3.1GB
0.2
Geo2Seq with Mamba
91.8M
2.2GB
100.0
108.4M
2.6GB
16.7
D.5. Results with Pretraining
To show the advantage of pretraining, compare the random generation performance on QM9 for models with and without pretraining Molecule3D (Xu et al., 2021c) dataset, which includes around 4M molecules. Specifically, we conduct experiments on an 8-layer GPT model and a 20-layer Mamba model. The models are pretrained for 20 epochs and then finetuned for 200 epochs. The results in Table 13 demonstrate the advantage of pretraining. Future studies could explore pretraining on larger datasets.
Table 13: Random generation performance on QM9 for models with and without pretraining on Molecule3 dataset.
Method
Atom Sta (%)
Mol Sta (%)
Valid (%)
Valid & Unique (%)
Geo2Seq with GPT
97.0
82.2
91.0
75.5
Geo2Seq with GPT pretraining
98.5
89.7
94.8
76.6
Geo2Seq with Mamba
97.4
86.8
93.0
78.8
Geo2Seq with Mamba pretraining
98.3
89.4
94.9
83.5
E. Extended Studies
E.1. Scaling Laws
Scaling law refers the relations between functional properties of interest, performance metrics in case, and properties of the architecture optimization process. In this section explore the scaling laws of our models, specifically regarding parameter size, since they provide typical insights for LMs. Scaling laws in 3D molecule generation appears similar to that in NLP. We provide experiments both GPT and Mamba in Table 14 and [5] respectively. As can be observed, LMs' performances on molecules grow significantly with parameter size increase, similar to the emergence abilities widely- recognized in NLP tasks. As known from NLP studies (Schaeifer et al., 2024), model capabilities grow consistently with model size, while emergence abilities are largely caused by nonlinear metrics. This matches our observations, since the chemical metrics are hardly linear.
Table 14: Scaling laws on Geo2Seq with GPT model.
Parameter size GPT
2556532
31309824
61650944
88012800
116342688
Atom sta ())
76.2
89.6
96.5
98.3
98.5
Mol sta (%)
5.1
42.4
81.3
89.1
90.6
Va)
45.5
73.1
90.9
94.3
95.1
Valid & Unique(%)
43.4
66.7
83.6
74.9
78.6
Note that we evaluate all models after 250 epochs for fairness concerns, while this fixed hyperparameter setting is not optimal for performances all parameter sizes. Other settings are the same as the ablation studies.
25

---

Geometry Informed Tokenization of Molecules for Language Model Generation
Table 15: Scaling laws Geo2Seq with Mamba model.
Parameter size Mamba
2180352
31458048
61631232
93088512
121977600
Atom sta(%)
81.6
95.7
97.4
97.8
97.9
Mol sta(%)
13.6
79.2
86.8
88.3
89.0
(%)
51.2
89.4
93
93.7
94.4
Valid & Unique(%)
49.6
78.7
78.8
82.6
83.5
E.2. Error Case Analysis
In the natural language domain, trained language models can produce error cases showing repetition or hallucinations. This also a problem that often arises with LLMs. In this section, we provide the analysis of some error cases to introduce more insights into the field.
Similarly NLP cases, our trained language models are showing repetition or hallucinations, especially when not trained to best convergence. This happens to both GPT and Mamba models. Below show some error cases from a 16-layer Mamba model trained 150 epochs on the QM9 dataset. The eπror case below shows typical repetition problem. The model generates repeated tokens for several periods, resulting in an invalid sample.
H 0.00 0 0.00 C 1.09 1.57 0.00 N 2.02 2. 0.00 3.39 1.99 -0.02 3.98 2.10 0.23 C.334 2.11 -0.35 H 4.09 4.09 4.09 4.09
4.43 2.38 -0.46 H 5.41 2.09 -0.29 H 4.29 1.96 -0.59 C 4.005 1.63 0.04 H 4.09 4.09 4.09 4.09
4.09 4.09 4.09 4.09 4.09 4.09 4.09 4.09 4.09 4.09 4.09 4.09 4.09 1.01 H 4.96 4.96 4.96 1.23 1.23 1.23 C.2.7
1.749 -0.03% H 4.02 1,83 -0.23 3.79 1.91 0.17 3.79 1.53 -0.03*
For hallucination, tokenization design actually prevents token-level hallucination by defining elements and whole- numerical-valuus as tokens, instead of using single characters. This prevents token-level hallucination, i.e., non-existent elements or numbers such as 'Hr' or 15'. However, there can still be sequence-level hallucinations, such as the error case below. The model generates distance values in the place the should be angle values (and vice versa).
H 0.00 0.00 0.00 N 1.01 1.57 0.00 H 1.70 2.14 0.00 2.06 1.13 -0.48 3.1.3 1.34 -0.42 N 2.49 0.62 -0.87% C 2.94 0.10 -1.64% H 3.20 0.39 3.91 H 2.81 0.33 -3.14 C..433 0.10 -1777 H5.15 0.22 0.22 H 2.78 0.21 -0.95
C 1.86 1.86 4.84 H 0.19 1.89 -2.06 8.28 1.94 -1.69 H 6.24 1,71 5.70 C 3.97 0.67 -0.90 H 5.02 0.68 -0.77"
H 4.155 0.91 -1.11 2.93 0.50 6.97 H 3.54 0.42 0.42 2.74 0.88 0.88*
These error cases will be rarer if the model well converges. When trained for 150 epochs, the model would generate 15% of invalid samples, including the above discussed syntax problems. When trained for 250 epochs, the model would generate 2% invalid samples.
F. Visualization Results
F.1. Visualization of Generated Molecules
In this section, provide visualizations of molecules generated from Geo2Seq with Mamba conditionally the property of Polarizability a in Figure 3 The Polarizability of molecule is the tendency to acquire an electric dipole moment when the molecule is subject to an external electric field. Large values usually correspond to less isometrically molecular geometries. This is consistent with our generated examples.
In addition, we provide visualizations molccules generated from Geo2Seq with Mamba trained on QM9 and DRUG in Figure and Figure respectively. These examples are randomly generated without any cherry pick. From the figures, we that the model generate realistic molecular geometries for both small and large size molecules. However, similar to previous methods (Hoogeboom al. 2022 Xu et al.) 2023a), there are disconnected components, especially for larger molecules. A possible future direction is to apply fragment based methods to reduce the sequence length, thus benefiting the training of language models.
26

---

Geometry Informed Tokenization of Molecules for Language Model Generation
67.72
70.82
74.94
76.62
78.83
83.57
90.12
91.78
Figure 3: Visualization of generated molecules condition the property of Polarizability a.
Figure 4: Visualization of molecules generated from Geo2Seq with Mamba trained on QM9.
27

---

Geometry Informed Tokenization of Molecules for Language Model Generation
Figure 5: Visualization of molecules generated from Geo2Seq with Mamba trained on GEOM-DRUGS.
28

---

In this section, provide UMAP visualizations of different (atom type, distance, and angle) token embeddings learned by Mamba models trained on QM9 and GEOM-DRUGS datasets. Patterns of the embeddings indicate that the model has successfully learned structure information from the sequence data, showcasing LMs' capabilities to understanding molecules precisely in 3D space. For example, Figure 8 shows that similar angle tokens (e.g., '1.41 and '1.42%") are placed next to each other and the overall structure of all angles is loop. Further, m-out-of-phase angles are placed near each other, such as "3.14", '-3,14", and '0"". For atom type tokens, the model appears to capture the structure of the periodic table, although the rows and columns are not perfect in Figure One reason the limited atom types in the datasets (5 in QM9 and 16 Geometry Informed Tokenization of Molecules for Language Model Generation
F.2. Visualization of Learned Token Embeddings
GEOM-DRUG) limiting the model's capabilities to learn chemical patterns from the entire periodic table, We provide analyses of the visualization results in the caption of each figure as Figure [6] Figure T0
13
80
5
17.75
70
17.50
14
Bi
6
3
60
17.25
4
15
17.00
S
B
16
40
0
16.75
30
16.50
20
2
1
16.25
1
H
17
10
0.5
1.0
1.5
2.0
2.5
3.0
3.5
4.0
Figure 6: UMAP visualization of element token embeddings learned by Mamba model trained on GEOM-DRUGS. Red groups indicate columns in the periodic table and blue groups indicate rows, which are both numbered. Points are colored by atomic weight. Overall, the model appears capture the structure of the periodic table. The column generally increases from top to bottom, and the row generally increases from left to right.
29

---

Gcometry Informed Tokenization of Molecules for Language Model Generation
9
-9.8
-10.0
-10.2
6
-10.4
5
-10.6
-10.8
N
3
-11.0
-5.75 -5.50 -5.25
-5.00
-475
-4.50
-4.25 -4.00
Figure 7: UMAP visualization of element token embeddings learned by Mamba model trained on QM9. Points are colored by atomic weight. Overall, the model appears to distinguish well between different elements. All different elements are distributed distantly from each other in the embedding space.
30

---

Geometry Informed Tokenization of Molecules for Language Model Generation
14
10
8
6
-1
4
-2
-4
-2
2
4
6
8
10
6
-1
10
-2
2
6
6
4
10
2
-3
Figure 8: 2D and 3D UMAP visualization of angle token embeddings learned by Mamba model trained on GEOM-DRUGS. It can be observed that similar tokens (e.g.. '1.41 and "1.429") are placed next to each other and the overall structure is loop. Further, o-out-of-phase angles are placed near each other, such as "3.14"* '-3.14"' and '0
31

---

Geometry Informed Tokenization of Molecules for Language Model Generation
15
10
10
0
-5
-5
0
5
10
15
10
15
5
15
10
2
5
5
15
Figure 9: 2D and 3D UMAP visualization of distance token embeddings learned by a Mamba model trained on QM9, Representations of distances lower than 6 form relatively distinct patterns. This is likely because these values are much more frequently seen in the training data, Values over 20 cluster into a clump, suggesting that they are also recognized by the model.
32

---

Geometry Informed Tokenization of Molecules for Language Model Generation
30
15
25
10
20
15
10
-5
-10
-5
0
5
10
15
30
15
25
10
5
20
0
-5
15
-10
10
0
10
15
-5
Figure 10: 2D and 3D UMAP visualization of distance token embeddings learned by Mamba model trained on GEOM- DRUGS. It is notable that the best and most distinct representations seem arise from between 5 and 20. This is likely because these values are much more frequently seen in the training data. Values over 20 form an indistinct clump. Interestingly, values > 20 are near values < 3, which is initially unintuitive; however, they are likely placed in a similar location in the embedding space since both small and large distances are rarely seen the data.
33


Reviews:
[{"review_id": "IpmkMlLcqJ", "reviewer": "Reviewer_39Ws", "summary": "The paper proposes Geo2Seq, which transforms molecular geometries into SE(3)-invariant discrete sequences for molecule generation. Existing language model-based molecule generation works do not consider the 3D molecular geometries in the tokenization process. The proposed paper address this limitations and show that the proposed Geo2Seq improves the performance in molecule gnerations.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer 39Ws,\n\nThank 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](https://anonymous.4open.science/r/geo2seq-rebuttal/Geo2Seq_rebuttal.pdf) and our responses to other reviewers.\n\n> Geo2Seq with larger or recent language models such as LLaMa\n\n- 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.\n\n| Property (Units)       | α (Bohr³) | Δε (meV) | ε_HOMO (meV) | ε_LUMO (meV) | μ (D)  | C_v (cal/mol·K) |\n|--|-|-|--|--|--|--|\n|Data|0.10| 64| 39 | 36| 0.043| 0.040|\n|Random |9.01| 1470| 645| 1457|1.616| 6.857|\n|GEOLDM|2.37| 587| 340| 522 | 1.108  | 1.025           |\n|Geo2Seq with Mamba| **0.46**| **98**| 57 | 71| 0.164 | **0.275**|\n|Geo2Seq with GPT| 0.53|102|**48**|**53**|**0.097**|0.325|\n|Geo2Seq with LLaMA| 0.71|102|51|98|0.324|0.357|\n\n- 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.\n\nWe 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!\n\nSincerely,\n\nAuthors", "claims_and_evidence": "The paper claims that the tokenization with preserving 3d molecular graph information improves the quality of molecule generation. The claim is well supported by theoretical analysis and experimental results.", "methods_and_evaluation_criteria": "The methods and evaluation criteria make sense for the problem from my perspective.", "theoretical_claims": "I've checked the correctness of the proofs in the paper and it seems to be correct from my side.", "experimental_designs_or_analyses": "The experimental designs and analyses are valid.", "relation_to_broader_scientific_literature": "The paper seems novel to me. Even though a recent paper proposes tokenization techniques considering 3D molecular  geometric information, it limits its scope to text generation. Different from it, Geo2Seq is designed to generate molecules with theoretical analysis. So, I think it is novel.", "essential_references_not_discussed": "I think that the essential references are discussed in this paper.", "other_strengths_and_weaknesses": "Strengths\n- The paper is well written and easy to follow.\n- The paper well demonstrates the effectiveness of the proposed method Geo2Seq with theoretical analysis and experimental results.\n- The visualization map in Figure 6 is interesting to me.\n\nWeaknesses\n- I cannot find crucial weaknesses in this paper.", "other_comments_or_suggestions": "- I'm wondering the performance of the proposed Geo2Seq with larger or recent language models such as LLaMa.", "questions_for_authors": "Please refer the above sections.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "WZXQg29Ejm", "reviewer": "Reviewer_hgAt", "summary": "The paper proposes a method called Geo2Seq to generate 3D molecules using language models. The authors convert each molecule into an SE(3)-invariant discrete sequence of tokens—one token per atom, with tokens containing both atom type and spherical-coordinate information. Once converted to a sequence, any language model can be trained to produce new molecular sequences, which are then mapped back to 3D structures.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer hgAt,\n\nThank 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](https://anonymous.4open.science/r/geo2seq-rebuttal/Geo2Seq_rebuttal.pdf) and our responses to other reviewers.\n\n> Claims And Evidence\n\n- 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.\n- 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.\n- 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. \n- 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.\n\n\n> Experimental Designs Or Analyses\n\nThanks 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.\n- `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**.\n- 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.\n- 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.\n- 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.\n\nWe 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!\n\nSincerely,\n\nAuthors", "claims_and_evidence": "Much of the theoretical proof (canonical labeling correctness, SE(3)-invariant spherical coordinates) is already known from prior literature on graph isomorphism and spherical transformations. The paper does not fully demonstrate how they guarantee stronger empirical generation beyond giving a valid, lossless tokenization. The paper would benefit from more ablation: e.g., do we see the same advantage if we do a simpler coordinate encoding?", "methods_and_evaluation_criteria": "The tokenization for unconditional and property-conditional generation does make sense: the method inserts property tokens into the same sequence, so an LM can learn to produce geometries with certain property values.", "theoretical_claims": "No obvious errors stand out.", "experimental_designs_or_analyses": "I believe it remains unclear whether the LM is capturing fundamental 3D chemistry knowledge or if it is mostly reproducing token patterns it has memorized. More analyses, like measuring internal geometry consistency or testing on molecules that deviate strongly from training data, would strengthen the claims.", "relation_to_broader_scientific_literature": "The paper was built based on prior works on 3D generation and proposed a geometry-aware tokenization to do generation tasks.", "essential_references_not_discussed": "NA", "other_strengths_and_weaknesses": "Strengths:\n1. The proposed method is model-agnostic and requires no special modifications to model architectures.\n2. The method preserves SE(3)-invariance.\n2. The paper is well-written and easy to follow.\n3. The paper conducts comprehensive experiments.", "other_comments_or_suggestions": "NA", "questions_for_authors": "Please refer to other sections.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "A9HETdel4e", "reviewer": "Reviewer_aNgc", "summary": "This paper explores the use of language models (LMs) for generating 3D molecules, a task that has previously been challenging due to the complex geometric structure of molecules. The paper proposes a novel tokenization method called Geo2Seq, which converts 3D molecular structures into SE(3)-invariant 1D discrete sequences that can be effectively processed by LMs. Results show that with various state-of-the-art 3D molecule generation methods, including diffusion-based models like EDM and GEOLDM. Geo2Seq achieves comparable or better results in terms of atom stability, molecule stability, and valid percentage, especially in controlled generation tasks.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer aNgc,\n\nThank you for your appreciation of our work and insightful comments! We address each point here and added experiments are in [this Link](https://anonymous.4open.science/r/geo2seq-rebuttal/Geo2Seq_rebuttal.pdf).\n>  Real number generalizability\n\nWe have conducted more studies on this point.\n- 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`.\n- We also analyze different discretization schemes. See experiments in `Appx C Table 5` and `[Link] Table 2`.\n- We can extend to additional datasets in future work.\n\n> Uniqueness of QM9\n- 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**.\n- 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%).\n\n> Error Case\n- We now include error case studies. Please see `[Link] Table 7`.\n- 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. \n\n> Comparison with graph-based methods\n- Previously we focus on comparing with SoTA methods. Here we compare with more graph-based methods, as in `[Link] Table 3`.\n- 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.\n\n> LMs like llama/qwen\n- We provide experiments extending Geo2Seq with LLaMA. Please see `[Link] Table 4` and `our Response to Reviewer 39Ws`.\n\n> More Metrics\n\nThanks for the advice. We already include these metrics in Appx. D.\n- `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.  \n\n> Def 3.4 Validity\n\nWe clarify its soundness regarding theoretical alignment with graph theory and physical relevance in molecular modeling.\n- 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.\n- 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.\n\n> Proof Completeness\n- 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.\n\n> Pretraining\n- We clarify that we do explore the impact of pretraining. See `Appx D.5 and Table 13`.\n\n> Limited Controlled Generation\n- We focus on quantum properties because they are available conditions relating 3D information. Here we provide more experiments exploring generalizability across conditions.\n- 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.\n\n> Limited Limitations\n- 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.\n\n> Figure Captions\n- We have included more informative captions. For Figure 2, see `[Link] Figure` for updated caption.\n\nThank 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.\n\nSincerely,\n\nAuthors", "claims_and_evidence": "1. Generalizability to Continuous Domains:\nThe paper acknowledges the limitations of the discrete tokenization approach in terms of generalization to the continuous domain of real numbers. However, the analysis provided is limited, and more evidence in various datasets is needed to demonstrate the extent of this limitation and potential solutions.\n\n2. Uniqueness of Generated Molecules:\n\nThe paper reports a lower uniqueness percentage for Geo2Seq compared to diffusion-based methods on the QM9 dataset. This could be due to several factors, including the tokenization approach and the size of the dataset. Further investigation is needed to understand the underlying reasons and explore ways to improve uniqueness.\n\n3. Error Case Analysis:\n\nThe provided error case analysis is limited and focuses on specific examples. A more comprehensive analysis of different types of errors (e.g., syntax errors, repetition, hallucinations) and their causes is needed to better understand the robustness and limitations of the approach.", "methods_and_evaluation_criteria": "1. Geo2Seq Tokenization: The use of canonical labeling and invariant spherical representations is a reasonable approach for converting 3D molecular structures into a format suitable for LMs. However, the paper lacks a comprehensive comparison with alternative tokenization methods, such as graph-based representations directly used with graph neural networks. A more thorough comparison would provide a better understanding of the advantages and limitations of Geo2Seq.\n2. Language Models: The choice of GPT and Mamba as LMs is reasonable, given their strong sequence modeling capabilities. However, the paper does not explore the potential of larger LMs, like llama and qwen.\n3. Evaluation Metrics: The paper focuses on atom stability, molecule stability, and valid percentage as primary evaluation metrics. While these metrics are important, they do not fully capture the quality and diversity of the generated molecules. Consider incorporating additional metrics, such as novelty, structural diversity, and property prediction accuracy, to provide a more comprehensive evaluation.", "theoretical_claims": "1. Validity of the 3D Graph Isomorphism Definition: The paper extends the concept of graph isomorphism to 3D graphs, which is not a standard definition. While the paper provides a definition, it is crucial to establish the validity and soundness of this definition in the context of 3D molecular structures. The proof should clearly justify the extension and demonstrate its consistency with existing graph theory concepts.\n\n2. Completeness of the Proof: The proof in the Appendix seems to focus on demonstrating the sufficiency of the Geo2Seq mapping (i.e., if two molecules are 3D isomorphic, their sequences are identical). However, it does not explicitly address the necessity (i.e., if two molecules have the same sequence, they must be 3D isomorphic). This is a crucial aspect of a bijective mapping, and a more comprehensive proof is needed to establish both sufficiency and necessity.", "experimental_designs_or_analyses": "1. Comparison with Alternative Tokenization Methods: The paper primarily focuses on comparing Geo2Seq with state-of-the-art 3D point cloud based methods. However, it lacks a comprehensive comparison with alternative tokenization methods, such as graph-based representations directly used with graph neural networks. This limits the understanding of the advantages and limitations of Geo2Seq compared to other approaches.\n2. Additional Metrics: The evaluation focuses on atom stability, molecule stability, and valid percentage. While these metrics are important, they do not fully capture the quality and diversity of the generated molecules. Incorporating additional metrics, such as novelty, structural diversity, and property prediction accuracy, would provide a more comprehensive evaluation of the generated molecules.\n3. Comparison with Pre-trained Models: The paper compares the performance of Geo2Seq with models that are not pre-trained on chemical datasets. Exploring the impact of pre-training on the performance of Geo2Seq would be valuable for understanding its effectiveness in leveraging knowledge from large chemical databases.", "relation_to_broader_scientific_literature": "Yes. By applying the proposed methods to generate a more accurate 3D coordinates of molecules, researchers could better find appropriate drug candidates.", "essential_references_not_discussed": "As far as I know, there is no essential reference not discussed.", "other_strengths_and_weaknesses": "#### Strengths:\n1. Originality: The paper presents a novel approach to 3D molecule generation using language models, which is a relatively unexplored area. The use of Geo2Seq for converting 3D molecular structures into 1D sequences is a creative combination of ideas from graph theory and language modeling.\n2. Potential Impact: The approach has the potential to significantly impact the field of 3D molecule generation, particularly in drug discovery and materials science. The ability to efficiently generate valid and diverse 3D molecules with desired properties could revolutionize these fields and enable the discovery of new drugs and materials.\n3. Clarity: The paper is generally well-written and provides a clear explanation of the approach and its benefits. The figures and tables are helpful in illustrating the key concepts and results.\n#### Weaknesses:\n1. Limited Comparison with Alternative Methods: The paper primarily focuses on comparing Geo2Seq with state-of-the-art 3D point cloud based methods. A more comprehensive comparison with alternative tokenization methods and graph-based approaches would provide a better understanding of the advantages and limitations of Geo2Seq.\n2. Limited Evaluation of Controlled Generation: The evaluation of controlled generation tasks is limited to specific quantum property values. Exploring the generalizability and robustness of controlled generation across different conditions would be valuable.\n3. Limited Discussion of Limitations: The paper acknowledges the limitations of the discrete tokenization approach but does not provide a detailed analysis of these limitations and potential solutions. A more thorough discussion of the limitations and their impact on the performance would be beneficial.", "other_comments_or_suggestions": "Figure Captions: Some figure captions could be more informative and provide a clearer explanation of the content. For example, Figure 2 would benefit from a more detailed description of the equivariant frame and invariant spherical representations.", "questions_for_authors": "Please see the above sections.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "9Zkt87nBGq", "reviewer": "Reviewer_CPUN", "summary": "This paper proposes Geo2Seq, a 3D molecule tokenization method for 3D molecular generation. The authors convert molecules (in 3D space) to 1D sequences while preserving SE(3) invariance, and then train a molecule generative model based on language model architecture. Geo2Seq equipped with various language models show superior performance in molecule generation tasks.\n\n## update after rebuttal\nI confirm that I have read the rebuttal and finalized my evaluation.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer CPUN,\n\nThank 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](https://anonymous.4open.science/r/geo2seq-rebuttal/Geo2Seq_rebuttal.pdf) and our responses to other reviewers.\n\n>  Some existing molecule 3D molecule tokenizing methods\n\n- Thank you for the advice. We have included the discussions in the paper revision and also briefly discuss here.\n- `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\nmolecule 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.\n- `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.\nCompared to our method, this work uses VQ-VAE to learn structure tokens, which lacks guarantee of structural completeness.\n- 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.\n\n> W1: Scalability of Geo2Seq\n\n- 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.\n- 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.\n\n> Paper Representation\n\n- 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.\n\n> Typos\n\n- Thank you for pointing out! We have revised the paper and corrected the typos.\n\n> Q1: train from scratch or pre-trained?\n\n- 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.\n\n> Q2: text-to-molecule generation tasks?\n\n- 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.\n\nWe 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!\n\nSincerely,\n\nAuthors", "claims_and_evidence": "The claims made in the submission are supported by clear and convincing evidence.", "methods_and_evaluation_criteria": "Proposed methods and/or evaluation criteria (e.g., benchmark datasets) make sense for the problem or application at hand.", "theoretical_claims": "The theoretical claims seem correct for me (at least at a high level).", "experimental_designs_or_analyses": "I checked the soundness/validity of the experimental designs or analyses.", "relation_to_broader_scientific_literature": "The key contribution of this paper is to relate 3D molecule generation and language models. Previous 3D molecule generation relies on 3D graph generation technique, however, such graph generation tasks show poor performance due to the under-explored 3D graph architectures. This paper shows such drawbacks can be alleviated by tokenizing 3D molecules to 1D sequences, which can be directly incorporated with well-developed language models.", "essential_references_not_discussed": "As far as I know, there exist some molecule 3D molecule tokenizing methods, e.g., 3D-MolT5: Towards Unified 3D Molecule-Text Modeling with 3D Molecular Tokenization [Pei et al., 2024] and Tokenizing 3D Molecule Structure with Quantized Spherical Coordinates [Gao et al., 2024], which can be discussed in this paper.", "other_strengths_and_weaknesses": "Strengths\n1. The problem of interest, 3D molecule tokenization, is important for molecular domain and poses a potential for real-world applications, e.g., drug discovery.\n2. Proposed method seems reasonable; tokenizing 3D spherical coordinates in an SE(3)-invariant manner.\n3. The improvements in controllable generation are impressive; Geo2Seq highly outperforms previous 3D graph-based molecule generation techniques.\n\nWeaknesses\n1. Scalability of Geo2Seq. Compared to graph-based models (which accepts continuous values), discretization technique in Geo2Seq may limit the performance when the molecules become large. The results in Table 1 also show that Geo2Seq works well on smaller molecules (QM9) but not quite well on larger molecules (GEOM).", "other_comments_or_suggestions": "Comments\n1. Representation can be further improved. Essential experimental ablations are deferred to the supplements, e.g., Table 4. I think the method section can be shortened, e.g., discussion about spherical coordinates.\n\nTypos\n\nL46, right column: \"subsequent LMs used. and can seamlessly\" -> \"subsequent LMs used, and can seamlessly\"\\\nL194, left column: \"can be be proved\" -> \"can be proved\"", "questions_for_authors": "1. Did the authors train GPT and Mamba from scratch or pre-trained checkpoints? \n2. Can Geo2Seq be applied to text-to-molecule generation tasks, e.g., ChEBI-20 dataset? Such results will further highlight the effectiveness of using language models.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: The paper introduces a new method, Geo2Seq, for molecule generation and appropriately provides theoretical analysis and comprehensive experimental evaluation. It is presented clearly and is intuitive. The addressed problem is very important in structural biology and drug design. The method is original and has the potential to be impactful. All reviewers recommend acceptance and the AC agrees. Thanks to the authors for this good submission. Please include the feedback in the rebuttal to the main paper.
Decision: Accept (poster)