Efficient Optimization of Variational Autoregressive Networks with Natural Gradient

Abstract

Estimating free energy is a fundamental problem in statistical mechanics. Recently, machine-learning-based methods, particularly the variational autoregressive networks (VANs) have been proposed to minimize variational free energy and to approximate the Boltzmann distribution. VAN enjoys notable advantages, including the exact computation of the normalized joint distribution and fast sampling, which are critical features often missing in Markov chain Monte Carlo algorithms. However, VAN also faces significant computational challenges. These include difficulties in the optimization of variational free energy in a complicated parameter space and slow convergence of learning. In this work, we introduce an optimization technique based on natural gradients to the VAN framework, namely ng-VAN, to enhance the learning efficiency and accuracy of the conventional VAN. The method has computational complexity cubic in the batch size rather than in the number of model parameters, hence it can be efficiently implemented for a large VAN model. We carried out extensive numerical experiments on the Sherrington-Kirkpatrick model, spin glasses on random graphs, and the two-dimensional Ising model. Our results indicate that compared with the conventional VAN, ng-VAN significantly improves the accuracy in estimating free energy and converges much faster with shorter learning time. This allows extending the VAN framework's applicability to challenging statistical mechanics problems that were previously not accessible.