Всем привет! Меня зовут Константин, и я занимаюсь вопросами глубокого обучения в естественных науках в AIRI.
Среди всех достижений глубокого обучения большие языковые модели — пожалуй, самые заметные. Однако помимо работы с текстами у нейросетей есть хорошие перспективы в области биологии, химии, физики и других наук. Можно вспомнить, например, последние Нобелевские премии за архитектуру AlphaFold.
Мы в AIRI активно интересуемся этим прогрессом, в особенности применением глубокого обучения для квантовой химии. Недавно нашу статью про предсказание электронной плотности с помощью нейросетей приняли в Journal of Cheminformatics [1], и мне бы хотелось рассказать, что именно мы там сделали.
Но в первой части своего рассказа я хотел бы сделать введение в тему, в частности, поговорить о том устроены целевые значения для предсказания свойств молекул, и что такое геометрические графовые нейронные сети.
Приятного чтения.
На 2025 год человечеству известно про примерно 219 миллионов органических химических соединений. При этом лекарств в мире порядка 10 000 (то есть, около 0.05% от всех органических веществ), а активно продающихся, без сильных побочных эффектов и широко доступных для лечения — около 5000.
Ещё до того, как новое соединение будет подвергнуто клиническим испытаниям, необходимо доказать наличие у него нужных физико‑химический свойств.
Исследование каждого вещества в лаборатории может стоить очень дорого, и очень хочется как‑то автоматически предсказывать свойства заранее. Хемоинформатика, квантовая и физическая химия — это как раз те теоретические разделы химии, которые решают эту задачу. Помимо фармакологии, у этих наук много других приложений: красители, пищевая промышленность, производство полимеров и нефтехимия.
219 миллионов органических соединений — это уже большие данные. Начиная где‑то с 2016 года датасайентсты поняли, что химические вещества могут быть хорошими целями для обработки с помощью нейронных сетей.
За очень короткое время появилось множество работ по теме. Например, SchNet [2] — одна из первых нейронных сетей для предсказания квантовых свойств молекул, увидевшая свет в 2017 году. Параллельно с ней Артур Кадурин (теперь он руководитель нашей группы в AIRI) вместе с коллегам предложил druGAN [3], одну из первых генеративных сетей, способных предлагать новые молекулы. А в 2018 году появилась знаменитая AlphaFold [4] — модель для предсказания структуры белков. Это позволило существенно ускорить исследования на стыке биологии и химии.
Ключевыми проблемами при работе с данными в химии и биологии являются представление объектов и выбор целевых свойств, которые требуется предсказать. Данные о том, в какую трёхмерную структуру уложены атомы в молекуле или аминокислоты в белке, имеют очевидно не строковый формат. Целевые значения для предсказания также могут сильно отличаться от таковых в задачах классического ML. Например, это могут быть матрицы, функции координат или положения объектов относительно друг друга. По этой причине многие задачи, связанные с предсказанием квантовых эффектов и свойств молекул, используют графовые нейронные сети.
Немного физики и химии
Прежде, чем двигаться дальше, поговорим немного о молекулах. Существует несколько способов записи и представления молекул: химические формулы, специальные строки в формате SMILES, двухмерные графы (как в учебниках по химии) и, наконец, трёхмерные структуры, заданные номерами атомов и их положением в пространстве (на полке в кабинете химии). Молекуле с одной химической формулой может соответствовать множество конформаций, то есть наборов положений атомов в трёхмерном пространстве. Малые молекулы обычно имеют от 5 до 50 конформаций при низких температурах.
В молекулах вокруг атомов летают электроны, взаимодействие части из них не даёт молекуле распасться и обеспечивает всё химическое разнообразие. Мозг человека очень плохо приспособлен для того, чтобы полностью понимать, что такое электрон, и сейчас мы объясним, почему.
В первую очередь электроны «не именованные». Представим, что на столе стоят две кружки. Если вы поменяете местами две кружки на столе, то вы придёте к совершенно иной физической системе. Но если вместо кружек у вас электроны, то такая замена систему не изменит. Это свойство электронов называется тождественностью.
Многие могли видеть картинки из школьных учебников, где нарисована планетарная модель атома. Ядро в центре, а вокруг него летают маленькие шарики — электроны. В реальности всё обстоит существенно иначе.
Сегодня вместо шариков мы представляем электроны как нечто, распределённое в пространстве. Строго говоря, мы должны описывать их с помощью комплексных волновых функций, но, чтобы эту распределённость хоть как‑то можно было визуализировать, часто переходят к квадрату её модуля. По физическому смыслу этот квадрат представляет собой плотность вероятности найти электрон. В молекулах же интеграл по пространству от плотности для всей электронной системы даст не единицу, как в теории вероятности, а количество электронов.
Конкретная форма волновых функций (их ещё часто называют орбиталями) определяется с помощью стационарного уравнения Шрёдингера — дифференциального уравнения второго порядка в частных производных. Как правило, они могут быть представлены с помощью сферических гармоник, домноженных на какую‑то функцию расстояния до ядра. Наконец, полная функция многоэлектронной системы складывается из отдельных орбиталей с помощью громоздкого детерминанта Слэтера — это необходимо, чтобы корректно учесть тождественность электронов.
Но это ещё не всё. Также у электронов есть обменно‑корреляционное взаимодействие. Это штука, для строгого объяснения которой опять нужно записывать интегралы, а также учитывать, что у электронов есть спин — он может «смотреть» вверх или вниз. Но мы попробуем описать его попроще.
Итак, каждую точку пространства всегда хотят занять два электрона. Важно, что квантовая механика запрещает это делать электронам с одинаковыми спинами — это приводит к их эффективному отталкиванию. Когда же один электрон со спином вверх, второй со спином вверх, они, наоборот, как бы стремятся притянутся друг к другу. Именно из‑за того, что электроны с разными спинами хотят находиться вместе, а не по одиночке, и появляются ковалентные связи в молекуле.
Вообще, можно совсем упростить и сказать, что вся молекула — это шарики (атомные остовы), погружённые в слайм из электронов. Если слайм начать растягивать, то появятся пустые места где будет не хватать электрона с определенным спином, из‑за этого молекула будет сопротивляться растягиванию.
Наконец, важно отметить, что электроны в атомах принято делить на внутренние (остовные) и валентные. Первые плотно кучкуются вокруг ядра и занимают там очень маленькое пространство — им соответствуют внутренние (остовные) орбитали. Такие электроны не участвуют в построении молекул, но они, например, ответственны за поглощение или испускание рентгеновских лучей. А вот электроны на валентных орбиталях очень сильно «размазаны» по пространству, и именно они образуют химические связи.
Другими словами, плотность электронов огромна по амплитуде в близи ядер, но очень мала в химических связях. Например, в молекуле бромоводорода на окрестность ядра брома приходится около 25,000 электронов на кубический Бор (единица измерения длины, равная боровскому радиусу), а в химических связях эта величина падает до 0.04. Такая разница создаёт большие сложности при применении в теоретической химии нейронных сетей, да и классических численных подходов тоже.
Классические подходы
Классическая задача квантовой химии формулируется следующим образом: предсказать состояние электронов на всех молекулярных орбиталях для данной конформации, заданной через положения атомов в пространстве и их тип. Звучит сложно, как мы уже отмечали выше, мозг человека плохо справляется с поведением квантовых объектов, потому что в обычной макроскопической жизни тела так себя не ведут.
По этой причине классические методы устроены так, чтобы выдавать понятные физикам и химикам величины. Например, полную энергию системы. Это одно число (скалярная величина), если оно отрицательное, значит атомам «удобно» находиться в таком состоянии, и вещество в принципе может существовать (хотя бы при нулевой температуре по Кельвину).
Помимо этого, есть множество других полезных и в то же время понятных вещей: например, силы, воздействующие на атомы (векторные величины, приписанные к каждому атому), тензоры поляризации (тензорные величины), полные гамильтонианы системы (симметричные матрицы) и распределения электронной плотности (функции в трёхмерном пространстве). Каждую их этих сущностей можно использовать для анализа и понимания каких‑то интересных свойств молекул.
Например, можно предсказать какого цвета будет вещество, какие у него будут спектры в разных аналитических приборах, будет ли оно биологически активно, как хорошо оно будет растворяться в воде и других веществах, будет ли у него активность против онкологических заболеваний. Важный момент здесь в том, что обычные нейронные сети редко предсказывают данные таких типов, поэтому здесь нужны специальные архитектуры.
И вот на этом моменте стоит пояснить, почему многие хотят перейти от классических подходов к машинному обучению. Классические методы квантовой химии — это численные методы. И, что немаловажно, это очень долго считающие численные методы.
Например, такой хороший и точный теоретический метод как связанные кластеры в варианте CCSD(T) имеет вычислительную сложность 
, где 
— это число электронов. Более простая и менее точная теория функционала плотности (density functional theory, DFT) считает несколько быстрее — за 
. Что это означает на практике?
Рассмотрим молекулу с одним кольцом (бензол), где уже имеется 42 электрона. Это значит, что, если у вас есть база хотя бы из одного миллиона веществ с 3–4 бензольными кольцами, то ресурсов всех компьютеров мира не хватит посчитать её за год на уровне теории CCSD(T). Во втором случае всё не так драматично: хватит и пары месяцев, но кластер из 20–30 машин с CPU или GPU вам всё же понадобится. Ну и конечно, если вы хотите тут использовать современные генеративный AI или RL, то у вас не получится это сделать из‑за долгой проверки сгенерированных объектов:((
Графовые нейронные сети
К счастью, люди придумали геометрические графовые нейронные сети (geometric graph neural network, G‑GNN). Почему их называют графовыми? Потому что они, в основном, оперируют абстракциями передачи данных по графу. Почему их называют геометрическими? В первую очередь потому что это часть геометрического глубокого обучения, то есть подход к построению специальных нейронных сетей, инвариантных и эквивариантных к определённым преобразованиям.
Так получается, что все физические законы от классической механики и до уравнений Шрёдингера эквивариантны к операциям вращения и переноса в пространстве (теория относительности и субатомные процессы — нет). Геометрическое машинное обучение разрабатывалось для обучения нейронных сетей именно с такими свойствами. То есть, при изменении данных в виде их переноса в пространстве и поворота, выходы сети сами по себе меняются в соответствии с этим преобразованием. Также геометрические графовые сети работают с геометрическим графом, который естественным образом появляется в пространстве, если мы считаем вершинами некоторые точки, а ребра строим на основании дистанций между этими точками.
G‑GNN предсказывают свойства молекул за 
, где 
— это количество атомов в системе. А если у вас малые молекулы, то за счёт массового параллелизма в GPU, в реальности скорость предсказание масштабируется как 
.
И вот тут мы приходим к главному. Классические решатели для квантовой химии — это устоявшийся класс программного обеспечения, они огромные и считают МНОГО нужных значений за раз. Глубокому обучению для химии чуть больше 12 лет. Методология классического глубокого обучения говорит нам, что нейронные сети должны принимать один элемент из датасета и предсказывать ОДИН результат на выходе.
Можно конечно написать сложную функцию потерь, но обычно это снижает точность предсказания каждого целевого значения по отдельности. В итоге, у нас есть один большой чёрный ящик, который делает много‑много магии (классическое программное обеспечение для квантовой химии). Мы его хотим разделить на много‑много нейронок, каждая из которых предсказывает своё целевое значение.
Хочется остановиться и ещё раз подчеркнуть разницу между глубоким обучением для квантовой химии и обычным машинным обучением. Машинное и глубокое обучение призвано находить паттерны в данных реального мира и предсказывать результаты. В мире квантовой химии уже существуют численные методы, которые очень хорошо повторяют некоторые свойства реальности (с оговорками и нюансами). И вот теперь нейронные сети учатся на синтетических данных этих методов, для того чтобы иметь возможность считать ещё больше случаев. То есть нейронные сети используются как метод ускорения обработки данных, а не как метод обработки данных сам по себе.
Для начала стоит разобраться, как работает предсказание простых свойств молекул, таких, например, как полная энергия. Для того, чтобы что‑то предсказывать, нужно сначала построить граф. Атомы будут вершинами графа, рёбрами станут дистанции между атомами в трёхмерном пространстве.
Существует два основных способа работать с таким графом. Первый, это использовать механизм внимания на графе (англ. Graph Attention). Он близок к классической архитектуре Трансформеров, но требует больше вычислительных ресурсов.
Второй способ, это использовать фреймворк нейронных сетей посылающих сообщения (англ. Message passing). Его суть в следующем:
В начале нужно приписать какие‑то данные в вершины атомов, будем называть их сообщениями. Это шаг кодирования или вложения данных (англ. Embedding).
Далее нужно обменяться информацией по всем ребрам графа. Обычно для этого сообщения суммируются по всем вершинам (англ. Aggregation).
Затем нужно выполнить какую‑то обработку в вершине. Это этап обновления (англ. Update). Операции посылки сообщения и обновления можно сделать несколько раз подряд в цикле.
Далее нужно собрать всю информацию с узлов графа вместе и сделать шаг обновления на некоторую абстрактную вершину, не находящуюся в трехмерном пространстве, которая обеспечивает получение нужного результата в нужной форме. Это шаг считывания (англ. Readout).
Теперь остаётся разъяснить всего несколько базовых терминов для понимания того, как это работает. Начальные данные в вершины графа (то есть, атомы), можно задать с помощью обычных слоёв Embedding из PyTorch, они будут конвертировать номер атома в набор фич из 32–1024 признаков. Также нужно передавать в нейронную сеть дистанцию между атомами. К сожалению, если кодирование атомов будет иметь размер 1024, а кодирование дистанций представлено лишь одним числом, то нейронная сеть не будет учиться на признаки дистанций.
Решение этой проблемы довольно простое: нужно сконвертировать дистанцию в несколько отдельных чисел с помощью явного инжиниринга признаков. Такие библиотеки как Pytorch Geometric, Deep Graph Library, SchNet Pack и e3nn имеют встроенные функции для задания такого кодирования. Обычно для этого используют набор гауссиан, функций Бесселя или синусоиды.
Также стоит поговорить о геометрическом графе. Если построить полный граф с бесконечной дистанцией отсечения, то для молекулы с 
атомами, будет 
рёбер для пересылки сообщения. Это немного по сравнению с классическими подходами, но всё равно неприятно. И вот здесь нам помогает физика: атомы и электроны, находящиеся на значительном удалении (12 Бор и больше), практически не взаимодействуют друг с другом. Поэтому достаточно построить геометрический граф только в окрестностях каждого атома и оставить дистанцию отсечения в 3, 6 или 10 Бор — это значительно ускорит и обучение нейронной сети, и её использование.
Таким образом, графовые нейронные сети выглядят очень перспективным инструментом для квантовой химии. По этой причине наша группа активно занимается этим направлением.
Как мы отмечали выше, графовые нейронные сети часто предсказывают целевые значения, которые не встречаются в обычном глубоком изучении: матрицы, функции или наборы тензоров. Год назад мы создали один из самых больших датасетов с веществами похожими на лекарства — nabla2DFT [5] (подробнее о нём на Хабре рассказывал мой коллега Кузьма Храбров) В нём есть оценки производительности 10 нейронных сетей для предсказания общей энергии (скаляр), сил на атомах (вектора), гамильтонианов электронных систем (скаляры).
В нашем новом исследовании мы разработали новую архитектуру нейронной сети и набор подходов к обучению, которые позволяют предсказывать электронную плотность в реальном пространстве, то есть преобразуют геометрический граф в функцию. При этом наша модель обладает эквивариантностью: если молекулу повернуть или перенести в пространстве, то функция изменится правильным образом.
Но подробнее про нашу работу я расскажу уже в следующем посте.
Материалы
TG‑канал автора статьи и текста Константина Ушенина: эйай фор сайенс.
TG‑канал Института AIRI.
Ссылки
Ushenin, Konstantin, et al. "LAGNet: better electron density prediction for LCAO-based data and drug-like substances." Journal of Cheminformatics 17.1 (2025): 65. https://doi.org/10.1186/s13321-025-01010-7
Schütt, Kristof, et al. "Schnet: A continuous-filter convolutional neural network for modeling quantum interactions." Advances in neural information processing systems 30 (2017).
Kadurin, Artur, et al. "druGAN: an advanced generative adversarial autoencoder model for de novo generation of new molecules with desired molecular properties in silico." Molecular pharmaceutics 14.9 (2017): 3098-3104.
Jumper, John, et al. "Highly accurate protein structure prediction with AlphaFold." nature 596.7873 (2021): 583-589.
Khrabrov, Kuzma, et al. "$\nabla^2$ DFT: A Universal Quantum Chemistry Dataset of Drug-Like Molecules and a Benchmark for Neural Network Potentials." arXiv preprint arXiv:2406.14347 (2024). https://arxiv.org/abs/2406.14347