PaperlayerView exact revision on arXiv

Source-linked AI summary

Learned End-to-End Guidance Schedules for Diffusion Models

Aneesh Barthakur, Mathias Niepert, Luiz F. O. Chamon

arXiv:2610.01502v1cs.LG

TL;DR

Guided diffusion’s slow inference makes balancing data quality and requirement satisfaction computationally costly. LEEGS learns time-dependent guidance schedules and is competitive with strong baselines across image and PDE tasks, matching the 1000-NFE DPS reference using 100 steps.

  • Problem

    Guided diffusion faces slow inference while balancing data quality with satisfaction of generation requirements.

    Source: §1 · pp. 1–2

  • Method

    LEEGS learns time-dependent guidance scales by minimizing the guidance function along the diffusion trajectory over a small set of examples.

    Source: §1 · p. 2, §3 · p. 3

  • Results

    Across image inpainting, noisy inverse problems, face-ID guidance, and PDE problems, LEEGS is competitive with strong task-specific baselines and matches the 1000-NFE DPS reference using 100 steps.

    Source: §6 · p. 9

  • Takeaways & Limitations

    Because useful schedule shapes are task-dependent, the authors conclude they should be learned rather than fixed by heuristic.

    Source: §6 · p. 9

    LEEGS training becomes costly at large step counts because its objective spans the full trajectory.

    Source: §6 · p. 9

Abstract

from arXiv · show

Diffusion models are a powerful generative paradigm used across multimedia and scientific applications. Guided diffusion methods impose requirements on the generation by adding the gradient of a differentiable loss (the guidance function) as a drift term during inference. The weight of this drift (the guidance scale) is critical for the trade-off between data quality and requirement satisfaction. To achieve both of these goals, guided diffusion must resort to small guidance scales and lengthy sampling, incurring high computational costs. This work proposes learned end-to-end guidance schedules (LEEGS) to achieve these objectives with fewer sampling steps. LEEGS trains a time-dependent schedule by minimizing the guidance function over a small set of examples using stochastic gradient descent. Backpropagating through guided sampling is computationally expensive, so LEEGS uses an approximation of the gradient that cuts training time by a factor of 4. We evaluate LEEGS on diverse guidance tasks, including (a) image inpainting, (b) noisy image inverse problems, (c) face-ID-guided generation, and (d) forward and inverse PDE problems, outperforming baselines at equal budget (50 or 100 NFEs), or matching constant guidance with only 10% of the steps.

1 Introduction

Guided diffusion adds the gradient of a differentiable violation loss as a drift to impose requirements, but balancing sample quality and requirement satisfaction can make inference costly. LEEGS learns time-dependent guidance schedules to improve this trade-off.

Introduction

  • Small guidance scales leave requirements unmet, while large scales produce unrealistic samples, so existing guided generation often uses small scales and many costly inference steps.The goal is to reduce sampling steps while balancing data quality and requirement satisfaction.Source: §1 · pp. 1–2
  • With 10 inference steps, LEEGS produces convincing inpainting, while constant guidance at w = 0.15 yields a plausible but incorrect face and w = 2.0 produces saturated artifacts.The constant scale w = 0.15 achieves optimal LPIPS on a held-out set of 100 images, but does not reconstruct the correct face.Source: §1 · p. 2
  • Prior time-dependent guidance schedules were studied for classifier-free guidance, but not for loss-based guidance, whose advantage is conditioning a pretrained model without retraining.Other approaches either still require choosing a guidance scale or require delicate tuning of a trust-region radius or penalty weight.Source: §1 · p. 2
  • LEEGS learns a time-dependent guidance schedule for a model, task, and sampler by minimizing the guidance function over a small example set.Gradient checkpointing bounds memory, while an approximate objective gradient reduces per-batch compute.Source: §1 · p. 2

2 Guided diffusion

Guided diffusion modifies reverse-time sampling to minimize a requirement-violation function, using time-dependent guidance scales. LEEGS learns those scales directly, addressing accumulated sampling error without relying on delicate constraint hyperparameters.

