Equilibrium Forcing: Adaptive Video Generation Without Noise Conditioning

Hansen Jin Lillemark* 1, Alex Rojas* 1, Zachary Novack1, Runqian Wang2,
Yilun Du3, Yian Ma1, Taylor Berg-Kirkpatrick1, Rose Yu1
*Equal Contribution
1 UC San Diego MIT logo 2 MIT Harvard logo 3 Harvard University
Paper arXiv Code (available on request)

Equilibrium Forcing (EqF) is a framework for video denoising generative models without noise level conditioning. EqF improves video generation quality and sampling speed through adaptive sampling, unifying and simplifying the theory and practice of noise-unconditional denoising models.

Abstract

Standard autoregressive video generation algorithms based on Diffusion and Flow Matching rely on rigid training objectives and static sampling schedules, limiting inference procedures from adapting to the data.

We introduce Equilibrium Forcing (EqF), a simplified framework for video denoising generative models without noise level conditioning. EqF pioneers modular training- and inference-time designs for noise-unconditional generation that decouple learning the denoising field from sampling.

This flexibility allows for inference-time algorithms that operate in a closed loop by adapting to feedback from the sample, improving video quality and consistency on challenging autoregressive video generation benchmarks. Extensive analysis elucidates exactly how removing the noise level conditioning enables EqF's data-dependent inference properties to surpass the performance of standard noise level-conditional denoising video methods.

Generated Video Samples

Below are video samples from the Minecraft, RealEstate10K, and DROID datasets. More are visible below in the experiments section of this webpage.

Minecraft Rollout #1
Re10K Rollout #1
DROID Rollout #1
Prompt: "A silver toaster sits on a table with a red and black checkered tablecloth. A robot arm pushes down the lever on the toaster."

Equilibrium Forcing Overview

Equilibrium Forcing (EqF) learns a denoising field for generating video samples without explicit noise level conditioning. Similar to Flow Matching, noisy training data is constructed from a simple interpolation with the corresponding velocity target:

\[ \mathbf{x}^{\boldsymbol{\sigma}} = (\boldsymbol{1}- \boldsymbol{\sigma})\odot\mathbf{x} + \boldsymbol{\sigma}\odot\boldsymbol{\epsilon}, \qquad \mathbf{v}=\boldsymbol{\epsilon}-\mathbf{x}. \]
\[ \begin{aligned} \mathbf{x}^{\boldsymbol{\sigma}} &= (\boldsymbol{1}- \boldsymbol{\sigma})\odot\mathbf{x} + \boldsymbol{\sigma}\odot\boldsymbol{\epsilon},\\ \mathbf{v}&=\boldsymbol{\epsilon}-\mathbf{x}. \end{aligned} \]

Flow Matching

\[ f_{\mathrm{FM}}(\mathbf{x}^{\boldsymbol{\sigma}},\boldsymbol{\sigma}) \rightarrow \mathbf{v} \]

Flow Matching receives both the noisy sample and a matched noise level. During inference, the sampler follows a prescribed sequence of noise levels to produce a sample from the target distribution.

Equilibrium Forcing

\[ f_{\mathrm{EqF}}(\mathbf{x}^{\boldsymbol{\sigma}}) \rightarrow \mathbf{v} \]

Instead of learning a chain of noise-conditional velocity fields, EqF learns a single denoising field over noisy samples, without explicit noise level conditioning.

Flow Matching vs. Equilibrium Forcing
To learn the correct denoising field, EqF naturally learns an internal estimator of the noise level, which is also a proxy for the distance to the data manifold. This internal estimate can be used to shape the sampling dynamics at inference time to have adaptive properties! Meanwhile, Flow Matching's rigid and inaccurate noise level conditioning produces incorrect denoising directions, especially in autoregressive video generation rollouts.

Modular Training and Inference Decisions with EqF

EqF separates the training of the denoising field from the shaping of the sampling dynamics. This is done by leveraging information that the model already learned as part of its training:

Estimate Sampling Progress

A lightweight readout model \(h_\omega\) easily learns to predict the current noise level from the intermediate EqF activations \(\mathbf{z}\) as EqF already has an internal estimate. This provides feedback about how far the current sample has progressed toward clean data.

\[ \begin{aligned} (\widehat{\mathbf{v}}^{(i)},\mathbf{z}^{(i)}) &= f_{\mathrm{EqF}}(\mathbf{x}^{(i)}),\\ \widehat{\boldsymbol{\sigma}}^{(i)}&=h_\omega(\mathbf{z}^{(i)}). \end{aligned} \]

Shape Sampling Dynamics

An inference-time warp \(\eta\) maps the estimated noise level to an update magnitude. Different choices of \(\eta\) can induce different sampling landscapes, introduced in the next section. Shown below is an attractor warp function that makes the sampling converge to a fixed point.

\[ \mathbf{x}^{(i+1)} = \mathbf{x}^{(i)} - \eta(\widehat{\boldsymbol{\sigma}}^{(i)}) \odot \hat{\mathbf{v}}^{(i)}. \]
Equilibrium Forcing overview
The same trained EqF model can support flexible inference procedures without changing the training objective. For example, a transport-style warp can intersect the data manifold with a nonzero velocity, while an attractor-style warp makes the update magnitude vanish near clean data.

Relation to Equilibrium Matching

Equilibrium Matching (EqM) proposes to remove noise conditioning for image generation, but it modifies the training target to \(c(\boldsymbol{\sigma}) \odot \mathbf{v}\) so that the learned field has baked in attractor-like sampling dynamics. This introduces costly training-time hyperparameters for the shape of \(c\) and its magnitude scale. By moving the sampling dynamics to an inference-time decision, we demonstrate that EqF with the warp can recover any trained EqM function without having to retrain. Additionally, we find that EqM's objective induces a suboptimal loss weighting; see the paper for more details.

