Multi-Lattice Sampling of Quantum Field Theories via Neural Operator-based Flows

Authors: bálint  máté
¸ois  franc
  fleuret
Venue: ICLR_cc_2024_Workshop_AI4DiffEqtnsInSci
Published: 2024-01-13 16:48:20

Abstract: We consider the problem of sampling discrete field configurations $\phi$ from the Boltzmann distribution $[d\phi] Z_1^{-1} e^{-S_1[\phi]}$, where $S_1$ is the  lattice-discretization of the continuous Euclidean action $\mathcal S_1$ of some quantum field theory. Since such densities arise as the approximation of the underlying functional density $[\mathcal D\phi(x)] \mathcal Z_1^{-1} e^{-\mathcal S_1[\phi(x)]}$, we frame the task as an instance of operator learning. In particular, we propose to approximate a time-dependent operator $\mathcal V_t$ whose time integral provides a mapping between the functional distributions of the free theory $[\mathcal D\phi(x)] \mathcal Z_0^{-1} e^{-\mathcal S_{0}[\phi(x)]}$ and of the target theory $[\mathcal D\phi(x)]\mathcal Z_1^{-1}e^{-\mathcal S_1[\phi(x)]}$. Once a particular lattice  is chosen, the operator $\mathcal V_t$ can be discretized to a finite dimensional, time-dependent vector field $V_t$ which in turn induces a continuous normalizing flow between finite dimensional distributions over the chosen lattice. This flow can then be trained to be a diffeormorphism between the discretized free and target theories $[d\phi] Z_0^{-1} e^{-S_{0}[\phi]}$, $[d\phi] Z_1^{-1}e^{-S_1[\phi]}$. We run experiments on the 2-dimensional $\phi^4$-theory to explore to what extent such operator-based flow architectures generalize to lattice sizes they were not trained on and show that pretraining on smaller lattices can lead to speedup over training directly on the target lattice size.

Main paper: 
 introduction: 
 Consider S as an action characterizing a quantum field theory, with S representing its discretization to a lattice. Albergo et al. (2019) suggest a method for sampling from the lattice quantum field theory described by S. This involves a normalizing flow parameterizing a density function q θ of discrete fields over the lattice, and optimizing the parameters θ until q θ closely approximates the probability density e -S Z , where Z is the normalizing constant of e -S .Operator learning promotes the viewpoint that the lattice/mesh is merely a computational tool, and model should capture the underlying continuous physics. Kovachki et al. (2021) term this property of models discretization invariance. 1 In this work, we apply the same idea to the task of sampling from lattice quantum field theories, motivated by the fact that lattice field theories also emerge as the discretization of continuous field theories.Suppose now that the field theory is defined on some domain D. Once a lattice, as a discretization of D, is chosen, one can construct a continuous normalizing flow Chen et al. ( 2018) as a time-dependent vector field V t that parametrizes the direction along which probability mass moves. Generalizing this idea, we propose to parametrize a time-dependent operator V t from the space of functions on D to itself that defines the direction in which functional probability mass moves. Such an operator can then be used to map the functional distributions [Dϕ(x)]Z -1 0 e -S0[ϕ(x)] , [Dϕ(x)]Z -1 1 e -S1[ϕ(x)] to one another. Computationally the operator V t can only be accessed by a choice of a lattice which induces a vector field V t as the discretization of V t .We then train this vector field to be a diffeomorphism between the discretized free and target theories, [dϕ]Z -1 0 e -S0 [ϕ] and [dϕ]Z -1 1 e -S1 [ϕ] . The upside of using a operator-based flow will be that a single model can be used to operate on multiple discretizations of the same underlying continuous system. Figure 1 provides a schematic overview of the objects and their relation in this paragraph.[Dϕ(x)] Z -1 0 e -S0[ϕ(x)] [Dϕ(x)] Z -1 1 e -S1[ϕ(x)] [dϕ] Z -1 0 e -S0[ϕ] [dϕ] Z -1 1 e -S1[ϕ]V t approximate approximateV t discretizeFigure 1 : Schematic overview of the probability distributions of interest. The top row shows the functional distributions of the free theory and the target theory connected by the time dependent operator V t . Descending to the bottom row means approximating all the objects of the top row on a lattice. In particular, in the bottom row all objects are finite dimensional, well-defined and can be numerically worked with.The structure of the paper is as follows:• Section §2 introduces the relevant background on the ϕ 4 quantum field theory.• In Section §3 we describe an operator-based normalizing flow architecture that we used in our experiments.• Section §4 documents our experiment on the ϕ 4 theory.• Section §A contains the relevant background on continuous normalizing flows, Boltzmann densities and neural operators.2 THE ϕ 4 (LATTICE) QUANTUM FIELD THEORY2 Let us now consider the Euclidean action on real valued scalar fields ϕ(x) with periodic boundary conditions on the D-dimensional hypercube of edge length L, ϕ : (R/LZ) D → R, for some constants m 2 and gS[ϕ] = (R/LZ) D d D x (∇ϕ) 2 + m 2 ϕ 2 + gϕ 4 (1)where we dropped the argument x of the field ϕ(x) for notational convenience. To estimate the expectation value of an observable O, we need to average over all field configurations that satisfy the boundary conditions, with each configuration weighted by the exponential of the negative actionEQUATIONThe action S corresponds to the energy function f of a Boltzmann density and the denominator Z = Dϕ e -S [ϕ] to the normalizing constant as introduced in Appendix §A.1.Equations ( 13) and ( 14) describe an infinite dimensional system. To tackle it numerically, one first needs to discretize it to a lattice. This comes at the cost of losing the information contained in the high-frequency components as the highest possible frequency of a periodic function on a lattice with edge length L with N nodes is 2πN L . The hope is that one can do the same on larger and larger lattices, and as the lattice approaches the continuum limit, the error due to discretization converges to zero.To discretise the action, we consider fields living on the points located at 0 N , L N , ..., (N -1)L N d forming a periodic lattice with cardinality N D and lattice spacing a = L/N . We then turn integrals into sums (R/LZ) D d D x → a D x and differentials into differences between nearest neighbors∂ i ϕ → 1 a ϕ(x + µ i ) -ϕ(x), where µ i is the generator of lattice along the i-th coordinate axis. After these substitutions we end up with the following discretised action on the lattice,S[ϕ] = a D 1 a 2 x,µ (ϕ x+µ -ϕ x ) 2 + x m 2 ϕ 2 x + gϕ 4 x (3)where x runs over the lattice sites and µ over the generators of the lattice. 
