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.

A modified FC-Gram approximation algorithm with provable error bounds

Gram quadrature: numerical integration with Gram polynomials

Extended Newton Methods for Conic Inequalities: Approximate Solutions and the Extended Smale Α-Theory

Improved Convergence Bounds For Operator Splitting Algorithms With Rare Extreme Errors

Sharper Bounds for Proximal Gradient Algorithms with Errors

Error Bounds for a Kernel-Based Constrained Optimal Smoothing Approximation

Error Bounds for Lanczos-Based Matrix Function Approximation

An efficient iterative procedure in hyperbolic space and application to non-linear delay integral equation

Scaling Optimized Hermite Approximation Methods

Asymptotic expansions for approximate solutions of boundary integral equations

Improved Laguerre Spectral Methods with Less Round-off Errors and Better Stability

An Accurate Numerical Algorithm Based on the Generalized Narayana Polynomials to Solve a Class of Caputo-Fabrizio and Liouville-Caputo Fractional-Order Delay Differential Equations

A posteriori error bounds for the block-Lanczos method for matrix function approximation

Computation of highly oscillatory integrals using a Fourier extension approximation

Interpolation by the Exact Inversion of the Gram Matrix

A note on improvement by iteration for the approximate solutions of second kind Fredholm integral equations with Green's kernels

Superconvergent Non-Polynomial Approximations

Guaranteed Approximation Bounds for Mixed-Precision Neural Operators

Improving the Accuracy of Exponentially Converging Quadratures

Convergence and optimality of an adaptive modified weak Galerkin finite element method

Effective Quadratic Error Bounds for Floating-Point Algorithms Computing the Hypotenuse Function