Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Yilie Huang, Wenpin Tang, Xunyu Zhou

TL;DR

This work tackles the challenge of timestep scheduling in score-based diffusion sampling, where uniform or handcrafted grids can waste compute under a fixed budget. It introduces Adaptive Reparameterized Time (ART), a time-warping framework that redistributes computation along the sampling trajectory while preserving the terminal time, and derives an optimization objective tied to the Euler discretization error. To solve ART, the authors formulate ART-RL, a continuous-time reinforcement learning auxiliary problem with Gaussian policies that yields an optimal mean control; they prove a value-shift relationship and policy-recovery results that connect ART-RL to the original ART solution. Empirically, ART-RL improves FID across budgets in CIFAR-10 and transfers effectively to AFHQv2, FFHQ, and ImageNet within the EDM pipeline, with a distilled time-only schedule offering a training-free drop-in replacement. The approach provides a principled, data-driven method for adaptive time discretization in diffusion sampling and opens avenues for solver-aware and cross-dataset generalization in diffusion-based generation.

Abstract

We consider time discretization for score-based diffusion models to generate samples from a learned reverse-time dynamic on a finite grid. Uniform and hand-crafted grids can be suboptimal given a budget on the number of time steps. We introduce Adaptive Reparameterized Time (ART) that controls the clock speed of a reparameterized time variable, leading to a time change and uneven timesteps along the sampling trajectory while preserving the terminal time. The objective is to minimize the aggregate error arising from the discretized Euler scheme. We derive a randomized control companion, ART-RL, and formulate time change as a continuous-time reinforcement learning (RL) problem with Gaussian policies. We then prove that solving ART-RL recovers the optimal ART schedule, which in turn enables practical actor--critic updates to learn the latter in a data-driven way. Empirically, based on the official EDM pipeline, ART-RL improves Fréchet Inception Distance on CIFAR-10 over a wide range of budgets and transfers to AFHQv2, FFHQ, and ImageNet without the need of retraining.

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

TL;DR

This work tackles the challenge of timestep scheduling in score-based diffusion sampling, where uniform or handcrafted grids can waste compute under a fixed budget. It introduces Adaptive Reparameterized Time (ART), a time-warping framework that redistributes computation along the sampling trajectory while preserving the terminal time, and derives an optimization objective tied to the Euler discretization error. To solve ART, the authors formulate ART-RL, a continuous-time reinforcement learning auxiliary problem with Gaussian policies that yields an optimal mean control; they prove a value-shift relationship and policy-recovery results that connect ART-RL to the original ART solution. Empirically, ART-RL improves FID across budgets in CIFAR-10 and transfers effectively to AFHQv2, FFHQ, and ImageNet within the EDM pipeline, with a distilled time-only schedule offering a training-free drop-in replacement. The approach provides a principled, data-driven method for adaptive time discretization in diffusion sampling and opens avenues for solver-aware and cross-dataset generalization in diffusion-based generation.

Abstract

We consider time discretization for score-based diffusion models to generate samples from a learned reverse-time dynamic on a finite grid. Uniform and hand-crafted grids can be suboptimal given a budget on the number of time steps. We introduce Adaptive Reparameterized Time (ART) that controls the clock speed of a reparameterized time variable, leading to a time change and uneven timesteps along the sampling trajectory while preserving the terminal time. The objective is to minimize the aggregate error arising from the discretized Euler scheme. We derive a randomized control companion, ART-RL, and formulate time change as a continuous-time reinforcement learning (RL) problem with Gaussian policies. We then prove that solving ART-RL recovers the optimal ART schedule, which in turn enables practical actor--critic updates to learn the latter in a data-driven way. Empirically, based on the official EDM pipeline, ART-RL improves Fréchet Inception Distance on CIFAR-10 over a wide range of budgets and transfers to AFHQv2, FFHQ, and ImageNet without the need of retraining.
Paper Structure (36 sections, 2 theorems, 50 equations, 7 figures, 7 tables, 1 algorithm)

This paper contains 36 sections, 2 theorems, 50 equations, 7 figures, 7 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Diffusion Models and ART
  3. Continuous-Time Score-Based Diffusion Models
  4. Diffusion sampling.
  5. Euler discretization.
  6. ART: Time Reparameterization via Control
  7. Randomized Control and RL Formulation
  8. ART-RL: Auxiliary Problem with Gaussian Policies
  9. Relationship between ART and ART-RL
  10. ART-RL Actor--Critic Algorithm
  11. Algorithm Design
  12. Implementable ART-RL Algorithm
  13. Numerical Experiments
  14. Datasets.
  15. Timestep schedules and baselines.
  16. ...and 21 more sections

Key Result

Theorem 3.1

If $V$ is a classical solution of eq_original_hjb, then is a classical solution of eq_auxiliary_hjb.

Figures (7)

  • Figure 1: Empirical mean (solid line) and 25–75 percent interquartile range (shaded region) of the learned control $\theta$ across time.
  • Figure 2: ImageNet samples under EDM and ART-RL schedules at increasing NFEs (top to bottom).
  • Figure 3: Empirical mean of the executed control $\theta$ and its 99 percent confidence interval, based on the last $10{,}000$ trajectories in the one–dimensional experiment.
  • Figure 4: CIFAR--10 samples across timesteps for the three schedules (Uniform, EDM, ART-RL). Each panel shows a grid where rows correspond to increasing NFEs.
  • Figure 5: CIFAR--10 samples across timesteps for interpolated and extrapolated grids (EDM and ART-RL). Each panel shows a grid where rows correspond to increasing NFEs.
  • ...and 2 more figures

Theorems & Definitions (4)

  • Theorem 3.1: Value function shift
  • Theorem 3.2: Recovery of the optimal ART control
  • proof : Proof of Theorem \ref{['thm_value_shift']}
  • proof : Proof of Theorem \ref{['thm_policy_recovery']}