flows parametrised by neural operators: 
 Let now ϕ ∈ R N ×....×N be a discretized scalar field on a lattice. The architecture then consists of the following sequence of steps, where the subscript θ denotes trainable parameters, 1. Use a per-node neural network f θ to embed the field values,ϕ emb = f θ (ϕ) ∈ R c×N ×....×N .2. Use a neural network to parametrize c-many continous spherically symmetric kernels K θ (r). Let then Kθ be the evaluation of the continous kernels on the lattice.3. Mask out the origin of the discrete kernel, i.e. set Kθ [:, 0] = 0.4. Perform a the channel-wise convolution ϕ emb ⋆ Kθ and denote the result by C ∈ R c×N ×....×N . Because of the previous step, C i is independent of ϕ i , and we will call it the conditioner (Chen & Duvenaud (2019)). 5. Apply a per-node neural network τ θ to the concatenation (C, ϕ emb ) with output Y = τ θ (C, ϕ emb ) ∈ R T ×N ×....×N . 6. Contract the first dimension of Y with a vector of length T that only depends on time. 7. Finally, denoting all the above steps as i, we set the output of the model to be V (ϕ, t) = 1 2 * (i(ϕ, t) -i(-ϕ, t)). This enforces the Z 2 symmetry of the system.To compute the divergence of the architecture one needs the Jacobians of the per-point operations f θ and τ θ , K θ does not have to be differentiated through.The normalizing flow architecture described in Section §A.1 requires an initial density from which samples can easily be drawn. Instead of sampling from a standard gaussian at every node, we choose a more physical initial density by setting g = 0 in the action ( 17). This results in the free theory with a gaussian Boltzmann density that becomes diagonal in momentum space. Position and momentum space are related by a discrete Fourier transformEQUATIONwhere p runs over -⌊N/2⌋ L , ...., 0 L , ..., ⌈(N -1)/2⌉ L D and the prefactor 1 √ N D makes the map {ϕ x } ↔ { φp } unitary. The covariance matrix of the free theory is diagonalized in the momentum basis 1 √ N D e 2πi⟨x,p⟩ with eigenvaluesEQUATIONTo constrain the sampling to real valued fields, we sample p from the hermitian symmetric subspace of real dimension N D of the Fourier-space of complex dimension N D . 
