Всем привет! На связи снова Константин Ушенин из AIRI, и мы продолжаем говорить о глубоком обучении в квантовой химии. В прошлом посте мы немного разобрались в том, что такое молекула, как её представлять в компьютере, и как работают графовые нейронные сети.

В этот же раз я расскажу о том, какие результаты в этой области получила наша команда. Речь пойдет о новой архитектуре для предсказания электронной плотности LAGNet, про которую у нас недавно вышла статья в Journal of Cheminformatics [1]. Мы применили несколько интересных усовершенствований к модели DeepDFT, что позволило в 8 раз снизить требование к объёму необходимых данных и в целом сделало выучивание плотности более эффективным.

Но обо всём по порядку.

Архитектура DeepDFT 

В отличие от простейшего предсказания скалярного значения, которое было разобрано в первой части, вычисление электронной плотности — это работа с функцией. То есть, теперь нейронная сеть должна преобразовывать граф в функцию, а полученная функция должна возвращать необходимое значение в любой точке пространства. Задать такую нейронную сеть можно несколькими способами. Конкретно наша работа опирается на концепцию, предложенную в 2020 году в архитектуре DeepDFT [2].

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

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

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

Архитектура LAGNet

Тут мы вплотную подходим к тому, что представляет собой предложенная нами архитектура LAGNet. Это улучшение DeepDFT, специально предназначенное для молекул, похожих на лекарства (англ. druglike molecules). Этот термин обозначает такие молекулы, которые потенциально имеют биологическую активность. Обычно они состоят из неметаллических элементов (точнее H, C, N, O, F, P, S, Cl, Se, Br, I), не являются сильными щелочами или кислотами, растворяются в воде или жире.

Вообще, датасетов по схожим с лекарствами молекулам не так уж и много. Наша команда активно стремится закрыть этот пробел. В 2022 году мы собрали датасет ∇DFT с информацией о более, чем миллионе молекул и пяти миллионах конформаций [3], а в прошлом году увеличили эти цифры до почти двух миллионов молекул и до 15,7 миллиона конформаций в ∇²DFT [4]. Подробности можно также найти в недавней статье на Хабре.

∇²DFT нам очень пригодился, чтобы заточить идеи, заложенные в DeepDFT, под указанный тип веществ. Фактически, LAGNet относится к DeepDFT примерно также, как 3D‑UNet или U‑Net++ относятся к основной архитектуре U‑Net. Это набор улучшений, которые увеличивают точность базовой архитектуре и делают её лучше в других аспектах.

Кодирование дистанций

Первое улучшение — это кодирование дистанций. Так сложилось, что специализированные пакеты для геометрических графовых нейронных сетей (SchNetPack, DGL, Pytorch Geometric, библиотека e3nn) по умолчанию предлагают для дистанций такие методы инжиниринга признаков, которые являются оптимальными для расстояния между атомами.

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

Проблема в том, что дистанция между атомом и точкой зондирования может быть совершенно любой. Получить значение плотности может понадобиться на дистанции 0.001, 0.1, 1 или 10 Бор от атома. Распределение всех дистанций между атомами и точками пробинга также отличается от распределения дистанций между атомами. Простое использование функций Бесселя приводит к огромным значениям на маленьких дистанциях, потому что эти функции имеют большой горб около нуля. Это одно из первых мест для улучшения DeepDFT, которое мы обнаружили.

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

Для остатка дистанции, если он есть, LAGNet использует вторую пересылку с синусным кодированием. В ходе наших экспериментов мы выяснили, что функции Бесселя хорошо работают при передаче сообщений на дистанции до 3 Бор, а синусоиды — на дистанциях более 8 Бор. Соответственно, мы объединили две эти идеи и сделали множественную независимую пересылку.

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

Семплирование с помощью решёток Лебедева

Следующей проблемой является семплирование точек пробинга. DeepDFT использует данные на равномерной трехмерной решетке. Это работает так: вокруг молекулы строится коробка с отступом от краёв, потом с маленьким шагом по каждой из трёх координат строятся точки. Из полученной решётки для каждой мини‑выборки случайно берётся 1000 точек.

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

Получается, что для обучения сети на функцию плотности нужно потратить огромный вычислительный ресурс, чтобы предподсчитать значения электронной плотности в огромном количестве точек в пространстве — и не все из них одинаково полезны. Слово за слово, оказывается, что один пример обучения для одной молекулы весит не как маленький короткий текстовый файлик в десяток килобайт, а 5–6 мегабайт, так как включает в себя данные примерно по 2–3 миллионам точек. А если в нашем датасете 12 миллионов drug‑like молекул, то все данные весят десятки терабайт (86.9 в нашем случае)!

Как вы понимаете, это очень много. А, главное, что это всё ещё нужно предподсчитывать перед обучением, пересылать по сети и где‑то хранить. И, конечно, у нас скорость обучения на таких больших данных линейно зависит от размера этих данных, просто по той причине, что узким местом становится их прокачка на карту. Такие объёмы делают технологию совершенно недоступной для практического использования. Что же делать?

Решение проблемы нашлось в трудах советского математика Вячеслава Ивановича Лебедева, которые тот опубликовал ещё в 70-х годах в СССР. Он изучал вопрос о том, как размещать наборы точек на равных расстояниях на поверхности сферы. Школьная математика говорит, что есть только пять фигур, которые задают множество эквидистантных точек на сфере — это вершины платоновых тел: тетраэдра, куба, октаэдра, додекаэдра, икосаэдра. И больше никак.

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

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

На самом деле, такой подход уже 30 лет используется в классических методах — там это называют стандартной решеткой (англ. the standard grid). Мы же решили это перенести в область работы с данными для глубокого обучения. В итоге размер одного примера сократился в восемь раз: с 5.8 Мб до 0.67 Мб, а полный размер датасета — с 86.9 до 10.19 Тб. Это конечно всё ещё много, и мы всё равно учили нейронки на подвыборке, но это делает технологию уже более доступной для обучения на огромных данных.

Подавление остовных орбиталей

И, наконец, третье серьёзное изменение. Как мы писали в прошлом посте, орбитали бывают внутренние (остовные) и внешние (валентные). Первые «упакованы» плотнее: электронная плотность рядом с атомом Br в молекуле бромводорода достигает около 25,000 электрон на кубический Бор, в то время как в химических связях эта величина падает до 0.04.

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

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

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

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

Заключение

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

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

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

Спасибо за внимание!

Материалы

• TG‑канал автора статьи и текста Константина Ушенина: эйай фор сайенс.

• TG‑канал Института AIRI.

Ссылки:

1. Ushenin, Konstantin, et al. "LAGNet: better electron density prediction for LCAO-based data and drug-like substances." Journal of Cheminformatics 17.1 (2025): 65.

2. Jørgensen, Peter Bjørn, and Arghya Bhowmik. "DeepDFT: Neural message passing network for accurate charge density prediction." arXiv preprint arXiv:2011.03346 (2020).

3. Khrabrov, Kuzma, et al. “nablaDFT: Large-Scale Conformational Energy and Hamiltonian Prediction benchmark and dataset.” Physical Chemistry Chemical Physics 24, 25853 (2022). 

4. 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).

5. Su, Jianlin, et al. "RoFormer: Enhanced Transformer with Rotary Position Embedding." arXiv preprint arXiv:2104.09864 (2021).