[論文レビュー] Multi-Grid Tensorized Fourier Neural Operator for High-Resolution PDEs

Kossaifi, Jean, Kovachki, Nikola+2 more|arXiv (Cornell University)|Jan 1, 2021

Model Reduction and Neural Networks被引用数 10

ひとこと要約

本稿では、マルチグリッドドメイン分割とフーリエドメインにおけるテンソル化パラメータ化を組み合わせることで、スケーラブルでメモリ効率の良い高分解能PDE用のニューラルオペレータ、Multi-Grid Tensorized Fourier Neural Operator (MG-TFNO) を提案する。入力ドメインを階層的なグローバルコンテキストを持つ局所的領域に分割し、重みを低ランクテンソル因子分解（例：Tucker）によって表現することで、標準FNOと比較して150倍以上のパラメータ圧縮と7倍の入力サイズ削減を達成し、乱流ナビエ-ストークス方程式において誤差の半分未満を実現する。

ABSTRACT

Memory complexity and data scarcity have so far prohibited learning solution operators of partial differential equations (PDEs) at high resolutions. We address these limitations by introducing a new data efficient and highly parallelizable operator learning approach with reduced memory requirement and better generalization, called multi-grid tensorized neural operator (MG-TFNO). MG-TFNO scales to large resolutions by leveraging local and global structures of full-scale, real-world phenomena, through a decomposition of both the input domain and the operator's parameter space. Our contributions are threefold: i) we enable parallelization over input samples with a novel multi-grid-based domain decomposition, ii) we represent the parameters of the model in a high-order latent subspace of the Fourier domain, through a global tensor factorization, resulting in an extreme reduction in the number of parameters and improved generalization, and iii) we propose architectural improvements to the backbone FNO. Our approach can be used in any operator learning setting. We demonstrate superior performance on the turbulent Navier-Stokes equations where we achieve less than half the error with over 150x compression. The tensorization combined with the domain decomposition, yields over 150x reduction in the number of parameters and 7x reduction in the domain size without losses in accuracy, while slightly enabling parallelism.

研究の動機と目的

• 高分解能PDEのためのオペレータ学習を制限する高いメモリ複雑性とデータ不足に対処する。
• 物理的およびパrameter空間における局所的・グローバル構造を活用することで、フーリエニューラルオペレータ（FNO）における次元の呪いとメモリ増大を克服する。
• ドメインを局所的領域に分割し、マルチグリッド階層的コンテキストを備えることで、入力サンプル間の並列処理を可能にする。
• フーリエドメインにおける高次元テンソル因子分解を用いて、パラメータ数とメモリフットプリントを削減しながら精度を損なわない。
• 気象予測や流体動力学などの実世界の科学計算問題における汎化性とスケーラビリティを向上させる。

提案手法

• 入力ドメインを局所的領域に分割するマルチグリッドドメイン分割戦略を適用し、関心領域からの距離に応じて段階的にダウンサンプリングされたグローバルコンテキストで各領域を強化する。
• 局所的領域を高分解率で処理し、遠くの領域を粗い分解率で処理する階層的・マルチスケール表現を採用し、入力サイズを最大7倍まで削減する。
• FNOの畳み込み重みをフーリエドメインで低ランクテンソル形式（例：Tucker, CP, TT）に表現するテンソル化フーリエオペレータ（TFNO）を導入し、パラメータを100倍以上圧縮する。
• 中程度の圧縮比で性能と圧縮のバランスを取るためにTucker分解を活用し、高圧縮を実現するCPを用い、極端な圧縮と安定性を確保するTTを用いる。
• マルチグリッド分割とテンソル化重みを統合した包括的アーキテクチャ（MG-TFNO）を構築し、入力領域間の完全な並列処理を可能にする。
• 局所的予測のステッチング後に高品質な再構成を保証するため、出力の滑らかさを促進する損失関数を用いる。

Figure 1: Comparison of the performance on the relative $L^{2}$ and $H^{1}$ test errors (lower is better) on a log-scale of our approach, compared with both our improved backbone ( FNO ) and the original FNO, on Navier-Stokes. Our approach enables large compression for both input and parameter, whil

実験結果

リサーチクエスチョン

• RQ1 マルチグリッドドメイン分割は、高分解能PDEにおける入力サイズとメモリフットプリントを削減しながら、精度を維持または向上させることができるか？
• RQ2 フーリエドメインにおけるFNO重みの高次元テンソル因子分解は、性能劣化を伴わずに極端なパラメータ圧縮を達成できるか？
• RQ3 ドメイン分割とテンソル化の併用効果は、一般化性、推論速度、メモリ効率にどのように影響を与えるか？
• RQ4 MG-TFNOは、乱流ナビエ-ストークスや気象予測などの実世界の高分解能科学的問題にどの程度スケーリング可能か？
• RQ5 テンソル化パラメータ化は正則化として機能し、限られたトレーニングデータでも一般化性能を向上させるか？

主な発見

• MG-TFNOはナビエ-ストークス方程式において150倍のモデルパラメータ削減と7倍の入力ドメインサイズ削減を達成し、精度に損失なしに実現した。
• ナビエ-ストークスベンチマークにおいて、MG-TFNOは相対$L^2$テスト誤差を0.85%まで低下させ、標準FNOの2.54%と比較して半分未満となった。
• Tuckerに基づくTFNOバージョンは10倍のモデル圧縮を達成し、550万パラメータで0.85%の誤差率を維持した。
• CP分解は482倍の圧縮を可能としたが、Tuckerと比較して誤差が高くなった（2.24%）ことから、圧縮と精度のトレードオフが明らかになった。
• TT分解は574倍の圧縮を達成し、11.7万パラメータと1.18%の誤差を実現し、極端な圧縮比においても安定性を示した。
• アブレーションスタディの結果、マルチグリッドドメイン分割とテンソル化重みの併用は、単独で用いる場合よりも優れた性能とメモリ節約効果を示した。

Figure 2: Overview of our approach . First (left), a multi-grid approach is used to create coarse to fine inputs that capture high-resolution details in a local region while still encoding global context. The resulting regions are fed to a tensorized Fourier operator (middle), the parameters of whic