[論文レビュー] Guaranteed Approximation Bounds for Mixed-Precision Neural Operators

Renbo Tu, Colin White+6 more|arXiv (Cornell University)|Jul 27, 2023

Model Reduction and Neural Networks被引用数 4

ひとこと要約

本稿では、Fourier Neural Operators (FNO) のための混合精度学習ルーチンを提案し、メモリ使用量を最大50%削減し、学習スループットを最大58%向上させながら、精度の損失を最小限に抑える。tanh前活性化と周波数モードの切り捨てを適用することで、複素数値のフーリエ領域における半精度学習が安定化し、モデル性能を損なわず、高解像度のPDE解の学習が効率的に行えるようになる。

ABSTRACT

Neural operators, such as Fourier Neural Operators (FNO), form a principled approach for learning solution operators for PDEs and other mappings between function spaces. However, many real-world problems require high-resolution training data, and the training time and limited GPU memory pose big barriers. One solution is to train neural operators in mixed precision to reduce the memory requirement and increase training speed. However, existing mixed-precision training techniques are designed for standard neural networks, and we find that their direct application to FNO leads to numerical overflow and poor memory efficiency. Further, at first glance, it may appear that mixed precision in FNO will lead to drastic accuracy degradation since reducing the precision of the Fourier transform yields poor results in classical numerical solvers. We show that this is not the case; in fact, we prove that reducing the precision in FNO still guarantees a good approximation bound, when done in a targeted manner. Specifically, we build on the intuition that neural operator learning inherently induces an approximation error, arising from discretizing the infinite-dimensional ground-truth input function, implying that training in full precision is not needed. We formalize this intuition by rigorously characterizing the approximation and precision errors of FNO and bounding these errors for general input functions. We prove that the precision error is asymptotically comparable to the approximation error. Based on this, we design a simple method to optimize the memory-intensive half-precision tensor contractions by greedily finding the optimal contraction order. Through extensive experiments on different state-of-the-art neural operators, datasets, and GPUs, we demonstrate that our approach reduces GPU memory usage by up to 50% and improves throughput by 58% with little or no reduction in accuracy.

研究の動機と目的

• 高解像度のPDEデータに対するFNOの学習に伴う高いメモリおよび計算リソースの要請に対処すること。
• 標準的な混合精度技術が適用できない複素数値のフーリエ領域における混合精度学習を可能にすること。
• 半精度FNO学習における数値的安定性と精度のトレードオフを検討すること。
• 精度を維持しつつリソース使用量を削減する、堅牢で高スループットの学習パイプラインを開発すること。
• Navier-StokesやDarcy流れといった実世界のPDEに対して、ゼロショットスーパーサンプルリゾリューションを含め、本手法の有効性を実証すること。

提案手法

• パフォーマンスボトルネックと潜在的な利得を特定するため、フル精度および混合精度学習のプロファイリング。
• フーリエ変換の前段階にtanh前活性化を適用し、半精度における数値の安定化を図る。
• 高周波数モードの切り捨てにより、半精度表現誤差の影響を低減する。
• 学習中に半精度から混合精度、最終的にフル精度へと移行する精度スケジュールの実装。
• ベクターキャストを用いた半精度複素数処理をサポートするcuFFTの利用による、効率的な計算。
• さまざまな解像度と精度設定で、Navier-StokesおよびDarcy流れデータセットを用いて手法の妥当性を検証。

Figure 1: Overview of the half-precision FNO block. We cast the input tensor to half precision and then use a tanh pre-activation for numerical stability. We use cuFFT to run the FFT, and we compute the tensor contraction in half precision by casting each complex number to a vector.

実験結果

リサーチクエスチョン

• RQ1 FNOの複素数値・関数空間演算という性質上、混合精度学習を安全に適用できるか？
• RQ2 半精度FNO学習において生じる数値的不安定性は何か？ それらはどのように緩和できるか？
• RQ3 周波数モードの切り捨ては、半精度FNOの精度と安定性にどのように影響するか？
• RQ4 モデル精度を損なわず、どれほどメモリ使用量と学習スループットを削減できるか？
• RQ5 提案手法は、ゼロショットスーパーサンプルリゾリューションにおいて、FNOの解像度不変性を保持するか？

主な発見

• 提案された混合精度FNO学習ルーチンにより、Navier-StokesおよびDarcy流れ問題においてGPUメモリ使用量が最大50%削減された。
• 最適化された混合精度パイプラインを用いることで、学習スループットが最大58%向上した。
• tanh前活性化と周波数モードの切り捨ての組み合わせが、半精度学習における数値オーバーフローおよびアンダーフローを効果的に防止した。
• 半精度FNOは、フル精度と比較してわずかな精度の低下（例：1024×1024のNavier-StokesにおいてL2誤差が0.00213から0.00236に増加）にとどまった。
• 精度スケジュール手法は純粋な半精度よりも顕著に高い精度を達成し、1024×1024のNavier-StokesにおいてL2誤差を0.00236から0.00170に低減した。
• 本手法は解像度不変性を保持しており、1024×1024解像度までの正確なゼロショットスーパーサンプルリゾリューションが、最小限の性能劣化で実現可能となった。

Figure 2: Training throughput and runtime as a function of the method, on different hardware. For mixed-precision FNO + AMP, we consistently observe an improvement of throughput up to 1.58X over the baseline, and a speedup of up to 36.8% . Our method also improves upon using only AMP in throughput b