Guaranteed Approximation Bounds for Mixed-Precision Neural Operators

TL;DR

This work tackles the memory and training-time barriers of learning neural operators for PDEs by introducing the first mixed-precision training framework tailored for Fourier-based operators. The authors prove theoretical bounds that quantify discretization and precision errors, showing that half-precision losses are asymptotically comparable to the discretization error, which justifies mixed-precision use. They develop a practical, memory-efficient contraction strategy and stabilize complex-valued spectral operations with tanh pre-activation, enabling robust training of FNO, TFNO, SFNO, and GINO across multiple datasets. Empirically, the method achieves up to 50% GPU memory reduction and up to 58% higher training throughput with minimal accuracy degradation, and it preserves discretization-convergence under zero-shot super-resolution with a precision-scheduling tactic. The approach significantly lowers the barrier to training high-resolution neural operators, with open-source code to facilitate adoption and reproducibility.

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.

Figures (16)

• Figure 1: Top: example data points from each dataset: Navier-Stokes, Darcy Flow, Spherical Shallow Water, Shape-Net Car CFD, and Ahmed-body CFD. Bottom: performance of our method compared to full-precision and AMP on each dataset. For each dataset, we plot test error (y-axis) and GPU memory (x-axis), and we annotate the maximum throughput (proportional to the area of each ball). All data are measured on the same hardware (RTX 3090 Ti) and the same virtual environment. Memory decreases by up to 50%, while $L^2$ loss increases by at most 0.28%.
• Figure 2: Overview of our Mixed-FNO pipeline.
• Figure 3: GPU memory usage reduction across different variants of neural operators on diverse tasks. We use an Nvidia RTX 3090 Ti GPU. Our method reduces memory by up to 50%, representing a super-linear combination of the two other methods because it avoids additional casting to and from full precision during the forward pass.
• Figure 4: Training throughput and runtime as a function of the method, on different GPUs. For mixed-precision FNO + AMP, we consistently observe an improvement of training throughput up to 1.58X over the baseline with the TFNO model on Navier Stokes, and up to 1.33X with the SFNO model on Spherical Shallow Water Equations (SWE). Our method also improves upon using only AMP in throughput by over 1.3X on Navier-Stokes and by over 1.2X on Spherical SWE. Batch sizes are selected to fully utilize each GPU.
• Figure 5: Test H1 error curves for FNO on the Navier-Stokes (top left) and Darcy flow (top right) datasets. Test L2 error curves for GINO on the Shape-Net Car (bottom left) and for SFNO on the Shallow Water Equation (bottom right) datasets. Each plot shows the mean of three random seeds and standard deviation as error bars. We compute and report the average difference between training curves in the legend. We also annotate the difference in final test errors with the memory savings of our mixed precision approach.

Theorems & Definitions (8)

• Theorem 3.1
• Theorem 3.2
• proof
• proof
• Theorem A.1
• proof
• Theorem A.2
• proof