Compressed diffusion models generate images faster with trajectory aware sampling
TamedFlow: Trajectory-Aware Distillation\\for Teacher-Assisted Diffusion Sampling
Computer Vision and Pattern RecognitionArtificial Intelligence
Summary
Diffusion Transformers create high-quality images but usually need many slow steps to finish. The authors show that speeding them up by combining smaller models and smart step-skipping often causes worse images. They create a new training and sampling method, called TamedFlow, that helps smaller models predict better sequences and decide when to ask a bigger model for help. This approach improves image quality and speeds up generation on several datasets.
What this means in practice
- •For image generation engineers: Build faster image generators that keep high quality by using trajectory-aware student-teacher models during sampling.
- •For mobile app developers: Deploy compressed diffusion-based image generation on devices with limited resources while maintaining output fidelity through adaptive sampling.$Commercial implications: Enables consumer apps to offer high-quality image generation with lower latency on phones and tablets.
Authors
Shuaiqi Cheng, Jiahao Zhang, Boyuan Gu, Siyu You, Yuxi Chen, Haiying Xu, Zexin Lu, Huiyu Zhou, Xuming Hu
Abstract
Diffusion Transformers (DiTs) deliver high-fidelity image and video generation, but inference requires many sequential model evaluations. Two acceleration strategies are increasingly important: model compression makes each evaluation cheaper, while adaptive inference avoids selected teacher computation through reuse, forecasting, or speculative execution. Their combination, however, is brittle: directly applying the evaluated adaptive methods to a compressed student can degrade quality, an observed failure mode we call the \textbf{Combination Gap}. We link this gap to a mismatch between pointwise distillation and short-horizon execution: a student trained at teacher states need not remain reliable when its own predictions are chained. We propose \textbf{TamedFlow}, which trains a compressed DiT for this deployed role through cross-architecture feature alignment, a near-identity Velocity Alignment Head, and randomized rollout supervision. At inference, a lightweight Learnable Trajectory Estimator routes completed drafts: accepted drafts bypass the teacher, while rejected drafts roll back and receive one teacher rescue step. Distilling DiT-XL/2 into DiT-B/2, TamedFlow reaches FID $2.58$ at $2.25\times$ measured speedup over DDIM-50 on ImageNet $256\times256$. Ablations support trajectory supervision, and routing controls favor LTE over random decisions at nearly matched latency. Evaluations on ImageNet $512\times512$, MS-COCO, and Latte/FaceForensics extend the empirical scope.