Flow-based generative models for Markov chain Monte Carlo in lattice field theory

TL;DR

This paper introduces a flow-based generative model for Markov chain Monte Carlo in lattice field theory, improving sampling efficiency and reducing critical slowing down without requiring prior samples.

Contribution

It presents a novel flow-based generative approach for MCMC that enhances autocorrelation times and can be trained without existing samples, outperforming traditional methods.

Findings

Reduces autocorrelation times in lattice simulations.

Performs comparably or better than HMC and Metropolis methods.

Effective even in critical slowing down regions.

Abstract

A Markov chain update scheme using a machine-learned flow-based generative model is proposed for Monte Carlo sampling in lattice field theories. The generative model may be optimized (trained) to produce samples from a distribution approximating the desired Boltzmann distribution determined by the lattice action of the theory being studied. Training the model systematically improves autocorrelation times in the Markov chain, even in regions of parameter space where standard Markov chain Monte Carlo algorithms exhibit critical slowing down in producing decorrelated updates. Moreover, the model may be trained without existing samples from the desired distribution. The algorithm is compared with HMC and local Metropolis sampling for $\phi^4$ theory in two dimensions.

Tables

Table 1. Table 1: Parameters { m 2 , λ } superscript 𝑚 2 𝜆 \{m^{2},\lambda\} of ϕ 4 superscript italic-ϕ 4 \phi^{4} theory on L × L 𝐿 𝐿 L\times L lattices used for numerical study of the flow-based MCMC algorithm proposed here. The coupling constants λ 𝜆 \lambda have been chosen to approximately maintain constant m p ​ L subscript 𝑚 𝑝 𝐿 m_{p}L as L 𝐿 L is varied.

	E1	E2	E3	E4	E5
$L$	$6$	$8$	$10$	$12$	$14$
$m^2$	$-4$	$-4$	$-4$	$-4$	$-4$
$\lambda$	$6.975$	$6.008$	$5.550$	$5.276$	$5.113$
$m_{p}L$	$3.96\left(3\right)$	$3.97\left(5\right)$	$4.00\left(4\right)$	$3.96\left(5\right)$	$4.03\left(6\right)$

Equations

$$p(\varphi )=e^{-S(\varphi )}/Z,\text{with}Z=\int \prod _{j=1}^{D}d{\varphi }_j{e}^{-S(\varphi )},$$

$$T(\varphi ,{\varphi }^{\prime })\ge 0\text{and}\int \prod _{j=1}^{D}d{\varphi }_j^{\prime }T(\varphi ,{\varphi }^{\prime })=1.$$

$$\exists n\text{ such that }{T}^n(\varphi ,{\varphi }^{\prime })>0\text{ for all }\varphi ,{\varphi }^{\prime },$$

$$\exists \varphi \text{ such that }T(\varphi ,\varphi )>0.$$

$$\int \prod _{j=1}^{D}d{\varphi }_{j}p(\varphi )T(\varphi ,{\varphi }^{\prime })=p({\varphi }^{\prime }).$$

$$A({\varphi }^{(i-1)},{\varphi }^{\prime })=\min \left(1,\frac{\tilde{p}({\varphi }^{(i-1)})}{p({\varphi }^{(i-1)})}\frac{p({\varphi }^{\prime })}{\tilde{p}({\varphi }^{\prime })}\right).$$

$$\tilde{p}(\varphi )>0,p(\varphi )>0\text{ for all }\varphi ,$$

$${\tilde{p}}_f(\varphi )=r(f(\varphi ))\left|\det \frac{\partial f(\varphi )}{\partial \varphi }\right|.$$

