[论文解读] A simple linear algebra identity to optimize Large-Scale Neural Network Quantum States

Riccardo Rende, Luciano Loris Viteritti+3 more|arXiv (Cornell University)|Oct 9, 2023

Quantum many-body systems参考文献 45被引用 4

一句话总结

本文提出了一种线性代数恒等式的创新应用，以在大规模神经网络量子态（NNQS）中实现精确的随机重构（SR）优化，通过将传统的 $P \times P$ 矩阵求逆瓶颈转化为 $M \times M$ 问题，从而克服该瓶颈。该方法在使用 267,720 个参数和仅 6,000 个蒙特卡洛样本的情况下，实现了 $J_1$-$J_2$ 赫巴德模型在 $J_2/J_1 = 0.5$ 时的最先进基态能量结果，该模型为 $10 \times 10$ 网格。

ABSTRACT

Neural-network architectures have been increasingly used to represent quantum many-body wave functions. These networks require a large number of variational parameters and are challenging to optimize using traditional methods, as gradient descent. Stochastic Reconfiguration (SR) has been effective with a limited number of parameters, but becomes impractical beyond a few thousand parameters. Here, we leverage a simple linear algebra identity to show that SR can be employed even in the deep learning scenario. We demonstrate the effectiveness of our method by optimizing a Deep Transformer architecture with $3 imes 10^5$ parameters, achieving state-of-the-art ground-state energy in the $J_1$-$J_2$ Heisenberg model at $J_2/J_1=0.5$ on the $10 imes10$ square lattice, a challenging benchmark in highly-frustrated magnetism. This work marks a significant step forward in the scalability and efficiency of SR for Neural-Network Quantum States, making them a promising method to investigate unknown quantum phases of matter, where other methods struggle.

研究动机与目标

• 克服在参数数量 $P$ 超过蒙特卡洛样本数量 $M$ 的大规模神经网络量子态（NNQS）中，随机重构（SR）计算不可行的问题。
• 为具有数十万参数的深度架构（如深度变换器）实现精确的 SR 优化，这些架构在标准 SR 下由于 $P \times P$ 矩阵求逆而难以处理。
• 为强阻挫的量子自旋系统（如 $J_1$-$J_2$ 赫巴德模型在 $J_2/J_1 = 0.5$ 时）提供高精度的基态能量估计，其他方法在此类系统中表现不佳。
• 证明在 NNQS 中高参数数量并不需要大量样本数 $M$ 即可实现精确的 SR，挑战了传统认为 $M \gg P$ 是必要条件的观点。
• 通过消除对波函数符号模式的假设，将 SR 的适用范围扩展到复杂且无符号结构的量子系统。

提出的方法

• 利用一个简单的线性代数恒等式，将标准 SR 中的 $P \times P$ 矩阵求逆转换为等价的 $M \times M$ 矩阵求逆，其中 $M$ 为蒙特卡洛样本数。
• 将该恒等式应用于 SR 框架中的费雪信息矩阵，实现无需近似的逆矩阵精确计算。
• 使用转换后的 $M \times M$ 矩阵计算神经网络参数的 SR 更新方向，保持 SR 的几何优化优势。
• 在变分蒙特卡洛框架中实现该方法，使用 JAX 进行可微分神经网络计算，并使用 mpi4jax 实现分布式并行化。
• 将修改后的 SR 更新与实值深度视觉变换器（ViT）架构结合，后接复值全连接输出层，以表示复值量子态。
• 在优化过程中逐步恢复物理对称性（平移、旋转、反射、自旋宇称），观察到能量持续下降。

实验结果

研究问题

• RQ1 能否在 $P \gg M$ 的大规模 NNQS 中精确应用随机重构（SR），且无需近似？
• RQ2 所提出的线性代数恒等式是否能够实现对具有数十万参数的深度神经网络在量子多体系统中的高效且稳定优化？
• RQ3 该方法能否在高度阻挫的量子自旋模型（如 $J_1$-$J_2$ 赫巴德模型在 $J_2/J_1 = 0.5$ 时）中实现最先进基态能量估计？
• RQ4 即使 $P \approx 267,720$，是否也能仅用 $M \approx 6,000$ 个样本实现 SR 的高精度？
• RQ5 在无符号先验的条件下，物理对称性的逐步恢复如何影响变分能量的优化过程？

主要发现

• 所提出的方法在仅使用 6,000 个蒙特卡洛样本的情况下，实现了对最多 267,720 个参数的 NNQS 的精确随机重构，成功绕过了 $P \times P$ 求逆瓶颈。
• 该方法在 $10 \times 10$ 网格的 $J_1$-$J_2$ 赫巴德模型在 $J_2/J_1 = 0.5$ 时，实现了每胞腔基态能量为 $-0.49575(3)$ 的新最先进结果。
• 在深度 ViT 架构中，逐步恢复物理对称性（平移、旋转、反射、自旋宇称）时，能量持续下降，证实了该方法的稳定性和物理一致性。
• 优化过程对初始种子不敏感，多次运行结果一致，表明其具有高度可靠性。
• 该方法在不假设基态符号结构的情况下表现良好，使其适用于广泛类别的量子系统，其中此类假设不可用。
• 该方法挑战了传统认为 $M \gg P$ 是实现精确 SR 所必需的观点，证明即使 $P \approx 267,720$，$M \approx 6,000$ 也已足够。