[論文レビュー] ART for Diffusion Sampling: A Reinforcement Learning Approach to Timestep Schedule

Yilie Huang, Wenpin Tang+1 more|arXiv (Cornell University)|Jan 26, 2026

Reinforcement Learning in Robotics被引用数 0

ひとこと要約

Adaptive Reparameterized Time (ART) を導入し、拡散サンプリングの適応タイムステップスケジュールを連続時間強化学習フレームワーク（ART-RL）を通じて学習；再訓練なしでデータセット間のサンプル品質と転移性を向上させる。

ABSTRACT

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.

研究の動機と目的

• 固定時間予算の下でスコアベース拡散モデルの principled な離散化を動機づける。
• Adaptive Reparameterized Time (ART) を提案し、サンプリング時間を再パラメータ化して計算を再分配する。
• 最適な時間スケジュールを学習するためのガウス政策を用いた連続時間強化学習アプローチとして ART-RL を開発する。
• ART と ART-RL の理論的関係を証明し、actor–critic 更新を導出する。
• CIFAR-10 での経験的利得を示し、再訓練なしで AFHQv2、FFHQ、ImageNet へ転送する。

提案手法

• 固定総時間 T の下で再パラメータ化時計 psi(t) と時間ワープ率 theta(t) を用いた拡散サンプリングを定式化する。
• 前方/後方ダイナミクスと、時間割り当てを導く Euler 離散化誤差代替量 Q(x,psi) を定義する。
• 分布をガウス政策 pi^(lambda) によるランダム化制御として導入し、分散を |Q(x,psi)| に結びつけて連続時間 RL の補助問題（ART-RL）を解く。
• ART の最適制御が最適ガウス政策の平均に対応することを示し、actor–critic 学習スキームを可能にする。
• V（批評家/価値ネットワーク）と actor/政策平均、時間制約のラグランジュ乗数を含む実装可能な ART-RL アルゴリズムを開発する。
• 学習された時間スケジュールを、効率的なサンプリングのための決定論的な時間のみのスケジュールとして蒸留するステップを提供する。

Figure 1 : Empirical mean (solid line) and 25–75 percent interquartile range (shaded region) of the learned control $\theta$ across time.

実験結果

リサーチクエスチョン

• RQ1 ART による適応時間再パラメトリゼーションは、固定時間予算下で拡散サンプラの離散化誤差を低減できるか。
• RQ2 ART-RL の補助問題を解くことで最適な ART 時間スケジュールを回復し、データ駆動のスケジュール学習を可能にするか。
• RQ3 ART-RL は均一スケジュールや手作業の EDM スケジュールと比較して、 diverse なデータセットと予算でどうなるか、学習済みスケジュールを再訓練なしで転送できるか。

主な発見

NFE	Uniform (FID)	EDM (FID)	ART-RL (FID)
3	280.29	465.83	152.86
9	213.13	35.54	32.13
13	191.69	6.79	5.44
19	168.87	2.54	2.45
35	118.02	1.85	1.85

• ART-RL は CIFAR-10 での Fréchet Inception Distance (FID) を改善し、再訓練なしで AFHQv2、FFHQ、ImageNet へ転移する。
• 一次元の実験では、ART-RL は全ての時間ステップで Uniform および EDM ベースのスケジュールを上回り、低予算で特に顕著な改善を示す。
• ART-RL は EDM パイプライン下の CIFAR-10 において Uniform および EDM に対して一貫した FID 改善をもたらし、特に小NFEsで顕著。
• 学習された時間スケジュールは補間・外挿されたタイムステップ数へ一般化し、データセット間で効果的に転移する。
• 蒸留された時間のみの ART-RL スケジュールは、各ステップのポリシー計算を必要としないドロップイン置換として機能し得る。

(a) EDM