GTMGC: Using Graph Transformer to Predict Molecule’s Ground-State Conformation
Authors: Guikun Xu
Yongquan Jiang
PengChuan Lei
Yan Yang
Jim Chen
Venue: ICLR 2024 spotlight (2024)
Keywords: molecular conformation prediction
molecule modeling
graph neural network
graph transformer
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
TL;DR: We propose a Graph-Transformer network which uses a novel self-attention mechanism we developed for molecular structure modeling, achieving SOTA performance in predicting a molecule’s ground-state 3D conformation from its 2D topology graph.

Abstract: The ground-state conformation of a molecule is often decisive for its properties. However, experimental or computational methods, such as density functional theory (DFT), are time-consuming and labor-intensive for obtaining this conformation. Deep learning (DL) based molecular representation learning (MRL) has made significant advancements in molecular modeling and has achieved remarkable results in various tasks. Consequently, it has emerged as a promising approach for directly predicting the ground-state conformation of molecules. In this regard, we introduce GTMGC, a novel network based on Graph-Transformer (GT) that seamlessly predicts the spatial configuration of molecules in a 3D space from their 2D topological architecture in an end-to-end manner. Moreover, we propose a novel self-attention mechanism called Molecule Structural Residual Self-Attention (MSRSA) for molecular structure modeling. This mechanism not only guarantees high model performance and easy implementation but also lends itself well to other molecular modeling tasks. Our method has been evaluated on the Molecule3D benchmark dataset and the QM9 dataset. Experimental results demonstrate that our approach achieves remarkable performance and outperforms current state-of-the-art methods as well as the widely used open-source software RDkit.

Published as a conference paper at ICLR 2024
GTMGC: USING GRAPH TRANSFORMER TO PREDICT MOLECULE'S GROUND-STATE CONFORMATION
Guikun Xu', Yongquan Jiang 2,33004 Pengchuan Lei', Yan Yangh,3, Jim X. Chen
Sccoool of Computing and Aritifical Intelligence, Southwest Jiaotong University, Chengdu, China
Institute of Aritifical Intelligence, Southwest Jiaotong University, Chengdu, China
Engineering Research Center of Sustainable Urban Intelligent Transportation, Ministry of Education, China
*Department of Computer Science, George Mason University. Fairfax, VA, USA ri :ichhu94455 90nmuu4550ga1.... com; lyqjiange yyange, my pc@my. swjtu.edu. en; mixxchen@gmall com
ABSTRACT
The ground-state conformation of a molecule often decisive for properties. However, experimental or computational methods, such as density functional the- ory (DFT), are consuming and labor-intensive for obtaining this conforma- tion. Deep learning (DL) based molecular representation learning (MRL) has made significant advancements molecular modeling and has achieved remark- able results in various tasks. Consequently, it has emerged promising ap- proach for directly predicting the ground state conformation of molecules. In this regard, introduce GTMGC, novel network based Graph-Transformee (GT) that seamlessly predicts the spatial configuration of molecules in 3D space from their 2D topological architecture an end-to-end manner. Moreover, we pro- pose novel self-attention mechanism called Molecule Structural Residual Self- Attention (MSRSA) for molecular structure modeling. This mechanism not only guarantees high model performance and easy implementation but also lends it- self well other molecular modeling tasks. Our method has been evaluated on the Molecule3D benchmark dataset and the QM9 dataset. Experimental results demonstrate that our approach achieves remarkable performance and outperforms current state-of oftte-of-hhe-att methods well the widely used open-source software RDkit.
INTRODUCTION
The molecular ground-state conformation refers the lowest energy state of a molecule on its potential energy surface. It represents the most stable configuration, requiring the least amount of energy to maintain. Additionally, plays a crucial role in determining the physical, chemical, and biological properties of the molecule in most cases.
Recently, MRL methods based on DL have made significant strides in the field of molecular modeling (Wigh al. 2022). This approach represents molecules graphs, where atoms and bonds are respectively depicted as nodes and edges with features. By utilizing graph neural net- works (GNNs) (Kipf & Welling, 2016) Veličković et al. 2017: al. 2020 Ying et al. 2021 Kim jet al. 2022) to learn the representation of molecules, a variety of tasks such as molecular property prediction, molecular generation, and molecular optimization can be accomplished (Gilmer al.
2017 Shi et al. 20200] Luo et al. [2022] Zhang et al. 2023b). Of particular note, some studies (Wu al. 2018 Townshend al. 2020 Liu al. [2021b] Luo al. 2022; Liao Smidi, [022) have attempted to incorporate the ground-state state 3D geometric information of molecules, such as atomic coordinates and interatomic distance matrix, as additional inputs the neural networks. Such rich and advantageous inductive bias has significantly enhanced the model's performance across wide
range of tasks.
It is clear that studying molecular ground-state conformation is of great significance. However, obtaining the conformations molecules challenging task. At present, the primary meth- ods include molecular dynamics simulations (De Vivo et al. 2016). rough approximations through manually hand-designed force fields (Rappé et al.) 1992; Halgren 1996). or computationally in- tensive (DFT) calculations (Parr al. 1979) and so experimental and computational resources or are too imprecise. Therefore, finding an efficient way to obtain molecular ground-state conformation is an important research topic. Due the success
These methods either demand substantial

---

