GNN для молекул и кристаллов

Коротко

Definition

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

В молекулах и кристаллах естественно есть структура:

• атомы;
• связи или соседства;
• расстояния;
• углы;
• локальные окружения;
• периодичность для кристаллов.

Поэтому такие объекты удобно представлять как графы.

Типичная идея:

$$\text{атомы и связи}\to \text{граф}\to \text{GNN}\to \text{свойство}$$

Интуиция

Молекула или кристалл — это не просто таблица признаков. Важно не только, какие атомы есть, но и как они соединены или расположены в пространстве.

Например, две молекулы могут иметь одинаковый состав, но разную структуру и разные свойства.

GNN учитывает это через message passing:

1. Каждый атом имеет начальное представление.
2. Атом получает сообщения от соседей.
3. Его представление обновляется.
4. После нескольких шагов атом содержит информацию о своём окружении.
5. Представления атомов агрегируются в представление всей молекулы или кристалла.
6. По этому представлению предсказывается свойство.

Интуитивно GNN учится отвечать на вопрос: как локальные атомные окружения влияют на глобальное свойство структуры.

Формальное описание

Химическую структуру можно представить графом:

$$G=(V,E)$$

где:

• $V$ — вершины, обычно атомы;
• $E$ — рёбра, химические связи или соседства;
• у вершин есть atom features;
• у рёбер есть bond или distance features.

Для каждой вершины $v$ есть скрытое представление:

$$h_v^{(l)}$$

где $l$ — номер слоя message passing.

Один шаг message passing можно записать так:

$$m_v^{(l)}=\sum _{u\in \mathcal{N}(v)}{M}_l(h_v^{(l)},h_u^{(l)},e_{uv})$$ $$h_v^{(l+1)}=U_l(h_v^{(l)},m_v^{(l)})$$

где:

• $\mathcal{N}(v)$ — соседи вершины $v$;
• $e_{uv}$ — признаки ребра между $u$ и $v$;
• $M_l$ — функция сообщения;
• $U_l$ — функция обновления.

После нескольких слоёв получают представления атомов, а затем агрегируют их в представление всей структуры:

$$h_G=READOUT(\{h_v{\}}_{v\in V})$$

По $h_G$ предсказывают свойство:

$$\hat{y}=f(h_G)$$

Молекулы как графы

Для молекулы обычно:

• вершины — атомы;
• рёбра — химические связи;
• признаки вершин — элемент, заряд, валентность, ароматичность, гибридизация;
• признаки рёбер — тип связи, сопряжённость, ароматичность, расстояние.

Пример:

Элемент графа	Химический смысл
Вершина	атом
Ребро	химическая связь
Node feature	тип атома, заряд, валентность
Edge feature	тип связи, расстояние
Graph target	токсичность, растворимость, энергия, активность

Для молекул GNN часто применяют к задачам:

• предсказание растворимости;
• токсичность;
• binding affinity;
• molecular property prediction;
• генерация молекул;
• классификация активности.

Кристаллы как графы

Кристаллы сложнее молекул, потому что у них есть периодическая структура.

Для кристалла обычно:

• вершины — атомы в элементарной ячейке;
• рёбра — соседства между атомами в радиусе cutoff;
• учитываются периодические копии ячейки;
• edge features часто включают расстояния и иногда направления.

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

Примеры задач:

• formation energy;
• band gap;
• elastic properties;
• stability;
• conductivity;
• catalytic properties.

Это напрямую связано с ML для материалов и проектом Materials ML.

Message passing

Message passing — центральная идея большинства GNN.

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

Интуитивно:

$$\text{атом}\leftarrow \text{соседи}\leftarrow \text{соседи соседей}$$

Это похоже на то, как химическое свойство атома зависит не только от самого атома, но и от его локального окружения.

Readout

После message passing нужно получить представление всей молекулы или кристалла.

Для этого используют readout:

• sum pooling;
• mean pooling;
• max pooling;
• attention pooling;
• set2set;
• более сложные global aggregation.

Пример:

$$h_G=\sum _{v\in V}{h}_v$$

или:

$$h_G=\frac{1}{|V|}\sum _{v\in V}{h}_v$$

Затем $h_G$ подают в prediction head, например полносвязную сеть.

Архитектуры

CGCNN

Crystal Graph Convolutional Neural Network — одна из классических архитектур для кристаллов.

Идея:

• представить кристалл как граф атомов;
• учитывать соседства атомов;
• обновлять atom embeddings через graph convolution;
• агрегировать представления атомов;
• предсказывать свойство кристалла.

CGCNN важна как базовая архитектура для понимания ML на кристаллических графах.

SchNet

SchNet — GNN для молекул и материалов, которая использует непрерывные фильтры и расстояния между атомами.

Она хорошо показывает идею, что геометрия и расстояния могут быть встроены в message passing.

