Geometric Transformer with Interatomic Positional Encoding
Authors: Yusong Wang
Shaoning Li
Tong Wang
Bin Shao
Nanning Zheng
Tie-Yan Liu
Venue: NeurIPS 2023 poster (2023)
Keywords: Geometric Deep Learning
Molecular Modeling
Positional Encoding

Abstract: The widespread adoption of Transformer architectures in various data modalities has opened new avenues for the applications in molecular modeling. 
Nevertheless, it remains elusive that whether the Transformer-based architecture can do molecular modeling as good as equivariant GNNs. 
In this paper, by designing Interatomic Positional Encoding (IPE) that
parameterizes atomic environments as Transformer's positional encodings,
we propose Geoformer, a novel geometric Transformer 
to effectively model molecular structures for various molecular property prediction. 
We evaluate Geoformer on several benchmarks, including the QM9 dataset and the recently proposed Molecule3D dataset. 
Compared with both Transformers and equivariant GNN models, Geoformer outperforms the state-of-the-art (SoTA) algorithms on QM9, and achieves the best performance on Molecule3D for both random and scaffold splits.
By introducing IPE, 
Geoformer paves the way for molecular geometric modeling based on Transformer architecture.
Codes are available at https://github.com/microsoft/AI2BMD/tree/Geoformer.

Geometric Transformer with Interatomic Positional Encoding
Yusong Wangl Shaoning Li2.3,4vf Tong Wang?; Bin Shao
Nanning Zheng', Tie-Yan Liu
National Key Laboratory of Human-Machine Hybrid Augmented Intelligence, National Engineering Research Center for Visual Information and Applications, and Institute of Artificial Intelligence and Robotics, Xi an Jiaotong University
Microsoft Research A14Science
Mila Québec AI Institute Université de Montréal
wangyusong22000ssuuu xjju. edu. cn, unzheng@mail xjtu. edu. cn shaoning.. 110umontreal.co {watong, binshao, tyliu)@microsoff...
Abstract
The widespread adoption of Transformer architectures in various data modalities has opened avenues for the applications in molecular modeling. Neverthe- less, it remains elusive that whether the Transformer based architecture can do molecular modeling as good equivariant GNNs. In this paper, by designing Interatomic Positional Encoding (IPE) that parameterizes atomic environments as Transformer's positional encodings, we propose Geoformer, a novel geometric Transformer to effectively model molecular structures for various molecular prop- erty prediction. We evaluate Geoformer several benchmarks, including the QM9 dataset and the recently proposed Molecule3D dataset. Compared with both Trans- formers and equivariant GNN models, Geoformer outperforms the state-of-the-art (SoTA) algorithms on QM9, and achieves the best performance on Molecule3D for both random and scaffold splits. By introducing IPE, Geoformer paves the way for molecular geometric modeling based on Transformer architecture. Codes are available https///chhhooooooocooooooA//old
Introduction
Transformer 44 has been a dominant architecture in modeling various data modalities such as natural language, images, and videos. As such, it is natural to generalize Transformers in molecule modeling. Molecules are represented as either 2D topology structures or 3D geometric structures. Prevailing algorithms for modeling 2D topology (22 attention, which treat the molecular structures as fully connected graphs and devise a diverse range of positional encoding schemes. In contrast topological graphs, geometric structures offer more comprehensive description of molecules by providing the information of 3D coordinates. There have been some recent attempts 23
26 20
employ Transformers with global
to integrate specific geometric features into Transformers. For example, Transformer-M 2 considers pairwise distances as learnable bias added to the attention weights supplement the 2D topology encoding used in Graphormer [2]. Nevertheless, they are insufficient to comprehensively encompass the intricacies of the entire 3D molecular space, as they solely rely on distance information for positional encoding, In
36
55
50
Work done during internship at Microsoft Research.
1 Equal contribution.
i *Correponndigg author.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

contrast, modern GNNs emphasize the significance of equivariance as crucial bias, and explore various methods to encode geometric information, including distances, angles, and dihedral angles 31 2 further utilize high-order geometric tensors to yield internal features models, ensuring equivariance with respect to the E(3)/SB(3) group. As a result, the performance of EGNNs in molecular property prediction significantly surpasses that of methods based on Transformers.
13
39
26
47 Several works 3
The success of EGNNs underscores the advantage of integrating directional geometric information into neural networks for molecular modeling [2 26 4r
while employing such information in the Transformer based architecture has yet to be developed. Intuitively, all geometric information is embedded in atomic coordinates, which can be naturally utilized as a bias for developing an effective positional encoding in Transformers. Based on atomic coordinates, atomic cluster expansion (ACE) a complete descriptor to represent the environment of centered atoms In this study, we first design interatomic positional encoding (IPE) by introducing cluster merging based on ACE theory. By incorporating IPE into traditional Transformers, we extend the capabilities of the Transformer, in terms of Geoformer, to effectively model molecular structures for molecular property prediction. We conduct a comprehensive evaluation of Geoformer using several benchmarks, which includes QM9 dataset
theory 1T0 21
K5 comprising 12 molecular properties and the recently proposed Molecule3D dataset [51], containing 3,899,60 ,647 molecules sourced from PubChemQC[B21 Our results demonstrate that Geoformer surpasses state-of-thes the-art algorithms on the majority of properties on QM9 dataset, and achieves the lowest mean absolute errors on Molecule3D for both random and scaffold splits. We also provide visualizations of the learned IPE, which shows that IPE can capture different positional information compared with PE that only encodes pairwise distances.
Our contributions can be summarized as follows:
We introduce a novel positional encoding method, i.e.. Interatomic Positional Encoding (IPE) to parameterize atomic environments in Transformer.
By incorporating IPE, we propose Geometric Transformer, in terms of Geoformer, which models valuable geometric information beyond pairwise distances for Transformer-based architecture.
The proposed Geoformer achieves superior performance on molecular property predictions compared with Transformers and EGNNs.
2 Preliminary
The Atomic Cluster Expansion (ACE) 10 is a complete descriptor of the local atomic chemical environments, represented by hierarchical many-body expansion. The key components of ACE are: a) ACE defines a complete set of basis functions for the environment of the centered atom (radial basis functions and spherical harmonics in practice); b) ACE significantly reduces the computational efforts for body order computing to linear time complexity scaling with the number of atoms within molecule. These advantages serve ACE as accurate, fast, and transferable theory framework for molecular modeling. In order to facilitate readers' comprehension of our interatomic positional encoding, we would present the core operations of ACE in several equations.
ACE includes a set of orthogonal basis functions (fis)) to describe the spatial relations between two atoms. f.j denotes the relative position pointing from atom i to atom j. and indicates the functions' polynomial degree. ACE focuses on modeling the potential energy of a system by focusing on collection of atoms, specifically atomic clusters. In this method, the centered atom, denoted as i, is surrounded by K neighboring atoms, which form the atomic cluster. The potential energy of the atomic cluster depends on the hierarchical interactions between the central atom i and its K neighboring atoms, known as many-body expansion. The expansion of atomic potential energy could be written by:
E, = E--E Go der (Pts)) JI 21
ΣΣ Goivz (fez) Des (frzz) +
(I)
J1j2 VL
with an unrestricted summation. G, indicates the expansion coefficients. However, the sum of surrounding neighbors within the cluster would scale to O(N*) with K neighbors and become

---

