[論文レビュー] UMA: A Family of Universal Models for Atoms

Brandon M. Wood, Misko Dzamba+16 more|ArXiv.org|Jun 30, 2025

Machine Learning in Materials Science被引用数 20

ひとこと要約

UMAは、分子、材料、触媒にまたがる約5億個の原子構造で学習した普遍的な機械学習型原子ポテンシャルのファミリーを提示します。モデルはMixture of Linear Experts (MoLE) を用いて容量を効率的にスケールさせ、タスク固有の微調整を行わずに競争力のある、または最先端の性能を達成します。

ABSTRACT

The ability to quickly and accurately compute properties from atomic simulations is critical for advancing a large number of applications in chemistry and materials science including drug discovery, energy storage, and semiconductor manufacturing. To address this need, Meta FAIR presents a family of Universal Models for Atoms (UMA), designed to push the frontier of speed, accuracy, and generalization. UMA models are trained on half a billion unique 3D atomic structures (the largest training runs to date) by compiling data across multiple chemical domains, e.g. molecules, materials, and catalysts. We develop empirical scaling laws to help understand how to increase model capacity alongside dataset size to achieve the best accuracy. The UMA small and medium models utilize a novel architectural design we refer to as mixture of linear experts that enables increasing model capacity without sacrificing speed. For example, UMA-medium has 1.4B parameters but only ~50M active parameters per atomic structure. We evaluate UMA models on a diverse set of applications across multiple domains and find that, remarkably, a single model without any fine-tuning can perform similarly or better than specialized models. We are releasing the UMA code, weights, and associated data to accelerate computational workflows and enable the community to continue to build increasingly capable AI models.

研究の動機と目的

• 材料、分子、触媒といった多様な化学分野に跨る高速で正確なDFT代替手法の必要性を動機づける。
• 単一の大規模モデルが微調整なしでタスク間を一般化できることを示す。
• 推論速度を犠牲にすることなく容量を増やすスケーラブルなアーキテ Archテクチャ（MoLE）を開発する。
• 速度とエネルギー保存精度のバランスを取る2段階の訓練手順を提案する。
• コミュニティ全体での利用と検証を可能にするため、コード、重み、データを公開する。）

提案手法

• 総電荷・スピン・DFTタスク入力を拡張したeSENベースの等変グラフニューラルネットワークアーキテクチャを採用する。
• 出力が線形専門家の密結合な組み合わせとなるMoLEを導入し、滑らかなエネルギー表面と回転等変性を保つ。
• システムレベルの埋め込みから小規模MLPを用いてエキスパート重みαを計算し、事前計算したW* = Σk αkWkを用いて推論を高速化する。
• 2段階スケジュールで訓練する。最初は力を直接予測し、次にエネルギー保存と応力のためにオートグラッドで微調整する。
• 事前学習にはBF16を用い、微調整時にはFP32に切り替えて精度を維持する。大規模MoLE構成へスケールさせるためにメモリ/グラフ並列性を活用する。
• 材料、分子、触媒からなる多様なデータセットの約5億個の原子構造で訓練し、異なるDFT設定間のマルチタスク学習のためにエネルギー参照スキームを採用する。）

実験結果

リサーチクエスチョン

• RQ1 単一の非微調整済みモデルが材料、分子、触媒にまたがる複数のDFTタスクで競争力のある精度を達成できるか？
• RQ2 モデルサイズ、データ量、計算コストがどのように相互作用して最適なUMA構成を決定するか？
• RQ3 MoLEアーキテクチャは、多タスクMLIPにおいて、特に長時間のMDのようなシミュレーションで、密モデルに対してどのような利点を提供するか？
• RQ4 統一モデルは、多様なタスクとデータセット全体でエネルギー保存と滑らかなポテンシャルエネルギー表面を維持できるか？

主な発見

• UMAは、材料、触媒作用、分子、分子結晶、MOFにおいて、タスク固有の微調整なしで競争力の、または最先端の性能を達成する。
• MoLEは大幅な効率性を提供し、同等スケールで密なモデルと比較して約2.5倍少ないアクティブパラメータで同様の損失を達成する（例：UMA-M）。
• UMA-S、UMA-M、UMA-LはMatbench Discoveryおよび吸着エネルギーベンチマークで強い成果を示し、触媒学習のAdsorbML成功率を25%改善した。
• 単一モデルで長時間のMDロールアウトを処理でき、推論速度とメモリ使用量が有利で、80GBのGPU1枚で100k atoms超のシミュレーションとマルチGPUスケーリングの可能性を可能にする。
• 2段階の訓練により効率的な学習とエネルギー保存を維持し、BF16の事前学習とFP32の微調整を組み合わせる。
• 経験的なスケーリング則は、モデル規模とデータ量に対して性能が対数線形にスケールすることを示唆し、計算最適化と推論最適化の訓練戦略を導く。