Equilibrium Matching: Generative Modeling with Implicit Energy-Based Models
Runqian Wang, Yilun Du
TL;DR
Equilibrium Matching (EqM) introduces a time-invariant equilibrium gradient over an implicit energy landscape as an alternative to non-equilibrium diffusion/flow dynamics. It trains a target gradient via corrupted interpolation x_γ and a decaying function c(γ), enabling gradient-descent sampling with adaptive step sizes and optimizers, and supports explicit energy variants. The authors prove that EqM learns the data manifold and that gradient-based sampling converges, while empirically achieving state-of-the-art-like results on ImageNet 256×256 (FID 1.90) and displaying strong scalability and flexible inference-time capabilities, including partial denoising, OOD detection, and compositional generation. By bridging flow-based and energy-based modeling, EqM offers a simple, optimization-driven inference paradigm with broad potential for scalable high-fidelity image generation.
Abstract
We introduce Equilibrium Matching (EqM), a generative modeling framework built from an equilibrium dynamics perspective. EqM discards the non-equilibrium, time-conditional dynamics in traditional diffusion and flow-based generative models and instead learns the equilibrium gradient of an implicit energy landscape. Through this approach, we can adopt an optimization-based sampling process at inference time, where samples are obtained by gradient descent on the learned landscape with adjustable step sizes, adaptive optimizers, and adaptive compute. EqM surpasses the generation performance of diffusion/flow models empirically, achieving an FID of 1.90 on ImageNet 256$\times$256. EqM is also theoretically justified to learn and sample from the data manifold. Beyond generation, EqM is a flexible framework that naturally handles tasks including partially noised image denoising, OOD detection, and image composition. By replacing time-conditional velocities with a unified equilibrium landscape, EqM offers a tighter bridge between flow and energy-based models and a simple route to optimization-driven inference.
