Обратимые преобразования глубокого обучения для эффективного генеративного моделирования

Натан С. Фрей

Соавтором этой публикации является Бхарат Рамсундар из Deep Forest Sciences.

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

К этому сообщению прилагается интерактивное руководство, которое можно запустить в Google Colab.

Зачем нам использовать нормализующие потоки в химии и биологии?

Существует много большой работы с использованием вариационных автоэнкодеров (VAE), рекуррентных нейронных сетей (RNN), графовых нейронных сетей (GNN) и генеративно-состязательных сетей (GAN) для генерации молекул. Нормализующие потоки (NF) можно использовать для создания молекулярных графов, как в статьях MoFlow и GraphNVP. Их также можно использовать для ускорения расчетов аффинности связывания между малыми молекулами и белками.

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

Что такое нормализующий поток?

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

Еще одно ключевое отличие состоит в том, что слои в НФ являются биективными преобразованиями — они обеспечивают однозначное сопоставление между входными и выходными данными, а не сжимают входные данные в скрытое пространство. Таким образом, NF полезны для любого приложения, которое требует вероятностной модели с расчетом плотности и выборкой либо с обоими. Я рекомендую эти замечательные сообщения в блогах от Eric Jang и Brad Saund для более подробной информации о NF и некоторых хороших примеров.

Подготовка набора данных строк SELFIES

Чтобы показать, как работают NF, мы попробуем смоделировать распределения из наборов данных QM9 и ChEMBL. Для представления молекул мы будем использовать самоссылающиеся встроенные строки (SELFIES) [1]. SELFIES — это стопроцентно надежное строковое представление молекул, а это означает, что даже совершенно случайные строки SELFIES являются химически действительными молекулами. Это отлично подходит для дизайнерских приложений, где модели глубокого обучения (например, NF) создают новые молекулы, не зная базовых правил химии, определяющих, является ли молекула действительной или нет. С функциями и примерами в репозитории селфи легко перевести строки Упрощенной системы молекулярного ввода (SMILES), предназначенные для наборов химических данных.

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

Обучение нормализации потока

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

Выбор базового распределения прост — мы выбираем многомерное нормальное распределение с тем же количеством измерений, что и наши однократные кодировки. Это настолько просто, насколько это возможно. Затем мы должны построить нашу НФ. Для FastFlows [2] мы используем слои Real NVP.

Вам может быть интересно, как слои в нейронной сети могут быть обратимыми; в общем их нет. Реальный NVP работает, ограничивая каждый элемент x_d+1 из преобразованного распределения, чтобы он зависел только от предыдущих элементов d из базового распределения для некоторого целого числа d. Alpha_i и mu_i — это скаляры, которые вычисляются путем передачи элементов базового распределения через нейронную сеть. Это дает простое преобразование масштаба и сдвига (аффинное), которое, тем не менее, может охватывать сложные целевые распределения, чувствительно завися от предыдущих переменных из базового распределения.

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

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

Для Real NVP мы должны указать количество слоев и количество скрытых единиц в каждом слое NVP. В этом примере мы будем использовать 8 слоев и [512, 512] скрытых единиц. Мы можем увеличить эти числа, чтобы сделать потоки более выразительными для захвата более сложных целевых распределений.

Создание новых молекул

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

 selfies.multiple_hot_to_selfies()

вернет нам представления SELFIES, которые мы можем расшифровать в SMILES и проанализировать с помощью программного обеспечения для хемоинформатики с открытым исходным кодом RDKit.

RDKit рассмотрит наши сгенерированные строки SMILES и проверит, являются ли они действительными молекулами. С помощью FastFlow мы можем генерировать 100 000 действительных молекул за 4,2 секунды на одном графическом процессоре! Если у вас есть опыт создания правильных строк SMILES с помощью авторегрессионных моделей, эта цифра впечатляет! Это скорость Real NVP и надежность представления SELFIES в действии. Теперь пришло время взглянуть на молекулы, которые придумал наш НФ.

Виртуальный досмотр с высокой пропускной способностью

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

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

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

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

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

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

Ограничения и дальнейшие действия

FastFlows обеспечивает быстрое обучение и выборку более 20 000 химически достоверных молекул в секунду после обучения на QM9 или почти 500 молекул в секунду после обучения на более сложных распределениях из ChEMBL, поэтому глубокое генеративное моделирование можно сочетать с столь же эффективным по времени апостериорным анализом. фильтры и многоцелевая оптимизация для эффективной молекулярной генерации. Но поскольку FastFlows полагается на полностью биективные преобразования, по мере увеличения размерности представления целевого распределения (~ 18 000 для селфи с горячим кодированием из ChEMBL), нормализующий поток страдает от низкой выразительности и требует увеличения глубины для захвата целевого распределения.

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

Ключевые выводы

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

Связь

Если вам понравился этот урок или у вас есть какие-либо вопросы, не стесняйтесь обращаться к Натану по электронной почте или связаться с LinkedIn и Twitter.

Эта работа была представлена ​​на семинаре ELLIS Machine Learning for Molecule Discovery Workshop 2021 года и доступна на arXiv.

Подробнее о проектах и ​​публикациях Натана можно узнать на его сайте.

Ссылки

[1] Кренн, Марио и др. «Самоссылающиеся встроенные строки (SELFIES): 100% надежное представление молекулярных строк». Машинное обучение: наука и технологии 1.4 (2020): 045024.

[2] Фрей, Натан С., Виджай Гадепалли и Бхарат Рамсундар. «FastFlows: модели на основе потоков для создания молекулярных графов». препринт arXiv arXiv:2201.12419 (2022 г.).