ART for Diffusion Sampling: A Reinforcement Learning Approach to Timestep Schedule

Technical Summary: ART for Diffusion Sampling

Problem Statement
Score-based diffusion models generate samples by integrating a learned reverse-time dynamic. A critical numerical choice in this process is the discretization of the time horizon into a grid of timesteps. Standard approaches typically employ uniform grids or hand-crafted schedules (e.g., EDM, DPM-Solver). The authors argue that these fixed schedules are often suboptimal given a fixed budget of function evaluations (NFE). Uniform grids allocate computational effort evenly, even when the reverse dynamics vary significantly along the trajectory, while hand-crafted grids may not adapt to specific data distributions or score function complexities. The goal is to systematically redistribute computational effort along the sampling trajectory to minimize the aggregate Euler discretization error.

Methodology
The paper introduces Adaptive Reparameterized Time (ART), a control-theoretic framework for optimal timestep scheduling.

1. Control Formulation (ART):

   • Instead of selecting grid points directly, ART treats the "clock speed" of the diffusion process as a control variable.
   • It introduces a reparameterized time variable $t$ mapped to the original diffusion time $\tau$ via a continuous mapping $\psi(t)$.
   • The control $\theta(t) = \dot{\psi}(t)$ represents the instantaneous rate of change of diffusion time.
   • The objective is to choose $\theta(t)$ to minimize the aggregate local Euler discretization error, subject to the constraint that the total time budget is preserved ($\int_0^T \theta(t) dt = T$).
   • The local error proxy is derived as proportional to $|Q(x, \psi)| \theta(t)^2$, where $Q$ depends on the score function and its derivatives.

2. Reinforcement Learning Solution (ART-RL):

   • The deterministic ART problem is intractable to solve directly due to the high-dimensional state space and the lack of a closed-form solution for the Hamilton-Jacobi-Bellman (HJB) equation.
   • The authors propose ART-RL, a continuous-time reinforcement learning (CTRL) approach.
   • They introduce a randomized auxiliary problem where the control $\theta$ is sampled from a Gaussian policy $\pi(\cdot | t, x, \psi) = \mathcal{N}(\mu, \lambda/|Q|)$.
   • This formulation allows the use of actor-critic algorithms (specifically adapting the framework of Jia and Zhou [14]) to learn the policy mean $\mu$ and value function via stochastic approximation.

3. Theoretical Bridge:

   • A core theoretical contribution is the establishment of a two-directional bridge between the deterministic ART optimum and the optimal Gaussian policy in the ART-RL auxiliary problem.
   • Sufficiency: The deterministic optimal control $\mu^*$ lifts to an optimal Gaussian policy whose mean is exactly $\mu^*$.
   • Necessity: Any optimal Gaussian policy for the auxiliary problem must have a mean that recovers the deterministic ART control $\mu^*$.
   • This result justifies parameterizing the actor as a Gaussian mean and deploying it deterministically after training, providing rigorous guarantees that the learned schedule is the true optimum for the original control problem.

4. Distillation:

   • After offline training, the stochastic policy is distilled into a deterministic, time-only grid. This removes the need for per-step actor network evaluations or variance calculations during inference, ensuring zero additional computational cost compared to standard schedulers.

Key Contributions

• Formulation: Proposes ART, the first control-theoretic framework for diffusion timestep scheduling via adaptive time reparameterization.
• Theory: Establishes the two-directional equivalence between a deterministic optimal control problem and a continuous-time actor-critic learning procedure with Gaussian policies. This transforms heuristic RL applications into a principled route for finding deterministic optima.
• Algorithm: Develops ART-RL, an actor-critic algorithm that learns adaptive time schedules data-drivenly.
• Generalization: Demonstrates that the learned schedule can be distilled and transferred across different datasets and timestep counts without retraining.

Experimental Results
The authors evaluate ART-RL within the official EDM pipeline across CIFAR-10, AFHQv2, FFHQ, and ImageNet.

• CIFAR-10: ART-RL consistently outperforms Uniform, DPM-Solver log-SNR, and standard EDM schedules across a wide range of NFE budgets.
  • At low budgets (e.g., NFE=3), ART-RL produces recognizable objects where baselines remain near noise.
  • At NFE=35, ART-RL achieves an FID of 1.82, improving upon EDM (1.85) and DPM-logSNR (1.89).
  • Statistical significance is confirmed via multiple seeds, with ART-RL's worst run beating EDM's best.
• Generalization:
  • Interpolation/Extrapolation: A schedule learned at $K=18$ steps on CIFAR-10 retains its performance advantage when interpolated or extrapolated to other step counts ($K \in \{4, 9, 12, 15, 20\}$).
  • Cross-Dataset Transfer: The CIFAR-10 distilled schedule is applied to AFHQv2, FFHQ, and ImageNet without retraining. It consistently improves FID over EDM and DPM-logSNR, particularly at low and moderate NFEs.
• One-Dimensional Analysis: In a setting with an analytical score function, ART-RL significantly reduces the Wasserstein distance ($W_2$) compared to hand-crafted schedules, confirming that the improvement stems from better time allocation rather than score estimation errors.

Significance and Claims
The paper claims to be the first to formulate diffusion timestep scheduling as a control problem via adaptive time reparameterization and the first to provide a rigorous two-directional bridge between a deterministic optimum and an actor-critic learning procedure.

The authors emphasize that their method is data-driven and amortized:

• Data-driven: The schedule is learned from the objective function and data dynamics, not hand-crafted.
• Amortized Transfer: The offline training cost (approx. 1–2 hours on a Colab T4 for CIFAR-10) is a one-time expense. The resulting distilled deterministic grid is reusable across different datasets and step counts with no extra inference cost, matching the runtime of standard schedulers like Uniform or EDM.

The work positions adaptive time reparameterization as a nascent but promising direction for systematic schedule design in generative diffusion models, offering a control-theoretic starting point for future research.