[论文解读] FAENet: Frame Averaging Equivariant GNN for Materials Modeling

Alexandre Duval, Victor Schmidt+5 more|arXiv (Cornell University)|Apr 28, 2023

Machine Learning in Materials Science被引用 7

一句话总结

FAENet 提出了一种新颖的轻量化图神经网络，在材料建模中通过几何数据的随机坐标系平均（SFA）实现 E(3)-等变性，无需架构约束。该方法在 OC20 和 QM9 数据集上实现了最先进水平的准确率与计算效率，优于以往的等变模型，在可扩展性和对称性保持方面表现更优。

ABSTRACT

Applications of machine learning techniques for materials modeling typically involve functions known to be equivariant or invariant to specific symmetries. While graph neural networks (GNNs) have proven successful in such tasks, they enforce symmetries via the model architecture, which often reduces their expressivity, scalability and comprehensibility. In this paper, we introduce (1) a flexible framework relying on stochastic frame-averaging (SFA) to make any model E(3)-equivariant or invariant through data transformations. (2) FAENet: a simple, fast and expressive GNN, optimized for SFA, that processes geometric information without any symmetrypreserving design constraints. We prove the validity of our method theoretically and empirically demonstrate its superior accuracy and computational scalability in materials modeling on the OC20 dataset (S2EF, IS2RE) as well as common molecular modeling tasks (QM9, QM7-X). A package implementation is available at https://faenet.readthedocs.io.

研究动机与目标

• 解决现有等变 GNN 在通过复杂架构设计强制对称性时所面临的表达能力与可扩展性限制。
• 开发一种灵活的框架，通过数据变换而非模型架构来保持 E(3) 对称性。
• 设计一种轻量化、高性能的 GNN——FAENet，可在自由处理几何信息的同时保持理论与实证的对称性不变性。
• 通过实证验证基于 SFA 的对称性强制相较于数据增强与无对称性强制在材料性质预测中的优越性。

提出的方法

• 提出随机坐标系平均（SFA）作为一种保持对称性的数据增强技术，将原子构型投影到规范坐标系中，以强制实现 E(3)-等变性。
• 引入 FAENet，一种处理相对原子位置而无需对称性保持架构约束的 GNN 架构，其不变性依赖于 SFA。
• 使用坐标系平均将输入图转换到规范坐标系，使任何模型都能对旋转、反射和位移保持等变性。
• 在 SFA 中采用可微的坐标系选择机制，以随机方式采样多个参考坐标系，提升泛化能力与对称性鲁棒性。
• 利用基于相对位置的几何消息传递机制，作为边特征，最大化表达能力，同时通过数据预处理保持对称性。
• 通过旋转不变性、反射不变性和力等变性等指标，在多个数据集上验证方法的对称性保持能力。

Figure 1: Overview of FAENet architecture making use of the Stochastic Frame Averaging framework. FAENet takes an input sample $D$ and makes an invariant $\Phi^{\text{{invar}}}_{{\mathcal{F}}}$ (e.g. energy) and/or equivariant prediction $\Phi^{\text{{equiv}}}_{{\mathcal{F}}}$ (e.g. atomic forces) b

实验结果

研究问题

• RQ1 是否可以通过数据变换而非设计对称性保持层，在 GNN 中无需架构约束地有效强制实现 E(3)-等变性？
• RQ2 与数据增强和无对称性强制相比，基于 SFA 的对称性强制在模型准确率与对称性保持方面表现如何？
• RQ3 轻量级 GNN 如 FAENet 在材料性质预测任务中是否比复杂且受对称性约束的模型具有更高的准确率与可扩展性？
• RQ4 与完整坐标系平均或标准数据增强相比，SFA 在多大程度上提升了对称性泛化能力？
• RQ5 FAENet 是否能在极少架构修改下有效泛化至包括固态晶体和分子在内的多样化材料数据集？

主要发现

• FAENet 搭配 SFA 在 S2EF 数据集上实现 464 meV 的测试 MAE，优于 SchNet（918 meV）与 ForceNet（655 meV），证明其在准确率与可扩展性方面的优越性。
• SFA 将 S2EF 上的旋转不变性误差降低至 6.94 meV，反射不变性误差降低至 6.96 meV，显著优于数据增强（8.88 和 9.12 meV）与无 FA（26.81 和 27.61 meV）。
• 在 IS2RE 上，FAENet 搭配 SFA 实现 569 meV 的 MAE，较 ForceNet 提高 12%，较 SchNet 提高 13%，同时保持了强大的对称性保持能力。
• SE(3)-SFA 变体的指标略优于完整 FA，表明从更少的坐标系中学习可提升泛化能力，同时不损失对称性。
• 无 FA 与无训练（No-Train）版本的对称性指标显著更差（如旋转不变性误差达 51.1 meV），证明对称性强制对鲁棒性能至关重要。
• SFA 在对称性保持与预测准确率方面均优于数据增强，尤其在三维对称性下，S2EF 上的旋转等变性误差降低 60%。

Figure 2: OC-20 Direct IS2RE performance / scalability trade-off. FAENet achieves the best MAE across methods while being much faster than the majority of baselines.