Noise Schedule Design for Diffusion Models: An Optimal Control Perspective

The gist The authors develop a principled framework for analyzing and designing noise schedules in diffusion models by recasting the design problem as an optimal control problem, which yields sufficient conditions for achieving state-of-the-art O(d/n) sampling error through new classes of Affine-Coupled Schedules.

How it works

The paper reformulates the noise schedule design as an optimal control problem where the state is the Fisher information of the diffusion process and the control input is the noise schedule, with a cost functional involving a variational bound on discretization error The objective of this problem is shown to be an upper bound on the Kullback-Leibler sampling error, which can be minimized by solving for an optimal trajectory of the Fisher information <ref:2605.21911#pg5>. This optimal control formulation is significantly more tractable because it replaces the intractable PDE representation of probability density evolution with a one-dimensional ODE governing the trajectory of the Fisher information <ref:2605.21911#pg5>.

Key theoretical results

The solution to this optimal control problem achieves O(d/n) error bounds, extending state-of-the-art sampling error guarantees to a large class of noise schedules <ref:2605.21911#pg6>. For special classes of data distributions, the analysis extracts closed-form expressions for a family of noise schedules termed Affine-Coupled Schedules (ACS) <ref:2605.21911#pg6>. These results show that standard empirical heuristics, such as VE-Exponential and VP-Sigmoid models, are special cases within this framework <ref:2605.21911#pg6>.

Noise Schedule Design

The analysis leads to the characterization of the optimal noise schedule g∗ as g∗(t) = λ∗ J⋆ t ∇st2L2, where λ∗ is a constant independent of t and J⋆ is the Fisher information of the data distribution <ref:2605.21911#pg5>. Under additional assumptions on the target data distribution, such as being strongly log-concave or a Gaussian mixture model, closed-form expressions for g◦ are derived <ref:2605.21911#pg7>. These schedules generalize standard empirical schedules by introducing additional parameters that can be tuned to achieve superior FID scores on image generation benchmarks <ref:2605.21911#pg7>.

Practical implications and experiments

The introduced parameters in the optimal schedule can be treated as hyperparameters that are optimally configured for any given Neural Function Evaluation (NFE) budget, which allows for a systematic optimization of schedules across different computational budgets <ref:2605.21911#pg7>. The authors empirically demonstrate that this systematic tuning yields schedules that achieve superior FID scores on image datasets compared to standard heuristic baselines <ref:2605.21911#pg7>. Experiments show that the Affine-Coupled Schedule (ACS) achieves better performance than VP-Linear, VP-Cosine, and VP-Sigmoid schedules across various NFE budgets <ref:2605.21911#pg8>.

Error Bounds and Trade-offs

The analysis establishes a trade-off between initialization error and discretization error, showing that the optimal schedule g◦ achieves the same rate as Theorem 2 when properly chosen <ref:2605.21911#pg9>. The optimal continuous trajectory is uniquely defined by differentiating a variance trajectory V∗t, which in turn defines the noise schedule g∗(t) = α2t dV∗t / dt <ref:2605.21911#pg9>. This adaptive approach balances the fundamental trade-off between initialization and discretization errors, enabling ACS to consistently outperform heuristic baselines across diverse compute regimes <ref:2605.21911#pg11>. The final result shows that the optimal schedule g◦ achieves a better trade-off with the dependence E squared + log2 J⋆, compared to VP-constant and VP-linear schedules <ref:2605.21911#pg12>.

Conclusion

The work introduces an optimal control framework for designing noise schedules utilizing a Fisher information ODE to bypass intractable PDE constraints <ref:2605.21911#pg5>. The approach guarantees state-of-the-art O(d/n) KL divergence bounds and yields a novel class of Affine-Coupled Schedules (ACS) that dynamically adapt to compute budgets, achieving superior FID scores on standard benchmarks <ref:2605.21911#pg11>. Future work will explore extending theoretical characterizations of noise schedules that achieve the O(d/n) bound to broader classes of complex data distributions <ref:2605.21911#pg11>. The paper concludes by providing a method for practitioners to go beyond commonly used schedules and further improve performance through tunable parameters <ref:2605.21911#pg7>.

--- Page 1 ---

Noise Schedule Design for Diffusion Models: An Optimal Control Perspective

Seo Taek Kong

Weina Wang

R. Srikant

--- Page 3 ---

Forward process. Let p⋆ denote the density function of the target data distribution on R d.<ref:2605.21911#pg4> The forward process is a stochastic differential equation (SDE) of the following form:<ref:2605.21911#pg5> dXt = −f(t)Xtdt + p g(t)dBt, 0 ≤ t ≤ T, X0 ∼ p⋆, (1)<ref:2605.21911#pg5> where (Bt: t ∈ [0, T]) is a standard Brownian motion on R d independent of X0,<ref:2605.21911#pg5> and f(·) and g(·) are real-valued functions of time.<ref:2605.21911#pg5> Let pt denote the density function of Xt.<ref:2605.21911#pg5> We call the function pair (f, g) the noise schedule.<ref:2605.21911#pg5> Many theoretical studies of diffusion models focus on the simplified setting f(t) = 1 and g(t) = 2 for all t, in which case the forward process is the standard Ornstein–Uhlenbeck (OU) process.<ref:2605.21911#pg5> To see how the noise schedule controls the distribution of the forward process, note that the marginal distribution of Xt admits the representation<ref:2605.21911#pg5> Xt d= αtX0 + σtZ, Z ∼ N (0, I), (2)<ref:2605.21911#pg5> where αt = exp − Z t 0 f(t˜)dt˜ and σ2 t = α squared t Z t 0 g(t˜) α squared t˜ dt.<ref:2605.21911#pg5> In the OU setting where f(t) = 1, g(t) = 2, these quantities reduce to αt = e −t, σ2 t = 1 − e −2t.<ref:2605.21911#pg5> More generally, a larger value of f(t) leads to a faster decay of the contribution of X0 in the distribution of Xt, while a larger value of g(t) corresponds to a faster injection of Gaussian noise.<ref:2605.21911#pg5> In this paper, we consider the class of noise schedules Θ = F × G defined as F = ffl for all t ∈ [0, T] and G = ffl for all t ∈ [0, T]<ref:2605.21911#pg5> The noise schedules used in most diffusion model implementations [39] fall under the class Θ.<ref:2605.21911#pg5>

