4Paradigm が分子特性予測のための生成型 3D 事前トレーニング済みモデルを開発

論文タイトル: 分子特性予測のための自動 3D 事前トレーニング

論文リンク: https://arxiv.org/abs/2306.07812

コードリンク: https://github.com/LARS-research/3D-PGT

この論文はKDD 2023に採択されました。

1. はじめに

分子特性の予測は、創薬や材料科学などの分野において重要な問題です。分子の3D構造情報は分子の特性と密接に関係していることから、近年では分子の3D構造情報とグラフ学習に代表されるさまざまな機械学習手法を組み合わせて分子特性の予測性能を向上させることが主流の研究ホットスポットとなっています。しかし、計算コストが高いため、大規模なデータセット上で分子の3D構造をリアルタイムで計算することはほぼ不可能である[1]。

図1: cc(=o)oc1=cc=cc=c1c(=o)oの2D分子図（左）と3D立体配座図（右）。

この課題に直面して、私たちは分子特性予測に特化した事前トレーニングパラダイムに基づく大規模モデルを提案します。事前学習パラダイムはAIGC[2]の重要な技術分野であり、大規模モデルの利点を最大限に活用して、大規模なデータセットを処理し、深い特徴を抽出することができます。分子分野では、本論文では、事前トレーニングに既存の3D分子立体配座を組み合わせ、実際の下流タスクでは分子の2D構造情報に基づいて特性を微調整して予測するだけで、高い効率を確保しながら下流タスクのパフォーマンスをさらに向上させます。

この研究では、自動化された 3D 事前トレーニング フレームワークである 3D-PGT を提案します。分子の化学結合長、結合角、二面角は、完全な分子の 3D 立体配座に対応する 3 つの基本的な幾何学的記述子であるという事実に基づいて、モデルが事前トレーニングを通じて 3D 幾何学的構造をエンコードできるように、それに応じて 3 つの生成的な事前トレーニング タスクを設計しました。 3 つの事前トレーニング タスクの重みを自動的に割り当てて、全体的な事前トレーニング目的関数に統合するために、総分子エネルギーに基づいて 3 つの事前トレーニング タスクの重み分布を自動的に検索する代替メトリックを設計しました。

図 2: 分子特性予測の核心課題の図解: 推論速度の高速化と推論誤差の低減の両方を実現する方法。

設計された事前トレーニング フレームワークの有効性を検証するために、決定された 3D 構造を持つ公開分子データセットに基づいて事前トレーニングを行い、8 つの量子化学特性予測タスクと薬理学および分子化学を含む 12 の下流タスクで微調整と実験検証を実行しました。結果は、3D-PGT が 3D 事前トレーニングを通じて大幅なパフォーマンス向上をもたらし、他の事前トレーニング ベースラインよりも優れていることを示しています。これは、分子特性予測の分野における複雑なタスクと大規模データの処理における大規模モデルの利点を改めて実証しています。同時に、3D-PGTベースの手法は、触媒シミュレーションチャレンジ「Open Catalyst Challenge 2022」で第3位、大規模分子特性予測チャレンジ「OGB-LSC@NeruIPS challenge 2022」で第8位も獲得しました。

2. 背景

分子の特性と挙動を記述する基本的な属性として、分子特性は、薬理学や材料化学などの研究分野や特定の応用において非常に重要です。しかし、分子の特性を決定するには、密度汎関数理論 (DFT) に代表される計算化学手法などのコストのかかる計算手法を使用する必要があることがよくあります。近年、大規模な量子化学コンピューティングとハイスループット実験の技術進歩に伴い、適切な誘導バイアスを備えた機械学習手法を使用し、それを測定された特性を持つ大量の分子データと組み合わせて効率的な分子特性予測を実現し、それを大規模な薬物分子スクリーニング、合成材料スクリーニング、触媒設計などの一連の実用的な下流アプリケーションに適用するという研究方向が、産業界と学界で新たに注目されています。

図3: 分子特性予測効率の比較の模式図。量子化学計算に基づく特定の分子特性の DFT 計算には数時間かかりますが、機械学習手法を使用すると 1 秒未満しかかかりません。この図は、分子の HOMO-LUMO エネルギーギャップを予測する OGB@NeurIPS 2022 チャレンジから得られたものです。

ここ数年、分子を2Dグラフ構造としてモデル化し、原子をノード、エッジを化学結合として扱い、分子特性予測全体をグラフレベルの予測タスクとしてモデル化し、DFT計算法を当てはめてGNN（グラフニューラルネットワーク）を適用し、分子の特性を予測するというアプローチが主流となってきました。しかし、現時点では、この方法は予測効率が高いだけであり、その予測性能は実用化には程遠い。

