Meta AI представляет UMA: революционные универсальные модели для атомных симуляций

Революция в вычислительной химии с UMA

Теория функционала плотности (DFT) является основой вычислительной химии и материаловедения, однако её высокая вычислительная стоимость ограничивает широкое применение. Машинное обучение межатомных потенциалов (MLIPs) предлагает перспективную альтернативу, приближая точность DFT при значительном ускорении — сокращая время вычислений с часов до долей секунды благодаря масштабированию O(n) вместо O(n³) у DFT. Однако создание MLIPs, способных хорошо обобщать задачи в разных областях химии, остаётся сложной задачей, так как традиционные методы опираются на небольшие специализированные датасеты, а не на масштабные данные, что стало ключом к успеху в языковых и визуальных моделях.

Проблемы создания универсальных MLIPs

Существующие попытки создания универсальных MLIPs сосредоточены на обучении на больших датасетах, таких как Alexandria и OMat24, что улучшает результаты на бенчмарках типа Matbench-Discovery. Вдохновляясь эмпирическими законами масштабирования в больших языковых моделях (LLMs), исследователи начали изучать взаимосвязь вычислительных ресурсов, объёма данных и размера модели для оптимального распределения ресурсов. Несмотря на успехи в языковом моделировании, подобный подход для MLIPs применялся редко до настоящего времени.

Представляем UMA: Универсальные модели для атомов

Учёные из FAIR Meta и Университета Карнеги-Меллона представили UMA — семейство универсальных моделей для атомов, направленных на максимальное повышение точности, скорости и способности к обобщению в химии и материаловедении. UMA использует беспрецедентный датасет примерно из 500 миллионов атомных систем и эмпирические законы масштабирования для выбора оптимального размера модели и стратегии обучения. В результате модель достигает или превосходит специализированные модели по точности и скорости вывода на широком спектре задач без необходимости дополнительной настройки.

Архитектура и обучение UMA

UMA построена на основе eSEN — эквивариантной графовой нейронной сети с модификациями для эффективного масштабирования и обработки дополнительных входов, включая общий заряд, спин и настройки DFT. Новая схема эмбеддингов интегрирует эти входы, создавая эмбеддинги той же размерности, что и сферические каналы. Обучение проводится в два этапа: сначала модель напрямую предсказывает силы для ускорения обучения, затем дообучается для сохранения сил и напряжений с помощью авто-градиента, обеспечивая сохранение энергии и гладкость потенциальной поверхности.

Производительность и масштабируемость

Модели UMA демонстрируют логарифмическое линейное масштабирование по диапазону FLOP, что указывает на необходимость увеличения мощности модели для лучшего соответствия датасету. Многоэкспертное обучение показывает значительное улучшение при увеличении числа экспертов с 1 до 8, меньшие приросты с 32 экспертами и практически отсутствующие при 128. Несмотря на большой размер, UMA остаётся эффективной при инференсе: модель UMA-S способна симулировать 1000 атомов с частотой 16 шагов в секунду и обрабатывать системы до 100 000 атомов на одном 80ГБ GPU.

Достижения и ограничения

UMA устанавливает новые стандарты на бенчмарках AdsorbML и Matbench Discovery, превосходя в задачах материалов, молекул, катализаторов, молекулярных кристаллов и металлоорганических каркасов. Ограничения включают неспособность обрабатывать дальнодействующие взаимодействия из-за стандартного отсечения в 6Å и использование отдельных эмбеддингов для дискретных значений заряда и спина, что снижает обобщаемость на новые значения. В будущем планируется преодолеть эти ограничения и развить по-настоящему универсальные MLIPs.

Дополнительные ресурсы

Подробности доступны в [статье], моделях на [Hugging Face] и на [GitHub]. Это достижение — значительный шаг к более эффективным и универсальным атомным симуляциям.

Все заслуги принадлежат исследователям проекта.