PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Yanjin He, Cuifeng Shen, Wanzhou Lei, Maohua Li, Tao Lan, Kan Liu, Hanlin Tang, Qing Qu

arXiv:2610.25844v1cs.LGcs.AIcs.CV

TL;DR

Diffusion models often rely on inference-time guidance, which adds sampling cost, and the paper asks whether weak-to-strong extrapolation can instead be moved into training. Swift trains on clean-sample targets extrapolated away from a frozen weak denoiser, then discards the guide. The paper reports gains across CIFAR-10 and ImageNet-256 and proves lower posterior-mean error than matched standard training under specified conditions.

  • Problem

    The paper asks whether weak-to-strong extrapolation can move from inference into training so the final denoiser benefits from the correction and the weak model can be discarded.

    Source: §1 · p. 2

  • Method

    Swift uses a frozen weak denoiser to extrapolate clean-sample regression targets, then discards the guide after main-model training.

    Source: §1 · pp. 2–3

  • Results

    Swift improves models across CIFAR-10 and ImageNet-256 representations, and a nonempty range of positive extrapolation strengths can yield strictly lower posterior-mean error than matched standard training.

    Source: §1 · pp. 2–3

  • Takeaways & Limitations

    Target extrapolation offers a training-time weak-to-strong correction while leaving a single denoiser and the original sampler at deployment.

    Source: §1 · p. 2, §6 · p. 10

    The theoretical comparison assumes matched population gradient-flow training in the same RKHS, with the weak guide taken from an earlier point on the ordinary training trajectory.

    Source: §4 · p. 6

Abstract

from arXiv · show

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.

1 Introduction

Swift moves weak-to-strong extrapolation from inference into training: a frozen weak denoiser guides extrapolated regression targets, then is discarded. The paper reports improvements across diffusion representations and a theoretical setting in which extrapolation can reduce posterior-mean error.

1 Introduction

  • 1.15 gFID without inference-time guidance and 0.99 with retuned internal guidance are reported for Swift on RAEv2/DINOv3-L.The figure also reports that Swift’s advantage persists throughout extended training in a matched RAE/DINOv2-B configuration.Source: Figure 1 · p. 2
  • For a nonempty range of positive extrapolation strengths, Swift can yield strictly lower posterior-mean error than matched standard training.The comparison is despite deliberately shifting the training target away from the posterior mean.Source: §1 · pp. 2–3
  • Across CIFAR-10 and ImageNet-256, Swift improves diffusion models spanning VAE latents, representation-autoencoder latents, learned latents, and pixel-space generation.The reported gains persist under longer training and after inference-time guidance is independently retuned.Source: §1 · pp. 2–3
  • With guidance independently retuned for both models, Swift achieves lower FID in every tested configuration.On RAEv2, the preferred internal-guidance scale falls from 1.78 for the baseline to 1.20 for Swift, while best guided gFID improves from 1.04 to 0.99.Source: §1 · p. 2
  • Swift uses a frozen weak denoiser to extrapolate clean-sample regression targets, then discards the guide at inference.The guide serves as a directional reference rather than a teacher to imitate.Source: §1 · pp. 2–3

2 Swift: Self-Guided Training of Diffusion Models

The preliminaries establish clean-sample denoising and the role of the trained denoiser in sampling. Swift changes the training target using a frozen weak guide, which is discarded before inference.

2.1 Preliminaries

  • The unrestricted population minimizer of clean-sample denoising is the posterior mean, while practical sampling uses the denoiser obtained after finite training.The sampler integrates the probability-flow ODE from high noise toward smaller noise levels.Source: §2.1 · p. 3

2.2 Proposed Method: Learning with a Frozen Weak Guide

  • Swift extrapolates the clean target away from a frozen weak denoiser’s prediction; α = 0 recovers standard training, while α > 0 extrapolates beyond the clean sample.The guide is trained ordinarily, frozen before the main run, and requires no backpropagation during main-model training.Source: §2.2 · pp. 3–4
  • With an early-stopped weak denoiser, theory establishes conditions under which Swift achieves strictly lower posterior-mean error than baseline training.The improvement depends on how directions where the weak guide remains inaccurate relate to the main model’s finite-training error.Source: §2.2 · p. 4
  • For a fixed training example and noise realization, Swift’s target stays fixed as the main denoiser updates; the guide is discarded and the original sampler is retained.Each training step adds one weak-denoiser forward pass, about 9% more FLOPs on DiffusionBench.Source: §2.2 · p. 4

3 From Autoguidance to Self-Guided Training

