Learning Local Equivariant Representations for Large-Scale Atomistic Dynamics

Что ж, сегодня Allegro от Гарварда. Наверное единственный метод, который классно скейлится между нодами GPU. Давайте разберемся как они смогли сделать симуляцию на 100 миллионов атомов 🐹

Собственно побуждение было вызвано плохим скейлиногом GNN из-за роста receptive field с количеством слоев, с чем также связано огромное количество вычислений между нодами. Allegro использует тензорные умножения и сферические функции - базу для сохранения векторных симметрий и скалярных инвариантностей. В статье используется обмен информацией между парами атомов через сумму взвешанных тензорных умножений между репрезентациями и сферическими тензорами. Свойство билинейности позволяет нам делать расчет более эффективно, разделяя операцию на фичи пары атомов и фичи соседей

Но, что супер важно, так это локальность операции агрегации. Все парные индексы имеют один и тот же центральный атом для конкретного атома. Это отличает Allegro от обычных GNN с их плохим скейлингом

Касательно экспериментов все очень необычно. У нас есть база типо QM9 / MD-17, но что интересно так это большие системы. Для того, чтобы пофлексить скейлингом авторы симулируют несколько крупных белков, а также твердотельный литий-ионный электролит на 100 миллиона атомов и оболочку ВИЧа на 44 миллиона соответственно

Хотя, как ни странно, для небольших систем уже сейчас существуют методы, которые стабильно бьют Allegro по точности. Оно не удивительно, ведь внутри стандратные MLP. Поэтому юзать есть смысл, только если у вас огромные молекулы или кристаллы

👀 LINK 1, LINK 2