How to Train a One‑Step Generative Model Without CFG, DMD, GAN, or Drifting

Generating images with a single forward pass is highly efficient, but reducing the sampling steps to one requires a stable and simple training signal that can bring the generated distribution close to the real distribution.

Existing routes each have trade‑offs: autoregressive and diffusion models need sequential or multi‑scale decoding; flow‑matching models repeatedly invoke the network along a noise‑to‑data path; GANs provide a single forward but rely on a discriminator in a sensitive min‑max game; consistency models and distribution‑matching distillation (DMD) need noise‑time supervision, teacher networks, or extra stabilisation tricks; drifting models require batch‑wise full pairing and multiple real samples per description, which does not fit “one prompt – one image” data.

TBSM reframes the problem as learning a per‑sample direction: for each generation condition (e.g., a prompt) a three‑sample event is formed—one real image, one generated image that retains gradients, and a second generated image with stopped gradients. The instantaneous three‑body scattering direction is computed as the gradient of an energy distance, consisting of a unit vector toward the real sample and a negative term pushing away the other generated sample. The Tracker network, a lightweight 30 M‑parameter model, learns the conditional expectation of these noisy directions, producing a stable Fake‑to‑Real flow field.

During training the generator regresses each generated sample to the position obtained by moving it along the Tracker‑predicted direction. The loss mixes the instantaneous scattering direction and the Tracker output using two weights (ρ and λ). Algorithm 1 describes the basic three‑sample regression, while Algorithm 2 shows the full training loop that updates both generator and Tracker parameters without any adversarial or teacher‑based supervision.

Experiments on ImageNet‑512 with the 800 M‑parameter PixelDiT‑XL/16 model demonstrate that, after TBSM training, a single‑forward generator achieves FID 1.92 (compared to FID 398.20 when the same model is forced to run with NFE = 1 without TBSM). Similar gains are reported on MNIST, Fashion‑MNIST, CIFAR‑10, and a 20 B‑parameter Qwen‑Image text‑to‑image model, where qualitative one‑step samples are shown. Ablation studies reveal that increasing the Tracker weight improves FID, that removing the generator‑generator interaction term still yields strong performance, and that a GAN‑like point (no generator‑generator term) also works well.

The authors discuss that TBSM provides a sample‑level distribution‑matching interface without an online teacher, making it a practical post‑training step for mature diffusion or flow‑matching models. Future directions include exploring random initialization at ImageNet scale and dynamically adjusting the weighting of instantaneous directions.

In conclusion, TBSM converts the energy‑distance between real and generated distributions into a per‑sample motion supervision signal; after training the auxiliary Tracker is discarded, leaving a generator that produces high‑quality images in a single forward pass.