PaperlayerView exact revision on arXiv

Source-linked AI summary

Adaptive Second-Order Solvers for Fast Stochastic Diffusion Sampling

Ella Kemperman, Luca Ambrogioni

arXiv:2610.03034v1cs.LGcs.CLcs.CV

TL;DR

Diffusion sampling is costly, and variable reverse-process difficulty makes time discretization consequential. The paper adapts PI step-size control for diffusion and aggregates adaptive paths into static schedules; results vary by task, with per-sample adaptation especially helpful in the 1D example.

  • Problem

    Diffusion sampling requires many neural-network evaluations, while reverse-process difficulty varies over time, making time discretization consequential to sampling cost and quality.

    Source: §1 · p. 1

  • Method

    The paper adapts PI step-size control with a diffusion-noise-normalized error estimator and aggregates per-sample trajectories into a static schedule.

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

  • Results

    Per-sample adaptation strongly benefits the 1D example, while image results show PI static schedules outperforming EDM in several settings and language results show lower perplexity than EDM-churn at all tested NFE.

    Source: §6.1 · p. 6, §6.2 · p. 6, §6.3 · p. 8

  • Takeaways & Limitations

    The benefit of per-sample adaptivity is problem-dependent, and image results show only a slight difference between the PI solver and stochastic Heun with a PI static schedule.

    Source: §6.1 · p. 6, §6.2 · p. 6

    The image and language experiments use VE models, the 1D experiments use a VP SDE, and no full search over τrel and τabs was conducted.

    Source: §7 · p. 9

Abstract

from arXiv · show

Diffusion models rely on numerical solvers requiring time-discretization, which has a large influence on the tradeoff between sampling cost and quality. However, the computational difficulty of the reverse process varies along the sampling trajectory and across data distributions, making the choice of discretization important. We adapt proportional-integral (PI) step-size control to diffusion, using our diffusion noise-normalised error estimator. Unlike existing adaptive methods in diffusion that respond only to the current error, the PI solver also incorporates the previous error, yielding smoother step adaptation. We further show that these per-sample trajectories exhibit shared structure and can be aggregated into a fixed schedule that retains much of the benefit of adaptive sampling. We evaluate both approaches on natural-image and language datasets, in terms of quality, measured by FID at a matched number of neural network evaluations (NFE), comparing them with widely used stochastic solvers and schedules. For images, our fixed discretization outperforms the commonly used EDM schedule in terms of sample quality when used with the stochastic Heun sampler, and with the EDM-churn sampler at low NFE. Additionally, our PI adaptive solver obtains better FID than most stochastic and adaptive baselines, although it does not beat the EDM-churn sampler at low NFE. Moreover, we find our solver outperforms both the EDM and the entropy schedule on language diffusion at low-to-medium NFE in terms of perplexity, with the drawback of lower token entropy. Lastly, we find that the benefit of per-sample adaptivity is problem-dependent. It is highly beneficial in 1D toy examples, while only marginal for image and language data, where the average schedule sometimes even outperforms the PI-adaptive solver. Code is available at https://github.com/ellakemperman/adaptive-second-order-diffusion-solvers

1 Introduction

Diffusion sampling is costly, and the reverse process has regions of differing difficulty, so uniform time-step allocation may be inefficient. This paper introduces PI-controlled adaptation and a way to aggregate per-sample paths into a static schedule.

1 Introduction

  • Diffusion sampling requires many neural-network evaluations, making time discretization important to the sampling-cost and quality tradeoff.The reverse process has critical windows where many generative decisions are made, suggesting effort should vary with local difficulty.Source: §1 · p. 1
  • Existing data-driven entropy schedules are reported to degrade with second-order methods, despite second-order samplers improving the NFE–quality tradeoff.The cited entropy schedule is described as effective only for first-order samplers.Source: §1 · pp. 1–2
  • Adaptive solvers can generate per-sample schedules based on local difficulty and are generally solver-order agnostic, but their use in diffusion remains limited.The cited diffusion-specific adaptive SDE solver adjusts step size using current local error.Source: §1 · p. 2
  • Adaptive PI control uses current and previous local errors, while a diffusion-noise-normalized error estimator improves its performance in diffusion sampling.The authors also aggregate shared structure across per-sample trajectories into a static schedule usable like a handcrafted schedule.Source: §1 · p. 2

