【Fourier Neural Operator】FNOを直感的に理解する

**本記事はFourier Neural Operator(FNO)の構造と設計意図について，特に「なぜFNOが必要とされるのか？」「従来のCNNやViTと何が違うのか？」という点を直感的に解説することを目的とします．**

コードや簡単な問題への実装は以下のリンクにまとめてありますので，もしよければご覧ください．
https://github.com/natsuki0720/neural_operator_comparison


# はじめに：なぜNeural Operatorのか？
近年，偏微分方程式（PDE）に基づく現象の高速予測を求められる場面が増えています．例えば，流体の時間発展や圧力分布の変化など，複雑な数値シミュレーションを一瞬で大体するような **サロゲートモデル** が必要とされています．  

こうした現象を深層学習で扱う場合，「テンソルを入力とし，テンソルを出力する」構造であるCNNやVision Transformer(ViT)を適用するのが自然なアプローチに思えます．PDEを離散化すれば，入力も出力も画像のようなテンソルになるためです．

しかし，ここには大きな落とし穴があります．

---
## PDEの本質は「関数の関数」：Operatorへ
PDEを通じてモデル化される現象は，実は「入力関数から出力関数への写像」であり，その背景にあるのは　**関数空間上の作用素（Operator）**　です．
例えば，以下のような写像関係を考えます．
$$
    \mathcal{G} : a(x) \to u(x)
$$
ここで， $a(x)$は入力関数（例：物質の分布や初期条件）， $u(x)$ はPDEの解として得られる出力関数です．  
このとき，学習すべき対象は有限次元のテンソル空間上の関数ではなく， **無限次元の関数空間上の写像＝作用素 $\mathcal{G}$** なのです．
CNNやViTといった手法は，離散化されたテンソルを入力とし，有限次元ベクトルとして出力を得ます：  
$$
    f_\theta : A \in \mathbb{R}^{n \times m} \to U \in \mathbb{R}^{n \times m}
$$
このアプローチでは，**解像度が変化した途端に再学習が必要となる** という根本的な限界を抱えています．なぜなら，関数そのものを扱っておらず，離散化された数値の羅列同士の関係を学習しているに過ぎないからです．

---
## Neural Operatorの着想
この問題を根本から解決するためのアプローチが **Neural Operator(NO)**です．

NOは関数空間上の作用素 $\mathcal{G}$を直接近似するニューラルネットワークであり，　**解像度非依存かつ構造的に関数を扱う** という性質を持っています．これにより，一度学習すれば異なるグリッドや解像度にも汎化可能なモデル構築が可能になります．

本記事では，Neural Operatorの中でも最も広く用いられているものの一つである **Fourier Neural Operator(FNO)** を取り上げ，
- どのように関数としての入力を処理するか
- CNNとの構造的な違いはどこか
などを順を追って解説していきます．

次章では，FNOが入力を関数として取り扱う方法を解説し，その後Fourier変換をベースとした演算処理について踏み込みます．

# FNOの処理構造
## 入力の関数として取り扱うとはどういうことか？
CNNやViTなどの深層学習モデルは，入力をテンソル（行列）として処理します．これは， **空間上の関数を離散的な点の集合として扱っている** に過ぎません．
例えば，連続な関数$a(x)$を有限個の点 $x_1,x_2, \cdots, x_n$でサンプリングし， $A = [a(x_1), a(x_2), \dots, a(x_n)]$ として扱うわけです．

このような処理では， **解像度（点の数）が変われば入力表現も変わってしまう** ため，同じ関数に対応するモデルの振る舞いが不安定になります．つまり， **入力を本質的に「関数」として扱っていない** のです．

## Neural Operatorにおける「関数のまま扱う」という発想
Neural Operatorの根本的な発想は， **関数$a(x)$をヒルベルト空間内の点として直接扱うこと**です． **ヒルベルト空間**は内積構造を持つ関数空間のことですが，詳細は後述します．

## フーリエ展開：関数を座標で表現する操作
では，どのようにして関数を「空間上の点」として表現するのでしょうか？
ここで登場するのが **Fourier展開** です．Fourier展開とは，関数を三角関数（または複素指数関数）の **正規直交規定** で展開し，重ね合わせによって表現する操作です．

