Information Routing in Atomistic Foundation Models: How Task Alignment and Equivariance Shape Linear Disentanglement

この論文は、Compositional Probe Decomposition (CPD) という手法を用いて、分子モデルにおける幾何学的・組成情報の線形分離性を定量化し、タスクの整合性、データ多様性、対称性に基づく情報経路が、モデルの表現における線形分離の度合いを決定づけることを明らかにしました。

要約

🧪 核心となる問題：「材料」と「形」の区別

分子を予測する AI は、2 つの重要な情報を持っています。

1. 材料（組成）： 水素が何個、炭素が何個入っているか（例：ケーキの材料）。
2. 形（幾何学）： それらがどう並んでいるか（例：ケーキのデコレーションや形）。

多くの AI は、この 2 つがごちゃ混ぜになっています。「材料さえ分かれば、形は大体想像できる」という癖があるため、「形そのもの」を正確に読み取るのが苦手な場合があります。

この論文は、**「AI が材料の情報を一度『消去』したあと、どれだけ『形』の情報を残して持っているか」**を測る新しい方法（CPD という名前）を開発しました。

---

🕵️‍♂️ 発見された 3 つのルール

10 種類の異なる AI モデルをテストしたところ、「形」の情報がどれだけ整理されているかに、驚くほどの差があることが分かりました。この差を生むのは、AI の「設計図（アーキテクチャ）」よりも、以下の 3 つの要因でした。

1. 勉強した科目が合っているか？（タスクアライメント）⭐最重要

これが一番大きな要因です。

• 例え話： 「ケーキの味（形）」を教えるために、AI に「材料の重さ（組成）」だけを勉強させた場合、AI は「形」の情報を整理して覚えられません。
• 結果： 「分子の形」に敏感な課題（電子の動きなど）を勉強させた AI は、形の情報を読み取るのが非常に得意でした。逆に、「エネルギー（重さ）」だけを勉強させた AI は、たとえ高性能な設計図を持っていても、形の情報を読み取るのが下手でした。
• 教訓： AI を使うときは、「何のために勉強させたか（学習目的）」が、その後の使いやすさを決めます。

2. 設計図の「対称性」は魔法の杖ではない（等変性）

最近の AI は、物理法則（回転しても変わらない性質）を守るように設計されています（等変性）。

• 例え話： 「回転しても壊れない頑丈な箱」を作る設計図（等変性）を持っているからといって、中身が整理されているとは限りません。
• 結果： 設計図が素晴らしい AI でも、間違った課題（重さの勉強）をさせていたら、中身はぐちゃぐちゃでした。逆に、シンプルな設計図でも、正しい課題（形の勉強）をさせていれば、中身はきれいに整理されていました。
• 教訓： 設計図よりも「何をするための AI か」の方が重要です。

3. 多様な経験が助けになる（データの多様性）

• 例え話： 1 つの料理しか作らない職人よりも、世界中の料理を学んだ職人の方が、新しい料理の「形」を直感的に理解しやすいかもしれません。
• 結果： 非常に多様な分子データで事前学習した AI は、特定の課題を勉強していなくても、ある程度は形を整理して覚えていました。ただし、これは「完璧な整理」には届かず、正しい課題を勉強した AI には勝てませんでした。

---

🧩 面白い発見：AI の「情報配送ルート」

特に「MACE」という AI について詳しく調べたところ、面白いことが分かりました。
この AI は、情報を**「スカラー（大きさ）」と「ベクトル（向き）」**という 2 つの異なるチャンネル（通路）に分けて運んでいます。

• スカラーの通路： 「分子の隙間の広さ」などの情報を運ぶ。
• ベクトルの通路： 「分子の電気の向き（双極子モーメント）」などの情報を運ぶ。

まるで、**「荷物の種類に合わせて、トラックとバイクを使い分けている」**ような、非常に効率的な整理方法です。しかし、別の AI（ViSNet）にはこの整理方法がなく、すべての荷物を同じトラックに放り込んでいました。

---

⚠️ 重要な注意点：「非线性」の罠

この研究で最も重要な方法論的な発見があります。
AI の中身を調べる際、「複雑な計算ができる探偵（非線形プローブ）」を使うと、嘘の結果が出ます。