2 Related Works

Related work improves diffusion sampling through sampler design, time-discretization design, and adaptive SDE methods. This paper focuses on adaptive step selection for second-order sampling rather than changing the integration method itself.

2 Related Works

  • Prior sampler work includes predictor-corrector frameworks and higher-order methods, with stochastic Heun and churn among the approaches described as improving low-NFE sample quality.The paper distinguishes its focus on time-step selection from work on the numerical integration method.Source: §2 · p. 2
  • Discretization methods include handcrafted schedules, solver- or dataset-optimized schedules, entropy spacing, learned predictions, and reinforcement-learning approaches.The cited reinforcement-learning method is shown to interpolate and extrapolate across NFE but only for first-order solvers.Source: §2 · p. 3
  • The proposed approach differs by using a second-order PI adaptive solver to determine per-sample, per-step discretizations and aggregate them into one static schedule.Prior diffusion adaptive work cited here uses current error only, whereas the PI approach incorporates previous error as well.Source: §2 · p. 3

3 Background

Diffusion modeling uses a forward process to corrupt data and a reverse process for inference. Because reverse diffusion generally requires numerical approximation, sampling depends on choices of solver and time discretization.

3 Background

  • The forward diffusion process transforms data toward a simple distribution, while the reverse process is used for inference to obtain samples from the target distribution.The simple distribution is generally chosen to be Gaussian.Source: §3 · p. 3
  • The forward process progressively adds noise and is governed by an Itô SDE; reverse diffusion progressively removes noise and is governed by a reverse-time SDE.The score is supplied by a trained neural network in the described setup.Source: §3 · p. 3
  • The probability flow ODE offers another model of the process, and its trajectories have the same marginal probability densities as the reverse SDE.The passage introduces this ODE as an alternative to the potentially slow reverse SDE.Source: §3 · p. 3
  • Numerical integration approximates the reverse SDE by discretizing time and advancing the sample with steps determined by the equation and integration scheme.Examples include Euler–Maruyama and stochastic Heun; the discretization has traditionally been chosen by hand.Source: §3 · p. 4

4 PI Controlled Adaptive Solver

The solver estimates local error with stochastic Heun and adjusts step sizes using PI control, with error scaling adapted to diffusion noise. It can also switch to predetermined probability-flow ODE steps at low noise levels.

4 PI Controlled Adaptive Solver

  • The error estimator uses the difference between first- and second-order estimates, with a tolerance term relative to diffusion noise to improve performance for variance-exploding diffusion.Stochastic Heun supplies the second-order estimate.Source: §4 · p. 4
  • The relative tolerance scales error based on the data and, specifically, the current noise level σ(t).The absolute tolerance scales error by a constant.Source: §4 · p. 4
  • Step rejection was disabled in final experiments because rejected steps generally worsened the NFE-quality tradeoff by wasting computation.The solver differs from the cited Ilie et al. implementation, which also uses rejection.Source: §4 · p. 4
  • The PI controller adjusts step size using both current and previous step errors.The controller is introduced as the step-size adaptation mechanism.Source: §4 · p. 4
  • The selected gains are ki = 0.3 and kp = 0.1; α is a safety factor, while αmin and αmax bound step-size increases per iteration.The authors report that this gain setting performed well after testing.Source: §4 · p. 4
  • At low NFE, the final few low-noise steps use a predetermined probability-flow ODE schedule with a Heun sampler to address grainy images.The authors attribute the graininess likely to stochastic noise that cannot be fully removed at low step counts.Source: §4 · p. 4