experiments: 
 One of the core dilemmas when choosing the lattice size lies in balancing the cost efficiency of a smaller lattice against the better approximation of the underlying continuous system of a larger one.Ideally, one would train on a small lattice and evaluate on a large one, but it is unreasonable to expect this approach to work well as the smaller lattices cannot capture the higher frequency components of the system. We thus explore the next best strategy: extensive pretraining on a small lattice followed by transferring the model to a larger lattice size and fine-tuning it with a reduced number of steps.We now work in D = 1 dimensions. Strictly speaking, a one dimensional lattice does not correspond to a quantum field theory, rather it describes the trajectory of a quantum mechanical particle in a potential. Nonetheless, it's the simplest setup in which we can experiment and serves as a good starting point. We also fix L = 4, m 2 = -4, g = 1 and train a single model for 5000 steps with mesh size uniformly sampled at each training step from N = L/a ∈ {4, 8, 16, ..., 128}. We then evaluate performance on lattices of size N = L/a up to 512 by sampling from the trained model to calculate the effective sample sizeESS = 1 N i w i 2 1 N i w 2 i (6)where w i is the importance weight p(ϕ i )/q θ (ϕ i ). We also estimate the two-point correlation function G(x, y)[ϕ] := ϕ(x)ϕ(y). In this experiment we consider the target D = 2, L = 12, m 2 = -4, g = 5.276, N = 64. Instead of training directly on the N = 64 lattice, we pretrain on a sequence of smaller lattices as they are significantly cheaper to work on. We start training on a 12 × 12 lattice for 2000 steps, after which we train on lattices of size 16 2 , 20 2 , 24 2 , 28 2 , 32 2 , 36 2 , 40 2 , 44 2 , 48 2 , 52 2 , 56 2 , 60 2 for 250 training steps each. Finally, we train on the target size 64 × 64 for 1000 steps. As a baseline, we also train the same architecture only on the target size for the same total number of steps (6000).While the performance, as measured by the effective sample size on target lattice, is comparable after training (Table 1 ), the training procedure that "trained through" the smaller lattices was ∼ 2.4times quicker to train(Figure 3 ). 2018) is a density estimator that operates by pushing forward a simple, usually Gaussian, initial density q 0 along a parametric, time-dependent vector fieldV θ : [0, 1] × R n → R n .Explicitly, the pushforward density q θ is given byEQUATION)where ∇ is the divergence operator in the spatial coordinates and x t is the integral curve of V θ that passes through x 1 at t = 1. In this work all normalizing flows will be continuous normalizing flows, and we will refer to them as normalizing flows or even just flows for brevity.The Boltzmann distribution of an energy function3 f : R n → R is a probability distribution with density functionEQUATIONwhere Z = dx e -f (x) is the normalizing constant ensuring that the density function integrates to 1. Boltzmann distributions appear in the context of the canonical ensemble, a statistical ensemble that describes a system in thermal equilibrium with an external heat reservoir. Such Boltzmann distributions describe the molecular systems in thermal equilibrium as well as Wick-rotated quantum field theories. Learning to sample from Boltzmann distributions using only the energy function (i.e. KL[q θ , p] = E x∼q θ [log q θ (x) -log p(x)] (10) = E x∼q θ [log q θ (x) + f (x)] + Z (11)where q θ is the density realized by the normalizing flow (8). Once a density q θ , approximating p = Z -1 e -f , is learnt, one can use importance sampling to correct for small inaccuracies of q θ when estimating the expected value of an observableEQUATIONA.2 THE ϕ4 (LATTICE) QUANTUM FIELD THEORY 4Let us now consider the Euclidean action on real valued scalar fields ϕ(x) with periodic boundary conditions on the D-dimensional hypercube of edge length L, ϕ : (R/LZ) D → R, for some constants m 2 and gS[ϕ] = (R/LZ) D d D x (∇ϕ) 2 + m 2 ϕ 2 + gϕ 4 (13)where we dropped the argument x of the field ϕ(x) for notational convenience. To estimate the expectation value of an observable O, we need to average over all field configurations that satisfy the boundary conditions, with each configuration weighted by the exponential of the negative actionEQUATIONThe action S corresponds to the energy function f of a Boltzmann density and the denominator [ϕ] to the normalizing constant as introduced in Section §A.1.Z = Dϕ e -SEquations ( 13) and ( 14) describe an infinite dimensional system. To tackle it numerically, one first needs to discretize it to a lattice. This comes at the cost of losing the information contained in the high-frequency components as the highest possible frequency of a periodic function on a lattice with edge length L with N nodes is 2πN L . The hope is that one can do the same on larger and larger lattices, and as the lattice approaches the continuum limit, the error due to discretization converges to zero.To discretise the action, we consider fields living on the points located at 0 N , L N , ..., (N -1)L N d forming a periodic lattice with cardinality N D and lattice spacing a = L/N . We then turn integrals into sums and differentials into differences between nearest neighborsEQUATIONEQUATION)where µ i denotes the generator of the lattice along the i-th coordinate axis. After these substitutions we end up with the following discretised action on the lattice,EQUATION)where x runs over the lattice sites and µ over the generators of the lattice. It is customary to absorb all the occurrences of a in the above formula by rescaling ϕEQUATIONThis results in an alternative form of the actionEQUATION)While the action (19) has the advantage of not being dependent on the lattice spacing a, we will continue working with (17) keeping the relation between different lattice sizes and to the underlying continuous setting explicit.Neural Operators Kovachki et al. (2021) are trainable function-to-function mappings. In particular, both their domains and codomains are infinite dimensional function spaces. In practice, one works with neural operators by choosing a mesh/lattice X ⊂ R n , representing functions by their evaluations on X and let the neural operator operate on this collection of evaluations. By design, neural operators can be evaluated on lattices of different size. Importantly, if a neural operator is applied to a sequence of meshes X i , approaching the continuum limit X i → R n , it converges to the underlying continuous operator. The main use case of neural operators is to approximate the solution of partial differential equations, i.e. learn the mapping from an initial condition to the time evolved state after some time ∆t (Figure 5 ), but have also been applied for multi-resolution generative modelling Voleti et al. ( 2021); Hagemann et al. (2023) . We will use them for parametrizing a flow, i.e. a vector field V t connecting the free theory (base density) to the ϕ 4 -theory (target density) in a way that can be evaluated at any mesh. Figure 5 : An operator that maps an initial condition u(0, x) (top row) to its time-evolved state u(∆t, x) (bottom row). The time evolution is given by the heat equation ∆u = ∂ t u. The blue dots denote the evaluation of u(0, x) on a discrete mesh, while the orange dots denote the output of the operator (a convolution in this case) evaluated on that same mesh. As the mesh gets denser, the operator becomes a better approximation of the map between the continuous u(0, x) (blue curve) and u(∆t, x) (orange curve). Discretization invariance means that the neural operator evaluated on finer and finer discretizations approximates the continuous operator. Thus, strictly speaking, it is not a requirement of invariance rather that of convergence.We recommend the book(Thijssen, 2007, Chapter 15) for further details on lattice field theories.ACKNOWLEDGEMENTThe authors were supported by the Swiss National Science Foundation under grant number CR-SII5 193716 -"Robust Deep Density Models for High-Energy Particle Physics and Solar Flare Analysis (RODEM)". We thank Samuel Klein for discussions.Assuming that exp(-f ) is integrable.We recommend the book(Thijssen, 2007, Chapter 15) for further details on lattice field theories. 
 
 Appendix:  


Decision: The work studies the application of neural operators for Quantum Field theory modeling. As mentioned by the reviewer, I also believe the experiments section needs clarification and better presentation. The authors are highly encouraged to do that for their final version. given they will improve their experiment presentation, I would be willing to accept this as poster.

Reviews: review: title: title: The paper proposes the use of operator learning to generalize to different lattice sizes corresponding to a given quantum field theory.
main_review: main_review: The paper proposes to use the idea of resolution invariant learning of operators in quantum field theory, thereby enabling the use of a single model on multiple discretization (lattice) of the same underlying system. 

While, the idea has been explored in many different use cases, I am unaware of it's specific use in quantum field theory.

The description of quantum field theory is adequate, given this might be new to readers of a ML venue.

However, I think the experimental section is very poorly written, and it is not easy to understand what the paper is trying to achieve in one go. I do suggest, more details on what the core experimental section is be added, plots can be moved out if needed. Without a clear message in the experimental section, the efforts made for the paper are futile. 

The experimental results are somewhat convincing.