[论文解读] A Self-Attention Ansatz for Ab-initio Quantum Chemistry

Ingrid von Glehn, James S. Spencer+1 more|arXiv (Cornell University)|Nov 24, 2022

Machine Learning in Materials Science被引用 30

一句话总结

Psiformer 是一种基于自注意力的波函数 Ansatz，用于从头算量子化学，其在地态能量方面优于 FermiNet 和 FermiNet+SchNet，尤其对于较大分子。

ABSTRACT

We present a novel neural network architecture using self-attention, the Wavefunction Transformer (Psiformer), which can be used as an approximation (or Ansatz) for solving the many-electron Schrödinger equation, the fundamental equation for quantum chemistry and material science. This equation can be solved from first principles, requiring no external training data. In recent years, deep neural networks like the FermiNet and PauliNet have been used to significantly improve the accuracy of these first-principle calculations, but they lack an attention-like mechanism for gating interactions between electrons. Here we show that the Psiformer can be used as a drop-in replacement for these other neural networks, often dramatically improving the accuracy of the calculations. On larger molecules especially, the ground state energy can be improved by dozens of kcal/mol, a qualitative leap over previous methods. This demonstrates that self-attention networks can learn complex quantum mechanical correlations between electrons, and are a promising route to reaching unprecedented accuracy in chemical calculations on larger systems.

研究动机与目标

• 在不依赖外部数据的情况下，开发一种基于自注意力的神经网络 Ansatz，用于求解多电子薛定谔方程。
• 相比现有神经网络 Ansatz，提升分子系统基态能量计算的准确性。
• 展示基于注意力的架构对更大分子规模的可扩展性。
• 通过 Jastrow 因子和包络函数引入 cusp 条件等物理约束。

提出的方法

• 将 Psi_theta 定义为 exp(J_theta) 乘以一个行列式之和，其中 J_theta 是类似 Jastrow 的前因子，Phi_theta^k 通过自注意力网络学习。
• 使用单流多头自注意力架构来建模电子-核相互作用输入特征，将自旋拼接进输入。
• 采用输入重标定以处理大分离几何，并通过指数包络来强制边界条件；应用残差连接和 tanh 非线性。
• 使用变分量子蒙特卡罗进行训练，采用 Rayleigh 商优化和 Metropolis-Hastings 采样；用 Kronecker-factored Approximate Curvature (KFAC) 进行优化。
• 将 Psiformer 与 FermiNet 和 FermiNet+SchNet 在小型和大型分子基准上进行比较，包括苯，以评估准确性提升。

实验结果

研究问题

• RQ1 基于自注意力的神经网络是否能够为从头量子化学中的多电子薛定谔方程提供具有竞争力或更优的 Ansatz？
• RQ2 相对于现有神经网络 Ansatz，Psiformer 是否提高了基态能量的准确性，尤其是对于更大的分子系统？
• RQ3 将自旋的引入、通过 Jastrow 处理 cusp，以及基于注意力的电子-核耦合如何影响不同系统尺寸下的学习和稳定性？
• RQ4 与先前方法相比，自注意力架构在拥有几十个电子的分子上的可扩展性在多大程度上能够实现？

主要发现

• Psiformer 在能量上显著优于 FermiNet 和 FermiNet+SchNet，尤其在更大系统上。
• 在小分子上，较大的 Psiformer 对于某些系统如臭氧，能量接近 CCSD(T)/CBS，差值约在 1 kcal/mol 之内。
• 对于 42–74 电子分子（苯到 CCl4），Psiformer 优于 FermiNet，甚至超过一些已发表的 DMC 能量，对苯和 CCl4 有显著提升。
• Psiformer 的能量随系统尺寸增大而改善，在大分子上相比传统基线，提升达到约 ~75 kcal/mol (100 mHa)。
• LayerNorm 及其他训练选择影响优化，Psiformer 通常在各种配置下保持更高的准确性。