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

TL;DR

This paper introduces ART, an adaptive time discretization method for diffusion models, optimized via reinforcement learning, leading to improved sample quality and transferability across datasets.

Contribution

The paper proposes ART and ART-RL, novel methods for adaptive time discretization in diffusion models, bridging deterministic control and reinforcement learning for improved sampling.

Findings

ART improves FID scores on CIFAR-10 across various budgets.

The deterministic ART schedule transfers effectively to other datasets without retraining.

ART-RL enhances diffusion sampling by optimizing timestep schedules via reinforcement learning.

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), which controls the clock speed of a reparameterized time variable to redistribute computation along the sampling trajectory while preserving the terminal time, with the objective of minimizing the aggregate Euler discretization error. We derive a randomized companion ART-RL that recasts ART as a continuous-time reinforcement learning problem with Gaussian policies, and prove a two-directional bridge between the two: the deterministic ART optimum lifts to an optimal Gaussian policy, and conversely any optimal Gaussian policy must recover the ART control through its mean. This bridge turns continuous-time actor--critic learning into a principled, rather than heuristic, route to the deterministic timestep optimum. Within the official EDM pipeline, ART-RL improves FID on CIFAR--10 across a wide range of budgets; after one-time offline training, the distilled deterministic schedule transfers without retraining to AFHQv2, FFHQ, and ImageNet at no extra inference cost.