5 Obtaining a Static Schedule

Because individual discretizations have different lengths, they cannot be aggregated directly. The method interpolates them onto a shared step-fraction scale and averages their time points to form a static schedule.

5 Obtaining a Static Schedule

  • Different trajectory lengths prevent direct aggregation, so linear interpolation samples each discretization at matching fractions of its total steps.This aligns the trajectories for aggregation despite their varying lengths.Source: §5 · p. 5
  • Normalizing each path's step index to a 0-to-1 fraction creates interpolated paths that can be averaged at corresponding fractions.The resulting average discretization uses the same number of points in each interpolated path.Source: §5 · p. 5
  • The average time points represent where the PI solver places evaluations; wider intervals indicate less allocated effort in that interval.Source: §5 · p. 5

6 Experiments

Across experiments, the PI solver’s per-sample adaptivity helps substantially in the 1D mixture but only marginally for images. Its fixed schedule improves on EDM in several image and language comparisons, with language perplexity gains accompanied by lower token entropy.

6.1 1D Multimodal Gaussian Diffusion

  • For the 1D Gaussian mixture, PI adaptation substantially outperforms other solvers at low NFE, while stochastic Heun with the PI static schedule degrades quickly.Individual effort allocations vary substantially, which the authors suggest explains the static schedule’s weaker performance.Source: §6.1 · p. 6, Figure 2 · p. 6

6.2 FFHQ and ImageNet

  • The PI static schedule beats EDM on FFHQ at every tested NFE and on ImageNet-64 at 49 NFE; PI adaptation generally beats other baselines except EDM-churn.The difference between the PI solver and stochastic Heun with its static schedule is generally slight.Source: §6.2 · p. 6, Table 1 · p. 7
  • For images, individual PI trajectories lie close to their average, which takes larger steps than EDM at high noise and smaller steps at low noise.This low per-sample variance is offered as a likely explanation for the slight difference between adaptive PI and stochastic Heun with the PI static schedule.Source: §6.2 · p. 7, Figure 3 · p. 7

6.3 Bitstream Language Diffusion

  • On LM1B, the PI solver beats EDM-churn in perplexity at all NFE and all configurations at 63 and 127 NFE, but has lower unigram entropy.The PI static schedule beats the entropy schedule with second-order EDM but loses to it with first-order DDIM.Source: Table 2 · p. 8, §6.3 · p. 8

6.4 Ablations

  • On FFHQ, configuration A outperforms the other PI solver ablations; alternate tolerance options perform substantially worse, while PI control adds a smaller gain over integral-only control.Table 3 reports FID across NFE for the solver configurations.Source: §6.4 · p. 8, Table 3 · p. 9
  • Native PI schedules outperform foreign schedules on both datasets, with the native FFHQ schedule giving a 3× improvement over the LM1B schedule at low NFE.Table 4 compares the FFHQ and LM1B schedules using FID on FFHQ and generative perplexity on LM1B.Source: Table 4 · p. 9, §6.4 · p. 9

7 Conclusion

The paper introduces a diffusion-noise-based error estimator and a method for aggregating adaptive PI trajectories into a static discretization. The approach supports second-order-compatible, data-driven schedules that improve on the handcrafted EDM schedule in several evaluated settings, with experimental scope limitations.

7 Conclusion

  • The aggregated PI trajectories yield second-order-compatible data-driven schedules that improve on the handcrafted EDM schedule in several evaluated settings.The method also introduces an error estimator based on the diffusion noise level.Source: §7 · p. 9
  • All image and language experiments used VE models, while the 1D experiments used a VP SDE and showed significant adaptivity benefits whose source—SDE or dataset—is unclear.Experiments covered stochastic sampling only; deterministic performance remains for future work.Source: §7 · p. 9

8 AI Use Statement