Published as a conference paper at ICLR 2024
of DL methods in the field of molecular modeling and the success of AlphaFold2 Jumper et al. 2021 in protein structure prediction, using deep neural networks (DNNs) to directly predict them is highly promising approach.
Related Work. In previous studies (Mansimov et al.;
2019
Simm & Hernández-Lobato 2019
Xu et al. 2021b c. Wu et al., 2022 Morehead & Cheng 2023 deep generative models have been employed to generate low-energy stable conformations of molecules. However, these generative methods mainly focused on producing a multitude of potential stable conformations, rather than specifically targeting the ground-state conformation. As a result, additional screening steps were required in real-world applications to identify the best conformation among the generated ensem- ble. Recently, a novel benchmark called Molecule3D (Xu et al. (021d) has emerged as a valuable resource for predicting the ground-state geometry of molecules. This benchmark comprises an ex- tensive dataset consisting of nearly 4 million molecules along with their corresponding ground-state conformations. In Sec. A more detailed related work has been organized.
Building upon the Molecule3D benchmark, we propose GT-based (Dwivedi & Bresson, 2020, Ying et al. 2022 Kim et al. 2022 Min et al., 2022) model for predicting molecule's ground conformation. Our model is capable of seamlessly transforming molecular 2D graph into its stable 3D ground-state conformation through end-to-end prediction, climinating the need for additional intermediate variables such predicting the interatomic distance matrix first and then recovering the 3D coordinates from it (Simm& & Hernández-Lobate 2019. Shi et al. 2021).
Implementation. In detail, The hole overview of our model's architecture is shown in Fig. Firstly, the newly introduced context-aware Mole-BERT Tokenizer (Xia et al.) (023) is creatively employed categorize identical atoms into distinct subclasses, mitigating the quantitative diver- gence between prevalent and scarce atoms. Each atom of molecules is tokenized into discrete values with chemical significance (e.g.. carbon to carbon in benzene ring or aldehyde.). serving as Input IDs. Then, the widely used Laplacian Positional Encoding (Dwivedi & Bresson, (2020) in the graph domain which can be solely computed from graph's 's adjacency matrix without any complex or exces- sive prior knowledge is applied to encode the positional information of each node within the entire graph. Furthermore, A novel self-attention mechanism, termed Molecule Structural Residual Self- Attention (MSRSA), has been proposed. This mechanism utilizes commonly available and computa- tionally inexpensive molecular structural information, adjacency matrix and creative row-subtracted interatomic distance matrix to compute residual bias terms upon self-attention scores (QK"), to enhance the self-attention mechanism's 's ability to model molccular structures. It largely retains the structure and elegance of the original Transformer (Vaswani al. 2017) network, unlike other works (Ying et al., 2202T] [Li et al. 2022) that necessitate complex design and computation. Finally, ford-forward network (FFN) head atoms' representations predict the molecular 3D> ground-state conformation, achieving end-to-end prediction from the molecule's s 2D graph.
Results and Contribution. (i). Our approach, GTMGC, that predicts molecular 3D ground-state conformation from its 2D graph topological representation in an end-to-end manner, has been val- idated on the recently proposed benchmark dataset Molecule3D (Xu et al. (021d) and the widely used QM9 (Ramakrishnan et al., (0144) dataset. Experimental results demonstrate that our approach outperforms the current best method (Xu et al.] (202dd) and surpasses the DG (Havell 1998) and ETKDG (Riniker & Landrum 2015) algorithms implemented in the most commonly used open- source software RDkit (Landrum et al. 2013). (ii). GTMGC employs a distinctive input form. We innovatively leverage the MoleBERT Tokenizer (Xia et al., 2023) to discretize each atom of the molecule into chemically significant tokens, serving as the model's input, rather than conventional atom type IDs or Ogb-style embeddings (Hu et al. 2021). This effectively helps the model predict molecule's 's ground-state conformation. (tx7). Furthermore, ablation studies strongly demonstrate the effectiveness of our method. (iv). More experiments in Sec. show that the proposed MSRSA mod- ule can be effectively applied to many other molecular property prediction tasks. (v). The source code of our method available at teltts::/gittee
com/Rich XGK/ XCKKXXXGCC
2 PRELIMINARY
2.1 NOTATION AND PROBLEM DEFINITION
Notation. For molecule with n atoms and m bonds, its 2D structure representation is denoted as g = E}. where V {e), U,...., 2n represents the set of n atoms and È (e1,e2.... cm represents the set of bonds. Each atom e and bond ej may possess its own features, such as
2

---

Published as a conference paper at ICLR 2024
Si = Q:K7
(Bmar D)
2D
3D
=
0 0EE+
+ yi
(a)
(b)
Figure 1: (a): The overview of the ground-state conformation prediction task; (b): Computation
process of the attention map in MSRSA module, viewed through the lens of attention mask.
atom type, bond type, etc. In the context of graph structure g, its adjacency matrix is denoted by A E Raxn where Aij = signifies the presence of edge connecting node vj vs while molecule
Aij = indicates the absence such an edge. The ground-state conformation denoted as G E I=*3 which represents a set of 3D Cartesian coordinates. Expanding on this, the interatomic distance matrix of molecule is defined as D Rexn, where Dij represents the Euclidean distance between atom bty and 0t-
Problem Definition. As depicted in F the ground-state conformation prediction task aims to predict the conformation of a molecule its ground state, [Rax solely based on its 2D molecular structure, G (V,E).
2.2 TRANSFORMER
In vanilla Transformer (Vaswani al. 2017), the multi-head sealed dot-product self-attention (MHSA) module is solely responsible for modeling relationships between different elements. By calculating the similarity between queries and keys, determines the importance of different ele- ments to a single element. The weights are then normalized and used to compute a weighted sum of the values, resulting in a new representation for the corresponding element. The use of multiple heads allows different attention heads to learn distinct information within separate subspaces.
Suppose the input fed into the MHSA module is X Ruxddodel, where n rep- resents the number of tokens and dmoods represents the hidden dimension of the model. W E Rddmode Xdg WK E RRomodei Xdz, and WV E X are three learnable linear mappings that map X to Q TR K E Banxd and V Raxdd respectively. The scaled dot-product attention
Self-Attention.
then computed as:
K K V XW = X K
(1)
Attention Q2 K, V) softmax QKT
vdk V
(2)
Multi-Head. After introducing heads, total of h w Rdmodax w E R and W/ X map X Qi R * K, Rax4 and V Rax 4x where
Rdmoda x
[1. M.. Then one head of sealed dot-product attention can be computed equation Let Wo Rdd Xdmodde be another learnable linear mapping to map the concatenated output of all heads to the output of the MHSA module:
QKT V
0 Attention (Q,, K,, Vi) softmax
(3)
Vdz
MHSA(X) = Concat (O,, Oz...., .., O₂ W W
Transformer Layer. After MHSA, a feed-forward network (FFN) consisting of two linear trans- formations with non-linear activation is applied each element respectively and identically. Both the output of MHSA and FFN are first passed through a residual connection (He et al. 2016) and then normalized by a LayerNorm (Ba et al. 2016). The output of Transformer layer can be com- puted as:
X' = LayerNorm + (MSA(XX)
(5)
Layer(X) X) LayerNorm(X FFN(X'))
(6)
3

---

Published as a conference paper at ICLR 2024
Malende Sinateral Roiieel Self Alteetaee (MSRSA))
Mulli-liadd Altndion
Predicte Graant GranndkSune Casferaniins
Sofwar
Fed Forurt
X/ Syore!
Lawer
N x
A Nert
add & Nera
Fead Parnar
(d)
tineceering
O snllytroor srop - 467
NIMSA with
492
Taieeber
(I)
(b)
Lapladinn Pocifonal
Erabwdding
Islcrttesk Dabanx
Figure 2: Overview of GTMGC. (a): Atoms in the molecule tokenized by MoleBERT Tok-
enizer (Xia et al. 2023 into chemically meaningful discrete codes as model's Input IDs. (b): The encoder of GTMGC consisting of N MSRSA blocks with only adjacency matrix residual bias term BA, (c): The Decoder of GTMGC consisting of MSRSA blocks with both adjacency matrix resid- ual bias term B and interatomic distance matrix residual bias term B (d): The main Molecule Structural Residual Self-Attention (MSRSA) module, using residual bias terms B and B to help pure dot doo-pprouuct self-attention mechanism model molecule's 's structure information.
3 METHOD
3.1 MODEL INPUT
Input IDs. Thanks to the convenience of (Xia et al. (023), we can now use the chemically mean- ingful IDs tokenized by the MoleBERT Tokenizer input IDs for our model, as shown in Fig. Aa). Since these tokens are assumed to already contain ample chemical information about atoms, we have ceased utilizing edge features to enhance the representation of molecular input features, thereby making the model more concise.
Positional Encoding. For graph structures, incorporating positional encoding becomes even more vital and challenging since elements in graph structures lack fixed positions. is necessary to differ- entiate the positions of various nodes within the graph structure, allowing the model to capture the relationships between nodes. To adhere to the principle of introducing minimal prior knowledge to the model and enhancing its generality, we opted for the Laplacian Positional Encoding as (Dwivedi & Bresson 2020) does, among using various graph structure positional encoding techniques (Zhang [dl.[200)) Ying et al., 202T Hussain et al. 2021; Park et al., 2022). This approach encodes posi- tional relationships between nodes using the eigenvectors of the graph Laplacian matrix. It can be solely computed from every molecular graph's 's adjacency matrix without any complex or excessive prior knowledge.
Thus, our final input, H E R Xdd is solely obtained by adding the Laplacian Positional Encod- ing vectors, L gaxn, to the feature vectors, E E [R Xdmoddd embedded from the MoleBERT- tokenized input IDs.
3.2
GRAPH STRUCTURAL RESIDUAL
(MSRSA)
we review in detail how a vanilla Transformer layer performs forward calculations. Its original dot-product self-attens self-attention mechanism models the relationships between all nodes in an undirected fully connected graph, similar to a GNN model (Kipf & Welling 2016 Veličković et al.) 2017 Gilmer al. 2017) that performs message passing
In Section 2.2
undirected fully connected graph.
For molecules, the local structure of atoms (such as functional groups) and the connectivity of edges (whether they have chemical bonds not) have huge impact on their chemical and structural prop- erties. Since the original self-attention mechanism is unable to effectively capture these information, propose an extended self-attention mechanism based molecular structure, named Molecule
4

---

Published as a conference paper at ICLR 2024
Structural Residual
(MSRSA). as can be observed in Fig. 2 d), which incorporates the adjacency matrix, E R=x", and interatomic distance matrix, E Rax" of the molecule as residual bias terms for self-attention scores, S Q:K?" R"x".
Global. Unlike the approach used by (Dwivedi & Bresson 2020) that employs the adjacency ma- trix as full attention mask to restrict the current node's 's foeus to neighboring nodes, method based on retaining complete global attention. This brings our model closer to the original Trans- former, we assume that the original self-attention mechanism captures global information by focusing on all atoms in the molecule. Therefore, the basic self-attention scores calculated by full self-attention in each head E [1, h) are:
S XW - Q.K?
Nearby. To capture something useful about the molecule's local structure, the adjacency matrix A has been introduced as residual bias term B, R=** to St, as shown in Fig. [ d). B, models atom's 's J-hop local structure information by attending to neighboring atoms connected to the current atom. For each head, first multiply the adjacency matrix A by learnable parameter 71" to automatically learn the requirement for local information. The result is then used in a hadamard product with Sy to obtain BA
= (7d x A)
Spatial. To enable the model to represent the spatial structure of molecules, introducing the in- teratomic distance matrix of molecules is direct and simple approach. Similarly, we use the in- teratomic distance matrix of molccules as another residual bias term B, E "=x", as shown in Fig. d). BP models molecule's s spatial structure by assigning more weight to atoms closer to the current atom, which introduces the simple distance relevance assumption: the greater the distance between atomic elements, the lower the interaction (Choukroun Wolf, 2021). For each head, first subtract the row-max value, 3Tmax i [1, n) which represents the furthest distance from other nodes to the current node v; from each row D to obtain row-subtracted distance ma- trix D mm--ub: representing the distance relevance assumption. It is worth noting that, as far know, this represents our biggest difference from others who use the interatomic distance matrix to introduce molecule's spatial structure information and its effectiveness supported by experimental evidence (Table 3 Table 4. Then another learnable parameter 7/ is multiplied to D row-suu to au- tomatically ascertain the necessity for spatial information. The result is finally used in a hadamard product with S/ to obtain B.:
D Dow
B Finax
D, Finax
max D vE[L,n]
BP = ② (TP D rew-sub.)
(10)
Thereafter, we sum up S,, B and B, together, where B and BP are residual bias terms of S and then scale the results by vdg, and finally normalize it using a softmax function to obtain the final output O of each head:
S =S + B + P
=S (1 x A) hy x (5max D)j
(11)
= -KKK @ [E, + n x A +
x (Jmax - D)]
S
0 softmax
V
(12)
dd
Analytically, as unveiled in Fig. (b), we can dissect equation I through the lens of implementing multiple attention masks on the self-attention scores S, Specifically, E,, E Raan initially executes an unobscured mask grounded on the original attention. Concurrently, A could be interpreted as an additional mask that confines the current node to solely focus on its neighboring nodes. D could be perceived as a spatial awareness mask that adjusts the level attention based on the distance between nodes. The significance of various masks is learned through parameters 7/" and and which are subsequently fused to influence Sj. By adopting this approach, we only need to introduce a achieve sub-
minimal amount of parameters based pure self-attention, compared to which, stantial performance enhancement. Otherwise, MSRSA shares certain similarities and associations with iMaziarka et al. 22020) Choukroun & Wolf 202I), as they also utilize A and D to enhance the self-ate attention mechanism. More details on the differences and connections between them have been provided in A.3
5

---

Published as conference paper at ICLR 2024
(A) Reference
ROKA (C-2MBDD)
(0.764)
(1,040)
(1.397)
(1.132)
(C) GPS
(C-RNISDI
(8.54))
(0:555)
(4) GTMGE (OHES)
(C-RMASR)
(0/667)
(0.210)
Figure 3: Examples inferred by RDkit, GPS and GTMGC (ours); (a): The ground truth, molecules'
ground-state 3D conformations; (b): The predictions RDkit; (c): The predictions of GPS; (c): The predictions of GTMGC. More examples are shown in Sec. C.I
3.3 ARCHITECTURE OF GTMGC
Encoder. As shown in Fig. (b). in the encoder with N blocks and H heads, only the adjacency matrix residual bias term B has been used enhance the original self-attention mechanism. Ini- tially, the input of the encoder only contains the node embeddings IH Raaddmdee and 2D) structural information, adjacency matrix, A E _axx of the molecule. After forward calculation in the encoder, we pass the output of the Encoder HУ pnxxdmoee through learnable weight-sharing two-layer FFN head (-) to obtain a relatively rough conformation prediction result G Geacke E [gax3 and then calculate the interatomic distance matrix of Geache: namely D cacke E R
Decoder. Instantly, we feed H' E Raxdmodei A E R and D Deache E R as inputs to the decoder with N blocks and H heads, depicted in Fig. (c). In the decoder, full MSRSA module with A and B. bias terms has been utilized to model the molecule's spatial structure. Finally, we pass the decoder's output HI" through the same FFN head f.r(-) used before return the final conformation G E |grx3 predicted result.
3.4 LOss
In order to keep the prediction results rotationally and translationally invariant, we use the Mean Absolute Error (MAE) between the interatomic distance matrix of the prediction, D and that of the ground truth, D, as the loss. In an innovative approach, we trained the model by combining the loss between its final output G and G with the loss between the Encoder's output Geeche and G. This enables the Encoder to learn more effective representation and provide the Decoder with more precise initialization space information cache. ultimately enhancing the model's 's performance. The final loss function is follows:
n
n Dij D
MAB(D, D*)
(13)
n
1=1 j=1
MAE(D, D*) + MAE(D, Deache)
(14)
EXPERIMENTS
4.1 DATASETS
Molecule3D. The first benchmark introduced by (Xu et al. 2021d) that aims to use DNNs to predict the ground-state 3D geometries of molecules based solely on their 2D graph structure. The proposed large-scale dataset comprises approximately 4 million molecules, each with its own 2D molecular graph, ground-state 3D geometric structure, and four additional quantum properties. employs splitting methods: random splitting according to the same probability distribution and scaffold splitting based on the molecule's 's core component. Both partitions use a 6:2:2 ratio for training, validation, and testing.
QM9, A small-scale quantum chemistry dataset (Ramakrishnan et al. 2014 Wu et al. 2018) that provides geometry, energy, electronic, and thermodynamic properties for nearly 130,000 organic

---

Published as a conference paper at ICLR 2024
molecules with 9 heavy atoms, containing molecular most stable conformation calculated by DFT. We adopt the identical data split as described in (Liao & Smidt (022), where 110k, 10k, and 11k molecules are allocated for training, validation, and testing, respectively.
4.2 METRICS
Followed (Xu et al. 2021d given dataset with totally N interatomic distances, the Mean Absolute Error (MAE) and Root Mean Square Error (RMSE) between the prediction, {dī Ti=1 and the ground truth, {di). are used to evaluate the performance at node-pair level:
N
D-MAE((4)) (d
d dil
(15)
1 N
D-RMSH (OMSB((4))))))) (4) (1
(d: di)²
(16)
Additionally, We followed previous generative works (Mansimov al. 2019 Wu al. 2022) to
calculate the Root Mean Square Deviation (RMSD) of n heavy atoms between the ground truth, G, and the prediction rigidly aligned to the ground truth by the Kabsch algorithm Kabsch 1978), G This metric can effectively measure the spatial difference between two conformations.
C-RMSD(G,G*) =
(17)
4.3 CONFORMATION PREDICTION
Setup. To validate the effectiveness and advancement of algorithm in predicting molecular ground-state conformation, we initially selected widely used DG and ETKDG algorithms imple- mented by RDkit as baselines. Additionally, the benchmark study (Xu et al. 2021d) employed DeeperGCN-DAGG (Liu et al. 2021a) to directly predict atom coordinates, achieving previously Mtate-of-the-are performance. Given the limited research in this specific area, we included GINE (Hu et al. 2019). GATv2 (Brody et al. 2021 and GPS (Rampášek et al., 2022) as comparative bench- marks. GINE and GATv2, known for their impressive capabilities, are robust 2D GNNs, and GPS
Table 1: The performance the Molecule3D and QM9 datasets (Å).
Validation
Test
D-MARJ
D-RMSEL
C-RMSDJ
D-MAEL
D-RMSEL C-RMSDJ
(a) Molcule3D Random Split
RDKit DG
0.581
0.930
1.054
0.582
0.932
1.055
RDKi ETKDG
0.575
0.941
0.998
0.576
0.942
0.999
DeeNeGGN-DAGGNN [uu et al. 2021d
0.509
0,571
0.961
GINE (u al. 2019
0.590
1.014
1.116
0.592
1.018
1.116
GATv2 al 2021
0.563
0.983
1,082
0.564
0.986
1.083
GPS (Rampasek et al. 2022
0.528
0.909
1.036
0.529
0.911
1.038
GTMGC (Ours)
0,432
0.719
0.712
0.433
0.721
0.713
(h) Molcule3D Scaffold Split
RDKit DG
0.542
0.872
1.001
0.524
0.857
0.973
RDKit ETKDG
0.531
0.874
0.928
0.511
0.859
(1898
DeepenGCN-DAGNT Xu et al.
2021d
0.617
0.930
0.763
1.176
GINE Hu al. 2019
0.883
1.517
1.407
1.400
2.224
1,960
GATv2 (Bodyy al.1
0.778
1.385
1.254
1.238
2.069
1.752
GPS (Rampasck et 2022
1.031
0.657
1.091
1.136
GTMGG (Ours)
0.406
0.675
0.678
0.400
0.679
0.693
(c) QMY
RDKit DG
0.358
0.616
0.722
0.358
0.615
0.722
ETKDG
0.355
0.355
0.689
GINE et 2019
0.357
0.673
0.685
0.357
0.669
0.693
G 021
0.339
0.663
0.661
0.339
0.659
0.666
GPS (Rampasek et 2022
0.326
0.644
0.662
0.326
0.640
0.666
GTMGC (Ours)
0.262
0.468
0.362
0.264
0.470
0.367
"The asterisk (*) indicates that the result this metric was reported
al. [222Td]

---

Published as a conference paper at ICLR 2024
Table 2: Ablation study on input format (A).
D-MAEL
D-RMSE4
C-RMSDD
Ogb-style embeddings (Hu et al. 2021
4204
.7162= 0003
.74314.0198
Atom Type IDs
443384.00z
.71954.00m
.72174.0002
MoleBERT Tokenized IDs (Ours)
.4300.000
.7213 0004
.7139 +.000
means the indicator we mainly focus on.
Table 3: Ablation study of MSRSA module for ground-state conformation prediction on
Molecule3D random split (A).
Methods
Index (Arrr)))
LPE
MHSA
B
rerader
B
B
deedder
Dimenassi divesler
D-MAEL
D-RMSEJ
C-RMSDL
(+0)
54644..0m
.009
99724_0111
2 (40)
2.004
.2377 2.008
.7388_ 0009
3 (+48)
/
4353 g.00
.217
.7213_
(+96)
4.000
.7216 0000
.7299_ 0111
5 (+144)
1325 20002
.2214
.7202 c057
6 (+144)
.4330
.72131.0004
.71391.0010
micans indicator we mainly focus on.
is a powerful GT-style network that was recently proposed. Our model, GTMGC, was compared against these baselines, and the results the Molecule3D and Qm9 datasets presented in Ta- ble (a), (b) and (c). Except for the results obtained from the original benchmark paper (Xu et al.) 2021d for DeeperGCN-DAGNN, all other findings presented in this study are derived from periments. It is important to note that the results GTMGC on QM9 were fine- fine-tuned based on its performance on Molecule3D. Finally, Fig. 3 illustrates a comparison between the predicted results of our model and the references on samples from the random split test set of Molecule3D.
Results and Analysis. (i). As shown in Table (a) and (b), in the Molecule3I) random and scaffold split, GTMGC improved 18. 34%, 20.86%, and 28. 28.63% and 21.71%, 20.79%, and 22.83% respectively on the D-MAE, D-RMSE, and C-RMSD metries in the test set compared to the previous top results. (ft). Additionally, in both splits, GTMGC's results on the validation set and test set are very close, indicating that our method has good generalization performance. (iit). Unlike (Xu al., 2021d). it has already overfitted on random split and performs worse scaffold split. GTMGC performs well on Scaffold split, indicating that can predict molecules with different scaffolds from those unseen during training. (iv). As demonstrated in Table (c), GTMGC also achieved superior results on the Qm9 test set, indicating its exceptional generalization performance across various datasets and its applicability to both large and small scale molecules. (v). Furthermore, as illustrated by the examples in Fig. the conformations predicted by GTMGC more closely aligns with the references.
4.4 EFFECTIVENESS STUDY
Setup. To evaluate the effectiveness of GTMGC, conducted ablation experiments on the Ran- dom Split of Molecule3D. Each result represents the mean (standard deviation, std) value, calculated after conducting the experiment 3 times. First, modified the input format of the model while keeping the network architecture constant to evaluate the impact our proposed input format on molecular ground-state conformation prediction. The results are presented in Table Moreover, we conducted experiments to evaluate the influence of individual components within the core MSRSA module. We initially implemented a pure self-attention form and progressively incorporated various components. The outcomes of these experiments are presented in Table 3 Finally, visualization of attention weights has been employed in Fig. 4 as means showcase the proficiency and ap- proach of our MSRSA module in the modeling of molecular structures. Specifically, the appendix (Sec. B Table 4 includes an extra ablation study designed to validate the contribution of each com- ponent when the MSRSA module is employed for molecular property prediction.
Results and Analysis. As shown in Table 2, our proposed input format, MoleBERT Tokenized IDs, outperforms the other two input formats, Ogb-style (Hu et al. (021) embeddings and Atom Type IDs, on the C-RMSD metric which best reflects the difference between conformations and mainly focus Although it falls slightly short in the D-MAE and D-RMSE indicators. the difference not significant. However, its improvement in C-RMSD is substantial. (t). This sug- gests that proposed input format is more effective for ground-state conformation prediction. We
8

---

Published as a conference paper at ICLR 2024
Moveculle Geaph
Nead
Movecuae Groph
Hewd O
630
6.0
608
cm
6AM
Nead
Nearr
Read
Headd
6.15
mase
0.55
0.re
3.0
0.58
(a)
(b)
Figure 4: attention weights visualization. The triangular node (A) represents the current node ((a): No.7, Br; (b): 14. S.). while the other nodes are shaded from dark to light based on the magnitude of their attention weights relative to the current node. More examples are shown Sec. C.2
hypothesize that the MoleBERT Tokenizer assigns distinct IDs atoms of the same type, result- ing in embeddings with unique chemical meanings. This addresses the issue of similar type nodes having close representations after being processed by the model. leading proximate outcomes (coordinates). Additionally, as shown in Table. 3 starting from pure selffattention, we first add
BA encoder the Encoder, followed by B decoder the Decoder. Finally, we incorporate B D(original) and B (vvos-sub) decoder into the Decoder respectively. (ii). As each component is integrated, the perfor- mance of the model almost gradually improves, particularly in the C-RMSD metric, which best flects differences between conformations and has seen improvements of 2.37%, 1.21%, and 3.37% respectively compared pure self-attention. (ivi). Notably, as each component incorporated, standard deviation of performance (C-RMSD) gradually decreases, indicating increased stability in our model. (iv). Moreover, by employing our proposed D pow-sub in place of the original inter- atomic distance matrix D, have achieved performance enhancement of 0.87% along with reduction in standard deviation. Regrettably, the Drow-sub we utilized is not sufficiently precise as is derived from preliminary rough estimates Geache By introducing the actual D jow-suub during molecular property prediction, superior outcomes are presented in Sec. B Table 4 (v). The attention weights learned are depicted in Fig. 4 illustrating the degree of focus specific atom has towards other atoms during its modeling. The figure reveals that the model discerns diverse attention patterns across different attention heads. For instance, in Fig- (a), the atom Br, represented by a triangle, primarily captures the molecule's global information in head O, paying substantial attention to the heavy atoms within the two rings. In heads and 4, distinctly showcases the acquisition of the molecule's is spatial structural information. As observed in the figure, moving from left to right (or vice versa), and from near to far, the attention of the Br atom towards other atoms progressively diminishes. This characteristic endows model with enhanced intuitive interpretability.
decoder
CONCLUSION
We propose novel Transformer-basee method, GTMGC, for end-to-end prediction of 3D ground- state conformations of molecules from their 2D topological structures. Our method achieves signif- icant performance improvement over the previous best methods, reaching the state-of-the-art level, and also surpasses the widely used open-source toolkit RDkit on various metrics. Moreover, we introduce novel and simple self-attention mechanism for molecular structure modeling, namely Molecule Structural Residual Self- Auff-Attetiion (MSRSA). It preserves the original self-attention struc- ture Transformer ture. Experiments show that MSRSA module not only provides great help for predicting the ground- state conformations of molecules, but also is easy to implement and generalize to other molecular modeling tasks, such molecular property prediction.
large extent, while also being able effectively model the molecular struc-
ACKNOWLEDGMENTS
This work is supported by the National Natural Science Foundation of China (No.61976) 61776247) and the Fundamental Research Funds for the Central Universities (N0.2682022TTPY055).

---

Published as a conference paper at ICLR 2024
REFERENCES
Uri Alon and Eran Yahav. On the bottleneck of graph neural networks and its practical implications. arXiv preprint arXiv:2006.0522 2020.
Brandon Anderson, Truong Son Hy, and Risi Kondor. Cormorant: Covariant molecular neural networks. Advances in neural information processing systems, 32, 2019.
Simon Axelrod and Rafael Hoonez-oobarrllli. Geom: Energy-annotated molecular conformations for property prediction and molecular generation. arXiv preprint arXiv:2006.05531, 2020.
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint 1rXvv:16070064500 2016.
Johannes Brandstetter, Rob Hesselink, Elise van der Pol, Erik J Bekkers, and Max Welling, Ge- ometric and physical quantities improve e (3) equivariant message passing. arXiv preprint 1r.iv:2100.02905, 2021.
Shaked Brody, Uri Alon, and Eran Yahav. How attentive are graph attention networks? arXiv preprint arXiv:2105.1 14491, 2021.
Deli Chen, Yankai Lin, Wei Li, Peng Li, Jie Zhou, and Xu Sun. Measuring and relieving the over- smoothing problem for graph neural networks from the topological view. In Proceedings of the AAAI conference artificial intelligence, volume 34, pp. 3438-3445, 2020.
Yoni Choukroun and Lior Wolf. Geometric transformer for end-to-end molecule properties predic- tion. arXiv preprint arXiv:2110.137211 2021.
Marco De Vivo, Matteo Masetti, Giovanni Bottegoni, and Andrea Cavalli. Role of molecular dy- namics and related methods in drug discovery. Journal of medicinal chemistry, 5909):4035-401,, 2016.
Jacob Devlin, Ming- Wei Chang. Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805. 2018.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x 16 words: Transformers for image recognition scale. arXiv preprint arXiv:2010.11929, 2020.
Vijay Prakash Dwivedi and Xavier Bresson. generalization of transformer networks to graphs. arXiv preprint arXhv:2012.09699 2020.
Zhiguang Fan, Yuedong Yang. Mingyuan Xu, and Hongming Chen. Ec-conf: A ultra-fast diffu- sion model for molecular conformation generation with equivariant consistency. arXiv preprint arXiv:2308.0023: 2023.
Mare Finzi, Samuel Stanton, Pavel Izmailov, and Andrew Gordon Wilson. Generalizing convolu- tional neural networks for equivariance to lie groups on arbitrary continuous data. In International Conference on Machine Learning, pp. 3165-3176. PMLR, 2020.
Fabian Fuchs, Daniel Worrall, Volker Fischer, and Max Welling. Se (rantrfssfss:mrs: 3d roto- translation equivariant attention networks. Advances in neural information processing systems, 33:1970-199 2020.
Octavian Ganea, Lagnajit Pattanaik, Connor Coley, Regina Barzilay, Klavs Jensen, William Green, and Tommi Jaakkola. Geomol: Torsional geometric generation of molecular 3d conformer en- sembles. Advances in Neural Information Processing Systems, 34:13757-13769, 2021.
Johannes Gasteiger, Shankari Giri, Johannes T Margraf, and Stephan Günnemann. Fast and uncertainty- directional message passing for non-equilibrium molecules. arXiv preprint 1rXi::2011.14155 2020.
10

---

Published as a conference paper at ICLR 2024
Justin Gilmer, Samuel S Schoenholz, Patrick F Riley, Oriol Vinyals, and George E Dahl. Neural message passing for quantum chemistry. In International conference on machine learning. pp- 1263-1272. PMLR, 2017.
Thomas A Halgren. Merck molecular force field. v. extension of mmff94 using experimental data, additional computational data, and empirical rules. Journal of Compulational Chemistry, 17(5-6): 616-641, 1996.
Timothy F Havel. Distance geometry: Theory, algorithms, and chemical applications. Encyclopedia of Computational Chemistry, 120:723-742, 1998.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp- 770-778, 2016.
Jonathan Ho. Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840-6851, 2020.
Weihua Hu, Bowen Liu, Joseph Gomes, Marinka Zitnik, Percy Liang, Vijay Pande, and Jure Leskovec. Strategies for pre-training graph neural networks. arXiv preprint arXiv:1905. 12265, 2019.
Weihua Hu, Matthias Fey, Hongyu Ren, Maho Nakata, Yuxiao Dong, and Jure Leskovec. Ogb-lsc: A large-scale challenge for machine learning on graphs. arXiv preprint arXiv:2 2021.
Md Shamim Hussain, Mohammed Zaki, and Dharmashankar Subramanian. Edge-augmented graph transformers: Global self-attention is enough for graphs. arXiv preprint arXiv:2/08.03345 2021.
Bowen Jing, Gabriele Corso, Jeffrey Chang, Regina Barzilay, and Tommi Jaakkola. Torsional dif- fusion for molecular conformer generation. Advances Neural Information Processing Systems, 35:24240-24253, 2022.
John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Zidek, Anna Potapenko, al. Highly accurate protein structure prediction with alphafold. Nature, 596(78733:583-58) 2021.
Wolfgang Kabsch. discussion of the solution for the best rotation relate two sets of vectors. Acia Crystallographica Section A: Crystal Physics, Diffruction, Theoretical and General Crys- tallography, 34(5):827-828, 1978.
Jinwoo Kim, Dat Nguyen, Seonwoo Min, Sungjun Cho, Moontae Lee, Honglak Lee, and Seunghoon Hong. Pure transformers are powerful graph learners. Advances Neurul Information Processing 2022.
Systems,
Thomas N Kipf and Max Welling, Sumi-super classification with graph convolutional net- works. arXiv preprint arXiv: 16X9:0060000207,, 2016.
Greg Landrum et al. Rdkit: A software suite for cheminformatics, computational chemistry, and predictive modeling. Greg Landrum, 8:31, 2013.
Tuan Le, Frank Noé, and AAm Clevert. Equivariant graph attention networks for molecular property prediction. arXiv preprint arXiv:2202.09891, 2022..
Gaichao Li, Jinsong Chen, and Kun He. Adaptive neghh-nneghborhhod attention based transformer for graph representation learning. arXiv preprint arXiv:22/1.0775 2022.
Guohao Li, Chenxin Xiong, Ali Thabet, and Bernard Ghanem. Deepergen: All you need to train
deeper gens. arXiv preprint 2rXii:2006.07739, 2020.
Yi-Lun Liao and Tess Smidt. Equiformer: Equivariant graph attention transformer for 3d atomistic graphs. arXiv preprint arXiv:2206. /1900, 2022.

---

Published as a conference paper at ICLR 2024
Meng Liu, Cong Fu, Xuan Zhang, Limei Wang, Yaochen Xie, Hao Yuan, Youzhi Luo, Zhao Xu, Shenglong Xu, and Shuiwang Ji. Fast quantum property prediction via deeper 2d and 3d graph networks. arXiv preprint arXh::2106008552 2021a.
Shengchao Liu, Hanchen Wang, Weiyang Liu, Joan Lasenby. Hongyu Guo, and Jian Tang. Pre- training molecular graph representation with 3d geometry. In International Conference Learn- ing Representations, 2021b.
Yi Liu, Limei Wang, Meng Liu, Yuchao Lin, Xuan Zhang, Bora Oztekin, and Shuiwang Ji. Spherical message passing for 3d molecular graphs. In International Conference on Learning Representa- tions, 202lc.
Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows. In Proceedings of the IEEE/CVF international conference computer vision, pp. 10012-10022, 2021d.
Shengjie Luo, Tianlang Chen, Yixian Xu, Shuxin Zheng, Tie- Yan Liu, Liwei Wang, and Di He. One transformer can understand both 2d & 3d molecular data. arXiv preprint arXiv:2210.01700 2022.
Elman Mansimov, Omar Mahmood, Seokho Kang, and Kyunghyun Cho. Molecular geometry pre- diction using deep generative graph neural network. Scientific reports, 911:20038,, 2019.
Eukasz Maziarka, Tomasz Danel, Slawomir Mucha, Krzysztof Rataj, Jacek Tabor, and Stanislaw Jastrzębski. Molecule attention transformer. arXiv preprint 2XXi:20022.08244 2020.
Eukasz Maziarka, Dawid Majchrowski, Tomasz Danel, Piotr Gaiński, Jacek Tabor, Igor Podolak, Pawel Morkisz, and Stanislaw Jastrzębski. Relative molecule self-attention transformer. arXiv preprint 10.05421 2021.
Grégoire Mialon, Dexiong Chen, Margot Selosse, and Julien Mairal. Graphit: Encoding graph structure in transformers. arXiv preprint MxXiv:2106.05660 2021.
Erxue Min, Runfa Chen, Yatao Bian, Tingyang Xu, Kangfei Zhao, Wenbing Huang, Peilin Zhao, Junzhou Huang, Sophia Ananiadou, and Yu Rong. Transformer for graphs: An overview from architecture perspective. arXiv preprint arXiv:2202.08455, 2022.
Alex Morehead and Jianlin Cheng. Geometry-complete diffusion for 3d molecule generation. In ICLR 2023-Machine Learning for Drug Discovery workshop. 2023.
Luis Müller, Mikhail Galkin, Christopher Morris, and Ladislav Rampášek. Attending to graph transformers. arXiv preprint arXiv:2302.04181, 2023.
Wonpyo Park, Woong-Gi Chang, Donggeon Lee, Juntae Kim, et al. Grpe: Relative positional encoding for graph transformer. In ICLR2022 Machine Learning for Drug Discovery, 2022.
Robert G Parr, Shridhar R Gadre, and Libero J Bartolotti. Local density functional theory of atoms and molecules. Proceedings of the National Academy of Sciences, 76(6):2522-255 2526, 1979.
Raghunathan Ramakrishnan, Pavlo O Dral, Matthias Rupp, and O Anatole Von Lilienfeld, Quantum chemistry structures and properties of 134 kilo molecules. Scientific data, 1(1):177, 2014.
Ladislav Rampášek, Michael Galkin, Vijay Prakash Dwivedi, Anh Tuan Luu, Guy Wolf, and Do- minique Beaini. Recipe for general, powerful, scalable graph transformer. Advances Neural Information Processing Systems, 35: 14501-14515. 2022.
Anthony K Rappé, Carla J Casewit, KS Colwell, William A Goddard III, and Mason Skiff. Uff. full periodic table foree field for molecular mechanics and molecular dynamics simulations. Journal of the American chemical society, 114225):10224100335 1992.
Sereina Riniker and Gregory A Landrum. Better informed distance geometry: using what know to improve conformation generation. Journal of chemical information and modeling. 55(12); 2562-2574, 2015.
12

---

Published as a conference paper at ICLR 2024
Yu Rong, Yatao Bian, Tingyang Xu, Weiyang Xie, Ying Wei, Wenbing Huang, and Junzhou Huang. Self-supervised graph transformer on large-scale molecular data. Advances in Neural Information Processing Systems, 33:12559-122777 2020.
Victor Garcia Satorras, Emiel Hoogeboom, and Max Welling. E (n) equivariant graph neural net- works. In International conference on machine learning, pp. 9323-9332. PMLR, 2021.
Kristof Schütt, Pieter-Jan Kindermans, Huziel Enoc Sauceda Felix, Stefan Chmiela, Alexandre Tkatchenko, and Klaus-Rober Müller. Schnet: continuous- filter convolutional neural network for modeling quantum interactions. Advances in neural information processing systems, 30, 2017.
Kristof Schütt, Oliver Unke, and Michael Gastegger. Equivariant message passing for the prediction of tensorial properties and molecular spectra. In International Conference on Machine Learning, pp. 917799377788 99888. PMLR, 2021.
Chence Shi, Minkai Xu, Zhaocheng Zhu, Weinan Zhang, Ming Zhang, and Jian Tang- Graphaf: a flow-based autoregressive model for molecular graph generation. arXiv preprint arXiv:2001.09382. 2020.
Chence Shi, Shitong Luo, Minkai Xu, and Jian Tang. Learning gradient fields for molecular con- formation generation. In International conference on machine learning, pp. 2021.
PMLR,
Gregor NC Simm and José Miguel Hemándezz-Lbbtoo A generative model for molecular distance geometry. arXiv preprint arXiv:1909. 1459, 2019.
Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances in neural information processing systems, 32, 2019.
Philipp Thölke and Gianni De Fabritüi.. Equivariant transformers for neural network based molecu- lar potentials. In International Conference on Learning Representations, 2021.
Nathaniel Thomas, Tess Smidt, Steven Kearnes, Lusann Yang. Li Li, Kai Kohlhoff, and Patrick Riley. Tensor field networks: Rotation-and translation-equivariant neural networks for 3d point clouds. arXiv preprint arXiv:1802.082/13 2018.
Raphael JL Townshend, Martin Vögele, Patricia Suriana, Alexander Derry, Alexander Powers, Yianni Laloudakis, Sidhika Balachandar, Bowen Jing. Brandon Anderson, Stephan Eismann, al. Atom3d: Tasks on molecules in three dimensions. arXiv preprint arXiv:20/2.04035. 2020.
Aaron Van Den Oord, Oriol Vinyals, et al. Neural discrete representation learning. Advances in neural information processing systems, 30, 2017.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural informa- fion processing systems, 30, 2017.
Petar Veličković, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pietro Lio, and Yoshua Bengio. Graph attention networks. arXiv preprint 1rXiv: 2017.
Limei Wang, Yi Liu, Yuchao Lin, Haoran Liu, and Shuiwang Ji. Comenet: Towards complete and efficient message passing for 3d molecular graphs. Advances in Neural Information Processing Systems, 35:650-664, 2022.
Daniel S Wigh, Jonathan M Goodman, and Alexei A Lapkin. review of molecular representa- tion in the age of machine learning. Wiley Interdisciplinary Reviews: Computational Molecular Science, 12(5)ee6033, 2022.
Fang Wu, Qiang Zhang, Dragomir Radev, Jiyu Cui, Wen Zhang, Huabin Xing, Ningyu Zhang, and Huajun Chen. 3d-transformer: Molecular representation with transformer in 3d space. 202 la.
Lemeng Wu, Chengyue Gong, Xingchao Liu, Mao Ye, and Qiang Liu. Diffusion-based molecule generation with informative prior bridges. Advances in Neural Information Processing Systems, 15:35533-36555. 2022.
13

---

Published as a conference paper at ICLR 2024
Zhanghao Wu, Paras Jain, Matthew Wright, Azalia Mirhoseini, Joseph E Gonzalez, and lon Stoica. Representing long-range context for graph neural networks with global attention. Advances in Neural Information Processing Systems, 34:13266-13279, 2021b.
Zhenqin Wu, Bharath Ramsundar, Evan N Feinberg, Joseph Gomes, Caleb Geniesse, Aneesh Pappu, Karl Leswing, and Vijay Pande. Moleculenet: a benchmark for molecular machine learn- ing. Chemical science, 9(2):113-530 2018.
Jun Xia, Chengshuai Zhao, Bozhen Hu, Zhangyang Gao, Cheng Tan, Yue Liu, Siyuan Li, and Stan Li. Mole-bert: Rethinking pre-training graph neural networks for molecules. In The Eleventh International Conference on Learning Representations, 2023.
Minkai Xu, Shitong Luo, Yoshua Bengio, Jian Peng, and Jian Tang. Learning neural generative dynamics for molecular conformation generation. arXiv preprint arXiv:2102.10244 2021a.
Minkai Xu, Wujie Wang, Shitong Luo, Chence Shi, Yoshua Bengio, Rafael Bombz--omaarrlli,, and Jian Tang. An end-to-end framework for molecular conformation generation via bilevel program- ming. International Conference on Machine Learning, pp. 11537 PMLR, 2021b.
Minkai Xu, Lantao Yu, Yang Song, Chence Shi, Stefano Ermon, and Jian Tang. Geodiff: geo-
metric diffusion model for molecular conformation generation. In International Conference Learning Representations, 2021c.
Minkai Xu, Lantao Yu, Yang Song, Chence Shi, Stefano Ermon, and Jian Tang. Geodiff: geo- metric diffusion model for molecular conformation generation. arXiv preprint arXiv:2203.02923, 2022.
Zhao Xu, Youzhi Luo, Xuan Zhang, Xinyi Xu, Yaochen Xie, Meng Liu, Kaleb Dickerson, Cheng Deng, Maho Nakata, and Shuiwang Ji. Molecule3d: A benchmark for predicting 3d geometries from molecular graphs. arXiv preprint I 2021d.
Chengxuan Ying. Tianle Cai, Shengjie Luo, Shuxin Zheng, Guolin Ke, Di He, Yanming Shen, and Tie- Yan Liu. Do transformers really perform badly for graph representation? Advances Neural Information Processing Systems, 34:28877-28888 2021.
Haotian Zhang, Shengming Li, Jintu Zhang, Zhe Wang, Jike Wang, Dejun Jiang. Zhiwen Bian, Yixue Zhang, Yafeng Deng, Jianfei Song, et al. Sdegen: learning to evolve molecular con- formations from thermodynamic noise for conformation generation. Chemical Science, 14(6): 1557-1568, 2023a.
Jiawei Zhang, Haopeng Zhang, Congying Xia, and Li Sun. Graph-bert: Only attention is needed for learning graph representations. arXiv preprint arXiv:2001.05140, 2020.
Ziqiao Zhang, Ailin Xie, Jihong Guan, and Shuigeng Zhou. Molecular property prediction by semantic-invariant contrastive learning. arXiv preprint 1xii:23033006000 2023b.
14

---

Published as a conference paper at ICLR 2024
A RELATED WORK
A. MOLECULE TORENIZATION
In natural language processing (NLP) pipelines (Vaswani et al. 2017 Devlin et al. 2018), tok- enization is the first step. This involves splitting text into individual tokens. These tokens are then converted into numerical values, known as input IDs, that can be embedded by the model. With this in mind, we wondered if we could extract meaningful tokens from molecules and convert them into input IDs for our model. The most straightforward approach is to use atomic types token IDs, but this results in a small and unbalanced dictionary size (only 118 classes of different atoms). This approach may not distinguish dominant and rare atoms in molecules, potentially causing the model to focus too much on primary atoms while ignoring rare atoms, leading to bias (Xia et al. 2023).
To our knowledge, Molc-BERT (Xia et al. 2023 is the first work to address this issuc. which leverages a context-aware tokenizer based on group VQ-VAE (Van Den Oord et al. 2017) that encodes atoms into discrete values with chemical significance. This distinguishes major elements such as carbon into carbon in different chemical environments (e.g., benzene ring and aldehyde carbons.). It greatly addresses the imbalance between major and rare elements while simultaneously expanding the vocabulary size (from 1 118 to 512).
A.2 MOLECULE GEOMETRY PREDICTION (GENERATION)
In recent years, there has been a surge in the application of generative models to the task of gen- erating multiple conformations of molecules. This work primarily utilizes the GEOM-QM9 and GEOM-Drugs datasets (Axelrod & Gomez-Bombaree exist in numerous potential stable low-energy conformations. For instance, CGCF (Xu al., 2021a) employs a Flow model to learn the distribution p p(D ()) of the interatomic distances given molecu- lar graph 9, and then generates potential conformations from p( G 9) and optimizes them using I Langevin Dynamics. Conv VAE (Xu et al., (021b) adopts a bilevel programming approach to split the task into a distance prediction problem and a distance geometry problem. The entire framework encodes the molecular graph into latent space using& conditional variational autoencoder (VAE). During inference, it first samples from the latent space, then obtains sampled samples of interatomic distances through the decoder, and finally obtains the generated conformation by optimizing the distance geometry problem. ConfGF (Shi et al.) (02T) suggests estimating the gradient field of the log density of interatomic distances initially, then calculating the gradient fields of the log density of atomic coordinates via the chain rule. The well-trained score network can estimate the gradi- ent field of atomic coordinates corresponding to different levels of noise. Utilizing this capability, the annealed Langevin dynamics (Song & Ermon 2019) algorithm can be employed to generate new conformations. However, all these methods require modeling the intermediate variable, inter- atomic distances, to obtain the coordinates of the generated conformation. This requirement has been identified as a contributing factor to their suboptimal performance. In an effort to circumvent the constraints posed by intermediate variables, GeoMol (Ganea et al. (02T)) models essential ge- ometric elements of molecules, such as torsion angles, bond distances, and bond angles. During inference, it outputs set of these geometric elements, enabling the comprehensive reconstruction of 3D conformations. Morcover, numerous studies have begun to leverage the increasingly popular diffusion models (Ho et al., 2020). These models aim learn the desired geometric distribution from a noise distribution via a reverse diffusion process. GeoDiff (Xu et al. (002) implements this diffusion and reverse diffusion process directly on atomic coordinates. It then recovers the desired real conformation from positions sampled from the noise distribution through the reverse diffusion process. TorDiff UTig et al. 2022), on the other hand, confines the diffusion process the tor- sion angles of molecular conformations. This approach effectively reduces the dimensionality of the sampling space. A growing body of research has been advancing in accelerating the sampling speed of diffusion models, resulting in significant enhancements in both the precision and speed of molecular conformation generation (Zhang et al. 2023a Fan et al. 2023).
2020). Each molecule in these datasets can
Prediction Methods. While generative models have demonstrated encouraging outcomes in the generation of molecular conformations, they are designed for a one-to-many task, aiming to gen- erate multiple potential stable low-energy conformations rather than necessarily the lowest-energy ground-state conformation. In practical applications, it becomes essential to further sift through the generated results to identify conformations of superior quality and lower energy. Moreover, the inherent sampling uncertainty during the inference phase of generative models could poten- tially lead to the generation of conformations that not practically feasible. Each molecule in
15

---

Published as a conference paper at ICLR 2024
the GEOM (Axelrod & conformations, rather than the most stable ground-state -tate conformation. Given these challenges, a benchmark task (Xu et al. 2021d) has been proposed which seeks to predict molecule's ground- state geometries directly from its 2D topological structure using a predictive paradigm network. A benchmark dataset called Molecule3D has been introduced which comprises nearly 4 million molecules, each with its ground-state geometric information and four quantum properties calculated at the B3LYP/6-31G level using Density Functional Theory (DFT). Building upon this dataset, the DeeperGCN-DAGNN model is employed as baseline by (Xu et al., dh for predicting the ground soound-state geometry of molecules. This involves predicting the interatomic distance matrix cor- responding to the molecule's ground-state conformation and directly predicting the 3D Cartesian coordinates of the molecule's ground-state conformation. Our work follows this latter approach, achieving an end-to-end prediction from a molecule's ' 2D graph to its 3D ground-state conforma- tion.
2020) dataset possesses multiple low-energy stable
A.3 (MOLECULE) GRAPH TRANSFORMER
Transformer (Vaswani et al. 2017) networks have made remarkable strides in diverse domains (De- vlin et al. 2018 Dosovitskiy ct al. 2020; Liu et al. 2021d). In recent years, numerous investiga- tions (Dwivedi & Bresson [020; Ying et al. 2021 Luo et al. 2022 Kim et al. 2022 Min et al.,
2022 have endeavored to employ transformer networks into graphs. Studies have demonstrated that transformer networks can effectively address prevalent challenges such as over- over-smoothing (Chen
jet al. 2020) and over-squashing (Alon & Yahav (020) encountered in GNNs. In molecular model- ing, MAT (to be discussed later) was proposed by (Maziarka et al. 2020) and R-MAT was proposed by (Maziarka al.] 202T) based on MAT. R-MAT calculates attention scores by employing distance embeddings and bond embeddings rather than using the original QK". (Choukroun & Wolf 2021) directly regulates the attention scores by incorporating a learnable gate function that considers the interatomic distances. (Wu et al. 2021a) employs two-dimensional convolution on interatomic dis- tances for position encoding to ensure roto- translatioool invariance, and makes the first attempt to incorporate motifs and knowledge of functional groups into Transformer network for 3D molecular representation learning. For a more comprehensive understanding of GTs, we recommend referring to (Min et al. 2022 Müller et al. 2023). These sources provide detailed overviews on the topic.
Discussion. Because A and D are the most prevalent prior knowledge for 3D molecules, nu- merous previous works have attempted to enhance the attention mechanism in Graph Trans- former networks by leveraging these inductive biases. The works most closely related to ours (aaziarka al. 2020, Choukroun & Wolf 202T). (Maziarka et al. 2020) directly incorporates K A and g(D) throughout the model by means of three hyperparameters, where
sof so tmaa vi
is kernel function that can either be softmax e-d. (Choukroun & Wolf 2021 utilizes f(D-1)
Q.KT") where f is a learnable fully-connected network. The
as gate function acting on so softmaz( /dd, crucial distinction between our work and theirs lies in the fact that we do not directly incorporate softmaz( Q.K[[ vdd ), A and D or impose distance-based gating constraints on so softmar( Q.K"" ). In- Vd, stead, we regard S = Q as a global attention and dynamically acquire the weights of the three masks E,,, A, and Zmax D (corresponding to global, nearby and spatial information) for each self-attention head through learnable parameters (1, nt and 7P)) We then adjust S by merging the three masks to obtain the overall mask M = E, + + x A + 7 (gmaxx D), and subsequently apply softmax to the adjusted S to obtain the final attention weights for updating node features as vdd Vi. Unlike (Maziarka et al. 2020), we are able automatically learn dif- ferent mask weights at every location in the model, thereby capturing distinct attention patterns. Unlike (Choukroun & Wolf (2011) we do not directly manipulate S based on distance information, rather we adjust S using the learned mask M. Furthermore, we introduce fewer learnable parame- ters. Our approach allows the model to learn different attention patterns in different self-attention heads, with some parts focusing more on global information and others on local or spatial informa- tion, thereby enhancing the expressive power of the model. A control experiment has been added to To to differenniate the variations between MAT (Maziarka et al. 2020) and our approach.
O = softmax( SOM
B MSRSA A FOR MOLECULAR PROPERTY PREDICTION
Setup. To assess the effectiveness and transferability of the MSRSA module (Sec. 3.2). We have implemented three Bert-like (encoder only) models (Devlin et al. 2018) of different sizes (small,
16

---

Published as a conference paper at ICLR 2024
base and large) with complete MSRSA modules, including both B and B Detailed model hy- perparameters are provided Table [ in Sec. [ Initially, an ablation study is conducted on our small-level model, utilizing the random split of the Molecule3D dataset. This is done validate the effectiveness of each individual component within the MSRSA when predicting molecular proper- ties. The results of these experiments are detailed in Table 4 Subsequently, a performance compar- ison was carried out with several baselines on the Molecule3D Random split for the prediction of the Homo-Lumo Gap. All baseline results are derived from (Wagg al. 2022). The results pre- sented in Table Finally, we conduct experiments on the molecular property prediction tasks using the QM9 dataset. The results of previous methods all borrowed from (Liao Smidt 2022). We use the same data split as in (Liao & Smidi (022), with 110k, 10k, and 11k molecules for train- ing, validation, and testing, respectively. Our results are based on the network that was trained in
Table 4: Ablation study of MSRSA module for molecular property prediction Molecule3D ran- dom split (eV).
Methods
Index (AParam)
MHSA
B
B (unigina)
(ouesaa)
EHOMO
ELUMO
A=
(+0)
.209
.239
2 (+0)
.111
.108
(448)
.092
.089
.104
(+96)
.034
.034
.053
(+96)
.031
.034
.045
MAT Maazirraa et al. 2020
.067
.070
.083
Table 5: MAE results of MSRSA module for molecular HOMO-LUMO gap prediction on Molecule3D random split (eV).
Time
Method
Ae
Train
Inference
GIN-Virtual (Hu al. 2021
15mmn
2min
.1036
SchNet (Schült
2017
15min
3min
.0428
DimeNet++ Gastegger et a
2020
133min
16min
.0306
SphereNet [Luu ct al. 2021c
182min
28min
.0301
ComENet Wang el 2022
22min
3min
.0326
GTMGC (Ours)
20min
3min
.0466
GTMGCpasc (Ours)
21min
3min
.0371
GTMGC (Ours)
25min
4min
.0316
Table 6: MAE results of MSRSA module for molecular property prediction QM9 test set.
Method
A (me
SHомO (mcV)
ELUMO (mc
C₁
(a
(D)
(cal/mol K)
NMP Gilmer al. 2017
.092
43
.030
.040
SchNet (Schütt ot al. 2017
.235
41
34
.033
.033
Cormorant (nderson
2019
.085
61
34
38
.038
.026
LicConv Finzi ct al.
20200
.084
49
30
25
.032
.038
Dete (Gatterger al. 2020
.044
33
25
20
.030
.023
TFN Thomas et 2018
.223
40
.064
.001
Fuchs et al. 2020
.142
53
35
33
.051
.054
EGNN Satorras et al. 2021
.071
29
PaiNN Schült ci al.
X12T
.045
46
28
20
.012
.024
TorchMDD-NET Tholke & De Fabritiis
2021
.059
36
20
18
.011
.026
SphereNet (Liu al. 2021c
.046
23
.026
.021
SEGNN Brandstetter et al.
2021
.060
42
24
21
.023
.031
EQGAT Le 2022
.053
32
20
16
.011
.024
Equiformer [Liao & Smidt
2022
.046
15
.011
.023
GTMGC (Ours)
.131
47
37
32
.055
.050
GTMGC (Ours)
.128
33
.052
.047
GTMGCtarge (Ours)
.117
39
29
26
.043
.043
denotes using different training, validation, testing data partitions.
17

---

Published as a conference paper at ICLR 2024
the previous experiment, with further simple fine-tuning performed the Qm9 dataset. The Mean Absolute Error (MAE) between the prediction and ground truth are shown in Table
Results and Analysis. (2). Table 4 robustly showcases the efficacy of cach constituent within MSRSA. Notably, the gradual incorporation of components B (Index 3) and B D(original) (Index 4), built on the foundation of pure self-attention (Index 2), has led to a stepwise enhancement in per- formance by 20.61 % and 59.54% respectively, yet the amount of additional parameters introduced nearly 0%. (iz). In specific advancement, the substitution of the original interatomic distance matrix D with innovative proposal D row-sub (Index 4), has further boosted the effectiveness by an additional 15.09%. This underlines that D Diw-sub offers superior assistance over in modeling the spatial configuration of molecules. (iii). As illustrated in Table [ our methodology delivers remarkable performance on large-scale molecules. Even though the large-level model encompasses significant number of parameters, the speed of inference is still notably satisfactory. This under- lines the considerable scalability of our approach when dealing with large molecules. (iv). The results presented in Table indicate that our method attains an above-average performance in pre- dicting three properties: Ac, CHоMO and ELUMO. Despite the relatively lower prediction results for other three properties, they remain competitive when compared numerous baseline methods. (v). All these results suggest that our proposed MSRSA module can be seamlessly applied to diverse molccular representation learning tasks and is characterized by its simplicity of implementation.
C MORE EXAMPLES OF VISUALIZATIONS
C.1 INFERENCE EXAMPLES VISUALIZATION
(b), RONE ETKDG
(B.02)
(1.332)
(0.115)
(0.155)
(1,356)
(1,44)
(0S51)
(U,533)
(0.550)
(0.533)
(0.830)
(0.585)
(0936)
(1.324)
(0.05))
(Chrs)
(B (0,090)
(0 03953
(0.529)
(1.288)
(0,005)
(06665)
60 240
(B (0.197)
C.2 ATTENTION WEIGHTS VISUALIZATION
0.0
200
0.8
Motecute
Head
Head
Head
O
0.7
520
0.00
0.6
0.2
12A
18

---

Published as a conference paper at ICLR 2024
Motecuie
Head
0.24
0.75
0.38
6122
0.00
0.55
C0TS
0.33
0000
0.02
Malecule 20 Graph
Head
Head
Read
0.150
6.125
0.15
0.05
5,100
5.000
Molccule 20 Groph
Head
Hoad
Hozd
E.175
0.175
6.125
5:100
0.4
E.05
CAGS
C.3 THE LEARNED "i' AND of
The values of y and 7p which we learned and exported from GTMGCsmali are illustrated the figure below. In the experiment, only 7/a and 7P are parameterized learning, while the weight of global information kept constant at The final weights 7 nt. and 7p are determined by normalizing the values of 2pt and 7pp
YR
block 3
Y/?
block
YI
Y!
44
0.3
block
block
02
01
D TRAINING DETAILS
MoleBERT Tokenizer Training. Since the MoleBERT tokenizer (Xia et al. 2023) is the first to tokenize atoms into set of chemically meaningful discrete tokens, should be re re-trained different datasets. However, once the tokenizer is trained, the tokenized results on the same dataset be used for different tasks without changing. is important note that the MoleBERT tokenizer
19

---

Published as a conference paper at ICLR 2024
Table Model hyperparameters for GTMGC on Molecule3D
Ground-state Conformation Prediction
Propertics Prediction
Encoder
Decoder
Small
Base
Lure
Epoch
20
60
Batchsize
100
100
Optimizer
Adamw
Adamw
Learning Rate
5e-5
5c-5
Warm Up
0.1
0.3
Lr Scheduler
linear
cosine
Num Transformer Blocks
6
6
8
Num Attention Heads
8
16
32
d_model
256
256
512
512
d_ffn
1024
768
2048
Dims in Prediction Head Use B Use B
[256, 768.
[256, 768, 1J
[512, 1536.
[512, 1536. 1]
Trainable Parameters
9,808,275 (whole model)
4,983,393
11,925,185
26, 10,,99
based on a small GNN model and VQ- VAE framework. We only use the tokenized tokens to train tasks without updating the tokenizer's 's parameters. This is different from some GA-style (GNN layer before, Transformer layer after) approaches (Rong et al. 2020; Wu et al. 2021b Mialon et al. 202T) that update both the GNN and Transformer layers' parameters simultaneously.
GTMGC Training. Table presents the crucial model and training hyperparameters for GT- MGC. For our primary task predicting the molecular ground-state conformation, both the encoder and decoder of GTMGC are configured with transformer blocks and 8 attention heads. Notably, we set our modee to 256, difn to 1024, and dhadden in the prediction head 768, resulting in lean model with only 9M parameters. During the training phase, we employ the AdamW optimizer with a learning rate of 5e-5 and batch size 100. We initially warm up the learning rate from 0 to Se-5, followed by a linear decay to O, over a total of 20 epochs. For our auxiliary task of molecular proper- ties prediction, we have constructed three versions of GTMGC: small, base, and large. The specific hyperparameters for these three versions training strategy across all three versions of the model, but increase the number of training epochs to 60, with larger learning rate warm up of 0.3 and cosine decay strategy. Our open-sourced code provides further details about the training process.
detailed in Table 7 For this task,
apply the same
20


Reviews:
[{"review_id": "0HwdezqDf1", "reviewer": "Reviewer_BbTi", "summary": "This paper introduces a novel graph transformer specifically designed for 3D ground state prediction. Notably, the proposed architecture is versatile and applicable to a wide range of 3D supervised tasks. The key contributions of this work include:\n\nA novel architectural proposal that elegantly extends the classical attention mechanism to 3D molecular graphs. This extension incorporates edge and interatomic distances as biases for the attention mechanism, enhancing its capabilities.\n\nA successful demonstration of the effectiveness of this architecture in the realm of 3D ground state prediction, as well as its application to predict various other 3D molecular properties.", "questions": "Was molecular property prediction approached as a single-task or multi-task endeavor?\n\nCould you clarify the rationale behind placing the molecular property prediction results in the appendix, especially considering that they do not outperform SOTA across the board?\n\nIt could be valuable to assess the scalability of your architecture across various graph sizes, thereby determining where it potentially outperforms existing methods.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Originality:\nDespite numerous unsuccessful attempts to construct graph transformers using principles akin to the original transformer, this work stands out as it elegantly achieves the intended goal with minimal architectural complexity. This work avoids unnecessary biases that often detract from the model's effectiveness, making it easier for most researchers to apply their existing intuitions from sequence transformers to this novel architecture.\n\nQuality:\nThe architectural design and its application in ground state conformation prediction are well executed, as reflected in the results, solidifying its position as a favorable solution compared to alternative methods. The iterative refinement of $G_{cache}$ within the decoder represents a notable innovation that enhances model performance in conformer prediction. Ablation studies further clarify the significance of each component within the network, facilitating an understanding of their contributions to this specific modeling task.\n\nClarity:\nThe paper is well-written and maintains a high level of clarity, making it easily comprehensible for readers.\n\nSignificance:\nWhile predicting the ground state of a molecule remains relatively underexplored due to its limited relevance in specific applications, this work serves as a foundational step that can be extended to tackle the broader challenge of full conformer generation. Such an extension holds are very significant, especially in the context of drug discovery.", "weaknesses": "The assertion regarding the innovative utilization of the MoleBERT Tokenizer might be overstated, especially in light of the results presented in Table 2. Previous molecular graph papers, such as the MolGPS paper, have explored various atomic featurizations that could potentially outperform the approach presented in this work.\n\nTo allocate more space for related works and experiments, it would be beneficial to consider shortening or omitting certain sections, such as those in the introduction (implementation) and preliminary sections (multi-head and transformer). \n\nThe related work should be integrated into the main text rather than relegated to the appendix. It is crucial to comprehensively cover the various attempts to construct graph transformers and elucidate why they are ill-suited for the tasks at hand.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "xxe0kbsP5M", "reviewer": "Reviewer_Po8N", "summary": "This work proposes GTMGC, a graph transformer model for ground-state molecular conformation prediction. GTMGC uses a novel self-attention module to achieve effective molecular structure modeling. Experiments show that the proposed GTMGC model achieves state-of-the-art performance in ground-state molecular conformation prediction benchmarks.", "questions": "No additional questions.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Originality: The proposed graph transformer model is novel, with many novel technical contributions in effectively capturing spatial structures by self-attention mechanism.  \nQuality: The effectiveness has been effectively demonstrated by experiments.  \nClarify: The writing and presentation of this paper is good and well-organized.  \nSignificance: The contribution of this work is very useful and meaningful to chemical and molecular biological science fields as the proposed method can significantly accelerate the computation of finding ground-state molecular conformations.", "weaknesses": "There are actually many prior studies about formulating the mapping from 2D molecular graphs to 3D molecular conformations as a generative problem. Though they are different from the problem studied in this work, these models can be trained on the used Molecule3D datasets and evaluated by generating only one molecular conformation. However, authors do not compare with any of these methods. Authors are recommended to compare with at least one molecular conformation generation method, such as [1].\n\n[1] Torsional Diffusion for Molecular Conformer Generation. NeurIPS 2022.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "lqEhvzchy1", "reviewer": "Reviewer_1WnG", "summary": "The paper proposes GTMGC, a transformer-based architecture for end-to-end prediction of 3D groundstate molecules' conformations from their 2D graph.\n\nThe method makes use of MoleBERT for initial embedding as well as LPE for positional encoding.\nThe inputs are then processed by a Transformer-based model, where the self-attention modules are augmented with the adjacency matrix and learned/predicted atomic distance matrix in a weighted sum fashion. \n\nThe loss is augmented by a regularization of the middle distance matrix prediction.,", "questions": "Currently, it seems the proposed contributions don't bring any advantage over a simple (large) transformer model.\n\n1. One needs to know the capacity of the model in order to assess the origin of the good performance. \nAccording to weakness 2, it seems the good performances are obtained *almost solely* from a large/powerful standard Transformer model. From Table 5, the model seems much bigger than other methods. Also, the discrepancies between the tables are disturbing (one single metric should be used for the model validation).\n\n2. Ablating the initial LPE. \n\n3. It would be beneficial to have a comparison performance with (at least one) other molecule transformer-based methods such as [4,5,6] or others at a similar capacity.\n\n4. It would be interesting to look at the learned weighting parameters (\\gamma) of the self-attention to better understand the contribution of each (maybe even adding a weighting to the global attention).\n\n\n[4] Relative molecule self-attention transformer.\n\n[5] 3dtransformer: Molecular representation with transformer in 3d space.\n\n[5] Geometric transformer for end-to-end molecule properties prediction.", "limitations": "", "rating": 3, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The paper is very clear and pleasant to read.\n2. The integration of various existing frameworks for the task of molecules' ground state conformation is original to some extent.\n3. The performances compared to the baselines are significant.", "weaknesses": "I believe there are two main weaknesses:\n\n1. Novelty: There are plenty of previous works (not cited) implementing Transformer architectures that implement the elements of the proposed self-attention, e.g., [1] uses the adjacency matrix, [2] makes use of the distance matrix (even cited in the manuscript), while the weighted summation is ubiquitous in many fields using Transformers [3]. The initial encoding is obviously not a technical contribution.\n\nThus, the contribution may be summarized as the integration of existing methods/approaches, augmented with the regularization loss on the distance matrix (the \\beta seems to bring very minor improvement).\n\n2. Strange results: I may be wrong, but according to Table 3 (ablation study) a *simple Transformer* architecture without any addition to the self-attention reaches 0.4395 (MAE) which is already better than all the other baselines. This is problematic. Also, the final improvement is only ~1%. Finally, the reported results seem to take the best-ablated model results for each metric, which is wrong.\n\nMoreover,  we have:\n\n3. The advantage of MoleBert over the standard atomic encoding is extremely shallow or even worse.\n4. Lack of comparison with other Transformer based methods.\n\n[1] *Molecule attention transformer.*\n\n[2] *Geometric transformer for end-to-end molecule properties prediction.*\n\n[3] *Axial-DeepLab*", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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:

Decision: Accept (spotlight)