эквивариантное обучение (GeomML) с байесовским мультизадачным лоссом и топологическим регуляризатором (TDA) на датасете QM9Spectra

1. Введение

В проеке попытка разработки квантово-топологического DL-калькулятора на стыке геометрического глубокого обучения и тополгического анплиза TDA для быстрого предсказания многомерных физических полей молекул на основе их 3D-скелета. В отличие от прежних попыток реализиации (линейных без inductive biases) разработанная модель напрямую реконструирует истинные пространственные физические объекты: 3D-вектор дипольного момента ($\vec{\mu}$, геометрический ранг 1o) и полную 3D-матрицу декартова тензора поляризуемости ($\alpha$) 0e+2e

В качестве датасета использован полный квантовохимический QM9Spectra (QM9S) объемом 129 817 молекул. Позволяет не тратить время на самостоятельный расчет обучающих данных как при использовании Alchemy.

2. Архитектура QM9S_net и геометрические приоры

Архитектура сети жестко ограничивает пространство поиска используя встроенные априорные индуктивные смещения:

1. Входной порядковый атомный номер элемента ($z$) посредством скалярной свертки разворачивается в 16-канальный скрытый изотропный "паспорт" внутренних квантовых свойств атома (электроотрицательность, валентная емкость оболочек, радиус экранирования ядра).
2. 3D-векторы связей проецируются на фиксированный базис неприводимых сферических гармоник группы вращений $SO(3)$ и унитарных симметрий $SU(2)$ рангов $L=1$ и $L=2$. На ребрах графа коэффициенты Клебша — Гордана выполняют роль пространственных линейных операторов. Информация распространяется каскадным наложением по двум слоям Message Passing, формируя непрерывное глубокое молекулярное поле.
3. Алгоритм giotto-tda батчами обсчитывает персистентную энтропию Шеннона фильтраций комплексов Вьеториса — Рипса по гомологиям $H_0$ и $H_1$. Полученные константные скаляры выступают "вертолетным снимком" глобального ландшафта молекулы (моменты рождения и гибели колец и циклов), накладывая геометрическое вето на финальную сборку полей.

3. Декомпозиция сферического тензора 3х3

Для полной математической строгости симметричная декартова матрица поляризуемости $3 \times 3$ (содержащая 6 уникальных чисел за счет зеркальности по осям $xy=yx, xz=zx, yz=zy$) разлагается средствами PyTorch на два неприводимых представления (Irreps) группы $SO(3)$:

• Изотропный сферический след (ранг 0e) $\to$ 1 число: Среднее диагонали, описывающее базовый объем электронного облака.
• Анизотропный девиатор (ранг 2e) $\to$ 5 чисел: Бесследовая матрица девиатора. Для идеального сопряжения с фильтрами Клебша-Гордана ребер, элементы переведены в канонический ортонормированный сферический базис по возрастанию квантового проекционного числа $m$ (от $-2$ до $+2$):
  1. $m = -2 \to \mathbf{xy}$
  2. $m = -1 \to \mathbf{yz}$
  3. $m = \ \ 0 \to \mathbf{z^2}$ (нормированная комбинация $(2zz - xx - yy)/\sqrt{3} \to -3(xx+yy)/\sqrt{3}$)
  4. $m = +1 \to \mathbf{xz}$
  5. $m = +2 \to \mathbf{x^2-y^2}$ (нормированная комбинация $(xx-yy)/2$)

Компоненты девиатора физически эквивалентны пяти волновым лепесткам квадрупольных $d$-орбиталей.

4. Динамическое байесовское взвешивание потерь

Вместо ручного подбора констант лосса внедрен метод Кендаллана на основе максимизации совместного гауссова правдоподобия многоцелевого пространства. В критерий ошибки full_train_loss внедрены 3 независимых обучаемых параметра логарифмов дисперсий ($\ln\sigma^2$) — по одному на каждое ортогональное поле: $$\mathcal{L}{\text{total}} = \exp(-s_1)\mathcal{L}{\mu} + \exp(-s_2)\mathcal{L}{\text{trace}} + \exp(-s_3)\mathcal{L}{\text{deviator}} + s_1 + s_2 + s_3$$ Параметры логарифмов дисперсий шума $s_i$ оптимизируются AdamW одновременно с весами сети. Логарифмический штраф служит математическим барьером, не позволяющим модели завышать неопределенность для снижения градиентного давления (защита от тривиальных паразитических решений). Многозадачный лосс удерживает систему в непрерывном балансе, распределяя приоритет градиентов по эпохам.

5. Итоговые результаты на тестовой выборке (components_R2_plots)

Многозадачное усреднение линейных ошибок MAE заставило сеть оптимизировать глобальную физическую адекватность всей макро-молекулы. На изолированном тестовом подмножестве полного датасета QM9S (который модель не видела в процессе обучения) коэффициенты детерминации ($R^2$) зафиксировали уровень сходимости:

Пространственное квантовое свойство	Неприводимое представление (Irreps)	Точность аппроксимации ($R^2$)
Компоненты вектора Диполя $\mu$ (X, Y, Z)	1x1o	*0.8955
След матрицы Поляризуемости $\alpha$ (0e)	1x0e	0.9347
Девиаторные d-компоненты матрицы $3\times3$	1x2e	0.0723

Анализ диаграмм рассеяния:

• Красный клинок диполя и синяя монолитная линия следа практически идеально легли на диагональ $Y=X$. Точность следа в 93.47% доказывает достаточную силу связки инвариантов 0e и топологического маяка TDA.
• Устойчивый положительный выход девиатора ($R^2 = 0.0723$) и визуальное стягивание (хоть и небольшое) красного облака к центру тренда подтверждают, что байесовские тиски Кендалла успешно зафиксировали высокочастотный шум $d$-гармоник, защитив макро-параметры от перекоса градиентов. При других попвтках без взвешеной ошибки результаты были более плоскогоризонтальны, принимали среднее значение в ходе обучения.

ИТОГ: По итогу данного проекта видна определенная уязвимость вычисления тензора поляризации из-за сложной структуры сферических гармоник и в отличии от диполя и скляра в данном случае TDA ландшафт не сильно способствует улучшению точности (отличие от других моделей на порядок, но так или иначе есть что улучшать). Заложенные симметрии и приоры помогают обеспечить в связке с TDA хорошую предсказательную силу для вектора диполя и скаляра внутренних характеристик (след тензора 3*3). Углубление структуры сети и более точная настройка гиперпараметров может улучшить данный результат в дальнейшем.