Swift moves weak-denoiser extrapolation from sampling into training, where it changes the regression target rather than adding a sampling-time correction. The section compares this approach with guidance methods and reports results on a branching toy distribution.

3 From Autoguidance to Self-Guided Training

  • Against GFT, the strongest baseline on the three reported metrics, Swift’s KL falls from 0.817 to 0.728, outlier rate from 4.97% to 3.66%, and log p rises from −0.40 to +0.08.The comparison uses a thin, branching Gaussian-mixture distribution; the highlighted regions show generated samples following the narrow branches.Source: §3 · p. 4, Figure 2 · p. 5

3.1 From Sampling-Time Guidance to Training Targets

  • Unlike autoguidance, which extrapolates a finite-time baseline denoiser during sampling, Swift uses the clean sample in the training target and discards the weak denoiser afterward.Under the stated assumptions, changing the target can improve approximation to the original posterior mean, but this is not guaranteed for every weak denoiser.Source: §3.1 · pp. 4–5

3.2 Comparison with Training-Based Guidance

  • Swift freezes a weak denoiser before main training, keeping each example’s correction fixed instead of using changing auxiliary predictions or supervising an additional output.Jointly optimizing weak and main denoisers provided no clear additional endpoint benefit in the CIFAR-10 comparisons; the weak guide still incurs training compute.Source: §3.2 · pp. 5–6

4 Why Swift Can Help

The analysis compares standard denoising with Swift after equal finite main-model training time at a fixed noise level. Under stated matched-training assumptions, a nonempty range of positive extrapolation strengths can strictly reduce posterior-mean error.

Analysis setting

  • The comparison uses population gradient flow at one noise level, with matched data law, kernel, initialization, and main-model training horizon.Guide preparation is excluded from the theorem's equal-time comparison.Source: §4 · pp. 6–7

Finite-time improvement

  • For a nonempty range of positive α, Swift achieves strictly lower posterior-mean risk than matched baseline training when the guide has nonzero response.The result concerns population denoising MSE at equal main-model optimization time, not FID.Source: §4 · pp. 6–7

Why extrapolation can help

  • An early-stopped guide has learned less of each learnable mode, so small positive extrapolation can partially cancel the baseline's finite-time residual, while larger α can overshoot.Improvement depends on how the guide's inaccurate directions relate to the main denoiser's finite-training error.Source: §4 · pp. 6–7

5 Experiments

Experiments evaluate Swift across ImageNet-256 configurations and compare it with matched baselines, including under retuned inference guidance and longer training. Gains persist across representations, while the results also identify useful choices for extrapolation strength and guide capacity.

5 Experiments

  • The evaluation compares baseline and Swift with and without inference-time guidance across 19 ImageNet-256 configurations, reporting FID and additional distribution and representation metrics where available.Architectures, schedules, samplers, and evaluation pipelines otherwise match the baseline; the ImageNet recipe uses α = 0.4 and a smaller frozen pretrained guide.Source: §5 · p. 7

5.1 Generation Quality and Transfer Across Representations

  • Across 18 configurations, Swift lowers unguided FID by 16.3%–49.4% (median 31.3%); RAE/DINOv2-B unguided FID falls from 2.09 to 1.75.The five-feature mean FDr for RAE/DINOv2-B falls from 4.13 to 3.69; on RAEv2/DINOv3-L, unguided FID falls from 1.56 to 1.15.Source: §5.1 · p. 7, Table 1 · p. 8

5.2 Combining Swift with Inference-Time Guidance

  • With independently retuned inference guidance, Swift gives the best FID in all five tested configurations; on RAEv2, the preferred scale shifts from 1.78 to 1.20.On RAEv2, best guided FID improves from 1.04 to 0.99.Source: §5.2 · p. 9, Figure 4 · p. 9

5.3 Longer Training and Training Cost

  • Swift surpasses the best DINOv2-B baseline FID observed through 320 epochs using approximately 0.45× the FLOPs needed to reach that baseline checkpoint.The matched Swift run remains ahead at all 36 evaluations from epochs 20 to 240.Source: Figure 5 · p. 9, §5.3 · p. 9
  • Baselines do not reach Swift's best FID within the observed training horizon, so their time-to-quality ratios are conservative.Baseline time-to-quality for those targets is not observed.Source: §5.3 · p. 9

5.4 How to Choose α and the Weak Guide

  • On CIFAR-10, moderate positive α and a competent half-width guide perform strongly; at α = 0.8, reducing guide width from 128 to 64 improves FID from 2.58 to 1.60.The fixed broad-experiment default is α = 0.4; the precise optimum varies by model.Source: §5.4 · p. 10, Figure 6 · p. 10

