Technical Summary: ART for Diffusion Sampling

Problem Statement

Diffusion models generate data by transforming noise into samples via a learned reverse-time dynamics. While the training of the score (or denoising) model is well-established, the sampling stage relies on numerically discretizing the reverse process on a finite grid of timesteps. Current standard practices utilize uniform grids or hand-crafted schedules (e.g., EDM, DPM-Solver). These approaches are typically ad hoc, fixed prescriptions that do not adapt to the specific numerical characteristics of the trajectory. Consequently, they may allocate computational resources suboptimally, leading to either significant discretization errors (if steps are too coarse) or wasted evaluations (if steps are too fine in regions where the dynamics are smooth). The core problem is to formulate a systematic, principled method for timestep allocation that optimizes sample quality under a fixed computational budget (number of function evaluations).

Methodology

1. Adaptive Reparameterized Time (ART)

The authors propose ART, a continuous-time optimal control framework that treats timestep selection as a control problem.

• Time Reparameterization: Instead of evolving the reverse process on a fixed physical time grid $\tau \in [0, T]$, ART introduces an auxiliary "sampling clock" $t \in [0, T]$. The physical time is a function of this clock, $\tau = \psi(t)$.
• Control Variable: The rate of progression in physical time relative to the sampling clock, $\theta(t) = \dot{\psi}(t)$, is treated as the control variable.
• Objective: The goal is to maximize an objective functional derived from the leading-order Euler discretization error. The local error of the Euler scheme is shown to be proportional to $|Q(x, \psi)| \theta^2$, where $Q$ captures the local stiffness of the probability-flow vector field. The control problem seeks to minimize the accumulated error subject to a total time budget constraint $\int_0^T \theta(t) dt = T$.
• Formulation: This results in a deterministic optimal control problem where the state includes the sample trajectory and the accumulated physical time, and the control $\theta$ dictates how fast or slow the sampler moves through physical time.

2. ART-RL: A Randomized Control Formulation

Solving the deterministic ART problem directly via the Hamilton–Jacobi–Bellman (HJB) equation is intractable due to the high dimensionality of the state space (e.g., image dimensions). To overcome this, the authors introduce ART-RL, an auxiliary randomized control problem.

• Gaussian Policies: The deterministic control $\theta$ is replaced by a stochastic policy $\pi(\lambda)$, specifically a Gaussian distribution $N(\mu, \lambda/|Q|)$. The variance is inversely proportional to the local stiffness $|Q|$, ensuring less randomness in stiff regions.
• Theoretical Equivalence: The paper proves that the mean of the optimal Gaussian policy in the ART-RL formulation is identical to the optimal control solution of the original deterministic ART problem. Thus, solving the randomized problem recovers the solution to the deterministic one.
• Actor–Critic Learning: Leveraging recent advances in continuous-time reinforcement learning (CTRL), the authors derive implementable actor–critic update rules.
  • Critic: Estimates the value function (policy evaluation) using martingale-based moment conditions.
  • Actor: Updates the policy mean $\mu$ to improve the value (policy improvement) based on the critic's feedback.
  • Lagrange Multiplier: A multiplier $\gamma$ is updated to enforce the total time budget constraint.

3. Implementation and Distillation

The learned policy is a feedback function of time, state, and accumulated time. However, empirical results show that in practice, the optimal control $\theta$ depends primarily on time. The authors propose a distillation step where the learned neural network policy is replaced by a fixed, deterministic time-only schedule (the empirical mean of the control). This allows the learned schedule to be deployed as a simple precomputed grid, incurring no additional inference-time overhead compared to hand-crafted schedules.

Key Contributions

1. Control-Theoretic Framework: The paper is the first to formulate timestep allocation in diffusion sampling as a continuous-time optimal control problem (ART), moving beyond heuristic grid designs.
2. Theoretical Bridge (ART-RL): It establishes a rigorous link between the deterministic control problem and a randomized reinforcement learning formulation. It proves that the mean of the optimal randomized policy solves the original deterministic problem, enabling the use of actor–critic algorithms.
3. Algorithmic Development: The authors derive specific policy evaluation and improvement characterizations for time reparameterization, leading to a practical, data-driven actor–critic algorithm (Algorithm 1) for learning optimal schedules.
4. Generalizability and Transfer: The work demonstrates that a schedule learned on a specific dataset (CIFAR-10) can be transferred without retraining to different datasets (AFHQv2, FFHQ, ImageNet), different resolutions (64x64 to 512x512), different pipelines (EDM to EDM2), and different representation spaces (pixel to latent).

Experimental Results

The authors evaluate ART-RL across diverse settings, keeping the score model, solver, and pipeline fixed, changing only the timestep grid:

• Low-Dimensional Analytical Case: Using a 1D diffusion with a known score function, ART-RL consistently outperforms Uniform, DPM, and EDM schedules in terms of Wasserstein-2 error, demonstrating that hand-crafted schedules designed for images may fail in other domains.
• CIFAR-10 (EDM Pipeline):
  • Heun Sampling: ART-RL achieves the best Fréchet Inception Distance (FID) across all budgets, including the strongest configuration (NFE=35), where it improves upon the state-of-the-art EDM schedule (1.82 vs 1.85).
  • Euler Sampling: Similar improvements are observed, confirming the benefit of learning the grid even with first-order solvers.
• MNIST (Small Model): Even with a lightweight, non-pretrained score model and an RK4 solver, ART-RL significantly outperforms baselines, showing the method's robustness to model capacity.
• Transfer Learning:
  • Across Step Counts: A schedule learned at $K=18$ transfers effectively to other step counts via log-linear interpolation/extrapolation.
  • Across Datasets: The CIFAR-10 learned schedule improves FID on AFHQv2, FFHQ, and ImageNet-64 without retraining.
  • Across Pipelines and Resolutions: The schedule transfers successfully to the EDM2 pipeline on ImageNet-512 (latent space, high resolution), outperforming both EDM and DPM baselines.

Significance and Claims

The paper claims that ART-RL provides a principled, data-driven alternative to the existing ad hoc, fixed heuristic grids used in diffusion sampling.

• Reusability: The primary practical significance is the ability to learn a reusable timestep schedule. Once trained and distilled on a single dataset (e.g., CIFAR-10), the schedule can be applied to a wide variety of downstream tasks (different datasets, resolutions, and pipelines) without retraining, amortizing the offline training cost.
• Drop-in Replacement: The method requires no changes to the score model architecture, the solver implementation, or the training protocol. It functions purely as a replacement for the timestep grid, making it a modular improvement for existing diffusion pipelines.
• Theoretical Grounding: Unlike previous heuristic approaches, ART is grounded in optimal control theory and continuous-time reinforcement learning, offering a systematic way to address the discretization error in diffusion sampling.

The authors modestly note that their current analysis is restricted to probability flow ODE sampling and that the distillation to time-only schedules removes state dependence, leaving open the question of whether richer state-conditioned schedules could yield further benefits. However, the current results establish a strong foundation for systematic timestep design in generative modeling.