図4: 既存の分子特性予測方法の一般的なプロセス。 DFT に基づく量子化学計算法は、まず分子の 3 次元構造を最適化して低エネルギーの立体配座を取得し、次にその立体配座に基づいて特定の分子特性を計算します。 2D モデルは DFT の全プロセスを近似し、入力された 2D 分子グラフを通じて分子特性を直接予測します。3D モデルでは、より優れた予測性能を得るために、2D 分子グラフに基づいて 3D 分子構造をさらに入力する必要もあります。

分子の 3D 構造は、原子と官能基間の相互作用と相対的な位置を反映するため、分子の特性を理解する上で非常に重要です。これを考慮して、一連の方法で分子の 3D 構造情報に基づいて 3D モデルを設計し、より優れた予測結果を実現します。しかし、3D 構造を取得するには、DFT などの量子化学計算方法に頼る必要があります。計算コストが高額なため、多くの実用的な下流タスクでは分子の 3D 構造が利用できないことがよくあります。

前述の 2 つの主流ルートの長所と短所を組み合わせた比較的新しい観点は、測定された 3D 構造情報を含むデータセットを事前トレーニング データセットとして使用し、3D 構造に基づいて事前トレーニング タスクを設計することで、モデルが 2D 分子グラフから 3D 構造を DFT 計算して最適化するプロセスを理解し、この事前情報を下流の 2D 分子特性予測タスクに転送し、効率とパフォーマンスの両方を実現することです。現在、GraphMVP[3]や3D Infomax[4]などの主流の現代研究では、分子の2Dビューと3Dビューのグラフレベル表現ベクトルを揃えることで事前トレーニングの目的関数を設計し、主流のベンチマークでの実験を通じて3D事前トレーニングスキームの有効性を証明しています。

要約すると、本論文は分子 3D 事前トレーニング パイプラインに沿って進み、3 つの基本的な生成事前トレーニング タスクを設計し、各事前トレーニング タスクの重みを自動的に検索するための分子エネルギー全体に基づく目的関数を設計し、それによって複数の事前トレーニング タスクの自動融合フレームワークを構築し、幅広い下流タスクで事前トレーニングの大きなメリットを実現します。

3. 本論文の方法

3.1 生成的事前トレーニングタスクの設計

分子の3D構造のDFT最適化の核心は、ポテンシャルエネルギー面上の局所最小値を見つけることです[5]。分子エネルギーがポテンシャルエネルギー面上の局所最小値にあるときの3D構造は、分子コンフォメーションとも呼ばれます。その中で、立体配座に含まれる 3D 幾何学情報は、2 つの原子間の化学結合長、2 つの化学結合間に形成される角度、および 3 つの化学結合によって形成される二面角という 3 つの記述子によって完全に記述できます。

生成的な事前トレーニング タスクにより、モデルは 2D トポロジから 3D ジオメトリまでの分子の DFT ベースの最適化プロセスを理解できるようになります。分子構造は電子の量子力学的挙動によって決定されるため、生成事前トレーニングタスクは、3D コンフォメーションの生成を学習することによって、量子化学的特性の予測を間接的に学習できます。これに基づいて、本論文では、以下に示すように、結合長、結合角、二面角の 3 つの記述子をそれぞれ生成するための 3 つの生成事前トレーニング タスクを設計します。

図5: 分子構造とポテンシャルエネルギー面の関係を示す模式図。ポテンシャル エネルギー面は、分子エネルギーを複数の座標軸の関数として定義します。分子の 3D 立体構造の最適化は、ポテンシャル エネルギー面の局所的最小値を見つけることです。3D 立体構造の幾何学的情報は、2 つの原子を結合する結合長、結合した 3 つの原子の結合角、および連続する 3 つの結合の二面角によって記述できます。

このうち、h は対応するノードの表現ベクトルを表し、f は対応する予測ネットワーク（ここでの設定は MLP）を表します。これら 3 つのタスクの核となるアイデアは、バックボーンによって抽出されたノード表現を使用して特定の回帰タスクを設計することであることがわかります。記述子に含まれる 3D 構造情報を解釈信号として使用することで、バックボーンは 3D 構造情報をエンコードできるようになり、この機能は事前にモデルに埋め込まれ、実際の下流タスクに移行されます。

