Хабр·Russia·Friday, October 16, 2026 at 08:18 AM UTC

Применение физически информированных нейронных сетей (PINN) в пучковых нейросетях на графах

Translate

xxxena6 минут назад

Привет, Хабр! Мы студенты мастерской по сервисам и платформам ИИ Инженерно-математической школы НИУ ВШЭ и VK: Ксения Белькова, Тимур Лиджиев и Мария Стеблий. В этой статье хотим рассказать об исследовании, в котором изучали, можно ли применять физически информированные нейронные сети (PINN) для аппроксимации диффузионных процессов в графовых нейронных сетях на пучках.

Исследование мы проводили совместно с Александром Таракановым, инженером машинного обучения AI VK и доцентом ФКН НИУ ВШЭ. В статье расскажем о подходе, проведенных экспериментах и полученных результатах.

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

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

Обозначим через представление вершины на слое . Тогда в общем виде message passing можно записать так:

$\begin{aligned}m_v^{t+1} &= \operatorname{AGGR}\left(\left\{h_u^t \mid u \in \mathcal{N}(v)\right\}\right),\\h_v^{t+1} &= \sigma\left(W^{t+1} \operatorname{CONCAT}\left(m_v^{t+1}, h_v^t\right)\right). \end{aligned}$

Когда слоёв становится больше, у такого подхода возникают проблемы — например, over-smoothing и over-squashing. В первом случае представления разных вершин постепенно становятся слишком похожими друг на друга. Во втором случае информацию от большого числа удалённых вершин приходится передавать через ограниченное число промежуточных представлений.

Рисунок 1. Пример графа с пятью вершинами и связями между ними.

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

На основе пучка (sheaf) строится sheaf Laplacian — оператор, обобщающий обычный графовый лапласиан. Если обычный лапласиан определяет взаимодействие признаков соседних вершин, то sheaf Laplacian дополнительно учитывает линейные restriction maps, связывающие локальные пространства признаков. Поэтому распространение признаков можно рассматривать как диффузионный процесс с дополнительной структурой локальных преобразований (Bodnar et al. 2022, Barbero et al. 2022).

Задачу распространения признаков можно записать в виде дифференциального уравнения

$\begin{equation} \frac{\partial u(t)}{\partial t} = -L_{\mathrm{sheaf}}u(t), \qquad u(0)=u_0,\end{equation}$

где $L_{\mathrm{sheaf}}$ — sheaf Laplacian, а — начальное состояние. В такой постановке sheaf Laplacian не только описывает взаимодействие признаков на графе, но и определяет непрерывную во времени диффузионную динамику.

Мы предлагаем использовать Physics-Informed Neural Network (PINN) для аппроксимации решения диффузионной задачи. Вместо того чтобы напрямую интегрировать уравнение численно, нейронная сеть обучается предсказывать состояние системы в зависимости от времени, а корректность получаемого решения контролирует невязка исходного дифференциального уравнения (Raissi et al. 2019). Диффузионное уравнение выступает как ограничение при обучении модели.

Дополнительная особенность предлагаемого подхода — использование представления Крылова. Мы не передаём в нейронную сеть плотное представление оператора $L_{\mathrm{sheaf}}$ , а используем последовательность его применений к начальному состоянию:

$\begin{equation} u_0,\; L_{\mathrm{sheaf}}u_0,\; L_{\mathrm{sheaf}}^2u_0,\; \ldots,\; L_{\mathrm{sheaf}}^K u_0.\end{equation}$

Такое представление сохраняет информацию о действии оператора на начальное состояние и позволяет использовать разреженность графового оператора (Saad 1992).

Метод объединяет обучаемое sheaf-представление графа и PINN, и их параметры оптимизируются совместно. Sheaf Laplacian определяет диффузионную динамику, Krylov-представление используется как компактное представление действия оператора, а PINN аппроксимирует зависящее от времени состояние системы. Обучение дополнительно учитывает начальное условие и невязку диффузионного уравнения.

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

Мы рассматриваем два варианта:

аппроксимация sheaf diffusion с использованием Krylov-представления
применение предварительно обученной PINN для аппроксимации действия матричной экспоненты внутри Bundle Neural Networks

Теория

Граф

Пусть задан граф , где — множество вершин, а — множество рёбер.

Каждой вершине $v\in V$ сопоставим вектор признаков $x_v \in \mathbb{R}^d$ .

Все признаки можно объединить в вектор $\begin{equation} x = \begin{bmatrix} x_1 \\ \vdots \\ x_{|V|} \end{bmatrix} \in \mathbb{R}^{|V| d}.\end{equation}$

