ofantomas 11 июл в 15:09

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

Средний

8 мин

1.8K

Блог компании AIRIBig Data*Искусственный интеллектФизикаХимия

Кейс

Привет! Меня зовут Артем Цыпин, я исследователь в Институте искусственного интеллекта AIRI. Наша команда занимается применением глубокого обучения в науках о жизни. В сферу наших интересов входят такие задачи как поиск новых лекарственных препаратов, дизайн материалов, анализ растворимости и другие.

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

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

Введение

Современные подходы для решения подобных задач работают со специальным представлением для молекулярных систем, называемом конформацией. Каждая конформация — мы будем обозначать их с помощью символа s — определяется типами атомов и их пространственным расположением. Атомы достаточно характеризовать набором зарядовых чисел

$z = \{ {z_1, . . . , z_N} \}, z_i ∈ N$

а положения в пространстве — набором троек координат

$X = \{ { \mathbf {x}_1, . . . , \mathbf {x}_N} \} , \mathbf {x}_i \in \mathbb{R}^{N \times 3}$

где N — количество атомов в системе.

Каждая конформация характеризуется потенциальной энергией E_s. Определить эту связь сложно, но можно с помощью громоздких, зато достаточно точных вычислений, например, в рамках теории функционала плотности (density functional theory, DFT). Молекулярные системы являются динамичными, и позиции атомов в пространстве постоянно меняются под действием межатомных сил

$\mathbf {F}_s = \{ { \mathbf {F}_{1, s}, . . . ,\mathbf {F}_{N, s} } \} , \mathbf {F}_{i, s} \in \mathbb{R}^{N \times 3}$

где F_i_,_s — это вектор силы, действующий на i-й атом конформации s. Эти векторы можно получать из потенциальной энергии, продифференцировав ее по координатам атомов системы (точнее, взяв антиградиент):

$\mathbf {F}_{i, s} = -\frac{\partial E_s}{\partial \mathbf {x}_i}.$

В равновесных условиях множеству всевозможных конформаций молекулярной системы можно сопоставить распределение Больцмана:

$p(s) \propto \exp \Big(-{\frac{E_s}{kT}}\Big).$

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

Оптимизация конформаций с помощью физических симуляторов

Традиционный подход к решению данной задачи подразумевает итеративную оптимизацию конформации с использованием межатомных сил в качестве анти-градиента:

$S_{t+1} = S_t + \alpha \mathbf{Opt}(\mathbf{F}_{s_t})$

Здесь S_t, S_t₊₁ — конформации на t и t+1 шаге оптимизационной траектории; α — гиперпараметр, отвечающий за темп оптимизации (learning rate); Opt — оптимизатор (в большинстве реализаций используется BFGS или L-BFGS).

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

Межатомные силы вычисляют с помощью физических симуляторов (оракулов). Сложность физических симуляторов ранжируется от линейной до экспоненциальной от количества электронов в системе N_e и влияет на их точность: чем точнее предсказания физического симулятора, тем более вычислительно сложны его операции. Популярным выбором являются DFT-симуляторы, так как они являются достаточно точными. Однако, такие методы имеют сложность порядка O(N_e⁴), что ограничивает их применение к большим системам в рамках задачи итеративной оптимизации.

Оптимизация конформаций с помощью нейронных потенциалов

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

$\begin{align} E_{s_t}^{NNP} &= E(s_t; \theta) \\ \mathbf{F}_{s_t}^{NNP} &= \mathbf{F}(s_t; \theta) = -\frac{\partial E(s_t; \theta)}{\partial X_t}, \end{align}$

где θ — это веса модели.

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

В наших экспериментах мы сосредоточились на задаче поиска низко-энергетических конформаций для молекул в вакууме. Мы показали, что оптимизация с использованием нейронного оракула в ~2000 раз быстрее (wall-time) чем оптимизация с помощью DFT-based симулятора.

Однако оказалось, что обученные на имеющихся в открытом доступе наборах данных (nablaDFT, SPICE) нейронные потенциалы, не могут быть использованы для задачи оптимизации без дообучения. Причиной этого является distribution shift, который возникает при использовании нейронных потенциалов для предсказания межатомных сил в процессе итеративной оптимизации. Дело в том, что датасеты типа nablaDFT (собранный, кстати, нашей командой!) и SPICE не содержат оптимизационных траекторий, из-за чего начинают ошибаться в предсказании сил по мере приближения к локальному минимуму. Это приводит к тому, что оптимизация сходится в «неправильную» конформацию или же расходится (молекулу разрывает или тяжелые атомы становятся слишком близкими друг к другу).

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

f^baseline — обучена на подмножестве D₀, |D₀| ≈ 10000 большого набора данных nablaDFT.
f^traj−¹⁰^k — обучена на конкатенации набора данных D₀ и множестве оптимизационных траекторий, полученных физическим симулятором, размером 10000 конформаций. Такой набор данных мы называем D_traj−₁₀_k, |D_traj−₁₀_k| ≈ 10000.
f^traj−¹⁰⁰^k — обучена на наборе данных D_traj−₁₀₀_k, |D_traj−₁₀₀_k| ≈ 110000.
f^traj−⁵⁰⁰^k — обучена на наборе данных D_traj−₅₀₀_k, |D_traj−₅₀₀_k| ≈ 510000.

