Guaranteed Approximation Bounds for Mixed-Precision Neural Operators
Authors: Renbo Tu
Colin White
Jean Kossaifi
Boris Bonev
Gennady Pekhimenko
Kamyar Azizzadenesheli
Anima Anandkumar
Venue: ICLR 2024 poster (2024)
Keywords: neural operators
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
TL;DR: We theoretically and empirically demonstrate the power of a new mixed-precision training pipeline for neural operators.

Abstract: Neural operators, such as Fourier Neural Operators (FNO), form a principled approach for learning solution operators for partial differential equations (PDE) 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.

Published conference paper at ICLR 2024
GUARANTEED APPROXIMATION BOUNDS FOR MIXED- PRECISION NEURAL OPERATORS
Renbo Tu Colin 2hite" Jean Kossaifi Boris Bonev", Gennady Pekhimenko',
Kamyar Azizzadenesheli? Anima Anandkumar2 University of Toronto," Caltech, NVIDIA
ABSTRACT
Neural operators, such Fourier Neural Operators (FNO), form principled ap- proach for learning solution operators for partial differential equations (PDE) and other mappings between function spaces. However, many real-world problems re- quire high-resolution training data, and the training time and limited GPU memory pose big barriers. One solution to train neural operators in mixed precision reduce the memory requirement and increase training speed. However, existing mixxedpprecison 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, 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 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. Specifi- cally, build on the intuition that neural operator learning inherently induces ground-truth
approximation error, arising from discretizing the input function, implying that training in full precision is not needed. We formalize this intuition by rigorously characterizing the approximation and precision eπrоrs 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 tensor contractions by greedily finding the optimal contraction order. Through extensive experiments on different state-of-the-art neural operators, datasets, and GPUs, demonstrate that approach reduces GPU memory usage by up 50% and improves throughput by 58% with little or no reduction accuracy.
intensive half- half- precision
INTRODUCTION
Real-world problems science and engineering often involve solving systems of partial differential equations (PDEs) (Strang, 2007). These problems typically require fine grid for numerical solvers to guarantee convergence, e.g., climate modeling and 3D fluid dynamics simulations, and the full-scale solutions these real-world systems are out of reach even for the world's largest supercomputers (Schneider et al., 2017).
To overcome the computational challenges of numerical solvers, fast surrogates using machine learning have been developed. Among them, neural operators are powerful data-driven technique for solving PDEs (Li al., 2020a; 2021a; Kovachki et al., 2023; Lu al., 2021). Neural operators learn maps between function spaces, and they can be used to approximate the solution operator of a given PDE. The input and output functions neural operators can at any resolution or on any mesh, and the output function can be evaluated at any point in the domain; therefore, neural operators are discretization convergent: once the neural operator is trained, it can be evaluated, without any retraining, at any resolution, and it converges to unique limit under mesh refinement (Kovachki et al., 2023). By learning from discretized data from input and solution functions, the trained models perform inference orders of magnitude faster than traditional PDE solvers (Kovachki et al., 2023). In particular, the Fourier Neural Operator (FNO) and its extensions have been successful in solving
Equul contribution. Email to: renbo. eu@ .utorrntoo.c

---

Published conference paper at ICLR 2024
Navter Stakes, TFNO
Darcy Row. TFNO
Spherical SWE SPMO
Shapeliet Car, GINO A mimmmayya pame: +8.00h
Anmed Body, GING
amenery- Suhrrp--luuticcmm orar +2.2%
Anamay-Js bama 4aai
Amemay -4 -4 ar. +6.2%
& o. aitt aaron va.s2%
C.T%. E.5m
Text U0
100% 6..%
9.%%
737GE
7:72
(68)
(Ours)
(Ours)
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 loss increases by at most 0.28%.
PDE-based problems with significant speedups (Li et al., 2021a; 2023; Bonev et al., 2023; Kossaifi al., 2023; Liu et al., 2022: Gopakumar et al., 2023).
Despite their successes, training of neural operators is still compute- and memory-intensive when faced with extremely high-resolution and large-scale problems. For conventional deep learning models, there is a wealth of knowledge on automatic mixed precision (AMP) training in order to reduce memory usage and increase computational throughput (Rakka et al., 2022). We find that their direct application to neural operators results in poor memory efficiency and numerical overflow; when applied to FNO, most memory-intensive operations are in the spectral domain, which complex- complex-valued valued and not handled well by the standard AMP,
Figure 1.
Conventional wisdom may suggest that reducing precision FNO without drastic accuracy degra- dation is not feasible. The presence of the Fourier transform in FNO appears be limitation reducing precision since it suffers from numerical instabilities on high-frequency modes under reduced precision. Hence, numerical solvers, such as pseudo-spectral spectral solvers, do indeed require very high precision to ensure numerical stability (Trefethen, 2000), and time-stepping leads to an accumulation of round-off errors under low precision that quickly explode. However, show both theoretically and empirically that this is not the case for FNO and other operator learning approaches.
Our approach: In this work, we devise the first -eeccd--reiision training method for neural operators and also derive approximation bounds that guarantee the expressivity of mixed-precision operators. We use the intuition that the Fourier transform within neural operators already approximated by the discrete Fourier transform, since the training dataset is a discrete approximation of the ground-truth continuous signal. Intuitively, since we already incur approximation error from discretization, there no need to run the discrete Fourier transform in full precision. Building on this, we show that, in single FNO layer, the round-off error due to lower precision is small. Therefore, the overall round-off eтror remains small since full FNO architecture is only made up of few much smaller than the number steps in a classical pseudo-spectrrg solver where errors accumulate and explode
We make the above intuitions concrete, and develop the following theoretical approximation bounds: we characterize the precision error and resolution error of the FNO block, proving asymptotic bounds of n - and 6, respectively, for mesh size n, dimension d, and dynamic range e. Therefore, justify using mixed precision since its eπror is comparable to the discretization error already present in FNO.
Motivated by these theoretical results, we introduce mixed-precision training method by optimizing the memory-intensire half-precision tensor contractions. We devise a simple and lightweight greedy strategy to find the optimal contraction order, which considers vectorization for each intermediate tensor. Unfortunately, naively training neural operators mixed precision leads to prohibitive numerical instability due overflows in the FNO block. To address this issue, study numerical
2

---

Published conference paper at ICLR 2024
stabilization techniques for neural operators, finding that tanh pre-activation before each FFT consistently avoids numerical instability. Further, we incorporate additional vectorization steps for complex-valued inputs, not present standard packages designed for real-valued networks.
We carry out extensive experiments to show significant reduction in memory usage and improved throughput without sacrificing accuracy; see Figure 1. We consider three variants of FNO: tensorized (TFNO) (Kossaffi et al., 2023), spherical (SFNO) (Bonev et al., 2023), and geometry-informed (GINO) (Li et al., 2023). Across different datasets and GPUs, our method results in up to 58% improvement in training throughput and 50% reduction in GPU memory usage with little reduction in accuracy.
Furthermore, show that method is discretization convergent via zero-shot super-resolution experiments. Finally, we propose a precision schedule training routine, which transitions from mixed to full precision during training: show this method achieves better than the baseline full-precision accuracy. Our mixed-precision training routine is lightweight and easily added to new neural operators. We release our codebase and all materials needed to reproduce our results https: // com/ hwtp://euaaapartoooooo wwwp:///wnnaaaparaeeeoooe
We summarize our main contributions below:
We introduce the first mixed-pprecisio training routine for neural operators, which optimizes the memory-intensinn tensor contraction operations in the spectral domain, and we propose using tanh pre-activations minimize numerical instability.
We theoretically ground our work by characterizing the precision and discretization errors of the FNO block, showing that these errors are comparable, proving that, done right, mixed- precision training of neural operators leads little or no performance degradation.
We empirically verify the superiority of our mixed- preecipoeciiion training approach on three state-of- the-art neural operators, TFNO, GINO, and SFNO, across four different datasets and GPUs. Our method uses half the memory and increases training throughput up to 58% % across different GPUs with little or no reduction in accuracy (< 0.1%).
We provide an efficient implementation of our approach in PyTTrcc, which we open-source, along with all data needed to reproduce our results.
2 BACKGROUND AND RELATED WORK
Neural Operators. Many real-world scientific and engineering problems rely on solving partial differential equations (PDEs). As such, there is recent focus using machine learning-base methods to solve PDEs (Gupta et al., 2021; Bhatmagar et al., 2019; Lu et al., 2019; Adler & Öktem, 2017). However, most of these methods use standard neural networks and are therefore limited to a mapping between fixed input and output grid. In other words, they operate on a fixed, regular discretization and cannot learn the continuous mapping between function spaces. Neural operators are a new technique that addresses this limitation by directly learning maps between function spaces (Li et al., 2021b; 2020a;b; Kovachki al., 2023). The input and output functions to neural operators can be in any resolution or mesh, and the output function can be evaluated at any point in the domain; therefore, neural operators are discretization convergent: once the neural operator trained, can be evaluated, without any retraining, at any resolution, and it converges to a unique limit under mesh refinement (Kovachki et al., 2023).
FNO and extensions. The Fourier neural operator, inspired by spectral methods, is highly successful neural operator (Li et al., 1a2 Gopakumar et al., 2023; Renn et al., 2023; Wen et al., 2022). Let A fa DA + R&A and 4 {u Du Ød denote the input and output function spaces, respectively. In this work, we consider the case where D = Du C Rdd for d E N. Given a dataset of pairs of initial conditions and solution functions {a,, Uy 1 which are consistent with operator G(a;) uy for all 1 S j S N, the goal is learn neural operator Go that approximates G. The primary operation in FNO is the Fourier convolution operator, (Kv+) F (R Tx (z)) Vr E D, where F and F-1 denote the Fourier transform and its inverse, R denotes a learnable transformation, Tx denotes a truncation operation, and denotes the function at the current layer of the neural operator. We use the discrete Fast Fourier Transform (FFT) and its inverse (iFFT) to implement this operator discrete data.
3

---

Published as a conference paper at ICLR 2024
Building on FNO, Kossaifi et al. (2023) introduced the Tensorized Fourier Neural Operator (TFNO). Building on tensor methods, which have proven very successful in deep learning Panagakis et al. (2021; 2024). TENO computes a tensor factorization (Kolda & Bader, 2009) of the weight tensors in the spectral domain, acting regularizer that boosts performance while decreasing parameter count. Additionally, two FNO-based neural operators have recently been proposed to improve performance in non-regular geometries. The Spherical Fourier Neural Operator (SFNO) (Bonev et al., 2023) is an extension of FNO to the spherical domain, which is achieved by the use of the spherical convolution theorem (Driscoll & Healy, 1994). The Eeoometry-Informen Neural Operator (GINO) (Li et al., 2023) is a highly efficient FNO-based architecture for solving PDEs with varying, irregular geometries. The architecture consists of an FNO, along with graph neural operators transform the irregular grid inputs into and from regular latent grids on which FNO can be efficiently applied (Li et al., 2020a).
Mixed-Precisiog Training. Mixed-precision training of neural networks consists of reducing runtime and memory usage by representing input tensors and weights (and performing operations) lower-than-ttandard precision. For example, PyTorch has a built-in pixee-precisioo mode called automatic mixed precision (A (AMP), which places all operations at float 1 6 rather than float 32, with the exception of reduction operations, weight updates, normalization operations, and specialized operations (Paszke et al., 2019). bfloat16 16 and TF32 are common mixed-precision methods, yet they do not support discrete Fourier transforms, which are essential to FNOs (Kalamkar et al., 2019; Dean et al., 2012). There is variety of work targeted at quantization for inference, yet these works do not reduce memory during training (Goel et al., 2020).
While well-studied for standard neural nets, mixed-precision training has not been studied for FNO (Zhao et al., 2022; De Sa et al., 2018; Micikevicius et al., 2017; Jia et al., 2018). The most similar work to ours is FourCastNet (Pathak et al., 2022), a large-scale climate model that uses mixed precision with Adaptive Fourier Neural Operators (Guibas et al., 2021). However, mixed precision is not applied to the FFT or complex- valuued multiplication operations, a key challenge the current work addresses. Very recently, another work studies the method of quantization for FNO at inference time (Dool et al., 2023). However, unlike work, they only study methods that improve memory usage at inference time, not training time, Another paper has appeared recently, which proposes to apply mixed precision to PINNs and DeepONet (Hayford et al., 2024).
3 GUARANTEED APPROXIMATION BOUNDS
In this section, to motivate our mixed-precision neural operator training routine in Section 4, we theoretically show that the inherent discretization error in the FNO block, from computing the discrete Fourier transform instead of the Fourier transform, is comparable to the precision error, from computing the FNO block in half precision rather than in full precision. We present our results in terms of forward Fourier transform, and we give the full details and discussion in Appendix A.
We start by formally defining the discretization error of the FNO block. Let D denote the closed unit cube [0, 1] for dimension d E N, subdivided into n md cubes Q1..... Qn with sidelength I/m and for each j. let E denote the point with the minimum value in each dimension, in Qj. Let V: D + D denote an intermediate function within the FNO. Recall from the previous section that the primary operation in FNO is the Fourier convolution operator, (Cox)(x)) We say the discretization error of F(v) is the absolute difference between the Fourier transform of v. and the discrete Fourier transform of u via discretization of u on @ = (fQ1...,,QA, Qn.. [E1,.... ..t........... Formally, given Fourier basis function (x) = 2T1((w.))
Disc(u, 2,w) = 1/ u (z)) - Σ Eu(S)eaG))1)
In other words, if we knew the true (infnite-dimensiooaa)) input function, we could run FNO using the (continuous) Fourier transform. However, for real-world applications, we only have access to discretization of the input, for example, on a 128 x 128 grid, so we incur a discretization error. We bound the discretization error follows. For the full proof details, see Appendix A.
Theorem 3.1. For any M 0 and I. ' 1 let K. C C(D) be the set of I-LLipschitz functions, ---iapcchitffunncionn,, bounded by ||oloo S M. Then all n, Q, there exist w, C1, 02 > 0 such that
civd Mn 3/2 < sup (Disc(r, w)) S 02 V Vd(ww) + L)Mn 1/4
ee
4

