[KongchangAI]
· 2 min read· 1,055 words

New Breakthrough in DMD Distillation: Critic Models Must Scale in Sync with Teacher Models

New Breakthrough in DMD Distillation: Critic Models Must Scale in Sync with Teacher Models

In DMD distillation, critics must scale with teachers and condition on generation steps to stably improve few-step generation quality.

This paper systematically investigates the capacity and information requirements of the critic in Distribution Matching Distillation (DMD). Controlled experiments find that when the teacher model grows while the critic's size stays fixed, student quality degrades and training can even collapse—showing that simply enlarging the teacher is not enough. The authors propose joint scaling to expand the critic in sync with the teacher, preventing collapse and letting students truly benefit from larger teachers. On the information side, traditional critics model a step-independent mixture distribution with severe information loss; the authors propose the Step-Aware Critic (STAC), which conditions on the student's generation step to estimate step-specific scores and provide more precise training signals. Combined, the two yield consistent improvements on class-conditional ImageNet, SDXL and FLUX.2 text-to-image, and Wan2.1 text-to-video, indicating that critic capacity and information design are general bottlenecks in DMD distillation.

Distribution Matching Distillation (DMD) has long been a key technique for few-step generative models. It trains a student model based on the difference between the teacher model's score and an online critic's score of the student distribution. But what capacity and information does the critic actually need to provide a reliable training signal? A recent study provides a systematic answer: the critic also needs to "keep up with the times."

DMD distillation research diagram

Insufficient Critic Capacity Leads to Training Collapse

The researchers conducted controlled experiments on ImageNet using seven SiT models of different scales and discovered a counterintuitive phenomenon: when the teacher model grows larger while the critic remains small and fixed, student model quality actually degrades, and training can even collapse. This means that simply enlarging the teacher model does not automatically yield better distillation results—if the critic cannot "keep up" with the teacher's representational capacity, the entire training process becomes imbalanced.

To address this, the team proposed a joint scaling strategy: scaling up the critic in sync with the teacher. Experiments show that this strategy prevents collapse and allows the student model to genuinely benefit from a larger teacher. This finding challenges the earlier simple assumption that "bigger teacher is always better" and emphasizes the importance of capacity matching among components in the distillation system.

SiT Model Overview

SiT (Scalable Interpolant Transformers) is a scalable diffusion/flow-matching generative model built on a Transformer backbone. It unifies diffusion model training under the interpolant framework, treating the forward noising process as an interpolation path between data and noise, and allows flexible choices of interpolation method, sampler, and time-step weighting. Because SiT offers a series of models at different parameter scales (from smaller to larger) with a unified architecture and training recipe, it is ideal for "controlled variable" scaling experiments. This paper leverages SiT's various scales to change only the capacity of the teacher and critic while keeping other conditions fixed, thereby clearly isolating the effects of "larger teacher" and "whether the critic scales in sync" on distillation performance.

Step-Aware Critic: Providing More Precise Score Estimates

Beyond the capacity issue, what information the critic needs is equally critical. In traditional DMD training, the critic is asked to model a step-independent mixture distribution of student outputs, which amounts to having the critic "blindly guess" the scores at different generation steps. The researchers argue that this approach suffers from severe information loss, and thus proposed the Step-Aware Critic (STAC) : conditioning the critic on the student's generation step to estimate step-specific scores.

This design enables the critic to distinguish distribution characteristics at different steps, thereby providing more precise training signals. Experiments show that STAC combined with joint scaling yields consistent improvements across multiple tasks, including class-conditional ImageNet generation, text-to-image generation based on SDXL and FLUX.2, and text-to-video generation based on Wan2.1.

Background on Score Matching and Score Functions

In diffusion models and distillation methods like DMD, "score" refers to the gradient of the log probability density with respect to the input, i.e., ∇ₓ log p(x). The reverse generation process of diffusion models essentially denoises at each step along the score direction corresponding to the current noise level, so the student model must learn the teacher's scoring behavior at each noise level. DMD does not directly regress the teacher's output; instead, it uses a learnable critic to estimate the score of the "student distribution" and then uses the difference between the teacher's score and this estimate as the gradient signal to update the student—this falls under the score distillation approach. The critic acts as a discriminator that is updated online during training: the more accurately it characterizes the student's distribution at each noise level, the closer the gradient signal is to the true distribution-matching direction. This also explains why the critic's capacity and conditioning information directly determine distillation success—it estimates distribution-level quantities, not regression targets for individual samples.

Multi-Task Validation: Generality from Images to Video

The study's validation covers three major categories of generative tasks, demonstrating the method's generality. In class-conditional ImageNet generation, the combination of joint scaling and STAC significantly improves few-step distillation; in text-to-image tasks, both mainstream models SDXL and FLUX.2 benefit; in text-to-video tasks, experiments on Wan2.1 likewise confirm the method's effectiveness. This cross-modal consistency indicates that the critic's capacity and information design are general bottlenecks in DMD distillation, not incidental phenomena specific to particular tasks.

Overview of SDXL, FLUX.2, and Wan2.1

SDXL is an enlarged text-to-image model in the Stable Diffusion series, using a larger UNet backbone and dual text encoders, and serves as a representative baseline for open-source text-to-image generation; FLUX.2 belongs to a new generation of text-to-image models based on flow matching and Transformer architecture, representing the current technical direction of high-quality open-source generative models; Wan2.1 is an open-source model series for text-to-video generation, which requires consistency along the temporal dimension and is significantly harder to distill than images. Applying DMD distillation to these three types of models means the method must simultaneously handle two backbones (UNet and Transformer), two training paradigms (diffusion and flow matching), and two data modalities (image and video). Achieving consistent improvements across these differences is key evidence for this paper's claim that "critic capacity and information design are general bottlenecks."

Implications for Distillation Practice

This study serves as a reminder for teams working on few-step generative models: when optimizing the distillation pipeline, one should not focus solely on the teacher model or student architecture—the critic's design also requires careful consideration. Joint scaling offers a simple and feasible capacity-matching solution, while STAC improves the quality of the training signal from the information perspective. Together, they are expected to become the standard configuration for future DMD distillation.

For resource-constrained teams, joint scaling means that when enlarging the teacher model, a corresponding compute budget must be allocated to the critic in sync—otherwise the results may backfire. Meanwhile, STAC has low implementation complexity yet delivers stable gains without adding much overhead, making it worth prioritizing.

Conclusion

DMD distillation performance depends not only on the size of the teacher model but, more importantly, on whether the critic can "match" the teacher. Joint scaling and the Step-Aware Critic (STAC) address this matching problem from both capacity and information dimensions, with effectiveness validated on image and video generation tasks. This work provides new design guidelines for training few-step generative models.

#DMD distillation#Distribution Matching Distillation#Step-Aware Critic#STAC#joint scaling#few-step generative models#score distillation#ImageNet#SDXL#FLUX.2#Wan2.1#text-to-video
Share:

Related articles