PaperlayerView exact revision on arXiv

Source-linked AI summary

Equilibrium Matching: Generative Modeling with Implicit Energy-Based Models

Runqian Wang, Yilun Du

arXiv:2510.02300v3cs.LGcs.AIcs.CV

TL;DR

Generative models based on diffusion and flow use time-conditional, non-equilibrium dynamics, while existing equilibrium approaches face quality or training limitations. Equilibrium Matching learns a time-invariant gradient over an implicit energy landscape and samples by optimization, achieving 1.90 FID on ImageNet 256×256 while supporting broader capabilities including denoising, OOD detection, and composition.

  • Problem

    Existing equilibrium approaches suffer from training instability, poor sample quality, or separate training stages, while removing time conditioning from diffusion can worsen generation quality.

  • Method

    Equilibrium Matching learns a time-invariant equilibrium gradient compatible with an implicit energy landscape and uses gradient-descent-based sampling with adaptive inference options.

  • Results

    Equilibrium Matching achieves 1.90 FID on ImageNet 256×256, outperforming diffusion and flow-based counterparts, while also supporting denoising, OOD detection, and image composition.

  • Takeaways & Limitations

    EqM combines energy-based and flow-based perspectives in a framework that supports interpretable landscapes and flexible optimization-driven inference.

Abstract

from arXiv · show

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.

1 INTRODUCTION

Equilibrium Matching replaces the time-conditional, non-equilibrium dynamics of diffusion and flow models with a single equilibrium gradient over an implicit energy landscape. It enables optimization-based sampling and achieves strong generation quality alongside denoising, OOD detection, and composition capabilities.

  • 1 INTRODUCTION: EqM learns a single time-invariant equilibrium gradient over an implicit energy landscape instead of time-conditional dynamics.The framework targets equilibrium behavior while addressing limitations of existing equilibrium approaches.
  • 1 INTRODUCTION: Gradient descent on the learned landscape supports adjustable step sizes, adaptive optimizers, and adaptive inference-time compute.Adaptive sampling can reduce function evaluations by up to 60%.
  • 1 INTRODUCTION: 1.90 FID on ImageNet 256×256, outperforming diffusion and flow-based counterparts in generation quality.Equilibrium Matching also exceeds its flow-based counterpart at all tested scales.
  • 1 INTRODUCTION: EqM extends beyond generation to partially noised image denoising, module-free OOD detection, and compositional image generation.These capabilities are presented as properties not found in traditional diffusion or flow models.

2 PRELIMINARIES: FLOW MATCHING

Flow Matching learns a time-dependent velocity field along a prescribed path from Gaussian noise to image samples. Removing timestep conditioning yields an equilibrium-oriented variant, but degrades generation quality without adding distinct capabilities.

  • 2 PRELIMINARIES: FLOW MATCHING: Flow Matching predicts conditional velocity along a linear noise-to-image path and integrates that velocity over a unit time horizon.The model receives both the corrupted sample and its timestep or noise level.
  • 2 PRELIMINARIES: FLOW MATCHING: Removing timestep conditioning makes the model noise-unconditional while leaving the original training and sampling procedures otherwise unchanged.The resulting objective differs only through the absence of t from the model input.
  • 2 PRELIMINARIES: FLOW MATCHING: The noise-unconditional variant degrades generation quality and does not exhibit properties distinct from the original Flow Matching model.This motivates a different construction for learning equilibrium dynamics.

3 EQUILIBRIUM MATCHING