---

Published conference paper at ICLR 2024
Intuitively, we bound the Riemann sum approximation the true integral by showing that in the case where the function is bounded and Lipschitz, then each of the n intervals contributes n -1++%/) error, up to parameters of the function. Note that the discretization error upper bound also scales linearly for the frequency w, although we empirically show in Section 4 that the energy concentrated the lower frequency modes. The lower bound is satisfied when v(x) 21 Idd Furthermore, note that given a function v, if n is not large enough, the discretization enror can become arbitrarily large due to aliasing. For example, the function v(x) = M sin (2 + w)x) has discretization error (2M)).
Next, say that the precision error of F(t) is the absolute difference between F(t) and F(v), computing the discrete Fourier transform half precision. Specifically, we define an (Au, c. T) precision system
mapping q R S for the set S {0} (ag f () 1==0- such that for all x E IR, ((x) = a argmmnges
y/- This represents a simplified version of the true mapping used by Python from IR to float32 float (we give further justification of this definition in Appendix A). Then, we define
Prec(u, Q.9.w)
Σ (E)))))))) Σ (al0s))9(92
Now, we bound the precision error as follows.
Theorem 3.2. For any M and L 2 1 let K. C. C(D) be the set of L-Lipschitz functions, bounded by |ol < M. Furthermore let q be an (4o, e, ())peeiinnriiii system. Then for all n, 2, w, there
exists e > 0 such that
sup (Prec(u, Q.).w)) S e. e.. veA
Once again, we show that each interval contributes s/n error up to the function's parameters. Taken together, Theorem 3.1 and Theorem 3.2 show that asymptotically, the discretization error can be as large as Mn (up to constants), while the precision error is always bounded by c.M. For example, for float precision (c = 10-4), our theory suggests that the precision error is comparable the discretization error for three-dimensional meshes up size 000 000. Given this motivation, present our mixed-precision training pipeline the next section.
EMPIRICAL STUDY OF OUR MIXED-PRECISIN NEURAL OPERATOR
We empirically study our proposed mixed-precision pipeline for FNO training and demonstrate significant memory usage reduction and large increases in training throughput on various GPUs. We then study failure modes of a naïve application of mixed precision and present our empirically justified rreaactivation-based stabilizer solution. Finally, we perform extensive experiments and ablations on our end-to-end training method across different architectures and datasets.
4.1 DATASETS AND EXPERIMENTAL SETUP
We summarize each of the four datasets we use; see Appendix B.2 for detailed descriptions.
Navier-Stokes. We consider the Navier-Stokes equations for viscous, incompressible fluid in vorticity form on the unit torus. We use the same dataset as Kossaifi et al. (2023), with a Reynolds number of 500, composed of 10000 training samples and 2000 test samples with resolution 128 128. Darcy Flow. We consider the steady-state 2D Darcy Flow equation, which models fluid flow through porous medium. We use the same dataset as Li et al. (2021a), with 5000 training samples and 1000 test samples at resolution 128 x 128. Spherical Shallow Water Equations. We use the dataset from Bonev
(2023), which generates random initial conditions on the sphere at resolution 256 x 512. At each epoch, 120 training samples and 20 validation samples are generated on the fly. Shape-Net Car and Ahmed-body. Our final two datasets are 3D real-world car dataset generated by prior work (Umetani & Bickel, 2018; Li et al., 2023), which consists of mesh points that represent unique 3D car, and the goal to predict the full 3D pressure field. We use 611 water water-tight shapes from surfaces from Shape-Ne- Net (Chang al., 2015), with 500 samples for training and the rest for the test set. For Ahmed-body body, we have 500 for training and 51 for test. The spatial resolution for both is 64 64 x 64.
Experimental Setup. We run the Navier Stokes and Darcy flow experiments on the TFNO archi- tecture (Kossaifi et al., 2023). For the Spherical SWE, we use SFNO (Bonev al., 2023) to handle
5

---

Published conference paper at ICLR 2024
the spherical geometry, and for Shape-Net Car and Ahmed-body, we use GINO (Li al., 2023) to handle the large-scale, irregular geometry. All models have the FNO as backbone. We use the official implementation and default hyperparameters for all models.
4.2 MIXED-PRECISION MODULE FOR COMPLEX-VALLEED TENSOR CONTRACTION
Now we introduce our theoretically and empirically motivated mixed- precision pipeline for the FNO block. We start by profiling FNO training workloads, identifying the complex- valued tensor contraction within the FNO block as the computational bot- tleneck, accounting for 4 out of the 5 most time-consuming GPU ker- nels in the entire training pipeline (see Appendix B.4 for the full pro- filing results). Furthermore, exist- ing mixed-precision tools such as Py- torch's Automatic Mixed Precision (AMP) (Paszke et al., 2019) leave the operator full precision since only autocasts real-valued modules in FNO.
Had-Praccserr FP16 ippor
Contract CPd2 with view, _a__ real
Hall-Prez
Half Preo
activation
Weightt
MUF Connection FP
Hast-
Hail-Precc Temp
Weights
Halkhhedsoo CPRR Packoreed "anaar Cortaction
Shep >
F
Hall- Proc
Hatt Phoo
Cuipui
Etep
Halk Hall-Precicerr
Choose Each Step to Minimze Size Temponary Results
FP16 Cuiput
We introduce simple and lightweight workaround: we break down each ten- contraction into sub-expressions consisting of at most two terms and perform each of these contractions by temporarily converting tensors to reals. To optimize the memory usage, next e nsum step that minimizes the intermediate tensor size (see Figure 2). The complex-valued operations, including the forward FFT, tensor contraction, and inverse FFT, are done in half-precision. This completes the half-precision FNO block as AMP manages non-FNO, real-valued operations.
Figure 2:
Overview of our Mixed-FNO pipeline.
simple greedy algorithm select the
Our simple half-precision FNO block module achieves up 38.6% reduction in memory and up to 50% reduction in memory when combined with Pytorch's 's native AMP; see Figure 3. Note that the reduction from AMP Half-Prec FNO is greater than the sum of its parts due to lack of the additional casting to and from half precissrcision. The memory usage reduction can then be used to train on larger batch size, thereby improving the overall training throughput and efficiency. On three different Nvidia GPUs, RTX 3090 Ti, V100, and RTX A6000, we demonstrate a consistent improvement in training throughput, in terms of numbers training samples processed per second, from 1.23X to 1.58X over the baseline in Figure 4.
4.3 NUMERICAL STABILITY VIA PRE-ACTIVATION
Mixed-precision training is prone to underflow and overflow because its dynamic range is signifi- cantly smaller than that of full precision (Rakka et al., 2022). Notably, for all four of the datasets in our study, naïvely running FNO in mixed precision results in training failure due to NaN outputs. Furthermore, we empirically show that many common solutions, including loss scaling, gradient clipping, normalization, and delaying updates, all fail to address the numerical instability of mixed- precision FNO (see Appendix B.6). To mitigate this overflow issue, we find that pre-activation before each forward FFT is very effective method for overcoming numerical instability. We also find that the tanh pre-activation is the Appendix B.6. Unlike other functions, tanh minimizes changes to small inputs, as it is approxi- mately the identity function near 0 and is smoothly differentiable. Furthermore, tanh preserves the discretization-conne convergget property of FNO. The strong performance of tanh is also predicted by our theoretical results in Section 3, since it decreases the Loe norm and the Lipschitz constant of the input function. In Appendix B.6, further show that the tnnh pre- are-avtivation minimally alters the frequency-domant signal in both amplitude and phase. Finally, we demonstrate through an ablation that tanh has negligible impact the model's final error (See Appendix Table 5).
operation that we considered according to

