Молекулярное моделирование с использованием машинного обучения, часть 2

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

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

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

Краткое резюме

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

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

Простейшая модель машинного обучения

При разработке модели машинного обучения полезно начать с самого основного вопроса: что в нее должно войти и что должно получиться.

Ну, мы хотим, чтобы модель вычисляла потенциал, представляющий собой одно число, зависящее от координат всех ядер. Итак, если у нас есть N ядер в трехмерном пространстве, модель должна взять 3N чисел и превратить их в одно число.

Простейшей возможной моделью машинного обучения, отвечающей этому требованию, было бы умножение на матрицу 1 на 3N W: M(x) = Wx. Это умножит каждый из входных данных 3N на вес и суммирует их все. Эти веса являются обучаемыми параметрами. Первоначально для них устанавливаются случайные значения, которые обновляются во время обучения путем исправления ошибок, которые делает модель. Эта часть останется неизменной.

Эта модель имеет много нежелательных свойств, наиболее проблематичным из которых является то, что она слишком проста. Более точно можно сказать, что это линейно. Это означает, что если мы умножим все входные данные на 2, выходные данные также будут умножены на 2. Или, если мы добавим 3, чтобы сказать вход x₁, выход будет исходным выходом плюс 3 w₁, где w₁ — соответствующий вес. В действительности, конечно, влияние этого изменения одной координаты на потенциал сложным образом зависит от его отношения к координатам других частиц. Эта модель никогда не сможет зафиксировать это, независимо от того, сколько данных вы ей подбросите.

Глубокое обучение

Снятие этого ограничения ведет нас от машинного обучения к глубокому обучению. Простейшая нейронная сеть представляет собой последовательность более простых моделей, представленных выше, смешанных с нелинейными функциями, называемыми функциями активации. (Без них конечный результат все равно был бы линейным!)

Самый простой пример этого класса моделей имеет один скрытый слой: он принимает 3N входных данных, чтобы сказать 100 чисел путем матричного умножения, и применяет функцию активации к каждому из них отдельно. Затем посредством второго матричного умножения он приводит их к одному числу, которое нам нужно в качестве вывода.

Существует множество различных функций активации, концептуально важно только то, что они нелинейны. Простым часто используемым примером является ReLU, который равен 0, если его вход отрицательный, а в противном случае просто передает ввод без изменений.

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

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

Чтобы уточнить это и сделать более точным, мы рассмотрим концепцию эквивариантности. Книгу и лекции по использованию структуры наших данных, которые являются гораздо более подробными, но также достаточно техническими, см. https://geometricdeeplearning.com/ и эту статью для приложений, специально предназначенных для молекул.

Эквивариантность

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

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

Частным случаем эквивалентности является инвариантность, когда подобное преобразование фактически является тождеством. Другими словами, если мы преобразуем ввод, вывод останется неизменным. Для простого примера этого, если бы гипотетическая модель на изображении выше была просто классификатором, указывающим, является ли изображение кошкой или собакой, мы ожидали бы, что она будет инвариантной: для обоих входных данных слева от изображения выходы справа должны быть «cat».

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

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

Перестановки

Превращая положения всех ядер в 3N-мерный входной вектор, мы должны были неявно определить порядок: какая частица первая, какая вторая и т. д. Ясно, что это просто выбор маркировки; если мы выберем другой порядок для точно такой же конфигурации частиц, потенциал не должен измениться. Итак, мы хотим, чтобы наша модель была инвариантной по отношению к перестановкам.

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

Есть несколько общих функций, которые не зависят от порядка их аргументов, т. Е. Инвариантны к перестановкам, например, сумма, среднее значение или максимум. Мы просто возьмем сумму, так как она сохраняет больше информации, чем максимум. Таким образом, самая простая модель, инвариантная к перестановкам, просто брала бы все положения ядер и суммировала их. (На самом деле это все равно оставило бы нам 3 координаты, поэтому нам понадобилась бы матрица 1 на 3 в конце, чтобы превратить ее в одно число.)

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

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

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

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

