介绍

本文将关注机器学习求解多体薛定谔方程的问题。具体的说，在物理上，我们关注电子-原子体系的基态能量和电子波函数的相关性质；在计算方法上，我们将介绍 AI 直接表示薛定谔方程中的电子波函数的思路，并结合变分方法（如 QMC）求解电子基态，获得相当高精度的结果。我们会介绍近期的主流相关方法（包括但不限于 FermiNet、PauliNet）的原理，并比较其优缺点；并且会上手实践 Deepmind 的 FermiNet 程序，在一些计算量较小的体系上求解基态能量。

问题背景

对于大部分关心的体系，通用的薛定谔方程可以写为 $$i\hslash \frac{\partial }{\partial t}\psi (\mathbf{R},\mathbf{r},t)=\hat{H}\psi (\mathbf{R},\mathbf{r},t)$$ 其中 $\mathbf{R}$ 表示原子中心的坐标，$\mathbf{r}$ 表示电子的坐标，$\hat{H}$ 为哈密顿量（Hamiltonian for everything） $$\hat{H}=-\sum _{i=1}^N\frac{{\nabla }_i^2}{2}-\sum _{I=1}^M\frac{{\nabla }_I^2}{2M_I}+\sum _{i=1}^N\sum _{j=i+1}^N\frac{1}{\left|{\mathbf{r}}_i-{\mathbf{r}}_j\right|}-\sum _{I=1}^M\sum _{i=1}^N\frac{Z_I}{\left|{\mathbf{r}}_i-{\mathbf{R}}_I\right|}+\sum _{I=1}^M\sum _{J=I+1}^M\frac{Z_I{Z}_J}{\left|{\mathbf{R}}_I-{\mathbf{R}}_J\right|}$$ 从左到右依次为：N 个电子动能项、M 个离子动能项、N 个电子间库伦势、电子-离子间库伦势、M 个离子间库伦势。

关于如何第一性的求解该薛定谔方程，此前的 Notebook 漫谈AI时代的科学计算与物理建模｜AI视角下的密度泛函建模与DeePKS 已经详细叙述了问题的背景、难点和几种常见的解决办法，例如 Hartree-Fock 等，这里不再赘述，本文的关注点在其中“AI 对电子波函数的直接变分求解”部分。我们这里先简要对比两种具有代表性的方法：密度泛函（DFT）和量子蒙特卡洛（QMC）。

DFT 将体系的能量直接写成由一个电子密度空间分布函数 $n(\mathbf{r})$ 所决定的泛函： $$E=E[n(\mathbf{r})]$$ 其可表示性由 Hohenberg-Kohn 定理所保证，然而该泛函并不能得到显式的表示，只能写为： $$E[n]=\sum _{\mathrm{i}=1}^N{\epsilon }_i-\frac{1}{2}\iint \frac{n(r)n\left(r^{\prime }\right)}{\left|r_1-r_2\right|}{d}^{3}r{d}^3{r}^{\prime }+E_{xc}[n]-\int n(r)V_{xc}(r)d^{3}r$$ 其中第一项为每个 Kohn-Sham 轨道的能量，而第三项 $E_{xc}[n]$ 则为无显式表示的交换关联项，尽管人们不断尝试优化这一项的表达，但这个非第一性给出的项大大限制了 DFT 的精度。

完全第一性的尝试是直接变分优化波函数的表示，得到基态能量 $$E=\underset{\psi }{\min }\frac{⟨\psi |H|\psi ⟩}{⟨\psi |\psi ⟩}=\underset{\psi }{\min }\frac{\int {\psi }^{\ast }(\mathbf{r})\hat{H}\psi (\mathbf{r})d\mathbf{r}}{\int {\psi }^{\ast }(\mathbf{r})\Psi (\mathbf{r})d\mathbf{r}}$$ 再使用各种 Monte Carlo 方法优化波函数， QMC 中的各种方法（如 Variational MC, Diffusion MC）由此而来。我们这里提供两种范式的对比，仅供参考：

