[论文解读] 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

一句话总结

本文提出了一种用于傅里叶神经算子（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训练中的数值稳定性与精度权衡问题。
• 开发一种鲁棒、高吞吐量的训练流水线，在降低资源消耗的同时保持精度。
• 在真实世界的PDE问题（如纳维-斯托克斯方程和达西流）上验证该方法的有效性，包括零样本超分辨率。

提出的方法

• 通过全精度和混合精度训练的性能分析，识别性能瓶颈与潜在优化空间。
• 在傅里叶变换前应用tanh预激活，以稳定半精度下的数值范围。
• 通过截断高频模态，减少半精度表示误差的影响。
• 采用从半精度到混合精度再到全精度的精度调度策略，逐步过渡训练过程。
• 利用cuFFT结合通过向量转换实现的半精度复数运算，提升计算效率。
• 在纳维-斯托克斯方程和达西流数据集上，针对不同分辨率和精度设置进行方法验证。

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训练方法在纳维-斯托克斯方程和达西流问题上，可将GPU显存使用量减少高达50%。
• 采用优化后的混合精度流水线后，训练吞吐量最高可提升58%。
• tanh预激活与频率模态截断的结合能有效防止半精度训练中的数值溢出与下溢。
• 与全精度相比，半精度FNO仅造成微小的精度下降（例如，在1024×1024纳维-斯托克斯问题上，L2误差从0.00213增加至0.00236）。
• 精度调度方法相比纯半精度训练显著提升了精度，1024×1024纳维-斯托克斯问题上的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