混合精度神经算子的保证近似界

神经算子，如 Fourier Neural Operators (FNO)，构成了一种学习 PDE 解算子以及函数空间之间其他映射的原则性方法。然而，许多现实问题需要高分辨率训练数据，而训练时间与有限的 GPU 内存构成了巨大障碍。一种解决方案是以混合精度训练神经算子，以降低内存需求并提升训练速度。然而，现有的混合精度训练技术是为标准神经网络设计的，我们发现将其直接应用于 FNO 会导致数值溢出和较差的内存效率。此外，乍看之下，FNO 中的混合精度似乎会导致严重的精度下降，因为降低傅里叶变换的精度在经典数值求解器中会产生糟糕的结果。我们表明情况并非如此；事实上，我们证明在以针对性方式完成时，降低 FNO 中的精度仍保证良好的近似界。具体而言，我们基于如下直觉：神经算子学习固有地引发一种近似误差，源于对无限维真实输入函数的离散化，这意味着无需以全精度训练。我们通过严格刻画 FNO 的近似误差与精度误差并为一般输入函数界定这些误差，使这一直觉形式化。我们证明精度误差在渐近意义上与近似误差相当。基于此，我们设计了一种简单方法，通过贪心寻找最优收缩顺序来优化内存密集的半精度张量收缩。通过对不同最先进神经算子、数据集和 GPU 的广泛实验，我们证明了我们的方法将 GPU 内存使用量减少高达 50%，并将吞吐量提升 58%，且精度几乎或完全没有降低。