Локальность, графовые нейронные сети

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

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

Физический принцип — локальность: то, что происходит в некоторой точке x, зависит только от вещей, близких к x. Более конкретно, в этом случае кулоновский потенциал между двумя частицами на расстоянии r убывает как 1/r. Это означает, что взаимодействие между двумя ядрами, находящимися очень далеко друг от друга, не вносит существенного вклада в потенциал.

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

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

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

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

Пространственные инвариантности

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

Если мы сдвинем («перенесем») все ядра в одном направлении на одинаковую величину, межатомный потенциал не изменится. То есть, проводим ли мы эксперимент в Амстердаме или точно такой же эксперимент в Нью-Йорке, результат должен быть одинаковым. И то же самое верно для вращений и отражений. (Если бы существовало какое-либо внешнее влияние, такое как гравитация, все было бы иначе, но на атомном уровне это незначительно.)

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

Самый простой способ обеспечить это — предоставить только входные данные модели, которые сами по себе инвариантны к этим преобразованиям. Ясно, что если входы не затронуты, выход также не затронут.

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

Это делает работу, но это снова довольно ограничительно. Обратите внимание на параллель с перестановками, которые мы обсуждали ранее: проще всего сделать все неизменным с самого начала, но это окажется излишне ограничивающим.

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

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

Высшие представления/орбитали

Однако это становится немного сложнее. Сложность в том, что есть много способов, которыми вещи могут трансформироваться при вращении. Химикам они известны как орбитали, физикам и математикам — как различные представления группы вращения.

Я дам идею, которая, хотя и не является математически строгой, отражает интуицию того, что здесь важно для нас. Представьте, что мы начинаем с двух векторов, представляющих положения двух ядер, которые находятся в так называемом представлении L=1 или p-орбитали; здесь вектор остается неизменным при повороте вокруг своей оси, но меняет свое положение при любом другом повороте.

Мы можем объединить эти два вектора, взяв длину их разности. Как упоминалось ранее, это инвариантно, оно называется представлением L=0 или s-орбиталью. Это единственное число, которое не меняется ни при каком вращении. Мы также можем просто добавить их, тогда он останется вектором.

Два вектора вместе определяют плоскость, двумерную поверхность, в которой они оба лежат (при условии, что они не параллельны). Или, более конкретно, представьте себе один вектор, указывающий из кабины реального самолета, и другой, направленный вдоль одного крыла. Плоскость, которую они определяют вместе, будет полом внутри фактической плоскости.

Грубо говоря, именно так два вектора могут быть объединены в более высокое представление, это будет представление L = 2 или d-орбиталь. Обратите внимание, что эта плоскость зависит от обоих векторов: если мы будем вращать вдоль крыла самолета, этот вектор не изменится, но вектор, указывающий на кабину, обязательно изменится, как и пол самолета.

Этот шаблон продолжается: каждый раз, когда мы добавляем в смесь новый вектор, мы можем либо уменьшить представление на единицу, оставить его прежним, либо увеличить его на единицу. Например, если мы начнем с плоскости и вектора, вы можете представить, что вектор проецируется на плоскость, что дает нам новый вектор, уменьшая представление на единицу.

Таким образом, важно то, что более сложные комбинации могут быть сведены к чему-то инвариантному, что нам требуется для вывода модели.

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

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

Это реализовано в библиотеке e3nn и используется, например, в моделях молекулярной динамики nequip и allegro.

Заключение

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

Сначала мы увидели, что нам нужно более одного слоя, чтобы изучать больше, чем просто линейные функции. Мы использовали инвариантность перестановок в сочетании с локальностью, чтобы получить графовую нейронную сеть. Кроме того, нам требовалась инвариантность по отношению к пространственным преобразованиям переносов, вращений и отражений. Это дает нам то, что называется E(3)-эквивариантной GNN (E(3) — название этих преобразований в трехмерном пространстве).

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