A Near-Optimal Gradient Flow for Learning Neural Energy-Based Models

TL;DR

This paper introduces a new Wasserstein gradient flow-based numerical scheme for training neural energy-based models, improving the approximation of data distributions and generating high-quality high-dimensional data.

Contribution

It proposes a second-order Wasserstein gradient flow approach for EBMs, addressing limitations of previous methods and enhancing distribution fitting and data generation.

Findings

Outperforms existing schemes in fitting complex distributions

Produces higher quality high-dimensional data

Demonstrates practical superiority through extensive experiments

Abstract

In this paper, we propose a novel numerical scheme to optimize the gradient flows for learning energy-based models (EBMs). From a perspective of physical simulation, we redefine the problem of approximating the gradient flow utilizing optimal transport (i.e. Wasserstein) metric. In EBMs, the learning process of stepwise sampling and estimating data distribution performs the functional gradient of minimizing the global relative entropy between the current and target real distribution, which can be treated as dynamic particles moving from disorder to target manifold. Previous learning schemes mainly minimize the entropy concerning the consecutive time KL divergence in each learning step. However, they are prone to being stuck in the local KL divergence by projecting non-smooth information within smooth manifold, which is against the optimal transport principle. To solve this problem, we derive a second-order Wasserstein gradient flow of the global relative entropy from Fokker-Planck equation. Compared with existing schemes, Wasserstein gradient flow is a smoother and near-optimal numerical scheme to approximate real data densities. We also derive this near-proximal scheme and provide its numerical computation equations. Our extensive experiments demonstrate the practical superiority and potentials of our proposed scheme on fitting complex distributions and generating high-quality, high-dimensional data with neural EBMs.