[論文レビュー] DPOT: Auto-Regressive Denoising Operator Transformer for Large-Scale PDE Pre-Training

Zhongkai Hao, Chang Su+7 more|arXiv (Cornell University)|Mar 6, 2024

Advanced Sensor and Control Systems被引用数 6

ひとこと要約

DPOTは、巨大で多様なデータセットからPDE解法演算子を学習するための、フーリエ注意機構を備えた自己回帰的ノイズ除去事前学習戦略を導入し、最先端の結果を達成するとともに、3Dデータを含む下流のPDEタスクへの強い転移を実現します。

ABSTRACT

Pre-training has been investigated to improve the efficiency and performance of training neural operators in data-scarce settings. However, it is largely in its infancy due to the inherent complexity and diversity, such as long trajectories, multiple scales and varying dimensions of partial differential equations (PDEs) data. In this paper, we present a new auto-regressive denoising pre-training strategy, which allows for more stable and efficient pre-training on PDE data and generalizes to various downstream tasks. Moreover, by designing a flexible and scalable model architecture based on Fourier attention, we can easily scale up the model for large-scale pre-training. We train our PDE foundation model with up to 0.5B parameters on 10+ PDE datasets with more than 100k trajectories. Extensive experiments show that we achieve SOTA on these benchmarks and validate the strong generalizability of our model to significantly enhance performance on diverse downstream PDE tasks like 3D data. Code is available at \url{https://github.com/thu-ml/DPOT}.

研究の動機と目的

• 多様なPDEタスクにおけるデータ効率を向上させるためのニューラルオペレータの事前学習を動機づける。
• PDEデータのための、スケーラブルな自己回帰型ノイズ除去事前学習戦略を提案する。
• PDE解法写像を効率的に学習するためのフーリエベースのトランスフォーマーバックボーンを設計する。
• 大規模なPDEファウンデーションモデルを構築し、3Dデータを含む下流タスクへの強力な一般化を示す。

提案手法

• 入力軌跡にガウスノイズを加えて次のタイムステップを予測することで、自己回帰型ノイズ除去事前学習を行う。
• データ前処理：データセット間で形状を統一するため（固定解像度へのパディング、チャネルパディング、幾何マスク）。
• トレーニング中のデータセット表現の公平性を確保するためのデータセットバランシング。
• 空間-時間PDE特徴を学習するためのパッチ化と時間的集約レイヤを備えたフーリエ注意バックボーン。
• 周波数ドメインで動作する共有MLPを備えたマルチヘッドフーリエミキサーが、カーネル風変換を実装。
• 理論的普遍性：フーリエ注意層はソボレフ空間間の連続作用素を近似できる。

実験結果

リサーチクエスチョン

• RQ1 自己回帰型ノイズ除去前訓練は、異種PDEファミリを横断して安定的で一般化可能なPDEファウンデーションモデルを可能にするか？
• RQ2 既存のニューラルオペレータと比較して、フーリエベースのトランスフォーマーバックボーンは大規模PDE事前学習でどのように機能するか？
• RQ3 事前学習とスケーリングは、3Dおよび高解像度データを含む多様な下流PDEタスクにおけるゼロショットおよびファインチューニング性能をどの程度向上させるか？

主な発見

• 自己回帰型ノイズ除去を用いた事前学習は、複数のPDEベンチマークで最先端の性能を達成し、特定のデータセットで最大52%の誤差削減を達成。
• DPOTモデルを0.5Bパラメータまでスケールアップすると、10+ PDEデータセットに対してゼロショットおよびファインチューニング性能が向上。
• フーリエ注意バックボーンは、PDEオペレータのスケーラブルで表現力のある学習を実現し、高次元で多様な下流タスク（例：3D Navier–Stokes、定常PDEs）への転移をサポートする。
• DPOTベースのファインチューニングは、ゼロからのトレーニングや一部の既存の事前学習ベースラインよりも下流パフォーマンスを一貫して改善する。
• このアプローチは、2D前訓練から3DPDE問題および高解像度データへの強い一般化を示す。