技術要約: ElemeNet

問題提起

ディープラーニングは分子特性予測に革命をもたらしたが、最先端（SOTA）のモデルは依然として、多様な化学空間に対する汎用性に欠ける、専門化されたコードベースに限定されていることが多い。Chempropのような既存の汎用ツールは、主に2D表現（SMILES文字列）と有機化学に限定されており、立体化学、有機金属における複雑な結合、重元素、または生物学的高分子をネイティブにサポートできていない。さらに、これらのツールは、3D等変アーキテクチャ、電荷やスピン状態への明示的な条件付け、および堅牢な不確実性定量化（UQ）へのサポートを欠いていることが多い。全周期表（原子番号1–100）をサポートし、数百万の分子へとスケールし、機械学習の深い専門知識を必要とせずに較正された不確実性推定を提供する、統一されたソフトウェアフレームワークが切実に求められている。

手法

著者らは、これらのギャップを埋たすべく設計された、分子機械学習のためのモジュール式汎用ソフトウェアパッケージであるElemeNetを紹介する。

1. 表記法と入力

• 元素の範囲: 原子番号1–100（水素からフェルミウム）をサポート。
• 入力フォーマット: 多様な2Dおよび3Dフォーマット（cif, cml, mol, mol2, sdf, SMILES, xyz）を受け入れる。
• グラフ構築:
  • 2D: トポロジカルな連結性から派生。座標は結合の解釈にのみ使用され、破棄される。
  • 3D: 原子座標から構築され、エッジは調整可能なカットオフ半径によって定義され、E(3)等変性を課すためにペアワイズ距離を利用する。
• 特徴量エンジニアリング:
  • ノード特徴量: 連続的な表形式の原子特性（質量、半径、電気陰影度など）と、カテゴリカルな特徴量（族、周期、価電子）の一熱エンコーディングからなる、バランスの取れた初期表現。
  • エッジ特徴量: 結合タイプ、推定結合長、および金属の含有をエンコード。
  • グラフ特徴量: 専用のニューラルネットワークを介して、分子の電荷とスピン多重度を明示的にエンコードし、モデルが電子状態の関数として特性を学習できるようにする。
  • 柔軟性: ユーザー提供の特徴量（例：部分電荷、Wiberg指数）や、暗黙的・明示的な水素処理をサポート。

2. モデルアーキテクチャ

ElemeNetは、モジュール式のエンコーダ・リードアウトアーキテクチャを採用している：

• エンコーダ:
  • 2D GNN: 有向メッセージパッシング（D-MPNN）、グラフ畳み込み、グラフ同型、およびグラフアテンションネットワークをサポート。
  • 3D 等変GNN (EGNN): 不変スカラー特徴量と3D座標のための別々のメッセージパッシング層を使用し、球関数を使用せずにE(3)等変性を保証するために、ガウス基底関数を介して原子間距離を投影する。
• リードアウト:
  • MLP: 標準的な多層パーセプトロン。
  • Transformer: グラフ構成要素の動的な重み付けのためのマルチヘッド自己注意メカニクス。
• ハイパーパラメータ最適化: ネットワークの形状（一定、増加、減少、砂時計型、ピラミッド型）、活性化関数、および畳み込み演算を含む広大な設計空間に対して、ベイズ最適化を行うOptunaを利用する。

3. マルチスケール学習タスク

本フレームワークは、単一のグラフ内での複数のスケールにおける予測をサポートする：

• グラフレベル: 分子特性（例：全エネルギー、毒性）。
• ノードレベル: 原子中心の特性（例：部分電荷）。
• エッジレベル: 結合特性（例：結合長）。
• モイエティレベル（新規）: ノードを中心とした特定のサブグラフ（k-hop近傍）に局在した特性を予測する新しい機能であり、損失計算中に残りの分子をマスクする。これは官能基や活性部位を対象とする。

4. 不確実性定量化 (UQ)

ElemeNetは、2つの階層のUQを実装している：

