SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Yusong Wang $^{1,2*}$ , Shaoning Li $^{2,3,4*}$ , Tong Wang $^{2*}$ , Bin Shao $^{2}$ , Nanning Zheng $^{1}$ , Tie-Yan Liu $^{2}$
$^{1}$  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  $^{2}$  Microsoft Research AI4Science
<sup>3</sup> Mila - Québec AI Institute <sup>4</sup> Université de Montréal  
wangyusong2000@stu.xjtu.edu.cn, nnzheng@mail.xjtu.edu.cn  
shaoning.li@umontreal.ca  
{watong, binshao, tyliu}@microsoft.com
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/A12BMD/tree/Geoformer.
Transformer [46] 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 [52, 36, 20] employ Transformers with global attention, which treat the molecular structures as fully connected graphs and devise a diverse range of positional encoding schemes. In contrast to topological graphs, geometric structures offer more comprehensive description of molecules by providing the information of 3D coordinates. There have been some recent attempts [28, 36, 55, 19, 29, 30, 50, 22, 8], to integrate specific geometric features into Transformers. For example, Transformer-M [28] considers pairwise distances as a learnable bias added to the attention weights as a supplement to the 2D topology encoding used in Graphormer [52]. 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
*Work done during an internship at Microsoft Research.
$\dagger$  Equal contribution.
$\ddagger$  Corresponding author.
contrast, modern GNNs emphasize the significance of equivariance as a crucial bias, and explore various methods to encode geometric information, including distances, angles, and dihedral angles [13, 12, 39, 26, 49, 47]. Several works [3, 31, 2] further utilize high-order geometric tensors to yield internal features in models, ensuring equivariance with respect to the E(3)/SE(3) group. As a result, the performance of EGNNs in molecular property prediction significantly surpasses that of methods based on Transformers.
The success of EGNNs underscores the advantage of integrating directional geometric information into neural networks for molecular modeling [12, 26, 47, 48], 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) theory [10, 21] is 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 [35] comprising 12 molecular properties and the recently proposed Molecule3D dataset [51], containing 3,899,647 molecules sourced from PubChemQC [32]. Our results demonstrate that Geoformer surpasses state-of-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.
- We introduce a novel positional encoding method, i.e., Interatomic Positional Encoding (IPE) to parameterize atomic environments in Transformer.
- By incorporating IPE, we propose a 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.
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 an 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  $\phi_v(\hat{r}_{ij})$  to describe the spatial relations between two atoms.  $\hat{r}_{ij}$  denotes the relative position pointing from atom  $i$  to atom  $j$ , and  $v$  indicates the functions' polynomial degree. ACE focuses on modeling the potential energy of a system by focusing on a 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 _ {i} = \sum_ {j _ {1}} \sum_ {v _ {1}} c _ {v _ {1}} \phi_ {v _ {1}} \left(\hat {r} _ {i j _ {1}}\right) + \sum_ {j _ {1} j _ {2}} \sum_ {v _ {1} v _ {2}} c _ {v _ {1} v _ {2}} \phi_ {v _ {1}} \left(\hat {r} _ {i j _ {1}}\right) \phi_ {v _ {2}} \left(\hat {r} _ {i j _ {2}}\right) + \dots \tag {1}
$$
with an unrestricted summation.  $c_{v}$  indicates the expansion coefficients. However, the sum of surrounding neighbors within the cluster would scale to  $\mathcal{O}(N^K)$  with  $K$  neighbors and become
Figure 1: Illustration of Interatomic Positional Encoding (IPE)  $C_{\eta}$ . Panel a depicts the extended Self-Attention with  $C_{\eta}$ .  $C_{\eta}$  contributes to the construction of attention weights while simultaneously being updated by atomic features; Panel b highlights the relationship between  $C_{\eta}$ , cluster  $\sigma_{i}$ , cluster  $\sigma_{j}$ , and the merged cluster  $\sigma_{ij}$ , as described in Theorem 1 and Equation 6 to Equation 11.
numerically expensive.  $N$  denotes the number of atoms within one molecule. ACE leverages the density trick to reduce the computational overhead. It defines atomic base  $A_{i,v}$  and  $A$ -basis  $A_{i,\mathbf{v}}$  as:
$$
A _ {i, v} = \sum_ {j \in N (i)} \phi_ {v} \left(\hat {r} _ {i j}\right) \tag {2}
$$
$$
A _ {i, \mathbf {v}} = \prod_ {t = 1} ^ {\epsilon} A _ {i, v _ {t}}, \quad \mathbf {v} = (v _ {1}, \dots , v _ {\epsilon}) \tag {3}
$$
where  $\epsilon$  denotes the order of body expansion, i.e.,  $(\epsilon + 1)$ -body expansion, and  $\mathbf{v}$  stands for the set of  $v$ . By firstly summing the neighboring basis functions and then applying multiplication, ACE can efficiently represent the atomic expansion scaling with the complexity  $\mathcal{O}(N)$ . Since  $A_{i,\mathbf{v}}$  is not rotationally invariant, we need additional Clebsch-Gordan coefficients  $C_v$  to construct fully permutation and isometry-invariant basis functions ( $B$ -basis), and describe the potential of cluster  $\sigma_i$  in Equation [1] as their linear combination with  $c$ -coefficients [21]:
$$
B _ {i, \mathbf {v}} = \sum_ {\mathbf {v} ^ {\prime}} C _ {\mathbf {v v} ^ {\prime}} A _ {i, \mathbf {v} ^ {\prime}} \tag {4}
$$
$$
E _ {i} = \sum_ {\epsilon} c _ {i, \mathbf {v}} B _ {i, \mathbf {v}} = \boldsymbol {c} _ {i} \cdot \boldsymbol {B} _ {i} \tag {5}
$$
In this section, we would introduce our Interatomic Positional Encoding (IPE) based on ACE theory in Section 3.1 and discuss how we integrate IPE into the Transformer architecture in 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 to encode in Transformers (termed as "positional encoding"). While most recent works have adopted pairwise distances between atoms as the relative PE, such a representation is often inadequate for capturing the complex interactions within molecules. As a result, it is essential to use a 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 1 Given two cluster  $\sigma_{i}$  and  $\sigma_{j}$  and their basis functions, there exists a set of invariant basis functions for the merged cluster  $\sigma_{ij}$  to describe integrated cluster potentials  $\tilde{E}_{ij}$ .
Proof. A merged cluster  $\sigma_{ij}$  can be represented as translating two clusters  $\sigma_i$  and  $\sigma_j$  such that their centered-atoms, i.e., atoms  $i$  and  $j$ , overlap as shown in Fig.  $\square (\mathrm{b})$ . All their neighbors are merged into a new cluster  $\sigma_{ij}$ , and the newly formed atomic base can be expressed as:
$$
\begin{array}{l} \tilde {A} _ {i j, v} = \left( \begin{array}{c c} A _ {i, v} & A _ {j, v} \end{array} \right) \otimes \left( \begin{array}{c c} A _ {i, v} & A _ {j, v} \end{array} \right) ^ {\top} \\ = \left( \begin{array}{l l} A _ {i, v} A _ {i, v} & A _ {i, v} A _ {j, v} \\ A _ {j, v} A _ {i, v} & A _ {j, v} A _ {j, v} \end{array} \right) \tag {6} \\ = \left( \begin{array}{c c} \sum_ {m _ {1} m _ {2}} \phi_ {v} (\hat {r} _ {i m _ {1}}) \phi_ {v} (\hat {r} _ {i m _ {2}}) & \sum_ {m n} \phi_ {v} (\hat {r} _ {i m}) \phi_ {v} (\hat {r} _ {j n}) \\ \sum_ {m n} \phi_ {v} (\hat {r} _ {j n}) \phi_ {v} (\hat {r} _ {i m}) & \sum_ {n _ {1} n _ {2}} \phi_ {v} (\hat {r} _ {j n _ {1}}) \phi_ {v} (\hat {r} _ {j n _ {2}}) \end{array} \right) \\ \end{array}
$$
with the product explicitly writing out.  $m,n$  are the neighbor atom symbols of atoms  $i,j$ , respectively.  $\otimes$  is the tensor product. We write the atomic base  $A_{i,v}A_{j,v}$  as  $A_{ij,v}$  (due to the permutational invariance, we have  $A_{ij,v} = A_{ji,v}$ ) which still follows the density trick. Therefore, we could construct a new  $A$ -basis for merged cluster  $\sigma_{ij}$  when taking product of  $\tilde{A}_{ij,v}$ :
$$
\begin{array}{l} \tilde {A} _ {i j, \mathbf {v}} = \tilde {A} _ {i j, v _ {1}} \odot \tilde {A} _ {i j, v _ {2}} \odot \dots \odot \tilde {A} _ {i j, v _ {\eta}} \\ = \left( \begin{array}{c c} A _ {i, \mathbf {v}} & A _ {j, \mathbf {v}} \end{array} \right) \otimes \left( \begin{array}{c c} A _ {i, \mathbf {v}} & A _ {j, \mathbf {v}} \end{array} \right) ^ {\top} \\ = \left( \begin{array}{l l} \prod_ {t = 1} ^ {2 \eta} A _ {i, v _ {t}} & \prod_ {t = 1} ^ {\eta} A _ {i j, v _ {t}} \\ \prod_ {t = 1} ^ {\eta} A _ {j i, v _ {t}} & \prod_ {t = 1} ^ {2 \eta} A _ {j, v _ {t}} \end{array} \right) \tag {7} \\ = \left( \begin{array}{c c} A _ {i (i), \mathbf {v}} & A _ {i j, \mathbf {v}} \\ A _ {j i, \mathbf {v}} & A _ {j (j), \mathbf {v}} \end{array} \right) \\ \end{array}
$$
where  $2\eta = \epsilon$  and  $\eta = 1,2,\ldots$ .  $\tilde{A}_{ij,\mathbf{v}}$  could describe  $(\epsilon +1)$ -body and  $(\epsilon +2)$ -body expansion simultaneously in  $\mathcal{O}(N)$ . To be concrete,  $A_{i(i),\mathbf{v}}$  contributes the  $(\epsilon +1)$ -body expansion, while  $A_{ij,\mathbf{v}}$  contributes the  $(\epsilon +2)$ -body expansion due to cluster merging. For instance, when taking  $\eta = 1$ ,  $A_{i,\mathbf{v}}$  and  $A_{i(i),\mathbf{v}}$  denote 2-body  $(im)$  and 3-body expansion  $(im_1m_2)$  in original cluster  $\sigma_i$ , and  $A_{ij,\mathbf{v}}$  denotes the 4-body expansion  $(mijn)$  in merged cluster  $\sigma_{ij}$ , respectively. Further explanation can be found in the following Remark.  $\tilde{A}_{ij,\mathbf{v}}$  exists when  $\epsilon \geq 2$ , which implies considering at least 4-body expansion within molecules. Then we could construct the corresponding matrix of  $B$ -basis  $\tilde{B}_{ij,\mathbf{v}}$  following Equation 4 and represent the potential of cluster  $\sigma_{ij}$  as:
$$
\tilde {E} _ {i j} = \sum_ {\eta} c _ {i j, v} \tilde {B} _ {i j, v} \tag {8}
$$
Remark. A straightforward illustration can be drawn by setting  $\upsilon = 1$  (Cartesian space) and  $\eta = 1$ . In this case, basis  $B_{i(i)}$  could be interpreted as the sum of cosine value of the surrounding angles within cluster  $\sigma_i$  [39, 43], i.e.,  $\sum_{m_1m_2}\cos \theta_{im_1m_2}$ , which represents the 3-body contributions in  $\mathcal{O}(N)$ . Similarly, basis  $B_{ij}$  could be treated as the sum of proper dihedral angles between two clusters  $\sigma_i$  and  $\sigma_j$  [49], i.e.,  $\sum_{mn}\cos \varphi_{mijn}$ , which represents 4-body contributions in  $\mathcal{O}(N)$ . The basis  $B_{ij}$  indeed serves as the contribution for torsion potential between two clusters in the original ACE, which primarily considers the contributions within a single cluster. The detailed proof could be found in Appendix B. As a result, incorporating  $A_{ij}$  effectively enhances the representation of interatomic relations by capturing the interactions between clusters, thereby offering a possible way for designing a geometric Transformer architecture.
Theorem 2 (Interactomic Positional Encoding (IPE)) Given one molecule with  $N$  atoms, there exists a positional encoding matrix  $\mathbf{C}_{\eta} \in \mathbb{R}^{N \times N}$ , which naturally describes the interatomic potentials.  $\eta$  denotes the orders of body expansion in Equation [7]. In particular,  $\mathbf{C}_{\eta}$  is directly multiplied with Query and Key before scaling. e.g., softmax, serving as the positional encoding in Transformer:
$$
\boldsymbol {\alpha} = \left(X W _ {Q}\right) \left(X W _ {K}\right) ^ {\top} \odot \boldsymbol {C} _ {\eta} \tag {9}
$$
where  $X \in \mathbb{R}^{N \times F}$  denotes the atomic features and  $W \in \mathbb{R}^{F \times F}$  denotes the learnable matrix.  $F$  is the hidden dimension.  $Q, K$  represent Query and Key, respectively.  $\odot$  is the Hadamard product.
(a) Geoformer Architecture
(b) Attention Block with IPE
(c) Embedding Layer
(d) Decoder
Figure 2: Geoformer Architecture. Geoformer consists of (c) an Embedding layer; (b) an elaborate Encoder incorporating  $L$  attention blocks with IPE  $C_{\eta}$  to extract geometric features and capture complex interatomic relationships within the molecular structure. The extended Self-Attention module with  $C_{\eta}$  is depicted in Fig. [1]; (d) a lightweight Decoder for predicting molecular properties of interest, such as energy and HOMO-LUMO gap.
Proof. As mentioned in Section 2 and Theorem 1, the basis functions  $A_{i,\mathbf{v}}$  in Equation 3 could represent the local chemical environment of cluster  $\sigma_{i}$  with density trick. Inspired by this, we first construct the  $A_{\mathbf{v}_{\tau}} = [A_{\mathbf{1},\mathbf{v}_{\tau}}, A_{\mathbf{2},\mathbf{v}_{\tau}}, \ldots, A_{\mathbf{N},\mathbf{v}_{\tau}}]^{\top}$  with integer  $\tau = [1, \eta]$  and  $A_{i,v_{\tau}} = \prod_{t=1}^{\tau} A_{i,v_t}$ . We then treat  $A_{\mathbf{v}_{\tau}}$  as the absolute positional encoding attached to the Query and Key, and attention matrix  $\alpha$  before scaling, e.g., softmax, could be modified as follows:
$$
\begin{array}{l} \boldsymbol {\alpha} _ {\tau} = \left(X W _ {Q} \circ \boldsymbol {A} _ {\mathbf {v} _ {\tau}}\right) \left(X W _ {K} \circ \boldsymbol {A} _ {\mathbf {v} _ {\tau}}\right) ^ {\top} \\ = \left(X W _ {Q}\right) \left(X W _ {K}\right) ^ {\top} \odot \left(\boldsymbol {A} _ {\mathbf {v} _ {\tau}} \boldsymbol {A} _ {\mathbf {v} _ {\tau}} ^ {\top}\right) \tag {10} \\ = (X W _ {Q}) (X W _ {K}) ^ {\top} \odot \sum_ {\mathbf {v} _ {\tau}} \tilde {A} _ {\mathbf {v} _ {\tau}} \\ \end{array}
$$
where  $\circ$  is the entry-wise Kronecker Product. The detailed proof could be found in Appendix B. Instructed by [10, 21], we add Clebsch-Gordan coefficients  $C_{\mathbf{v}_{\tau}}$  to ensure the rotationally invariance, and therefore obtain  $\tilde{B}_{\mathbf{v}_{\tau}} = \sum_{\mathbf{v}_{\tau}^{\prime}}C_{\mathbf{v}\mathbf{v}_{\tau}^{\prime}}\tilde{A}_{\mathbf{v}_{\tau}^{\prime}}$ . Then we can write the linear expansion:
$$
\boldsymbol {\alpha} = \left(X W _ {Q}\right) \left(X W _ {K}\right) ^ {\top} \odot \left(\sum_ {\tau = 1} ^ {\eta} W _ {\tilde {B}} \tilde {B} _ {\mathbf {v} _ {\tau}}\right) \tag {11}
$$
where  $W_{\tilde{B}}$  is a learnable weight matrix denoting  $c$ -coefficients, which is similar to the message construction in MACE [2]. Such operation could be done by a tensor broadcast. Eventually, our interatomic positional encoding  $C_{\eta}$  could be represented as:
$$
C _ {\eta} = \sum_ {\tau = 1} ^ {\eta} W _ {\tilde {B}} \tilde {B} _ {\mathbf {v} _ {\tau}} \tag {12}
$$
We can further modify the Equation 11 as Equation 9 and complete the proof.
3.2 Geoformer: Geometric Transformer for Molecules
Overall Design. The comprehensive structure of Geoformer is depicted in Figure 2. This design comprises an Embedding layer, an elaborate Encoder for extracting geometric features, and a lightweight Decoder for predicting molecular properties. Within the Geoformer's Encoder,  $L$  attention blocks are integrated, each having hidden dimension  $F$  and employing the proposed interatomic positional encoding  $C_{\eta}$ . Notably,  $C_{\eta}$  is updated by the atomic features  $X \in \mathbb{R}^{N \times F}$  within each block. Geoformer takes the atom type  $Z \in \mathbb{R}^{N}$  and atomic coordinates  $R \in \mathbb{R}^{N \times 3}$  from one molecule with  $N$  atoms as inputs, and produces the corresponding molecular properties as outputs.
**Embedding Layer.** The Embedding layer maps the atom type  $Z$  to  $X^0 \in \mathbb{R}^{N \times F}$ :
$$
X ^ {0} = \operatorname {L a y e r N o r m} (\operatorname {e m b e d} (Z)) \tag {13}
$$
and initialize the IPE  $C_{\eta}^{0}\in \mathbb{R}^{N\times N\times F}$  with 2-body expansion, i.e., radial basis functions (RBF) [45]:
$$
\mathrm {g} _ {k} (\hat {R}) = \phi (\| \hat {R} \|) \cdot \exp \left(- \beta_ {k} \left(\exp (- \| \hat {R} \|) - \mu_ {k}\right) ^ {2}\right) \tag {14}
$$
$$
C _ {\eta} ^ {0} = \mathbf {g} (\hat {R}) W _ {\mathrm {R B F}} \tag {15}
$$
where  $\hat{R} \in \mathbb{R}^{N \times N \times 3}$  denotes the relative position between two atoms,  $\|\cdot\|$  is the vector norm,  $\beta_{k}, \mu_{k}$  are optional learnable parameters that specify center and width of  $\mathrm{g}_{k}(\hat{R})$ , and  $\phi(\cdot)$  is a smooth cosine cutoff function.  $\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.  $W_{\mathrm{RBF}} \in \mathbb{R}^{K \times F}$  is a learnable matrix mapping basis functions to the hidden size.
Attention block with IPE. The extended attention block is illustrated in Fig. [1] (a). In contrast to the traditional Transformer, a learnable IPE matrix  $C_{\eta}$  is introduced to each attention block. Following TorchMD-NET [43], the softmax function is substituted with the SiLU activation to enhance accuracy, and the attention weight is scaled by a smooth cutoff:
$$
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) \odot C _ {\eta} ^ {l}\right)\right) \cdot \phi (\| \hat {R} \|) \tag {16}
$$
where  $*$  denotes the batched tensor product, i.e.,  $(X^l W_Q^l) * (X^l W_K^l)^\top \in \mathbb{R}^{N \times N \times F}$ .  $l$  indicates  $l$ -th attention block. In Fig. [L], 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 [5]. Then we produce the weighted values per atoms after self-attention to update  $C_{\eta}^l$ :
$$
\operatorname {A t t n} _ {V} \left(X ^ {l}\right) = A \left(X ^ {l}\right) \odot X ^ {l} W _ {V} ^ {l} \tag {17}
$$
Then under the instruction of Equation 7, we construct the  $A$ -basis for all merged cluster:
$$
\boldsymbol {A} _ {\mathbf {v} _ {\tau}} ^ {l} = \prod_ {t = 1} ^ {\tau} \sum_ {j} \operatorname {A t t n} _ {V} \left(X ^ {l}\right) W _ {A t t n} ^ {l} Y _ {l ^ {*}, m} (\hat {R} / \| \hat {R} \|) \tag {18}
$$
$$
\tilde {A} _ {\mathbf {v} _ {\tau}} ^ {l} = \left(A _ {\mathbf {v} _ {\tau}} ^ {l} W _ {\tilde {A} _ {\mathbf {v} _ {t} ^ {(1)}}} ^ {l}\right) \otimes \left(A _ {\mathbf {v} _ {\tau}} ^ {l} W _ {\tilde {A} _ {\mathbf {v} _ {t} ^ {(2)}}} ^ {l}\right) ^ {\top} \tag {19}
$$
where  $W_{Attn}^{l}$  is a learnable matrix,  $Y_{l^{*},m^{*}}(\hat{R} / \| \hat{R}\|)$  denotes the spherical harmonics with order  $l^*$  and degree  $m^{*}$ .  $W_{\tilde{A}_{v_t}^{(1)}}^l$  and  $W_{\tilde{A}_{v_t}^{(2)}}^l$  are two learnable matrix without bias to ensure equivariance [39]. We could further construct a new form of residual IPE within each block following Equation 8 and 11:
$$
\delta C _ {\eta} ^ {l} = \sum_ {\tau = 1} ^ {\eta} W _ {\tilde {B}} ^ {l} \sum_ {\mathbf {v} _ {\tau} ^ {\prime}} C _ {\mathbf {v v} _ {\tau} ^ {\prime}} \tilde {A} _ {\mathbf {v} _ {\tau} ^ {\prime}} \tag {20}
$$
where  $W_{\tilde{B}}^{l}$  is a learnable matrix. Finally we apply the residual connection to compute IPE for the next block:
$$
\boldsymbol {C} _ {\eta} ^ {l + 1} = \operatorname {S i L U} \left(\boldsymbol {C} _ {\eta} ^ {l} W _ {\boldsymbol {C}} ^ {l}\right) \odot \delta \boldsymbol {C} _ {\eta} ^ {l} + \boldsymbol {C} _ {\eta} ^ {l} \tag {21}
$$
where  $W_{C}^{l}$  is a learnable matrix and  $\mathrm{SiLU}(C_{\eta}^{l}W_{C}^{l})$  plays a role of gated filter. The update of atomic features still follows the traditional procedures:
$$
\operatorname {A t t n} \left(X ^ {l}\right) = \sum_ {j} \operatorname {A t t n} _ {V} \left(X ^ {l}\right) + X ^ {l} \tag {22}
$$
$$
\operatorname {F F N} \left(X ^ {l}\right) = \operatorname {S i L U} \left(X ^ {l} W _ {1} ^ {l}\right) W _ {2} ^ {l} + X ^ {l} \tag {23}
$$