6 Conclusion

Swift treats the regression target as a design choice, moving a weak-to-strong correction into training. Despite using a deliberately biased target, finite training can yield a better predictor and better samples.

Target extrapolation

  • Swift intervenes in diffusion training by carrying the clean target past itself and away from a frozen weak model.This moves a weak-to-strong correction from inference into training.Source: §6 · p. 10

Effect and deployment

  • The modified target is deliberately biased, yet finite training on it can produce a better final predictor and better samples.Deployment retains a single denoiser and the original sampler, with about 9% additional main-training FLOPs.Source: §6 · p. 10

Implication

  • The results suggest treating the regression target as a design interface for writing useful corrections into training.The conclusion raises the question of what other corrections could be expressed this way.Source: §6 · p. 10

Reproducibility statement

The paper specifies implementation, evaluation, compute accounting, experiments, and theory details across its supplementary sections. The implementation is stated to be made public upon publication.

Implementation

  • Swift changes only the regression target; the respective framework's remaining training loop is retained.Section D.1 gives architectures, freezing points, coefficients, budgets, hardware, and the velocity-target implementation.Source: Reproducibility statement · p. 11

Evaluation and results

  • Sections D.2 and D.3 specify evaluation choices and report GPU hours and FLOPs accounting.Section E gives ImageNet-256 results across configurations and metrics, while Section F reports CIFAR-10 experiments.Source: Reproducibility statement · p. 11

Theory and availability

  • The paper provides the matched Gaussian-mixture comparison and assumptions and proofs in Sections C and A, respectively.The implementation is stated to be made publicly available upon publication.Source: Reproducibility statement · p. 11

AI use statement

The appendix states that LLMs assisted with editing, drafting, figure and table code, and summary-statistic computation, with authors reviewing the content. It then outlines proofs and extensions of the finite-time analysis, covering target identities, kernel training responses, spectral comparisons, and finite-step results.

AI use statement

  • LLMs edited and drafted text, generated LaTeX and TikZ for tables and figures, and computed summary statistics; the authors reviewed and revised all assisted content and take responsibility.The appendix overview lists proofs, guidance comparisons, a Gaussian-mixture comparison, implementation and compute details, ImageNet-256 and CIFAR-10 results, and additional samples.Source: AI use statement · p. 11, Appendix · p. 14

A.1 Target identity and the quantity being compared

  • The target identity analysis shows that Swift's target shift changes the training response, not the unrestricted population optimum; posterior-mean risk and clean-target MSE rank predictors identically.The unrestricted population minimizer of the Swift regression loss is gα.Source: §A.1 · p. 14

A.2 Kernel training response

  • Population gradient flow in prediction space depends on the conditional-mean target, while the target's conditional variance affects loss value but not the trajectory.The analysis develops the operator properties and training-response identities used in the spectral proof.Source: §A.2 · p. 15

A.3 A general commuting-response theorem

  • For commuting population responses, the posterior-mean risk is a quadratic in extrapolation strength, and it is lower than baseline risk exactly when 0 < α < 2α⋆PM.When A > 0, the unique minimizer is α⋆PM = A/C > 0.Source: §A.3 · p. 16

A.4 Spectral comparison and proof of Theorem 1

  • With a strictly weaker early-stopped guide and a nonzero ordinary predictor, the spectral comparison gives a nonempty range of positive extrapolation strengths that improves posterior-mean risk.The stated boundary is R(α) < R(0) exactly when 0 < α < 2A/C; an exact orthogonal-projection response instead gives A = 0 and no improvement.Source: §A.4 · p. 17

A.5 Gaussian-mixture posterior and learnability

  • For Gaussian-mixture data with a Gaussian kernel, the kernel operator is injective and ordinary training yields a nonzero predictor whenever the posterior mean is nonzero.This verifies the theorem's nonzero-predictor condition without requiring a deliberately underpowered function class.Source: §A.5 · pp. 18–19

A.6 Reduced-rank and spectrally nested guides

  • The reduced-rank guide guarantee applies when the narrower guide corresponds to a nested spectral subspace; arbitrary architectural width reduction is not covered without an additional approximation argument.Under the stated nonzero-response condition, Swift improves posterior-mean risk for 0 < α < 2A/C.Source: §A.6 · pp. 19–20

A.7 Exact label-affine response for finite-step training

  • With fixed features and label-independent initialization and step sizes, finite-step training produces parameters affine in the labels; label-dependent feature learning is not covered without further assumptions.The argument also covers pre-specified minibatch schedules, fixed linear preconditioners, weight decay, and momentum after augmenting the state.Source: §A.7 · p. 20

