ART for Diffusion Sampling: Continuous-Time Control and Actor-Critic Learning

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.

• The paper presents ART, formulating timestep allocation as a continuous control problem using variational calculus and error surrogates.
• It embeds ART into a reinforcement learning framework via Gaussian policy randomization, achieving effective actor-critic updates.
• Empirical evaluations demonstrate ART-RL’s superior sample quality, robust transferability, and efficiency across diverse datasets and solvers.

Adaptive Reparameterized Time (ART): A Control-Theoretic and RL Approach to Diffusion Sampling

Background: Diffusion Sampling and the Timestep Allocation Problem

In score-based diffusion generative models, data is transformed through stochastic forward and learned reverse diffusion processes, generating samples that match a target distribution. During inference, the reverse process is discretized over a finite set of timesteps. The selection of this timestep grid fundamentally impacts both sample quality and computational efficiency, as each step incurs a (potentially expensive) score model evaluation. Existing methods typically apply uniform grids or rely on fixed heuristic schedules (e.g., EDM, DPM), which are not universally optimal and lack principled justification, especially when model or data properties vary.

Figure 1: Illustration of the forward noising process and the corresponding backward generative process in a score-based diffusion model.

ART: Time Reparameterization as Deterministic Control

The paper introduces Adaptive Reparameterized Time (ART), treating timestep allocation as a continuous-time optimal control problem. By introducing a time-reparameterization $\psi: [0,T] \to [0,T]$, ART models the progression in diffusion time as a controllable rate $\theta(t) = \dot{\psi}(t)$. A uniform grid on this "learned clock" maps non-uniformly to diffusion time, adaptively concentrating function evaluations where necessary according to trajectory stiffness.

Error analysis based on leading-order Euler discretization motivates the use of state- and time-dependent error surrogates to guide optimal allocation. Crucially, the ART formalism does not fix $\psi$ to specific functional forms but adopts a variational calculus perspective, optimizing time allocation subject to a global time budget constraint.

Figure 2: ART as a time change between two clocks, with $\theta(t) = \dot{\psi}(t)$ governing non-uniform allocation.

ART-RL: Gaussian Policy Randomization and Continuous-Time RL

Directly solving the high-dimensional deterministic ART control problem is intractable for generative modeling regimes (e.g., high-dimensional images). To address this, the authors introduce ART-RL, an auxiliary problem wherein $\theta$ is sampled from Gaussian policies whose variance is inversely proportional to local stiffness ($|Q|$), effectively embedding ART into the recently developed continuous-time reinforcement learning framework.

A rigorous equivalence is established: the deterministic ART control is recovered as the mean of the optimal ART-RL Gaussian policy. This alignment is characterized through the HJB equations of the original and randomization-augmented control problems. The mean and value functions are updated via a derived continuous-time actor-critic scheme, exploiting martingale identities to yield trajectory-level, implementable moment conditions for policy evaluation and improvement.

Algorithmic Methodology

ART-RL employs coupled neural network parameterizations for the actor (policy mean) and the critic (value function). Updates are derived from stochastic approximation of the RL-induced moment identities, with Lagrange multipliers enforcing global time budget constraints. The practical implementation employs time-discretized rollouts of the controlled diffusion process, with policy and value function updates carried out per batch, and an explicit distillation step to yield deterministic (time-only) schedules as a plug-in for inference-time sampling.

Empirical Evaluation and Analysis

Extensive experiments demonstrate the efficacy and universality of ART-RL. Empirical tests cover known-score synthetic data, small-model settings (MNIST), and state-of-the-art high-dimensional image generation benchmarks (CIFAR-10, FFHQ, AFHQv2, ImageNet-64/512) across multiple solvers (Euler, Heun, RK4) and pipelines (EDM, EDM2). All baselines (Uniform, EDM, DPM) are matched in function evaluation count and kept identical in all components except for the timestep schedule.

ART-RL consistently achieves superior sample quality (lowest FID and Wasserstein-2 errors) at matched evaluation budgets, outperforming hand-crafted schedules, including at the largest budgets where existing heuristics are strongest. For example, on CIFAR-10 in the EDM pipeline, ART-RL achieves FID = 1.82 at 35 NFE, compared to 1.85 for EDM and 1.89 for DPM.

Figure 3: Visual overview of ART-RL across key experiments, showing improved sample quality at matched evaluation budgets.

The empirical control $\theta$ quickly collapses to a smooth, nearly deterministic function of time, with negligible state dependence in 1D and realistic settings. This justifies a schedule distillation step: post-training, one can discard the policy NN and reuse a precomputed deterministic schedule, reducing runtime overhead.

Figure 4: Empirical mean (solid) and interquartile range of $\theta$ under ART-RL schedules, highlighting deterministic collapse.

Figure 5: Empirical mean and 99% confidence interval for $\theta$ on CIFAR-10, supporting distillation to time-only schedules.

Transfer, Generalization, and Plug-and-Play Schedules

A key claim, strongly supported by experiments, is the robust transfer and generalization ability of ART-RL. A schedule learned on one dataset (e.g., CIFAR-10) at a chosen step budget remains competitive, outperforming baselines when reused—without retraining—across:

• Different step counts (using interpolation/extrapolation of the learned $\psi$)
• Other datasets (AFHQv2, FFHQ, ImageNet-64) under the same EDM pipeline
• New architectures, latent space representations, and higher resolutions (EDM2 on ImageNet-512)

This generality ensures that ART-RL's (modest) one-off training cost is amortized over many practical deployments: once distilled, an ART-RL schedule is fully analogous in deployment to existing hand-crafted grids.

Theoretical Implications

The paper advances the perspective that timestep allocation in generative diffusion models is best understood as a control problem, driven by error surrogates derived from numerical analysis of SDE/ODE integration. Embedding this into continuous-time RL enables tractable solution in high-dimensional spaces, leveraging policy gradient and actor-critic machinery to achieve systematic, data-driven optimization of sampling clocks.

The equivalence between the deterministic ART control and the mean of the optimal ART-RL policy provides a rigorous foundation for policy randomization in this context, distinct from standard RL exploration: randomization here is purely analytic, enabling functional optimization.

Limits and Future Directions

The current formulation is built on probability flow ODE sampling and error surrogates motivated by the Euler method, which may not directly correspond to the numerical characteristics of higher-order solvers. Extending ART principles to generic SDE samplers and integrating error surrogates tailored to advanced integrators (e.g., Heun, Runge-Kutta) are promising directions. Moreover, although distillation to time-only schedules is highly effective empirically, theoretical exploration of state-conditioned or adaptive schedules could uncover additional gains, especially in non-stationary or data-dependent regimes.

Conclusion

This work provides the first theoretically principled, control-theoretic, and RL-amenable framework for timestep allocation in generative diffusion sampling (2607.02137). The ART and ART-RL methodology systematically optimizes sampling clocks, achieving state-of-the-art performance and exhibiting broad generalizability across diverse models, solvers, and datasets. By aligning the error control objective with continuous-time stochastic control and modern RL techniques, the approach represents a significant step toward automated, optimal design of diffusion sampling protocols in generative modeling.