MEGNet

MEGNet использует graph networks и может учитывать:

• node features;
• edge features;
• global state attributes.

Это удобно для материалов, где глобальные условия или свойства структуры тоже могут быть важны.

M3GNet и MatGL

M3GNet и MatGL-подходы связаны с современными pretrained models для материалов.

Они полезны, когда нужно использовать уже обученные представления и делать transfer learning на новую задачу.

Equivariant GNN

Equivariant GNN учитывают симметрии пространства.

Для 3D-структур важно, чтобы модель корректно реагировала на:

• перенос;
• поворот;
• отражение, если применимо;
• перестановку атомов.

Это связано с геометрическим ML.

Какие свойства предсказывают

GNN для молекул и кристаллов используют для разных задач.

Регрессия

Примеры:

• formation energy;
• band gap;
• total energy;
• forces;
• solubility;
• binding affinity;
• adsorption energy;
• catalytic activity.

Это регрессионные задачи, оцениваемые через Метрики качества регрессоров.

Классификация

Примеры:

• активна / не активна;
• токсична / нетоксична;
• стабильна / нестабильна;
• проводник / полупроводник / изолятор.

Это задачи классификации.

Генерация

GNN могут быть частью генеративных моделей для молекул или материалов.

Но генерация химически валидных структур сложнее обычного property prediction, потому что нужно учитывать валентность, стабильность, синтезируемость и физические ограничения.

Когда использовать

GNN стоит использовать, если:

• объект естественно является графом;
• важны связи между атомами;
• доступна структура молекулы или кристалла;
• табличных descriptors недостаточно;
• нужно учитывать локальные окружения;
• есть достаточно данных или pretrained model;
• задача связана с molecular/material property prediction.

В materials ML GNN особенно полезны, когда структура материала важнее, чем один только состав.

Когда не использовать

GNN может быть плохим выбором, если:

• данных очень мало;
• структуры ненадёжные или отсутствуют;
• достаточно простых табличных descriptors;
• нужна максимально простая и интерпретируемая модель;
• вычислительные ресурсы ограничены;
• baseline на градиентном бустинге уже решает задачу достаточно хорошо;
• graph construction плохо определён.

Перед сложной GNN полезно построить baseline на табличных данных.

Типичные проблемы

Построение графа

Для молекул связи часто известны. Для кристаллов рёбра обычно строятся по cutoff или nearest neighbors.

Если правило построения графа выбрано плохо, модель будет получать искажённую структуру.

Периодичность

Кристаллы имеют periodic boundary conditions. Нужно учитывать соседей из соседних ячеек.

Если периодичность игнорировать, модель потеряет важную часть структуры.

Инвариантности и эквивариантности

Свойства материала не должны зависеть от произвольного поворота, переноса или перестановки одинаковых атомов.

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

Data leakage

Похожие материалы могут попасть и в train, и в test.

Это завышает качество и создаёт иллюзию хорошего обобщения.

Нужны аккуратные splits: по составам, chemical systems, scaffold-like groups или источникам данных.

Малые датасеты

Многие materials datasets малы. GNN может переобучиться, особенно если модель большая.

В таких случаях помогают:

• pretrained models;
• transfer learning;
• cross-validation;
• регуляризация;
• простые baseline.

Интерпретация

GNN сложнее интерпретировать, чем модели на ручных descriptors.

Важно дополнительно анализировать:

• ошибки;
• attention или contribution scores, если они есть;
• чувствительность к структуре;
• поведение на химических семействах.

Минимальный пример

Пусть есть молекула воды.

Граф:

Вершина	Атом
0	O
1	H
2	H

Рёбра:

Откуда	Куда
0	1
0	2

GNN обновляет представление атома кислорода на основе соседних атомов водорода, а представления водородов — на основе кислорода.

После message passing модель агрегирует атомные представления:

$$h_G=READOUT(h_O,h_{H1},h_{H2})$$

Затем prediction head предсказывает свойство молекулы:

$$\hat{y}=f(h_G)$$

Например, растворимость или энергию.

Практические замечания

Хороший workflow для GNN в materials ML:

1. Определить target property.
2. Проверить качество структур.
3. Выбрать способ построения графа.
4. Задать node и edge features.
5. Сделать честный train/validation/test split.
6. Построить baseline на табличных descriptors.
7. Обучить GNN.
8. Сравнить качество с baseline.
9. Проанализировать ошибки по химическим семействам.
10. Проверить, насколько модель обобщает на новые структуры.

Для проекта Materials ML важно не только обучить GNN, но и объяснить, почему графовое представление подходит к выбранному датасету и target property.

Связанные понятия

Что знать перед этим

• Графовые данные
• Табличные данные
• Задачи регрессии
• Метрики качества регрессоров

Связанные заметки

• Трансформер