FermiNet: Квантовая физика и химия с первых принципов

Применение методов глубокого обучения для решения ключевых задач в вычислительной квантовой химии и изучения взаимодействия вещества с светом

В публикации в журнале Physical Review Research продемонстрировано, как глубокое обучение позволяет решать фундаментальные уравнения квантовой механики для реальных систем. Это не только значимая научная проблема, но и потенциальный инструмент для будущих приложений, где ученые смогут моделировать новые материалы и химические синтезы на компьютере перед лабораторными экспериментами.

Архитектура нейронной сети FermiNet (фермионная нейронная сеть) идеально подходит для описания квантового состояния больших групп электронов, которые формируют основу химических связей. Код этого исследования опубликован открыто, чтобы сообщества вычислительной физики и химии могли развивать его и применять к разнообразным задачам.

FermiNet стал первым примером использования глубокого обучения для расчета энергии атомов и молекул с первых принципов с достаточной точностью для практического применения, а Psiformer, наша инновационная архитектура на базе механизма самоаттеншена, остается наиболее точным методом на основе ИИ на сегодняшний день.

Мы стремимся к тому, чтобы инструменты и концепции из исследований в области искусственного интеллекта способствовали решению фундаментальных научных вопросов. FermiNet дополняет наши проекты по свертыванию белков, динамике стекол, решеточной квантовой хромодинамике и многим другим, воплощая эту идею в жизнь.

Краткая история квантовой механики

Упоминание квантовой механики чаще вызывает недоумение, чем ясность. Оно вызывает ассоциации с парадоксальной кошкой Шрёдингера, которая одновременно жива и мертва, и с фундаментальными частицами, которые каким-то образом являются волнами.

В квантовых системах частица, такая как электрон, не имеет точного положения, как в классическом описании. Вместо этого ее позиция представлена вероятностным облаком — она размазана по всем разрешенным местам. Эта интуитивно противоречивая ситуация побудила Ричарда Фейнмана заявить: «Если вы думаете, что понимаете квантовую механику, значит, вы ее не понимаете».

Несмотря на эту загадочную странность, суть теории сводится к нескольким простым уравнениям. Самое известное из них — уравнение Шрёдингера — описывает поведение частиц на квантовом уровне аналогично тому, как законы Ньютона описывают движение объектов на привычном нам человеческом масштабе. Хотя интерпретация этого уравнения может вызывать бесконечные размышления, математика гораздо проще в работе, что приводит к распространенному совету преподавателей «замолчи и рассчитывай», когда студенты задают сложные философские вопросы.

Эти уравнения достаточны для описания поведения всей знакомой нам материи на уровне атомов и ядер. Их неинтуитивная природа приводит к экзотическим явлениям: сверхпроводники, сверхтекучие жидкости, лазеры и полупроводники возможны только благодаря квантовым эффектам. Даже простая ковалентная связь — основа химии — возникает из квантовых взаимодействий электронов.

После того как эти правила были разработаны в 1920-х годах, ученые осознали, что впервые получили детальную теорию химических процессов. В принципе, можно было настроить эти уравнения для разных молекул, вычислить энергию системы и определить, какие молекулы стабильны, а какие реакции протекают спонтанно. Однако при попытке реальных расчетов они смогли точно решить уравнения только для самого простого атома (водорода) и почти ничего больше. Все остальное оказалось слишком сложным.

Многие приняли вызов Дирака, и вскоре физики создали математические приемы, которые могли приближенно описывать качественное поведение молекулярных связей и других химических явлений. Эти методы начинались с приближенного описания поведения электронов, знакомого из вводных курсов химии.

В этом описании каждый электрон приписывается к определенной орбитали, которая указывает вероятность нахождения одного электрона в любой точке около атомного ядра. Форма каждой орбитали зависит от средней формы всех остальных орбиталей. Поскольку это «среднеполевое» описание приписывает каждый электрон только одной орбитали, оно дает очень неполную картину реального поведения электронов. Тем не менее, этого достаточно для оценки общей энергии молекулы с ошибкой около 0,5%.

К сожалению, ошибка в 0,5% все еще недостаточна для практической пользы в химии. Энергия молекулярных связей составляет лишь малую долю от общей энергии системы, и правильное предсказание стабильности молекулы часто зависит от 0,001% общей энергии, или примерно 0,2% оставшейся «корреляционной» энергии.

Например, общая энергия электронов в молекуле бутадиена составляет почти 100 000 килокалорий на моль, но разница в энергии между разными возможными формами молекулы — всего 1 килокалория на моль. Это значит, что для точного предсказания естественной формы бутадиена требуется такая же точность, как измерение ширины футбольного поля с миллиметровой точностью.