2 Guided diffusion

  • Diffusion models generate samples by numerically integrating learned reverse-time dynamics from Gaussian noise, forming a Markov trajectory whose final distribution approximates the data distribution.The reverse step may use a stochastic or deterministic integrator, such as DDIM or DDPM.Source: §2 · p. 3
  • Guidance uses a nonnegative function h(x0, c) to measure requirement violations and seeks samples that minimize it, ideally to zero.For image inpainting, h measures the difference between the generated image and the condition at observed pixels.Source: §2 · p. 3
  • Following Chung et al. (2023), guidance evaluates h on the one-step denoised estimate and modifies reverse dynamics using its gradient and time-dependent nonnegative scales.Computing the gradient requires differentiating through the diffusion model because the denoised estimate depends on its noise-prediction network.Source: §2 · p. 3
  • As guided marginals diverge from unguided ones, sampling error accumulates and can push final samples away from the data distribution, producing artifacts.Prior trust-region or proximal constraints introduce a delicate radius or penalty-weight hyperparameter; this work instead learns scales to optimize h directly.Source: §2 · p. 3

3 Learned end-to-end guidance schedules

LEEGS trains a nonnegative, time-dependent guidance schedule by minimizing guidance-function values along the sampling trajectory, rather than only at the terminal sample. To reduce the cost of backpropagation, it approximates the gradient by stopping gradients through guidance gradients while retaining their effect through the sampler Jacobian.

Objective

  • Including guidance-function terms at every trajectory step restores training signal where a terminal-only objective has gradients that decay for early steps.The objective averages over conditions from a fixed distribution, with the pretrained model held fixed.Source: §3 · p. 3

Gradient computation

  • Ignoring the second-derivative term yields a roughly 4× cheaper gradient approximation that still propagates schedule effects through the sampler Jacobian.The approximation applies a stop-gradient to the guidance gradient during sampling; the true gradient requires repeated second-derivative computations.Source: §3 · p. 4

Training implementation

  • Projected stochastic gradient descent clips negative schedule entries to zero, and the trained schedule adds no cost at sampling time.Gradient checkpointing makes peak memory independent of the number of steps by recomputing activations during the backward pass.Source: §3 · p. 4

Regularization

  • Driving the guidance function too low can cause reward hacking and unrealistic samples, mitigated by early stopping or penalizing schedule magnitude.The passage notes this typically appears at high step counts such as S = 100, rather than lower counts such as S = 25.Source: §3 · p. 4

Parametrization

  • LEEGS optimizes schedule values freely; restricting them to a small parametric family trades flexibility for fewer parameters and smoother schedules, but is unnecessary in practice.The passage attributes this trade-off to LEEGS’s cost being dominated by guided sampling and notes that a parametric form must be chosen in advance.Source: §3 · p. 4

4 Experiments

LEEGS is evaluated on image-generation and PDE tasks, where learned guidance schedules improve results at constrained sampling budgets. The analyses examine schedule specialization, sampler flexibility, inference-step scaling, and training-data requirements.

4.1 Benchmark tasks

  • The benchmarks cover inpainting, noisy image restoration, face-identity guidance, and forward/inverse Helmholtz, Poisson, and Darcy PDE problems.Schedules are trained separately for each subtask and NFE budget on 30 examples; LEEGS results are means of three runs.Source: §4.1 · pp. 5–6, §4.2 · p. 7

4.2 Main results

  • On FFHQ and ImageNet noisy inverse tasks, LEEGS is best or second-best in nearly every column at 50 and 100 NFE, and beats 1000-NFE DPS on both metrics for most subtasks.On ImageNet at 100 NFE, it is best on 4 of 5 subtasks for both LPIPS and FID.Source: Table 1 · p. 6, Table 2 · p. 6, §4.2 · p. 7
  • At a matched 200-NFE budget, LEEGS improves over RePaint on every inpainting mask type across CelebA-HQ and ImageNet.RePaint with 12× the budget remains ahead on most masks.Source: Table 3 · p. 6, §4.2 · p. 7
  • At 100 NFE, LEEGS performs best on both Face-ID guidance metrics, improving over DSG and constant guidance.At 50 NFE, all three methods are similar except DSG has significantly worse FID.Source: Table 4 · p. 7, §4.2 · p. 7
  • LEEGS attains the lowest relative error on 4 of 6 PDE problems at 50 steps and 5 of 6 at 100, using half the NFE of constant guidance.DiffusionPDE retains an advantage with 20× more steps; the best fixed schedule shape varies by problem.Source: Table 5 · p. 7, §4.2 · p. 8