The AI use statement describes feedback and debugging assistance that was manually checked, with no AI-generated code in the final codebase. The appendix explains how the EDM probability-flow formulation is rewritten to obtain a compatible reverse-SDE drift and diffusion.

8 AI Use Statement

  • Generative AI provided feedback on methodology and manuscript style and helped identify code bugs; all feedback was manually checked, and no AI-generated code entered the final codebase.The authors state that AI was not used for other purposes and take responsibility for the final content.Source: §8 · pp. 9–10

A Reparametrizing the EDM Probability Flow ODE

  • The EDM probability-flow formulation is rewritten to obtain drift and diffusion compatible with the forward SDE, allowing its use with the reverse-SDE formulation.The resulting SDE can be used with variance-exploding denoising networks from EDM and EDM2.Source: A Reparametrizing the EDM Probability… · p. 12
  • In this formulation, s(t) is a scaling function and σ(t) indicates the noise at time t; the derivation uses the score to obtain drift and diffusion.For the EDM framework, the text sets s(t) = 1 and σ(t) = t before deriving those terms.Source: A Reparametrizing the EDM Probability… · p. 12

B.1 Error Estimator

The error estimator uses diffusion noise to set a mixed tolerance, which generally achieves similar sample quality with fewer neural-network evaluations than absolute or non-mixed tolerances. The experiments also disable step rejection after finding it wastes evaluations, while batching masks samples that have finished.

B.1 Error Estimator

  • The absolute-value-based estimator worked in 1D VP experiments but struggled to achieve low NFE and quality on image datasets using VE models.A relative tolerance was introduced because the problem scale varies substantially in the VE formulation.Source: §B.1 · p. 13
  • The mixed noise tolerance generally requires fewer NFE for similar sample quality than absolute or non-mixed tolerance, so it is used in all experiments.It is defined as τ(t) = τabs + τrelσ(t), using the known diffusion noise level as a scale for normalization.Source: §B.1 · p. 13

B.2 Step Rejection

  • Disabling step rejection produced the same or better quality, improved 1D results, and avoided NFE wasted on rejected steps.The rejection mechanism was therefore disabled in the final experiments.Source: §B.2 · p. 13

B.3 Batches

  • Batch elements can finish at different times, so masking removes finished elements from the batch before score-network evaluation.Source: §B.3 · p. 13

C.1 1D Validation

The 1D tests show that the PI solver’s NFE–Wasserstein tradeoff depends strongly on its initial step size, while several controller parameters have limited influence. The later tuning sections describe how the solver and a comparison solver were adjusted for image and language experiments.

C.1 1D Validation

  • The solver was insensitive to α, kp, and most ki values, while ki < 0.05 was detrimental; τrel was best kept at 0 for this 1D experiment.The authors attribute τrel’s effect partly to limited variation in the variance-preserving SDE’s scale.Source: §C.1 · pp. 14–15
  • h0 = 0.15 reduced Wasserstein distance quickly as NFE rose and reached the same plateau as smaller h0 values; h0 = 0.2 attained a worse minimum.The authors selected h0 = 0.15 for the 1D experiments and identify h0 as an important tradeoff parameter.Source: §C.1 · p. 15, Figure 6 · p. 15

C.2.1 PI Solver

  • Image experiments transferred the 1D settings as a starting point but required substantial changes to τabs, τrel, and h0 to bring sampling from 750 NFE toward 50–100 NFE.The final 20% of the PI trajectory used ODE steps with the EDM schedule, with step counts set according to total NFE.Source: §C.2.1 · pp. 15–16
  • When τabs + τrel became sufficiently large, image quality degraded as NFE fell; the follow-up search selected τabs = 0.005.The authors tuned h0 alongside τrel to reach target NFE, choosing the largest h0 that retained high-quality results.Source: §C.2.1 · p. 16, Figure 7 · p. 16

