OmniMol: Transferring Particle Physics Knowledge to Molecular Dynamics with Point-Edge Transformers

本論文は、Point-Edge トランスフォーマー・アーキテクチャと高エネルギー物理学からの知識転移を活用し、最小限の微調整と極めて高速な推論で卓越した性能を達成する、小分子向けの最先端の機械学習型原子間ポテンシャル「OmniMol」を紹介する。

要約

2 つの非常に異なる世界を想像してみてください。1 つは素粒子物理学の世界（科学者たちが原子を衝突させて飛び出すものを観察する、混沌とした高速の世界）であり、もう 1 つは分子化学の世界（原子が結合して医薬品、材料、そして生命を形成する、複雑で粘着質な世界）です。

長らく、これらの分野の科学者たちはそれぞれの世界を理解するために全く異なるツールを用いてきました。しかし、この論文では、著者たちがOmniMolという新しいツールを紹介しています。これは、すでに構築した「基盤モデル」を用いて、素粒子物理学の専門家たちに化学を理解させることを試みるものです。

以下に、彼らがどのように行い、何を発見したのかを簡潔にまとめます。

1. 「巨匠シェフ」の比喩

Omnilearnedと呼ばれる元のモデルを、何年も粒子ジェットを用いて料理を続けてきた巨匠シェフと想像してください。

• 材料: 素粒子物理学において、「ジェット」とは衝突から飛び出す陽子や中性子などの亜原子粒子の噴霧のことです。
• 技能: このシェフは、これらの噴霧の中のパターンを認識することを学びました。粒子がどのように相互作用し、どのようにクラスター化し、次に何が起きるかを予測する方法を知っています。彼らは10 億もの異なる粒子噴霧で訓練されました。

ここで、著者たちは問いかけます：この同じシェフは、分子の料理を作れるでしょうか？

• 新しい材料: 亜原子粒子の代わりに、「材料」は分子内の原子（炭素、酸素、水素など）です。
• 課題: 原子は亜原子粒子とは異なる振る舞いをしますが、類似した構造を持っています。つまり、特定のタイプを持つ空間上の点に過ぎないのです。

2. 「万能翻訳機」（アーキテクチャ）

これを機能させるために、彼らはゼロから新しいシェフを構築したわけではありません。既存の「巨匠シェフ」（Omnilearned）を持ち出し、新しい道具を与えました。

• ポイント・エッジ・トランスフォーマー（PET）: シェフが皿に盛られた料理を見ていると想像してください。このツールは、1 つの材料だけを順に見るのではなく、すべての材料を一度に見て、それぞれの材料が他のすべての材料とどのように関連しているかを把握できるようにします。
• 「物理的バイアス」: これが秘密のソースです。モデルには「ねえ、これらの 2 つの粒子/原子は互いに近いから、お互いにより注意を払うべきだ」と伝える組み込みの「ルールブック」があります。これにより、モデルはノイズに混乱することなく、最も重要な関係性に焦点を当てることができます。

3. 実験：ファインチューニング

著者たちは、この素粒子で訓練されたモデルを、oMol（数百万の分子のコレクション）というデータセットを用いて化学の「集中講座」にかけました。

• 目標: モデルを**機械学習による原子間ポテンシャル（MLIP）**として機能させることでした。平易に言えば、このモデルは任意の原子のグループに対して 2 つのことを予測する必要があります。
  1. エネルギー: それらを結びつけている「のり」の強さはどれくらいか？
  2. 力: 1 つの原子を押すと、どれくらい強く押し返してくるか？

4. 結果：高速で驚くほど優れている

この論文は、いくつかの興奮すべき発見をもたらしました。

• 「少ショット」のスーパーパワー: 通常、コンピュータに化学を教えるには膨大な量のデータが必要です。しかし、OmniMol は素粒子物理学の「知識」から出発したため、化学を非常に迅速に学びました。比較的小さな量の新しいデータ（例えば 10 万個の分子）であっても、数百万で訓練されたモデルとほぼ同等のパフォーマンスを発揮しました。これは、味と熱の基礎を理解している巨匠シェフが、わずか数種類のレシピで新しい料理を習得できるようなものです。
• 速度: OmniMol は信じられないほど高速です。他のモデルが分子の動きを計算するのに長い時間を要するのに対し、OmniMol は瞬時に行います。著者たちは、計算時間 1 時間あたり、OmniMol は競合するモデルの3 倍多くの分子をシミュレーションできると指摘しています。
• トレードオフ: 彼らが膨大な量のデータ（数百万の分子）を持っていた場合、素粒子物理学の知識から始めることの利点は少し薄れました。これは、「素粒子物理学の知識」が強力なスタートダッシュとして機能することを示唆していますが、ゼロからモデルを訓練するのに十分な時間とデータがあれば、そのスタートダッシュの重要性は低くなることを意味します。

