Zatom-1: A Multimodal Flow Foundation Model for 3D Molecules and Materials

Zatom-1 は、分子と材料の両方の 3D 構造を統合的に学習し、生成と予測の両方のタスクで専門的なモデルに匹敵する性能を発揮しながら、推論時間を大幅に短縮する初のオープンソース・マルチモーダル・フロー・ファウンデーションモデルです。

要約

ZATOM-1：化学の世界を「創造」し「予測」する万能な AI 助手

こんにちは！今日は、科学の分野で大きな一歩を踏み出した新しい AI、**「ZATOM-1（ザトーム・ワン）」**について、難しい専門用語を使わずに、わかりやすくお話しします。

想像してみてください。化学の世界には、無数の「分子（薬の材料など）」と「材料（スマホの電池や新しい金属など）」が存在します。これまで、科学者はこれらの新しいものを作るために、何十年もかけて実験を繰り返してきました。しかし、ZATOM-1 は、この「試行錯誤」を AI が一瞬でシミュレーションし、さらに「どんな性質を持つのか」も予測してしまう、画期的な存在なのです。

🧩 従来の AI と ZATOM-1 の違い：「料理人」vs「万能シェフ」

これまでの AI は、まるで**「料理人」**のようでした。

• A さん（分子用 AI）： 寿司（分子）を作るのは天才だけど、パスタ（材料）は作れない。
• B さん（材料用 AI）： パスタは絶品だけど、寿司は作れない。
• C さん（予測用 AI）： 料理の味見は得意だけど、自分で料理は作れない。
• D さん（生成用 AI）： 料理は作れるけど、味見は苦手。

つまり、それぞれが「得意分野」しか持っていなかったのです。

一方、**ZATOM-1 は「万能シェフ」**です。

• 寿司もパスタも作れる（生成）： 新しい分子も、新しい材料も、ゼロから生み出せます。
• 味見も得意（予測）： 作った料理が美味しいか（性質が良いか）、カロリーはどれくらいか（エネルギー）を即座に予測できます。
• 何より速い： 従来の AI が 1 時間かかる作業を、ZATOM-1 は数分で終わらせてしまいます。

🏗️ ZATOM-1 がどうやって働くのか？3 つの魔法

ZATOM-1 がこれほどすごいのは、3 つの「魔法」を使っているからです。

1. 3D の「粘土細工」をそのまま扱う（直接生成）

多くの AI は、複雑なものを理解するために、いったん「暗号（潜在空間）」に変換してから作業します。それは、粘土を一度溶かして型に流し込み、また固めてから形を作るようなものです。
しかし、ZATOM-1 は**「粘土を直接手で捏ねる」**ように、原子の位置や種類を 3D 空間で直接扱います。

• メリット： 変換の手間がないので、超高速で、より自然な形を作れます。

2. 「混ぜて、混ぜて、混ぜて」から「元に戻す」学習（フローマッチング）

ZATOM-1 の学習方法は、**「カオスな状態から秩序ある状態へ」**を学ぶことに似ています。

• イメージ： きれいに並べられたレゴブロック（完成品）を、一度全部バラバラにして、箱の中で激しく揺らします（ノイズ）。
• 学習： AI は、「このバラバラの状態から、どうすれば元のきれいなレゴに戻せるか？」を練習します。
• すごい点： この練習を「分子」と「材料」の両方で行うことで、AI は化学の「共通言語」を身につけます。

3. 化学の「共通言語」を学ぶ（転移学習）

これが ZATOM-1 の最大の強みです。

• 従来の方法： 分子の勉強をしたら、材料の勉強をゼロからやり直す必要がありました。
• ZATOM-1 の方法： 分子と材料を同時に勉強します。
  • 例え話： 「日本語（分子）」と「スペイン語（材料）」を同時に勉強すると、両方の言語に共通する「文法（化学の法則）」が深く理解できるようになります。
  • 結果： 材料の勉強をしたおかげで、分子の予測精度が向上しました！これは、これまで見たことのない「化学分野を超えた学習」です。

🚀 どれくらい速くてすごいのか？

