Gradient estimators for normalising flows

TL;DR

This paper introduces a new gradient estimator for training normalizing flows in neural MCMC, achieving faster convergence and more accurate free energy estimates for the $\

Contribution

It presents a novel gradient estimator that reduces variance and speeds up training of normalizing flows in neural MCMC methods, especially for complex models.

Findings

Achieves same precision in half the time compared to standard methods.

Provides better free energy estimates for the $\

Demonstrates lower variance of the new estimator improves training efficiency.

Abstract

Recently a machine learning approach to Monte-Carlo simulations called Neural Markov Chain Monte-Carlo (NMCMC) is gaining traction. In its most popular form it uses neural networks to construct normalizing flows which are then trained to approximate the desired target distribution. In this contribution we present new gradient estimator for Stochastic Gradient Descent algorithm (and the corresponding \texttt{PyTorch} implementation) and show that it leads to better training results for $\phi^4$ model. For this model our estimator achieves the same precision in approximately half of the time needed in standard approach and ultimately provides better estimates of the free energy. We attribute this effect to the lower variance of the new estimator. In contrary to the standard learning algorithm our approach does not require estimation of the action gradient with respect to the fields, thus has potential of further speeding up the training for models with more complicated actions.