5. 全体像

この論文は、OmniMolが、ある科学分野（素粒子物理学）向けに構築された「基盤モデル」が、全く異なる分野（化学）へ成功裏に転用された最初の事例であると結論付けています。

彼らは、ある分野で空間内の点がどのように相互作用するかを理解するスマートなモデルがあれば、それを別の分野における空間内の点の相互作用を理解するように適応させることができ、時間と計算資源を節約できることを証明しました。

要約すると: 著者たちは、高エネルギーの粒子衝突で訓練された超スマートな AI を、粒子ではなく原子を理解するようにその脳を調整し、それが特にデータが不足している場合に、分子の挙動を予測するための超高速かつ高精度なツールとなったことを発見しました。

技術概要

技術的サマリー：OmniMol

問題定義

機械学習（ML）は、特に素粒子物理学と分子化学において、複雑な物理系の表現とシミュレーションを変革してきた。これらの分野はエネルギー規模において大きく異なるが、共通する基本的なデータ構造を持っている：位相空間における可変サイズの粒子（または原子）の集合であり、実質的に構造化された点雲を形成している。

対処される主要な課題は、効率的な機械学習型原子間ポテンシャル（MLIP）の開発である。密度汎関数理論（DFT）などの従来手法は計算コストが高く、大規模かつ長期的な分子動力学（MD）シミュレーションを制限している。MLIP は、このコストの断片でポテンシャルエネルギー曲面と力を近似することを目指している。しかし、堅牢な MLIP を訓練するには、通常、膨大なデータセットと莫大な計算資源が必要となる。本論文は、素粒子物理学（特に粒子ジェット）の点雲で事前学習されたファウンデーションモデルを分子動力学へ転用することで、最適化を加速し、低データ領域における精度を向上させる可能性があると仮説を立てている。

手法

アーキテクチャ：Point-Edge Transformer (PET)

OmniMol は、高エネルギー物理学（HEP）における粒子ジェットの分類と生成のために設計されたファウンデーションモデル「Omnilearned」を適応させることで構築されている。中核となるアーキテクチャは、$k$-近傍に対する局所アテンションとグローバルな全対全トランスフォーマーブロックを結合した**Point-Edge Transformer (PET)**である。

主要なアーキテクチャ構成要素は以下の通り：

1. 入力埋め込み: 原子は、位置情報（$\vec{r}$）、離散的な原子番号（$Z$）、および追加の特徴量（電荷、スピン）を組み合わせたトークン空間に埋め込まれる。
2. 局所アテンションブロック: 各原子について、$K$-近傍（分子の場合は $K=15$、ジェットの場合は $K=10$）を使用して局所近傍を構築する。距離項、距離の逆べき乗、および原子埋め込みの学習済み関数を含む、ペアごとの物理的特徴量が計算される。これらは小さな局所トランスフォーマーによって処理され、局所埋め込みベクトルが作成される。
3. 相互作用バイアスを伴うグローバルアテンション: グローバル自己アテンション機構には、ペアごとの物理的特徴量から導出された明示的なバイアスが組み込まれる。アテンションロジットは $A^*_{ij} = A_{ij} + B_{ij}$ として修正され、ここで $B_{ij}$ は MLP 埋め込みされたバイアス項である。この「相互作用行列アテンションバイアス」は、ペアごとの物理的事前知識をトランスフォーマーに直接注入し、表現力を犠牲にすることなく、ネットワークを物理的に意味のある近傍へと誘導する。
4. 出力ヘッド: Omnilearned の生成ヘッドは、2 つのタスクのために再利用される：
   • 力予測: 原子ごとの力を予測する置換等変換ヘッド。
   • エネルギー予測: 原子ごとのエネルギー補正を予測するヘッドであり、これらは総分子エネルギーを生成するために合計され、広大性の事前知識を保持する。

不変性と保存則の制約

物理的制約を満たすために、著者は 2 つの要件に対処する：

• エネルギー保存: 力は直接予測されるのではなく、エネルギー出力の逆伝播を通じて計算される（$\vec{F}_i = \nabla_{\vec{r}_i} E$）。これにより厳密なエネルギー保存が保証されるが、訓練中の計算コスト（二重逆伝播が必要）が増加する。したがって、この制約は「小」モデル変種にのみ適用される。
• 回転等変換性: 標準的なアーキテクチャは、生の座標差が MLP に入力されるため、本質的に等変換性を持たない。これを補うため、著者は「等変換かつ保存則を満たす」変種を導入する。このバージョンは、ペアごとの特徴量から直接の座標差項を除去し、代わりに隣接原子間のベクトルが形成する角度の余弦（コサイン）といった角度情報を局所ブロックに組み込む。この修正により、座標項を除去することに伴う性能低下を大幅に回復させつつ、等変換性を維持する。

