Steering Masked Discrete Diffusion Models via Discrete Denoising Posterior Prediction

要約

離散データの生成モデリングは、ChatGPT のようなテキストベースのエージェントから、タンパク質配列におけるまさに生命の構成要素の設計に至るまで、重要なアプリケーションの基礎となります。
ただし、アプリケーション ドメインは、指定されたプロパティ、報酬、またはアフィニティ メトリックを満たすために、通常は RLHF を介して生成プロセスを操作することにより、生成されたデータを制御する必要があります。
この論文では、従来の自己回帰モデルに代わる強力な代替手段を提供する最近の離散拡散モデルであるマスク拡散モデル (MDM) のステアリングの問題を研究します。
ターゲットの事後ベイズからサンプリングすることを学習することで、事前トレーニングされた MDM を操作するタスクを確率的推論の問題としてキャストする新しいフレームワークである、離散ノイズ除去事後予測 (DDPP) を紹介します。
私たちの DDPP フレームワークは、すべてシミュレーション不要で、一般的な微分不可能な報酬関数に適用しながらスケーラブルな 3 つの新しい目標のファミリーにつながります。
経験的に、クラス条件付きのピクセルレベルの画像モデリング、テキストベースの報酬を使用した MDM の RLHF ベースのアライメント、およびより多様な二次構造と短いタンパク質を生成するタンパク質言語モデルの微調整を実行するように MDM を操作することにより、DDPP をインスタンス化します。
私たちは、報酬を最適化したタンパク質配列の一過性発現を観察するウェットラボ検証を通じて設計を実証します。

要約(オリジナル)

Generative modeling of discrete data underlies important applications spanning text-based agents like ChatGPT to the design of the very building blocks of life in protein sequences. However, application domains need to exert control over the generated data by steering the generative process – typically via RLHF – to satisfy a specified property, reward, or affinity metric. In this paper, we study the problem of steering Masked Diffusion Models (MDMs), a recent class of discrete diffusion models that offer a compelling alternative to traditional autoregressive models. We introduce Discrete Denoising Posterior Prediction (DDPP), a novel framework that casts the task of steering pre-trained MDMs as a problem of probabilistic inference by learning to sample from a target Bayesian posterior. Our DDPP framework leads to a family of three novel objectives that are all simulation-free, and thus scalable while applying to general non-differentiable reward functions. Empirically, we instantiate DDPP by steering MDMs to perform class-conditional pixel-level image modeling, RLHF-based alignment of MDMs using text-based rewards, and finetuning protein language models to generate more diverse secondary structures and shorter proteins. We substantiate our designs via wet-lab validation, where we observe transient expression of reward-optimized protein sequences.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google