TriForces: Augmenting Atomistic GNNs for Transferable Representations

TriForces は、自己教師あり学習と分離された構成および構造表現を組み合わせるモデル非依存の 3 ストリームフレームワークであり、機械学習ポテンシャルのための原子グラフニューラルネットワークの転移性とデータ効率を大幅に向上させる。

要約

ロボットシェフに料理を教えることを想像してください。

問題：「万能型」シェフ
現在、科学者たちは強力なAIモデル（MLIPsと呼ばれる）を用いて、原子の挙動を予測しています。例えば、物質が持つエネルギー量や、原子を動かす難易度などです。これらのモデルは、スーパーコンピュータ（DFT）から得られた膨大なデータで訓練されています。

しかし、これらのモデルには欠点があります。それは、特定の料理の正確な味を暗記したものの、なぜその味になったのかを忘れたシェフのようなものです。少し異なる料理（新しい種類の物質）を頼むと、彼らは苦労します。彼らは材料（組成）と鍋の形（構造）を混同してしまいます。材料を変えると形について混乱し、その逆もまた然りです。これにより、特に教えるためのデータがあまりない場合、新しいタスクを素早く学習することが苦手になります。

解決策：TriForces（3つのストリームを持つキッチン）
著者たちは、これらのAIシェフを構築する新しい方法としてTriForcesを紹介しています。一度にすべてを覚えようとする巨大な脳ではなく、脳を3つの専門的な「ストリーム」または部門に分けます。

1. 材料ストリーム（組成）： この部門は、鍋の中に何が入っているか（例えば、「水素が2つ、酸素が1つある」）だけを見ます。形は完全に無視します。化学を学びます。
2. 形状ストリーム（構造）： この部門は、原子が空間内でどのように配置されているか（例えば、「三角形になっている」）だけを見ます。原子が実際には何であるかは無視します。幾何学を学びます。
3. 相互作用ストリーム： これはメインのシェフで、材料部門と形状部門からのメモを受け取り、それらを組み合わせて最終結果（エネルギーや力）を予測します。

秘密のソース：自己教師あり学習
モデルが特定の性質を予測するように求められる前に、著者たちは「自己教師あり学習」と呼ばれるゲームを用いてこれを訓練します。これは、AIが以下のことを行う練習セッションのようなものです。

• ノイズ除去： 分子の少し壊れたりノイズの混じった画像を見て、それを修正する。
• マスキング： 材料を隠して、隣接するものに基づいてそれが何だったかを推測する。
• マッチング： 同じ分子のわずかに異なる2つのバージョンを見て、それらが同じものであると認識する。

この訓練により、AIは知識を整理して組織化することを強制されます。材料と形状を混ぜてしまうのではなく、材料は1つのフォルダに、形状は別のフォルダに収めることを学びます。

なぜこれが重要なのか（結果）
この論文は、この新しい「3ストリーム」キッチンが、古い「1つの脳」キッチンよりもはるかに優れていることを示しています。

• 高速な学習： 新しいデータ（数百万ではなく2万の例など）を少量与えられた場合、TriForcesははるかに速く学習し、間違いを減らします。まるで、1000回調理する必要があるのではなく、1度味見しただけで新しいレシピを覚えるシェフのようです。
• 優れた記憶： AIは学んだことを忘れません。混乱することなく、ある種類の物質から別の種類の物質へ知識を転移できます。
• 検索可能な知識： AIが「材料」と「形状」を分離して保持しているため、同じ見た目だが異なる材料を持つ物質、あるいは同じ材料だが異なる形状を持つ物質を見つけるように依頼できます。古いモデルは知識が混雑しすぎていたため、これをできませんでした。

まとめ
TriForcesは、原子を理解するという複雑な仕事を、材料を知る、形状を知る、そしてそれらがどのように相互作用するかを知るという3つの単純な仕事に分解するフレームワークです。AIにこれらの仕事を分離して維持させ、「推測ゲーム」（自己教師あり学習）で練習させることで、このモデルは新しい物質を発見するための、はるかに柔軟で効率的かつ正確なツールとなります。