---

Published conference paper at ICLR 2024
Comparison GPU Memory Consumption, Datasets
0.0%
GPU (MB)
15000
0.0%
74.5%
12500
-8..%%
0.07%
10000
0.0%
1.00%
0.4%
-14.000
350 13
7500
25.09%
50.3%
Flow
Spherical
Body
Dataset
(Ours)
FNG (Ouns)
Figure 3: GPU memory usage reduction across different variants of neural operators on diverse tasks. We use Nvidia RTX 3090 Ti GPU. Our method reduces memory by up 50%, representing a super-linear combination of the two other methods because avoids additional casting to and from full precision during the forward pass.
Throughputs Different Navier Stokes
Throughputs Different GPUs, Spherical
Training Throughout (samptess))
3m5
AMP
Mari PMD AMP joun0
Ad*
aso
Movd 7R (Dunn)
300
30
Hardware
Figure 4: Training throughput and runtime as a function of the method, on different GPUs. For mixed-precision FNO + AMP, consistently observe improvement of training throughput up to 1.58X over the baseline with the TFNO model on Navier Stokes, and 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 fully utilize each GPU.
4.4 ERROR COMPARISON WITH
Having resolved the critical issue of numerical stability, demonstrate that mixed-precision approach achieves errors within 1% of the full-precision baseline across the four datasets. Additionally, we propose a greeision-scheduiie training, which performs better than full precision in zero-shot super-resolution inference.
technique that transitions from mixed to full precision during
Mixed- vs. Full-Precision Training Curves. Figure illustrates that our mixed-precision approach achieves test eтrors on par with full precision throughout the training process, remaining within 1% of the preesionn baseline. We ensure apples-to-apples comparison by keeping all hyperparameters constant across the two precision settings. These hyperparameter configurations are originally optimized for full precision training, so we show that they are directly transferable to mixed precision.
Precision Scheduling and Zero-Shot Inference. An important property of neural operators is their discretization convergence, meaning that they can be trained on one resolution and tested on higher resolution (zero-shot super-resolution) (Kovachki et al., 2023). To achieve the best result, propose a precision schedule, which the first 25% of training is in mixed precisiie the middle 50% applying only AMP, and the final 25% in full precision. This follows simple intuition: in the early stages training, it is okay for gradient updates to be coarser, since the gradient updates are larger overall. However, in the later stages of training, the average gradient updates much smaller, so

---

Published conference paper at ICLR 2024
L2 Loss Epech. Navier Stokes
L2 Loss vs. Epoch, Darcy Flow
PNO
Full FNO
Full FNO
Average Difference: 0.00
412
Average Difference: 0.02
6:00
1:2
1055
6004
tnail Enee 3306-
6:00
Mamory Restatioon 30.3%
ffrennce Memony Rodudtonn 25.0%
Epoch
Epoch
1.2 Loss vs. Epoch, Spherical SWE
12 Loss vs. Epoch, ShapeNet Car
Mixed (Curs)
Mixed FNO (Curs) FNG Average Diffeence: 0.
Full
Average Difference: 0.00
0:22
4.4
fhaal Emir i ittrence: 4.000e-03 Mernery Reduction 28.7%
Emer 6.940c
Homory Roductions 45.0%
Epoch
Figure 5: Test H1 error curves for FNO on the Navier-Stokes (top left) and Darcy flow (top right) datasets. Test 1.2 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.
128x128
256x256
512x512
1024x1024
H²
1.2
H1
1,2
H
1,2
H²
1,2
Full FNO
0.00557
0.00213
0.00597
0.00213
0.00610
0.00213
0.00616
0.00213
Mixed FNO (Ours)
0.00624
0.00236
0.00672
0.00228
0.00688
0.00226
0.00693
0.00226
Precision schedule (Ours)
0.00503
0.00170
0.00542
0.00170
0.00555
0.00170
0.00558
0.00170
Table 1: Zero-shot super resolution. We test zero-shot super-resolution by training each model on
128 x 128 resolution for 19 hours. Mixed precision has small decrease in accuracy compared to full precision, and using a precision schedule achieves significantly better accuracy.
full precision is more important. We use the Kavier-Stokes dataset, trained on 128 x 128 resolution (the same setting and model as Figure 5), and tested on 256 x 256, 512 x 512, and 1024 x 1024 resolutions; see Table 1. We find that half- half-precision has small decrease in accuracy compared to full precision, and using a precision schedule achieves significantly better generalization.
4.5 COMPARISON AGAINST U-NETS
Despite being orignally designed for computer vision tasks, U-Nets have recently been used as PDE surrogates (Ronneberger et al., 2015). Here, compare FNOs against the U-Net baseline on the Darcy Flow and Navier-Stokes datasets. As shown in Table 2, FNO outperforms UNet: our mixed-precision approach yields higher memory reduction compared AMP applied to U-Nets.
4.6 ABLATION STUDIES
Here, we perform ablations of our mixed -eeecs-preiision procedure on different parameterizatioos of FNOs, regularization via reducing frequency modes, and training with other numerical systems.
Decomposition of FNO Weights. On the Navier-Stokes and Darcy Flow datasets, adopt Canonical-Polyadii (CP) factorization of the FNO weights using Kossaifi et al. (2019) to ensure better final error. For generic problem, FNO weights are usually not factorzied and are saved in
8

---

Published conference paper at ICLR 2024
Navier-Stokes
Darcy Flow
Model
Error
Memory Reduction
Error
Memory Reduction
Full ENO
0.003
0.01
50.4%
25.8%
Mixed FNO (Ours)
0.004
0.007
Full U-Net
0.024
20.9%
24.9%
U-Net + AMP
0.111
0.022
Table Comparison of L error and memory reduction with U-Nets, FNO consistently shows better final error, and our mixed-precision approach yields significantly more memory reduction on Navier-Stokes than AMP applied to U-Nets.
H1 Loss vs. Runtime, Navier Stokes
H1 Loss vs. Runtime, Darcy Flow
Maad IND (Dun), Derse
Mixed FNO (Ours),
Full IND, Denve
0.93
3827
Mixed PND [Ours), Dense
Full
(Ours).
3.35
0.81
3334
3322
035
0.5
Ria
63.3
42%
40.0
500
125
126
12.
348
Time per Epoch (seconds)
Time per Epoch Eseconds)
Figure For Navier-Stokes (Left) and Darcy Flow (Right), our method accelerates training, whether with CP factorized weights Dense unfactorized weights, without losing more than 1% in error.
their original, dense form. As shown in Figure 6, we experiment with both scenarios and show that mixed-precision approach improves runtime without sacrificing accuracy.
Number of Frequency Modes. Recall that the FNO architecture, after the FFT, we truncate to a fixed fraction of frequency modes to improve generalization, typically 1/3 to 2/3. We run ablation study on the number of frequency modes used the FNO architecture. We run frequency modes {16, 32, 64, 128} on the Darcy flow dataset in full and half precision; Figure 14. We find that using two few frequency modes hurts accuracy substantially while using too many frequency modes increases runtime substantially. There is not significant difference between half-precision and full precision for all frequencies. In addition, we show in Appendix B. 10 that for synthetic data, the precision error from half precision higher for higher frequencies relative the amplitude.
Other Mixed Precision Options. We also experimented with BrainFloat16 (BF16) and Tensor- Float32 (TF32). However, PyTorch does not support BF16 for discrete Fourier transforms, which are essential to FNOs. Even when applied to the rest of the network, BF16 suffers from error degradation on Navier Stokes, possibly due to having fewer precision bits than FP16. On the other hand, TF32 is not as efficient our approach, even on optimized hardware such the A100 GPU. Moreover, we simulated FP8 training via clipping. See Appendix B. for additional details.
5 CONCLUSIONS AND FUTURE WORK
In this work, we introduced the first mixed- preeeisin training method for neural operators. We derived approximation bounds for mixed -eeecds-reiiion neural operator and rigorously characterized the approximation and precision errors of FNO-based operators, proving that the precision error comparable to the approximation error. Using this insight, we show, practice, a significant reduction of memory usage and improvement in throughput, without sacrificing accuracy, solving critical issues from standard neural net- bet-based training methods. Through extensive experiments on different state-of-the- neural operators, datasets, and hardware, we demonstrate that approach reduces GPU memory usage by up to 50% with little or no reduction in accuracy.
Overall. calf-precision FNO makes it possible train significantly larger data points with the same batch size. Going forward, we plan to apply this to real-world applications that require super-resolution enable larger-scale training.

---

