Tsinghua's MGFlow & Elastic Forcing: Distributional Training for One-Step Images & Streaming Video

Introduction: Few-Step Generation for Real-Time Interaction

Multimodal interaction demands visual generation that keeps pace with user actions. Long waiting times break the experience, making few-step generation — compressing generation to one or a few steps — a critical capability for real-time multimodal interaction and world models. Two works from Tsinghua University's Meow Lab, MGFlow and Elastic Forcing, address this challenge through a shared principle: distributional training (样本分布训练), which supervises a batch of generated samples by comparing their overall distribution with a reference distribution, rather than requiring a unique reference answer for each sample.

MGFlow: A Unified Distributional Optimization Framework

Problem: From Per-Sample to Collective Supervision

Generative tasks have no single correct answer; a model must learn the richness of plausible outputs. Multi-step diffusion corrects iteratively, but a one-step generator must map noise to image in a single forward pass. Distributional training asks whether a batch of generated results can directly learn "what the overall distribution should look like" from a batch of reference results. It freezes a visual encoder, maps both generated and reference samples to a feature space, and compares the two feature distributions. If the reference set contains multiple modes but the generator only covers one, the model has not learned the full distribution even if individual samples look realistic.

Unifying FD-Loss and Drifting via Wasserstein Gradient Flow

Existing methods exhibit collective supervision in different forms: FD-Loss minimizes the Fréchet Distance between mean and covariance of two Gaussian-fitted distributions; Drifting uses attraction to reference samples and repulsion among generated samples to construct update directions. MGFlow unifies these by separating two design choices: how to model the distribution (single Gaussian, Gaussian mixture, kernel density estimation) and what criterion to compare distributions (optimal transport, KL divergence, etc.). The connection is the Wasserstein gradient flow : imagining generated features as particles, a distribution-level objective yields a per-particle velocity field, linking "overall discrepancy" with "per-sample updates".

Single Gaussian + optimal transport → FD-Loss.

Gaussian kernel density estimation + KL → update field matching Gaussian-kernel Drifting up to a bandwidth scale factor.

This unified view turns ad-hoc methods into a systematic design space . For example, keeping a single Gaussian but swapping optimal transport for KL yields another analytic update; adopting a Gaussian mixture allows trading off between global statistics and local sample structure.

Coupling Sample Assignment with Update Direction

Modeling multiple modes does not guarantee samples reach the correct modes. If the reference distribution has two equally weighted modes but the generator favors one, global statistics may improve while the proportion remains wrong. MGFlow refines the distribution model with multiple Gaussian components and couples capacity-constrained sample assignment with pairwise updates , ensuring "which mode to assign" and "where to move" stay consistent. At inference, only a single network evaluation is needed.

Experimental Results

ImageNet 256×256 one-step generation: MGFlow-KL on pMF-H and JiT-H backbones achieves FDr⁶ = 1.45 and 1.64 (lower is better), a ~23% and ~38% reduction over FD-Loss. FDr⁶ aggregates distribution discrepancy across six representation spaces.

Notably, MGFlow-KL improves FDr⁶ without directly optimizing FD loss, and also improves on three encoders not used in training, showing generalization beyond the training representation.

Text-to-image on FLUX.2 [klein] 4B: Post-training for one-step generation with joint image-text matching raises GenEval from 0.794 (4-step original) to 0.900 and PickScore from 21.85 to 21.98 , surpassing the original 4-step model and all same-scale baselines.

Image-only matching yields GenEval 0.856; adding joint image-text matching boosts it to 0.900, demonstrating that supervising the conditional alignment directly improves text faithfulness.

Elastic Forcing: Scaling Distributional Training to 14B Video Models

Motivation: Streaming Video Generation on a Single Node

