A learned clock helps diffusion models spend sampling steps more wisely

The sampling clock is part of the algorithm

A diffusion model’s learned network is only one part of image generation. At sampling time, a numerical solver follows a reverse process from noise toward an image, evaluating the network at a series of time points. A uniform grid spreads those evaluations evenly; established alternatives use preset, hand-crafted grids. Either way, the grid determines where a fixed compute budget is spent.

The authors propose Adaptive Reparameterized Time (ART), which treats the speed of the sampling clock as a control. Imagine two clocks: the original diffusion clock and a new clock whose ticks are evenly spaced. By learning how quickly the new clock advances relative to the original one, ART turns its even ticks into an adaptive, generally uneven set of sampling times. More evaluations can thus be placed in parts of the trajectory where the method’s error proxy suggests that smaller steps are useful.

That proxy comes from a leading-order analysis of Euler discretization error. It is a reasoned guide to allocating steps, not a direct measure of final image quality. The authors then introduce ART-RL, a randomized formulation intended to make the control problem tractable in high-dimensional settings. It uses Gaussian policies and an actor–critic learning procedure. Here, the reinforcement-learning framing is a mathematical tool for learning the schedule; it does not mean the sampler is exploring an unknown real-world environment.

What they tested

The study moves from a one-dimensional example with a known score function—which helps isolate timestep effects—to image-generation experiments. These include MNIST with a deliberately small score model and image benchmarks using EDM and EDM2 pipelines. The researchers compare ART-RL with uniform, EDM and DPM schedules at matched numbers of model evaluations. In the controlled comparisons, they keep the score model and sampling setup fixed and change the timestep grid, so the comparison tests scheduling rather than a new model architecture.

The authors report that ART-RL improves the reported quality metric over those schedules across the tested budgets in their experiments. For images, they use Fréchet Inception Distance (FID), or LeNet-FID for MNIST; lower values indicate closer agreement under those metrics, not a complete measure of visual quality. In one reported CIFAR-10 Heun comparison at 35 model evaluations, ART-RL’s FID is 1.82 versus 1.85 for EDM. The gap in that example is small, and the paper also reports three additional matched runs supporting the result. Results in the one-dimensional setting likewise favor ART-RL on the reported Wasserstein distance, while MNIST results favor it on LeNet-FID.

A further set of experiments reuses a schedule learned on CIFAR-10 without retraining. The authors report favorable comparisons after adapting it to different step counts, other image datasets, and an EDM2 ImageNet-512 setup. These tests suggest that the learned allocation may transfer beyond the configuration used to learn it. They do not establish that it will work equally well for every model, dataset or sampling task.

What the result does—and doesn’t—show

The practical attraction is that the learned schedule can be distilled into a fixed list of times. At deployment, the sampler then uses that list without an extra actor-network evaluation at each step. But learning it is not free: the paper reports a one-time CIFAR-10 schedule-training cost of about one to two hours on a Colab T4 GPU. Whether that cost pays off depends on how often the schedule can be reused and what compute resources are available.

There are also important limits. ART’s objective is motivated by a leading-order Euler error surrogate, while the experiments include higher-order solvers such as Heun and RK4. The authors identify alternative objectives better aligned with higher-order integrators as future work. Their analysis is restricted to probability-flow ordinary differential equation sampling; extending the approach to stochastic differential equation samplers remains open. Although ART is formulated with state-dependent control, the experiments find that distilled time-only schedules work well; whether richer state-conditioned schedules would help is unresolved.

The paper describes its preliminary version as an arXiv preprint and says the work was accepted to NeurIPS 2026. This article relies on the supplied paper text and abstract, not an independent replication or assessment of peer review. The reported results should therefore be read as evidence from the authors’ experiments, not as a guarantee of improvement in other settings.