SWIFT: Self-guided Weak-to-strong dIFfusion models via Target extrapolation

SWIFT. Self-guided Weak-to-strong dIFfusion models via Target extrapolation. By Yanjin He and seven others, from the University of Michigan, Ann Arbor, Alibaba Group, and two other institutions. Published on the nineteenth of October. No code repository is linked to it.

Training a high-quality diffusion model is computationally expensive. And the guidance applied at inference time, to improve samples, often adds further sampling cost on top. So the cost lands twice: once to train the model, and again every time it generates. The question is whether a stronger model can be trained in a way that needs nothing extra when it is used.

Swift is a self-guided training method that uses a frozen weak denoiser to improve a stronger model, without adding inference overhead. For every noisy input, Swift extrapolates the clean regression target away from the weak denoiser's prediction. The effect is to amplify the components that are still underfit. Where the weak model's prediction already sits, the target is pushed further on, so the stronger model is asked to go past what the weaker one learned rather than merely to match it. The guidance happens to the target during training, not to the sampler during generation.

The weak guide is cheap to obtain. It can be an early checkpoint from training. It needs no backpropagation while the main model trains, and it is thrown away afterward. The resulting model keeps the original sampler. It supports both conditional and unconditional generation, without label dropout and without an unconditional branch. Once training ends, nothing of the weak guide remains in the deployed model.

The authors explain the mechanism through a spectral analysis of finite-time population gradient flow in an RKHS. With an early-stopped weak guide, they prove that there is a nonempty range of extrapolation strengths giving strictly lower posterior-mean error than matched standard training. That holds even though the regression target has been shifted away from the posterior mean. Moving the target off the quantity being estimated, and still estimating it better, is the counterintuitive part the analysis is there to explain.

Experiments on CIFAR-10 and ImageNet-256 show 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, counting guide preparation and training overhead. On RAEv2, it reaches an FID of one point one five without inference-time guidance, and zero point nine nine with guidance.

Aiming slightly past the right answer sounds like a mistake, and the proof is there to show it is not. Target extrapolation is offered as a route to weak-to-strong diffusion training in which the help comes during training, and inference stays a single model.