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

Yilie Huang; Wenpin Tang; Xunyu Zhou

cs.LGJan 26, 2026

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

Authors: Yilie Huang, Wenpin Tang, Xunyu Zhou

Organizations: Department of Industrial Engineering and Operations Research, Columbia University, New York, NY 10027, USA. · Department of Industrial Engineering and Operations Research & Data Science Institute, Columbia University, New York, NY 10027, USA.

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.

Figures & tables

Figure 1 : Quantitative overview: ART-RL Pareto-dominates Uniform, DPM-logSNR, and EDM on CIFAR–10 across NFE budgets, and the same distilled grid retains the advantage on AFHQv2, FFHQ, and ImageNet without retraining. DPM-logSNR is the DPM-Solver uniform log-SNR grid.

NFE	3	9	13	19	35
Uniform	280.29	213.13	191.69	168.87	118.02
AYS	465.83	157.29	86.82	35.30	3.60
DPM-logSNR	465.83	52.29	8.67	2.76	1.89
EDM	465.83	35.54	6.79	2.54	1.85
ART-RL	152.86	32.13	5.44	2.45	1.82

Table 1 : CIFAR–10 FID vs. NFE. The analytic baselines (DPM-logSNR, EDM, AYS) all fix

\sigma_{\max}

and

\sigma_{\min}

as schedule boundaries and prescribe interior values analytically; at

\mathrm{NFE}=3

the schedule is essentially specified by these two boundary values, so the analytic baselines collapse to the same noise-level FID. ART-RL learns its schedule from data and is not pinned by an analytic boundary prescription at this budget.

NFE	7	17	23	29	39
DPM-logSNR	185.63	3.52	2.19	1.94	1.88
EDM	85.80	3.11	2.06	1.88	1.85
ART-RL	33.73	2.57	2.00	1.84	1.82

Table 2 : CIFAR–10 interpolation/extrapolation FID.

Panel of Figure 2 — Figure 2 : ImageNet samples across NFEs.

Dataset	Method	3	9	13	19	35
AFHQv2	DPM-logSNR	375.76	67.64	9.77	3.44	2.15
	EDM	375.76	27.88	7.56	2.99	2.11
	ART-RL	243.48	20.48	6.12	2.85	2.07
FFHQ	DPM-logSNR	466.76	113.87	15.94	5.25	2.66
	EDM	466.76	57.13	15.87	5.26	2.73
	ART-RL	305.97	35.73	11.08	4.31	2.57

Table 3 : Cross-dataset FID using the CIFAR–10 ART-RL schedule.

Explore similar work

Adaptive Reparametrized Time for Score-Based Diffusion Sampling
Jul 2, 2026Yilie Huang, Wenpin Tang, Xun Yu ZhouDiffusion Sampling Timestep
Department of Applied Mathematics, The Hong Kong Polytechnic University, Hong Kong. · Department of Industrial Engineering and Operations Research, Columbia University, New York, NY 10027, USA. · Department of Industrial Engineering and Operations Research and Data Science Institute, Columbia University, New York, NY 10027, USA.
Temporal Difference Learning for Diffusion Models
Jun 13, 2026Qizhen Ying, Yangchen Pan, Victor Adrian Prisacariu +1Denoising Trajectory Generative Models
Department of Engineering Science, University of Oxford, Oxford, United Kingdom · School of Computer Science, Carleton University, Ottawa, Canada.
ReDiF: Resource-Efficient Few-Step Diffusion Distillation via Reinforcement Learning
Dec 28, 2025Amirhossein Tighkhorshid, Zahra Dehghanian, Hamid R. RabieeDiffusion-Based Reinforcement Learning Methods Diffusion Sampling
Computer Engineering Sharif University of Technology

Jul 2, 2026, cs.LGJ/K move · Enter open · S save

Adaptive Reparametrized Time for Score-Based Diffusion Sampling

Yilie Huang, Wenpin Tang, Xun Yu Zhou

Department of Applied Mathematics, The Hong Kong Polytechnic University, Hong Kong. · Department of Industrial Engineering and Operations Research, Columbia University, New York, NY 10027, USA. · Department of Industrial Engineering and Operations Research and Data Science Institute, Columbia University, New York, NY 10027, USA.

Diffusion Sampling Timestep Score-Based Diffusion Model Soft Actor-Critic Reparameterization

Open paper AskPDF arXiv

We study timestep allocation for score-based diffusion sampling, where a learned reverse-time dynamics is discretized on a finite grid. Uniform and hand-crafted schedules are standard choices, but they rely on fixed prescriptions and can therefore be suboptimal. To address this limitation, we propose Adaptive Reparameterized Time (ART), a continuous-time control formulation that learns a time change by treating the speed of the sampling clock as the control, so that a uniform grid on the learned clock induces adaptive timesteps in the original diffusion time. Based on a leading-order Euler error surrogate, ART provides a principled objective for allocating timesteps along the sampling trajectory. To solve this deterministic control problem, we introduce ART-RL, an auxiliary randomized formulation with Gaussian policies that turns schedule learning into a continuous-time reinforcement learning problem. We prove that the randomized ART-RL formulation is equivalent to ART at the optimizer level, in the sense that its optimal Gaussian policy recovers the optimal ART time-warping rate through its mean. We further establish policy evaluation and policy improvement characterizations and derive trajectory-based moment identities that yield implementable actor--critic updates for learning the schedule. Across experiments ranging from controlled low-dimensional settings to image generation, ART-RL can be plugged into existing diffusion samplers by changing only the timestep grid, consistently improving sample quality over strong baseline schedules at matched budgets while leaving the rest of the sampling pipeline unchanged. The learned schedules also exhibit broad generalization, transferring without retraining across sampling budgets, datasets, solvers, pipelines, and representation spaces.

Abstract

Explore similar work

Adaptive Reparametrized Time for Score-Based Diffusion Sampling

Temporal Difference Learning for Diffusion Models

ReDiF: Resource-Efficient Few-Step Diffusion Distillation via Reinforcement Learning