4.3 Analysis

  • Learned schedule shapes are similar across budgets but differ by problem: increasing schedules help noisy inverse tasks, while decreasing schedules help PDE tasks.Among five FFHQ inverse problems, each task’s learned schedule is best on that task in nearly every case.Source: Figure 3 · p. 8, §4.3 · p. 8
  • LEEGS also applies without modification to reverse-diffusion sampling interleaved with renoising, performing competitively with vanilla LEEGS on FFHQ inverse problems and better on box and random inpainting.This scheme learns hill-shaped schedule segments, with unchanged learning cost.Source: §4.3 · p. 8
  • More inference steps generally improve LEEGS with diminishing returns; performance and schedules change little beyond 30 training examples.A schedule trained at small S can be interpolated and rescaled for larger S, reducing training cost at large budgets.Source: §4.3 · p. 8

5 Related work

Prior work varies or learns CFG guidance scales, while LEEGS is presented as the first learned schedule for loss-guided diffusion. Related constrained-diffusion methods use projection, thresholding, or consensus, whereas LEEGS applies an offline schedule during ordinary guided sampling.

Guidance scheduling for CFG

  • CFG schedules have used fixed trajectories, limited intervals of nontrivial scaling, and adaptive scales derived from noise predictions, likelihoods, or greedy search.Kynkäänniemi et al. (2024) showed that a nontrivial scale is needed only within a limited interval; other works proposed fixed or adaptive schedules.Source: §5 · p. 9

Learned guidance schedules

  • LEEGS is, to our knowledge, the first method to learn guidance schedules for loss-guided diffusion, unlike prior learned CFG-scale methods using labels or constant-guidance references.Prior methods learn scales from the intermediate sample and time, or use evolutionary search supervised by a constant-guidance process requiring many 1000-step rollouts.Source: §5 · p. 9

Constrained diffusion

  • Projection methods can require inner optimization and renoising, while proximal and trust-region methods introduce a delicate threshold or radius hyperparameter.Projection steps may deviate substantially from the diffusion-induced distribution, and renoising lengthens sampling.Source: §5 · p. 9
  • Consensus methods split inference into projected and unguided trajectories, anneal them together, and can set the implicit guidance scale by dual ascent.LEEGS instead learns one schedule offline and uses ordinary guided sampling without projection, renoising, or extra inference cost.Source: §5 · p. 9

6 Conclusion

LEEGS learns time-dependent guidance schedules for loss-guided diffusion, achieving strong results across diverse tasks with fewer sampling steps. The conclusion also identifies limitations in scaling training and adapting schedules to conditions or initial noise.

6 Conclusion

  • 100 steps match the 1000-NFE DPS reference across image inpainting, noisy inverse problems, face-ID guidance, and PDE problems, while LEEGS is competitive with strong task-specific baselines.The evaluation spans the tasks reported in Section 4.Source: §6 · p. 9
  • LEEGS learns time-dependent guidance schedules for loss-guided diffusion using as few as 30 examples by backpropagating guidance loss through the sampler.The method is described as simple and general.Source: §6 · p. 9
  • A learned schedule beats constant guidance at a matched budget, and useful schedule shapes depend on the task, favoring learning over fixed heuristics.Source: §6 · p. 9
  • LEEGS works with both stochastic DDPM and deterministic Euler samplers.Source: §6 · p. 9
  • Training becomes costly at very large S; the gradient approximation may be less effective when guidance is highly sensitive to x_s, and adapting schedules to conditions or initial noise remains open.Source: §6 · p. 9

AI use statement

LLM tools supported code, analysis, literature search, and manuscript preparation, while authors reviewed AI-assisted work and took responsibility for the final content.

AI use statement

  • LLM tools helped implement and modify LEEGS code, samplers, guidance functions, baselines, experiment scripts, and plotting and analysis notebooks, and assisted in interpreting results.They also helped search for relevant papers, edit the main text, and draft several appendix sections.Source: AI use statement · p. 11
  • The authors did not use generative AI to create synthetic datasets, develop theoretical models or conceptual frameworks, or clean and reformat datasets.Mathematical claims, proof assistance, and translation were not applicable to this work.Source: AI use statement · p. 11
  • Authors reviewed all AI-assisted work, manually reviewed critical method and experiment code and AI-generated text, and accepted responsibility for the paper’s final content and artifacts.Source: AI use statement · p. 11

