[論文レビュー] Bridging Model-Based Optimization and Generative Modeling via Conservative Fine-Tuning of Diffusion Models

Masatoshi Uehara, Yulai Zhao+6 more|arXiv (Cornell University)|May 30, 2024

Advanced Mathematical Modeling in Engineering被引用数 4

ひとこと要約

本稿では、報酬モデルと拡散モデルを組み合わせることで、オフライン設計最適化を改善する、保守的微調整フレームワークBRAIDを提案する。二重に保守的なアプローチ——報酬モデルにおける分布外領域のペナルティ付与とKLダイバージェンス制約の導入——により、拡散モデルがオフラインで観測された最良の設計を上回るが、無効または敵対的生成を回避する。

ABSTRACT

AI-driven design problems, such as DNA/protein sequence design, are commonly tackled from two angles: generative modeling, which efficiently captures the feasible design space (e.g., natural images or biological sequences), and model-based optimization, which utilizes reward models for extrapolation. To combine the strengths of both approaches, we adopt a hybrid method that fine-tunes cutting-edge diffusion models by optimizing reward models through RL. Although prior work has explored similar avenues, they primarily focus on scenarios where accurate reward models are accessible. In contrast, we concentrate on an offline setting where a reward model is unknown, and we must learn from static offline datasets, a common scenario in scientific domains. In offline scenarios, existing approaches tend to suffer from overoptimization, as they may be misled by the reward model in out-of-distribution regions. To address this, we introduce a conservative fine-tuning approach, BRAID, by optimizing a conservative reward model, which includes additional penalization outside of offline data distributions. Through empirical and theoretical analysis, we demonstrate the capability of our approach to outperform the best designs in offline data, leveraging the extrapolation capabilities of reward models while avoiding the generation of invalid designs through pre-trained diffusion models.

研究の動機と目的

• 報酬モデルが分布外領域でモデルを誤導する可能性がある、オフラインモデルベース最適化における過剰最適化の課題に対処する。
• 静的オフラインデータセットから学習された報酬関数を用いて、事前学習済み拡散モデルの微調整により、生成モデルとモデルベース最適化を統合する。
• オフラインデータで観測された最良の設計を越えて外挿を可能にしつつ、有効な設計空間への忠実度を維持する手法を開発する。
• 微調整中に不確実性を考慮したペナルティとKL正則化を組み込むことで、分布シフトに対するロバストネスを確保する。

提案手法

• オフラインデータを用いて、分布外入力に対してペナルティを増加させる不確実性評価項を組み込んだ保守的報酬モデルを訓練する。
• ソフトエントロピー正則化付きのマルコフ決定過程フレームワークを用いて、微調整目的関数への保守的正則化の組み込みを理論的に正当化する。
• 強化学習を用いて保守的報酬モデルを最適化することで、事前学習済み拡散モデルを微調整し、生成を有効な設計空間内に制約するためのKLペナルティを追加する。
• 訓練プロセスを二段階に分解する：(1) 保守的正則化を伴うオフライン報酬モデルの訓練、(2) 保守的報酬モデルを用いた拡散モデルの微調整。
• 学習済み報酬と視覚言語モデル（例：LLaVA）からの一貫性信号を組み合わせた複合報酬関数を定義し、画像生成におけるプロンプト適合性を保証する。
• 生物的配列設計（DNA/UTR/エンハンサー）および画像生成タスクの両方に対して、EnformerおよびStable Diffusionを事前学習済みモデルとして適用する。

実験結果

リサーチクエスチョン

• RQ1 報酬モデルが限定的かつ静的データセットから学習される場合、保守的微調整アプローチは、オフラインモデルベース最適化における過剰最適化を緩和できるか？
• RQ2 報酬モデルの外挿能力を活用しながら、生成設計の有効性を損なわず、どのように統合できるか？
• RQ3 報酬モデルの学習段階で不確実性を考慮したペナルティを組み込むことで、観測されたデータ分布を超えて一般化性能が向上する程度はどの程度か？
• RQ4 保守的報酬モデルとKL正則化の組み合わせが、配列および画像生成タスクの両方でよりロバストかつ忠実な生成を実現するか？
• RQ5 報酬性能と分布的忠実度の観点から、提案手法は条件付き微調整および標準的なRL微調整と比較して、どのように異なるか？

主な発見

• BRAIDは、DNA配列および画像生成タスクの両方で、オフラインデータに含まれる最良の設計を上回り、観測されたデータ分布を超えて一般化する能力を示した。
• 画像生成において、BRAIDが生成する画像はLLaVAおよびLAION予測子による美的スコアが高く、プロンプト適合性と視覚的品質を維持している。これに対してSTRLは、プロンプトから逸脱した画像を生成する。
• 学習済み報酬（$\hat{r}$）に関する訓練曲線では、BRAIDは安定的かつ一貫した改善を示したが、STRLは報酬の急激な上昇を示したが、一般化性能と分布シフトに劣っていた。
• 実験的結果から、BRAID-bootおよびBRAID-bonusは報酬性能においてSTRLを著しく上回り、保守的正則化の有効性を裏付けた。
• 理論的分析により、本手法はソフトエントロピー正則化付きMDP設定で定式化され、保守的微調整がレギュレートを低減し、オフラインベースラインを上回る性能を実現できることを示した。
• アブレーションスタディの結果、保守的報酬モデルとKLペナルティの両方が不可欠であることが確認された。両者のいずれかを除去すると、性能が低下し、分布外生成が増加した。