ART for Diffusion Sampling: Continuous-Time Control and Actor-Critic Learning
- URL: http://arxiv.org/abs/2607.02137v2
- Date: Fri, 03 Jul 2026 02:14:28 GMT
- Title: ART for Diffusion Sampling: Continuous-Time Control and Actor-Critic Learning
- Abstract summary: We study timestep allocation for score-based diffusion sampling, where a learned reverse-time dynamics is discretized on a finite grid.<n>We propose a continuous-time control formulation that learns a time change by treating the speed of the sampling clock as the control.<n>We introduce ART-RL, an auxiliary randomized formulation with Gaussian policies that turns schedule learning into a continuous-time reinforcement learning problem.
- Score: 8.67719968902777
- License: http://creativecommons.org/licenses/by/4.0/
- 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 papers
- Adversarial Learning of Classifier-Free Guidance Schedules [74.70485037604884]
We learn the guidance schedule as a function of diffusion time, conditioning and the current noisy sample.<n>We frame this as a density ratio estimation problem.<n>A lightweight generator network predicts the optimal, state-dependent guidance scale.
arXiv Detail & Related papers (2026-08-14T07:30:36Z) - LC-GRPO: Bridging Train-Inference Gap for Flow-Based GRPO with Langevin Correction [59.32834240526981]
Flow-based generative models are typically sampled by solving a deterministic ordinary differential equation (ODE)<n>Existing GRPO methods for flow models replace the Euler-time ODE with aotropic differential equation (SDE) during training.<n>We introduce LC-GRPO, a flow-based GRPO framework with Langevin correction targeting the marginal distribution at the resulting timestep.
arXiv Detail & Related papers (2026-08-06T04:47:57Z) - From Atoms to Entropy: Optimal Noise Allocation for Diffusion Training in the Convex Regime [52.459513910242485]
We develop a general statistical framework for studying optimal noise-level allocation in diffusion training.<n>We show that the optimized training schedule admits an atomic minimizer, concentrated on finitely many noise levels.<n>We then evaluate the entropic schedule in larger-scale experiments, where full schedule optimization is currently intractable.
arXiv Detail & Related papers (2026-07-10T20:59:06Z) - Precise: SDE-Consistent Stochastic Sampling for RL Post-Training of Flow-Matching Models [56.67321805551389]
Reinforcement learning (RL) has become an effective way to improve prompt alignment and perceptual quality in diffusion and flow-matching generators.<n>The sampler, controlling the exploration behavior and denoising dynamics, is thus part of the policy.<n>We propose a new sampler that balances effective exploration with stability.
arXiv Detail & Related papers (2026-05-22T11:37:22Z) - RAVEN: Real-time Autoregressive Video Extrapolation with Consistency-model GRPO [53.38929612273108]
Causal autoregressive video diffusion models support real-time streaming generation by extrapolating future chunks from previously generated content.<n>We introduce the Real-time Autoregressive Video Extrapolation Network (RAVEN), a training-time test framework that repacks each self rollout into an interleaved sequence of clean historical endpoints.
arXiv Detail & Related papers (2026-05-14T17:59:30Z) - Inference-Time Attribute Distribution Alignment for Unconditional Diffusion [54.28942858023809]
We formalize the inference-time attribute distributional alignment problem for pretrained unconditional diffusion models.<n>To address this, we cast inference-time attribute distributional alignment as an optimal control problem over the reverse diffusion process.<n>We solve for the perturbations using an optimal-control-based algorithm to optimize a differentiable distribution-matching objective.
arXiv Detail & Related papers (2026-05-08T09:02:47Z) - ART for Diffusion Sampling: A Reinforcement Learning Approach to Timestep Schedule [8.67719968902777]
We consider time discretization for score-based diffusion models.<n>We derive a randomized control companion, ART-RL, and formulate time change as a continuous-time reinforcement learning problem.<n> Empirically, based on the official EDM pipeline, ART-RL improves Fréchet Inception Distance on CIFAR-10 over a wide range of budgets and transfers to AFHQv2, FFHQ, and ImageNet without the need of retraining.
arXiv Detail & Related papers (2026-01-26T16:56:40Z) - Inference-Time Scaling of Diffusion Language Models with Particle Gibbs Sampling [70.8832906871441]
We study how to steer generation toward desired rewards without retraining the models.<n>Prior methods typically resample or filter within a single denoising trajectory, optimizing rewards step-by-step without trajectory-level refinement.<n>We introduce particle Gibbs sampling for diffusion language models (PG-DLM), a novel inference-time algorithm enabling trajectory-level refinement while preserving generation perplexity.
arXiv Detail & Related papers (2025-07-11T08:00:47Z) - Data-Driven Exploration for a Class of Continuous-Time Indefinite Linear--Quadratic Reinforcement Learning Problems [6.859965454961918]
We study reinforcement learning for continuous-time linear-quadratic (LQ) control problems.<n>We propose a model-free, data-driven exploration mechanism that adaptively adjusts entropy regularization by the critic.<n>Our method achieves a sublinear regret bound that matches the best-known model-free results for this class of LQ problems.
arXiv Detail & Related papers (2025-07-01T01:09:06Z) - Score-Optimal Diffusion Schedules [29.062842062257918]
An appropriate discretisation schedule is crucial to obtain high quality samples.<n>This paper presents a novel algorithm for adaptively selecting an optimal discretisation schedule.<n>We find that our learned schedule recovers performant schedules previously only discovered through manual search.
arXiv Detail & Related papers (2024-12-10T19:26:51Z) - Truncating Trajectories in Monte Carlo Policy Evaluation: an Adaptive Approach [51.76826149868971]
Policy evaluation via Monte Carlo simulation is at the core of many MC Reinforcement Learning (RL) algorithms.
We propose as a quality index a surrogate of the mean squared error of a return estimator that uses trajectories of different lengths.
We present an adaptive algorithm called Robust and Iterative Data collection strategy Optimization (RIDO)
arXiv Detail & Related papers (2024-10-17T11:47:56Z) - When to Sense and Control? A Time-adaptive Approach for Continuous-Time RL [37.58940726230092]
Reinforcement learning (RL) excels in optimizing policies for discrete-time Markov decision processes (MDP)
We formalize an RL framework, Time-adaptive Control & Sensing (TaCoS), that tackles this challenge.
We demonstrate that state-of-the-art RL algorithms trained on TaCoS drastically reduce the interaction amount over their discrete-time counterpart.
arXiv Detail & Related papers (2024-06-03T09:57:18Z) - Time-series Generation by Contrastive Imitation [87.51882102248395]
We study a generative framework that seeks to combine the strengths of both: Motivated by a moment-matching objective to mitigate compounding error, we optimize a local (but forward-looking) transition policy.
At inference, the learned policy serves as the generator for iterative sampling, and the learned energy serves as a trajectory-level measure for evaluating sample quality.
arXiv Detail & Related papers (2023-11-02T16:45:25Z)
This list is automatically generated from the titles and abstracts of the papers in this site.
This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.