Speeding up Fourier Neural Operators via Mixed Precision

解决问题：本论文旨在解决利用傅里叶神经算子（Fourier neural operator, FNO）学习偏微分方程（PDE）解算子的过程中，高分辨率数据点下的训练时间和内存使用问题。同时，论文还通过研究混合精度训练的数值稳定性，提出了一种训练方法，可以显著减少训练时间和内存使用（最多可减少34%），同时保持准确性。

关键思路：本文的关键思路是利用混合精度训练（mixed-precision training）来加速FNO的训练，从而减少训练时间和内存使用。与标准神经网络不同，FNO操作在复数域和函数空间中，因此无法直接应用于实数数据类型和有限维空间的混合精度训练技术。然而，由于傅里叶变换已经是一种近似方法，我们不需要以完整的精度执行操作。通过研究混合精度训练的数值稳定性，本文提出了一种训练方法，可以显著减少训练时间和内存使用，同时保持准确性。

其他亮点：本文的实验结果表明，提出的混合精度训练方法可以显著减少训练时间和内存使用（最多可减少34%），同时保持准确性，特别是在Navier-Stokes和Darcy流方程中。此外，本文还提出了张量化FNO的概念，与该方法相结合后，模型的性能得到了进一步提升。本文提出的混合精度训练方法和张量化FNO的思想值得进一步研究和探索。

关于作者：本文的主要作者包括Colin White、Renbo Tu、Jean Kossaifi、Gennady Pekhimenko、Kamyar Azizzadenesheli和Anima Anandkumar。他们分别来自卡内基梅隆大学、加州大学洛杉矶分校和加州理工学院。Anima Anandkumar是机器学习领域的知名学者，在深度学习、无监督学习、分布式优化等方面有很多代表性工作。例如，她曾经提出了一种基于张量分解的分布式优化算法，并应用于神经网络的训练（Anandkumar et al., 2014）。

相关研究：近期的相关研究包括以下论文：1）"Fourier Neural Networks for Approximation of High-Dimensional Functions"（作者：Zongyi Li等，机构：加州大学伯克利分校）；2）"Learning to Solve Partial Differential Equations via Deep Neural Networks"（作者：E. Haber等，机构：加州大学圣巴巴拉分校）；3）"Deep learning for universal linear embeddings of nonlinear dynamics"（作者：J. Nathan Kutz等，机构：华盛顿大学）。

论文摘要：本文介绍了一种名为傅里叶神经算子（Fourier neural operator，FNO）的技术，用于学习偏微分方程（PDE）解算符的替代映射。然而，在许多需要高分辨率数据点的实际应用中，训练时间和内存使用是重要的瓶颈。虽然标准神经网络有混合精度训练技术，但那些技术只适用于有限维度的实值数据类型，因此不能直接应用于FNO，后者至关重要地在（复值）傅里叶域和函数空间中运作。另一方面，由于傅里叶变换本身就是一种近似（由于离散化误差），我们不需要以完全精度执行该操作。本文（i）对完全精度和混合精度训练的FNO进行内存和运行时间分析，（ii）对FNO混合精度训练的数值稳定性进行了研究，（iii）设计了一种训练程序，大大减少了训练时间和内存使用（高达34%），并且在Navier-Stokes和Darcy流方程中几乎没有降低准确性。结合最近提出的张量化FNO（Kossaifi等人，2023），得到的模型性能更好，同时比原始FNO快得多。

内容中包含的图片若涉及版权问题，请及时与我们联系删除