Для нескольких каналов признаков будем использовать матрицу $X\in\mathbb{R}^{|V|d\times f}$ .

Клеточный пучок

Клеточный пучок $\mathcal{F}$ над графом сопоставляет каждой вершине векторное пространство $\mathcal{F}(v)$ , каждому ребру пространство $\mathcal{F}(e)$ , а каждой паре «вершина — инцидентное ребро» линейное отображение $\begin{equation} \mathcal{F}_{v\subset e}: \mathcal{F}(v)\rightarrow\mathcal{F}(e),\end{equation}$ называемое $\textit{restriction map}$ .

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

Предположим, что все stalk имеют размерность , поэтому restriction maps — матрицы размера $d\times d$ .

Рисунок 2. Схематическое представление клеточного пучка на графе.

На рисунке показано, что значения в вершинах сначала принадлежат своим локальным пространствам $\mathcal{F}(v)$ , а restriction maps переводят и в пространство соответствующего ребра $\mathcal{F}(e)$ . Поэтому признаки соседних вершин можно сравнивать после их отображения в общее пространство.

0- и 1-коцепи

Пространство значений на вершинах задаётся как $\begin{equation} C^0(G,\mathcal{F}) = \bigoplus_{v\in V}\mathcal{F}(v),\end{equation}$

а пространство значений на рёбрах: $\begin{equation} C^1(G,\mathcal{F}) = \bigoplus_{e\in E}\mathcal{F}(e)\end{equation}$ .Таким образом, элемент $x\in C^0(G,\mathcal{F})$ — это набор векторов , по одному для каждой вершины.

Coboundary operator

Выберем ориентацию каждого ребра. Для ребра определим оператор $\begin{equation} \delta: C^0(G,\mathcal{F}) \rightarrow C^1(G,\mathcal{F})\end{equation}$ как $\begin{equation} (\delta x)_e = \mathcal{F}_{v\subset e}x_v - \mathcal{F}_{u\subset e}x_u\end{equation}$ . Таким образом, $\delta$ сравнивает значения на концах ребра после их отображения в общее пространство $\mathcal{F}(e)$ .

Sheaf Laplacian

Sheaf Laplacian определяется как $\begin{equation} L_{\mathcal{F}} = \delta^{\top}\delta.\end{equation}$ Для любого $x\in C^0(G,\mathcal{F})$ $\begin{equation} x^{\top}L_{\mathcal{F}}x = \|\delta x\|^2 \geq 0\end{equation}$ , поэтому $L_{\mathcal{F}}$ — симметричная положительно полуопределённая матрица.

Диагональный блок, соответствующий вершине , имеет вид $\begin{equation} L_{\mathcal{F}}(v,v) = \sum_{e\ni v} \mathcal{F}_{v\subset e}^{\top} \mathcal{F}_{v\subset e}\end{equation}$ . Для соседних вершин и , соединённых ребром , $\begin{equation} L_{\mathcal{F}}(v,u) = - \mathcal{F}_{v\subset e}^{\top} \mathcal{F}_{u\subset e}\end{equation}$ . В тривиальном случае, когда все stalk одномерны, а restriction maps тождественны, эта конструкция сводится к обычному графовому лапласиану.

Sheaf diffusion

Для обычного графа диффузионный процесс задаётся уравнением $\begin{equation} \dot{X}(t) = -LX(t)\end{equation}$ , где — графовый лапласиан. Это уравнение описывает распространение признаков между соседними вершинами. Оно служит дискретным аналогом уравнения теплопроводности.

Sheaf Laplacian обобщает эту конструкцию: он заменяет обычный графовый лапласиан матричным оператором, учитывающим restriction maps. Поэтому аналогичный диффузионный процесс принимает вид $\begin{equation} \dot{X}(t) = -\Delta_{\mathcal{F}}X(t)\end{equation}$ .

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

Bundle Neural Networks

Эту идею развивают Bundle Neural Networks (BuNN), предложенные в работе Bamberger et al. В BuNN каждой вершине сопоставляется локальное векторное пространство, а переход между локальными пространствами соседних вершин задаётся ортогональным преобразованием. Для соседних вершин и его можно записать как $\begin{equation}O_{vu}=\mathcal{F}_{v \trianglelefteq e}^{\top}\mathcal{F}_{u \trianglelefteq e}\end{equation}$ , причём $\begin{equation}O_{vu}\in O(d),\qquad O_{vu}^{\top}O_{vu}=I\end{equation}$ . Тогда сообщение для вершины можно представить в виде $\begin{equation}m_v^t =\operatorname{AGGR}_{u\in\mathcal{N}(v)}\left(h_v^t-O_{vu}h_u^t\right)\end{equation}$ .

