[论文解读] 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 Science被引用 6

一句话总结

本论文提出 SE(3) 不变和 SO(3) 变换等变的晶体图表示，以及两种图变换器变体 iComFormer 和 eComFormer，在三个晶体基准数据集上取得了最先进的结果。

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).

研究动机与目标

• 解决捕捉周期性模式并区分手性晶体的几何完整晶体表示的挑战。
• 开发基于晶格的晶体图，编码完整几何信息，同时尊重晶体被动对称性。
• 提出两种可扩展的图变换器变体（iComFormer 和 eComFormer），利用不变性和等变表示。
• 在三个标准晶体基准（JARVIS、Materials Project、MatBench）上进行实证验证，达到最先进的结果。

提出的方法

• 定义晶体的被动对称性（晶胞的 SE(3) 不变性、晶胞的 SO(3) 变换等变性、周期性不变性），并将晶体的几何完整性形式化。
• 通过在晶格中对原子进行复制并将晶格向量和夹角编码为边特征来构建 SE(3) 不变的晶体图；在晶格构建中确保周期性不变性。
• 使用带球谐基特征的边向量以及基于张量积的消息传递来保持旋转等变性，从而构建 SO(3) 变换等变的晶体图。
• 实现 iComFormer，采用 SE(3) 不变图和 O(nk) 变换器架构（节点级和边级更新），以实现高效计算。
• 实现 eComFormer，使用张量积层将向量边特征和节点特征以旋转等变方式融合，构建 SO(3) 变换等变图。
• 展示完整性与可扩展性，其复杂度为 O(nk)，其中 n 为晶胞中的原子数，k 为平均邻居数。

实验结果

研究问题

• RQ1 SE(3) 不变和 SO(3) 变换等变的晶体图能否实现晶体材料的几何完整性？
• RQ2 在大规模、中等规模和小规模晶体基准上，iComFormer 和 eComFormer 是否达到最先进的预测准确性？
• RQ3 将基于晶格的周期性和手性敏感性纳入预测，如何影响材料性质预测？
• RQ4 与现有晶体 GNNs 相比，所提出模型的计算效率如何？
• RQ5 所提出的表示是否能区分手性晶体，并在晶胞周期变换下保持鲁棒性？

主要发现

方法	形成能（meV/原子）	带隙(OPT)（eV）	E_total（meV）	带隙(MBJ)（eV）	E_hull（meV）
CGCNN	63	0.20	78	0.41	170
SchNet	45	0.19	47	0.43	140
MEGNET	47	0.145	58	0.34	84
GATGNN	47	0.170	56	0.51	120
ALIGNN	33.1	0.142	37	0.31	76
Matformer	32.5	0.137	35	0.30	64
PotNet	29.4	0.127	32	0.27	55
eComFormer	28.4	0.124	32	0.28	44
iComFormer	27.2	0.122	28.8	0.26	47

• iComFormer 在 JARVIS 上比 PotNet 在形成能、E_total 和 E_hull 分别提升 8%、10% 和 15%。
• eComFormer 在 E_hull 上提升 20%，在 JARVIS 的其他任务上为第二名。
• 在 Materials Project 上，iComFormer 在四个任务中的三个任务上提供最佳表现；eComFormer 在形成能上表现最佳。
• 在 MatBench 上，iComFormer 在 e_form-132752 和 jdft2d-636 任务中达到或超过强基线，且 RMSE 具有竞争力。
• ComFormer 变体在大规模 e_form-132752、小到中等规模基准上达到最先进的结果，并在有限训练数据下表现稳健（如 Matformer/MP 比较）。
• 该方法在晶胞大小和周期表示各异的晶体上扩展时，保持几何完整性。