深度学习哈密顿蒙特卡罗

摘要

我们用一堆神经网络层来推广哈密顿蒙特卡罗算法，并评估其从二维晶格规范理论中的不同拓扑进行采样的能力。 我们证明我们的模型能够成功地混合不同拓扑的模式，从而显着降低生成独立规范场配置所需的计算成本。 我们的实现可在 https://github.com/saforem2/l2hmc-qcd 获取。

1简介

在Lévy等人（2017）中，作者提出了学习哈密顿蒙特卡罗（L2HMC）算法，并证明了其在各种目标分布上优于传统哈密顿蒙特卡罗（HMC）的能力。 他们表明，经过训练的 L2HMC 采样器能够通过探索传统 HMC 无法访问的相空间区域来在模式之间进行混合。 在本文中，我们提出了一种通用采样器，使用深度神经网络进行自训练，以提出新的马尔可夫链状态，该状态在 Metropolis-Hastings （Hastings，1970） 算法的帮助下变得精确。 我们的深度学习哈密顿蒙特卡罗 (DLHMC) 算法的目标是模拟晶格规范理论，其中最先进的模拟具有数十亿个自由度，并在具有数千个节点的超级计算机上运行，​​但在接近连续体时会出现指数级减速物理学（Schaefer等人，2009；Cossu等人，2017）。

2 主要贡献

我们提出 DLHMC 算法作为广义 HMC 算法，使用一堆不同的神经网络层代替连续的蛙跳步骤。 每个不同的神经网络层都带有 $N_{\mathrm{LF}}$ 蛙跳层的离散索引 $k=0,1,\mathrm{\dots },N_{\mathrm{LF}}-1$，执行增强的蛙跳方程（方程 1 和方程 2)，它们由不同的神经网络参数化。

我们将所提出的方法应用于在具有周期性边界条件的方晶格上定义的$2$维$U(1)$晶格规范理论。 我们专门为我们的物理系统设计了一个损失函数，鼓励拓扑特性的隧道效应。 我们发现，通过拓扑电荷的积分自相关时间来衡量，使用 DLHMC 的计算成本显着降低。 我们将各种轨迹长度和耦合常数的结果与传统 HMC 进行比较，结果表明我们训练的模型始终优于传统 HMC（见图 3)。 在较小的耦合常数上训练的模型，当应用于较大的耦合常数时，只需要最少的重新训练。 我们发现，经过训练的 DLHMC 采样器的隧道拓扑扇区效率可以通过物理系统通过深度神经网络时的行为来解释，见图 5(a) 和图 5(b)。

3相关工作

最近，人们对应用生成机器学习技术来构建更智能、更高效的科学模拟越来越感兴趣。 随着 RealNVP (Dinh 等人, 2017) 架构的发展，已经提出了许多旨在利用其可逆结构的技术。 特别是，格子规理论和格子 QCD 的模拟将从这种新方法中获益匪浅，这一点可以从网络架构 （Dinh 等人，2017 年；Favoni 等人，2020 年；Toth 等人，2020 年）、生成模型 （Albergo 等人，2019 年；2021 年；Medvidovic 等人，2020 年；Boyda 等人，2020 年；Kanwar 等人，2020 年；Wehenkel & Louppe, 2020 年；Tompe, 2020 年） 等在这方面迅速增长的研究成果中得到证明、2019；2021；Medvidovic 等人，2020；Boyda 等人，2020；Kanwar 等人，2020；Wehenkel & Louppe，2020；Tomiya & Nagai，2021），以及新型 MCMC 方法 （Pasarica & Gelman，2010；Tanaka & Tomiya，2017；Hoffman et al、2019; Neklyudov & Welling, 2020; Neklyudov et al.

4DLHMC算法

我们回顾了通用哈密顿蒙特卡罗 (HMC) 算法，并在 A.1 节中设置了一些相关符号。 在模拟系统 $x$ 时，我们使用理论上给定的概率密度 $p(x)$ 和可能难以实现的归一化，用目标分布 $p(\xi )=p(x,v,d)=p(x)p(v)p(d)$ 将马尔科夫链状态增强为 $\xi =(x,v,d)$ 。 HMC 算法中通常使用的共轭动量 $v$ 具有已知且易于采样的分布。 二进制方向变量，如 L2HMC (Lévy 等人, 2017)、$d\in 𝒰(+,-)$ 中介绍的，表示我们更新的“方向”（向前/向后）。