• 速さ： 従来の AI が 10 時間かかる生成作業を、ZATOM-1 は1 時間以下で終わらせます。さらに、最新のモデルでは10 倍〜100 倍の速さです。
• 精度： 薬の候補となるような複雑な分子を作っても、化学的に「ありえない形」を作ってしまうことがほとんどありません（99% 以上が正しい形）。
• オープンソース： この「万能シェフ」のレシピ（コードと重み）は、世界中の誰でも無料で使えます。

🌟 まとめ：未来の科学を加速する

ZATOM-1 は、単なる計算ツールではありません。

• 新しい薬を設計する
• 環境に優しい電池材料を見つける
• 超軽量で強い素材を発明する

これらを、これまでよりもはるかに短時間で、かつ正確に行える可能性を秘めています。まるで、化学の設計図を描くための「魔法のペン」を手に入れたようなものです。

この AI が普及すれば、私たちが待ち望んでいた「画期的な新薬」や「クリーンエネルギー」が、もっと早く、もっと安く実現するかもしれません。ZATOM-1 は、科学の未来を「予測」するだけでなく、実際に「創造」し始める最初のステップなのです。

技術概要

ZATOM-1: 3D 分子および材料のためのマルチモーダルフロー基盤モデル

技術的概要（日本語）

本論文は、3D 分子と材料の両方に対する**生成（Generative）と予測（Predictive）**のタスクを統合した、世界初のエンドツーエンドかつフルオープンソースの基盤モデル「ZATOM-1」を提案しています。既存の AI 手法は通常、分子か材料のどちらか一方のドメイン、あるいは生成か予測のどちらか一つのタスクに特化しており、表現の共有や転移学習が制限されていました。ZATOM-1 はこの課題を解決し、化学分野における汎用的な基盤モデルの実現を示しました。

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

---

1. 問題定義 (Problem)

化学システム（分子や材料）の設計空間を探索することは、創薬や新材料開発において極めて重要ですが、従来の実験や計算手法では時間とコストがかかり、エラーも発生しやすい課題がありました。

• ドメインの分断: 既存のモデルは「分子（非周期的）」と「材料（周期的）」を別々に扱っており、両者を統一的にモデル化する基盤モデルが存在しませんでした。
• タスクの分断: 生成モデル（新しい構造の生成）と予測モデル（物性、エネルギー、力の予測）が分離されており、生成モデルを事前学習として利用して予測タスクの精度を向上させるアプローチが不足していました。
• スケーラビリティと速度: 従来の生成モデル（拡散モデルや潜在拡散モデル）は、トレーニングやサンプリングに時間がかかり、大規模なモデルへのスケーリングが困難でした。

2. 手法 (Methodology)

ZATOM-1 は、**マルチモーダル・フローマッチング（Multimodal Flow Matching）**を基盤としたトランスフォーマーアーキテクチャを採用しています。

2.1. 入力表現とモダリティ

原子レベルの粒度でシステムを表現し、以下の 5 つの統一されたモダリティを使用します：

1. 原子種 (Atom Types): 離散的なモダリティ（整数値）。
2. 3D 座標 (3D Coordinates): 連続的なモダリティ（Å単位）。
3. 分数座標 (Fractional Coordinates): 単位胞内の位置（[0, 1) 範囲）。
4. 単位胞の長さ (Unit Cell Lengths): 格子定数 $a, b, c$。
5. 単位胞の角度 (Unit Cell Angles): 格子角度 $\alpha, \beta, \gamma$。

• 非周期的な分子の場合、単位胞関連の入力はマスクされます。

2.2. アーキテクチャ: Trunk-based Flow Transformer (TFT)

• 標準化されたトランスフォーマー: 特殊なグラフニューラルネットワークや手作業の生成事前知識に依存せず、標準的なトランスフォーマー（Query-Key Normalization, Flash Attention, SwiGLU を採用）を使用します。
• 条件付きフローマッチング (Conditional Flow Matching):
  • 離散的な原子種と連続的な幾何学モダリティ（座標、格子など）を同時にモデル化します。
  • 確率流（Probability Flow）を近似するベクトル場を学習し、ソース分布からターゲット分布への軌跡を ODE（常微分方程式）として定義します。
  • 非自己回帰的: 生成プロセスは自己回帰的ではなく、並列処理が可能であるため、サンプリングが高速です。
