Matching final images helps four-step generators improve
eDMD: Endpoint Distribution Matching for Few-Step Diffusion Distillation · arXiv:2611.17194
In plain words
A diffusion image generator can be trained to make an image in just a few steps, using a slower generator as its teacher. Endpoint Distribution Matching Distillation trains a comparison model on the fast generator’s finished images, then uses it to correct each step. It also improves some teacher guidance.
In tests on four image-generation models, the resulting four-step generators had the highest reported image-preference scores among the compared students. The result applies to those tested models and image evaluations.
The research
A fast image generator may take only a few denoising steps, so its intermediate predictions can differ from a teacher’s even when its finished images match. This result makes the training reference consequential: matching final outputs may avoid corrections that pull an already-correct result away.
Endpoint Distribution Matching Distillation trains its comparison model on a generator’s final outputs, aligning each step’s correction with the images it delivers.
During training, eDMD fits one shared fake model only to the student’s final outputs, then uses it to correct every step. It also sharpens low-noise teacher targets and selects intermediate query levels using teacher consistency.
Across four image-generation backbones, eDMD had the highest reported peak image-preference scores among compared four-step students. That pattern supports endpoint-focused distillation in these tests; the comparison sets varied by backbone.
The image tests used 200 text prompts with five seeds per prompt, and the comparison set varied by backbone. The result is bounded to those prompts and scoring protocols.
What the evidence shows
- Match what gets delivered
Endpoint Distribution Matching Distillation (eDMD) is a method for training a generator to produce images in a few denoising steps, using a slower pretrained generator as its teacher.
It trains a shared comparison model only on the student’s final outputs, then uses that model to correct each step. The choice of comparison model matters because the desired result is the student’s final image distribution, even when its intermediate predictions differ from the teacher’s.
- Intermediate targets can misdirect
Some multi-step methods train their comparison model on intermediate predictions or states. The paper’s theoretical analysis identifies an extra correction when that model’s reference distribution differs from the student’s final-output distribution.
Under the stated assumptions, the extra correction disappears at every noise-level exactly when the reference matches those final outputs. That gives a specific reason to align the training reference with the images the generator ultimately produces.
- A toy test shows the mismatch
A two-dimensional synthetic test compared endpoint, mixture, and stage references for two students whose intermediate steps differed from the teacher. One student’s final outputs already matched the teacher; in that case, the endpoint reference added no correction, while mixture and stage references still pushed the student away.
This comparison demonstrates the mismatch the method is designed to avoid, in a synthetic setting rather than an image test.
- Four backbones favor eDMD
The image evaluations tested four-step students on DrawBench-1000, a set of generated images made from 200 text prompts with five seeds per prompt. HPSv3++ is an image-preference score.
Across four teacher backbones, eDMD had the highest reported peak score among the compared four-step distribution-matching students. This supports the method on those evaluations, although the set of methods compared differed by backbone.
- Removing the endpoint costs
A component test compared the full method with a version that replaced its endpoint reference with intermediate-state training. On SD3-Medium, the HPSv3++ image-preference score was 11.171 for the full method and 10.122 without the endpoint reference.
This was the largest overall drop among component removals on that backbone. The test points to the endpoint reference as a major part of the measured result, within this component-test setup.
Deep
- Original paper authors
- Tong Zhao, Xianfang Zeng, Yucheng Han, Ruoyu Wang, Mingkun Lei
- Submitted to arXiv
- License
- CC BY 4.0