DLHMC 使用单独的神经网络进一步概括了 L2HMC 中的蛙跳步骤，这些神经网络由权重 $\theta$ 共同参数化。 DLHMC 中的每个蛙跳步骤都是深度神经网络的一层，它将第 $k$ 层的输入 ${\xi }_k=(x_k,v_k,d_k)$ 转换为输出 $(x_k^{\prime \prime },v_k^{\prime \prime },d_k)$，其中 $d_k$ 保持不变，选择向前或向后的蛙跳层。 随后，第 $k+1$ 层使用 ${\xi }_{k+1}=(x_k^{\prime \prime },v_k^{\prime \prime },d_k)$ 作为其输入。 为了简单起见，我们考虑向前的$d=+1$方向，并引入符号：

$v_k^{\prime }$ $\equiv {\mathrm{\Gamma }}_k^{+}(v_k;{\zeta }_{v_k})=v_k\odot \exp \left(\frac{{\epsilon }_v^k}{2}{s}_v^k({\zeta }_{v_k})\right)-\frac{{\epsilon }_v^k}{2}\left[{\partial }_{x}S(x_k)\odot \exp \left({\epsilon }_v^k{q}_v^k({\zeta }_{v_k})\right)+t_v^k({\zeta }_{v_k})\right],$ (1)

$x_k^{\prime }$ $\equiv {\mathrm{\Lambda }}^{+}(x_k;{\zeta }_{x_k})=x_k\odot \exp \left({\epsilon }_x^k{s}_x^k({\zeta }_{x_k})\right)+{\epsilon }_x^k\left[v_k^{\prime }\odot \exp \left({\epsilon }_x^k{q}_x^k({\zeta }_{x_k})\right)+t_x^k({\zeta }_{x_k})\right]$ (2)

其中：(1.) ${\zeta }_{v_k}=(x_k,{\partial }_{x}S(x_k))$和${\zeta }_{x_k}=({\overline{m}}^k\odot x_k,v_k)$,${\zeta }_{x_k^{\prime }}=(m^k\odot x_k^{\prime },v_k^{\prime })$表示（分别为$x$,$v$)网络的输入，其中 $\odot$ 表示逐元素乘法； (2.) 我们通过${\mathrm{\Gamma }}_k^{\pm },{\mathrm{\Lambda }}_k^{\pm }$上的上标$\pm$来指示更新的方向； (3.) $k\in \{0,1,\mathrm{\dots },N_{\mathrm{LF}}-1\}$ 表示当前沿着轨迹的蛙跳步。 我们在图 1 中展示了用于公式 1 和公式 2 更新的 $k$ 第三跃迁层的网络结构。

图1： 公式 1 和公式 2 中使用的网络架构示意图。

使用这种表示法，我们可以编写一个完整的蛙跳更新（向前 $d=+1$ 方向） 作为：

1. 半步动量更新：$v_k^{\prime }={\mathrm{\Gamma }}_k^{+}(v_k;{\zeta }_{v_k})$

2. 全步半位置更新： $x_k^{\prime }={\overline{m}}^k\odot x_k+m^k\odot {\mathrm{\Lambda }}_k^{+}(x_k;{\zeta }_{x_k})$

3. 全步半位置更新：$x_k^{\prime \prime }={\overline{m}}^k\odot {\mathrm{\Lambda }}_k^{+}(x_k^{\prime };{\zeta }_{x_k^{\prime }})+m^k\odot x_k^{\prime }$

4. 半步动量更新：$v_k^{\prime \prime }={\mathrm{\Gamma }}_k^{+}(v_k^{\prime };{\zeta }_{v_k^{\prime }})$

总的来说，我们将这一系列更新称为单个leapfrog层，它执行单个更新$\xi \to {\xi }^{\prime }$。 请注意，为了保持更新的可逆性，我们引入了一个二进制掩码 ${\overline{m}}^k=\mathrm{𝟙}-m^k$，将 $x$ 更新拆分为两个子更新，该掩码会按顺序更新 $x$ 的一半组件。

与在 HMC 中一样，我们通过按顺序执行 $N_{\mathrm{LF}}$ 蛙跳步骤，然后执行 Metropolis-Hastings 接受/拒绝步骤（如方程 7 中所述）来形成完整的轨迹。 然而，与 HMC 的表达式不同，我们必须考虑更新 $\xi \to {\xi }^{\prime }$ 中的整体雅可比因子，可以轻松计算为 $\left|\frac{\partial v_k^{\prime \prime }}{\partial v_k}\right|=\exp \left(\frac{1}{2}{\epsilon }_v^k{s}_v^k({\zeta }_{v_k})\right)$、$\left|\frac{\partial x_k^{\prime \prime }}{\partial x_k}\right|=\exp \left({\epsilon }_x^k{s}_x^k({\zeta }_{x_k})\right)$ 。 我们保持个体自由度的雅可比因子易于处理。