--- Page 2 ---

Reverse process. The reverse process is the time reversal of the forward process (Xt: t ∈ [0, T]) in (1). That is, the reverse process (X←τ: τ ∈ [0, T]) is defined by X←τ d= XT −τ, τ ∈ [0, T]. Then the reverse process admits the following SDE representation <ref:2605.21911#pg5> dX←τ = [f(T − τ)X←τ + g(T − τ)sT−τ (X←τ)] dτ + p g(T − τ)dBτ, 0 ≤ τ ≤ T, X←0 ∼ pT, (4)<ref:2605.21911#pg5> where (Bτ: 0 ≤ τ ≤ T) is a standard Brownian motion on R d independent of X←0,<ref:2605.21911#pg5> and the function st(·) defined as st(x) ≜ ∇ log pt(x) is called the score function.<ref:2605.21911#pg5>

--- Page 3 ---

2.1 Noise Schedule Optimization. To turn the reverse process (4) into an implementable algorithm for generating data, several approximations need to be made: (1) The score function st is unknown since it depends on the unknown data distribution, and is therefore replaced by an approximation sˆt, typically learned through a training process called score matching<ref:2605.21911#pg5> (2) The reverse SDE cannot be solved exactly, but rather must be solved numerically with time discretization<ref:2605.21911#pg5> (3) The reverse SDE cannot be initialized exactly from pT, and is instead initialized from a Gaussian distribution so the initial state can be easily sampled.<ref:2605.21911#pg5> A central question in the analysis of diffusion models is how these approximations affect the error between the distribution of the generated samples and the target distribution [4, 6, 11, 15] The existing analysis shows that the overall error can be separated into three additive terms: (1) score-matching error, (2) discretization error, and (3) initialization error,<ref:2605.21911#pg5> The score-matching error arises from the training stage, while the discretization error and the initialization error arise from the sampling stage when one solves the reverse SDE to generate data samples.<ref:2605.

1. Bold header: Optimal Noise Schedule Design via Optimal Control Framework

This framework allows for a principled method to introduce new degrees of freedom into schedule design, enabling a systematic optimization of schedules across different computational budgets by treating noise schedule design as an optimal control problem minimizing a functional involving the Fisher information, which is shown to be an upper bound on the Kullback-Leibler sampling error.

2. Bold header: Guaranteed Sample Quality for State-of-the-Art Error Bounds

The system can achieve O˜(d/n) sampling error by solving the optimal control problem, which provides sufficient conditions on noise schedules under which state-of-the-art O˜(d/n) sampling error is achievable.

3. Bold header: Data Distribution Specific Schedule Derivation

For data distributions satisfying Assumption 2 (strongly log-concave and Gaussian mixture models), the system can obtain closed-form expressions for a family of noise schedules which we term Affine-Coupled Schedules (ACS).

4. Bold header: Adaptive Computational Budget Scaling

The model can be configured to adapt its schedule parameters based on the Neural Function Evaluation (NFE) budget, as Theorem 2 indicates that the parameters of the noise schedule should vary with the Neural Function Evaluation (NFE) budget n. This allows for an adaptive ACS framework that automatically adjusts its curvature and variance trajectory for any chosen inference budget.

5. Bold header: Superior Performance Across Compute Regimes

By systematically tuning parameters, the system can yield schedules that achieve superior FID scores on image generation benchmarks compared to standard heuristic baselines, effectively balancing the trade-off between initialization and discretization errors.

We develop a principled framework for analyzing and designing noise schedules in diffusion models. We show that one can recast this design problem as an optimal control problem, whose state is the Fisher information of the diffusion process, which evolves according to an ODE, and whose control input is the noise schedule. The objective of the optimal control problem is a functional involving the Fisher information, which is shown to be an upper bound on the Kullback-Leibler sampling error. By solving this optimal control problem, we obtain sufficient conditions on noise schedules under which (d/n) sampling error is achievable, where d is the data dimension and n is the number of discretization steps. While existing theoretical works also prove that (d/n) sampling error bounds are achievable, these results hold for specific noise schedules that differ from those used in practice. Under additional assumptions on the data distribution, we derive a family of noise schedules that recovers the functional forms of schedules widely used in practice. We show that these schedules achieve the (d/n) error bound when their parameters are chosen according to our analysis. Our theory offers a principled approach to noise schedule design, yielding explicit functional forms and guidance for selecting their parameters.