Equilibrium Matching: Generative Modeling with Implicit Energy-Based Models

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.

• The paper introduces Equilibrium Matching, merging energy-based and flow-based methods through a learned equilibrium gradient field over an implicit energy landscape.
• It employs optimization-based sampling with flexible step sizes and adaptive compute, achieving state-of-the-art results on class-conditional ImageNet generation.
• Empirical evaluations highlight improved out-of-distribution detection and compositional generation, offering scalable and robust performance for generative tasks.

Equilibrium Matching: Generative Modeling with Implicit Energy-Based Models

Overview and Motivation

Equilibrium Matching (EqM) introduces a generative modeling paradigm that unifies energy-based and flow-based approaches by learning a time-invariant equilibrium gradient field over an implicit energy landscape. Unlike diffusion and flow models, which rely on non-equilibrium, time-conditional dynamics and require explicit noise or time conditioning, EqM discards these constraints and instead learns a gradient field that is compatible with an underlying energy function. This enables optimization-based sampling at inference, where samples are generated via gradient descent on the learned landscape, supporting flexible step sizes, adaptive optimizers, and adaptive compute allocation.

The EqM framework is motivated by the limitations of existing generative models: diffusion and flow models achieve high sample quality but are restricted by their non-equilibrium design, while energy-based models (EBMs) offer equilibrium dynamics but suffer from training instability and poor sample quality. EqM addresses these issues by constructing a single equilibrium gradient field, theoretically guaranteeing that ground-truth samples are local minima and empirically demonstrating superior generation quality and scalability.

Figure 1: Conceptual comparison of Flow Matching (left) and Equilibrium Matching (right) in 2D. EqM learns a time-invariant gradient field converging to ground-truth data points.

Theoretical Foundations

EqM is formulated by defining a corruption scheme that interpolates between data and noise via a factor $\gamma \in [0,1]$, producing intermediate samples $x_\gamma = \gamma x + (1-\gamma)\epsilon$. The model is trained to predict a target gradient $(\epsilon - x)c(\gamma)$, where $c(\gamma)$ controls the gradient magnitude and is designed to vanish at the data manifold ($c(1) = 0$). This ensures that ground-truth samples are stationary points of the learned energy landscape.

Key theoretical results include:

• Learned Gradient at Ground-Truth Samples: Under perfect training, EqM assigns approximately zero gradient to ground-truth samples, ensuring they are local minima.
• Property of Local Minima: All local minima of the learned landscape correspond to ground-truth data points in high-dimensional settings.
• Convergence of Gradient-Based Sampling: Gradient descent sampling on the EqM landscape converges to the data manifold at a rate of $O(1/N)$, where $N$ is the number of steps.

These results establish that EqM learns a valid energy landscape and supports optimization-driven inference.

Training and Implementation

EqM is implemented by adapting transformer-based backbones (e.g., SiT) and removing time/noise conditioning. The training objective is a mean squared error between the model's predicted gradient and the target gradient $(\epsilon - x)c(\gamma)$. Several choices for $c(\gamma)$ are explored, including linear decay, truncated decay, and piecewise functions, with empirical results favoring truncated decay with a constant segment before decaying to zero.

Pseudocode for training and sampling is straightforward:

def training_loss(f, eps, x, g): xg = (1-g)*eps + g*x target = (eps-x)*c(g) loss = (f(xg) - target)**2 return loss def generate(f, st, eta, N): xn = st for i in range(N): xn = xn - eta*f(xn) return xn

EqM also supports explicit energy modeling via two formulations: dot product ($g(x_\gamma) = x_\gamma \cdot f(x_\gamma)$) and squared $L_2$ norm ($g(x_\gamma) = -\frac{1}{2}||f(x_\gamma)||_2^2$), with the dot product variant exhibiting better stability and performance.

Sampling and Inference-Time Flexibility

EqM enables optimization-based sampling, where samples are generated by gradient descent on the learned landscape. This approach supports:

• Flexible Step Sizes: EqM is robust to a wide range of step sizes, unlike flow models which require precise scheduling.
• Adaptive Optimizers: Techniques such as Nesterov Accelerated Gradient (NAG-GD) can be employed, yielding improved sample quality, especially with fewer steps.

Figure 2: NAG-GD sampling achieves better sample quality than vanilla GD, with the gap increasing for fewer steps.

• Adaptive Compute: EqM can allocate different numbers of sampling steps per sample, terminating when the gradient norm falls below a threshold, reducing compute by up to 60% without significant degradation in sample quality.

Figure 3: EqM scales favorably with training epochs, parameter count, and patch size, outperforming Flow Matching at all tested scales.

Figure 4: EqM produces realistic images earlier in the sampling process compared to FM and generalizes beyond memorization, as shown by nearest neighbor analysis.

Empirical Results

EqM demonstrates strong empirical performance on class-conditional ImageNet 256$\times$256 generation, achieving an FID of 1.90, surpassing StyleGAN-XL, VDM++, DiT-XL/2, and SiT-XL/2. EqM exhibits superior scaling behavior across model size, training length, and patch size.

Figure 3: Curated samples from EqM-XL/2 and scalability plots showing EqM's consistent outperformance over Flow Matching.

Ablation studies reveal that truncated decay for $c(\gamma)$ with a gradient multiplier $\lambda=4$ yields optimal results. Explicit energy modeling via the dot product variant is preferred due to better stability and performance.

EqM's sampling process is robust to step size variations, and NAG-GD further improves sample quality. Adaptive compute allocation enables efficient inference.

Unique Properties and Applications

EqM exhibits several properties not supported by traditional diffusion/flow models:

• Partially Noised Image Denoising: EqM can denoise partially noised images directly, with generation quality improving as input noise decreases, unlike flow models which degrade when not starting from pure noise.

Figure 2: EqM is robust to a wide range of step sizes, while FM only functions properly at a specific step size.

• Out-of-Distribution Detection: EqM inherently supports OOD detection via energy values, achieving the best average AUROC across tested datasets compared to PixelCNN++, GLOW, and IGEBM.

Figure 5: EqM achieves strong OOD detection performance and supports compositional generation by summing energy landscapes.

• Compositional Generation: EqM supports compositionality by adding gradients from multiple models, enabling generation of images conditioned on multiple labels, similar to EBMs but with greater stability and scalability.

Figure 5: Compositional samples generated by EqM-XL/2 using two ImageNet labels per sample.

Practical Implications and Future Directions

EqM provides a principled framework for generative modeling that unifies flow-based and energy-based perspectives. Its equilibrium dynamics enable flexible, optimization-driven inference, supporting adaptive compute and compositionality. EqM's superior empirical performance and scalability suggest its suitability for large-scale generative tasks.

Potential future directions include:

• Extending EqM to other modalities (e.g., text, audio, video) by leveraging its equilibrium landscape.
• Investigating more advanced optimization techniques for sampling, such as adaptive learning rates or second-order methods.
• Exploring compositionality for multi-modal and multi-task generative modeling.
• Further analysis of the learned energy landscape for interpretability and controllability.

Conclusion

Equilibrium Matching offers a robust, scalable, and flexible generative modeling framework by learning equilibrium dynamics over an implicit energy landscape. It achieves state-of-the-art generation quality, supports optimization-based sampling, and enables unique capabilities such as adaptive compute, OOD detection, and compositional generation. EqM represents a significant step toward unifying energy-based and flow-based generative modeling, with promising implications for future research and applications in AI.