• 例え話： 「材料」の情報を消したはずの箱から、探偵が「材料の匂い」を嗅ぎ取って「材料が入っている！」と報告してしまうようなものです。
• 結果： 複雑な探偵は、消したはずの情報を「推測」して復活させてしまい、AI が実は整理できていないのに「できている」と誤解させます。
• 解決策： この研究では、**「単純な探偵（線形プローブ）」**を使うことで、嘘のない真実を測ることができました。

---

💡 私たちへのメッセージ

この研究が私たちに教えてくれることはシンプルです。

1. AI を選ぶときは「何のために作られたか」を見ること。 設計図の良さよりも、学習させた課題が目的に合っているかが重要です。
2. AI の「脳」は、人間が思っている以上に整理されていないかもしれない。 複雑な AI でも、必要な情報が取り出しやすいように整理されているとは限りません。
3. 新しい分析方法（CPD）を使えば、AI の中身がどう整理されているかを正しく見極められる。 これにより、薬の発見や新素材の開発で、より適切な AI を選べるようになります。

つまり、**「AI の性能は、設計図の複雑さではなく、何を学ばせたかで決まる」**というのが、この論文が伝える最大のメッセージです。

技術概要

1. 研究の背景と問題提起

分子物性予測モデル（MACE, SchNet, PaiNN など）は、密度汎関数理論（DFT）に近い精度でエネルギーや力を予測できるようになりました。しかし、これらのモデルの中間表現が「分子が何でできているか（組成）」と「原子がどのように配置されているか（幾何構造）」をどのようにエンコードしているかは不明確です。

• 問題点: 従来のプロービング（内部表現に対する予測タスク）では、組成と幾何構造が相関しているため、モデルが真に幾何情報を学習しているのか、それとも組成のみに依存しているのかを区別できません。
• 既存手法の限界: 残留表現（Residuals）に対して非線形なプローブ（勾配ブースティング木など）を使用すると、投影されたはずの組成情報が非線形相互作用を通じて再構築され、過大評価されたスコア（偽陽性）が得られることが発見されました。

2. 提案手法：Compositional Probe Decomposition (CPD)

著者は、組成信号を除去し、残りの幾何情報が線形にアクセス可能かを測定する新しい手法CPDを提案しました。

• プロトコル:
  1. 組成特徴量の定義: 分子ごとの元素比率と原子数からなるベクトル $Z$ を作成。
  2. 線形投影（除去）: 凍結されたモデルの表現 $X$ に対して、$Z$ に対する最小二乗法（OLS）を適用し、組成成分 $\hat{X}_{comp} = Z\hat{\beta}$ を計算。
  3. 残留表現の取得: 幾何情報を含むと仮定される残留表現 $X_{geom} = X - \hat{X}_{comp}$ を算出（クロスバリデーションの各フォールドごとに独立して投影係数を算出することで、情報リークを防ぐ）。
  4. プロービング: 残留表現 $X_{geom}$ に対して**リッジ回帰（線形プローブ）**を適用し、ターゲット物性（例：HOMO-LUMO ギャップ）の予測精度（$R^2_{geom}$）を測定。
• 重要な知見（非線形プローブの危険性）: 平均原子質量（純粋に組成依存のターゲット）に対して、残留表現から勾配ブースティング木（GBT）を使用すると、$R^2$ が 0.68〜0.95 と高く評価されてしまいます。これは GBT が非線形性を利用して投影された組成情報を再構築しているためであり、残留表現の解析には線形プローブのみが信頼できることを示しています。

3. 実験設定

• データセット: QM9（最大 9 個の重原子を持つ有機分子）および Materials Project（結晶）。
• モデル: 5 つのアーキテクチャファミリーに属する 10 種類のモデル（MACE, PaiNN, ViSNet, SchNet, DimeNet++, ANI-2x など）。
• 変数: 等変性（Equivariance）の有無、トレーニング目的（HOMO-LUMO ギャップ予測 vs エネルギー予測）、トレーニングデータの多様性（QM9 のみ vs 大規模事前学習）。

4. 主要な結果

4.1 線形アクセス可能性の勾配（Linear Accessibility Gradient）

