TR2-D2: 離散拡散モデルのための軌道認識型微調整をガイドする木探索

要旨

確率的最適制御を用いた強化学習は、拡散モデルのファインチューニングにおいて有望なフレームワークを提供します。ここでは、事前学習済みの拡散モデルを最適化し、報酬に傾いた分布へと導く経路を生成します。これらのアプローチは、最適分布からの明示的なサンプルにアクセスすることなく最適化を可能にしますが、現在のファインチューニングされたモデル下でのロールアウトに基づいて訓練を行うため、報酬が低い非最適な軌道を強化してしまう可能性があります。この課題を克服するため、我々はTRee Search Guided TRajectory-Aware Fine-Tuning for Discrete Diffusion (TR2-D2)を提案します。これは、報酬に基づく離散拡散軌道を最適化し、軌道を意識したファインチューニングのためのリプレイバッファを構築する新しいフレームワークです。これらのバッファはモンテカルロ木探索（MCTS）を用いて生成され、その後、確率的最適制御の目的関数の下で事前学習済みの離散拡散モデルをファインチューニングするために使用されます。我々は、生物学的配列拡散モデルの単一および多目的ファインチューニングにおいてこのフレームワークを検証し、離散配列生成における信頼性の高い報酬に基づくファインチューニングのためのTR2-D2の全体的な有効性を強調します。

English

Reinforcement learning with stochastic optimal control offers a promising framework for diffusion fine-tuning, where a pre-trained diffusion model is optimized to generate paths that lead to a reward-tilted distribution. While these approaches enable optimization without access to explicit samples from the optimal distribution, they require training on rollouts under the current fine-tuned model, making them susceptible to reinforcing sub-optimal trajectories that yield poor rewards. To overcome this challenge, we introduce TRee Search Guided TRajectory-Aware Fine-Tuning for Discrete Diffusion (TR2-D2), a novel framework that optimizes reward-guided discrete diffusion trajectories with tree search to construct replay buffers for trajectory-aware fine-tuning. These buffers are generated using Monte Carlo Tree Search (MCTS) and subsequently used to fine-tune a pre-trained discrete diffusion model under a stochastic optimal control objective. We validate our framework on single- and multi-objective fine-tuning of biological sequence diffusion models, highlighting the overall effectiveness of TR2-D2 for reliable reward-guided fine-tuning in discrete sequence generation.

TR2-D2: 離散拡散モデルのための軌道認識型微調整をガイドする木探索

TR2-D2: Tree Search Guided Trajectory-Aware Fine-Tuning for Discrete Diffusion

要旨

Support