Title: UMA: A Family of Universal Models for Atoms
Venue: NeurIPS 2025
Accepted: yes
Keywords: Chemistry

Materials Science

Machine Learning Interatomic Potential

Density Functional Theory

Graph Neural Networks
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.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS AND WEAKNESSES:\nOverall, 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?\n\nQUESTIONS:\n* 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?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 4, "DATE": "2025-06-23"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\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.\n\nQUESTIONS:\nQuestions:\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?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 5, "CLARITY": 4, "IMPACT": 4, "ORIGINALITY": 3, "DATE": "2025-06-30"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\n**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.\n\nQUESTIONS:\n**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.\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 4, "ORIGINALITY": 2, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\n* 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.\n\nQUESTIONS:\n1. 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.\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-07-03"}