同時に、近傍数が|N|の場合、分子の結合角と二面角の計算量はΟ(|N|^2)とΟ(|N|^3)増加するため、大規模な3D事前トレーニングを実装するにはコストがかかりすぎることを考慮すると、この論文では、RGC (Runtime Geometry Calculation) に基づいて、結合角と二面角の目的関数を再設計し、すべての結合角と二面角の予測を、各原子に関係する結合角と各化学結合の設計された二面角の合計に置き換え、計算の複雑さを線形レベルまで削減します。

3.2 自動化されたマルチ事前トレーニングタスク融合フレームワーク

複数の事前トレーニングタスクを同時に定義し、各事前トレーニングタスクの損失関数はモデルパラメータに対して異なる勾配最適化方向を持つため、各事前トレーニングタスクの重み割り当ては考慮する必要がある問題です。下流のタスクでより優れたパフォーマンス向上を実現するために、複数の事前トレーニング タスクの重みを自動的かつ最適に割り当てる方法が、複数の事前トレーニング タスク融合フレームワークを設計する上での中核となります。

結合長、結合角、二面角はすべてローカル記述子であるため、ローカルエンコーディング能力ではなく、全体的な分子 3D 構造に対する事前トレーニング済みバックボーンのエンコーディング能力を評価するには、明確な代替メトリックが必要です。この点に関して、本論文では、分子の全エネルギーと分子の 3D 構造との対応を考慮し、分子の全エネルギーに基づいて目的関数を設計し、これに基づいて 3 つの生成事前トレーニング タスクの重みを検索するための 2 レベルの油絵フレームワークを設計します。最後に、全体的な事前トレーニング パイプラインを下の図に示します。事前トレーニング段階では、まず事前トレーニング分子の低エネルギー立体配座に基づいて 3 つの生成事前トレーニング タスクを設計し、分子の低エネルギー立体配座に対応する総分子エネルギーに基づいて代理メトリックを設計し、3 つの事前トレーニング タスクのそれぞれの重みを検索しました。下流の微調整段階では、事前トレーニング段階で分子形状の事前条件を導入したため、微調整によって実際の下流タスクでパフォーマンスの向上が得られます。

図6: 3D-PGTの事前トレーニングフレームワークの概略図

4. 実験結果

本論文では、3D 事前トレーニング技術ルートと、本論文で設計した事前トレーニング タスクおよび自動融合フレームワークの有効性を実証するために、一連の実験を具体的に設計しました。予測に使用できるのが 2D 分子グラフのみの場合、各分子の 3D コンフォメーションを生成するための膨大な計算コストを回避でき、微調整によって大幅なパフォーマンス向上が得られます。

本稿では、主に3D構造情報を含む3つのデータセットに対して事前トレーニングを行い、12の下流タスクに対して微調整と効果検証を主に行います。データセットの詳細な統計は次のとおりです。

図7: データセットの統計詳細

4.1 量子化学特性予測の性能

この論文ではまず、一般的なベンチマークデータセット QM9 上で 3D-PGT の量子化学特性予測機能を評価します。まず、事前トレーニング用に、QM9 に含まれる 134,000 個の単一分子コンフォメーション サンプルからコンフォメーション情報を持つ 50,000 個の分子サンプルをランダムに選択し、残りのサンプルから 50,000 個の分子サンプルを選択して 3D 構造情報をマスクし、微調整と評価を行いました。結果は次のとおりです。

図8: QM9データセット実験の比較結果（評価指標はMAE（平均絶対誤差））

実験結果から、次の現象を観察しました。最初に、3D-PGTはGPSを前訓練と微調整のためにバックボーンとして使用しますが、3D-PGTは8つの量子化学特性を摂取すると、8量子化学物質を摂取することで平均17.7％減少します3D構造情報を使用せずにトレーニング前にもたらされる3）2つの最先端の3Dプリトレーニング方法と比較して、私たちのトレーニング前のフレームワークはパフォーマンスの利点があります。 3Dプリトレーニングパラダイムの効率とパフォーマンスの利点を反映した8つの予測タスク。

図9: HOMO予測タスクにおけるさまざまな手法の推論時間とパフォーマンスランキング

4.2 その他の幅広い下流タスクにおける一般化能力

分子の 3D 構造に密接に関連する量子化学特性の予測に加えて、下流のタスクを薬理学、物理学、生物学などの分野に拡張します。これらのタスクの入力には 2D 分子グラフのみが含まれます。事前トレーニング用に、GEOM データセットから単一の立体配座を持つ 50,000 個の分子をランダムに選択し、8 つの主流の下流分子特性分類タスクと 6 つの回帰タスクで微調整しました。これらの下流タスクに含まれるデータの量は比較的少ないです。具体的な結果は次のとおりです。

