FREED++. Ускоряем поиск новых лекарств с помощью нейросетей

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

О том, что сегодня для поиска новых соединений используют нейросети, слышали многие. Взять хотя бы нашумевший AlphaFold 3 от DeepMind, решающий задачу генерации трехмерной структуры разнообразных молекулярных комплексов. Существуют и другие задачи, в которых нейросети преуспели над классическими численными методами. Ярчайший пример — генерация молекул‑лекарств. Одним из самых заметных подходов к этой задаче стал фреймворк генерации молекул‑лекарств на основе методов обучения с подкреплением FREED. Но и он оказался далёк от идеала.

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

Разработка лекарственных препаратов

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

Как правило, такая разработка происходит по схеме «лекарство‑мишень» в три этапа:

1. Сначала определяется мишень‑белок — ключевая молекула, вовлеченная в конкретный метаболический или сигнальный путь, связанный с определенным болезненным состоянием.

2. Определяется трехмерная структура белка (обычно методами кристаллографии).

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

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

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

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

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

Обучение с подкреплением

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

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

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

Направленная генерация молекулярного графа

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

Функцию награды можно ввести несколькими способами:

или 

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

Схематически такой подход изображен на Рис.1.

Отметим, что вплоть до появления модели FREED (то есть, конца 2021 года), в подавляющем большинстве работ по молекулярной генерации (CVAE, ORGAN, REINVENT, GraphVAE, MolGAN, JT‑VAE, GCPN) рассматривалась задача безусловной генерации, такой, в которой не принимается во внимание белок‑мишень (хотя были и исключения — LiGAN). Такие модели не являются оптимальными, поскольку терапевтический эффект молекулы существенно определяется её энергией связывания с целевым белком.

Метод FREED

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

1. Фрагментная генерация существенно сокращает пространство возможных действий и состояний, позволяя эффективно обучать RL‑агента;

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

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

• в существующей реализации присутствуют множественные баги; 

• протокол сравнения неконсистентен между предложенной моделью и бейзлайнами;

• количество выбранных белков-мишеней для оценки качества недостаточно; 

• модель переусложнена и эффекты от ее отдельных частей недоисследованы.

От FREED к FFREED и FREED++

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

Перед тем, как описывать изменения внесенные в исходную модель, нужно более детально описать её саму. Однако поскольку исходная модель содержит в себе баги и неконсистентные блоки, это может ввести читателя в заблуждение. Поэтому я сначала опишу, как устроена исправленная модель FFREED (Fixed FREED), а затем приведу пример одного из найденных багов в модели.

FFREED

Рассмотрим, как устроен во FFREED выбор действия (архитектура политики, актор) и функция ценности состояния‑действия (критик).

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

Процесс выбора действия изображен на Рис. 2a.

Шаг 1: текущее состояние обрабатывается графовой сетью, и его эмбеддинг объединяется с эмбеддингами точками присоединения. После выбирается одна из точек присоединения , и ее эмбеддинг (более точно будет сказать эмбеддинг, обусловленный на эмбеддинг состояния) используется на следующем шаге.

Шаг 2: объединяется с эмбеддингами доступных фрагментов. Один из фрагментов выбирается как , и его эмбеддинг используется на следующем шаге (на самом деле выбор фрагмента устроен чуть более сложно, детальнее это описано в нашей статье).

Шаг 3: эмбеддинги точек присоединения выбранного фрагмента объединяются с . Наконец, одна из точек присоединения фрагмента выбирается как .

Критик: Эмбеддинги всех действий конкатенируются с эмбеддингом состояния и обрабатываются критиком (Рис. 2b).

Ниже мы опишем один из самых существенных багов, связанный с параметризацией критика (функции ценности состояния‑действия).

В оригинальной имплементации критик параметризован многослойным перцептроном, который принимает на вход конкатенацию 4 векторов:

1) эмбеддинга молекулы, полученного графовой нейронной сетью;

2) вектора вероятностей, ассоциированных с точками присоединения на молекуле;