С появлением цифровых вычислений после Второй мировой войны ученые разработали множество методов, выходящих за рамки этого среднеполевое описание электронов. Хотя эти методы обозначаются разными аббревиатурами, они все лежат на оси, где точность обменивается на эффективность. На одном конце — практически точные методы, которые масштабируются хуже экспоненциально с числом электронов, делая их непрактичными для всех, кроме самых малых молекул. На другом конце — методы с линейным масштабированием, но низкой точностью. Эти вычислительные подходы оказали огромное влияние на химию — Нобелевская премия по химии 1998 года была присуждена создателям многих из этих алгоритмов.

Фермионные нейронные сети

Несмотря на разнообразие существующих инструментов вычислительной квантовой механики, мы сочли необходимым новый метод для эффективного представления. Есть причина, почему даже самые приближенные квантовохимические расчеты ограничиваются десятками тысяч электронов, в то время как классические химические техники, такие как молекулярная динамика, справляются с миллионами атомов.

Состояние классической системы легко описать — нужно отслеживать только позицию и импульс каждой частицы. Представление состояния квантовой системы гораздо сложнее. Вероятность должна быть присвоена каждой возможной конфигурации позиций электронов. Это закодировано в волновой функции, которая присваивает положительное или отрицательное число каждой конфигурации электронов, а квадрат волновой функции дает вероятность нахождения системы в этой конфигурации.

Пространство всех возможных конфигураций огромно — если представить его сеткой с 100 точками по каждой размерности, то число возможных конфигураций электронов для атома кремния превысит количество атомов во Вселенной. Именно здесь мы увидели потенциал глубоких нейронных сетей.

В последние годы произошли значительные прорывы в представлении сложных высокомерных вероятностных распределений с помощью нейронных сетей. Мы научились эффективно и масштабируемо обучать эти сети. Мы предположили, что, поскольку эти сети уже доказали способность подстраиваться под высокомерные функции в задачах ИИ, они могут представлять и квантовые волновые функции.

Исследователи, такие как Джузеппе Карлео, Матиас Тройер и другие, показали, как современное глубокое обучение может решать идеализированные квантовые проблемы. Мы хотели применить глубокие нейронные сети к более реалистичным задачам в химии и физике конденсированного вещества, что подразумевало включение электронов в расчеты.

Есть одно осложнение при работе с электронами. Электроны подчиняются принципу запрета Паули, который не позволяет им занимать одно и то же пространство одновременно. Это потому, что электроны — тип частиц, известных как фермионы, включающие строительные блоки большей части материи: протоны, нейтроны, кварки, нейтрино и т.д. Их волновая функция должна быть антисимметричной. Если поменять местами позиции двух электронов, волновая функция умножается на -1. Это значит, что если два электрона совпадают, волновая функция (и вероятность такой конфигурации) равна нулю.

Поэтому нам пришлось создать новый тип нейронной сети, антисимметричную относительно входов, который мы назвали FermiNet. В большинстве методов квантовой химии антисимметрия вводится с помощью функции, называемой детерминантом. Детерминант матрицы имеет свойство, что при обмене двух строк выход умножается на -1, как и волновая функция для фермионов.

Таким образом, можно взять набор функций для одного электрона, вычислить их для каждого электрона в системе и упаковать результаты в одну матрицу. Детерминант этой матрицы — правильно антисимметричная волновая функция. Главное ограничение этого подхода в том, что результирующая функция — слэтеровский детерминант — не очень универсальна.

Волновые функции реальных систем обычно гораздо сложнее. Обычный способ улучшения — взять большую линейную комбинацию слэтеровских детерминантов — иногда миллионы или больше — и добавить простые поправки на основе пар электронов. Даже тогда этого может не хватить для точного расчета энергий.

Глубокие нейронные сети часто гораздо эффективнее представляют сложные функции, чем линейные комбинации базисных функций. В FermiNet это достигается тем, что каждая функция, входящая в детерминант, зависит от всех электронов (см. примечание). Это выходит далеко за пределы методов, использующих только функции одного и двух электронов. FermiNet имеет отдельный поток информации для каждого электрона. Без взаимодействия между этими потоками сеть была бы не выразительнее обычного слэтеровского детерминанта.

Чтобы выйти за эти рамки, мы усредняем информацию из всех потоков на каждом слое сети и передаем ее в каждый поток на следующем слое. Таким образом, эти потоки обладают нужными симметричными свойствами для создания антисимметричной функции. Это похоже на то, как графовые нейронные сети агрегируют информацию на каждом слое.

В отличие от слэтеровских детерминантов, FermiNet — универсальные аппроксиматоры функций, по крайней мере, в пределе, когда слои нейронной сети достаточно широки. Это значит, что если мы правильно обучим эти сети, они смогут подогнаться под почти точное решение уравнения Шрёдингера.

Мы обучаем FermiNet, минимизируя энергию системы. Для точного этого нужно вычислить волновую функцию во всех возможных конфигурациях электронов, поэтому мы делаем это приближенно. Мы выбираем случайный набор конфигураций электронов, вычисляем энергию локально для каждого расположения, суммируем вклады и минимизируем это вместо истинной энергии. Это известный как метод Монте-Карло, поскольку напоминает игрока, бросающего кубики снова и снова. Хотя это приближение, для повышения точности мы всегда можем бросить кубики еще раз.

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

Основная идея, известная как вариационный метод Монте-Карло (или VMC), существует с 1960-х годов и считается дешевым, но не очень точным способом расчета энергии системы. Заменив простые волновые функции на основе слэтеровских детерминантов на FermiNet, мы резко повысили точность этого подхода для всех рассмотренных систем.

Чтобы убедиться, что FermiNet представляет прорыв в состоянии искусства, мы начали с изучения простых, хорошо исследованных систем, таких как атомы первого ряда периодической таблицы (от водорода до неона). Это малые системы — 10 электронов или меньше — и достаточно простые, чтобы их можно было обработать наиболее точными (но экспоненциально масштабируемыми) методами.

FermiNet значительно превосходит сравнимые расчеты VMC — часто сокращая ошибку относительно экспоненциальных методов вдвое или больше. На более крупных системах экспоненциальные методы становятся нерешаемыми, поэтому мы используем метод связанной кластерной теории как базу. Этот метод хорошо работает на молекулах в стабильной конфигурации, но испытывает трудности, когда связи растягиваются или разрываются, что критично для понимания химических реакций. Хотя он масштабируется гораздо лучше экспоненциально, используемый нами вариант все равно масштабируется как седьмая степень числа электронов, так что подходит только для молекул среднего размера.

Мы применили FermiNet к progressively более крупным молекулам, начиная с гидрида лития и доходя до бициклобутана, самой большой системы в нашем исследовании, с 30 электронами. На самых малых молекулах FermiNet захватил ошеломляющие 99,8% разницы между энергией связанной кластерной теории и энергией от одного слэтеровского детерминанта. На бициклобутане FermiNet все еще захватывал 97% или больше этой корреляционной энергии, что огромно для такого простого подхода.

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

Мы рассмотрели две эталонные растянутые системы: молекулу азота (N2) и цепочку из 10 атомов водорода (H10). Связь в азоте особенно сложна, поскольку каждый атом азота вносит три электрона. Цепочка водорода интересна для понимания поведения электронов в материалах, например, предсказания, будет ли материал проводить электричество.

На обеих системах методы связанной кластерной теории преуспели в равновесии, но имели проблемы при растяжении связей. Обычные расчеты VMC показали плохие результаты повсеместно, но FermiNet был среди лучших методов независимо от длины связи.

Новый подход к расчету возбужденных состояний

В августе 2024 года мы опубликовали следующий этап этой работы в журнале Science. Наше исследование предлагает решение одной из самых трудных задач в вычислительной квантовой химии: понимание переходов молекул в возбужденные состояния и обратно под стимуляцией.

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

Возбужденные состояния фундаментальны для понимания взаимодействия вещества со светом. Точное количество поглощаемой и испускаемой энергии создает уникальный отпечаток для разных молекул и материалов, влияя на эффективность технологий от солнечных панелей и светодиодов до полупроводников, фотокатализаторов и других. Они также играют ключевую роль в биологических процессах с участием света, таких как фотосинтез и зрение.

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

Мы разработали инновационный подход к расчету возбужденных состояний, более надежный и общий, чем предыдущие. Он применим к любой математической модели, включая FermiNet и другие нейронные сети. Метод работает, находя основное состояние расширенной системы с дополнительными частицами, так что существующие алгоритмы оптимизации используются с минимальными изменениями.

Мы проверили эту работу на широком спектре эталонов с многообещающими результатами. На малой, но сложной молекуле — димере углерода — мы достигли средней абсолютной ошибки (MAE) в 4 мэВ, что в пять раз ближе к экспериментальным результатам, чем предыдущие золотые стандарты с 20 мэВ. Мы также протестировали метод на некоторых из самых сложных систем в вычислительной химии, где одновременно возбуждаются два электрона, и обнаружили, что находимся в пределах около 0,1 эВ от наиболее требовательных сложных расчетов, выполненных на сегодняшний день.