Energy-Guided Distribution Matching Distillation for Diverse Few-Step Visual Generation

Diffusion distillation enables high-quality generation in only a few sampling steps, but aggressive acceleration often substantially reduces sample diversity. In particular, Distribution Matching Distillation (DMD) optimizes a reverse-KL objective whose mode-seeking behavior can collapse the teacher's rich multimodal distribution. Existing diversity-preserving methods mitigate this issue through trajectory-level regularization, forcing individual student trajectories to reproduce teacher trajectories under a prescribed noise correspondence. We argue that such pairing is unnecessarily restrictive: diversity should instead be preserved at the distribution level. We introduce Energy-Guided Distribution Matching Distillation (EG-DMD), which complements DMD with population-level teacher supervision for improved mode coverage. At each training step, EG-DMD jointly rolls out multiple student and teacher trajectories to an intermediate diffusion time and matches their empirical distributions through a Sinkhorn-divergence energy. Unlike paired regression, this objective optimizes the coupling between student and teacher populations, preserving teacher modes without enforcing a particular noise-to-sample correspondence. Across diffusion models with diverse architectures, EG-DMD consistently improves the diversity–quality frontier over DMD and existing diversity-preserving baselines while retaining four-step inference. On SDXL, EG-DMD improves ImageReward over the strongest baseline by 8.9%, while increasing DINO and CLIP diversity by 13.3% and 23.6%, respectively. These results establish population-level energy guidance as a simple, data-free, and architecture-agnostic principle for high-fidelity, diverse few-step generation.