訓練と微調整戦略

モデルはoMolデータセット（具体的には oMol-25、oMol-4M、oMol-100M、および oMol-140M サブセット）で微調整される。2 つの戦略が検討される：

1. LoRA（低ランク適応）: 事前学習された PET バックボーンの重みは凍結される。低ランクアダプターは、トランスフォーマー本体の行列（$W_Q, W_K, W_V, W_O, W_{MLP}$）に対してのみ導入され、分子入力エンコーダ、バイアス MLP、およびタスクヘッドの訓練と併せて行われる。「埋め込み適応」レイヤーも、学習済み埋め込みを変更するために追加される。
2. フル微調整: 本体と入力エンコーダのすべての重みが凍結解除され訓練され、タスクヘッドはゼロから訓練される。

訓練目的関数は、エネルギーと力の平均絶対誤差（MAE）の和を最小化するものであり、力はより重み付けされる（$\lambda_F = 10$）。

主要な結果

oMol における性能

• フル訓練: 大規模データセット（oMol-4M および oMol-100M/140M）で訓練された場合、OmniMol は最先端の MLIP と競合する性能を達成する。例えば、oMol-140M において、OmniMol-large モデルは原子あたりエネルギー MAE で 1.04 meV、原子あたり力 MAE で 13.59 meV/Å を達成する。
• 低データ領域: 訓練データが制限されている場合に、最も顕著な改善が観察される。10 万分子のみで微調整した場合、または oMol-4M 上で非常に少ないエポック（2 パス）で微調整した場合、事前学習済み OmniMol 変種は、ゼロから訓練されたモデルを大幅に上回る。
  • 10 万のサブセットにおいて、事前学習により、中規模モデルのエネルギー MAE が最大29.4%、力 MAE が**26.9%**改善された。
  • oMol-4M 上で 2 エポックのみで訓練した場合、中規模モデルは、事前学習を行わない対応モデルと比較して、エネルギー MAE で54.6%、力 MAE で**56.9%**の改善を示した。
• 等変換/保存則変種: 等変換かつ保存則を満たすモデル変種は、低データ領域で（特に力において）著しく改善された性能を示すが、この利点はデータセットサイズが増加するにつれて減少する。

スケーリングと推論速度

• スケーリング: OmniMol はモデルサイズに対してクリーンなべき乗則スケーリングに従い、10 億パラメータまで飽和の兆候は見られず、トランスフォーマーベースの MLIP に関する最近の知見と一致している。
• 推論速度: 大規模なパラメータ数にもかかわらず、OmniMol はトランスフォーマー向けのハードウェア最適化により、極めて高速な推論速度を示す。A100 GPU 上で約 100 原子のシステムにおいて、OmniMol-medium は、同等のグラフニューラルネットワーク（GNN）ベースライン（eSEN-md-d および AllScAIP-md）と比較して約 3 倍高速であり、かつ競合する精度を維持している（AllScAIP-md よりもエネルギー誤差が約 0.7 meV/atom 高いのみ）。

意義と主張

本論文は、科学的な点雲ファウンデーションモデルにおける分野横断的な転用の最初の実証を提示すると主張している。高エネルギー物理学の粒子ジェットで事前学習されたモデルを分子動力学へ適応させることで、著者は以下のことを実証している：

1. 分野横断的転用の実現可能性: 素粒子物理学向けに構築されたファウンデーションモデルが分子化学へ効果的に転用可能であり、点雲構造の基礎が、広範に異なる物理スケールにわたって学習可能な特徴を共有していることを示唆している。
2. 帰納的バイアスによる学習の加速: 事前学習は強力な帰納的バイアスとして機能する。データが不足している場合に等変換性が役立つのと同様に、事前学習という「苦い教訓」は、訓練データが制限されている場合の迅速な最適化と精度向上を可能にする。
3. 効率性: アーキテクチャの転用により、極めて高速な推論速度が可能となり、小分子創薬など設計空間の迅速な探索を必要とする応用にとって重要である。

著者は、本研究が MLIP に焦点を当てているが、点雲ファウンデーションモデルに関する教訓は、相互作用する物体の順序不定な集合として記述されるシステムを扱う科学分野全体で広く有用である可能性があると結論付けている。すべての領域ですべての既存手法に対して普遍的に優位であると主張するものではないが、低データシナリオと推論速度における特定の利点を強調している。