Kento Nishio, Kiyou Shibata, Teruyasu Mizoguchi Institute of Industrial Science, The University of Tokyo DS06: Integrating Machine Learning with simulations for accelerated material modeling Prediction of forces and energy based on rotational invariant spherical representations

2 Graph Neural Networks (GNNs) for Material Science J. Gasteiger et al., arXiv abs/2003.03123 (2020). J. Gasteiger et al., arXiv abs/2106.08903 (2021). K. T. Schütt et al., arXiv abs/2102.03150 (2021). S. Batzner et al., Nat. Commun. 13, 2453 (2022). Spectra prediction (PaiNN) MD simulation (NequIP, Allegro) High prediction accuracy and efficiency. Differentiable and can be used as a force field. Need to explicitly compute invariant geometric features for 3-body or 4-body combinations. Material design (CGCNN) MD simulation (SchNet) Property prediction (DimeNet, GemNet) T. Xie and J. C. Grossman, Phys. Rev. Lett. 120, 145301 (2018). K. T. Schütt et al., J. Chem. Phys. 148, 241722 (2018). A. Musaelian et al., Nat. Commun. 14, 579 (2023). 2 Efficiently incorporate higher-order geometric features. High prediction accuracy for equivariant properties. (ex: Forces) Limitation on nonlinear transformations. High computational cost of equivariant tensor operations. Energy predictions (invariance) Force predictions (equivariance)

3 SCN Accuracy improvement efforts C. Lawrence et al., arXiv abs/2206.14331 (2022). S. Passaro et al., arXiv abs/2302.03655 (2023). 3 ü Improved expressiveness by relaxing non-linearity limitations. ü Approximation to equivariant operations on lower order groups to reduce cost. eSCN Improved expressiveness. Achieved state-of-the-art accuracy in OC20 dataset. Models do not have strict invariance and equivariance to transformations. Models with excellent accuracy and strict physical constraints are required.

4 Purpose of this study Construction of a strictly invariant GNN model for the prediction of equivariant material properties with high expressive power and low computational cost. ◎ No restrictions on non-linear transformations = High expressive power. ◎ Predicts strictly equivariant properties by extending invariant outputs. J. Gasteiger et al., arXiv abs/2106.08903 (2021). ◎ The proposed method can efficiently consider dihedral angles while being strictly invariant. Invariant GNN Model & Expressiveness Computational cost C. W. Park et al.,. npj Computational Materials. 7, 1–9 (2021). 4

5 Converts the relative vector into the invariance 1. Calculate the relative vector ⃗ 𝑟!" , ⃗ 𝑟!#! , ⃗ 𝑟!#" 2. Create the new coordinate axis with ⃗ 𝑟!#! & ⃗ 𝑟!#" . 3. Coordinate transformation: ⃗ 𝑟!" to ⃗ 𝑟!" !,%. 4. Repeat this process in various combinations in order of ploximity and make multiple coordinate axis and get transformed vectors 𝑟!" !,&. ※ In 3-dimensional case, a third coordinate axis is set by taking the outer product. Conversion Procedure (2D example) 𝑢! : 𝑘 nearest neighbor atom from the central atom 𝑖. ⃗ 𝑟"# : Vector from central atom 𝑖 to neighbor atom 𝑗. ⃗ 𝑟"# ",% : Transformed relative vector from central atom 𝑖 to 𝑗. ,1 ⃗ 𝑟"# ",% is a strictly rotationally invariant relative vector. 5 𝑢 # 𝑢 $ 𝑟"# ",& 𝑢 # 𝑢% 𝑟"# ",' 𝑢% 𝑢 $ 𝑟"# ",( ・・・ ・・・ 𝑟"# ",% 𝑢 & ! 𝑢 & "

6 Message passing with spherical representation 𝑴"# ' = 𝜙()**+,) 𝒎"# ' , ⃗ 𝑟"# , ⃗ 𝑟"- , 𝜃#"-, ⃗ 𝑟"# ",. 𝒉"# '/. = 𝜙012+3) 𝒉 𝒉" ', 𝒉# ', 𝑴"# ' 𝒎"# '/. = 𝜙012+3) 𝒎 𝒉" ', 𝒉# ', 𝒎"# ' , 𝑴"# ' 𝒆'() !",& = 𝑗* ⃗ 𝑟!" 𝑌*+ ⃗ 𝑟!" !,&/ ⃗ 𝑟!" !,& J. Gasteiger et al., arXiv abs/2106.08903 (2021). 𝒉" ) : node embedding (invariant) 𝒎"# ) : edge embedding (invariant) ⃗ 𝑟"# ",% : Transformed relative vector from central atom 𝑖 to 𝑗. ü𝜙,-../0- → Bilinear product between 𝒆'() !",& and 𝒎!" * . Maintain invariance by adding all transformed vectors. ü𝜙123/4- → same as GemNet Generate Spherical representation 𝒆 𝐒𝐁𝐅 𝒊𝒋,𝒕 Message & Update functions 6

