UMA: Uniwersalne modele dla atomów zmieniają oblicze chemii obliczeniowej

Współczesna chemia obliczeniowa i materiałoznawstwo w dużej mierze opierają się na teorii funkcjonału gęstości (DFT), która, choć precyzyjna, jest niezwykle kosztowna obliczeniowo. Ogranicza to jej zastosowanie w modelowaniu złożonych układów, gdzie czas symulacji może wynosić godziny. Potencjały międzyatomowe oparte na uczeniu maszynowym (MLIPs) stanowią obiecującą alternatywę, zdolną do osiągnięcia dokładności zbliżonej do DFT przy znacznym przyspieszeniu obliczeń – z godzin do ułamków sekundy. Skalowalność MLIPs, charakteryzująca się złożonością O(n) w porównaniu do O(n³) dla DFT, jest kluczową przewagą.

Jednakże, mimo potencjału, dotychczasowe metody uczenia MLIPs borykały się z problemem generalizacji. Trening na małych, problemowo-specyficznych zbiorach danych uniemożliwiał pełne wykorzystanie skalowalności, która napędzała postęp w modelach językowych i wizyjnych. Dotychczasowe próby stworzenia uniwersalnych MLIPs, choć poprawiły wydajność na benchmarkach takich jak Matbench-Discovery dzięki większym zbiorom danych (np. Alexandria, OMat24), nadal miały ograniczenia.

Naukowcy z FAIR (Fundamental AI Research) w Meta oraz z Carnegie Mellon University podjęli to wyzwanie, prezentując rodzinę Uniwersalnych Modeli dla Atomów (UMA). Celem było przetestowanie granic dokładności, szybkości i generalizacji pojedynczego modelu w szerokim spektrum zastosowań chemicznych i materiałoznawczych. Kluczowym elementem ich podejścia było opracowanie empirycznych praw skalowania, które powiązały moc obliczeniową, rozmiar danych i rozmiar modelu. Pozwoliło to na optymalne określenie architektury modelu i strategii treningowej, co miało fundamentalne znaczenie biorąc pod uwagę bezprecedensowy zbiór danych około 500 milionów układów atomowych.

Architektura UMA, bazująca na ekwiwarianckich grafowych sieciach neuronowych eSEN, została zmodyfikowana w celu efektywnego skalowania i przetwarzania dodatkowych danych wejściowych, takich jak całkowity ładunek, spin czy parametry emulacji DFT. Wprowadzono również nowe osadzenie, pozwalające modelom UMA na integrację zadań związanych z ładunkiem, spinem i DFT. Proces treningu przebiega dwuetapowo: początkowo model bezpośrednio przewiduje siły, co przyspiesza trening, a następnie, po usunięciu warstwy przewidującej siły, model jest dostrajany do przewidywania zachowawczych sił i naprężeń z wykorzystaniem automatycznej dyferencjacji (auto-grad), co zapewnia zachowanie energii i gładkie powierzchnie energii potencjalnej.

Wyniki demonstracji UMA są imponujące. Modele wykazują logarytmiczno-liniowe skalowanie w testowanych zakresach FLOPs, co wskazuje na potrzebę większej pojemności modelu do pełnego wykorzystania zbioru danych UMA. Obserwuje się znaczną poprawę w redukcji strat przy przejściu z jednego do ośmiu „ekspertów” w treningu wielozadaniowym, z mniejszymi zyskami przy 32 ekspertach i marginalnymi przy 128. Pomimo dużej liczby parametrów, modele UMA charakteryzują się wyjątkową efektywnością wnioskowania – UMA-S jest w stanie symulować 1000 atomów z prędkością 16 kroków na sekundę i radzić sobie z układami liczącymi do 100 000 atomów na pojedynczej karcie graficznej o pojemności 80 GB.

W kierunku uniwersalnych symulacji

Podsumowując, UMA osiąga wyniki na poziomie lub przewyższającym specjalizowane modele w dokładności i prędkości wnioskowania na szerokiej gamie benchmarków dla materiałów, molekuł i katalizy, bez konieczności dostrajania do konkretnych zadań. Uzyskano nowe najlepsze wyniki na ustalonych benchmarkach, takich jak AdsorbML i Matbench Discovery. Należy jednak zaznaczyć, że model ma swoje ograniczenia: standardowy limit odcięcia 6Å utrudnia obsługę oddziaływań dalekiego zasięgu, a oddzielne osadzenia dla dyskretnych wartości ładunku lub spinu mogą ograniczać generalizację na nieznane wartości. Przyszłe badania mają na celu dalszy rozwój uniwersalnych MLIPs, co otworzy nowe możliwości w symulacjach atomowych, jednocześnie podkreślając potrzebę tworzenia bardziej wymagających benchmarków napędzających postęp.