論文レビュー: DeepONet
概要及び要約
- リファレンス、数式の番号、表記法などはできるだけ論文に従う。
アクセシビリティのためにジャーナルに掲載されたバージョンではなく、アーカイブにアップロードされているバージョンを基準にレビューする。実験の部分で扱われている問題は少し異なるが、結局のところ、中心は実験結果や性能ではなく、DeepONet技法そのものの説明にある。
DeepONetは作用素を学習するために提案されたディープラーニング技法だ。作用素とは関数を関数へと対応付ける関数である（本文で詳しく説明する）。すなわち、ある関数$u$に対して、作用素$G$とは以下のようになる。
$$ G : u \mapsto G(u) $$
このとき、$u$も関数であり、$Gu = G(u)$も関数だ。「DeepONetは作用素を学習する」というのが第一の重要な点であり、第二に「$Gu$を級数として近似する」ということである。ある適切な関数空間$X$があり、その基底を$\left\{ \phi_{k} \right\}$としよう。このとき、$Gu \in X$は次のように表現される。
$$ Gu = \sum_{k=1}^{\infty} c_{k}\phi_{k} $$
DeepONetが学習するのは$c_{k}$と$\phi_{k}$であり、係数$c_{k}$を学習する部分をブランチネットワークbranch network、基底$\left\{ \phi_{k} \right\}$を学習する部分をトランクネットワークtrunk networkと呼ぶ。
実装
- PyTorchで実装する
- ジュリアで実装する
1 Introduction
普遍近似定理the universal approximation theoremはニューラルネットワークが任意の連続関数を近似できることを保証する。これは人工ニューラルネットワークとディープラーニング技法がうまく機能することの理論的基盤となり、これを基にさまざまな分野でディープラーニング技法が活躍している。しかし、これよりもさらに驚くべき結果は、人工ニューラルネットワークがすべての非線形汎関数functional及び（非線形）作用素operatorまで近似できるということだ。
数学に馴染みのない読者のために関数、汎関数、作用素について少し説明しよう。この3つは基本的には大きなくくりで言えば関数（定義域にある1つの元を共域にある正確に1つの元に対応させること）である。しかしあえて汎関数や作用素という言葉を使う文脈ではこれらは少し特別な意味を持つ。まず、特に説明がなければ関数は数（あるいはベクトル）を数（あるいはベクトル）に対応させることを指す。多項式、三角関数など一般的によく扱う関数がすべてこの文脈での関数である。
$$ \text{function}: \mathbb{R}^{n} \to \mathbb{R}^{m} $$
関数を数（スカラー）に対応させる関数を特に汎関数と呼ぶ。具体的には定積分がある。$I_{[a,b]}$という汎関数を$\displaystyle I_{[a,b]}(f) = \int_{a}^{b} f(x)dx$と定義すると、これは$f$が与えられるたびに区間$[a, b]$での$f$のグラフ下の面積を対応させる関数である。$X$を適切な関数空間とすると、汎関数は次のようになる。
$$ \text{functional}: X \to \mathbb{R} $$
作用素とは、関数を関数に対応させる関数を指す。例としては不定積分、微分などがある。
$$ \text{operator}: X \to X $$ 関数$f$に対して$D$という作用素を$D(f) = \dfrac{df}{dx}$と定義すると、これは与えられた関数をその関数の導関数に対応させる微分作用素になる。$I$という作用素を$\displaystyle I(f) = \int f(x) dx$と定義すると、これは与えられた関数をその関数の不定積分に対応させる作用素になる。
ここでいう関数、汎関数、作用素は上述の説明と同じである。本格的な話に入る前に、論文全体で通用される表記法を紹介する。$G$は変数が関数$u$である作用素を指す。
$$ G : u \mapsto G(u) $$
$G$が作用素なので、$G$の関数値である$G(u)$も1つの関数であり、その変数を$y$と表記する。
$$ G(u) : y \mapsto G(u)(y) $$
したがって、$y$と$G(u)(y)$はどちらも実数である。
$$ y, G(u)(y) \in \mathbb{R} $$
本論文では作用素を学習することが目標なので、以下のように$u$と$y$を入力として受け取り、$G(u)(y)$を出力するニューラルネットワークを考慮する。
$$ \text{network} : (u, y) \mapsto G(u)(y) $$
理論的には作用素$G$は$u$という関数自体を変数として持つが、コンピュータシミュレーションのためには離散化が必要であり、$u$の代わりに有限個の関数値$u(x_{1})$、$u(x_{2})$、$\dots$、$u(x_{m})$をニューラルネットワークの入力として扱う。このとき、$\left\{ x_{1}, x_{2}, \dots, x_{m} \right\}$を論文ではセンサーsensorsと呼ぶ。つまり提案するニューラルネットワークは以下のような構造を持つ（図1A）。
図1A
定理1（作用素に関する普遍近似定理） $\sigma$を連続する非多項式関数non-polynomial functionとする。$X$をバナッハ空間、$K_{1} \subset X$、$K_{2} \subset X$をコンパクト集合とする。$V \subset C(K_{1})$をコンパクト集合、$G : V \to C(K_{2})$を非線形連続作用素とする。
すると、任意の$\epsilon > 0$に対して、正の整数$n$、$p$、$m$と定数$c_{i}^{k}$、$\xi_{ij}^{k}$、$\theta_{i}^{k}$、$\zeta_{k} \in \mathbb{R}$、$w_{k} \in \mathbb{R}^{d}$、$x_{j} \in K_{1}$（$i = 1,\dots,n$、$k = 1,\dots,p$、$j = 1,\dots,m$）が存在し、次のことが成り立つ。
$$ \left| G(u)(y) - \sum\limits_{k=1}^{p} \underbrace{\sum\limits_{i=1}^{n} c_{i}^{k}\sigma\left( \sum\limits_{j=1}^{m} \xi_{ij}^{k}u(x_{j}) + \theta_{i}^{k} \right)}_{branch} \underbrace{\sigma(w_{k} \cdot y + \zeta_{k})}_{trunk} \right| < \epsilon \quad \text{for all } u \in V, y \in K_{2} \tag{1} $$
本論文では、上記の定理の近似を大きく2つの部分に分けてブランチbranchとトランクtrunkと呼ぶ。
上述の近似定理はニューラルネットワークが非線形作用素を学習できるだろうということを示唆するが、これを実際にどのように効果的に学習できるかについては示していない。実際に普遍近似定理によれば、任意の[MLP]が任意の連続関数を近似できるはずだが、画像に関連する作業ではCNNやその他のニューラルネットワーク構造がよりよく機能する。有用なネットワークは訓練しやすく、一般化性能generalization errorが良くなければならない。著者たちはこれが可能な新しい方法論を提案しようとしている。
提案する方法が非線形作用素を学習するのに適していることを示すために、データに対する制約constraintsを非常に緩やかに設定した。この条件は具体的には入力データの$u_{i}$が同じセンサーを持つべきであるということである。もちろん、センサーが均一な格子上にある必要はなく、変数$y$に対する制約はない。この条件は次の図（図1B）によく示されている。
図1B
著者たちは提案する構造をDeepONet（DeepOperatorNetwork）と呼び、これは入力関数（$u(x_{1}), \dots, u(x_{m})$）に対するブランチネットbranch netと出力関数の変数（$y$）に対する関数トランクネットtrunk netで構成されている。これについては第2章で詳しく説明する。
論文では2種類の作用素、常微分方程式(ODE)で表現される動的システムと偏微分方程式を考慮する。
2 Methodology
2.1 Deep operator networks (DeepONets)
著者たちは一般的な状況での作用素の学習に焦点を当て、これのための制約条件は唯一、入力関数（$u$）が同じセンサーを持つべきであるということである。提案するニューラルネットワークの入力値は大きく2つの部分に分かれ、上記の図1Aに示すように、$[u(x_{1}), \dots, u(x_{m})]$と$y$である。ニューラルネットワークの構造には制約がなく、論文では性能の優秀さを自慢するために最も簡単なニューラルネットワークである[fully-connected neural networks]（FNNs）を使用した。希望に応じて[CNN]、RNNなどの構造を適用でき、アテンションメカニズムを適用することもできると説明している。
まず、トランクネットワークは$y$を入力として受け取り、$[t_{1}, t_{2}, \dots, t_{p}]^{T} \in \mathbb{R}^{p}$を出力する。$p$個のブランチネットワークは$[u(x_{1}), \dots, u(x_{m})]$を入力として受け取り、それぞれが$b_{k} \in \mathbb{R}$を出力する（$k = 1,2,\dots,p$）。これらを数式$(1)$のように合わせると以下のようになる。
$$ G(u)(y) \approx \sum_{k=1}^{p} b_{k}t_{k} = \sum_{k=1}^{p} b_{k}([u(x_{1}), u(x_{2}), \cdots, u(x_{m})]) t_{k}(y) $$
注目すべきは、トランクネットの最後のレイヤーにも活性化関数が適用されている点だ。上記の数式ではあまり表れていないが、このようなアプローチは「$Gu = G(u)$という関数を級数で近似すること」と解釈できる。適切な関数空間$X$が与えられており、その基底を$\left\{ \phi_{k} \right\}$とする。このとき、$Gu \in X$に対して次のように表現できる。
$$ Gu = \sum_{k=1}^{\infty} b_{k}\phi_{k} $$
言い換えれば$t_{k} = \phi_{k}(y)$と考えると、級数の$t_{k}$は基底部分、$b_{k}$は級数の係数部分である。すなわち、DeepONetは「$Gu$を直接近似するのではなく、級数に分解して近似する」ということだ。Theorem 1では必要のない部分だが、以下のようにバイアス（定数項）を追加することで一般化性能が向上する。
$$ G(u)(y) \approx \sum_{k=1}^{p} b_{k}t_{k} + b_{0} $$
実際に実装する際には$p$は最低でも10個以上であり、$p$が増加するほど計算コストが増加する。そこで論文では、それぞれの$b_{k}$を学習するブランチネットワークを個別に持つ形式（図1C）であるStacked DeepONetと、1つのネットワークで全ての$b_{k}$を学習する形式（図1D）であるUnstacked DeepONetを紹介する。DeepONetに関するすべてのコードはhttps://github.com/lululxvi/deepxdeで確認できるが、著者の他の研究で提案されたすべてのコードが含まれており、希望する部分を見つけるのが少し難しい。
図1Cと図1D
2.2 Data generation
論文では2つの関数空間、ガウスランダム場Gaussian random field(GRF)と[直交多項式空間]を扱う。著者は平均が$0$であるGRFを使用した。
$$ u \sim \cal{G}(0, k_{l}(x_{1}, x_{2})) $$
ここで$k_{l}(x_{1}, x_{2}) = \exp (- \| x_{1} - x_{2} \|^{2} / 2l^{2})$は共分散カーネルcovariance kernelである。直交多項式空間にはチェビシェフ多項式を選んだ。$M > 0$であり、$T_{i}$を第1種チェビシェフ多項式とすると、
$$ V_{\text{poly}} = \left\{ \sum\limits_{i=0}^{N-1} a_{i} T_{i}(x): |a_{i}| \le M \right \} $$
データセットは$a_{i} \in [-M, M]$をランダムサンプリングして生成した。このように生成されたデータセットで、それぞれの$u$についてルンゲ・クッタ法でODEシステムを解き、有限差分法で2次PDEを解いてリファレンスソリューションを求めた。
3 Number of sensors for identifying nonlinear dynamic systems
このセクションでは、DeepONetで非線形力学システムを解く際に任意の精度$\varepsilon$を達成するためにどのくらいの数のセンサーが必要かを議論する。
4 Simulation results
このセクションでは、まず最も簡単な線形問題ですらFNNよりDeepONetが優れた性能を示すことを確認し、3つの非線形ODEおよびPDE問題に対する結果を示す。すべての問題でオプティマイザは学習率が$0.001$のAdamを使用し、特に明記されていない限りネットワークのサイズは以下の表に示されているものと同じである。
表1と表2
4.1 A simple 1D dynamic system
一次元力学システムは次のように表される。
$$ \begin{align*} \dfrac{ds(x)}{dx} &= g(s(x), u(x), x), \qquad x\in[0, 1] \\ s(0) &= 0 \end{align*} $$
この問題で目標は任意の$u$に対して、ソリューション$s(x) \text{ on } [0,1]$を見つけることである。
4.1.1 Linear case: $g(s(x), u(x), x) = u(x)$
まず非常に簡単な場合を考えてみる。
$$ \begin{align*} \dfrac{ds(x)}{dx} &= u(x), \qquad x\in[0, 1] \\ s(0) &= 0 \end{align*} $$
この場合、作用素$G : u \mapsto s$は次のような不定積分作用素である。
$$ G : u(x) \mapsto s(x) = \int_{0}^{x} u(\tau)d\tau $$
まず比較のために深度と幅を調整しながらFNNで$G$を学習するようにした。深度は増やしても性能に大きな影響を与えず、幅が増えると訓練誤差が減少するが、依然として一般化性能(テスト誤差)が改善されることはない（図2）。
図2
一方でDeepONetは訓練誤差とテスト誤差の差がほとんどない（図3A）。バイアス$b_{0}$を追加すると性能が若干向上する。また、Unstacked DeepONetがStacked DeepONetに比べて訓練誤差は大きいが、より重要なテスト誤差は低い。Unstacked DeepONetはパラメーター数が少ないため、より少ないメモリを使用してより高速に学習することができる。
図3
4.1.2 Nonlinear case: $g(s(x), u(x), x) = −s^{2}(x) + u(x)$
この場合では、Unstacked DeepONetとStacked DeepONetを比較することに焦点を当てた。訓練誤差とテスト誤差の相関を見ると、Unstacked DeepONetがより強いことがわかる（図4A）。さまざまな学習率や初期値に対してもUnstackedの方が強い相関を示した（図4B）。
図4
4.2 Gravity pendulum with an external force
このサブセクションでは次のような外力がある振り子運動を扱う。
$$ \begin{align*} \dfrac{ds_{1}}{dt} &= s_{2} \\ \dfrac{ds_{2}}{dt} &= -k \sin s_{1} + u(t) \\ s_{1}(0) &= 0, \quad s_{2}(0) = 0 \end{align*} $$
続く内容では、DeepONetがこの問題に対してうまく機能することを示し、センサーの数、誤差収束などについて扱う。
4.3 Diffusion-reaction system with a source term
次のような拡散-反応方程式を扱う。
$$ \dfrac{\partial s}{\partial t} = D \dfrac{\partial^{2} s}{\partial^{x}} + ks^{2} + u(x),\qquad x\in [0,1], t\in [0,1] $$ $$ \text{with zero initial/boundary conditions} $$
前の例と異なり、$u(x)$の変数は1次元だが、$s(x, t)$の変数は2次元である。この場合でもDeepONetがうまく機能することを示す。1つの$u$に対するトレーニングデータは以下のような式である。
$$ \begin{align*} \big( (u, (x_{1}, t_{1})), s(x_{1}, t_{1}) \big) \\ \big( (u, (x_{2}, t_{2})), s(x_{2}, t_{2}) \big) \\ \vdots \\ \big( (u, (x_{p}, t_{p})), s(x_{p}, t_{p}) \big) \end{align*} $$
$(u, (x_{i}, t_{i}))$はDeepONetの入力値であり、$s(x_{i}, t_{i})$は最終出力値である。具体的に、$u$はブランチ、$(x_{i}, t_{i})$はトランクの入力である。このように構成されたデータが異なる$u$ごとに生成され、学習に使用される。
5 Conclusion
本論文では非線形作用素を学習するためのDeepONetを提案した。DeepONetはブランチとトランクで構成されており、ブランチは係数を、トランクは基底を学習すると解釈できる。本論文には、テスト誤差に関するさまざまな要因（センサーの数、最大予測時間、入力関数空間の複雑さ、訓練データセットのサイズ、ネットワークのサイズ）の影響に関する分析が含まれている。また、f近似誤差がさまざまな要因にどのように影響を受けるかを理論的に導出し、結果が計算と一致することを示した。
しかし、DeepONetそのものに関する理論的な分析についてはまだ研究すべきことが多く残されている。また、論文ではFNNしか使用していないが、CNNやアテンションメカニズム、あるいはその他の異なるニューラルネットワーク構造／技法との連携についても研究の余地が残されている。