Towards Controllable Diffusion Models via Reward-Guided Exploration

Towards Controllable Diffusion Models via Reward-Guided Exploration

拡散モデルは、データサンプルの形成をマルコフノイズ除去過程として定式化することで、様々なタスクにおいて最先端の性能を実現している。近年、制御されたサンプル生成を可能にするために、拡散モデルの多くのバリエーションが提案されている。これらの手法の多くは、制御情報をノイズ近似器の入力（条件表現）として定式化するか、テストフェーズで事前に訓練された分類器を導入して、ランジュバンダイナミクスを条件ゴールへ導くものである。しかし、前者は制御情報が条件付き表現として定式化できる場合にのみ有効であり、後者は事前に訓練された誘導分類器が微分可能であることが必要である。本論文では、強化学習（RL）により拡散モデルの学習段階を誘導するRGDM（Reward-Guided Diffusion Model）と名付けられた新しいフレームワークを提案する。提案する学習フレームワークは、重み付け対数尤度と最大エントロピーRLの目的を橋渡しし、政策そのものからではなく、指数関数的にスケールされた報酬に比例する報酬分布からのサンプルによって政策勾配を計算することを可能にするものである。このような枠組みは、高い勾配分散を緩和し、拡散モデルが逆プロセスで高い報酬を得られるサンプルを探索することを可能にします。3次元形状および分子生成タスクの実験では、既存の条件付き拡散モデルに対する大幅な改善が見られた。