AI 方法怎样引入 QMC

QMC 中最常见的变分蒙特卡洛方法（VMC）的关键点在于如何人为选取一组表示 3N 维电子多体波函数的基，以选取合适的试探波函数（这也是其主要误差来源）。当然，这些基的表示首先应该满足泡利不相容原理的约束，即电子波函数应当是交换反对称的： $$\psi \left(r_1\dots r_i\dots r_j\dots r_N\right)=-\psi \left(r_1\dots r_j\dots r_i\dots r_N\right)$$ 最常规的 QMC 的选取是 Slater-Jastrow ansatz， $$\psi \left(r_1,r_2,\dots ,r_N\right)=e^{J(\mathbf{q})}D\left[{\varphi }_i\left(q_j\right)\right]$$ 其中，第一组人为选取的参数是坐标变换 $\mathbf{q}(\mathbf{r})$，或通常称为 backflow transformation（下面只是一种人为选取的方式）： $$\begin{array}{rl}& q_i=r_i+\sum _{j\ne i}^N\eta \left(r_{ij}\right)\left(r_i-r_j\right)\\ & \eta \left(r_{ij}\right)=f\left(r_{ij};L_{\eta }\right)\sum _{k=0}^{N_{\eta }}{c}_k{r}_{ij}^k\end{array}$$ 这里 $f\left(r_{ij};L_{\eta }\right)={\left(\frac{L_{\eta }-r_{ij}}{L_{\eta }}\right)}^C\Theta \left(L_{\eta }-r_{ij}\right)$ 是一个截断函数，多项式系数 $c_k$ 也是 QMC 中待优化的参数。由此可以看出 backflow 变换的目的：使用变换后的坐标可以用于描述电子间在截断距离内相互作用引起的实时关联（截断距离之外回归到原有形式），这是 Slater 行列式所不包括的： $$D\left[{\varphi }_i\left(q_j\right)\right]=\det \left|\begin{array}{cccc}{\varphi }_1\left(q_1\right)& {\varphi }_2\left(q_1\right)& \cdots & {\varphi }_N\left(q_1\right)\\ {\varphi }_1\left(q_2\right)& {\varphi }_2\left(q_2\right)& \cdots & {\varphi }_N\left(q_2\right)\\ \cdots & \cdots & \cdots & \cdots \\ {\varphi }_1\left(q_N\right)& {\varphi }_2\left(q_N\right)& \cdots & {\varphi }_N\left(q_N\right)\end{array}\right|$$ 其中 ${\varphi }_i$ 表示单电子的轨道波函数。 在 backflow 变换的基础上，再乘以一个 Jastrow factor $J(\mathbf{q})$ 的指数项就是选取的试探波函数，Jastrow factor 的形式是 $$J(\mathbf{q})=\sum _i^N\chi \left(q_i\right)-\frac{1}{2}\sum _i^N\sum _{j\ne i}^{N}u\left(q_i,q_j\right)$$ 其中，二体关联项一般写为 $$u\left(q_i,q_j\right)=\frac{A}{q_{ij}}\left(1-e^{-q_{ij}/F}\right)$$ 一体项 $\chi (q_i)$ 则用于描述电子-原子核关联，没有固定形式，通常写成一组基的线性叠加形式（例如在周期条件下写成平面波展开），其系数也是待 QMC 优化的参数，接下来只需要一直优化 $E=⟨\psi |H|\psi ⟩/⟨\psi |\psi ⟩$ 一直取到最小值即可。