5 $U(1)$ 格子规范理论

图2： 匾额

晶格规范理论是量子场论与规范场的一种表述，例如。 量子电动力学（QED）和量子色动力学（QCD）。 它是可以进行数值计算的主要方法，并广泛用于模拟亚原子粒子和原子核。 在这里，我们描述了 $U(1)$ 规范 (QED)，并将我们的 DLHMC 应用在二维系统中。 让$U_{\mu }(n)=e^{i{x}_{\mu }(n)}\in U(1)$和$x_{\mu }(n)\in [-\pi ,\pi ]$表示链接变量，其中$x_{\mu }(n)$是面向$\widehat{\mu }$的链接-方向位于站点$n$。我们可以根据威尔逊动作$S_{\beta }(x)$将我们的目标分布写为

$$p_t(x)=\frac{1}{Z}{e}^{-{\gamma }_t{S}_{\beta }(x)},\text{with}{S}_{\beta }(x)=\beta \sum _{P}1-\cos (x_P)$$ (3)

其中 $Z$ 是棘手的标准化； ${\gamma }_t$ 是在训练期间缓慢退火的缩放因子 ($\Vert {\gamma }_t\Vert \le 1$)(B 节）； $x_P\equiv x_{\mu }(n)+x_{\nu }(n+\widehat{\mu })-x_{\mu }(n+\widehat{\nu })-x_{\nu }(n)$是围绕$1\times 1$基本模板的链接变量的组合，如图2所示；总和是格子上所有这些小匾额的总和。 这里，$\beta =2/g_0^2$是耦合常数，$\beta \to \mathrm{\infty }$恢复了理论的连续极限。

物理上，每个晶格配置都有其拓扑电荷 ${𝒬}_{\mathbb{Z}}\in \mathbb{Z}$，定义为

$${𝒬}_{\mathbb{Z}}(x)\equiv \frac{1}{2\pi }\sum _P\lfloor x_P\rfloor ,\text{where}\lfloor x_P\rfloor =x_P-2\pi \lfloor \frac{x_P+\pi }{2\pi }\rfloor .$$ (4)

当前的方法在不同拓扑之间混合的能力受到严重限制，当模拟接近连续极限时，这种现象称为拓扑冻结。 因此，我们希望构建一个损失函数，鼓励我们的采样器探索不同的拓扑部分。 为此，我们引入了拓扑电荷的连续版本 ${𝒬}_{ℝ}\in ℝ$，通过替换方程 4 中的投影来给出：

$${𝒬}_{ℝ}\equiv \frac{1}{2\pi }\sum _P\sin (x_P).$$ (5)

这个量的优点是连续可微，这对于训练深度神经网络很重要。 我们的损失函数定义为

$$ℒ(\theta )={𝔼}_{p_t(\xi )}\left[-\delta ({\xi }^{\prime },\xi )A({\xi }^{\prime }|\xi )\right]$$ (6)

其中 $\delta ({\xi }^{\prime },\xi )\equiv {\left({𝒬}_{ℝ}(x^{\prime })-{𝒬}_{ℝ}(x)\right)}^2$ 和 $A({\xi }^{\prime }|\xi )$ 由公式 7 给出。 在实践中，我们用一批 MC 配置的样本来近似精确的损失函数。

6结果

我们将 DLHMC 算法与 C 节中的训练过程和 B 节中的退火计划应用到二维 $U(1)$ 规范理论，参数逐渐接近连续极限。 然后，我们使用经过训练的模型进行推理来生成马尔可夫链，并将其效率与 HMC 进行比较。 我们使用 Horovod （Sergeev & Balso，2018） 和 TensorFlow （Abadi 等人，2016） 在阿贡领导计算设施的 ThetaGPU 超级计算机上训练我们的模型。 在 1 个节点（$8\times$ NVIDIA A100 GPU）上运行一个典型的训练，使用批量大小 $M=2048$、每个 $N_{\mathrm{LF}}=10$ 跃迁层的隐藏层形状 $=[256,256,256]$，在 $16\times 16$ 网格上运行 $5\times {10}^5$ 训练步骤，大约需要 $24$ 小时才能完成。

图3： 集成自相关时间 ${\tau }_{\mathrm{int}}^{{𝒬}_{ℝ}}$ 与 $\beta$ 的估计，按 $N_{\mathrm{LF}}$ 缩放以考虑模拟成本。

为了衡量我们模型的改进，我们评估了拓扑电荷的积分自相关时间 ${\tau }_{\mathrm{int}}^{{𝒬}_{ℝ}}$，它可以粗略地解释为绘制独立样本之前所需的轨迹数量（平均）。 为了更准确地捕获两种方法之间的计算量，我们通过跳跃步数 $N_{\mathrm{LF}}{\tau }_{\mathrm{int}}^{{𝒬}_{ℝ}}$ 来缩放对积分自相关时间的估计。 我们可以在图3中看到，经过训练的模型在$\beta =2,3,\mathrm{\dots },7$范围内始终优于通用HMC，并且对于研究的最大晶格耦合$\beta =7$，标准HMC将采用大约一百倍的更新才能生成与 DLHMC 一样的独立配置。 为了确保我们结果的有效性，每次训练运行都会与 $N_{\mathrm{LF}}$ 和 $\epsilon$ 值范围内的多次 HMC 运行进行比较。 我们在图 4 中看到，对于大部分模拟，HMC 仍停留在特定的 ${𝒬}_{\mathbb{Z}}$ 值，而经过训练的采样器在值之间快速跳转。 为了更好地理解驱动这种改进行为的机制，我们在系统通过 $N_{\mathrm{LF}}$ 蛙跳层期间评估了系统的不同量。 图5(a)和图5(b)显示了通过$N_{\mathrm{LF}}=10$ 越级层数。 我们的采样器在轨迹的前半段人为地增加了物理系统的能量密度，然后返回到其原始物理值。 这不是规定的行为，而是在训练中自学的。 我们相信，这种在轨迹期间改变能量的能力有助于采样器克服拓扑扇区之间的能量壁垒，而 HMC 仍然停滞不前。 从图5(c)可以看出，连续拓扑电荷在深度神经网络的中间表现出混合行为。

图4： HMC（黑色）和训练模型（蓝色）的拓扑电荷 ${𝒬}_{\mathbb{Z}}$ 与 MC 步骤。

(a) Difference in the average plaquette, $⟨x_P-x_P^{\ast }⟩$ at each leapfrog layer.

(b) The adjusted energy, $\mathscr{H}-\sum \log |𝒥|$ at each leapfrog layer.

(c) The continuous topological charge, ${𝒬}_{ℝ}$ at each leapfrog layer.

图5： 该图显示了各种兴趣量在通过广义蛙跳更新中的每个蛙跳层时如何转换。

7讨论和进一步研究

我们提出 DLHMC 作为 MCMC 的有效算法，并观察到模拟 2D $U(1)$ 规范理论的显着改进。 能够有效地隧道拓扑扇区是朝着有效模拟描述我们宇宙的晶格规范理论迈出的重要的第一步。 我们计划继续开发这种方法，目标是最终扩展到 QCD 的 4D $SU(3)$ 规范理论。

致谢

这项研究得到了百亿亿级计算项目 (17-SC-20-SC) 的支持，该项目是美国能源部科学办公室和国家核安全管理局的合作项目。 这项研究是使用阿贡领导计算设施 (ALCF) 的资源进行的，该设施是能源部科学用户设施办公室，根据合同 DE_AC02-06CH11357 提供支持。 这项工作描述了客观的技术结果和分析。 作品中可能表达的任何主观观点或意见并不一定代表美国能源部或美国政府的观点。 本研究中提出的结果是使用 Python （Van Rossum & Drake，2012）、编程语言及其许多数据科学库（Hunter，2007；Harris 等人，2020；Abadi 等）获得的人，2016；Pérez & Granger，2007；

附录A附录

A.1 哈密尔顿蒙特卡罗 (HMC)

哈密​​顿蒙特卡罗算法是一种广泛使用的技术，它允许我们通过构建状态链 $\{x^{(0)},x^{(1)},\mathrm{\dots },x^{(n)}\}$ 从分析已知的目标分布 $p(x)$ 中进行采样，使得 $x^{(n)}\sim p(x)$ 在限制 $n\to \mathrm{\infty }$ 内。 出于我们的目的，我们假设我们的目标分布可以表示为玻尔兹曼分布 $p(x)=\frac{1}{𝒵}{e}^{-S(x)}\propto e^{-S(x)}$，其中 $S(x)$ 是我们理论的作用，并且 $𝒵$ 是归一化因子（配分函数）。 在这种情况下，HMC 首先使用虚拟动量变量 $v$ 来扩充状态空间，该变量呈独立于 $x$ 的正态分布，即 $v\sim 𝒩(0,\mathrm{𝟙})$。 我们的联合分布可以写为 $p(x,v)=p(x)p(v)\propto e^{-S(x)}{e}^{-\frac{1}{2}{v}^{T}v}=e^{-\mathscr{H}(x,v)}$，其中 $\mathscr{H}(x,v)$ 是联合 (x, v) 系统的哈密顿量。 该系统遵循哈密顿方程：$\dot{x}=\frac{\partial \mathscr{H}}{\partial v}$、$\dot{v}=-\frac{\partial H}{\partial x}$，可以沿着$\mathscr{H}=\text{const.}$的等概率轮廓进行数值积分。 明确地，对于步长$\epsilon$和初始状态$\xi =(x,v)$，蛙跳积分器通过执行以下一系列更新来生成提案配置${\xi }^{\prime }\equiv (x^{\prime },v^{\prime })$：

1. 半步动量更新：$v^{1/2}\equiv v\left(t+\frac{\epsilon }{2}\right)=v-\frac{\epsilon }{2}{\partial }_{x}S(x)$

2. 全步位置更新：$x^{\prime }\equiv x(t+\epsilon )=x+\epsilon v^{1/2}$

3. 半步动量更新：$v^{\prime }\equiv v(t+\epsilon )=v^{1/2}-\frac{\epsilon }{2}{\partial }_{x}S(x^{\prime })$

然后，我们可以通过按顺序执行$N_{\mathrm{LF}}$个蛙步来构建长度为$\lambda =\epsilon N_{\mathrm{LF}}$的完整轨迹。 在我们的轨迹结束时，我们根据 Metropolis-Hastings 接受标准接受或拒绝提案配置，

(7)

蛙跳积分器是辛的，因此雅可比因子减小为$\left|\frac{\partial {\xi }^{\prime }}{\partial {\xi }^T}\right|=1$。

附录 B退火时间表

为了帮助我们的采样器克服隔离模式之间的巨大能量障碍，我们在训练阶段($N$训练步骤）${\{{\gamma }_t\}}_{t=0}^N=\{{\gamma }_0,{\gamma }_1,\mathrm{\dots },{\gamma }_{N-1},{\gamma }_N\}$引入了退火计划，其中,${\gamma }_{t+1}-{\gamma }_t\ll 1$。 只要我们在训练结束时使用 $\gamma \equiv 1$ 恢复真实分布并在没有此因素的情况下评估我们训练的模型，我们就可以在初始训练阶段自由改变 $\gamma$ 。 明确地说，对于 ，此缩放因子有助于降低能垒的高度，使我们的采样器更容易探索以前无法访问的相空间区域。 就这个额外的退火计划而言，我们的目标分布选择一个额外的索引 $t$ 来表示我们在训练阶段的进度，可以明确地写为

$$p_t(x)\propto e^{-{\gamma }_{t}S(x)},\text{for}t=0,1,\mathrm{\dots },N$$ (8)

附录C训练算法

input : 1. Loss function, ${ℒ}_{\theta }({\xi }^{\prime },\xi ,A({\xi }^{\prime }|\xi ))$ 2. Batch of initial states, $x$ 3. Learning rate schedule, ${\{{\alpha }_t\}}_{t=0}^{N_{\mathrm{train}}}$ 4. Annealing schedule, ${\{{\gamma }_t\}}_{t=0}^{N_{\mathrm{train}}}$ 5. Target distribution, $p_t(x)\propto e^{-{\gamma }_t{S}_{\beta }(x)}$

Initialize weights $\theta$

for :

resample $v\sim 𝒩(0,\mathrm{𝟙})$

resample $d\sim 𝒰(+,-)$

construct ${\xi }_0\equiv (x_0,v_0,d_0)$

for :

propose (leapfrog layer) ${\xi }_k^{\prime }\leftarrow {\xi }_k$

compute $A({\xi }^{\prime }|\xi )=\min \{1,\frac{p({\xi }^{\prime })}{p(\xi )}\left|\frac{\partial {\xi }^{\prime }}{\partial {\xi }^T}\right|\}$

update $ℒ\leftarrow {ℒ}_{\theta }({\xi }^{\prime },\xi ,A({\xi }^{\prime }|\xi ))$

backprop $\theta \leftarrow \theta -{\alpha }_t{\nabla }_{\theta }ℒ$

assign

Algorithm 1 Training procedure