Technical Summary: Noise Schedule Design for Diffusion Models: An Optimal Control Perspective

Problem Statement

Diffusion models have emerged as a leading framework for approximate sampling, formulating generation as the simulation of a time-reversed stochastic differential equation (SDE) driven by a learned score function. While the training objective (score matching) is well-understood, the design of noise schedules—the functions $f(t)$ and $g(t)$ governing the forward and reverse processes—remains largely heuristic.

Existing theoretical work establishes sampling error bounds (measured in Kullback-Leibler divergence) for specific, often constant, noise schedules. However, these bounds typically scale sub-optimally with the data dimension $d$ and problem constants, often appearing as $\tilde{O}(d^2/n)$ or $\tilde{O}(\max\{J^\star, d\}/n)$, where $J^\star$ is the Fisher information of the data distribution and $n$ is the number of discretization steps. Since $J^\star$ can be significantly larger than $d$, these guarantees suffer from severe problem-dependent penalties. Conversely, state-of-the-art theoretical bounds of $\tilde{O}(d/n)$ exist but rely on specific noise schedules that are rarely deployed in practice, creating a gap between diffusion model theory and empirical implementation.

Methodology

The authors propose a principled framework to recast noise schedule design as an optimal control problem. The core innovation lies in circumventing the intractable optimization of probability density evolution (governed by the Fokker-Planck partial differential equation) by focusing on the Fisher information trajectory.

1. Optimal Control Formulation:

   • State: The Fisher information $J_t = J(p_t)$ of the diffusion process, which evolves according to a one-dimensional ordinary differential equation (ODE) derived via Bochner's formula:
     $$\dot{J}_t = 2f(t)J_t - g(t)\|\nabla s_t\|_{L^2}^2$$
   • Control: The noise schedule functions $(f, g)$.
   • Objective: Minimize a functional involving the Fisher information that serves as an upper bound on the KL sampling error. The authors utilize tools from probability and information theory, including Girsanov's theorem, Bochner's formula, the concavity of entropy power, and the Blachman-Stam inequality, to derive this bound.

2. Derivation of Closed-Form Schedules:

   • Under additional assumptions that the target data distribution is either Strongly Log-Concave (SLC) or a Gaussian Mixture Model (GMM), the authors derive a closed-form characterization for the optimal noise schedule.
   • They introduce a class of schedules termed Affine-Coupled Schedules (ACS), where $f(t)$ and $g(t)$ satisfy an affine relationship: $f(t) = \theta g(t) + \omega$.
   • Solving the resulting ODEs yields explicit expressions for $f(t)$ and $g(t)$ parameterized by $(\theta, \omega, g_0, \lambda^*)$.

3. Adaptive Parameterization:

   • The framework treats the derived parameters as hyperparameters that can be tuned based on the computational budget (Neural Function Evaluations, or NFE, denoted as $n$).
   • The optimal decay rate $\lambda^*$ is shown to scale dynamically with $n$ via the Lambert-W function, allowing the schedule to automatically balance initialization error and discretization error.

Key Contributions

1. Optimal Control Framework: The paper formulates noise schedule design as an optimal control problem where the state is the Fisher information trajectory. This replaces the intractable PDE constraints with a tractable ODE, enabling the derivation of sufficient conditions for optimal error bounds.
2. $\tilde{O}(d/n)$ Error Bounds: The authors prove that the solution to their optimal control problem achieves a sampling error bound of $\tilde{O}(d/n)$ in KL divergence. This extends state-of-the-art theoretical guarantees to a large class of noise schedules, including those used in practice.
3. Affine-Coupled Schedules (ACS): They derive a new parametric class of noise schedules (ACS) that generalizes standard empirical schedules:
   • VE-Exponential: Recovered by setting parameters $(\theta, \omega) = (0, 0)$.
   • VP-Sigmoid: Recovered by setting $(\theta, \omega) = (1/2, 0)$.
   • The ACS framework introduces additional tunable parameters ($g_0, c$) that allow for systematic optimization beyond these fixed heuristics.
4. Empirical Validation: Experiments on CIFAR-10 and ImageNet (64x64) demonstrate that systematically tuning the ACS hyperparameters yields superior Fréchet Inception Distance (FID) scores compared to standard baselines (VP-Linear, VP-Cosine, VP-Sigmoid) across various NFE budgets.

Results

• Theoretical: The paper establishes that for SLC and GMM distributions, the proposed ACS framework achieves the optimal $\tilde{O}(d/n)$ convergence rate. It further analyzes why constant or linear schedules (common in practice) may suffer from sub-optimal trade-offs, particularly when the data distribution is highly anisotropic ($J^\star \gg d$).
• Empirical: On unconditional CIFAR-10 and class-conditional ImageNet, the adaptive ACS framework consistently outperforms standard baselines. For instance, at an NFE of 50 on ImageNet, ACS achieved an FID of 5.30, outperforming VP-Linear (7.46), VP-Cosine (5.67), and VP-Sigmoid (6.77). The results indicate that ACS adapts its curvature and variance trajectory effectively to the computational budget, whereas fixed schedules degrade in quality as the step size changes.

Significance

The paper claims to bridge the gap between diffusion model theory and practice by providing a principled method for noise schedule design. Rather than relying on empirical sweeps or fixed heuristics, the authors offer a theoretical justification for why certain schedules work and a mechanism to generate new, optimized schedules.

The significance lies in:

1. Generalization: Showing that standard empirical schedules (Exponential, Sigmoid) are special cases of a broader, theoretically optimal class (ACS).
2. Adaptability: Providing a method to dynamically adjust schedules based on the available compute budget ($n$) and intrinsic data properties (Fisher information), rather than using static schedules.
3. Performance: Demonstrating that this theoretically grounded approach leads to tangible improvements in generation quality (FID scores) on standard benchmarks, validating the utility of the optimal control perspective in generative modeling.

The authors note that while the closed-form derivation assumes SLC or GMM distributions, the resulting schedules perform well on complex real-world data, suggesting the theoretical insights generalize beyond the specific assumptions made for derivation. Future work is identified as extending these characterizations to broader classes of complex data distributions.