TensorNet: Cartesian Tensor Representations for Efficient Learning of Molecular Potentials
Authors: Guillem Simeon
Gianni De Fabritiis
Venue: NeurIPS 2023 poster (2023)
Keywords: Neural network interatomic potentials
Equivariant graph neural network
Message passing neural network

Abstract: 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.

TensorNet: Cartesian Tensor Representations for Efficient Learning of Molecular Potentials
Guillem Simeon
Gianni De Fabritiis Computational Science Laboratory Icrea, Universitat Pompeu Fabra, Acellera g w.deebriiiis@@aaal com
Computational Science Laboratory Universitat Pompeu Fabra guillem. simeon@gmail.c com
Abstract
The development of efficient machine learning models for molecular systems representation is becoming crucial in scientific research. We introduce Tensor- Net, innovative 0(3) eqquivaria that leverages Cartesian tensor representations. By using Cartesian tensor atomic embeddings, feature mixing is simplified through matrix product operations. Fur- thermore, 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-an performance with significantly fewer pa- rameters. For small molecule potential energies, this can be achieved even with single interaction layer. As a result of all these properties, the model's com- putational 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 's framework opens up a new space for the design state-of-the-an equivariant models.
cesssage-psssing neural network architecture
Introduction
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
g Current state-of offie--o-the-arr the-art art methodologies rely on equivariant
D 10] and message- message-passing neural network (MPNNs) frameworks [2] 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 0(3) which encompasses reflections
graph neural networks (GNNs)
and rotations in 3D space- in the form of spherical tensors [13 [4] Although lower-rank Cartesian
representations (scalars and vectors) have been employed 15 [6], their success has been limited
compared to -tate-of-the-art spherical models I7 19). MPNNs typically necessitate substantial number of cesssage-pasiigg iterations, and models based on irreducible representations are generally computationally demanding due to the need compute tensor products, even though some successful alternative has been put forward (20).
The pursuit of computationally efficient approaches for incorporating higher-rank equivariance essential. In this paper, we introduce novel -3)eequuivari 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 reduced computational cost. This efficiency is realized
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

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 3x3 matrix products. Additionally, our model requires fewer message-passoe 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.
2 Background and related work
Equivariance. A function f between two vector spaces X and Y, f X + Y, is said to be equivariant to the action of some abstract group if fulfills
(I) for all elements E G and x E X, where Dx [g] and Dy[g] denote the representations of 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 O(3) in R. the latter one being comprised by the rotation group (3) and reflections, and regarded as a whole as the Euclidean group E3).
Dylglf(z)= = [DDx(g)=).
Cartesian tensors and irreducible tensor decomposition. Tensors are algebraic objects that generalize the notion of vectors. In the same way, vectors change their components with respect to some basis under the action a rotation R E SO(3), v' = Rv, a rank-k Cartesian tensor T can be multidimensional array with & indices, where each index transforms
(very) informally regarded 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 T" = R TR" where RT 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 R can be rewritten in the following manner (1
1 (xx)dd 3
2 xx
2 - xx) + - (x()dd))
1
X
3
where Tx() E, X,; is the trace operator and 1d 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 symmetric traceless contribution. It can be shown that expression 2 is decomposition into separate representations that are not mixed under the action of the rotation group [21]. In particular, the first component x Tr(X))l has only 1 degree of freedom and is invariant under rotations, that is, it a scalar; the second term A [ - X has 3 independent components since is skew-symmetric tensor, which can be shown to rotate as a vector; and 5 [ +X X 2 Tr())dd rotates like 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 as a full tensor and to the components x Ax S scalar, vector, and tensor features, respectively.
Message passing neural network potentials. Message passgge-painnggggg neural networks (MPNNs) have been successfully applied to the prediction of molecular potential energies and forces [1] Atoms are represented by graph nodes, which are embedded 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 very complex way the states of other atoms, and therefore MPNNs can be regarded as some learnable approximation to the sany-bodyy potential energy function. However, these neural networks have typically needed substantially large amount of seassage-passing steps (up to 6 in some cases) [7]
diree-dimensionl Euclidean space, and
Equivariant models. Initially, since the potential energy is scalar quantity, atomic features were built using geometric information which invariant to translations, reflections, and rotations, such as in SchNet (22), DimeNet (24 225, 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 22) ified way under some group action. Molecules are physical systems embedded in three-dimensional Fuclidean space, and their properties display well-defned 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 O(3), that is, rotations and reflections of the set of atoms in 3D space.
28]. In equivariant GNNs, internal features transform in a spec-
In models such as NewtonNet 29 EGNN B0), PaiNN [5], the Equivariant Transformer [6 and SAKE BTI Cartesian vector features are used on top of invariant features. These vector features are built using relative position vectors between input atoms, such a way that when the input atomic Cartesian coordinates are transformed under the action of some R E O(3) represented by a 3x3 matrix, r - Rr, internal vector features and vector outputs v transform accordingly, v -> Rv.. Other models such as Cormorant [33, Tensor Field Networks (32) NequlP [17]. Allegro [8], BOTNet 19 and MACE (20), work directly with internal features that are irreducible representations of the group (3), which can be labeled by some / N (including / 0), and with dimensions 2/ + 1. The representations = and / correspond to scalars and vectors, respectively. In this case, under transformation R E (3) of the input coordinates r + Rr, internal features hum. (r) transform as him (Rt) Em D D'm (R) (r). where D' R2) E [2(2+1)x(2(+11) is an order 1 Wigner D-matrix. In this case, features are rank-l spherical tensors or pscudotensors, depending on their parity. The decomposition of Cartesian tensor described in (2) and the irreducible representations in terms of spherical tensors are directly related by a change basis [2[]. To generate new features that satisfy (33--vvivarriane,, these are built by means of tensor products involving Clebsch- coefficients and parity selection rules. In particular, models with features Allegro, BOTNet, and MACE have achieved state-of oftee----hhhrrrr 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.
1 such as NequlP,
3. TensorNet's architecture
3.1 Operations respecting 0(3)-qquivariance
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 O(3): equivariance under 0(3) instead of the subgroup of rotations 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; 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 decomposi- tion of a tensor in Eq. i is with respect to the rotation group 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 (v", v², v²), one can build well-behaved tensor X under rotations by composing X = I A+S.
0
2
S vV - - Ir(v
1
1 1= =[f([[]]] A
21
(3)
3

---

where f is some function and vv" denotes the outer product of the vector with itself. In this case, under parity the vector transforms as v + -v, and is explicit that I and S remain invariant, while A + A = AT, and the full tensor X transforms as X 7 + A + S - X' = T A + = X", 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 -1+ S by multiplying invariant quantities to the components, X' = = fAA+ + fsS, where fA and fs 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 fA and fs 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 Ins Ag. S3 with j E {0,1,... I.., 1}. one can build full tensors X;
X = Σ
Σ w j=0
w S
(4)
where why why wi can be learnable weights, in which case the transformation reduces to the application of three different linear layers without biases to inputs 1,, Aj, Sy.
Matrix product. Consider two tensors, X and Y, and some rotation matrix R E (O))). Under the transformation R. the tensors become RXR R and R RYRR RT. The matrix product of these tensors gives new object that also transforms like tensor under the transformation, XY - RX R. RRRRRRRRRRRRRR RY
RXR RY HT = R(XY)R" since for any rotation matrix R. R = R- Taking into account their irreducible decomposition X fx + AX + S and Y - AY + SY, the matrix product XY consists of several matrix products among rotationally independent sectors (IX + A + SX + +AY + +.S)). These products will contribute to the different parts of the irreducible decomposition XY = 1 + AAXY + 5XY, 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 akew-symmetric parts are initialized from vectors, this matrix product mixes components with different parities, and resulting components J AAXY, ,XX would not have well-defined behavior under parity (see Appendix, section A.2). To achieve 0(3) -quiivvarve we propose the use of the matrix products XY + YX. Under parity X + X ,. + Y and one can show that
XXTyT=yTX* = XX A X
= - X SXTYT
= -XX (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) 0(3)-equivariant operation. We finally note that one can produce 0(3) ((3)-invariant quantities from full tensor representations or their components by taking their Frobenius norm Tr(X = Tr(XXT) = Eij |X,]
3.2 Model architecture
In this work we propose model that learns set of Cartesian full tensor representations ( (3x3 matrices) for every atom (i), from which atomic or molecular properties can be predicted, using as inputs atomic numbers 24 and atomic positions r;. 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 (() can be decomposed at any point into scalar, vector and tensor contributions I(9), A 5(() via 2 and TensorNet can be regarded as operating with physical inductive bias akin to the usual decomposition of interaction energies in terms of monopole, dipole and quadrupole moments [3.. We refer the reader Figure and the Appendix (section A. for diagrams of the methods and the architecture.
Embedding. By defining a cutoff radius r., we obtain vectors rú = r r; between central atom i and neighbors j within distance Te. We initialize per-edge scalar features using the identity

---

matrix (0 = Id. and per-edge vector and tensor features using the normalized edge vectors fu = Tzy / = (fig Pi fij). We create symmetric traceless tensor from the outer product fy with itself. S = f 3 Tr Id, and vector features are initialized by identifying the independent components of the skew-symmetric contribution with the components of fij as
denoted in (3 getting for every edge (ij) initial irreducible components ( A (
(( To
a) Embedding
b) Interaction
1
+I 11 ,
k
1 A g(t)
Tij
X ik
i
4
k
y(
10
400
0
M 10
o(ru)
12 AMA
M (0))
4
A
S
r
i
0
k
5200
M (4))
M (s2)
600 Zup(ra))
((k) (ik)
sum
M (4)
X
Key steps, from top to bottom, in the embedding and interaction modules for some central atom and neighbors j and & found within the cutoff radius. a) Relative position vectors 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.
Figure 1:
encode interatomic distance and atomic number information in the tensor representations we use an embedding layer that maps the atomic number of every atom 2; to n invariant features Z,, and expand interatomic distances Tij to d invariant features by means of an expansion in terms of exponential radial basis functions
(6) where Sg and M4 are fixed parameters specifying the center and width of radial basis function k. The a vector initialized with values equally spaced between exp(-ra) and 1, and 3 initialized as (2d-1(1 cxp ( (-r=))) for all & as proposed in (2)). After creating n identical copies of initial
RBF (ru) exp (-3u (exp(-ra)) -
components (5)) (5) ((1) (n feature channels), for every edge (ij) we map with a linear layer the
concatenation of Z; and Z, to pair-wise invariant representations Zij, and the radial basis functions are further expanded to n scalar features by using three different linear layers to obtain
ff nBr
NBF (ry)) bb, f2 (e RBF (re)) + + f ( (BB (rw)) + + (7)
5

---

X (or (35)
(65) +22 + 135
(8)
where the cutoff function is given by o((i)) = = {ccos M3s +1 when Tiy S ro and 0 otherwise.
That is, edge-wise tensor representations X( 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 || = Tr(X+ X) of atomic representations
x are computed and fed normalization layer, multilayer perceptron, and a SiLU activation to obtain three different ((3))invariant functions per channel,
5 12 ,5 SiL SiLU(MLPPLaaeette (LayerNorr(a ((())
which, after the decomposition of tensor embeddings into their irreducible representations, are then used to modify component-wer wise linear combinations to obtain the final atomic tensor embeddings
(
S W A
(1) g$
(10)
Interaction and node update. We start by normalizing each node's tensor representation X X(XX| + 1) and decomposing this representation into scalar, vector, and tensor features. We next transform these features A
S(() , by computing independent linear combinations, ( W ( + W ( + g()) In parallel, edge distances' radial basis expansions are fed to multilayer perceptron and a SiLU activation to transform them into tuples of three invariant functions per channel weighted with the cutoff function ((ri)),
((1)) (13) ( o(ri))((LUU)))))) (SLLU(MLP(( RBF (ri;))).
((0))
(11)
At this point, after decomposition of node features y define the messages sent from neighbors ( (t) + ((5) A0) + EjeN(o M We use the irreducible decomposition of matrix products ( M + (
j to central atom as (i) = ,(مر
8 ((1)) sti), which get aggregated into M
between node embeddings and aggregated messages to generate new atomic scalar, vector, and tensor features. New features are generated in this way to guarantee the preservation of the original parity of scalar, vector, and tensor features. These new representations [(0, A()), S(()) are individually normalized dividing by || + A + +500 || + and are further used to compute independent linear combinations to get ( + W 760 + W A + g(5 A residual update ( for original
embeddings (+) is computed with the to eventually obtain updated representations (+) (6 + (
matrix polynomial Ax y + (x)())
Scalar output. The Frobenius norm Tr(x"X) of full tensor representations and components in TensorNet is O(3)-invariant, For molecular potential predictions, total energy U computed from atomic contributions (1() which are simply obtained by using the concatenated final norms of every atom's scalar, vector, and tensor features || [1,
u = MLP(LayerNorrm [1. || [[1001]
(12)
obtaining forces via backpropagation.
Vector output. Since interaction and update operations preserve the parity of tensor components, the skew-symmetric part of any full tensor representation X in TensorNet is guaranteed to be vector, not a pscudovector. Therefore, from the antisymmetrization Ax, one can extract vectors v = (Uas Ug Ue) by means of the identification given in
Tensor output. Taking into account that -2 tensors have even parity and the -eem-rymmetric part A in TensorNet is vector, not pseudovector, one might need to produce pseudovector features before rank-2 tensor predictions can be built by combining irreducible representations. This can be easily done by obtaining two new vector features with linear layers, AO) W(1) A and (22) (22) A, and computing (A()) 4(2) (A()_{(2) TT ), which is skew-symmetric, rotates like vector, and is invariant under parity, the simultaneous transposition of (11) and (
4 Experiments and results
We refer the reader to the Appendix (section A.3) for further training, data set and experimental details

---

QM9: Chemical diversity. To assess TensorNet's accuracy in the prediction of energy-related molecular properties with a training set of varying chemical composition we used QM9 [3]. We trained TensorNet to predict: Uo. the internal energy of the molecule at O K; U, the internal energy 298. 15 K: H, the enthalpy, also at 298.15 K; and G, the free energy at 298. 15 K. Results can be found in Table 1, which show that TensorNet outperforms Allegro [8]. and MACE [34] on Up. U and H. Remarkably, this achieved with 23% Allegro's parameter count. Furthermore, TensorNet uses only scalar, vector and rank-2 tensor features, opposed to Allegro, which uses also their parity counterparts, and without the need of explicitly taking into account many-body terms as done in MACE.
Table I: QM9 results. Mean absolute error energy-related molecular properties from the QM9 dataset, in meV, averaged over different splits. Parameter counts for some models are found between parentheses.
DimeNet++
ET
PaiNN
Allegro
MACE
TensorNet (4.0M)
Property
25
16
[5
(17.9M)
34
Un
6.3
6.2
5.9
4.7
4.1
3.9(1)
U
6.3
6.3
5.7
4.4
4.1
3.9(1)
H
6.5
6.5
6.0
4.4
4.7
4.0(1)
G
7.6
7.6
7.4
5.7
5.5
5.7(1)
rMD17: Conformational diversity. We also benchmarked TensorNet on rMD17 35 the revised version of MD17 [36] 337. a data set of small organic molecules in which energies and forces were obtained by running molecular dynamics simulations with DFT. We report the results in Table 2. In the case of energies, TensorNet with two interaction layers (2L) is the model that achieves state-of- the-art accuracy for the largest number of molecules (6 out of 10), outperforming all other spherical models for benzene, with a parameter count of 770k. Energy errors are also within the range of other spherical models, except for ethanol and aspirin, and reach state-of-the-ase accuracy for the case of toluene, with just one interaction layer (1L) and parameter count of 535k. Force errors for 2L are also mostly found within the ranges defined by other spherical models, except for ethanol, aspirin, and salicylic acid, which case these are slightly higher. However, for one interaction layer, force errors are increased and in most cases found outside of the range of accuracy of the other spherical models. We note that the smallest spherical models have approximately 2.8M parameters, and therefore TensorNet results are achieved with reductions of 80% and 70% in the number of parameters for IL and 2L, respectively. Also, TensorNet is entirely based at most on rank-2 tensors.
SPICE, ANIIX, C0MP6: Compositional and conformational diversity. To obtain general- purpose neural network interatomic potentials, models need to learn simultancously compositional and conformational degrees of freedom. In this case, data sets must contain wide range of molecular systems as well as several conformations per system. To evaluate TensorNet's out-of-the tit-hhe-or performance without hyperparameter fine-tume tne-tuning, we trained the light model with two interaction layers used on rMD17 on the SPICE 388 and ANIIx (39 (40 data sets using the proposed Equivariant Transformer's SPICE hyperparameters [1 (for ANIIx, in contrast to the SPICE model, we used 32 radial basis functions instead of 64, and a cutoff of 4.5A instead of 10Å), and further evaluated ANIIx-trained models on the COMP6 benchmarks 39
For SPICE, with a maximum force filter of 50.94 ev/Â & Ha/Bohr, TensorNet's 's mean absolute error in energies and forces are 25.0 meV and 40.7 meÀ/ respectively, while the Equivariant Transformer achieves 31.2 meV and 49.3 me V/À. In this case, both models used cutoff of 10Å. Results for ANI1x and model evaluations on COMP6 are found in Table 3. We note that for ANII x training, which contains molecules with up to 63 atoms, TensorNet used a cutoff of 4.5A The largest rMD17 17 molecule is aspirin with 21 atoms. The light TensorNet model shows better generalization capabilities across all COMP6 benchmarks.
Scalar, vector and tensor molecular properties for ethanol in vacuum. We next tested TensorNet performance for the simultaneous prediction of scalar, vector, and tensor molecular properties: potential energy, atomic forces, molecular dipole moments μ,, molecular polarizability tensors and nuclear-shielding tensors a, for the ethanol molecule in vacuum [5] 42 We trained TensorNet to generate atomic tensor representations that can be used by different output modules to predict the
7

---

Table 2: rMD17 results. Energy (E) and forces (F) mean absolute errors meV and meV/Å, averaged over different splits.
TensorNet
TensorNet
NequlP
Allegro [18]
BOTNet
MACE
Molecule
IL (535k)
2L (770k)
[17
[9
22
2.7
2.4
2.3
2.3
2.3
2.2
Aspirin 10.2(2)
8.9(1)
8.2
7.3
8.5
6.6
E
0.9
0.7
0.7
1.2
0.7
1.2
Azobenzene
3.8
3.1
2.9
2.6
3.3
3.0
0.03
0.02
0.04
0.3
0.03
0.4
Benzene
0.3
0.3
0.3
0.2
0.3
0.3
0.5
0.5
0.4
0.4
0.4
0.4
Ethanol
3.9(1)
3.5
2.8
2.1
3.2
2.1
0.8
0.8
0.8
Malonaldehyde
5.8(1)
5.4
5.1
3.6
5.8
4.1
0.3
0.2
0.9
0.2
0.2
0.5
Naphthalene
1.9
1.6
1.3
0.9
1.8
1.6
E
1.5
1.3
1.4
1.5
1.3
1.3
Paracetamol
6.9
5.9(1)
5.9
4.9
5.8
4.8
E
0.9
0.8
0.7
0.9
0.8
0.9
Salicylie acid
5,4(1)
4.6(1)
4.0
2.9
4.3
3.1
0.3
0.3
0.3
0.4
0.4
0.5
Toluene
2.0
1.7
1.6
1.8
1.9
1.5
0.5
0.4
0.4
0.6
0.4
0.5
Uracil
F
3.6(1)
3.1
3.1
1.8
3.2
2.1
Table 3: ANIIx and COMP6 results. Energy (E) and forces (F) mean absolute errors in meVV and
me V/Á for Equivariant Transformer and TensorNet models trained on ANIIx and evaluated on the ANIIX test set and the C0MP6 benchmarks, averaged over different training splits. 43 meV
kcal/mol.
Model
ANIIx
ANI-MD
GDB7-9
GDB10-13
DrugBank
Tripeptides
S66x8
E
21.2
249.7
17.8
51.0
95.5
57.9
30.7
ET
42.0
50.8
29.2
57.4
47.7
37.9
19.0
E
17.3
69.9
14.3
36.0
42.4
40.0
27.1
TensorNet
F
34.3
35.5
23.1
41.9
32.6
26.9
14.3
desired properties. The specific architecture of these output modules can be found in the Appendix (section A.3).
Table 4: Ethanol in vacuum results. Mean absolute enror for the prediction of energies (E), forces (F), dipole moments (μ), polarizabilities (a), and chemical shifts for all elements (Tau), averaged over different splits, with corresponding units between parentheses.
Model
E (kcal/mol)
F
Tall (ppm)
(kcal/mol/A)
(D)
(Bohr3)
PaiNN L5
0.027
0.150
0.003
0.009
NieldSchNet
42
0.017
0.128
0.004
0.008
0.169
TensorNet 0.008(1)
0.058(3)
0.003(0)
0.007(0)
0.139(4)
Results from Table 4 show that TensorNet can learn expressive atomic tensor embeddings from which multiple molecular properties can simultaneously predicted. particular, TensorNet's energy and force errors are approximately factor of two and three smaller when compared to FieldSchNet 42 and PaiNN [15), respectively, while increasing the prediction accuracy for the other target molecular properties, with the exception of the dipole moment.

---

Equivariance, interaction and cutoff ablations. TensorNet can be straightforwardly modified such that features are -quivaaana and scalar predictions are SE(3) SE33a-invriiant by modifying the matrix products in the interaction mechanism. An interaction product between node features and aggregated messages 2 y M() instead of y ( + ((), gives vector and tensor representations which are combinations of even and odd parity contributions. We refer the reader to the Supplementary Material for detailed derivations. Furthermore, the norms Tr(X (XTXX X) used to produce scalars will only be invariant under rotations, not reflections. This flexibility in the model allows us to study the changes in prediction accuracy when considering ((3) or SO(3) equivariant models. We also evaluated the impact on accuracy for two rMD17 molecules, toluene and aspirin, when modifying the receptive field of the model by changing the cutoff radius and the number of interaction layers, including the case of using the embedding and output modules alone, without interaction layers (OL), with results in Table 5.
Table 5: Equivariance, interaction and cutoff ablations results. Energy (E) and force (F) mean absolute errors in meV and me V/À for rMD17 toluene and aspirin, averaged over different splits, varying the number of interaction layers, the cutoff radius, and the equivariance group.
TensorNetOL.
TensorNet 1L.
Tensor
TensonNet IL.
TensorNet 21.
0(3)
O(3)
O(3)
SO(3)
SO(3)
Molecule
4.5A
9A
4.5A
9A
4.5A
9A
4.5A
4.5A
E
3.3
2.0
0.33
0.36
0.26
0.32
0.50
0.42
Toluene
F
15.7
11.5
2.0
2.2
1.7
2.0
2.9
2.4
E
9.8
7.8
2.7
2.9
2.4
2.8
3.7
3.4
Aspirin
F
32.7
28.3
10.1
11.0
8.9
10.5
13.1
11.8
The inclusion or exclusion of equivariance and energy invariance under reflections has a significant impact accuracy. The consideration of the full orthogonal group O(3), and therefore the physical symmetries of the true energy function, leads to higher accuracy for both energy and forces. Fur- thermore, the use of interaction products produces a drastic decrease in errors (note that TensorNet IL 4.5A and TensorNet OL 9 have the same receptive field). In line with rMD17 results, second interaction layer in the case of re = 4.5A gives an additional but more limited improvement in both energy and force errors. For forces, the use of a second interaction layer with r. 9A encompassing the whole molecule provides a smaller improvement when compared to r. = 4.5A. We note that for OL., when the model can be regarded as just a learnable aggregation of local atomic neighborhoods, TensorNet with both cutoff radii achieves for aspirin (the rMD17 molecule on which the model
performs the worst) lower mean absolute errors than ANI (16.6 me and 40.6 V/À) 7 43 and
SchNet (13.5 meV and 33.2 V/À) 22 44).
Computational cost. We found that TensorNet exhibits high computational efficiency, even higher than an equivariant model using Cartesian vectors such as the Equivariant Transformer [16] in some cases. We provide inference times for single molecules with varying numbers of atoms in Table 6, and in Table 7 we show training steps per second when training on the ANI1x data set, containing molecules with up to 63 atoms.
For molecules containing up to ~20 atoms, TensorNet 1L and 2L are faster or similar when compared to the ET, even when its number of message passing layers is reduced (ET optimal performance for MD17 was achieved with 6 layers, found to be one of the fastest neural network potentials in the literature [16] 45)), meaning that energy and forces on these molecules can be predicted with rMD17 state-of-the-art TensorNet models with a lower or similar computational cost than the reduced ET. For larger molecules with thousands of atoms, TensorNet 2L becomes significantly slower. However, TensorNet IL. which still exhibits remarkable performance on rMD17 (see Table 1), performs on par the reduced ET in terms of speed even for Factor IX, containing 5807 atoms. For training on ANIIx, TensorNet IL and 2L are faster or comparable to the ET up to a batch size of 64, being the speed for the 2L model being significantly slower for a batch size of 128. Nevertheless, the model with 1 interaction layer is still comparable to the reduced ET.
TensorNet's efficiency is given by properties that are in contrast to state-of-the-ar equivariant spherical models. In particular, the use of Cartesian representations allows one to manipulate full tensors or their decomposition into scalars, vectors, and tensors at one's me's convenience, and Clebsch-

---

Table 6: Inference time. Inference time for energy and forces for single molecules (batch size of 1),
ms, on an NVIDIA GeForce RTX 4090.
Molecule
N
TensorNet OI.
TensorNet II.
TensorNet 21.
ET 41.
ET 51.
Alanine dipeptide
22
10.0
22.1
26.5
27.2
29.0
Chignolin
166
10.4
22.5
26.9
26.8
28.9
DHFR
2489
27.3
66.9
106.7
52.4
67.7
Factor IX
5807
53.1
149.8
248.6
110.6
136.4
Table 7: Training speed. Number batch training steps per second for ANIIx dataset on an NVIDIA
GeForce RTX 4090.
Batch size
TensorNet OL
TensorNet 1L
TensorNet 2L
ET 4L
ET 5L
32
20.5
13.2
10.1
9.5
8.4
64
19.1
12.9
9.1
9.4
8.3
128
17.3
8.9
5.9
9.1
8.0
Gordan tensor products are substituted for simple 3x3 matrix products. As detailed in the model's architecture, state-of-the-ee performance can be achieved by computing these matrix products after message aggregation (that is,
node level) and using full tensor representations, without having to individually compute products between different irreducible components. When considering average number of neighbors per atom M controlled by the cutoff radius and the density, given that matrix products performed after aggregation
neighbors, these do not scale with M. This is in contrast to spherical models, where tensor products are computed on edges, and therefore display a worse scaling with the number of neighbors M, that is, worse scaling when increasing the cutoff radius at fixed density. Also, the use of higher-order many-body messages or many message-passingg steps is not needed.
5 Conclusions and limitations
We have presented TensorNet, a novel 033)-equivariant message- message-passing architecture leveraging Cartesian tensors and their irreducible representations. We showed that even though the model is limited to the use of rank-2 tensors, in contrast to other spherical models, it achieves state-of-the offte-of-theart performance on QM9 and rMD17 with reduced number of parameters, few pesssage-passing steps, and it exhibits a low computational cost. Furthermore, the model is able to accurately predict vector and rank-2 tensor molecular properties on top potential energies and forces. Nevertheless, the prediction of higher-rank quantities is directly limited by our framework. However, given the benefits of the formalism for the construction of machine learning potential, TensorNet can be used as an alternative for the exploration of the design space of efficient equivariant models. TensorNet can be found in httph:tt //githu.. com//trrcmm//tt cooandddtohhhmddm net.
Acknowledgments and Disclosure of Funding
We thank Raul P. Pelaez for fruitful discussions. GS is financially supported by Generalitat de Catalunya's Agency for Management of University and Research Grants (AGAUR) PhD grant FI- 2-00587. This project has received funding from the European Union's 's Horizon 2020 research and innovation programme under grant agreement No. 823712 (RG, AV, RF); the project PID2020- 116564GB-10W has been funded by MCIN AEI 10. 30.39/0// 10001 1033. Research reported in this publication was supported by the National Institute of General Medical Sciences (NIGMS) of the National Institutes of Health under award number GM140090. The content solely the responsibility of the authors and does not necessarily represent the official views of the National Institutes of Health,
10

---

References
[1] Frank Noć, Alexandre Tkatchenko, Klaus-Robert Müller, and Cecilia Clementi. Machine leam- ing for molecular simulation. Annual Review of Physical Chemistry. 71(1):61-330 April 2020. doi: 10.1 10/1n44//nnureev 10.146/aamuryy-pisccee-202222 Sunurev-physchenn 042018-052331
pwwschem-042018--0000
URL https://ddot http:////oooog////////// 1146/
[2] Justin S. Smith, Adrian E. Roitberg, and Olexandr Isayev. Transforming computational drug discovery with machine learning and AI. ACS Medicinal Chemistry Lettery, 9111)1066--106, October 2018. doi: 10. http://csmeedcceeetttt00000555
httpatta::///o org/10. 0022/acaamdchemmettt.. 8b00437
[3] Jessica Vamathevan, Dominic Clark, Paul Czodrowski, Ian Dunham, Edgardo Ferran, George Lee, Bin Li, Anant Madabhushi, Parantu Shah, Michaela Spitzer, and Shanrong Zhao. Applications of machine learning in drug discovery and development. Nature Reviews Drug Discovery, 18(6):6334477 477, April 2019. org/10. 1008//411333009900000000
doi: (0038//1573-019-00222
[4] Albert Musaelian, Anders Johansson, Simon Batzner, and Boris Kozinsky. Scaling the leading accuracy of deep equivariant models biomolecular simulations of realistic size, 2023. URL htttp:://arii... abs/2304. 10066
[5] Jörg Behler and Michele Parrinello. Generalized neural-network representation of high-dimensional potential-enene surfaces. Physical Review Letters, 98(14), April 2007. doi: 98.1100//pyyse1eee.884444000 URL https://://// org/10...rr1.... www.//////yeeeeett.t..oo 98. 146401
[6] Kristof Schüitt, Farhad Arbabzadah, Stefan Chmiela, Klaus Müller, and Alexandre Tkatchenko. Quantum-chemical insights from deep tensor neural networks. Nature Communications, 8(1), January
2017. doi: 10.1038/ncomms 13890. URL https://:///dii http://doooorrr/// 10. 1038/ wwt/////comms1130 13890
[7] J. S. Smith, O. Isayev, and A. E. Roitberg. ANI-I: am extensible neural network potential with DFT accuracy at force field computational cost. Chemical Science, 8(4):3192 3203, 2017. doi: 10.1039/c6se05520a. URL httph/tttt///ooll// org/ 10. 1039/ c636c22
[8] Michael Gastegger and Philipp Marquetand, Molecular dynamics with neural network potentials. In Machine Learning Meets Quantum Physics, pages 233 252. Springer International Publishing. 2020. doi: 00000/9788-3030-42200
o /
-30-0-4-0000000 1000--0225-777 7__2
[9] Jie Zhou, Ganqu Cui, Shengding Hu, Zhengyan Zhang, Cheng Yang, Zhiyuan Liu, Lifeng Wang. Changcheng Li. and Maosong Sun. Graph neural networks: A review of methods and applications. 2018. URL httptst//axxiii www./oggg/////////o///..o/...ee/ aa/aaaaa///1// /tte//812.0000 08434
[10] Peter W. Battaglia, Jessica Hamrick, Victor Bapst, Alvaro Sanchez- Gonzalzz Vinicius Zambaldi, Mateusz Malinowski, Andrea Tacchetti, David Raposo, Adam Santoro, Ryan Faulkner, Caglar Gulcehre, Francis Song, Andrew Ballard, Justin Gilmer, George Dahl, Ashish Vaswani, Kelsey Allen, Charles Nash, Victoria Langston, Chris Dyer, Nicolas Heess, Daan Wierstra, Pushmeet Kohli, Matt Botvinick, Oriol Vinyals, Yujia Li, and Razvan Pascanu. Relational inductive biases, deep learning, and graph networks, 2018. URL htttps://arratt http:://axxxvvvv.
[11] Michael M. Bronstcin, Joan Bruna, Taco Cohen, and Potar Veličković. Geometric deep lcarning: Grids, groups, graphs, geodesics, and gauges, 2021. htt.L/ttpt///aavv.ro/aaaa/////// 13478
[12] Justin Gilmer, Samuel S. Schoenholz, Patrick F. Riley, Oriol Vinyals, and George E. Dahl. Neural message passing for quantum chemistry, 2017. URL https::/arxxv.oortt http://aaaav..orrra///////////// 01220
[13] Brandon Anderson, Truong-Son Son Hy, and Risi Kondor. Cormorant: Covariant molecular neural networks. In Proceedings of the 33rd International Conference on Neural Information Processing Systems, Red Hook, NY, USA, 2019. Curran Associates Inc.
[14] Johannes Brandstetter, Rob Hesselink, Elise van der Pol, Erik Bekkers, and Max Welling, Geometric and physical quantities improve e(3) cquivariant message passing, 2021. F110.02900
[15] Kristof T. Schütt, Oliver Unke, and Michael Gastegger. Equivariant message passing for the prediction of tensorial properties and molecular spectra, 2021. URL https:////// org/ abs abs/2102 2102.03150
[16] Philipp Thölke and Gianni De Fabritiis. Torchmd-net Equivariant transformers for neural network based molecular potentials, 2022. URL httts::/aaxiis wwweeorr//raae http//bbb/222....e 02541
11

---

[17] Simon Batzner, Albert Musaelian, Lixin Sun, Mario Geiger. Jonathan P. Mailoa, Mordechai Kornbluth, Nicola Molinari, Tess E. Smidt, and Boris Kozinsky. E())-quiivarian graph neural networks for data- efficient and accurate interatomic potentials. Nature Communicationss 13(1), May 2022. doi: 10.1038/ 01167-022-29999--2 URL https: h/tp:://ooooo..// org/ 10. 1008/ s41467 022 29939-993955
[18] Albert Musaclian, Simon Batzner, Anders Johansson, Lixin Sun, Cameron Owen, Mordechai Kornbluth, and Boris Kozinsky. Learning local equivariant representations for large-scale atomistic dynamics. Nature Communications, 14(1), February 2023. doi: 10. 11.10038/441467-02336622222 10. -4038//44467022225 -0--34439-----33-9--9-
URL htRlhhtps://doo org//
[19] Ilyes Batatia, Simon Batzner, Dávid Péter Kovács, Albert Musaelian, Gregor N. C. Simm, Ralf Drautz, Christoph Ortner. Boris Kozinsky, and Gábor Csányi. The design space of equivavain alom-centered interatomic potentials, 2022.
[20] Ilyes Batatia, David Peter Kovacs, Gregor N. C. Simm, Christoph Ortner, and Gabor Csanyi. MACE: Higher order equivariant message passing neural networks for fast and accurate force fields. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances Neural Information Processing Systems, 2022. URL URL bbttps://openreveee.eeet wwt./foene/iiideee YPo.dd-
[21] Maurice Weiler, Mario Geiger, Max Welling, Wouter Boomsma, and Taco Cohen. 3d stecrable enns: Learning rotationally equivariant features in volumetric data, 2018. URL htttss///axii: /bs// 1204702547
[22] K. T. Schütt, H. E. Sauceda, P.-J. Kindermans, A. Tkatchenko, and K.-R. Müller. SchNet deep learning architecture for molecules and materials. The Journal of Chemical Physics, 148(24):241722, June 2018. doi:
URL http::tt/cci...ooooo org/ 10. 5019779
[23] Johannes Gasteiger, Florian Becker. and Stephan Günnnemann. Gemnet: Universal directional graph neural networks for molecules, 2021. URI
08903
[24] Johannes Gasteiger, Janek GroB, and Stephan Günnemann. Directional message passing for molecular graphs, 2020. URL (RRtihttt:///rrrtt org/ wbwpb003 03123
[25] Johannes Gasteiger, Shankari Giri, Johannes T. Margraf, and Stephan Günnemana. Fast and uncertainty- aware directional message passing for non-equilibrimm molecules, 2020. URL lttps:////// org/ abs/2011. 14115
[26] Oliver T. Unke and Markus Meuwly. PhysNet: A neural network for predicting energies, forces, dipole moments, and partial charges. Journal of Chemical Theory and Computation, (566)3678-369) May 2019. doi: 10.1022/acss,jtt::bb000
URL httpps:///do /oi. org. 10. 1021/ acs
[27] Yi Liu, Limei Wang, Meng Liu, Xuan Zhang, Bora Oztekin, and Shuiwang Ji. Spherical message passing for 3d graph networks, 2021. UtL.:ttts: www..o/g///......ooooo...oesee/e wtwpa/////.aaeeeeeeeeeeeeeeteeee 2101305013
[28] Fabian B. Fuchs, Daniel E. Worrall, Volker Fischer, and Max Welling. teans-rransssommrss:: 3d roto- translation equivariant attention networks, 2020. https:////// http:://wwwwv.////////// hwtp://.rg//b/////////e// 10503
[29] Mojtaba Haghighatlari, Jie Li, Xingyi Guan, Oufan Zhang, Akshaya Das, Christopher Stein, Farna Heidar-Zadeh, Meili Liu, Martin lead-Gordon, Luke Bertels. Hongxia Hao, Itai Leven, and Teresa Head- Gordon. NewtonNet: a newtonian message passing network for deep learning of interatomic potentials and forces. Digit. Discov. 1(3):333-343 June 2022.
[30] Victor Garcia Satorras, Emiel Hoogehoom, and Max Welling. E(n) equivariant graph neural networks, 2021. URL. htt.hhttt:///artti org/ abs/2102.... 211.209922201
[31] Yuanqing Wang and John Chodera. Spatial attention kinetic networks with e(n)-equivariance, 2023. URL //axiiv./////iv///aa org/ httg//bbb2331..... 088933
[32] Nathaniel Thomas, Tess Smidt. Steven Kearnes. Lusann Yang, Li Li, Kai Kohlhoff, and Patrick Riley. Tensor field networks: Rotation and translation -awaaaaiioonuuiriiain neural networks for 3d point clouds, 2018.
URL http:::://rxiiis org/ a8802 1822/
[33] Raghunathan Ramakrishnan, Pavlo O. Dral, Matthias Rupp, and O. Anatole von Lilienfeld. Quantum chemistry structures and properties of 134 kilo molecules. Scientific Data 1(1), August 2014. doi: (038//data.2014.222 URL https://do/: org/ 10. 1038/sdata. sdata.daaaaa... 2014.22
[34] Dávid Péter Kovács, Ilyes Batatia, Eszter Sára Arany, and Gábor Csányi Evaluation of the MACE force field architecture: From medicinal chemistry to materials science. The Journal of Chemical Physics, 159 (4), July 2023. doi:
httphhttt/:/ttt org/ 100 1063/5. 0155322
12

---

https://figshant
[36] Stefan Chmiela, Huziel E. Sauceda, Klaus-Robert Müller, and Alexandre Tkatchenko. Towards exact molec- ular dynamics simulations with machine-learned force fields. Nature Communications, 9(1), September 2018. doi: 01105/444667-1-06661666
URL https://doi org/ /tg/10..//// -18-://184670011-----8
[37] Stefan Chmiela, Alexandre Tkatchenko, Huziel E. Sauceda, Igor Poltavsky, Kristof T. Schütt, and Klaus- Robert Müller. Machine learning of accurate energy-conserrvngg molecular force fields. Science Advances, 3(5), May 2017. doi: 10. 1116//ciadv.
URL httpttss:////o org/10 10. 1126/sciady 1603015
[38] Peter Eastman, Pavan Kumar Behara. David Dotson, Raimondas Galvelis, John Herr Josh T. Horton, Yuezhi Mao, John Chodera, Benjamin P. Pritchard. Yuanqing Wang, Gianni Fabritiis, and Thomas E. Markland. SPICE, a dataset of drug-like molecules and peptides for training machine learning potentials. Scientific Data, 10(1), January 2023, doi:
URL
https://do1. www/or.g//....... 1038/ #41597 01882--
[39] Justin S. Smith. Roman Zubatyuk. Benjamin Nebgen, Nicholas Lubbers, Kipton Barros, Adrian E. Roitherg, Olexandr Isayev, and Sergei Tretiak. The ANI-Iccx and ANI-Ix data sets, coupled-cluster and density functional theory propertics for molecules. Scientific Data, 7(1), May 2020. doi: 10.
-41597002-- URL http:////aaa...g11000000//111111 -44-33
[40] Justin Smith, Ben Nebgen, Nicholas Lubbers, Olexandr Isayev, Adrian Roitberg, Less more: Sampling chemical space with active learning. The Journal of Chemical Physics, 148((4))241733, June 2018. doi: 10.1063/1 StRRahtttaa///. //do:pa://di...o org/ 10. 1.533/00
[41]
chand net
[42] Michael Gastegger, Kristof T. Schütt, and Klaus-Robert Müller. Machine learning of solvent effects on molecular spectra and reactions. Chemical Science, 12(34):1 14733-11147-1148 2021. doi: 10. sc2777.
URL httptttts///....... http:///oo...rr///// 10. 1039/ /at///12sc22222e
[43] Dávid Péter Kovács, Cas van der Oord, Jiri Kucera, Alice A. Allen, Daniel Cole, Christoph Ortner, and Gábor Csányi. Linear atomic cluster expansion force fields for organic molecules: Beyond RMSE, Journal of Chemical Theory and Computation, 17(12)7696-7711, November 2021. doi: 10. 102 1021//cs.jcte. 1c00447. httppt://///oii.... orgg 1021/ acs. http://cc..ccc....oo 1c00647
[44] Kristof T. Schürt, Stefaan S. P. Hessmann, Niklas W. A. Gebauer, Jonas Lederer, and Michael Gastegger. SchNeiPack 2.0: A neural network toolbox for atomistic machine learning. The Journal of Chemical Physics, (58(14):14480, April 2023. doi: 10.1063/5.01
URL https://doi. org/10. 1063//
0138367
[45] Vaibhav Bihani. Utkarsh Pratiush, Sajid Mannan, Tao Du, Zhimin Chen, Santiago Miret, Matthieu Micoulaut, Morten M Smedskjaer. Sayan Ranu, and N M Anoop Krishnan. Egraffbench: Evaluation of equivariant graph neural network force fields for atomistic simulations. 2023. URL https://arxiv
org/ abs///aaa.22./// 02428
[46] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan. Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Köpf. Edward Yang. Zach Vetito. Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang. Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high- tigh-ppeffrrancc deep learning library, 2019. URL
https://arxxx wwt.:/rr//bbb///.1// 01703
[47] http: http://qannuuuuuum hagaaatuuummcciii..orrattttaeae
13

---

A Appendix
A.1
Architecture diagrams
-IIII
X
|
RDP
4
-1-1+=X
W +-8
oubilt
W 4-8
4
5
++6
Nelghborn
Elmbodd
W.
W.
W
concas
1 A g()
1
18
Embod(a(0)
Σ
JAN10
Inseroction
((
4
51()
7
(4)
5(6)
||
W.
LayerNorm
W.
W.
MLP
1/
SILU
(1)
x (9)
(
(a) Embedding
(b) Interaction and node update
Figure A.1: Diagrams of the embedding and interaction and node update modules in TensorNet, descibed in Section 3.2. Models built from an initial embedding module, the concatenation of several interaction layers, and the output module. The dotted line in Interaction and node update indicates the point where, after node level transformations, neighboring atom's features are used to create messages.
A.2 Proofs and derivations
Parity and interaction products
Consider some tensor X, with its corresponding decomposition x + A + S and initialized as in TensorNet (see Section 3.2), that is, with vector features obtained directly by the identification of vector components with matrix entries (3). When considering the parity transformation, the components of the vector used to build the initial features flip sign, and for any full tensor initialized in this way
x =
x + X
+ x + X'
= X
(A.1)
Consider now another tensor Y = 1 + A SY initialized in the same way, and the product XY
x +5* +A" + + ). One can compute the resulting decomposition XY TTxXYTTd
= (XY - YT x and S.YY (XY + (XY)T 3 (I(())) = (XY + YTXT TT(XY Id) from the individual decompositions of X and Y. Taking into account that transposition is equivalent to sign reversal of skew-symmetric components, one obtains after some matrix algebra
AXY = ((XX-(XX))
14

---

Y
Ir + A + S* + Tr(A S + ) +
3
scalar
m
+Tr(I* 1 +
+ ld.
(A.2)
AxY + A 1 (AXS (A (A - (AV
vector
+ (A A (A - 1 - s* - 5T)
(A.3)
pseudovector
= - - MId+
1 + 2 (A* + ( Tm( St jld) + (S* A - AS*AA"" 2
Tr(s A IId)
3
paeudotensor
+ + + xx f +1**
2 Th (d) [18 S + 50**0²"
2 This s styd), (id)
tensor
(A.4)
where we have made explicit that the product XY gives rise to both even parity (transposing A and AY does not flip the sign of the expression) and odd parity (transposing A and A does flip the sign of the expression) skew-symmetric and symmetric traceless contributions. However, when considering XY Y YX, one gets
XX
Tr(21* PROI*I++AAAA A (A*AY)) S + * (SX ( ld (A.5)
scalar
AXY+YX= - 21 12 + S*
(AX (A (A* +
vector
1 AA A * AX A*A
(5x - S s + SK 5x - s*
(A.6)
that is, the undesired pscudovector contributions cancel out. For the symmetric traceless part pseudotensor contributions also cancel out, getting eventually
xXY+Yx= Sx + 25x + (A* A + (A*AA))
-2 Tr(A* A
* + (S _ TrSS 5 ITd)),
(A.7)
which has even parity. These results can be summarized expression (5) in the main text
xx XX xxx --x++X SXTY++x*x* SXX+XX
(A.8)
15

---

Invariance of Frobenius norm
The Frobenius norm of some tensor representation X in TensorNet is invariant under the full orthogonal group O(3). This follows from the cyclic permutation invariance of the trace. As previously shown, parity induces the transposition of X, which amounts to cyclic permutation of the matrix product inside the trace operator, Tr(X* x) = Tr(Xx*), When considering some rotation R. Tr(X TXTXX X) + Tr((RX R$ Y ( R = Tr(RX"
T=(RX"R=RXXXP" TT(RXTX"") = Tr(XTRR R) = Tr TTXXTXXXXXXXX Tx(Xxx)
0(3) to SO(3) symmetry breaking
The embedding module in TensorNet always produces irreducible representations that are scalars, vectors and tensors, that is, preserves the parity of the initial components. As previously shown, the computation of the Frobenius norm is (33)-invariant in this case, making the normalization an operation that also preserves initial parities, since can be regarded as the modification of the irreducible representations by means of invariant weights. However, depending on the interaction product, ='u-vvvvvaar can broken to -euuieqquvariiane in the first interaction layer.
Input irreducible representations for the first interaction layer, which are the outputs of the embedding module, have dell-defined parity. The first normalization that takes place makes use of the Frobenius norm, which at this point is O(3) D(3)-ime -invariant, giving ((3) (3))equrvarraant representations with the desired parity. Nevertheless, if one considers an interaction product proportional to ( (+ and computes its irreducible decomposition to generate new features, seemmmymettric and symmetric traceless representations will receive both even and odd parity contributions, proved above (expressions A.3 and A.4). At this point, one can write some node's full tensor embedding X as
- A + + 5 +
(A.9)
where + and denote even and odd parity, respectively. After obtaining new features through the interaction product, another normalization takes place. The computation of the Frobenius norm Tr(XT X) with current node representations X used normalize features is longer One can write
= -T(( A + A + + + + At + + + +
- - A + + S A+ + A + +S))),
(A.10)
where we have made use of the symmetry or symmeeyyyy of the matrices under transposition, and considering the action of parity A + S S +
S-
Tr((T -+ A + S - ST) A+ A- -
(A.11)
which manifestly different from A.10 and cannot be reduced to a cyclic permutation of A.10. Thus, the Frobenius norm of representations X this point of the architecture is not invariant under parity, and therefore normalization amounts to the modification of components with weights 1/([|X]] 1) which not invariant under the full O(3). but just under SO(3). From this moment, all features generated in TensorNet will no longer be O(3) ('uiiiarirr but the computation of any Frobenius norm (such the used in the scalar output module to predict energies) will only be SO(3) (30(3)-inarria)) invariant.
and
A.3 Data sets and training details
TensorNet was implemented within the TorchMD-NET framework 41 using PyTorch 1.11.0 46 PyTorch Geometric 2.0.3 and PyTorch Lightning 1.6.3.
QM9
The QM9 33 47
data set consists of 1 130,831 optimized structures of molecules that contain up to 9 heavy elements from C, N, and F. On top of the structures, several muantum-chemican properties computed at the DFT previous work, we used 1 100000 structures for training, which were shuffled after every epoch, 10,000 for validation and the remaining ones were used for testing. We used four random splits, initializing the model with different random seeds. Reported results are the average errors and the standard deviation between parentheses of the last significant digit over these splits.
SELLY/6-311((d6)
level of theory are provided. To be consistent with
16

---

Models trained with QM9 had 3 interaction layers, 256 hidden channels, used 64 non- non-trainable radial basis functions and a cutoff of 5A. The MLP for the embedding part had 2 linear layers, mapping 256 hidden channels as [256, 512, 768J. with SiLU activation. The MLPs in the interaction layers had 3 linear layers, mapping the 64 radial basis functions as [64, 256, 512, 768], also with SiLU non-linearities. The scalar output MLP mapped [768. 256, 128, 1] with SiLU activations, giving an atomic contribution to the energy which is added to the atomic reference energies from the data set. All layers were initialized using default PyTorch initialization, except for the last two layers of the output MLP which were initialized using a Xavier uniform distribution and with vanishing biases. Layer normalizations used the default PyTorch parameters. For training, we used a batch size of 16, an initial learning rate of le-4 which was reached after 1000-step linear warm-up, the Adam optimizer with PyTorch default parameters and the MSE loss. The learning rate decayed using an on-plateau scheduler based on the validation MSE Ioss, with patience of 15 and decay factor of 0.8. We did not use an exponential moving average for the validation loss. Inference batch size was 128. Training automatically stopped when the learning rate reached le-7 or when the validation MSE loss did not improve for 150 epochs. We used gradient clipping at norm 40. Training was performed on two NVIDIA GeForce RTX 3090 with float32 precision.
rMD17
The rMD17 35 data set contains 100,000 recomputed structures of 10 molecules from MD17 36
371. a data set of small organic molecules obtained by running molecular dynamics simulations. The DFT PBE/def2-SVV level of theory, a very dense DFT integration grid and very tight SCF convergence were used for the recomputations. We used 950 and 50 random conformations for training and validation, respectively, and evaluated the error on all remaining conformations. We used five random splits, initializing the model with different random seeds. Reported results are the average errors and the standard deviation between parentheses (when different from zero) of the last significant digit over these splits.
Models trained with rMD17 had and interaction layers, 128 hidden channels, used 32 non non-trainable radial basis functions and a cutoff of 4.5A. The MLP for the embedding part had 2 linear layers, mapping 128 hidden channels as [128, 256, 384], with SiLU activation. The MLPs in the interaction layers had linear layers, mapping the 32 radial basis functions as [32, 128, 256, 384], also with SiLU non-linearities. The scalar output MLP mapped [384, 128, 64, 1J with SiLU activations, giving an atomic contribution the energy which after addition for all atoms is scaled by the training data standard deviation and shifted with the training data mean. All layers were initialized using default PyTorch initialization, except for the last two layers of the output MLP which were initialized using a Xavier uniform distribution and with vanishing biases. Layer normalizations used the default Py Torch parameters. For training, we used batch size of 8, an initial learning rate of 1e-3 which was reached after a 500-step linear warm-up, the Adam optimizer with PyTorch default parameters and weighted MSE losses of energy and forces with both weights equal to 0.5. The learning rate decayed using an on-plateau scheduler based on the total weighted validation MSE losses, with a patience of 25 and a decay factor of 0.8. We used an exponential moving average for the energy MSE validation loss with weight 0.99. Inference batch size was 64. Training automatically stopped when the learning rate reached le-8 or when the total weighted validation MSE losses did not improve for 300 epochs. We used gradient clipping at norm 40. Training was performed on single NVIDIA GeForce RTX 2080 Ti with float32 precision.
SPICE
SPICE (38 is a data set with an emphasis on the simulation of the interaction of drug-like small molecules and proteins. It consists of a collection of dipeptides, drug-like small molecules from PubChem, solvated aminoacids, monomer and dimer structures from DES370K and ion pairs, with a varying number of systems and conformations in each subset, and computed at the WB97M- D3BI/de12-TZVPPD level of theory. After filtering molecules containing forces higher than 50.94 ev/À ~ 1 Hartree/Bohr, SPICE (v 1.1.3) contains approximately IM data points. We used three rraining/vlidationntee splits, initializing the model with different random seeds. Reported results are the average errors over these splits.
random 80%/10%
Models trained with SPICE had 2 interaction layers, 128 hidden channels, used 64 -tairiiiaelle radial basis functions and cutoff of 10Å. The MLP for the embedding part had 2 linear layers,
17

---

mapping 128 hidden channels as [128, 256, 384], with SiLU activation. The MLPs in the interaction layers had linear layers, mapping the 64 radial basis functions as [64, 128, 256, 384], also with SiLU non-linearities. The scalar output MLP mapped [384, 128, 64, 1] with SiLUU activations, giving an atomic contribution to the energy, which after addition for all atoms, the model was trained to match the reference energy of the molecule with subtracted atomic reference energies. All layers were initialized using default PyTorch initialization, except for the last two layers of the output MLP which were initialized using Xavier uniform distribution and with vanishing biases. Layer normalizations used the default PyTorcl parameters. For training, we used a batch size of 16, an initial learning rate of le-4 which reached after a 5t0-step linear warm-up, the Adam optimizer with Py Torch default parameters and weighted MSE losses of energy and forces with both weights equal to 0.5. The learning rate decayed using an on-plateau scheduler based the total weighted validation MSE losses, with a patience of 5 and a decay factor of 0.5. We did not use an exponential moving average for the validation loss. Training automatically stopped when the learning rate reached 1e-7 or when the total weighted validation MSE losses did not improve for 50 epochs. Inference batch size was 16. We used gradient clipping at norm 100. Training was performed on 4 NVIDIA GeForce RTX 2080 Ti with float32 precision. The Equivariant Transformer model used for comparison purposes had 4 interaction layers, used 64 non-trainable radial basis functions, 8 attention heads and a cutoff of ι0Å.
ANIIx, COMP6
ANIIx 39 40 was built using an involved active learning process on the ANII data set, giving approximately 5M data points containing organic molecules composed of C, N. and H. The COMP6 [39] is benchmarking data set consists of five benchmarks (GDB07t009, GDB 10to13, Tripeptides, DrugBank, and ANI-MD) that cover broad regions of organic and biochemical space containing also the elements C, N. O and H, and a sixth one from the previously existing S66x8 noncovalent interaction benchmark. Energies and forces for all non-equilibrium conformations presented were computed at wBB7x55/6-31 random 80%/10%//1 traiming/valle weapiing//aliiation/teese splits, initializing the model with different random seeds. Reported results on A and COMP6 are the average eπrors over these splits.
level of theory. We trained the model on ANIIx by using three
Models trained with ANIIX had interaction layers, 128 hidden channels, used 32 non-trainable radial basis functions and a cutoff of 4.5A. The MLP for the embedding part had 2 linear layers, mapping 128 hidden channels as [128, 256, 384], with SiLU activation. The MLPs in the interaction layers had linear layers, mapping the 64 radial basis functions as [64, 128, 256, 384], also with SiLU non-linearities. The scalar output MLP mapped [384, 128, 64, 1] with SiLU activations, giving an atomic contribution to the energy, which after addition for all atoms, the model was trained to match the reference energy of the molecule with subtracted atomic reference energies. All layers were initialized using default PyTorch initialization, except for the last two layers of the output MI.P which were initialized using Xavier uniform distribution and with vanishing biases. Layer normalizations used the default PyTorch parameters. For training, we used a batch size of 64, an initial learning rate of lc-4 which was reached after 1000-step linear warm-up, the Adam optimizer with PyTorch default parameters and weighted MSE losses of energy and forces with weights and 100. The learning rate decayed using on-plateau scheduler based on the total weighted validation MSE losses, with a patience of 4 and a decay factor of 0.5. We did not use an exponential moving average for the validation loss. Training automatically stopped when the learning rate reached 1e-7 or when the total weighted validation MSE losses did not improve for 30 epochs. Inference batch size was 64. We used gradient clipping at norm 100. Training was performed on 4 NVIDIA GeForce RTX 2080 Ti with float32 precision. The Equivariant Transformer [16] model used for comparison purposes had 4 interaction layers, used 64 non-trainable radial basis functions, 8 attention heads and a cutoff of 5Å.
Scalar, vector and tensor properties of ethanol
We used the reference data for ethanol in vacuum provided together with the FieldSchNet paper 42 47), in which energies, forces, molecular dipole moments, polarizability tensors and nuclear shielding tensors were computed with the PBEO functional. To be consistent with results from FieldSehNet [42] and PaiNN [13), 8000 conformations for training and 1000 for both validation and testing were considered. We used three random splits, initializing the model with different random seeds. Reported results are the average errors and the standard deviation between parentheses of the last significant digit over these splits.
18

---

As mentioned in the main text, after generating atomic full tensor embeddings X we used different output modules to process their decompositions _ , and predict the desired properties.
Energies and forces were predicted as described in the scalar output subsection. For the remaining properties:
Molecular dipole moments were predicted by identifying the n skew-symmetric parts A with n vectors v-1) and mapping with first linear layer these n vectors to n/2 vectors, and then with a second linear layer from n//2 vectors to a single vector y per atom. In parallel, the norms [|
are fed to a two-layer MLP with SiLU activation to obtain a single scalar per atom ||xx00|] Then, the molccular dipole moment prediction is obtained by summing over all atoms μ = E(s) lI (
For the prediction of polarizability tensors we used only 7(0 and 5 to directly enforce their symmetry, We fed these components to two sets of two linear layers. to produce two distinct single atomic predictions a and ag Simultaneously, the norms ||7(*) + |(0) were processed by a two-layer MLP with SiLU activation to produce two scalars per atom ||a, || and ||o ||og)]] Then, (8) (4)
polarizability was predicted by adding all per-atom contributions a E() llo ||
+ (9 llog"
Eventually, for the prediction of nuclear shielding tensors, we generated a pseudovector skew- symmetric contribution per atom (8) as described in the Tensor output subsection. We applied three sets of two linear layers to the components ((0) , ((),, stt) to obtain three single predictions || A4 5 were fed to two-layer MLP give
per atom of (3)
S The norms
(i)
101 used to predict per-atom nuclear shieldings as r(1) (lllss)) (8)) where wi() are element-dependee weights. We used 1 to account for the different magnitudes of the nuclear - 1/0.167, w 1/0.022, w.
three scalars per atom ||0 11. 11 A,
(8)
w (1l=)) 10,/ (1) + 110 (6)
A, 11 (2)
A, + 11o 0
shielding tensors, as suggested in
Models trained with ethanol in vacuum had 2 interaction layers, 128 hidden channels, used 32 non-trainable radial basis functions and a cutoff of 4.5A. The MLP for the embedding part had 2 linear layers, mapping 128 hidden channels as [128, 256, 384], with SiLU activation. The MLPs in the interaction layers had 3 linear layers, mapping the 32 radial basis functions as [32, 128, 256, 384), also with SiLU non-linearities. The scalar output MLP mapped [384, 128, 64, 1] with SiLU activations, giving an atomic contribution to the energy which after addition for all atoms is scaled by the training data standard deviation and shifted with the training data mean, on top of vector and tensor predictions obtained as described above. All layers were initialized using default PyTorch initialization, except for the last two layers of the output MLP which were initialized using Xavier uniform distribution and with vanishing biases. Layer normalizations used the default Py Torch parameters. For training. we used a batch size of 8, an initial learning rate of le-3 without warm-up, the Adam optimizer with PyTorch default parameters and weighted MSE losses of energy, forces, dipole moments, polarizabilities and nuclear shieldings, with weights of 0.5 2)", 1115.111111111111 (the squared conversion factors from Ha to kcal/mol and Ha/Bohr to kcal/mol/A), 100, 100 and 100, respectively. The learning rate decayed using on-plateau scheduler based on the total weighted validation MSE losses, with a patience of 30 and a decay factor of 0.75. We used exponential moving average for the energy MSE validation loss with weight 0.99. Training automatically stopped when the learning rate reached 1e-8 or when the total weighted validation MSE losses did not improve for 300 epochs. Inference batch size was 64. We used gradient clipping at norm 100. Training was performed used the original data atomic units, performing unit conversion at test stage, using two NVIDIA GeForce RTX 2080 Ti with float32 precision. Errors in chemical shifts were computed as MAEs of mean traces of predicted and target nuclear shielding tensors.
Equivariance, interaction and cutoff ablations
We used the rMD17 data set and training details for both aspirin and toluene, varying only the number of layers, the cutoff radius and the interaction products, as detailed in the experiment. We used three random splits, initializing the model with different random seeds. Reported results are the average errors over these splits.
19

---

Computational cost
TorchScript code optimization and PyTorch (2.0) compilation were not used for the experiments. For TensorNet, we used the model with rMD17 architectural specifications. The ET used 128 hidden channels, 32 radial basis functions and 8 attention heads. For inference time, we used the blocked autorange benchmarking functionality in PyTorch with a minimum runtime of 10. For training step and epoch timings, we waited until the value stabilized. All experiments were performed with floar32 precision.
20


Reviews:
[{"review_id": "iQG6vm8BVt", "reviewer": "Reviewer_1mQm", "summary": "The paper introduces Tensor3 Net, a message-passing neural network architecture designed for molecular systems representation. Tensor3 Net leverages rank-2 Cartesian tensor representations and O(3)-equivariance. The tensors are decomposed into rotation group irreducible representations, enabling separate processing of scalars, vectors, and tensors when necessary.", "questions": "See weaknesses. ", "limitations": "The authors do not discuss the limitations of their approach. ", "rating": 2, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 1, "strengths": "The paper is well-written and the idea of building a O(3) equivariant model with Cartesian Tensor is valid, given the computation complexity of models such as e3nn, based on spherical harmonics. ", "weaknesses": "The one main weakness is that the author only model rank-2 tensors, which is relatively arbitrary. What it should be done is to introduce a network with general rank cartesian tensors (these will then be 3 x 3 x 3 x ... x 3 tensors). I see the value of adding the l=2 components in a cartesian way too marginal for publication at Neurips. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the feedback on the model in terms of its validity compared to e3nn. Here we address the main point raised.\n \n**The one main weakness is that the author only model rank-2 tensors, which is relatively arbitrary. What it should be done is to introduce a network with general rank cartesian tensors (these will then be 3 x 3 x 3 x ... x 3 tensors). I see the value of adding the l=2 components in a cartesian way too marginal for publication at Neurips.**\n\nTensorNet proposes a novel idea on how to achieve O(3)-equivariance using simple and fast matrix operations. These operations are commonly executed in neural networks and fast on GPUs. \n\nEven though we are limiting the operations to rank-2 Cartesian tensors, the model achieves state-of-the-art accuracies or performs better than models based on e3nn or using up to rank-3 spherical tensors (and pseudotensors). The use of rank-3 spherical tensors renders those models more computationally expensive, and in fact, no neural network potentials have used a rank higher than 3, while they have the possibility of using arbitrary higher ranks (of course, at extra cost). TensorNet achieves the same accuracy as models using higher-rank tensors (and pseudotensors). \n\nThe results are important because if we want to enable molecular dynamics simulations using neural network potentials at quantum-level of accuracy, we need to compute forces for millions and even billions of time steps and fast potentials are critical for that.  We foresee TensorNet’s approach to be important in the quest to replace molecular mechanics potentials commonly used in bio-molecular dynamics.\nIn fact, any equivariant model based on Cartesian vectors (such as PaiNN [15] and the ET [16]), can benefit from the idea and be rewritten using TensorNet’s formalism, by identifying scalar features with tensor features (matrix features) proportional to the identity matrix, and vector features with skew-symmetric tensors. Given the current implementation of Cartesian vector equivariant models, separated into scalar and vector pathways, one could consider the incorporation of a rank-2 feature pathway. In this regard, TensorNet provides a framework that allows working in a unified way with ‘general’ geometrical objects (full tensors) and their decomposition into scalars, vectors, and tensors, as opposed to designing differentiated pathways for different-rank features. \n\nThe decomposition of arbitrary rank Cartesian tensors into irreducible representations of the rotation group is highly non-trivial. For the case 3x3x3 (rank 3), it can be shown (see for example the paper *‘Decomposition of third-order constitutive tensors’* from Y. Itin and S. Reches, available on arXiv, subsection 4.2, and Figure 1 on page 24) that the 3x3x3 = 27-dimensional representation can be decomposed into 1 + 3 + 3 + 3 + 5 + 5 + 7, that is a scalar, three vectors, two quadrupoles, and an octupole (as opposed to rank-2 decomposition, 3x3 = 9 = 1 + 3 + 5). The decomposition is very complex (subsection 4.2) and involves a significant amount of operations. These decompositions for arbitrary ranks would be unfeasible in terms of memory and computational efficiency. In contrast, for rank-2 Cartesian tensors, the decomposition only requires computing a trace and a matrix transpose and already achieves SOTA.\n\nIn terms of limitations, constraining TensorNet to rank-2 restricts the prediction of quantities built on top of higher-rank tensors, like functions expanded in terms of spherical harmonics, e.g. electronic densities. However, TensorNet can still predict the density as a scalar on each point, the limitation comes from the expansion in terms of higher-rank tensors. We want to emphasize that our ultimate intention is to predict energies and forces in a fast and accurate way, and therefore this would not strictly be a requirement. Also, the prediction of these quantities of rank higher than two is uncommon. Examples of physical quantities that can correctly be predicted up to rank-2 apart from energies and forces are molecular or atomic dipoles, polarizability tensors, nuclear-shielding tensors, and quadrupole moments, which account for the vast majority of quantities that are used in molecular settings.\n\nIn summary, we believe that TensorNet deserves publication here because it is a novel idea that will be very useful to be integrated by any models aiming for O(3) equivariance using fast operations. The rank-2 restriction does not limit the applicability of the model due to SOTA accuracy and the wide range of physical quantities to which it can be applied. \n"}, {"review_id": "AqenxURJMq", "reviewer": "Reviewer_3Tca", "summary": "This paper proposes a cartesian tensor representation for efficient learning of molecular potentials. It enables the feature mixing process to be a simple matrix product operation. In addition, the matrix product operation is simplified by cost-effective decomposition techniques. Experimental results demonstrate that the proposed method can effectively reach a comparable performance with a much smaller number of parameters. ", "questions": "It seems the proposed architecture is an extension on TorchMD-net. Is it possible to apply this refinement to other equivariant models? If not, then this limitation will significantly lower its significance. ", "limitations": "NA. ", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "(1) The proposed method is very technically solid and this reviewer also thinks the efficient computation of equivariant architectures becomes increasingly important; \n\n(2) The extension of torchMD-net with efficient decomposition techniques is reasonably motivated, just like low-rank decomposition techniques for large language models;\n\n(3) Experimental settings are very solid, covering enough number of experimental settings, which fully support the effectiveness of the proposed method. ", "weaknesses": "(1) The presentation is not that easy to follow and the notation is a little bit complicated, which brings additional hardness for readers to understand the core idea. In addition, although the paper is more about mathematical techniques utilization, no figure illustration provided is still very tough for readers to quickly capture the core idea; \n\n(2) The proposed method is more like an application of efficient tensor decomposition techniques to the existing equivariant network architectures. Without molecular domain-specific insights somehow lowers the significance of the proposed method. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the nice feedback on the strengths of TensorNet. We clarify here the other points raised.\n\n**The presentation is not that easy to follow and the notation is a little bit complicated, which brings additional hardness for readers to understand the core idea. In addition, although the paper is more about mathematical techniques utilization, no figure illustration provided is still very tough for readers to quickly capture the core idea**\n\nWhile the mathematics can be a bit overwhelming, we tried to make it as simple as possible using the supplementary information for the more technical parts. We do provide a figure illustration in supplementary, however, we will introduce another more illustrative figure (see general rebuttal’s pdf) and, space permitting, include the one found in supplementary in the main manuscript. Also, we will use the extra space to further clarify the explanation of the mathematical techniques.\n\n**The proposed method is more like an application of efficient tensor decomposition techniques to the existing equivariant network architectures. Without molecular domain-specific insights somehow lowers the significance of the proposed method**\n\nTo our understanding, TensorNet, apart from being a mathematical framework, contains physical inductive biases which are meaningful in molecular domain-specific settings. A common approach in molecular physics is to expand some interaction energy in terms of charges (scalars), dipoles (vectors), quadrupoles (rank-2 tensors) and so on, see for example reference [13] (sections 2 and 3). In TensorNet, atomic features are learnable full tensors that depend on neighboring atoms and that can be decomposed precisely into scalars, vectors, and tensors which interact with other atoms’ features by means of matrix products, giving rise to new scalar, vector and tensor features. These products can be regarded as computing all possible combinations (interactions) between scalars, vectors and tensors, akin to scalar/dipole, dipole/dipole, dipole/quadrupole and quadrupole/quadrupole interactions (and so on) in a unified way, contributing to the predicted potential energy of the system. \n\nIn fact, we also show that for ethanol in a vacuum, TensorNet can simultaneously and accurately predict physical quantities of different geometrical nature from shared atomic features, pointing to the fact that these features are physically meaningful. These facts will be reinforced in the manuscript.\n\n**It seems the proposed architecture is an extension on TorchMD-net. Is it possible to apply this refinement to other equivariant models? If not, then this limitation will significantly lower its significance.**\n\nTensorNet is not an extension of the Equivariant Transformer, but a completely new model built on the same codebase. TorchMD-NET is a highly optimized PyTorch-based library for neural network potentials, in which the Equivariant Transformer (ET) [16] (probably the model the reviewer refers to when mentioning TorchMD-NET) is one of several existing models. Other models include a graph neural network similar to SchNet, an invariant transformer, and currently TensorNet. \n\nIn fact, any equivariant model based on Cartesian vectors (such as PaiNN [15] and the ET [16]), can be rewritten using TensorNet’s formalism, by identifying scalar features with tensor features (matrix features) proportional to the identity matrix, and vector features with skew-symmetric tensors. As the reviewer mentions, these models lack the ‘refinement’ of the incorporation of rank-2 features. Given the current implementation of Cartesian vector equivariant models, separated into scalar and vector pathways, one could consider the incorporation of a rank-2 feature pathway. In this regard, TensorNet provides a framework that allows working in a unified way with ‘general’ geometrical objects (full tensors) and their decomposition into scalars, vectors, and tensors, as opposed to designing differentiated pathways for different-rank features. \n"}, {"review_id": "fK8WBpRRpi", "reviewer": "Reviewer_u3os", "summary": "This paper proposes TensorNet, an O(3)-equivariant neural network architecture for molecules. Using the decomposition of a 3x3 matrix into a scalar, vector, and matrix shown in Eq. (2), TensorNet efficiently computes the interaction of O(3)-equivariant features up to l=2, where l is the degree (frequency) of the O(3) representation. The performance is evaluated on several standard benchmarks, such as qm9, which show that TensorNet achieves equivalent or better performance than baselines. ", "questions": "n/a", "limitations": "Limitations are not explicitly addressed. No particular concern for potential negative societal impact. ", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. TensorNet is a novel method. I have never seen the construction of an equivariant net based on the decomposition (2).\n2. The performance is evaluated with several different datasets and in terms of different metrics (e.g., prediction error, computation speed).\n3. The performance is comparable to or better than existing approaches. ", "weaknesses": "1. The paper has room to improve in terms of presentation. I'm unfamiliar with the chemistry (molecule) domain, and some parts seem challenging to understand without expertise. For example, a vector r_ij is defined on line 175, but the mathematical definition is not described. Also, the meaning of the cutoff radius is unexplained. Another point is that there is no reference nor citation to Eq. (2), the core equation of this paper. These are not well known in the machine learning community, and it could be better to introduce them in plain words. \n2. The limitations of the proposed method are not explicitly discussed. One limitation is that TensorNet cannot capture higher degree (l>2) information of O(3). \n3. It is argued that some existing methods have a downside: \"the computation of tensor products in most of these models containing higher-rank tensors and pseudotensors can be expensive\" (lines 111-112). However, only one method (ET) is compared with the proposed method in terms of computational cost. Also, there is no theoretical evaluation of the complexity e.g., using big-O notation. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their positive feedback on the novelty of the approach, completeness of the validation, and results. We here answer the further points raised.\n \n**The paper has room to improve in terms of presentation. I'm unfamiliar with the chemistry (molecule) domain, and some parts seem challenging to understand without expertise. For example, a vector r_ij is defined on line 175, but the mathematical definition is not described. Also, the meaning of the cutoff radius is unexplained. Another point is that there is no reference nor citation to Eq. (2), the core equation of this paper. These are not well known in the machine learning community, and it could be better to introduce them in plain words.**\n\nWe understand the concern that the reviewer mentions about some definitions not being explicitly addressed for researchers not familiar with the molecular domain. We will address this and improve it in the new version of the manuscript. It will also be clarified in the additional figure that would be incorporated into the manuscript (see general rebuttal’s pdf).\n\nThe edges between nodes (atoms) are built by defining a cutoff radius, where atoms $i$ and $j$ are connected by an edge $ij$ if they are at a distance smaller than the cutoff radius. At the same time, for every edge, we can consider a relative position vector $r_{ij}$, computed from 3D coordinates of atoms $r_i$ and $r_j$, as $r_{ij}$ = $r_j$ - $r_i$. \n\nRegarding the decomposition (2), the reference is [21] (page 3, first and second paragraphs), even though it is true that it is found two or three lines below the expression and it may not be immediately clear that it refers to (2). Therefore, we will add an additional reference and citation to clarify.\n\n**The limitations of the proposed method are not explicitly discussed. One limitation is that TensorNet cannot capture higher degree (l>2) information of O(3).**\n\nRegarding limitations, we will include them as an independent section at the end of the manuscript. \n\nTensorNet is not using tensors of rank higher than 2 to favor computational efficiency. In terms of accuracy for lower-rank quantities, this does not seem to be a problem, as results show that TensorNet performs competitively or better compared to models using higher ranks. TensorNet cannot choose the maximum rank of the tensors being used,  in contrast to spherical models, where the rank can be seen as a hyperparameter, and the architectures are built in a general way that can consistently include higher ranks. The main limitation comes from the prediction of quantities built on top of higher-rank tensors, like functions expanded in terms of spherical harmonics, e.g. electronic densities. However, TensorNet can still predict the density as a scalar on each point, the limitation comes from the representation in terms of higher-rank tensors (bearing in mind that our ultimate intention is to predict energies and forces in a fast and accurate way, and therefore this is not a requirement). The prediction of these quantities of rank higher than two is uncommon. Examples of physical quantities that can correctly be predicted by TensorNet up to rank-2 are energies, forces, molecular or atomic dipoles, polarizability tensors, nuclear-shielding tensors, and quadrupole moments, which account for the majority of quantities that are used in molecular systems.\n\n**It is argued that some existing methods have a downside: \"the computation of tensor products in most of these models containing higher-rank tensors and pseudotensors can be expensive\" (lines 111-112). However, only one method (ET) is compared with the proposed method in terms of computational cost. Also, there is no theoretical evaluation of the complexity e.g., using big-O notation.**\n\nThe computational cost of models based on spherical tensors is well-established and we will add more comparisons in the manuscript. For example, in reference [20] (MACE), a fast model is built that explicitly addresses this issue, and one can see in the reference in Table 2, at the end of page 8, times to compute energy and forces are compared for NequIP [17], BOTNet [19] and MACE on the 3BPA molecule with 27 atoms. MACE is 4x faster than previous models when considering the full model (L=2, the SOTA model for rMD17). TensorNet 2L (our SOTA model for rMD17) with a batch size of 32 is ~4x faster than MACE on the same molecule and batch size, on a NVIDIA RTX 2080 Ti, and therefore it is by far faster than NequIP and BOTNet. \n\nWe compared our speed to the ET [16], a Cartesian vector equivariant model, as opposed to the other higher-rank spherical models, because Cartesian vector representations are found to be fast (also, notice that the SOTA ET model had 6 layers, and we compare to 4 and 5 layers). Our intention is to emphasize that TensorNet, with accuracies comparable to SOTA spherical models, is as computationally efficient (or even more computationally efficient in some cases) than a Cartesian vector model for small molecules.\n\nFinally, the complexity of TensorNet is linear regardless of the rank with the number of atoms N (~O(N)), which is the main scaling parameter. The number of neighbors per atom M is controlled by the cutoff radius and the density, which are usually fixed across models and systems, therefore it is not usually a scaling factor to take into account, but rather a constant factor that is included in the overall speed. In any case, since matrix products are performed after aggregation over neighbors (we compute matrix products node-wise, not edge-wise), these do not scale with M. This is in contrast to spherical models, where tensor products are computed on edges, and therefore display a worse scaling with the number of neighbors M, that is, a worse scaling when increasing the cutoff radius at fixed density. We will add this to the manuscript.\n"}, {"review_id": "lmmUVB9M3F", "reviewer": "Reviewer_QwGa", "summary": "The paper presents TensorNet, a novel O(3)-equivariant message-passing neural network for efficient representation of molecular systems in scientific research. This model utilizes Cartesian tensor atomic embeddings to simplify feature mixing via matrix product operations. By decomposing tensors into rotation group irreducible representations, it enables independent processing of scalars, vectors, and tensors when required. TensorNet outperforms higher-rank spherical tensor models in performance, utilizing fewer parameters, even with a single interaction layer for small molecule potential energies. Additionally, TensorNet can accurately predict vector and tensor molecular quantities on top of potential energies and forces, greatly reducing the model's computational cost. Therefore, TensorNet provides a promising framework for developing state-of-the-art equivariant models with enhanced efficiency and computational affordability.", "questions": "1. The proposed architecture (illustrated in Fig. A1) appears quite intricate. Could you provide more insight or intuitive reasoning behind the specific design choices in the model? What fundamental principles or considerations have influenced the complexity and uniqueness of this architecture?\n\n2. Could you delve deeper into the primary constraints or drawbacks of the TensorNet model compared to the existing methods in the field? What are the potential areas where other models might still hold an advantage, and how does TensorNet aim to address these challenges in future iterations or improvements?", "limitations": "N/A", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. This paper adeptly presents TensorNet, a new learning model that not only establishes state-of-the-art performance but does so with a remarkable reduction in the number of parameters utilized. This marks a significant leap in model efficiency without sacrificing performance quality, setting a new benchmark in the field.\n\n2. The majority of the empirical outcomes display a marked enhancement over existing methods, with the implementation of TensorNet consistently yielding superior results. The experimental evidence provided substantiates the model's efficacy, reinforcing the robustness and applicability of this innovative approach in real-world scenarios.", "weaknesses": "The paper's exposition of the model architecture can be somewhat challenging to comprehend due to its complex nature. A potential improvement would be the inclusion of intuitive diagrams or visual aids within the main body of the text, not just in the Appendix. Simplified illustrations, possibly even a step-by-step visual guide, could greatly enhance the reader's understanding of the architecture and make the methodology more accessible to a broader audience.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their positive feedback on the model. We now address the questions raised.\n\n**The proposed architecture (illustrated in Fig. A1) appears quite intricate. Could you provide more insight or intuitive reasoning behind the specific design choices in the model? What fundamental principles or considerations have influenced the complexity and uniqueness of this architecture?**\n\nWe will address this issue in terms of clarity of the figure and the model by adding an additional figure (see general rebuttal’s pdf).\n\nOne can justify the design choices by taking into account the precise operations that allow TensorNet to be O(3) equivariant and taking them as the fundamental building blocks:\n\n1) Linear combinations of irreducible components: mixing independently different rank features (different linear layers for I, A and S) is presumed to give more flexibility to the model when compared to just applying a linear layer to X (= I+A+S). \n\n2) Multiplication with invariant quantities: again, we assume it is preferable to modify independently I, A, S compared to modifying X alone. We can distinguish subcases:\n\n- In the case of the embedding module: We need to generate scalars, vectors and tensors from relative position vectors, that is, edge-wise. Assuming more model flexibility when modifying independently I, A, S (compared to just modifying X) by multiplying invariant quantities, it makes sense to encode edge interatomic distance with different learnable functions $f^{(ij)}_I$, $f^{(ij)}_A$, $f^{(ij)}_S$. After aggregation of edge-wise features into node-wise features, we use again the modification with invariants by taking the (normalized) norm and obtaining new invariants through a MLP that will modify the node’s I, A, S, enabling further processing of the direct neighborhood that has been aggregated. \n\n- In interaction layers, when considering the message-passing framework, we find that encoding interatomic distance (weighted with a cutoff) separately to incoming I, A, S features from neighbors goes in line with previous message-passing models for molecular potentials while exploiting the flexibility of independent modifications of different rank features. \n\n3) Interaction via matrix product: We prove that using a particular sum of products the model exhibits O(3) equivariance. Importantly, these products are performed at the node level, after aggregation, which is more efficient. There are an infinite amount of expressions in terms of matrix products that allow O(3) equivariance: for example, w * YM + w * MY, where w can be a learnable vector of weights with hidden channels length and * is the element-wise product. The most simple one, 1 * YM + 1 * MY, is the one currently used.\n\n4) We can consider separately the normalization operations that take place. There are two normalizations:\n\n- Using LayerNorm: It can be expected that when taking the norms of I, A and S, their magnitude can be quite different. LayerNorm would equilibrate the norms coming from I, A and S. In fact, the LayerNorm used in the output MLP giving energies gave a substantial benefit in terms of accuracy.\n\n- Normalizing X tensors by means of X -> X / (||X||+1): we found this to stabilize training (without this normalization, for some molecules the loss became NaN), +1 ensures numerical stability (though any other number could be used). This normalization also enforces a non-linearity. Since ||X|| = || I + A + S || is used (as opposed to using separately || I ||, ||A||, ||S||) the components I, A, S receive non-linear contributions from the other components.\n\nWe hope this explanation helps in clarifying the architectural choices, given that the use of the different components I, A, S along the architecture makes the diagram convoluted.\n\n**Could you delve deeper into the primary constraints or drawbacks of the TensorNet model compared to the existing methods in the field? What are the potential areas where other models might still hold an advantage, and how does TensorNet aim to address these challenges in future iterations or improvements?**\n\nThe main limitation comes from the prediction of quantities built on top of higher-rank tensors, like functions expanded in terms of spherical harmonics, e.g. electronic densities. However, TensorNet can still predict the density as a scalar on each point, the limitation comes from the representation in terms of higher-rank tensors (bearing in mind that our ultimate intention is to predict energies and forces in a fast and accurate way, and therefore this is not a requirement). The prediction of these quantities of rank higher than two is very uncommon. Examples of physical quantities that can correctly be predicted up to rank-2 apart from energies and forces are molecular or atomic dipoles, polarizability tensors, nuclear-shielding tensors, and quadrupole moments, which account for the vast majority of quantities that are used in molecular settings.\n\nWe do think that TensorNet will be primarily used to predict neural network potentials (molecular energies and forces), and molecular properties which is the reason for which it has been designed and where it holds its main advantages. e3nn based models will still be used where speed is less important or somehow the higher rank tensors are required, e.g. for representing functions in terms of spherical harmonics.\n\nFor the future, we are currently exploiting the simple matrix operations of TensorNet further to achieve even greater performance speed-ups, mainly using Cuda graphs and other standard techniques. Integration with a widely used biomolecular dynamics package is already in progress with the final aim to be used in practical applications.\n"}]

Author rebuttal:
The questions raised by the reviewers have been addressed in their individual rebuttals. However, since some of them expressed concerns about the clarity of exposition, we use this general rebuttal to attach a figure that will be added to the manuscript, which we hope provides clarification on both the notation and the methods.

Meta-review: The paper proposes a neural network architecture that is invariant to SO(3) transformations. The key idea is to represent matrices in a rotation invariant name (named 'rank-2 tensors', but for ML community such notation is not very standard). The reviewers agree that paper is well-written and the results are quite convincing. The main concern (I also share it) is that the paper only works with matrices, not going into higher-dimensional tensors. The authors, however, have provided a rebuttal where they explain the theory behind that. Unfortunately, the reviewer have not engaged back into the discussion, thus I base my decision on my own reading and understanding of the paper, as well as other reviews, which are positive.
Decision: Accept (poster)