Рисунок 3. Локальная передача сообщений на графах в сравнении с глобальной диффузией сообщений на пучках.

В отличие от обычного локального message passing, в BuNN для распространения признаков используется диффузия. Это позволяет учитывать информацию не только от соседей вершины.

Диффузия в BuNN

Как и в случае sheaf diffusion, распространение признаков в BuNN можно рассматривать как диффузионный процесс. Диффузию признаков можно описать уравнением $\begin{equation}\frac{d\mathbf{u}}{dt}=-L\mathbf{u}\end{equation}$ , где $L$ — оператор Лапласа. Для постоянного решение имеет вид $\begin{equation}\mathbf{u}(t)=e^{-tL}\mathbf{u}(0)\end{equation}$ .

Чтобы реализовать диффузию, нужно вычислить действие матричной экспоненты $e^{-tL}$ на начальный вектор признаков.

Рисунок 4. Пример реализации механизма диффузии сообщений на графе с четырьмя вершинами и четырьмя рёбрами.

Обзор существующих методов

Существующие sheaf-based методы в основном используют sheaf Laplacian как основу для дискретного распространения признаков. В Neural Sheaf Diffusion sheaf обучается вместе с моделью и помогает формировать диффузию на гетерофильных графах. В Sheaf Neural Networks with Connection Laplacians restriction maps вычисляются заранее на основе геометрического выравнивания локальных пространств, что позволяет снизить вычислительные затраты. Другие работы рассматривают совместные диффузионные процессы и оптимизацию спектральных свойств sheaf Laplacian.

Близкий подход предлагает работа Bundle Neural Networks for Message Diffusion on Graphs. Диффузию на векторных пучках после дискретизации можно интерпретировать как частный случай SNN.

Существующие sheaf-based подходы можно условно разделить на две группы:

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

Мы предлагаем другой способ решения диффузионного уравнения — с помощью Physics-Informed Neural Network (PINN).

Physics-Informed Neural Network

Вместо непосредственного численного решения диффузионного уравнения $\begin{equation}\dot{u}(t) = -Lu(t),\qquad u(0)=u_0\end{equation}$ , можно аппроксимировать его решение нейронной сетью $\begin{equation}\hat{u}(t) = f_{\theta}(t, X, u_0)\end{equation}$ , где $\theta$ — параметры сети, задаёт структуру оператора, а — начальное состояние.

Ключевая особенность PINN — сеть обучается не только по известному решению, но и с учётом самого дифференциального уравнения. Для этого вводится невязка $\begin{equation}r_{\theta}(t)=\frac{\partial \hat{u}(t)}{\partial t}+L\hat{u}(t)\end{equation}$ , которая должна быть близка к нулю.

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

В статье мы рассмотрим две реализации нейросетевой модели для классификации вершин графа. В первой PINN обучается совместно со всеми слоями модели под конкретную задачу. Во второй применяется подход, где используется предварительно обученная сеть (pre-trained PINN). Веса модели инициализируются результатами предварительного обучения, после чего замораживаются.

Предобученная PINN для действия матричной экспоненты

Рассмотрим подробнее, как устроена и обучается PINN, которая во втором варианте заменяет вычисление $e^{-tX}u_0$ внутри BuNN. Ниже — её упрощённое описание.

Какую задачу решает сеть

Сеть обучается не для одной конкретной матрицы, а для целого семейства задач Коши $\begin{equation} \dot u(t) = -A\,u(t), \qquad u(0)=u_0\end{equation}$ , где $A\in\mathbb{R}^{n\times n}$ — кососимметричная ( $A^{\top}=-A$ ) ленточная матрица с ненулевыми наддиагоналями. Точное решение $u(t)=e^{-tA}u_0$ .

Экспонента кососимметричной матрицы ортогональна, поэтому $\|u(t)\| = \|u_0\|$ при любом : решение изменяет направление вектора, но не его длину. Мы используем это свойство при обучении PINN: оно даёт устойчивую динамику и сохраняет норму решения.

Ортогональные матрицы, возникающие как экспоненты кососимметричных матриц при обучении PINN, — не то же самое, что ортогональные преобразования в BuNN. В BuNN такие преобразования используются при построении bundle Laplacian, который затем определяет диффузионный процесс.

Кроме того, для кососимметричной матрицы выполняется $\langle e, Ae\rangle = 0$ . Поэтому в энергетической оценке этот член не влияет на изменение нормы ошибки. Аналогичную оценку через невязку можно получить и для рассматриваемой системы.