可以看出，QMC 的关键在于选取合适的试探波函数的 ansatz，例如上面人为选取的 Jastrow factor 和 backflow transformation，至少比单独只用 Slater determinant 做优化要精确很多，如果需要无限度提高精度，就选取多个 Slater-Jastrow factor 组合起来优化 $$\psi =e^{J_1\left({\mathbf{q}}_1\right)}\times D\left[\mathbf{\varphi }({\mathbf{q}}_1)\right]+e^{J_2\left({\mathbf{q}}_2\right)}\times D\left[\mathbf{\varphi }({\mathbf{q}}_2)\right]+\cdots$$ 但完全精确的全组态量子蒙特卡洛（FCIQMC）是极度浪费的，需要 $O(e^N)$ 的计算复杂度，对于多体问题是不现实的。因此如何人为选取合适的试探波函数进行优化，是决定 VMC 精度和效率的关键。用神经网络尝试表示波函数，此前有三种思路：

1. 直接用神经网络表示整个多体波函数

一个较早的尝试是用深度神经网络（DNN）直接表示多体波函数（可见韩劼群、张林峰、鄂维南的DeepWF方法（2018年）），而不借助任何先验的假定（例如轨道波函数的 Slater 行列式），把波函数分解成 $$\psi (\mathbf{r})=S(\mathbf{r})A^{\uparrow }\left({\mathbf{r}}^{\uparrow }\right)A^{\downarrow }\left({\mathbf{r}}^{\downarrow }\right)$$ 其中 $\uparrow ,\downarrow$ 表示电子的自旋，泡利不相容原理要求 $S(\mathbf{r})$ 对称（这一项与 Jastrow factor 起的作用类似），对应的 $A^{\uparrow }\left({\mathbf{r}}^{\uparrow }\right),A^{\downarrow }\left({\mathbf{r}}^{\downarrow }\right)$ 反对称。这样每个函数都可以用 DNN 来尝试表示，例如反对称函数可以这样构造： $A(\mathbf{r})=$ ${\prod }_{i\ne j}a\left({\mathbf{r}}_{\mathbf{i}},{\mathbf{r}}_{\mathbf{j}}\right)$，其中 $a\left({\mathbf{r}}_{\mathbf{i}},{\mathbf{r}}_{\mathbf{j}}\right)=\mathrm{Net}\left({\mathbf{r}}_{\mathbf{i}},{\mathbf{r}}_j,r_{ij}\right)-\mathrm{Net}\left({\mathbf{r}}_j,{\mathbf{r}}_{\mathbf{i}},r_{ij}\right)$。这样我们就可以不依赖人为选取的 ansatz 训练试探波函数求解电子基态，复杂度也可以降到 $O(N^4)$。

2. 仅用神经网络表示 Jastrow factor 和 backflow tranformation

另一个“完全相反”的方法是保留原有的 Jastrow-Slater ansatz 的整体结构，只是使用 DNN 来表示本来人为选取的 Jastrow factor $J(\mathbf{q})$ 和 backflow tranformation $f(\mathbf{r})$。Frank Noé 组的 PauliNet 工作采用了这种思路，即待优化的试探波函数为 $$\begin{array}{c}{\psi }_{\mathbf{\theta }}(\mathbf{r})={\mathrm{e}}^{\gamma (\mathbf{r})+J_{\theta }(\mathbf{r})}\sum _p{c}_p\det \left[{\tilde{\mathbf{\varphi }}}_{\mathbf{\theta },{\mu }_{p}i}^{\uparrow }(\mathbf{r})\right]\det \left[{\tilde{\mathbf{\varphi }}}_{\mathbf{\theta },{\mu }_{p}i}^{\downarrow }(\mathbf{r})\right]\\ {\tilde{\mathbf{\varphi }}}_{\mu i}(\mathbf{r})={\varphi }_{\mu }\left({\mathbf{r}}_i\right)f_{\mathbf{\theta },\mu i}(\mathbf{r})\end{array}$$ PauliNet 保留了行列式中的反对称性和分子轨道波函数 ${\varphi }_{\mu }({\mathbf{r}}_i)$，这些不再需要构建 DNN 去满足，而只有 $J_{\theta }(\mathbf{r})$ 和 $f_{\mathbf{\theta },\mu i}(\mathbf{r})$ 需要神经网络表示。这种方式在保留较高精确度的情况下，只需要训练两个 DNN，节省了大量计算资源，同等花销下精度可以显著高于 DeepWF。

