[论文解读] Point Edge Transformer model trained on highly distorted methane configurations

Sergey N. Pozdnyakov, Michele Ceriotti|arXiv (Cornell University)|May 30, 2023

Image Processing and 3D Reconstruction被引用 15

一句话总结

本文提出 Point Edge Transformer (PET) 以及 Equivariant Coordinate System Ensemble (ECSE)，在非等变骨干之上实现严格的旋转等变性，并在多项原子级基准测试中达到最先进的结果。

ABSTRACT

Point clouds are versatile representations of 3D objects and have found widespread application in science and engineering. Many successful deep-learning models have been proposed that use them as input. The domain of chemical and materials modeling is especially challenging because exact compliance with physical constraints is highly desirable for a model to be usable in practice. These constraints include smoothness and invariance with respect to translations, rotations, and permutations of identical atoms. If these requirements are not rigorously fulfilled, atomistic simulations might lead to absurd outcomes even if the model has excellent accuracy. Consequently, dedicated architectures, which achieve invariance by restricting their design space, have been developed. General-purpose point-cloud models are more varied but often disregard rotational symmetry. We propose a general symmetrization method that adds rotational equivariance to any given model while preserving all the other requirements. Our approach simplifies the development of better atomic-scale machine-learning schemes by relaxing the constraints on the design space and making it possible to incorporate ideas that proved effective in other domains. We demonstrate this idea by introducing the Point Edge Transformer (PET) architecture, which is not intrinsically equivariant but achieves state-of-the-art performance on several benchmark datasets of molecules and solids. A-posteriori application of our general protocol makes PET exactly equivariant, with minimal changes to its accuracy.

研究动机与目标

• 说明在原子级机器学习中对精确旋转对称性需求以避免模拟中的伪影。
• 提出一个通用的对称化协议（ECSE），使任意骨架架构在事后实现精确的旋转等变性，同时不损害其他对称性属性。
• 开发 Point Edge Transformer (PET)，利用边缘聚焦的变换器式消息传递以提升表达能力，同时保持局部性与效率。
• 在失真甲烷配置、小分子和具有多种原子种类的周期性晶体上展示 PET 的最先进性能。

提出的方法

• 引入 Equivariant Coordinate System Ensemble (ECSE)，对由 cutoff 内邻居对定义的所有坐标系进行预测平均，并通过权重确保平滑性。
• 为原子环境定义自适应内 cutoff R_in(A_i)，以在确保至少存在一个定义良好的坐标系的同时限制坐标系数量。
• 对非等变骨干事后应用 ECSE，以获得精确的旋转等变性，同时保持其它不变性（置换、平移、平滑性）。
• 提出 Point Edge Transformer (PET)，一种基于变换器的、以边特征为中心的消息传递架构，处理中心原子及其邻居的 token，并使用置换不变的自注意力。
• 通过 f_c 在注意力中实现平滑 cutoff，确保可微分，避免邻居在 cutoff 出现/消失时的不连续性。
• 在数据集包括液态水、COLL、MnO、HME21、HEA 和 QM9 双极子等上对 PET 进行有 ECSE 与无 ECSE 的基准测试。

实验结果

研究问题

• RQ1 一个通用的对称化协议能否在不损害其他对称性属性的前提下，在任意骨架模型之上强制实现精确的旋转等变性？
• RQ2 在包括液体、分子与含多种元素类型的固体的多样原子数据集上，经过 ECSE 加持的 PET 架构是否实现最先进的精度？
• RQ3 通过增加模型深度/块数，性能提升的程度如何？ECSE 对推断成本与精度有何影响？
• RQ4 与领先的等变模型（如 NEQUIP、MACE）在标准基准与失真配置上相比，PET 的表现如何？

主要发现

dataset	MAE f (COLL)	MAE E /at. (COLL)	RMSE f (MnO)	RMSE E /at. (MnO)	MAE f (HME21)	MAE E /at. (HME21)	MAE f (HEA)	MAE E /at. (HEA)
SOTA	26.4 [17]	47 [86]	125 [88]	1.11 [88]	138 [89]	15.7 [89]	190 [90]	10 [90]
PET ( y0 )	23.1	12.0	22.7	0.312	140.5 ± 2.0	17.8 ± 0.1	60.2	1.87
PET ( yS )	23.1	11.9	22.7	0.304	141.6 ± 1.9	17.8 ± 0.1	60.1	1.87
PET (ens.)				128.5	16.8

• 相对于最先进的等变模型 NEQUIP，PET 在相关基准的液态水配置上将误差降低约 30% 。
• 以 cutoff 4.25 Å 为阈值时，PET 在近液态水数据中的力的 MAE 为 14.4 meV/Å，展示了对长程相互作用的强处理能力。
• 在 COLL 数据集上，PET 力的 MAE 提升约 13%，原子化能误差降低约 4 倍，相较于先前的最佳方法。
• 在 MnO、HME21、HEA 数据集上，PET 及其 y0/yS 变体达到与最先进水平相当的性能，且 PET yS 在所报告的指标上几乎与 PET y0 相同。
• 一个 PET 模型集成（ens.）进一步降低误差，在某些指标上（如集成行中的 16.8 的报告值）优于单模型变体。
• ECSE 实现了精确的旋转等变性（且协变输出），对精度的影响极小，使 PET 适用于具有严格对称性要求的原子级模拟。