Геометрическое ML в химии

Коротко

Definition

Геометрическое ML в химии — это применение моделей, которые учитывают структуру молекул, кристаллов и материалов как графов или 3D-объектов, а также встраивают физические симметрии прямо в архитектуру.

В химии и материаловедении данные часто имеют геометрическую структуру:

• молекула — это атомы, связи и 3D-координаты;
• кристалл — это атомы, координаты, периодичность и решётка;
• материал — это не просто таблица признаков, а объект с внутренней структурой.

Геометрическое ML пытается использовать эту структуру напрямую, вместо того чтобы сводить всё к плоскому вектору признаков.

Ключевые идеи:

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

Интуиция

Обычная табличная модель видит объект как набор признаков:

$$x=(x_1,x_2,\dots ,x_m)$$

Но молекула или кристалл — это не просто список чисел.

Например, если поменять порядок атомов в файле, сама молекула не изменится. Если повернуть молекулу в пространстве, её энергия тоже не изменится. Если перенести молекулу на другой участок координатной системы, её внутренние свойства не должны поменяться.

Хорошая модель для химии должна понимать такие вещи.

Геометрическое ML встраивает это в архитектуру:

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

Так модель получает правильные ограничения ещё до обучения. Это повышает data efficiency и уменьшает риск физически бессмысленных предсказаний.

Основные идеи

Симметрии в химических данных

В химии важны преобразования, которые меняют описание объекта, но не меняют сам физический объект.

Группа	Что делает	Пример
$T(3)$	Трансляции в 3D	Перенос молекулы не меняет её энергию
$SO(3)$	Вращения в 3D	Поворот молекулы не меняет скалярные свойства
$E(3)$	Вращения + трансляции	Полное движение твёрдого тела
$S_n$	Перестановки	Порядок атомов в списке не должен менять молекулу

Для молекул и кристаллов особенно важны:

• permutation invariance;
• translation invariance;
• rotation invariance;
• rotation equivariance;
• periodicity для кристаллов.

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

Инвариантность означает, что выход не меняется при преобразовании входа:

$$f(g\cdot x)=f(x)$$

Пример: энергия молекулы. Если повернуть молекулу, её энергия должна остаться той же.

Эквивариантность означает, что выход меняется согласованно с входом:

$$f(g\cdot x)=g\cdot f(x)$$

Пример: силы. Если повернуть молекулу, векторы сил должны повернуться вместе с ней.

Свойство	Тип реакции модели	Пример
Энергия	Инвариантность	Поворот молекулы не меняет $E$
HOMO-LUMO gap	Инвариантность	Нумерация атомов не меняет свойство
Силы	Эквивариантность	Поворот молекулы поворачивает $\vec{F}$
Координатные обновления	Эквивариантность	Новые координаты должны двигаться вместе со старыми

Молекулы и кристаллы как графы

Молекулу можно представить как граф:

• узлы — атомы;
• рёбра — связи или соседства;
• признаки узлов — тип атома, заряд, гибридизация;
• признаки рёбер — расстояние, тип связи, направление, периодический сдвиг.

Кристалл тоже можно представить как граф, но с дополнительными особенностями:

• периодические граничные условия;
• элементарная ячейка;
• координаты атомов;
• соседства между атомами с учётом периодических копий;
• параметры решётки.

Такой подход связан с Графовые данные и GNN для молекул и кристаллов.

Message Passing Neural Networks

Базовый подход для графов — message passing.

На каждом слое узлы обмениваются сообщениями с соседями.

Общая схема:

$$m_{ij}={\varphi }_e(h_i,h_j,e_{ij})$$ $$h_i^{new}={\varphi }_h\left(h_i,\sum _{j\in N(i)}{m}_{ij}\right)$$

где:

• $h_i$ — признаки узла $i$;
• $e_{ij}$ — признаки ребра между $i$ и $j$;
• $m_{ij}$ — сообщение от узла $j$ к узлу $i$;
• $N(i)$ — соседи узла $i$.

Если модель использует только расстояния:

$$\Vert x_i-x_j\Vert$$

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

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

E(3)-эквивариантные модели

E(3)-эквивариантные модели учитывают не только графовую структуру, но и геометрию в 3D-пространстве.

В таких моделях обновляются:

• скрытые признаки атомов;
• иногда координаты или векторные признаки;
• сообщения с учётом расстояний и направлений.

Пример идеи EGNN:

$$m_{ij}={\varphi }_e(h_i,h_j,\Vert x_i-x_j{\Vert }^2)$$

Обновление координат:

$$x_i^{new}=x_i+\sum _{j\ne i}(x_i-x_j){\varphi }_x(m_{ij})$$

Обновление признаков:

$$h_i^{new}={\varphi }_h\left(h_i,\sum _j{m}_{ij}\right)$$

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

Физически корректное предсказание энергии и сил

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

Физическая связь:

$${\vec{F}}_i=-\frac{\partial E}{\partial {\vec{x}}_i}$$

То есть сила — это антиградиент потенциальной энергии по координатам атома.

Можно обучать модель так:

1. Подать молекулу или кристалл в модель.
2. Предсказать скалярную энергию $\hat{E}$.
3. Через autodiff вычислить градиент энергии по координатам.
4. Получить силы:

$${\hat{\vec{F}}}_i=-\frac{\partial \hat{E}}{\partial {\vec{x}}_i}$$

5. Считать loss и по энергии, и по силам:

$$L={\lambda }_E\Vert E-\hat{E}{\Vert }^2+{\lambda }_F\Vert \vec{F}-\hat{\vec{F}}{\Vert }^2$$

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

Топологический анализ данных

Иногда локального message passing недостаточно, чтобы увидеть глобальную форму молекулы или материала.

Топологический анализ данных, TDA, может добавлять признаки о глобальной структуре:

• связные компоненты;
• циклы;
• полости;
• устойчивые топологические элементы.

Через persistent homology отслеживают, когда топологический объект появляется и исчезает при изменении масштаба.

Примеры:

Объект	Смысл
$H_0$	Связные компоненты
$H_1$	Циклы, кольца
$H_2$	Полости

Результат persistent homology часто нужно векторизовать, чтобы подать в ML-модель.

Варианты:

• persistence image;
• persistence landscape;
• persistence entropy;
• статистики persistence diagram.

Итоговая модель может использовать объединение:

$$Input=[GNN\_Embedding]\oplus [TDA\_Vector]$$

Это полезно, если глобальная форма объекта важна для свойства.

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

Геометрическое ML стоит использовать, когда:

• объект имеет графовую или 3D-структуру;
• важны атомы, связи, расстояния, углы или координаты;
• нужно предсказывать свойства молекул;
• нужно предсказывать свойства кристаллов;
• важны физические симметрии;
• данных не так много, и хочется встроить inductive bias;
• нужно предсказывать энергии, силы, band gap, formation energy или другие свойства материалов;
• простая табличная featurization теряет важную структурную информацию.

Типичные задачи:

• molecular property prediction;
• crystal property prediction;
• interatomic potentials;
• force field learning;
• materials screening;
• prediction of formation energy;
• prediction of band gap;
• molecular dynamics surrogate models.

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

Геометрическое ML может быть избыточным, если:

• есть только готовые табличные дескрипторы;
• структура объекта неизвестна или ненадёжна;
• данных очень мало для обучения сложной модели;
• задача хорошо решается простым baseline;
• важна максимальная интерпретируемость;
• нет ресурсов на построение графов, соседств и 3D-представлений;
• свойство почти полностью определяется простыми composition-based признаками.

Для первого baseline в materials ML часто полезно сравнить:

• простую табличную модель;
• composition-based descriptors;
• Градиентный бустинг;
• GNN;
• pretrained модели для материалов.

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

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

import torch
from torch import nn
 
class ToyEnergyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.mlp = nn.Sequential(
            nn.Linear(1, 16),
            nn.SiLU(),
            nn.Linear(16, 1),
        )
 
    def forward(self, coordinates):
        distances = torch.cdist(coordinates, coordinates)
        upper = distances.triu(diagonal=1)
        feature = upper.sum().view(1, 1)
        energy = self.mlp(feature)
        return energy.squeeze()
 
coordinates = torch.randn(5, 3, requires_grad=True)
 
model = ToyEnergyModel()
energy = model(coordinates)
 
forces = -torch.autograd.grad(
    energy,
    coordinates,
    create_graph=True,
)[0]
 
print("energy:", energy)
print("forces shape:", forces.shape)

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

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

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

• Графовые данные
• GNN для молекул и кристаллов
• Метрики качества регрессоров

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

• ML для материалов
• Materials ML
• Регуляризация

Связанные проекты

• Materials ML