人工知能技術は医薬品開発プロセスを大きく変革しており、分子特性予測と分子生成という二つの中核タスクは、それぞれ独立した技術的発展経路を辿って長らく発展してきました。分子特性予測は、分子の構造情報に基づいて多様な化学的・生物学的特性を予測し、医薬品スクリーニングを加速することを目指しています。一方、分子生成は、分子データの分布を推定し、原子間相互作用や立体配座情報を学習することで、化学的に妥当な新規分子の新規生成を可能にし、医薬品設計の可能性を拡大することを目指しています。近年、これらの分野で広範な研究が行われているにもかかわらず、それらは主に独立して発展しており、二つの重要なステップの間に相乗効果をもたらす道筋を効果的に確立できていません。

この問題に対処するため、清華大学と中国科学院の研究チームはUniGEMモデルを提案し、拡散モデルに基づく2つのタスクの相乗的な強化を初めて実現しました。研究チームは、分子生成と特性予測は高度に相関しており、効果的な分子表現に依存していると指摘しました。チームは革新的な2段階生成プロセスを提案し、従来の共同学習における矛盾を克服し、分子生成と特性予測の分野に新たな道を開きました。「UniGEM：分子の生成と特性予測への統一的アプローチ」と題されたこの研究は、ICLR 2025に採択されました。

論文の宛先:

https://openreview.net/pdf?id=Lb91pXwZMR

QM9量子化学データセット:

https://go.hyper.ai/zgPLs

GEOM-Drugs 3D分子コンフォメーションデータセット:

https://go.hyper.ai/FOtLx

オープンソース プロジェクト「awesome-ai4s」は、200 を超える AI4S 論文の解釈をまとめ、膨大なデータセットとツールを提供します。

https://github.com/hyperai/awesome-ai4s

生成と予測のタスクを統合する動機

研究チームは、生成と予測の両方のタスクの本質は分子表現の学習にあると考えています。一方で、様々な分子事前学習法の有効性は、分子特性予測が堅牢な分子表現を基盤として依存していることを示しています。他方、分子生成には、生成プロセス中に適切な表現を作成するために、分子構造の深い理解が必要です。

最近の研究はこの見解を裏付けています。例えば、コンピュータービジョンの研究では、拡散モデル自体が効果的な画像表現を学習する能力を持っていることが示されています。分子領域では、生成的な事前学習によって分子特性予測タスクを強化できることが研究で示されていますが、これらの手法では最適な予測性能を得るためには通常、追加の微調整が必​​要です。さらに、予測器は分類器誘導法を通じて分子生成を導くことができますが、予測器の学習が直接的に生成性能を向上させるかどうかは依然として不明です。

したがって、既存の研究では、生成タスクと予測タスクの関係が完全に解明されておらず、 「生成タスクと予測タスクの相乗的な強化を実現する統一モデルを構築できるのか」という重要な疑問が生じます。

従来の方法の失敗理由の分析

これら2つのタスクを組み合わせる最も簡単な方法は、従来のマルチタスク学習フレームワークを用いることです。このフレームワークでは、モデルが生成損失と予測損失の両方を同時に最適化します。しかし、研究チームの実験では、このアプローチは生成タスクと特性予測タスクの両方のパフォーマンスを著しく低下させることが示されました（生成安定性が6%も急激に低下し、予測誤差が2倍以上に増加しました）。生成モデルの重みを固定し、特性予測タスク用に別のヘッドを追加して生成パフォーマンスを維持した後でも、研究者たちはゼロから学習した場合と比較して特性予測パフォーマンスの向上を確認できませんでした。

研究者たちは、従来の手法が満足のいく結果につながらない原因として、生成タスクと予測タスクの間に固有の不一致を挙げています。拡散生成においては、分子構造は無秩序なノイズから微細構造へと徐々に再構築されます。一方、予測タスクにおいては、基本的な分子構造が確立された後にのみ、意味のある分子特性を定義できます。そのため、単純なマルチタスク最適化手法を単純に適用すると、拡散初期段階で、高度に無秩序な分子構造と特性ラベルとの関連付けに誤りが生じ、分子生成と特性予測の両方に悪影響を及ぼす可能性があります。