Отметим, что оптимизационные траектории были получены путем оптимизации конформаций из набора данных D₀; средняя длина оптимизационной траектории ≈ 100 конформаций. Далее мы оценили полученные нейронные потенциалы на задаче оптимизации и построили график зависимости ошибки предсказания межатомных сил от шага оптимизации:

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

Как видно из графика, ошибка постепенно уменьшается с увеличением количества оптимизационных траекторий в обучающей выборке. Более того, нейронный потенциал f^traj−⁵⁰⁰^k имеет сравнимое с физическим симулятором качество оптимизации (то есть ошибка в среднем оказывается меньше химической точности — 1 kcal/mol)! Однако, для того чтобы собрать и обсчитать дополнительные 500 000 конформаций, мы потратили приблизительно 9 CPU-лет вычислений. Для более сложных атомных систем таких как адсорбент-адсорбат, молекула в растворе, молекула в белке, собрать такое количество данных и вовсе не представляется возможным.

Как уменьшить количество необходимых дополнительных данных при обучении нейронного потенциала?

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

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

Чтобы решить эту проблему, мы предлагаем использовать неточный, но фактически «бесплатный» физический симулятор, который мы называем суррогатным оракулом. В его основе лежит использование простой эмпирической модели, известной как модель молекулярной механики или молекулярных силовых полей. В ней мы каждую химическую связь представляем в виде классического осциллятора. Соответственно, всё, что нам нужно, чтобы вычислить энергию, это параметры этих осцилляторов и схема их построения. В нашей работе мы выбрали молекулярное силовое поле Мерка (Merck molecular force field, MMFF).

Построенный таким образом суррогатный оракул оценивает для нас потенциальную энергию конформации E_sₜ₊₁^MMFF на каждом шаге оптимизации:

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

Если энергия E_sₜ₊₁^MMFF уменьшилась, то оптимизация продолжается. Если же энергия увеличилась, мы считаем, что нейронный потенциал «плохо» предсказал силы для конформации s_t. Далее мы вычисляем истинные силы и энергию для конформации s_t с помощью дорогостоящего физического симулятора и добавляем в обучающую выборку. После этого нейронный потенциал дообучается на расширенной обучающей выборке.

Только наиболее релевантные конформации обсчитываются с помощью дорогостоящего симулятора и добавляются в обучающую выборку

Мы стартуем с нейронного потенциала, обученного на наборе D₀, для предсказания межатомных сил и построения оптимизационной траектории, а сам процесс продолжается пока не исчерпается бюджет на дополнительные вычисления. Получившийся фреймворк мы назвали постепенным выучиванием оптимизации (Gradual Optimization Learning Framework, GOLF).

Оценка качества

Многие современные подходы для молекулярного моделирования используют в качестве метрики RMSD (root mean squared distance). Эта метрика подсчитывает среднее расстояние между атомами получившейся в результате оптимизации с нейронным потенциалом конформации и целевой конформации. Однако, метрики, основанные на расстоянии не являются оптимальными для данной задачи. Например, две конформации могут иметь высокий RMSD и при этом быть одинаково выгодными с точки зрения потенциальной энергии. Поэтому в данной работе мы используем метрики, основанные на энергии.

$\begin{align} \overline{\operatorname{pct}}_t &= \frac{1}{|\mathcal{D}_{\text{test}}|} \sum_{s \in \mathcal{D}_{\text{test}}} 100\% * \frac{E_{s_0}^{DFT} - E_{s_t}^{DFT}}{E_{s_0}^{DFT} - E_{s_{\mathbf{opt}}}^{DFT}}, \\ \end{align}$

где E_sₒₚₜ^DFT — энергия целевой конформации, полученной с помощью физического симулятора. Средний процент оптимизированной энергии

$\begin{align} \overline{E^{\text{res}}}_t &= \frac{1}{|\mathcal{D}_{\text{test}}|} \sum_{s \in \mathcal{D}_{\text{test}}} E_{s_t}^{DFT} - E_{s_{\mathbf{opt}}}^{DFT}. \\ \end{align}$

Средняя разность энергии конформации, полученной оптимизацией с нейронным оракулом, и целевой конформации, полученной оптимизацией с физическим симулятором:

$\begin{align} \operatorname{pct}_{\text{success}} &= \frac{1}{|\mathcal{D}_{\text{test}}|} \sum_{s \in \mathcal{D}_{\text{test}}} I\left[E_{s_t}^{DFT} - E_{s_{\mathbf{opt}}}^{DFT} < 1\right]. \end{align}$

это процент конформаций, для которых разность энергий меньше химической точности (1 kcal/mol).

Далее мы сравнили нейронные потенциалы, дообученные на разном количестве данных, на подмножестве набора данных nablaDFT D_test, |D_test| ≈ 20000. Также мы сравнили нейронные потенциалы с генеративными подходами (TD, ConfOpt, Uni-Mol+):

Из таблицы можно видеть, что нейронный потенциал, дообученный с помощью GOLF на 10000 конформаций имеет сравнимое качество с нейронным потенциалом, дообученном на 500000 конформаций

А вот как ведёт себя ошибка NNP, обученного с помощью GOLF, по сравнению с обучением на 500k данных:

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

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

Этот результат был получен целой командой исследователей из AIRI, ФИЦ ИУ РАН, МФТИ и Университета Констрактор в Бремене. Совсем недавно мы представили его на конференции ICLR 2024, которая прошла в Вене, детали можно найти в соответствующей статье. Весь код лежит в отрытом доступе у нас на GitHub, будем рады вашим звёздочкам и форкам!

Теги:

Хабы: