Neural Operators

ニューラルオペレーターとは、固定サイズのベクトル間ではなく、関数間の写像を学習するニューラルネットワークです。入力は、初期速度場、岩石層の浸透率マップ、今日の大気などの関数で、出力は数秒後の流れ、結果として生じる圧力場、明日の天気などの別の関数です。学習したオペレーターは連続領域上で定義されるため、粗いグリッド、細かいグリッド、不規則なメッシュでサンプリングした入力を、同じ重みで処理できます。数値ソルバーの解で学習すると、ニューラルオペレーターは偏微分方程式（PDE）群のシミュレーションを繰り返す代わりに、1回の高速な順伝播で処理します。

オペレーターと一般的なサロゲートの違い#

64 x 64のシミュレーションでサロゲートとして学習した畳み込みニューラルネットワーク（CNN）は、そのピクセル間隔に結び付いたカーネルを学習します。一方、ニューラルオペレーター層は積分オペレーターを近似し、領域全体にわたって積分した学習済みカーネルと点ごとの線形項を用いて各点を更新します。グリッドは積分をどの細かさで評価するかだけを決めます。グリッドを細分化すると予測が収束する性質は離散化不変性と呼ばれます。そのため、低コストの粗いシミュレーションで学習したモデルをより高解像度で実行できます。Fourier Neural Operatorの論文では、これをゼロショット超解像と呼んでいます。

新しい解像度で同じ重みを使うことは、ビジョンモデルでもおなじみです。Ultralytics YOLO26は、予測モードであらゆるimgszを受け付け、最大ストライドの32ピクセルの倍数に切り上げます。これは畳み込みカーネルが画像サイズに依存しないためです。ただし、各3 x 3カーネルが対象とするピクセル数は固定なので、imgszを2倍にすると、各カーネルが捉えるシーンの幅は半分になります。

ニューラルオペレーターと混同されやすい関連手法が2つあります。物理情報ニューラルネットワーク（PINN）は、係数と初期条件が固定された単一のPDEインスタンスを満たすように1つのネットワークを学習させるため、条件が変わると再学習が必要です。Neural ODEは、固定長の状態ベクトルの時間微分を学習し、ソルバーがその値を積分します。

Fourier Neural Operator#

オペレーター学習は、Lu Lu、George Karniadakisらが2019年に公開したDeepONetと、Zongyi Li、Nikola Kovachki、Anima Anandkumar、Andrew StuartらがCaltechで開発したニューラルオペレーターフレームワークを機に発展しました。ICLR 2021で発表された彼らのFourier Neural Operator（FNO）は、周波数空間での畳み込みとしてカーネル積分を計算します。高速フーリエ変換（FFT）で各チャネルをフーリエモードに変換し、学習可能な複素行列で低次モードのチャネルを混合して残りを除外し、逆FFTでグリッドに戻します。1つの層で出力の各点が入力全体に依存するため、グローバルな受容野を持ちます。

スペクトル経路は数行のPyTorchで実装できます。ここでは、同一の重みセットで、64点と256点でサンプリングした同じ4つの入力関数を処理します。

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個の複素重みは変化しません。層がグリッド点ではなく基礎となる関数のフーリエモードに作用するため、同じ物理座標での出力は小数点以下3桁まで一致します。FNOの論文では、乱流状態のNavier-Stokes流、Burgers方程式、Darcy流において、従来のPDEソルバーより最大3桁高速な予測が報告されました。

その他のアーキテクチャ#

2021年にNature Machine Intelligenceで発表されたDeepONetは、Tianping ChenとHong Chenによる1995年のオペレーターの普遍近似定理を基礎としています。ブランチネットワークは固定されたセンサー位置での値から入力関数をエンコードし、トランクネットワークは出力を求める座標をエンコードします。その内積が座標における出力値になります。グラフニューラルオペレーターは、複雑な形状周辺の不規則なメッシュを扱うため、グラフニューラルネットワーク（GNN）のメッセージパッシングを取り入れています。オペレーターTransformerは、領域内でサンプリングした点やパッチをトークンとして扱います。2024年にETH Zurichが発表したPoseidonは、流体力学で事前学習され、PDE向けの基盤モデルとして、事前学習データに含まれていない物理分野を含む15の下流タスクでファインチューニングされました。オープンソースのneuraloperatorライブラリはFNOと多くの派生手法を実装しています。

用途#

学習データは従来のソルバーから得られるため、各学習ペアは合成データであり、生成コストは1回だけ発生します。学習済みのオペレーターは多くのクエリに高速で応答できます。たとえば、翼型や原子炉の形状に関する設計の比較検討、数千回の実行が必要な不確実性定量化の研究、完全なシミュレーションより高速な更新が求められる機械やプラントのデジタルツインなどです。物理情報ニューラルオペレーターは、損失関数にPDE残差を加え、学習データが少ない場合でも予測が支配方程式と整合するようにします。

天気予報では、NVIDIAと学術機関の共同研究者が2022年に開発したFourCastNetが、適応型Fourierニューラルオペレーターを採用しています。0.25°の解像度で全球予報を生成し、1週間分の予報を2秒未満で作成します。短期予報の大規模変数についてECMWF統合予報システムと同等の精度を実現しながら、桁違いに高速です。