SlowGuess/NeurIPS_2023_OCR huggingface dataset:
where  $\sum_{j}$  denotes the sum of atomic values weighted by the attention score.  $W_1^l$  and  $W_2^l$  are two learnable matrix in feed-forward layer. The output of each step is fed into a Layer Normalization (LayerNorm). It is important to emphasize that while the theoretical derivations of the Geoformaler architecture above may appear complex, in practice, the model utilizes simplified settings to reduce computational complexity. Specifically, we employ  $\eta = 1$  for body expansion and  $l^{*} = 1$  for spherical harmonics. It streamlines the complex tensor contraction and Clebsch-Gordan product calculations, making the model more efficient and easier to implement. Despite this, Geoformaler achieves or
surpasses state-of-the-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). It consists of a two linear layers with SiLU activation and an aggregation module  $\sum$  to predict the specific molecular property. The lightweight Decoder ensures that the Geoformer remains computationally efficient while maintaining its ability to accurately predict molecular properties based on the geometric features captured by the Encoder. More details on Decoder for specific properties are shown in Appendix G.
Geoformer is evaluated on both QM9 dataset [35] that consists of 12 molecular properties and a large-scale Molecule3D dataset [51] derived from PubChemQC [32] with ground-state structures and the corresponding properties calculated at DFT level. All results are measured by mean absolute error (MAE) on 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.
4.2 QM9
Table 1: Mean absolute errors (MAE) of 12 kinds of molecular properties on QM9 compared with state-of-the-art algorithms. The best one in each category is highlighted in bold.
Target
Unit	μ
mD	α
ma03	εHOMO
meV	εLUMO
meV	Δε
meV	(R2) ma02	ZPVE
meV	U0
meV	U
meV	H
meV	G
meV	Cv
meal
mol K
NMP [14]	30	92	43	38	69	180	1.50	20	20	17	19	40
SchNet [38]	33	235	41	34	63	73	1.70	14	19	14	14	33
Cormorant [1]	38	85	34	38	61	961	2.03	22	21	21	20	26
LieConv [11]	32	84	30	25	49	800	2.28	19	19	24	22	38
DimeNet++ [13]	30	44	25	20	33	331	1.21	6.32	6.28	6.53	7.56	23
EGNN [37]	29	71	29	25	48	106	1.55	11	12	12	12	31
PaiNN [39]	12	45	28	20	46	66	1.28	5.85	5.83	5.98	7.35	24
TorchMD-NET [42]	11	59	20	18	36	33	1.84	6.15	6.38	6.16	7.62	26
GNS + NoisyNode [15]	25	52	20	19	29	700	1.16	7.30	7.57	7.43	8.30	25
SphereNet [25]	25	45	23	19	31	268	1.12	6.26	6.36	6.33	7.78	22
SEGNN [4]	23	60	24	21	42	660	1.62	15	13	16	15	31
EQGAT [23]	11	53	20	16	32	382	2.00	25	25	24	23	24
PaxNet [54]	11	45	23	19	31	249	1.17	5.90	5.92	6.04	7.14	23
ComENet [47]	25	45	23	20	32	259	1.20	6.59	6.82	6.86	7.98	24
Equiformer [24]	11	46	15	14	30	251	1.26	6.59	6.74	6.63	7.63	23
AMP [44]	12	67	26	23	45	93	4.10	11.3	11.4	11.3	12.4	32
Molformer [50]	28	41	25	26	39	350	2.05	7.52	7.46	7.38	8.11	25
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 [8]	26.4	51	27.5	20.4	36.1	157	1.24	7.35	7.55	7.73	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 is associated with 12 targets covering its energetic, electronic, and thermodynamic properties. We randomly split them in to 110,000 samples for training, 10,000 samples for validation and the remains for testing following the prior work [43]. The evaluation results on QM9 are shown in Table I with the upper section displaying EGNNs and the lower showcasing Transformer-based methods. When compared with other SoTA methods, Geoformer achieves the state-of-the-art results on 8 kinds of properties and shows comparable results on the remaining properties, which underscores that our IPE can help Transformers learn useful positional information to better model molecular structures. Specifically, we conduct a comparison between Geoformer and the previously best-performing Transformer-based model, Transformer-M, which incorporates pairwise distances as PE. As Transformer-M has been pretrained on the large-scale PCQM4Mv2 dataset [17], targeting the prediction of HOMO-LUMO gaps, it exhibits comparable performance in terms of related properties. Nevertheless, a 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-M in 9 kinds of properties without pretraining.
4.3 Molecule3D
The Molecule3D dataset consists of 3,899,647 molecules, each with the corresponding groundstate structures and quantum properties calculated 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 a distribution shift among different subsets. We focus our analysis on the prediction of the HOMO-LUMO gap, in comparison with ComENet [47]. Table 2 displays the results of our experiments on Molecule3D,
Table 2: 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 [38]	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
indicating that Geoformer achieved a 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 our Geoformer serves as an 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. 3, 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.
To verify the effectiveness of our IPE, we conduct comprehensive ablation studies with model variants on four properties  $U_{0}$ ,  $U$ ,  $H$  and  $G$  in QM9. First, we remove the residual connection for IPE, which leads to a Transformer that exclusively encodes pairs of distances using the initial non-updated IPE (Non-updated
Table 3: Ablation study on four properties  $U_0$ ,  $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
U0	5.63	7.62	4.43
U	5.87	7.66	4.41
H	6.21	7.93	4.39
G	7.24	9.01	6.13
IPE). More specifically, the RBF feature is constructed as Equation [14], where  $\mathbf{g}(\hat{R}) = [\mathrm{g}_1(\hat{R}),\dots,\mathrm{g}_K(\hat{R})]^\top \in \mathbb{R}^{N\times N\times K}$  is composed of the values of  $K$  radial basis functions and the PE undergoes a transformation in the  $l$ -th layer through a distinct linear layer that is not shared across layers with the form:
$$
C _ {\eta} ^ {l} = \mathbf {g} (\hat {R}) W _ {\mathrm {R B F}} ^ {l} \tag {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.
Figure 3: Visualization of IPE  $C_{\eta}$  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_{\eta}$  (second row) encodes different positional information beyond pairwise distance. In some atomic positions, IPE effectively emphasizes the atomic relationships, providing a more comprehensive 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. [10]), multiplication emerges as a natural approach, whereas the addition of PEs as attention bias prevails in the previous Transformer. To verify the effectiveness of our design, we also introduce our IPEs into the Transformer in an additive manner (Addition), which we replace multiplication to addition as follows:
$$
A \left(X ^ {l}\right) = \operatorname {S i L U} \left(\sum_ {F} \left(\left(\left(X ^ {l} W _ {Q} ^ {l}\right) * \left(X ^ {l} W _ {K} ^ {l}\right) ^ {\top}\right) + C _ {\eta} ^ {l}\right)\right) \cdot \phi (\| \hat {R} \|) \tag {25}
$$
where  $C_{\eta}^{l}$  updated in the same way as the original Geoformer, only  $\odot$  has been replaced with  $+$ . As 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 is important.
The performance also drops when using the addition of PEs as the attention bias, which demonstrates the multiplication is a proper way to incorporate IPE into Transformers. Furthermore, the results for both variants outperform Transformer-M, which uses pairwise distance information as the attention bias. This observation further highlights the significance of these components to the overall performance gains.
5.1 Positional Encoding in Transformers
The Transformer architecture contains a series of Transformer blocks as well as a 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 be categorized into absolute PE 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-sequence functions [53], it tends to exhibit inferior generalization capabilities
in comparison to relative PE, particularly when dealing with longer sequences[33]. The relative PE further considers the pairwise relationship between two tokens. Shaw [41], T5 [34], DeBERTa [16], and Transformer-XL [9] 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 [27] 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 [28] and Uni-Mol [55] have incorporated pairwise distances as relative positional encoding to capture positional information within 3D space. TorchMD-Net [43] 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 an comprehensive overview of the available PE methods employed in graph Transformers. The MAT [29] and R-MAT [30] 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 [22] 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 [8]. 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 a rotation-invariant 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 molecules. Invariant and equivariant graph neural networks (EGNNs) are representative methods in GDL. SchNet [38], DimeNet++, [13], GemNet [12], SphereNet [25], ComENet [47] gradually explicitly incorporate more geometric information including distances, angles and dihedrals. Some works like PaiNN [39], TorchMD-Net [43], ViSNet [49] adopt vector embedding and implicitly extract the above geometric information with lower consumption. Another mainstream approach such as NequIP [3], MACE [2], Allegro [31] and Equiformer [24] guarantee equivariance through group representation theory, which can achieve higher accuracy leveraging high-order geometric tensors.
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 embedded 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 a 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-based 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.
We thank the reviewers for their valuable comments. Yusong Wang and Nanning Zheng were supported in part by NSFC under grant No. 62088102.
[1] B. Anderson, T. 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:11423-11436, 2022.
[3] S. Batzner, A. Musaelian, L. Sun, M. Geiger, J. P. Mailoa, M. Kornbluth, N. Molinari, T. E. Smidt, and B. Kozinsky. E (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 e (3) equivariant message passing. arXiv preprint arXiv:2110.02905, 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. E. Sauceda, K.-R. Müller, and A. Tkatchenko. Towards exact molecular dynamics simulations with machine-learned force fields. Nature communications, 9(1):1–10, 2018.
[7] S. Chmiela, A. Tkatchenko, H. E. Sauceda, I. Poltavsky, K. T. Schütt, and K.-R. Müller. Machine learning of accurate energy-conserving molecular force fields. Science advances, 3(5):e1603015, 2017.
[8] Y. Choukroun and L. Wolf. Geometric transformer for end-to-end molecule properties prediction. arXiv preprint arXiv:2110.13721, 2021.
[9] Z. Dai, Z. Yang, Y. Yang, J. Carbonell, Q. V. Le, and R. Salakhutdinov. Transformer-xl: Attentive language models beyond a fixed-length context. arXiv preprint arXiv:1901.02860, 2019.
[10] R. Drautz. Atomic cluster expansion for accurate and transferable interatomic potentials. Physical Review B, 99(1):014104, 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. In International Conference on Machine Learning, pages 3165-3176. PMLR, 2020.
[12] J. Gasteiger, F. Becker, and S. Gunnemann. 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. Gunnemann. Fast and uncertainty-aware directional message passing for non-equilibrium molecules. arXiv preprint arXiv:2011.14115, 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 on machine learning, pages 1263-1272. PMLR, 2017.
[15] J. Godwin, M. Schaarschmidt, A. Gaunt, A. Sanchez-Gonzalez, Y. Rubanova, P. Velicković, J. Kirkpatrick, and P. Battaglia. Simple gnn regularisation for 3d molecular property prediction & beyond. arXiv preprint arXiv:2106.07971, 2021.
[16] P. He, X. Liu, J. Gao, and W. Chen. Deberta: Decoding-enhanced bert with disentangled attention. arXiv preprint arXiv:2006.03654, 2020.
[54] S. Zhang, Y. Liu, and L. Xie. Efficient and accurate physics-aware multiplex graph neural networks for 3d small molecules and macromolecule complexes. arXiv preprint arXiv:2206.02789, 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.