7 Prediction of material properties 𝑓"# = 𝜙603103 7 𝒎"# 8 , 𝒆 9:; "# 𝑭" = - #∈𝒩* 𝑓"# 𝒓"# 𝒓"# ü Sum between edges in the cutoff sphere and impose equivariance by multiplying the relative vector. ü The output interatomic forces 𝑭! are strictly equivariant. 𝐸" = - #∈𝒩* 𝜙603103 > 𝒎"# 8 , 𝒆9:; "# 𝐸 = - "∈𝒢 𝐸" ü Sum between edges in the cutoff sphere. ü The output 𝐸! and 𝐸 are strictly invariant. The model output was strictly invariant and equivariant. Energy prediction (invariant) Extend to force prediction (equivariant) 7

8 Experiment with 3 different Datasets I. E. Castelli et al., Energy Environ. Sci. 5, 9034–9043 (2012). J. Gasteiger et al.,. arXiv 2011.14115 (2020). L. Chanussot et al., ACS Catal. 11, 6059–6072 (2021). Cubic Perovskite Dataset COLL Dataset OC20 Dataset ü Formation energies of cubic perovskite type compounds. ü 56 elemental species with atomic numbers from 3 to 83. ü Energies and forces from MD simulations of molecular collisions. ü QM9 extension with many variations in structures. ü Energies and forces of surface adsorption dynamics. ü More than 1.2 million adsorption structures consisting of 82 adsorption molecules and 55 different surfaces. Training condition ü Optimizer: Adam ü GPU: NVIDIA A100 80GB PCIe D. P. Kingma and J. Ba, arXiv 1412.6980 (2014). 8

9 Results| Comparison of accuracy with previous models COLL Dataset Prediction accuracy of energy for COLL Dataset. (MAE [meV / atom]) 9 Energy Forces Formation Energy Perovskite Dataset

10 Results| Comparison of accuracy with previous models 9 OC20 2M Dataset Energy Forces l The predictions were as accurate as or better than state-of-the-art models particularly in predicting energy (invariance). l Achieved high accuracy in both energy and force for diverse datasets of different materials and stability. l Comparing the previous models (DimeNet++ & GemNet-dT), which consider effects up to 3-body interactions, the proposed model achieved higher accuracy in both energy and forces. l The transformed invariant relative vectors are effective in improving accuracy. l Incorporating higher order geometry.

11 Results| Effect of the number of coordinate axis The maximum number of coordinate axis per node (Perovskite). l Prediction accuracy improved up to 𝑡 = 8 in perovskite dataset but up to 𝑡 = 5 in COLL dataset. l It is not necessary to consider all dihedral angles. 10 𝑢 # 𝑢 $ 𝑟"# ",& 𝑢 # 𝑢% 𝑟"# ",' 𝑢% 𝑢 $ 𝑟"# ",( ・・・ ・・・ 𝑟"# ",% 𝑢 & ! 𝑢 & " ü The accuracy was verified by increasing the maximum number of coordinate axes for each node. It is important to consider the appropriate number of dihedral angles for the system The maximum number of coordinate axis per node (COLL).

12 Results| Comparison of training cost SchNet DimeNet++ GemNet-dT GemNet-dQ This study (t=5) min / epoch 3.8 7.2 8.4 23.1 16.1 Forces MAE 172.0 40.0 43.1 38.1 42.1 Time per epoch in training process with COLL dataset. Computational Cost l Faster than GemNet-dQ, which considers all 4-body combinations within the cutoff. l The model can train on OC20 dataset but GemNet-dQ cannot. 11

