EBM-WGF: Training energy-based models with Wasserstein gradient flow

Area of Science:

• Machine Learning
• Artificial Intelligence
• Computational Statistics

Background:

• Energy-based models (EBMs) are effective for density estimation but face computational challenges with Markov Chain Monte Carlo (MCMC) sampling.
• Existing EBMs utilizing minimax games improve efficiency but suffer from unstable energy estimation and generator optimization.
• Instability in minimax EBMs stems from inaccurate KL divergence minimization via standard gradient flow.

Purpose of the Study:

• To improve the stability and efficiency of energy-based models for density estimation.
• To address the computational expense of MCMC sampling in traditional EBMs.
• To overcome the instability issues in minimax game-based EBMs.

Main Methods:

• Leveraging Wasserstein gradient flow (WGF) to correct generator optimization direction in minimax games.
• Pulling back WGF to parameter space and solving with a variational scheme for bounded solution error.
• Developing a novel EBM incorporating WGF, avoiding MCMC sampling.

Main Results:

• The proposed EBM with WGF demonstrates enhanced stability in generator optimization.
• The WGF solution in this approach is equivalent to Langevin dynamics used in MCMC sampling, offering computational benefits.
• Empirical validation on toy and natural datasets confirms the approach's effectiveness.

Conclusions:

• The novel EBM with WGF provides a stable and computationally efficient alternative to traditional methods.
• This approach effectively resolves the instability issues inherent in minimax game-based EBMs.
• The method offers a promising direction for advancing density estimation techniques in machine learning.