来源论文: https://arxiv.org/abs/2605.13807v1 生成时间: May 14, 2026 07:29

0. 执行摘要

神经量子态（Neural-Network Quantum States, NQS）自 2017 年问世以来，已成为求解量子多体问题最强有力的变分工具之一。然而，在计算效率与扩展性（Scalability）方面，领域内一直存在一种“二元对立”：基于 Transformer 的架构虽然具备强大的并行化训练能力，但往往需要复杂的显式位置编码来捕捉物理系统的距离相关性；而循环神经网络（RNN）虽具有天然的归纳偏置（Inductive Bias），能够直观地编码物理距离，却受限于本质上的序列化（Sequential）计算，难以利用现代 GPU 的大规模并行能力，导致其在超大体系上的应用进展缓慢。

本研究通过引入“并行扫描循环神经网络量子态”（Parallel Scan Recurrent Neural Quantum States, PSR-NQS），彻底打破了这一瓶颈。作者证明，通过现代序列模型（如 SSM 和 LRU）中的并行扫描（Parallel Scan）算法，RNN 的前向和反向传播可以从 $O(N)$ 的序列深度降低至 $O(\log N)$，从而在保持 RNN 优良物理特性的同时，获得了足以媲美甚至超越 Transformer 的计算速度。在实际基准测试中，该方法成功处理了高达 $52 \times 52$（2704 个自旋）的二维海森堡模型，这在神经量子态领域是前所未有的规模。更重要的是，PSR-NQS 仅需约 1400 GPU 小时的适度资源即可完成全部训练，证明了该架构在量子化学和凝聚态物理研究中的巨大应用潜力。

1. 核心科学问题，理论基础，技术难点，方法细节

核心科学问题：循环结构的“并行化悖论”

在变分蒙特卡洛（VMC）框架中，变分波函数 $\Psi_\theta(\sigma)$ 的表达能力和采样效率决定了模拟的上限。RNN 因其自回归（Autoregressive）特性，可以实现无拒绝采样的直接配置生成，这在处理符号问题（Sign Problem）和复杂关联时具有天然优势。然而，传统的 RNN（如 LSTM, GRU）在处理长度为 $N$ 的序列（即 $N$ 个格点）时，必须依次计算隐藏状态 $h_t$，其计算深度为 $O(N)$。在格点数达到数千个时，这种序列依赖性成为计算时间的绝对杀手。

理论基础：状态空间模型（SSM）与线性循环单元（LRU）

本研究的理论支柱源于近年来深度学习领域对“线性递归”的重新发现。作者首先回顾了连续时间的状态空间模型（Equation 1 & 2）：

$$\frac{d}{dt}h(t) = \bar{A}h(t) + \bar{B}x(t)$$

$$y(t) = \text{Re}\{Ch(t)\} + Dx(t)$$

通过零阶保持（Zero-Order Hold）离散化后，递归公式变为（Equation 3）：

$$h_t = Ah_{t-1} + Bx_{t-1}$$

其中 $A = \exp(\Delta \bar{A})$。关键的观察点在于：如果 $A$ 与输入 $x$ 无关，且算子具有结合律，那么这个线性递归过程就可以被视为一个前缀和（Prefix Sum）问题。通过并行扫描算法（Blelloch, 1990），可以将计算复杂度在并行设备上压缩至 $O(\log N)$。

技术难点：从 1D 线性到 2D 复杂的泛化

  1. 线性化的代价：为了实现并行化，必须牺牲传统 RNN 中的非线性激活（如 $\tanh$）在递归路径中的应用。如何在线性递归的前提下保持波函数的非线性表达能力？作者采用了在递归层外部嵌套非线性门控单元（GLU）和 MLP 的策略（Sec II.D）。
  2. 二维拓扑的映射：二维晶格（$L \times L$）通常需要按照蛇形（Snake）顺序展开。在 PSR-NQS 中，作者设计了二维 minGRU 单元。对于站点 $(i, j)$，隐藏状态接收来自水平（$h_H$）和垂直（$h_V$）两个方向的因果预测信息（Equation 23, 24）。虽然 2D 递归无法实现全局单一并行扫描，但作者巧妙地实现了“逐行并行（Row-by-row Parallel Scan）”，将序列深度从 $O(L^2)$ 降低到 $O(L \log L)$。