A.8 Finite-sample bias–variance extension

The finite-sample extension accounts for label-noise amplification as well as residual alignment when assessing target extrapolation. Its improvement condition is conditional, and the analysis has explicit assumptions and scope limits.

A.8 Finite-sample bias–variance extension

  • Under conditional centering, the finite-sample risk reduction is exact, but centering is not automatic when the guide was trained on the same clean labels.The argument also assumes a label-affine predictor and square-integrability.Source: §A.8 · pp. 20–21
  • With finite label noise, a beneficial extrapolation strength exists only when residual alignment exceeds variance inflation, A > V.When C + V > 0, the risk-minimizing coefficient is conditional on the realized inputs and guide; if C + V = 0, risk is constant in α.Source: §A.8 · p. 21

A.9 Scope of the analysis

  • The finite-sample result requires conditional centering and integrability, while the finite-step result assumes fixed features.The population argument needs square-integrable targets, not Gaussianity.Source: §A.9 · p. 22
  • The analysis compares fixed-noise denoising MSE at equal main-model optimization time; it does not guarantee improved FID or iterative sampling.Guide preparation and forward passes add compute, and neural feature learning need not meet the analysis conditions.Source: §A.9 · p. 22

B Additional Comparisons with Diffusion Guidance Methods

This block compares diffusion-guidance methods by their training references and sampling requirements, then tests adapted objectives against the class-A posterior mean. The reported Gaussian-mixture comparison favors Swift in the matched setup, with scope and evaluation conditions specified.

B Additional Comparisons with Diffusion Guidance Methods

  • CFG and autoguidance retain ordinary denoising targets and extrapolate predictions at sampling, whereas Swift puts the clean sample into the extrapolated training target.Autoguidance uses a separately trained weak model at sampling; Swift discards its fixed reference after training.Source: B Additional Comparisons with… · pp. 22–23

C Distance to the optimal denoiser on a Gaussian mixture

  • The Gaussian-mixture experiment scores each method’s implied posterior-mean estimate, converting Swift’s raw output because its target is shifted away from the posterior mean.The selected MG and GFT references are frozen, and the tested IG arm uses auxiliary supervision rather than IG’s target-correction variant.Source: C Distance to the optimal denoiser on… · p. 23
  • At d=16, Swift has the lowest distance to the class-A posterior mean at all five noise levels, reducing baseline distance by 12.2–17.2%.At d=8, the selected Swift row is also lowest at all five levels; the small GFT margins and five-seed comparison are specific to this matched setup, not a ranking of the original implementations.Source: C Distance to the optimal denoiser on… · p. 25, Table 4 · p. 25

D Experimental Setup

  • Matched Swift comparisons change the training target while keeping architecture, optimizer, schedule, sampler, and evaluator fixed.The frozen guide receives the same noisy input, noise level, and class condition, and is used without gradients.Source: D Experimental Setup · p. 25

D.1 Training Configurations

  • The ImageNet-256 DiffusionBench sweep uses α=0.4 with a separate half-width guide pretrained for half the standard training duration and then frozen.CFG is tuned independently for the baseline and Swift; ImageNet velocity-prediction configurations implement the equivalent clean-target extrapolation in velocity space.Source: §D.1 · p. 26

D.2 Sampling and Evaluation Settings

  • Evaluations use 50,000 generated samples per checkpoint, and guidance strengths are swept separately for each model in guided re-evaluations.RAEv2’s headline epoch-140 values are not a same-epoch comparison; its matched comparison is at epoch 160.Source: §D.2 · pp. 26–27

D.3 Compute Accounting

Swift improves unguided and guided FID across the reported ImageNet-256 configurations, while its guide adds training compute. The block also distinguishes published comparisons from re-evaluated results and details RAEv2 performance.

D.3 Compute Accounting

  • For one 80-epoch DiffusionBench Swift run, guide and main training together cost about 98 baseline-equivalent epochs, or 1.22 times an 80-epoch baseline.The guide adds about 9% FLOPs per update; once pretrained, each additional main run incurs only that per-update increase.Source: §D.3 · p. 27
  • For a 50-epoch RAEv2 Swift model, guide preparation and training together cost about 71 baseline-equivalent epochs, or 0.44 times a 160-epoch baseline.The guide adds about 12% FLOPs per main-training update, and its pretraining costs about 15 baseline-equivalent epochs.Source: §D.3 · p. 27