3. 用神经网络代替单电子轨道波函数

Google DeepMind 开发的 FermiNet 选择了第三条路径，即将原先的 Slater 行列式 $$\begin{array}{r}\det \left[{\varphi }_i\left({\mathbf{x}}_j\right)\right]=\left|\begin{array}{ccc}{\varphi }_1\left({\mathbf{x}}_1\right)& \dots & {\varphi }_1\left({\mathbf{x}}_n\right)\\ ⋮& & ⋮\\ {\varphi }_n\left({\mathbf{x}}_1\right)& \dots & {\varphi }_n\left({\mathbf{x}}_n\right)\end{array}\right|\end{array}$$ 中的单电子轨道波函数 ${\varphi }_i({\mathbf{x}}_j)$ 全部写成 general 的形式：${\varphi }_i\left({\mathbf{x}}_j;\left\{{\mathbf{x}}_{/j}\right\}\right)$，这里 $\left\{{\mathbf{x}}_{/j}\right\}$ 表示除了 ${\mathbf{x}}_j$ 之外的其他电子坐标。原先 N 个电子轨道波函数改用 N 个神经网络代替，FermiNet 的 ansatz 为 $$\begin{array}{rl}\psi \left({\mathbf{r}}_1^{\uparrow },\dots ,{\mathbf{r}}_{n\downarrow }^{\downarrow }\right)=& \sum _k{\omega }_k(\det \left[{\varphi }_i^{k\uparrow }\left({\mathbf{r}}_j^{\uparrow };\left\{{\mathbf{r}}_{/j}^{\uparrow }\right\};\left\{{\mathbf{r}}^{\downarrow }\right\}\right)\right]\\ & \times \det \left[{\varphi }_i^{k\downarrow }\left({\mathbf{r}}_j^{\downarrow };\left\{{\mathbf{r}}_j^{\downarrow }\right\}\right);\left\{{\mathbf{r}}^{\uparrow }\right\};\right]).\end{array}$$ 即不需要再乘上 Jastrow factor 或作用 backflow 变换（已经包含在单电子波函数中），而是写为 $$\begin{array}{rl}{\varphi }_i^{k\alpha }\left({\mathbf{r}}_j^{\alpha };\left\{{\mathbf{r}}_{/j}^{\alpha }\right\};\left\{{\mathbf{r}}^{\bar{\alpha }}\right\}\right)=& \left({\mathbf{w}}_i^{k\alpha }\cdot {\mathbf{h}}_j^{L\alpha }+g_i^{k\alpha }\right)\\ & \times \sum _m{\pi }_{im}^{k\alpha }\exp \left(-\left|{\mathbf{\Sigma }}_{im}^{k\alpha }\left({\mathbf{r}}_j^{\alpha }-{\mathbf{R}}_m\right)\right|\right),\end{array}$$ 其中 ${\mathbf{h}}_j^{L\alpha }$ 就表示神经网络第 L 层的输出。网络的结构是 $$\begin{array}{c}{\mathbf{h}}_i^{\ell +1\alpha }\leftarrow \tanh \left(\mathrm{matmul}\left({\mathbf{V}}^l,{\mathbf{f}}_i^{\ell \alpha }\right)+{\mathbf{b}}^l\right)+{\mathbf{h}}_i^{\ell \alpha }\\ {\mathbf{f}}_i^{\ell \alpha }=\left({\mathbf{h}}_i^{\ell \alpha },\frac{1}{n^{\uparrow }}\sum _{j=1}^{n^{\uparrow }}{\mathbf{h}}_j^{\ell \uparrow },\frac{1}{n^{\downarrow }}\sum _{j=1}^{n^{\downarrow }}{\mathbf{h}}_j^{\ell \downarrow },\frac{1}{n^{\uparrow }}\sum _{j=1}^{n^{\uparrow }}{\mathbf{h}}_{ij}^{\ell \alpha \uparrow },\frac{1}{n^{\downarrow }}\sum _{j=1}^{n^{\downarrow }}{\mathbf{h}}_{ij}^{\ell \alpha \downarrow }\right)\end{array}$$ 在多层神经网络下，FermiNet 一般只需要一个或者有限几个行列式（即 N 个神经网络，或者 N 的整数倍个），即可达到很高的精度。