著者たちはコードと事前学習済みモデルを公開しており、他の科学者たちがこの「3ストリームキッチン」を用いて、材料科学のためのより良いAIを構築できるようにしています。

技術概要

技術的サマリー：TriForces

問題提起

幾何学的グラフニューラルネットワーク（GNN）に基づく機械学習原子間ポテンシャル（MLIPs）は、大規模な密度汎関数理論（DFT）データセットで訓練された場合、原子論的性質の予測において高い精度を達成しています。しかし、その実用的有用性は、主に 2 つの制限によって阻害されています：

1. 一貫性のない転移性（Transferability）： MLIPs は、特に小規模で高価なタスク固有のデータセットでの微調整において、新しい化学種やタスクへ効果的に転移できず、失敗することが多いです。事前学習済みモデルは、結晶系の同定や主要元素の特定などの単純な診断タスクで頻繁に苦戦し、ドメイン間での「破滅的忘却」や不安定な転移性能を示します。
2. 表現の再利用性の低さ： 現在の MLIPs は、一般的な再利用ではなく、特定の予測ターゲット（エネルギーと力）に対して表現を最適化しています。その結果、これらの表現は組成（化学）と構造（幾何学）が絡み合っており、探索的分析、最近傍検索、または分解された類似性検索には不適切です。標準的な教師あり目的関数は、回帰に十分な表現を促しますが、アクセス可能な組成および構造情報を保持するように組織化された表現を促すものではありません。

手法：TriForces フレームワーク

著者らは、既存の幾何学的 GNN を 3 つのストリームアーキテクチャと多目的自己教師あり学習（SSL）で拡張する、モデル非依存のフレームワーク「TriForces」を提案します。中核的な革新は、原子論的表現を 3 つの明確な成分に明示的に因数分解することにあります：

1. 3 つのストリームアーキテクチャ

単一の潜在ベクトルの代わりに、TriForces はノードレベルの表現 $h_i$ を 3 つの連結されたストリームに分解します：

• 組成ストリーム（$h^{comp}$）： 座標なしで化学情報をエンコードします。これは、カウント重み付きアテンションを持つ Transformer を用いて、固有の原子元素のセットとその化学量論的カウントを処理します。このストリームは、幾何学に依存せずにシステムサイズとエネルギー規模をエンコードするために、絶対的な元素カウントを保持します。
• 構造ストリーム（$h^{struct}$）： 元素の同一性（タイプ非依存）なしで幾何学をエンコードします。これは、SOAP（原子位置の平滑重なり）に着想を得た回転不変の局所記述子を構築し、半径基底関数、球面調和関数、およびマルチスケールのカットオフを利用します。このストリームは、不変なメッセージパッシングによって強化された、再利用可能な幾何学的モチーフとトポロジカルなパターンを捉えます。
• 相互作用ストリーム（$h^{int}$）： 組成と幾何学の間の結合を捉え、元のアーキテクチャの表現力を保持する標準的なベース幾何学的 GNN（例：MACE、eSEN、Orb-v3）です。

2. 自己教師あり事前学習

潜在空間を組織化し、転移性を向上させるために、TriForces は確率的拡張（位置ノイズ、元素マスク、グラフ変異、回転）を用いた多目的 SSL 事前学習戦略を採用します。このフレームワークは、3 つの相補的な目的を組み合わせます：

• 非再構成（LeJEPA）： 同じ構造の 2 つの拡張ビューからの埋め込みを、ノードレベルおよびグラフレベルの両方で整合させます。これにより、拡張に対する不変性が強制され、ストップグラディエントやモーメントエンコーダを必要とせずにグローバルな潜在空間が組織化されます。
• ノイズ除去： 模型にノイズのある入力からクリーンな原子位置を復元させるように訓練します。これにより幾何学的表現が安定化され、暗黙的に回転拡張が提供されます。
• マスク化： 周囲の幾何学と組成に基づいて、マスクされた原子タイプを予測します。これにより、模型は組成パターンと文脈を学習することが促されます。