b
Update
Scale
my
Mask (opt.)
B Bi(().
Bij.v
8y0*
MatMul
Bjiv
B Byuv Byu
Cluster
Linear
Lineer
Lincar
Cluster
Interatomic Positional Encoding (IPE) Cy
Merged oj
Merge
Figure 1: Illustration of Interatomic Positional Encoding (IPE) Cyp Panel depicts the extended Self-Attention with Cy. C contributes to the construction of attention weights while simultaneously being updated by atomic features; Panel b highlights the relationship between C.,, cluster Ø; cluster Jj. and the merged cluster Jij. as described in Theorem and Equation to Equation
numerically expensive. N denotes the number of atoms within molecule. ACE leverages the density trick to reduce the computational overhead. It defines atomic base Ai, and A-basis A as:
A
Σ (f(g)) JEN(()
(2)
Aiv II Ai, 1=1
== (UI,...,
(3)
where denotes the order of body expansion, i.c., (e + 1)-body expansion, and stands for the set of By firstly summing the neighboring basis functions and then applying multiplication, ACE can efficiently represent the atomic expansion scaling with the complexity O(N). Since A, is rotationally invariant, we need additional Rebbcch-Gordas coefficients C. to construct fully permutation and isometry-invariant basis functions (B-basis), and describe the potential of cluster o; Equation M as their linear combination with c-coefficients 21
Biv
Ecrr
(4)
E, = Σ Civ = =c,
(5)
Methods
In this section, we would introduce our Interatomic Positional Encoding (IPE) based on ACE theory
Section 3.1 and discuss how we integrate IPE into the Transformer architecture Section 3.2
3.1
Positional Encoding for Geometric Molecule Modeling
In the context of geometric molecule modeling, the positions of atoms are the most intuitive positional information encode in Transformers (termed as "positional encoding"). While most recent works have adopted pairwise distances between atoms as the relative PE, such representation is often inadequate for capturing the complex interactions within molecules. As a result, it essential to use more comprehensive and appropriate PE for geometric data. Motivated by ACE, we propose an Interatomic Positional Encoding (IPE) to efficiently describe the many-body contributions in Transformers for geometric molecule modeling. The distinction from ACE is that IPE further takes the interactions between atomic clusters into account. More details on the IPE method and its integration into the Transformer architecture will be provided in the subsequent sections.

---

Theorem Given two cluster and o and their basis functions, there exists set of invariant basis functions for the merged cluster oij to describe integrated cluster potentials Esj.
Proof. A merged cluster oij can be represented as translating two clusters og and Øj such that their centered-atoms, i.e., atoms and overlap shown in Fig. (b). All their neighbors merged into a new cluster ijj and the newly formed atomic base can be expressed as:
Aij.0 (Ai, Aju)) @ (Ai, Az,0)
Ai,dAio Aj,oAi
AivAjo
AjuAjo)
(6)
E Emima Do (o
Emn
Em ( (fim)
Enins Gu(fins )o((fjn))
with the product explicitly writing out. It are the neighbor atom symbols of atoms i, respectively, is the tensor product. We write the atomic base Ai, as Agj (due the permutational
invariance, we have Aij. = Afi,)) which still follows the density trick. Therefore, we could construct new A-basis for merged cluster oi when taking product of Ago::
Âgg.v = Aijan O Aij.rz O .... 2 Asjon
(Ar, AJ,) 2 (Ai, Ay.))
I1 Aiv
Il² Aijm
II Agi.n
I
27
Ajm
(Asi).v
Agj.v Az().v)
Aji,y
where 21 and 1) 1,2,.... A.j.v could describe (c (c+ 11-boyy and (e 2)-body expansion simultaneously O(N). To be concrete, Ast)),v contributes the (c+1)) expansion, while Aij, contributes the (e + 2bbbd expansion due to cluster merging. For instance, when taking v = 1, A1, and As(o,v denote 2-body (im) and 3-body expansion (im1m²) in original cluster ø;, and A.,.v denotes the 4-body expansion (mijn) merged cluster osj. respectively. Further explanation can be found in the following Remark. Aij.v exists when 2 2, which implies considering at least 4-body expansion within molecules. Then we could construct the corresponding matrix of B-basis Beg.v following Equation and represent the potential of cluster oij as:
Esg Σ
Cij. Bij.v
(8)
Remark. A straightforward illustration can be drawn by setting v = (Cartesian space) and 1. In this case, basis Buy) could interpreted as the sum of cosine value of the surrounding angles within cluster o; 39 [3 i.c., Em cos om which represents the 3-body contributions in O(N). Similarly, basis B:y could be treated as the sum of proper dihedral angles between two clusters and of 499. i.e., Σ mm COS Pmijn. which represents 4-body contributions in O(N). The basis B3 indeed serves as the contribution for torsion potential between two clusters the original ACF, which primarily considers the contributions within single cluster. The detailed proof could be found in Appendix As a result, incorporating Aú effectively enhances the representation of interatomic relations by capturing the interactions between clusters, thereby offering possible way for designing a geometric Transformer architecture.
Theorem 2 (Interactomic Positional Encoding (IPE)) Given molecule with N atoms, there exists a positional encoding matrix C, XN,, which naturally describes the interatomic potentials. n denotes the orders of body expansion in Equation / In particular, C, is directly multiplied with Query and Key before scaling. e.g., softmax, serving as the positional encoding in Transformer:
a (XW))(XXWk) C,,
where E X denotes the atomic features and W E FxF denotes the learnable matrix. F is the hidden dimension. Q, K. represent Query and Key, respectively. @ is the Hadamard product.

---

X
Decoder
LayerNorm
Type
Coordinates
Attention Block
Attention Block
Layorhorm
Encoder
XL
LayerNorm
X
Cy
(c) Embedding Layer
Attention Block
Self-Attention
X
x
Lineer
S1
Linear
priperty
Embedding
Ch
(a) Geoformer Architecture
(b) Attention Block with IPE
(d) Decoder
Geoformer Architecture. Geoformer consists of (c) Embedding layer; (b) elaborate Encoder incorporating L attention blocks with IPE C to extract geometric features and capture complex interatomic relationships within the molecular structure. The extended Self-Attention module with C, is depicted in Fig. [ (d) a lightweight Decoder for predicting molecular properties of interest, such as energy and HOMO-LUMO gap.
Figure 2:
Proof. As mentioned in Section 2 and Theorem the basis functions A,v in Equation 3 could represent the local chemical environment of cluster a, with density trick. Inspired by this, we first construct the Av, [A1.v- A2.v..... AN,v. with integer т = [1, n] and A II A We then treat Av, as the absolute positional encoding attached to the Query and Key, and attention matrix a before scaling, e.g., softmax, could be modified as follows:
0r (XWQ Av,)(XW, Av.) (XWx) & AT.
(10)
(XWx) & E. Ave
where o is the entry-wise Kronecker Product. The detailed proof could be found in Appendix Instructed by [10 21], we add Clebsch-Gordee Gordan coefficients Cv, to ensure the rotationally invariance, and therefore obtain B -vv; Cwv, Ave Then we can write the linear expansion:
WBBv.) a (XWQ) & (E"_WBBu,
(11)
where Wa is a learnable weight matrix denoting c-coefficients, which is similar to the message construction in MACE 22) Such operation could be done by tensor broadcast. Eventually, our interatomic positional encoding C could be represented as:
==
W WgB.,
(12)
We further modify the Equation |
Equation 9 and complete the proof.
3.2 Geoformer: Geometric Transformer for Molecules
Overall Design. The comprehensive structure of Geoformer depicted in Figure 2 This design comprises an Embedding layer, an elaborate Encoder for extracting geometric features, and lightweight Decoder for predicting molecular properties. Within the Geoformer's 's Encoder, attention blocks are integrated, each having hidden dimension F and employing the proposed interatomic positional encoding C,,. Notably, C, is updated by the atomic features X R xF' within each block. Geoformer takes the atom type Z E R and atomic coordinates R E Nxx from one molecule with N atoms as inputs, and produces the corresponding molecular properties outputs.

---

Embedding Layer. The Embedding layer maps the atom type Z to X ERN:
x LayerNorm (embed
(13)
and initialize the IPE C. E RNNNNF with 2-body expansion, i.e., radial basis functions (RBF) 45
-3h (exp (-11kil) p4k) ))
gx(R) ([] exp
(14)
- g())
(15)
where R RNNNN3 denotes the relative position between two atoms, ||-I| is the vector norm, SR. Pa are optional learnable parameters that specify center and width of ge (R), and (b()) is a smooth cosine cutoff function. g R) = [ga .).....,g.. (R)]T E N K is composed of the values of K radial basis functions. W RBF E R
learnable matrix mapping basis functions to the hidden size.
Attention block with IPE. The extended attention block is illustrated in Fig. aa). In contrast the traditional Transformer, learnable IPE matrix C, is introduced to each attention block. Following TorchMD-NET + the softmax function is substituted with the SiL.U activation to enhance accuracy, and the attention weight is scaled by smooth cutoff::
A(X') = SiLU ΣE, (((X'W)) * (X'W)) O c.)) ((1
(16)
where denotes the batched tensor product, i.e., (X'W)) (X' E indicates [-th attention block. In Fig. [ the mask operation corresponds to the implementation of an alternative attention mask. It effectively filters out atoms with excessive distance, concentrates attention and improves the performance of the attention mechanism Then we produce the weighted values per atoms after self-attention to update C²:
Attny(X') A(X') @ X'W,
(17)
Then under the instruction of Equation
7 we construct the A-basis for all merged cluster:
A I1
Attny Attn Yi- ml ( (t/1] R())
(18)
1=1 5
E (A W (0n) (A,, W
AL A )
(19)
where W Attn is learnable matrix, V--mm (R/|[[[[)) denotes the spherical harmonics with order 1* and W are two learnable matrix without bias to ensure equivariance (39
and degree m Wi AlD A& We could further construct new form of residual IPE within each block following Equation g and
=
W B Ev, Crv, Av
(20)
where Wi is a learnable matrix. Finally apply the residual connection compute IPE for the next block:
B
c = SiLU (C,,')) @ sc, + + where WL. is learnable matrix and SILUCC,,W, plays role of gated filter. The update of atomic features still follows the traditional procedures:
(21)
Attn(X')= = E, Attny + x
(22)
FFN(X') SiLU(X' +
(23)
where E, denotes the sum of atomic values weighted by the attention score. W and W/2 are two learnable matrix in ferward layer. The output of each step fed into Layer Normalization (LayerNorm). It is important to emphasize that while the theoretical derivations of the Geoformer architecture above may appear complex, practice, the model utilizes simplified settings to reduce computational complexity. Specifically, we employ 7 1 for body expansion and 1* 1 for spherical harmonies. It streamlines the complex tensor contraction and Tiebsch-GGrddan product caleulations, making the model more efficient and casier to implement. Despite this, Geoformer achieves or

---

surpasses state-of -ae---f-thhhaaaaa the-art art prediction results on several benchmark datasets, as demonstrated in the subsequent section.
Decoder: The Geoformer utilizes a lightweight Decoder, as depicted in Fig. 2 d). consists of a two linear layers with SiLU activation and aggregation module > predict the specific molecular property. The lightweight Decoder ensures that the Geoformer remains computationally efficient while maintaining its ability accurately predict molecular properties based the geometric features captured by the Encoder. More details on Decoder for specific properties are shown in Appendix
4 Experiments
4.1 Experimental Setup
Geoformer is evaluated on both QM9 dataset 351 that consists of 12 molecular properties and large-scale Molecule3D dataset [51] derived from PubChemQC [2] with ground-state structures and the corresponding properties calculated at DFT level. All results measured by mean absolute eлror (MAE) test sets and baseline results are directly taken from the corresponding papers. All models are trained using the AdamW optimizer, and we use the learning rate decay if the validation loss stops decreasing. We also adopt the early stopping strategy to prevent over-fitting, The optimal hyperparameters such as learning rate and batch size are selected on validation sets. More detailed hyperparameters setting for Geoformer are provided in Appendix Table
4.2 QM9
Table I: Mean absolute errors (MAE) of 12 kinds of molecular properties on QM9 compared with state-of-the-art algorithms. The best one each category is highlighted bold.
Target
J
Q mag
EHGMO meV
ELUMO meV
Ae
(R²)
ZPVE E me
U
H
G
C2
Unit
mD
meV
mag
me
me
me
meV
meal
SchNet
15
33
235
41
34
63
73
1.70
14
19
14
14
33
Cormorant
38
85
2.03
20
LieConv |
32
84
30
25
49
800
2.28
19
19
24
22
38
DimeNet++
[K]
30
44
20
1.21
6.32
6.28
6.53
7.56
23
EGNN 17
29
71
25
106
1.55
11
12
31
PaiNN 39
12
45
28
20
46
66
1.28
5.85
5.83
5
7.35
24
TorchMD-NET
122
1.84
7.62
GNS NoisyNode
25
52
20
19
29
700
1.16
7,30
7.57
7.43
8.30
25
SphereNet
25
45
268
6.26
7.78
SEGNN 59
23
60
24
21
42
660
1.62
15
13
16
15
31
EQGAT
PaxNet
11
45
23
19
31
249
1.17
5.90
5.92
6.04
7.14
23
ComENet
25
45
20
32
1.20
6.59
6.82
6.86
7.98
24
Equiformer
A1
11
46
15
14
30
251
1,26
6.59
6.74
6.63
7.63
23
AMP
12
67
4.10
11.3
11.4
11.3
12.4
32
Molformer
28
2.05
7.52
GeoT
22
29.7
52.7
25.0
20.2
43.9
300,8
1.73
11.1
11.7
11.3
11.7
27.6
Geometric Transformer
26.4
27.5
20.4
157
1.24
7.35
8.21
28.0
Transformer-M 28
37
41
17.5
16.2
27.4
75
1.18
9.37
9.41
9.39
9.63
22
Geoformer
10
40
18.4
15.4
33.8
27.5
1.28
4.43
4.41
4.39
6.13
22
QM9 dataset consists of 130,831 small organic molecules with up to 9 heavy atoms. Each molecule associated with 12 targets covering its energetic, electronic, and thermodynamic properties. We randomly split them to 10,000 samples for training, 10,000 samples for validation and the remains for testing following the prior work [ The evaluation results on QM9 are shown in Table ! with the upper section displaying EGNNs and the lower showcasing Transformer-basee methods. When compared with other SoTA methods, Geoformer achieves the state-of-the-ap results on kinds of properties and shows comparable results on the remaining properties, which underscores that our IPE can help Transformers learn useful positional information better model molecular structures. Specifically, we conduct comparison between Geoformer and the previously best- performing Transformer-basel model, Transformer-M. which incorporates pairwise distances as PE. As Transformer-N M has been pretrained on the large-scale PCQM4Mv2 dataset [17], targeting the prediction of HOMO-LUMO gaps, exhibits comparable performance terms of related propertics. Nevertheless, noticeable performance disparity persists when evaluating other properties

---

in comparison to the EGNNs. By integrating more directional information beyond distances into Geoformer, it significantly surpasses Transformer-N M in 9 kinds of properties without pretraining.
4.3
Molecule3D
The Molecule3D dataset consists of 1899,647 molecules, each with the corresponding ground- state structures and quantum properties calcu- lated by Density Functional Theory (DFT). The dataset is split into train, validation, and test set with the ratio of 6:2:2. The official GitHub repository of Molecule3D provides both random and scaffold splits, which are both employed in our experiments. The random split ensures that the training, validation, and test sets are sampled from the same distribution, while the scaffold split introduces distribution shift among dif- ferent subsets. We focus our analysis on the prediction of the HOMO-LUMO gap, in com- parison with ComENet 47
Table Mean absolute errors (MAE) of HOMO- LUMO gap (eV) on Molecule3D test set for both random and scaffold splits compared with state-of- the-art algorithms.
Model
Random
Scaffold
GIN-Virtual
18
0.1036
0.2371
SchNet
0.0428
0.1511
DimeNet++
13
0,0306
0.1214
SphereNet
25
0.0301
0.1182
ComENet
47
0.0326
0.1273
Geoformer
0.0202
0.1135
Table 2 displays
the results of our experiments on Molecule3D.
indicating that Geoformer achieved reduction of 32.56% and 3.98% in test MAE on the random and scaffold splits, respectively. These results highlight the superiority of our approach compared to invariant GNNs on a large-scale dataset.
4.4 Analysis on Interatomic Positional Encoding
The exceptional performance demonstrated by Geoformer serves evidence that Interatomic Positional Encoding (IPE) can effectively guide Transformers in modeling the geometry of molecular systems. In order to analyze the distinctions between our learned IPE and other PEs that solely utilize pairwise distances, we visualize the positional encoding for different molecules by IPE and other PEs, respectively. As shown in Fig. IPE exhibits different positional encoding information compared with PEs that solely utilize pairwise distances. Specifically, for some positions that exhibit strong signals shown in the PEs only with distances, IPE further enhances such signals and shows significant distinction from background signals.
4.5 Ablation Study
To verify the effectiveness of our IPF, we conduct com- prehensive ablation studies with model variants on four properties Up. U, H and G' in QM9, First, we re- move the residual connec- tion for IPE, which leads to Transformer that exclu- sively encodes pairs of dis- tances using the initial non- updated IPE (Non-updated
Table 3: Ablation study on four properties Uo. U, H and G in QM9 test set for model variants. The best one in each property is highlighted in bold.
Property
Non-updated IPE
Addition
Geoformer
Uo
5.63
7.62
4.43
5.87
7.66
4.41
H
6.21
7.93
4.39
7.24
9.01
6.13
IPE). More specifically, the RBF feature is constructed Equation 14 where g(ft)
igi (ft), (f),.... (f))]= RNNNNXKK is composed of the values of K radial basis functions and the PE undergoes transformation in the 1-th layer through distinct linear layer that is not shared
across layers with the form:
= -(Rg(WWWL RBF
(24)
The remaining operations are consistent with those found in the original Geoformer implementation. This should be highly similar to the previous Transformer, albeit with a slightly different pairwise distance encoding approach.
8

---

Distances
0.035
5
10
10
0.030
15
10
0.025
20
10
0.020
0.015
IPE Cr
5
10
0.010
10
0.005
15-
15
20
0.000
10
Figure 3: Visualization of IPE C, on molecules GDB65488, GDB101712, GDB87153 and GDB56373 in QM9 test set. Other PEs (first row) only encode pairwise distances by RBF, and the learned IPE C, (second row) encodes different positional information beyond pairwise distance. In some atomic positions, IPE effectively emphasizes the atomic relationships, providing prehensive representation of the underlying geometry within the molecular structure. This enhanced encoding of geometric information enables the Transformer-based model to capture the intricate interactions between atoms and better predict molecular properties.
Second, we discover that the way of incorporating IPE into the attention weights is crucial. In our derivation (Eq. I0., multiplication emerges as natural approach, whereas the addition PEs as attention bias prevails in the previous Transformer. To verify the effectiveness of our design, also introduce our IPEs into the Transformer in an additive manner (Addition), which we replace multiplication addition follows:
A(X') SiLU (E, (((X'W)) (X'Wk) - - ((1)))
(25)
where C updated in the same way the original Geoformer, only has been replaced with +. As a show in Table 3 the performance of Geoformer without the updated IPE is worse than the original version, which shows that the additional directional information important.
The performance also drops when using the addition of PEs as the attention bias, which demonstrates the multiplication is proper way incorporate IPE into Transformers. Furthermore, the results for both variants outperform Transformer-M, which uses pairwise distance information as the atten- tion bias. This observation further highlights the significance of these components the overall performance gains.
Related Work
5.1
Positional Encoding in Transformers
The Transformer architecture contains a series of Transformer blocks as well as positional encoding (PE) to model the sequence data. Since the self-attention modules in Transformer blocks are invariant to the sequence orders, the positional encoding plays an essential role in injecting positional information from sequences into the Transformer. Recent PEs can categorized into absolute and relative PE. The original Transformer model incorporates absolute PE by directly adding the positional encoding to the tokens 46.. Although absolute PE has been demonstrated to approximate any continuous sequence-to-seqqeeee functions [53] it tends to exhibit inferior generalization capabilitiees

---

in comparison to relative PE, particularly when dealing with longer sequences 33 The relative PE further considers the pairwise relationship between two tokens. Shaw [1.. T5 DeBERTa [T.. and Transformer-Xl. Y have developed various relative PE. approaches for parameterizing relative positions. The success of relative PE in natural language processing has inspired its applications in other domains. For instance, the Swin Transformer 2 employs relative PE to model relationships between image patches, while Graphormer (52) utilizes both absolute PE (centrality encoding) and
relative PE (spatial encoding and edge encoding) to model the graph topology. Recently, Transformer- M PS and Uni-Mol [55] have incorporated pairwise distances as relative positional encoding to capture positional information within 3D space. TorchMD-Net - included the radial basis functions (RBF) as distance filter to the attention matrix. MUformer 19 further extended the distance filter by incorporating additional 2D structural information, resulting in improved performance and applicability to molecule 2D-3D co-generation. Several works have further explored the integration of different types of positional encoding in Transformers. GPS (36] offers comprehensive overview of the available PE. methods employed in graph Transformers. The MAT (2 and R-MAT (0
approaches methodologically introduce inter-atomic distances and chemical bond information as
domain-specific inductive biases in Transformer architecture. Molformer [50] employs Adaptive PE
to model molecules of varying sizes; GeoT p forgoes softmax in favor of distance matrices as
scaling factor, and the application of multiplication for PE has been previously tried in the Geometric Transformer In this study. our objective is to develop a relative PE for modeling molecular geometry. Drawing upon the atomic cluster expansion theory, which will be discussed in Section 2 we derive rotation-inwe invvariant relative PE that incorporates additional positional information beyond pairwise distances.
5.2 Geometric Deep Learning for Molecules
Geometric deep learning (GDL) has emerged as a promising approach to modeling molecular geometry and predicting the properties of molecules, playing an important role in fields such as drug discovery, materials science, and computational chemistry. By leveraging the inherent geometric structures and incorporating symmetry in architecture design, GDL approaches offer an effective and efficient representation for molccules. Invariant and equivariant graph neural networks (EGNNs) are representative methods in GDL. SchNet [38], DimeNet++ [13,, GemNet [12. SphereNet (25), ComENet 477 gradually explicitly incorporate more geometric information including distances, angles and dihedrals. Some works like PaiNN 3y, TorchMD 4%). ViSNet A adopt vector embedding and implicitly extract the above geometric information with lower consumption. Another mainstream approach such as NequlP B) MACE (2) Allegro [3I and Equiformer [24] guarantee equivariance through group representation theory, which can achieve higher accuracy leveraging high-order geometric tensors.
6 Discussion
In this paper, we propose a novel Transformer architecture in the field of molecular modeling. This innovative approach, incorporating the novel Interatomic Positional Encoding (IPE), effectively captures complex geometric information and interatomic relations beyond pairwise distances embeded in the molecular structures. The extensive results on QM9 and Molecule3D dataset elucidate the capability of Geoformer compared with Transformers and EGNNs. Further research can explore its applicability to broader range of systems such as materials and polymers. Moreover, the concept Interatomic Positional Encoding may inspire the development of more advanced encoding schemes in Transformers.
Limitation and Societal Impacts: Like other Transformer architectures, Geoformer suffers from common training instabilities, necessitating the use of a relatively small learning rate during the training process. Effective molecular geometry modeling, as provided by the Geoformer, significantly benefits the materials science and pharmaceutical industries. However, it is essential to acknowledge that the same technology could be misused for illicit activities, such as the manufacturing of illegal drugs or the development of biochemical weapons.
10

---

Acknowledgments and Disclosure of Funding
We thank the reviewers for their valuable comments. Yusong Wang and Nanning Zheng were supported in part by NSFC under grant No. 62088102.
References
[1] B. Anderson, S. Hy, and R. Kondor. Cormorant: Covariant molecular neural networks. Advances in neural information processing systems, 32, 2019.
[2] I. Batatia, D. P. Kovacs, G. Simm, C. Ortner, and G. Csányi. Mace: Higher order equivariant message passing neural networks for fast and accurate force fields. Advances in Neural Information Processing Systems, 35:1 1423-1122111 1136, 2022.
[3] S. Batzner, A. Musaelian, L. Sun, M. Geiger, J. P. Mailoa, M. Kornbluth, N. Molinari, T. E. Smidt, and B. Kozinsky. (3)-equivariant graph neural networks for data-efficient and accurate interatomic potentials. Nature communications, 13(1):2453, 2022.
[4] J. Brandstetter, R. Hesselink, E. van der Pol, E. J. Bekkers, and M. Welling, Geometric and physical quantities improve (3) equivariant message passing. arXiv preprint arXiv:21 2021.
[5] T.-C. Chi, T.-H. Fan, and A. I. Rudnicky. Receptive field alignment enables transformer length extrapolation. arXiv preprint arXiv:2212. 10356, 2022.
[6] S. Chmiela, H. Sauceda, K.-R. Müller, and A. Tkatchenko. Towards exact molecular dynamics simulations with machine-leamed force fields. Nature communicattonns, 9(1):1-10, 2018.
[7] S. Chmiela, A. Tkatchenko, H. E. Sauceda, I. Poltavsky, T. Schütt, and K.-R. Müller. Machine learning of accurate energy-conserving molecular force fields. Science advances, 3(5):e 2017.
[8] Y. Choukroun and L. Wolf. Geometric transformer for end-to-end molecule properties prediction. arXiv preprint arXiv:2110.13722 2021.
[9] Z. Dai, Z. Yang, Y. Yang. J. Carbonell, Q. V. Le, and R. Salakhutdinov. Transformer-xll: Attentive language models beyond fixed-length context. arXiv preprint 2019.
[10] R. Drautz. Atomic cluster expansion for accurate and transferable interatomic potentials. Physical Review 99(1):01440 2019.
[11] M. Finzi, S. Stanton, P. Izmailov, and A. G. Wilson. Generalizing convolutional neural networks for equivariance to lie groups on arbitrary continuous data. International Conference on Machine Learning, pages 3165-3176. PMLR, 2020.
[12] J. Gasteiger, F. Becker, and S. Günnemann. Gemnet: Universal directional graph neural networks for molecules. Advances in Neural Information Processing Systems, 34:6790-6802, 2021.
[13] J. Gasteiger, S. Giri, J. T. Margraf, and S. Günnemann. Fast and uncertainty- directional message passing for non-equilibrium molecules. arXiv preprint arXiv:2011.14111, 2020.
[14] J. Gilmer, S. S. Schoenholz, P. F. Riley, O. Vinyals, and G. E. Dahl. Neural message passing for quantum chemistry. In International conference machine learning, pages 1263-1272 PMLR, 2017.
[15] J. Godwin, M. Schaarschmidt, A. Gaunt, A. Sanchez-Gonzalez, Y. Rubanova, P. Veličković, J. Kirkpatrick, and P. Battaglia. Simple gnn regularisation for 3d molecular property prediction beyond. arXiv preprint arXiv:2106.00 07971, 2021.
[16] He, X. Liu, Gao, and W. Chen. Deberta: Decoding-enhanced bert with disentangled attention. arXiv preprint ar wwX.x:200660363:3 2020.

---

[17] W. Hu, M. Fey. H. Ren, M. Nakata, Y. Dong, and J. Leskovec. Ogb-Isc: A large-scale challenge for machine learning on graphs. arXiv preprint arXiv:2103.09450, 2021.
[18] W. Hu, M. Fey, H. Ren, M. Nakata, Y. Dong, and J. Leskovec. Ogb-Isc: A large-scale challenge for machine learning on graphs. arXiv preprini arXiv:2103.094500 2021.
[19] C. Hua, S. Luan, M. Xu, R. Ying. J. Fu, S. Ermon, and D. Precup. Mudiff: Unified diffusion for complete molecule generation. arXiv preprint arXiv:2304.14621, 2023.
[20] M. S. Hussain, M. J. Zaki, and D. Subramanian. Global self-attention as a replacement for graph convolution. In Proceedings of the 28th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, pages
2022.
[21] D. P. Kovács, C. v. d. Oord, J. Kucera, A. E. Allen, D. J. Cole, C. Ortner, and G. Csányi. Linear atomic cluster expansion force fields for organic molecules: beyond rmse. Journal of chemical theory and compulation, 17(12)7696-7711, 2021.
[22] B. Kwak, J. Jo. Lee, and S. Yoon. Geometry-aware transformer for molecular property prediction. arXiv preprint arXiv:2106. 2021.
[23] Le, F. Noé, and D.-A. Clevert. Equivariant graph attention networks for molecular property prediction. arXiv preprint arXiv:2202.098)) 2022.
[24] Y.-L. Liao and T. Smidt. Equiformer: Equivariant graph attention transformer for 3d atomistic graphs. arXiv preprint 1rXiv:2206.11999, 2022.
[25] Y. Liu, L. Wang, M. Liu, Y. Lin, X. Zhang, B. Oztekin, and S. Ji. Spherical message passing for
3d molecular graphs. In International Conference on Learning Representations (ICLR), 2022.
[26] Y. Liu, L. Wang, M. Liu, X. Zhang. B. Oztekin, and S. Ji. Spherical message passing for 3d graph networks. arXiv preprint arXiv:2102.050133 2021.
[27] Z. Liu, Y. Lin, Y. Cao, H. Hu, Y. Wei, Z. Zhang, S. Lin, and B. Guo. Swin transformer: Hierarchieal vision transformer using shifted windows. In Proceedings of the IEEE/CVF international conference on computer vision, pages 10012-10022, 2021.
[28] S. Luo, T. Chen, Y. Xu, S. Zheng, T.-Y. Liu, L. Wang, and D. He. One transformer can understand both 2d & 3d molecular data. arXiv preprint 1rXiv:2210.01165, 2022.
[29] E. Maziarka, T. Danel, S. Mucha, K. Rataj. J. Tabor, and S. Jastrzebski. Molecule attention transformer. arXiv preprint arXiv:2002. 082644 2020.
[30] F.. Maziarka, D. Majchrowski, T. Danel, P. Gainski, J. Tabor, I. Podolak, P. Morkisz, and Jastrzebski. Relative molecule self-attention transformer. arXiv preprint 1r.iv22110.08844) 2021.
[31] A. Musaelian, Batzner, A. Johansson, L. Sun, C. J. Owen, M. Korbbluth, and B. Kozin- sky. Learning local equivariant representations for large-scale atomistic dynamics. Nature Communications, 14(1):579, 2023.
[32] M. Nakata and T. Shimazaki. Pubchemqe project: a large-scale first-principle electronic structure database for data-driven chemistry. Journal of chemical information and modeling. 5766):1300-1306 1308, 2017.
[33] O. Press, N. Smith, and M. Lewis. Train short, test long: Attention with linear biases enables input length extrapolation. arXiv preprint arXiv:2 /2409, 2021.
[34] C. Raffel, N. Shazeer, A. Roberts, K. Lee, S. Narang, M. Matena, Y. Zhou, W. Li, and P. J. Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. The Journal of Machine Learning Research, 21 (1):54885-555) 2020.
[35] R. Ramakrishnan, P. O. Dral, Rupp, and O. A. Von Lilienfeld. Quantum chemistry structures and properties of 134 kilo molecules. Scientific data, 1(1):-77, 2014.
12

---

[36] L. Rampášek, M. Galkin, V. Dwivedi, A. T. Luu, G. Wolf, and D. Beaini. Recipe for a general, powerful, scalable graph transformer. Advances in Neural Information Processing Systems, 35:14501-14515, 2022.
[37] V. G. Satorras, E. Hoogeboom, and M. Welling. E (n) equivariant graph neural networks. In International conference on machine learning, pages 9323-9332. PMLR, 2021.
[38] K. Schütt, P.-J. Kindermans, H. E. Sauceda Felix, S. Chmiela, A. Tkatchenko, and K.-R. Müller. Schnet: A continuous- filter convolutional neural network for modeling quantum interactions. Advances in neural information processing systems, 30, 2017.
[39] K. Schütt, O. Unke, and M. Gastegger. Equivariant message passing for the prediction of tensorial properties and molecular spectra. In International Conference on Machine Learning, pages 9377-9388. PMLR, 2021.
[40] K. T. Schütt, F. Arbabzadah, S. Chmiela. K. R. Müller, and A. Tkatchenko. Quantum-chemical insights from deep tensor neural networks. Nature communications, 8(1):188, 2017.
[41] P. Shaw, J. Uszkoreit, and A. Vaswani. Self-attention with relative position representations. arXiv preprini arXiv: 1803.022555 2018.
[42] Thölke and G. De Fabritiis. Equivariant transformers for neural network based molecular potentials. In International Conference on Learning Representations, 2022.
[43] Thölke and G. De Fabritiis. Torchmd-net: equivariant transformers for neural network based molecular potentials. arXiv preprint arX1v:2200.022411, 2022.
[44] M. Thürlemann and S. Riniker. Anisotropic message passing: Graph neural networks with directional and long-range interactions. In The Eleventh International Conference on Learning Representations, 2023.
[45] O. T. Unke and M. Meuwly. Physnet: moments, and partial charges. Journal of chemical theory and computation, (5(6):367-3369) 2019.
A neural network for predicting energies, forces, dipole
[46] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, t.. Kaiser, and I. Polosukhin. Attention is all you need. Advances neural information processing systems, 30, 2017.
[47] L. Wang, Y. Liu, Y. Lin, H. Liu, and S. Ji. Comenet: Towards complete and efficient message passing for 3d molecular graphs. Advances in Neural Information Processing Systems, 2022.
[48] T. Wang, X. He, M. Li, Y. Wang, Z. Wang, S. Li, B. Shao, and T.-Y. Liu. Ai2bmd: efficient characterization of protein dynamics with ab initio accuracy. bioRxiv, pages 2023-07, 2023.
[49] Y. Wang, S. Li, X. He, M. Li, 7. Wang, N. Zheng, B. Shao, T. Wang, and T.-Y. Liu. Visnet: a scalable and accurate geometric deep learning potential for molecular dynamics simulation. arXiv preprint arXiv:2210.1633
2022.
[50] Wu, D. Radev, and S. Z. Li. Molformer: Motif-based transformer on 3d heterogencous molecular graphs. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 37, pages 531 1-220, 2023.
[51] Z. Xu, Y. Luo, X. Zhang, X. Xu, Y. Xie, M. Liu, K. Dickerson, C. Deng, M. Nakata, and S. Ji. Molecule3d: A benchmark for predicting 3d geometries from molecular graphs. arXiv preprint 11000:2110.01717, 2021.
[52] C. Ying, T. Cai, S. Luo, S. Zheng, G. Ke, D. He, Y. Shen, and T.-Y. Liu. Do transformers really perform badly for graph representation? Advances in Neural Information Processing Systems, 34:28877-28888 2021.
[53] C. Yun, Bhojanapalli, A. S. Rawat, S. J. Reddi, and S. Kumar. Are transformers universal approximators of sequence-to-sequeee functions? arXiv preprint
13

---

[54] S. Zhang, Y. Liu, and L. Xie. Efficient and accurate physics aware multiplex graph neural net- works for 3d small molecules and macromolecule complexes. arXiv preprint 022702900600789, 2022.
[55] G. Zhou, Z. Gao, Q. Ding, H. Zheng, H. Xu, Z. Wei, L. Zhang, and G. Ke. Uni-mol: A universal 3d molecular representation learning framework. In International Conference on Learning Representations, 2023.
14


Reviews:
[{"review_id": "qo7LYbYO11", "reviewer": "Reviewer_M4bW", "summary": "The paper proposes a geometric Transformer called Geoformer for geometric molecular modeling. It designs Interatomic Positional Encoding (IPE), taking 3D geometric information into account to parametrize atomic environment for positional encoding in Transformer architecture. Geoformer adopts a learned Interatomic Positional Encoding (IPE) to the attention block of the Transformer. Geoformer is evaluated on QM9 and Molcule3D datasets with various tasks and shows superior performance compared with Transformers and EGNN baselines. ", "questions": "- In line 163, is A_{i, v_{tau}} defined in the same way as equation 3? If so, the bold font should be unified.  \n- Is C_{vv^{prime}_{tau}} the same as C_{vv^{prime}} in the equation 4?\n- In subsection 4.5 and figure 3, is there any task-related or domain-specific ground truth where IPE focuses more on the atomic relationship in learned IPE? It is necessary for the claim that \"IPE further enhances signals and shows significant distinction\" to be persuasive.\n", "limitations": "The limitations are well addressed in the paper.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "- Geoformer effectively integrates domain knowledge from physics, specifically Atomic Cluster Expansion (ACE), into the machine learning field for molecular property prediction. The Interatomic Positional Encoding (IPE) with ACE offers a robust method for incorporating geometric priors from other scientific disciplines into the positional encoding of Transformers.\n- IPE reflects geometric information and interatomic relations within the Transformer architecture. Leveraging IPE, Geoformer demonstrates strong performance on 3D molecular property prediction tasks across QM9 and Molecule3D datasets over various baselines.\n", "weaknesses": "Overall, the representation of this paper could be polished further to improve accessibility. Considering this paper is targeted at a machine learning conference, more detailed explanations about the background (ACE) would be beneficial for potential readers from this community. If the space matters, the appendix would be a good place for the additional information. \n- For example, the term \"cluster\" is used in Theorem 1 without a clear definition. I presume this is the concept introduced in ACE as well.\n- Although there are some introductory sentences, Section 3.1 only consists of two theorems without any proper explanation of why these theorems are needed and how they can be used.  For example, Theorem 1 is introduced without any prior explanation of its necessity. It would be better to explain why Theorem 1 is needed and its main messages beforehand.\n- The attention block introduced through Equations 18 to 23 seems like an important contribution of this paper. However, with the given content, it is unclear what intuition leads to the specific form of learnable IPE in Eq 20, 21. \n- Preliminary and related work sections need to be separated for better explanation. Currently, section 2.1 and 2.3 is more like related work. I suggest moving these to the appendix and providing more information on ACE in the preliminary.\n- It is questionable whether Theorem 2 is indeed a theorem; it seems like a matrix IPE matrix C was defined for positional encoding. More explanation would be appreciated.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "## Response\n\nThanks for your valuable comments. We provide point-to-point responses in the following:\n\n### Weakness 1\n\nThe term **cluster** is a concept in the original Atomic Cluster Expansion (ACE) theory, which represents the local chemical environment of centered atoms. A cluster $\\alpha$ contains one centered atom $i$, $K$ neighbor atoms (elements, denoted as $j$) with $K$ bonds. We’ve added it in the Appendix.\n\n### Weakness 2\n\nWe apologize for the confusion. Theorem 1 serves as the foundation for Theorem 2 and the architecture of Geoformer in Section 3. It describes the proposed interatomic positional encoding (IPE) from a physics perspective. Specifically, we demonstrate how to construct the IPE (Theorem 2) **in theory**, representing the potential of merged cluster (Theorem 1). Next, we illustrate how to leverage the IPE (Theorem 2) within traditional Transformers (Section 3.1) **in practical**. In summary, we aim to demonstrate the proposed IPE both theoretically and practically.\n\n### Weakness 3\n\nWe apologize for the oversight. The concepts behind Eq. 20 and 21 originate from AlphaFold2, which utilizes pairwise embedding $z$ with an activation function as a learnable gate, combined with the previous residual. We recommend referring to the *row/column-wise gated self-attention* in AlphaFold2 supplementary. We have also added this reference to the manuscript.\n\n### Weakness 4\n\nFollowing your suggestion, we separate the sections into **Preliminaries** and **Related Work** and provide a more detailed introduction of ACE.\n\n### Weakness 5\n\nIn Theorem 2, we aim to prove why the interatomic positional encoding should be multiplied with the Key and Query in the traditional Transformers, as opposed to other operations such as attention bias or distance filter. Therefore, we consider this construction and proof as a theorem. Additionally, we have included the complete proof in the Appendix.\n\n### Question 1\n\nThanks for reaching this out. We have removed the bold font $i$ in line 163.\n\n### Question 2\n\nRight. They both represent the Clebsch-Gordan coefficients to ensure the permutation and isometry-invariance.\n\n### Question 3\n\nWe appreciate your comments. To testify “IPE further enhances signals and shows significant distinction”, in Figure 3, we compare the IPE with the distance signal, which the most of previous work adopted in the original manuscript to show how IPE functions. Based on your suggestion, we are exploring for some domain-specific ground truth, e.g., electron density, to further substantiate the effectiveness of the proposed IPE. More results would be added if we find something interesting."}, {"review_id": "p81X35bJPg", "reviewer": "Reviewer_nHVJ", "summary": "The paper introduced a Transformer-based model with interatomic positional encoding (IPE) for molecular modeling, which the authors termed \"Geoformer\". Geoformer incorporated novel positional encoding derived from empirical physical knowledge (atomic cluster expansion) and was able to achieve comparable or even superior performance on multiple benchmarks.", "questions": "### Notations about the merged cluster.\n\n1. The relation between Fig.1 and Eq. 7 is confusing. In Eq. 7, the author shows the matrix of $A$-basis to represent the new basis for the merged cluster, while in Fig. 1, it turns out to be the matrix of $B$-basis to describe the proposed positional encoding. Their relation needs to be further explained in the manuscript.\n2. It is hard to follow the relation between the basis of the merged cluster with the proposed interatomic positional encoding for those readers who are not familiar with ACE. Also, only after reading several times could I realize that the matrix of $A$-basis in Theorem 1 is the additional (right) term in Eq.7. The relation between Theorem 1 and 2 should be clarified.\n3. What is the difference between $A_{i(i)}$ and $A_{i}$?\n4. Following Q3, what is the difference between $A_{i(i)}$ in the merged cluster with the original $A$-basis in (linear) ACE? \n\n### Complexity\n\n5. The author mentioned they only utilized reduced settings, i.e., body-expansion=1 and order of spherical harmonics=1. What is the time/memory complexity for the algorithm?\n6. Compared with the explicit calculation like GEM [1], and those approaches like SphereNet [2] which employ the CG-product, how much time/memory would Geoformer reduce?\n7. What is the model size of Geoformer compared with other EGNNs?\n\n### Other minor problems\n\n8. The index of interatomic positional encoding (C) in the left corner, Fig. 1 is wrong.\n9. how is the molecule graph being padded?\n\n\n\n[1] Liu, Lihang, et al. \"GEM-2: Next Generation Molecular Property Prediction Network by Modeling Full-range Many-body Interactions.\" (2022).\n\n[2] Liu, Yi, et al. \"Spherical message passing for 3d graph networks.\" *arXiv preprint arXiv:2102.05013* (2021).\n", "limitations": "The authors have adequately addressed the limitation and potential negative societal impact in Section 5.", "rating": 8, "confidence": 5, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "1. To the best of my knowledge, the proposed approach for injecting physical knowledge, i.e., the atomic cluster expansion into the positional encoding in Transformers is novel in this field, which is dominated by EGNNs. The extensive experiments and ablation studies also demonstrated the expressiveness of the proposed architecture.\n\n2. The idea of merging clusters further considers the interaction/potential between different clusters, while the original ACE only considers one individual cluster.\n\n3. The proposed method has advantages over current EGNNs such as incorporating angles and dihedrals, and seems more computationally efficient because the calculation between clusters is still under linear time complexity.\n\n4. It revisits the Transformer architecture in the field of molecular modeling and outperforms most of the current sota models in QM9 and Molecule3D benchmark.", "weaknesses": "1. Some notations and figures about the merged cluster are a bit confusing and can be improved (see the following question section).\n\n2. The authors claimed the proposed model was computationally more efficient and memory-saving. The authors can provide quantitative data for a clearer justification (see the following question section).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "## Response\n\nThanks for your comments and acknowledgment of our work. We provide the point-to-point responses in the following.\n\n### Question 1\n\nWe appreciate your suggestion and have provided more explanation of the construction of the B-basis in Theorem 2 and the Appendix for better clarity. Indeed, the merged B-basis remains in matrix form, similar to the merged A-basis matrix. It is additionally multiplied with CG-coefficients to ensure permutation and isometry-invariance.\n\n### Question 2\n\nThank you for pointing this out. We understand that readers unfamiliar with ACE or the physics of harmonics might find it challenging to follow. We have added more explanation of the connection between Theorem 1 and Theorem 2. The central point is that \"the potential of the merged cluster in Theorem 1 is the physical explanation of the interatomic positional encoding in Theorem 2.\" In Section 3.1, we demonstrate how to incorporate IPE within the Transformer in practice. We aim to introduce the proposed IPE both theoretically and practically.\n\n### Question 3 and 4\n\n$A_i$ is the original A-basis in ACE. Since we conduct the cluster merging, $A_{i(i)}$ denotes the multiplication of two $A_i$, describing the $\\epsilon+1$ body expansion in the manscript line 137. Together, the matrix in Eq. 7 could represent the $\\epsilon+1$ and $\\epsilon+2$ body expansion simutaneously.\n\n### Question 5 and 6\n\nDue to the linear ACE theory, the complexity of constructing of interatomic positional encoding is $\\mathcal{O}(N)$ with $N$ denoting the number of atom with one molecule. In contrast, the explicit extraction method such as GEM, GemNet and SphereNet have a complexity of $\\mathcal{O}(N^{3})$ due to extracting dihedral angles for 4-body interaction. The main computational consumption lies in the Attention block, with a complexity of $\\mathcal{O}(N^{2})$. We cannot replace this part with the prevailing linear attention at this moment because we have to update learnable IPE. We would explore these possibilities in the future work. Compare to higher-order of spherical harmonics, which require the pre-computed CG-coefficients and products, we streamline this part by reduced setting. Since the number of spherical harmonics is $(l+1)^2$ with degree $l$, it would reduce consumption (tensor paths) by approximately $Nl$.\n\n### Question 7\n\nModern EGNN models are indeed relatively small (50.7M vs 10M), however, due to the  CG-product or graph operations in their implementations, Transformer-based approach has a similar training speed to them (see the below table for details).\n\n|  | Model Size | Overall Training Time (GPU-hours) | MAE on U0 | MAE on U | MAE on H | MAE on G |\n| --- | --- | --- | --- | --- | --- | --- |\n| SEGNN | 1.03M | 81 | 15 | 13 | 16 | 15 |\n| TorchMD-NET | 6.86M | 92 | 6.15 | 6.38 | 6.16 | 7.62 |\n| Equiformer | 3.53M | 61 | 6.59 | 6.74 | 6.63 | 7.63 |\n| Transformer-M | 47.4M | - | 9.37 | 9.41 | 9.39 | 9.63 |\n| Geoformer | 50.1M | 55 | 4.43 | 4.41 | 4.39 | 6.13 |\n| Geoformer-S | 6.4M | 20 | 5.20 | 5.12 | 5.19 | 6.78 |\n\n### Minor Problem 1\n\nThanks for reaching this out. We have modified the $C_{\\gamma}$ with $C_{\\eta}$ in Fig. 1.\n\n### Minor Problem 2\n\nIn one batch, we select the the maximum number of molecular atom, and pad other molecules with $0$, detaching the gradient at these padding indices in *Embedding* layer and add *key_padding_mask* for computing attention correctly."}, {"review_id": "ADi6D80mTz", "reviewer": "Reviewer_z5Cf", "summary": "The paper introduces Geoformer, a new geometric transformer for molecular property prediction. The authors argue that while transformers have been dominant in various data modalities, their application to molecular modeling has been limited. To address this, the authors propose Interatomic Positional Encoding (IPE) based on atomic cluster expansion (ACE) theory, which captures complex interactions within molecules. They evaluate Geoformer on the QM9 dataset and the Molecule3D dataset, demonstrating its superior performance compared to the existing methods.", "questions": "Please refer to the weaknesses for my concerns about the potential applications of this method.", "limitations": "The method could be computationally expensive as the presented positional encoding involves additional learnable parameters.\n\n\n", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "I am unfamiliar with this topic, but the idea of capturing complex geometric features by introducing Interatomic Positional Encoding seems novel. The motivation is well explained, and the authors also provide some mathematical proofs. The presented positional encoding may contribute to this field. \n\nThe results on the QM9 dataset and the Molecule3D dataset show that Geoformer outperforms other competitors in terms of most metrics. This highlights the effectiveness of the proposed method in capturing and utilizing valuable geometric information.\n", "weaknesses": "I did not see a major weakness, but I have some minor concerns about the experiments. \n\nAlthough this paper handles the problem of molecular property prediction, it seems that the presented geometric transformer block could potentially serve as a plugin in other machine learning and computer vision tasks. For example, the task of point cloud analysis could also benefit from describing complex geometric features. It would be beneficial and convincing if the authors could show the effectiveness of the presented method in such a context.\n\nThe presented positional encoding includes some learnable parameters, which is different from the traditional encoding strategy, and the encoding is dynamically updated. Therefore, the ablation study towards the effectiveness of the presented method is crucial. However, the authors put some results in the appendix instead of the main paper, which downgrades the importance. It would be better if the authors could shed more light on the ablation study in the main paper.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "## Reponse\n\nThanks for raising some valuable concerns. We provide the point-to-point response in the following.\n\n### Weakness 1\n\nAfter conducting a literature review on point cloud analysis, we discovered related work which incorporates geometric information in the point cloud modeling [1-5]. Though few of these studies consider *equivariance* as in molecules, there are several works like TFN [7] are famous, and the most relavent work would be by Qin et al. [6], who propose *Pair-wise Distance Embedding* and *Triplet-wise Angular Embedding* in Transformers as attention bias to *Key*. Therefore, we believe our proposed interatomic positional encoding (IPE) have the potential to generalize to the field of point cloud analysis beyond the molecule.\n\nDue to the time limitation, we cannot present more results in such a context, but would investigate its effectiveness in the future. Thanks for your valuable insights again!\n\n[1] Yu, Xumin, et al. \"Pointr: Diverse point cloud completion with geometry-aware transformers.\" *Proceedings of the IEEE/CVF international conference on computer vision*. 2021.\n\n[2] Ma, Xu, et al. \"Rethinking network design and local geometry in point cloud: A simple residual MLP framework.\" *arXiv preprint arXiv:2202.07123* (2022).\n\n[3] Chen, Zhi, et al. \"Sc2-pcr: A second order spatial compatibility for efficient and robust point cloud registration.\" *Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition*. 2022.\n\n[4] Yu, Hao, et al. \"Rotation-invariant transformer for point cloud matching.\" *Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition*. 2023.\n\n[5] Hou, Ji, et al. \"Mask3D: Pre-training 2D Vision Transformers by Learning Masked 3D Priors.\" *Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition*. 2023.\n\n[6] Qin, Zheng, et al. \"Geometric transformer for fast and robust point cloud registration.\" *Proceedings of the IEEE/CVF conference on computer vision and pattern recognition*. 2022.\n\n[7] Thomas, Nathaniel, et al. \"Tensor field networks: Rotation-and translation-equivariant neural networks for 3d point clouds.\" arXiv preprint arXiv:1802.08219 (2018).\n\n### Weakness 2\n\nDue to the page limitation, we had to put some ablation experiments in the Apendix. However, we highlight the analysis of ablation study in the manuscript as you suggested.\n\n### Limitation 1\n\nIndeed, our approach introduces additional computational overhead. Compared with Transformer-M using the same setting (the number of layers and hidden dimensions), our model increases by less than 8% parameters.  However, considering the improvement in performance, we believe the extra computational overhead is acceptable."}, {"review_id": "Eeq0xGNxLw", "reviewer": "Reviewer_efGJ", "summary": "The authors propose an Interatomic Positional Encoding (IPE), and introduce the Geoformer. IPE is motivated by Atomic Cluster Expansion (ACE), which describes the many-body contributions in transformers for molecular modeling. The authors first derive integrated cluster potentials for atom pair i,j in a molecule, then leverage the interatomic potentials to propose the interactomic positional encoding. Further, the authors introduce the Geoformer architecture that implements IPE. \n\nExperimentally, the authors compare Geoformer with popular molecular baselines on two datasets, QM9 and Molecule3D, and achieve the best results on most of the tasks. Additionally, the authors conduct ablation studies to varify the effectiveness of interatomic positional encoding.\n\nOverall, the paper proposes a novel approach, and the experimental results show a better model performance than exisiting methods. ", "questions": "1. I would like to hear from authors about how ACE/IPE surpass the simplest positional encoding in [1] and [2] in terms of expressive power and calculation time. Moreover, how do the authors choose the trade-off between efficiency and effectiveness?\n\n2. The authors mention that the architecture suffers from training instabilities. I wonder if there could be other techniques for stable training other than just using small learning rates.\n\n3. I would like to see the comparison of model size (number of parameters) and training time, as well as the standard deviation over experimental runs.\n\n4. I would like to see the comparison of model performance on the MD17 dataset, as MD17 provide more valid molecules. I also wonder why the authors do not perform experiments on MD17 in the beginning, as QM9 and MD17 are two of the most standard datasets for molecular simulation.\n\n\n[1] Thölke, Philipp, and Gianni De Fabritiis. \"Torchmd-net: equivariant transformers for neural network based molecular potentials.\" arXiv preprint arXiv:2202.02541 (2022).\n\n[2] Hua, Chenqing, et al. \"MUDiff: Unified Diffusion for Complete Molecule Generation.\" arXiv preprint arXiv:2304.14621 (2023).", "limitations": "N/A", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The approach inspired by Atomic Cluster Expansion is novel. Indeed, the positional encodings for modelling interatomic interactions have always been one of important factors in designing effective graph transformers for molecular data. Inspired by atomic cluster expansion, the authors propose a novel interatomic positional encoding that 'successfully' describes the interactions between atoms in a molecule, with theoretical justification. Experimentally, Geoformer surpass existing methods by a certain margin, which demonstrates the effectiveness of their approach. Overall, the paper presents a good quality of work with clear writing.", "weaknesses": "First, I find this paper is not so well-motivated. TorchMD-Net[1] model interatomic interactions by distance filters, and Muformer[2] leverage it and use additional 2D structures in the distance filters to model interatomic interactions. So, employing such information in Transformers is not completely underdeveloped. And it is unclear to me how ACE/IPE surpass previous methods in terms of effectiveness and efficiency. \n\nIn addition, the overall architecture suffers from numerical and training instabilities, for which the authors have also mentioned the limitation in section 5. Also, the authors do not compare the overall training time and model parameters in the experimental section, it is important to see the comparison of model sizes. Also, the error bars (standard deviations) are not shown in the tables. The authors do not perform sufficient experiments on proper datasets, it is also important to include experiments on MD17.\n\n[1] Thölke, Philipp, and Gianni De Fabritiis. \"Torchmd-net: equivariant transformers for neural network based molecular potentials.\" arXiv preprint arXiv:2202.02541 (2022).\n\n[2] Hua, Chenqing, et al. \"MUDiff: Unified Diffusion for Complete Molecule Generation.\" arXiv preprint arXiv:2304.14621 (2023).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "## Response\n\nThanks for your valuable comments. We provide the point-to-point responses at length in the following.\n\n### Weakness 1\n\n1. The distance filter (Radial Basis Fuctions, RBF) in TorchMD-Net is a sub-class of Atomic Cluster Expansion (ACE), which could be considered as the 2-body expansion. We mentioned it in Section 2.2, line 88. Here we add a detailed explanation. The basis function in ACE is a combination of radial basis functions (RBF) and spherical harmonics (SHs),  $\\phi(r)=\\sqrt{4\\pi}R_{nl}(\\Vert r \\Vert)Y_l^m(r)$ in the original paper, with $n$ denoting the degree of RBF, $l$ denoting the degree of SHs, and $m$ denoting the index of SHs. If we set the $l=0$, the SH degrades to scalar $1$. Therefore, the basis functions remain as RBF, which is adopted by TorchMD-Net as the distance filter.\n2. Compared to MUformer, which treated the 3D structure information as attention bias in attention block, we provide the theorical proof in Theorem 2, explaining the reason we directly multiply the proposed interatomic positional encoding (IPE) to the *Key* and *Query* from the perspective of the ACE. In a nutshell, we believe that employing the geometric information itself is not novel, but how to effectively utilize them is the most essential question. Concerns about instabilitiy, the comparison of model size and training as well as the standard deviations are replied in Question 2 and 3.\n\n### Weakness 2\n\nThank you for raising this concern. The main issue for us is that, for the MD17 experiment, it is required to use 0.7% of the data for training because they are generated from the same trajectory. This setup aims to test whether the machine learning force field is data-efficient. However, we think this is not suitable for Transformers, which prefer a large amount of data. Instead, we chose the prevailing QM9 and Molecule3D benchmarks, containing 130,831 and 3,899,647 molecules, respectively. However, as you suggested, we have conducted experiments on MD17 and display the energy and force error. The results show that, equipped with IPE, Transformer can also be applied to problems in limited-data scenarios.\n\n**Energy MAE**\n\n| Molecule | Geoformer | NequIP | TorchMD-Net | GemNet | PaiNN |\n| --- | --- | --- | --- | --- | --- |\n| Aspirin | **0.116** | 0.131 | 0.123 | - | 0.167 |\n| Ethanol | **0.051** | **0.051** | 0.052 | - | 0.064 |\n| Malonaldehyde | **0.074** | 0.076 | 0.077 | - | 0.091 |\n| Naphthalene | 0.087 | 0.113 | **0.085** | - | 0.116 |\n| Salicylic Acid | **0.093** | 0.106 | **0.093** | - | 0.116 |\n| Toluene | 0.078 | 0.092 | **0.074** | - | 0.095 |\n| Uracil | **0.095** | 0.104 | **0.095** | - | 0.106 |\n\n**Forces MAE**\n\n| Molecule | Geoformer | NequIP | TorchMD-Net | GemNet | PaiNN |\n| --- | --- | --- | --- | --- | --- |\n| Aspirin | **0.169** | 0.184 | 0.253 | 0.217 | 0.338 |\n| Ethanol | **0.063** | 0.071 | 0.109 | 0.085 | 0.224 |\n| Malonaldehyde | **0.115** | 0.129 | 0.169 | 0.155 | 0.319 |\n| Naphthalene | 0.043 | **0.039** | 0.061 | 0.051 | 0.077 |\n| Salicylic Acid | **0.088** | 0.090 | 0.129 | 0.125 | 0.195 |\n| Toluene | **0.044** | 0.046 | 0.067 | 0.060 | 0.094 |\n| Uracil | **0.066** | 0.076 | 0.095 | 0.097 | 0.139 |\n\n### Question 1\n\nAs discussed in Theorem 1, Remark, line 145, IPE with $v=1$ and  $\\eta=1$ could be interpreted as introducing more geometric information like angle and dihedral with linear time complexity. In conjunction with the response to Weakness 1, we believe this explains why ACE/IPE outperforms those simple positional encoding, i.e., distance filters and the sum of Gaussian distances.\n\n### Question 2\n\nThe training instability is s common issue in EGNNs, like PaiNN and TorchMD-Net (e.g., see https://github.com/shehzaidi/pre-training-via-denoising/issues/3#issuecomment-1324959415) as well as other Transformers. However, to ensure the **equivariance** of the basis functions, we **cannot** directly apply *LayerNorm* to them. Here attempt to linearly shrink the norm of the basis over the last dimension to stablize training.\n\n### Question 3\n\nWe list the standard deviations when we repeated an additional 2 trials on the QM9 dataset, and from the results it can be seen that the model is not particularly sensitive to the partition of the dataset and the initial parameters.\n\n|  | $\\mu$ | $\\alpha$ | HOMO | LUMO | gap | $R^2$ | ZPVE | U0 | U | H | G | Cv |\n| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |\n| Std | 0.38 | 1.65 | 0.51 | 0.73 | 0.76 | 0.69 | 0.05 | 0.15 | 0.14 | 0.13 | 0.13 | 0.41 |\n\nWe directly used the model sizes and training times reported in Equiformer and Transformer-M, and experimented with Geoformer on a NVIDIA V100 GPU. It can be seen that in contrast to the Transformer-based method, our method increases the number of parameters by about 8%(47.4M vs 50.1M). Comparing to the EGNNs, although our number of parameters of the model is much larger than theirs, we can achieve faster training speed because some operators in EGNNs slow down training. In addition, to verify that the improvement in model performance does not only come from the larger model size, we also experimented with the model using a similar size as EGNNs, and although the performance drops compared to the larger model, it is still generally better and significantly faster than the other methods.\n\n|  | Model Size | Overall Training Time (GPU-hours) | MAE on U0 | MAE on U | MAE on H | MAE on G |\n| --- | --- | --- | --- | --- | --- | --- |\n| SEGNN | 1.03M | 81 | 15 | 13 | 16 | 15 |\n| TorchMD-NET | 6.86M | 92 | 6.15 | 6.38 | 6.16 | 7.62 |\n| Equiformer | 3.53M | 61 | 6.59 | 6.74 | 6.63 | 7.63 |\n| Transformer-M | 47.4M | - | 9.37 | 9.41 | 9.39 | 9.63 |\n| Geoformer | 50.1M | 55 | 4.43 | 4.41 | 4.39 | 6.13 |\n| Geoformer-S | 6.4M | 20 | 5.20 | 5.12 | 5.19 | 6.78 |\n\n### Qustion 4\n\nSame as Weakness 2."}, {"review_id": "SGXS6GFEmW", "reviewer": "Reviewer_mN7Q", "summary": "The paper proposes a Transformer based architecture incorporating interatomic information upon learned Atomic Cluster Expansion, integrated into the self-attention and refined residually over the network.\nThe method obtains state-of-the-art performance on a majority of prediction tasks on two popular molecular datasets.", "questions": "1) From the ablation study, we can observe that variant 1 has the smallest impact on accuracy decrease among the two variants.\nTo better understand the impact of ACE/IPE, it would be nice to see the ablations of Variant 1 coupled with the residual procedure proposed in the paper (it should not be worse at least).\n\n2) The Transformer literature (including molecular) provides many possible integrations of the IPE. One should provide more ablations rather than only with the standard sum (and besides the statements in line 194-)\n\n3) Can you provide a capacity comparison with the other methods? The (significant) improvement in accuracy may also come from a higher model capacity and better and/or longer training. Also, regarding reproducibility, some details regarding the training are missing (e.g. number of epochs)\n\n4) How do you explain the employment of a minimal number of expansions and harmonics still leads to such substantial results? \nAlso, in case such minimal hyperparameters are chosen, I believe one may simplify the cumbersome general case notation into a clear and easier-to-implement formulation.\n\n5) The Weaknesses Section above should be addressed (e.g., novelty, comparisons).\n\nGiven the very good performance of the model, addressing these points may certainly rise the paper's impact and rating.", "limitations": "No specific/significant limitations are presented.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The paper is well written. \nThe main novelty lies in the integration of learned ACE-based information into the model and its residual refinement.\nThe performances seem extremely significant.\n", "weaknesses": "I believe the most significant weakness of the paper is the lack of a literature review, which decreases the novelty level of the suggested method.\n\nTransformer-based architectures for molecular property prediction is a field of research of at least three years old.\nThe manuscript cites only two related papers and compares its performance to a very low-accuracy Transformer model. \\\nTo name a few of the firsts:\\\nMolecule attention transformer. arXiv:2002.08264, 2020\\\nRelative molecule self-attention transformer. arXiv:2110.05841, 2021\\\n3dtransformer: Molecular representation with transformer in 3d space. arXiv:2110.01191, 2021\\\nGeometry-aware transformer for molecular property prediction. arXiv:2106.15516, 2021\\\nGeometric transformer for end-to-end molecule properties prediction. arXiv:2110.13721, 2021\n\nThus, the proposed different embeddings and the integration of information into the self-attention are then not that novel.\\\n For example, the initial IPE is similar to Physnet, and the integration of pairwise (also multi-scale) information into the transformer is similar to the references above (or other Transformer based propositions applied to other modalities).\nAlso, the proposed solution may be suboptimal compared to the literature and has not been ablated enough (e.g., learned or choice of handcrafted cut-off functions, the integration modality of the IPE into self-attention (beyond the addition presented in the Appendix's ablation); please see Questions).\n\nIncorporating information from molecular mechanics/force field into neural networks has been proven as a powerful tool ([31,37,10,32 and many others]) by including beneficial inductive bias.\nThus, and as stated in the strengths, besides the very good results, I see the main contributions of the paper in the integration of augmented ACE information and also in the residual refinement of the IPE as the \"neural\" novelty.\n\nFinally, the proposed model seems to have at least 20% higher capacity than the competing transformers and probably more compared to the others.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "## Response\n\nThanks for your valuable comments. We provide the point-to-point responses in the following.\n\n### Weakness 1\n\n1. Sorry for the lack of literature review torward the Transformer-based architecture. The works you mentioned are pioneers in the molecular modeling. We introduce them in the Introduction and Related Work, and obtain valuable insights from them to design some variants of positional encoding (see Question 5 for the novelty and comparison).\n2. The key idea of different methods is try to design more effective positional encoding. According to your suggestion, we add these parts in the manuscript and conduct more ablation study (see Question 1 and 2).\n3. Right. Thanks for your acknowledgement. In addition, the prevailing molecular mechanics/force field methods are based on EGNNs. Here we revisit the power of Transformers with the proposed IPE and we believe it is one of the novelty in this paper.\n4. The introduction of IPE does introduce additional parameters. When we used the same settings as Transformer-M, such as the number of layers and hidden layer dimensions, we found that the parameters of the model went up a bit (about 8% higher capacity). Modern EGNN models are indeed relatively small (< 10M), however, due to the  CG-product or graph operations in their implementations, Transformer-based approach has a similar training speed to them (see Question 3 for details).\n\n### Question 1\n\nThanks for your comments. This suggested experiment does provide further evidence of the validity of our proposed framework for IPE. As you suggested, we update the initialized IPE using residuals that come from the transformation of the initial IPE itself, with no further information introduced, i.e. $C_{\\eta}^{l+1} =  \\operatorname{SiLU} \\left(C_{\\eta}^{l}W_{C}^l \\right)  + C_{\\eta}^{l}$\n\nThe experimental results are as you expected, the residual update brings a slight performance gains (Non-updated IPE **0.563**; Residual IPE without additional information **0.551**; Geoformer **0.443**), while the main performance improvement comes from the additional information $\\delta C_{\\eta}^l = \\sum_{\\tau=1}^{\\eta} W_{\\tilde{B}}^l \\sum_{v_{\\tau}'} C_{vv_{\\tau}'} \\tilde{A_{vv_{\\tau}'}}$  introduced to the IPE module.\n\n### Question 2\n\nThanks again for pointing these work. We add additional 3 variants (besides Question 1) brought from the ideas of above 4 papers:\n\n- for **Varaint 1** from *Molecule attention transformer*, we modify the IPE and attention block as:\n\n$$\nA=\\left(\\lambda_a \\operatorname{SiLU}(\\frac{QK^{T}}{\\sqrt{d_k}})+\\lambda_c\\bf{C}_{\\eta}\\right)V\n$$\n\nsetting $\\lambda_a=\\lambda_c=0.33$. \n\n- for **Variant 2** from *Geometry-aware transformer for molecular property prediction*, we modify the IPE following its paper:\n\n$$\nA=\\frac{(QK^{T}\\odot\\bf{C}_{\\eta})}{\\sqrt{d_k}} \\cdot V\n$$\n\n- for **Variant 3** from *Geometric transformer for end-to-end molecule properties prediction*, we modify the IPE following its paper:\n\n$$\nA=\\operatorname{Softmax}(\\frac{QK^{T}}{\\sqrt{d_k}})\\odot \\bf{C}_{\\eta}\n$$\n\nDue to the resource and time limitation, we only conduct the experiments on energy $U_0$:\n\n|  | Varaint 1 | Variant 2 | Variant 3 |\n| --- | --- | --- | --- |\n| U0 | 7.43 | 5.02 | 5.34 |\n\nWe discovered that the variants involving the positional encoding **multiplicated** with *Key* and *Query* perform better than those using summation, implicitly supporting our theorem.\n\n### Question 3\n\nWe have provided the model size and overall training time in the table below, and to answer your question, we have also provided a model using the same size as EGNNs, and you can see that although the performance of the smaller model drops a bit compared to the larger model, the model is still SOTA when compared to the other methods. Following the previous studies, the maximum number of epochs is set to 600. We will add detailed settings in the appendix.\n\n|  | Model Size | Overall Training Time (GPU-hours) | MAE on U0 | MAE on U | MAE on H | MAE on G |\n| --- | --- | --- | --- | --- | --- | --- |\n| SEGNN | 1.03M | 81 | 15 | 13 | 16 | 15 |\n| TorchMD-NET | 6.86M | 92 | 6.15 | 6.38 | 6.16 | 7.62 |\n| Equiformer | 3.53M | 61 | 6.59 | 6.74 | 6.63 | 7.63 |\n| Transformer-M | 47.4M | - | 9.37 | 9.41 | 9.39 | 9.63 |\n| Geoformer | 50.1M | 55 | 4.43 | 4.41 | 4.39 | 6.13 |\n| Geoformer-S | 6.4M | 20 | 5.20 | 5.12 | 5.19 | 6.78 |\n\n### Question 4\n\nIn MACE [1], with $l=2$ could surpass most of the methods with higher order geometric tensors. Since the basis function matrix in Eq. 7 could represent the $\\epsilon+1$ and $\\epsilon+2$ body expansion simutaneously as well as the power of Transformers, we could use lower degree to achieve the comparable results. To obtain the rigor and universal derivation, we consider the higher degree of spherical harmonics.\n\n[1] Batatia, Ilyes, et al. \"MACE: Higher order equivariant message passing neural networks for fast and accurate force fields.\" *Advances in Neural Information Processing Systems* 35 (2022): 11423-11436.\n\n### Question 5\n\nAll five studies endeavor to incorporate distance matrices for representing geometrical models, with certain investigations introducing the adjacency matrix (MAT) or chemical bond information (R-MAT) to further characterize geometries. Moreover, several studies examine the incorporation of attention bias; for instance, Molformer employs Adaptive PE to model molecules of varying sizes, GeoT forgoes softmax in favor of distance matrices as a scaling factor, and the application of multiplication for PE has been previously tried in the Geometric Transformer. Contrasting these  pioneer efforts, our methodology is grounded in the principles of ACE, providing a theoretical foundation for the empirical effectiveness of multiplication. Furthermore, as you mentioned, we introduced IPE based on ACE beyond pairwise distances, continuously refining it and bridging the performance gap between transformer-based approaches and EGNN."}]

Author rebuttal:

Meta-review: The paper tries to improve application of popular transformer architecture for molecular property prediction. In this regards, the authors incorporate domain knowledge from physics, specifically Atomic Cluster Expansion (ACE), into the transformer attention mechanism and positional embeddings. Empirically, it bridges the performance gap between EGNN and Transformer. Reviewers found the paper to be technically solid and demonstrates strong results, but also raised concerns regarding novelty, model complexity, and presentation issues like notation, background etc. 

We thank the authors and reviewers to actively engage during the discussion phase to improve the paper including providing more ablation experiments, which were also intriguing.

- The "neural" contribution/novelty is limited and the main contribution seems to be the integration of the ACE into existing (very overparameterized) Transformer architecture, similar to what was previously done with other pairwise metrics or handcrafted descriptors but unlike those, ACE seems to have such a large impact on the accuracy. Thus contribution of the paper should be updated.

- In terms of writing, several reviewers pointed out issues with background knowledge, and missing related work. After the rebuttal, the authors provided revised paragraphs in response to them. Please incorporate these including comparison with other transformers and continue to refine the manuscript.
Decision: Accept (poster)