Research map / Learning theory and algorithms / Combinatorial and search optimization
Guaranteed Approximation Bounds for Mixed-Precision Neural Operators
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes the first mixed-precision training method for neural operators (FNO, TFNO, SFNO, GINO), using a greedy tensor-contraction ordering, tanh pre-activation for stability, and a theoretical bound showing half-precision error is comparable to discretization error.
Problem
Training neural operators such as FNO is memory- and compute-intensive at high resolution; standard automatic mixed precision (AMP) causes overflow and poor memory savings for the complex-valued spectral operations, and it is unclear whether reduced precision harms approximation accuracy.
Approach
Shows theoretically that the precision error of computing the discrete Fourier transform in half precision is asymptotically comparable to the inherent discretization error of FNO, so half precision is justified; then runs spectral tensor contractions in half precision via a greedy contraction order and stabilizes with tanh pre-activation before each FFT.
Key results
- Mixed-precision FNO reduces GPU memory by up to 50% and improves training throughput by up to 1.58X over full precision on TFNO/Navier-Stokes.
- Test errors remain within 1% of full-precision baselines across four datasets, with little or no accuracy loss.
- A precision schedule (mixed then AMP then full precision) achieves better zero-shot super-resolution error than full precision on Navier-Stokes (H1 0.00503 vs 0.00557 at 128x128).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Train FNO with mixed precision instead of full precision | FNO · Navier-Stokes zero-shot 128x128 | H1 error: 0.00557 -> 0.00624 | Hurt |
| Train FNO with mixed precision instead of full precision | FNO · Navier-Stokes zero-shot 128x128 | L2 error: 0.00213 -> 0.00236 | Hurt |
| Use a mixed-to-AMP-to-full-precision schedule instead of full precision throughout | FNO · Navier-Stokes zero-shot super-resolution | H1 error: 0.00557 -> 0.00503; 0.00597 -> 0.00542; 0.00610 -> 0.00555; 0.00616 -> 0.00558 | Helped |
| Use a mixed-to-AMP-to-full-precision schedule instead of full precision throughout | FNO · Navier-Stokes zero-shot super-resolution | L2 error: 0.00213 -> 0.00170 | Helped |
| Use mixed-precision FNO contractions with AMP | TFNO · Navier-Stokes | training throughput: 1.58X | Helped |
| Use mixed-precision FNO contractions with AMP | SFNO · Spherical SWE | training throughput: 1.33X | Helped |
| Use half-precision FNO contractions together with AMP | FNO · GPU training | GPU memory reduction: 50% | Helped |
| Add tanh pre-activation before each forward FFT (ablation) | Mixed-precision FNO · Mixed-precision FNO training | numerical stability | Helped |
| Add tanh pre-activation (ablation) | FNO · Final model error | test error | No clear effect |
| Apply loss scaling, gradient clipping, or delayed updates rather than tanh pre-activation (ablation) | Mixed-precision FNO · Darcy Flow | training stability | Hurt |
| Use too few frequency modes (ablation) | FNO · Darcy Flow | accuracy | Hurt |
| Use too many frequency modes (ablation) | FNO · Darcy Flow | runtime | Hurt |
Surprising findings
- Naive mixed-precision FNO training produces NaN outputs on all four datasets, and common stabilizers such as loss scaling and gradient clipping fail to prevent divergence.
- BF16 degrades accuracy on Navier-Stokes, and PyTorch does not support BF16 for discrete Fourier transforms.
Limitations
- The theoretical bounds are worst-case over Lipschitz functions and the lower bound is shown only for frequency omega=1.
- Reported results are on specific architectures (TFNO, SFNO, GINO) and datasets; GINO on ShapeNet Car has batch size fixed at 1 so throughput gains are not measured there.
- Precision schedule and zero-shot results come from a single 19-hour training setting.
Experimental setup
Datasets: Navier-Stokes Darcy Flow Spherical Shallow Water Equations Shape-Net Car Ahmed-body CFD
Benchmarks: Navier-Stokes Darcy Flow Spherical Shallow Water Equations Shape-Net Car Ahmed-body CFD
Metrics: L2 error H1 error GPU memory reduction training throughput
Baselines: Full-precision FNO PyTorch AMP U-Net U-Net + AMP TF32 BF16
Builds on: Fourier Neural Operator Tensorized Fourier Neural Operator Geometry-Informed Neural Operator Spherical Fourier Neural Operator automatic mixed precision
Closest papers by meaning
- Amortized Fourier Neural Operators (NeurIPS 2024 · similarity 0.90)
- Optimization for Neural Operators can Benefit from Width (ICML 2025 · similarity 0.89)
- Maximum Update Parametrization and Zero-Shot Hyperparameter Transfer for Fourier Neural Operators (ICML 2025 · similarity 0.89)
- Component Fourier Neural Operator for Singularly Perturbed Differential Equations (AAAI 2024 · similarity 0.89)
- Neural Operators with Localized Integral and Differential Kernels (ICML 2024 · similarity 0.88)
- MgNO: Efficient Parameterization of Linear Operators via Multigrid (ICLR 2024 · similarity 0.88)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.