Published conference paper at ICLR 2024
ACKNOWLEDGMENTS
AA is supported by the Bren Foundation and the Schmidt Sciences through the AI 2050 senior fellow program. This project and GP were supported, in part, by the Canada Foundation for Innovation JELF grant, NSERC Discovery grant, AWS Machine Learning Research Award (MLRA), Facebook Faculty Research Award, Google Scholar Research Award, and VMware Early Career Faculty Grant. We would like to express our sincere thanks Nikola Kovachki for insightful suggestions and discussions, which were invaluable for completing this work. We thank members of the EcoSystem lab, especially Shang Wang, for their constructive feedback.
REFERENCES
Jonas Adler and Ozan Öktem. Solving ill-posed inverse problems using iterative deep neural networks. Inverse Problems, 33(12):124007 2017.
Syed R Ahmed, Ramm, and Gunter Faltin. Some salient features of the time-averaged ground vehicle wake. SAE transactions, pp. 473-503, 1984.
Saakaar Bhatnagar, Yaser Afshar, Shaowu Pan, Karthik Duraisamy, and Shailendra Kaushik. Predic- tion of aerodynamic flow fields using convolutional neural networks. Computational Mechanics, 64:525-545, 2019.
Boris Bonev, Thorsten Kurth, Christian Hundt, Jaideep Pathak, Maximilian Baust, Karthik Kashinath, and Anima Anandkumar. Spherical fourier neural operators: Learning stable dynamics on the sphere. In Proceedings of the International Conference Machine Learning (ICML), 2023,
Gary Chandler and Rich R Kerswell. Invariant recurrent solutions embedded in turbulent two- dimensional kolmogorov flow. Journal of Fluid Mechanics, 722:554-595, 2013.
Angel Chang, Thomas Funkhouser, Leonidas Guibas, Pat Hanrahan, Qixing Huang, Zimo Li, Silvio Savarese, Manolis Savva, Shuran Song, Hao Su, et al. Shapenet: An information-rich 3d model repository. arXiv preprint arXiv:1512.03012 2015.
Christopher De Sa, Megan Leszezynski, Jian Zhang, Alana Marzoev, Christopher R Aberger, Kunle Olukotun, and Christopher Ré. High-accuracy low-precision training. arXiv preprint arXiv:1803.03383 2018.
Jeffrey Dean, Greg Corrado, Rajat Monga, Kai Chen, Matthieu Devin, Mark Mao, Marc' aureloo Ranzato, Andrew Senior, Paul Tucker, Ke Yang, et al. Large scale distributed deep networks. Advances in neural information processing systems, 25, 2012.
Winfried van den Dool, Tijmen Blankevoort, Max Welling, and Yuki Asano, Efficient neural pde-solvers using quantization aware training. arXiv preprint arXiv:2308. 07350, 2023.
J.R. Driscoll and D.M. Healy. Computing fourier transforms and convolutions on the 2- sphere. Advances in Applied Mathematics, 15:202-250, 6 1994. ISSN 01968858. doi: 10. 10.1006/aama. 1994. URL pii/S01968858888
e elsevier http://ww....eeeeee..o.eerrrreoo wtt.:/trie/eeeeeeeeeeeeeeeeeeeee
Abhinay Goel, Caleb Tung, Yung-Hsiang Lu, and George K Thiruvathukal. A survey of methods for low-power deep learning and computer vision. In 2020 IEEE 6th World Forum Internet of Things (WF-loT), pp. 1-6. IEEE, 2020.
Vignesh Gopakumar, Stanislas Pamela, Lorenzo Zanisi, Zongyi Li, Anima Anandkumar, and MAST Team. Fourier neural operator for plasma modelling. arXiv preprint arXiv:2302.06542, 2023.
John Guibas, Morteza Mardani, Zongyi Li, Andrew Tao, Anima Anandkumar, and Bryan Catanzaro. Adaptive fourier neural operators: Efficient token mixers for transformers. In Proceedings of the International Conference Learning Representations (ICLR), 2021.
10

---

Published conference paper at ICLR 2024
Gaurav Gupta, Xiongye Xiao, and Paul Bogdan. Multiwavelet-base operator learning for differential equations. Proceedings of the Annual Conference on Neural Information Processing Systems (NeurIPS), 34:24048-2400 2021.
Joel Hayford, Jacob Goldman- -eetzer,, Eric Wang, and Lu Lu. Speeding up and reducing memory usage for scientific machine learning via mixed precision. arXiv preprint arXiv:2401. 1tti::2401.6664 16645, 2024.
Xianyan Jia, Shutao Song, Wei He, Yangzihao Wang, Haidong Rong, Feihu Zhou, Liqiang Xie, Zhenyu Guo, Yuanzhou Yang, Liwei Yu, al. Highly scalable deep learning training system with mixed-precision Training imagenet in four minutes. arXiv preprint arXiv: 180..1107..... 2018.
Dhiraj Kalamkar, Dheevatsa Mudigere, Naveen Mellempudi, Dipankar Das, Kunal Banerjee, Sasikanth Avancha, Dharma Teja Vooturi, Nataraj Jammalamadaka, Jianyu Huang, Hector Yuen, et al. study of bfloat for deep learning training. arXiv preprint arXiv: I 1905. 10009995522222 /2322, 2019.
Tamara G Kolda and Brett W Bader. Tensor decompositions and applications. SIAM review, 51(3): 455-500, 2009.
Jean Kossaifi, Yannis Panagakis, Anima Anandkumar, and Maja Pantic. Tensorly: tensor learning in python. The Journal of Machine Learning Research, 20(1):925-930, 2019.
Jean Kossaifi, Nikola Kovachki, Kamyar Azizzadenesheli, and Anima Anandkumar. Multi-gria tensorized fourier neural operator for high-resolution pdes. arXiv preprint arXiv:2310.00120 2023.
Nikola Kovachki, Zongyi Li, Burigede Liu, Kamyar Azizzadenesheli, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Neural operator: Learning maps between function spaces with applications to pdes. Journal of Machine Learning Research, 24(89):1-97, 2023. URL html.
http: / httmmll.orr//oooeemooooooooeeoee wttp://paprs///////111-111
Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Neural operator: Graph kernel network for partial differential equations. arXiv preprint arXiv:2003.03485, 2020a.
Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Andrew Stuart, Kaushik Bhattacharya, and Anima Anandkumar. Multipole graph neural operator for parametric partial differential equations. Proceedings of the Annual Conference on Neural Information Processing Systems (NeurIPS), 0066755567666, 2020b.
Zongyi Li, Nikola Borislavov Kovachki. Kamyar Azizzadenesheli, Kaushik Bhattacharya, Andrew Stuart, Anima Anandkumar, et al. Fourier neural operator for parametric partial differential equations. In Proceedings of the International Conference on Learning Representations (ICLR), 202
Zongyi Li, Hongkai Zheng, Nikola Kovachki, David Jin, Haoxuan Chen, Burigede Liu, Kamyar Azizzadenesheli, and Anima Anandkumar. -nnfores- neural operator for learning partial differential equations. arXiv preprint arXiv:2111.03399 2021b.
Zongyi Li, Nikola Borislavov Kovachki, Chris Choy, Boyi Li, Jean Kossaifi, Shourya Prakash Otta, Mohammad Amin Nabian, Maximilian Stadler, Christian Hundt, Kamyar Arizzadenesheli, and Anima Anandkumar. Geometry-informer neural operator for large-scale 3d pdes. arXiv preprint arXiv:2309. 00583, 2023.
Burigede Liu, Nikola Kovachki, Zongyi Li, Kamyar Azizzadenesheli, Anima Anandkumar, An- drew M Stuart, and Kaushik Bhattacharya. A learning-based multiscale method and its application inelastic impact problems. Journal of the Mechanics and Physics of Solids, 158:104665 2022.
Lu Lu, Pengzhan Jin, and George Em Karniadakis. Deeponet: Learning nonlinear operators for identifying differential equations based on the universal approximation theorem of operators. arXiv preprint arXiv:1910.00 1910.03 193, 2019.

---

Published conference paper at ICLR 2024
Lu Lu, Pengzhan Jin, Guofei Pang, Zhongqiang Zhang, and George Em Karniadakis. Learning nonlinear operators via deeponet based on the universal approximation theorem of operators. Nature machine intelligence, 3(3):218-229, 229, 2021.
Paulius Micikevicius, Sharan Narang, Jonah Alben, Gregory Diamos, Erich Elsen, David Garcia, Boris Ginsburg, Michael Houston, Oleksii Kuchaiev, Ganesh Venkatesh, et al. Mixed precision training. arXiv preprint arXiv:/710.037440 2017.
Paulius Micikevicius, Dusan Stosie, Neil Burgess, Marius Cornea, Pradeep Dubey, Richard Grisenth- waite, Sangwon Ha, Alexander Heinecke, Patrick Judd, John Kamalu, et al. Fp8 formats for deep learning. arXiv preprint 22Xxv:22000.0433, 2022.
Yannis Panagakis, Jean Kossaifi, Grigorios G. Chrysos, James Oldfield. Mihalis A. Nicolaou, Anima Anandkumar, and Stefanos Zafeiriou. Tensor methods computer vision and deep learning. Proceedings of the IEEE, 109(5):863-888) 2021.
Yannis Panagakis, Jean Kossaifi, Grigorios G. Chrysos, James Oldfield, Taylor Patti, Mihalis A. Nicolaou, Anima Anandkumar, and Stefanos Zafeiriou. Chapter 15 tensor methods in deep learning. In Paulo S.R. Diniz (ed.), Signal Processing and Machine Learning Theory, pp. 1009-1048. Academic Press, 2024. ISBN
0rt.:/...r1/00.11111011 doi: https://oi.org//1... B978-0-32-391172--8.. -0002-1. URL hetps::/www. www.:wscieneiieeeeoo wttpe//re....cccc.ceeeeeeeeeeeee www.ssiie/neeaecoooo.cocco....oe
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high-performance deep learning library. In Proceedings of the Annual Conference on Neural Information Processing Systems (NeurIPSS, 2019.
Jaideep Pathak, Shashank Subramanian, Peter Harrington, Sanjeev Raja, Ashesh Chattopadhyay, Morteza Mardani, Thorsten Kurth, David Hall, Zongyi Li, Kamyar Azizzadenesheli, et al. Fourcast- net: A global data-driven high-resolution weather model using adaptive fourier neural operators. arXiv preprint arXiv:2202. 11214, 2022.
Mariam Rakka, Mohammed Fouda, Pramod Khargonekar, and Fadi Kurdahi. Mixed-precision neural networks: A survey, arXiv preprint arXiv:2208.06000 2022.
Peter Renn, Cong Wang, Sahin Lale, Zongyi Li, Anima Anandkumar, and Morteza Gharib. Fore- casting subcritical cylinder wakes with fourier neural operators. arXiv preprint arXiv:2301.08299 2023.
Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. Nassir Navab, Joachim Hornegger, William M. Wells, and Alejandro Frangi (eds.), Medical Image Computing and Computer-Assisted Intervention MICCAI 2015, pp. 2,1-241, Cham, 2015. Springer International Publishing. ISBN 978-3-319-24574-4,
Tapio Schneider, João Teixeira, Christopher S Bretherton, Florent Brient, Kyle G Pressel, Christoph Schär, and Pier Siebesma. Climate goals and computing the future of clouds. Nature Climate Change, 701):3-,,200
Gilbert Strang. Computational science and engineering. Optimization, 551(663):571-586 2007.
Lloyd N Trefethen. Spectral methods in MATLAB. SIAM, 2000.
Nobuyuki Umetani and Bernd Bickel. Learning three-dimensional flow for interactive aerodynamic design. ACM Transactions on Graphics (TOG), 37(4):1-10, 2018.
Gege Wen, Zongyi Li, Qirui Long, Kamyar Azizzadenesheli, Anima Anandkumar, and Sally M Benson. Accelerating carbon capture and storage modeling using fourier neural operators. arXiv preprint arXiv:22100 17051, 2022.
David C Wilcox et al. Turbulence modeling for CFD, volume 2. DCW industries La Canada, CA. 1998.
12