C.2.2 Gotta Go Fast Solver

  • Gotta Go Fast’s initial settings caused step-size and local-error oscillations; reducing r to 0.1 improved stability, and lowering α to 0.7 helped further.Increasing r slightly reintroduced instability, and the dataset required a larger τrel than the original paper.Source: §C.2.2 · pp. 16–17

C.3 Bitstream Diffusion

  • For Bitstream diffusion, τrel showed a tipping point: lower values could make the solver require infinitely many steps to converge, though it could still be tuned to target NFE.Source: §C.3 · p. 17

D Experiment Details

  • The experiments used a new codebase inspired by Song et al. [1] and Karras et al. [2], integrating and extending Batzolis et al.’s [35] code.Runs used 10 Quadro RTX 6000 GPUs with PyTorch.Source: D Experiment Details · p. 17

D.1.1 Experiment Setup

The appendices specify synthetic, image, and text-generation setups, including how schedules, metrics, and solver settings were configured. They also describe PI-solver ablations and cross-modality schedule tests.

D.1.1 Experiment Setup

  • The 1D setup used a 10-mode Gaussian mixture with varied weights, means, and variances, and compared solver estimates against the original samples using Wasserstein distance.The score was analytic, with 10^-13 added to avoid NaN values outside the distribution.Source: §D.1.1 · pp. 17–18

D.1.2 Hyperparameters

  • The 1D evaluation swept Euler–Maruyama and Heun from 10 to 200 NFE and varied PI τabs from 0.001 to 1.The tolerance sweep used exponential spacing; the full PI settings are referenced in Table 6.Source: §D.1.2 · p. 18, Table 6 · p. 18

D.2.1 Image Generation

  • Image experiments used EDM2-based generation with ImageNet-64 and FFHQ checkpoints, matched starting noise and labels across solvers, and adapted the denoiser output into a score.EDM-churn used its repository sampler, while other listed solvers used the drift and diffusion of the specified SDE.Source: §D.2.1 · pp. 18–19

D.2.2 Metric Calculation

  • Image metrics used three 50,000-sample generations per configuration, and adaptive NFE counted batch evaluations while excluding samples already finished by masking.The metric pipeline was checked against EDM2 and produced identical FID results; second-order non-adaptive evaluations used odd NFE values.Source: §D.2.2 · p. 19

D.2.3 Hyperparameters

  • Image baselines used the EDM schedule, and stochastic Heun and Euler–Maruyama switched to ODE steps below t = 0.05; EDM-churn used Karras et al.’s [2] churn settings.PI hyperparameters were mostly shared across ImageNet-64 and FFHQ, with τrel and h0 slightly altered between datasets.Source: §D.2.3 · p. 19, Table 7 · p. 20

D.3.1 Experiment Details / D.3.2 Hyperparameters

  • The Bitstream setup extended Batzolis et al.’s [35] framework with PI-solver support and self-conditioning; its text-sampling SDE’s noise function controls added noise and is zero for the probability-flow ODE.For PI configurations, τrel and h0 were tuned by NFE, and static schedules were generated from 1,024 samples at the target NFE.Source: §D.3.1 · p. 20, §D.3.2 · pp. 20–21

D.4.1 PI Solver Ablations

  • The FFHQ ablations varied proportional control, the tolerance form, and the error estimator relative to the PI solver configuration.Configuration B sets kp = 0; configurations C and D use alternative error-estimation choices.Source: §D.4.1 · p. 21, Table 9 · p. 21

D.4.2 Schedule Across Modalities

  • Cross-modality tests applied an LM1B PI static schedule to FFHQ and an FFHQ PI static schedule to LM1B.The LM1B schedule was generated at 255 NFE and the FFHQ schedule at 99 NFE.Source: §D.4.2 · p. 21

E Samples

The appendix provides generated ImageNet-64 and FFHQ sample panels from different SDE solvers at three function-evaluation budgets. The captions identify datasets, checkpoints, and NFE, but do not state a quality comparison.

E.1 ImageNet-64

E.2 FFHQ

Loading 2610.03034v1…