Действие матричной экспоненты $e^{-tA}$ на начальный вектор можно вычислять разными численными методами. В эксперименте ниже базовый метод — приближение усечённым рядом Тейлора порядка 16. Идея PINN-подхода — перенести основную часть вычислений на этап предварительного обучения, а на инференсе получать приближение решения с помощью прямого прохода нейронной сети.

Вход модели

Плотная матрица содержит элементов, но независимых среди них только . Поэтому на вход подаются наддиагоналей $a^{(1)},\dots,a^{(k)}$ (дополненные нулями до длины ) вместе с начальным состоянием — как каналов одномерного сигнала:

$\begin{equation} x_{\text{matrix}} = \bigl[a^{(1)};\,\dots;\,a^{(k)};\,u_0\bigr] \in \mathbb{R}^{(k+1)\times n}\end{equation}$ .

Для трёхдиагональной матрицы из BuNN () это всего два канала.

Архитектура FiLM-PINN

Общая схема модели показана на рис. 5. Сеть состоит из двух потоков, которые обрабатывают время и матрицу по отдельности, и блоков слияния.

Временной поток. Обычный MLP плохо выучивает высокочастотные функции (эффект spectral bias), а на длинных интервалах решение сильно осциллирует. Поэтому скалярное время сначала отображается в признаки Фурье $\begin{equation} \gamma(t) = \bigl[\cos(2\pi B t),\ \sin(2\pi B t)\bigr], \qquad B_i \sim \mathcal{N}(0,\sigma^2)\end{equation}$ , а затем двухслойным MLP переводится в эмбеддинг $z_t\in\mathbb{R}^{d_h}$ .

Матричный поток. Вход $x_{\text{matrix}}$ проходит через свёртку $1\times1$ (поднятие до каналов) и несколько остаточных блоков из одномерных свёрток.

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

Слияние через FiLM. Самое важное — как время влияет на обработку матрицы. Вместо простой конкатенации признаков используется Feature-wise Linear Modulation: эмбеддинг времени генерирует для каждого блока поканальные коэффициенты масштаба $\gamma_l(z_t)$ и сдвига $\beta_l(z_t)$ :

$\begin{equation} h_{l+1} = h_l + \operatorname{Conv}_2\!\Bigl( \sigma\bigl((1+\gamma_l(z_t))\odot\operatorname{Conv}_1(\sigma(h_l)) + \beta_l(z_t)\bigr)\Bigr)\end{equation}$ .

Фактически временной энкодер работает как гиперсеть, которая перенастраивает свёрточные фильтры под конкретный момент времени.

Параметризация $1+\gamma_l$ делает модуляцию близкой к тождественной в начале обучения.

Рисунок 6. Остаточный блок с FiLM-модуляцией: коэффициенты и вычисляются по эмбеддингу времени . — Рисунок 6. Остаточный блок с FiLM-модуляцией: коэффициенты $\gamma_l$ и $\beta_l$ вычисляются по эмбеддингу времени .

Жёсткое начальное условие. Последняя свёртка $1\times1$ выдаёт вектор $\Delta_\theta\in\mathbb{R}^n$ — не само решение, а его приращение. Время нормируется на отрезок , причём $t_{\min}=-1$ соответствует , и предсказание строится как

$\begin{equation} \hat u_{\text{raw}} = u_0 + (t-t_{\min})\,\Delta_\theta(t, x_{\text{matrix}}), \qquad \hat u = \hat u_{\text{raw}}\, \frac{\|u_0\|}{\|\hat u_{\text{raw}}\|+\epsilon}\end{equation}$ .

Первое равенство гарантирует $\hat{u}(0)=u_0$ по построению (hard constraint), поэтому слагаемое $\mathcal{L}_{IC}$ в функции потерь не нужно. Нормировка явно сохраняет норму решения, как и положено динамике с кососимметричной матрицей .

Обучение

Схема обучения показана в нижней части рис. 5.

Данные. Эталонные решения для обучения не нужны. Элементы диагоналей и начальные состояния семплируются из $\mathcal{N}(0,1)$ , моменты времени — из $\mathcal{U}(0,T)$ . Точные значения $e^{-tA}u_0$ используются только на валидации.

Невязка. Производная по времени вычисляется не автоматическим дифференцированием, а центральной разностью с шагом $\varepsilon=10^{-3}$ :

$\begin{equation} R_\theta(t) = \frac{\hat u_\theta(t+\varepsilon)-\hat u_\theta(t-\varepsilon)}{2\varepsilon} + A\,\hat u_\theta(t)\end{equation}$ .

Это стоит двух дополнительных прямых проходов, но дешевле и стабильнее, чем дифференцирование по входу через модулированные слои. Аналогичный гибридный подход используется в Chiu et al. 2022.

