Adaptive Reparametrized Time for Score-Based Diffusion Sampling
Abstract
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.
Related
- Score Mismatch and Set-Constrained Sampling in Zero-Shot Conditional Discrete Diffusion
- Serial Depth of Adaptive Parallel Sampling: Lower Bounds and Schedules for Masked Diffusion Models
- Rollcast: Proper-Score Gated Rolling Anchors for Adaptive Probabilistic Time-Series Forecasting
- Approximate Matrix Multiplication through Joint Leverage Score Sampling
- Post-Corrected Raw-Score Martingale Posterior Sampling for von Mises-Fisher Models
BibTeX
@misc{huang2026adaptive,
title = {{Adaptive Reparametrized Time for Score-Based Diffusion Sampling}},
author = {Yilie Huang and Wenpin Tang and Xun Yu Zhou},
howpublished = {arXiv},
year = {2026},
month = sep,
doi = {10.48550/arxiv.2607.02137},
eprint = {2607.02137},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2607.02137},
note = {CC BY 4.0},
}
RIS
TY - UNPB
TI - Adaptive Reparametrized Time for Score-Based Diffusion Sampling
AU - Yilie Huang
AU - Wenpin Tang
AU - Xun Yu Zhou
PY - 2026
DA - 2026/09/30
PB - arXiv
DO - 10.48550/arxiv.2607.02137
UR - https://arxiv.org/abs/2607.02137
AB - 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.
ER -