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

📄 arXiv: 2601.18681v1 📥 PDF

作者: Yilie Huang, Wenpin Tang, Xunyu Zhou

分类: cs.LG, cs.AI, eess.SY, math.OC

发布日期: 2026-01-26

备注: 17 pages, 7 figures

💡 一句话要点

提出ART-RL，通过强化学习优化扩散模型采样的时间步长调度，提升生成质量。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 扩散模型 时间步长调度 强化学习 自适应重参数化时间 连续时间控制

📋 核心要点

  1. 扩散模型采样过程中，固定或手工设计的时间步长调度并非最优，限制了生成质量。
  2. 提出自适应重参数化时间（ART）框架，通过强化学习动态调整时间步长，优化采样轨迹。
  3. 实验表明，ART-RL在多个数据集上提升了生成质量，且具有良好的迁移能力。

📝 摘要（中文）

本文研究了基于分数的扩散模型的时间离散化问题，旨在通过学习到的逆时间动态在有限网格上生成样本。均匀或手工设计的网格在给定的时间步数预算下可能并非最优。为此，我们引入了自适应重参数化时间（ART），它控制重参数化时间变量的时钟速度，从而在保持终端时间不变的情况下，实现时间变化和采样轨迹上的不均匀时间步长。目标是最小化离散欧拉方案产生的累积误差。我们推导出一个随机控制伴随，即ART-RL，并将时间变化建模为一个具有高斯策略的连续时间强化学习（RL）问题。然后，我们证明了解ART-RL可以恢复最优的ART调度，进而实现actor-critic更新，以数据驱动的方式学习后者。实验结果表明，基于官方EDM流程，ART-RL在CIFAR-10上改进了Fréchet Inception Distance，且在各种预算下均优于基线方法，并可迁移到AFHQv2、FFHQ和ImageNet，无需重新训练。

🔬 方法详解

问题定义：扩散模型通过逆扩散过程从噪声中生成数据，该过程需要离散化为一系列时间步。均匀或手工设计的时间步长分配方式可能并非最优，尤其是在计算资源有限的情况下，如何找到最优的时间步长调度策略，以在给定的时间步数预算下最大化生成质量，是本文要解决的核心问题。现有方法缺乏自适应性，无法根据数据分布动态调整时间步长。

核心思路：本文的核心思路是将时间步长调度问题建模为一个连续时间强化学习问题。通过引入自适应重参数化时间（ART），将时间变量进行重参数化，并控制其“时钟速度”，从而实现对时间步长的动态调整。利用强化学习算法学习最优的“时钟速度”控制策略，使得采样轨迹上的时间步长分配更加合理，从而提高生成质量。

技术框架：整体框架包含以下几个主要部分：1) 自适应重参数化时间（ART）：定义了一个重参数化的时间变量，其“时钟速度”由一个控制策略决定。2) 随机控制伴随（ART-RL）：将时间步长调度问题转化为一个连续时间强化学习问题，推导出一个随机控制伴随，用于指导策略学习。3) Actor-Critic算法：使用Actor-Critic算法学习最优的控制策略，Actor网络输出“时钟速度”，Critic网络评估当前策略的价值。

关键创新：最重要的创新点在于将时间步长调度问题建模为一个连续时间强化学习问题，并提出了ART-RL算法。与传统方法相比，ART-RL能够根据数据分布动态调整时间步长，从而获得更优的采样轨迹。此外，该方法还具有良好的可迁移性，无需针对每个数据集进行重新训练。

关键设计：在ART-RL中，Actor网络使用高斯策略，输出“时钟速度”的均值和方差。Critic网络用于评估当前策略的价值。损失函数包括两部分：一部分是基于随机控制伴随推导出的损失函数，用于指导策略学习；另一部分是正则化项，用于防止策略过于激进。具体实现细节和超参数设置在论文中有详细描述。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，ART-RL在CIFAR-10数据集上显著提升了生成质量，Fréchet Inception Distance (FID) 指标优于基线方法。更重要的是，该方法具有良好的迁移性，可以直接应用于AFHQv2、FFHQ和ImageNet等数据集，无需重新训练，这大大降低了应用成本。

🎯 应用场景

该研究成果可广泛应用于图像生成、音频生成、视频生成等领域。通过优化时间步长调度，可以提高生成模型的效率和生成质量，降低计算成本。此外，该方法还可应用于其他需要序列决策的任务，例如机器人控制、金融交易等。

📄 摘要（原文）

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) that controls the clock speed of a reparameterized time variable, leading to a time change and uneven timesteps along the sampling trajectory while preserving the terminal time. The objective is to minimize the aggregate error arising from the discretized Euler scheme. We derive a randomized control companion, ART-RL, and formulate time change as a continuous-time reinforcement learning (RL) problem with Gaussian policies. We then prove that solving ART-RL recovers the optimal ART schedule, which in turn enables practical actor--critic updates to learn the latter in a data-driven way. 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.