Balanced Training of Energy-Based Models with Adaptive Flow Sampling

TL;DR

The paper tackles the intractable normalization of energy-based models (EBMs) and multimodal density estimation by coupling EBMs with normalizing flows (NFs) into a jointly trained framework. By tilting the EBM with a base NF density and employing FlowMC for calibrated, independent-flow proposals, the authors achieve accurate gradient estimates and fast, multimodal sampling throughout training. Empirical results on 2D multimodal tasks, high-dimensional Gaussian mixtures, and CIFAR-10 show improved mode-weights accuracy and chain mixing compared to ULA-based methods and related approaches, though CIFAR-10 remains challenging with room for flow and energy improvements. Overall, the NF-assisted, jointly trained EBM framework addresses key challenges of statistical density accuracy, sampling efficiency, and mode coverage, enabling more reliable evaluation of relative mode weights in multimodal data.

Abstract

Energy-based models (EBMs) are versatile density estimation models that directly parameterize an unnormalized log density. Although very flexible, EBMs lack a specified normalization constant of the model, making the likelihood of the model computationally intractable. Several approximate samplers and variational inference techniques have been proposed to estimate the likelihood gradients for training. These techniques have shown promising results in generating samples, but little attention has been paid to the statistical accuracy of the estimated density, such as determining the relative importance of different classes in a dataset. In this work, we propose a new maximum likelihood training algorithm for EBMs that uses a different type of generative model, normalizing flows (NF), which have recently been proposed to facilitate sampling. Our method fits an NF to an EBM during training so that an NF-assisted sampling scheme provides an accurate gradient for the EBMs at all times, ultimately leading to a fast sampler for generating new data.

Figures (8)

• Figure 1: Comparison of PCD and flowMC EBM training on toy 2d mixture with 2:1 weight ratio. The EBM learned with flowMC (f) captured correctly the relative weights, unlike the EBM trained with ULA (b), as is also clear from the conditional densities along the axis going trough the centroids (d,h). This statistical accuracy is promoted by the fast mixing NF-assisted sampling (g). (Right) Estimated weight of the top right mode during the training of ULA-EBM and flowMC-EBM (Details in \ref{['app:motivating_example']})
• Figure 2: Estimated energies using different algorithms on toy 2D experiments - (Top) 8 Gaussians (Bottom) Rings
• Figure 3: Estimated energies using different algorithms on the extended motivating example.
• Figure 4: Space partitioning for histograms - the level lines corresponds to $p^{\star}$ and the four colors correspond to each zone.
• Figure 5: Four different negative chains viewed in the latent space of NT-EBM trained on the extended motivating example