13 Summary l We constructed a strictly invariant & equivariant physical property prediction model that efficiently incorporates higher- order geometric features in message passing function. l We proposed the method to efficiently capture the dihedral angle using coordinate transformations. l Compared to previous models, we achieved better or the same accuracy paritcuraly on invariant properties. l We found that preparing multiple coordinate axis improves accuracy. It was also found that there may be an appropriate number of 4-body combinations depending on the system. l Our model was more computationally efficient than prior models that explicitly compute between all 4-body combinations in the cutoff sphere and can be learned on OC20. 11 ,1 𝑢 . 𝑢 @ 𝑟12 1,3 𝑢 . 𝑢A 𝑟12 1,4 ・・・ ・・・ 𝑟12 1,5 𝑢 - + 𝑢 - ,

14 GNNs for materials informatics. CGCNN T. Xie and J. C. Grossman, Phys. Rev. Lett. 120, 145301 (2018). SchNet K. T. Schütt et al., J. Chem. Phys. 148, 241722 (2018). ü Nodes, edges, and global attributes are represented as fixed-length vectors. ◎ High flexibility and improved accuracy archived by incorporating more inductive bias. Modeling graph structure with Graph Neural Network (GNN) GNNs for materials P. W. Battaglia et al., arXiv abs/1806.01261 (2018). ü Node: each atom in system. ü Edge: bond information. Message Passing architecture models interactions between atoms. 12

15 Example of the coordination number = 6. To capture the many-body effect more efficiently ü When multiple coordinate axes are taken, the axes are taken mechanically in the order of the nearest distance. If the first coordination number is 6, then the axis must be ; C< = 15 or more to be considered up to the second coordination. If we can make the coordinate axes and between atoms of the second and third coordination with less t, the accuracy will be improved with more efficiently. 𝑢 # 𝑢 $ 𝑟"# ",& 𝑢 # 𝑢% 𝑟"# ",' ・・・ ・・・ 𝑟"# ",% 𝑢 & ! 𝑢 & " 13

16 Training conditions for Cubic Perovskite Dataset Model parameters • Embedding size of 𝒉! * : 128 • Embedding size of 𝒎!" * : 128 • Embedding size of 𝒆=() !" : 128 • Embedding size of 𝒆>() !"? : 128 • Embedding size of 𝒆'() !",& : 128 • Embedding size of T-MP convolution : 64 • Embedding size of D-MP convolution : 64 • The number of rbf : 6 • The number of sbf : 294 • The number of interaction block : 4 Learning prameters • Batch size : 128 • Initial LR : 0.001 • Patience of ReduceLROnPlateau : 10 • Decay factor of ReduceLROnPlateau : 0.6 • Patience of early stopping : 80 14

17 Training conditions for COLL Dataset Model parameters • Embedding size of 𝒉! * : 128 • Embedding size of 𝒎!" * : 256 • Embedding size of 𝒆=() !" : 32 • Embedding size of 𝒆>() !"? : 32 • Embedding size of 𝒆'() !",& : 64 • Embedding size of T-MP convolution : 64 • Embedding size of D-MP convolution : 32 • The number of rbf : 6 • The number of sbf : 294 • The number of interaction block : 4 Learning prameters • 𝜌 of loss function : 0.99 • Batch size : 32 • Initial LR : 0.001 • Patience of ReduceLROnPlateau : 10 • Decay factor of ReduceLROnPlateau : 0.5 • Patience of early stopping : 30 𝐿 = 1 − 𝜌 MAE 𝐸, 8 𝐸 + 𝜌MAE 𝐹" , 8 𝐹" 15

18 Training conditions for OC20 Datasets Model parameters • Embedding size of 𝒉! * : 256 • Embedding size of 𝒎!" * : 512 • Embedding size of 𝒆=() !" : 16 • Embedding size of 𝒆>() !"? : 16 • Embedding size of 𝒆'() !",& : 32 • Embedding size of T-MP convolution : 64 • Embedding size of D-MP convolution : 32 • The number of rbf : 6 • The number of sbf : 294 • The number of interaction block : 4 Learning prameters • 𝜌 of loss function : 0.99 • Batch size : 48 • Initial LR : 0.0005 • Decay step of LR : 3 • Decay factor of LR : 0.3 • Max epochs : 20 16 𝐿 = 1 − 𝜌 MAE 𝐸, 8 𝐸 + 𝜌MAE 𝐹" , 8 𝐹"

19 Results| Force Cosine on OC20 2M validation Force Cosine values for OC20 2M Validation S2EF Dataset. Target SchNet DimeNet++ GemNet-dT SCN eSCN This study ⃗ 𝐹 0.109 0.217 0.557 0.650 0.662 0.568 17