この点をさらに説明するために、研究者らは、拡散学習中のノイズ除去ネットワークにおける中間表現と標的分子間の相互情報量を理論的に解析しました。そして、拡散モデルが中間表現と標的分子間の相互情報量の下限を暗黙的に最大化することを理論的に証明し、その表現学習能力を実証しました。しかし、中間表現と標的分子間の相互情報量は単調に減少する傾向を示し、より大きな時間ステップではゼロに近づきます。これは、無秩序段階の中間表現では効果的な予測をサポートできないことを意味します。したがって、直感と理論の両方から、生成タスクと予測タスクは、より小さな時間ステップ、つまり分子が比較的高い秩序を維持している場合にのみ整合できることが示唆されます。

二段階拡散生成メカニズム

研究チームは、上記の分析に基づいて、下図に示すように、分子特性の予測と生成を統合することを目指した新しい2段階生成方法を提案しました。

キャプション：UniGEMの2段階生成プロセス

研究者は分子生成プロセスを「分子核形成段階」と「分子成長段階」の 2 つの段階に分けます。この区分は物理学における結晶形成プロセスにヒントを得たものです。

分子核形成段階では、分子は完全に無秩序な状態から骨格を形成し、その後、この骨格に基づいて完全な分子へと成長します。これらの2つの段階は「核形成時間」によって区切られています。研究者たちは、この2つの段階を記述するために、新たな分子生成メカニズムを導入しました。「核形成時間」の前段階では、拡散モデルは分子座標を徐々に生成します。核形成後、モデルは分子座標を調整し続けながら、特性と原子の種類に対する損失予測を最適化します。

従来の生成モデルでは原子の種類と座標の同時拡散が一般的ですが、この革新的なアプローチは座標拡散のみに焦点を当て、原子の種類を独立した予測タスクとして扱います。これは、研究者が、形成された分子座標から原子の種類を推測できることが多いことに気づいたためです。具体的には、核形成前は拡散プロセスの目的は座標を再構築することであり、核形成後は、原子の種類と特性の予測損失を統一された学習フレームワークに統合します。

UniGEMトレーニング戦略

従来の結合拡散法との比較を容易にするため、研究者らはEGNNをネットワーク構造のバックボーンとする等拡散モデル（EDM）を使用しました。成長段階は学習プロセス全体の約1%を占めるに過ぎません。標準的な拡散学習手順に従い、時間ステップを均一にサンプリングした場合、予測タスクの反復回数は学習プロセス全体のわずか1%に過ぎず、このタスクにおけるモデルのパフォーマンスは大幅に低下します。そのため、予測タスクに十分な学習量を確保するため、研究者らは成長段階の時間ステップをオーバーサンプリングしました。

しかし、研究者たちは、オーバーサンプリングはタイムステップ全体にわたってトレーニングの不均衡につながり、生成プロセスの品質に影響を与える可能性があることに気付きました。この問題に対処するため、彼らはマルチブランチネットワークアーキテクチャを提案しました。このネットワークは、浅い層ではパラメータを共有しますが、深い層では2つのブランチに分岐し、それぞれが独立したパラメータセットを持ちます。これらのブランチは、トレーニングの異なる段階でアクティブ化されます。1つのブランチは核形成段階に焦点を当て、もう1つのブランチは成長段階を処理します（下図を参照）。この設計により、予測タスクと生成タスクが互いに影響を与えることなく効果的にトレーニングできるようになります。

キャプション: UniGEMのトレーニングプロセス

UniGEMの推論プロセス

キャプション：UniGEMの分子形成プロセス

UniGEMでは、分子生成は逆拡散プロセスによって原子座標を再構築し、生成された座標に基づいて原子の種類を予測します（図を参照）。特性予測では、ネットワーク入力の時間ステップを0に固定し、特性予測ヘッドを使用します。注目すべきは、この手法では生成タスクと予測タスクのいずれにおいても追加の計算オーバーヘッドは発生せず、推論時間全体はベースラインと同じです。

