Learned End-to-End Guidance Schedules for Diffusion Models
Guided diffusion forces a costly trade-off between enforcing requirements and preserving sample quality, usually resolved by small guidance scales and long sampling trajectories. This paper introduces LEEGS, which learns a time-dependent guidance schedule end-to-end by backpropagating the guidance loss through the sampling trajectory. A cheap gradient approximation—treating the guidance gradient as a stop-gradient—cuts training time by roughly 4$\times$, letting the method train on only 30 examples and match or beat strong baselines across image inverse problems, face-ID guidance, and PDEs with as few as 50–100 NFEs.
LEEGS is an elegant and empirically strong contribution: treating per-step guidance scales as differentiable trajectory parameters is well-motivated, and the paper shows convincingly that the optimal schedule shape is task-dependent. Still, the theoretical justification for dropping second-order trajectory terms is thin, the abstract's claim of matching constant guidance at 10\% steps is not directly demonstrated in the main results tables, and the need for per-task retraining limits off-the-shelf utility. The authors candidly note that the approximation 'may be less effective when the guidance field is highly sensitive to $x_s$' (Section 6).
The paper's greatest strength is empirical breadth: LEEGS is evaluated on noiseless inpainting, noisy inverse problems, face-ID guidance, and forward/inverse PDEs, often outperforming DSG, CDIM, DAPS, and constant-guidance baselines at matched budgets (Tables 1–5). The training objective in Eq. (3) is a thoughtful design; by summing $h(\hat{x}_{0|s},c)$ over intermediate steps rather than the terminal sample alone, it restores gradient signal to early trajectory steps. The ablations are thorough: the approximation is shown to be $\sim$4$\times$ faster than the true gradient with comparable or better stability, and only $\sim$30 training examples are needed. Finally, the paper demonstrates clearly that schedule shape is task-dependent: 'The schedules are similar across different step budgets for a particular problem, but differ between problems, highlighting the need for learning methods' (Section 4.3).
The gradient approximation drops the Hessian-like term $\partial g_s / \partial x_s$ from Eq. (5), and while the authors honestly flag sensitivity to this omission, no quantitative criterion or bound is offered for when it fails. More pressingly, the abstract claims LEEGS can match constant guidance 'with only 10\% of the steps,' yet the main results tables only report comparisons at 50, 100, and 200 NFEs; the 10$\times$ reduction explicitly stated in the text refers to beating the DPS reference at 1000 NFEs with 100 NFEs (Section 4.2), not to the constant-guidance baseline. Schedules also fail to transfer across tasks: 'the matched schedule is best on that task in nearly every case' (Appendix C.4), so users must retrain per task and budget. Finally, Figure 2 shows reward hacking—where the terminal constraint improves while reconstruction LPIPS degrades—unless early stopping or weight decay is applied.
The core claim—that learned schedules dominate constant or heuristic linear schedules at fixed NFE budgets—is well supported by Tables 1–5. On FFHQ and ImageNet noisy inverse tasks, LEEGS is best or second-best in almost every column at 50 and 100 NFEs. Some comparisons are less clean, however. CDIM's average NFEs exceed the nominal budget (e.g., 73.4 NFEs for a nominal 50-NFE setting on ImageNet box inpainting; Table 6), making the comparison imperfectly matched. In the PDE domain, DiffusionPDE at 2000 steps retains a clear accuracy advantage over LEEGS at 100 steps (Table 5), so the claim of drastic step reduction is strongest for image tasks. The heuristic linear schedules also expose an inconsistency: linear-increasing is strong on images while linear-decreasing is strong on PDEs, reinforcing the value of learning, but also showing that no single heuristic works everywhere.
The paper provides Algorithm 1 with PyTorch-style pseudocode and detailed hyperparameters in Appendix B (e.g., 'learning rate 0.3, 25 epochs for FFHQ super-resolution at $S=50$'; Table 7). Training is lightweight—about 1–4 hours on a single A100 for most schedules (Table 11)—and uses only 30 examples per subtask. No public code release is mentioned, which is a notable gap. Independent reproduction would also need to navigate the interaction between step budget, learning rate, epochs, and weight decay to avoid reward hacking (Figure 2), and the paper notes that gradient checkpointing is required to keep peak memory bounded as $S$ grows. The reliance on existing codebases is disclosed, but without released integration code, reproducing the exact training pipeline would require non-trivial effort.
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.
Pick a starting point or write your own. Challenges run in the background, so you can keep reading while the AI investigates.
No challenges yet. Disagree with the review? Ask the AI to revisit a specific claim.