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

技术摘要：用于扩散采样的 ART

问题陈述
基于分数的扩散模型通过积分学习到的反向时间动态来生成样本。在此过程中，一个关键的数值选择是将时间范围离散化为时间步网格。标准方法通常采用均匀网格或手工设计的调度方案（例如 EDM、DPM-Solver）。作者认为，鉴于固定的函数评估预算（NFE），这些固定调度方案通常并非最优。均匀网格即使在与轨迹显著变化的反向动态不一致时，也会均匀分配计算资源，而手工设计的网格可能无法适应特定的数据分布或分数函数的复杂性。目标是系统地重新分配采样轨迹上的计算资源，以最小化总体欧拉离散化误差。

方法论
本文提出了自适应重参数化时间（ART），这是一种用于最优时间步调度的控制理论框架。

1. 控制公式化（ART）：

   • ART 不直接选择网格点，而是将扩散过程的“时钟速度”视为控制变量。
   • 它引入一个重参数化时间变量 $t$，通过连续映射 $\psi(t)$ 映射到原始扩散时间 $\tau$。
   • 控制量 $\theta(t) = \dot{\psi}(t)$ 表示扩散时间的瞬时变化率。
   • 目标是在保持总时间预算不变（$\int_0^T \theta(t) dt = T$）的约束下，选择 $\theta(t)$ 以最小化总体局部欧拉离散化误差。
   • 局部误差代理被推导为与 $|Q(x, \psi)| \theta(t)^2$ 成正比，其中 $Q$ 取决于分数函数及其导数。

2. 强化学习求解（ART-RL）：

   • 由于高维状态空间以及哈密顿 - 雅可比 - 贝尔曼（HJB）方程缺乏闭式解，确定性的 ART 问题难以直接求解。
   • 作者提出了ART-RL，一种连续时间强化学习（CTRL）方法。
   • 他们引入了一个随机化辅助问题，其中控制量 $\theta$ 从高斯策略 $\pi(\cdot | t, x, \psi) = \mathcal{N}(\mu, \lambda/|Q|)$ 中采样。
   • 该公式允许使用演员 - 评论家（actor-critic）算法（特别是借鉴 Jia 和 Zhou [14] 的框架），通过随机近似来学习策略均值 $\mu$ 和价值函数。

3. 理论桥梁：

   • 核心理论贡献在于建立了确定性 ART 最优解与 ART-RL 辅助问题中的最优高斯策略之间的双向桥梁。
   • 充分性： 确定性最优控制 $\mu^*$ 提升为一个最优高斯策略，其均值恰好为 $\mu^*$。
   • 必要性： 辅助问题的任何最优高斯策略，其均值必须能够恢复确定性 ART 控制量 $\mu^*$。
   • 这一结果证明了将演员参数化为高斯均值并在训练后确定性部署的合理性，提供了严格的保证，即学习到的调度方案是原始控制问题的真正最优解。

4. 蒸馏：

   • 在离线训练后，随机策略被蒸馏为一个确定性的、仅依赖时间的网格。这消除了在推理过程中进行每步演员网络评估或方差计算的需求，确保与标准调度方案相比零额外计算成本。

主要贡献

• 公式化： 提出了 ART，这是首个通过自适应时间重参数化解决扩散时间步调度的控制理论框架。
• 理论： 建立了确定性最优控制问题与带有高斯策略的连续时间演员 - 评论家学习过程之间的双向等价性。这将启发式的强化学习应用转化为寻找确定性最优解的原则性途径。
• 算法： 开发了 ART-RL，这是一种数据驱动地学习自适应时间调度的演员 - 评论家算法。
• 泛化性： 证明了学习到的调度方案可以在不同数据集和时间步数之间进行蒸馏和迁移，而无需重新训练。

实验结果
作者在官方 EDM 流程中评估了 ART-RL，涵盖 CIFAR-10、AFHQv2、FFHQ 和 ImageNet 数据集。

• CIFAR-10： 在广泛的 NFE 预算范围内，ART-RL 始终优于均匀调度、DPM-Solver log-SNR 和标准 EDM 调度。
  • 在低预算下（例如 NFE=3），ART-RL 能生成可识别的物体，而基线方法仍接近噪声。
  • 在 NFE=35 时，ART-RL 实现了 1.82 的 FID，优于 EDM（1.85）和 DPM-logSNR（1.89）。
  • 通过多次随机种子确认了统计显著性，ART-RL 的最差运行结果甚至优于 EDM 的最佳结果。
• 泛化性：
  • 插值/外推： 在 CIFAR-10 上于 $K=18$ 步学习到的调度方案，在插值或外推到其他步数（$K \in \{4, 9, 12, 15, 20\}$）时，仍保持其性能优势。
  • 跨数据集迁移： 将 CIFAR-10 蒸馏后的调度方案应用于 AFHQv2、FFHQ 和 ImageNet，无需重新训练。它始终在 FID 上优于 EDM 和 DPM-logSNR，特别是在低和中等 NFE 下。
• 一维分析： 在具有解析分数函数的设置中，与手工设计的调度方案相比，ART-RL 显著降低了 Wasserstein 距离（$W_2$），证实了改进源于更好的时间分配而非分数估计误差。

意义与主张
本文声称是首个通过自适应时间重参数化将扩散时间步调度公式化为控制问题的工作，也是首个在确定性最优解与演员 - 评论家学习过程之间提供严格双向桥梁的工作。

作者强调，他们的方法是数据驱动且摊销的：

• 数据驱动： 调度方案是从目标函数和数据动态中学习得到的，而非手工设计。
• 摊销迁移： 离线训练成本（CIFAR-10 在 Colab T4 上约 1–2 小时）是一次性支出。生成的蒸馏确定性网格可在不同数据集和步数间复用，无额外推理成本，其运行时间与均匀或 EDM 等标准调度方案相当。

该工作将自适应时间重参数化定位为生成式扩散模型中系统性调度设计的一个新兴但充满希望的方向，为未来研究提供了控制理论的起点。