分子生成タスクにおいて、研究者らはUniGEMと従来の結合生成手法との生成誤差の違いも分析しました。まず、UniGEMでは、結合生成における原子種予測損失による誤差が、原子種ノイズ除去生成損失による誤差よりも小さいことが分かりました。次に、結合生成では、座標生成が原子種予測結果の振動の影響を受け、誤差が増大します。最後に、結合生成手法では、初期分布および離散化誤差も大きくなります。これらの要因が相まって、UniGEMが優れた生成結果を達成する理由を説明しています。

実験結果: 分子生成と特性予測の両方のタスクにおいてベースライン モデルを上回ります。

分子生成：UniGEMがベンチマークモデルを上回る

研究者らはまず、EDMの改良版であるUniGEMと、QM9データセットおよびGEOM-DrugsデータセットにおけるEDMバリアントを比較しました。下図に示すように、UniGEMはほぼすべての評価指標においてベースラインモデルを上回りました。特に、 UniGEMは事前知識に依存せず、追加のオートエンコーダトレーニングも必要としないため、他のEDMバリアントよりも大幅にシンプルでありながら、EDM-BridgeやGeoLDMを上回るパフォーマンスを示しており、UniGEMの優位性が際立っています。

キャプション：UniGEMの生成効果

さまざまな生成アルゴリズムに適応する UniGEM の柔軟性を実証するために、研究者は UniGEM をベイジアン フロー ネットワーク (BFN) に適用し、座標と原子タイプを共同で生成する GeoBFN を上回り、QM9 データセットで最先端 (SOTA) の結果を達成しました。

さらに、研究者らは条件付き生成タスクにおける UniGEM のパフォーマンスをテストし、サンプリング プロセス中にモデルに組み込まれたプロパティ予測モジュールをガイドとして使用することで、条件付き生成モデルを再トレーニングする必要性を回避しました。

分子特性予測: UniGEMはほとんどの事前学習法を上回る

キャプション：UniGEMの特性予測性能

研究者らは、QM9データセットにおける特性予測におけるUniGEMの性能を、テストセットの平均絶対誤差（MAE）を評価指標として評価しました。図に示すように、 UniGEMはゼロから学習したEGNNを大幅に上回り、統合モデリングの有効性を実証しました。驚くべきことに、これらの最先端の事前学習手法は追加の大規模事前学習データセットを用いているにもかかわらず、UniGEMは依然としてほとんどの事前学習手法を凌駕しています。これは、統合された生成・予測モデルの利点を浮き彫りにしており、追加のデータや事前学習手順を必要とせずに、生成プロセス中に分子表現を学習する能力を効果的に活用しています。

結論

UniGEMモデルは、分子生成と特性予測のタスクを統合し、双方の性能を大幅に向上させます。UniGEMの性能向上は、堅実な理論解析と包括的な実験研究によって裏付けられています。この革新的な二段階生成プロセスとそれに対応するモデルは、分子生成フレームワークの開発に新たなパラダイムを提供し、より高度な分子生成フレームワークの開発を促し、より具体的な応用分野における分子生成に有益なものとなると考えています。

この研究はATOM Labが主導しました。チームは、分子事前学習、分子生成、タンパク質構造予測、仮想スクリーニングの分野で多くの研究成果を上げています。今後の展開にご期待ください。

ATOM Lab ホームページへようこそ:
https://atomlab.yanyanlan.com/

著者について:

* 藍燕燕は清華大学知能産業研究所（AIR）の教授です。彼女の研究分野はAI4Science、機械学習、自然言語処理です。

* 馮世坤は清華大学知能産業研究所（AIR）の博士課程に在籍しています。研究分野は、表現学習、生成モデル、AI4Scienceです。

* 倪宇燕は、中国科学院数学・システム科学研究所（AMSS）の博士課程学生です。彼女の研究分野は、生成モデル、表現学習、AI4Science、深層学習理論です。

本論文の主著者である馮世坤博士と倪宇燕博士は現在、求職活動を行っています。ご興味のある方は、お気軽にご連絡ください。

* 馮世坤のメールアドレス: mailto:[email protected]

* Ni Yuyan のメールアドレス: mailto:[email protected]