Improved Training Technique for Shortcut Models

TL;DR

The paper tackles why shortcut models underperform by formalizing issues from guidance compounding to trajectory curvature. It proposes iSM, a framework combining Intrinsic Guidance, Multi-Level Wavelet Loss, Scaling Optimal Transport, and Twin EMA to stabilize training and improve sample fidelity across one-step to multi-step generation. On ImageNet 256×256, iSM delivers significant FID improvements, closing the gap with larger diffusion and GAN systems while enabling flexible inference budgets from a single network. The results suggest shortcut models can be a practical, competitive class of generative models with broad applicability and efficiency advantages.

Abstract

Shortcut models represent a promising, non-adversarial paradigm for generative modeling, uniquely supporting one-step, few-step, and multi-step sampling from a single trained network. However, their widespread adoption has been stymied by critical performance bottlenecks. This paper tackles the five core issues that held shortcut models back: (1) the hidden flaw of compounding guidance, which we are the first to formalize, causing severe image artifacts; (2) inflexible fixed guidance that restricts inference-time control; (3) a pervasive frequency bias driven by a reliance on low-level distances in the direct domain, which biases reconstructions toward low frequencies; (4) divergent self-consistency arising from a conflict with EMA training; and (5) curvy flow trajectories that impede convergence. To address these challenges, we introduce iSM, a unified training framework that systematically resolves each limitation. Our framework is built on four key improvements: Intrinsic Guidance provides explicit, dynamic control over guidance strength, resolving both compounding guidance and inflexibility. A Multi-Level Wavelet Loss mitigates frequency bias to restore high-frequency details. Scaling Optimal Transport (sOT) reduces training variance and learns straighter, more stable generative paths. Finally, a Twin EMA strategy reconciles training stability with self-consistency. Extensive experiments on ImageNet 256 x 256 demonstrate that our approach yields substantial FID improvements over baseline shortcut models across one-step, few-step, and multi-step generation, making shortcut models a viable and competitive class of generative models.

Figures (9)

• Figure 1: Visualizing the accumulated guidance problem in standard shortcut models (left) versus our integrated‐guidance approach (right). Top row shows multi-step generation; bottom row shows one-step generation. In the original shortcut model, the fixed guidance strength learned at training time effectively compounds across the many implicit sub-steps, leading to oversaturated colors and blurriness. By conditioning the network on the guidance scale and training it to apply guidance explicitly at each step, our method produces consistent, artifact-free images even in the extreme one-step and few-step regimes.
• Figure 2: Efficacy of Scaling Optimal Transport (sOT) Matching in improving shortcut model training, demonstrated by varying the OT scaling factor $K$. (a) On a bimodal target, forward trajectories (top row) without OT exhibit frequent intersections (red), compelling the reverse generative process (bottom row) to follow high-curvature paths that initially average the target modes (blue, green). Our sOT approach, by increasing the effective OT scaling $K$, progressively disentangles these forward couplings, yielding substantially straighter reverse trajectories. (b, c) The benefits of this trajectory straightening are reflected in the training losses: both self-consistency (b) and velocity & guidance (c) losses are consistently lower and more stable with larger $K$, especially after the interval guidance point $t_{\mathrm{interval}}=0.3$ (left-side of the red dashed line). Note for (a): To avoid overly dense trajectory, only intersecting forward paths and sampling trajectories with high curvature are highlighted. Zoom for detail.
• Figure 3: Effect of applying different numbers of wavelet loss layers ($L$) on image quality under one and few-step inference. The traditional low-level objective (leftmost column) results in noticeable degradation and artifacts. In contrast, using multi-level wavelet loss, especially with $L=5$ (green box, right), produces high-quality, consistent images across all inference settings.
• Figure 4: Uncurated samples on ImageNet-$256\times256$ with 1-step sampling.
• Figure 5: Uncurated samples on ImageNet-$256\times256$ with 2-step sampling.

Theorems & Definitions (1)

• Proposition 1