So3krates：面向分子系统任意尺度相互作用的等变注意力机制
机器学习
2023-01-10 v3
摘要
机器学习方法在量子化学中的应用使得众多传统ab-initio方法在计算上难以处理的化学现象得以研究。然而，分子和材料的某些量子力学性质依赖于非局域电子效应，由于其高效建模的困难，这些效应常被忽略。本工作提出了一种适配底层物理的改进注意力机制，能够恢复相关的非局域效应。具体而言，我们引入球谐坐标（SPHCs）以反映分子中每个原子的高阶几何信息，从而在SPHC空间中实现非局域的注意力表述。我们提出的模型So3krates——一种基于自注意力的消息传递神经网络——将几何信息从原子特征中解耦，使它们可独立地适用于注意力机制。由此我们构建了球滤波器，将欧几里得空间中连续滤波器的概念拓展至SPHC空间，并作为球自注意力机制的基础。我们表明，与其他已发表的方法不同，So3krates能够描述任意长度尺度上的非局域量子力学效应。此外，我们发现证据表明纳入高阶几何关联提升了数据效率并改善了泛化能力。So3krates在流行基准上达到或超越最先进性能，值得注意的是，与其他模型相比，其所需参数量显著更少（0.25-0.4倍），同时带来大幅加速（训练6-14倍，推理2-11倍）。
引用
@article{arxiv.2205.14276,
title = {So3krates: Equivariant attention for interactions on arbitrary length-scales in molecular systems},
author = {J. Thorben Frank and Oliver T. Unke and Klaus-Robert Müller},
journal= {arXiv preprint arXiv:2205.14276},
year = {2023}
}