$$g_i(\varphi ):=\{\begin{array}{l}{z}_a={\varphi }_a\\ z_b={\varphi }_b\odot e^{s_i({\varphi }_a)}+t_i({\varphi }_a),\end{array}$$

$$g_i^{-1}(z):=\{\begin{array}{l}{\varphi }_a=z_a\\ {\varphi }_b=(z_b-t_i(z_a))\odot e^{-s_i(z_a)}.\end{array}$$

$$\left|\det \frac{\partial g_i(\varphi )}{\partial \varphi }\right|=\prod _{j=1}^{D/2}{e}^{{\left[s_i({\varphi }_a)\right]}_j},$$

$$f(\varphi )=g_1(g_2(\dots g_n(\varphi )\dots )).$$

$$L({\tilde{p}}_f)$$

$$=\int \prod _{j}d{\varphi }_j{\tilde{p}}_f(\varphi )\left(\log {\tilde{p}}_f(\varphi )-\log p(\varphi )-\log Z\right)$$

$$=\int \prod _{j}d{\varphi }_j{\tilde{p}}_f(\varphi )\left(\log {\tilde{p}}_f(\varphi )+S(\varphi )\right).$$

$$\widehat{L\left({\tilde{p}}_f\right)}=\frac{1}{M}\sum _{i=1}^M\left(\log {\tilde{p}}_f({\varphi }^{(i)})+S({\varphi }^{(i)})\right).$$

$$p_{\tau \text{rej}}\equiv ⟨\prod _{i=1}^{\tau }\text{mathbbm}{1}_{\text{rej}}(i)⟩=\rho (\tau )/\rho (0),$$

$${\widehat{\rho (\tau )/\rho (0)}}_{\text{acc}}=\frac{1}{N-\tau }\sum _{j=1}^{N-\tau }\prod _{i=1}^{\tau }\text{mathbbm}{1}_{\text{rej}}(i+j).$$

$${\widehat{\rho (\tau )/\rho (0)}}_{\mathcal{O}}=\frac{\frac{1}{N-\tau }{\sum }_{i=0}^{N-\tau -1}({\mathcal{O}}_i-\bar{\mathcal{O}})({\mathcal{O}}_{i+\tau }-\bar{\mathcal{O}})}{\frac{1}{N}{\sum }_{i=0}^{N-1}{\left({\mathcal{O}}_i-\bar{\mathcal{O}}\right)}^2},$$

$$\rho (\tau )/\rho (0)=\underset{\varphi \sim p}{\mathbb{E}}\left[p_{\text{rej}}^{\tau }(\varphi )\right]\ge {\left(\underset{\varphi \sim p}{\mathbb{E}}\left[p_{\text{rej}}(\varphi )\right]\right)}^{\tau }=(1-a{)}^{\tau }.$$

$${\tau }_{\mathcal{O}}^{\text{int}}=\frac{1}{2}+\underset{{\tau }_{\text{max}}\to \infty }{\lim }\sum _{\tau =1}^{{\tau }_{\text{max}}}\frac{{\rho }_{\mathcal{O}}(\tau )}{{\rho }_{\mathcal{O}}(0)}.$$

$$S(\varphi )=\sum _x\left(\sum _y\varphi (x)\square (x,y)\varphi (y)+m^2\varphi (x{)}^2+\lambda \varphi (x{)}^4\right),$$

$$\sum _y\square (x,y)\varphi (y)=\sum _{\mu }\left(2\varphi (x)-\varphi (x-\hat{\mu })-\varphi (x+\hat{\mu })\right).$$

$$G_c(x)$$

$${\tilde{G}}_c(\vec{p},t)=\frac{1}{L^{d-1}}\sum _{\vec{x}}{e}^{i\vec{p}\cdot \vec{x}}{G}_c(\vec{x},t),$$

$$m_p=-{\partial }_t\log ⟨{\tilde{G}}_c(0,t)⟩,$$

$${\chi }_2=\sum _x{G}_c(x).$$

$$E=\frac{1}{d}\sum _{1\le \mu \le d}{G}_c(\hat{\mu }),$$

$$r(z)\propto \prod _i{e}^{-z_i^2/2}.$$

$$m_p^{\text{eff}}(t)=\text{arccosh}\left(\frac{{\tilde{G}}_c(0,t-1)+{\tilde{G}}_c(0,t+1)}{2{\tilde{G}}_c(0,t)}\right).$$

$$\underset{N\to \infty }{\lim }\mathbb{E}\left[{\widehat{\rho (\tau )/\rho (0)}}_O\right]$$

Full text

††preprint:

MIT-CTP/5114

I Introduction

A key problem in lattice field theory and statistical mechanics is the evaluation of integrals over field configurations, referred to as path integrals. Typically, such integrals are evaluated via a Markov chain Monte Carlo (MCMC) approach: field configurations are sampled from the desired probability distribution, dictated by the action of the theory, using a Markov chain. A significant practical concern is the existence of correlations between configurations in the chain. Critical slowing down Wolff (1990) refers to the divergence of the associated autocorrelation time as a critical point in parameter space is approached. This behavior drastically increases the computational cost of simulations in these parameter regions Del Debbio et al. (2004); Meyer et al. (2007). For some models, algorithms have been found which significantly reduce or eliminate this slowing down Kandel et al. (1988); Bonati and D’Elia (2018); Hasenbusch and Schaefer (2018); Swendsen and Wang (1987); Wolff (1989); Prokof’ev and Svistunov (2001); Kawashima and Harada (2004); Bietenholz et al. (1995), enabling efficient simulation. For field theories, a number of methods have been proposed to circumvent critical slowing down by variations of Hybrid Monte Carlo (HMC) techniques Ramos (2012); Gambhir and Orginos (2015); Cossu et al. (2018); Jin and Osborn (2019), multi-scale updating procedures Endres et al. (2015); Detmold and Endres (2016, 2018), open boundary conditions or non-orientable manifolds Luscher and Schaefer (2013); Mages et al. (2017); Burnier et al. (2018), metadynamics Laio et al. (2016), and machine learning tools Tanaka and Tomiya (2017); Shanahan et al. (2018). In important classes of theories, however, critical slowing down remains limiting; for example, in lattice formulations of Quantum Chromodynamics (QCD, the piece of the Standard Model describing the strong nuclear force) it is a major barrier to simulations at the fine lattice spacings required for precise control of the continuum limit.

Here, a new flow-based MCMC approach is proposed and is applied to lattice field generation. The resulting Markov chain has autocorrelation properties that are systematically improvable by an optimization (training) step before sampling. In this method, samples $z$ are drawn from a simple distribution and then transformed by a change-of-variables (or “flow”) $\phi=f^{-1}(z)$, resulting in samples $\phi$ with a new effective distribution $\tilde{p}_{f}$. The mapping $f^{-1}$ is chosen to be efficient to compute, making it easy to draw samples $\phi$, and is optimized within a variational family to produce a distribution $\tilde{p}_{f}$ close to the desired one. To guarantee asymptotic exactness of sampling, a Markov chain is constructed using Metropolis-Hastings steps with $\tilde{p}_{f}$ taken as a proposal distribution. Since proposed samples are independent of the previous samples in the chain, the autocorrelation time and acceptance rate are coupled; the autocorrelation time drops to zero as the acceptance rate approaches $1$. This is true even in regions of parameter space where standard algorithms exhibit critical slowing down. Under mild conditions (detailed in Section II), this approach is guaranteed to generate samples from the desired probability distribution in the limit of a large number of updates.

This method has several features that make it attractive for the evaluation of path integrals in lattice field theories:

The autocorrelation time of the Markov chain can be systematically decreased by training the model; 2. 2.

Each step of the Markov chain requires only the model evaluation and an action computation; 3. 3.

Each update proposal is independent of the previous sample, thus proposals can be generated in parallel and efficiently composed into a Markov chain; 4. 4.

The model is trained using samples produced by the model itself, without the need for existing samples from the desired probability distribution.

Several other machine learning approaches have been applied to MCMC, for statistical mechanics systems, synthetic distributions, and simple lattice quantum field theories. Self-learning Monte Carlo (SLMC) methods have been applied fairly successfully to one- to three-dimensional Ising and fermionic systems. These methods construct, by a variety of techniques, an effective Hamiltonian for a theory that can be more easily sampled than the original Hamiltonian Huang and Wang (2017); Liu et al. (2017, 2016); Nagai et al. (2017); Shen et al. (2018). The effective Hamiltonian is learned using supervised learning techniques based on training data drawn from a combination of existing MCMC simulations, randomly-mutated samples, and the accelerated Markov chain itself (hence the term “self-learning”). Flow-based methods have been used for Monte Carlo sampling in the two-dimensional Ising model Li and Wang (2018), many-body systems Noé and Wu (2018), and synthetic distributions Song et al. (2017); Levy et al. (2017), and generative adversarial methods have been applied to two-dimensional scalar field theory Urban and Pawlowski (2018); Zhou et al. (2019).

In contrast to these approaches, the method proposed here focuses on directly generating samples from a close approximation to the true distribution in such a way that the exact likelihood of each produced sample is known. The direct generation allows self-learning as in SLMC, and the known likelihood allows use of a Metropolis-Hastings acceptance step to ensure exactness.

The proposed flow-based MCMC algorithm is detailed in Section II. A numerical study of its effectiveness in the context of two-dimensional $\phi^{4}$ theory is presented in Section III. Finally, Section IV outlines the further development and scaling of the approach that will be required for applications to theories defined in a larger number of spacetime dimensions and to more complicated field theories such as QCD.

II A flow-based Markov chain Monte Carlo algorithm

In lattice field theory, a Markov chain Monte Carlo (MCMC) process is an efficient way to generate field configurations $\phi\in\mathbb{R}^{D}$ distributed according to a target probability distribution

[TABLE]

where $j$ indexes the $D$ components of $\phi$, $S(\phi)$ is the action that defines the theory and $Z$ is the partition function. Here, $\phi$ is defined to be a vector of $D$ real components representing the combined internal $(\alpha)$ and spacetime $(x)$ degrees of freedom of the field $\phi(x,\alpha)$ evaluated on a finite, discrete spacetime lattice (generalizations to gauge fields are discussed in Section IV). A MCMC process generates a chain $\phi^{(0)}\rightarrow\phi^{(1)}\rightarrow\ldots\phi^{(N)}$ by steps through configuration space starting with an arbitrary configuration $\phi^{(0)}$. The steps are stochastic and are determined by the probabilities $T(\phi,\phi^{\prime})$ associated with each possible transition $\phi\rightarrow\phi^{\prime}$. These probabilities must be non-negative and normalized:

[TABLE]

They must also satisfy the conditions of ergodicity and balance to ensure that samples in the chain are drawn from a distribution that converges to $p(\phi)$ after thermalization. For the chain to be ergodic, it must be possible to transition from a starting configuration $\phi$ to any other configuration $\phi^{\prime}$ in a finite number of steps, i.e.,

[TABLE]

and the chain must not have a period, for which it is sufficient that a single state has non-zero self-transition probability, i.e.,

[TABLE]

Balance is the condition that $p(\phi)$ is a stationary distribution of the transition:

[TABLE]

Any procedure which satisfies these conditions will, in the limit of a sufficiently long Markov chain, produce field configurations $\{\phi^{(i)}\}$ distributed according to $p(\phi)$.

II.1 Metropolis-Hastings with generative models

Given a model that allows sampling from a known probability distribution $\tilde{p}(\phi)$, a Markov chain for a desired probability distribution $p(\phi)$ can be constructed via the independence Metropolis sampler, a specialization of the Metropolis-Hastings method Tierney (1994). For each step $i$ of the chain, an update proposal $\phi^{\prime}$ is generated by sampling from $\tilde{p}(\phi)$, independent of the previous configuration. This proposal is accepted with probability

[TABLE]

If the proposal is accepted, $\phi^{(i)}=\phi^{\prime}$, otherwise $\phi^{(i)}=\phi^{(i-1)}$. This procedure defines the transition probabilities of the Markov chain.

The general Metropolis-Hastings algorithm has been proven to satisfy balance Hastings (1970) for any proposal scheme. For the independence Metropolis sampler, under the further condition that every state $\phi$ has non-zero proposal density and non-zero desired density,

[TABLE]

the Markov chain is also ergodic and thus guaranteed to converge to the desired distribution Tierney (1994).

This Markov chain can be intuitively considered a method to correct an approximate distribution $\tilde{p}(\phi)$ to the desired distribution $p(\phi)$. The accept/reject statistics of the Metropolis-Hastings algorithm serve as a diagnostic for closeness of the approximate and desired distributions; if the distributions are equal, proposals are accepted with probability 1 and the Markov chain process is equivalent to a direct sampling of the desired distribution. This is made precise in Section II.3.

II.2 Sampling using normalizing flows

Here, a normalizing flow model is used to define a proposal distribution $\tilde{p}(\phi)$ for a generative Metropolis-Hastings algorithm. Normalizing flows Rezende and Mohamed (2015) are a machine learning approach to the task of sampling from complicated, intractable distributions. They do so by learning a map from an input distribution that is easy to sample to an output distribution that approximates the desired distribution. Normalizing flow models produce both samples and their associated probability densities, allowing the acceptance probability in Eq. (6) to be calculated.

A normalizing flow enacts the transformation between distributions by a change-of-variables111The convention of using $f^{-1}$ for the change-of-variables stems from typical applications of normalizing flows.: a smooth, bijective function, $f^{-1}:\mathbb{R}^{D}\rightarrow\mathbb{R}^{D}$ maps samples $z$ from a prior distribution $r(z)$ to $\phi=f^{-1}(z)$. This mapping defines an output distribution $\tilde{p}_{f}(\phi)$, by the change-of-variables formula

[TABLE]

Typically, the prior distribution is a simple and analytically-understood distribution (e.g., a normal distribution). While the desired distribution $p(\phi)$ is often complicated and difficult to sample from directly, optimizing the function $f$ allows one to generate samples from $\tilde{p}_{f}(\phi)\approx p(\phi)$. The function $f$ is chosen to have a tractable Jacobian such that the probability density $\tilde{p}_{f}(\phi)$ can be computed exactly according to Eq. (8).

To encode a map from a simple distribution $r(z)$ to a complicated distribution $\tilde{p}_{f}(\phi)$, the map $f$ must be highly expressive while also being invertible and having a computable Jacobian. Here, the real non-volume-preserving (real NVP) flow Dinh et al. (2016) machine learning approach is used: $f$ is constructed by the composition of affine coupling layers that scale and offset half of the components of the input at a time; the choice of which components of the data are transformed is part of the layer definition. Splitting the $D$-dimensional vector $\phi$ into ($D/2$)-dimensional pieces $\phi_{a}$ and $\phi_{b}$ according to this choice, a single coupling layer $g_{i}$ transforms $\phi$ to $z=g_{i}(\phi)$ via

[TABLE]

where $s_{i}$ and $t_{i}$ are neural networks mapping from $\mathbb{R}^{D/2}$ to $\mathbb{R}^{D/2}$ and $\odot$ denotes element-wise multiplication. Importantly, each layer $g_{i}$ is invertible without inverting the neural networks $s_{i}$ or $t_{i}$:

[TABLE]

The Jacobian matrix is lower-triangular and its determinant can be easily computed. For coupling layer $g_{i}$:

[TABLE]

where $j$ indexes the $D/2$ components of the output of $s_{i}$. Stacking many coupling layers $g_{1},\dots,g_{n}$ which alternate which half of the data is transformed, the function $f$ is defined as

[TABLE]

Using the chain rule, the determinant of the Jacobian of $f$ is a product of the contributions from each $g_{i}$. By increasing the number of coupling layers and the complexity of the networks $s_{i}$ and $t_{i}$, $f$ can systematically be made more expressive and general. Figure 1 depicts how composing many coupling layers incrementally modifies a prior distribution which is easy to sample into a more complex output distribution that approximates a distribution of interest.

For a fixed initial distribution $r(z)$, the neural networks within each affine coupling layer of $f$ can be trained to bring $\tilde{p}_{f}(\phi)$ close to the desired distribution $p(\phi)$. This training is undertaken by minimizing a loss function. Here, the loss function used is a shifted Kullback-Leibler (KL) divergence222This training paradigm is a specific instance of Probability Density Distillation Oord et al. (2017); Li and Wang (2018). between the target distribution of the form $p(\phi)=e^{-S(\phi)}/Z$ and the proposal distribution $\tilde{p}_{f}(\phi)$:

[TABLE]

This loss function has been successfully applied in related generative approaches to statistical lattice models Zhang et al. (2018); Li and Wang (2018). The formal shift by $\log{Z}$ in Eq. (13) eliminates the need to compute the true partition function, and does not affect the gradients or location of the minima. By non-negativity of the KL divergence, the lower bound on the loss is $-\log{Z}$, and this minimum is achieved exactly when $\tilde{p}_{f}=p$. In practice, the loss is stochastically estimated by drawing batches of $M$ samples from the model $\left\{\phi^{(i)}\sim\tilde{p}_{f}\right\}$ and computing the sample mean:

[TABLE]

The loss minimization can then be undertaken using stochastic optimization techniques such as stochastic gradient descent or momentum-based methods including Adam and Nesterov Kingma and Ba (2015); Nesterov (1983).

By construction, the flow model allows sampling from $\tilde{p}_{f}$ efficiently. The training process can thus be performed by drawing samples from the model itself, rather than using existing samples from the desired distribution as training data. This self-training is a key feature of the proposed approach to Monte Carlo sampling for field theories, where samples from the desired distribution are often computationally expensive to obtain. If samples do exist, they can be used to ‘pre-train’ the network, although in the numerical studies undertaken here this was not found to be markedly more efficient in network optimization than using only self-training.

Given a trained model with distribution $\tilde{p}_{f}(\phi)\approx p(\phi)$, samples from $\tilde{p}_{f}$ can be used as proposals to advance a Markov chain using the generative Metropolis-Hastings algorithm described above. This forms the basis for the flow-based MCMC algorithm proposed here:

A flow-based generative model (here, a real NVP model) is trained using the shifted KL loss given in Eq. (13) to have output distribution $\tilde{p}_{f}(\phi)\approx p(\phi)$; 2. 2.

$N$ proposals $\left\{\phi^{\prime(i)}\sim\tilde{p}_{f}\right\}$ are produced by sampling from the flow-based model (this can be done in parallel) and the associated action $S(\phi^{\prime})$ is computed for each proposal; 3. 3.

Starting from an arbitrary initial configuration, each proposed sample is successively accepted or rejected using the Metropolis-Hastings algorithm given in Eq. (6) to build a Markov chain of length $N$.

When the prior distribution $r(z)$ is strictly positive, the invertibility and continuity of $f$ guarantees that the generated distribution $\tilde{p}_{f}(\phi)$ is also strictly positive. For all models with finite action, and thus $p(\phi)>0$, the resulting Markov chain is then ergodic by the arguments detailed in Section II.1.

II.3 Autocorrelation time for generative Metropolis-Hastings

For any Markov chain constructed via a generative Metropolis-Hastings algorithm (with independent update proposals), an observable-independent estimator for autocorrelation time can be defined from the accept/reject statistics of the chain. This serves both as a measure of the similarity between the proposal and desired distributions and enables proper error estimation for lattice observables Wolff (2004).

Precisely, the autocorrelation at Markov chain separation $\tau$, for all observables, is given by the probability of $\tau$ rejections in a row,

[TABLE]

where $\mathbbm{1}_{\text{rej}}(i)$ is an indicator variable taking value $1$ when the proposed step from $i-1$ to $i$ was rejected in the Metropolis-Hastings algorithm, and [math] otherwise. In practice, for a near-equilibrium, finite Markov chain with length $N$, a finite-sample estimator provides a good approximation to $\rho(\tau)/\rho(0)$:

[TABLE]

This measure of autocorrelation is consistent with the usual definition; it is shown in Appendix A that the standard estimator for the autocorrelation of any given observable $\mathcal{O}$,

[TABLE]

also converges to $p_{\tau\text{rej}}$ in the limit $N\rightarrow\infty$.

The qualitative relation between acceptance rate and autocorrelations gives a convenient measure of the autocorrelation characteristics of a Markov chain. Precisely, the autocorrelation at distance $\tau$ can be bounded in terms of the average acceptance rate $a=1-\mathbb{E}\left[p_{\text{rej}}\right]$:

[TABLE]

Increasing the acceptance rate of an independence Metropolis sampler is thus a necessary condition to reduce autocorrelations. Additionally, $a=1$ exactly when the proposal and desired distributions are equal. In this case, $p_{\text{rej}}(\phi)=0$ for each $\phi$, and there are no autocorrelations. While bringing $a$ close to $1$ does not provide an upper bound on autocorrelation, stochastically improving a loss function that measures distance between distributions is expected to reduce autocorrelations on average. In practice, autocorrelations should be evaluated as a test metric alongside the training loss to confirm improvement over the course of training the model. The correspondence between loss minimization, acceptance rate, and autocorrelations is studied in the context of $\phi^{4}$ theory in Section III, where a clear correlation between $a$ and $\rho(\tau)/\rho(0)$ is observed.

II.4 Critical slowing down

When a distribution is sampled using a Markov chain with large autocorrelation time, many updates are required to produce decorrelated samples. Critical slowing down (CSD) is defined as the divergence of the autocorrelation time of Markov chain sampling as a critical point in parameter space is approached Wolff (1990). A numerically-stable definition of the characteristic autocorrelation time of a Markov chain is the integrated autocorrelation time:

[TABLE]

As a critical point is approached, analysis of standard local-update algorithms for lattice models suggests $\tau_{\mathcal{O}}^{\text{int}}$ is typically well-described by a power law in the lattice spacing, or for fixed physical volume, a power law in the lattice sites per dimension, $L$. A dynamical critical exponent $z_{\mathcal{O}}$ is thus defined by a fit to $\tau_{\mathcal{O}}^{\text{int}}=\alpha_{\mathcal{O}}L^{z_{\mathcal{O}}}$ along a line of constant physics. An update algorithm for which the critical exponent is zero is unaffected by CSD.

In any generative Metropolis-Hastings simulation, the autocorrelation time is completely fixed by the expected accept/reject statistics, which in turn result from the structure of the proposal and desired distributions. For models trained with a target value of the integrated autocorrelation time used as a stopping criterion, CSD associated with the Markov chain sampling is thus trivially removed at the expense of up-front training costs. The difficulty of CSD is in essence shifted to the training of the model, i.e., to the optimization of the proposal distribution. The cost and scaling of this optimization task for $\phi^{4}$ theory, as studied here, and the prospects of scaling this approach to more complicated theories, are discussed in Sec. III.4.

III Application of flow-based MCMC to $\phi^{4}$ theory

The theory with a massive scalar field $\phi(x)$ and a quartic self-interaction is one of the simplest interacting field theories that can be constructed. It is thus a convenient testing ground for new algorithms for lattice field theory, such as the flow-based MCMC approach proposed here.

In a $d$-dimensional Euclidean spacetime, a discretized formulation of $\phi^{4}$ theory can be defined on a lattice with sites $x_{\mu}=an_{\mu}$, where $a$ denotes the lattice spacing, $\mu\in\left\{1,\dots,d\right\}$ labels spacetime dimension, and $n_{\mu}\in\mathbb{Z}^{d}$. Here, a finite lattice volume $V=(aL)^{d}$ is considered, with periodic boundary conditions in all dimensions. The lattice action (in units where $a=1$) can be expressed as

[TABLE]

where the parameters $m^{2}$ and $\lambda$ are the bare mass squared and bare coupling, respectively, and the lattice d’Alembert operator is defined by

[TABLE]

By taking expectation values over the distribution $p(\phi)=e^{-S(\phi)}/Z$, observables in the theory can be estimated.

The observables studied here are the connected two-point Green’s function

[TABLE]

and its momentum-space representation

[TABLE]

where $x_{\mu}=(\vec{x},t)$, as well as the corresponding pole mass

[TABLE]

and the two-point susceptibility

[TABLE]

In the limit $\lambda\rightarrow\infty$, with $m^{2}/\lambda<0$ fixed, scalar $\phi^{4}$ theory reduces to an Ising model. Another observable of interest is therefore the average Ising energy density Vierhaus (2010), defined by

[TABLE]

where the sum runs over single-site displacements in all dimensions.

The action of $\phi^{4}$ theory is invariant under the discrete symmetry $\phi(x)\rightarrow-\phi(x)$. Depending on the value of the parameters $m^{2}$ and $\lambda$, this symmetry can be spontaneously broken. The theory thus has two phases: a symmetric phase and a broken-symmetry phase.

III.1 Model definition and training

For this proof-of-principle study, the flow-based MCMC algorithm detailed in Section II was applied to $\phi^{4}$ theory in two dimensions with $L=\{6,8,10,12,14\}$ lattice sites in each dimension. The parameters $m^{2}$ and $\lambda$ were chosen to fix $m_{p}L\approx 4$ for each lattice size; their numerical values are given in Table 1. For simplicity in this initial work, all parameters were chosen to lie in the symmetric phase. In principle, the flow-based MCMC algorithm can be applied with identical methods to the broken-symmetry phase of the theory, but it remains to be shown that models can be trained for such choices of parameters.

For each set of parameters, real NVP models were defined using 8–12 affine coupling layers (see Sec. II.2). The coupling layers were defined to update half of the lattice sites in a checkerboard pattern; successive layers alternately updated the odd and even sites. The neural networks $s_{i}$ and $t_{i}$ used in coupling layer $g_{i}$ (see Eq. (9)) were constructed from two to six fully-connected layers, each defined as multiplication by a rectangular matrix followed by pointwise application of a nonlinear function (here, a leaky rectified linear unit Nair and Hinton (2010)). Intermediate vectors (hidden units) had sizes ranging between 100–1024. The prior distribution $r(z)$ was chosen to be an uncorrelated Gaussian distribution

[TABLE]

The models were trained to minimize the shifted KL loss between the output distribution $\tilde{p}_{f}(\phi)$ and the desired distribution $p(\phi)=e^{-S(\phi)}/Z$ using gradient-based updates with the Adam optimizer Kingma and Ba (2015), a specific variety of gradient descent with momentum. A mean absolute error loss, defined in Appendix B, was optimized before training in the case of the $14^{2}$ model where it was found to accelerate convergence to the KL loss minimum.

An exhaustive study of the optimal choice of prior distribution $r(z)$, model depth, architecture and initialization of the neural networks, and of the mode of coupling of the affine layers, is beyond the scope of this proof-of-principle study. The parameters used here, however, proved to define sufficiently expressive models such that the Metropolis-Hastings algorithm applied to output from the trained models easily achieved acceptance rates of well over 50%. With further investment in hyperparameter optimization, higher rates of acceptance could be achieved. In any Markov chain using the Metropolis-Hastings algorithm, there is a tradeoff between computational cost and correlations resulting from low acceptance rates. The optimal acceptance rate minimizes the cost per decorrelated sample from the chain. Here, the cost of training, and not just model evaluation, must be considered, and the optimal level of training in future applications will depend on many factors, such as the desired ensemble size.

For each set of parameters studied, instances of the model were trained to reach both 50% and 70% average Metropolis acceptance. Figure 2 shows histograms of the number of updates between accepted configurations for models at both levels of training. Models trained to reach the higher acceptance rate are seen to have shorter runs of consecutive rejections. Because autocorrelation is related to rejections by $\rho(\tau)/\rho(0)=p_{\tau\text{rej}}$ for independence Metropolis sampling, a reduced frequency of rejection runs with length longer than $\tau$ directly implies a reduction in $\rho(\tau)/\rho(0)$. Implications for critical slowing down of the generation of decorrelated configurations are discussed in Section III.3.

For comparison, ensembles of $10^{6}$ lattice configurations were generated using the machine-learned models in flow-based MCMC as well as standard local Metropolis Wolff (1996) and Hybrid Monte Carlo (HMC) Duane et al. (1987) algorithms at matched parameters. The local Metropolis algorithm employed a fixed order of sequential updates to each site, with proposed updates to $\phi(x)$ sampled uniformly from the interval $[\phi(x)-\delta,\phi(x)+\delta]$ followed by a Metropolis-Hastings accept/reject step; for all parameters considered, the width $\delta$ was tuned to achieve a 70% acceptance rate. The HMC method was implemented using a leapfrog integrator with a fixed division of trajectory length $\tau$ into $10$ steps; the trajectory length $\tau$ was also tuned to achieve a 70% acceptance rate. In both the local Metropolis and HMC methods, samples were saved after every 10th update.

III.2 Tests: physical observables and error scaling

Since the flow-based MCMC algorithm satisfies ergodicity and balance, it is guaranteed to produce samples from the desired probability distribution in the limit of an infinite chain. To test the performance of the algorithm for a finite number of samples, each of the physical observables defined above was computed on ensembles of configurations at the parameters of Table 1, generated both using standard HMC and local Metropolis methods, as well as with the trained flow-based MCMC algorithm. Figures 4–6 compare the observables computed on ensembles generated using all three methods.

To estimate the pole mass $m_{p}$, an effective mass is defined based on the zero-momentum Green’s functions at various time separations:

[TABLE]

For all observables, the values computed using the flow-based MCMC ensembles are consistent within statistical uncertainties with those computed using the standard methods. Moreover, Figure 6 shows that the statistical uncertainties of the observables scale as $1/\sqrt{N}$ with the number of samples $N$, as expected for decorrelated samples.

III.3 Critical slowing down

For $\phi^{4}$ theory, a number of algorithms have been developed that mitigate CSD to various extents, such as worm algorithms Vierhaus (2010), multigrid methods Goodman and Sokal (1989), Fourier-accelerated Langevin updates Batrouni and Svetitsky (1987) and cluster updates via embedded Ising dynamics Brower and Tamayo (1989). The path towards generalizing those algorithms to more complicated theories such as QCD, however, is not clear. Algorithms such as HMC and local Metropolis, which are also used for studies of QCD and pure gauge theory, exhibit CSD for $\phi^{4}$ (as well as more complicated theories) as the continuum limit is approached.

The parameter sets chosen for the study of $\phi^{4}$ theory in this work (Table 1) correspond to a critical line with constant $m_{p}L$ as $L\rightarrow\infty$. For the flow-based MCMC approach proposed here, as well as for ensembles generated using the HMC and local Metropolis algorithms, the autocorrelation times of the set of physical observables discussed previously were fit to leading-order power laws in $L$ to determine the dynamical critical exponents $z_{\mathcal{O}}$ for that observable. Figure 7 shows the autocorrelation times for each observable for each approach to ensemble generation. The absolute values of $\tau_{\text{int}}$ are not directly comparable between methods because the cost per update differs. The scaling with lattice size, on the other hand, indicates the sensitivity of each method to critical slowing down. For both HMC and local Metropolis, the critical behavior and consequently the performance of the algorithm was found to depend on the observable. In each case, the critical exponent was $0.3\lesssim z_{\mathcal{O}}\lesssim 2.0$. In comparison, for the flow-based MCMC ensembles at a fixed acceptance, the critical exponent was found to be consistent with zero, with the autocorrelation time observable-independent and in agreement with the acceptance-based estimator defined in Section II.3.

Since the the mean acceptance rate was used as the stopping criterion for training these models, it was not guaranteed a priori that the measured integrated autocorrelation time would be constant across the different models used. The results in Figure 7, however, suggest that beyond the simple lower bound from Eq. (18) there is a strong correlation between the mean acceptance rate and integrated autocorrelation time for models trained using a shifted KL loss. This is further confirmed by the similarity of the rejection run histograms across lattice sizes for flow-based MCMC, as shown in Figure 2.

III.4 Training costs

While CSD in the sampling step for the flow-based MCMC is eliminated, training the generative model introduces an additional up-front cost, as discussed in Section II.4. Since this cost is amortized over the ensemble, this approach will naturally be computationally advantageous in the limit of generating a large number of samples. For a finite target ensemble size, the potential acceleration offered depends crucially on the training time.

In this work, all models were trained using one to two GPU-weeks, with the larger lattices incurring the most computational cost. For the simple fully-connected architecture used in this work, the scaling of both the sampling and training time is controlled by dense matrix-vector multiplications which require $O(V^{2})$ floating point operations each. The number of epochs used to train the largest lattice was also roughly $10\times$ that of the smallest lattice. This asymptotic scaling is a result of the simple model architecture used in this proof-of-principle study. For related methods applied to image generation, using convolutional neural networks and a multi-scale architecture reduced training and sampling costs significantly and improved scaling to $O(V)$ Dinh et al. (2016). There are physical grounds to expect these tools to apply equally well to the present application. Convolutional networks use only local information to update values in each layer, exploiting locality in the system, and use identical weights for each point on the lattice, manifestly preserving translational invariance. A multi-scale architecture learns coarse-grained distributions and fine-graining procedures in separate layers; this is an effective division of tasks for renormalizable quantum field theories, where simple coarse-grained descriptions are expected to arise. Generative models, and in particular flow-based models, are also rapidly evolving towards more efficient representation capacity. Complex coupling layers have been implemented Dinh et al. (2014, 2016), as have generalized convolutions Kingma and Dhariwal (2018); Hoogeboom et al. (2019) and transformations with continuous dynamics that are not dependent on restricted coupling layers Grathwohl et al. (2019). These developments allow models to better capture a distribution within a given number of training steps.

For complex applications, it is also critical that larger models with many coupling layers can be trained without exceeding memory bounds. The algorithm proposed here can be trained with constant memory cost as the number of layers is increased Li and Grathwohl (2018), alleviating the storage limitations that can arise in gradient-based optimization. Memory costs can be further reduced by distributing samples within each training batch across many machines.

Finally, typical applications seek to produce ensembles at many different choices of parameters, and often require parameter tuning. Training costs can therefore by amortized further; models trained with respect to an action at a given set of parameter values can either be used to initialize training or as a prior distribution for models targeting that action at nearby parameter values.

IV Summary

This work defines a flow-based MCMC algorithm to sample lattice field configurations from a desired probability distribution:

A real NVP flow model is trained to produce approximately the desired distribution; 2. 2.

Samples are proposed from the trained model; 3. 3.

Starting from an arbitrary configuration, each proposal is accepted or rejected to advance a Markov chain using the Metropolis-Hastings algorithm.

The approach is shown to define an ergodic and balanced Markov chain, thus guaranteeing convergence to the desired probability distribution in the limit of a long Markov chain. In essence, the flow-based MCMC algorithm combines the expressiveness of normalizing flows based on neural networks with the theoretical guarantees of Markov chains to create a trainable and asymptotically-correct sampler. Since these flows are applicable for arbitrary configurations with continuous, real-valued degrees of freedom, one can generically apply this method to any of a broad class of lattice theories. Here, the algorithm is implemented in practice for $\phi^{4}$ theory, and is demonstrated to produce ensembles of configurations that are indistinguishable from those generated using standard local Metropolis and HMC algorithms, based on studies of a number of physical observables.

A key feature of the approach is that models trained to a fixed acceptance rate do not experience critical slowing down in the sampling stage. In particular, the autocorrelation time for all observables is dictated entirely by the accuracy with which the flow model has been trained; perfect training corresponds to decorrelated samples and 100% acceptance in the Metropolis-Hastings step of the MCMC process. Nevertheless, the efficiency with which the training step of this approach can be scaled to larger model sizes, and to more complicated theories such as QCD, remains to be studied. Recent advances in the training and scaling of flow models provide reasons for optimism on this front. Further, incorporating symmetries generally improves data efficiency of training, and implementing spacetime and gauge symmetries Cohen et al. (2019) may be a natural next step to practically train these flow models for lattice gauge theories like QCD.

In moving towards lattice gauge theories such as QCD, several theoretical developments are also required. The real NVP model chosen to parameterize the normalizing flows here is described in terms of vectors of variables $\phi\in\mathbb{R}^{D}$. Gauge configurations, however, live in a compact manifold arising from the Lie group structure. Extending this method will require a normalizing flow model that can act on this manifold while remaining sufficiently expressive. The choice of prior likewise will need to be extended to a distribution over the manifold of lattice gauge configurations which can be easily sampled. A uniform distribution, for example, may be a candidate for a prior, but this choice must be tested in the context of a specific flow model.

If the flow-based MCMC algorithm proposed here can be implemented for a complex theory such as QCD, the advantages would be significant; arbitrarily large ensembles of field configurations could be generated at minimal cost. The independence of the proposal step from any previous configuration allows parallel generation of proposals, and the continually-improving support in hardware and software for neural network execution suggests future practical gains for this style of ensemble-generation. Given efficient sample generation from a trained model, ensembles would not need to be stored long-term. Moreover, a model trained for one action could either be re-trained or used as a prior for another flow model targeting an action with nearby parameter values. This would allow efficient tuning of parameters and generation of additional ensembles interpolating between and extrapolating from existing models.

Acknowledgements

We thank J.-W. Chen, K. Cranmer, W. Detmold, R. Melko, D. Murphy, A. Pochinsky, and B. Trippe for helpful discussions. This work is supported in part by the U.S. Department of Energy, Office of Science, Office of Nuclear Physics under grant Contract Number DE-SC0011090. This research used resources of the Argonne Leadership Computing Facility, which is a DOE Office of Science User Facility supported under Contract DE-AC02-06CH11357, under the ALCF Aurora Early Science Program. Some work was undertaken at the Kavli Institute for Theoretical Physics, supported by the National Science Foundation under Grant No. NSF PHY-1748958. PES is supported by the National Science Foundation under CAREER Award 1841699, GK is supported by the U.S. Department of Energy under the SciDAC4 award DE-SC0018121, and PES and MSA are supported in part by NSERC and the Perimeter Institute for Theoretical Physics. Research at the Perimeter Institute is supported by the Government of Canada through the Department of Innovation, Science and Economic Development and by the Province of Ontario through the Ministry of Research and Innovation.

Appendix A Acceptance rate estimator for autocorrelation

Here it is shown that the standard estimator for the autocorrelation of an observable $\mathcal{O}$ converges in the limit of infinite path length to $p_{\tau\mathrm{rej}}$, as claimed in Section II.3. The standard estimator is defined by:

[TABLE]

where the final equality is true assuming the Markov chain is initialized with a sample from the stationary distribution $p(\phi)$ (i.e., it is assumed that enough prior iterations were discarded that the chain is thermalized). The expectation value can then be split by cases and, conditioning on the fixed accept/reject pattern, the expectation values can be computed by identifying the distributions of observables $\mathcal{O}_{i}$ and $\mathcal{O}_{i+\tau}$:

[TABLE]

The limit $N\rightarrow\infty$ is used to drop biases arising from conditioning on behavior within the region $[i,i+\tau]$.

Appendix B Mean absolute error loss

The mean absolute error (MAE) loss optimized before training some models is defined by:

[TABLE]

It is bounded below by the KL divergence $D_{KL}(\tilde{p}_{f}||p)$ and has global minima exactly where the KL loss does, when $\tilde{p}_{f}=p$.

In practice, the loss is stochastically estimated by drawing batches of $M$ samples from the model $\left\{\phi^{(i)}\sim\tilde{p}_{f}\right\}$ and computing the sample mean:

[TABLE]

To employ this loss, the partition function must either be estimated ahead of time, or initialized as a trainable parameter. In this study, a multistage method Valleau and Card (1972) was used to estimate and fix the partition function value used while optimizing $L_{\text{MAE}}$.

This loss is appealing due to the point-by-point potential driving the distribution towards the correct one. Any errors in computing $\log{Z}$, however, result in a minimum at which the model distribution $\tilde{p}_{f}(\phi)$ does not necessarily agree with the desired distribution $p(\phi)$. This loss was therefore only used prior to training with the shifted KL loss.