Произведение $A\hat u_\theta$ вычисляется сдвигами по диагоналям без формирования плотной матрицы: $\begin{equation} (A u)_i = \sum_{j=1}^{k}\Bigl(a^{(j)}_i\,u_{i+j} - a^{(j)}_{i-j}\,u_{i-j}\Bigr)\end{equation}$ , где выходящие за границы индексы считаются нулевыми.

Причинное взвешивание. В эволюционных уравнениях PINN склонна подгонять невязку на поздних временах раньше, чем решение выучено на ранних. Из-за этого она сходится к неверному решению. Чтобы этого избежать, можно использовать causal weighting: батч сортируется по и делится на последовательных блоков с невязками , а итоговая потеря равна

$\begin{equation} \mathcal{L}_{\text{PDE}} = \frac{1}{C}\sum_{c=1}^{C} w_c\,\mathcal{L}_c, \qquad w_c = \exp\Bigl(-\varepsilon_{\text{causal}}\sum_{j<c}\mathcal{L}_j\Bigr)\end{equation}$ .

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

Оптимизация. Используется Adam с косинусным расписанием скорости обучения ( $2,5\cdot10^{6}$ итераций) и экспоненциальным скользящим средним весов (EMA, коэффициент 0,999), которое и применяется на инференсе. Размерность скрытого пространства .

Точность и обобщение

В табл. 1 есть сравнение FiLM с более простыми способами объединения временного и матричного потоков при и $t\in[0,1]$ : конкатенацией признаков (Baseline) и аддитивным добавлением эмбеддинга времени внутри остаточных блоков (ResNet).

Таблица 1. Относительная ошибка аппроксимации $e^{-tA}u_0$ на $t\in[0,1]$ , , среднее по пяти запускам.

Модель	Параметры	Rel. L_2, k=1	Rel.L_2, k=2
Baseline (конкатенация)	13,76 тыс.	$1,7\cdot10^{-3}$	$8,3\cdot10^{-2}$
ResNet (сложение)	29,25 тыс.	$1,6\cdot10^{-4}$	$5,0\cdot10^{-3}$
FiLM Small	8,54 тыс.	$8,3\cdot10^{-2}$	—
FiLM	32,42 тыс.	$1,2\cdot10^{-4}$	$4,6\cdot10^{-3}$

Даже у FiLM Small, где параметров меньше, чем у Baseline, ошибка почти на порядок ниже. Значит, выигрыш даёт способ учитывать время, а не размер модели. На коротком интервале предсказание FiLM-PINN визуально совпадает с точным решением (рис. 7).

Рисунок 7. Точное решение (сплошная линия) и предсказание FiLM-PINN (пунктир) на интервале для четырёх случайных пар , , . Показаны первые и последние три компоненты вектора . — Рисунок 7. Точное решение $e^{-tA}u_0$ (сплошная линия) и предсказание FiLM-PINN (пунктир) на интервале $t\in[0,1]$ для четырёх случайных пар , , . Показаны первые и последние три компоненты вектора .

Длинные интервалы. При решение сильно осциллирует, и обучать сеть на независимых случайных точках мало: решение расходится при . Лучше семплировать батч полными траекториями (много моментов для одной пары ). Тогда сеть в целом воспроизводит форму и фазу колебаний. Но к концу интервала накапливаются заметные отклонения по амплитуде и фазе (рис. 8).

Рисунок 8. Прямая аппроксимация на интервале : FiLM-PINN, обученная на полных траекториях при (, ). — Рисунок 8. Прямая аппроксимация на интервале $t\in[0,10]$ : FiLM-PINN, обученная на полных траекториях при (, ).

Рисунок 8. Прямая аппроксимация на интервале $t\in[0,10]$ : FiLM-PINN, обученная на полных траекториях при (, ).

Гораздо лучше работает другой приём: обучить сеть на коротком интервале и применять её авторегрессионно. Предсказание $\hat u(1)$ становится новым начальным условием для следующего отрезка. Это корректно благодаря свойству $e^{-(t+s)A}=e^{-sA}e^{-tA}$ . Относительная ошибка на при этом составляет 0,0009 против 0,37 при прямой аппроксимации. Авторегрессионная траектория практически совпадает с точным решением (рис. 9).

Рисунок 9. Сравнение прямой аппроксимации на (зелёный пунктир) и авторегрессионного применения модели, обученной на (красный пунктир), с точным решением (синяя линия).

