[论文解读] Better, Faster Fermionic Neural Networks

James S. Spencer, David Pfau+2 more|arXiv (Cornell University)|Nov 13, 2020

Machine Learning in Materials Science参考文献 15被引用 7

一句话总结

该论文通过增加网络宽度和行列式数量，改进了费米子神经网络（FermiNet），实现了对氩等第二周期原子的化学精度。通过迁移到JAX并简化架构，训练时间减少了一个数量级，在包括双环丁烷和环丁二烯互变异构化在内的复杂体系上实现了最先进性能。

ABSTRACT

The Fermionic Neural Network (FermiNet) is a recently-developed neural network architecture that can be used as a wavefunction Ansatz for many-electron systems, and has already demonstrated high accuracy on small systems. Here we present several improvements to the FermiNet that allow us to set new records for speed and accuracy on challenging systems. We find that increasing the size of the network is sufficient to reach chemical accuracy on atoms as large as argon. Through a combination of implementing FermiNet in JAX and simplifying several parts of the network, we are able to reduce the number of GPU hours needed to train the FermiNet on large systems by an order of magnitude. This enables us to run the FermiNet on the challenging transition of bicyclobutane to butadiene and compare against the PauliNet on the automerization of cyclobutadiene, and we achieve results near the state of the art for both.

研究动机与目标

• 将FermiNet的精度扩展到第二周期元素，包括此前版本未达到化学精度的氩。
• 显著降低大规模多电子体系的训练时间和计算成本。
• 通过架构改进和优化训练协议，提升收敛性和能量精度。
• 在多参考体系（如环丁二烯互变异构化和双环丁烷过渡态）上，将FermiNet与PauliNet进行基准比较。
• 评估基于神经网络的量子化学波函数Ansatz中精度与训练速度之间的权衡。

提出的方法

• 增加FermiNet架构中单电子流的宽度和行列式数量，以提升表达能力和精度。
• 将FermiNet实现从TensorFlow迁移至JAX，实现更快计算和更优的自动微分。
• 通过移除完整协方差包络并采用各向同性权重初始化，简化网络结构，降低复杂度。
• 将MCMC采样步数从10次增加到50次，以改善平衡化并减少训练过程中的噪声。
• 将训练轮数从200,000次延长至300,000–400,000次，以确保在大体系上的收敛性。
• 采用变分蒙特卡洛方法，结合随机梯度下降和KFAC优化，最小化能量期望值。

实验结果

研究问题

• RQ1 仅通过增大FermiNet架构规模，能否在氩等第二周期原子上实现化学精度？
• RQ2 在不损失精度的前提下，架构简化与JAX迁移能在多大程度上减少训练时间？
• RQ3 与PauliNet相比，FermiNet在双环丁烷至丁二烯和环丁二烯互变异构化等多参考体系上的表现如何？
• RQ4 在具有挑战性的有机反应中，FermiNet计算基态与过渡态之间相对能量的相对精度如何？
• RQ5 为何FermiNet的训练速度慢于PauliNet？这一差距能否在不损失精度的前提下缩小？

主要发现

• 增加行列式数量和单电子流宽度后，FermiNet在除硫以外的所有第二周期元素上实现了化学精度（误差在1 kcal/mol以内）。
• 对于氩原子，FermiNet的能量达到-527.5405(1) $E_h$，与精确值在化学精度范围内一致。
• 训练时间从TensorFlow的11,520 GPU小时减少至JAX简化架构下的1,104 GPU小时，提升十倍。
• 在双环丁烷至丁二烯的转化过程中，FermiNet的相对能量差与最佳可用方法相比处于化学精度范围内。
• 在环丁二烯互变异构化过程中，FermiNet的能量比PauliNet低约70 m$E_h$，相对能量为10.3 ± 0.1 kcal/mol，与MR-CC结果相当。
• 尽管每次迭代的计算成本较高（约为PauliNet的5倍），FermiNet在多参考体系上仍展现出更优的绝对能量精度。