ADDITIONAL MATERIAL · A Pseudocode for LEEGS training

The appendix pseudocode trains a time-dependent guidance schedule by differentiating an objective over sampled guided trajectories, with checkpointing to limit memory and an optional gradient approximation that detaches the guidance gradient.

A Pseudocode for LEEGS training

  • LEEGS initializes a schedule with independent uniform weights, then uses Adam updates over epochs and minibatches while projecting the schedule onto nonnegative values after each update.Minibatches contain training conditions.Source: A Pseudocode for LEEGS training · p. 17
  • For each minibatch, COMPUTEGRADIENT samples a fresh trajectory, accumulates intermediate and terminal guidance values, and returns their objective’s gradient with respect to the schedule.Intermediate terms evaluate guidance on denoised iterates; the terminal term evaluates it on the final sample.Source: A Pseudocode for LEEGS training · p. 17
  • Each GUIDEDSTEP evaluates guidance on the denoiser prediction, computes its gradient, and subtracts the weight-scaled gradient from the unguided reverse-diffusion update.The step returns the updated iterate and that step’s guidance value.Source: A Pseudocode for LEEGS training · p. 17
  • Gradient checkpointing recomputes internal activations during backward passes, keeping peak memory independent of the number of sampling steps.The implementation uses PyTorch’s built-in checkpointing functionality.Source: A Pseudocode for LEEGS training · pp. 17–18
  • The sampler primitive returns a denoiser prediction and an unguided reverse step; for noise-predicting models, the prediction is recovered with Tweedie’s formula, while direct denoised-sample predictors need no conversion.The guidance function is evaluated on a denoised sample.Source: A Pseudocode for LEEGS training · pp. 17–18
  • With the LEEGS approximation enabled, the guidance gradient is detached, omitting its Hessian contribution while propagating the sampler Jacobian across steps.Without the approximation, the guidance gradient retains its graph and its Hessian contributes to the schedule gradient.Source: A Pseudocode for LEEGS training · pp. 17–18

B Additional experimental details … C.3 Effect of regularization

The appendix details task-specific training and evaluation protocols, then tests schedule variability, the detached-gradient approximation, and weight-decay regularization. LEEGS schedules vary less with larger training sets, the approximation is faster and more stable than the true gradient on ImageNet, and regularization effects depend on task and step budget.

B Additional experimental details; B.1 Noisy inverse problems

  • LEEGS uses Adam with batch size 1 and 30 randomly drawn training examples per task and NFE budget; noisy inverse-problem tests use five measurement operators.The noisy inverse-problem operators include 4× bicubic superresolution, Gaussian and motion deblurring, and box or random inpainting.Source: B Additional experimental details · p. 19, §B.1 · p. 19, Figure 4 · p. 19

B.2 Image inpainting; B.3 Face-ID guidance

  • Image-inpainting experiments use Lugmayr et al.’s masks; Face-ID guidance targets each input image’s identity, measured with ArcFace.For CelebA-HQ inpainting, test and training pools are disjoint, and most mask types are fixed across images. Face-ID FID is computed against the CelebA-HQ training split.Source: §B.2 · p. 21, §B.3 · pp. 21–22

B.4 PDE problems

  • For PDE tasks, a pretrained joint model represents coefficient-solution pairs, while guidance enforces consistency with sparse observations for forward or inverse recovery.The forward problem recovers the solution from the coefficient, and the inverse problem recovers the coefficient from the solution; observations cover 500 randomly chosen locations, approximately 3% of pixels.Source: §B.4 · p. 22

B.5 Computational resources

  • Training wall-clock times are measured on a single A100 without validation logging, with ranges reported when tasks at a fixed NFE budget use different epoch counts.The experiments used NVIDIA A40 and A100 GPUs and an RTX 3090 workstation; validation was for monitoring, not early stopping or model selection.Source: §B.5 · p. 23, Table 11 · p. 23

