Exascale Multi-Task Graph Foundation Models for Imbalanced, Multi-Fidelity Atomistic Data

Frontier 超算を用いた大規模マルチタスク学習により、従来法では数年を要する計算を 50 秒で完了させる原子構造グラフ基礎モデルを開発し、広範な化学設計空間の高速探索を可能にした。

要約

この論文は、**「超巨大なスーパーコンピューターを使って、新しい素材（材料）を見つけるための『万能な AI 先生』を作った」**という画期的な成果について書かれています。

専門用語を抜きにして、日常の言葉と面白い例え話で説明しますね。

1. 何をしたの？（大きな絵）

これまで、新しい電池や太陽光パネル、薬などの素材を見つけるには、科学者が一つ一つ実験したり、複雑な計算（量子力学など）でシミュレーションしたりしていました。これは**「一つ作るのに数年かかる」**ような、とても時間とコストがかかる作業でした。

この研究では、**「16 種類の異なる教科書（データセット）」をすべて同時に読み込ませて、「5 億 4,400 万個もの分子の形」を学習させた AI を作りました。
その結果、この AI は「11 億個もの分子の候補」**を、たったの 50 秒でチェックし終えることができました。

• 昔のやり方： 11 億個の候補を調べるのに、同じコンピューターで6 年半かかっていたはず。
• 今回のやり方： 50 秒で完了。

まるで、**「11 億冊の辞書を手で引いて意味を探すのに 6 年半かかるのを、AI が 50 秒で全部読み上げて意味を教えてくれた」**ようなものです。

2. 何がすごいのか？（3 つのポイント）

① 「ごちゃ混ぜ」の勉強ができる（マルチタスク学習）

通常、AI は「有機物（生き物など）」のデータだけ、あるいは「無機物（金属など）」のデータだけで勉強させると、どちらか一方しか得意になりません。
でも、この AI は**「16 種類も違う教科書（有機物、無機物、触媒など）」を同時に勉強**させました。

• 例え話： 16 人の異なる先生（有機化学、無機化学、触媒など）が同時に教えてくれる教室です。AI は「共通の基礎知識（分子の動き）」を学びつつ、先生ごとに「答えの出し方（予測の頭）」を少し変えることで、どんな問題にも対応できるようになりました。

② 「超高速」で選別する（インフラの工夫）

5 億個ものデータを処理するには、普通の読み込みでは遅すぎてしまいます。そこで、**「データを一時的にメモ帳（NVMe）にコピーして、計算機に近づける」**という工夫をしました。

• 例え話： 図書館（ハードディスク）から本を取りに行くのが遅いので、「机の上（メモ帳）」に本を全部並べておいて、計算する人がその場で即座に取れるようにしたような状態です。これにより、通信の待ち時間がなくなり、爆速で計算できました。

③ 「精度」と「速さ」のバランス調整（ミックス精度）

AI の計算には「100% 正確な計算（FP64）」と「少し近似した計算（BF16/FP32）」があります。

• 勉強中（学習）： 100% 正確な計算で、基礎をガッチリ固めました。
• テスト中（推論）： 速さを優先して、少し近似した計算も試しました。
• 結果： 「速さを重視しても、科学実験に必要な精度は保てた」ことが証明されました。これは、**「スピード重視のスポーツカーでも、レースのルール（科学の法則）を破らずに走れる」**ことを意味します。

3. 具体的な成果（数字で見る）

• 学習データ： 16 種類のデータセット、5 億 4,400 万個の分子構造、85 種類以上の元素。
• 使用したコンピューター： アメリカの「Frontier（フロンティア）」という世界最強クラスのスーパーコンピューター。
• 処理速度： 11 億個の分子を 50 秒でチェック。1 秒間に約 2,180 万個の分子を処理。
• 応用： 学習した AI は、新しい素材を探すだけでなく、「データが少ない分野（例えば、新しい薬の候補）」でも、少量のデータで高い精度を発揮することがわかりました。

4. この研究がもたらす未来

この技術は、**「材料開発のスピードを劇的に変える」**可能性があります。