E.1 All configurations

  • Swift lowers unguided FID in all 18 DiffusionBench configurations by 16.3–49.4% (median 31.3%) and improves RAEv2 from 1.56 to 1.15.The same α=0.4 recipe is used without per-configuration coefficient tuning.Source: §E.1 · p. 28, Table 6 · p. 29
  • Swift has lower guided FID in all 19 configurations: by 0.10–1.31 (6.7–33.0%, median 21.9%) across DiffusionBench, and from 1.04 to 0.99 on RAEv2.Guidance is independently tuned for re-evaluated pairs; published guided baselines retain their reported settings.Source: §E.1 · p. 28, Table 7 · p. 30

E.2 Published results in Figure 1

  • Table 8 reports published ImageNet-256 FID-50K values from the corresponding papers rather than re-evaluations, and their evaluation procedures differ.Some papers sample 50 images per class and others sample class labels at random, a difference reported to change FID by about 0.1 [38, Table 14].Source: Table 8 · p. 31, §E.2 · p. 31

E.3 RAEv2

  • RAEv2’s canonical results are 1.56→1.15 without guidance and 1.04→0.99 with independently tuned internal guidance.The unguided 1.15 checkpoint is from epoch 140 within the 160-epoch budget; Swift is already below the released epoch-80 checkpoint by epoch 40.Source: §E.3 · p. 31

E.4 Training trajectories and the long schedule

The training-trajectory analyses test when Swift’s FID advantage appears and whether it persists across guides, configurations, and longer training. The reported matched comparisons show Swift ahead from epoch 40 onward, including through 240 epochs.

E.4 Training trajectories and the long schedule

  • Across 12 DiffusionBench configurations with training-time baseline evaluations, Swift is ahead at every evaluation from epoch 40 onward.Figure 9 shows four trajectories spanning baselines with roughly an order of magnitude in FID.Source: §E.4 · p. 32, Figure 9 · p. 32
  • On the matched 80-epoch DINOv2-B ablation, α=0.4 and 0.8 with the default guide lag at epoch 20 but lead from epoch 40 onward; an earlier guide leads throughout.The default-guide arms lead by 0.28–0.48 FID from epoch 40, while the earlier-guide α=0.4 arm leads by 0.47–0.59 FID throughout.Source: §E.4 · p. 32
  • In the matched 240-epoch DINOv2-B pair, Swift has lower FID at all 36 evaluations from epoch 20 onward, with a 0.26 FID gap at epoch 240.The comparison uses the fixed α=0.4 default frozen-guide recipe.Source: §E.4 · p. 33, Table 9 · p. 33, Figure 10 · p. 33

F Additional CIFAR-10 experiments

Additional CIFAR-10 experiments vary guide designs and extrapolation schedules under a fixed evaluation setup, while accompanying figures show class-conditional ImageNet samples from Swift on RAEv2/DINOv3-L.

F Additional CIFAR-10 experiments

  • The additional CIFAR-10 experiments use a class-conditional 32 × 32 setup with a DDPM++ student and evaluate FID-50k using the 18-step EDM sampler.The student uses channel base 128, batch size 512, and 200k kimg; unless varied, the frozen DDPM++ guide uses channel base 64 at 200k.Source: F Additional CIFAR-10 experiments · p. 34
  • Guide variants replace the DDPM++ guide, co-train it with the student, replay checkpoints of different training ages, or change guide capacity during training.Ahead-of-student checkpoint variants require pre-trained future checkpoints; capacity schedules include 16 → 32 → 64 and an optional width-96 stage.Source: F Additional CIFAR-10 experiments · p. 34
  • Coefficient variants change extrapolation strength over training or noise level, including piecewise profiles, negative-α prefixes, random coefficients, and online loss-probe adjustments.The tested training-time profiles include middle-only pulses and late-onset peak/hold schedules.Source: F Additional CIFAR-10 experiments · p. 34
  • Combined methods multiply a training-time coefficient profile by a noise-dependent factor, and the accompanying table reports each method’s lowest completed-run FID at the common 200k endpoint.The results comprise 156 run labels, 152 scored; online adaptive rules are preliminary, and the table entries are not run averages.Source: F Additional CIFAR-10 experiments · p. 34, Table 10 · p. 34

G Additional samples

  • Figures 11–14 show 28 samples per class for 32 ImageNet classes from Swift on RAEv2/DINOv3-L, generated with retuned internal guidance and a 50-step Euler ODE sampler.The examples span birds and mammals, marine life, vehicles and objects, and fruit, food, and plants.Source: G Additional samples · p. 35, Figure 11 · p. 35, Figure 12 · p. 36

Related papers

Loading 2610.25844v1…