Благодаря свёрточной архитектуре модель переносится на другие размерности без дообучения. Сеть, обученная при , на матрицах размера от до сохраняет относительную ошибку на уровне $1,2\text{--}1,25\cdot10^{-4}$ . Именно эта модель с (трёхдиагональная кососимметричная матрица) используется дальше в эксперименте с BuNN. После предобучения её веса замораживаются. PINN используется для аппроксимации действия матричной экспоненты $e^{-tX}$ на начальный вектор — в базовом варианте его вычисляют с помощью усечённого ряда Тейлора.

Обзор методов

Мотивация метода с представлением Крылова

Для линейного диффузионного уравнения $\begin{equation}\frac{\partial u(t)}{\partial t}=-Lu(t),\qquadu(0)=u_0\end{equation}$ , точное решение имеет вид $\begin{equation}u(t)=e^{-tL}u_0\end{equation}$ . Основная вычислительная задача связана с действием функции матрицы $e^{-tL}$ на начальное состояние. Чтобы эффективно представить такое действие, используют подпространства Крылова $\begin{equation}\mathcal{K}_K(L,u_0)=\operatorname{span}\left\{u_0,Lu_0,\ldots,L^{K-1}u_0\right\}\end{equation}$ . Krylov-методы широко применяют, чтобы приближать действия матричной экспоненты и решать эволюционные задачи. Они позволяют заменить работу с исходным оператором на операции в подпространстве, порождённом последовательным применением оператора к начальному состоянию (Saad 1992).

В нашей реализации эта идея — не отдельный численный решатель, а представление действия L для нейронной сети. В PINN передаётся последовательность $\begin{equation}u_0,\ L_{\mathrm{sheaf}}u_0,\ L_{\mathrm{sheaf}}^2u_0,\ldots,\ L_{\mathrm{sheaf}}^K u_0\end{equation}$ . Это позволяет сохранять разреженность оператора и не подавать в сеть его плотное представление.

Пусть — точное решение дискретного диффузионного уравнения

$\begin{equation}\frac{\partial u_n(t)}{\partial t}+L_nu_n(t)=0,\qquad u_n(0)=u_0\end{equation}$ , а $u_\theta(t)$ — его нейросетевое приближение. Определим невязку PINN:

$\begin{equation}r_\theta(t)=\frac{\partial u_\theta(t)}{\partial t}+L_nu_\theta(t)\end{equation}$ .

Обозначим ошибку приближения через $\begin{equation}e_\theta(t)=u_\theta(t)-u_n(t)\end{equation}$ . Тогда из уравнений для $u_\theta$ и следует

$\begin{equation}\frac{\partial e_\theta(t)}{\partial t}+L_ne_\theta(t)=r_\theta(t).\end{equation}$

Если — симметричный положительно полуопределённый оператор, то, умножая последнее уравнение скалярно на $e_\theta(t)$ , получаем оценку

$\begin{equation}\frac{1}{2}\frac{d}{dt}|e_\theta(t)|^2+\langle e_\theta(t),L_ne_\theta(t)\rangle=\langle e_\theta(t),r_\theta(t)\rangle.\end{equation}$

Поскольку для симметричного положительно полуопределённого выполняется $\begin{equation}\langle e_\theta,L_ne_\theta\rangle\geq0,\end{equation}$ верна оценка $\begin{equation}|e_\theta(t)|\leq|e_\theta(0)|+\int_0^t|r_\theta(s)|,ds.\end{equation}$

Точность решения PINN контролируют две величины: ошибка начального условия и невязка дифференциального уравнения. Поэтому минимизация $\begin{equation}\mathcal{L}_{\mathrm{PDE}}=\frac{1}{T}\int_0^T|r_\theta(t)|^2,dt\end{equation}$ совместно с $\begin{equation}\mathcal{L}_{\mathrm{IC}}=|u_\theta(0)-u_0|^2\end{equation}$ напрямую уменьшает ошибку приближения решения диффузионного уравнения.

Рисунок 10. Архитектура Krylov-based модели.

PINN для аппроксимации экспоненты с использованием представления Крылова

Модель объединяет обучаемое sheaf-представление графа, предложенное в работе Y. Choi, J. Choi, T. Ko, J. Kim и C.-K. Kim «Sheaf Graph Neural Networks via PAC-Bayes Spectral Optimization», с физически информированной нейронной сетью (PINN), аппроксимирующей решение уравнения диффузии.

В работе Choi et al. (2026) диффузионная ветвь реализована на основе слоя $\textit{SVR+AFM}$ , который совмещает диффузию со стохастическим уменьшением дисперсии (Stochastic Variance Reduced, SVR) и адаптивное частотное смешивание (Adaptive Frequency Mixing, AFM). В нашей модели SVR-компонента уступает место PINN, а конфигурация обозначена как $\textit{PINN+AFM}$ .

