[论文解读] So3krates: Equivariant attention for interactions on arbitrary length-scales in molecular systems

J. Thorben Frank, Oliver T. Unke+1 more|arXiv (Cornell University)|May 28, 2022

Machine Learning in Materials Science被引用 15

一句话总结

So3krates 引入了一种新颖的等变注意力机制，通过学习作为动态、注意力驱动表示的球谐坐标（SPHCs），在分子系统中建模非局部电子效应。通过将几何信息与原子特征解耦，并在 SPHC 空间中应用球面自注意力机制，该模型在参数量减少 0.25–0.4 倍、训练与推理速度提升 2–14 倍的情况下，实现了最先进水平的精度，从而实现了跨任意长度尺度的高效长程相关性学习。

ABSTRACT

The application of machine learning methods in quantum chemistry has enabled the study of numerous chemical phenomena, which are computationally intractable with traditional ab-initio methods. However, some quantum mechanical properties of molecules and materials depend on non-local electronic effects, which are often neglected due to the difficulty of modeling them efficiently. This work proposes a modified attention mechanism adapted to the underlying physics, which allows to recover the relevant non-local effects. Namely, we introduce spherical harmonic coordinates (SPHCs) to reflect higher-order geometric information for each atom in a molecule, enabling a non-local formulation of attention in the SPHC space. Our proposed model So3krates - a self-attention based message passing neural network - uncouples geometric information from atomic features, making them independently amenable to attention mechanisms. Thereby we construct spherical filters, which extend the concept of continuous filters in Euclidean space to SPHC space and serve as foundation for a spherical self-attention mechanism. We show that in contrast to other published methods, So3krates is able to describe non-local quantum mechanical effects over arbitrary length scales. Further, we find evidence that the inclusion of higher-order geometric correlations increases data efficiency and improves generalization. So3krates matches or exceeds state-of-the-art performance on popular benchmarks, notably, requiring a significantly lower number of parameters (0.25 - 0.4x) while at the same time giving a substantial speedup (6 - 14x for training and 2 - 11x for inference) compared to other models.

研究动机与目标

• 为解决在传统机器学习力场中由于计算和架构限制而常被忽略的分子系统中非局部电子效应建模挑战。
• 开发一种可微分、可学习的高阶几何相关性表示，以捕捉超出局部邻域的长程相互作用。
• 通过适配 3D 分子系统的 SO(3) 对称性的自注意力机制，实现对任意长度尺度上分子相互作用的高效、可扩展建模。
• 通过将球谐坐标（SPHCs）作为可学习特征显式建模几何相关性，提升数据效率和泛化能力。
• 在保持或超越最先进性能的同时，降低模型复杂度和推理成本。

提出的方法

• 该模型将球谐坐标（SPHCs）作为原子几何信息的可学习、动态表示，替代固定或手工设计的几何特征。
• 它将不变的原子特征与等变的 SPHCs 解耦，使注意力机制可独立作用于两者，从而在 SPHC 空间中实现局部与非局部注意力。
• 通过将连续滤波器从欧氏空间推广至球谐空间的球形滤波器，构建了保持 SO(3) 等变性的球面自注意力机制。
• 该架构通过在 SPHC 上进行注意力的消息传递，实现远距离信息传播，使模型能够捕捉非局部电子效应。
• 模型通过能量与力的联合损失进行训练，当仅监督力时，通过最小二乘拟合计算能量积分常数。
• SPHCs 在训练过程中被更新，使模型能够学习反映分子结构中高阶相关性的最优几何表示。

实验结果

研究问题

• RQ1 可学习的、基于注意力的球谐坐标表示能否提升分子系统中非局部电子效应的建模效果？
• RQ2 通过 SPHCs 将几何信息与原子特征解耦，是否能提升机器学习力场中的数据效率与泛化能力？
• RQ3 球面自注意力机制能否有效捕捉分子系统中任意长度尺度上的长程相关性？
• RQ4 通过 SPHCs 引入高阶几何相关性，在降低模型复杂度的同时，能在多大程度上维持或提升预测精度？
• RQ5 与现有最先进模型相比，So3krates 在参数效率、训练速度和推理速度方面表现如何？

主要发现

• So3krates 在多个基准测试中达到最先进水平，包括 QM7-X250 和 MD17，相比先前 SOTA 模型，模型参数量减少 0.25–0.4 倍。
• 尽管精度更高，该模型在训练阶段相比其他基于注意力的模型提速 6–14 倍，在推理阶段提速 2–11 倍。
• So3krates 有效捕捉了聚炔中氢旋转子的长程相关性，忠实地重现了从头算基态，而局部模型则无法做到。
• 通过 SPHCs 引入高阶几何相关性显著提升了数据效率与泛化能力，体现在小样本数据下的性能更优。
• 结果表明，将 SPHCs 作为架构的一部分进行学习，可实现有效的非局部注意力，优于使用固定或手工设计几何特征的模型。
• 当仅以力进行训练时，模型通过最小二乘拟合计算能量积分常数，确保无需额外监督即可实现精确的能量预测。