Speeding up Fourier Neural Operators via Mixed Precision

要約

フーリエ ニューラル演算子 (FNO) は、偏微分方程式 (PDE) 解演算子の代理マップを学習するための強力な手法です。
高解像度のデータ ポイントを必要とする多くの実世界のアプリケーションでは、トレーニング時間とメモリ使用量が大きなボトルネックとなります。
標準的なニューラル ネットワーク用の混合精度トレーニング手法はありますが、それらは有限次元の実数値データ型に対して機能するため、(複素数値の) フーリエ領域および関数空間で重要に動作する FNO に直接適用することはできません。
一方、フーリエ変換はすでに近似値であるため (離散化誤差により)、完全な精度で演算を実行する必要はありません。
この研究では、(i) 完全精度および混合精度トレーニングを使用した FNO のメモリと実行時間をプロファイリングし、(ii) FNO の混合精度トレーニングの数値安定性に関する研究を実施し、(iii) 実質的なトレーニング ルーチンを考案します。
Navier-Stokes および Darcy の流れ方程式の精度をほとんどまたはまったく低下させることなく、トレーニング時間とメモリ使用量 (最大 34%) を削減します。
最近提案されたテンソル化 FNO (Kossaifi et al., 2023) と組み合わせると、結果として得られるモデルは、元の FNO よりも大幅に高速であると同時に、パフォーマンスがはるかに向上します。

要約(オリジナル)

The Fourier neural operator (FNO) is a powerful technique for learning surrogate maps for partial differential equation (PDE) solution operators. For many real-world applications, which often require high-resolution data points, training time and memory usage are significant bottlenecks. While there are mixed-precision training techniques for standard neural networks, those work for real-valued datatypes on finite dimensions and therefore cannot be directly applied to FNO, which crucially operates in the (complex-valued) Fourier domain and in function spaces. On the other hand, since the Fourier transform is already an approximation (due to discretization error), we do not need to perform the operation at full precision. In this work, we (i) profile memory and runtime for FNO with full and mixed-precision training, (ii) conduct a study on the numerical stability of mixed-precision training of FNO, and (iii) devise a training routine which substantially decreases training time and memory usage (up to 34%), with little or no reduction in accuracy, on the Navier-Stokes and Darcy flow equations. Combined with the recently proposed tensorized FNO (Kossaifi et al., 2023), the resulting model has far better performance while also being significantly faster than the original FNO.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google