ICML 2026poster0 citations

Diversity-Preserved Distribution Matching Distillation for Fast Visual Synthesis

Tianhe Wu, Ruibin Li, Lei Zhang, Kede Ma

Abstract

Distribution matching distillation (DMD) aligns a multi-step generator with its few-step counterpart to enable high-quality generation under low inference cost. However, DMD tends to suffer from mode collapse, as its reverse-KL formulation inherently encourages mode-seeking behavior, for which existing remedies typically rely on perceptual or adversarial regularization, thereby incurring substantial computational overhead and training instability. In this work, we propose a role-separated distillation framework that explicitly disentangles the roles of distilled steps: the first step is dedicated to preserving sample diversity via a target-prediction (e.g., v-prediction) objective, while subsequent steps focus on quality refinement under the standard DMD loss, with gradients from the DMD objective blocked at the first step. We term this approach Diversity-Preserved DMD (**DP-DMD**), which, despite its simplicity—no perceptual backbone, no discriminator, no auxiliary networks, and no additional ground-truth images—preserves sample diversity while maintaining visual quality on par with state-of-the-art methods in extensive text-to-image experiments. Code and models will be released publicly.

OptimizationRobustnessVisionRetrieval
BibTeX
@inproceedings{
wu2026diversitypreserved,
title={Diversity-Preserved Distribution Matching Distillation for Fast Visual Synthesis},
author={Tianhe Wu and Ruibin Li and Lei Zhang and Kede Ma},
booktitle={Forty-third International Conference on Machine Learning},
year={2026},
url={https://openreview.net/forum?id=R3JzLU2qZj}
}