C Further experiments and ablations; C.1 Ablation on training dataset size

  • Schedule and reconstruction variability decrease as training-set size grows, while reconstruction metrics plateau at larger sizes; the mean schedule remains relatively stable, especially for FFHQ motion deblur.Schedule variability continues to decrease with size for Helmholtz, possibly because its sampler is deterministic; the stochastic FFHQ sampler adds variability. The mean-schedule trend suggests smaller training sets may have more variance and less bias relative to larger-set schedules.Source: Figure 6 · p. 24, §C.1 · p. 24, Figure 7 · p. 24

C.2 Comparison of LEEGS vs true gradient

  • ∼4–5× faster: LEEGS training iterations cost less than true-gradient iterations; on ImageNet, the true gradient has 4–6× higher LPIPS standard deviation and ∼3× higher schedule distances.On FFHQ, the two gradient methods are comparable. The timing comparison covers FFHQ box inpainting and Helmholtz forward; relative to plain sampling, true-gradient costs are ≈20× and ≈16×, respectively, while LEEGS costs ≈4× for both.Source: §C.2 · pp. 25–26, Table 12 · p. 26, Figure 10 · p. 26

C.3 Effect of regularization

  • At S = 100, superresolution held-out LPIPS improves through γ = 5 × 10−4, then worsens at 10−3; weight decay flattens the training elbow, while box inpainting is largely insensitive except at 10−3.At S = 50, superresolution’s held-out LPIPS optimum is γ = 10−4; at S = 10 and 25, it worsens monotonically with weight decay. For S = 10, box inpainting remains flat even at γ = 10−3. Early stopping before the unregularized superresolution elbow is another fix, though its training LPIPS does not reach the pre-elbow value.Source: §C.3 · pp. 26–27, Figure 11 · p. 27

C.4 Generalization across tasks · C.5 Generalization across number of steps · C.6 Ablation on trajectory objective vs terminal objective

Across task and step-budget generalization studies, LEEGS schedules retain some transferable structure, but task-specific magnitude and shape matter; rescaling helps adapt schedules to larger budgets. The trajectory-level objective also outperforms terminal-only training, especially as the number of diffusion steps grows.

C.4 Generalization across tasks

  • Learned schedules differ in magnitude by up to a factor of four and in shape; random inpainting has the largest guidance, while superresolution and random inpainting rise earliest.Gaussian deblur rises roughly linearly over the second half of the trajectory; random inpainting drops sharply at the final steps, and superresolution has a gradual, hill-shaped profile.Source: §C.4 · pp. 27–28, Figure 12 · p. 28
  • 4 of 5 matched schedules perform best in their task columns at both S; gaussian deblur narrowly beats motion deblur’s schedule for motion deblur.The motion-deblur comparisons are 0.209 vs. 0.212 at S = 50 and 0.177 vs. 0.179 at S = 100.Source: §C.4 · p. 27, Figure 13 · p. 29
  • After magnitude matching, random inpainting remains selective: random-inpainting and superresolution schedules perform best there, while box inpainting retains its own schedule as best.At S = 50, the two best random-inpainting schedule results are 0.196 and 0.217; at S = 100, they are 0.173 and 0.171.Source: §C.4 · pp. 27–28, Figure 13 · p. 29

C.5 Generalization across number of steps

  • Training at a smaller step budget and extrapolating a schedule to a larger budget can improve LEEGS when training compute is limited.The study evaluates this strategy on FFHQ box inpainting and the Helmholtz forward PDE problem, using interpolation with and without magnitude rescaling.Source: §C.5 · pp. 28–30
  • Across FFHQ box inpainting and Helmholtz, schedules learned at different S have similar shapes, while their mean guidance decreases monotonically as S increases.A linear fit to mean guidance across trained budgets supplies the rescaling factors used for interpolation.Source: Figure 14 · p. 29, §C.5 · p. 30
  • At Stgt = 100, plain interpolation from Ssrc ∈{10, 20, 30, 50} underperforms trained schedules across both tasks, while rescaling performs very close in every table row.For smaller gaps, interpolation often matches or beats the trained schedule, and can be marginally better when little magnitude mismatch needs correction.Source: §C.5 · p. 30, Table 13 · p. 30, Table 14 · p. 30

C.6 Ablation on trajectory objective vs terminal objective

  • Terminal-only training has worse training LPIPS than the full trajectory objective, with a gap that widens as S grows; held-out LPIPS shows the same widening gap.Terminal-only training also shows a reward-hacking elbow, while its held-out LPIPS curve is flatter and less able to benefit from additional steps; terminal gradients decay for steps far from the trajectory’s end.Source: §C.6 · p. 31, Figure 15 · p. 31, Figure 16 · p. 31

