UMA: A Family of Universal Models for Atoms
Authors: Brandon M Wood
Misko Dzamba
Xiang Fu
Meng Gao
Muhammed Shuaibi
Luis Barroso-Luque
Kareem Abdelmaqsoud
Vahe Gharakhanyan
John R. Kitchin
Daniel S. Levine
Kyle Michel
Anuroop Sriram
Taco Cohen
Abhishek Das
Sushree Jagriti Sahoo
Ammar Rizvi
Zachary Ward Ulissi
C. Lawrence Zitnick
Venue: NeurIPS 2025 spotlight ()
Keywords: Chemistry
Materials Science
Machine Learning Interatomic Potential
Density Functional Theory
Graph Neural Networks
Primary area: machine_learning_for_sciences
TL;DR: We introduce a Universal Model for Atoms (UMA), trained on nearly half a billion 3D atomic structures, designed to push accuracy, speed, and generalization in atomic simulations across chemistry using a single model without fine-tuning.

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, we present 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 $\sim$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 build increasingly capable AI models.

UMA: A Family of Universal Models for Atoms
Brandon M. Wood 1,*,
Misko Dzamba 1,*
Xiang Fu 1,* Meng Gao 1,*
Muhammed Shuaibi 1,*
Luis Barroso-Luque
Kareem Abdelmaqsoud 2
Vahe Gharakhanyan
John R. Kitchin
Daniel S. Levine
Kyle Michel Anuroop Sriram
Taco Cohen
Abhishek Das
Ammar Rizvi!
Sushree Jagriti Sahoo
Zachary W. Ulissi
C. Lawrence Zitnick L.
FAIR at Meta
²CMU
Abstract
The ability to quickly and accurately compute properties from atomic simulations is critical for advancing large number of applications in chemistry and materials sci- ence including drug discovery, energy storage, and semiconductor manufacturing. To address this need, present a family of Universal Models for Atoms (UMA), designed to push the frontier of speed, accuracy, and generalization. UMA models are trained on half billion unique 3D atomic structures (the largest training runs to date) by compiling data across multiple chemical domains, c.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 novel architectural design refer to 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 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 accelerate com- putational workflows and enable the community to build increasingly capable Al models.
Introduction
Density Functional Theory (DFT) models the interaction of atoms from first principles through the estimation of their electronic structure. It serves as the foundation of modern computational chemistry and materials science, and has provided insights into many applications including drug discovery [70]
78), energy storage [73 27 72), and semiconductors [62 (76). Despite DFT's widespread adoption,
its considerable computational expense limits its usage.
Machine learning models offer the potential to accurately approximate DFT while being dramatically faster (O(n) vs. (n²²) for DFT, where the number of atoms); reducing computation time from hours to less than second. Ideally, these models referred to as Machine Learning Interatomic Potentials (MLIPs)- would generalize across the many tasks that utilize DFT. We refer to "task" as chemical domain (e.g. materials, catalysts, molecules) plus the specialized DFT settings required by domain. For each task, we can explore different applications of DFT, such as molecular dynamics, relaxations, vibrational analysis, and mechanical response. Training MLIPs that generalize across such tasks remains an open problem.
"Co-first Author
Co-corresponding Author
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Molecules 100M
Materials 100M
Catalysis 230M
OMC 25M
ODAC 29M
UMA
Figure I: Visualization of the different datasets used for training. The 2D plots (bottom) illustrate the number of pairwise interactions contained in each dataset for every combination of elements. Note their combination covers nearly the entire chemical space with the exception of the radioactive elements. Model accuracies have improved with training dataset size (upper right), and this paper explores the limits of this scaling.
The scaling of datasets and model sizes has led to major breakthroughs in language and vision models, enabling their generalization across diverse data distributions and tasks 2 24). A similar scaling of atomic datasets is more challenging due to their computational cost, resulting in models typically being trained on smaller problem-speciffe datasets (57 18 17. potential solution is to pool data across tasks to allow for the creation of an exceptionally large dataset for training multi-task models. Recently, several large domain-soe sppcciic datasets including catalysts [0] 72], materials [5] 67], and molecules [4] have been released that, when combined, total nearly half billion atomic systems. This combined dataset defines a new paradigm in terms of the diversity of interactions and chemical environments (Figure D. allowing for the learning of general-purport models.
In this paper, we present a family of Universal Models for Atoms (UMA) designed to test the limits of accuracy, speed, and generalization for a single model across chemistry and materials science. The unprecedented amount of data (~500M atomic systems, Figure poses new challenges in balancing accuracy and efficiency with multi-task training. To explore this space, we develop empirical scaling laws, relating compute, data, and model size, to determine the model size required to fit the UMA dataset and to define compute-optimal and inference-optimal training strategies. To accommodate the growing demand for model capacity while maintaining speed, we introduce a novel Mixture of Linear Experts (MoLE) architecture that efficiently scales the model size without increasing inference times for applications such as Molecular Dynamics (MD). For efficient model training, we propose a novel two-stage training schedule that efficiently pre-trains a model that is then refined in the second stage to conserve energy at higher precision.
We demonstrate that UMA, without fine-tuning to specific tasks, performs similarly or better in both accuracy and inference-speed/memory -effeniiccccccccyyyye..e than specialized models on wide-range of material, molecular and catalysis benchmarks. Highlights include state-of-the-an results on the popular Matbench Discovery leaderboard [59], a 25% improvement successful adsorption energy calculations for catalysis 42), and accuracy sufficient for practical applications (e.g. ligand-strain energy) structure- based drug-design [43). All code and data used for training UMA is publicly available. UMA model weights are available with a commercially permissive license (with some geographic and acceptable use restrictions).
Approach
An MLIP acts surrogate for DFT, and so it requires the same inputs DFT: atom positions, their atomic numbers, and optionally spin and charge information. As outputs, MLIPs estimate the energy of an atomic structure from which other properties may be computed through the use of derivatives, such as per-atom forces, stress, etc. For many applications, such as molecular dynamics or performing relaxations, an MLIP is used to calculate the atomic forces to run simulations that can require thousands or even millions of iterations. For this reason, MI.PPs must be both accurate and computationally efficient.

---

noa-trainin Phase
Input
Forces, Strees
1.6
Single-aek
Nerma Energy MAE
Multi-task
Ingus Embedding
Grobal Embeeciing
MoE
Multi MoLE
1.0
Edpesitss
Router
0.8
SD3
cheenta
cC20
DMear5
QDAC23
Nodewaie Feod Forward
Embedding
MoLE
Darse
Validiation Loss.
8 Etpprts
W
W
120 Exparta
Energy Output Head
Router
Energy
Training Steps
Figure 2: (left) Overview of UMA model architecture. The SO2 convolution is made up of a set of linear operations and each one of these operations is replaced with MoLE. (middle) Illustration of MoE and MoLE. The embedding used for routing, which estimates the expert weights a. only depends on global information making possible to merge before the model forward pass (middle, bottom), which has substantial benefits for applications that require long roll outs such molecular dynamics. (right) Bar plot of UMA-S trained with MoLE for multi-task and without MoLE for single and multi-task. Note the MoLE model outperforms non-MoLE models. (right, bottom) Loss plots when varying the number experts from 128 for UMA-S.
Table I: Summary of UMA models.
Inferences per second for 1 Atoms
Max Atoms per 80GB GPU
Model
Total Parameters
Active Parameters
Conservative
UMA-S
150M
16
100k+
UMA-M
1.4B
10k+
UMA-L
700M
700M
1.6
1k+
X
Inference spead and max atoms measured on Nvidia H100 with pesiodic system that has es 50 neighbors per atom withas 6A, see Appendix
2.1 UMA: Universal Model for Atoms
In this paper, we describe family of models that can be categorized into three sizes: small (S), medium (M), and large (L.). summary of the models can be found Table Additional details on the different versions of UMA can be found in Appendix E.1 All the models have different accuracy and speed attributes and as result are best suited for distinct use cases. UMA-S aims to strike balance between accuracy and efficiency and is suitable for computationally intensive applications such as long molecular dynamics simulations. The most general-purpose model family is UMA-M, which is more accurate and be used as DFT surrogate for relaxations or vibrational analysis as well as for molecular dynamics. UMA-L is a highly accurate DFT surrogate and is intended limits of what is currently feasible.
proof-of-principp to help understand scaling behavior and to demonstrate the
We begin by providing basic introduction to the UMA architecture. Next, we highlight our use of Mixture of Linear Experts (MoLE), an efficient way to scale capacity and add flexibility for training models across different DFT datasets and tasks. Lastly, present overview of the training procedure used for UMA models.
2.2 Architecture
The UMA architecture is based on eSEN 21 an equivariant graph neural network, with number of important modifications to enable the model to efficiently scale and handle additional inputs, such as total charge and spin, and the DFT settings desired for emulation. The standard eSEN architecture

---

takes 3D atomic positions and atomic numbers as inputs and returns the total energy, per-atom forces, and, optionally, stress. The network operates by updating a spherical harmonic node embedding for each atom through a series of message passing layers. The embeddings are initialized based on the atom's atomic number. Messages are passed between neighboring atoms that are less than a predefined distance (6 Å) from each other. Each message passing layer consists of an edge wise block followed by node-wise feed-forward block with residual connections. Normalization is performed after each layer. The central component of the edgewise block is the eSCN convolution [54]. After message passing, there is a single task-independent node-wise feed-forward block to predict the outputs.
2.2.1 Charge, Spin, and DFT Task Inputs
In addition to the imputs handled by eSEN, we expand the model to incorporate information about the system's 's total charge, total spin multiplicity, and the DFT task. The charge and spin are indicated using an integer value for each and allow the MLIP to model structures which may be electrically charged (have extra or missing electrons) or have unpaired electrons. The DFT task indicates which of the five training dataset DFT settings the model should attempt to replicate. single dataset is specified and random vector corresponding to it is fed into the network. This is necessary since different datasets use different plane-wave or localized orbital DFT calculators (VASP 39 and ORCA 50) and levels of theory.
We introduce a new embedding to UMA models that enables the inclusion of charge, spin, and DFT task. Each of these inputs generates an embedding with the same dimension as the number of spherical channels used, which is concatenated and then fed through a I-layer feed-forward network. The result is added to the node embeddings at each layer for the spherical harmonics coefficients of degree 0 (L = 0). This embedding is also used compute the global embedding used for MoLE routing, which is described below.
2.3 Mixture of Linear Experts
A common approach to learning general-purpose models is to increase the amount and diversity of their training data. As datasets increase in size, scaling laws suggest model sizes must also be scaled to optimally reduce losses [24, 2]. However, this presents a tradeoff, which may lead to accurate models that too computationally expensive to in practical applications. One method (MoEs) to increase the number of model parameters while minimizing the additional computational cost. In MoE model, the outputs of a block are calculated by a set of experts, each with their own individual set of weights. Typically, sparse weighted combination of their outputs is combined and passed to the next block. This approach has been shown to work well for LLMs to improve both their efficiency and their ability to generalize 19
to improve these trade-offs is to use Mixture of Experts 32 31
60
The application of MoEs to MLIPs requires additional considerations. Contrary to language modeling, estimating the potential energy surface is regression task whose outputs should vary smoothly to ensure energy conservation (11. In addition, the estimated forces should be equivariant to rotation. This implies that the use of experts should not introduce discontinuities on the energy surface, and care must be taken to ensure equivariance is maintained. Another practical consideration for MIJPs is that the task and of atoms is commonly held constant during long simulations. Ideally, an MoE approach for MLIPs would take advantage of this to improve efficiency in sequential inference applications.
We propose using a Mixture of Linear Experts (MoLE) for MLIPs. An MoLE combines a set of linear experts 46 31:
y = Σ as (War)
where each expert i has a set of weights W₁ and contribution 0g E [0, 1]. While this was one of the original approaches proposed for MoEs [1I, MoEs that use sparse sets of mon-linear experts which compete for attention [32 are much more commonly used in modern networks [9 60 15 63
However, MoLEs offer distinct advantages for MLIPs. First, the MIIP can learn functions that share information and vary smoothly between tasks by encouraging the dense use of all experts without enforcing sparseness. Second, since MoLE is a mixture of linear experts, they maintain

---

rotational equivariance when used within the eSCN convolution [54]. Finally, if the expert weights are only dependent on tineiinvariant global information such as element composition, the network weights may be precomputed before running simulations [46 75 The precomputed weights W* are calculated by moving the summation in Equation which results in MoE. inference times that are similar to non-MoE models:
y W x
where W
Σ CAe W4
The contribution an of each expert is calculated as function of system-level features, including element composition, charge, spin, and task information. Embeddings are calculated for each of these properties, concatenated, and passed through 3-layer MLP followed by a softmax estimate akk Σx 04 = 1. We specifically exclude information when calculating a's ''s that may vary during the course of relaxations or MD simulations, such as relative atom positions or other neighborhood information.
2.4 Training Procedure
Training models large datasets across numerous tasks is challenging due to the computational resources required. This is especially true for conservative models, which require an additional backward pass calculate forces stress. To improve training efficiency, we implemented two-stage approach described in detail in the supplementary material (Appendix [A]. As proposed by 7 train model that directly predicts forces in the first stage. In the second stage, remove the force head and fine-tune the model to predict conserving forces and stresses using auto-grad. This approach takes advantage of the faster training enabled by direct models, while providing energy conservation and smooth potential energy landscapes required by many applications.
Several other novel improvements were made to each stage's 's training to further improve efficiency. Half-precision is critical to training efficiently on modern GPUs. Similar to other domains such as LLM training, we found that 6F16 is significantly more stable compared to FP16 with automatic mixed precision. However, unlike LLMs, MLIPs are significantly more sensitive to numerical precision; we observed that BF16 alone will degrade accuracy by much as 20- 50% depending on the task and the property being computed. We found that pre- pre-training with BF16 and switching to FP32 for fine-tuning recovers the loss in accuracy. Finally, training the models with large number of MoLE experts is challenging due to memory constraints. We make further optimization to help bound and amortize memory usage (Appendix AA). In addition, models were trained with combination of graph parallelism [66] and fully-sharded data parallelism for large MoLE layers and activation checkpointing (Appendix A). When combined, this allows us reliably scale up model training up to 1OB total parameters.
The energies for each dataset can vary significantly due the use of different DFT settings, which makes multi-task training difficult. To account for this, an energy referencing scheme was employed as described in Appendix AA66) A..6 This approach allowed for the use of single energy head across all tasks without the need for task-specific heads.
3 Datasets
To train a model capable of generalizing across DFT tasks, ideal training dataset would include materials, molecules, and their interactions. The Open Molecules 2025 (OMol25) 43 and Open Materials 2024 (OMar24) [5] datasets cover both molecules and materials respectively. The Open Catalyst 2020 (OC20) [T]] and OpenDAC 2025 (ODAC25) [68] datasets are useful for modeling the interactions of molecules and materials. Finally, the Open Molecular Crystals 2025 (OMC25) 23 dataset specializes in modeling the interaction between molecules in periodic structures. In combina- tion, these datasets contain close to 500 million training examples with over 30 billion atoms. As visualized in Figure the combination of these datasets contains nearly all pairwise interactions be- tween atoms of different elements. We focused on large datasets, because of the challenges of mixing different DFT tasks and since the inclusion of much smaller datasets was unlikely to significantly impact the model weights.
While all DFT calculations yield energy and force labels by determining ground-state electronic structures, different chemical systems require domain-specifi, DFT settings. For instance, the OMat24
5

---

1.4
Dense ba7.07s
3e-18 FLOPs
[cl
(0)
Val Loss vs Active Params
2.0
3a ELORs
00119 FLOPS
20+20 FLOPS
Benus
MulE Opc, FI
(6)
V0l Loos
Votl befisg
Trannin Computt is Eutt
3ost8 FLCPs
(d)
Compute Opt. Ftt
FUCPS
error
4
29 FUOPs
Ge+18 FLOPS
20 PUOPS
HOLE Con pute Opt.
Canas OuL
std error
3e-19 lsoFLOPs
(b)
1e-20 Lsof
11P
307
UP
120
100
WP
16'
10
20
Activ Narumaters
Active Parameters
Figure 3: Empirical scaling measurements of dense (blue) MoLE (red) model architectures.
FLOPs vs. validation loss for (a) dense and (b) MoLE (8-expert) models. Experiment sets are performed by holding FLOPss constant and varying model size and training data. Diamonds represent the compute optimal frontier. (c) Training compute vs. parameters. Vertical green dotted line represents our estimated training budget of (1122)) FL.OPs, horizontal dotted lines are corresponding dense and MoLE model sizes. (d) Compute vs. dataset size (atoms) with the green dotted line representing the training FLOPs required for 1 epoch of UMA training data (50B atoms) (e) Overlay of dense vs. MoLE compute optimal frontiers from (a-b) and the fitted power law of validation loss as a function of parameters. compute optimal MoLE model with A 2.5x fewer active parameters can achieve an equivalent loss. Fitting details and parameters are described in Appendix
dataset uses the PBE functional and is generated with the plane-ward wave code VASP, whereas OMol25 employs the WB97M-V functional with the localized orbital code ORCA. Designing single model that performs well across such diverse settings a non-trivial challenge. Additionally, the datasets differ in size and information content, so we adjust their sampling ratios: 4 for OMat24 and OMol25, 1 for OC20 and ODAC25, and 2 for OMC25. Further details on the datasets and our unified modeling strategy can be found in Appendix B and A respectivelly
4 Results
4.1 Model and Data Scaling
In various domains such language, empirical power-law relationships have been successfully used 228 24 35) to predict the optimal model size and quantity of training data given a fixed compute budget. However, for modeling MLIPs, models have been trained the size of UMA's dataset and very few models have been trained to > >100M parameters [5]; we do not know whether such relationships exist or whether the models will continue to improve with data and parameters. In this study, we show that UMA models do show log-linear scaling behavior as seen in other domains in the FLOP ranges we tested, indicating that greater model capacity is required to fit the UMA dataset. We used these scaling relationships to select the appropriate model size for our dataset and demonstrate the advantage of MoLE over the dense architecture.
we first examine the behavior of training compute in FLOPs C' as a function of model size N and dataset size D. To determine this, used total validation loss (computed in the same way as training loss) as the observable quantity. We trained a series of models by sweeping the model parameter size while fixing the training compute budget in FLOPs. This is done by varying the amount of training data for each model, i.e., smaller models receive more training data, while larger models receive less. This is performed for both the dense and MoLE (8 expert) models classes, producing "Iso-FLe FLOPS" curves Figure 3 b).. The minima of the FLOPs represent the compute optimal model, or the best achievable loss for the given compute budget (diamonds in Figure
Following (28] 24 35
Optimal model size. We fit the Iso-FL.PPs minima using power laws 35 28
to estimate the Figure c.d)
optimal model and dataset sizes given fixed compute budget (detailed in Appendix

---

indicates that the dense and MoLE models display convincing log-linear scaling behavior at the scales of our experiments (10' 1020 FLOPs). Given that our training dataset was approximately 50 billion atoms (1 epoch) and our estimated compute budget for pretraining was ((102)) FL.OPs (extrapolated from our preliminary training runs), Figure suggests the largest compute optimal dense model we should train is ~ 700M parameters corresponding to the UMA-L. In practice, we used the compute optimal model size as a starting point and continued training beyond the compute optimal regime to produce lower losses while minimizing parameter sizes. However, we observed that training on too many epochs can lead to overfitting and significantly deviate from log-linear scaling behavior, prompting to keep training within 2-3 epochs.
Dense vs MoLE. If we combine Figures a..b) and fit the Iso-FLOPs minima as a function of loss (Figure Be). we can compare the model sizes needed by the dense and MoLE models to achieve a fixed loss. Effectively, an optimal Mol.E. model can achieve the same loss as an optimal dense model that is A times larger, e.g., A x 2.5 + 0.2 for UMA-M. We observe that this advantage is reduced at larger model sizes; training a 700M active parameter (8 expert 5.6B parameter total) MoLE model had marginal improvement over the dense version. This effect can be seen in overlaid IsoFLOPs in Figure (e.. e). As parameters increase, dense and MoLE performance converge. We hypothesize that we are limited by our dataset's size, regardless of MoLE or dense architectures.
4.2 Multi-task vs. Single-Tast Task
To explore the benefits of MoLE models for multi-task training, we compare UMA-S (with MoLE) to non-MoLE dense models trained for single/multi-task in Figure 2 (right, top). In this limited parameter regime, multi-task models trained without MoLE always perform worse than single-task (ST) models. With the use of MoLE, UMA-S can achieve results comparable to task specialized small models. We also experimented with varying the number of experts for UMA-S in Figure 2 (right, bottom). A significant improvement in loss is observed when moving from expert 8. A smaller gain is found with 32 experts, and the additional gain from 32 to 128 experts is negligible. In our experiments, UMA-S uses 32 experts. For large models, multi-task training offers benefits even without MoLE, as shown in Appendix F where multi-task models avoid overfitting and generally achieve better performance.
4.3 Inference Efficiency
The ability to increase the length scale (number of atoms) and time scale of molecular dynamics (MD), relaxations and other types of long running simulations is critical to our design. Leveraging our MoLE strategy, we can pre-merge all the weights and pay additional penalty in inference time or memory when running a long time-scale simulations on a single system. Here we show that UMA models, despite having very large parameter counts and achieving SOTA accuracy across many domains, are extremely competitive in simulation settings in both speed and raw number of atoms that can fit in memory compared to specialized single-domain models. For example, the UMA-S can simulate 1000 atoms at 16 steps per second (1.4 ns-per-day) and fit system sizes up 100,000 or more atoms in memory on a single 80GB GPU. Lastly. our models are designed be scaled with multi-GPU parallel inference using the Graph Parallelism strategy [66] described in training; giving the potential to provide orders of magnitude speed-ups when running simulations at the 100k+ atoms scale and unlock new science that is not possible today. We will leave multi-GPU inference outside the scope of this paper.
4.4 Evaluation
Our evaluation contains two main components: (1) diverse set of held-out test splits (Table 2 and (2) a suite of practically important benchmarks [Table(3. To cover all chemical domains studied: materials, catalysis, molecules, molecular crystals, and metal-organic frameworks (MOFs), we propose new test sets and benchmarks, in addition to established benchmarks in existing literature. We highlight selected results below, with more detailed discussion in Appendix [E]
Materials. Table 2 we report the test-set performance on two OOD test sets: WBM 74 and high entropy alloy (HEA). The HEA test set is introduced in this work (Appendix E.2 UMA-M performs on-par with SOTA task specialized models, such as eSEN-30M--Mat and V--Matt The degraded performance of UMA-L on energy of the HEA dataset (Table II] suggests overfitting,
7

---

Table 2: Test MAE results held test splits for materials [59], catalysis I molecules 43 molecular crystals 23 and MOF' 667. All energies are in meV, forces are in meÀ and stresses are in VeÀ3 Results for UMA are compared against the SOTA literature results. Target accuracies for practical utility are provided as an approximate guide for reference.
Catalyns
crystols
Adh Barae
0OD Doth
As. Enrge
OUD OuD-Comp Comp Farggi
EnergeSttoe
OMC Text
Eaerge Alsr
OOD&LI
WBM
Porees
Sbrak
Ruvss
Rvvus
Rinos
Renus
Streze
Rumon
Model
10
UMA
UMA-S-LI
24.9
85.7
3.5
51.5
24.1
68.8
30.7
8.64
0.34
15.47
1,0
5.04
0.95
UMA-M-I.I
3.5
10.34
0.54
0.90
Literature
SEEN-00M-MM-
H0
16.2
49.6
20.0
59.5
3.2
OMM 5
14.9
46.3
3.6
20.3
47.9
2.7
OC20
[4
149.1
Gem
201
163.5
4p2--DAC
67
316.0
eSEN
1.39
0.83
12.72
Target
Practical Utilit
10-20
10-20
100 100
1-3 1-3
1-3
100
Table 3: Evaluation results on Matbench-Discovar [59], MDR phonon
45 elastic tensor 16
and AdsorbML benchmarks [42]. Results are also provided for a diverse set of molecule 43 and molecular crystal [29] 23 benchmarks. NVE MD [2i] tests whether energy conservation is observed
when running the model for molecular dynamics. SOTA results from literature are reported where available. Additionally, for the materials evaluations, UMA models were fine-tuned on MPtrj [17
and sAlex 61 5 to be consistent with the benchmark's DFT settings.
Canalvais
Motecwlor
Molecaiar Crynnts
Bankon LO
Malhemnh
Fiee Deoggy Waee
KGA (GPg
NVE. MD
AdveMML Secura Bute
Eigotel shtate
NVE MD AE
CSP Tagss Conurm
RMCOO AA
Ga
(ee')
POB
Model
UMA
0.913
0.064
66.80%
4,56
127.7
194.7
Literaurre
COEN-OMM-OA
0.955
M
0.901
0.064
0.021
0.317
GRACE-2--AM
0.580
01667
0035
0.764
MACE-MPA-
0.652
0.073
0,28
0.422
0p22-0C20
60.80%
cSEN-en-ons,
4.56
147.3
21.6
GEN-S.OMC
6.15
0.74
0.18
and demonstrates the test set's challenging nature. Table we report benchmark results on the prediction of materials' thermodynamic stability, thermal conductivity, and vibrational properties. Note for these results, the UMA models are fine-tuned on the MPTrj [77 and sAlex 61 5 to be consistent with the DFT settings of the public benchmarks. On the popular benchmark Matbench- Discovery (MBD) [59, UMA-M achieves the highest Fl score to date. All UMA models of different sizes show strong performance, which suggests that further scaling of model sizes is unlikely to result better performance on the MBD benchmark. The conserving UMA models and M) also excel at phonon-related properties, which are reflected by metrics including KSRME, free energy MAF, and shear/bulk clasticity MAE.
Catalysis. We evaluate our models on two established catalysis benchmarks: OC20 S2EF [I] and AdsorbML 42]. These benchmarks focus on adsorption energy predictions that is, the models predict the change in energy as a molecule, known as an adsorbate, comes in contact with surface. Most previous models on the OC20 benchmark directly predict the adsorption energy given the adsorbate's 's position on the surface. Since our model predicts total energy, we make total energy prediction on the same structure, and another on a clean surface without the adsorbate II. These two values are subtracted (along with the energy of the adsorbate in isolation) to predict the adsorption energy. shown in Table 2 UMA models significantly outperform previous SOTA models on OC20 S2EF adsorption energy prediction reducing the errors by around 80%. On the more realistic
8

---

benchmark of AdsorbML, which evaluates a model's capability to predict the global minimum adsorption energy, all UMA models outperform the previous SOTA (EquiformerV2). In particular, UMA-L
achieves achivees a 25% improvement in the success rate (Table 12).
trained on a preview subset ( 70%) of the OMol25 dataset, since portions of the dataset were still being calculated when UMA model training began (results in Appendix E Both UMA-1 and UMA-1.1 were trained on the full OMol25 dataset. In Table 2 we report the prediction performance on two challenging test splits of OMol25: OOD-Comp and OOD PDB-TM. Our results show that UMA-S achieves performance comparable to the domain-spens -ppecii eSEN-sm-cons, while UMA-M significantly outperforms both models. In Table 3 we evaluate ligand strain and pocket-liganee interaction energy, which are quantities relevant to structure- discovery. In particular, for ligand strain (and, similarly, for conformers in Table [9]. UMA-M achieves errors approaching those of the DFT reference, demonstrating its ability to act as a practical surrogate for these calculations. On the distance scaling benchmarks (short and long range) in Table 3
Molecules. Initial UMA models
drug
UMA-S slightly surpasses -SEN-m--co despite having the same receptive field.
Molecular Crystals. In contrast to materials, molecules, and catalysis, the development and use of universal MLIPs for molecular crystals has been relatively underexplored. In addition the OMC25 test set, we evaluate our UMA's capability to accurately predict lattice energies, rank, and match structures of the most recent 7th Crystal Structure Prediction (CSP) Blind Test [29] In all evaluations. UMA-S outperforms the task-specific eSEN-S-OME baseline, which was trained on the OMC25 dataset 23). The high accuracy of lattice energy predictions ($ kJ/mol) indicates that UMA can be reliable substitute for DFT in many crystal structure prediction applications.
Metal Organic Frameworks. Computing the adsorption energy of CO2 for MOF sorbent is important for direct-air carbon capture applications and an established benchmark [67]. Similar to OC20, UMA models use total energy predictions compute the adsorption energy. UMA models perform on par with previous SOTA models on the ID test set (Appendix E) while achieving the best performance on the hardest ODAC OOD test set (Table 2), suggesting improved generalization.
5 Related Work
5.1 Universal MLIPs
The field of MLIP' has been rapidly improving due in part the availability of larger datasets. A recent example in the materials space the release of the Alexandria [61] and OMat24 [5] datasets, which quickly led to improved performance on the Matbench-Discover MLIPs trained on these materials datasets are often referred to as "universal" because they include nearly all the elements in the periodic table 112 6 20
leaderboard (59]. While
77 53 they may generalize well to
One reason for this is the chemical, structural, and elemental distribution shifts between materials and other domains such molecules. Another reason is the level of DFT theory that is considered accurate differs between domains, e.g. PBE for materials and WB97M-V for molecules. As we demonstrate, training across domains can lead to unified representations and help generalization.
other domains such as molecules (37) or surfaces
While there have been number of promising works exploring training MLIPs across multiple domains, it remains an open challenge to demonstrate high-accuracy zero-shot performance. One approach is pre-train model using large corpus of data (> 100M) and fine-tuning for specific tasks [65] (7).. The fine-tuned models are shown to perform significantly better than models trained from scratch. Nevertheless, removing the need for specialization would make these models substan- tially more useful, Recent studies suggest that achieving zero-shot performance across two DFT tasks may be possible 33 64
5.2 Scaling Relations
Empirical scaling laws provide wealth of insight into the relationships between compute, data, and model size. In LLMs, scaling laws motivated the community to train on more tokens with larger models because performance improvements became predictable (28 24 35) Additionally, scaling
relations help practitioners decide on how to best (compute-optimal) allocate resources such as dataset and model size. These types of relations have been explored in MLIPs to compare the asymptotic

---

scaling behavior of different model architectures and training paradigms. For example, [10] showed that equivariant models have different scaling behavior compared non-equivariant models.
6 Limitations
While UMA represents a significant step forward, there still limitations and areas for improvement. Currently, the UMA models are limited their ability to model long-range interactions. Our small and medium models use a standard MLIP cutoff distance 6Å, the actual receptive field is much larger due to message passing, but this can present problem for inputs where sets of atoms separated by more than 6Å. For example, if an adsorbate starts at 7Á from catalyst's surface the model views these as two independent non-interacting structures. Improvements may also be made in how charge and spin are incorporated into the model [17] Currently, each discrete charge spin uses separate embedding, which limits its ability to generalize to unseen spins or charges.
Our work aims to accelerate scientific discovery in chemistry and materials science using MLIPs; however, potential misuses exist, and we train on datasets designed for upplications beneficial to society to mitigate these risks.
Discussion and Conclusion
We explore techniques for training MLIPs across diverse DFT tasks using nearly 500 million training examples gathered over numerous datasets. By taking advantage of scaling relationships between dataset size, model capacity and loss, we are able to train models that are competitive with-or or better than-task specialized models. Naive approaches accurate but slow model. To address this, we introduced Mixture of Linear Experts, a method to increase model capacity while maintaining inference efficiency. Among family of models, UMA- S offers a favorable balance between speed and accuracy, capable of performing MD simulations of 1,000 atoms at 1.4ns per day single 80GB GPU.
scaling model parameters would result in
evaluated UMA across wide-range of benchmarks for materials, molecules, catalysts, molecular crystals, and metal organic frameworks, demonstrating strong performance across all. For well established benchmarks, such as AdsorbML and Matbench Discovery, we achieve new state-of-the. art results. Looking ahead, the development of more challenging benchmarks will be crucial for driving progress the field.
Our findings suggest that single model can achieve sufficient accuracy for practical research applications across a broad spectrum chemistry and materials science applications. This paves the way for universal MLIPs and opens new opportunities for atomic simulations at scale.
References
[1] Kareem Abdelmaqsoud, Muhammed Shuaibi, Adeesh Kolluru, Raffaele Cheula, and John R. Kitchin. Investigating the error imbalance of large-scale machine learning potentials in catalysis. Catal. Sci. Technol, 14:5899-5908 2024. doi: 10. 10.1039/D4CY0001155
[2] Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. Gpt-4 technical report. arXiv preprint arXiv:2303.00 08774, 2023.
[3] Qianxiang Ai, Vinayak Bhat, Sean M. Ryno, Karol Jarolimek. Parker Sorbberger, Andrew Smith, Michael M. Haley, John E. Anthony, and Chad Risko. OCELOT: An infrastructure for data-driven research to discover and design crystalline organic semiconductors. The Journal of Chemical Physics, 154(17):174705, May 2021. ISSN 0021-9606. doi: 10. 11163/5500487 14.
[4] Luis Barroso-Luque, Julia H. Yang, Fengyu Xie, Tina Chen, Ronald L. Kam, Zinab Jadidi, Peichen Zhong, and Gerbrand Ceder. SMOL: A Python package for cluster expansions and beyond. Journal of Open Source Softwre, ((77):5004, 2022. ISSN 2475-9066. doi: 10.21 105/ 045040
10

---

[5] Luis Barroso- Muhammmd Shuaibi, Xiang Fu, Brandon M Wood, Misko Dzamba, Meng Gao, Ammar Rizvi, C Lawrence Zitmick, and Zachary W Ulissi. Open materials 2024 (omat24) inorganic materials dataset and models. arXiv preprint arXiv:2410. 1n::::210.12771,, ,
[6] Ilyes Batatia, Philipp Benner, Yuan Chiang, Alin M Elena, Dávid P Kovács, Janosh Riebesell, Xavier R Advincula, Mark Asta, Matthew Avaylon, William Baldwin, et al. A foundation model for atomistic materials chemistry. arXiv preprint arXiv:2401 00006, 2023.
[7] Filippo Bigi, Marcel Langer, and Michele Ceriotti. The dark side of the forces: assessing non- conservative force models for atomistic machine learning. arXiv preprint arXiv:2412.11569, 2024.
[8] Anton Bochkarev, Yury Lysogorskiy, and Ralf Drautz. Graph atomic cluster expansion for semilocal interactions beyond equivariant message passing. Physical Review X, 14(2):021036, 2024.
[9] Stanislav S. Borysov, R. Matthias Geilhufe, and Alexander V. Balatsky, Organic materials database: An open-access online database for data mining. PLOS ONE, 12(2)=a0171001, February 2017. ISSN 1932-6203.
doi: 10.137 hoa.:://jomaallpoo..00177
[10] Johann Brehmer, Sönke Behrends, Pim de Haan, and Taco Cohen. Does equivariance matter at scale? arXiv preprini arXiv:2410.23179, 2024.
[11] Lowik Chanussot, Abhishek Das, Siddharth Goyal, Thibaut Lavril, Muhammed Shuaibi, Mor- gane Riviere, Kevin Tran, Javier Heras- Poras-Domingo, Caleb Ho, Weihua Hu, et al. Open catalyst 2020 (oc20) dataset and community challenges. Acs Catalysis, 1(100)6059-607) 6072, 2021.
[12] Chi Chen and Shyue Ping Ong. A universal graph deep learning interatomic potential for the periodic table. Nature Computational Science, 2(11))718 (0::18-882888 -788, 2022.
[13] Yongchul G. Chung, Jeffrey Camp, Maciej Haranczyk, Benjamin J. Sikora, Wojciech Bury, Vaiva Krungleviciute, Taner Yildirim, Omar K. Farha, David S. Sholl, and Randall Q. Snurr. Computation-ready experimental organic frameworks: A tool to enable high- throuhhhoughhut screening of nanoporous crystals. Chem. Mater, 26(21):6185 6192, 2014. doi: 10.1021/ cm502594j.
[14] Yongchul G. Chung, Emmanuel Haldoupis, Benjamin J. Bucior, Maciej Haranczyk, Seulchan Lee, Hongda Zhang, Konstantinos D. Vogiatzis. Marija Milisavljevic, Sanliang Ling, Jeffrey S. Camp, Ben Slater, J. Ilja Siepmann, David S. Sholl, and Randall Q. Snurr. Advances, updates, and analytics for the computation-reade experimental metal-organic framework database: Core mof 2019. J. Chem. Eng. Data, 64( 64(12):5985-599 5998, 2019. doi: 10.1021/acs.je 9000000 9b00835.
[15] Damai Dai, Chengqi Deng, Chenggang Zhao, RX Xu, Huazuo Gao, Deli Chen, Jiashi Li, Wangding Zeng, Xingkai Yu, Yu Wu, et al. Deepseekmoe: Towards ultimate expert specializa- tion in mixture-of-expeeee
language models. arXiv preprint arXiv:2401.06066 2024.
[16] Maarten de Jong, Wei Chen, Thomas Angsten, Anubhav Jain, Randy Notestine, Anthony Gamst, Marcel Sluiter, Chaitanya Krishna Ande, Sybrand van der Zwaag, Jose J. Plata, Cormac Toher, Stefano Curtarolo, Gerbrand Ceder, Kristin A. Persson, and Mark Asta. Charting the complete elastic properties of inorganic crystalline compounds. Scientific Data, 2(1):150009, March 2015. ISSN 2052-4463. doi: 10. 00.103/sdata.2015.99
[17] Bowen Deng, Peichen Zhong, KyuJung Jun, Janosh Riebesell, Kevin Han, Christopher Bartel, and Gerbrand Ceder. Chgnet as pretrained universal neural network potential for charge-informed atomistic modelling. Nature Machine Intelligence, 5(9): 2023.
[18] Peter Eastman, Pavan Kumar Behara, David Dotson, Raimondas Galvelis, John Herr, Josh Horton, Yuezhi Mao, John D Chodera, Benjamin P Pritchard, Yuanqing Wang, et al. Spice, dataset drug-like molecules and peptides for training machine learning potentials. Scientific Data, 10(1):11, 2023.
[19] William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. Journal of Machine Learning Research, 23 (120):1-39, 2022.

---

[20] Bruno Focassio, Luis Paulo M. Freitas, and Gabriel R Schleder. Performance assessment of universal machine learning interatomic potentials: Challenges and directions for materials' surfaces. ACS Applied Materials & Interfaces, 2024.
[21] Xiang Fu, Brandon M Wood, Luis Barroso-Luque Daniel S Levine, Meng Gao, Misko Dzamba, and C Lawrence Zitnick. Learning smooth and expressive interatomic potentials for physical property prediction. arXiv preprint arXiv:2502. 12147, 2025.
[22] Johannes Gasteiger, Muhammed Shuaibi, Anuroop Sriram, Stephan Günnemann, Zachary Ulissi, C Lawrence Zitnick, and Abhishek Das. Gemnet-oe: developing graph neural networks for large and diverse molecular simulation datasets. arXiv preprint arXiv:2204.02782, 2022.
[23] Vahe Gharakhanyan, Luis Barroso-Luque, Yi Yang, Muhammed Shuaibi, Kyle Michel, Daniel S. Levine, Misko Dzamba, Xiang Fu, Meng Gao, Xingyu Liu, Haoran Ni, Keian Noori, Brandon M. Wood, Matt Uyttendacle, Arman Boromand, C. Lawrence Zitnick, Noa Marom, Zachary W. Ulissi, and Anuroop Sriram. Open molecular crystals 2025 (omc25) dataset and models, 2025.
URL /ttps://///a (erg/abs/25000 2508..2250..0 022651
[24] Aaron Grattafiori, Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Alex Vaughan, et al. The llama 3 herd of models. arXiv preprint arXiv:2407.21783 2024.
[25] Stefan Grimme, Jens Antony, Stephan Ehrlich, and Helge Krieg. A consistent and accurate ab initio parametrization of density functional dispersion correction (DFT-D) for the 94 elements H-Pu. The Journal of Chemical Physics, 132(15):154104, April 2010. ISSN 0021-9606. doi: 10.1063/1.3382344.
[26] Bjørk Hammer, Lars Bruno Hansen, and Jens Kehlet Nørskov. Improved adsorption energetics within density-functione theory using revised perdew--urke-eemzerrof functionals. Physical Review B, 59(1 1):7413, 1999.
[27] Qiu He, Bin Yu, Zhaohuai Li, and Yan Zhao. Density functional theory for battery materials. Energy & Environmental Materials, 2(4): 27944422792200 2019.
[28] Jordan Hoffmann. Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbi, Aidan Clark, al. Training compute-optimal large language models. arXiv preprint arXiv:2203. /5556, 2022.
[29] Lily M Hunnisett, Nicholas Francia, Jonas Nyman, Nathan S Abraham, Srinivasulu Aitipamula, Tamador Alkhidir, Mubarak Almehairbi, Andrea Anelli, Dylan M Anstine, John E Anthony, et al. The seventh blind test of crystal structure prediction: Structure ranking methods. Acta Crystallographica Section B: Structural Science, Crystal Engineering and Materials, 80(6): 548-574, December 2024. ISSN 2052-5206. doi: 10. 10.1107/220525222440066777
[30] Lily M Hunnisett, Jonas Nyman, Nicholas Francia, Nathan S Abraham, Claire Adjiman, Srini- vasulu Aitipamula, Tamador Alkhidir, Mubarak Almehairbi, Andrea Anelli, Dylan M Anstine, et al. The seventh blind test of crystal structure prediction: Structure generation methods. Acta Crystallographica Section B: Structural Science, Crystal Engineering and Materials, 80(6): 517-547, December 2024. ISSN 2052-5206. doi: 10. 00.1077/520225222400049999
[31] Robert A Jacobs, Michael I Jordan, and Andrew G Barto. Task decomposition through com- petition in modular connectionist architecture: The what and where vision tasks. Cognitive science, 15(2):2192250, 1991.
[32] Robert A Jacobs, Michael Jordan, Steven J Nowlan, and Geoffrey E Hinton. Adaptive mixtures of local experts. Neural computation, 3(1):79-8) 07,779977771 1991.
[33] Albert Q Jiang, Alexandre Sablayrolles, Antoine Roux, Arthur Mensch, Blanche Savary, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Emma Bou Hanna, Florian Bressand, al. Mixtral of experts. arXiv preprint wxXiv:2401.04088 2024.
[34] Aaron D. Kaplan, Runze Liu, Ji Qi, Tsz Wai Ko, Bowen Deng, Janosh Riebesell, Gerbrand Ceder, Kristin A. Persson, and Shyue Ping Ong. A Foundational Potential Energy Surface Dataset for Materials, March 2025.
12

---

[35] Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprini arXiv:200100861,, 2020.
[36] Jaesun Kim, Jisu Kim, Jachoon Kim, Jiho Lee, Yutack Park, Youngho Kang, and Seungwu Han. Data-efficient multifidelitty training for high-fidelity machine learning interatomic potentials. Journal of the American Chemical Sociery, (4711):002-1154, 2024.
[37] Dávid Péter Kovács, J Harry Moore. Nicholas J Browning, Ilyes Batatia, Joshua T Hor- ton, Venkat Kapil, William C Witt, Joan-Bogdan Magdäu, Daniel J Cole, and Gábor Csányi. Mace-off23: Transferable machine learning force fields for organic molecules. arXiv preprint 1.21 2023.
[38] Georg Kresse and Jürgen Furthmüller. Efficient iterative schemes for ab initio total-energy calculations using a plane-wave basis set. Physical review B, 54(16):11169, 1996.
[39] Georg Kresse and Jürgen Hafner. Ab initio molecular dynamics for liquid metals. Physical review B, 47(1):558, 1993.
[40] Georg Kresse and Jurgen Hafner. Norm-conserving and ultrasoft pseudopotentials for first-row and transition elements. Journal of Physics: Condensed Matter, 6(40))8245, 1994.
[41] Georg Kresse and Daniel Joubert. From ultrasoft pscudopotentials to the projector augmented- wave method. Physical review b, 1758) 1999.
[42] Janice Lan, Aini Palizhati, Muhammed Shuaibi. Brandon M Wood, Brook Wander, Abhishek Das, Matt Uyttendaele, C Lawrence Zitnick, and Zachary W Ulissi. Adsorbml: leap in efficiency for adsorption energy calculations using generalizable machine learning potentials. npj Computational Materials, 911: 2023.
[43] Daniel S. Levine, Muhammed Shuaibi, Evan Walter Clark Spotte-Smith, Michael G. Taylor, Muhammad R. Hasyim, Kyle Michel, Ilyes Batatia, Gábor Csányi, Misko Dzamba, Peter Eastman, Nathan C. Frey, Xiang Fu, Vahe Gharakhanyan, Aditi S. Krishnapriyan, Joshua A. Rackers, Sanjeev Raja, Ammar Rizvi, Andrew S. Rosen, Zachary Ulissi, Santiago Vargas, C. LLwrrence Zitnick, Samuel M. Blau, and Brandon M. Wood. The open molecules 2025 (omol25) dataset, evaluations, and models, 2025. URL https:///rxiv www.aawwwwwwaaaaaaeeeeaaa wrgg/br/////oooeoooooe 08762
[44] Yi-Lun Liao, Brandon Wood, Abhishek Das, and Tess Smidt. Equiformerv2: Improved equivari- ant transformer for scaling to higher-degree representations. arXiv preprini arXiv:2306. 2tti::3306/12553 12059, 2023.
[45] Antoine Loew, Dewen Sun, Hai-Chen Wang, Silvana Botti, and Miguel A. L. Marques. Universal Machine Learning Interatomic Potentials are Ready for Phonons, December 2024.
[46] Andrea Madotto, Zhaojiang Lin, Chien-Sheng Wu, Jamin Shin, and Pascale Fung. Attention over parameters for dialogue systems. NeurIPS Conversational Al Workshops, 2020.
[47] Saeed Masoudnia and Reza Ebrahimpour. Mixture of experts: literature survey, Artificial Intelligence Review, 42:275-295 2014.
[48] Arslan Mazitov, Filippo Bigi, Matthias Kellner, Paolo Pegolo, Davide Tisi, Guillaume Fraux, Sergey Pozdnyakov, Philip Loche, and Michele Ceriotti. Pet-mad, lightweight universal interatomic potential for advanced materials modeling, 2025. URL
org/
abs/2503. 14118
[49] Fukasz Mentel. mendeleev A Python package with properties of chemical elements, ions, isotopes and methods to manipulate and visualize periodic table., March 2021. URI.httpp //github. wwt.:/.commmmmmmm com.//mmeentell/eedeeeeeee
[50] Frank Neese, Frank Wennmohs, Ute Becker, and Christoph Riplinger. The orea quantum chemistry program package. The Journal of chemical physics, 152(22), 2020.
13

---

[51] Mark Neumann, James Gin, Benjamin Rhodes, Steven Bennett, Zhiyi Li, Hitarth Choubisa, Arthur Hussey, and Jonathan Godwin, Orb: A fast, scalable neural network potential. arXiv preprint arXiv:2410.22570, 2024.
[52] Shyue Ping Ong, William Davidson Richards, Anubhav Jain, Geoffroy Hautier, Michael Kocher, Shreyas Cholia, Dan Gunter, Vincent L. Chevrier, Kristin A Persson, and Gerbrand Ceder. Python materials genomics (pymatgen): A robust, open-source python library for materials analysis. Computational Materials Science, (8:314-339, 2013..
[53] Yutack Park, Jaesun Kim, Seungwoo Hwang, and Seungwu Han. Scalable parallel algorithm for graph neural network interatomic potentials in molecular dynamics simulations. Chem. Theory Comput., 20(1 1):44574444866 2024. doi: 10.
[54] Saro Passaro and C. Lawrence Zitnick. Reducing SO(3) convolutions to SO(2) for efficient equivariant GNNs. In Proceedings of the 40th International Conference on Machine Learning, volume 202. PMLR, 2023.
[55] John Perdew, Kieron Burke, and Matthias Ernzerhof. Generalized gradient approximation made simple. Physical review letters, 77(18):3866 1996.
[56] Balázs Póta, Paramvir Ahlawat, Gábor Csányi, and Michele Simoncelli. Thermal Conductivity Predictions with Foundation Atomistic Models, September 2024.
[57] Raghunathan Ramakrishnan, Pavlo O Dral, Matthias Rupp, and O Anatole Von Lilienfeld. Quantum chemistry structures and properties of 134 kilo molecules. Scientific data, 1(1):1-7, 2014.
[58] Benjamin Rhodes, Sander Vandenhaute, Vaidotas Šimkus, James Gin, Jonathan Godwin, Tim Duignan, and Mark Neumann. Orb-v3: atomistic simulation at scale. arXiv preprint arXiv:2504.06231, 2025.
[59] Janosh Ricbesell, Rhys EA Goodall, Philipp Benner, Yuan Chiang, Bowen Deng, Alpha A Lee, Anubhav Jain, and Kristin A Persson. Matbench discovery-a framework to evaluate machine learning crystal stability predictions. arXiv preprint arXiv:2308.14999 2023.
[60] Carlos Riquelme, Joan Puigcerver, Basil Mustafa, Maxim Neumann, Rodolphe Jenatton, André Susano Pinto, Daniel Keysers, and Neil Houlsby. Scaling vision with sparse mixture of experts. Advances in Neural Information Processing Systems, 34:8583 8595, 2021.
[61] Jonathan Schmidt, Tiago FT Cerqueira, Aldo H Romero, Antoine Loew, Fabian Jäger, Hai-Chen Wang. Silvana Botti, and Miguel AL Marques. Improving machine-learning models in materials science through large datasets. Materials Today Physics, 48: 48:101560, 2024.
[62] MD Segall, Philip JD Lindan, MJ al Probert, Christopher James Pickard, Philip James Hasnip, SJ Clark, and MC Payne. First-principles simulation: ideas, illustrations and the castepcode. Journal of physics: condensed matter, 14(1 14(11): 2002.
[63] Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoe Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. arXiv preprint arxiv: 1177700065550 2017
[64] Tomoya Shiota, Kenji Ishihara, Tuan Minh Do, Toshio Mori, and Wataru Mizukami. Taming multi-domain, -fidelity data: Towards foundation models for atomistic scale simulations, 2024.
http.:: org/ absg/bbs/ae 2412. 13088
[65] Nima Shoghi, Adeesh Kolluru, John R Kitchin, Zachary W Ulissi. C Lawrence Zitnick, and Brandon M Wood. From molecules to materials: Pre-training large generalizable models for atomic property prediction. arXiv preprint ar arXiv:2310. 16802, 2023.
[66] Anuroop Sriram, Abhishek Das, Brandon M Wood, and C. Lawrence Zitnick. Towards train- ing billion parameter graph neural networks for atomic simulations. In International Con- ference on Learning Representations, 2022. URI. ltttpa://operreviee fet/ffoummiid
OjP2n0YFmkt
14

---

[67] Anuroop Sriram, Sihoon Choi, Xiaohan Yu, Logan M Brabson, Abhishek Das, Zachary Ulissi, Matt Uyttendacle, Andrew Medford, and David S Sholl. The open dac 2023 dataset and challenges for sorbent discovery in direct air capture, 2024.
[68] Anuroop Sriram, Logan M. Brabson, Xiaohan Yu, Sihoon Choi, Kareem Abdelmaqsoud, Elias Moubarak, Pim de Haan, Sindy Löwe, Johann Brehmer, John R. Kitchin, Max Welling, C. Lawrence Zitnick, Zachary Ulissi, Andrew J. Medford. and David S. Sholl. The open dac 2025 dataset for sorbent discovery in direct air capture, 2025. URL https:
org/
abs/2508.03162
[69] Annika Stuke, Christian Kunkel, Dorothea Golze, Milica Todorović, Johannes T. Margraf, Karsten Reuter, Patrick Rinke, and Harald Oberhofer. Atomic structures and orbital energies of 61, 61,489 crystal-forming organic molecules. Scientific Data, 7(1):58, February 2020. ISSN 2052-4463. doi: 10. 10.1008/441597-0000088----
[70] Hiteshi Tandon, Tanmoy Chakraborty, and Vandana Suhag. A brief review on importance of dft drug design. Res. Med. Eng. Sci, 7(4):791-795, 2019.
[71] Rithwik Tom, Timothy Rose, Imanuel Bier, Harriet O' O'Brien, Álvaro Vázquez-Mayagoitaa and Noa Marom. Genarris 2.0: random structure generator for molecular crystals. Computer Physics Communications, 100:107170, 2020.
[72] Richard Tran, Janice Lan, Muhammed Shuaibi, Brandon M Wood, Siddharth Goyal, Abhishek Das, Javier Heras- Drera--Domino,, Adeesh Kolluru, Ammar Rizvi, Nima Shoghi, et al. The open catalyst 2022 (oc22) dataset and challenges for oxide electrocatalysts. ACS Catalysis, 13(5): 3066-3084, 2023.
[73] Alexander Urban, Dong-Hwa Seo, and Gerbrand Ceder. Computational understanding of li-ion batteries. npj Computational Materials, 2(1):1-13, 2016.
[74] Hai-Chen Wang, Silvana Botti, and Miguel A. Marques. Predicting stable crystalline compounds using chemical similarity. npj Computational Materials, 7(1):1-9, January 2021. ISSN 2057-3960.
doi: 10. 0e00444111---20-040400
[75] Mitchell Wortsman, Gabriel Ilharco, Samir Ya Gadre, Rehecca Roelofs, Raphael Gontijo-Lopes, Ari S Morcos, Hongseok Namkoong, Ali Farhadi, Yair Carmon, Simon Kornblith, et al. Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time. In International conference machine learning, pages 23965-23998, PMLR, 2022.
[76] Hai Xiao, Jamil Tahir- Kheli, and William A Goddard III. Accurate band gaps for semiconductors from density functional theory. The Journal of Physical Chemistry Letters, 2(3):212-217, 2011.
[77] Han Yang, Chenxi Hu, Yichi Zhou, Xixian Liu, Yu Shi, Jielan Li, Guanzhi Li. Zekun Chen, Shuizhou Chen, Claudio Zeni, Matthew Horton, Robert Pinsler, Andrew Fowler, Daniel Zügner, Tian Xic, Jake Smith, Lixin Sun, Qian Wang, Lingyu Kong, Chang Liu, Hongxia Hao, and Ziheng Lu. Mattersim: A deep learning atomistic model across elements, temperatures and pressures, 2024. URL /top:tttt///axii.rrr//////// 220g/ 04967
[78] Naike Ye, Zekai Yang, and Yuchen Liu. Applications of density functional theory in covid-19 drug modeling. Drug Discovery Today, 27(5):1411-141, 2022.
[79] Duo Zhang, Xinzijian Liu, Xiangyu Zhang, Chengqian Zhang, Chun Cai, Hangrui Bi, Yiming Du, Xuejian Qin, Anyang Peng, Jiameng Huang, et al. Dpa-2: large atomic model as a multi-task learner. npj Computational Materials, 10(1):293, 2024.
[80] Alex Zunger, S.-H. Wei, L. G. Ferreira, and James E. Bernard. Special quasirandom structures. Physical Review Letters, 65(3):353-355 July 1990. doi: 10.103/.
15

---

Appendix Table of Contents
Training Details and Hyperparameters
17
A.1
Two-stage training
17
A.2
Parallelisms/
17
A.3
Model
17
MPuj and sAlex Fine-tuning
18
A.5
Training Compute Resources
18
A.6
Referencing and Normalization
18
Training Data
19
B.
Materials
20
2 Molecules
20
B.3
Catalysis
20
B.
Molecular Crystals
20
B.5
Metal-organic frameworks (MOFs)
20
Scaling Laws Methods
21
FLOP counting
21
C.2
Compute optimal fits
Fitting loss vs. parameters
D Inference
22
Eval vvaluatoons
Additional
CMA Results
24
2
Materials
24
Catalysis
25
Molecules
25
5.5
Molecular Crystals
27
E.6
Metal Organic Frameworks
27
Single-task vs Multi-task
28
Additional MoLE Analysis
28
G.1
Expert Analysis
28
G.2
Generalization Across Architectures
28
16

---

A Training Details and Hyperparameters
A.1 Two-stage training
While conservative models have been found to provide reliable performance in diverse physical property prediction tasks [21] the backward pass required for force/stress prediction significantly increases inference costs. Models with direct force prediction are much more efficient, and have been found to be effective as pre-training strategy to save compute when subsequently fine-tuned as conservative model [21 7. The UMA-S and UMA-M both follow this procedure. In addition, we introduce low-precision training, max-atom batching, and max neighbor switching to further enhance the scalability and efficiency of our training process. These novel strategies are discussed in turn below. Detailed hyper- parameersssssrss are summarized in Table 4
Precision. For pretraining, we used BF16 numerical format, commonly used in training LLMs but uncommon in MLIP models due high numerical precision requirements. In our experiments, we found that BF16 is significantly more stable than AMP-FP16 (automatic mixed precision) especially our multi-modal setting where data distributions can vary dramatically, frequently causing gradient and loss spikes that would destabilize AMP training. However, it suffers accuracy drop compared to AMP-IPI6 and FP32. We found the degradation can be nearly completely recovered after very small number of finetuning steps in FP32 (<1% % of data).
Max-atom Batching. Due the large differences in system topology and number of atoms/edges system, using a fixed number of systems as batch size is infcasible. Instead we chose to use max-atom batching scheme where randomly pack batches that contain as close an upper bound (max atoms) as possible without going over to guarantee upper bound on memory usage.
Max Neighbors. For training efficiency purposes, bors per atom during pretraining and found that it has no effect on the final performance, energy conservation, and smoothness properties of the model after finetuning with effectively "infinite" max neighbors.
significantly smaller number of neigh-
A.2 Parallelisms
Although our models are designed with inference efficiency in mind, training models with large number of MoLE experts is memory intensive, In particular, for the finetuning stages, a combination of infinite neighbors, FP32 precision, and autograd forces puts significant constraints on memory and training speed. We used combination three parallelism training techniques summarized as follows:
Partitioning graphs across GPUs during message passing layers is used when scaling up large number of atoms large model sizes. Graph partitioning is only used within a node with fixed graph parallel rank size (2 or 4) during conservative fine-tuning stages.
Graph parallelism (GP) [66]
Fully-sharded data parallel (FSDP): We use the Pytorch FSDPv1 implementation on MoLE expert layers only for models with total parameter count exceeding 1B during conservative fine-tuning when memory is scarce. Parameters are sharded within a node and replicated across nodes.
Distributed Data Parallel (DDP): We use the standard PyTorch DDP implementation, with modifications for per-atom loss averaging and compatibility with graph parallelism.
Furthermore, we leveraged Pytorch's Distributed Checkpointing framework ensure saving and loading extremely large checkpoints is efficient and stable across different node configurations. Exponential moving average (EMA) is used for stable validation performance.
A.3 Model
One model for all tasks. UMA is designed for multi-task learning under diverse DFT settings. For two inputs with exactly the same atomic numbers and positions, DFT labels will be different when different DFT settings are used. Such DFT settings include the level of theory and system total charge/spin. These task specifications are global information of the atomic system, and we process them through initial embedding layers. In this paper, five levels of theories involved
17

---

Table 4: Summary of main training-related typer-parameters for the pre-training and fine-tuning
stages. These hyper-parament parametera are shared among model sizes.
Pre-training
Finetuning
Precision
BFI6
FP32
Radius cutoff
6
Max neighbors
30
300
Force prediction
Direct
Autograd
Stress prediction
None
Autograd
Optimizer
Adam W
Adam W
Learning rate scheduling
Cosine
Cosine
Maximum learning rate
8 10
4 X 10
Warmup epochs
0.01
0.01
Warmup factor
0.2
0.2
Gradient clipping norm threshold
100
100
Model EMA decay
0.999
0.999
Weight decay
x 10-3
1 x 10-3
Energy loss coefficient
10
20
OMol energy loss coefficient
30
Foree loss coefficient
30
Stress loss coefficient
Table 5: Hyper-parameters for UMA models of different sizes.
Hyper-parameter
UMA-S
UMA-M
UMA-L
Number of MoLE experts
32
32
Dense
Number of layer blocks
4
10
16
Maximum degree Lmnax
2
Maximum order Maax
2
Number of channels Nehannel
128
128
256
Number of radial basis functions
611
128
256
Global batch size (atoms)
88k
44k
44k
Number pre-training steps
1.68M
2.08M
2.58M
Number of fine-tuning steps
IM
545k
350k
OMat24, OC20, OMol25, OMC25, and ODAC25 all use different DFT levels of theory. OMol25 further contains systems with non-neutral charge/spin.
For this iteration of the UMA models, these global information embedded follows:
Furthermore,
single model for all tasks, is crucial to normalize the labels such that targets from different datasets fall into similar numerical ranges. We specifically design referencing scheme that brings the diverse datasets to similar level, detailed in Appendix The model hyperparameters are shown in Table 6
A.4 MPtrj and sAlex Fine-tuning
For materials evaluations, UMA models were fine-tuned on the MPTrj 17 and sAlex 61 datasets to ensure consistent DFT settings. The fine-tuning procedure is the same as 30M-0AM as documented in [21]
A.5
Training Compute Resources
The resources used to train UMA models are described in Table 7
A.6
Referencing and Normalization
While each dataset used in this work comes with its own specific set DFT settings, we wanted referencing scheme that provides way to make energy magnitudes comparable across datasets. We
18

---

N Table 6:
for UMA models of different sizes.
Slype--arameters.
UMA-S
UMA-M
UMA-L
Number of MoLE experts
32
32
Dense
Number of layer blocks
4
10
16
Maximum degree Liuxx
2
Maximum order Miax
2
Number of channels Nehannel
128
128
256
Number of radial basis functions
611
128
256
Global batch size (atoms)
88k
44k
44k
Number pre-training steps
1.68M
2,08M
2.58M
Number fine-tuning steps
1M
545k
350k
Table 7: Training Times for UMA models.
Stage
in Parillel
Training Days
GPU-Type
UMA-S S
Direct Pre-train
128
H200 140GB
UMA-S
Conserve Fine-tune
256
5
H200 140GB
128
UMA-M
Conserve Fine-tune
256
14
H200 140GB
UMA-L
Direct Pre-train
128
H100 80GB
UMA-L
Stress Fine-tune
128
H100 80GB
UMA-L
FP32 Fine-tune
128
H100 80GB
do this through "heat of formation" (HOF) reference that is applied to the energies:
Fref = EDPT Σ
FI AH
Where EDPT corresponds the total DFT energy of the system, is the atom number, N is the total number of atoms in the system, DFD is the DFT energy an isolated atom in a box, and AHJ, f. is the heat of formation of atomic number as taken directly from Mendeleev 449 D was calculated using the DFT settings for each of the unique datasets in this work. Additionally, we apply linear reference to the above energies help with the convergence and training stability of our models. We follow the same protocol as described in the OC22 Appendix 722.
We a custom normalization x' ž= for all targets (energy, forces, stress), where the shift and the scale term Force root mean square (RMS). For the combined dataset the Force RMS is computed the weighted average (based on the number of systems) of all individual dataset Force RMS values.
B Training Data
summary of the five datasets used to train UMA is shown Table 8 In total, the dataset has 459 million training examples, containing up to 350 atoms. The average number of atoms varies based on the dataset, from 19 for OMat24 to 178 for ODAC25.
Table 8: Overview of the five datasets used to train UMA. For each dataset various statitises are provided alongside the sampling ratio used for training.
Dataset
Domain
Training Size
Labels
# Elements
Avg Size
Force RMS
Sampling ratio
OMat24
Materials
100.824.585
E.F.S
89
19
2.83
OMul25
Molecules
75.889.983
E,F
83
52
0.985
OC20++
Catalysis
229,054,045
E,F
56
77
0.624
OMC25
Molecular Crystals
24,870,226
E.F.
12
130
0.103
ODAC25
MOFs
28,517,826
E,F
70
178
0.046
Total
459.156.663
19

---

B.1 Materials
The field of inorganic bulk materials moving at an incredibly fast pace. Here we train on the Open Materials (OMat24) dataset (100M) [5], one of the largest and most diverse datasets in the community. All DFT calculations from this domain were run with VASP 33
41] and used the PBE (55] functional. Due to the differences in psuedopotential version and different pseudopotentials
38
40
for certain elements in OMat24 and Materials Project [3] calculation settings used for the data in most third party benchmarks, finetuning also performed on MPuj [[7] and subsampled Alexandria (sAlex) [6]] to ensure consistent evaluation on the materials benchmarks.
B.2 Molecules
The community has seen dozens of molecular datasets spanning different scales for a variety of applications []. However, the varying levels of DFT theory and quality makes it challenging to unify under a single model. The release of the Open Molecules 2025 (OMol25) dataset [43] helps address this by providing the largest single dataset (100M+) spanning 80+ elements covering metal- complexes, biomolecules, electrolytes, and several existing datasets under single, high-quality level of theory. All DFT calculations were performed using Orca [50] (w:B9MM- M/d-VV/df222-ZZDDD
At the time of training, only 75M samples from OMol25 were available for use. and we refer to this as Splits were constructed to ensure that this snapshot of the dataset is consistent with the full dataset release. All ablations and results were trained with this OMol-preview, unless stated otherwise. Released models, however, will be retrained with the full OMol25 dataset to ensure the best models are accessible by the community.
B.3 Catalysis
The Open Catalyst (OC20) dataset [1I] provides the largest adsorbate+surfacc dataset in the com- munity. OC20 enumerates 1M+ unique surface + adsorbate combinations, spanning 55 elements, and runs local geometry optimizations. Here, we train the OC20 All (133M), MD (38M), and Rattled (17M) datasets. Unlike prior work, we also leverage OC20's clean surface data (14M) since models here are trained total energies. One limitation of OC20 is that it only contains single adsorbates that interact with surface. To address this, we introduce the DC20-Multi-Adsorba (mAds) dataset (22M) better capture coverage effects and adsorbate-adsorbat interactions. All with the RPBE exchange-corellaton
DFT calculations were performed using VASP 39 38 40 functional functional (26).
B.4 Molecular Crystals
The most recent 7th Crystal Structure Prediction (CSP) Blind Test organized by Cambridge Crystallo- graphic Data Center (CCDC) demonstrated the effectiveness of tailored machine learning interatomic potentials (MLIPs) in predicting, filtering, and ranking molecular crystal structures [30 29). However, despite the widespread applications molecular crystals, there has been limited focus on developing universal MLIPs for molecular crystals, mostly because of the scarcity of publicly available datasets. Currently, publicly available datasets of molccular crystals are scarce, with at most 60,000 materials represented [3] 19. To address this data gap, use the Open Molecular Crystals (OMC25) dataset, which comprises 25 million molecular crystal structures. The dataset includes multiple relaxation trajectories of various molecular crystal packings generated with Genarris (71 starting from organic molecules from the OE62 dataset [69] The dataset includes 12 elements (all elements from OE62, excluding Li, As, Se, and Te) and the maximum number of atoms is capped 300. The dataset is computed using VASP (3 38 40 41 with the PBE exchange-correlation functional (55) and D3 dispersion correction 25 The OMC25 dataset, along with in-depth details on data and polymorph structure generation, will be released in an upcoming publication (2
B.5 Metal-organic frameworks (MOFs)
The Open Direct Air Capture 2025 (ODAC25) dataset represents the largest MOF dataset (78M) for
DAC applications [6] ODAC25 extends the earlier ODAC23 67 dataset, and is derived from the
COREMOF [3 14 dataset. ODAC25 focuses on the adsorption and co-adsorption CO2 and H₂O
in MOFs among other adsorbates, which is critical for DAC performance evaluation. This work uses
20

---

the subset of ODAC25 that overlaps with ODAC23, where all DFT- adsorption energies have been upgraded higher k-points sampling grid density, providing improved accuracy over ODAC23 (29M calculations). All DFT calculations were performed using VASP with the PВE exchange-correlation functional [55] and D3 dispersion correction
39
38
40 41
25
C Scaling Laws Methods
The scaling law experiments were only performed on pretraining, due to compute constraints, we did not study the effect finetuning with energy conservation. We used an 8-expert MoLE version of the model with the UMA-M settings (lmax = 4, mmaz = 2, and Rneighhors = 300 for consistency.
C.1 FLOP counting
We approximate our training FLOPs by
C(N,D) 2 RND,
where N is the total number of parameters in the network, and D is the number of inputs (tokens for LLMs and atoms or edges for MLIPs) computed on. This relationship holds for any network that dominated by linear layers where к indicates how many times a parameters is re-used on an input. For single forward pass of linear network,
2. full training cycle for such network with single backward pass brings K = 6 as we need to compute the gradient with respect to both the parameters and inputs. Thus, for most LLMs K = FLLOr//prameterrtoknn
[35 for single forward
and backwards pass where D is in units of tokens.
For our edge-based SO2 equivariant networks [54 K is function of bmax and Mmaa spherical harmonic orders and the number of edges per atom neneihhorrs For the scaling law experiments, we used Imax Mmas 2, and Reeighhors 30 which corresponds the settings of UMA-M model, resulting in k z= 270 L..Ps/paaameterratoono (or 9 FLOPs/parameter/edge per training step, which can be computed or experimentally determined. As long as the number input edges and parameters are sufficiently large (which holds for UMA models), this flop approximation holds as contributions from all other operators are marginal. We also verified this assumption with direct FLOP counting in our model code.
Overall, parameter reuse is significantly higher in ENui compared to LLMs, and hence the flop count is 2-3 orders of magnitude higher for a similar parameter-sized LLM network.
C.2 Compute optimal fits
The compute optimal model and dataset sizes can then be fitted to power laws
35
28
log( N (N)(C)) log(C) + A log (C) = 33 (o) (D((C)++)) +
(4)
(C' represents the optimal model size (in parameters) as a function of compute, and D (C) represents the optimal dataset size (in units of atoms). N (C') is determined by finding the minima of fitted parabolas for each Isoflop curve. The 10% and 90% percentile bootstrap errors are shown in Figure ee). a and / are the scaling coefficients w.r.t. model size and dataset size respectively. A and B are offset constants of the fit. Fit coefficients and bootstrap errors are shown in Table 9
Where C' is the compute in FLOPs described in Appendix C.I
C.3 Fitting loss vs. parameters
To understand the minimum achievable loss for dense vs MoLE models we can fit the more general parameterized ansatz of L(N, D) proposed by 35
À
(5)
No
D8
21

---

This maps the power law coefficients from Equation those in Equation by using a+B and = a+B3" Here we can either minimize the L(N, D) directly by fitting the 5 parameters
A, B, â, 3 with a iterative minimization procedure such LBFGS 28 10 or by examining the loss power relationship of N using
log L(N") log([V") +
(6)
where 7 - log([1 + 914) with E = We found both methods yielded similar results but the minimization of L. was more sensitive to the choice of hyperparameters.
Table 9: Power Law Coefficients determined from fitting Equations and 6 Error bounds are determined by bootstrap sampling 1000 times and taking the 10th and 90th percentile values, quoted in brackets.
Parameter
Dense
MoL.
0.61 (0.57,0.65)
0.56 (0.49,0.59)
B
0.39 (0.35, 0.43)
0.44 (0.39, 0.43)
A
-4.5 (-3.8, -5.3)
-3,8 (-2.56, -4.65)
B
3.6 (2.9, 4.4)
2.9 (1.6, 3.7)
â
-0.29 (-0.27, -0.31)
-0.25 (-0.2, -0.3)
2.16 (2.02, 2.34)
1.82 (1.61, 2.12)
D Inference
For inference benchmarking we use periodic fee carbon system with lattice constant a 3.&Å. This results in a fixed density of approximately 50 edges per atom within 6Å. For UMA models, we use a combination of torch.com compile, cuda graphs and pre-merged MoLE experts for inference speed. For large number of atoms (> 1000), we use edge-based activation checkpointing trade off memory for some speed, allowing to fit 100k+ atoms for the UMA-S S into memory. We checked all our optimizations chosen does not degrade simulation accuracy, equivariance or energy conservation properties. While our benchmarks do not include graph generation, our internal CUDA based graph generation algorithm is very fast and decreases throughput by no more than 10% even for the largest systems tested. In the case of non-MoLE merging, we found the inference speed was comparable but the parameters require more GPU memory store.
Table 10: Single-GPU simulation speeds for energy- conrry-conservvtive models in steps per second: compar- ing conservative UMA models the top two models (eSEN and OrbV3) on the Matbench Discovery leaderboard [59] and the MACE materials and molecules models. Benchmarks are run, excluding graph generation, on single Nvidia H100 80GB GPU using FP32 (TF32-high precision) and torch compile when possible. Test systems are standard periodic atomic system that have == 50 neighbors per atom with a 6Å cutoff. OOM indicates the model ran out of memory. Refer to Sec D for more
details.
Atoms
UMA-S
UMA-M (50M)
SSEN-30M- OAM (30M)
Orb-v3 conservative. inf-omat (25M)
MACE- MPA-0 (9M)
MACE-
(6.6M)
OFF23-1
(4.7M)
100
44
21
8
77
38
89
1.000
16
1.7
30
24
20
10,000
1.6
0OM
3.7
2.9
OOM
50,000
0OM OOM
0OM
OOM
OOM
OOM
100,000
0.1
0OM
OOM
00M
0OM
For fair comparisons against other models, used pytorch2.6.0. cuda12.4, python3. 12 and TF-32 precision universally on H100 80GB GPU. We use standard torch.compile settings whenever possible (only MACE-MPA-0 failed to compile). Different models have different radius cutoffs and
22

---

max neighbors settings. We made sure that all models was receiving roughly 50 neighbors per atom for the same number of atoms.
E Evaluations
Section E.1 provides results for all versions of UMA.
Sections E.2-E..6 provides additional results for the original version of UMA, which trained preview subset of the OMol25 dataset.
Table 11: Test MAE results on held out test splits for materials [59], catalysis T. molecules 43), molecular crystals [23] and MOFs [67]. All energies are in meV, forces are V/À and stresses are V/À Results for UMA are compared against the SOTA literature results. Target accuracies for practical utility are provided as an approximate guide for reference.
Muterials
Catafyais
Molecules
Molecialor crystals
MOFs
Energghhion
Ebergeiton
DXI DKID-8lth
Both
DXDD DXID-Comp
PDB-TM
Energ Energgniion
OMC-Tea
Energnidion
OOD-UT
Adh. Energg
WDM
Forces
Stocs
HEA
Focces
Forees
Slocah
Fvere
Model
UMA
UMA-S
20.0
60.8
4.4
22.0
72.8
3.1
52.1
24.3
70.2
30.9
3.64
10.80
0.88
16.12
0.91
4.77
0.97
292.4
16.0
3.2
16.0
10:37
UMA-L
17.6
45.5
3.8
24.8
48.3
2.8
32.4
12.2
43.5
15.9
2.33
5.19
0.81
8.76
0.59
2:28
0.10
291.1
6.5
19.4
62.2
21,9
73.5
3.1
24.5
70.4
31.2
8.25
0.53
15.56
0.93
5.15
1.01
287.1
13.6
UMA-S-111
20.2
62.8
24,9
3.5
51.5
24.1
68.8
30.7
0.95
8.64
0.84
15:47
1.00
5.04
0.93
289.9
13.3
UMA M-1.1
18.2
50.7
4,2
21.9
69.0
3.5
31.8
15.5
45.5
20.2
0.74
5.44
0.50
10.14
0.84
2.83
0.90
294.1
10.3
Literature
20.0
vyV2-OMal
14.9
46.3
3.6
20.3
47.0
2.7
31V2-0000
149
11.6
15.7
GemNes-OC21
163.5
16.3
343.3
25.1
eSEN
7.39
0.83
12.72
oSEN-S-OO0 SEN-S-OMC
1,05
5.39
0.94
eqv2-ODAC
31600
7.2
Target
Practccll Utilit
10:20
J0
100
100
MDR phonon 45], elastic tensor and AdsorbML benchmarks 42). Results are also provided for diverse set of molecule 43 and molecular crystal 22 23 benchmarks. NVE MD [2] tests whether energy conservation is observed when running the model for molecular dynamics. SOTA results from literature are reported where available. For the materials evaluations, UMA models were fine-tuned MPtrj [17] and sAlex 6 to be consistent with the benchmark's DFT settings.
Table 12: Evaluation results on Matbench-Discovers 59
16
Materials
Catalisis
Molecuier
Molecutar Crystals
Mutoena 2D
MAB (VVinonn
Farng
Elustshy M
Kia [GPs]
AdsoMMl: +
Lipond srram /mcc?)
Dosd LR (ow)))
CSP Targes 20
Energ [[d0ml]]
Kendalf Rack
RMSD
Prome
Phonons
(K]
Fire
Cou
NYE ND
DMMus
PDD
[mee]]
Din-SR
NVB XD
Conurr
RMSD
Model
UMA
UMA-S
0916
0.064
0.020
0.203
17.59
5.00
8.54
4.96
68.35%
4.39
150.3
67.6
432.1
2.685
0.82
0.12
UMA-M
0930
13.91
585.7
UMA-L
0928
0065
0.18
0.671
73.30
18.20
20.56
14.48
744.41%
3.37
71.7
16.5
296.1
2.488
0.34
0.12
UMA-S
0916
17.59
UMA-S-1
0914
0064
0.20
0.211
18.65
5.51
13.72
5.19
64.85%
5.19
158.3
23.8
2.57
081
0.13
0913
18.82
5.16
66.80%
127.7
26.7
256.7
2.13
0.86
0.13
UMA-M-LI
0.929
0.01
0018
0.176
14.81
3.87
8.57
4.78
72.25%
76.8
21.5
214.8
3.24
0.82
0.14
Literature
6SEN- 30M-0AM
21
0925
0.061
0018
0.170
15.00
4.00
9.13
5.73
0.90
0.200
SovenNar-M-orree
20
0.01
0064
0.021
03177
MACE-MPA-O
0.52
0.073
0.088
0.412
oqy2
+
$4.899 88%
eSEN
147.3
28.6
268.6
eSEN-S-OMC
6.18
0.74
0.18
23

---

E.1 Additional UMA Results
In this section, we provide results for the UMA, UMA-1, and UMA-1.1 models. The original UMA models (S, M, L) were trained on preview subset of the OMol25 dataset (~2 70%), since portions of the dataset were still being calculated when UMA model training began. Both UMA-1 and UMA-1.1 were trained on the full OMol25 released 05/14/2025. The other datasets used for training were unchanged, with the exception of the sampling ratios (Table ( where the OMat24 coefficent was changed from 4 extensive bug later discovered. These models were fine-tuned on our 1.0 models with the fix in place. Tables [[]andd T2 correspond to Tables a i the main text. Additional UMA I series models will be included in the arviv version of the paper
for both UMA-1 and UMA-1.1. UMA-1.1 resolved size
E.2 Materials
Table 13: Materials validation and test evaluations from OMat24 S and HEA, All energies are in Ve forces are in m/À and stresses are in me V/A
NaY S
Joer
Composition
Eoegynnn
Poice Coiiee
EigggAAn
Annn
Shne
Teces
Ponces
UMA
UMA-S
11.3
57.1
2,9
0.98
20.0
60.8
4.4
11.5
57.0
3.0
9.9
56.9
2,6
22.0
12.8
3.1
10.0
10.2
8.5
LMA-1
0.7
Literature
300 OMe
10.7
0.59
16.2
4.1
10.7
9.0
20.0
59.5
a/V2-66M-OMB
1000
Table 14: Materials evals results. We note that UMA models are fine-tuned on the MPTrj and sAlex 61
5 datasets to be consistent with the DFT settings of the benchmarks.
Matbench SY
Kappa 101
MDR
Elantcirs
Elantiiry
NYE MD
Energ.
MAE
Maag MAE [A]
Estray MAB
C.. MAE
Gum- MAB
K.a- MAE
Cade MAR 1GPg
K.a- MAE
Cunservaive
ProcusoSE
DAP
MAE
MAE
[GPu
[GPul
Model
[K3
UMA
UMA-S
0.92
6.00
0.92
0.97
0.02
0.07
0.87
0.09
0,20
17.59
741
13.59
3.49
5.00
8.54
4,96
8.57
5.25
UMA-M
0.93
6.08
0.98
0.02
0.07
0.09
0.20
13.91
5.11
2,66
3.39
8.40
4.76
7.07
4.75
UMAL
0:93
6.09
0.53
0.98
0.02
0,07
0.86
0.45
0.67
78.50
27.68
43.04
15.85
18.20
20.56
14,48
21.95
17.01
Literature
CSEN-3IM-GAM GAM
6.07
0.98
0.07
1021
3,00
4.00
9.13
5.73
5.73
c4V2 B6M OAM
1
0.92
6.05
0.92
0.98
0.02
0.07
0.85
1.82
1.94
840.33
377.96
426.79
102.72
251.14
19.60
26.25
22.02
26.50
Table both validation and test results following OMar24 (5) along with the new high entropy alloy HEA test introduced in this paper. The HEA dataset contains relaxation trajectories for over 5000 alloys with atomic configuration disorder. Input structures were generated by sampling metallic element combinations of up to 6 different unique elements and using the special quasirandom structure method [80] 4 to decorate
Full results on materials' benchmarks are in Tables 14
face-centered cubic, body-centered cubic and hexagonal close packed structures. DFT relaxations
were carried out following the settings used in the OMat24 dataset [
In Table 14 we show full results for Matbench discovery [59 MDR phonon, elastic tensors, high entropy alloy IS2RE and NVE MD conservation benchmarks. The Matbench-Discovee benchmark evaluates a model's ability to predict ground-state thermodynamic stability by optimizing geometry and predicting energy. The thermal conductivity prediction task demands accurate modeling of harmonic and anharmonic phonons, which are crucial for precise predictions of thermal transport. The MDR Phonon benchmark assesses a model's 's performance in predicting phonon and vibrational thermodynamic properties. The MP elastic constant benchmark tests model's 's accuracy in predicting
24

---

bulk and shear moduli, requiring precise calculations of stress tensors and their derivatives with respect cell deformations.
E.3
Catalysis
For catalysis, show full validation and test results for OC20 11 in Table 15 the dataset contain molecules, called adsorbates, interacting with surfaces. The goal is to estimate the adsorption energy, which is the change in energy as the adsorbates come into contact with the surface, and the forces on the atoms. Force MAEs are comparable across UMA-M and UMA-L to other state-of-the-an art models. Adsorption energies for UMA are calculated by subtracting two total energy calculations as described in the main text, which improves results over prior models.
The structures in
In addition energy and forces MAEs, practically relevant task finding the global minima adsorption energy 42). One limitation of the originally proposed AdsorbML pipeline is that it requires DFT evaluation of the ML-identified global minima structure. To make this benchmark more accessible to those without access to DFT, propose slightly altered version of this benchmark that only requires MI.. Here, follow the same AdsorbMI. pipeline but allow the use of the MI.-predicted global minima energy, but success is now only considered if the energy is within 0.1 eV of the DFT minima. Previously, success did not enforce lower bound so long a DFT evaluation confirmed that energy prediction is realized. Without DFT, we also set lower bound of 0 eV provide some flexibility finding better global minima than DFT. We show that metrics are still highly correlated when you perform the original evaluation the one proposed here.
AdsorbML to evaluate the performance the
Table 15: Catalysis validation and test results on OC20 metrics. All energies are in meV and forces are in me V/A.À.
Eaoggt
Eeergi 80D
000
Medil
UMA
UMA-5
61A
31.1
0.63
H17
0355
57.1
26.1
212
984
UMA-M
43.1
15.8
0.73
9000
19.2
0.35
33.4
160
46.
21,
LM44L
32.6
12.0
0.77
49.
145
0.39
43.5
153
Liternture
SVV2-0C20
1091
11.6
106.5
1578
1823
15:33
143.3
2311
E,4 Molecules
Table 16: Open Molecule 2025 43
validation evaluations across biomolecules, electrolytes, metal complexes, neutral organics and OOD-comp. All energies are in meV and forces are in meV/À. All models are trained with preview OMol25.
Biomolecales
Electrolytes
Menul Compleses
Neatral Organicr
OOD-Comp
Enegy/Amm
EnngyWAao
EneryWAtom Alem
Enery/
Energy/
Force
Fore
Fowe
Model
UMA-S
0.53
5.89
2.89
11.65
4.63
37.85
1.00
13.15
3.62
12.02
UMA-M
0.44
10.21
UMA-L
0.33
2.90
1.13
4.52
3.35
24.85
0.65
5.02
2.39
3.83
Baseline
0.54
6.06
12.78
We report results on molecules following the Open Molecules 2025 43 These include energy and force estimates for validation and test splits, as well as, numerous other tasks. The validation and test results are shown in Tables a7] respectively. Note that the eSEN-S-OMol model is only trained on the preview OMol25 dataset, which is 70% of the full OMol25 dataset for fair
25

---

Table 17: Open Molecule 2025 43 test evaluations biomolecules, electrolytes, metal com- plexes, neutral organics and OOD-comp, metal ligand, PDB-TM, reactivity, COD and anions. All energies are in meV and forces are in meV/Á. All models are trained with preview OMol25.
Biomolecules
Eleetrolpeez
Mosal Complexes
Nestral Organies
O0D-Cummp
Metal Ligund
PDB-TM
Reacibitty
COD
Anious
EnagylAnee
Ener maagg/Aoo
EnsgyiAn
Enegy/Aron
Enery/AAoo
Enegy/Acco
Eneey
Rvce
Pocce
Encr
Fooce
Feea
Force
Fooe
Focce
UMA
UMA-S
0.51
5.70
3.80
13.95
3.07
33.56
1.49
20.35
3.64
10.80
1.23
17.71
0.88
16.12
4.82
61.80
2.92
29.82
0.66
10.85
UMA-M
3.89
13.19
24.85
10.83
3.26
0,99
47.75
20.11
8.03
UMA-L
21.67
2,33
1.05
41.97
15.44
Baseline
S-ONMl
0,52
6.06
3.52
15.25
2,86
33.30
1.35
20.06
3.67
11.56
1,18
17.49
0.79
14.11
4.89
61.16
2.0
24.12
0.47
10.38
comparison with the UMA models that were trained on the same subset. In general, the UMA-S and cSEN-S-OMol models provide comparable results.
Table 18:
single point evaluations for protein-ligand, IE/EA, spin gap and distance scaling. All energies are in me and forces are in me V/Å. All models are trained with preview OMol25.
Open Molecule 2025 48
Frotede Bggnn
IEAA
Spin pap
Distance scaling
A Eonryy (SR)
Finos (SR)
Tenon
A. Fepoes
voaler als.
MAE
MAE
MAE
MBEE
Model
UMA
150.25
5.09
336.8
80.60
0.77
665-75
112.09
0.60
67.60
4.11
432.14
5.54
66.25
41.60
Baselines
<SEN-S-IMol
73.02
Table 19: Open Molecule 2025 4 optimization evaluations including ligand-straie conformer prediction, and protonation states. Results are reported across a variety energy and structure based metrics for each task. All energies in meV. All models are trained with preview OMo125.
Ligand stainn
Conformers
Protonation
Strain evergy
arcublle JL
(neV))
Elngy rapp (outty
ansmmbe (4
MAE (me
Jine/1
RMSD
(4 JL
RMSD
RMSD
butt, A
A Energy
ROSD
rejel 1u
RMSD
Energyy
RMSD
rape M4
A Energy
Model
UMA
UMA-S
4.39
0.28
0.06
0.05
5.76
0.02
3.06
0.13
40.42
0.04
18.35
UMA-M M
UMA-L
3.37
0.25
0.06
0.05
4.97
0.01
2.27
0.11
30.31
0.02
12.54
Buselines
6SPN-sDDM
0.06
0.00
32.82
OMol25 [43] provides numerous other tasks for evaluation. These include evaluations which only require the estimation of a single point DFT calculation and evaluations that require optimizations. The comparison for single point DFT calculations is shown in Table 18 Similar to the test results, the UMA-S and eSEN-S OMol models perform similarly with the UMA-M and UMA-L performing significantly better. Table 19 shows the results for tasks that require optimizations. which require repeated calls to the model update atoms positions until local energy minima is found. The small models report similar performance, and the UMA-M demonstrates the highest accuracies. It is likely that UMA-M outperforms UMA-I due to UMA-M being energy conserving and better behaved during optimization tasks.
26

---

Table 20: Open Molecular Crystals 2025 forces are in me V/À and stresses are in meV/A
23 validation and test table. All energies in meV,
Val
Test
Enogy/Ann
EnegytA Abo
Force Costre
Forza
Stres
UMA
0.98
UMA
0.95
UMA-L
0,6
23
0.3
0.6
0.5
2.3
0.1
0.96
Basclines
CSEN-S-OMC
1.06
5.58
0.96
0.2
1.08
5.39
0.94
0.92
Table 21: Open Molecular Crystals 2025 Phase of CCDC 7th CSP Blind Test
23 evaluation for polymorphs from the Structure Ranking
29
All lattice energies are per molecule basis in kJ/mol.
CCCC Polysorphs
Kank condumi
Morch nore
Model
UMA
UMA-M
2.66
3.71
0.60
0.81
0.1
0.13
0.07
0.99
UMAL
2.49
3.70
0.81
0.84
0.95
0.12
0.07
1.00
Baselines
SEXSSOMC
6188
7.38
0.07
0.74
0.7
0.18
0.08
0.91
E.5
Molecular Crystals
Open Molecular Crystals (OMC25) 23 evaluates whether model can predict the packing of molecule into crystal structures. This task requires the accurate estimation of inter-molecular forces. Results on validation and test splits are shown Table 20 is notable that all sizes of UMA models outperform the eSEN-S-OMC model trained on only OMC25. This indicates that the other datasets, such as OMol25, provide useful complementary information for task. One important and real- world task for molecular crystals is to predict the lowest energy packing, called polymorph, for molecule. Results for this task for a subset of molecular crystal polymorphs from the most recent 7th Crystal Structure Prediction (CSP) Blind Test (29 shown in Table 21 The pymatgen's [52] StructureMatcher class with default settings is used to match DFT and UMA-relaxed polymorphs, and root mean square deviation (RMSD) is computed for matches. Similar to the test metries, the UMA models outperform the eSEN-S-OMC trained on only OMC25.
E.6 Metal-Organic Frameworks
Table 22: OpenDAC 67 val and test table. All energies are in meV and forces are in meV/á.
Var rowel Kiweg))
TorriAt Eoregy
000-U/T
Depp
EL00
Medel UMA
59
2924
38
290.2
35
Literature
0189
3188
The results on OpenDAC 67 are shown in Table 22 The OpenDAC dataset contains Metal-Organic Frameworks (MOFs) with CO2 and water molecules. The goal is to estimate the change in energy in the presence with and without the CO2 and water molecules. These adsorption energies are computed in the same manner for catalysts. The use of total energies leads to significantly better adsorption
27

---

Single-task
Multi-task
Single-task
Multi-task
Dirua OMCET
Nomaltone Eneegy MAB
Norm Energy MAC
Nomal MIE
Home Bod Force MYE
0MCC2
0.
03
0.0
0.
0.0
0.2
00
Normalzedd Epoch
Normal lood Eooch
Normalzod Epoch
Norma Lood Epoci
Figure 4: Pre-training curves of UMA-L for both single-task and multi-task models. Errors are normalized based single-task performance. Note single-task models can overfit (forces right), and the multi-task model generally converges to lower errors.
energy estimates, similar to catalysis. The forces of UMA-M and UMA-L are similar to the SOTA cqV2-ODAC [67] model.
F
Single-task vs Multi-task
For large models, multi-task training offers benefits even without MoLE. In Figure 4 plot the direct-force pre-training curves of UMA-Ig and single-task models with the same model architecture and size. All metrics are normalized to those achieved with the models trained on single tasks to easily compare their relative performance with multi-task models. We observe that single task models frequently overfit to forces (OMat overfits upon further training), while the multi-task UMA model does not. Furthermore, UMA achieves lower losses most cases. The one exception is OMol forces, for which errors are already small (< 10 me V/À) for both models.
G Additional MoLE Analysis
G.1 Expert Analysis
Using a limited validation set consisting of 10,000 OMat24, 5,000 OC20, 20,000 OMol25, 10,000 OMC25, and 5,000 ODAC23 samples, we calculate the mean expert coefficient for each element- expert pair across all systems where the pair appears. We visualize these results using 32 periodic tables, each representing
of the 32 experts in UMA-S (Figure
Additionally we visualize the expert cofficient mean and variance across datasets (tasks). Many experts utilized in OC20 are also used in OMat24. In contrast, the experts associated with OMol25 show minimal overlap with other datasets, aside from a small subset shared with OM25C. Lastly ODAC23 and OMC25 utilize the fewest experts, and these two share single expert between them
G.2 Generalization Across Architectures
Table 23: Testing the generalization capability MoLE layers. We applied 8-expert MoLE layers to both e2EN[21] and EquiformerV2[A. model architectures and measured their relative performance on direct pretraining against the versions without no MoLE, For each row, we show the relative improvement over the baseline.
OMol25 Val subset
eSEN 6M
EquiformerV2 6M
Metal Complexes
+10%
+10%
Spice
+5%
+5%
Neutral organics
0%
+3%
Biomolecules
+25%
+21%
Electrolytes
+12%
+ 10%
Mean across subsets
+10 %
+10 %
We performed additional analysis to ablate the benefit of the MoLE layers. Since our MoL. implement very general and can be applied any linear layers inside model, we ran experiments adding MoLE
28

---

Expert 1
Expert
Expert 3
Expert
Expert
Expert
Expert 7
Expert
-10
Expert
Expert
Expert
Expert 12
-15
Expert
Expert
Expert
Expert
-25
Expert 17
Expert
Expert 19
Expert
-30
Expert 21
Expert 22
Expert 23
Export 24
-35
-4.
Expert 25
Expert
Expert 27
Expert 28
Expert
Expert
Figure 5: Log mean expert coefficient across element-expert pairs.
both eSEN and Equiformers architectures, controlling for model size (6M base parameters) and 8 training epochs on Omol dataset, direct pretraining only.
23 shows that despite EquiformersV2 and eSEN being very different architectures, the relative benefit of adding MoLE is fairly consistent both cases. This suggests that MoLE is simple and general approach boost the capacity of MLIP models without incurring inference speed and memory overhead.
29

---

Variance (log scale)
0.04
oc20/id
omat/rattled-300
Variance
omol/val
0.00
omc/val
odac/val
0.00
20
Expert
Mean (log scale)
0.50
oc20/id
omat/rattle-3300
0.17
omol/val
0.06
Mean
omc/val
0.02
odac/val
0.01
15
20
25
Expert
Figure 6: Log mean expert coefficient across element-expert pairs.
30

---

NeurIPS Paper Checklist
The checklist is designed to encourage best practices for responsible machine learning research, addressing issues of reproducibility, transparency, research ethics, and societal impact. Do not remove the checklist: The papers not including the checklist will be desk rejected. The checklist should follow the references and follow the (optional) supplemental material. The checklist does NOT count towards the page limit.
Please read the checklist guidelines carefully for information on how to answer these questions. For each question in the checklist:
You should answer [Yes], [No]. or [NA].
[NA] means either that the question is Not Applicable for that particular paper or the relevant information is Not Available.
Please provide short (1-2 sentence) justification right after your answer (even for NA).
The checklist answers are an integral part of your paper submission. They are visible to the reviewers, area chairs, senior area chairs, and ethics reviewers. You will be asked to also include (after eventual revisions) with the final version of your paper, and its final version will be published with the paper.
The reviewers of your paper will be asked to use the checklist as one of the factors in their evaluation. While "[Yes] is generally preferable to "[No]" it is perfectly acceptable to answer "[No] provided proper justification is given (e.g., "error bars are not reported because it would be too computationally expensive" or "we were unable to find the license for the dataset we used"). In general, answering "[No] or "[NA] is not grounds for rejection. While the questions are phrased in binary way, we acknowledge that the true answer is often more nuanced, please just use your best judgment and write justification to elaborate. All supporting evidence can appear either in the main paper or the supplemental material, provided in appendix. If you answer [Yes] to question, in the justification please point to the section(s) where related material for the question can be found.
IMPORTANT, please:
Delete this instruction block, but keep the section heading "NeurIPS Paper Checklist",
Keep the checklist subsection headings, questions/answers and guidelines below.
Do not modify the questions and only use the provided macros for your answers.
1. Claims
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Question:
Answer: [Yes]
The main claims made in the abstract and introduction accurately reflect the paper's s contributions and scope. The abstract provides short summary of the innovations and key findings presented the paper, while the introduction outlines the research problem, objectives, and the significance of the work. Both sections are aligned with the results discussed in the subsequent sections of the paper.
Justification:
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
31

---

Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: We discuss some of the limitations of our work in the discussion section, including long-range interactions and incorporation of charge and spin.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex system might not be used reliably provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [NA]
Justification: [NA]
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
32

---

Justification: We describe the architecture, especially the novel features, in detail as part of the main text. We provide descriptions of the training procedure and the datasets used in main text with additional details in the appendix. For established evaluations we cite the standard procedure and the evaluation new we describe it. These cover the main components needed to reproduce the main results of the paper.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is a dataset and/or model, the authors should describe the steps taken make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (c.g., in the case of large language model), releasing of a model checkpoint, other means that appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm. (b) the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors welcome deseribe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification:
All code and data will be available upon publication.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.co cc/
for more details.
While encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https //nips. http:///wwpp./eeeeeee////eeeeaaa
for more details.
33

---

The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data, The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments should state which ones are omitted from the script and why.
reproducible, they paper) is recommended, but including URLs to data and code is permitted.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible supplemental material (appended to
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We provide full details about the training procedure, hyperparameters, datasets, and evaluations in the appendix.
Guidelines:
The answer NA means that the paper does include experiments.
The experimental setting should be presented in the core the paper to a level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, appendix, as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [No]
Justification: Given the computational expense, we did not train ensembles models.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
Question: For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
34

---

Answer: [Yes]
Justification: We provide details about the resources used for training our main family of models (small, medium, large) in the appendix. This includes breakdown of the various training stages.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental rums well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code of ethies
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics http:://wweuappssccccclpe/////ee
Answer:
[Yes]
Justification:
We respect the NeurIPS code of Ethics and our best to preserve anonymity,
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No. they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (c.g.. there a special consid- eration due to laws regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative socictal impacts of the work performed?
Answer: [Yes]
Justification: We discuss the positive societal impact in the introduction/resulls and discuss potential misuse and how that risk is mitigated in the discussion.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g.. deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
35

---

11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [Yes]
Justification:
The model will have appropriate terms of use when released.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [No]
Justification: We have cited any assets used (to the best of our ability), but we do not explicitly state the license and terms of use for every asset.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include a URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms use in the package should be provided. For popular datasets, paperssi thcodee com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re- re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors encouraged to reach out to the asset's 's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: When we introduce new assets such data/code/models we make sure they are well documented. Upon publication, we plan to include URLs to access all new assets including detailed license information. These have not been included in the initial submission due to the difficulty in remaining anonymous.
Guidelines:
The answer NA means that the paper does not release new assets.
36

---

Researchers should communicate the details of the dataset/code/mmodel as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: [NA]
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: [NA]
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if' applicable), such the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Answer: [NA]
Justification: [NA]
Guidelines:
The answer NA means that the method development in this research does not involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy (https /ttpap:////uri.... cct.e///eeeeeeecc/222555.. for what should or should not be described.
37


Reviews:
[{"review_id": "mXg2GU7N7C", "reviewer": "Reviewer_tgt6", "summary": "This paper presents UMA, a novel family of Universal Models for Atoms, aiming to achieve state-of-the-art speed, accuracy, and generalization in atomic simulations across diverse chemical domains. The authors leverage an exceptionally large dataset (half a billion 3D atomic structures) and introduce a Mixture of Linear Experts (MoLE) architecture to efficiently scale model capacity. The work demonstrates that a single UMA model can perform comparably to or better than specialized models without fine-tuning across various applications in chemistry and materials science.", "questions": "1. The manuscript frequently refers to an \"Appendix\" (e.g., Appendix A, B.6, C, D, E, F). However, these appendices appear to be missing from the provided PDF document. Could the authors please ensure that the complete supplementary material, including all referenced appendices, is included in the final submission?\n2. While the paper introduces the Mixture of Linear Experts (MoLE) architecture and applies it to the eSEN base, the novelty of MoLE is very limited.  Could the authors provide a clearer summary of the specific model or theoretical innovations developed in this paper beyond the large-scale data training? Additionally, a more comprehensive ablation study comparing MoLE against alternative architectural choices or existing MoE implementations would strengthen the claim of its efficiency and effectiveness. In its current form, the paper leans more towards a technical report on large-scale training rather than highlighting fundamental architectural breakthroughs.\n3. The authors claim that UMA models are designed to push the frontier of speed, accuracy, and generalization, and that a single model can perform similarly or better than specialized models without fine-tuning. However, upon reviewing Table 2, it appears that UMA models do not consistently outperform all specialized baselines. For instance, in the \"WBM Energy\" column, the eqV2-OMat model (14.9 meV) shows a lower error than even the largest UMA model, UMA-lg (17.6 meV). Similar observations can be made for other metrics. Could the authors clarify these discrepancies and elaborate on the specific contexts or metrics where UMA's generalization truly translates to superior performance compared to highly specialized models?\n4. Following up on Table 2, it is observed that UMA-lg, despite having the largest number of parameters, does not always exhibit the best performance among the UMA family (e.g., UMA-md often performs comparably or better). This seems to deviate from the expected behavior suggested by the empirical scaling laws presented in the paper (e.g., Figure 3). Could the authors explain why the scaling laws might not fully apply to these specific evaluation tests, or if there are other factors contributing to this non-monotonic performance with increasing model size?\n5. Regarding the training data, the paper mentions combining datasets from various chemical domains. It is known that different DFT calculators and levels of theory (e.g., PBE functional with plane-wave codes for materials vs. ORCA software with Gaussian basis sets for molecules) result in different reference zero-energy lines. While an \"energy referencing scheme\" is briefly mentioned (line 161), but the appendix B.6 is not given in the pdf. Could the authors provide a more detailed explanation of how these disparate setups across different domains were reconciled during data compilation and model training to ensure consistency and prevent performance degradation?\n6. The \"Related Work\" section (Section 5, starting line 293) provides a good overview but appears to omit several highly relevant and recent works in the field of universal MLIPs or large-scale atomistic modeling. Specifically, works such as Orb, PET-MAD, MatterSim, DPA, and SevenNet also contribute significantly to this field. Including these references would provide a more comprehensive overview of the current landscape and better contextualize UMA's contributions.", "limitations": "The authors did not discuss the limitations of their model in the paper. I would strongly suggest the authors to answer my questions above and put some of the answers as part of the limitation discussion in the main text.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "* Strengths:\n\n1. The paper addresses a critical challenge in MLIPs by aiming for a single, generalizable model across a wide range of chemical domains (molecules, materials, catalysis).\n2. Very large training data size, probably largest to date.\n3. The authors demonstrate scaling laws\n\n* Weakness:\n\n1. This paper looks like a technical report, instead of a research paper. \n2. The paper only contains very limited theoretical innovation and it only trained models on top of existing architecture on a large dataset.", "was_revised": true, "final_justification": "Thank you for your response. Your rebuttal has addressed most of my concerns. While I still have some reservations regarding the novelty of the work, I appreciate the clarifications provided and have accordingly updated my evaluation to a borderline accept.", "rebuttal": "We would like to thank reviewer tgt6 for their time and the constructive feedback! We address each of the reviewer’s comments/questions below. \n\n**Questions**\n\n>Q1: The manuscript frequently refers to an \"Appendix\" (e.g., Appendix A, B.6, C, D, E, F). However, these appendices appear to be missing from the provided PDF document. …\n\nOur appendix was included in the Supplementary Material, which is standard practice for NeurIPS and was available for reviewers to download (see link on review page). We apologize it wasn’t easier to find, but at least one other reviewer referenced data from the appendix, so it was accessible at the time of review. \n\n>Q2: While the paper introduces the Mixture of Linear Experts (MoLE) architecture and applies it to the eSEN base, the novelty of MoLE is very limited. ...\n\nBelow we have listed the major novel contributions in this paper:\n\n1. MoLE architecture \n    - To the best of our knowledge, MoLE (or any type of MoE) is a novel design choice for a top performing MLIP. Standard MoE architectures borrowed from LLMs cannot be easily applied to MLIPs as it is difficult to preserve equivariance, ensure smoothness of the potential energy surface (PES), and make them memory and compute efficient at inference time.\n\n     - Our MoLE routing is specifically designed to depend on static features of the atomic structures (charge, spin, composition, etc.) rather than dynamic latent space features compared to standard MOEs. This allows us to pre-merge our experts weights before running a sequential simulation, making it fast and memory efficient. If the routing depended on the positions of atoms it could change throughout a sequential simulation and introduce a discontinuity in the PES (see Fu et al., 2025). During development, we tested dynamic routing schemes but ultimately rejected them as they didn't meet our listed requirements.\n\n2. Handling total charge/spin\n\n    - Most universal MLIPs e.g. MACE, Orb, Mattersim, DPA, SevenNet, PET-MAD etc. do not include charge and spin, but it is particularly important for molecules. \n\n3. Handling different DFT levels of theory\n\n   - We are not aware of another top performing universal MLIP that can generalize to multiple diverse tasks (for example in materials + molecules space) with different DFT levels of theory without relying on fine-tuning or regressing on performance.\n\n4. Scaling laws\n\n    - We provide the first systematic study of empirical scaling relationships including compute, data, and model size in MLIPs.\n\n5. Training innovations\n\n    - In section A of the Appendix, we highlight several innovations that made training at this scale possible including mixed BF16 (previously never used in MLIP training)/FP32 training, max-atom batching, and max neighbor switching.\n\nA number of these contributions and their novelty were specifically noted by the other reviewers:\n- Reviewer iEKm: “The introduction of an inference-optimized MLIP that can be run efficiently. In particular, the idea of contracting weight ahead of time — and amortizing this cost over the simulation — is novel to this application and seems to yield very promising results.”\n- Reviewer wJPq: “The paper provides a novel, performant model architecture that combines atomistic data from different sources and includes a way to different the source of different DFT calculations.”\n- Reviewer isJH: “It takes total charge and total spin multiplicity as additional inputs, enabling a potential extension to more diverse quantum chemical settings.” \n\n>Q2 cont: Additionally, a more comprehensive ablation study comparing MoLE against alternative architectural choices or existing MoE implementations would strengthen the claim of its efficiency and effectiveness. ...\n\nBased on your comment, we ran additional experiments to test the generalization ability of the MoLE architecture and applied it to EquiformerV2 on the OMol25 dataset. We found that the relative energy improvement on a similar sized EquiformerV2 (6M) and eSEN (6M) with MoLE is very comparable (see table below). This gives us more confidence that MoLE is a general approach to improve model capacity for MLIPs while preserving equivariance and simulation speed. We will include the results in the appendix of the camera-ready paper.\n\n| **OMol25 Val subset** | **eSEN 6M** | **EquiformerV2 6M** |\n| :- | :- | :- |\n| Metal Complexes | +10% | +10% |\n| Spice | +5% | +5% |\n| Neutral organics | 0% | +3% |\n| Biomolecules | +25% | +21% |\n| Electrolytes | +12% | +10% |\n| **mean across subsets** | **+10%** | **+10%** |\n\nIn addition, in the main text we demonstrated with ablations that naively training multi-task models in the small parameter regime without MoLE severely degrades performance (up to 20-30%) on nearly all tasks (Figure 2 top right). This provides further evidence that multi-task training is nontrivial and the MoLE architecture is a necessary component for training smaller models, which are critical to the community.\n\n>Q2 cont: In its current form, the paper leans more towards a technical report on large-scale training rather than highlighting fundamental architectural breakthroughs.\n\nWhile we do believe our work introduces a number of novel techniques for the model and for training/inference (as listed above), we note that this is not the only criteria for acceptance. Novel contributions related to model scaling have also been rewarded at NeurIPS, one famous example is the GPT-3 paper which won an outstanding paper award in 2020 and was focused on how scaling up LLMs improves performance without introducing any architectural updates. \n\n>Q3: The authors claim that UMA models are designed to push the frontier of speed, accuracy, and generalization ...\n\nWe agree that UMA does not outperform all specialized baselines on every single metric — we tried to be transparent about this fact. However, what is impressive is that a single model is on par or in a majority of cases better than SOTA across multiple chemistry areas (materials, molecules, catalysis, etc.). We would also like to point out that no previously published result has shown that multi-task models can even perform close to on-par with single-task variants in the space of MLIPs.\n\nNote that for WBM energy, both eqv2 and UMA perform very well and are within the threshold of practical utility (10-20 meV/atom). The strength of UMA is that while it does solid on single-point energy metrics like WBM, it also performs well on physical properties like thermal conductivity (kSMRE in Table 3) where models like eqv2 cannot be used.  \n\nA few examples where UMA’s generalization shines is catalysis OOD Both energy, MOFs OOD-L/T energy (the hardest OOD splits from OC20 and ODAC25 respectively), and unseen CSP target lattice energy.\n\n>Q4: Following up on Table 2, it is observed that UMA-lg, despite having the largest number of parameters, does not always exhibit the best performance among the UMA family …\n\nWe expect the empirical scaling laws to hold around the order of magnitudes of compute that we measured (10^18 - 10^20 flops) and similar to other scaling law studies (LLAMA3, Chinchilla, etc.) we expect our results to extrapolate out 2-3 orders of magnitude beyond what we measured.\n\nThe deviation of UMA-lg’s performance from other models can be explained from several confounding factors:\n\n\n- UMA-lg is a direct-force model and is generally not comparable with the performance of UMA-sm/md which are gradient-force models. There can be quantitative differences between direct-force and gradient-force models, for example direct models tend to do better in single-point calculations (Table 2) while gradient-force models perform better with longer rollout tasks or ones that depend on having smooth gradients with respect to the energy (Table 3). Ideally, we would show the conserving lg model performs better however it is too computationally expensive. \n- The scaling laws we measured here are based on validation loss which only provides a rough guiding metric for model performance in the MLIP space. There are many examples where improving validation loss does not translate to improved downstream performance:\n  - Benchmarks that are close to saturation, for example in Matbench Discovery F1 scores.\n  - Benchmarks that are limited by other factors such as smoothly varying potential surface, i.e. Matbench discovery kSRME metric/phonons etc. \n- We observed empirically (and noted in the main text) that scaling beyond UMA-lg with MoLE or increased dense parameters (up to 2-10B parameters) leads to very little further improvement even in validation loss which suggests that at these model sizes we are limited by the current dataset size and diversity. We did not have the compute budget to further explore this question in this paper.\n\n>Q5: Regarding the training data, the paper mentions combining datasets from various chemical domains. ...\n\nAs mentioned above, we included the appendix B.6 (Referencing) in the supplementary material. To briefly answer your question here: we reference all energies to isolated atomic energies computed with the respective level of theory (corrected by the heat of formation), followed by a linear referencing scheme (described in B.6) to help normalize the energies. This may help align the zero-energy line but importantly it enables the model to make predictions in a similar numerical range which helps with stability during training. We give the model a task embedding that allows the model to learn the energy landscape given the level of theory of the training data.   \n\n>Q6: Related work ...\n\nThanks for the suggestion, we will include these references for the camera-ready version of the paper. \n\n**Limitations**\n\nWe did provide a limitations paragraph in the main text of the paper (see lines 319-328), this was acknowledged by all the other reviewers. Given your comments/suggestions, we plan to add additional content on the limitations of scaling laws to the camera-ready version of the paper."}, {"review_id": "GAuAGtJS08", "reviewer": "Reviewer_isJH", "summary": "This paper introduces a family of foundation machine learning (ML) potential models, called Universal Models for Atoms (UMA). These models are trained on a large-scale dataset that combines multiple publicly available materials datasets, encompassing a wide range of atomic structures.\n\nThe main characteristics of the model are as follows:\n* It adopts a mixture of linear experts architecture, inspired by the mixture of experts technique used in large language models.\n* It introduces a two-stage training procedure designed to improve the overall training efficiency: the first stage predicts forces and stresses directly, and the second stage fine-tunes the model using autograd to produce energy-conserving outputs.\n* It takes total charge and total spin multiplicity as additional inputs, enabling a potential extension to more diverse quantum chemical settings.\n\nThe paper presents extensive experiments on prediction accuracy, inference efficiency, and empirical scaling laws. Moreover, the full details of the model, including hyperparameters, code, and pretrained checkpoints, are publicly released, ensuring reproducibility.", "questions": "**Question:**\n\nI am curious about the intended directions for future extensions of this work. Will the authors prioritize expanding the range of properties that the model can predict (e.g., spectra, magnetism), or will they focus on further scaling up the dataset used in training? In the latter case, extending the model to all-atom protein structures also seems like a promising possibility.\n\n**Suggestion 1:**\n\nIn the ML potential field, the evaluation of model quality is gradually moving beyond simple energy and force prediction accuracy toward more sophisticated assessments of stability and reliability. Several recent works have proposed additional evaluation metrics. It would be valuable for the authors to consider including some of these metrics in their evaluation. Belows are relevant references for such discussions:\n* Raja et al., “Stability-Aware Training of Machine Learning Force Fields with Differentiable Boltzmann Estimators.” TMLR (2025).\n* Yi et al., “Towards Physically Reliable Molecular Representation Learning.” UAI (2023).\n* Fu et al., “Forces are not Enough: Benchmark and Critical Evaluation for Machine Learning Force Fields with Molecular Simulations.” TMLR (2023).\n\n**Suggestion 2 (minor):**\n\nIn Figure 3, it would be better to align the y-axis ranges for panels (a) and (b), as well as for (c) and (d). While the current setup does not hinder interpretation of general trends, it makes direct comparisons more difficult.", "limitations": "Yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 4, "originality": 2, "strengths_and_weaknesses": "**Strengths:**\n* This paper presents a unified ML potential model that covers a wide range of material types, including molecules and crystals. By releasing both the checkpoints and code, the work makes a significant contribution to the field.\n* The paper conducts an empirical analysis of scaling laws using such a large-scale model, which is an experiment that most researchers would find difficult to carry out. This provides valuable insights into this domain.\n\n**Weaknesses:**\n\nWhile the scope of materials handled by the baseline models may vary, the current manuscript omits quite a few results in the comparison tables. More comprehensive baseline comparisons would strengthen the paper.\n\n**Recommendation:**\n\nAlthough the comparison with baseline models is somewhat limited, the release of a foundational ML potential model applicable across diverse material structures, along with the scaling law analysis, represents a substantial contribution to the field. In light of this, I recommend accepting this paper.", "was_revised": true, "final_justification": "After reviewing the authors' rebuttal and considering other reviewers' comments along with the authors' responses to them, I find that the authors have adequately addressed my concerns. Therefore, I maintain my original assessment.", "rebuttal": "We would like to thank reviewer isJH for their time and the constructive feedback! We address each of the reviewer’s comments/questions below.\n\n**Questions**\n\n>Q1: I am curious about the intended directions for future extensions of this work. Will the authors prioritize expanding the range of properties that the model can predict (e.g., spectra, magnetism), or will they focus on further scaling up the dataset used in training? In the latter case, extending the model to all-atom protein structures also seems like a promising possibility.\n\nWe are quite excited about the future of these types of models! One of the benefits of open sourcing the work is that the community can build off of it and extend our work. There is a lot of opportunity to fine-tune our models to predict other properties such as band gaps and for distillation to generate specialized models that are even faster. We are particularly excited about the possibility of multi-GPU inference, as mentioned in the paper, that enables much larger simulations such as you mentioned for all-atom protein structures. \n\nAs for scaling up the dataset, it will likely be more targeted in the future. One of the exciting observations from this work is that in a number of cases the models are already at or very close to chemical accuracy so scaling further isn’t needed. That is not true across the board and there are clearly areas that would benefit from more data. Understanding what data is needed and what data has the highest information content will be the focus of future work.\n\n>Q2: In the ML potential field, the evaluation of model quality is gradually moving beyond simple energy and force prediction accuracy toward more sophisticated assessments of stability and reliability. Several recent works have proposed additional evaluation metrics. It would be valuable for the authors to consider including some of these metrics in their evaluation. ...\n\nWe agree there is a lot of value in assessing downstream evaluations in addition to the standard energy and force metrics. In our paper, Table 2 focuses on the standard energy and force metrics whereas Table 3 focuses on downstream evaluations that require multiple energy/force evaluations. To highlight a few:\n\n- kSMRE (thermal conductivity), which requires both second and third order derivatives of the potential computed through finite differences (Póta et al., 2024).\n- AdsorbML, global (not just local) optimization problem for adsorption energy (Lan et al., 2023).\n- NVE Molecular Dynamics (MD) energy drift test, this evaluation requires running 100 ps of NVE MD while monitoring the energy drift, which is a proxy for the smoothness of the MLIP (Fu et al., 2025). If the total energy drift is below a defined threshold, we say the model is conservative. UMA-lg does not pass this test because it is a direct force model. \n\n>Q3: In Figure 3, it would be better to align the y-axis ranges for panels (a) and (b), as well as for (c) and (d). While the current setup does not hinder interpretation of general trends, it makes direct comparisons more difficult.\n\nThanks for this feedback. In the camera-ready version we will make sure any common axes (e.g. val loss) use the same scale.\n\n**Perceived weaknesses**\n\n>PW1: While the scope of materials handled by the baseline models may vary, the current manuscript omits quite a few results in the comparison tables. More comprehensive baseline comparisons would strengthen the paper.\n\nWe tried our best to include baseline results where they existed and in cases where there were none we trained our own. In Tables 2 and 3 you see a bunch of blanks “-” for baseline models because we only evaluated single-task baselines on the task they were trained on.\n\nAre there specific baselines that have been reported in the literature you think we missed? We would be happy to include them in the camera-ready version of the paper. \n\nWe acknowledge that there is substantial room for improvement in benchmarking these kinds of universal MLIPs both in terms of the depth and breadth of evaluations, but in this work we went well beyond what any single paper has done in the past in terms of benchmarks, marking a significant effort to understand model performance. In some cases we created our own benchmarks as part of this work e.g. HEAs for materials."}, {"review_id": "WtPZEoNPhs", "reviewer": "Reviewer_iEKm", "summary": "This paper present a “Universal Model for Atoms” with major contributions in both pre-training model architecture and scaling of datasets to create an efficient yet highly performant series of models. The architecture is based on Mixture of Experts approaches; however, rather than standard sparse selection of experts, the authors introduce the MoLE (Mixture of Linear Experts). Unlike language models, MLFF present a regression task where the output must vary smoothly, especially with respect to expert selection. The author therefore argue for a weighted contraction of the expert weight matrices. This leads to a “fewer” number of activated parameters by the effective contraction of the weight matrix. Note, this is particularly promising for inference procedures where the contraction can be performed ahead of time and repeated through the course of the network. For training, UMA introduces a number of new techniques, including low precision training (bfloat16) and following by increasing the precision in a second stage. For datasets, UMA enables the ability to combine a number of previously independent datasets including OC20, ODAC23, and OMC25 covering an extremely wide variety of atomic interactions. Other novel architectural components are the inclusion of charge and spin within the resulting DFT task.\n\nIn the results section, the authors discuss how model and data scaling laws were used to design levels of sparsity and scale models. The Iso-FLOP training curves in particular are particular helpful for determining an optimal model given a fixed compute budget. The authors then explore how UMA—as a multi-task model—compares to models trained on individual datasets. The multi-task model (UMA) often remains competitive with top single-task models and can even improve performance (for example, on OMol) — highlighting the benefits from the transfer learning. For standard benchmarks such as Matbench-Discovery, MDR phonon, and others, UMA achieve state of the art performance. The authors highlight that this strong performance is across Materials, Catalysis, Molecules, and Molecular Crystals — indicating the generality of the methods.", "questions": "Questions:\n1) From the appendix, it appears that the many of the experts have a very low element-expert pair utilization (expert 19, 23, and so on).  Were there load-balancing losses that were tried in order to motivate expert utilization? Similarly, the expert analysis in appendix G highlights how certain experts are correlated to certain elements of the periodic table. However, a core focus of the paper is on multi-task learning. Was similar analysis performed for each of the datasets — e.g. whether certain datasets re-use similar experts or choose different ones?\n2) Is there a specific linear operation within the SO2 convolution that is replaced by the MoLE method or all of them? In particular, the integration of MoLE into the SO2 convolution is barely discussed in the core body of the paper [mainly the caption of Figure 2]. Additional details here would be appreciated. Moreover, there is a specific linear operation in the node-based feed forward as well. Was there an ablation for this network layer as well?\n3) How do the speeds compare between the two stages of pretraining? Are there applications where the model at the end of stage 1 is possible to use or is the bfloat16 degradation 20%+ across the board?", "limitations": "Yes", "rating": 6, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 4, "originality": 3, "strengths_and_weaknesses": "Strengths:\n1) The paper introduces a highly ‘universal’ MLIP that is able to combine a wide variety of datasets into a single model. The universality of this model, from materials, catalysis, molecules, molecular crystals, to ODAC (MoFs) means that this model can have wide-scale applicability. In particular, the transfer learning between different domains (e.g. with the improvements in Molecular Crystals which seems to benefit from multi-task learning) highlights how the model has enough capacity to learn the best of each dataset.\n2) The introduction of an inference-optimized MLIP that can be run efficiently. In particular, the idea of contracting weight ahead of time — and amortizing this cost over the simulation — is novel to this application and seems to yield very promising results. The UMA method introduces a new way to scale MLIPs in a way that will remain optimized for inference.\n3) The paper is extremely well-written with a number of ablations that really help understand the benefits of the UMA model. For example, the scaling laws in terms of the number of active parameters really help show how the sparse models improve the performance for a fixed number of parameters. Also, the IsoFLOP curves help understand how compute should be split to train similar models under fixed compute budgets. All of these detailed findings provide a very explicit understanding of the models and what can be discovered.\n\nLimitations:\n1) The Mixture-of-Linear Experts (MoLE) method [to the best of my understanding] has been introduced previously in Mixture of Experts literature. For example, the method is highly similar to dense mixture of expert strategies — with particular similarity to SMEAR [https://arxiv.org/pdf/2306.03745] and Lory [https://arxiv.org/pdf/2405.03133]. Both also incorporate the idea of a pre-merged FFN for inference procedures of a learned router. This does not detract significantly from the originality of this work as I believe that it is the first of its-kind to apply this method to graph networks. However, it would be worth mentioning the other method works and potentially toning down the claims of originality.\n2) The implementation of UMA method in training and inference is insufficiently discussed in the core body of the paper. For example, standard issues such as expert capacity would likely be of concern when the data points of the atomistic systems are of variable sizes. The inference speeds of the methods suggest some form of acceleration but details of the implementation as well as further exploration into the efficiency would be worth discussing.", "was_revised": true, "final_justification": "The authors have answered my questions in the review and have made the commitment to add the details to the final version of paper. I have read through the other reviews and despite minor limitations in the form style of benchmarking (i.e. specific to two model families), I believe the questions brought up by other reviewers has been well-answered.\n\nHighlights from the paper\n1) A universal model that covers a wide variety of materials\n2) An inference-time speed up with details about the implementation\n3) Empirical understanding of scaling laws\n\nHighlights from the rebuttals:\n1) Discussion around load-balancing and shared experts between datasets highlights interpretability \n2) Training innovations clearly stated\n\nAll of these points are positives for the paper, so I recommend acceptance.", "rebuttal": "We would like to thank reviewer iEKm for their time and the constructive feedback! We address each of the reviewer’s comments/questions below. \n\n**Questions**\n\n>Q1: From the appendix, it appears that the many of the experts have a very low element-expert pair utilization (expert 19, 23, and so on). Were there load-balancing losses that were tried in order to motivate expert utilization? Similarly, the expert analysis in appendix G highlights how certain experts are correlated to certain elements of the periodic table. However, a core focus of the paper is on multi-task learning. Was similar analysis performed for each of the datasets — e.g. whether certain datasets re-use similar experts or choose different ones?\n\nWe did not try load balancing, but you’re right there are a few experts with low utilization and it's possible this could be improved. This experiment will take a bit of time but we will try to include a preliminary result (to inform future work, we won’t retrain all models) in the appendix for the camera-ready version of the paper. \n\nBased on your comment, we conducted a comparable analysis focusing on the expert weightings across different datasets (i.e. tasks). We describe it here in words per NeurIPS rules on images, but we will include the figure in the expert analysis section of the appendix. The results indicate that OC20, OMat, and OMol employ the largest number of experts. Notably, many experts utilized in OC20 are also used in OMat. In contrast, the experts associated with OMol show minimal overlap with other datasets, aside from a small subset shared with OMC. Lastly, the datasets ODAC and OMC utilize the fewest experts, and these two share a single expert between them.\n\n>Q2: Is there a specific linear operation within the SO2 convolution that is replaced by the MoLE method or all of them? In particular, the integration of MoLE into the SO2 convolution is barely discussed in the core body of the paper [mainly the caption of Figure 2]. Additional details here would be appreciated. Moreover, there is a specific linear operation in the node-based feed forward as well. Was there an ablation for this network layer as well?\n\nThe SO2 convolution consists of 1 linear layer for M=0 and two linear layers for other values of M. We replace each one of these linear layers with a MoLE. In fact our code is written in a general way that searches inside specific modules and then replaces linear layers with MoLE layers. We do this for all SO2 convolution operations in the network (2 x Num layers). We do not replace the linear operations inside the node-based feed forward layers. Given the camera-ready version can be 1 page longer than the original submission we can expand the description of MoLE in the main text.  \n\n>Q3: How do the speeds compare between the two stages of pretraining? Are there applications where the model at the end of stage 1 is possible to use or is the bfloat16 degradation 20%+ across the board?\n\nThe conservative fine-tuning stage is much slower than the direct pre-training stage, \\~3.5x slower than the UMA-sm model. The BF16 degradation is fairly universal, but this can be fixed with a short amount of FP32 fine-tuning (~ 50k steps). After the small amount of FP32 fine-tuning the models would be usable but they would be direct force models and these have a number of known limitations. \n\n**Perceived weaknesses**\n\n>PW1: The Mixture-of-Linear Experts (MoLE) method [to the best of my understanding] has been introduced previously in Mixture of Experts literature. For example, the method is highly similar to dense mixture of expert strategies — with particular similarity to SMEAR [https://arxiv.org/pdf/2306.03745] and Lory [https://arxiv.org/pdf/2405.03133]. Both also incorporate the idea of a pre-merged FFN for inference procedures of a learned router. This does not detract significantly from the originality of this work as I believe that it is the first of its-kind to apply this method to graph networks. However, it would be worth mentioning the other method works and potentially toning down the claims of originality.\n\nOur understanding is inline with yours that MoLE-like strategies have been previously discussed in the MoE literature. We included the citations we knew about on line 122, but we agree with the ones you mentioned as well. We will include these citations in the camera-ready version of the paper.  \n\n> PW2: The implementation of UMA method in training and inference is insufficiently discussed in the core body of the paper.\n\nWe agree, however, due to the 9-page limit of the NeurIPS format, we had to move most of the technical details of training and inference to the appendix. Given we get an extra page for the camera-ready version of the paper we can try to incorporate more of these details into the main text.\n\n>PW2 cont: For example, standard issues such as expert capacity would likely be of concern when the data points of the atomistic systems are of variable sizes. \n\nThe routing of MoLE experts in our implementation does not depend on the size of the atomic system, only on charge, spin, task, and composition. Furthermore, the composition embedding is based on the reduced atomic composition. For example a system with 1 carbon and 100 carbons would see the same composition embedding and therefore the same expert routing.\n\n>PW2 cont: The inference speeds of the methods suggest some form of acceleration but details of the implementation as well as further exploration into the efficiency would be worth discussing.\n\nWe discussed briefly in supplementary materials but we agree with the reviewer that we should add more details to the final paper. In particular, there were significant innovations used to speed up the model and improve its memory efficiency.\n\n1. Chunked activation checkpointing\n\n    - During inference with a large number of atoms, the activations used in our SO2 convolutions dominate GPU-memory as they scale linearly with the number of atoms, edges and embedding size. Standard activation checkpointing only provides slight improvements to memory because it still scales with the number of atoms. However, if we break up or \"chunk\" the input and compute them sequentially during SO2 convolutions at inference time and checkpoint the chunked activations, this allowed us to scale UMA-sm from 2k atoms to 100k with just ~30% inference time penalty.\n\n2. Pre-merge of MoLE experts\n\n    - This has been discussed as one of the central themes of MoLE. It allows us to skip the router call and expert parameter memory entirely and turn UMA into eSEN essentially at inference time. This is particularly important for large models like UMA-md where the parameter count can be brought down from 1.4B to 50M, dramatically reducing the footprint during inference time.\n\n3. Specialized GPU graph construction\n\n    - We have an in-house GPU graph construction algorithm that makes generating graphs for large systems fast. We evaluated many 3rd systems such as Pymat Gen, cuML KNN, openMM NNOps etc and none of them was fast enough or as simple to use for the scale that we want.\n\n4. Other optimizations\n\n     - Based on profiling, we eliminated many inference time bottlenecks that made eSEN particularly slow such as inefficient slicing operations, un-fused operators etc. We also added optimizations like cuda-graphs for Wigner matrix generation (collapsed many small operations), TF32 and torch compile to get further speed-ups"}, {"review_id": "Oqz6vNaO6V", "reviewer": "Reviewer_WjpQ", "summary": "The paper introduce a new set of models, Universal Models for Atoms (UMA), for atomistic structures spanning periodic crystals, molecules and surfaces interacting with molecules. In addition to the new models and architectures, the paper also outlines a study of scaling laws for atomistic models that aims to provide insight on the optimal relationship between data size, model size and compute budget. The paper starts by motivating the need for Machine Learning Interatomic Potentials (MLIPs) to accelerate the computation of accurate energies and forces of atomic systems based on Density Functional Theory (DFT) data. The introduction also poses the question of whether MLIPs can be trained across different types of DFT types and whether scaling of data and model sizes can unlock breakthroughs for atomistic modeling with machine learning. The proposed UMA models are trained on a combination of five datasets spanning 500M atomic systems and include scaling laws for different size UMA models, including UMA-sm, UMA-md and UMA-lg. \n\nSection 2 describes the overall approach starting with the UMA architecture. The UMA architecture is based of eSEN, which leverages message passing in spherical harmonic embeddings, and is augmented by global information that can represent particularities of the how the underlying DFT data was obtained (e.g., charge and spin). Furthermore, UMA uses a mixture of linear expert (MoLE) that combine experts based on the sum of a set of weights multiplied by a contribution. The paper claims that the MoLE approach enables smooth variation of prediction while also maintaining rotational equivariance. The training procedure undergoes two stages: The first stage performs direct force prediction and the second stage without a force head that uses conservative energy predictions via autograd. Given the large of training, the paper outlines additional efficiency improvements, including: pretraining with BF16 followed by FP32 fine-tuning and amortization of GPU memory usage for the MoLE layers. Section 3 briefly describe the datasets used for training UMA models.\n\nSection 4 describes the main results focusing on model and data scaling, optimal model size and modeling performance. The scaling laws analysis based on a fixed compute FLOP budget provides new insights into large-scale MLIP training, including a power law for UMA model training that serves the size of the different models. The results also include a discussion on overfitting challenges during training leading to deviations from the fitted power law. The scaling also leads to the observation that MoLE can be more parameter efficient compared to a dense model, thereby enabling more efficient inference for downstream applications. The second set of results compared single task and multi-task performance with MoLE models. The results in Table 2 generally show competitive performance of the UMA models while single task models still outperform in their specific task. The most compelling results for UMA are shown in Table 3 where the models show outperformance across different tasks and benchmarks for the atomistic systems covered in the training set. Section 5 provides a discussion of related work on MLIPs and scaling relation highlighting the need to use different DFT settings for different systems and the utility of scaling laws for model design. The paper ends with a discussion of limitations focusing on the limitations of the cutoff distance and the need to explore better ways to embed relevant DFT parameters into the model input for greater generalizability.", "questions": "* Can you describe in more detail why MoLE preserves rotational equivariance?\n* When you say pretraining happens with BF16 does that refer to the direct force pretraining or to both stages?\n* Can you describe why UMA-lg is not conservative?\n* Where the models in Table 3 fine-tuned for the task in the benchmark? Can you provide more details on the performance difference between Table 2 and Table 3?", "limitations": "yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 3, "significance": 3, "originality": 4, "strengths_and_weaknesses": "Overall, the paper provides a valuable contribution to the field of MLIPs with new, performant model architectures and insights on scaling laws for data, model size and compute.\n\n**Strengths:**\n* The paper provides a novel, performant model architecture that combines atomistic data from different sources and includes a way to different the source of different DFT calculations. The architecture shows promising scalability for downstream applications given the MoLE design. [Originality, Significance]\n* The study of scaling laws represents a useful insight for the MLIP domain given that few studies have provided this. [Significance, Quality]\n* The paper presents a compelling case for scaling MLIPs with more and lays out the methods and experiments with good clarity and relevant detail. [Clarity]\n\n**Weaknesses:**\n* While the paper has a strong technical foundation, it could provide a more detailed analysis on the consequences and limitations of its scaling study [Significance]. Some questions to consider: \n    * Do the authors expect this type of scaling to hold for other types of architectures or only for UMA?  \n    * How important is the MoLE to the observed scaling compared to the eSCN convolution? This could have consequences for future model design. The dense vs MoLE isoflops provide some good data for this, so a better discussion would be interesting.\n    * Given the increased scale, how important are engineering optimizations to achieve greater scale? Are there additional challenges, such as mitigating overfitting?", "was_revised": true, "final_justification": "I maintain my support and original score as the authors clarified many of details asked in my review. While some questions remain unresolved (e.g., scaling applicability to other methods), they constitute future work in my opinion.", "rebuttal": "We would like to thank reviewer WjpQ for their time and the constructive feedback! We address each of the reviewer’s comments/questions below.\n\n**Questions**\n\n>Q1: Can you describe in more detail why MoLE preserves rotational equivariance?\n\nOur router is invariant to rotations as it only depends on the reduced atomic composition (not atomic positions), global charge/spin, and the task. Adding MoLE swaps one linear operation with another linear operation during the SO2 convolution. Since these linear operations are used in the same manner as the original SO2 convolution, the block maintains its equivariance as described in the eSCN paper. \n\n>Q2: When you say pretraining happens with BF16 does that refer to the direct force pretraining or to both stages?\n\nThe direct pre-training stage utilizes BF16 whereas the conservative fine-tuning stage uses FP32.\n\n>Q3: Can you describe why UMA-lg is not conservative?\n\nWe opted to make UMA-lg a direct model because it can still run inference on a single 80GB GPU for smaller systems i.e. < 1k atoms. In contrast, a conservative model at that size would require model parallelism, making it too computationally expensive for most of the MLIP community. Additionally, the extra training overhead associated with conservative models (~3.5x) is impractical at the lg model size. \n\n>Q4: Where the models in Table 3 fine-tuned for the task in the benchmark? Can you provide more details on the performance difference between Table 2 and Table 3?\n\nThe UMA models in Table 3 were fine-tuned for the materials evaluations not for catalysis, molecules, or molecular crystal evaluations. The reason the models were fine-tuned (on MPtraj and sAlex) for the materials evaluations is to ensure a consistent level of theory (DFT settings). UMA models are trained on OMat24 data which has a slightly different level of theory compared to the Matbench Discovery benchmark and other materials evaluations that are common in the field. In practice, we recommend using the OMat24 level of theory (i.e. the standard UMA OMat task), so this is just an issue for evaluation. None of the models in Table 2 are fine-tuned, this was for materials evaluations in Table 3. \n\n**Perceived weaknesses**\n\n>PW1: Do the authors expect this type of scaling to hold for other types of architectures or only for UMA?\n\nWe believe it is likely that the log-linear scaling behavior holds for a larger family of models due to UMA’s similarity to other architectures that exhibit such scaling behaviors (Kaplan et al., 2020, Brehmer et al., 2024). In particular since UMA’s learnable parameters almost entirely reside in linear layers, akin to LLMs, UMA’s training compute (flops) scales linearly with the number of parameters N and size of dataset D, where a LLM token can be thought as an edge in UMA (see Appendix C). There are also other studies in equivariant GNNs that show this behavior (e.g. Brehmer et al., 2024). However, we do not have enough evidence currently to support or reject this hypothesis. We hope this encourages more independent scaling law studies in this space in the future to better understand these questions!\n\n> PW1 cont: How important is the MoLE to the observed scaling compared to the eSCN convolution? This could have consequences for future model design. The dense vs MoLE isoflops provide some good data for this, so a better discussion would be interesting.\n\nWe believe MoLE is independent of the scaling behavior compared to the eSCN convolution. In Figure 3c-e, we found the power laws coefficients with and without MoLE look similar (with the main distinction being a constant offset). The offset can be explained by the fact that MoLE just gives the eSCN convolution a reservoir of additional parameter capacity (~2.5x effective parameter increase from Figure 3e) at a cost of more memory and compute (which can be alleviated by pre-merging the MoLE experts during inference).\n\n>PW1 cont: Given the increased scale, how important are engineering optimizations to achieve greater scale? Are there additional challenges, such as mitigating overfitting?\n\nIn order to scale effectively, we need to optimize both for training time and GPU-memory.\nMany optimizations are important to reduce training time across all scales such as two-stage direct/conserved training, mixed BF16/FP32 precision etc.\n\nFor memory, different optimizations need to be applied depending on the dominant consumer of GPU memory. There are typically 2 main types of memory consumers:\n- Parameters - when total parameter count is close or >1B, they start to become a significant bottleneck on 80GB GPUs (optimizer states, EMA, other copies etc can take up 10s GBs quickly), parameter sharding techniques like FSDP are important to enable further scaling\n- Activations - activations scale with the number of layers, amount of operations and number of input atoms/edges and atom representation sizes. To perform conservative batched training with large systems, we need techniques like activation checkpointing and graph parallelism to reduce and distribute activation memory across GPUs. For example, conservative fine-tuning with Uma-md was very difficult without graph parallelism. \n\nFurthermore, dynamic batching in MLIP training with a large variety of system sizes can make training very inefficient (due to frequent allocations and memory fragmentation) and lead to OOMs. We found our max-atom batching scheme was critical to make sure training is memory stable and efficient throughout the run.\n\nWe observed overfitting frequently when training large models (ie: UMA-lg) on smaller or less diverse datasets (Section F of the Appendix). However, we found when training on large mixed datasets (i.e. the full UMA training set), it took much longer to observe overfitting and when it occurred the effect was much less pronounced. We also add expert dropout and standard regularization to help mitigate overfitting."}]

Author rebuttal:

Meta-review: The paper presents a family of foundation models that can be used to predict physical properties of atom structures. The models are trained on a very large-scale dataset curated by the authors by combining several earlier datasets. The proposed models exploit Mixture of Experts (MoE) techniques to increase computational efficiency. Extensive experiments illustrate the performance of the models in various tasks.

The main strengths include extensive scale of the dataset, carefully conducted experiments, (arguably minor) architectural novelties, and overall high quality of presentation. Weaknesses and room for improvement were pointed out in terms of limited theoretical novelty and acknowledgement of the limitations and scope of the scaling properties of the models, as pointed out by the reviewers.

Overall, the paper was found to be highly relevant to a critical problem in atom-level physical modeling applications and a major step forwards in practical application of foundation models for atom structures in physics,

The authors were able to respond to most of the critical comments and questions posed by the reviewers in a satisfactory manner.
Decision: Accept (spotlight)