1次元の例で言えば：
$$
    f(t) = \frac{a_0}{2} + \sum_{k=1}^\infty \{a_k \cos(k\omega t) + b_k \sin(k \omega t)\}
$$
ここで，$[\cos(k\omega t), \sin(k\omega t)]_{k=1}^\infty$はヒルベルト空間上の正規直交基底であり，$(a_0, a_1, b_1, a_2, b_2, \dots)$ は関数 $f(t)$ の**空間内の座標**に対応します．

>言い換えれば，Fourier展開とは「関数を無限次元ベクトルに変換する操作」なのです．

## FNOにおけるFourier展開の役割
FNOは，このFourier展開のアイディアを拡張し，**空間的な入力をFourier空間（周波数ドメイン）に変換し，スペクトル領域での演算** を行います．

これにより，
- 入力を「関数空間の点」として解釈できる（明示的に関数として扱える）
- 空間局所性に縛られず， **関数全体のグローバルな構造に基づいた演算が可能** （周波数特性を扱うため）
- 解像度に依存しない処理が可能（グリッド数が変わっても意味が変わらない）
といった，従来のCNNやViTにはない構造的な利点を獲得できる．

### 補足：Hilbert空間としての関数空間
- 関数空間$L^2([0,1])$ は以下のような内積が定義される空間で，正規直交基底としてフーリエ基底が存在します
$$
    \langle f,g\rangle = \int_0^1 f(x)g(x)dx
$$
- Fourier展開とは，この基底に対して関数を展開し，**無限次元空間における「座標表現」**を与える操作です
- これにより，関数$f(x)$を$L^2$空間の**点=無限次元ベクトル**として扱うことが可能になります

### 補足：実装上のFourier展開の扱い
Fourier展開は理論的には無限次元の基底展開ですが，FNOにおいては**有限個の周波数モードのみを取り出して処理**します．
つまり，厳密には関数空間全体への作用ではなく，**低次元の周波数空間への射影を通じた近似的な作用**となっています．

この処理は，「関数空間の全体構造を低次元空間に落として操作する」という意味で，**Hilbert空間からの低次元射影**と解釈できます．
よってFNOの枠組みで適切に処理できる関数は厳密な意味での「全ての関数」ではなく，**高周波成分が小さい，すなわち低周波成分で十分表現可能な滑らかな関数に限られる**という制約があります．


# FNOの全体構造
本章では，FNOの全体構造についての解説を行います．
![モデル構造](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/4052711/d4824e80-607c-48dc-8e8c-803dee705482.png)


## 位置符号化（Positional Encoding）について
FNOでは，入力関数$a(x)$をグリッド上に離散化する際に**座標成分をチャネルとして連結する**Positional Encodingを必須の前処理としています．  
2次元グリッド$(x_i, y_i)$の場合，元のスカラ値$a_{ij} = a(x_i,y_i)$を  
$$
    [a_{ij}, \frac{x_i}{L_x}, \frac{y_j}{L_y}]
$$
という3チャネル構成に拡張します．$(L_x, L_y)$は領域長で正規化係数です．
領域長を正規化係数として用いることで，FNOの解像度ロバスト性を維持したまま，位置符号化を行うことが可能になります．

### なぜ位置符号化が不可欠か？
1. 非周期境界の位相ずれ補正
    Fourier基底は周期関数を暗黙に仮定するため，非周期性のある境界では位相情報が失われます．座標を明示することで，ネットワーク自身が境界位置を条件として学習でき，誤差を大幅に減少します． 
1. 高周波モード打ち切りの補完
FNOは有限個の周波数モードしか保持しません．そのため，高周波成分は失われがちです．座標チャネルは局所畳み込み層を通じて高周波情報を補完するため，モード制限による入力の滑らかさに関する制約を緩和します．



## リフト
続いてはリフトと呼ばれる操作について解説します．  
FNOの入力は「スカラー＋座標」のわずか3チャネルですが，このままでは後続の処理が十分な表現力を獲得できません．

そこで，**カーネルサイズが $1 \times 1$の畳み込み層や線型層**を用いて，入力を複数チャネルに変換する以下の操作を行います．  
$$
    \mathbb{R}^{H \times W \times C_{in}} \to \mathbb{R}^{H \times W \times d_{lift}}
$$

* $1 \times 1$カーネルは空間情報を変えず，チャネル方向だけの線形変換となります