Reducing generation steps is even more crucial for autoregressive video: if generation speed exceeds playback speed, real-time streaming ("generate while playing") becomes possible. However, scaling collective supervision to video is challenging — videos are longer, models larger, and batch comparisons more expensive. Prior Self Forcing methods use asymmetric DMD to distill bidirectional video diffusion into few-step autoregressive models, effective on small models but unstable and memory-prohibitive at 10B+ parameters.

Method: MMD-Based Distributional Supervision with Decoupled Estimation and Backprop

Elastic Forcing retains the autoregressive generation pipeline but replaces the auxiliary score networks (pretrained diffusion teacher + online score model) with Maximum Mean Discrepancy (MMD) between generated and reference video features in a frozen feature space. This eliminates the need to maintain two extra networks during post-training.

MMD involves cross-sample interactions, typically requiring huge batches. Elastic Forcing decouples the problem into two layers:

Collective direction estimation: Use a large set of reference features (via precomputed statistics and random sampling) and a full batch of generated features to determine the update direction, preserving cross-sample relationships.

Controlled backpropagation: Replay the computation graph separately, backpropagate only through a subset of samples with gradient scaling. Samples not backpropagated still influence the overall direction.

Thus, the scale of distribution estimation and the backprop budget become independently tunable.

Results

1.3B controlled comparison (aligned architecture, initialization, generation flow): Elastic Forcing raises VBench total score from 83.80 to 84.64 while maintaining 17 FPS generation speed.

14B model post-training: Completed on 8 H200 GPUs in 23.2 hours . Self Forcing would OOM even on 64 GPUs. Quality surpasses industrial model Krea-Realtime.

Human blind evaluation (42 participants, 20 video pairs, 1–5 scale): Elastic Forcing 14B scores 3.354 vs. Krea Realtime 14B 3.054 , covering human motion, animal movement, object interaction.

Single-node budget enables rapid iteration on new data, styles, characters — critical for product development.

Reference Distribution Defines What the Model Learns

Because the target is defined by reference samples, changing the reference set changes the post-training objective. Elastic Forcing demonstrates adaptation to specific characters (e.g., "Nailong"), styles (black-and-white), and compositions (panoramic projection) on the 1.3B model without re-adapting a diffusion teacher . Reference sources are flexible: standard experiments use Wan-generated videos; qualitative results also show adaptation to real videos from OpenVid.

MGFlow's joint image-text matching concatenates image and text features, embedding the correspondence directly into the supervision. In text-to-image ablation, image-only matching gives GenEval 0.856; joint matching raises it to 0.900. This shows that supervising "does this text match this image" improves alignment, and provides a natural entry for adding action, state, or other conditions.

Future Directions: World Models & Action Models

When few-step generation incorporates action conditions, the line connects to world models (predicting environment dynamics) and world action models (linking actions to consequences). Each prediction feeds the next decision, so generation speed directly limits observation-planning-response frequency. Few-step generation serves as fast prediction; autoregressive streaming provides continuous visual state rollout. For action-driven simulation, the next step is to organize current observation, action condition, and future visual trajectory into a joint distribution, letting supervision cover both future plausibility and action-condition correspondence.

This aligns with distributional training's premise: the same scene and action often have multiple valid futures; the model must learn their overall regularities and proportions. MGFlow supplies the methodological framework for designing such distributional supervision; Elastic Forcing demonstrates a feasible training path for scaling direct sample supervision to large video models. Together, they point toward interactive video, game environments, and embodied intelligence: fewer steps for faster feedback, joint distributions for conditioning, and affordable post-training for continuous data absorption.

Paper & Project Resources

MGFlow

Paper: https://arxiv.org/abs/2609.35763 Project & samples: https://shihaoyang0423.github.io/MGFlow-website/ Code: https://github.com/shihaoyang0423/MGFlow Model:

https://huggingface.co/shy0423/MGFlow

Elastic Forcing

Paper: https://arxiv.org/abs/2609.35491 Project & videos: https://video-examples-m8r2v6.pages.dev/ Code:

https://gitee.com/liu-yueyi/elastic-forcing