3) one‑hot представления выбранного фрагмента;

4) вектора вероятностей, ассоциированных с точками присоединения на выбранном фрагменте.

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

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

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

FFREED++

На следующем шаге мы значительно ускорили модель FFREED и уменьшили количество обучаемых параметров в ней. Получившуюся модель мы называли FREED++. Расскажу по порядку, что именно мы сделали.

Архитектура политики. Рассмотрим более детально, как устроена архитектура политики в моделях FREED и FFREED. Выбор действия происходит авторегрессионно, и каждая компонента действия выбирается отдельной нейронной сетью. Устроены они однотипным образом, схематически изображенным на Рис. 4.

Каждая компонента политики вычисляет распределение вероятностей на множестве доступных опций при некотором условии. Например, для вычисляется вероятность выбора точки присоединения, при условии, что точка выбирается на молекуле, собранной к текущему шагу. Сначала эмбеддинг, ассоциированный с каждой опцией, сливается с эмбеддингом условия, например, эмбеддинги конкатенируются (на схеме функция слияния обозначена ). Затем эмбеддинг обусловленной опции обрабатывается многослойным перцептроном : , чтобы получить соответствующий логит. Наконец, условное распределение над опциями вычисляется оператором SoftMax (на самом деле, используется трюк Gumbel-Softmax, так как агент обучается методом SAC).

Функции слияния. Функция слияния (fusing functions) нужна для объединения нескольких эмбеддингов в один. В моделях FREED и FFREED в качестве функций слияния используются Multiplicative Interactions (MI). Пусть , представляют эмбеддинги, которые нужно объединить. Тогда, MI-слой примененный к и будет выглядеть следующим образом:

где  — это обучаемый трехмерный тензор, ,  — обучаемые матрицы, и  — смещение.

Отмечу, что MI‑слой содержит в себе вычисление билинейной формы, которое является дорогим с точки зрения времени и памяти.

Ускорение FFREED. Анализируя время вычисления различных частей политики, мы выяснили, что основное время тратится на шаге вычисления распределения вероятностей на множестве фрагментов. Это объясняется 3 факторами:

1) в качестве функций слияния используются MI‑взаимодействия;

2) входные эмбеддинги имеют высокую размерность (1024);

3) функцию слияния нужно вычислять раз, где  — размерность словаря фрагментов (в оригинальной статье ).

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

В модели FREED++ мы заменили MI-слой на вычислительно более дешевую и часто используемую комбинацию конкатенации линейного слоя:

Финальная модель FREED++, полученная после всех упрощений, оказалась быстрее в 8,5 раза и в 22 раза менее затратной по памяти, чем модель FFREED. Важно, что модель получилось ускорить без каких‑либо потерь в качестве генерации.

Сравнение с бейзлайнами

После исправления ошибок в исходном коде и оптимизации модели мы сравнили наш результат с другими существующими моделями на задаче генерации молекул с высокой аффинностью. В оригинальной работе протокол сравнения между бейзлайнами и моделью FREED неконсистентен, а набор тестовых белков недостаточен для аккуратного сравнения. В нашей работе мы расширили набор тестовых белков с 3 до 6, выполнили перебор гиперпараметров для бейзлайнов, добавили 2 метода для сравнения (комбинаторный генератор и Pocket2Mol, который и сегодня считается сильным бейзлайном), а также увеличили вычислительный бюджет для бейзлайн моделей (в оригинальной статье модели REINVENT и MolDQN были существенно недоучены).

Наши результаты показывают, что модели FFREED и FREED++ способны генерировать молекулы с более высокими значениями аффинности, чем существующие подходы.

Заключение

Таким образом, у нас, кажется, получилось исправить, упростить и улучшить модель FREED, причём довольно неплохо. Однако будет лукавством заявлять, что мы достигли предела совершенства в задаче о генерации молекул‑лекарств методами глубокого обучения.

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

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

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

Реализация FFREED и FREED++ доступна на нашем github, будем рады вашим звездочкам и форкам!