用强化学习优化扩散模型采样时间步,提升生成质量。
ART for Diffusion Sampling: A Reinforcement Learning Approach to Timestep Schedule
- 将时间步调度建模为连续强化学习问题,动态分配计算资源。
- 在CIFAR-10上显著降低FID,且可跨数据集无成本迁移。
- 理论证明最优策略与确定性调度等价,方法更可信可靠。
我们研究基于得分的扩散模型的时间离散化,以在有限时间步数下从学习到的逆时序动态中生成样本。均匀或手工设计的时间网格在给定步数预算下可能表现不佳。本文提出自适应重参数化时间(ART),通过控制重参数化时间变量的时钟速度,在保持终态时间不变的前提下重新分配采样轨迹上的计算量,目标是最小化整体欧拉离散误差。我们推导出随机版本ART-RL,将ART重述为具有高斯策略的连续时间强化学习问题,并证明二者之间的双向关联:确定性ART最优解可映射为最优高斯策略,反之任意最优高斯策略的均值必须恢复ART控制。这一桥梁使连续时间演员-评论家学习成为通向确定性时间步最优解的合理路径,而非启发式尝试。在官方EDM流程中,ART-RL在多种预算下改进了CIFAR-10的FID;经一次离线训练后,提炼出的确定性调度可无需再训练直接迁移至AFHQv2、FFHQ和ImageNet,推理成本无增加。
原文摘要 · Abstract (English)
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), which controls the clock speed of a reparameterized time variable to redistribute computation along the sampling trajectory while preserving the terminal time, with the objective of minimizing the aggregate Euler discretization error. We derive a randomized companion ART-RL that recasts ART as a continuous-time reinforcement learning problem with Gaussian policies, and prove a two-directional bridge between the two: the deterministic ART optimum lifts to an optimal Gaussian policy, and conversely any optimal Gaussian policy must recover the ART control through its mean. This bridge turns continuous-time actor--critic learning into a principled, rather than heuristic, route to the deterministic timestep optimum. Within the official EDM pipeline, ART-RL improves FID on CIFAR--10 across a wide range of budgets; after one-time offline training, the distilled deterministic schedule transfers without retraining to AFHQv2, FFHQ, and ImageNet at no extra inference cost.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。