Liftにより，「スカラー＋座標」だった入力を**周波数領域で処理しやすい高次元テンソルへ変換**できました．  
次節ではこのテンソルを用いて，FNOの肝となるSpectral Convによって作用素近似を行う仕組みを解説します．  

## Spectral Convについて  
いよいよFNOの中核を担うSpectral Convの解説です． 


### 演算フロー
Spectral Convの処理フローの手順は以下のとおりです．
1.  **高速フーリエ変換（FFT）** により，各チャネルを周波数空間（スペクトル表現）へ写像
1.  **同一周波数モード$(k_x,k_y)$ごとに，チャネル方向の線形変換**を適用
1.  **逆フーリエ変換（iFFT）** により座標空間（元のテンソル表現に復元）

#### フーリエ変換
前節にて，「スカラー値＋座標成分」かならる入力を
$$
    a(x) \in \mathbb{R}^{H \times W \times 3} \to z(x) \in \mathbb{R}^{H \times W \times d_{\text{lift}}}
$$
へとリフトしました．
この$d_{\text{lift}}$個のチャネルそれぞれに対して，高速フーリエ変換（FFT）を適用します．

> 補足：FFTは入力サイズが$2^n$の時に最も効率良く動作します．一般サイズの時には，ライブラリでzero-paddingなどを通じて処理されますが，厳密には離散フーリエ変換（DFT）への切り替えも選択肢になります．ただし，計算量は大幅に悪化します．

FFTによって得られたフーリエ係数$\widehat{z}_k$は，フーリエ基底に対する展開係数であり，**元のチャネルを関数空間上の点とみなすための座標**と解釈できます．  
この表現を得ることで，離散テンソル表現だった$z(x)$を **関数空間の低次元射影（≒有限モードでの近似）**　として扱えるようになります．

---

#### 線形変換：モード毎の作用
得られたスペクトル表現に対して，**周波数モード$(k_x, k_y)$ごとに**学習可能な線形写像$W_k$を適用します：  

$$
    \widehat{z}_k \in \mathbb{C}^{d _{\text{lift}}} \to \tilde{z}_k  = W_k \widehat{z}_k \quad (W_k \in \mathbb{C}^{d _{\text{lift}} \times d _{\text{lift}}})
$$

これは単なるテンソルの数値的変換ではなく， **周波数基底で展開された関数に対する「空間全体にわたる作用」** として解釈できます．
各モードに個別の$W_k$を与えることで，**空間的に広がった（＝非局所な）構造に基づく演算** をモデルに内在させることが可能になります．  

---

#### iFFTによる復元と非線形変換の順序
最後に，線形変換された$\tilde{z}_k$に逆フーリエ変換（iFFT）を適用し，座標空間での出力を得ます．  
ここで重要なのは：
-  **非線形活性化関数(GELUなど)はiFFT後に適用される**
-  これは，**FFTによるスペクトル表現が複素数値であるため**，ReLUなどの標準的非線形関数で安全な変換が行えないためです

そのため，FNOでは**線形処理はスペクトル空間で，非線形処理は座標空間で**という役割分担が為されています．

## その後の処理
ここまでは，FNOの中核である**Spectral Conv**に焦点を当て，周波数空間上での処理がどのように作用素的な写像を実現するのかを見てきました．
本節では，その後に続く処理構造の流れを整理します．

###  **1.残差接続：Spectral Conv の補完としての直結**
FNOの処理は，単に周波数変換を行っているだけでなく，**スペクトルブロックの出力に元の入力を加える残差構造**を取っています：  
$$
    \text{output}_{\text{block}} = \text{SpectralConv}(f) + f
$$
これは**高周波モードの打ち切りによる情報損失や過剰な変換の抑制**を目的としており，ResNetと同様に学習安定性や性能向上に寄与します．

###  2.非線形変換
SpectralConvの変換後，出力テンソルに対してチャネル方向のMLP(全結合層)が適用されます．  
この操作は各空間点$(x_i, y_i)$ごとに独立に適用され，**チャネル間の非線形混合**を実現します．
SpectralConvが空間的（非局所）構造をとらえるのに対し，MLPは**チャネル次元における非線形表現の補強**を担っています．
通常は，これらのSpectralConvの残差接続や非線形変換を一つのブロックとして，複数回繰り返すことで，空間構造・周波数構造の両方を深くとらえることが可能となる．

