[Paper Review] Equivariant Graph Neural Networks for Prediction of Tensor Material Properties of Crystals

Alex Heilman, Claire Schlesinger+1 more|arXiv (Cornell University)|Jun 5, 2024

Material Properties and Failure MechanismsMaterials Science3 citations

TL;DR

This paper proposes a novel SE(3)-equivariant graph neural network framework for direct, component-wise prediction of tensorial material properties—such as dielectric, piezoelectric, and elasticity tensors—from crystalline structures. By leveraging spherical harmonic decomposition and SO(3) equivariant message passing, the model predicts tensor components with rotational invariance, achieving state-of-the-art performance on tensor prediction tasks, though transfer learning across tasks showed limited gains due to data scarcity and model overparameterization.

ABSTRACT

Modern E(3)-Equivariant networks may be used to predict rotationally equivariant properties, including tensorial quantities. Three such quantities: the dielectric, piezoelectric, and elasticity tensors, are computationally expensive to produce ab initio for crystalline systems; however, with greater availability of such data in large material property databases, we now have a sufficient target space to begin training equivariant models in the prediction of such properties. Here we explicitly develop spherical harmonic decompositions of these tensorial properties using their general symmetries. We then apply three distinct E(3)-equivariant convolutional structures to the prediction of the components of these decompositions, allowing us to predict the aforementioned tensorial quantities in an equivariant manner and compare performance. We further report results testing the transferability of these predictive models between different tensorial target sets.

Motivation & Objective

• To enable direct, component-wise prediction of coordinate-dependent tensor properties (e.g., dielectric, piezoelectric, elasticity) from crystal structures.
• To address the limitation of existing models that predict only scalar, invariant properties by incorporating SO(3) equivariance.
• To explore transfer learning from scalar pretraining tasks to improve generalization for low-data tensor prediction tasks.
• To investigate whether chaining pretraining across multiple tensor tasks improves performance on downstream tensor predictions.
• To evaluate the expressiveness and practicality of SO(3) equivariant models in capturing crystal symmetry constraints for tensor prediction.

Proposed method

• The model uses SE(3)-equivariant graph neural networks where node features are represented as spherical harmonics indexed by rotational order ℓ and azimuthal order m.
• It employs SO(3) equivariant convolutions defined via Clebsch-Gordan coefficients to combine features and filters across different ℓ and m channels while preserving rotational symmetry.
• Tensor components are predicted by decomposing target tensors into harmonic subspaces using spherical harmonic decomposition, enabling direct readout from model outputs.
• The framework integrates self-interaction layers and nonlinearities that are equivariant under SO(3) transformations to maintain symmetry across layers.
• Transfer learning is explored via pretraining on scalar properties (formation energy, band gap) and chained pretraining across multiple tensor tasks.
• Model performance is evaluated using mean absolute error (MAE) across individual tensor components, with ablation on data sparsity and pretraining order.

Experimental results

Research questions

• RQ1 Can SO(3)-equivariant GNNs effectively predict component-wise tensorial material properties such as dielectric, piezoelectric, and elasticity tensors from crystal structures?
• RQ2 Does pretraining on scalar properties (e.g., formation energy, band gap) improve downstream performance on low-data tensor prediction tasks?
• RQ3 Does chaining pretraining across multiple tensor tasks (e.g., elastic → dielectric → piezoelectric) lead to better generalization than single-task pretraining?
• RQ4 Why do models fail to improve performance on the piezoelectric tensor despite pretraining, and is this due to data scarcity or model overparameterization?
• RQ5 To what extent does the SO(3) framework’s overparameterization—generating many unused spherical harmonic combinations—hinder practical tensor prediction in crystals?

Key findings

• The SE(3)-equivariant model achieved state-of-the-art performance on tensor component prediction, with MAE of 7.65 (log10 GPa) for the elasticity tensor using SETransformer.
• Pretraining on scalar targets (formation energy, band gap) led to minimal improvement in tensor prediction accuracy, with MAE values remaining nearly unchanged (e.g., 0.170 for piezoelectric tensor across models).
• Chained pretraining across multiple tensor tasks did not improve performance on the piezoelectric tensor, with MAE values consistently at 0.170–0.171 across all models.
• The dielectric tensor prediction performance degraded when pretraining included elastic tensor data, with MAE increasing from 4.82 to 4.67 (SEGNN) under chained pretraining, indicating negative transfer.
• The study suggests that the SO(3) framework may be overly expressive for crystal symmetry constraints, generating many redundant spherical harmonic combinations that are not physically meaningful.
• The authors conclude that restricting convolutions to crystal point group symmetries—rather than full SO(3)—could improve efficiency and accuracy by eliminating zero-valued tensor components a priori.