Concepts de base
ニューラルオペレータの混合精度学習手法を提案し、理論的に近似誤差を保証する。これにより、メモリ使用量を最大50%削減し、スループットを最大58%向上させることができる。
Résumé
本研究では、ニューラルオペレータの混合精度学習手法を提案している。ニューラルオペレータは、偏微分方程式の解を学習する強力な手法であるが、高解像度のデータを扱うには計算リソースが大きな障壁となる。
提案手法では、ニューラルオペレータの中核となるフーリエ変換の演算を混合精度で行うことで、メモリ使用量を最大50%削減し、スループットを最大58%向上させることができる。理論的には、フーリエ変換の精度誤差が離散化誤差と同程度であることを示し、混合精度でも十分な近似精度が得られることを証明している。
具体的には、以下の取り組みを行っている:
- 複素数テンソル演算の最適化: メモリ集約的な複素数テンソル演算を、部分的に実数演算に置き換えることで効率化する。
- 数値安定化: フーリエ変換前にtanh関数を適用することで、混合精度での数値不安定性を解消する。
- 精度スケジューリング: 学習初期は混合精度、後期は単精度に切り替えることで、高精度な推論を実現する。
提案手法は、Navier-Stokes、Darcy Flow、球面浅水波、車体CFDなどの4つのデータセットで検証され、既存手法と比べて大幅な性能向上を示している。
Stats
混合精度FNOはNavier-Stokesデータセットで、メモリ使用量を50.3%削減し、スループットを1.58倍向上させた。
混合精度FNOはDarcy Flowデータセットで、メモリ使用量を25.0%削減し、スループットを1.32倍向上させた。
混合精度FNOはSphereical SWEデータセットで、メモリ使用量を45.9%削減し、スループットを1.33倍向上させた。
混合精度GINOはShapeNet Carデータセットで、メモリ使用量を28.7%削減した。
混合精度GINOはAhmed Bodyデータセットで、メモリ使用量を38.4%削減した。