Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Training normalizing flows with computationally intensive target probability distributions

Piotr Bialas, Piotr Korcyl, Tomasz Stebel

TL;DR

The paper tackles the challenge of training normalizing flows when the target distribution involves computationally intensive gradients, such as lattice field theories with fermionic determinants. It adapts a REINFORCE-based gradient estimator to avoid differentiating the action, and compares it to the standard reparameterization trick on the 2D Schwinger model with Wilson fermions. The RE estimator delivers substantial practical benefits, including up to ~10x faster wall-clock times, reduced memory usage (up to ~30%), and robust single-precision performance, primarily due to avoiding backpropagation through the determinant. The study also provides a detailed breakdown of the computational graph and identifies the determinant assembly as a key bottleneck for RT, offering insights for extending the approach to other models with expensive target probabilities. Overall, the RE approach shows strong potential to accelerate NMCMC training in fermionic lattice theories and other domains with costly target gradients.

Abstract

Machine learning techniques, in particular the so-called normalizing flows, are becoming increasingly popular in the context of Monte Carlo simulations as they can effectively approximate target probability distributions. In the case of lattice field theories (LFT) the target distribution is given by the exponential of the action. The common loss function's gradient estimator based on the "reparametrization trick" requires the calculation of the derivative of the action with respect to the fields. This can present a significant computational cost for complicated, non-local actions like e.g. fermionic action in QCD. In this contribution, we propose an estimator for normalizing flows based on the REINFORCE algorithm that avoids this issue. We apply it to two dimensional Schwinger model with Wilson fermions at criticality and show that it is up to ten times faster in terms of the wall-clock time as well as requiring up to $30\%$ less memory than the reparameterization trick estimator. It is also more numerically stable allowing for single precision calculations and the use of half-float tensor cores. We present an in-depth analysis of the origins of those improvements. We believe that these benefits will appear also outside the realm of the LFT, in each case where the target probability distribution is computationally intensive.

Training normalizing flows with computationally intensive target probability distributions

TL;DR

The paper tackles the challenge of training normalizing flows when the target distribution involves computationally intensive gradients, such as lattice field theories with fermionic determinants. It adapts a REINFORCE-based gradient estimator to avoid differentiating the action, and compares it to the standard reparameterization trick on the 2D Schwinger model with Wilson fermions. The RE estimator delivers substantial practical benefits, including up to ~10x faster wall-clock times, reduced memory usage (up to ~30%), and robust single-precision performance, primarily due to avoiding backpropagation through the determinant. The study also provides a detailed breakdown of the computational graph and identifies the determinant assembly as a key bottleneck for RT, offering insights for extending the approach to other models with expensive target probabilities. Overall, the RE approach shows strong potential to accelerate NMCMC training in fermionic lattice theories and other domains with costly target gradients.

Abstract

Machine learning techniques, in particular the so-called normalizing flows, are becoming increasingly popular in the context of Monte Carlo simulations as they can effectively approximate target probability distributions. In the case of lattice field theories (LFT) the target distribution is given by the exponential of the action. The common loss function's gradient estimator based on the "reparametrization trick" requires the calculation of the derivative of the action with respect to the fields. This can present a significant computational cost for complicated, non-local actions like e.g. fermionic action in QCD. In this contribution, we propose an estimator for normalizing flows based on the REINFORCE algorithm that avoids this issue. We apply it to two dimensional Schwinger model with Wilson fermions at criticality and show that it is up to ten times faster in terms of the wall-clock time as well as requiring up to less memory than the reparameterization trick estimator. It is also more numerically stable allowing for single precision calculations and the use of half-float tensor cores. We present an in-depth analysis of the origins of those improvements. We believe that these benefits will appear also outside the realm of the LFT, in each case where the target probability distribution is computationally intensive.
Paper Structure (9 sections, 37 equations, 5 figures, 6 tables, 1 algorithm)

This paper contains 9 sections, 37 equations, 5 figures, 6 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Neural Markov Chain Monte Carlo
  3. Gradient estimators
  4. Schwinger model
  5. Results
  6. Timing
  7. Memory
  8. Summary
  9. Implementation

Figures (5)

  • Figure 1: Schematic picture of two algorithms for gradient estimation discussed in the paper: a) reparametrization trick b) REINFORCE. Double line arrows represent the flow: upward-pointing arrows represent forward propagation, and downward-pointing arrows represent reversed propagation. Dashed arrows denote propagation which does not require gradient calculations.
  • Figure 2: Training history for the Schwinger model on a $16\times 16$ lattice at criticality $\beta=2.0$, $\kappa=0.276$. Each gradient step was calculated on a batch of $3\times512$ samples (the batch was split into three parts to fit on the GPU). Left: the effective sample size (ESS) defined in eq. (\ref{['eq:ess-def']}) as a function of the number of gradient steps for two gradient estimators. Red curves were obtained using REINFORCE estimator with automatic mixed precision (amp) which enabled the use of tensor cores on GPU using half-float precision. We present the history of two different runs for each estimator. Right: the acceptance rate of the MCMC algorithm calculated for a particular state of the network after a given number of gradient steps.
  • Figure 3: An excerpt from the Monte Carlo history for $\sigma$ (upper panel) and chiral condensate (lower) for $L=16$. Series of non-accepted configurations (bridges) that are over 100 in length are marked in red.
  • Figure 4: Training history for the Schwinger model on a $24\times 24$ lattice at criticality $\beta=2.0$, $\kappa=0.276$. Each gradient step was calculated on a batch of $4\times384$ samples (the batch was split into four parts to fit on the GPU). Left: the effective sample size (ESS) defined in eq. (\ref{['eq:ess-def']}) as a function of the number of gradient steps for the REINFORCE gradient estimator. Red curves were obtained using REINFORCE estimator with automatic mixed precision (amp) which enabled the use of tensor cores on GPU using half-float precision. We present the history of two different runs for each estimator. Right: the acceptance rate of the MCMC algorithm calculated for a particular state of the network after a given number of gradient steps.
  • Figure 5: An excerpt from the Monte-Carlo history for $\sigma$ and chiral condensate for $L=24$. Series of non-accepted configurations (bridges) that are over 100 in length are marked in red.