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.