• トレーニング戦略:
  • ステージ 1（事前学習）: MP20（材料）と QM9/GEOM-Drugs（分子）のデータセットを用いた生成タスクでモデルを事前学習します。これはラベル不要の自己教師あり学習として機能します。
  • ステージ 2（微調整）: 事前学習された重みを凍結し、物性、エネルギー、力の予測タスクに対して追加のヘッドとバックボーンを微調整します。

2.3. 損失関数

連続モダリティには MSE 損失、離散モダリティ（原子種）には交差エントロピー損失を使用し、重み付けパラメータ $\lambda_{discrete}$ で統合されたマルチオブジェクティブ損失を最適化します。

3. 主要な貢献 (Key Contributions)

1. 初の統合基盤モデル: 分子と材料の両方に対応し、生成（新規構造の創出）と予測（物性、エネルギー、力の予測）を単一のモデルで実行可能にしました。
2. 高性能な生成と予測:
   • 生成: 材料生成ベンチマーク（LeMat-GenBench）および分子生成ベンチマーク（QM9, GEOM-Drugs）において、既存の最先端モデル（SOTA）と同等かそれ以上の性能を達成しました。
   • 予測: 事前学習された生成モデルを転移学習することで、分子の物性予測精度が向上し、化学ドメイン間での正の転移学習（材料の学習が分子の予測精度を高める）を実証しました。
3. 圧倒的な高速性とスケーラビリティ:
   • 潜在拡散モデル（Latent Diffusion）のベースラインと比較して、GPU トレーニング時間が3 倍短縮されました。
   • サンプリング速度は、5 億パラメータの潜在拡散モデルベースラインに対して12.5 倍高速（3 億パラメータの ZATOM-1 が 4 分未満で 1 万サンプル生成）です。
   • モデルサイズを増やすことで、トレーニング損失や生成の妥当性（Validity）が予測可能な形で改善されることを示しました（スケーリング則の追従）。
4. オープンソース化: コードと重みが公開され、研究の再現性と発展を促進します。

4. 実験結果 (Results)

• 材料生成 (MP20 / LeMat-GenBench):
  • 80M パラメータの ZATOM-1 は、47M パラメータの MatterGen や 7B パラメータの PLaID++ などの既存モデルと比較して、妥当性（Validity）、一意性（Uniqueness）、新規性（Novelty）において競争力のある結果を達成しました。
  • 特に、緩和（Relaxation）を行わない設定でも、安定した結晶構造を生成できることが示されました。
• 分子生成 (QM9 / GEOM-Drugs):
  • QM9 において、PoseBusters（物理的妥当性チェック）のパス率が約 99% に達し、既存のモデル（ADiT など）を上回る SOTA を達成しました。
  • GEOM-Drugs（大規模分子）でも、94% の妥当性を達成し、薬物様分子の生成能力を証明しました。
• 予測タスク (QM9 / Matbench):
  • 分子と材料を共同で事前学習したモデルは、分子の物性予測において SOTA 性能を達成しました。
  • 転移学習の検証: 材料データで事前学習したモデルを分子の予測に適用した際、分子データのみで事前学習した場合よりも精度が向上しました。これは、化学ドメイン間の共通表現が学習されていることを示唆しています。
• エネルギーと力の予測:
  • 材料の力予測において、最先端モデル Orb-v1 と同等以上の性能を示しました。

5. 意義と将来展望 (Significance)

• 科学 AI (SciML) のパラダイムシフト: 従来の「生成」と「予測」の分離、および「分子」と「材料」の分離を超え、単一の基盤モデルで化学空間全体を扱う新たなアプローチを確立しました。
• 効率性の向上: 高速なトレーニングとサンプリングにより、大規模な化学空間の探索や、反復的な設計プロセス（創薬、材料開発）への実用化が現実的なものになりました。
• 転移学習の可能性: 生成モデルを事前学習として利用することで、ラベル付きデータが不足する領域でも高精度な予測が可能になることを示しました。
• オープンサイエンス: フルオープンソースであるため、世界中の研究者が基盤を共有し、化学発見の加速に貢献できます。

総じて、ZATOM-1 は、3D 分子と材料のモデルリングにおいて、スケーラブルで高速、かつ高精度な新しい標準を提示した画期的な研究です。