SlowGuess/NeurIPS_2023_OCR huggingface dataset:
The Frobenius norm of some tensor representation  $X$  in TensorNet is invariant under the full orthogonal group  $\mathrm{O}(3)$ . This follows from the cyclic permutation invariance of the trace. As previously shown, parity induces the transposition of  $X$ , which amounts to a cyclic permutation of the matrix product inside the trace operator,  $\operatorname{Tr}(X^{\mathrm{T}}X) = \operatorname{Tr}(XX^{\mathrm{T}})$ . When considering some rotation  $R$ ,  $\operatorname{Tr}(X^{\mathrm{T}}X) \to \operatorname{Tr}((RXRT)^{\mathrm{T}}(RXRT)) = \operatorname{Tr}(RXRTRXRRT) = \operatorname{Tr}(RX^{\mathrm{T}}R^{-1}RXRT) = \operatorname{Tr}(RX^{\mathrm{T}}XR^{\mathrm{T}}R) = \operatorname{Tr}(X^{\mathrm{T}}XR^{\mathrm{T}}R) = \operatorname{Tr}(X^{\mathrm{T}}XR^{-1}R) = \operatorname{Tr}(X^{\mathrm{T}}X)$ .
O(3) to SO(3) symmetry breaking
The embedding module in TensorNet always produces irreducible representations that are scalars, vectors and tensors, that is, it preserves the parity of the initial components. As previously shown, the computation of the Frobenius norm is O(3)-invariant in this case, making the normalization an operation that also preserves initial parities, since it can be regarded as the modification of the irreducible representations by means of invariant weights. However, depending on the interaction product, O(3)-equivariance can be broken to SO(3)-equivariance in the first interaction layer.
Input irreducible representations for the first interaction layer, which are the outputs of the embedding module, have a well-defined parity. The first normalization that takes place makes use of the Frobenius norm, which at this point is  $\mathrm{O}(3)$ -invariant, giving  $\mathrm{O}(3)$ -equivariant representations with the desired parity. Nevertheless, if one considers an interaction product proportional to  $Y^{(i)}M^{(i)}$  and computes its irreducible decomposition to generate new features, skew-symmetric and symmetric traceless representations will receive both even and odd parity contributions, as proved above (expressions A.3 and A.4). At this point, one can write some node's full tensor embedding  $X$  as
$$
X = I + A ^ {+} + A ^ {-} + S ^ {+} + S ^ {-}, \tag {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  $\mathrm{Tr}(X^{\mathrm{T}}X)$  with current node representations  $X$  used to normalize features is no longer O(3)-invariant. One can write
$$
\begin{array}{l} \operatorname {T r} \left(X ^ {\mathrm {T}} X\right) = \operatorname {T r} \left(\left(I + A ^ {+} + A ^ {-} + S ^ {+} + S ^ {-}\right) ^ {\mathrm {T}} \left(I + A ^ {+} + A ^ {-} + S ^ {+} + S ^ {-}\right)\right) = \\ = \operatorname {T r} \left(\left(I - A ^ {+} - A ^ {-} + S ^ {+} + S ^ {-}\right) \left(I + A ^ {+} + A ^ {-} + S ^ {+} + S ^ {-}\right)\right), \tag {A.10} \\ \end{array}
$$
where we have made use of the symmetry or skew-symmetry of the matrices under transposition, and considering the action of parity  $A^{-} \to -A^{-}, S^{-} \to -S^{-}$ ,
$$
\operatorname {T r} \left(\left(I - A ^ {+} + A ^ {-} + S ^ {+} - S ^ {-}\right) \left(I + A ^ {+} - A ^ {-} + S ^ {+} - S ^ {-}\right)\right), \tag {A.11}
$$
which is manifestly different from A.10 and cannot be reduced to a cyclic permutation of A.10. Thus, the Frobenius norm of representations  $X$  at 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 are 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)-equivariant, but SO(3)-equivariant, and the computation of any Frobenius norm (such as the one used in the scalar output module to predict energies) will only be SO(3)-invariant.
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 130,831 optimized structures of molecules that contain up to 9 heavy elements from C, N, O and F. On top of the structures, several quantum-chemical properties computed at the DFT B3LYP/6-31G(2df,p) level of theory are provided. To be consistent with previous work, we used 110,000 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.
Models trained with QM9 had 3 interaction layers, 256 hidden channels, used 64 non-trainable radial basis functions and a cutoff of  $5\AA$ . The MLP for the embedding part had 2 linear layers, mapping 256 hidden channels as [256, 512, 768], 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 1e-4 which was reached after a 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 loss, with a patience of 15 and a 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 1e-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; 37], a data set of small organic molecules obtained by running molecular dynamics simulations. The DFT PBE/def2-SVP level of theory, a very dense DFT integration grid and a 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 1 and 2 interaction layers, 128 hidden channels, used 32 non-trainable radial basis functions and a cutoff of  $4.5\AA$ . 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. 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 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 1e-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 a single NVIDIA GeForce RTX 2080 Ti with float32 precision.
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  $\omega$ B97M-D3BJ/def2-TZVPPD level of theory. After filtering molecules containing forces higher than 50.94 eV/Å ≈ 1 Hartree/Bohr, SPICE (v 1.1.3) contains approximately 1M data points. We used three random  $80\% / 10\% / 10\%$  training/validation/test splits, initializing the model with different random seeds. Reported results are the average errors over these splits.
Models trained with SPICE had 2 interaction layers, 128 hidden channels, used 64 non-trainable radial basis functions and a cutoff of  $10\AA$ . 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 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 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 1e-4 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 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  $10\AA$ .
AN1x, COMP6
ANI1x [39; 40] was built using an involved active learning process on the ANI1 data set, giving approximately 5M data points containing organic molecules composed of C, N, O and H. The COMP6 [39] is a benchmarking data set consists of five benchmarks (GDB07to09, GDB10to13, 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  $\omega$ B97x59/6-31G(d) level of theory. We trained the model on ANI1x by using three random  $80\% / 10\% / 10\%$  training/validation/test splits, initializing the model with different random seeds. Reported results on ANI1x and COMP6 are the average errors over these splits.
Models trained with AN1x had 2 interaction layers, 128 hidden channels, used 32 non-trainable radial basis functions and a cutoff of  $4.5\AA$ . 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 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 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 64, an initial learning rate of 1e-4 which was reached after a 1000-step linear warm-up, the Adam optimizer with PyTorch default parameters and weighted MSE losses of energy and forces with weights 1 and 100. The learning rate decayed using an 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\AA$ .
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 PBE0 functional. To be consistent with results from FieldSchNet [42] and PaiNN [15], 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.
As mentioned in the main text, after generating atomic full tensor embeddings  $X^{(i)}$ , we used different output modules to process their decompositions  $I^{(i)}, A^{(i)}, S^{(i)}$  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^{(i)}$  with  $n$  vectors  $\mathbf{v}^{(i)}$  and mapping with a 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  $\mu^{(i)}$  per atom. In parallel, the norms  $\| A^{(i)}\|$  are fed to a two-layer MLP with SiLU activation to obtain a single scalar per atom  $\| \mu^{(i)}\|$ . Then, the molecular dipole moment prediction is obtained by summing over all atoms  $\mu = \sum_{(i)}\| \mu^{(i)}\| \mu^{(i)}$ .
For the prediction of polarizability tensors we used only  $I^{(i)}$  and  $S^{(i)}$  to directly enforce their symmetry. We fed these components to two sets of two linear layers, to produce two distinct single atomic predictions  $\alpha_{I}^{(i)}$  and  $\alpha_{S}^{(i)}$ . Simultaneously, the norms  $\| I^{(i)} + S^{(i)}\|$  were processed by a two-layer MLP with SiLU activation to produce two scalars per atom  $\| \alpha_{I}^{(i)}\|$  and  $\| \alpha_{S}^{(i)}\|$ . Then, polarizability was predicted by adding all per-atom contributions  $\alpha = \sum_{(i)}\| \alpha_{I}^{(i)}\| \alpha_{I}^{(i)} + \| \alpha_{S}^{(i)}\| \alpha_{S}^{(i)}$ .
Eventually, for the prediction of nuclear shielding tensors, we generated a pseudovector skew-symmetric contribution per atom  $A_{p}^{(i)}$  as described in the Tensor output subsection. We applied three sets of two linear layers to the components  $I^{(i)}, A_{p}^{(i)}, S^{(i)}$  to obtain three single predictions per atom  $\sigma_{I}^{(i)}, \sigma_{A_{p}}^{(i)}, \sigma_{S}^{(i)}$ . The norms  $||I^{(i)} + A_{p}^{(i)} + S^{(i)}||$  were fed to a two-layer MLP to give three scalars per atom  $||\sigma_{I}^{(i)}||, ||\sigma_{A_{p}}^{(i)}||, ||\sigma_{S}^{(i)}||$ , used to predict per-atom nuclear shieldings as  $\sigma^{(i)} = w^{(i)}(||\sigma_{I}^{(i)}||\sigma_{I}^{(i)} + ||\sigma_{A_{p}}^{(i)}||\sigma_{A_{p}}^{(i)} + ||\sigma_{S}^{(i)}||\sigma_{S}^{(i)})$ , where  $w^{(i)}$  are element-dependent weights. We used  $w^{C} = 1 / 0.167$ ,  $w^{O} = 1 / 0.022$ ,  $w^{H} = 1$  to account for the different magnitudes of the nuclear shielding tensors, as suggested in [42].
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.5\AA$ . 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 a Xavier uniform distribution and with vanishing biases. Layer normalizations used the default PyTorch parameters. For training, we used a batch size of 8, an initial learning rate of 1e-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 \times (627.5)^{2}$ ,  $0.5 \times (1185.82117)^{2}$  (the squared conversion factors from Ha to kcal/mol and Ha/Bohr to kcal/mol/Å), 100, 100 and 100, respectively. The learning rate decayed using an 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 an 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 using 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.
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 float32 precision.