Equilibrium Matching trains a time-invariant gradient field whose magnitude vanishes on the data manifold, creating an implicit energy landscape. It samples by optimization and can also learn explicit energies, use accelerated or adaptive procedures, and support related differential-equation samplers.

  • 3 EQUILIBRIUM MATCHING: EqM learns a time-invariant gradient compatible with an energy function, with zero gradient at data samples and increasing gradient toward noise.Ground-truth samples become stationary points of the learned equilibrium landscape.
  • 3 EQUILIBRIUM MATCHING: EqM training corrupts images with implicit interpolation factors and fits the model to target gradients, while inference iteratively optimizes predicted gradients.The provided procedure combines gradient descent, Nesterov acceleration, and adaptive compute.
  • 3.1 TRAINING: The target gradient points from noise toward data and uses c(1) = 0 so its magnitude vanishes at the data manifold.When c(γ) = 1, the EqM objective is the negation of Flow Matching’s objective.
  • 3.1 TRAINING: EqM supports linear, truncated, piecewise, and multiplier-controlled gradient magnitude schedules for constructing the equilibrium landscape.Truncated decay keeps the gradient constant far from data before linearly reducing it toward zero.
  • 3.2 LEARNING EXPLICIT ENERGY: Explicit-energy EqM matches the gradient of a scalar energy model and assigns low energy to real samples and high energy to noise.The scalar energy can be constructed from an existing EqM model through a dot product or squared L2 norm.
  • 3.3 SAMPLING: Sampling is formulated as gradient descent, allowing adaptive step sizes, optimization methods such as Nesterov acceleration, integration-based samplers, and per-sample stopping.Adaptive compute terminates when the gradient norm falls below a threshold.

4 ANALYSIS

The analysis argues that EqM learns ground-truth samples as local minima and that gradient-based sampling converges under stated assumptions. These results provide a theoretical basis for generating samples close to the data manifold.

  • Under perfect training and high-dimensional assumptions, ground-truth samples have approximately zero learned gradient, while local minima with zero model output are approximately ground-truth samples.Together, these statements identify the learned local minima with samples from the ground-truth dataset.
  • Equilibrium Matching therefore learns ground-truth images as local minima during training.
  • With an L-smooth energy bounded below, gradient descent sampling has a convergence rate of O(1/N) under the stated step-size condition.N denotes the total number of sampling steps, while the convergence statement is formalized for iterates after K optimization steps.
  • Under the given assumptions, the analysis concludes that Equilibrium Matching produces samples close to ground truths before empirical validation.

5 EXPERIMENTS

Experiments evaluate EqM on ImageNet generation, scaling, sampling, ablations, and properties beyond standard generation. EqM achieves strong generation quality while supporting flexible samplers, adaptive compute, generalization, denoising, OOD detection, and composition.

  • 5.1 IMAGE GENERATION: EqM scales across training length, model size, and patch size, consistently outperforming Flow Matching under all tested configurations.The scalability evaluation varies each axis while controlling the others as specified in the experimental setup.
  • 5.1 IMAGE GENERATION: 1.90 FID on class-conditional ImageNet 256×256 surpasses tested diffusion and flow-based methods, while EqM consistently exceeds Flow Matching during training.Table 1 reports the headline ImageNet result; training curves further show a persistent advantage over the Flow Matching baseline.
  • 5.3 ABLATION STUDY: Ablations show that NAG-GD improves across tested step counts, the truncated-decay gradient uses a = 0.8 and λ = 4 by default, and the dot-product energy variant is preferred.The L2 norm energy formulation degrades performance and is sensitive to initialization; the target gradient helps only in the noise-unconditional setting.
  • 5.2 INFERENCE-TIME EXPERIMENTS: EqM-generated samples differ from their nearest training-set neighbors, indicating generalization beyond memorizing training images.The nearest neighbors are measured by mean squared distance and shown as top-3 matches.
  • 5.2 INFERENCE-TIME EXPERIMENTS: EqM supports faster and more flexible inference: NAG-GD improves FID, EqM converges earlier, and adaptive compute reduces total computation to 40% of fixed sampling.EqM also remains effective across tested step sizes, whereas Flow Matching requires a specific step size.
  • 5 EXPERIMENTS: EqM supports partially noised-image denoising, energy-based OOD detection, and composition by adding model gradients for multiple conditions.These capabilities extend the framework beyond ordinary generation and include composition examples from paired ImageNet class conditions.

6 RELATED WORK

