Improved Training Technique for Shortcut Models
Anh Nguyen, Viet Nguyen, Duc Vu, Trung Dao, Chi Tran, Toan Tran, Anh Tran
Shortcut models represent a promising pathway for generative modeling, supporting one-step, few-step, and many-step sampling without adversarial training. Yet their practical adoption has been hindered by several key issues: the need to fix its classifier-free guidance scale at training time, which limits flexibility at inference; high variance from joint training on random noise–data pairs, which decelerates and destabilizes convergence; and reliance on low-level distance on direct domain that bias reconstructions toward low frequencies and degrade sample quality. Moreover, we uncover a previously overlooked problem of accumulation scale in classifier-free guidance and a subtle conflict between EMA updates and the self-consistency objective. To address these challenges, we introduce a unified training framework for shortcut models that (1) parameterizes guidance scales to support dynamic guidance sampling at inference, (2) mitigates frequency bias with a Multi-level Wavelet Loss, (3) incorporates interval guidance directly into the loss, (4) reduces training variance via Scaling Optimal Transport Matching, and (5) preserves self-consistency alongside training stability through a Twin EMA strategy. Extensive experiments on ImageNet $256\times256$ demonstrate that our approach yields substantial FID improvements over baseline shortcut models in one-step, few-step, and multi-step generation.