神经网络波函数的非随机优化策略

近日，清华大学胡憾石课题组联合字节跳动ByteDance Research团队提出实现了分子体系神经网络波函数的一种非随机优化策略，显著提高了优化过程的效率、精度与稳定性，为神经网络在电子结构计算领域的应用提供了一种有前景的新算法。

基于神经网络方法来求解分子薛定谔方程的核心思想是把神经网络模型作为通用的函数模拟器，通过优化神经网络参数来拟合分子体系的目标电子波函数。借助神经网络处理分子高维波函数有望突破电子结构计算中的“指数墙”困难。以量子化学中高精准FCI波函数为例，它由大量电子组态线性组合而成，其涉及的组态数目随体系大小呈指数增长，实际过程中计算量极大。如果使用神经网络波函数算法，则能够直接拟合出任意电子组态与其对应组态系数之间的映射函数，从而有望加速获得组态系数。

图1.基于受限玻尔兹曼机模型的神经网络电子波函数，输入值是电子组态的轨道占据数，输出值是电子组态的系数。（图片来源于原文）

为了让神经网络波函数逼近真实的FCI波函数，通常是以变分原理为依据，向波函数能量降低的方向优化神经网络参数。这离不开蒙特卡洛过程，即通过随机采样得到一部分对波函数贡献较大的电子组态来近似计算波函数能量及其梯度。然而该基于随机过程的优化策略仍面临计算效率低、产生的电子组态少、优化过程不稳定等挑战。该文章一作李相同学在研究中发现，神经网络波函数的能量计算过程与选择组态相互作用（SCI）方法中的迭代组态挑选过程具有相似的结构特征。基于SCI方法的组态挑选策略，在计算神经网络波函数能量的同时便挑选出了全部波函数值绝对值较大（即对波函数贡献较大）的电子组态，从而避免了原来的蒙特卡洛过程。这一非随机的组态挑选过程无需额外计算开销，不仅提升了计算效率，还提高了用于能量与梯度计算的电子组态样本的质量。

图2.非随机组态挑选与波函数能量计算同时进行。（图片来源于原文）

相比依赖蒙特卡洛过程的随机优化策略，利用该组态挑选方案的非随机优化策略在19个已有文献报导的分子体系中使神经网络波函数达到了与文献结果相近或更优的能量精度。此外，以不同键长的N₂分子为例，非随机的优化策略展现出了更快的收敛速度、更稳定的优化过程以及更高的能量精度。由于非随机的优化策略避免了耗时的蒙特卡洛过程，计算结果还展现出了显著的效率提升。

图3.不同键长N₂分子神经网络波函数基于随机优化策略（MC）与非随机优化策略（SC）的优化过程。（图片来源于原文）

图4.不同键长N₂分子神经网络波函数基于随机优化策略（MC）与非随机优化策略（SC）优化所需的CPU耗时。（图片来源于原文）

这一成果近期发表在学术期刊Journal of Chemical Theory and Computation 上，文章的第一作者是清华大学化学系博士研究生李相，清华大学胡憾石副教授为通讯作者。

原文（扫描或长按二维码，识别后直达原文页面，或点此查看原文）：

A Nonstochastic Optimization Algorithm for Neural-Network Quantum States

Xiang Li, Jia-Cheng Huang, Guang-Ze Zhang, Hao-En Li, Chang-Su Cao, Dingshun Lv, and Han-Shi Hu*

J. Chem. Theory Comput., 2023, 19, 8156–8165, DOI: 10.1021/acs.jctc.3c00831

通讯作者简介

胡憾石，清华大学化学系副教授，博导，国家优秀青年基金获得者。2006、2011年先后获清华大学学士、博士学位。2011-2017年分别在清华物理系、美国西北太平洋国家实验室和英国曼彻斯特大学从事博士后研究工作。2017获得国家级青年人才支持，并于同年加入清华大学化学系。主持国家自然科学青年、面上和优青项目基金，担任科技部国家重点研发计划课题负责人等。

研究方向：主要研究领域为f-区元素化合物、团簇和固体材料的相对论量子化学研究和从头算波函数量子化学算法开发等。已在Nat. Chem., Nat. Commun., J. Am. Chem. Soc., Angew. Chem. Int. Ed. 等国际重要学术期刊发表论文80余篇。

https://www.x-mol.com/groups/HanshiHU