UMA: A Family of Universal Models for Atoms

We would like to thank reviewer WjpQ for their time and the constructive feedback! We address each of the reviewer’s comments/questions below.

Questions

Q1: Can you describe in more detail why MoLE preserves rotational equivariance?

Our 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.

Q2: When you say pretraining happens with BF16 does that refer to the direct force pretraining or to both stages?

The direct pre-training stage utilizes BF16 whereas the conservative fine-tuning stage uses FP32.

Q3: Can you describe why UMA-lg is not conservative?

We 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.

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?

The 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.

Perceived weaknesses

PW1: Do the authors expect this type of scaling to hold for other types of architectures or only for UMA?

We 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!

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.

We 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).

PW1 cont: Given the increased scale, how important are engineering optimizations to achieve greater scale? Are there additional challenges, such as mitigating overfitting?

In order to scale effectively, we need to optimize both for training time and GPU-memory. Many optimizations are important to reduce training time across all scales such as two-stage direct/conserved training, mixed BF16/FP32 precision etc.

For memory, different optimizations need to be applied depending on the dominant consumer of GPU memory. There are typically 2 main types of memory consumers:

• 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
• 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.

Furthermore, 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.

We 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.

We would like to thank reviewer iEKm for their time and the constructive feedback! We address each of the reviewer’s comments/questions below.

Questions

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?

We 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.

Based 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.

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?

The 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.

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?

The 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.

Perceived weaknesses

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.

Our 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.

PW2: The implementation of UMA method in training and inference is insufficiently discussed in the core body of the paper.

We 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.

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.

The 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.

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.

We 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.

1. Chunked activation checkpointing

   • 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.

2. Pre-merge of MoLE experts

   • 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.

3. Specialized GPU graph construction

   • 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.

4. Other optimizations

   • 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

I have read through all of the reviewers comments as well as the rebuttals from the authors.

My review was focused mainly on additional details from individual sections (e.g. router utilization), and the answers in rebuttals suggest that these details will be added to the main body of the paper. In addition, the inference speedups as well as custom implementations (i.e. of graph creation) seem like impressive technical improvements. In particular, an only 30% slowdown at inference seems like a significant win -- allowing for drastically different systems in a single setting.

I maintain my support for the paper and keep my score/confidence.

We would like to thank reviewer isJH for their time and the constructive feedback! We address each of the reviewer’s comments/questions below.

Questions

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.

We 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.

As 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.

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. ...

We 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:

• kSMRE (thermal conductivity), which requires both second and third order derivatives of the potential computed through finite differences (Póta et al., 2024).
• AdsorbML, global (not just local) optimization problem for adsorption energy (Lan et al., 2023).
• 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.

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.

Thanks for this feedback. In the camera-ready version we will make sure any common axes (e.g. val loss) use the same scale.

Perceived weaknesses

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.

We 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.

Are 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.

We 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.

We would like to thank reviewer tgt6 for their time and the constructive feedback! We address each of the reviewer’s comments/questions below.

Questions

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. …

Our 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.

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. ...

Below we have listed the major novel contributions in this paper:

1. MoLE architecture

   • 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.

   • 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.

2. Handling total charge/spin

   • 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.

3. Handling different DFT levels of theory

   • 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.

4. Scaling laws

   • We provide the first systematic study of empirical scaling relationships including compute, data, and model size in MLIPs.

5. Training innovations

   • 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.

A number of these contributions and their novelty were specifically noted by the other reviewers:

• 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.”
• 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.”
• Reviewer isJH: “It takes total charge and total spin multiplicity as additional inputs, enabling a potential extension to more diverse quantum chemical settings.”

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. ...

Based 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.

In 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.

Q2 cont: In its current form, the paper leans more towards a technical report on large-scale training rather than highlighting fundamental architectural breakthroughs.

While 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.

Q3: The authors claim that UMA models are designed to push the frontier of speed, accuracy, and generalization ...

We 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.

Note 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.

A 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.

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 …

We 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.

The deviation of UMA-lg’s performance from other models can be explained from several confounding factors:

• 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.
• 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:
  • Benchmarks that are close to saturation, for example in Matbench Discovery F1 scores.
  • Benchmarks that are limited by other factors such as smoothly varying potential surface, i.e. Matbench discovery kSRME metric/phonons etc.
• 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.

Q5: Regarding the training data, the paper mentions combining datasets from various chemical domains. ...

As 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.

Q6: Related work ...

Thanks for the suggestion, we will include these references for the camera-ready version of the paper.

Limitations

We 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.