Better critics improve training of student models in AI generation tasks
A Better Critic Matters for Distribution Matching Distillation
Computer Vision and Pattern RecognitionMachine Learning
Summary
Training smaller AI models to mimic larger ones can sometimes fail if the helper model, called a critic, is too simple. The authors find that making this critic bigger alongside the teacher model helps avoid training problems. They also show that letting the critic know about the specific step in the student's generation process improves its guidance. These improvements work across different image and video generation tasks, helping smaller models learn more effectively from big, complex ones.
What this means in practice
- •For ai model developers: Train smaller generative models more effectively by scaling the critic alongside the teacher and using step-aware critics for better score guidance.
- •For multimedia generation teams: Improve training of models that generate images or videos from texts by applying joint scaling and step-aware critics in the distillation process.
Authors
Zongze Wu, Xianbang Wang, Zhengyang Geng, Mingdeng Cao, Zhoutong Zhang, Richard Zhang, Eli Shechtman, Ruiming Cao
Abstract
Distribution Matching Distillation (DMD) trains a few-step student using the difference between a teacher's score and an online critic's estimate of the student distribution's score. We investigate what the critic needs, in capacity and in information, to provide a reliable training signal. On the capacity side, controlled experiments across seven SiT model sizes on ImageNet show that increasing teacher size while keeping a small critic fixed can degrade student quality and cause training collapse. Scaling the critic alongside the teacher, a strategy we call joint scaling, prevents collapse in our experiments and allows the student to benefit from larger teachers. On the information side, rather than forcing the critic to model a step-agnostic mixture of student outputs, we introduce a step-aware critic (STAC) that conditions on the student's generation step to estimate step-specific scores. Together, joint scaling and STAC consistently improve few-step distillation across class-conditional ImageNet generation, text-to-image generation with SDXL and FLUX.2, and text-to-video generation with Wan2.1.