Revisiting the Canonicalization for Fast and Accurate Crystal Tensor Property Prediction

TL;DR

This work tackles the challenge of predicting high-order crystal tensor properties while preserving $O(3)$-equivariance. It introduces GoeCTP, a canonicalization-based framework that uses polar decomposition to obtain an invariant form and a corresponding orthogonal transformation to recover equivariant outputs, avoiding heavy equivariant architectures. The approach yields competitive accuracy across dielectric, piezoelectric, and elastic tensors and delivers up to 13x speedups with minimal overhead, acting as a plug-and-play enhancement for existing crystal graph models. The results demonstrate robust $O(3)$-equivariant tensor prediction and practical efficiency, with discussion of limitations related to 2D/space-group cases and potential integration of symmetry priors for further gains.

Abstract

Predicting the tensor properties of crystalline materials is a fundamental task in materials science. Unlike scalar property prediction, which requires invariance, tensor property prediction requires maintaining O(3) group tensor equivariance. Achieving such equivariance typically demands specialized architectural designs, which substantially increase computational cost. Canonicalization, a classical technique for geometry, has recently been explored for efficient learning with symmetry.In this work, we revisit the problem of crystal tensor property prediction through the lens of canonicalization. Specifically, we demonstrate how polar decomposition, a simple yet efficient algebraic method, can serve as a form of canonicalization and be leveraged to ensure equivariant tensor property prediction. Building upon this insight, we propose a general O(3)-equivariant framework for fast and accurate crystal tensor property prediction, referred to as GoeCTP. By utilizing canonicalization, GoeCTP achieves high efficiency without requiring the explicit incorporation of equivariance constraints into the network architecture.Experimental results indicate that GoeCTP achieves the high prediction accuracy and runs up to 13 times faster compared to existing state-of-the-art methods, underscoring its effectiveness and efficiency.

Figures (3)

• Figure 1: The illustration of $O(3)$-equivariance for crystal tensor prediction. The visualization of the crystal structures on the left was generated using VESTA momma2011vesta, while the visualization of the crystal tensor properties on the right follows the method of yanspace and VELAS ran2023velas.
• Figure 2: The Illustration of GoeCTP. To begin with, (1) the R&R module rotates and reflects the input crystal structure, which may have an arbitrary orientation, to the canonical form of this crystal. Next, (2) Crystal Graph Construction module organizes the adjusted input into a crystal graph, followed by (3) the Node & Edge Feature Embedding module, which encodes the features of the crystal graph. Subsequently, (4) the Prediction module leverages these embedded features to predict the canonical form of tensor properties corresponding to the canonical form of input crystal. Finally, (5) the Reverse R&R module applies the orthogonal matrix $\mathbf{Q}$, obtained from the R&R module, to ensure the equivariance of the output tensor properties.
• Figure 3: The detailed architectures of the node-wise transformer layer and node-wise equivariant updating layer, adapted from yancomplete.

Theorems & Definitions (14)

• Definition 1: Orbit.
• Definition 2
• Definition 3
• Proposition 1
• Theorem 1
• Proposition 2
• Lemma 1
• proof
• Lemma 2
• proof