A General Tensor Prediction Framework Based on Graph Neural Networks

Graph neural networks (GNNs) have been shown to be extremelyflexibleand accurate in predicting the physical properties of molecules andcrystals. However, traditional invariant GNNs are not compatible withdirectional properties, which currently limits their usage to theprediction of only invariant scalar properties. To address this issue,here we propose a general framework, i.e., an edge-based tensor predictiongraph neural network, in which a tensor is expressed as the linearcombination of the local spatial components projected on the edgedirections of clusters with varying sizes. This tensor decompositionis rotationally equivariant and exactly satisfies the symmetry ofthe local structures. The accuracy and universality of our new frameworkare demonstrated by the successful prediction of various tensor propertiesfrom first to third order. The framework proposed in this work willenable GNNs to step into the broad field of prediction of directionalproperties.

Automated summary

In this study, we propose a new framework called edge-based tensor prediction graph neural network that addresses the incompatibility of traditional invariant GNNs with directional properties. By expressing tensors as linear combinations of local spatial components projected on the edge directions of clusters with varying sizes, our framework is rotationally equivariant and satisfies the symmetry of local structures. We demonstrate the accuracy and universality of our framework by successfully predicting various tensor properties from first to third order. This work enables GNNs to step into the broad field of prediction of directional properties.