図10: 4つの分子特性予測タスク(RMSE)と2つの薬物タンパク質標的親和性予測タスク(MSE)の実験結果

3D-PGT はほとんどのダウンストリーム タスクで他のベースラインよりも優れていることがわかります。また、GPS と比較すると、3D-PGT は事前トレーニングによって安定したパフォーマンスの向上を得ることができます。これは、私たちが設計した事前トレーニング フレームワークが、3D 構造に密接に関連する量子化学特性の予測に限定されず、幅広い下流タスクにパフォーマンスの向上を一般化できること (分子特性予測自体に加えて、2 つの薬物および標的タンパク質親和性予測タスクである Davis と KIBA でもパフォーマンスの向上が達成されました)、および小さなサンプル データでの微調整でもパフォーマンスの向上を実現できることを意味します。

4.3 大規模データチャレンジにおけるパフォーマンス

OGB-LSC のグラフレベル予測タスクトラックでは、PCQM4Mv2 は 374 万分子を含む大規模な分子データセットであり、そのうち 337 万のトレーニング サンプルには DFT によって計算された 3D 幾何学情報が含まれています。大規模な仮想スクリーニング シナリオにアプローチするために、このチャレンジでは検証セットとテスト セットに 3D コンフォメーションを提供しておらず、15 万分子の HOMO-LUMO ギャップの推論を 4 時間以内に完了するには単一の GPU を使用する必要があります。つまり、モデルの推論フェーズ中にすべてのテスト サンプルの幾何学的構造を計算することは現実的ではありません。このデータセットでの 3D-PGT のパフォーマンスは、次のようにリーダーボードのトップ メソッドと比較されます。

図11: PCQM4Mv@検証セットでの実験結果

3D構造情報を考慮しない既存のGNNやTransformerシリーズの手法と比較すると、3D-PGTは生成的な事前トレーニングタスクを導入することで大幅なパフォーマンスの向上を実現していることがわかります。 GPS ベースラインでは、3D-PGT は設計された自動事前トレーニング フレームワークを通じて MAE を 10.6 削減し、単一モデルのパフォーマンスは OGB@NeruIPS 2022 チャンピオン ソリューション GPS++ よりも優れていました。

一方、Open Catalyst Challenge 2022 は、分子シミュレーションを使用して高性能触媒を発見することに特化したチャレンジです。このコンテストの中心的な課題の 1 つは、機械学習モデルを使用して触媒分子の活性を予測することです。トレーニング セットには 200 万個の分子が含まれており、各分子サンプルの 3D 構造と総エネルギーが記録されます。最終的に、AutoGraph のアルゴリズムと組み合わせることで、3D-PGT はチャレンジで 3 位を獲得し、事前トレーニング フレームワークの有効性と競争力を再び証明しました。

図12: Open Catalyst Challenge 2022 リーダーボードの結果

5. 結論

本研究では、実際の 3D 構造が利用できない分子特性予測タスクのパフォーマンスを向上させることに重点を置いた、自動化された 3D 事前トレーニング フレームワークである 3D-PGT を提案します。

3D-PGT は、幾何学的事前分布を微調整段階に持ち込むことができる複数の生成事前トレーニング タスクを設計します。同時に、これらの事前トレーニング タスクをより適切に自動的に統合し、その利点を普遍的にするために、各事前タスクの適応重みを検索するための総分子エネルギーに基づく事前トレーニング済みの代替メトリックを設計しました。この論文では、既存の主流の量子化学データセットと幅広い下流タスクに基づいて実験を設計し、実験結果から、3D-PGT が事前トレーニングを通じて潜在的な幾何学的事前確率を導入することが示されており、これは量子化学特性の予測だけでなく、薬理学、物理化学、生物物理学などの分野での予測にも有益です。さらに、OGB リーダーボードでは、3D-PGT は大規模な分子予測におけるトップソリューションのすべてのベースラインを上回っています。

今後の研究では、NLP 分野における GPT 事前トレーニング パラダイムの成功を考慮して、このフレームワークと 3D 分子事前トレーニング パラダイムを、再生可能エネルギーを貯蔵するための触媒の開発などの実際の下流アプリケーションに拡張します。

注: この論文は KDD 2023 に採択されました。ご質問がある場合は、[email protected] までお問い合わせください。