技術要約：拡散モデルのテスト時スケーリングのためのLookahead Sample Reward Guidance (LiDAR)

問題提起

拡散モデルは、画像、ビデオ、言語において驚異的な生成性能を達成していますが、生成されたサンプルが人間の意図に完全に沿わないことが頻繁にあります。人間の好みを捉える報酬関数は存在しますが、それらをサンプリングプロセスに統合することには大きな課題があります。テスト時スケーリングの核心的な目的は、報酬によって傾斜付けされた分布 $p_\theta^r(x_0|c) \propto p_\theta(x_0|c) \exp(\lambda \cdot r(x_0, c))$ からサンプリングすることです。

これには、逆拡散プロセスをガイドするために、中間粒子 $x_t$ に対する期待将来報酬 (Expected Future Reward: EFR)、$r_t^\lambda(x_t, c)$ を計算する必要があります。既存の手法には以下の決定的な限界があります：

1. バックワード・ロールアウト (Backward Rollout): EFRを推定するために $x_t$ から $x_0$ へ直接サンプリングすることは、各タイムステップでの反復的なデノイジングを必要とするため、計算コストが極めて高くなります。
2. Tweedieに基づく近似: Gradient GuidanceやSequential Monte Carlo (SMC) といった手法は、EFRを推定するために一次テイラー近似（Tweedieの公式）に依存しています。これらは、特に報酬信号の強度（$\lambda$）が高い場合にバイアスを導入し、またニューラルネットワークのバックプロパゲーションを必要とする（Gradient Guidance）、あるいは高次元空間でサンプル崩壊（sample collapse）に陥る（SMC）ことがあります。

手法：LiDAR サンプリング

著者らは、中間粒子 $x_t$ とEFRの間のニューラルな依存関係を切り離し、閉形式（closed-form）かつ微分不要なガイダンスを可能にするフレームワークである LiDAR (Lookahead Sample Reward Guidance) を提案しています。

1. 期待将来報酬の再定式化

論文では、EFRの新しい定式化（定理3.1）を導出し、EFRの期待値を、ニューラルネットワークを介して現在の粒子 $x_t$ に条件付けるのではなく、学習済み拡散モデルからの周辺サンプル (marginal samples) と前方摂動カーネルのみを用いて表現します：
$$r_t^\lambda(x_t, c) = \log \mathbb{E}_{p_\theta(x_0|c)} \left[ \frac{p(x_t|x_0)}{\mathbb{E}_{p_\theta(x_0|c)}[p(x_t|x_0)]} \exp(\lambda \cdot r(x_0, c)) \right]$$
この「フォワード・ロールアウト」アプローチにより、任意の $x_t$ に対して、ガイダンスステップ中にスコアネットワーク $s_\theta$ や報酬モデルのバックプロパゲーションを行うことなく、事前生成された周辺サンプルの集合 $\{\hat{x}_0^i\}$ を用いてEFRを計算できます。

2. 効率的なルックアヘッド・サンプリング

すべてのプロンプトに対して周辺サンプルを生成するコストを軽減するため、著者らは効率的なルックアヘッド・サンプリング (lookahead sampling) 戦略を導入しています。低速な完全拡散プロセスを使用する代わりに、高速なジェネレータ $q(x_0|c)$（例：数ステップのODEソルバーであるDPM-Solverや蒸留モデル）を用いて、一連の $n$ 個のルックアヘッド・サンプル $\{\hat{x}_0^i\}_{i=1}^n$ を生成します。これらのサンプルには報酬値 $r(\hat{x}_0^i, c)$ が付与されます。

3. 閉形式の微分不要なガイダンス

有限のルックアヘッド・サンプルを用いて、経験的ルックアヘッド報酬の勾配の閉形式の式（定理3.3）を導出します。このガイダンス項は、現在の粒子 $x_t$ を高報酬のルックアヘッド・サンプルへと導き、低報酬のサンプルからは遠ざけます：
$$\nabla_{x_t} \hat{r}_t^\lambda(x_t, c) = \sum_{i=1}^n (w_i^r - w_i) \frac{\hat{x}_0^i}{\sigma_t^2}$$
ここで、$w_i^r$ と $w_i$ は、報酬の値と $x_t$ とルックアヘッド・サンプル $\hat{x}_0^i$ の距離に基づくソフトマックス重みです。この定式化は、ニューラル微分を必要としません。

4. 理論的なスケーリング則

論文では、本手法のスケーリング則を確立しています：

• ルックアヘッド精度 ($\delta$): ターゲット分布とLiDAR生成分布の全変動距離は、ルックアヘッド・ソルバーの離散化ステップ数 $\delta$ に対して $O(1/\sqrt{\delta})$ で減少します。
• サンプルサイズ ($n$): 経験的報酬推定の分散は $O(1/\sqrt{n})$ で減少します。これは、ルックアヘッド・サンプルの数を増やすことで近似精度が向上することを意味します。

主な結果

実験は、SD v1.5、SDXL、FLUX (flow matching)、およびUDLM (discrete diffusion) を用い、ImageReward、CLIP、HPS、GenEvalの指標を用いて行われました。

• 性能と効率性: SDXLにおいて、LiDARは最新の勾配ガイダンス手法（DATE）に匹敵するGenEval性能を達成しつつ、推論時間を9.5倍高速化しました。
• メモリ効率: バックプロパゲーションのために約3倍のメモリを必要とする勾配ガイダンス手法とは異なり、LiDARのメモリ使用量はバニラなサンプリングとほぼ同等です。これは、ターゲット・サンプリング段階（アルゴリズム2）において勾配計算を必要としないためです。
• 報酬ハッキングへの堅牢性: 勾配ガイダンス手法は、ImageRewardを最適化しようとするとCLIPやHPSなどの指標が低下する（報酬ハッキング）ことが多いですが、LiDARはすべての指標を同時に改善します。これは、ルックアヘッド・サンプル内におけるImageRewardと他の指標との相関により、サンプラーが自然にバランスの取れた方向へと導かれるためです。
• 直交的な改善: LiDARは、DPO (Direct Preference Optimization) で既に微調整されたモデルに適用した場合でも追加の性能向上をもたらし、学習ベースのアライメントに対する直交性を示しています。
• スケーラビリティ: 本手法は、ルックアヘッド・サンプルの数 ($n$) およびソルバーのステップ数 ($\delta$) と共に効果的にスケールし、SMCやBest-of-Nといった他の手法とも組み合わせることが可能です。

意義と主張

本論文は、LiDARが既存のテスト時スケーリング手法における効率性と報酬ハッキングの限界を体系的に解決することを主張しています。EFRの計算をフォワード・ロールアウトと周辺サンプルに依存するように再定式化することで、著者らは、計算効率が高く、理論的根拠に基づいた微分不要のガイダンスを可能にしました。

著者らは、LiDARが既存のベースラインと比較して優れた効率と性能のトレードオフを実現していることを強調しています。具体的には、ニューラルなバックプロパゲーションによる計算オーバーヘッドや、SMCによるサンプル崩壊の問題を回避しつつ、テイラー近似に基づく手法よりも正確なEFR推定を提供します。この研究は、報酬信号を中間粒子のニューラルな依存関係から切り離すことが、スケーラブルで人間と整合した生成モデリングへの有効な道であることを示唆しています。