Adaptive Fixed-point Inference

EqF's framework enables inference procedures with natural feedback mechanisms to adjust to the state of the sample. We use this flexibility in three complimentary ways:

Closed-Loop Feedback

EqF estimates the current state of the sample via the noise level estimate, rather than following a predetermined trajectory. The warp adjusts using this noise level estimate to adapt to progress that is faster or slower than expected.

Attractor Sampling + NAG

Warps with \(\eta(\sigma)\rightarrow 0\) near clean data turns sampling into fixed-point search. This enables gradient-based acceleration algorithms such as Nesterov Accelerated Gradient (NAG).

Budget-Adaptive Sampling

In addition to the denoising progress, the sampler can use the number of steps remaining to continuously reallocate the overall inference budget, rather than committing to a schedule before generation begins.

Budget-adaptive sampling results
The model uses feedback from the current sample to accelerate toward a fixed point, and dynamically redistribute the sampling dynamics under a limited inference budget. These capabilities, uniquely enabled by EqF's framework, produce gains across Minecraft, Re10K, and Droid datasets, as described in the following section.

Video Generation Experiments

We train and evaluate EqF across three video generation settings: action-conditioned Minecraft, camera-conditioned RealEstate10K, and language prompt-conditioned robotic manipulation on Droid. The Minecraft model is trained from scratch, with all baselines trained under matched budgets. Re10K models are finetuned from open source Flow Matching model Wan 2.1 with 1.3B params, a process we term Equilibrium Finetuning. Droid is trained at an even larger scale at 5B params.

These results demonstrate the scalability and broad applicability of EqF's framework. Results are visualized below, with more experimental details and results in the paper.

Minecraft Rollouts

We compare the generated EqF videos with Equilibrium Matching (EqM) and FM baselines on the Minecraft dataset. Each rollout follows a set of action conditions either going left, right, or straight. Rollouts are conditioned on 25 frames of history, and generate 275 frames into the future with a model self-attention window size of 50.

RealEstate10K Rollouts

We compare videos of rollouts on RealEstate10K between the finetuned EqF-FT and FM-FT models. Each rollout follows camera conditions that describe the movement of the camera throughout the trajectory. Rollouts are conditioned on 37 frames of history, and generate 152 frames into the future with a model self-attention window size of 49.

DROID Rollouts

We also provide example videos of EqF-FT's rollouts on the Droid dataset, with the text prompt below the video. These videos are non-autoregressively generated, generated from 21 frames of context to generate 28 future frames.

Prompt: "A white robot arm moves down and grasps a white tea towel on a white table, then moves the tea towel to wipe the table surface."
Prompt: "A robot arm descends and picks up the lid from a black pot on a kitchen counter and then sets the lid down on the counter next to the pot."
Prompt: "A white robot arm approaches a table with a red pepper, a silver salt shaker, and a yellow corn plush. The arm's gripper picks up the corn plush and moves it to the left on the table. The camera remains stationary."
Prompt: "A robot arm moves a folded gray towel to the right on a brown desk."

Analysis: Why Is Fixed Noise Conditioning an Issue?

Flow Matching inference assumes that the scheduled noise level accurately describes the current state of the generated sample. During training, Flow Matching always receives the exact pair of a noisy sample and the exact noise level that the sample is at. During autoregressive generation, However, this assumption does not hold in practice, as the forward and reverse processes are not the same. Furthermore, integration and generation errors over autoregressive rollouts make this assumption even less sturdy.

We probe this with an auxiliary model trained to estimate the effective noise level of the current sample and compare it with the scheduled noise level. Indeed, we find that the scheduled and true noise levels has a nonzero divergence that increases over autoregressive rollouts (left). We posit that this is a primary cause for the degrading performance of Flow Matching for long rollouts (right). An attempt to close the loop by passing in the estimated noise level can help bridge the gap, but it still is unable to match EqF's performance.

Noise-level divergence and autoregressive generation quality
Since Flow Matching is trained to predict the denoising direction given a noise level that it assumes is correct, it is unable to adapt during inference time when the two signals have a mismatch. This mismatch is a contributing factor to Flow Matching's performance degradation over autoregressive video generation rollouts. See the paper for further analysis.

Analysis: How EqF Learns Without Noise Conditioning

Removing \(\boldsymbol{\sigma}\) initially appears to make the denoising task ambiguous. Flow Matching explicitly receives the noise level, while EqF must infer the appropriate denoising direction from the noisy sample alone.

However, for high dimensional data, the posterior \(p(\boldsymbol{\sigma}\mid\mathbf{x}^{\boldsymbol{\sigma}})\) concentrates around a single noise level. The noisy sample therefore contains enough information to identify its own effective noise level, enabling EqF to actually learn the same optimal denoising field as Flow Matching, which gets the noise level information directly.

We find that in practice across datasets, this decomposition holds: EqF learns the same denoising field as Flow Matching, and it does so by having an internal estimate of the noise level. Together, these render the explicit noise level conditioning superfluous.

Flow Matching and Equilibrium Forcing loss difference and residual field gap
The residual between the Bayes-optimal EqF and Flow Matching losses vanishes across training, as predicted by our analysis. Together, this provides an explanation for how noise-unconditional denoising models can improve over noise-conditional ones: the underlying optimal denoising fields are equally optimal, but EqF can avoid incorrect noise level conditioning altogether, and it can further take advantage of adaptive algorithms.

BibTeX

TODO