# Guaranteed Approximation Bounds for
Mixed-Precision Neural Operators

Renbo Tu 14, Colin White∗2, Jean Kossaifi3, Boris Bonev3, Nikola Kovachki3,
Gennady Pekhimenko14, Kamyar Azizzadenesheli3, Anima Anandkumar2
1 University of Toronto, 2 Caltech, 3 NVIDIA, 4 Vector Institute Equal
contribution. Email to<EMAIL_ADDRESS>
###### 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.

##  1 Introduction

Real-world problems in science and engineering often involve solving systems
of partial differential equations (PDEs) (Strang, 2007). These problems
typically require a fine grid for numerical solvers to guarantee convergence,
e.g., climate modeling and 3D fluid dynamics simulations, and the full-scale
solutions to 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
a powerful _data-driven_ technique for solving PDEs (Li et al., 2020a; 2021a;
Kovachki et al., 2023; Lu et 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 in neural operators can be 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 a unique limit
under mesh refinement (Kovachki et al., 2023). By learning from discretized
data from input and solution functions, the trained models can 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 PDE-based problems with significant
speedups (Li et al., 2021a; 2023; Bonev et al., 2023; Kossaifi et 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 is complex-valued and not handled
well by the standard AMP, as seen in Figure 1.

Conventional wisdom may suggest that reducing precision in FNO without drastic
accuracy degradation is not feasible. The presence of the Fourier transform in
FNO appears to be a limitation in reducing precision since it suffers from
numerical instabilities on high-frequency modes under reduced precision.
Hence, numerical solvers, such as pseudo-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, we show both theoretically and empirically that this
is not the case for FNO and other operator learning approaches.

Figure 1:  Top: example data points from each dataset: Navier-Stokes, Darcy
Flow, Spherical Shallow Water, Shape-Net Car CFD, and Ahmed-body CFD. Bottom:
performance of our method compared to full-precision and AMP on each dataset.
For each dataset, we plot test error (y-axis) and GPU memory (x-axis), and we
annotate the maximum throughput (proportional to the area of each ball). All
data are measured on the same hardware (RTX 3090 Ti) and the same virtual
environment. Memory decreases by up to 50%, while $L^{2}$ loss increases by at
most 0.28%.

Our approach: In this work, we devise the first mixed-precision 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 is 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 is no need to run
the discrete Fourier transform in full precision. Building on this, we show
that, in a single FNO layer, the round-off error due to lower precision is
small. Therefore, the overall round-off error remains small since a full FNO
architecture is only made up of a few layers—much smaller than the number of
steps in a classical pseudo-spectral solver where errors accumulate and
explode.

We make the above intuitions concrete, and we develop the following
theoretical approximation bounds: we characterize the precision error and
resolution error of the FNO block, proving asymptotic bounds of
$n^{-\nicefrac{{2}}{{d}}}$ and $\epsilon$, respectively, for mesh size $n$,
dimension $d$, and dynamic range $\epsilon$. Therefore, we justify using mixed
precision since its error is comparable to the discretization error already
present in FNO.

Motivated by these theoretical results, we introduce a mixed-precision
training method by optimizing the memory-intensive 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, naïvely training neural operators in mixed precision
leads to prohibitive numerical instability due to overflows in the FNO block.
To address this issue, we study numerical 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 in standard packages designed for
real-valued networks.

We carry out extensive experiments to show a significant reduction in memory
usage and improved throughput without sacrificing accuracy; see Figure 1. We
consider three variants of FNO: tensorized (TFNO) (Kossaifi 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 or no reduction in accuracy.

Furthermore, we show that our 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; we 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 at https://github.com/neuraloperator/neuraloperator.

We summarize our main contributions below:

  * •

We introduce the first mixed-precision training routine for neural operators,
which optimizes the memory-intensive tensor contraction operations in the
spectral domain, and we propose using tanh pre-activations to 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 to little or no performance degradation.

  * •

We empirically verify the superiority of our mixed-precision 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 PyTorch, 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 a recent focus on using
machine learning-based methods to solve PDEs (Gupta et al., 2021; Bhatnagar 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; 2020b; Kovachki et 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 is trained, it 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 a highly successful neural operator (Li et al., 2021a; Gopakumar et al.,
2023; Renn et al., 2023; Wen et al., 2022). Let
$\mathcal{A}:\\{a:D_{\mathcal{A}}\rightarrow\mathbb{R}^{d_{\mathcal{A}}}\\}$
and
$\mathcal{U}:\\{u:D_{\mathcal{U}}\rightarrow\mathbb{R}^{d_{\mathcal{U}}}\\}$
denote the input and output function spaces, respectively. In this work, we
consider the case where $D_{\mathcal{A}}=D_{\mathcal{U}}\subset\mathbb{R}^{d}$
for $d\in\mathbb{N}$. Given a dataset of pairs of initial conditions and
solution functions $\\{a_{j},u_{j}\\}_{j=1}^{N}$, which are consistent with an
operator $\mathcal{G}(a_{j})=u_{j}$ for all $1\leq j\leq N$, the goal is to
learn a _neural operator_ $\mathcal{G}_{\theta}$ that approximates
$\mathcal{G}$. The primary operation in FNO is the Fourier convolution
operator, $(\mathcal{K}v_{t})(x)=\mathcal{F}^{-1}(R\cdot
T_{K}(\mathcal{F}v_{t}))(x),$ $\forall x\in D$, where $\mathcal{F}$ and
$\mathcal{F}^{-1}$ denote the Fourier transform and its inverse, $R$ denotes a
learnable transformation, $T_{K}$ denotes a truncation operation, and $v_{t}$
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 on discrete data.

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), TFNO computes a
tensor factorization (Kolda & Bader, 2009) of the weight tensors in the
spectral domain, acting as a 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 Geometry-Informed 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 to transform the irregular grid
inputs into and from regular latent grids on which FNO can be efficiently
applied (Li et al., 2020a).

#### Mixed-Precision Training.

Mixed-precision training of neural networks consists of reducing runtime and
memory usage by representing input tensors and weights (and performing
operations) at lower-than-standard precision. For example, PyTorch has a
built-in mixed-precision mode called automatic mixed precision (AMP), which
places all operations at float16 rather than float32, with the exception of
reduction operations, weight updates, normalization operations, and
specialized operations (Paszke et al., 2019). bfloat16 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 a 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-valued 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 our 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 a forward Fourier transform, and we give the
full details and discussion in Appendix A.

We start by formally defining the discretization error. Let $D$ denote the
closed unit hypercube $[0,1]^{d}$ for dimension $d\in\mathbb{N}$. Let
$Q_{1},\dots,Q_{n}$ denote the unique (up to $d$) partitioning of $D$, such
that each $Q_{j}$ is a hypercube with sidelength $\nicefrac{{1}}{{m}}$. For
each $1\leq j\leq n$, let $\xi_{j}\in Q_{j}$ denote the vertex that is closest
to the origin; formally, we have $Q_{j}=\prod_{k=1}^{d}[s_{j,k},t_{j,k}]$, and
we define $\xi_{j}=(s_{j,1},\dots,s_{j,d})$.

Let $v:D\rightarrow\mathbb{R}$ denote an intermediate function within the FNO.
Recall from the previous section that the primary operation in FNO is the
Fourier convolution operator, $(\mathcal{K}v_{t})(x)$. We say the
_discretization error_ of $\mathcal{F}(v)$ is the absolute difference between
the Fourier transform of $v$, and the discrete Fourier transform of $v$ via
discretization of $v$ on
$\mathcal{Q}_{d}=(\\{Q_{1},\dots,Q_{n}\\},\\{\xi_{1},\dots,\xi_{n}\\})$.
Formally, given Fourier basis function $\varphi_{\omega}(x)=e^{2\pi
i\langle\omega,x\rangle}$,

$\texttt{Disc}(v,\mathcal{Q}_{d},\omega)=\Big{|}\int_{D}v(x)\varphi_{\omega}(x)dx-\sum_{j=1}^{n}v(\xi_{j})\varphi_{\omega}(\xi_{j})|Q_{j}|\Big{|}.$
(1)

In other words, if we knew the true (infinite-dimensional) input function, we
could run FNO using the (continuous) Fourier transform. However, for real-
world applications, we only have access to a discretization of the input, for
example, on a $128\times 128$ grid, so we incur a discretization error. We
bound the discretization error as follows. For the full proof details, see
Appendix A.

######  Theorem 3.1.

For any $D=[0,1]^{d}$, $M>0$, and $L\geq 1$, let $\mathcal{K}\subset C(D)$ be
the set of L-Lipschitz functions, bounded by $||v||_{\infty}\leq M$. Then
there exists constants $c_{1},c_{2}>0$ such that for all $n,d,\omega$, we have

$c_{1}\sqrt{d}\cdot
Mn^{-\nicefrac{{2}}{{d}}}\leq\sup_{v\in\mathcal{K}}\left(\texttt{Disc}(v,\mathcal{Q}_{d},1)\right)\text{
and
}\sup_{v\in\mathcal{K}}\left(\texttt{Disc}(v,\mathcal{Q}_{d},\omega)\right)\leq
c_{2}\sqrt{d}(|\omega|+L)Mn^{-\nicefrac{{1}}{{d}}}.$

Note that the upper bound is true for all $\omega$, while the lower bound is
shown for $\omega=1$. It is an interesting question for future work to show
the lower bound for $\omega>1$.

Intuitively, we bound the Riemann sum approximation of the true integral by
showing that in the case where the function $v$ is bounded and Lipschitz, then
each of the $n$ intervals contributes $n^{-(1+\nicefrac{{1}}{{d}})}$ error, up
to parameters of the function. Note that the discretization error upper bound
also scales linearly for the frequency $\omega$, although we empirically show
in Section 4 that the energy is concentrated in the lower frequency modes. The
lower bound is satisfied when $v(x)=x_{1}\cdots x_{d}$. Furthermore, note that
given a function $v$, if $n$ is not large enough, the discretization error can
become arbitrarily large due to aliasing. For example, the function
$v(x)=M\sin\left(2\pi(n+\omega)x\right)$ has discretization error $\Omega(M)$.

Next, we say that the _precision error_ of $\mathcal{F}(v)$ is the absolute
difference between $\mathcal{F}(v)$ and $\overline{\mathcal{F}(v)}$, computing
the discrete Fourier transform in half precision. Specifically, we define an
$(a_{0},\epsilon,T)$-_precision system_ as a mapping $q:\mathbb{R}\rightarrow
S$ for the set
$S=\\{0\\}\cup\\{a_{0}(1+\epsilon)^{i}\\}_{i=0}^{T}\cup\\{-a_{i}(1+\epsilon)^{i}\\}_{i=0}^{T}$,
such that for all $x\in\mathbb{R}$, $q(x)=\texttt{argmin}_{y\in S}|x-y|$. This
represents a simplified version of the true mapping used by Python from
$\mathbb{R}$ to float32 or float16 (we give further justification of this
definition in Appendix A). Then, we define

$\texttt{Prec}(v,\mathcal{Q}_{d},q,\omega)=\Big{|}\sum_{j=1}^{n}v(\xi_{j})\varphi_{\omega}(\xi_{j})|Q_{j}|-\sum_{j}q(v(\xi_{j}))q(\varphi_{\omega}(\xi_{j}))|Q_{j}|\Big{|}.$
(2)

Now, we bound the precision error as follows.

######  Theorem 3.2.

For any $D=[0,1]^{d}$, $M>0$, and $L\geq 1$, let $\mathcal{K}\subset C(D)$ be
the set of L-Lipschitz functions, bounded by $||v||_{\infty}\leq M$.
Furthermore let $q$ be an $(a_{0},\epsilon,T)$-precision system. There exists
$c>0$ such that for all $n,d,\omega$,

$\sup_{v\in\mathcal{K}}\left(\texttt{Prec}(v,\mathcal{Q}_{d},q,\omega)\right)\leq
c\cdot\epsilon M.$

Once again, we show that each interval contributes
$\nicefrac{{\epsilon}}{{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^{-\nicefrac{{2}}{{d}}}$ (up to
constants), while the precision error is always bounded by $\epsilon M$. For
example, for float16 precision ($\epsilon=10^{-4}$), our theory suggests that
the precision error is comparable to the discretization error for three-
dimensional meshes up to size 1 000 000. Given this motivation, we present our
mixed-precision training pipeline in the next section.

##  4 Empirical Study of our Mixed-Precision 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 pre-
activation-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 Section B.2 for detailed
descriptions.

Navier-Stokes. We consider the Navier-Stokes equations for a 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 10 000 training samples and 2000 test samples with resolution $128\times
128$. Darcy Flow. We consider the steady-state 2D Darcy Flow equation, which
models fluid flow through a porous medium. We use the same dataset as Li et
al. (2021a), with 5000 training samples and 1000 test samples at resolution
$128\times 128$. Spherical Shallow Water Equations. We use the dataset from
Bonev et al. (2023), which generates random initial conditions on the sphere
at resolution $256\times 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 a unique 3D car, and the goal is to predict the full 3D pressure
field. We use 611 water-tight shapes from car surfaces from Shape-Net (Chang
et al., 2015), with 500 samples for training and the rest for the test set.
For Ahmed-body, we have 500 for training and 51 for test. The spatial
resolution for both is $64\times 64\times 64$.

Experimental Setup. We run the Navier Stokes and Darcy flow experiments on the
TFNO architecture (Kossaifi et al., 2023). For the Spherical SWE, we use SFNO
(Bonev et al., 2023) to handle the spherical geometry, and for Shape-Net Car
and Ahmed-body, we use GINO (Li et 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-Valued Tensor Contraction

Figure 2:  Overview of our Mixed-FNO pipeline.

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 bottleneck, accounting for 4 out of the 5 most time-consuming
GPU kernels in the entire training pipeline (see Section B.4 for the full
profiling results). Furthermore, existing mixed-precision tools such as
Pytorch’s Automatic Mixed Precision (AMP) (Paszke et al., 2019) leave the
operator in full precision since it only autocasts real-valued modules in FNO.

We introduce a simple and lightweight workaround: we break down each tensor
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, we use a simple greedy algorithm to select the next
einsum 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.

Our simple half-precision FNO block module achieves up to 38.6% reduction in
memory and up to 50% reduction in memory when combined with Pytorch’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-precision. The memory usage reduction can then be used to train on a
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 of training samples processed per second, from 1.23X to 1.58X over the
baseline in Figure 4.

Figure 3:  GPU memory usage reduction across different variants of neural
operators on diverse tasks. We use an Nvidia RTX 3090 Ti GPU. Our method
reduces memory by up to 50%, representing a super-linear combination of the
two other methods because it avoids additional casting to and from full
precision during the forward pass.

Figure 4:  Training throughput and runtime as a function of the method, on
different GPUs. For mixed-precision FNO + AMP, we consistently observe an
improvement of training throughput up to 1.58X over the baseline with the TFNO
model on Navier Stokes, and up to 1.33X with the SFNO model on Spherical
Shallow Water Equations (SWE). Our method also improves upon using only AMP in
throughput by over 1.3X on Navier-Stokes and by over 1.2X on Spherical SWE.
Batch sizes are selected to fully utilize each GPU.

###  4.3 Numerical Stability via Pre-Activation

Mixed-precision training is prone to underflow and overflow because its
dynamic range is significantly 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 Section B.6). To
mitigate this overflow issue, we find that pre-activation before each forward
FFT is a very effective method for overcoming numerical instability. We also
find that the tanh pre-activation is the highest-performing operation that we
considered according to Table 3. Unlike other functions, tanh minimizes
changes to small inputs, as it is approximately the identity function near 0
and is smoothly differentiable. Furthermore, tanh preserves the
discretization-convergent property of FNO. The strong performance of tanh is
also predicted by our theoretical results in Section 3, since it decreases the
$L_{\infty}$ norm and the Lipschitz constant of the input function. In Section
B.6, we further show that the tanh pre-activation minimally alters the
frequency-domain signal in both amplitude and phase. Finally, we demonstrate
through an ablation that tanh has negligible impact on the model’s final error
(See Appendix Table 5).

###  4.4 Error Comparison with Full-Precision

Having resolved the critical issue of numerical stability, we demonstrate that
our mixed-precision approach achieves errors within 1% of the full-precision
baseline across the four datasets. Additionally, we propose a precision-
scheduling technique that transitions from mixed to full precision during
training, which performs better than full precision in zero-shot super-
resolution inference.

Figure 5:  Test H1 error curves for FNO on the Navier-Stokes (top left) and
Darcy flow (top right) datasets. Test L2 error curves for GINO on the Shape-
Net Car (bottom left) and for SFNO on the Shallow Water Equation (bottom
right) datasets. Each plot shows the mean of three random seeds and standard
deviation as error bars. We compute and report the average difference between
training curves in the legend. We also annotate the difference in final test
errors with the memory savings of our mixed precision approach.

Mixed- vs. Full-Precision Training Curves. Figure 5 illustrates that our
mixed-precision approach achieves test errors on par with full precision
throughout the training process, remaining within 1% of the full-precision
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 a higher resolution (zero-shot super-
resolution) (Kovachki et al., 2023). To achieve the best result, we propose a
precision schedule, in which the first 25% of training is in mixed-precision,
the middle 50% applying only AMP, and the final 25% in full precision. This
follows a simple intuition: in the early stages of 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 are
much smaller, so full precision is more important. We use the Navier-Stokes
dataset, trained on $128\times 128$ resolution (the same setting and model as
Figure 5), and tested on $256\times 256$, $512\times 512$, and $1024\times
1024$ resolutions; see Table 1. We find that half-precision has a small
decrease in accuracy compared to full precision, and using a precision
schedule achieves significantly better generalization.

| 128x128 | 256x256 | 512x512 | 1024x1024
---|---|---|---|---
| $H^{1}$ | $L^{2}$ | $H^{1}$ | $L^{2}$ | $H^{1}$ | $L^{2}$ | $H^{1}$ | $L^{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\times 128$ resolution for 19 hours. Mixed
precision has a small decrease in accuracy compared to full precision, and
using a precision schedule achieves significantly better accuracy.

###  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, we
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 to AMP applied to U-Nets.

Model | Navier-Stokes | Darcy Flow
---|---|---
Error | Memory Reduction | Error | Memory Reduction
Full FNO | 0.003 | 50.4% | 0.01 | 25.8%
Mixed FNO (Ours) | 0.004 | 0.007
Full U-Net | 0.111 | 20.9% | 0.024 | 24.9%
U-Net + AMP | 0.111 | 0.022

Table 2: Comparison of $L^{2}$ 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.

###  4.6 Ablation studies

Here, we perform ablations of our mixed-precision procedure on different
parameterizations 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, we
adopt a Canonical-Polyadic (CP) factorization of the FNO weights using
Kossaifi et al. (2019) to ensure better final error. For a generic problem,
FNO weights are usually not factorzied and are saved in their original, dense
form. As shown in Figure 6, we experiment with both scenarios and show that
our mixed-precision approach improves runtime without sacrificing accuracy.

Figure 6:  For Navier-Stokes (Left) and Darcy Flow (Right), our method
accelerates training, whether with CP factorized weights or Dense unfactorized
weights, without losing more than 1% in error.

Number of Frequency Modes. Recall that in the FNO architecture, after the FFT,
we truncate to a fixed fraction of frequency modes to improve generalization,
typically $\nicefrac{{1}}{{3}}$ to $\nicefrac{{2}}{{3}}$. We run an ablation
study on the number of frequency modes used in the FNO architecture. We run
frequency modes $\\{16,32,64,128\\}$ on the Darcy flow dataset in full and
half precision; see 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 a significant difference between half-
precision and full precision for all frequencies. In addition, we show in
Section B.10 that for synthetic data, the precision error from half precision
is higher for higher frequencies relative to the amplitude.

Other Mixed Precision Options. We also experimented with BrainFloat16 (BF16)
and TensorFloat32 (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 as our approach, even on optimized hardware such as the A100 GPU.
Moreover, we simulated FP8 training via clipping. See Section B.11 for
additional details.

##  5 Conclusions and Future Work

In this work, we introduced the first mixed-precision training method for
neural operators. We derived approximation bounds for mixed-precision neural
operator and rigorously characterized the approximation and precision errors
of FNO-based operators, proving that the precision error is comparable to the
approximation error. Using this insight, we show, in practice, a significant
reduction of memory usage and improvement in throughput, without sacrificing
accuracy, solving critical issues from standard neural net-based training
methods. Through extensive experiments on different state-of-the-art neural
operators, datasets, and hardware, we demonstrate that our approach reduces
GPU memory usage by up to 50% with little or no reduction in accuracy.

Overall, half-precision FNO makes it possible to train on significantly larger
data points with the same batch size. Going forward, we plan to apply this to
real-world applications that require super-resolution to enable larger-scale
training.

## 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 thank
members of the EcoSystem lab, especially Shang Wang, for their constructive
feedback.

## References

  * Adler & Öktem (2017) Jonas Adler and Ozan Öktem.  Solving ill-posed inverse problems using iterative deep neural networks.  _Inverse Problems_ , 33(12):124007, 2017.
  * Ahmed et al. (1984) Syed R Ahmed, G Ramm, and Gunter Faltin.  Some salient features of the time-averaged ground vehicle wake.  _SAE transactions_ , pp. 473–503, 1984.
  * Bhatnagar et al. (2019) Saakaar Bhatnagar, Yaser Afshar, Shaowu Pan, Karthik Duraisamy, and Shailendra Kaushik.  Prediction of aerodynamic flow fields using convolutional neural networks.  _Computational Mechanics_ , 64:525–545, 2019.
  * Bonev et al. (2023) 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 on Machine Learning (ICML)_ , 2023.
  * Chandler & Kerswell (2013) Gary J Chandler and Rich R Kerswell.  Invariant recurrent solutions embedded in a turbulent two-dimensional kolmogorov flow.  _Journal of Fluid Mechanics_ , 722:554–595, 2013.
  * Chang et al. (2015) Angel X 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.
  * De Sa et al. (2018) Christopher De Sa, Megan Leszczynski, Jian Zhang, Alana Marzoev, Christopher R Aberger, Kunle Olukotun, and Christopher Ré.  High-accuracy low-precision training.  _arXiv preprint arXiv:1803.03383_ , 2018.
  * Dean et al. (2012) Jeffrey Dean, Greg Corrado, Rajat Monga, Kai Chen, Matthieu Devin, Mark Mao, Marc’aurelio Ranzato, Andrew Senior, Paul Tucker, Ke Yang, et al.  Large scale distributed deep networks.  _Advances in neural information processing systems_ , 25, 2012.
  * Dool et al. (2023) Winfried van den Dool, Tijmen Blankevoort, Max Welling, and Yuki M Asano.  Efficient neural pde-solvers using quantization aware training.  _arXiv preprint arXiv:2308.07350_ , 2023.
  * Driscoll & Healy (1994) 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.1006/aama.1994.1008.  URL https://linkinghub.elsevier.com/retrieve/pii/S0196885884710086.
  * Goel et al. (2020) Abhinav 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 on Internet of Things (WF-IoT)_ , pp. 1–6. IEEE, 2020.
  * Gopakumar et al. (2023) 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.
  * Guibas et al. (2021) 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 on Learning Representations (ICLR)_ , 2021.
  * Gupta et al. (2021) Gaurav Gupta, Xiongye Xiao, and Paul Bogdan.  Multiwavelet-based operator learning for differential equations.  _Proceedings of the Annual Conference on Neural Information Processing Systems (NeurIPS)_ , 34:24048–24062, 2021.
  * Hayford et al. (2024) Joel Hayford, Jacob Goldman-Wetzler, Eric Wang, and Lu Lu.  Speeding up and reducing memory usage for scientific machine learning via mixed precision.  _arXiv preprint arXiv:2401.16645_ , 2024.
  * Jia et al. (2018) Xianyan Jia, Shutao Song, Wei He, Yangzihao Wang, Haidong Rong, Feihu Zhou, Liqiang Xie, Zhenyu Guo, Yuanzhou Yang, Liwei Yu, et al.  Highly scalable deep learning training system with mixed-precision: Training imagenet in four minutes.  _arXiv preprint arXiv:1807.11205_ , 2018.
  * Kalamkar et al. (2019) Dhiraj Kalamkar, Dheevatsa Mudigere, Naveen Mellempudi, Dipankar Das, Kunal Banerjee, Sasikanth Avancha, Dharma Teja Vooturi, Nataraj Jammalamadaka, Jianyu Huang, Hector Yuen, et al.  A study of bfloat16 for deep learning training.  _arXiv preprint arXiv:1905.12322_ , 2019.
  * Kolda & Bader (2009) Tamara G Kolda and Brett W Bader.  Tensor decompositions and applications.  _SIAM review_ , 51(3):455–500, 2009.
  * Kossaifi et al. (2019) 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.
  * Kossaifi et al. (2023) Jean Kossaifi, Nikola Kovachki, Kamyar Azizzadenesheli, and Anima Anandkumar.  Multi-grid tensorized fourier neural operator for high-resolution pdes.  _arXiv preprint arXiv:2310.00120_ , 2023.
  * Kovachki et al. (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 http://jmlr.org/papers/v24/21-1524.html.
  * Li et al. (2020a) 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.
  * Li et al. (2020b) 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)_ , 33:6755–6766, 2020b.
  * Li et al. (2021a) 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)_ , 2021a.
  * Li et al. (2021b) Zongyi Li, Hongkai Zheng, Nikola Kovachki, David Jin, Haoxuan Chen, Burigede Liu, Kamyar Azizzadenesheli, and Anima Anandkumar.  Physics-informed neural operator for learning partial differential equations.  _arXiv preprint arXiv:2111.03794_ , 2021b.
  * Li et al. (2023) Zongyi Li, Nikola Borislavov Kovachki, Chris Choy, Boyi Li, Jean Kossaifi, Shourya Prakash Otta, Mohammad Amin Nabian, Maximilian Stadler, Christian Hundt, Kamyar Azizzadenesheli, and Anima Anandkumar.  Geometry-informed neural operator for large-scale 3d pdes.  _arXiv preprint arXiv:2309.00583_ , 2023.
  * Liu et al. (2022) Burigede Liu, Nikola Kovachki, Zongyi Li, Kamyar Azizzadenesheli, Anima Anandkumar, Andrew M Stuart, and Kaushik Bhattacharya.  A learning-based multiscale method and its application to inelastic impact problems.  _Journal of the Mechanics and Physics of Solids_ , 158:104668, 2022.
  * Lu et al. (2019) 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.03193_ , 2019.
  * Lu et al. (2021) 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, 2021.
  * Micikevicius et al. (2017) 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:1710.03740_ , 2017.
  * Micikevicius et al. (2022) Paulius Micikevicius, Dusan Stosic, Neil Burgess, Marius Cornea, Pradeep Dubey, Richard Grisenthwaite, Sangwon Ha, Alexander Heinecke, Patrick Judd, John Kamalu, et al.  Fp8 formats for deep learning.  _arXiv preprint arXiv:2209.05433_ , 2022.
  * Panagakis et al. (2021) Yannis Panagakis, Jean Kossaifi, Grigorios G. Chrysos, James Oldfield, Mihalis A. Nicolaou, Anima Anandkumar, and Stefanos Zafeiriou.  Tensor methods in computer vision and deep learning.  _Proceedings of the IEEE_ , 109(5):863–890, 2021\.
  * Panagakis et al. (2024) 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 978-0-323-91772-8.  doi: https://doi.org/10.1016/B978-0-32-391772-8.00021-1.  URL https://www.sciencedirect.com/science/article/pii/B9780323917728000211.
  * Paszke et al. (2019) 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 (NeurIPS)_ , 2019.
  * Pathak et al. (2022) Jaideep Pathak, Shashank Subramanian, Peter Harrington, Sanjeev Raja, Ashesh Chattopadhyay, Morteza Mardani, Thorsten Kurth, David Hall, Zongyi Li, Kamyar Azizzadenesheli, et al.  Fourcastnet: A global data-driven high-resolution weather model using adaptive fourier neural operators.  _arXiv preprint arXiv:2202.11214_ , 2022.
  * Rakka et al. (2022) Mariam Rakka, Mohammed E Fouda, Pramod Khargonekar, and Fadi Kurdahi.  Mixed-precision neural networks: A survey.  _arXiv preprint arXiv:2208.06064_ , 2022.
  * Renn et al. (2023) Peter I Renn, Cong Wang, Sahin Lale, Zongyi Li, Anima Anandkumar, and Morteza Gharib.  Forecasting subcritical cylinder wakes with fourier neural operators.  _arXiv preprint arXiv:2301.08290_ , 2023.
  * Ronneberger et al. (2015) Olaf Ronneberger, Philipp Fischer, and Thomas Brox.  U-net: Convolutional networks for biomedical image segmentation.  In Nassir Navab, Joachim Hornegger, William M. Wells, and Alejandro F. Frangi (eds.), _Medical Image Computing and Computer-Assisted Intervention – MICCAI 2015_ , pp. 234–241, Cham, 2015. Springer International Publishing.  ISBN 978-3-319-24574-4.
  * Schneider et al. (2017) Tapio Schneider, João Teixeira, Christopher S Bretherton, Florent Brient, Kyle G Pressel, Christoph Schär, and A Pier Siebesma.  Climate goals and computing the future of clouds.  _Nature Climate Change_ , 7(1):3–5, 2017.
  * Strang (2007) Gilbert Strang.  Computational science and engineering.  _Optimization_ , 551(563):571–586, 2007.
  * Trefethen (2000) Lloyd N Trefethen.  _Spectral methods in MATLAB_.  SIAM, 2000.
  * Umetani & Bickel (2018) Nobuyuki Umetani and Bernd Bickel.  Learning three-dimensional flow for interactive aerodynamic design.  _ACM Transactions on Graphics (TOG)_ , 37(4):1–10, 2018.
  * Wen et al. (2022) 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:2210.17051_ , 2022.
  * Wilcox et al. (1998) David C Wilcox et al.  _Turbulence modeling for CFD_ , volume 2.  DCW industries La Canada, CA, 1998.
  * Zhao et al. (2022) Jiawei Zhao, Steve Dai, Rangharajan Venkatesan, Brian Zimmer, Mustafa Ali, Ming-Yu Liu, Brucek Khailany, William J Dally, and Anima Anandkumar.  Lns-madam: Low-precision training in logarithmic number system using multiplicative weight update.  _IEEE Transactions on Computers_ , 71(12):3179–3190, 2022.

##  Appendix A Additional Details from Section 3

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 $D$ denote the closed unit hypercube $[0,1]^{d}$ for dimension
$d\in\mathbb{N}$. Let $Q_{1},\dots,Q_{n}$ denote the unique (up to $d$)
partitioning of $D$, such that each $Q_{j}$ is a hypercube with sidelength
$\nicefrac{{1}}{{m}}$. For each $1\leq j\leq n$, let $\xi_{j}\in Q_{j}$ denote
the vertex that is closest to the origin; formally, we have
$Q_{j}=\prod_{k=1}^{d}[s_{j,k},t_{j,k}]$, and we define
$\xi_{j}=(s_{j,1},\dots,s_{j,d})$.

Let $v:D\rightarrow\mathbb{R}$ denote an intermediate function within the FNO.
Recall from the previous section that the primary operation in FNO is the
Fourier convolution operator, $(\mathcal{K}v_{t})(x)$. We say the
_discretization error_ of $\mathcal{F}(v)$ is the absolute difference between
the Fourier transform of $v$, and the discrete Fourier transform of $v$ via
discretization of $v$ on
$\mathcal{Q}_{d}=(\\{Q_{1},\dots,Q_{n}\\},\\{\xi_{1},\dots,\xi_{n}\\})$.
Formally, given Fourier basis function $\varphi_{\omega}(x)=e^{2\pi
i\langle\omega,x\rangle}$,

$\texttt{Disc}(v,\mathcal{Q}_{d},\omega)=\Big{|}\int_{D}v(x)\varphi_{\omega}(x)dx-\sum_{j=1}^{n}v(\xi_{j})\varphi_{\omega}(\xi_{j})|Q_{j}|\Big{|}.$

###  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 $D=[0,1]^{d}$, $M>0$, and $L\geq 1$, let
$\mathcal{K}\subset C(D)$ be the set of L-Lipschitz functions, bounded by
$||v||_{\infty}\leq M$. Then there exists constants $c_{1},c_{2}>0$ such that
for all $n,d,\omega$, we have_

$c_{1}\sqrt{d}\cdot
Mn^{-\nicefrac{{2}}{{d}}}\leq\sup_{v\in\mathcal{K}}\left(\texttt{Disc}(v,\mathcal{Q}_{d},1)\right)\text{
and
}\sup_{v\in\mathcal{K}}\left(\texttt{Disc}(v,\mathcal{Q}_{d},\omega)\right)\leq
c_{2}\sqrt{d}(|\omega|+L)Mn^{-\nicefrac{{1}}{{d}}}.$

###### Proof.

Let $\varphi_{\omega}^{r}(x)=\sin(2\pi\langle\omega,x\rangle)$ denote the real
part of the Fourier base of $v$ at frequency $\omega$. Define

$\texttt{Disc}^{r}(v,Q_{d},\omega)=\left|\int_{D}v(x)\varphi_{\omega}^{r}(x)dx-\sum_{j=1}^{n}v(\xi_{j})\varphi_{\omega}^{r}(\xi_{j})|Q_{j}|\right|.$
(3)

Recall that as defined above,
$\mathcal{Q}_{d}=(\\{Q_{1},\dots,Q_{n}\\},\\{\xi_{1},\dots,\xi_{n}\\})$, where
$\\{Q_{1},\dots,Q_{n}\\}$ is the unique (up to $d$) partitioning of $D$, and
for each $1\leq j\leq n$, $\xi_{j}\in Q_{j}$ denotes the vertex that is
closest to the origin; formally, we have
$Q_{j}=\prod_{k=1}^{d}[s_{j,k},t_{j,k}]$, and we define
$\xi_{j}=(s_{j,1},\dots,s_{j,d})$. First we prove the upper bound. We have:

$\displaystyle\texttt{Disc}^{r}(v,Q_{d},\omega)$
$\displaystyle=\left|\int_{D}v(x)\varphi_{\omega}^{r}(x)dx-\sum_{j=1}^{n}v(\xi_{j})\varphi_{\omega}^{r}(\xi_{j})|Q_{j}|\right|$
(4)
$\displaystyle=\left|\int_{D}v(x)\varphi_{\omega}^{r}(x)dx-\sum_{j=1}^{n}\int_{Q_{j}}v(\xi_{j})\varphi_{\omega}^{r}(\xi_{j})dx\right|$
(5)
$\displaystyle=\left|\sum_{j=1}^{n}\int_{Q_{j}}\left(v(x)\varphi_{\omega}^{r}(x)-v(\xi_{j})\varphi_{\omega}^{r}(\xi_{j})\right)dx\right|$
(6)
$\displaystyle=\left|\sum_{j=1}^{n}\int_{Q_{j}}\left(v(x)(\varphi_{\omega}^{r}(x)-\varphi_{\omega}^{r}(\xi_{j}))+(v(x)-v(\xi_{j}))\varphi_{\omega}^{r}(\xi_{j})\right)dx\right|$
(7)
$\displaystyle\leq\sum_{j=1}^{n}\int_{Q_{j}}\left(|v(x)|\cdot|\varphi_{\omega}^{r}(x)-\varphi_{\omega}^{r}(\xi_{j})|+|v(x)-v(\xi_{j})|\cdot|\varphi_{\omega}^{r}(\xi_{j})|\right)dx$
(8)
$\displaystyle\leq\sum_{j=1}^{n}\int_{Q_{j}}\left(M\cdot\left(|\omega|\frac{\sqrt{d}}{m}\right)+\left(L\cdot\frac{\sqrt{d}}{m}\right)\cdot
1\right)dx$ (9)
$\displaystyle=\sum_{j=1}^{n}\int_{Q_{j}}\left(\frac{\sqrt{d}}{m}\left(M|\omega|+L\right)\right)dx$
(10)
$\displaystyle=n\cdot\left(\frac{1}{m^{d}}\right)\sqrt{d}\left(M|\omega|+L\right)n^{-\nicefrac{{1}}{{d}}}$
(11) $\displaystyle=\sqrt{d}\left(M|\omega|+L\right)n^{-\nicefrac{{1}}{{d}}}.$
(12)

Bounding $\texttt{Disc}^{c}(v,Q_{d},\omega)$, the complex part of the Fourier
base, follows an identical argument. Setting $c_{2}=2$ concludes the proof of
the upper bound, which is true for all $\omega$. Also note that we only needed
the fact that $\xi_{j}\in Q_{j}$ for all $1\leq j\leq n$.

Now we move to the lower bound. We set $v(x)=x_{1}\cdots x_{d}$ and
$\omega=1$, and we will show that

$\displaystyle\left|\int_{D}v(x)\sin(2\pi
x)dx-\sum_{j=1}^{n}v(\xi_{j})\sin(2\pi\xi_{j})\right|=\frac{d}{3\cdot
2^{d}\pi^{d-2}}\cdot n^{\nicefrac{{1}}{{2d}}}.$

First, in one dimension, $\int_{0}^{1}x_{1}\sin(2\pi
x_{1})dx_{1}=\frac{1}{2\pi}.$ Therefore, for $d$ dimensions, we have

$\int_{D}v(x)\sin(2\pi x)dx=\left(\int_{0}^{1}x_{1}\sin(2\pi
x_{1})dx_{1}\right)^{d}=(2\pi)^{-d}.$

Next, we will show a lower bound on
$\sum_{j=1}^{n}v(\xi_{j})\sin(2\pi\xi_{j})$. Here, we will need the explicit
definition of the $\xi_{j}$’s defined at the start of this section.
Specifically, for each $1\leq j\leq n$, we have
$Q_{j}=\prod_{k=1}^{d}[s_{j,k},t_{j,k}]$, and we define
$\xi_{j}=(s_{j,1},\dots,s_{j,d})$. Now we note that by construction, since the
unit hypercube $D$ is partitioned uniformly across each dimension in segments
of $\nicefrac{{1}}{{m}}$, we can parameterize each $Q_{j}$ by a unique
$d$-tuple $(i_{1},\dots i_{d})$, where each $i_{k}$ is an integer from $0$ to
$m-1$, and the corresponding
$\xi_{j}=(s_{j,1},\dots,s_{j,d})=(\nicefrac{{i_{1}}}{{m}},\dots\nicefrac{{i_{d}}}{{m}})$.
Also, each $d$-tuple of integers from $0$ to $m-1$ defines a unique $Q_{j}$.
In other words, we are defining a different parameterization of the $Q_{j}$’s
by using the fact that the $\xi_{j}$’s form a uniform lattice across the unit
hypercube. This parameterization is convenient for the next part of the proof.

Recall that $v(x)=x_{1}\cdots x_{d}$. We have

$\displaystyle\sum_{j=1}^{n}\left(v(\xi_{j})\sin(2\pi\xi_{j})|Q_{j}|\right)$
$\displaystyle=\frac{1}{m^{d}}\left(\sum_{j=1}^{n}\left(\prod_{k=1}^{d}s_{j,k}\sin(2\pi
s_{j,k})\right)\right)$ (13)
$\displaystyle=\frac{1}{m^{d}}\left(\sum_{i_{1}=1}^{m}\cdots\sum_{i_{d}=1}^{m}\left(\prod_{k=1}^{d}\left(\frac{i_{k}}{m}\cdot\sin\left(2\pi\frac{i_{k}}{m}\right)\right)\right)\right)$
(14)
$\displaystyle=\frac{1}{m^{d}}\left(\sum_{i_{1}=1}^{m}\frac{i_{1}}{m}\sin\left(2\pi\frac{i_{1}}{m}\right)\right)\cdots\left(\sum_{i_{d}=1}^{m}\frac{i_{d}}{m}\sin\left(2\pi\frac{i_{d}}{m}\right)\right)$
(15)
$\displaystyle=\frac{1}{m^{d}}\left(\sum_{j=1}^{m}\left(\frac{j}{m}\right)\cdot\sin\left(2\pi\frac{j}{m}\right)\right)^{d}$
(16)
$\displaystyle=m^{-2d}\left(\sum_{j=1}^{m}\left(j\cdot\sin\left(2\pi\frac{j}{m}\right)\right)\right)^{d}$
(17)
$\displaystyle=m^{-2d}\left(-\frac{m}{2}\text{cot}\left(\frac{\pi}{m}\right)\right)^{d}$
(18) $\displaystyle\geq 2^{-d}\cdot
m^{-d}\left(-\frac{m}{\pi}+\frac{1}{3}\cdot\frac{\pi}{m}\right)^{d}$ (19)
$\displaystyle=(2\pi)^{-d}\cdot\frac{(m+\frac{\pi^{2}}{3m})^{d}}{m^{-d}}$ (20)
$\displaystyle\geq(2\pi)^{-d}\left(1+d\cdot\frac{\pi^{2}}{3}\cdot
m^{-2}\right).$ (21)

Finally, since $n=m^{d}$, we have

$\left|(2\pi)^{-d}-(2\pi)^{-d}\left(1+d\cdot\frac{\pi^{2}}{3}\cdot
m^{-2}\right)\right|=\frac{d}{3\cdot 2^{d}\pi^{d-2}}\cdot
n^{\nicefrac{{1}}{{2d}}},$

which concludes the proof. ∎

Note that the upper bound is true for all $\omega$, while the lower bound is
shown for $\omega=1$. It is an interesting question for future work to show
the lower bound for $\omega>1$.

Now we bound the precision error of $\mathcal{F}(v)$. Recall from Section 3
that the _precision error_ of $\mathcal{F}(v)$ is the absolute difference
between $\mathcal{F}(v)$ and $\overline{\mathcal{F}(v)}$, computing the
discrete Fourier transform in half precision. Specifically, we define an
$(a_{0},\epsilon,T)$-_precision system_ as a mapping $q:\mathbb{R}\rightarrow
S$ for the set
$\\{0\\}\cup\\{a_{0}(1+\epsilon)^{j}\\}_{j=0}^{T}\cup\\{-a_{j}(1+\epsilon)^{j}\\}_{j=0}^{T}$,
such that for all $x\in\mathbb{R}$, $q(x)=\texttt{argmin}_{y\in S}|x-y|$.

This represents a simplified version of the true mapping used by Python from
$\mathbb{R}$ to float32 or float16. Recall that these datatypes allocate some
number of bits for the mantissa and for the exponent. Then, given a real
number $x\cdot 2^{y}$, its value in float32 or float16 would be
$(x+\epsilon_{1})\cdot 2^{(y+\epsilon_{2})}=x\cdot 2^{y}+\epsilon_{2}x\cdot
2^{y}+\epsilon_{1}\cdot
2^{y}+\epsilon_{1}\epsilon_{2}2^{y}=(1+\nicefrac{{\epsilon_{1}}}{{x}}+\epsilon_{2})x\cdot
2^{y}$. Therefore, our definition above is a simplified but reasonable
approximation of floating-point arithmetic.

Now, we define

$\texttt{Prec}(v,\mathcal{Q}_{d},q,\omega)=\Big{|}\sum_{j=1}^{n}v(\xi_{j})\varphi_{\omega}(\xi_{j})|Q_{j}|-\sum_{j}q(v(\xi_{j}))q(\varphi_{\omega}(\xi_{j}))|Q_{j}|\Big{|}.$

Now we bound the precision error of $\mathcal{F}(v)$.

Theorem 3.2 (restated). _For any $D=[0,1]^{d}$, $M>0$, and $L\geq 1$, let
$\mathcal{K}\subset C(D)$ be the set of L-Lipschitz functions, bounded by
$||v||_{\infty}\leq M$. Furthermore let $q$ be an
$(a_{0},\epsilon,T)$-precision system. There exists $c>0$ such that for all
$n,d,\omega$,_

$\sup_{v\in\mathcal{K}}\left(\texttt{Prec}(v,\mathcal{Q}_{d},q,\omega)\right)\leq
c\cdot\epsilon M.$

###### Proof.

Let $\varphi_{\omega}^{r}(x)=\sin(2\pi\langle\omega,x\rangle)$ denote the real
part of the Fourier base of $v$ at frequency $\omega$. Define

$\texttt{Prec}^{r}(v,\mathcal{Q}_{d},q,\omega)=\left|\sum_{j=1}^{n}v(\xi_{j})\varphi_{\omega}^{r}(\xi_{j})|Q_{j}|-\sum_{j=1}^{n}q(v(\xi_{j}))q(\varphi_{\omega}^{r}(\xi_{j}))|Q_{j}|\right|$
(23)

We prove the upper bound as follows. We have

$\displaystyle\texttt{Prec}^{r}(v,\mathcal{Q}_{d},q,\omega)$
$\displaystyle=\left|\sum_{j=1}^{n}v(\xi_{j})\varphi_{\omega}^{r}(\xi_{j})|Q_{j}|-\sum_{j=1}^{n}q(v(\xi_{j}))q(\varphi_{\omega}^{r}(\xi_{j}))|Q_{j}|\right|$
$\displaystyle=\left|\frac{1}{n}\sum_{j=1}^{n}\left(v(\xi_{j})\varphi_{\omega}^{r}(\xi_{j})-q(v(\xi_{j}))q(\varphi_{\omega}^{r}(\xi_{j}))\right)\right|$
$\displaystyle\leq\left|\frac{1}{n}\sum_{j=1}^{n}\left(v(\xi_{j})(\varphi_{\omega}^{r}(\xi_{j})-q(\varphi_{\omega}^{r}(\xi_{j})))+(v(\xi_{j})-q(v(\xi_{j})))q(\varphi_{\omega}^{r}(\xi_{j}))\right)\right|$
$\displaystyle\leq\frac{1}{n}\sum_{j=1}^{n}\left(|v(\xi_{j})|\cdot|\varphi_{\omega}^{r}(\xi_{j})-q(\varphi_{\omega}^{r}(\xi_{j}))|+|v(\xi_{j})-q(v(\xi_{j}))|\cdot|q(\varphi_{\omega}^{r}(\xi_{j}))|\right)$
$\displaystyle\leq\frac{1}{n}\sum_{j=1}^{n}\left(M\cdot(\epsilon\cdot
1)+(\epsilon\cdot M)\cdot 1\right)$ $\displaystyle=\frac{1}{n}\cdot
n\left(2\epsilon M\right)$ $\displaystyle=2\epsilon M.$

Bounding $\texttt{Prec}^{c}(v,\mathcal{Q}_{d},q,\omega)$, 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 $\mathcal{F}(v)$, a function $v$ times
the Fourier basis function.

######  Theorem A.1.

For any $D=[0,1]^{d}$, $M>0$, and $L\geq 1$, let $\mathcal{K}\subset C(D)$ be
the set of L-Lipschitz functions, bounded by $||f||_{\infty}\leq M$. Then for
all $n,d,\omega$, we have

$2^{-d+1}\cdot d\cdot
n^{-\nicefrac{{1}}{{d}}}\leq\sup_{f\in\mathcal{K}}\left(\texttt{Disc}(f,\mathcal{Q},\omega)\right)\leq
L\sqrt{d}\cdot n^{-\nicefrac{{1}}{{d}}}.$

###### Proof.

First, we define

$\texttt{Disc}^{r}(f,\mathcal{Q}_{d},\omega)=\left|\int_{D}f(x)dx-\sum_{j=1}^{n}f(\xi_{j})|Q_{j}|\right|.$
(24)

Recall that
$\mathcal{Q}_{d}=(\\{Q_{1},\dots,Q_{n}\\},\\{\xi_{1},\dots,\xi_{n}\\})$, where
the $n=m^{d}$ hypercubes of side length $\nicefrac{{1}}{{m}}$ subdivide $D$.
Now we prove the upper bound. We have:

$\displaystyle\texttt{Disc}^{r}(f,\mathcal{Q}_{d},\omega)$
$\displaystyle=\left|\int_{D}f(x)dx-\sum_{j=1}^{n}f(\xi_{j})|Q_{j}|\right|$
(25)
$\displaystyle=\left|\int_{D}f(x)dx-\sum_{j=1}^{n}\int_{Q_{j}}f(\xi_{j})dx\right|$
(26)
$\displaystyle=\left|\sum_{j=1}^{n}\int_{Q_{j}}\left(f(x)-f(\xi_{j})\right)dx\right|$
(27)
$\displaystyle\leq\sum_{j=1}^{n}\int_{Q_{j}}\left|f(x)-f(\xi_{j}\right|dx$
(28)
$\displaystyle\leq\sum_{j=1}^{n}\int_{Q_{j}}\left(L\cdot\frac{\sqrt{d}}{m}\right)dx$
(29)
$\displaystyle=n\cdot\frac{1}{m^{d}}\left(L\cdot\frac{\sqrt{d}}{m}\right)$
(30) $\displaystyle=L\sqrt{d}\cdot n^{-\nicefrac{{1}}{{d}}}$ (31)

This concludes the proof of the upper bound.

For the lower bound, first, in one dimension,
$\int_{0}^{1}x_{1}dx_{1}=\frac{1}{2}.$ Therefore, for $d$ dimensions, we have

$\int_{D}f(x)dx=\left(\int_{0}^{1}x_{1}dx_{1}\right)^{d}=(2)^{-d}.$

Next, we will use the same explicit definition and reparameterization of the
$\xi_{j}$ as we used in the lower bound of Theorem 3.1: for each $1\leq j\leq
n$, we have $Q_{j}=\prod_{k=1}^{d}[s_{j,k},t_{j,k}]$, and we define
$\xi_{j}=(s_{j,1},\dots,s_{j,d})$. Now we note that by construction, since the
unit hypercube $D$ is partitioned uniformly across each dimension in segments
of $\nicefrac{{1}}{{m}}$, we can parameterize each $Q_{j}$ by a unique
$d$-tuple $(i_{1},\dots i_{d})$, where each $i_{k}$ is an integer from $0$ to
$m-1$, and the corresponding
$\xi_{j}=(s_{j,1},\dots,s_{j,d})=(\nicefrac{{i_{1}}}{{m}},\dots\nicefrac{{i_{d}}}{{m}})$.
Also, each $d$-tuple of integers from $0$ to $m-1$ defines a unique $Q_{j}$.
In other words, we are defining a different parameterization of the $Q_{j}$’s
by using the fact that the $\xi_{j}$’s form a uniform lattice across the unit
hypercube. This parameterization is convenient for the next part of the proof.
Recall that $f(x)=x_{1}\cdots x_{d}$. We have

$\displaystyle\sum_{j=1}^{n}\left(f(\xi_{j})|Q_{j}|\right)$
$\displaystyle=\frac{1}{m^{d}}\left(\sum_{j=1}^{n}\left(\prod_{k=1}^{d}f(s_{j,k})\right)\right)$
(32)
$\displaystyle=\frac{1}{m^{d}}\left(\sum_{i_{1}=1}^{m}\cdots\sum_{i_{d}=1}^{m}\left(\prod_{k=1}^{d}f\left(\frac{i_{k}}{m}\right)\right)\right)$
(33)
$\displaystyle=\frac{1}{m^{d}}\left(\sum_{i_{1}=1}^{m}\frac{i_{1}}{m}\right)\cdots\left(\sum_{i_{d}=1}^{m}\frac{i_{d}}{m}\right)$
(34) $\displaystyle=\frac{1}{m^{d}}\left(\sum_{j=1}^{m}\frac{j}{m}\right)^{d}$
(35) $\displaystyle=m^{-2d}\left(\frac{m(m+1)}{2}\right)^{d}$ (36)
$\displaystyle=2^{-d}\left(1+\frac{1}{m}\right)^{d}$ (37) $\displaystyle\geq
2^{-d}\left(1+2dm^{-1}\right)^{d}$ (38)

Finally, since $n=m^{d}$, we have

$\left|2^{-d}-2^{-d}\left(1+2dm^{-1}\right)\right|=2^{-d+1}\cdot d\cdot
n^{\nicefrac{{1}}{{d}}},$

which concludes the proof. ∎

Now we similarly bound the precision error of $f$.

######  Theorem A.2.

For any $D=[0,1]^{d}$, $M>0$, and $L\geq 1$, let $\mathcal{K}\subset C(D)$ be
the set of L-Lipschitz functions, bounded by $||f||_{\infty}\leq M$.
Furthermore let $q$ be an $(a_{0},\epsilon,T)$-precision system. Then for all
$n,d,\omega$, there exists $c>0$ such that

$\nicefrac{{1}}{{4}}\cdot\epsilon
M\leq\sup_{f\in\mathcal{K}}\left(\texttt{Prec}(f,\mathcal{Q}_{d},q,\omega)\right)\leq\epsilon
M.$

###### Proof.

Define

$\texttt{Prec}^{r}(f,\mathcal{Q}_{d},q,\omega)=\left|\sum_{j=1}^{n}f(\xi_{j})|Q_{j}|-\sum_{j=1}^{n}q(f(\xi_{j}))|Q_{j}|\right|$
(40)

We prove the upper bound as follows. We have

$\displaystyle\texttt{Prec}^{r}(f,\mathcal{Q}_{d},q,\omega)$
$\displaystyle=\left|\sum_{j=1}^{n}f(\xi_{j})|Q_{j}|-\sum_{j=1}^{n}q(f(\xi_{j}))|Q_{j}|\right|$
$\displaystyle=\left|\frac{1}{n}\sum_{j=1}^{n}\left(f(\xi_{j})-q(f(\xi_{j}))\right)\right|$
$\displaystyle\leq\frac{1}{n}\sum_{j=1}^{n}\left|f(\xi_{j})-q(f(\xi_{j}))\right|$
$\displaystyle\leq\frac{1}{n}\sum_{j=1}^{n}\left(\epsilon M\right)$
$\displaystyle=\epsilon M$

This concludes the proof of the upper bound.

For the lower bound, given the definition of an $(a_{0},\epsilon,T)$-precision
system, it follows that there exists $y$ such that $\nicefrac{{M}}{{2}}<y<M$
and $\nicefrac{{1}}{{2}}\cdot\epsilon y<|y-q(y)|.$ Then, for the lower bound
we set $f(x)=y$, and we have

$\displaystyle\texttt{Prec}^{r}(f,\mathcal{Q}_{d},q,\omega)$
$\displaystyle=\left|\sum_{j=1}^{n}f(\xi_{j})|Q_{j}|-\sum_{j=1}^{n}q(f(\xi_{j}))|Q_{j}|\right|$
$\displaystyle=\left|\frac{1}{n}\sum_{j=1}^{n}\left(f(\xi_{j})-q(f(\xi_{j}))\right)\right|$
$\displaystyle\geq\left|\frac{1}{n}\sum_{j=1}^{n}\left(\frac{1}{2}\cdot\epsilon
y\right)\right|$ $\displaystyle=\frac{1}{2}\cdot\epsilon y$
$\displaystyle\geq\frac{1}{4}\cdot\epsilon M$

This concludes the proof. ∎

###  A.3 Plotting Theoretical Bounds

In this section, we run additional experiments by plotting our theoretical
bounds alongside the true empirical discretization and resolution errors of
the Darcy flow dataset. Specifically, we plot both our bounds assuming a
Fourier basis (Theorem 3.1, Theorem 3.2) and our bounds for general functions
(Theorem A.1, Theorem A.2) compared to the true Darcy flow dataset, measured
after 10 epochs, at the start of the FNO block (just before the forward FFT).
The true errors are calculated using the definitions of precision and
discretization error in the previous section, and we use the true difference
in precision between float32 and float16 when computing the precision error.
See Figure 7. 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 a 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 7.

Figure 7:  Discretization and precision errors of Darcy flow 1D and 3D,
compared to our Fourier basis and general bounds. Note that our bounds are
worst-case bounds (we showed there exists a 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.

##  Appendix B Additional Details from Section 4

In this section, we give additional details on the datasets and experimental
setup, as well as additional experiments.

###  B.1 Ahmed-Body CFD Experiment

Similar to Shape-Net Car, we train the Geometry-Informed Neural Operator on
the Ahmed-body dataset, using the same hyperparameters as the original
implementation (Li et al., 2023). The memory, error, and throughput results
for Ahmed-body are appended to Figure 1. We add the L2 error curves here
following the format of Figure 5.

Figure 8:  Test L2 error curves for GINO on 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 $\mathbb{T}^{2}\cong[0,2\pi)^{2}$,

$\displaystyle\partial_{t}\omega+\nabla^{\bot}\phi\cdot\omega=\frac{1}{\text{Re}}\Delta\omega+f,\quad$
$\displaystyle\text{for }x\in\mathbb{T}^{2},~{}t\in(0,T]$
$\displaystyle-\Delta\phi=\omega,\quad\int_{\mathbb{T}^{2}}\phi=0,\quad$
$\displaystyle\text{for }x\in\mathbb{T}^{2},~{}t\in(0,T]$ (41)

where $f\in L^{2}(\mathbb{T}^{2},\mathbb{R})$ is a forcing function and
$\text{Re}>0$ is the Reynolds number, for the initial condition
$\omega(0,\cdot)=0$. The goal is to learn the non-linear operator
$\mathcal{G}^{\dagger}:f\mapsto\omega(T,\cdot)$ for $T=5$ and $\text{Re}=500$.
We use the dataset from prior work (Kossaifi et al., 2023), which sets
$\text{Re}=500$ and generates forcing functions from the Gaussian measure,
$\mathcal{N}(0,27(-\Delta+9I)^{-4})$, 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\times 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 $a\in
L^{\infty}((0,1)^{2};\mathbb{R}_{+})$, and the forcing function $f\in
L^{2}((0,1)^{2};\mathbb{R})$, the pressure $u$ satisfies

$\displaystyle-\nabla\cdot(a(x)\nabla u(x))=f(x),\quad$
$\displaystyle\text{for }x\in D$ (42) $\displaystyle u(x)=0,\quad$
$\displaystyle\text{for }x\in\partial D$ (43)

We seek to learn the non-linear operator $\mathcal{G}^{\dagger}:a\mapsto 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\equiv 1$
and generates 5000 training and 1000 test samples with resolution $128\times
128$.

#### Spherical Shallow Water Equations.

The spherical Shallow Water Equations on the rotating sphere are a model for a
variety of geophysical flow phenomena such as atmospheric flows, tsunamis and
storm surges. We consider the evolution of two state variables $\varphi$, $u$
(geopotential height and tangential velocity of the fluid column), governed by
the equations

$\displaystyle\partial_{t}\varphi+\nabla\cdot(\varphi u)=0,\quad$
$\displaystyle\text{for }x\in\mathbb{S}^{2},t\in(0,T]$ (44)
$\displaystyle\partial_{t}(\varphi u)+\nabla\cdot F=S,\quad$
$\displaystyle\text{for }x\in\mathbb{S}^{2},t\in(0,T]$ (45)

for initial conditions $\varphi(\cdot,0)=\varphi_{0}$, $u(\cdot,0)=u_{0}$,
where $\mathbb{S}^{2}$ is the unit sphere, $F$ the momentum flux tensor with
$F^{ij}=\varphi u^{i}u^{j}+\frac{1}{2}\varphi^{2}$, and $S$ a source term
$S=-2\Omega x\times(\varphi u)$, which models the Coriolis force due to the
rotation of the sphere with angular velocity $\Omega$. We use the dataset from
prior work (Bonev et al., 2023), which generates random initial conditions on
the sphere at resolution $256\times 512$ and solves the PDE using the spectral
solver in the torch-harmonics package (Bonev et al., 2023).

#### Shape-Net Car

We evaluate on 3D real-world car dataset generated by prior work (Umetani &
Bickel, 2018; Li et al., 2023) using the Reynolds Averaged Navier Stokes
(RANS) equations (Wilcox et al., 1998). Each data point consists of mesh
points representing a unique 3D car, and the goal is to predict the full 3D
pressure field, given an inlet velocity of 20m/s.

#### Ahmed-Body CFD

Including another 3D real-world dataset for evaluation, we the dataset from
(Li et al., 2023), which is based on Ahmed-body shapes from (Ahmed et 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 to those of Shape-Net Car, but
Ahmed-body has many more mesh points (100k) on the surface.

###  B.3 Model Details

We use the official implementation of each model, with the default
hyperparameters, unless otherwise specified.
111https://github.com/neuraloperator/neuraloperator Furthermore, we release
all of our code at https://github.com/neuraloperator/neuraloperator.

For the case of Shape-Net Car and Ahmed Body on GINO, we note that the default
batch size is 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-
geometry datasets are too large to fit in memory, so sub-sampling is used (Li
et 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 Ti, 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 PyTorch modules and by individual GPU
kernels. We observe that complex-valued operations are accountable for the
majority of the runtime.

Figure 9:  For Navier-Stokes, we use the PyTorch 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 4 amongst the 5 most costly kernels.

###  B.5 Numerical Stability: Post-Forward Methods

Most commonly, numerical underflows and overflows that occur during mixed-
precision training are handled after the forward pass using loss scaling,
gradient clipping, and other methods (Micikevicius et 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 a no-stabilizer baseline on the Darcy Flow experiment
setting. In Figure 10, we 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.

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).

###  B.6 Numerical Stability: Pre-FFT Methods

Here, we give additional details and experiments for our study on numerical
stability of mixed-precision FNO.

A simple idea for addressing overflow in FP16 is to scale down all values by
adding a 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 a 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:
$10^{1},10^{2},10^{3}$. At $10^{4}$, the network completely fails to learn due
to vanishing gradients. We also find that adding a 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.

Table 3 compares several pre-FFT stabilization methods that can address
numerical instability with acceptable computational costs. hard-clip forces
maximum and minimum values in the input to a default range. 2$\sigma$-clip
similarly performs clipping but based on the data’s mean and standard
deviation. Out of these methods, tanh is founded to be the highest-performing.
After adopting tanh as our pre-activation function, we assess its impact on
the signal using a trained model, as shown in Figure 11. We find that the loss
of signal information is minimal in amplitude and phase.

Fourier | AMP | Pre- | Runtime | Train
---|---|---|---|---
Precision | (T/F) | Activation | per epoch | loss
Full | F | N/A | 44.4 | 0.0457
Full | T | N/A | 42.3 | 0.0457
Half | F | N/A | N/A | N/A
Half | T | hard-clip | 37.1 | 0.0483
Half | T | 2$\sigma$-clip | 40.0 | 0.0474
Half | T | tanh | 36.5 | 0.0481
Table 3: Comparison of different pre-activation functions used for numerical
stability. tanh achieves the fastest runtime without significantly
compromising on loss. Since our focus now is on numerical stability, we
compare only the average train loss over the 5 final epochs (we compare the
test losses in Section 4.4).

Figure 11:  Impact of pre-activation on frequency domain signals, for a fully-
trained full FNO model (Left) and a mixed FNO model (Right),  in terms of mean
magnitude and phase differences, on a minibatch of the Navier-Stokes dataset.
Pre-activation only changes a extremely small fraction of frequencies, and the
two signals are well-aligned in phase.

###  B.7 FNO Block Precision Ablation Study

As shown in Figure 2, our FNO block has three complex-valued operations that
are 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 on FNO training. This creates a 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 on the
Darcy Flow dataset for each setting. The empirical results demonstrates the
acorss-the-board advantage of our approach, which keeps all three operations
in half-precision in the forward pass.

Forward FFT | Contraction | Inverse FFT | Runtime per epoch | Memory | Train Error
---|---|---|---|---|---
(F/H) | (F/H) | (F/H) | (sec) | (MB) | (L2 loss)
F | F | F | 17.06 | 8870 | 9.00
F | F | H | 16.55 | 8908 | 8.71
F | H | F | 17.11 | 8614 | 8.76
F | H | H | 16.96 | 8658 | 8.15
H | F | F | 17.64 | 7824 | 9.13
H | F | H | 16.81 | 8004 | 7.51
H | H | F | 16.57 | 7558 | 8.75
H | H | H | 15.63 | 7550 | 7.49
Table 4: Performance of each setting on Darcy Flow trained for 30 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 is
in half-precision.

###  B.8 Tanh Ablation Study

A natural question from using pre-activation is to ask how a hyperbolic
tangent would affect the full precision FNO model with no other changes. In
Table 5, we answer this question by running the full precision FNO model on
the Navier-Stokes dataset. We find that there is no noticeable change in the
test losses, and the runtime-per-epoch increases by 0.8 seconds, or 1.6%. This
concludes that tanh is a reasonable and practical choice to improve numerical
stability for half precision methods.

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 tanh is a practical choice to improve numerical stability in low precision methods.  | $H^{1}$ | $L^{2}$ | time-per-epoch (sec)
---|---|---|---
Full precision | 0.0121 | 0.00470 | 51.72
Full precision + tanh | 0.0122 | 0.00465 | 52.56

###  B.9 Comparing the $H^{1}$ Loss

We present additional experiments similar to the results in Section 4, but
using the $H^{1}$ metric instead of $L^{2}$. In Figure 12, we plot an ablation
study on the number of frequency modes for Darcy flow, similar to Figure 14
with $L^{2}$ loss. We see that the overall trends are similar, but test
$L^{2}$ loss is noisier than test L1 loss; this makes sense, because the
training loss is $H^{1}$. In Figure 13, we plot the performance-vs-time and
Pareto frontier plots for Navier Stokes and Darcy flow, similar to Figure 5
with $L^{2}$ loss. As before, we see largely the same trends as with $H^{1}$
loss, but the results are noisier. In Table 6, we plot the final performance
from the models run in Figure 5.

Figure 12:  Comparison of the full-precision and mixed-precision
(AMP+HALF+TANH) FNO with different frequency modes, on the Darcy flow dataset.

Figure 13:  Test $H^{1}$ error curves for FNO on the Navier-Stokes (top left) and Darcy flow (top right) datasets. Pareto frontier for FNO on the Navier-Stokes (bottom left) and Darcy flow (bottom right) datasets, for Canonical-Polyadic factorization (CP) or no factorization (Dense). We train each model for 500 epochs, and we 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 are the average over 3 seeds trained to 19 hours each (which is the time it takes full precision FNO to reach 500 epochs). | $H^{1}$ | $L^{2}$ | time-per-epoch (sec)
---|---|---|---
Full FNO | $.00536\pm 2.1e-5$ | $.00214\pm 3.5e-5$ | 121.4
Mixed FNO (Ours) | $.00645\pm 6.6e-5$ | $.00212\pm 1.4e-5$ | 80.2
Precision schedule (Ours) | $.00515\pm 8.3e-5$ | $.00812\pm 4.1e-5$ | 80.2, 83.8, 121.4
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 to full precision) is higher than for 32, 64, or 128 modes.

###  B.10 Frequency Mode Ablation Study

We run an experiment on synthetic data to demonstrate that the error caused by
half precision is higher for higher frequencies, relative to the amplitude. We
create a signal based on sine and cosine waves with frequencies from 1 to 10,
with randomly drawn, exponentially decaying amplitudes. Then we plot the
Fourier spectrum in full and half precision, as well as the absolute error of
the half-precision spectrum as a 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.

Figure 15:  Synthetic signal (top), its frequency modes (middle), and the error due to half precision, as a percentage of the amplitude (bottom). The percentage error increases for higher frequencies.  | Navier-Stokes | Darcy Flow
---|---|---
FNO + TF32 | 57.44 | 14.12
Mixed FNO (Ours) | 53.72 | 13.52
Table 7: Training time per epoch on a Nvidia A100 GPU.

###  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. In addition to suffering from 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 A100 GPU, on which TF32 is highly
optimized. We report the time per epoch when training on Navier-Stokes and
Darcy Flow datasets in Table 7. Our method is still more efficient in
comparison.

In addition, we simulate 8-bit floating point (FP8) training using our method
on the Darcy 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 as shown in
Figure 16.

Furthermore, the worse result of FP8, is actually predicted Theorem 3.2: while
the precision constant for FP16 corresponds to $\epsilon=10^{-4}$, it
corresponds to $\epsilon>10^{-2}$ for FP8 (even when using E4M3), and we can
no longer claim that the precision error is lower than the resolution error.

Figure 16:  BF16 has significantly higher error 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.

###  B.12 Ablation Study on Contract and View-as-real Usage

This ablation study shows the benefit of (1) decomposition of einsum into sub-
equations and appropriate view-as-real usage, (2) our einsum path caching, (3)
our greedy tensor contraction procedure, and (4) both weights and inputs in
half-precision.

The results are as follows:

  1. 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 in one
equation (naive);

     * •

Option B: View-as-real on two tensors at a time, and compute step-by-step
einsums with two-tensor sub-equations (optimized);

     * •

Option C: View-as-real only on high-dimensional einsums, and contracting low-
dimension sub-equations in complex format (ours, optimal).

Results on the Navier-Stokes can be found in Table 8.

  2. 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 we 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. 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 to allow higher-resolution PDEs. Results on Ahmed-body
and Shape-Net Car could be found in Table 10. On memory-consuming 3D datasets,
our method could save up to 12%.

  4. 4.

Approximate only weights in half-precision (naive) vs. inputs and weights both
half-precision (ours):

If only weights were in half-precision, the memory reduction would have been
significantly limited for the contract operation since the inputs are in full-
precision. Results on Darcy Flow and Navier-Stokes can be found in Table 11.
Note that in the case of Navier-Stokes, pytorch uses a much less memory-
efficient kernel if inputs were in full-precision. Hence, it is critical to
cast both weights and inputs to half-precision.

Table 8: Comparison of tensor contraction implementations on 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 vs. Caching contract paths Dataset | Time to calculate paths (s) | Time to compute einsums (s) | Path/einsum (%)
---|---|---|---
Navier-Stokes | 0.000573 | 0.000751 | 76.3
Darcy Flow | 0.000441 | 0.000716 | 61.6
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 vs. inputs and weights Dataset | Half-Prec Memory (MB) | Inputs-Full Memory (MB) | Reduction (%)
---|---|---|---
Darcy Flow | 7550 | 8166 | 7.54
Navier-Stokes | 4832 | 9380 | 48.46