Fine-Tuning Discrete Diffusion Models via Reward Optimization with Applications to DNA and Protein Design

要約

最近の研究では、自然言語から生物学的配列生成まで、ドメイン全体の離散シーケンス上の拡散モデルの強力な経験的パフォーマンスが実証されています。
たとえば、タンパク質の逆折りたたみ式タスクでは、条件付き拡散モデルは、元の構造に折り返す天然のようなシーケンスを生成する印象的な結果を達成しました。
ただし、実用的な設計タスクでは、条件付き分布をモデル化するだけでなく、特定のタスク目標を最適化する必要があることがよくあります。
たとえば、安定性の高いタンパク質配列を好む場合があります。
これに対処するために、自然なシーケンスを生成できる事前に訓練された離散拡散モデルがあるシナリオと、シーケンスをタスク目的にマッピングするモデルに報いることを検討します。
次に、自然性を維持するために、補強された拡散モデルに対するKLの発散を最小限に抑えながら、補強学習（RL）に類似した離散拡散モデル内の報酬最大化問題を策定します。
このRLの問題を解決するために、Gumbel-Softmaxのトリックを使用して微分可能になっている元々拡散性の軌跡を作成することにより、拡散モデルによって生成された軌跡全体を通して報酬の直接的なバックプロパゲーションを可能にする新しいアルゴリズムであるドレイクを提案します。
私たちの理論分析は、私たちのアプローチが自然なようなシーケンスを生成し、高い報酬をもたらすことができることを示しています。
連続ドメインの拡散モデルで同様のタスクが最近検討されていますが、我々の研究は、ブラウン運動ではなく連続時間マルコフチェーンで基礎から生じる離散拡散モデルに特有の独自のアルゴリズムと理論的課題に対処します。
最後に、遺伝子療法とタンパク質ベースの治療薬の重要なタスクをそれぞれ最適化するDNAおよびタンパク質配列の生成におけるドレイクの有効性を実証します。

要約(オリジナル)

Recent studies have demonstrated the strong empirical performance of diffusion models on discrete sequences across domains from natural language to biological sequence generation. For example, in the protein inverse folding task, conditional diffusion models have achieved impressive results in generating natural-like sequences that fold back into the original structure. However, practical design tasks often require not only modeling a conditional distribution but also optimizing specific task objectives. For instance, we may prefer protein sequences with high stability. To address this, we consider the scenario where we have pre-trained discrete diffusion models that can generate natural-like sequences, as well as reward models that map sequences to task objectives. We then formulate the reward maximization problem within discrete diffusion models, analogous to reinforcement learning (RL), while minimizing the KL divergence against pretrained diffusion models to preserve naturalness. To solve this RL problem, we propose a novel algorithm, DRAKES, that enables direct backpropagation of rewards through entire trajectories generated by diffusion models, by making the originally non-differentiable trajectories differentiable using the Gumbel-Softmax trick. Our theoretical analysis indicates that our approach can generate sequences that are both natural-like and yield high rewards. While similar tasks have been recently explored in diffusion models for continuous domains, our work addresses unique algorithmic and theoretical challenges specific to discrete diffusion models, which arise from their foundation in continuous-time Markov chains rather than Brownian motion. Finally, we demonstrate the effectiveness of DRAKES in generating DNA and protein sequences that optimize enhancer activity and protein stability, respectively, important tasks for gene therapies and protein-based therapeutics.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google