Gradient-Free Noise Optimization for Reward Alignment in Generative Models

技術的サマリー：ZeNO（ゼロ次ノイズ最適化）

問題定義
生成モデル（拡散モデルやフローモデルなど）に対する既存の報酬アライメント手法は、通常、多段階の確率的軌跡の最適化に依存しています。これらの設定に対しては効果的ですが、同様の軌跡構造を自然に露出しない決定論的または単一ステップの生成器（一貫性モデル、蒸留拡散モデルなど）へは、そのようなアプローチを拡張することが困難です。自然な代替案として、入力ノイズを直接最適化することが考えられます。しかし、標準的な勾配ベースのノイズ最適化には、生成器と報酬関数の両方を通じた逆伝播が必要です。これにより、微分可能なパイプラインへの適用性が制限され、勾配が意味のある意味論的改善ではなく、偽の局所的な感度を反映している場合に信頼性が低下する可能性があります。さらに、多くの実用的な報酬関数（タンパク質の設計性、構成的なテキスト - 画像アライメントなど）は微分不可能またはブラックボックスであり、勾配ベースの手法を実行不可能にします。

手法：ZeNO
著者らは、ノイズ最適化を経路積分制御（PIC）問題として定式化する、勾配フリーのフレームワークであるZeNO（ゼロ次ノイズ最適化）を提案します。その核心となる考え方は、入力ノイズの進化を制御可能な確率過程として扱い、ゼロ次報酬評価のみを使用してそれを最適化することです。

1. 経路積分制御の定式化:
   この手法は、ノイズ最適化を制御された確率微分方程式（SDE）としてモデル化します。目的は、制御努力を正則化しつつ、ノイズ軌跡を高い報酬を持つ終端状態をもたらす領域へ誘導することです。最適な制御 $u^*(z, t)$ は、経路積分形式で導出されます：
   $$u^*(z, t) = \frac{\mathbb{E}_{z:t \to T} [\exp(r(G_\theta(z_T))/\lambda) d w_t]}{\mathbb{E}_{z:t \to T} [\exp(r(G_\theta(z_T))/\lambda)]}$$
   重要なのは、この推定量が終端コスト評価（報酬）のみに依存し、生成器 $G_\theta$ や報酬関数 $r$ を通じた微分を必要としない点です。

2. 参照ダイナミクス（オルンシュタイン・ウーレンベック）:
   更新されたノイズが事前学習された生成器と互換性を持つことを保証するため、ZeNO は参照ダイナミクスとしてオルンシュタイン・ウーレンベック（OU）過程を採用します。分散発散型 SDE とは異なり、OU 過程はガウス定常分布 $N(0, I)$ を保持し、これはほとんどの事前学習済み生成器の事前分布と一致します。この選択により、ノイズ更新は、報酬傾斜分布 $p^*(z) \propto p(z) \exp(r(G_\theta(z))/\lambda)$ を標的とするランジュバンダイナミクスと暗黙的に結びつきます。

3. 線形化推定量と分散低減:
   標準的な経路積分推定量は、指数重み付けにより高い分散に悩まされます。著者らは、平均報酬 $\bar{r}$ を基準として減算し、小さな報酬分散を仮定して指数項を線形化することで、線形化近似（命題 1）を導入します：
   $$u^*(z_t, t) \approx \mathbb{E}[(r(G_\theta(z_T)) - \bar{r}) d w_t] / \lambda$$
   この推定量は、自然進化戦略で用いられるスコア関数推定量と一致し、経験的に指数型の変種よりも優れた性能を示します。

4. 単一ステップの再帰的ホライズン:
   長ホライズンのロールアウトに伴う計算コストを回避するため、ZeNO は単一ステップの再帰的ホライズン近似（$H=1$）を採用します。各反復において、$N$ 個の摂動ノイズ候補が参照ダイナミクスからサンプリングされ、報酬モデルによって評価されます。その後、線形化された制御推定量を用いて現在のノイズ状態が更新されます。これにより、生成器を変更することなく、粒子数（$N$）または反復回数（$M$）を増やすことで推論時のスケーリングが可能になります。

主な貢献

• 勾配フリーのアライメント: ZeNO は、生成器または報酬モデルを通じた逆伝播を必要とせずに、決定論的および単一ステップの生成器に対する報酬アライメントを可能にします。
• ブラックボックスとの互換性: この手法は、微分不可能およびブラックボックスの報酬関数に適用可能です。これは、既存の勾配ベースのノイズ最適化手法（ORIGEN、ReNO など）が失敗する設定です。
• 理論的関連性: このフレームワークは、ゼロ次ノイズ最適化と、報酬傾斜分布を標的とするランジュバンサンプリングとの間の理論的リンクを確立し、更新規則に対する原理的な正当性を提供します。
• 推論時のスケーリング: ZeNO は、計算予算（$N$および$M$を介して）と報酬性能のトレードオフを可能にする柔軟なメカニズムを導入します。

実験結果
著者らは、多様な生成器（SDXL-Turbo、DMD2、LCM、DFGAN）および報酬関数（審美性スコア、PickScore、GenEval、タンパク質設計性）において ZeNO を評価しました。

• 画像生成: ZeNO は、テストされたすべての生成器および報酬タイプにおいて、一貫して報酬スコアを向上させました。勾配ベースのベースライン（ORIGEN、$\Psi$-Sampler）および Best-of-N サンプリングと比較して、競争力のあるまたは優れた性能を達成しました。特に LCM 生成器において、勾配ベースの手法は制御されていないベースライン以下に性能が低下しましたが、ZeNO は改善を維持しました。これは、局所的な勾配感度に対するゼロ次手法の堅牢性を浮き彫りにしています。
• 微分不可能な報酬: 報酬が微分不可能な GenEval タスク（物体の計数と位置決め）において、ZeNO は計数タスクにおいて Best-of-N を上回り、反復的なノイズ最適化が、独立したサンプリングよりも同じ計算予算からより多くの価値を抽出することを示しました。
• タンパク質構造生成: 単一ステップのリーマン平均フロー（RMF）タンパク質バックボーン生成器に適用したところ、ZeNO は設計性（自己整合性 RMSD によって測定）を大幅に向上させました。単一のフロー評価のみを使用して、5 ステップ RMF ベースラインと同等の設計性レベルを達成し、ノイズ最適化を通じて減少したサンプリングステップを補う能力を実証しました。
• 多様性: Vendi スコアを用いた分析は、ZeNO がサンプルの多様性を維持しながら報酬を向上させ、勾配ベースのアプローチでよく見られる単一モードへの収束を回避していることを示しています。

意義と主張
本論文は、ZeNO が、効率的で決定論的かつ単一ステップの生成モデルの成長するクラスを報酬関数と整合させるための、実用的かつ理論的に裏付けられた解決策を提供すると主張しています。その主な意義は、ブラックボックス報酬モデルや複雑な生物学的設計タスクなど、勾配が利用できないか信頼できない設定で動作できる点にあります。ノイズ最適化を経路積分制御問題として定式化することで、ZeNO は、再学習や基礎となる生成モデルの修正を必要としない、スケーラブルで生成器に依存しない推論時のアライメントフレームワークを提供します。著者らは、この手法には限界があることを指摘しています。すなわち、各更新ステップで複数の報酬評価が必要であり、報酬モデル自体が高価な場合、計算集約的になる可能性があるという点です。