$\begin{equation}\frac{\partial u_\theta(t)}{\partial t}+L_{\mathrm{sheaf}}u_\theta(t)= 0.\end{equation}$

Оператор $L_{\mathrm{sheaf}}$ строится из обучаемых sheaf-представлений, поэтому его параметры оптимизируются совместно с параметрами PINN. На вход PINN получает начальное состояние , оператор $L_{\mathrm{sheaf}}$ и время . Действие оператора передаётся через Krylov-представление $\begin{equation}u_0, L_{\mathrm{sheaf}}u_0,L_{\mathrm{sheaf}}^2u_0,\ldots,L_{\mathrm{sheaf}}^K u_0\end{equation}$ . Это позволяет работать с разреженным оператором без формирования его плотного представления.

Физическое ограничение задаётся через невязку $\begin{equation}R_\theta(t)=\frac{\partial u_\theta(t)}{\partial t}+L_{\mathrm{sheaf}}u_\theta(t),\end{equation}$ а начальное условие контролируется отдельно: $\begin{equation}\mathcal{L}_{\mathrm{IC}}=\operatorname{MSE}(u_\theta(0),u_0).\end{equation}$

Обучение включает одновременно задачу классификации и ограничения, задаваемые диффузионным уравнением: $\begin{equation}\mathcal{L}=\mathcal{L}_{\mathrm{cls}}+\lambda_{\mathrm{PDE}}\mathcal{L}_{\mathrm{PDE}}+\lambda_{\mathrm{IC}}\mathcal{L}_{\mathrm{IC}}.\end{equation}$

Результаты на тестовой выборке и время инференса. Значения представлены как среднее ± стандартное отклонение по пяти запускам с различными random seed. Метод SVR+AFM приведён из работы Choi et al. 2026. Модели мы отбирали по валидационной выборке.

Таблица 2. Время для этого метода мы замеряли на GPU NVIDIA GeForce RTX 4060, а вычисления запускали через CUDA. Поэтому характеристики времени в таблицах 2, 4 относятся именно к вычислениям на GPU, а не к CPU.

Dataset	Method	Test Accuracy (%)	Inference Time (ms)
Cora	PINN+AFM	81,42 ± 1,25	3,76 ± 0,02
	SVR+AFM	81,26 ± 1,27	6,17 ± 0,05
Citeseer	PINN+AFM	71,30 ± 0,88	4,14 ± 0,02
	SVR+AFM	71,36 ± 0,83	6,51 ± 0,12
Pubmed	PINN+AFM	75,46 ± 2,38	8,16 ± 0,16
	SVR+AFM	77,36 ± 0,50	11,61 ± 0,02
Chameleon	PINN+AFM	49,34 ± 1,28	3,43 ± 0,31
	SVR+AFM	50,31 ± 1,43	10,26 ± 0,18
Squirrel	PINN+AFM	34,49 ± 1,47	7,22 ± 0,11
	SVR+AFM	34,70 ± 0,87	15,96 ± 0,46
Actor	PINN+AFM	37,42 ± 0,98	2,95 ± 0,06
	SVR+AFM	36,39 ± 0,49	9,85 ± 0,22
Cornell	PINN+AFM	76,76 ± 2,42	2,64 ± 0,40
	SVR+AFM	78,38 ± 1,91	4,30 ± 0,64
Texas	PINN+AFM	79,46 ± 1,48	2,65 ± 0,14
	SVR+AFM	82,16 ± 4,10	3,11 ± 0,14
Wisconsin	PINN+AFM	80,39 ± 1,39	2,61 ± 0,04
	SVR+AFM	80,78 ± 1,64	5,05 ± 0,70

PINN для аппроксимации матричной экспоненты в BuNN

Кроме аппроксимации sheaf diffusion через Krylov-представление, мы рассмотрели отдельный вариант использования PINN внутри Bundle Neural Networks. Здесь мы берём заранее обученную PINN, которая аппроксимирует действие матричной экспоненты: $f_{\theta}(t,X,u_0)\approx e^{-tX}u_0$ .

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

В качестве baseline мы взяли приближение экспоненты усечённым рядом Тейлора порядка 16:

$e^A\approx I+A+\frac{A^2}{2!}+\ldots+\frac{A^{16}}{16!}$ .

В эксперименте мы сравниваем два способа вычисления действия одной и той же матричной экспоненты на вектор:

аппроксимация усечённым рядом Тейлора порядка 16;
предварительно обученная PINN.

Низкоразмерная параметризация