FermiNet 是目前最常用的 VMC 神经网络方法之一，并且还有后续的相关工作对 FermiNet 的效率进行了提升，例如字节的 Forward Laplacian 工作，简化了动能部分拉普拉斯算符的计算（即计算 $⟨\psi |{\nabla }^2|\psi ⟩$ 的步骤），原先需要正反向传播计算 Hessian 矩阵以得到动能，改使用一次前向传播直接计算，将 FermiNet 的效率提高了约一个数量级，并且也被 DeepMind 集成在了最近更新的 FermiNet 项目中。也有若干工作拓展了 FermiNet 的应用场景，下面的海报中的信息可以作为一个参考。

物理模拟实践

DeepMind 的 FermiNet 项目长期维护，适合用于 Neural Network + Variational Monte Carlo 的代码实践。我们将首先配置 FermiNet 的环境：

需要模拟某个体系，我们可以在 ferminet/configs 中使用相应的 config 文件，例如对于氢分子，我们在该路径下已经创建了相应的 h2_config.py 文件，我们把键长设为稳定基态时的位型：1.4 个玻尔半径，批量大小设为 256，FermiNet会先使用 Hartree-Fock 方法预训练 100 次：

from ferminet import base_config
from ferminet.utils import system

# Settings in a config files are loaded by executing the the get_config
# function.
def get_config():
  # Get default options.
  cfg = base_config.default()
  # Set up molecule
  cfg.system.electrons = (1,1)
  cfg.system.molecule = [system.Atom('H', (0, 0, -0.7)), system.Atom('H', (0, 0, 0.7))]

  # Set training hyperparameters
  cfg.batch_size = 256
  cfg.pretrain.iterations = 100

  return cfg

接下来使用如下命令训练，作为例子，我们已在 ferminet/base_config.py 中将训练次数由默认的 1000000 次调小到了 10000 次，训练需要约 3 小时：

运行完成后，在路径下生成标有完成时间的文件夹，包含所有所需输出。我们也把目前的输出结果放在了数据集 QMC2 中，可供未运行的读者直接参考。

根据路径名称打开最新的输出，例如这里输出了训练过程的断点 checkpoint 的神经网络参数，以及训练过程的能量等变化数据 train_stats.csv

我们直接读取FermiNet的training数据，并对蒙卡得到的随机性能量波动做一个滚动平均，观察训练过程中的能量变化。对于氢分子的电子基态能量，benchmark得到的精确值为 1.17448 a.u.，我们观察 FermiNet 在 10000 步之内的训练的结果：

可以看出能量在约4000步前收敛，结果是相当精确的。我们计算一下后面4000-10000步的值：

误差在0.01%内，FermiNet提供了正确的结果。

FermiNet 目前是一个开发中的预版本，在 notebook 运行的范围内，我们能最直接做的操作就是观察能量的训练数据是否收敛；如果需要更多可观测量的数据，需要在源代码中做修改，使用神经网络表示的波函数来手动计算。作为一个方法的对比和检验，我们先展示了使用 FermiNet 表示的波函数在基态能量上得到了精确的值。FermiNet 最近的更新还支持了激发态计算，但无论是基态还是激发态都只提供了少数的可观测量，做更多电子结构的分析需要在理解源代码的基础上自行编程。