UMA: A Family of Universal Models for Atoms
Brandon Wood, Misko Dzamba, Xiang Fu, Meng Gao, Muhammed Shuaibi, Luis Barroso-Luque, Kareem Abdelmaqsoud, Vahe Gharakhanyan, John Kitchin, Daniel Levine, Kyle Michel, Anuroop Sriram, Taco Cohen, Abhishek Das, Sushree Sahoo, Ammar Rizvi, Zachary Ulissi, Larry Zitnick
The ability to quickly and accurately compute properties from atomic simulations is critical for advancing a large number of applications from drug discovery to energy storage to semiconductor manufacturing. We introduce a family of Universal Models for Atoms (UMA), designed to push the frontier of speed, accuracy, and generalization across chemistry and materials science. 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 catalysis. 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. All code, weights, and associated data are released with permissive licenses.