Swift method improves diffusion model training without extra inference cost
SWIFT: Self-guided Weak-to-strong dIFfusion models via Target extrapolation
Machine LearningArtificial IntelligenceComputer Vision and Pattern Recognition
Summary
Training advanced image generation models called diffusion models can take a lot of computing power, and making them better during use can slow them down. The authors created Swift, a way to improve these models during training by using a simpler, frozen guide model to nudge the main model without slowing down later. Swift adjusts the training target to focus on the parts the simple guide misses, producing better results without extra steps when generating images. They tested this on popular image datasets and found it improves quality and efficiency.
What this means in practice
- •For machine learning engineers: Train diffusion models faster and with less compute by using a simple guide model to boost main model quality without changing sampling steps.
- •For computer vision developers: Improve image generation quality in conditional and unconditional diffusion pipelines without adding inference overhead or complex sampling strategies.
- •For creative ai tool builders: Build AI-powered image generation tools that deliver higher quality outputs needing less training compute and maintain efficient, single-model inference.$Commercial implications: Enables improved image synthesis models that can be sold as AI creative services or embedded software with lower operational costs.
Authors
Yanjin He, Cuifeng Shen, Wanzhou Lei, Maohua Li, Tao Lan, Kan Liu, Hanlin Tang, Qing Qu
Abstract
Training high-quality diffusion models is computationally expensive, and inference-time guidance often adds further sampling cost. In this work, we introduce Swift, a simple yet effective self-guided training method that uses a frozen weak denoiser to improve a stronger model without adding inference overhead. For each noisy input, Swift extrapolates the clean regression target away from the weak denoiser's prediction, amplifying components that remain underfit. The weak guide can be an early training checkpoint, requires no backpropagation during main-model training, and is discarded afterward. The resulting model retains the original sampler and supports both conditional and unconditional generation without label dropout or an unconditional branch. We explain this mechanism through a spectral analysis of finite-time population gradient flow in an RKHS. With an early-stopped weak guide, we prove that a nonempty range of extrapolation strengths yields strictly lower posterior-mean error than matched standard training, despite shifting the regression target away from the posterior mean. Experiments on CIFAR-10 and ImageNet-256 demonstrate improvements across model scales and representations. In one ImageNet-256 configuration, Swift surpasses the baseline's best FID using roughly half the total training compute, including guide preparation and training overhead. On RAEv2, it achieves FID 1.15 without inference-time guidance and 0.99 with guidance. These results establish target extrapolation as an effective approach to weak-to-strong diffusion training with efficient single-model inference.