最終的な事前学習損失は、これら 3 つの目的の加重和です。

主要な貢献

1. アーキテクチャの分解： 組成、構造、相互作用を明示的に分離する 3 つのストリーム設計により、両方の因子が絡み合うのではなく、設計上保持されることを保証します。
2. ハイブリッド事前学習戦略： 埋め込み空間をより良い下流転移のために構造化するために、再構成ベースの目的（ノイズ除去、マスク化）と潜在予測学習（LeJEPA）を組み合わせた自己教師ありアプローチ。
3. 解釈可能な検索： 組成、構造、または結合埋め込み空間においてターゲットとした類似性検索を実行する能力により、特定の基準（例：化学のみ、または構造のみ）に基づく材料比較を可能にします。
4. 実証的検証： 複数のアーキテクチャ（Orb-v3、eSEN、MACE）とベンチマーク（OMat24、MatBench、QM9）にわたる広範な実験により、データ効率、転移性能、および表現の質の向上を実証しました。

結果

• 転移性能（OMat24）： 限られたデータ領域において、TriForces はベースラインを大幅に上回ります。2 万サンプルにおいて、ベースモデルと比較してエネルギーの平均絶対誤差（MAE）を 57% 削減します。また、すべてのサンプルサイズで力の MAE を改善し、応力誤差を削減します。
• データ効率： TriForces は、すべてのデータセットサイズ（2 万から 200 万サンプル）でより低い誤差を達成し、特に低データ設定で最も顕著な改善が見られます。
• ベンチマーク性能：
  • MatBench： TriForces のバリエーションは、8 タスクのうち 6 つで全体的に最良の結果を達成し、自己教師ありおよび DFT ラベル付き事前学習ベースラインの両方を上回ります。例えば、フォノンの MAE は 57.8 から 19.5 cm$^{-1}$ に改善されました。
  • MatBench Discovery： TriForces eSEN-sm は、60% 少ないパラメータを使用し、最大 5 倍高速に訓練しながら、はるかに大きな eSEN-30M-OAM モデルと同等のエネルギー MAE を達成します。
  • QM9： 多様な化学入力（バルク＋分子）での事前学習は、バルクのみまたは SSL なしベースラインと比較して、一貫して MAE を低下させます。
• 表現の質： 凍結された埋め込みに対する線形プロービングは、TriForces が標準的な MLIPs が失う基本的な情報（結晶系、主要元素、配位数）を保持することを示しています。TriForces は結晶系および主要元素の分類で 96–100% の精度を達成する一方、ベースラインは苦戦します（55–73%）。
• 検索： このフレームワークは効果的な k-NN 検索を可能にし、組成ストリームは元素セットのリコールに優れ、構造ストリームは空間群のリコールに優れています。これは単一ストリームモデルには欠けている能力です。

意義と主張

本論文は、TriForces を単なる自己教師あり手法ではなく、その表現が SSL によってさらに強化されるアーキテクチャフレームワークとして位置づけています。

• 領域依存性： 著者らは、ストリームの因数分解が大規模な教師あり設定において支配的な利益をもたらす一方、SSL は低データ転移、表現の組織化、および検索タスクにおいて最も価値があると主張しています。
• 脱結合： 組成と構造を分離することで、TriForces は現在の基盤モデルの「脆弱な転移」および「再利用が困難」という問題に対処します。これにより、模型は予測に加えて、分析（例：プロービング、検索）のために組織化された表現を学習することが可能になります。
• 実用性： このフレームワークはモデル非依存であり、プラグイン可能であるため、既存または新しい原子論的アーキテクチャに即座に適用可能です。著者らは、下流の原子論的モデリングでの再利用を促進するために、事前学習済みチェックポイントとコードを公開しています。

この研究は、将来の原子論的基盤モデルが、化学系の明確な物理的要因を明示的に保持する因数分解された表現へと移行し、単一ストリームの予測アーキテクチャを超えていくべきであることを示唆しています。