• 昔： 「新しい電池材料を見つける」→ 何十年もかかるかもしれない。
• 今： 「AI が 11 億個の候補を 50 秒でチェック」→ 有望な候補をすぐに絞り込み、人間が実験に集中できる。

まるで、**「広大な森（化学物質の宇宙）で、目的の宝石を見つける作業」を、「手探りで探す」から「ドローンで上空から一瞬でスキャンする」**ように変えたようなものです。

まとめ

この論文は、**「スーパーコンピューターと AI を組み合わせて、科学者が何十年もかけて行うはずだった『素材探し』を、50 秒で終わらせる」**という、人類の科学技術における大きな飛躍を報告したものです。これにより、未来のエネルギー問題や環境問題の解決策を、もっと早く見つけられるようになるでしょう。

技術概要

論文要約：Exascale Multi-Task Graph Foundation Models for Imbalanced, Multi-Fidelity Atomistic Data

この論文は、Oak Ridge 国立研究所（ORNL）および AMD などの研究チームによって執筆され、Frontier などのエクサスケールスーパーコンピュータを用いた、原子構造データのための大規模マルチタスク・グラフ・ファウンデーションモデル（GFM）の開発と実証について述べています。

以下に、問題定義、手法、主要な貢献、結果、および意義について詳細にまとめます。

1. 背景と問題定義

材料科学や化学の分野において、密度汎関数理論（DFT）などの第一原理計算は高精度ですが、計算コストが極めて高く、広範囲な材料スクリーニングには現実的ではありません。機械学習ポテンシャル（MLIP）やグラフ・ファウンデーションモデル（GFM）はこのボトルネックを解消する有望な手段ですが、以下の課題が存在しました。

• データの不均一性と不均衡: 有機物、無機物、ハイブリッド系など、化学的に多様なデータソース（16 のオープンソースデータセット）を統合して学習させる際、データ量の偏り（不均衡）や、異なる近似理論（異なる機能や基底セット）によるラベルの不一致（マルチフィデリティ）が最適化を不安定にします。
• アーキテクチャの選択: 単一のモデル構造がすべての化学領域で最適とは限らず、大規模なハイパーパラメータ最適化（HPO）が必要ですが、従来の手法では計算リソースの制約から包括的な探索が困難でした。
• スケーラビリティと移植性: 異なるハードウェアアーキテクチャ（AMD MI250, Intel XPU, NVIDIA A100 など）間で、データパイプラインやトレーニング効率を維持したまま大規模にスケーリングするワークフローの欠如。
• 推論の速度: 数十億規模の候補構造をスクリーニングするには、第一原理計算に匹敵する精度を保ちつつ、極めて高速な推論が必要です。

2. 提案手法と技術的革新

本研究は、HydraGNN、ADIOS2、DDStore を統合した包括的なエクサスケールワークフローを構築しました。

A. 大規模マルチタスク学習（MTL）アーキテクチャ

• 共有メッセージパッシングと専用ヘッド: 16 の異なるデータセット（5 億 4400 万超の構造、85 元素以上）を同時に学習します。共有されたメッセージパッシング層で普遍的な原子間相互作用の特徴を学習し、各データセット固有のフィデリティやラベルのオフセットを吸収するための「タスク固有の出力ヘッド」を設けることで、不均衡なデータでの安定した学習を実現しました。
• 組成条件付きブランチ重み付け: 推論時に、化学組成ベクトルに基づいて MLP（多層パーセプトロン）が各データセット固有の専門家（ブランチ）からの予測を重み付けして統合します。これにより、訓練データに存在しない新しい化学系に対しても柔軟に予測できます。

B. システム設計とデータパイプライン

• ADIOS2/DDStore パイプライン: 大規模な異種データセットの効率的な読み込みと、NVMe ストレージを活用したデータステージングにより、ファイルシステムのボトルネックを回避し、2,048 ノード規模の継続的トレーニングを可能にしました。
• シャード分散トレーニング: HydraGNN の拡張により、データセットのサイズに応じて MPI ランクを割り当てるタスク並列性を導入し、小規模なデータセットを持つノードのアイドル時間を削減しました。

C. エクサスケール規模のハイパーパラメータ最適化（HPO）

