SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Guillem Simeon
Computational Science Laboratory
Universitat Pompeu Fabra
guillem.simeon@gmail.com
Gianni De Fabritiis
Icrea, Universitat Pompeu Fabra, Acellera
g.defabritiis@gmail.com
The development of efficient machine learning models for molecular systems representation is becoming crucial in scientific research. We introduce TensorNet, an innovative O(3)-equivariant message-passing neural network architecture that leverages Cartesian tensor representations. By using Cartesian tensor atomic embeddings, feature mixing is simplified through matrix product operations. Furthermore, the cost-effective decomposition of these tensors into rotation group irreducible representations allows for the separate processing of scalars, vectors, and tensors when necessary. Compared to higher-rank spherical tensor models, TensorNet demonstrates state-of-the-art performance with significantly fewer parameters. For small molecule potential energies, this can be achieved even with a single interaction layer. As a result of all these properties, the model's computational cost is substantially decreased. Moreover, the accurate prediction of vector and tensor molecular quantities on top of potential energies and forces is possible. In summary, TensorNet's framework opens up a new space for the design of state-of-the-art equivariant models.
Interatomic potential modeling using neural networks is an emerging research area that holds great promise for revolutionizing molecular simulation and drug discovery pipelines [1; 2; 3; 4]. The conventional trade-off between accuracy and computational cost can be bypassed by training models on highly precise data [5; 6; 7; 8]. Current state-of-the-art methodologies rely on equivariant graph neural networks (GNNs) [9; 10] and message-passing neural network (MPNNs) frameworks [11; 12], where internal atomic representations incorporate well-defined transformation properties characteristic of physical systems. The integration of equivariant features into neural network interatomic potentials has led to remarkable improvements in accuracy, particularly when using higher-rank irreducible representations of the orthogonal group  $\mathrm{O}(3)$ —which encompasses reflections and rotations in 3D space—in the form of spherical tensors [13; 14]. Although lower-rank Cartesian representations (calvars and vectors) have been employed [15; 16], their success has been limited compared to state-of-the-art spherical models [17; 18; 19]. MPNNs typically necessitate a substantial number of message-passing iterations, and models based on irreducible representations are generally computationally demanding due to the need to compute tensor products, even though some successful alternative has been put forward [20].
The pursuit of computationally efficient approaches for incorporating higher-rank equivariance is essential. In this paper, we introduce a novel O(3)-equivariant architecture that advances the integration of Cartesian representations by utilizing Cartesian rank-2 tensors, represented as 3x3 matrices. We demonstrate that this method achieves state-of-the-art performance comparable to higher-rank spherical models while having a reduced computational cost. This efficiency is realized
through the cheap decomposition of Cartesian rank-2 tensors into their irreducible components under the rotation group and the ability to mix features using straightforward  $3 \times 3$  matrix products. Additionally, our model requires fewer message-passing steps and eliminates the need for explicit construction of many-body terms. The architecture also facilitates the direct prediction of tensor quantities, enabling the modeling of molecular phenomena where such quantities are relevant. In summary, we propose an alternative framework for developing efficient and accurate equivariant models.
Equivalence. A function  $f$  between two vector spaces  $X$  and  $Y$ ,  $f: X \to Y$ , is said to be equivariant to the action of some abstract group  $G$  if it fulfills
$$
D _ {Y} [ g ] f (x) = f \left(D _ {X} [ g ] x\right), \tag {1}
$$
for all elements  $g \in G$  and  $x \in X$ , where  $D_X[g]$  and  $D_Y[g]$  denote the representations of  $g$  in  $X$  and  $Y$ , respectively. Equivariant neural networks used in neural network potentials focus on equivariance under the action of translations and the orthogonal group  $\mathrm{O}(3)$  in  $\mathbb{R}^3$ , the latter one being comprised by the rotation group  $\mathrm{SO}(3)$  and reflections, and regarded as a whole as the Euclidean group  $E(3)$ .
Cartesian tensors and irreducible tensor decomposition. Tensors are algebraic objects that generalize the notion of vectors. In the same way, as vectors change their components with respect to some basis under the action of a rotation  $R \in \mathrm{SO}(3)$ ,  $\mathbf{v}' = R\mathbf{v}$ , a rank- $k$  Cartesian tensor  $T$  can be (very) informally regarded as a multidimensional array with  $k$  indices, where each index transforms as a vector under the action of a rotation. In particular, a rank-2 tensor transformation under a rotation can be written in matrix notation as  $T' = RTR^{\mathrm{T}}$ , where  $R^{\mathrm{T}}$  denotes the transpose of the rotation matrix  $R$ . In this paper, we will restrict ourselves to rank-2 tensors. Moreover, any rank-2 tensor  $X$  defined on  $\mathbb{R}^3$  can be rewritten in the following manner [21]
$$
X = \frac {1}{3} \operatorname {T r} (X) \mathrm {I d} + \frac {1}{2} \left(X - X ^ {\mathrm {T}}\right) + \frac {1}{2} \left(X + X ^ {\mathrm {T}} - \frac {2}{3} \operatorname {T r} (X) \mathrm {I d}\right), \tag {2}
$$
where  $\operatorname{Tr}(X) = \sum_{i} X_{ii}$  is the trace operator and  $\operatorname{Id}$  is the identity matrix. The first term is proportional to the identity matrix, the second term is a skew-symmetric contribution, and the last term is a symmetric traceless contribution. It can be shown that expression (2) is a decomposition into separate representations that are not mixed under the action of the rotation group [21]. In particular, the first component  $I^{X} \equiv \frac{1}{3} \operatorname{Tr}(X) \operatorname{Id}$  has only 1 degree of freedom and is invariant under rotations, that is, it is a scalar; the second term  $A^{X} \equiv \frac{1}{2} (X - X^{\mathrm{T}})$  has 3 independent components since it is a skew-symmetric tensor, which can be shown to rotate as a vector; and  $S^{X} \equiv \frac{1}{2} (X + X^{\mathrm{T}} - \frac{2}{3} \operatorname{Tr}(X) \operatorname{Id})$  rotates like a rank-2 tensor and has 5 independent components, since a symmetric tensor has six independent components but the traceless condition removes one degree of freedom. In terms of representation theory, the 9-dimensional representation (a 3x3 matrix) has been reduced to irreducible representations of dimensions 1, 3, and 5 [21]. We will refer to  $X$  as a full tensor and to the components  $I^{X}, A^{X}, S^{X}$  as scalar, vector, and tensor features, respectively.
Message passing neural network potentials. Message-passing neural networks (MPNNs) have been successfully applied to the prediction of molecular potential energies and forces [11]. Atoms are represented by graph nodes, which are embedded in three-dimensional Euclidean space, and edges between nodes are built according to their relative proximity after the definition of some cutoff radius. The neural network uses atomic and geometric information, such as distances, angles or relative position vectors, to learn useful node representations by recursively propagating, aggregating, and transforming features from neighboring nodes [22; 23; 15]. In the case of neural network potentials, after several rounds of message passing and feature transformations, node features are mapped to single per-atom scalar quantities which are atomic contributions to the total energy of the molecule. These energy contributions depend in a very complex way on the states of other atoms, and therefore MPNNs can be regarded as some learnable approximation to the many-body potential energy function. However, these neural networks have typically needed a substantially large amount of message-passing steps (up to 6 in some cases) [16; 17].
Equivariant models. Initially, since the potential energy is a scalar quantity, atomic features were built using geometric information which is invariant to translations, reflections, and rotations, such as in SchNet [22], DimeNet [24; 25], PhysNet [26], SphereNet [27] and GemNet [23]. Nevertheless,
it has been shown that the inclusion of equivariant internal features leads to substantially better performances and data efficiency [14; 28]. In equivariant GNNs, internal features transform in a specified way under some group action. Molecules are physical systems embedded in three-dimensional Euclidean space, and their properties display well-defined behaviors under transformations such as translations, reflections, and rotations. Therefore, when predicting molecular properties, the group of interest is the orthogonal group in three dimensions  $\mathrm{O}(3)$ , that is, rotations and reflections of the set of atoms in 3D space.
In models such as NewtonNet [29], EGNN [30], PaiNN [15], the Equivariant Transformer [16] and SAKE [31], Cartesian vector features are used on top of invariant features. These vector features are built using relative position vectors between input atoms, in such a way that when the input atomic Cartesian coordinates  $\mathbf{r}$  are transformed under the action of some  $R\in \mathrm{O}(3)$  represented by a 3x3 matrix,  $\mathbf{r}\rightarrow R\mathbf{r}$ , internal vector features and vector outputs  $\mathbf{v}$  transform accordingly,  $\mathbf{v}\rightarrow R\mathbf{v}$ . Other models such as Cormorant [13], Tensor Field Networks [32], NequIP [17], Allegro [18], BOTNet [19] and MACE [20], work directly with internal features that are irreducible representations of the group  $\mathrm{O}(3)$ , which can be labeled by some  $l\in \mathbb{N}$  (including  $l = 0$ ), and with dimensions  $2l + 1$ . The representations  $l = 0$  and  $l = 1$  correspond to scalars and vectors, respectively. In this case, under a transformation  $R\in \mathrm{O}(3)$  of the input coordinates  $\mathbf{r}\rightarrow R\mathbf{r}$ , internal features  $h_{lm}(\mathbf{r})$  transform as  $h_{lm}(R\mathbf{r}) = \sum_{m'}D_{m'm'}^l (R)h_{lm'}(\mathbf{r})$ , where  $D_{m'm'}^l (R)\in \mathbb{R}^{(2l + 1)\times (2l + 1)}$  is an order  $l$  Wigner D-matrix. In this case, features are rank- $l$  spherical tensors or pseudotensors, depending on their parity. The decomposition of a Cartesian tensor described in (2) and the irreducible representations in terms of spherical tensors are directly related by a change of basis [21]. To generate new features that satisfy O(3)-equivariance, these are built by means of tensor products involving Clebsch-Gordan coefficients and parity selection rules. In particular, models with features  $l > 1$  such as NequIP, Allegro, BOTNet, and MACE have achieved state-of-the-art performances in benchmark datasets in comparison to all other MPNNs. However, the computation of tensor products in most of these models containing higher-rank tensors and pseudotensors can be expensive, especially when computing them in an edge-wise manner.
3 TensorNet's architecture
3.1 Operations respecting O(3)-equivariance
In this work, we propose the use of the 9-dimensional representation of rank-2 tensors (3x3 matrices). TensorNet operations are built to satisfy equivariance to the action of the orthogonal group  $\mathrm{O}(3)$ : equivariance under  $\mathrm{O}(3)$  instead of the subgroup of rotations  $\mathrm{SO}(3)$  requires the consideration of the differences between tensors and pseudotensors. Tensors and pseudotensors are indistinguishable under rotations but display different behaviors under parity transformation, i.e. a reflection of the coordinate system through the origin. By definition, scalars, rank-2 tensors and in general all tensors of even rank do not flip their sign, and their parity is said to be even; on the contrary, vectors and tensors of odd rank have odd parity and flip their sign under the parity transformation. Pseudoscalars, pseudovectors, and pseudotensors have precisely the opposite behavior. Necessary derivations for the following subsections can be found in the Appendix (section A.2).
Composition from irreducible representations. The previously described irreducible decomposition of a tensor in Eq. 2 is with respect to the rotation group  $\mathrm{SO}(3)$ . Building a tensor-like object that behaves appropriately under rotations can be achieved by composing any combination of scalars, vectors, tensors, and their parity counterparts. However, in neural network potential settings, the most direct way to produce tensors is by means of relative position vectors and, in general, it is preferred for the neural network to be able to predict vectors rather than pseudovectors. One has the possibility to initialize full tensor representations from the composition of scalars, vectors encoded in skew-symmetric matrices, and symmetric traceless tensors. For instance, if one considers some vector  $\mathbf{v} = (v^{x}, v^{y}, v^{z})$ , one can build a well-behaved tensor  $X$  under rotations by composing  $X = I + A + S$ ,
$$
I = f (\left\| \mathbf {v} \right\|) \mathrm {I d}, \quad A = \left( \begin{array}{c c c} 0 & v ^ {z} & - v ^ {y} \\ - v ^ {z} & 0 & v ^ {x} \\ v ^ {y} & - v ^ {x} & 0 \end{array} \right), \quad S = \mathbf {v} \mathbf {v} ^ {\mathrm {T}} - \frac {1}{3} \operatorname {T r} \left(\mathbf {v} \mathbf {v} ^ {\mathrm {T}}\right) \mathrm {I d}, \tag {3}
$$
where  $f$  is some function and  $\mathbf{v}\mathbf{v}^{\mathrm{T}}$  denotes the outer product of the vector with itself. In this case, under parity the vector transforms as  $\mathbf{v} \rightarrow -\mathbf{v}$ , and it is explicit that  $I$  and  $S$  remain invariant, while  $A \rightarrow -A = A^{\mathrm{T}}$ , and the full tensor  $X$  transforms as  $X = I + A + S \rightarrow X' = I + A^{\mathrm{T}} + S = X^{\mathrm{T}}$ , since  $I$  and  $S$  are symmetric matrices. Therefore, one concludes that when initializing the skew-symmetric part  $A$  from vectors, not pseudovectors, parity transformation produces the transposition of full tensors.
Invariant weights and linear combinations. One can also modify some tensor  $X = I + A + S$  by multiplying invariant quantities to the components,  $X' = f_{I}I + f_{A}A + f_{S}S$ , where  $f_{I}, f_{A}$  and  $f_{S}$  can be constants or invariant functions. This modification of the tensor does not break the tensor transformation rule under the action of rotations and preserves the parity of the individual components given that  $f_{I}, f_{A}$  and  $f_{S}$  are scalars (learnable functions of distances or vector norms, for example), not pseudoscalars. Also, from this property and the possibility of building full tensors from the composition of irreducible components, it follows that linear combinations of scalars, vectors, and tensors generate new full tensor representations that behave appropriately under rotations. Regarding parity, linear combinations preserve the original parity of the irreducible components given that all terms in the linear combination have the same parity. Therefore, given a set of irreducible components  $I_{j}, A_{j}, S_{j}$  with  $j \in \{0, 1, \dots, n - 1\}$ , one can build full tensors  $X_{i}'$
$$
X _ {i} ^ {\prime} = \sum_ {j = 0} ^ {n - 1} w _ {i j} ^ {I} I _ {j} + \sum_ {j = 0} ^ {n - 1} w _ {i j} ^ {A} A _ {j} + \sum_ {j = 0} ^ {n - 1} w _ {i j} ^ {S} S _ {j}, \tag {4}
$$
where  $w_{ij}^{I}, w_{ij}^{A}, w_{ij}^{S}$  can be learnable weights, in which case the transformation reduces to the application of three different linear layers without biases to inputs  $I_{j}, A_{j}, S_{j}$ .
Matrix product. Consider two tensors,  $X$  and  $Y$ , and some rotation matrix  $R \in \mathrm{SO}(3)$ . Under the transformation  $R$ , the tensors become  $RXR^{\mathrm{T}}$  and  $RYR^{\mathrm{T}}$ . The matrix product of these tensors gives a new object that also transforms like a tensor under the transformation,  $XY \to RXR^{\mathrm{T}}RYR^{\mathrm{T}} = RXR^{-1}RYR^{\mathrm{T}} = R(XY)R^{\mathrm{T}}$ , since for any rotation matrix  $R$ ,  $R^{\mathrm{T}} = R^{-1}$ . Taking into account their irreducible decomposition  $X = I^{X} + A^{X} + S^{X}$  and  $Y = I^{Y} + A^{Y} + S^{Y}$ , the matrix product  $XY$  consists of several matrix products among rotationally independent sectors ( $I^{X} + A^{X} + S^{X})(I^{Y} + A^{Y} + S^{Y})$ ). These products will contribute to the different parts of the irreducible decomposition  $XY = I^{XY} + A^{XY} + S^{XY}$ . Therefore, one can regard the matrix product as a way of combining scalar, vector, and tensor features to obtain new features. However, when assuming that the skew-symmetric parts are initialized from vectors, this matrix product mixes components with different parities, and resulting components  $I^{XY}, A^{XY}, S^{XY}$  would not have a well-defined behavior under parity (see Appendix, section A.2). To achieve O(3)-equivariance, we propose the use of the matrix products  $XY + YX$ . Under parity  $X \to X^{\mathrm{T}}, Y \to Y^{\mathrm{T}}$ , and one can show that
$$
I ^ {X ^ {\mathrm {T}} Y ^ {\mathrm {T}} + Y ^ {\mathrm {T}} X ^ {\mathrm {T}}} = I ^ {X Y + Y X}, A ^ {X ^ {\mathrm {T}} Y ^ {\mathrm {T}} + Y ^ {\mathrm {T}} X ^ {\mathrm {T}}} = - A ^ {X Y + Y X}, S ^ {X ^ {\mathrm {T}} Y ^ {\mathrm {T}} + Y ^ {\mathrm {T}} X ^ {\mathrm {T}}} = S ^ {X Y + Y X}, \tag {5}
$$
that is, the scalar and symmetric traceless parts have even parity, and the skew-symmetric part has odd parity. The irreducible decomposition of the expression  $XY + YX$  preserves the rotational and parity properties of the original components and, therefore, it is an O(3)-equivariant operation. We finally note that one can produce O(3)-invariant quantities from full tensor representations or their components by taking their Frobenius norm  $\operatorname{Tr}(X^{\mathrm{T}}X) = \operatorname{Tr}(XX^{\mathrm{T}}) = \sum_{ij} |X_{ij}|^2$ .
3.2 Model architecture
In this work we propose a model that learns a set of Cartesian full tensor representations  $X^{(i)}$  (3x3 matrices) for every atom  $(i)$ , from which atomic or molecular properties can be predicted, using as inputs atomic numbers  $z_{i}$  and atomic positions  $\mathbf{r}_i$ . We mainly focus on the prediction potential energies and forces, even though we provide in Section 4 experiments demonstrating the ability of TensorNet to accurately predict up to rank-2 physical quantities. Atomic representations  $X^{(i)}$  can be decomposed at any point into scalar, vector and tensor contributions  $I^{(i)}, A^{(i)}, S^{(i)}$  via (2), and TensorNet can be regarded as operating with a physical inductive bias akin to the usual decomposition of interaction energies in terms of monopole, dipole and quadrupole moments [13]. We refer the reader to Figure 1 and the Appendix (section A.1) for diagrams of the methods and the architecture.
**Embedding.** By defining a cutoff radius  $r_c$ , we obtain vectors  $\mathbf{r}_{ij} = \mathbf{r}_j - \mathbf{r}_i$  between central atom  $i$  and neighbors  $j$  within a distance  $r_c$ . We initialize per-edge scalar features using the identity
matrix  $I_0^{(ij)} = \mathrm{Id}$ , and per-edge vector and tensor features using the normalized edge vectors  $\hat{r}_{ij} = \mathbf{r}_{ij} / ||\mathbf{r}_{ij}|| = (\hat{r}_{ij}^x, \hat{r}_{ij}^y, \hat{r}_{ij}^z)$ . We create a symmetric traceless tensor from the outer product of  $\hat{r}_{ij}$  with itself,  $S_0^{(ij)} \equiv \hat{r}_{ij}\hat{r}_{ij}^{\mathrm{T}} - \frac{1}{3}\mathrm{Tr}(\hat{r}_{ij}\hat{r}_{ij}^{\mathrm{T}})\mathrm{Id}$ , and vector features are initialized by identifying the independent components of the skew-symmetric contribution with the components of  $\hat{r}_{ij}$  as denoted in (3), getting for every edge  $(ij)$  initial irreducible components  $I_0^{(ij)}, A_0^{(ij)}, S_0^{(ij)}$ . To
Figure 1: Key steps, from top to bottom, in the embedding and interaction modules for some central atom  $i$  and neighbors  $j$  and  $k$  found within the cutoff radius. a) Relative position vectors are used to initialize edge-wise tensor components, modified using edge-wise invariant functions, and summed to obtain node-wise full tensors. b) Node full tensors are decomposed and weighted with edge invariant functions to obtain pair-wise messages, and summed to obtain node-wise aggregated messages, which will interact with receiving node's full tensors via matrix product.
encode interatomic distance and atomic number information in the tensor representations we use an embedding layer that maps the atomic number of every atom  $z_{i}$  to  $n$  invariant features  $Z_{i}$ , and expand interatomic distances  $r_{ij}$  to  $d$  invariant features by means of an expansion in terms of exponential radial basis functions
$$
e _ {k} ^ {\mathrm {R B F}} \left(r _ {i j}\right) = \exp \left(- \beta_ {k} \left(\exp \left(- r _ {i j}\right) - \mu_ {k}\right) ^ {2}\right), \tag {6}
$$
where  $\beta_{k}$  and  $\mu_{k}$  are fixed parameters specifying the center and width of radial basis function  $k$ . The  $\mu$  vector is initialized with values equally spaced between  $\exp(-r_c)$  and 1, and  $\beta$  is initialized as  $(2d^{-1}(1 - \exp(-r_c)))^{-2}$  for all  $k$  as proposed in [26]. After creating  $n$  identical copies of initial components  $I_0^{(ij)}, A_0^{(ij)}, S_0^{(ij)}$  ( $n$  feature channels), for every edge ( $ij$ ) we map with a linear layer the concatenation of  $Z_i$  and  $Z_j$  to  $n$  pair-wise invariant representations  $Z_{ij}$ , and the radial basis functions are further expanded to  $n$  scalar features by using three different linear layers to obtain
$$
f _ {I} ^ {0} = W ^ {I} (e ^ {\mathrm {R B F}} (r _ {i j})) + b ^ {I}, f _ {A} ^ {0} = W ^ {A} (e ^ {\mathrm {R B F}} (r _ {i j})) + b ^ {A}, f _ {S} ^ {0} = W ^ {S} (e ^ {\mathrm {R B F}} (r _ {i j})) + b ^ {S}, \quad (7)
$$
$$
X ^ {(i j)} = \phi \left(r _ {i j}\right) Z _ {i j} \left(f _ {I} ^ {0} I _ {0} ^ {(i j)} + f _ {A} ^ {0} A _ {0} ^ {(i j)} + f _ {S} ^ {0} S _ {0} ^ {(i j)}\right), \tag {8}
$$
where the cutoff function is given by  $\phi (r_{ij}) = \frac{1}{2}\bigl (\cos \bigl (\frac{\pi r_{ij}}{r_c}\bigr) + 1\bigr)$  when  $r_{ij}\leq r_c$  and 0 otherwise. That is,  $n$  edge-wise tensor representations  $X^{(ij)}$  are obtained, where the channel dimension has not been written explicitly. Then, we get atom-wise tensor representations by aggregating all neighboring edge-wise features. At this point, the invariant norms  $\| X\| \equiv \mathrm{Tr}(X^{\mathrm{T}}X)$  of atomic representations  $X^{(i)}$  are computed and fed to a normalization layer, a multilayer perceptron, and a SiLU activation to obtain three different O(3)-invariant functions per channel,
$$
f _ {I} ^ {(i)}, f _ {A} ^ {(i)}, f _ {S} ^ {(i)} = \operatorname {S i L U} (\operatorname {M L P} (\operatorname {L a y e r N o r m} (| | X ^ {(i)} | |))), \tag {9}
$$
which, after the decomposition of tensor embeddings into their irreducible representations, are then used to modify component-wise linear combinations to obtain the final atomic tensor embeddings
$$
X ^ {(i)} \leftarrow f _ {I} ^ {(i)} W ^ {I} I ^ {(i)} + f _ {A} ^ {(i)} W ^ {A} A ^ {(i)} + f _ {S} ^ {(i)} W ^ {S} S ^ {(i)}. \tag {10}
$$
Interaction and node update. We start by normalizing each node's tensor representation  $X^{(i)} \gets X^{(i)} / (||X^{(i)}|| + 1)$  and decomposing this representation into scalar, vector, and tensor features. We next transform these features  $I^{(i)}, A^{(i)}, S^{(i)}$  by computing independent linear combinations,  $Y^{(i)} = W^{I}I^{(i)} + W^{A}A^{(i)} + W^{S}S^{(i)}$ . In parallel, edge distances' radial basis expansions are fed to a multilayer perceptron and a SiLU activation to transform them into tuples of three invariant functions per channel weighted with the cutoff function  $\phi(r_{ij})$ ,
$$
f _ {I} ^ {(i j)}, f _ {A} ^ {(i j)}, f _ {S} ^ {(i j)} = \phi \left(r _ {i j}\right) \text {S i L U} \left(\operatorname {M L P} \left(e ^ {\mathrm {R B F}} \left(r _ {i j}\right)\right)\right). \tag {11}
$$