[論文レビュー] Understanding Reinforcement Learning-Based Fine-Tuning of Diffusion Models: A Tutorial and Review

Masatoshi Uehara, Yulai Zhao+2 more|arXiv (Cornell University)|Jul 18, 2024

Traffic control and management被引用数 4

ひとこと要約

このチュートリアルでは、タンパク質安定性や分子活性など、下流の報酬関数を最適化するために事前学習済み拡散モデルを微調整する統一された強化学習（RL）ベースのフレームワークを提示する。ノイズ除去プロセスをマークフ・意思決定過程（MDP）として定式化することで、PPO や報酬重み付きMLE、パス整合性学習といったRLアルゴリズムを用いて、サンプル品質や一般化性能を維持したまま、望ましい指標を直接最大化することが可能になる。

ABSTRACT

This tutorial provides a comprehensive survey of methods for fine-tuning diffusion models to optimize downstream reward functions. While diffusion models are widely known to provide excellent generative modeling capability, practical applications in domains such as biology require generating samples that maximize some desired metric (e.g., translation efficiency in RNA, docking score in molecules, stability in protein). In these cases, the diffusion model can be optimized not only to generate realistic samples but also to explicitly maximize the measure of interest. Such methods are based on concepts from reinforcement learning (RL). We explain the application of various RL algorithms, including PPO, differentiable optimization, reward-weighted MLE, value-weighted sampling, and path consistency learning, tailored specifically for fine-tuning diffusion models. We aim to explore fundamental aspects such as the strengths and limitations of different RL-based fine-tuning algorithms across various scenarios, the benefits of RL-based fine-tuning compared to non-RL-based approaches, and the formal objectives of RL-based fine-tuning (target distributions). Additionally, we aim to examine their connections with related topics such as classifier guidance, Gflownets, flow-based diffusion models, path integral control theory, and sampling from unnormalized distributions such as MCMC. The code of this tutorial is available at https://github.com/masa-ue/RLfinetuning_Diffusion_Bioseq

研究の動機と目的

• エントロピー正則化付きMDPにおける逐次的意思決定問題として、拡散モデルのRLベースの微調整を形式化すること。
• PPO、報酬重み付きMLE、価値重み付きサンプリング、パス整合性学習といったさまざまなRLアルゴリズムを、拡散モデルの微調整において比較・対比すること。
• 報酬フィードバックの可用性（オンライン対オフライン、既知対未知の報酬）に基づいて微調整の状況を分類し、アルゴリズム選択に与える影響を分析すること。
• RLベースの微調整の形式的目的を明確にし、オフライン学習におけるターゲット分布と分布シフトのリスクを明らかにすること。
• 分類器ガイドランス、GflowNets、フローに基づくモデル、パス積分制御理論といった関連概念との関連を確立すること。

提案手法

• 各ノイズ除去ステップを報酬の累積を最大化する意思決定として扱う、拡散ノイズ除去プロセスのMDPとしての形式化。
• Proximal Policy Optimization（PPO）、微分可能な報酬バックプロパゲーション、報酬重み付き最大尤度推定（MLE）といった市販のRLアルゴリズムを用いてポリシーを最適化する。
• 高次元かつマルチモーダルな分布において特に有効な、価値重み付きサンプリングとパス整合性学習（PCL）を用いて、サンプル品質と学習安定性を向上させる。
• 事前学習済み拡散モデルを参照SDE（確率微分方程式）として活用することで、事前分布を再学習する必要がなくなり、推論速度が向上する。
• オフラインRLにおける分布シフトに起因するリスクを分析し、分布外のサンプルに対する過剰最適化の危険性を特定し、緩和戦略を提案する。
• 関連フレームワークとの統合：分類器ガイドランス（価値重み付きサンプリングの特殊ケースとして）、GflowNets（フローに基づくモデリング）、未正規化分布のサンプリングのためのMCMC。

Figure 1 : Illustrative examples of RL-based fine-tuning, aimed at optimizing pre-trained diffusion models to maximize downstream reward functions.

実験結果

リサーチクエスチョン

• RQ1 拡散モデルにおける逐次的ノイズ除去プロセスは、どのように自然に強化学習問題として定式化できるか？
• RQ2 PPO、報酬重み付きMLE、PCLといった異なるRLアルゴリズムの、下流の報酬最大化のための微調整における相対的な長所と短所は何か？
• RQ3 フィードバック収集の状況（オンライン、オフライン、既知対未知の報酬）は、RLベースの微調整におけるアルゴリズム設計と性能にどのように影響を与えるか？
• RQ4 RLベースの微調整の形式的目的とは何か？また、ターゲット分布や分布シフトとはどのように関係しているか？
• RQ5 RLベースの微調整手法は、分類器ガイドランス、GflowNets、パス積分制御理論といった既存の概念とどのように関連しているか？

主な発見

• RLベースの微調整により、タンパク質安定性、RNA翻訳効率、分子QEDスコアといった下流の報酬関数を直接最適化でき、ターゲット生成において非RLベースラインを上回る性能を発揮する。
• 報酬重み付きMLEとPPOは有効なベースラインであり、PPOは高次元かつマルチモーダルな設定でも安定性とポリシー最適化フレームワークのおかげで優れた性能を示す。
• パス整合性学習（PCL）と価値重み付きサンプリングは、安定した学習ダイナミクスを提供し、特に複雑な報酬ランドスケープにおいてモード崩壊を低減する。
• 事前学習済み拡散モデルを参照SDEとして活用することで、学習時間を著しく短縮でき、MCMCベースの手法とは異なり事前分布の推定が不要になる。
• オフラインRLの文脈では、フィードバックカバレッジが限られているため、微調整済みモデルは分布外のサンプルに対して過剰最適化されやすくなる。これは、分布ロバスト性の重要性を示唆している。
• 本フレームワークは、RLベースの微調整と分類器ガイドランスとの正式な関係を確立し、分類器ガイドランスが特定の報酬形状を施した価値重み付きサンプリングの特殊ケースであることを示している。

Figure 4 : Offline scenario