IKNO для Навье — Стокса точен при поворотах и в 5,6 раза меньше TFN

Нейросетевой суррогат для уравнений в частных производных может хорошо ложиться на обучающие данные и всё равно выдавать разные ответы для одного и того же физического состояния, если повернуть систему координат. Новая работа на arXiv предлагает архитектуру, у которой этой проблемы нет по построению. Invariant-Conditioned Isotropic Kernel Neural Operator, или IKNO, прогнозирует трёхмерную динамику Навье — Стокса на нерегулярно расположенных точках. Согласованность при поворотах держится до численной точности. По точности модель сравнима с общей вращательно-осведомлённой Tensor Field Network, а параметров у неё в 5,6 раза меньше.

В чём ломаются обычные суррогаты

Задача выглядит так. Есть облако точек в 3D, в каждой записана скорость. Модель должна предсказать, как скорость изменится на следующем шаге. Физика изотропна: если повернуть и положения, и скорости, то изменение скорости обязано повернуться так же. Это эквивариантность относительно группы SO(3).

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

Как устроен IKNO

Идея в разделении ролей. Локальные взаимодействия между соседними точками модель строит из двух типов величин:

• скаляры, не меняющиеся при повороте (например, расстояния и скалярные произведения);
• направления-векторы, которые поворачиваются вместе с данными.

Скаляры определяют силу взаимодействия, а векторы задают, куда оно направлено. Ядро изотропно, так что выход по построению поворачивается вместе со входом. Это свойство заложено в архитектуру, а не выучивается из данных.

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

Что показали измерения

Оценка шла на отложенной выборке, которую зафиксировали после того, как дизайн модели был завершён. Это снижает риск подгонки архитектуры под тест.

• Согласованность при поворотах у IKNO: до численной точности.
• Неограниченные графовые модели, обученные с аугментацией поворотами: зависимость снижается, но остаётся.
• Точность прогноза: на уровне Tensor Field Network.
• Размер: в 5,6 раза меньше параметров, чем у TFN.
• Против графового симулятора с тем же числом параметров: IKNO выигрывает.

Последний пункт важен для честного сравнения. Выигрыш получен не за счёт размера модели, а за счёт вшитой симметрии.

Где это пригодится и где границы

Для тех, кто строит суррогаты в CFD, метеорологии, моделировании частиц или робототехнике, вывод практический. Если у задачи есть известная симметрия, её выгоднее закладывать в архитектуру, чем надеяться, что аугментация её выучит. Гарантия точной эквивариантности также упрощает валидацию: не нужно прогонять модель по сетке ориентаций, чтобы проверить, что ответ не плывёт.

Границы тоже видны из абстракта. Результаты получены на одной семье задач, трёхмерной динамике Навье — Стокса. Заявление о компактных PDE-специализированных моделях пока подтверждено на ней. Как архитектура поведёт себя с граничными условиями, анизотропными средами или другими уравнениями, из описания сказать нельзя. Там, где физика сама нарушает вращательную симметрию (например, из-за гравитации или стенок), чистая SO(3)-эквивариантность потребует доработки.

Общий вектор понятен. Для научных суррогатов специализированная и маленькая модель с правильной симметрией может заменить универсальную и тяжёлую без потери качества.