方法细节：PSR-NQS 的架构设计

  • 1D LRU 架构：采用复数域的对角矩阵 $A$。通过显式参数化衰减（Decay）和相位（Phase），使得模型能自然捕捉量子态中的振荡行为和长程关联。每个块包含 LRU 层、GLU 映射和残差连接。
  • 2D minGRU 架构：去除了传统 GRU 中复杂的重置门，仅保留更新门 $z_t$，使其计算更加轻量。针对 2D 系统，作者使用了 $2 \times 2$ 的 Patch 策略，将物理自旋打包输入，进一步降低了序列长度并提升了局部关联捕捉能力。
  • 变分优化：利用 Adam 优化器最小化能量期望值 $E(\theta)$。通过随机梯度估算（Stochastic Estimates）进行参数更新。值得注意的是，为了稳定超大规模体系的训练，作者采用了“迭代再训练（Iterative Retraining）”技术：先在小体系（如 $6 \times 6$）上训练，再将参数平滑迁移至更大体系。

2. 关键 benchmark 体系，计算所得数据，性能数据

1. 一维横场伊辛模型 (TFIM)

在临界点 $h=1$ 处，TFIM 具有极强的量子纠缠，是检验 NQS 表达能力的试金石。作者在 $N=6$ 到 $N=256$ 的体系上进行了测试。

  • 精度表现：表 I 显示，在所有尺寸下，PSR-NQS 的相对误差（Relative Error）始终保持在 $10^{-4}$ 以下，多个尺寸甚至达到了 $10^{-6}$ 级别。
  • 热力学极限外推：通过对有限尺寸能级进行拟合，外推得到的热力学极限能量密度 $e_\infty = -1.2731999(8)$，与精确值 $-1.27323954$ 的相对误差仅为 $3.11 \times 10^{-5}$。

2. 二维平方格点海森堡模型 (Heisenberg Model)

这是本研究最引人注目的部分，主要针对具有开口边界条件（OBC）的反铁磁体系。

  • 10x10 与 16x16 对比：如表 II 所示，2D minGRU 在 10x10 体系上的能量优于常规的 PEPS 和 PixelCNN，接近最先进的 2D TRNN。在 16x16 体系上，2D minGRU 给出了当前文献中最低的变分能量（-0.643504），逼近了量子蒙特卡洛（QMC）的参考值。
  • 52x52 史诗级模拟：这是 NQS 首次在如此大规模的 2D 体系上展现出与 QMC 高度一致的结果。图 2 的有限尺寸缩放（Finite-size scaling）显示，随着格点数增加，PSR-NQS 的能量变化趋势与 QMC 完全吻合。在 52x52 体系下，相对误差仍能控制在 $10^{-4}$ 数量级，而此时总自旋数已达到 2704 个。

3. 计算性能数据 (Runtime Benchmarks)

  • 加速比：图 1 展示了并行扫描相对于传统序列评估的压倒性优势。在 1D 情况（N=1024）下，单步训练时间缩短了约一个数量级。在 2D 情况（L=32）下，加速效果同样显著。
  • 硬件效率：所有的 52x52 计算均在单块 NVIDIA A100 GPU 上完成。1D TFIM (N=256) 的单步训练时间约为 0.9 秒；2D Heisenberg (52x52) 虽然需要约 400 秒每步，但考虑到体系规模，这已经是极高的效率。总计算耗时约 1400 GPU 小时，远低于大规模 Transformer 动辄上万小时的开销。

代码实现架构

作者团队已将代码开源，主要基于 Python 与高性能计算库（推测为 JAX 或 PyTorch，鉴于并行扫描在 JAX 中的成熟支持）。

  • GitHub 地址：https://github.com/ParallelScan-RNNs/PSR-NQS
  • 核心模块：
    • models/lru.py: 实现 1D 线性循环单元，包含复数参数初始化和并行扫描算子。
    • models/mingru_2d.py: 2D minGRU 的单元定义，包含水平与垂直状态的融合逻辑。
    • vmc/sampler.py: 实现自回归采样，支持利用并行扫描进行高效的概率对数计算。

