Швейцарский нож устроен просто: один корпус, много лезвий, каждое под свою задачу. Достал нужное — и работаешь. В материаловедении графовые нейросети долго напоминали разрозненные инструменты: MEGNet лежал в ящике TensorFlow, M3GNet — в ящике PyTorch, CHGNet — вообще в отдельной коробке со своими скриптами. Чтобы перенести модель на новые данные или сравнить архитектуры, приходилось изобретать переходники. Библиотека MatGL собирает эти лезвия в один корпус.

Внутри — пять архитектур: MEGNet, M3GNet, CHGNet, TensorNet и SO3Net. Они построены из общих модулей — слоёв, пулинга, функций активации. Это как нож с общим механизмом: меняется только форма лезвия. Обучение идёт через PyTorch Lightning, что позволяет масштабировать расчёты на несколько GPU. Для межатомных потенциалов — функций, предсказывающих энергию по позициям атомов, — реализован класс Potential. Он автоматически учитывает нормировку энергий и считает силы, напряжения и гессианы через автодифференцирование. Интерфейсы к ASE и LAMMPS дают запускать релаксацию структуры и молекулярную динамику одной командой.

Бенчмарки на четырёх наборах данных показывают, что MatGL воспроизводит точность оригинальных моделей. На QM9 ошибка предсказания свободной энергии у M3GNet — 0.025 эВ, у TensorNet — 0.027 эВ, у SO3Net — 0.025 эВ; MEGNet отстаёт с 0.037 эВ. На Matbench для энергии формирования M3GNet, TensorNet и SO3Net дают ошибку 0.020–0.024 эВ/атом, а MEGNet — 0.037 эВ/атом. Для потенциалов на ANI-1x эквивариантный SO3Net достигает ошибки по силам 0.035 эВ/Å против 0.075 эВ/Å у M3GNet с переносом обучения. На кристаллах из Materials Project CHGNet обходит M3GNet и TensorNet: ошибка по энергии 19.9 мэВ/атом против 23.0 и 30.3 мэВ/атом. Скорость вывода у TensorNet на QM9 — 14.9 секунды на 6541 молекулу на одной RTX 3090, у MEGNet — 12 секунд, но с худшей точностью.

Раньше каждая графовая модель для материалов существовала в своём репозитории со своим форматом данных и своим способом обучения. Теперь появилась единая библиотека, где пять архитектур, предобученные потенциалы и интерфейсы к симуляторам работают через общий API. Собрали модульные слои и конвейеры данных поверх DGL и PyTorch Lightning, а затем прогнали все модели через одинаковые бенчмарки. Материаловеды могут сразу запускать предсказания свойств и молекулярную динамику, не разбираясь в деталях каждой архитектуры. Это снижает порог входа для экспериментаторов и ускоряет воспроизводимость. В отличие от соседних работ, которые фокусируются на одной архитектуре или на обучении с подкреплением, MatGL предлагает сразу пять моделей и связывает их с ASE и LAMMPS, что превращает библиотеку в платформу для полного цикла — от данных до молекулярной динамики.

Следующий шаг — поддержка многопроцессорного обучения на нескольких GPU и работа с базами данных, превышающими миллионы структур. Авторы планируют использовать memory-mapped базы для хранения графов, чтобы снять ограничение по памяти. Это позволит обучать универсальные потенциалы на всём объёме Materials Project и запускать молекулярную динамику систем с сотнями тысяч атомов. Швейцарский нож не лежит в ящике — он раскрывается под новые задачи.