Source-linked AI summary
Unlocking Few-Step Diffusion for Faithful Previews
Jing Jia, Sifan Liu, Guanyang Wang
TL;DR
Repeatedly sampling candidates can make image search costly, and ordinary few-step samplers may not match their full-step counterparts. The paper learns corrections to the input noise or denoising trajectory; experiments report improved reference fidelity and candidate selection, with input correction reusable across sampling budgets without retraining.
Problem
Image search can require repeated costly sampling, while standard few-step outputs may differ substantially from corresponding full-step images.
Source: §1 · pp. 1–2Method
Results
Takeaways & Limitations
Corrected few-step previews let users screen candidates and reserve full-step generation for promising ones, while input corrections can transfer across sampling budgets without retraining.
Source: §1 · p. 2, §5 · p. 9Learned correctors remain less accurate than target-aware oracle noise optimization, and the findings need evaluation on larger models and broader transfer settings.
Source: §5 · p. 9
Abstract
from arXiv · showhide
Sampling latency compounds in diffusion workflows, where users generate and discard many candidates before keeping one. Surprisingly, the poor outputs of standard few-step samplers do not reflect a lack of reconstruction capacity: by optimizing only the initial noise, frozen 3-4-step samplers can closely reproduce their corresponding full-step outputs. Building on this finding, we learn corrections to the initial noise and denoising updates using endpoint supervision, improving correspondence with full-step outputs generated from the same noise and prompt. The resulting previews allow users to screen candidates cheaply and reserve full-step generation for promising ones. Input correction also transfers across sampling budgets without retraining. Experiments show substantial improvements in reference fidelity, including 53-78% lower reconstruction MSE than retrained LD3 on unconditional benchmarks, alongside improved ranking preservation and candidate selection on SD1.5, SDXL, and FLUX.1-dev.
1 INTRODUCTION
Few-step previews could reduce the cost of searching for images, but standard samplers often fail to match their corresponding full-step outputs. The paper learns corrections that improve this instance-wise correspondence and support candidate screening.
1 INTRODUCTION
- Because image search repeats sampling across candidates, users could screen corrected few-step previews and reserve full-step generation for images they choose to keep.A single image can require tens to hundreds of denoising steps.Source: §1 · pp. 1–2
- Standard three- or four-step DDIM often produces blurry outputs unlike full-step counterparts, while sharper distilled samples do not necessarily match their corresponding references.The target is instance-wise correspondence for the same initial noise and prompt.Source: §1 · pp. 1–2
- Oracle optimization of only the initial noise reduces reconstruction MSE by 99.1–99.5% across four unconditional models with three-step sampling and by 98.6% for four-step FLUX.1-dev.The optimization uses full-step reference images and frozen samplers.Source: §1 · p. 2
- Learned input corrections improve fidelity to full-step references, generalize across sampling budgets without retraining, and trajectory corrections further improve correspondence on text-to-image models.Trajectory corrections use low-rank adapters while keeping pretrained base weights frozen.Source: §1 · p. 2
- Experiments on SD1.5, SDXL, and FLUX.1-dev show gains in ranking preservation and candidate selection over evaluated distilled models, learned preview solvers, and probe-based predictors.Corrected few-step previews let users screen candidates before reserving full-step generation for promising ones.Source: §1 · p. 2
- Unlike acceleration methods focused on fast, high-quality generation, this work targets instance-wise fidelity to a full-step teacher for previews.The paper contrasts its endpoint-supervised corrections with prior noise refinement and distribution-level objectives.Source: §1 · p. 2
2 METHOD
The method trains corrections to make a few-step sampler match a fixed full-step reference on the same input. It supports changes to the starting noise or denoising trajectory and uses corrected outputs as previews for candidate selection.
2 METHOD
- The training objective minimizes image reconstruction error between the corrected few-step output and the full-step reference for the same input.Training samples (z, c), computes F_ref(z, c), and optimizes correction parameters through the full few-step sampler while freezing pretrained backbone weights.Source: §2 · p. 3
- Input correction predicts a residual added to the initial noise, while trajectory correction modifies denoising updates using step-specific low-rank adapters.Trajectory adapters are jointly trained with final-output supervision, requiring no intermediate reference states and treating the reference sampler as a black box.Source: §2 · p. 3
- For text-to-image diffusion, the correction model takes classifier-free guidance scale as an input because changing it changes the generation trajectory.The scale is encoded with Fourier features and a learned projection for continuous modulation.Source: §2 · p. 3
- A corrector trained at k steps can be reused at K steps without retraining by scaling its predicted noise displacement, for example by α_K = (k/K)^2.A three-step corrector can use 0.36 of its predicted displacement for five-step sampling.Source: §2 · pp. 3–4
- Corrected few-step outputs serve as previews for choosing or discarding candidates, with full-step generation run only for a selected candidate using its original input.Selection can use a task-specific metric or human judgment because training minimizes image reconstruction error rather than a particular selection metric.Source: §2 · pp. 3–4
3 INSTANCE-WISE CORRESPONDENCE THROUGH FEW-STEP SAMPLERS
Frozen few-step samplers can closely reproduce corresponding full-step outputs through noise optimization, motivating learned input and trajectory corrections. Across unconditional and text-to-image evaluations, these corrections improve teacher correspondence, and input correction remains useful at other sampling budgets without retraining.
3.1 HIDDEN RECONSTRUCTION CAPACITY OF FROZEN FEW-STEP SAMPLERS
- Noise optimization raises mean PSNR from 14.62–17.00 dB to 37.41–39.98 dB across four unconditional models, and from 14.54 to 37.49 dB for four-step FLUX.1-dev.These target-aware experiments optimize only the initial noise of frozen samplers and motivate learning input correction without target access or per-image inference optimization.Source: §3.1 · p. 4
3.2 INPUT CORRECTION
- Without retraining, transferred input correction improves PSNR over raw DDIM at every tested budget and also reduces MSE and LPIPS while increasing SSIM, HPS, and ImageReward.At transferred budgets, p = 2.5 gives the highest PSNR, while p = 1.7 gives the highest HPS; larger exponents favor teacher fidelity and smaller ones perceptual quality.Source: §3.2 · pp. 5–6
- Ablations support noise-norm normalization, sample-specific correction directions, and reduced correction magnitudes when transferring to more sampling steps.Source: §3.2 · p. 6
- In matched CIFAR-10 and SD1.5 comparisons, trajectory correction raises PSNR from 26.60 to 29.79 dB and from 20.60 to 23.25 dB, respectively, over input correction.Both correction strategies substantially outperform uncorrected three-step sampling.Source: §3.2 · p. 6
3.3 TRAJECTORY CORRECTION
- Across SD1.5 and SDXL at CFG scales 1, 3, and 5, trajectory correction reduces LPIPS by 18.3%, increases SSIM by 18.7%, and improves PSNR by 2.32 dB against the best reported baselines.Stronger classifier-free guidance makes trajectories harder to approximate under step compression, while trajectory correction mitigates the resulting degradation.Source: §3.3 · p. 6
3.4 TRANSFER ACROSS BACKBONES
- Cross-backbone input correction improves CelebA-HQ teacher-reconstruction PSNR from 14.78 to 22.82 dB and exceeds the full-step teacher in HPS, ImageReward, and PickScore.The authors hypothesize that backbone mismatch regularizes correction toward perceptually preferred structure.Source: §3.4 · p. 6
3.2 INPUT CORRECTION
- 78.0%, 53.2%, and 71.6% lower MSE than retrained 4-step LD3 on CIFAR-10, Church, and CelebA-HQ, respectively, with input correction improving all four metrics over evaluated baselines.The method uses three denoiser evaluations and one lighter corrector evaluation; examples show closer matches in object shape, architecture, and facial details.Source: §3.2 · p. 5, Table 2 · p. 7
4 PRESERVATION OF CANDIDATE RANKINGS AS PREVIEW
The section evaluates whether few-step previews preserve candidate rankings and support candidate selection relative to full-step references. On SD1.5 and SDXL, ranking agreement improves over evaluated baselines; on FLUX.1-dev, preview-based selection leads the non-oracle methods on all five metrics.
4 PRESERVATION OF CANDIDATE RANKINGS AS PREVIEW
- The evaluation compares preview and teacher rankings across eight noises per prompt using automated scores, with Spearman correlation measuring ranking agreement.It reports regret and selected teacher score alongside Spearman correlation; regret is the gap between the best teacher score and the selected candidate’s teacher score.Source: §4 · p. 7, §4.1 · pp. 7–8
4.1 STABLE DIFFUSION FAMILY
- Mean Spearman correlation rises from 0.298 for ConSolver to 0.387 on SD1.5, and from 0.319 for LD3 to 0.421 on SDXL.The section reports consistent preservation across guidance scales and model families, with similar trends in agent-based evaluation.Source: §4.1 · p. 8, Table 3 · p. 8
4.2 IMAGE SELECTION ON FLUX.1-DEV
- On FLUX.1-dev, metric-specific selection chooses among eight previews for 512 COCO prompts, with probes, ConSolver, TDD-FLUX, random selection, and a full best-of-eight oracle as comparisons.The final candidates use 28-step FLUX.1-dev generation; all methods select from the same fixed candidate pool.Source: §4.2 · pp. 8–9
- 1.1417 ImageReward is the preview method’s mean, versus 1.0331 for TDD-FLUX and 1.0085 for ConSolver; it leads non-oracle selection on all five metrics.The paired ImageReward gain over TDD-FLUX is 0.1086 (95% prompt-bootstrap confidence interval [0.0654, 0.1509]); the preview workflow uses 60 NFE versus 224 for full best-of-eight.Source: §4.2 · p. 9, Table 4 · p. 9
5 DISCUSSION, LIMITATIONS, AND FUTURE DIRECTIONS
The discussion presents corrected few-step samplers as a way to screen candidates before full-step generation, while identifying accuracy and generality as open issues.
5 DISCUSSION, LIMITATIONS, AND FUTURE DIRECTIONS
- Preview–select–render can potentially reduce the time to obtain a satisfactory image by reserving full-step generation for selected candidates.The paper identifies image editing, inverse problems, and helping smaller models approximate larger teachers as possible future directions.Source: §5 · p. 9
- Learned correctors remain less accurate than target-aware oracle noise optimization, and their generality and transfer conditions need broader evaluation.The paper also calls for mechanistic investigation of when transfer succeeds and which properties of learned corrections enable it.Source: §5 · p. 9
A CORRECTION OBJECTIVES AND METHOD DETAILS
The appendix details input and trajectory corrections, guidance conditioning, and transfer across sampling budgets. For unconditional input correction, the corrected noise preserves the original norm; trajectory adapters are trained jointly through the rollout.
A.1 INPUT CORRECTION DETAILS
- For unconditional input correction, the predicted residual is added to the initial noise and the result is normalized to preserve the original noise norm.Normalization is per sample with a small denominator floor; training uses α = 1, and gradients pass through the frozen few-step sampler toward the original-noise reference target.Source: §A.1 · p. 13
A.2 TRAJECTORY CORRECTION DETAILS
- Trajectory correction uses independent LoRA adapters at each denoising step, trained jointly through the complete student rollout.Supervision covers both final output and terminal latent, requires endpoint reference latents and images, and does not require intermediate reference states.Source: §A.2 · p. 13
A.3 CLASSIFIER-FREE GUIDANCE CONDITIONING
- Stable Diffusion conditions its correction on Fourier-feature guidance-scale embeddings, while FLUX uses its native guidance embedding.The FLUX reference uses guidance 3.5 and the trained student uses 1.0, without a separate unconditional branch.Source: §A.3 · p. 13
A.4 TRANSFER ACROSS SAMPLING STEPS
- A frozen input corrector trained with k steps transfers to K steps by scaling its predicted noise displacement by α = (k/K)^p before norm normalization.The reported settings use k = 3, K ∈ {3, 4, 6, 8}, and p ∈ {1.7, 2.0, 2.5}.Source: §A.4 · p. 13
B EXPERIMENTAL SETUP AND IMPLEMENTATION DETAILS
The experimental setup uses pretrained unconditional and text-to-image diffusion backbones. It specifies oracle noise-optimization settings and the correctors and timestep choices used for input-correction training.
B EXPERIMENTAL SETUP AND IMPLEMENTATION DETAILS
- Experiments use pretrained unconditional DDPM backbones for CIFAR-10, LSUN-Church, CelebA-HQ, and LSUN-Bedroom, and SD1.5, SDXL, and FLUX.1-dev for text-to-image.Source: B EXPERIMENTAL SETUP AND… · p. 13
B.1 ORACLE NOISE OPTIMIZATION
- Oracle noise optimization uses 200 Adam updates per image and 100 seeds per unconditional dataset with a frozen three-step DDIM sampler.References use 1,000-step DDIM for CIFAR-10 and 20-step DDIM for LSUN-Church, CelebA-HQ, and LSUN-Bedroom; the FLUX oracle uses four prompt–noise pairs, four-step students, and 2,560 updates.Source: §B.1 · p. 14
B.2 INPUT CORRECTION TRAINING
- Input correction uses a residual U-Net for LSUN-Church and adds attention blocks for CelebA-HQ; all students use DDIM timesteps [666, 333, 0].Reference budgets are 1,000 steps for CIFAR-10 and 20 steps for Church and CelebA-HQ.Source: §B.2 · p. 14
B.3 TRAJECTORY CORRECTION TRAINING
The section describes trajectory-correction training, an agent-based visual-quality evaluation protocol, and FLUX.1-dev image-selection baselines. It also specifies training configurations for SD models and FLUX.
B.3 TRAJECTORY CORRECTION TRAINING
- SDXL and SD1.5 use three-step sampling with independent rank-64 LoRA adapters at each step, while CFG scale controls continuous rank-wise gates.Training uses 4,000 prompts at five CFG scales; the FLUX student uses four rank-64 adapter sets and 36,864 training pairs.Source: §B.3 · p. 14
B.4 AGENT-BASED VISUAL-QUALITY EVALUATION
- The visual-quality evaluator scores and ranks eight candidates per prompt, with candidate order randomized and preview and reference groups evaluated separately.It receives the prompt and images, but not method, sampling budget, guidance scale, original position, or noise seed.Source: §B.4 · p. 14
B.5 FLUX.1-DEV IMAGE SELECTION
- Diffusion Probe scores 0.9731 ImageReward versus 0.9483 for random selection on the 512 × 512 FLUX.1-dev benchmark; retraining did not yield a statistically significant improvement.The highest retraining score was 0.9785, with a paired gain of 0.0054 and 95% confidence interval [−0.0392, 0.0506].Source: §B.5 · p. 15
- Probe-Select is adapted by training its released prediction-head architecture on FLUX features, while ConSolver adapts a flow-matching implementation to four-step FLUX.1-dev generation.Probe-Select uses 35,332 image–score pairs with the FLUX backbone frozen; ConSolver trains a time-conditioned coefficient MLP while keeping its backbone, VAE, and reward model frozen.Source: §B.5 · p. 15
C.1 ADDITIONAL SD FAMILY RESULTS
Trajectory correction improves SD1.5 and SDXL results at higher guidance scales. Input correction trained at three steps also improves metrics when transferred to larger sampling budgets without retraining.
C.1 ADDITIONAL SD FAMILY RESULTS
- Trajectory correction consistently leads the reported baselines across four image-level similarity metrics for SD1.5 and SDXL at CFG 7 and 9, with especially pronounced improvements at CFG 9.The reported metrics are MSE, LPIPS, SSIM, and PSNR.Source: §C.1 · pp. 15–16
C.2 EXTENDED TRANSFER ACROSS SAMPLING STEPS
- At transferred budgets K ∈ {3, 4, 6, 8}, every tested correction improves all six metrics over raw DDIM at the same step count.With 20-step DDIM outputs as references, p = 2.5 gives the lowest MSE and LPIPS and highest PSNR, while p = 1.7 gives the highest HPS.Source: §C.2 · p. 16, Table 5 · p. 16
C.3 CONTROLLED COMPARISON OF INPUT AND TRAJECTORY CORRECTION
Under matched training conditions, both input and trajectory correction improve fidelity over raw sampling, with trajectory correction achieving higher fidelity in the reported comparison. Cross-backbone transfer to CelebA-HQ improves preference scores, while in-domain correction yields closer teacher agreement.
C.3 CONTROLLED COMPARISON OF INPUT AND TRAJECTORY CORRECTION
- Trajectory correction improves PSNR over input correction by 3.19 dB on CIFAR-10 and 2.65 dB on SD1.5, while reducing LPIPS from 0.025 to 0.011 and from 0.408 to 0.189, respectively.Both correction methods improve all four fidelity metrics over raw sampling; the learned methods share the pretrained backbone, paired data, reference sampler, endpoint objective, and update budget.Source: §C.3 · pp. 16–18, Table 6 · p. 17
- Input correction adds one full U-Net call, whereas trajectory correction operates within the existing three calls.Input correction predicts a residual adjustment to the initial noise; trajectory correction uses separate adapters at each denoising step.Source: §C.3 · pp. 17–18
C.4 CROSS-BACKBONE TRANSFER
- The Bedroom-trained corrector transferred to CelebA-HQ improves PickScore, HPS, and ImageReward over raw DDIM3 by approximately 1.34, 10.81, and 1.52, respectively.It was trained using cached CelebA 20-step teacher targets; the in-domain CelebA corrector achieves closer teacher agreement across all four fidelity metrics.Source: §C.4 · p. 18, Table 7 · p. 18
C.5 ADDITIONAL FLUX EVALUATIONS
Additional FLUX evaluations find that corrected previews improve reference fidelity and candidate selection against evaluated alternatives. These gains also appear under equal-budget screening, at higher resolution, and on a different prompt collection.
C.5 ADDITIONAL FLUX EVALUATIONS
- Our four-step FLUX.1-dev previews achieve 18.954 dB PSNR versus 11.357 dB for Schnell, and preview-based selection yields final reference-image scores of 1.1361 versus 0.9514.Schnell has higher ImageReward on the previews themselves, but the proposed method more closely matches the 28-step reference.Source: §C.5 · p. 18
- At equal 60-NFE budgets, eight corrected four-step previews yield selected-image ImageReward of 1.1361 versus 1.0573 for four raw eight-step previews.The paired gain is 0.0788, with 95% confidence interval [0.0292, 0.1308].Source: §C.5 · pp. 18–19
- At 768 × 768 without retraining, the frozen corrector achieves 18.45 dB PSNR and selected-image ImageReward of 1.247, above the raw control and native eight-step Euler.The comparisons are 13.96 and 15.35 dB PSNR, and 1.125 and 1.154 ImageReward, respectively.Source: §C.5 · p. 19
- On PartiPrompts without retraining, previews achieve 19.082 dB PSNR and 0.380 LPIPS versus 14.052 dB and 0.542 for native four-step Euler.ImageReward-based selection improves final-image score from 1.3232 to 1.4102 and Spearman rank correlation from 0.079 to 0.357.Source: §C.5 · p. 19
C.6 ABLATION STUDIES
The ablations support preserving the corrected noise norm and using sample-specific correction directions, while reducing correction magnitude when transferring from three to eight steps. Runtime measurements show modest overhead, and the figures document cross-backbone comparisons, transfer across sampling steps, and candidate-selection examples.
C.6 ABLATION STUDIES
- Removing noise norm normalization reduced three-step PSNR from 26.58 to 23.59 dB on LSUN-Church and from 29.35 to 26.64 dB on CelebA-HQ.The unnormalized outputs were also blurrier, with less distinct fine details.Source: §C.6 · p. 19, Figure 8 · p. 21
- Reversing the predicted residual or substituting another sample’s direction yielded lower PSNR than raw three-step DDIM, supporting sample-specific correction directions.The controls preserved residual magnitude and normalized the corrected noise.Source: §C.6 · p. 19, Figure 9 · p. 22
- From 11.50 to 29.09 dB on LSUN-Church and from 11.78 to 31.94 dB on CelebA-HQ, attenuation improved PSNR when the frozen three-step corrector was reused for eight-step DDIM.The compared attenuation was α = (3/8)^2.5; full-strength correction substantially degraded fidelity.Source: §C.6 · p. 19, Figure 10 · p. 23
C.7 RUNTIME AND COMPUTATIONAL COST
- On LSUN-Church, three-step DDIM, input-corrected sampling, and four-step DDIM took 87.6, 102.2, and 116.8 ms; on CelebA-HQ, they took 90.5, 109.4, and 121.4 ms.Input correction cost less than an additional denoising step on both datasets.Source: §C.7 · p. 19
- For FLUX.1-dev at 512 × 512 and batch size eight, trajectory correction changed four-step denoising time from 1.739 s to 1.740 s.The adapter sets total 418.38 M parameters, or 3.52% of the frozen 11.90 B-parameter FLUX transformer.Source: §C.7 · p. 24
D.3 CROSS-BACKBONE TRANSFER
- The additional seed-comparison figures cover two living-room prompts and an eight-seed selection example, without stating a comparative outcome.Source: Figure 17 · p. 30, Figure 18 · p. 31, Figure 19 · p. 32
- Selected examples show a frozen three-step corrector reused across step counts with exponent 2.5 and noise norm normalization; some remain close to the teacher, while others approach it as steps increase.Source: Figure 21 · p. 34
- Cross-backbone figures compare three-step samplers with 60-step references on SDXL and Stable Diffusion 1.5 at CFG 1, 3, and 5; CelebA-HQ examples also show correction transferred from a Bedroom backbone.The CelebA-HQ triplets show raw three-step DDIM, the paired 20-step teacher, and transferred input correction; training used CelebA teacher targets.Source: Figure 12 · p. 25, Figure 13 · p. 26, Figure 14 · p. 27