SWIFT: Self-guided Weak-to-strong dIFfusion models via Target extrapolation
1University of Michigan, Ann Arbor 2Alibaba Group 3University of California, Berkeley 4Nanjing University
♯Work done during internship at Alibaba †Project lead
TL;DR. With a finite training budget, a diffusion model falls short of the posterior mean, and models of the same kind trained the same way fall short in the same direction. SWIFT reads that direction off a frozen weak model and extrapolates the regression target against it, so the trained denoiser lands closer to the posterior mean, with nothing changed at inference. On ImageNet-256 it reaches FID 1.15 without guidance and 0.99 with guidance.
Overview. Left: ImageNet-256 FID without vs. with guidance for published models; SWIFT (RAEv2 / DINOv3-L) sits in the lower-left corner. Right: on the matched DINOv2-B setup, SWIFT stays ahead of the baseline throughout training.
Finite training misses the posterior mean, in a predictable direction.
A diffusion model is trained by regressing its denoiser \(D_\theta(\bm{x}_t,\sigma_t)\) onto the clean sample \(\bm{x}_0\). The optimum is the \(D^\star = \mathbb{E}[\bm{x}_0 \mid \bm{x}_t]\), and the sampler uses the trained denoiser as if it were \(D^\star\). But training stops after a finite budget, so the denoiser we sample with falls short of \(D^\star\).
With \(\bm{x}_t = \bm{x}_0 + \bm{\varepsilon}_t\), \(\bm{\varepsilon}_t \sim \mathcal{N}(\bm{0},\sigma_t^2\bm{I})\), the minimizer of \(\mathbb{E}\,\lambda(t)\|D(\bm{x}_t,\sigma_t)-\bm{x}_0\|^2\) over all functions is \(D^\star(\bm{x}_t,\sigma_t) = \mathbb{E}[\bm{x}_0 \mid \bm{x}_t]\). Tweedie's formula turns it into the score, \(\nabla_{\bm{x}_t}\log p_t(\bm{x}_t) = (D^\star - \bm{x}_t)/\sigma_t^2\), which is what the sampler integrates.
This shortfall is not random. Models of the same kind, trained the same way, learn the same parts of \(D^\star\) slowly, so they fall short in the same direction. A weaker such model, smaller or stopped earlier, falls further along that direction, so its prediction shows where training will miss. We can then aim past the target in the opposite direction and land near \(D^\star\).
SWIFT: aim past the target.
SWIFT keeps the loss, the weighting, the architecture and the sampler, and changes only what the denoiser is asked to predict:
Here \(D_{\mathrm{weak}}\) is a frozen weak denoiser and \(\alpha \ge 0\) sets the extrapolation strength; \(\alpha = 0\) recovers ordinary training. The target is pushed past the clean sample, away from the weak prediction, that is, against the direction in which a model like ours falls short. Conditioned on \(\bm{x}_t\), it averages to \((1+\alpha)D^\star - \alpha D_{\mathrm{weak}}\).
Target extrapolation, in two dimensions. Training starts at 0 and stops short of what it aims at (dotted). Aimed at the posterior mean \(D^\star\), it lands at \(D_{\mathrm{weak}}\), off by \(v_1\). SWIFT aims past \(D^\star\), away from \(D_{\mathrm{weak}}\). Because the same kind of model is trained the same way, the shortfall \(v_2\) toward the new target is close to \(v_1\) (in a linear model of training they differ only by a term proportional to \(\alpha\)), so the learned \(D_{\mathrm{SWIFT}}\) lands much closer to \(D^\star\).
The guide is an ordinary denoiser trained with the standard loss, for example an early checkpoint of ordinary training, and frozen before the main run. In our default ImageNet setup it is smaller, with about a quarter of the main model's parameters, and trained for half as long. It runs without gradients, adding about per training step, and is thrown away after training, so sampling uses one network and the original sampler. The method is self-guided: the correction comes from a weaker denoiser on the same task, with no external teacher, no auxiliary head and no label dropout. It works for conditional and unconditional generation, and velocity-predicting models apply it to the .
The DiffusionBench guide's forward pass costs about 26% of the main model's (87.3 vs. 337.5 GFLOPs per sample). A baseline update is one forward and one backward pass, about three forward-equivalents; SWIFT adds one guide forward pass and no backward pass through it, so an update costs about \(1 + 0.26/3 \approx 1.09\) times the baseline. Pretraining the guide costs about 10.4 baseline-equivalent epochs, paid once and reusable across runs; charged to a single 80-epoch run, the total is about 1.22 times an 80-epoch baseline.
The ImageNet models use linear flow matching, \(\bm{x}_t = (1-t)\,\bm{x}_0 + t\,\bm{\epsilon}\), and most of them predict the velocity \(\bm{v} = \bm{\epsilon} - \bm{x}_0\). At a fixed noisy input, the velocity is affine in the clean prediction, so extrapolating the clean target is the same as extrapolating the velocity target: \(\bm{v}_{\mathrm{SWIFT}} = (1+\alpha)\,\bm{v} - \alpha\,\bm{v}_{\mathrm{weak}}\), with \(\bm{v}_{\mathrm{weak}}\) the frozen guide's prediction.
The guide stays frozen, so for a given clean sample, noise draw and noise level the target never changes during training. This sets SWIFT apart from training-time guidance methods whose reference . The same algebraic form also appears in (Karras et al., 2024), which extrapolates a strong model's prediction away from a weak one's at every sampling step; SWIFT arrives at it from the training side, as a correction of where finite training lands, and needs nothing extra at inference.
Autoguidance evaluates \(D_{\mathrm{AG}} = (1+s)\,D_{\mathrm{base}} - s\,D_{\mathrm{weak}}\) at every sampling step, where \(s\) is the guidance scale. The mean of the SWIFT target, \((1+\alpha)D^\star - \alpha D_{\mathrm{weak}}\), has the same form with the ideal denoiser \(D^\star\) in place of \(D_{\mathrm{base}}\); SWIFT uses it as a regression target instead of a sampling rule.
What the correction looks like.
On a two-class, tree-shaped 2D distribution with thin branching support (after the fractal example of Karras et al.), SWIFT follows the narrow branches more closely than classifier-free guidance, autoguidance and four other guidance methods, each with independently tuned settings. It has the lowest KL divergence (0.728 vs. 0.817 for GFT, the strongest baseline), the lowest outlier rate (3.66% vs. 4.97%) and the highest log-likelihood (+0.08 vs. −0.40).
A 2D toy distribution. Ground truth and samples from seven methods; the boxes zoom into branching regions. Lower KL and outlier rate and higher \(\log p\) are better.
Why a deliberately biased target gives a better denoiser.
The SWIFT target is biased on purpose: its mean is \((1+\alpha)D^\star - \alpha D_{\mathrm{weak}}\), not \(D^\star\), so with unlimited training SWIFT would end up worse than ordinary training. The gain comes from finite training.
We model training as , for any data distribution with finite second moment. After time \(T\), an eigen-direction of the kernel with eigenvalue \(\kappa\) is learned only up to a fraction \(1-e^{-T\kappa}\), so ordinary training misses \(D^\star\) by \(e^{-T\kappa}\) along it. A guide stopped earlier, at \(T_g \lt T\), misses by more, \(e^{-T_g\kappa}\), along the same directions, so the extrapolation is largest exactly where training is still short. Trained toward the SWIFT target for time \(T\), the model's error along each direction is
$$ \Big(\underbrace{-\,e^{-T\kappa}}_{\text{finite-time shortfall}} \;+\; \underbrace{\alpha\,(1-e^{-T\kappa})\,e^{-T_g\kappa}}_{\text{SWIFT correction}}\Big)\, d^\star_\kappa , $$with \(d^\star_\kappa\) the component of \(D^\star\) along that direction. The correction has the opposite sign of the shortfall, so a small \(\alpha \gt 0\) reduces the error along every direction at once (Theorem 1); too large an \(\alpha\) overshoots, consistent with the U-shape in \(\alpha\) we see in experiments. The guide does not need to be good, only wrong in the same directions, and the guarantee concerns distance to the posterior mean rather than FID directly.
Write \(L\) for the kernel integral operator on \(L^2(P_{\bm{x}_t})\) and \(H_s = I - e^{-sL}\) for the response of zero-initialized gradient flow after time \(s\). Then \(D_{\mathrm{weak}} = H_{T_g}D^\star\), \(D_{\mathrm{base}} = H_T D^\star\) and \(D^{(\alpha)}_{\mathrm{SWIFT}} = H_T\big[(1+\alpha)D^\star - \alpha D_{\mathrm{weak}}\big]\). Along an eigen-direction of \(L\) with eigenvalue \(\kappa\), \(H_s\) multiplies by \(1-e^{-s\kappa}\), which gives the formula above; directions the kernel cannot learn receive no correction. The result is stated at one noise level and equal main-model training time, with the guide's own training not counted. The paper's appendix also treats a smaller, spectrally nested guide, finite-step training and finite samples.
Results on ImageNet-256.
On ImageNet, SWIFT uses one fixed recipe: \(\alpha = 0.4\) and a smaller pretrained guide, frozen before main-model training. Across 18 DiffusionBench configurations, spanning VAE latents, representation-autoencoder latents, end-to-end learned latents and pixel space, this recipe lowers unguided FID in every case, by 16–49% (median 31%), and at the best available guidance settings it lowers guided FID in all 18 as well. On the stronger RAEv2 model, SWIFT reaches FID 1.15 without guidance and 0.99 with guidance.
| Representation | FID-50K, no guidance | FID-50K, with guidance | ||
|---|---|---|---|---|
| Baseline | SWIFT | Baseline | SWIFT | |
| DiffusionBench · 615M · 80 epochs · CFG | ||||
| RAE: DINOv2-B | 2.09 | 1.75−16% | 1.98 | 1.59−20% |
| VAE: SD-VAE | 5.49 | 3.77−31% | 2.32 | 1.81−22% |
| End-to-end: Qwen-VAE | 4.94 | 3.23−35% | 1.50 | 1.40−7% |
| Pixel: PixNerd | 20.37 | 10.31−49% | 4.17 | 3.38−19% |
| RAEv2 · 875M · up to 160 epochs · internal guidance | ||||
| RAE: DINOv3-L | 1.56 | 1.15−26% | 1.04 | 0.99−5% |
Matched pairs. Same architecture, schedule and sampler; only the regression target differs. RAEv2's SWIFT values use its best checkpoint (epoch 140) within the 160-epoch budget.
Samples from the 0.99-gFID model (RAEv2 / DINOv3-L), 50-step ODE sampler with retuned internal guidance.
SWIFT stacks with inference-time guidance.
Guidance at sampling time still helps a SWIFT model: with autoguidance on the DiffusionBench models and internal guidance on RAEv2, SWIFT plus retuned guidance gives the best FID in all five configurations. It also needs less of it: on RAEv2 the best guidance scale drops from 1.78 to 1.20, while guided FID improves from 1.04 to 0.99.
SWIFT with inference-time guidance. (a) FID with autoguidance on the DiffusionBench models and internal guidance on RAEv2. (b) On RAEv2, SWIFT shifts the preferred internal-guidance scale from 1.78 to 1.20.
Training cost, and longer training.
SWIFT adds 9–12% FLOPs per training step for the frozen guide, plus a one-off 10–15 baseline-equivalent epochs to pretrain it. Counting both, on DINOv2-B it beats the baseline's best FID over 320 epochs (1.53) by epoch 120 (1.50), with 0.45× the training FLOPs; on DINOv3-L it beats the epoch-160 baseline (1.56) by epoch 50 (1.34), with 0.44× the training FLOPs. The lead also persists over training: on DINOv2-B through epoch 240, and from epoch 40 onward in every one of the 12 DiffusionBench configurations we tracked.
Training compute on DINOv2-B. Guide pretraining and the per-step overhead are both charged to SWIFT; the x-axis is in baseline-epoch equivalents.
How to choose \(\alpha\) and the guide.
CIFAR-10 ablations. (a) The coefficient \(\alpha\), with a half-width guide. (b) The guide's width, with the student at channel base 128. (c) How long the guide is trained before it is frozen.
- \(\alpha\): moderate positive values work, and the curve is U-shaped; \(\alpha\) between 0.4 and 0.8 works well on both CIFAR-10 and ImageNet. All ImageNet results use \(\alpha = 0.4\) without tuning.
- Guide: weaker than the student, but competent. A half-width guide works best (CIFAR-10 FID 2.58 → 1.60 compared with a full-width one); much smaller guides hurt, and past a competence threshold the guide's training length matters little.
- Keep it simple: a frozen guide and a fixed \(\alpha\). Co-training the guide gave no clear additional benefit.
Conclusion.
Finite training misses the posterior mean in a direction that a weaker model of the same kind reveals, and SWIFT turns this into a better regression target: aim past the clean sample, away from a frozen weak model. It costs about 9% more training compute, leaves inference untouched, and lowers FID in every ImageNet configuration we tried. More broadly, the regression target is a design choice in its own right, not a fixed quantity.
BibTeX
@article{he2026swift,
title={{SWIFT}: Self-guided Weak-to-strong Diffusion Models via Target Extrapolation},
author={He, Yanjin and Shen, Cuifeng and Lei, Wanzhou and Li, Maohua and Lan, Tao and Liu, Kan and Tang, Hanlin and Qu, Qing},
journal={arXiv preprint arXiv:2610.25844},
year={2026}
}