• Frontier 上で 6 つの異なる MPNN（メッセージパッシングニューラルネットワーク）バックボーン（EGNN, SchNet, DimeNet, MACE, PaiNN, PNAEq）を対象に、DeepHyper を用いた大規模 HPO を実施しました。
• 精度だけでなく、「時間-to-エポック」やリソース効率も考慮し、実用的なモデルを選択するプロセスを自動化しました。

D. 精度と推論の最適化

• 精度感度分析: BF16、FP32、FP64 の各精度での推論と微調整を比較し、科学的な信頼性（FP64）とスループットのトレードオフを定量化しました。
• 推論パイプラインの最適化:
  1. エンコーダ再利用: 共有エンコーダの計算を 1 回のみ実行し、結果をキャッシュ。
  2. ブランチスキップ: 重みが小さいブランチの計算を省略。
  3. 融合勾配（Fused Gradient）: 16 のデコーダからのエネルギー予測を重み付け合計し、力の計算を単一の自動微分呼び出しで実行。
  4. torch.compile: コンパイルによるカーネル融合。
     これらの最適化により、ベースラインと比較して最大 33 倍の推論速度向上を実現しました。

3. 主要な結果

A. トレーニングとモデル選択

• Frontier 上での HPO により、PaiNN ベースのモデルが最も優れた性能と計算効率のバランスを示し、リードモデルとして選定されました。
• このモデルは約 1,210 万パラメータ、92MB のメモリ使用量と軽量でありながら、100 エポックのトレーニングで高い予測精度を達成しました。

B. スケーラビリティ

• クロスシステムスケーリング: Frontier (AMD)、Aurora (Intel)、Perlmutter (NVIDIA) の 3 つのリーダーシップクラスシステム間で、強いスケーリング（Strong Scaling）と弱いスケーリング（Weak Scaling）の両方で高い効率を維持しました。
  • Perlmutter: 最大 2,048 GPU までほぼ線形スケーリング。
  • Aurora: 最大 6,144 GPU まで。
  • Frontier: 最大 1,024 GPU まで。

C. 微調整（Fine-tuning）性能

• 事前学習済みモデルを 12 の多様な下流タスク（有機分子、無機結晶、凝縮相など）で微調整した結果、ゼロから学習する場合やバックボーンを凍結する場合と比較して、検証誤差が最大 1 桁改善しました（特にエネルギー表面（PES）に関連するタスクで顕著）。
• 少量データ（例：Wiggle150 の 150 構造）の領域でも、事前学習済みモデルは優れた汎化性能を示しました。

D. 推論スループットと大規模スクリーニング

• Frontier の 9,300 ノード（約 74,400 GPU）を用いた大規模推論実験において、50 秒間で 11 億個の原子構造をスクリーニングすることに成功しました。
• これは、第一原理計算（1 構造あたり約 30 分と仮定）で同等のタスクを行う場合、約 6.7 年かかる計算量を、数秒で完了させたことを意味します。
• 精度と速度のトレードオフ：FP64 は最高精度ですが、FP32 や BF16 を用いることでさらに高速化が可能であり、用途に応じた選択が可能であることを示しました。

4. 意義と貢献

1. 科学的 ML の新たなパラダイム: 単一のデータセットや均一なフィデリティに依存せず、異種で不均衡な大規模データを統合して学習する「マルチフィデリティ・マルチタスク」アプローチの実証的成功です。
2. 材料発見の加速: 11 億規模のスクリーニングを 50 秒で完了させる能力は、これまで第一原理計算では不可能だった広大な化学設計空間の探索を現実的なものにし、新材料発見のサイクルを劇的に短縮します。
3. ポータブルなエクサスケールワークフロー: 異なるハードウェアアーキテクチャ間で、データパイプラインからトレーニング、推論までを再設計なしに移植可能にする手法を確立しました。
4. 実用性の高いモデル: 事前学習済みモデルが軽量（92MB）であり、エッジデプロイや複雑な物理シミュレーションへの直接統合が可能である点も重要です。

この研究は、エクサスケールコンピューティングと科学機械学習を融合させ、材料科学における「計算の壁」を打破する具体的な道筋を示した画期的な成果と言えます。