### 3.出力層
最後に得られた最終表現に対し，**線形写像**を適用し，目的の出力関数$u(x)$を得ます． 
この出力$u(x)$は，PDEの解（例：圧力場や流速場）などの物理量を表現します．

> この線形変換によって，高次元に拡張していたチャネルを線形変換によって集約し，最終的なチャネル数に変換されます

# Darcy Flowへの適用例
さて，これまでFNOの理論や実装についての話を繰り返し行ってきましたが，ここでは実際にDarcyFlowの問題への適用例をダイジェストで紹介していきます．
詳細については，GitHubのリポジトリにまとめてあるので，興味のある方はそちらも併せてご覧ください．

## CNNとFNOの比較
行列$\to$行列の変換を考えた時にCNNの適用を考えるのはごく普通です．
そこでCNNとFNOの比較を行ったので以下に結果を示します．

### 推論時間・評価指標（順問題, 16x16解像度）

| Model     | Params    | Inference Time [ms]|     MSE ↓    　 |   PSNR ↑   |   SSIM ↑   |  HF_MSE ↓    |
|-----------|-----------|------------------------|----------------|------------|------------|--------------|
| CNN      | ~416K     | **0.76**                   | 3.03e-03 | 25.19  | 0.979 | 1.75e-01   |
| FNO       | ~204K     | 0.91                   | 1.79e-01     | 21.05      | 0.946     | 4.53e+01   |

### 推論時間・評価指標（順問題, 32x32解像度）

| Model     | Params    | Inference Time [ms]|     MSE ↓　　     |   PSNR ↑   |   SSIM ↑   |   HF_MSE ↓   |
|-----------|-----------|------------------------|----------------|------------|------------|----------------|
| CNN      | ~416K     | **1.60**                   | 2.26e-01     | 6.46       | 0.6209     | 1.06e+00 |
| FNO       | ~204K     | 1.98                   | 1.72e-01 | 19.80* | 0.9183 | 1.79e+02     |

## 位置符号化の影響の検証
以下に，位置符号化の有無によって，推論精度にどのような影響を与えるかを示したものです．
各評価指標の意味などについては，上述のリポジトリで解説しているので，そちらも併せて参照ください．

#### FNO構造比較：座標情報（coord）あり vs なし（順問題, 16x16解像度）

| Model           |     MSE ↓      |   PSNR ↑   |   SSIM ↑   |  HF_MSE ↓      |
|-----------------|----------------|------------|------------|----------------|
| FNO with coord  |  1.8002e-01    | **20.96**  | **0.9432** | 4.5290e+01     |
| FNO no coord    | **1.7464e-01** | 12.49      | 0.7131     | **4.4261e+01** |

合わせて，実際の推論結果のサンプルもいくつか示します．

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/4052711/f1de4a34-f706-4580-b5f3-5877e6006533.png)

## FNO構造比較：座標情報（coord）あり vs なし（順問題, 32x32解像度）

| Model          |     MSE ↓       |   PSNR ↑   |   SSIM ↑   |  HF_MSE ↓      |
|----------------|-----------------|------------|------------|----------------|
| FNO with coord | 1.7411e-01      | **19.47**  | **0.9122** | 1.7880e+02     |
| FNO no coord   | **1.7161e-01** | 12.54      | 0.7059     | **1.7685e+02** |

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/4052711/e5df9ff6-e59c-401a-aff9-7be15489f989.png)

### 位置符号化の影響について
Darcy Flowの透水圧分布の予測を例題として位置符号化の有無による差異を精度指標の平均値といくつかのサンプルを示しました．
$16 \times 16$と$32 \times 32$の二つの解像度についていくつかのサンプルを示しました．
位置情報を付与することで，境界部での精度が大きく向上していることがわかります．

# まとめ
本記事では，FNOの理論解説を中心に行いました．
かなり難解な内容になりましたが，理解いただければ幸いです．
もし内容に不備があれば，遠慮なくお申し付けください．
詳細は本記事最上部にリンクを載せたGitHubリポジトリにありますので，もしよろしければご覧ください．
本記事がNewral Operatorに関心を持つ方の読者の一助になれば幸いです．