В этой реализации оператор задаёт кососимметричная трёхдиагональная матрица

$X =\begin{pmatrix}0 & a_1 & 0 & \cdots & 0 \\-a_1 & 0 & a_2 & \ddots & \vdots \\0 & -a_2 & 0 & \ddots & 0 \\\vdots & \ddots & \ddots & \ddots & a_{n-1} \\0 & \cdots & 0 & -a_{n-1} & 0\end{pmatrix}$ .

Поэтому вместо передачи всех элементов матрицы достаточно задать вектор $a=(a_1,\ldots,a_{n-1})$ .

В наших экспериментах , поэтому матрицу определяют 63 параметра.

Эксперимент с BuNN

Все эксперименты проводились на CPU MacBook Air с процессором Apple M4 без использования GPU. Поэтому время работы ниже соответствует вычислениям на CPU.

Перед запуском экспериментов на графовых датасетах мы отдельно проверили, насколько хорошо PINN воспроизводит действие матричной экспоненты, полученное с помощью Taylor-16.

При размере батча 2708 время вычисления составило 263,19 мс для PINN и 438,18 мс для Taylor-16, что соответствует ускорению в 1,66 раза.

При этом относительная ошибка составила undefined $\mathrm{CosineSimilarity}=0,999999$ , а cosine similarity между результатами двух методов: $\mathrm{CosineSimilarity}=0,999999$ .

В этом тесте PINN практически воспроизводит результат Taylor-16, но требует меньше времени на вычисление.

Таблица 3. Сравнение PINN и Taylor-16 внутри BuNN.

Dataset	PINN acc.	Taylor acc.	PINN, ms	Taylor, ms	Speedup
Cora	0,455	0,312	360,5	484,5	1,34
CiteSeer	0,388	0,192	400,1	595,1	1,49
PubMed	0,627	0,632	3361,7	5357,8	1,59

Результаты и выводы

Мы рассмотрели подход к моделированию sheaf diffusion с помощью Physics-Informed Neural Network. В отличие от дискретного применения диффузионного оператора, наша модель аппроксимирует зависящее от времени решение диффузионного уравнения. Корректность этого решения контролируют невязка уравнения и начальное условие.

Важная часть метода — использование Krylov-представления, в котором PINN получает последовательность применений sheaf Laplacian к начальному состоянию. Это позволяет передавать информацию о действии разреженного оператора без формирования его плотного представления.

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

Зато измерения времени инференса показывают более устойчивое различие между двумя подходами. Диффузионная ветка PINN работает быстрее соответствующей SVR-ветки на всех рассмотренных датасетах.

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

Также мы рассмотрели вариант, где предварительно обученная PINN аппроксимирует действие матричной экспоненты внутри Bundle Neural Networks. Отдельный тест показал, что результаты PINN и Taylor-16 близки: относительная ошибка составила 0,001261, а cosine Similarity — 0,999999.

Внутри BuNN PINN оказался быстрее Taylor-16 на всех трёх рассмотренных датасетах. Полученное ускорение составило от 1,34 до 1,59 раза. При этом accuracy на всех трёх датасетах осталась сопоставимой либо увеличилась.

Две серии экспериментов показывают, что PINN подходит не только для прямой аппроксимации sheaf diffusion. Она может аппроксимировать действия матричной экспоненты в рассмотренной реализации Bundle Neural Networks.

Перспективное направление для дальнейшей работы — исследовать влияние размерности Krylov-представления, весов $\lambda_{\mathrm{PDE}}$ и $\lambda_{\mathrm{IC}}$ , а также параметров PINN на точность и вычислительную стоимость модели.

Таблица 4. Время инференса диффузионной ветки для PINN и SVR. Значения представлены как среднее $\pm$ стандартное отклонение по пяти запускам и указаны в миллисекундах

Dataset	PINN Diffusion (ms)	SVR Diffusion (ms)
Cora	0,671 ± 0,006	3,152 ± 0,032
Citeseer	0,684 ± 0,006	3,098 ± 0,042
Pubmed	1,674 ± 0,046	5,225 ± 0,021
Chameleon	0,567 ± 0,088	7,558 ± 0,126
Squirrel	0,595 ± 0,014	9,354 ± 0,317
Actor	0,795 ± 0,013	7,764 ± 0,153
Cornell	0,633 ± 0,092	2,497 ± 0,644
Texas	0,631 ± 0,050	1,320 ± 0,135
Wisconsin	0,615 ± 0,020	3,232 ± 0,708

View the original on Хабр →

KioskNews shows a cleaned-up reading view extracted from the publisher’s page — the original always lives on their site, not ours.