組成を除去した後の幾何情報の線形アクセス可能性（$R^2_{geom}$）はモデル間で6.6 倍もの差がありました（0.081 〜 0.533）。この勾配を説明する 3 つの要因は以下の通りです。

1. タスクの整合性（Task Alignment）の支配的役割:

   • トレーニング目的とプロービング対象が一致するモデル（例：HOMO-LUMO ギャップ予測でトレーニングされたモデル）は、一致しないモデル（エネルギー予測でトレーニングされたモデル）よりも $R^2_{geom}$ が約 0.25 高い値を示しました。
   • アーキテクチャが同じでも（例：PaiNN）、トレーニング目的をエネルギーから HOMO-LUMO ギャップに変えるだけで、性能が劇的に向上しました（$\Delta R^2 \approx 0.22$）。
   • 結論: 幾何情報の線形アクセス可能性を決定する最も重要な要因はアーキテクチャではなく、トレーニング目的です。

2. 等変性（Equivariance）の条件付き効果:

   • 等変性アーキテクチャ（MACE など）自体が自動的に幾何情報を線形に分離するわけではありません。
   • エネルギー予測でトレーニングされた等変性モデル（MACE QM9）は、不変（Invariant）モデル（SchNet）よりも幾何情報のアクセス可能性が低くなりました。
   • 等変性が有効になるのは、タスク整合性のあるトレーニング目的と組み合わさった場合のみです。

3. データの多様性による部分的な補償:

   • 大規模な多様なデータ（MPTraj など）で事前学習されたモデルは、QM9 のみでトレーニングされたモデルよりも高い幾何アクセス可能性を示しましたが、タスク整合性のあるモデルには及びませんでした。

4.2 情報ルーティングと等変性チャネル（MACE の分析）

MACE の内部構造を解析したところ、不可約表現（Irreducible Representations）ごとの情報ルーティングが確認されました。

• スカラーチャネル（L=0）: 電子的なスカラー量（HOMO-LUMO ギャップなど）を優先的にエンコード。
• ベクトルチャネル（L=1）: 双極子モーメントなどのベクトル量を優先的にエンコード。
• 対照的な ViSNet: ViSNet もスカラー/ベクトルストリームを持っていますが、最終的な表現ではベクトルチャネルがほとんど情報を担っておらず、この「チャネルごとの特化」は MACE のようなテンソル積ベースのアーキテクチャに特有の現象であることが示唆されました。

4.3 構造的異性体ベンチマークによる検証

組成が同一で構造のみが異なる「異性体」ペアに対して、組成成分と幾何残留成分を分類タスクに適用しました。

• 組成成分は偶然レベル（52.5%）の精度しか出ませんでした。
• 幾何残留成分は、PaiNN で 94.6% の高い精度を達成し、CPD が組成を正しく除去しつつ、構造情報を保持していることを実証しました。

5. 貢献と意義

1. 手法論的貢献:

   • CPD の提案: 組成と幾何を分離するための検証済みのプロービング手法。
   • 非線形プローブの警告: 残留表現に対する非線形プローブが組成情報を再構築してスコアを過大評価することを示し、線形プローブの使用を推奨しました。

2. 理論的知見:

   • タスク整合性の重要性: 分子モデルにおいて、幾何情報の線形分離性を高めるためには、アーキテクチャの複雑さよりも、トレーニング目的の選択が重要であることを実証しました。
   • 情報の線形アクセス可能性: 幾何情報が存在するかどうかではなく、「線形に読み取り可能な形で組織化されているか」が、少サンプル転移学習や解釈可能性において重要であることを示しました。

3. 実用的意義:

   • 下流タスク（触媒設計、創薬など）において、幾何構造に敏感な物性を予測する場合、同じアーキテクチャでも「幾何構造に敏感な目的でトレーニングされたモデル」を選択することが、事前学習済みモデルの選定において最も重要であることを示唆しています。
   • 大規模な事前学習はタスクミスマッチを部分的に補償しますが、完全には解決しません。

結論

この研究は、分子基礎モデルの表現学習において、「何のために学習したか（タスク）」が「どのように表現が組織化されるか（幾何情報の線形分離性）」を支配するという重要な原則を明らかにしました。また、非線形プローブの危険性を指摘し、分子表現の解釈可能性を高めるための厳密な評価手法（CPD）を確立しました。