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

Fuente: arXiv
Guardado en:
Detalles Bibliográficos
Autores principales: Huang, Yilie, Tang, Wenpin, Zhou, Xunyu
Formato: Preprint
Publicado: 2026
Materias:
Acceso en línea:
Etiquetas: Agregar Etiqueta
Sin Etiquetas, Sea el primero en etiquetar este registro!
_version_ 1866910199874322432
author Huang, Yilie
Tang, Wenpin
Zhou, Xunyu
author_facet Huang, Yilie
Tang, Wenpin
Zhou, Xunyu
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2601_18681
institution arXiv
publishDate 2026
record_format arxiv
spellingShingle ART for Diffusion Sampling: A Reinforcement Learning Approach to Timestep Schedule
Huang, Yilie
Tang, Wenpin
Zhou, Xunyu
Machine Learning
Artificial Intelligence
Systems and Control
Optimization and Control
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.
title ART for Diffusion Sampling: A Reinforcement Learning Approach to Timestep Schedule
topic Machine Learning
Artificial Intelligence
Systems and Control
Optimization and Control
url https://arxiv.org/abs/2601.18681