[論文レビュー] Aligning Target-Aware Molecule Diffusion Models with Exact Energy Optimization

Siyi Gu, Minkai Xu+7 more|arXiv (Cornell University)|Jul 1, 2024

Machine Learning in Materials ScienceMaterials Science被引用数 3

ひとこと要約

本論文は、正確なエネルギー選好最適化（E²PO）を用いて、事前学習済みのターゲットに依存する分子拡散モデルを望ましい機能的特性に整合させる新規フレームワークAliDiffを提案する。バインディング親和性と構造的妥当性のユーザー定義報酬関数を学習することで、CrossDocked2020ベンチマークにおいて、最高水準のバインディングエネルギー（最高-7.07平均Vinaスコア）を達成しながら、良好な分子特性と多様性を維持した。

ABSTRACT

Generating ligand molecules for specific protein targets, known as structure-based drug design, is a fundamental problem in therapeutics development and biological discovery. Recently, target-aware generative models, especially diffusion models, have shown great promise in modeling protein-ligand interactions and generating candidate drugs. However, existing models primarily focus on learning the chemical distribution of all drug candidates, which lacks effective steerability on the chemical quality of model generations. In this paper, we propose a novel and general alignment framework to align pretrained target diffusion models with preferred functional properties, named AliDiff. AliDiff shifts the target-conditioned chemical distribution towards regions with higher binding affinity and structural rationality, specified by user-defined reward functions, via the preference optimization approach. To avoid the overfitting problem in common preference optimization objectives, we further develop an improved Exact Energy Preference Optimization method to yield an exact and efficient alignment of the diffusion models, and provide the closed-form expression for the converged distribution. Empirical studies on the CrossDocked2020 benchmark show that AliDiff can generate molecules with state-of-the-art binding energies with up to -7.07 Avg. Vina Score, while maintaining strong molecular properties. Code is available at https://github.com/MinkaiXu/AliDiff.

研究の動機と目的

• 既存のターゲットに依存する拡散モデルには制御性が欠けているという問題に対処する。これにより、高いバインディング親和性や望ましい分子特性を優先しない分子が生成される。
• バインディングエネルギーが低く、構造的に妥当性の高い領域に、ターゲット条件付きの化学的分布をシフトさせる一般化された整列フレームワークを開発する。
• 閉形式の収束を備えた正確なエネルギー最適化手法を導入することで、報酬関数を用いた整列における過学習を克服する。
• 高い親和性を持つ分子を生成する際にも、ドラッグライクネス（QED）、合成可能性（SA）、多様性といった良好な分子特性を維持する。
• 異なる拡散ベースの構造基盤ドラッグデザインモデルに一般化可能であることを示す。

提案手法

• フレームワークは、離散的化学的タイプと連続的3次元座標の両方の尤度を整列化するための統一された変分下界を用いる。
• 正確なエネルギー選好最適化（E²PO）を導入し、ユーザー定義報酬関数下での収束分布を解析的に計算することで、標準的な選好最適化における過学習を回避する。
• E²POは更新された条件付き分布の閉形式表現を提供し、拡散モデルの正確かつ効率的な整列を可能にする。
• 報酬関数を統合し、バインディングエネルギーと構造的妥当性をエンコードすることで、より高い親和性と安定性を持つリガンドへとモデルを誘導する。
• IPDiff や TargetDiff などの事前学習モデルを微調整することで、順序付けられた分子ペアからの選好信号を用いてアプローチを適用する。
• ハイパーパrameter β はバインディング親和性と分子特性の間のトレードオフを制御し、異なるドラッグ開発ニーズに合わせたカスタマイズ生成を可能にする。

Figure 1: High-level illustration of AliDiff . For a protein target, we can have multiple candidate ligands and rank the preference by certain reward functions, e.g . , binding energy. We align the target-aware molecule diffusion model with these preferences by adjusting the conditional likelihoods.

実験結果

リサーチクエスチョン

• RQ1 ターゲットに依存する分子拡散モデルを、分子特性を損なわせることなく、より高いバインディング親和性へと効果的に制御することは可能か？
• RQ2 報酬関数を用いて拡散モデルを整列させる際、選好最適化における過学習をどのように回避できるか？
• RQ3 エネルギーに基づく選好最適化に対して、正確な閉形式解が標準的手法よりも優れた性能をもたらすか？
• RQ4 提案された整列フレームワークは、異なる拡散ベースの構造基盤ドラッグデザインモデルに一般化可能か？
• RQ5 バインディング親和性と分子特性のトレードオフは何か？そして、1つのハイパーパrameterで制御可能か？

主な発見

• AliDiffは、CrossDocked2020ベンチマークで最高水準の平均Vinaスコア-7.07を達成し、ベースラインモデルを著しく上回った。
• AliDiff-E²PO（E²POを用いたバージョン）は、直接選好最適化ベースライン（AliDiff-DPO）と比較して、優れたバインディング親和性指標を達成した一方で、QED、SA、多様性スコアも競争的であった。
• AliDiffを用いてTargetDiffを微調整したAliDiff-Tは、元のモデルと比較して平均Vinaスコアが6.2%向上、QEDが16.6%向上、SAが2.8%向上、多様性が2.8%向上した。
• アブレーションスタディの結果、β はバインディング親和性と分子特性のトレードオフを制御しており、β=1でバランスの取れた性能が得られた。
• β の低い値（例：0.01）は、バインディング親和性を顕著に向上させるが、QEDと多様性を低下させる。一方、β の高い値はドラッグライクネスを向上させるが、わずかに親和性が低下する。
• 本フレームワークは一般化可能であり、IPDiff や TargetDiff などの複数の拡散モデルにおいて一貫した改善が確認された。

Figure 2: Overview of AliDiff . This workflow can be summarized as 1) For each protein target (pocket) ${\mathbf{p}}$ in the training set, we retrieve two candidate ligands ${\mathbf{m}}$ ; 2) Label the two ligands as wining sample ${\mathbf{m}}^{w}$ and losing sample ${\mathbf{m}}^{l}$ by desirable