[论文解读] Scalable variational Monte Carlo with graph neural ansatz

Li Yang, Wen-Jun Hu+1 more|arXiv (Cornell University)|Nov 25, 2020

Model Reduction and Neural Networks参考文献 42被引用 11

一句话总结

本论文提出了一种基于图卷积网络（GCNs）的图神经网络变分 ansatz（GNA），以实现对量子多体系统的通用、几何无关的变分蒙特卡洛（VMC）模拟。通过利用参数共享和在128个TPU核心上的可扩展分布式训练，该方法成功实现了对2D凯文格、三角形及随机连接晶格（最大432个格点）的基态能量精确预测，并利用小系统预训练模型的迁移能力，实现了在大规模系统上的快速收敛。

ABSTRACT

Deep neural networks have been shown as a potentially powerful ansatz in variational Monte Carlo for solving quantum many-body problems. We propose two improvements in this direction. The first is graph neural ansatz (GNA), which is a variational wavefunction universal to arbitrary geometry. GNA results in accurate ground-state energies on 2D Kagome lattices, triangular lattices, and randomly connected graphs. Secondly, we design a distributed workflow on multiple accelerators to scale up the computation. We compute Kagome lattices with sizes up to 432 sites on 128 TPU cores. The parameter sharing nature of the GNA also leads to transferability across different system sizes and geometries.

研究动机与目标

• 开发一种适用于任意晶格几何结构（包括非周期性和不规则结构）的通用变分ansatz，用于量子多体系统。
• 克服以往神经网络ansatz（如CNN和RBM）因受限于晶格对称性或基于网格的表示方式而带来的局限性。
• 利用分布式加速器（如TPU）实现大规模量子系统的可扩展VMC模拟，突破内存和计算瓶颈。
• 展示训练好的GNA模型在不同系统尺寸和几何结构之间的迁移能力，从而缩短训练时间并提升收敛性。
• 提供一个统一且可扩展的框架，用于模拟强关联量子系统，包括未来扩展至复值波函数和费米子系统。

提出的方法

• 提出一种基于图卷积网络（GCNs）的图神经ansatz（GNA），用于编码多体波函数，其中每个格点作为节点，通过邻接图上的消息传递机制更新其可学习嵌入。
• 采用包含残差连接、ReLU非线性激活和层归一化的GCN架构，通过归一化邻接矩阵和跨层共享权重矩阵来更新节点特征。
• 通过单个全连接层将最终聚合的节点嵌入映射为波函数的标量对数振幅，实现变分能量的可微分优化。
• 在多个加速器（如TPU）上实现重要性采样梯度优化（ISGO）的分布式版本，支持大规模系统上的小批量采样与参数更新。
• 通过在邻接矩阵中连接对边，实现晶格构建中的周期性边界条件，支持环面几何结构的模拟。
• 通过在小系统（如12个格点）上预训练GNA，并在大系统（如432个格点）上微调，实现迁移学习，利用跨几何结构的共享参数。

实验结果

研究问题

• RQ1 基于图神经网络的ansatz是否能在无需架构修改的情况下，泛化至包括凯文格、三角形和随机连接图在内的多种晶格几何结构？
• RQ2 GNA中的参数共享在多大程度上实现了从小系统到大系统的迁移学习，从而提升训练收敛性和稳定性？
• RQ3 能否利用TPU等分布式加速器将GNA-VMC框架扩展至大规模系统（如432个格点），从而克服单设备的内存限制？
• RQ4 在不同相互作用强度（t/V）和晶格类型下，GNA-VMC方法在基态能量估计方面与精确对角化相比的准确性如何？
• RQ5 在不同几何结构（如三角形或随机图）上预训练的模型，是否仍能加速目标凯文格晶格的训练收敛？

主要发现

• GNA在2D凯文格和三角形晶格上的基态能量预测结果与精确对角化结果相差不超过0.1%，即使在较大的t/V比下依然保持高精度。
• 对于12×12×3的凯文格晶格（共432个格点），该方法成功利用128个TPU核心完成基态能量计算，该系统规模此前在单个加速器上难以实现。
• 在12个格点的小系统上预训练GNA（仅为目标系统大小的1/36），可将收敛时间缩短至仅5次采样迭代，而随机初始化的模型即使在40次迭代后仍无法收敛。
• GNA展现出显著的迁移能力：在三角形或随机连接图上预训练的模型，仍能显著加速目标凯文格晶格的训练，表明其在不同几何结构间具有稳健的泛化能力。
• 可扩展的ISGO实现通过分区系统并在多个加速器间聚合梯度，实现了高效的分布式训练，使大规模VMC模拟成为可能。
• 该方法在包括非规则和不规则图在内的多种几何结构中均保持高精度，证实了GNA作为变分ansatz的通用性。