Diffusion and flow models use time-conditioned non-equilibrium dynamics, while energy-based models learn equilibrium landscapes but face instability and sample-quality problems. Recent hybrid efforts have not fully resolved these limitations.

  • Diffusion models generate samples through noise-conditioned noising and denoising dynamics, while flow models formulate sampling as differential-equation integration.
  • Energy-based models learn equilibrium energy landscapes and support tasks such as OOD detection, but often suffer from training instabilities and poor sample quality.
  • Removing noise conditioning from diffusion models worsens generation quality, while Energy Matching combines flow and Langevin-based training in two stages.

7 CONCLUSION

The paper proposes Equilibrium Matching as a simple generative model that learns equilibrium dynamics. It combines energy-based and flow-based advantages while supporting interpretable, flexible sampling.

  • Equilibrium Matching learns equilibrium dynamics in a simple and effective generative modeling framework.
  • The method combines energy-based and flow-based advantages, achieves strong generation quality, and provides an interpretable energy landscape for diverse sampling methods.
  • The authors argue that EqM’s equilibrium dynamics could inspire more effective and scalable inference algorithms.

A EXPERIMENTAL SETTING

The experimental setting reports the training configuration and sampler choices used for Equilibrium Matching and comparison models. Sampler selection varies across the paper’s experiments.

  • The training setting for Equilibrium Matching is presented in Table 7.
  • The experiments use NAG-GD for EqM in ImageNet generation, while other sections use vanilla gradient descent for EqM and Euler sampling for SiT.

B.1 CIFAR-10 EXPERIMENTS

On CIFAR-10 with U-Net architectures, Equilibrium Matching improves over noise-unconditional Flow Matching but does not surpass standard Flow Matching. Additional metrics show relatively good sFID and Inception Score.

  • Equilibrium Matching outperforms the noise-unconditional Flow Matching baseline on CIFAR-10 but fails to improve over standard Flow Matching.The authors attribute the gap to extensive optimization of the standard baseline’s noise and sampling schedules, which they describe as an unfair comparison.
  • Equilibrium Matching achieves relatively good sFID and Inception Score compared with other generative methods.

C.1 LEARNED GRADIENT AT GROUND-TRUTH SAMPLES

Under perfect training, Equilibrium Matching assigns approximately zero gradient to ground-truth samples and makes its local minima approximately coincide with training data. Its gradient-based sampler also has a standard convergence guarantee and recovers explicit Euler sampling as a special case.

  • C.1 LEARNED GRADIENT AT GROUND-TRUTH SAMPLES: Ground-truth images receive approximately zero gradient under perfect training in high-dimensional settings.
  • C.1 LEARNED GRADIENT AT GROUND-TRUTH SAMPLES: The proof uses exponentially vanishing Gaussian density at exact training points for γ < 1, while γ = 1 retains mass on the finite dataset.
  • C.2 PROPERTY OF LOCAL MINIMA: All local minima are approximately samples from the ground-truth dataset under perfect training in high-dimensional settings.
  • C.3 CONVERGENCE OF GRADIENT-BASED SAMPLING: With an L-smooth energy bounded below, gradient descent satisfies a convergence bound on the minimum squared gradient norm across its iterates.
  • D.1 ODE SAMPLING AS A SPECIAL CASE OF OPTIMIZATION: When the velocity is conservative, explicit Euler sampling exactly matches gradient descent with step size η = 1/N on a unit time horizon.

D.2 GENERAL INTEGRATION SAMPLERS

Viewing a conservative velocity as an energy descent direction turns integration-based samplers into optimization updates. This permits adaptive step sizes without changing the underlying direction field.

  • D.2 GENERAL INTEGRATION SAMPLERS: Any time integrator induces an optimization update when the velocity is the negative gradient of an energy.
  • D.2 GENERAL INTEGRATION SAMPLERS: Adaptive step sizes can replace the uniform ODE grid while retaining the same direction field and adapting existing integration samplers to Equilibrium Matching.
Loading 2510.02300v3…