Neural Operators

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

Операторы и обычные суррогатные модели#

Свёрточная нейронная сеть (CNN), обученная как суррогат на симуляциях 64 x 64, использует ядра, привязанные к этому шагу пиксельной сетки. Слой нейронного оператора вместо этого аппроксимирует интегральный оператор: обновляет каждую точку на основе интеграла обученного ядра по всей области и поточечного линейного слагаемого, а сетка лишь определяет точность вычисления интеграла. По мере измельчения сетки предсказания сходятся — это свойство называют инвариантностью к дискретизации. Благодаря ему модель, обученную на дешёвых симуляциях с грубой сеткой, можно запрашивать с более высоким разрешением; в статье о Fourier Neural Operator этот подход назван сверхразрешением без дообучения.

Работа одних и тех же весов с новым разрешением знакома по моделям компьютерного зрения. Ultralytics YOLO26 принимает любые imgsz в режиме predict, округляя их вверх до кратного максимальному шагу в 32 пикселя: свёрточные ядра не зависят от размера изображения, но каждое ядро 3 x 3 охватывает фиксированное число пикселей, поэтому при удвоении imgsz ширина сцены, охватываемой каждым ядром, уменьшается вдвое.

Нейронные операторы часто путают с двумя родственными методами. Физически информированная нейронная сеть (PINN) обучается удовлетворять одному экземпляру PDE с фиксированными коэффициентами и начальными условиями; при их изменении её нужно обучать заново. Нейронное ОДУ обучается вычислять производную по времени для вектора состояния фиксированной длины, а затем решатель интегрирует её.

Операторы Фурье#

Обучение операторов получило развитие благодаря DeepONet, опубликованной Лу Лу, Джорджем Караниадисом и коллегами в 2019 году, и системе нейронных операторов Зонъи Ли, Николы Ковачки, Анимы Анандкумар, Эндрю Стюарта и коллег из Caltech. Их Fourier Neural Operator (FNO), опубликованный на ICLR 2021, вычисляет интеграл ядра как свёртку в частотной области: быстрое преобразование Фурье (FFT) преобразует каждый канал в моды Фурье, обучаемые комплексные матрицы смешивают каналы на низших модах и отбрасывают остальные, а обратное FFT возвращает данные на сетку. Каждая выходная точка зависит от всего входа — один слой обеспечивает глобальное рецептивное поле.

Спектральный путь можно реализовать всего в нескольких строках PyTorch. Здесь одни и те же веса обрабатывают четыре входные функции, заданные с 64 и 256 точками:

import torch

torch.manual_seed(0)
channels, modes = 4, 8
# Learned complex weights act on the 8 lowest Fourier modes, whatever the grid size
weights = torch.randn(channels, channels, modes, dtype=torch.cfloat) / channels

for points in (64, 256):
    x = torch.arange(points) / points  # the same periodic domain [0, 1), sampled coarsely or finely
    field = torch.stack([torch.sin(2 * torch.pi * k * x) + x * (1 - x) for k in range(1, channels + 1)])
    spectrum = torch.fft.rfft(field)  # one row of Fourier coefficients per channel
    mixed = torch.zeros(channels, points // 2 + 1, dtype=torch.cfloat)
    mixed[:, :modes] = torch.einsum("im,iom->om", spectrum[:, :modes], weights)
    output = torch.fft.irfft(mixed, n=points)  # back to the grid the input came from
    print(f"{points} points: output shape {tuple(output.shape)}, u(0.25) = {output[0, points // 4]:.4f}")

Вывод содержит тензор формы (4, 64) со значением u(0.25) = -1.0908 и тензор формы (4, 256) со значением u(0.25) = -1.0909. 128 комплексных весов не меняются, а выходы в одной и той же физической точке совпадают до третьего знака после запятой, поскольку слой работает с модами Фурье исходной функции, а не с точками сетки. В статье о FNO сообщается, что для турбулентного течения Навье — Стокса, уравнения Бюргерса и течения Дарси модель давала предсказания до трёх порядков величины быстрее традиционных решателей PDE.

Другие архитектуры#

DeepONet, опубликованная в Nature Machine Intelligence в 2021 году, опирается на теорему об универсальной аппроксимации операторов, сформулированную Тяньпином Ченом и Хуном Ченом в 1995 году: ветвь сети кодирует входную функцию по значениям в фиксированных точках измерения, стволовая часть кодирует координату, для которой нужен выход, а их скалярное произведение даёт значение выхода в этой точке. Графовые нейронные операторы используют передачу сообщений из графовых нейронных сетей (GNN), чтобы работать с нерегулярными сетками сложной геометрии. Операторные трансформеры рассматривают отсчёты или фрагменты области как токены. Обученная на гидродинамике Poseidon, созданная в ETH Zurich в 2024 году, служит фундаментальной моделью для PDE; её дообучили на 15 последующих задачах, включая физические явления, отсутствовавшие в обучающих данных. Библиотека neuraloperator с открытым исходным кодом реализует FNO и многие из этих вариантов.

Применение#

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

В прогнозировании погоды FourCastNet, созданная NVIDIA и академическими партнёрами в 2022 году, основана на адаптивном нейронном операторе Фурье и строит глобальные прогнозы с разрешением 0,25°, создавая прогноз на неделю менее чем за 2 секунды. Она на порядки быстрее ECMWF Integrated Forecasting System и при этом обеспечивает сопоставимую точность краткосрочного прогноза крупномасштабных переменных.