FNO-MP: a mixed-precision training acceleration framework for FNO
TL;DR
FNO–MP achieves acceleration through three synergistic designs: a progressive multi-resolution and modality growth strategy that incrementally reduces FFT overhead; a stabilized early stopping mechanism adapted to FNO’s training oscillations to adaptively terminate each stage; and mixed precision as a foundational accelerator working in concert with the above strategies.
Abstract
Neural operators like the Fourier Neural Operator (FNO) have demonstrated remarkable success in solving partial differential equations (PDEs) but suffer from high training costs due to fast Fourier Transform (FFT) operations. Targeting the unique computational bottleneck of the FFT in FNO, FNO–MP achieves acceleration through three synergistic designs: a progressive multi-resolution and modality growth strategy that incrementally reduces FFT overhead; a stabilized early stopping mechanism adapted to FNO’s training oscillations to adaptively terminate each stage; and mixed precision as a foundational accelerator working in concert with the above strategies. On the Navier- Stokes and Darcy flow equations, FNO–MP delivers 1.92× and 2.38× speedups, respectively, while maintaining comparable or slightly improved test errors. These results demonstrate that the synergistic alignment of precision scheduling, architectural adaptation, and optimization control can effectively unlock training efficiency in neural operators, providing a practical solution for accelerating FNO without compromising predictive accuracy.