[Paper Review] Complete and Efficient Graph Transformers for Crystal Material Property Prediction

Keqiang Yan, Cong Fu+3 more|arXiv (Cornell University)|Mar 18, 2024

Machine Learning in Materials Science6 citations

TL;DR

This paper introduces SE(3) invariant and SO(3) equivariant crystal graph representations and two graph transformer variants, iComFormer and eComFormer, achieving state-of-the-art results on three crystal benchmarks.

ABSTRACT

Crystal structures are characterized by atomic bases within a primitive unit cell that repeats along a regular lattice throughout 3D space. The periodic and infinite nature of crystals poses unique challenges for geometric graph representation learning. Specifically, constructing graphs that effectively capture the complete geometric information of crystals and handle chiral crystals remains an unsolved and challenging problem. In this paper, we introduce a novel approach that utilizes the periodic patterns of unit cells to establish the lattice-based representation for each atom, enabling efficient and expressive graph representations of crystals. Furthermore, we propose ComFormer, a SE(3) transformer designed specifically for crystalline materials. ComFormer includes two variants; namely, iComFormer that employs invariant geometric descriptors of Euclidean distances and angles, and eComFormer that utilizes equivariant vector representations. Experimental results demonstrate the state-of-the-art predictive accuracy of ComFormer variants on various tasks across three widely-used crystal benchmarks. Our code is publicly available as part of the AIRS library (https://github.com/divelab/AIRS).

Motivation & Objective

• Address the challenge of geometrically complete crystal representations that capture periodic patterns and distinguish chiral crystals.
• Develop lattice-based crystal graphs that encode complete geometric information while respecting crystal passive symmetries.
• Propose two scalable graph transformer variants (iComFormer and eComFormer) leveraging invariant and equivariant representations.
• Empirically validate on three standard crystal benchmarks (JARVIS, Materials Project, MatBench) achieving state-of-the-art results.

Proposed method

• Define crystal passive symmetries (unit cell SE(3) invariance, unit cell SO(3) equivariance, periodic invariance) and formalize geometric completeness for crystals.
• Construct SE(3) invariant crystal graphs by duplicating atoms across the lattice and encoding lattice vectors and angles as edge features; ensure periodic invariance in lattice construction.
• Construct SO(3) equivariant crystal graphs using edge vectors with spherical-harmonics-based features and tensor-product based message passing to preserve rotation equivariance.
• Implement iComFormer with SE(3) invariant graphs and an O(nk) transformer architecture (node-wise and edge-wise updates) for efficient computation.
• Implement eComFormer with SO(3) equivariant graphs using tensor-product layers to fuse vector edge features and node features in a rotation-equivariant manner.
• Demonstrate completeness and scalability, with complexity O(nk) where n is the number of atoms per unit cell and k is the average number of neighbors.

Experimental results

Research questions

• RQ1 Can SE(3) invariant and SO(3) equivariant crystal graphs achieve geometric completeness for crystalline materials?
• RQ2 Do iComFormer and eComFormer achieve state-of-the-art predictive accuracy across large-scale, medium-scale, and small-scale crystal benchmarks?
• RQ3 How does incorporating lattice-based periodicity and chiral sensitivity influence material property predictions?
• RQ4 What is the computational efficiency of the proposed models compared to existing crystal GNNs?
• RQ5 Do the proposed representations distinguish chiral crystals and remain robust under unit cell periodic transformations?

Key findings

• iComFormer achieves 8% improvement on formation energy, 10% on E_total, and 15% on E_hull over PotNet on JARVIS.
• eComFormer achieves 20% improvement on E_hull and second-best on other JARVIS tasks.
• On Materials Project, iComFormer delivers best performance on three of four tasks; eComFormer best on formation energy.
• On MatBench, iComFormer matches or surpasses strong baselines across e_form-132752 and jdft2d-636 tasks, with competitive RMSEs.
• ComFormer variants achieve state-of-the-art results on large-scale e_form-132752, small-to-medium scale benchmarks, and robust performance with limited training data (e.g., Matformer/MP comparisons).
• The methods maintain geometric completeness while scaling to crystals of varying unit cell sizes and periodic representations.