D Extension to sampling with renoising

LEEGS learns time-dependent guidance schedules for a sampler that alternates guided reverse diffusion with unguided renoising, without increasing the NFE budget. On FFHQ noisy inverse tasks, LEEGS-renoise is comparable to standard LEEGS overall and improves both inpainting results.

Extension to sampling with renoising

  • With S=50 and K=5, each of five outer steps applies 10 guided reverse steps, then renoises without guidance to the next start time; this adds no network evaluations.The learned schedule has one weight per guided step, and the final sample comes from the last outer step.Source: D Extension to sampling with renoising · p. 32

Results

Learned schedules

  • Every outer-step schedule except the last has a hill shape; Gaussian deblur’s first-step peak is about 85, versus about 37 in later steps.The last segment increases roughly linearly except for random inpainting, which drops sharply at its final step; peak guidance grows across outer steps for superresolution and random inpainting.Source: D Extension to sampling with renoising · p. 33, Figure 17 · p. 33

Conclusion

  • Learning guidance schedules also applies beyond DPS-style sampling, where renoising makes hand-tuning per-outer-step schedule segments harder.LEEGS has the same cost for either sampler when the sampling trajectory uses the same number of NFEs.Source: D Extension to sampling with renoising · p. 33

E Comparison at S = 10 NFEs (FFHQ) … I Qualitative results

At 10 NFEs on FFHQ, LEEGS leads the noisy-inverse comparison, while the remaining appendices document variability, supplementary metrics, learned schedules, baseline searches, and qualitative examples.

E Comparison at S = 10 NFEs (FFHQ)

  • At 10 NFEs, LEEGS is best on four of ten columns and best or second-best on eight, while DAPS and CDIM are the strongest competitors.The comparison uses the 950-image FFHQ test set; CDIM averages 12.4 NFEs across five tasks, and Linear increase is not best on any task at this low budget.Source: E Comparison at S = 10 NFEs (FFHQ) · p. 34, Table 18 · p. 34

F Additional metrics for main experiments; F.1 Variance of LEEGS across draws

  • LEEGS results in Tables 1–5 average three independent training-set draws; Tables 19–23 report their standard deviations, while baselines have no repeated draws.The variance appendix includes FFHQ and ImageNet noisy-inverse results, image inpainting, Face-ID guidance, and PDE problems.Source: §F.1 · p. 38, Table 19 · p. 38, Table 20 · p. 38

F.2 Noisy inverse problems

  • For noisy inverse problems, the appendix reports test-sample LPIPS variance, PSNR, and SSIM for FFHQ and ImageNet, omitting the DPS reference row.PSNR and SSIM are pixel-level metrics that can favor smooth reconstructions over sharp, slightly misaligned details; higher is better for both.Source: §F.2 · pp. 38–39, Table 24 · p. 39, Table 25 · p. 39

F.3 Image inpainting problems

  • For 200-NFE image-inpainting results, the appendix adds LPIPS variance, PSNR, and SSIM, omitting the approximately 2400-NFE RePaint reference.The metrics are computed as for the noisy-inverse problems.Source: §F.3 · p. 41, Table 30 · p. 41, Table 31 · p. 41

G Additional learned guidance schedules; G.1 Noisy inverse problems; G.2 Image inpainting; G.3 Face-ID guidance; G.4 PDE problems

H Hyperparameter searches for the baselines; H.1 Noisy inverse problems; H.2 Image inpainting; H.3 Face-ID guidance; H.4 PDE problems

  • Baseline hyperparameter searches are grouped by task: LPIPS is used for noisy-inverse and inpainting selection, mean FaceID score for Face-ID, and task-specific relative error or error rate for PDEs.The noisy-inverse appendix also adds a 10-NFE FFHQ search for constant guidance, linear increase, DSG, and DAPS.Source: H Hyperparameter searches for the… · p. 47, §H.1 · p. 47, §H.2 · p. 51

I Qualitative results

  • Qualitative results for the main experiments are supplied as separate supplementary material.The appendix directs readers to large image panels hosted separately.Source: I Qualitative results · p. 55

Related papers

Loading 2610.01502v1…