复现指南

  1. 环境准备：建议使用 CUDA 11.8+ 环境，安装最新版本的 JAX (或 PyTorch)。并行扫描需要底层关联算子的支持（如 jax.lax.associative_scan）。
  2. 小体系预训练：
    • 首先运行 python train.py --model lru --L 6 --task tfim 建立冷启动基准。
    • 观察能量收敛情况，确保相对误差在 $10^{-4}$ 以内。
  3. 体系扩展（Iterative Retraining）：
    • 使用 --load_path 加载上一级尺寸的模型参数。
    • 设置 --L 为新尺寸（如从 6 扩展到 8, 10…）。
    • 参考附录 Tab III 和 Tab V 调整学习率。通常随着尺寸增大，学习率需适当调小（如从 $5\times 10^{-4}$ 降至 $1\times 10^{-4}$）。
  4. 2D Patch 设置：对于二维体系，确保设置 patch_size=(2,2)，这对于降低显存占用和提升 2D 关联捕捉至关重要。

4. 关键引用文献，以及对这项工作局限性的评论

关键引用文献

  1. Carleo & Troyer (2017): NQS 的奠基之作，定义了变分神经波函数的基本框架。
  2. Hibat-Allah et al. (2020): 首次在大规模系统上应用 RNN 量子态，本研究的 2D 蛇形路径参考了该工作。
  3. Gu, Dao et al. (2021/2024): SSM 和 Mamba 的提出者，提供了并行化递归的理论支撑。
  4. Sandvik (2026/2017): QMC 基准数据的来源，为验证本工作精度提供了标准答案。

局限性评论

尽管 PSR-NQS 取得了显著进展，但在实际科研应用中仍存在以下局限：

  • 显存与带宽的权衡：并行扫描虽然降低了时间复杂度，但它要求同时存储整个序列的隐藏状态以便计算梯度。在极长序列下，GPU 显存会成为比计算速度更早遇到的瓶颈。
  • 二维因果律的局限性：虽然“逐行扫描”提升了速度，但蛇形排列本质上还是将 2D 拓扑强行一维化。这可能导致在处理高度受挫（Frustrated）的量子磁体时，模型难以捕捉跨行之间的精细纠缠，即便有垂直状态传递，其表达能力仍逊色于真正的 2D 算子。
  • 线性递归的表达上限：为了并行化而去掉递归内部的非线性，意味着波函数的“深度非线性”只能依靠堆叠更多的层。对于某些具有非平凡拓扑序（Topological Order）的相，LRU 的线性核可能需要极大的隐藏维度才能模拟其纠缠熵特性。
  • 费米子符号问题：目前的 benchmark 主要集中在自旋系统。在涉及格点费米子（Hubbard Model）时，由于交换反对称性导致的剧烈符号波动，PSR-NQS 是否依然能保持如此高的效率和精度尚待验证。

5. 其他必要补充：为什么这是量子化学的利好？

对量子化学模拟的启示

对于从事量子化学（Ab-initio Quantum Chemistry）的研究者来说，本工作的意义不仅在于自旋模型。在分子轨道的全配置交互（FCI）空间中，波函数同样可以映射为长序列的占据数表征。以往基于 Transformer 的方法在处理数以千计的轨道时，自注意力机制的 $O(N^2)$ 复杂度（或 $O(N)$ 的近似版）往往难以处理轨道间复杂的长程库仑相互作用。

PSR-NQS 提供了三个核心优势：

  1. 天然的轨道排序敏感性：化学体系中轨道的排列顺序对关联能计算至关重要。RNN 的顺序处理特性与化学家对“活性空间”的理解契合度更高。
  2. 轻量化部署：如文中所示，PSR-NQS 在单卡 GPU 上即可完成数千格点的模拟，这意味着量子化学实验室无需构建超大规模的计算集群即可开展复杂的关联电子计算。
  3. 距离归纳偏置：在实空间格点化学模拟中，PSR-NQS 能够自然地编码原子间的空间距离，无需像 Transformer 那样繁琐地调试位置编码参数。

总结与展望

PSR-NQS 的出现标志着循环架构在神经量子态领域的“文艺复兴”。通过借鉴现代 AI 领域的算法创新（并行扫描），RNN 不再是“慢”的代名词，而是成为了大规模量子模拟中“轻快、精准、高效”的代表。未来，随着该框架扩展到费米子系统和动力学模拟，我们有望看到其在催化剂设计、高温超导机理探索等前沿化学领域发挥关键作用。