---

Published conference paper at ICLR 2024
Jiawei Zhao, Steve Dai, Rangharajan Venkatesan, Brian Zimmer, Mustafa Ali, Ming- Yu Liu, Brucek Khailany, William J Dally, and Anima Anandkumar. Lns-madam: Low- Low-precision training in logarithmic number system using multiplicative weight update. IEEE Transactions on Computers, 71(12):3179-3190 2022.
13

---

Published conference paper at ICLR 2024
A ADDITIONAL DETAILS FROM SECTION
In this section, we give the full proofs, details, and discussions from Section 3.
First, for convenience, we restate the definition of the discretization error. Let denote the closed unit cube [0, 1J for dimension d E N, subdivided into n m cubes QQ..... Qn with sidelength 1/m. For each j. let ç denote the point with the minimum value in each dimension, in Q, Let V: - D denote an intermediate function within the FNO. Recall from the previous section that the primary operation in FNO the Fourier convolution operator, (KCx)(2)) We say the discretization error of F(v) is the absolute difference between the Fourier transform of v. and the discrete Fourier {..
Formally, given Fourier transform of via discretization of on = ((QI,.... Qn): ; (E,,...,6..)
basis function (x) = 2m)(wa))
Disc(e, 2,w) = 1/ v(r))u(z)dx - Σ (((S)
j=1
A.1 DETAILS FOR RESOLUTION AND PRECISION BOUNDS
In this section, we give the full details for Theorem 3.1 and Theorem 3.2.
Theorem 3.1 (restated). For any M > 0 and L 2 let K. C C(D) be the set of L-Lisschitz functions. bounded by ||o||00 S M. Then for all n, Q, there exist w. C1s c2 such that
viv. Mn 3/2 S sup (Disc(r, 2,w)) S e2va(1ul+ + L)Mn
sek
Proof. Let (x) sin(2r(w,x)) denote the real part of the Fourier base of frequency w. Define
Disc (v, Q,w) = / (()))((()) - Σ (e(
(3)
Recall that as defined above, Q = ((Q...... Q.). fE....., En }), where the n = mdd hypercubes of
side length 1/m subdivide D. First prove the upper bound. We have:
14

---

Published conference paper at ICLR 2024
/ ((1)) Σ
Disc Q,)
(S))
(4)
(())
- Σ 5=1
(5)
D
Σ 9q
(u(x) - ((S)) da
(6)
Σ
(u())(<L(()) - +L(E))) + (v(x) - =(E))) dæ
(7)
/=1
Q
(|o(x)) 1PC((=) (E,E) + 12(z) x(E,)) (P dx
(8)
j=1
Vd
Vd . 1 dx
Σ
M
|w| m
+
(9)
j=1 Q1
Σ /
Vd
(MWl + D) dr
(10)
==1 R
n
md
(M + L) 1/u
(11)
(M + n//
(12)
Bounding Disc w), the complex part of the Fourier base, follows identical argument. Setting e = 2 concludes the proof of the upper bound, which is true for all w.
For the lower bound, recall that D [0, 1] and that for each j. Ej
the minimum value in each dimension, in Q2. Setting v(z) = æ1 ... Fd and w = 1, we will show that
S o(x) sin sin(2rx))de )xr Σ v(t)) sin sin (2
d
nn 1/24
3. 2d
First, in one dimension, So' σ1 sin(2Tx))dez 2= Therefore, for dimensions, we have
d J. u(x) sin(2rx))d (f, xn sin (2 )d21 =(1m)
15

---

Published conference paper at ICLR 2024
Next, let & denote the kth component of E,. for S and S $
We have
n Σ (w(Es) sin (2
v (La sin (i (in(2(mW(m))
(13)
md
==1
j. sin (2x m)
(14)
m
m
10
Σ j=1
(j
sin 2 (2Tj/mm)) /m])
(15)
m
2d
cot
(16)
m
d
2-d
-d
(17)
π
m
(m +
yd
(2T)
3m
(18)
m--
1 d π2
z (2T) d
m
(19)
3
(20)
Finally, since n md we have
x²
d
(2T) (2T) (2T)-
(1 1+
3. 2d n 1/28
which concludes the proof.
Note that the upper bound is true for all while the lower bound is shown for 1. It is an interesting question for future work to show the lower bound for w > 1.
Now bound the precision error of F(v). Recall from Section that the precision error of F(v) is the absolute difference between F(v) and Flv), computing the discrete Fourier transform in half precision. Specifically, we define an (ag. Preeeciiionn system as mapping q R - S for the set (0) U {ao(1 c> ) u1-a((1+ + <) !==0 such that for all R, q(x) ar gminnesll y). )
This represents simplified version of the true mapping used by Python from R to float 32 or float16. Recall that these datatypes allocate some number of bits for the mantissa and for the exponent. Then, given a real number æ 29, its value in float.32 or float1 6 would be (z + ()) 2(x++e))
2 62æ 2 + 6 2" + 61622 (1 + 9/2 22)x= 2", Therefore, our definition above is simplified but reasonable approximation of floating -pinnt arithmetic.
Now, we define
Prec(e, Q.g.w) = Σ (())) Σ (u(L)))((
Now we bound the precision error of F(u).
Theorem 3.2 (restated). For any M > 0 and L 2 let K. c. C(D) be the set of L-Lipschitz functions, bounded by ||oi]oo < M. Furthermore let be an (0o, T)-precision system. Then for all n, Q w. there exists > 0 such that
sup (Prec(u, Q. S eM.
VEA
16

---

Published conference paper at ICLR 2024
Proof. Let (z) sin(27(w,x)) 2)) denote the real part of the Fourier base of frequency w. Define
m
Prec" (v, Q.q.w)
Σ
(E)))))))))
Σ (lo(E)))
(21)
We prove the upper bound as follows. We have
Prec"(e, Q. w)
(o(S)) ((15))))))(
(L(5)(()()) - ((E(())) (o(E) (
([e(E)) (P((E)) 9 ((E -
n
(M (e- () + (--
n j=1
nn (2cM)
n
M.M.
Bounding P Pre Q, q, w). the complex part of the Fourier base, follows an identical argument. Setting c 4 concludes the proof.
A.2 GENERAL BOUNDS
Now, we give results similar to Theorem 3.1 and Theorem 3.2, but with a general function f, rather than for F(v), function times the Fourier basis function.
Theorem A.1. For any M > 0 and L 2 let K. C(D) be the set L-Lipschitr,ffunciooss bounded by ll I M. Then for all Q,w, we have
2-1+1 d. n 1/4 < sup (Disc(f. Q,w)) & Lvd Lvd-n"" 4uu
jEK'
Proof. First, we define
1 ((od)d Σ s())
Disc'(f.Q.u)
(22)
Recall that Q (fQ....., Qn P. (Ei,..., En.) where the n = md hypercubes of side length 1/m
subdivide D. Now we prove the upper bound. We have:
17

---

Published conference paper at ICLR 2024
(f,,"():Q,W) =
f(zdda Σ Q
(23)
J [ac)da -E
T1
[(E))do
(24)
$=1 Q,
Σ
(f(x) f(Ex)) dx
(25)
1=1
Qg
1f(x) - f(E5) dr
(26)
jml
vd
$
L m
dx
(27)
j=l
Qg
L dd
nl
md
(28)
LVd n
(29)
This concludes the proof of the upper bound.
For the lower bound, recall that for each j. Ej is set to the minimum value in each dimension in Qj. We set (x) I1 22 Fdd-
First, in one dimension, To x xıdrı
2. Therefore, for d dimensions, we have
J, f(z)do (%, 21dz1) =(2)-d
Next, let Ej, denote the kth component of E,. for s k S d and 1 S I n. We have
(((E)));))
1
II
(30)
md
k=
j
(31)
md
m(m 1)
d
m
2d
(32)
| 1
2-d
(33)
m
> 2 d (1 2dm -1)
(34)
(35)
Finally, since n md have
22-d 2- (1+ + 2dm 2dm= 2-1+1 d. 1/4
which concludes the proof.
Now similarly bound the precision enror of f.
18

---

Published conference paper at ICLR 2024
Theorem A.2. For any M > 0 and L 2 llet C C(D) be the set L-ipschhhttt bounded by || < M. Furthermore let q be an (0o, pprecision system. Then for all n, Q, w, there exists C > such that
1/1 eM S sup (Preclf, 0, S eM.
SEK
Proof. Define
f(())10)) Σ (((E)))
(36)
We prove the upper bound as follows. We have
Prec" w)) =
1E) - Σ aU(EM)
5=1
(S(E)) (((())))
((E)) a(f(E)))
" n (cM)
cM
This concludes the proof the upper bound.
For the lower bound, given the definition of an (ao, e, -rrrciiio system, it follows that there exists y such that 4122 < y < M and 1/2 ey < ly g(y)). Then, for the lower bound we set f(x) = y. and
we have
Prec" (f.Q,q,w))
((5)) Σ (((E)))))
(f(E)) {(f(E))))
Σ 1 cy
=
2 cM
This concludes the proof.
A.3 PLOTTING THEORETICAL BOUNDS
In this section, we run additional experiments by plotting theoretical bounds alongside the true empirical discretization and resolution errors of the Darey flow dataset. Specifically, we plot both our bounds assuming Fourier basis (Theorem 3.1, Theorem 3.2) and our bounds for general functions (Theorem A. Theorem A.2) compared to the true Darcy flow dataset, measured after 10 epochs, the start of the FNO block (just before the forward FFT). The true errors calculated using the
19

---

Published conference paper at ICLR 2024
Darcy Hlow Error Analysis, Fourier Bound
Analysis,
3071
Entr
10"
LB)
1hm 3.1 (Res LB)
Thm A.l [Res UB)
Thm 3.1 (Res UB)
1hm A-2 (Prec
307
(Prec
True Disc error
Tue Prec arrar
True Prez error
Thm A.2 (Prec LB)
Darcy Analysis, 30, Bound
Analysis. 3D,
Thm 3.1 (Res LB)
Thm (Res UB)
True Disc emror
Irue Prec error
E11
Thm A.1 [Res LB)
UB)
Thm A.2 |Prec UB)
True Disc error
Thm A2 (Proc LB)
Discretization and precision errors of Darcy flow ID and 3D, compared to our Fourier basis and general bounds. Note that our bounds are worst-case bounds (we showed there exists function from the class of bounded L-Lipschitz functions which achieves at least the desired error), meaning that the true Darcy flow error may be lower than the worst-case bounds.
Figure 7:
definitions of precision and discretization error in the previous section, and we use the true difference precision between float 32 and floatll 1 6 when computing the precision error. See Figure We find that the discretization error is higher than the precision error, as expected by our theory. Furthermore, the Darcy flow discretization and precision errors are always lower than their respective upper and lower bounds. Note that the lower bounds in our theorems are worst-case lower bounds (we showed there exists function from the class of bounded L--Lipschitz functions which achieves at least the desired error), which means that the true errors can be lower than the bounds, which is the case in Figure
B ADDITIONAL DETAILS FROM SECTION 4
In this section, give additional details the datasets and experimental setup, as well additional experiments.
B. SHMED-BODY CFD EXPERIMENT
Similar to Shape-Net Car, we train the coometry-Informed Neural Operator on the Ahmed-boue dataset, using the same hyperparameters the original implementation (Li et al., 2023). The memory, error, and throughput results for thmed-body are appended to Figure We add the L2 error curves here following the format of Figure
20

---

Published conference paper at ICLR 2024
L2 Loss vs. Epoch, Ahmed Body
Mixed FNO (Ours)
Full FNO
Average Difference: 0.01
0.14
L2 Loss
Final Error Difference: 8.526e-04 Memory Reduction: 38.4%
0.08
Epoch
Figure 8: Test L2 error curves for GINO the Ahmed-body CFD dataset on three random seeds and standard deviation as error bars.
B.2 DATASET DETAILS
Navier-Stokes. We consider the 2D Navier-Stokes equation for a viscous, incompressible fluid in vorticity form on the unit torus T _ (0,
1
daw V ④
Aw + f f T2, te (0,
Re
Ad = w, [ - for LE (0,
(37)
where E LL is a forcing function and Re is the Reynolds number, for the initial condition (0,-) = = 0. The goal is to learn the non-lineat operator Gf : f + (T,,) for T = 5 and Re 500. We use the dataset from prior work (Kossaifi et al., 2023), which sets Re 500 and generates forcing functions from the Gaussian measure, (V 27(-A + 9/)- and computes the solution function via a pseudo-spectral solver (Chandler & Kerswell, 2013). There are 10 000 training and 2000 test samples with resolution 128 128.
Darcy Flow. We consider the steady-state 2D Darcy Flow equation, which models fluid flow through a porous medium. Given D (0, 1)2, the diffusion coefficient E L (); and the forcing function f E ((0 1)2; R), the pressure u satisfies
V (a (a(z) (u(x)) = f(x), for x E D
(38)
u(x) 0, for E OD
(39)
We seek to learn the non-linear operator g+ a + u, the mapping from the diffusion coefficient a to the solution function u. We use the dataset from prior work (Li et al., 2021a), which fixes f 1 and generates 5000 training and 1000 test samples with resolution 128 x 128.
Spherical Shallow Water Equations. The spherical Shallow Water Equations on the rotating sphere are model for variety of geophysical flow phenomena such atmospheric flows, tsunamis and storm surges. We consider the evolution of two state variables yp, u (geopotential height and tangential velocity of the fluid column), governed by the equations
0:p + V -(pu) 0, for E S², E (0,
(40)
(p()))) F=F=S.
F- for f (0,
(41)
21

---

Published conference paper at ICLR 2024
for initial conditions (p(, So. u(-,0) uo, where S is the unit sphere, F the momentum flux tensor with Fd you'n" u2 + 7 30² and S a source term 5 = -22r (ypu), which models the Coriolis force due to the rotation of the sphere with angular velocity Ω. We use the dataset from prior work (Bonev al., 2023), which generates random initial conditions the sphere resolution 256 x 512 and solves the PDE using the spectral solver in the corch-harmon 2023).
package (Bonev et al.,
Shape-Net Car We evaluate on 3D real-worid car dataset generated by prior work (Umetani & Bickel, 2018; Li et al., 2023) using the Reynolds Averaged Navier Stokes (RANS) equations (Wilcox al., 1998). Each data point consists of mesh points representing unique 3D car, and the goal is to predict the full 3D pressure field, given an inlet velocity of 20m/s.
Including another 3D real-world dataset for evaluation, we the dataset from (Li al., 2023), which is based on Ahmed-body shapes from (Ahmed al., 1984), for which steady-state aerodynamic simulations were run using the GPU-accelerated Open-FOAM solver, with varying inlet velocity from 10m/s to 70m/s. The input and output data formats are similar those of Shape-Net Car, but Ahmed- Mimded-bdyy has many more mesh points (100k) on the surface.
Ahmed-Body CFD
B.3 MODEL DETAILS
We the official implementation of each model, with the default hyperparameters, unless otherwise specified. Furthermore, we release all of our code at https: //gtt///ggttbbb//// com/ wwwuraloperato/neeeee howralopeaatooo...
For the case of Shape-Net Car and Ahmed Body on GINO, we note that the default batch size 1, in fact, the only allowed batch size is 1, since each car geometry is unique. As shown in Figure 1, throughput is identical (in fact, our memory saving is enough to have batch size 2, but this is not possible). Although each individual Shape-Net Car fits in memory, other irregular-goometry datasets are too large to fit in memory, so sub-sampling is used (Li er al., 2023). Therefore, mixed-precision GINO allows sampling significantly more mesh points than the full-precision default.
B.4 FNO PROFILING
Using Nvidia's ' RTX 3090 Tí, we start by profiling full-precision FNO training on the Navier-Stokes dataset with the PyTorch profiler. Figure 9 provides an overview of runtime breakdown by Py PyTorch modules and by individual GPU kernels. We observe that complex-valued operations are accountable for the majority of the runtime.
B.5 NUMERICAL STABILITY: POST-FORWARD METHOOS
Most commonly, numerical underflows and overflows that occur during mixed-precision training handled after the forward pass using loss scaling, gradient clipping, and other methods (Micikevicius al., 2017). For experimentation, we implemented the following: (1) automatic loss scaling via AMP: (2) gradient clipping to a default value (5.0); (3) delayed updates via gradient accumulation every 3 batches. These methods are compared to our tanh approach and no-stabilizer baseline on the Darcy Flow experiment setting. In Figure 10, show that all three baselines diverge during the first epoch of training, while tanh remains stable numerically. It is noteworthy that while loss scaling almost does not diverge, its scale decreases drastically with each update and becomes infinitesimal. The common problem for global stabilizers is that they do not directly address the numerical overflow of forward FFT operations within each FNO block. In the next subsection, we discuss local methods that acts on the FNO block itself.
B.6 NUMERICAL STABILITY: FFE-FFT METHODS
Here, we give additional details and experiments for our study on numerical stability of mixed- precision FNO.
hhtpps:////t thub. httpp////sstttt.com///reereeeerr
22

---

Published conference paper at ICLR 2024
Total Time
?
kernel-128,
6.4%
at: nativeeeggu kernel, Impl-at native::dr.... megeeeee.kkttttttteee..rr.
7.7%
13. 13.8%
void attrnativeveecoorized elementwise
kemelc4, Sennaiee.GelBBaawaarrlloo.
9.1%
sopere_geemm_com 64432_ ampere_cgemme
64x32_m
ampere secere_cemmmm__m
htt___eem____xxx_____ 64x64_nt _nnt
9.24
void at:nattee:ollmmettttt. kernel<128....
11%
void at:raatieevvcctorized ellmeetttieeeee____....
10%
void at:rative.vectorized elementwise....
10.49
void at:retive.vectorized elementwise_...
void vector 2d II-128u, 128u, 16u, 16u...
Figure 9: For Navier-Stokes, we use the Py Torch profiler to visualize the runtime breakdown of FNO training by both modules (top) and GPU kernels (bottom). Spectral convolutions (complex-valued tensor contraction) accounts for most of the runtime, and correspondingly. kernels associated with complex numbers are amongst the most costly kernels.
Norms of Activations for Global Stablizers
5000
Norm Activation
Scale
3000
2000
Scate llog
20
Grad 0r0d staror
Minibatch
600
Scale Updates
33 90 60
75 16
Figure 10: All three global methods diverge during the first epoch of training, with loss scaling completing the most batches without NaN values (Left). However, the scale quickly becomes infinitesimal during training (Right).
simple idea for addressing overflow in FP16 is to scale down all values by adding fixed pointwise division operation before the FFT. However, this is an sub-optimal solution because all the values are scaled down equally, removing the data outliers but simultaneously squashing normal data. This forces all numbers into very small range, which half precision cannot distinguish, preventing the model from converging to an acceptable performance. We find that accuracy deteriorates as the division factor grows: 1 102, 10 At 10 the network completely fails learn due vanishing gradients. We also find that adding normalization layer before the Fourier transform (essentially equivalent to scaling) does not fix the numerical instability in the long term. Additionally, normalization layers present extra GPU memory overhead because their statistics have to be tracked for gradient computation.
23

---

Published conference paper at ICLR 2024
Fourier
AMP (T/F)
Runtime per epoch
Train
Precision
Activation
loss
Full
N/A
44.4
0.0457
Full
T
N/A
42.3
0.0457
Half
F
N/A
N/A
N/A
Half
hard-clip
37.1
0.0483
Half
20-clip
40.0
0.0474
Half
tanh
36.5
0.0481
Table 3: Comparison different pre-activation functions used for numerical stability. tanh achieves the fastest runtime without significantly compromising loss. Since focus is numerical stability, we compare only the average train loss over the 5 final epochs (we compare the test losses in Section 4.4).
Magnitude (1aft) and Phase (ighh)] Difference from Pre-Activation
Magnitude (leff) and Phase (Right) Change from Pre-Activation mean: 5.2005
meare 5 15
meen:-1 20c00
0414
meat: -5200 64
A/85
1000
4
46
FRouendis
9/85
A4
-024
79
20
Figure 11: Impact of pre-activation frequency domain signals, for fully-trained full FNO model (Left) and a mixed FNO model (Right), in terms of mean magnitude and phase differences, on minibatch of the Navier-Stokes dataset. Pre-activation only changes extremely small fraction of frequencies, and the two signals are well-aligned in phase.
Appendix B.6 compares several pre-FFT stabilization methods that address numerical instability with acceptaale computational costs. hard-clip forces maximum and minimum values in the input default range. 2a-clip similarly performs clipping but based the data's mean and standard deviation. Out of these methods, tanh is founded to be the adopting tanh as our pre-activation function, we assess its impact on the signal using trained model, as shown in Figure 11, We find that the loss of signal information is minimal in amplitude and phase.
After
B.7 FNO BLOCK PRECISION ABLATION STUDY
As shown in Figure 2, FNO block has three complex-valued operations that performed in half-precision: the forward FFT, the tensor contraction, and the inverse FFT. We investigate the effect of altering the precision setting of each FNO training. This creates total of eight settings since each operations have two precision configurations: full or half. In Table 4, we show error, runtime, and memory results from training for 30 epochs the Darcy Flow dataset for each setting. The empirical results demonstrates the acorss-the-boan advantage of our approach, which keeps all three operations in half half-precision in the forward pass.
B.8 TANH ABLATION STUDY
A natural question from using pre-activation is to ask how hyperbolic tangent would affect the full precision FNO model with no other changes. In Table 5, we answer this question by rumning the full precision FNO model the Navier-Stokes dataset. We find that there is no noticeable change in the test losses, and the cuntime-per-epore increases by 0.8 seconds, or 1.6%. This concludes that tanh reasonable and practical choice to improve numerical stability for half precision methods.
24

---

Published conference paper at ICLR 2024
Forward FFT
Contraction
Inverse FFT
Runtime per epoch
Memory
Train Error
(F/H)
(F/H)
(F/H)
(sec)
(MB)
(L2 loss)
F
F
17.06
8870
9.00
F
H
16.55
8908
8.71
F
F
17.11
8614
8.76
F
H F
16.96
8658
8.15
H
F
17.64
7824
9.13
H
H
16.81
8004
7,51
H
H
F
16.57
7558
8.75
H
H
15.63
7550
7.49
Table Performance each setting on Darcy Flow trained for epochs. The fully half-precision setting (our method) is advantageous across all the metrics. Note that the numerical stabilizer is only used when forward FFT half-precision.
Loss vs. Runtime, Darcy Flow
Mixed FNO (Ours), 26 Modes
Pull 16 Modes
(Curs),
Full PNO, 32 Modes
Mixed FNO (Curs), 64 Modes Modes
3.13
L2 Lo00
Pull PNO, 128 Modes
Mixed FNO (Curs), 128 Modes
0.12
Time Epoch (seconds)
Figure 12: Comparison of the full-precision and mixed-precision (AMP (AMP+HALF+TAN)) FNO with different frequency modes, on the Darcy flow dataset.
B.9 COMPARING THE H Loss
We present additional experiments similar to the results in Section 4, but using the H' metric instead of T.2 In Figure 12, we plot an ablation study on the number of frequency modes for Darcy flow, similar to Figure 14 with L loss. We see that the overall trends are similar, but test L loss is noisier than test L1 loss; this makes sense, because the training loss is H² In Figure 13, we plot the cotoomance-vs-tinee and Pareto frontier plots for Navier Stokes and Darcy flow, similar Figure with L loss. As before, we see largely the same trends as with H' loss, but the results are noisier. In Table 6, plot the final performance from the models run in Figure 5.
1. FREQUENCY MODE ABLATION STUDY
We run an experiment on synthetic data demonstrate that the error caused by half precision higher for higher frequencies, relative to the amplitude. We create a signal based on sine and cosine waves with frequencies from 10, with randomly drawn, exponentially decaying amplitudes. Then we
Table 5: Ablation study on full- precision FNO with and without tanh on the Navier-Stokes dataset. There is no noticeable change in accuracy, showing that t tanh is a practical choice to improve numerical stability low precision methods.
H'
L
time-per-epoch (sec)
Full precision
0.0121
0,00470
51.72
Full precision + tanh
0.0122
0.00465
52.56
25

---

Published conference paper at ICLR 2024
H1 Loss Wallclock Time, Nevier Stokes
Hi Loss Wallclock Time,, Dancy Flow
Vised Froddon
Hoed Predecen
Thme nuDeN)
(minutes))
L2 Loss vs. Runtime, Navier Stokes
L2 Loss vs. Runtime, Darcy Flow
FNO (urst,
Mixed FNO (Ours), OP
Full FNO. CP
LOS5
8.090
8612
Moad PNO (Curs), Dense (Curs), Full FNO, CP
Dense
1404
273
60.0
02.3
4.55
67.3
2000
1244
L2.
id4
24
14.
Time per Epoch (seconds)
Time per Epoch (seconds)
Figure 13: Test H error curves for FNO on the Navier-Stokes (top left) and Darcy flow (top right) datasets. Pareto frontier for FNO the Navier-Stee Stokes (bottom left) and Darcy flow (bottom right) datasets, for Canonical Cooonica-PPolyddcc factorization (CP) or no factorization (Dense). We train each model for 500 epochs, and plot the standard deviation across 3 trials for each setting.
Table 6: Results for FNO with full precision, mixed precision, and precision schedule. All results the average over seeds trained to 19 hours each (which is the time it takes full precision FNO to reach 500 epochs).
H
L
time pece-eer-ppcchh (sce)
Full FNO
.00536 = 2.1e
.00214
3.5e
121.4
Mixed FNO (Ours)
.00645 = 6.6e
.00212 =
80.2
Precision schedule (Ours)
.00515 = 8.3e
.00812 + 1.1e
80.2, 83.8, 121.4
plot the Fourier spectrum in full and half precision, well the absolute error of the half-precision spectrum as percentage of the true amplitudes. See Figure 15; we find that the percentage error exponentially increases. Since in real-world data, the energy is concentrated in the lowest frequency modes, and the higher frequency modes are truncated, this gives further justification for our half precision method.
B.11 OTHER NUMERIC SYSTEMS
We compare training curves of full FNO, our mixed FNO pipeline, and FNO trained with AMP using Brain Float 16 (BF16) on the Navier-Stokes dataset in Figure 16. addition to suffering from
Navier-Stokes
Darcy Flow
FNO TF32
57.44
14.12 13.52
Mixed FNO (Ours)
53.72
Table 7: Training time per epoch on a Nvidia A100 GPU.
26

---

Published conference paper at ICLR 2024
H1 Loss Runtime, Darcy Flow
Mixed (Ours),
Full FND, 16 Modes
Mixed PNO (Ours), 32 Modes
Mixed FNO (Ours), 64 Modes
0.0425
Full FNO, 64 Modes
Mixed FNO (Ours), 128 Modes
H1 LO55
FND, 128 Modes
0.000
0.0375
0.055
0.025
Time per Epoch (seconds)
Figure 14:
Comparison of the Full-FNO and Mixed-FNO with different frequency modes on the Darcy Flow dataset. For 16 frequency modes, the half precision error (compared full precision) is higher than for 32, 64, or 128 modes.
Frequency
66.
03
Broor fom no. oroddiinn frquuenyy medos, porcentage of tud:
Frequency
Figure 15: Synthetic signal (top), its frequency modes (middle), and the error due to half precision, percentage of the amplitude (bottom). The percentage error increases for higher frequencies.
27

---

Published conference paper at ICLR 2024
12 Loss vs. Epoch, Navier Stakess
12 Loss vs. Epech, Darcy Flow
Nixed RND (Durs)
Full IND
04
AMPIBF16)
Moxed Fho (Duns)
hull PND
FPB
Epoce
Spach
Figure 16: BF16 has significantly higher enror compared to the other two approaches, possibly due to reduced bits for precision compared to FP16. FP8 training diverges due to the absence of precision bits in FP16.
reduced precision, BF16 is not supported for most FFT-related operations in PyTorch, which makes its application on FNOs extremely difficult.
Furthermore, we compare the efficiency of our mixed precision approach with TensorFloat 32 (TF32) on an Nvidia A 100 GPU, on which TF32 is highly optimized. We report the time per epoch when training on Navicr-Stokes and Darcy Flow datasets Table Our method still more efficient in comparison.
In addition, we simulate 8-bit floating point (FP8) training using our method on the Darey Flow dataset via clipping out-of-range values to the maximum and minimum representable under the E5M2 format, which has a higher dynamic range than the E4M3 format (Micikevicius et al., 2022). However, the GINO network diverges during training due to less available precision bits shown in Figure 16.
Furthermore, the worse result of FP8, is actually predicted Theorem 3.2: while the precision constant for FP16 corresponds to c = 10-4 corresponds to c > 10-2 for FP8 (even when using E4M3), and we can no longer claim that the precision error is lower than the resolution error.
B. ABLATION STUDY ON CONTRACT AND VIEW-AS-REAL USAGE
This ablation study shows the benefit of (1) decomposition of cinsum into sub-equations and appro- priate view-as-real usage, (2) our einsum path eaching, (3) our greedy tensor contraction procedure, and (4) both weights and inputs in half-precision.
The results are follows:
1. The usage of view-as-real is different for different tensor contract implementations:
Option A: View-as-real on all tensors, and compute a single einsum one equation (naive);
Option View-as-real on two tensors a time, and compute step-by-step einsums with two-tensor sub-equations (optimized);
Option C: View-as-real only on high- -ish-iimensooall einsums, and contracting low- dimension sub-equations in complex format (ours, optimal).
Results on the Navier-Stokes be found in Table 8.
2. Re-computing contract paths (naive) vs. Caching contract paths (ours):
Since tensor shapes are static, we avoid repeated path computation in the default contract implementation. Path computation could take up to 75% of the cost of the contract operation in the forward pass, which try to avoid.
Results on the Navier-Stokes and Darcy Flow datasets can be found in Table 9. Our method reduces time to calculate paths to almost zero during each iteration of einsum in FNOs.
3. Opt-einsum's FLOP-optimal path (naive) vs. memory-greedy path (ours):
By default, opt-einsum minimizes FLOPs for the contraction, but our method optimizes memory usage allow higher-resolution PDEs. Results Ahmed-body and Shape-Net
28

---

Published conference paper at ICLR 2024
Table 8: Comparison of tensor contraction implementations the Navier-Stokes dataset
Metric
Option A
Option B
Option C, Ours
Reduction from Option A (%)
Time per Epoch (s)
1730
101.7
92.59
94.65
Memory Usage (MB)
10310
5048
4832
53.12
Contract Forward Time (s)
0.178
0.0084
0.0076
95.73
Table 9: Re-computing contract paths Caching contract paths
Dataset
Time to calculate paths (s) Time to compute einsums (s)
Path/einsum (%)
Navier-Stoker Stokes
0.000573
0.000751
76.3
Darcy Flow
0.000441
0.000716
61.6
Car could be found in Table 10. On cemori-nonsumingg 3D datasets, our method could save up to 12%.
4. Approximate only weights in half-precision (naive) vs. inputs and weights both half- precision (ours):
If only weights were in half half-precision, the memory reduction would have been significantly limited for the contract operation since the inputs are in full-precision. Results on Darcy be found in Table 1. Note that in the case of Navier-Stokes,
Flow and Navier-Stokes pytorch uses much less memory-efficiee kernel if inputs were in full- full-precision. Hence, is critical to cast both weights and inputs half- half-precision.
29

---

Published conference paper at ICLR 2024
Table 10: FLOP-optimal path vs. memory-greedy path
Dataset
Greedy Memory (MB) Default Memory (MB)
Reduction (%)
Shape-Net Car
7906
8662
8.72
Ahmed-body
11990
13610
11.91
Table 11: Approximate precision in weights inputs and weights Half-Pree Memory (MB) Inputs-Full Memory (MB)
Dataset
Reduction (%)
Darcy Flow
7550
8166
7.54
Navier-Stokes
4832
9380
48.46
30


Reviews:
[{"review_id": "16hqpwuhZg", "reviewer": "Reviewer_XXPr", "summary": "The paper introduces a mixed-precision training method for neural operators, focusing on Fourier Neural Operators (FNO) used in solving partial differential equations and other mappings between function spaces. The paper discusses the challenges of high-resolution training data, limited GPU memory, and long training times in the context of neural operators. It emphasizes the need for mixed-precision training to mitigate these issues. The paper demonstrates that, contrary to our expectations, mixed-precision training in FNO does not lead to significant accuracy degradation. It presents rigorous theoretical characterization of approximation and precision errors in FNO, highlighting that the precision error is comparable to the approximation error. The paper introduces a method for optimizing memory-intensive tensor contractions using mixed precision, reducing GPU memory usage by up to 50% and improving throughput by 58% while maintaining accuracy.\n\n These findings have the potential to advance the efficiency and scalability of neural operators in various downstream applications.", "questions": "The work is indeed very interested and provides a strong contribution to the research community. My questions are:\n\n1. Why is it that the standard mixed precision training used for training ConvNets and ViTs is ineffective for training Fourier Neural Operators?\n2. Similarly, why is is that the common solutions (loss scaling, gradient clipping, normalization, delaying updates) fail to address the numerical instability of mixed-precision FNO?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "The strengths of the paper are as follows:\n\n1. Mixed-Precision Training for Neural Operators: The paper introduces the first mixed-precision training routine tailored specifically for neural operators. This novel approach optimizes the memory-intensive tensor contraction operations in the spectral domain and incorporates the use of tanh pre-activations to address numerical instability. \n\n2. Theoretical Results; The paper provides a strong theoretical foundation for its work by characterizing the precision and discretization errors of the FNO block. It demonstrates that these errors are comparable and proves that, when executed correctly, mixed-precision training of neural operators does not lead to significant performance degradation. \n\n3. Experimental Evaluation: The paper conducts thorough empirical validation of its mixed-precision training approach on three state-of-the-art neural operators (TFNO, GINO, and SFNO) across four different datasets and GPUs. The results indicate that the method significantly reduces memory usage (using half the memory) and increases training throughput by up to 58% across various GPUs, all while maintaining high accuracy with less than 0.1% reduction. \n\n4. Open-Source Code: The paper provides an efficient implementation of its approach in PyTorch, making it open-source and providing all the necessary data to reproduce the results.", "weaknesses": "While the method of using tanh pre-activation before each FFT seems to avoid numerical instability, it would have been better if some theoretical justification was given (even for simplistic cases). I believe similar theory should have been provided for the learning rate schedule. However, the paper indeed contributes significantly in theoretical aspects (in Section 3) by characterizing the precision and discretization errors of the FNO block and showing that these errors are comparable. Given the strong theoretical contributions, the above is not significant. Moreover, a strong ablation study for tanh is provided in Appendix B.5", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "XPXJGHz94b", "reviewer": "Reviewer_dRFB", "summary": "In this paper the authors provide a way to apply mixed-precision training to Fourier Neural Operators. The authors point out that while some previous approaches have used mixed-precision training with FNO, they have done so only on the real-valued parameters (i.e, the linear weights, and biases) and not on the FFT/DFT that is applied to the input in each layer. \n\nThey first show that mixed-precision training should not cause extra error in the training due to the Fourier transform, since the discretization error from FFT→ DFT is orders of magnitude bigger than the DFT→ half-precision DFT. \n\nFor the application of mixed-precision on complex tensors (the output after an FFT/DFT application) the authors apply mixed-precision by converting each “tensors to real” (this is something that the authors don’t elaborate upon enough). \n\nThe authors show through their experiments, that they are indeed able to train various FNO based architectures with low-memory and hence are also able to increase the throughput of their runs while only taking a less than 1 percent hit in the performance as compared to non-mixed-precision baselines.\n\nThe authors also point to the use of a tanh based pre-activation that helps in mitigating the mixed-precision based overflow that usually occurs when trained FNO based architectures.", "questions": "The authors mention that neural operators are discretization convergent, however have not cited any relevant work around it. While empirically we know that for some FNO based architecture we can get zero shot super-resolution, are there any works that prove it? In general, adding relevant citation to that claim would be useful.", "limitations": "", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The paper is well-written and easy to follow.\n\nThe experimental results of the paper are impressive. The authors are able to increase the training throughput for navier stokes $1.41$ times, by achieving an almost 50 percent memory reduction. \n\nThe use of tanh activation as a pre-activation to the neural operator block is a good technique to reduce the overflow issue, and an important finding. \n\nThe authors provide a theoretical proof for why the error for the mixed-precision training would be negligible when compared to the discretization error for DFT, while the results there are pretty standard its a good addition to have.", "weaknesses": "I think that the primary methodology as to how the mixed-precision is applied to the Fourier Kernel is not clear. From the looks of it the mixed-precision approximation is applied to the weights in the complex domain (that are used in the kernal operator).\n\nHowever, this is something that seems to be different from what they say in the introduction, where they mention that they enable some mixed precision in the entire FNO block (which I assumed will try to also enabled mixed precision in the DFT algorithm). \n\nFrom the looks of it the primary contribution seems to be the addition of tanh, and the application of mixed precision in the complex domain by treating the real and the imaginary components as two distinct real tensors. Together, the overall contribution does not seem very novel in of itself.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "9nXn9Ng4z0", "reviewer": "Reviewer_NM9S", "summary": "The authors demonstrated mixed-precision training of FNOs on GPUs.  They derived theoretic bounds on the effect of rounding error, improved AMP to support complex arithmetic, addressed an instability by introducing a pre-activation thah, and conducted experiments on practical datasets.", "questions": "See minor requests above.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- Application is of significance.\n- Well written. \n- Code released.", "weaknesses": "> It would be helpful to include an illustration of empirical characterization of rounding error in direct comparison with the theoretical bound scaling law.   \n> Mixed-precision with FP8 has been proposed, how does the method fare with FP8?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "NMygitKrb0", "reviewer": "Reviewer_5SPZ", "summary": "The paper intends to reduce the cost of training and inference of FNO. The method used here include training by mixed-precision to reduce memory and time cost. The mixed-precision training technique proposed here is specialized for FNO, where bounds for precision error and approximation error are theoretically analyzed and guaranteed. The experiment results show that the proposed method greatly improves the speed and throughput of FNO while not sacrificing accuracy.", "questions": "None.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "**Originality:** The novelty of this paper lies in leveraging the analysis of precision error in Fourier transform, which successfully reduced the computational cost and maintained the accuracy of FNO.\n\n**Quality:** The quality of the paper is high in both theoretical analysis and experiment conduction.\n\n**Clarity:** The paper delivers its idea in a quite clear and detailed way.\n\n**Significance:** For people who intend to have a large scale of implementations of FNO, this work can be very interesting to them.", "weaknesses": "The major doubt for this paper is regarding the significance. While it is dedicated to improving FNO, an important architecture of neural operators, FNO is not the only one and arguably the best one. If FNO is yet to be a dominant model for large scale applications of neural operators, the significance of this work is limited.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)