• ヒューリスティックな尺度: 潜在空間の距離（ユークリッド距離およびコサイン類似度）および分類のための情報エントロピー。これらは低コストであるが、誤差の大きさに対して正式に較正されているわけではない。
• アンサンブルベースの尺度:
  • 回帰: エンコーダトランクを共有しながら、最終的なリードアウト層を$N$回複製する**浅いアンサンブル（shallow ensembles）**を使用する。モデルは、アンサンブルの平均と分散によってパラメータ化されたガウス分布の負の対数尤度（NLL）を最小化する。
  • 分類: **平均確率（MoP）**スキームを用いて浅いアンサンブルを拡張する。活性化関数（シグモイド/ソフトマックス）を適用した後、ロジットを各ヘッド間で平均化し、確率空間におけるアンサンブル標準偏差を不確実性の指標として使用する。

主な結果

著者らは、4つの多様なデータセットにおいて、ElemeNetを文献のベースラインと比較検証した：

1. QM9 (有機化学):

   • 3D EGNNモデルは2D GNNを大幅に上回り、特に双極子モーメントにおいて（MAEが6倍減少）、幾何学的に依存する特性に対する3D情報の必要性を裏付けた。
   • ElemeNetモデル（2Dおよび3D）は、Chemprop v1/v2の性能に匹敵するか、それを上回った。
   • 浅いアンサンブルは、予測精度を大きく損なうことなく、不確実性の較正（MAEとのSpearman順位相関）を改善した。

2. tmQMg (無機/遷移金属化学):

   • EGNNモデルは、HOMO-LUMOギャップ、双極子、および分極率において2D GNNを上回った。
   • 最良のElemeNetモデルは、トランスフォーマー・リードアウトにデフォルトのハイパーパラメータを使用しているにもかかわらず、文献でベンチマークされた既存の3Dアーキテクチャの平均を上回った。
   • マルチタスクEGNNモデルは、3つのターゲットすべてに対して競争力のある性能を同時に達成した。

3. pydentate (配位化学):

   • ElemeNetは、配位原子および配位数予測において、公開されている2Dモデルと同等の精度を達成した。
   • トランスフォーマー・リードアウトは、配位原子予測のバランス精度を向上させた（MLPの92.4%に対し96.1%）。
   • 以前は報告されていなかった、エッジレベル（結合長）およびサブグラフレベル（金属-配位子結合長）の予測に関する新しいベンチマークを確立した。

4. GEMS (生物化学、270万分子):

   • 数百万件のエントリを持つデータセットへのスケーラビリティを実証した。
   • デフォルトのハイパーパラメータで訓練されたモデルは、電子エネルギーおよび双極子の大きさに対して競争力のあるMAEを達成した。
   • 電荷とスピンの明示的なエンコーディングが極めて重要であることが示され、電荷を持つ種を中性として誤ラベル付けすると、予測誤差が著しく増大することが示された。

重要性と主張

論文は、ElemeNetが、周期表全体にわたって高性能なディープラーニングモデルを訓練するための、実用的で統一されたプラットフォームを提供すると主張している。その主な意義は以下の通りである：

• 民主化: 簡潔なコマンドラインインターフェースと自動化されたハイパーパラメータ最適化により、非計算科学研究者の参入障壁を下げる。
• 汎用性: 有機化学を超えて、複雑な結合と電子状態を持つ有機金属、無機物、および生物学的システムをサポートする。
• マルチスケール能力: モイエティレベルの予測を導入し、単一のフレームワーク内でノード、エッジ、グラフ、およびサブグラフのターゲットをサポートする。
• 堅牢性: 浅いアンサンブルを通じて、よく較正された不確実性推定を提供し、これは能動学習や信頼できる展開のための重要な前提条件である。
• 性能: 最小限のユーザー努力と適度な計算予算で、多様な化学ドメインにわたってSOTAまたは競争力のある性能を達成する。

著者らは、ElemeNetが、化学および物理科学のあらゆる分野における分子の発見と最適化のために、現代的なディープラーニング手法を活用することを研究者に可能にする、と結論付けている。