Multi-Grid Tensorized Fourier Neural Operator for High- Resolution PDEs
Authors: Jean Kossaifi
Nikola Borislavov Kovachki
Kamyar Azizzadenesheli
Anima Anandkumar
Venue: Accepted by TMLR ()

Abstract: Memory complexity and data scarcity have so far prohibited learning solution operators of partial differential equations (PDE) at high resolutions. We address these limitations by introducing a new data-efficient and highly parallelizable operator learning approach with reduced memory requirement and better generalization, called multi-grid tensorized neural operator (MG-TFNO). MG-TFNO scales to large resolutions by leveraging local and global structures of full-scale, real-world phenomena, through a decomposition of both the input domain and the operator’s parameter space. Our contributions are threefold: i) we enable parallelization over input samples with a novel multi-grid-based domain decomposition, ii) we represent the parameters of the model in a high-order latent subspace of the Fourier domain, through a global tensor factorization, resulting in an extreme reduction in the number of parameters and improved generalization, and iii) we propose architectural improvements to the backbone FNO. Our approach can be used in any operator learning setting. We demonstrate superior performance on the turbulent Navier-Stokes equations where we achieve less than half the error with over 150× compression. The tensorization combined with the domain decomposition, yields over 150× reduction in the number of parameters and 7× reduction in the domain size without losses in accuracy.

Published in Transactions on Machine Learning Research (08/2024)
Multi-Grid Tensorized Fourier Neural Operator for High- Resolution PDEs
Jean Kossaifi'
jkossoff@nvidia.com com
NVIDIA
Nikola Kovachki
nkovachii@nridiaacom
NVIDIA
Kamyar Azizzadeneheeli NVIDIA
kamyoro@nvidia..cm
Anima Anandkumar
anima@caltech.. edu
Caltech
Reviewed on OpenReview:
https: openreview. net/ forum? id= AWaD1063bH
Abstract
Memory complexity and data scarcity have so far prohibited learning solution operators of partial differential equations (PDE) at high resolutions. We address these limitations by introducing new data efficient and highly parallelizable operator learning approach with reduced memory requirement and better generalization, called multi-grid lensorized neural operator (MG-TFNO). MG-TFNO scales to large resolutions by leveraging local and global structures of full-scale, real-world phenomena, through docomposition of both the input domain and the operator's parameter space. Our contributions are threefold: i) we enable parallelization over input samples with novel multi-grid-basse domain decomposition, ii) we represent the parameters of the model in a high-order latent subspace of the Fourier do- main, through a global tensor factorization, resulting in an extreme reduction in the number of parameters and improved generalization, and iii) we propose architectural improvements to the backbone FNO. Our approach can be used in any operator learning setting. We demonstrate superior performance on the turbulent 2D Navier-Stokes equations where achieve less than half the error with over 150x compression compared to the FNO baseline. The tensorization combined with the domain decomposition, yields over 150x reduction in the number of parameters and 7x reduction in the domain size without losses in accuracy.
1 Introduction
Real-world scientific computing problems often time require repeatedly solving large-scale and high-resolution ren partial differential equations (PDEs). For instance, in weather forecasts, large systems of differential equa- tions are solved to forecast the future state of the weather. Due to internal inherent and alcatoric uncertain- ties, multiple repeated runs carried out by meteorologists every day to quantify prediction uncertainties. Conventional PDE solvers constitute the mainstream approach used to tackle such computational problems. Such methods provide convergence guarantess and explicit error bounds to the true solution and have long served as the backbone for tackling many engireening problems. However, these methods are known to be slow and memory-intensive. They require an immense amount of computing power, are unable to learn and adapt based on observed data, and oftentimes require sophisticated tuning (Siinoo & Palmer 2011 Leutbecher Palmer 2008 Blanusa et al. 2022
These authors contributed equally this work

---

Published in Transactions on Machine Learning Research (08/2024)
5.0%
4.5%
4.0%
3.5%
3.0%
h1 (lower is better)
FNO [Ours]
67.14M
2.0%
1.5%
67.18M
1.0%
TFNO [Ours]
3.56M
0.6%
1.0%
1.2% 1.4%
12 (lower is better)
Figure 1: Comparison of the performance on the relative L and H' test crrors (lower is better) log-scale of our approach, compared with both our improved backbone (FNO) and the original FNO, on Navier-Stokes. Our approach enables large compression for both input and parameter, while outperforming regular FNO.
Neural operators are a new class of models that aim at tackling these challenging problems (Li al, 2020a). They maps between function spaces whose trained models emulate the solution operators of PDEs Kovachki ct al., 2021b). In the context of PDEs, these deep learning models are orders of magnitude faster than conventions ional solvers, can easily learn from data, incorporate physically relevant information, and recently enabled solving problems deemed to be unsolvable with the current state of available PDE methodologics (Liu et al. 2022 Li ct al. 2021c Among neural operator modols, Fouricr neural operators (FNOs), in particular, have seen successful application scientific computing for the task learning the solution operator to PDFs as well as in computer vision for classification, in-painting, and segmentation 2021b
(Li al. successfully advanced frontiers in weather forecasts, carbon storage, and scismology (Pathak al. 2022
Kovachki al.) 2021a Guibas et al. 2021). By leveraging spectral theory, FNOs have
Wen et al.
2022
Yang et al. 2021
While FNOss have shown tremendous speed-up over classical numerical methods, their efficacy can be limited due to the rapid growth in memory needed to represent complex operators. In the worst case, large memory complexity required and, fact, unavoidable due to the need for resolving fine-scale features globally. However, many real-world problems, possess local structure not currently exploited by neural operator methods. For instance, consider weather forecast where predictions for the next hour are heavily dependent the weather conditions in local regions and minimally global weather conditions. Incorporating and learning this local structure of the underlying PDEs is the key to overcoming the curse of memory complexity.
In this work, we propose a new, scalable neural operator that addresses these issues by leveraging the structure in both the domain space and the parameter space, Figure Specifically, we introduce the multi- grid tensor operator (MG-TFNO), a model that exploits locality in physical space by a novel multi-grid domain decomposition approach compress the input domain size by up to 7x while leveraging the global interactions of the model parameters to compress them by over 100x without any loss of accuracy.
2

---

Published in Transactions on Machine Learning Research (08/2024)
F
Tensor Operator Block
Tensor Operator Block
Level
Multi-Grid Domain Decomposition
Joint Parameter Decomposition
Overview of our approach. First (left), multi-grid approach is used to create coarse fine inputs that capture high-resolution details in a local region while still encoding global context. The resulting regions are fed to a tensorized Fourier operator (middle), the parameters of which are jointly represented in single latent space via a low-rank tensor factorization (here, a Tucker form). Here F denotes Fourier transform. Finally, the outputs (right) are stitched back together to form the full result. Smoothness in the output ensured via the choice of the loss function.
Figure 2:
In the input space, predict the solution in any region of the domain, MG-TENO decomposes the input domain into small local regions to which hierarchical levels of global information are added in a multi-grid fashion. Since a local prediction depends most strongly on its immediate spatial surroundings, the farther field information is downsampled to lower resolutions, progressively, based on its distance from the region of interest. Thus, MG-TFNO allows parallelization over the input domain as relies on high-resolution data only locally and coarso-resolutio data globally. Due its mato-of-the-art performance on PDE problems and efficient FFT-based implementation, we use the FNO the backbone architecture for our method. It is worth noting that the multi-g multi-grid approach is readily amendable to neural network settings and, moreover, any other neural operator architecture can be used in place of FNO as backbone.
In the parameter space, we exploit the spatiolemporal structure of the underlying PDE solution operator by parameterizing the convolutional weights within the Fourier domain with a low-rank tensor factorization. Specifically, we impose a coupling betwoen all the weights in the Fourier space by jointly parameterizing them with a single tensor, learned factorized form such as Tucker or Canonical- Polyadii (Kolda Bader 2009). This coupling allows to limit the number of parameters in the model without limiting its expressivity. On the contrary, this low-rank regularization Intuitively, our method can be thought of as a fully-learned implicit scheme capable of converging in a small, fixed number of iterations. Due to the global nature of the integral kernel transform, the FNO avoids the only a few steps (Courant et al. 1928). Our weight coupling ensures maximum communication between the steps, mitigating possible redundancies in the learned kernels and imposing low-rank regularizing of the network Panagakis et al. (0021b).
the model mitigates over-fitting and improves generalization.
(CFL) condition plaguing explicit schemes, allowing convergence in
In summary, we make the following contributions:
We propose architectural improvements to the backbone which we validated through thor- ough ablations.
We propose MG-TFNO, a neural operator parameterized the spectral domain by single low- rank factorized tensor, allowing its size to grow linearly with the size of the problem. It achieves better performance with fraction of the parameters: outperform FNO on solving the turbulent
3

---

Published in Transactions on Machine Learning Research (08/2024)
Navier Stokes equations with more than 400x weight compression ratio, Figure 7 Our method overfits less and does better in the low-data regime. In particular, it outperforms FNO with less than half the training samples, Figure 9
We introduce a novel multi-grid domain decomposition approach, a technique which allows the operator to predict the output only on local portions of the domain, thus reducing the memory usage by an order of magnitude with performance degradation. Combining tensorization with multi-grid domain decomposition leads to MG-TFNO, which is more efficient in terms of task per- formance, computation, and memory. MG-TENO achieves 2.5x lower error with 10x model weight compression, and 1.8x domain compression.
We release an open-source implementation of our proposed method to run all variations of FNO and MG- TFNO, well as the Navier-Stokes data used in this paper, along with the paper.
2 Background
Here, we review related works and introduce the background necessary to explain our approach.
Many physical phenomena are governed by PDEs and wide range of scientific and engineering computation problems are based on solving these equations. In recent years, new perspective to PDEs dictates to formulate these problems as machine learning problems where solutions to PDEs are learned. Prior works mainly focused on using neural networks to train for the solution map of PDEs (Guo et al. 2016 Zhu & Zabaras 2018 Adler & Oktem 2017 Bhatnagar et al. 2019; Gupta et al. 2021]. The use of neural networks in the prior works limits them to a fixed grid and narrows their applicability to PDEs where maps between function spaces are desirable. Multiple attempls have been made to address this limitation. For example
mesh free methods are proposed that locally output mesh-free solution (Lu et al. 2019 Esmaeilzadeh et al. 2020), but they are still limited to fixed input gird.
A new deep learning paradigm, neural operators, are proposed as maps between function spaces (ii et al. 2020a; Kovachki et al., 2021b). They are discretization invariants maps. The input functions to neural operators can be presented in any discretization, mesh, resolution, or basis. The output functions can be evaluated at any point in the domain, either on regular grid (e.g. FFT based architectures) or for any ar- bitrary geometry (e.g- graph based architectures). Variants of neural operators deploy a variety of Nyström approximation to develop new neural operator architecture. Among these, multi-pole neural operators (Li et al.] 2020b utilize the multi-pole approach to develop computationally elficient neural operator archi- tecture. Inspired by the spectral method, Fourier-based neural operators show significant applicability in
practical applications şLi et al.) 2021b Yang et al., 2021 Wen et al., 2022 Rahman ct al. 2022a). and the architectures have been used in neural networks for vision and text tasks Guibas et al. 2021 Dao el al. 2022). Principle component analysis and u-shaped methods are also considered (Bhattacharya et al. 2020 Liu et al. 2022: |ahmman al., 2022b Yang et al.) (022). It is also shown that neural operators can be additionally trained using PDEs, resulting in physics-ifformed neural operators, opening new venues for hybrid data and equation methods (Li et. al. 2021c) to tackle problems in scientific computing.
Decomposing the domain in smaller subdomains is at the core of many methods in computational sci- ences (ryjaa & Widlund 1990 Chan & Mathew 1994) and extensively developed in deep learning (Dosovit-] skiy et al. 2020). Prior deep learning methods on neural networks propose to decompose the input finite dimension vector to multiple patchos, accomplish local operations, and aggregate the result of such process in the global sense (Dosovitskiy et al., 22000; Guibas et al., 2021). ). Such methods do not decompose the output domain and directly predict the entire output vector. In contrast, MG-TENO works function spaces, and not only decomposes the input domain, but also decomposes the domain of the output functions, and separately predicts the output at each subdomain.
For the purposes of this paper, tensors are multi-dimensional arrays and generalize the concept of matrices to more than 2 modes (dimensions). For instance, RGB images are encoded as third-order (threo-dimensional) tensors, videos are 4th order tensors and so on and so forth. Tensor methods generalize linear algebraic methods to these higher-order structures. They have been very successful in various applications in computer
4

---

Published in Transactions on Machine Learning Research (08/2024)
vision, signal processing, data mining and machine learning Panagakis et al.
2021b Janzamin et al. 2019b;
Sidiropoulos et al. 2017 Papalexakis et. al.
2016).
Janzamin et al. 2019a previous works have been able to compress and improve deep networks for vision tasks Panagakis et al. (021a). Either a weight matrix is tensorized and factorized Novikov et al.
Using tensor decomposition Kolda & Bader
(2009);
directly to the convolutional kernels before fine-tuning to recover-for lost accuracy, which also allows for an efficient reparametrization of
(2015).
or tensor decomposition
the network (Lebedev et al., 2015 Kim et al.) 2016 Gusak et al. 2019). There is tight link between efficient convolutional blocks and tensor factorization and factorized higher-order structures (Kossaifi et al., 2020). Similar strategies have been applicd to multi-task learning (Bulat et al. 2020a) and NLP (Papadopoulos et al. 2022 Cordonnier et al. 2020). Of all these prior works, none has been applied to neural operator. In this work, we propose the first application of tensor compression to learning operators and propose a Tensorized Fourier Neural Operator (TFNO).
3 Methodology
Here, we briefly review operator learning as well as the Fourier Neural Operator, on which we build to introduce our proposed Tensor OPerator (TFNO) as well as the Multi-Grit Domain Decomposition, which together form proposed MG-TFNO.
3.1 Operator Learning
Let := {a DA - RRda) and U == (u Du Rdy) denote two input and output function spaces respectively. Each function a, in the input function space A, is a map from a bounded, open set DA C Rd to the dx-dimmesional Euclidean space. Any function in the output function space U is map from bounded open set Du Rd to the dy-dimensional Euclidean space. In this work we consider the case D D Du C IRd.
We aim to learn an operator A + U which is a mapping between the two function spaces. In particular, given a dataset of N samples (e.g. PDE solves) {(a,, uy)) where the pair (aj, uy) are functions satisfying g(a,) = uj, we build an approximation of the operator G As a backbone operator learning model, we use neural operators as they consistent and universal learners in function spaces. We refer the reader to Kovachki et al. (2021b) for overview of theory and implementation. We specifically use the FNO and give details in the forthcoming section (Li ct al. 2021b
3.2 Notation
We summarize the notation used throughout the paper in Table
3.3
Fourier Neural Operators
For simplicity, we will work on the d-dimensional unit torus T and first describe single, pre-activation FNO layer mapping valuue functions to -vauudd functions. Such layer constitutes the mapping G L (TTTRRR + L (Td, RR") defined
(F))= F(o),, Ve I.
(1)
where x (T4, (TT:TTax)) Baxm) ) is function constituting the layer parameters and F,J are the Fourier transform and its inverse respectively. The Fourier transform of the function is parameterized directly by some fixed number of Fourier nodes denoted E N.
To implement equation I F,F-1 are replaced by the discrete fast Fourier transforms f.f-1 Let E X=-xxmm denote the evaluation the function y on a uniform grid discretizing T with 8, E N points in cach direction. We replace F(k) with weight tensor T Co X--xxxxm consisting of the Fourier modes of x which are parameters to be learned. To ensure that k is parameterized Raludvalued function with fixed, maximum amount of wavenumbers
5

---

Published in Transactions on Machine Learning Research (08/2024)
Variable
Meaning
Dimensionality
T
Tensor of weights in the Fourier domain
Coy xxxxmxxx
W
Weight tensor parameterizing the entire operator
Covax (oyaxxxXxxxaa Xaxxxxx2*
A
Input function space
Infinite
4
output function space
Infinite
a
Input function Output function
Infinite Infinite d d
DA
Domain of function a
Du
Domain of function
A
Dimension of the co-domain of the input functions Dimension of the co-domain of the output functions Fourier transform
du
F
Infinite Infinite In Between and Infinite
Fourier transform
Number of integral operation layers
Layer index
Point-wise activation operation
b
Bias vector
Function at each layer
Infinite Between 1 and 1 min(s1, ,d)
Number of kept frequencies in Fourier space
Table 1: Table of notation
2 min ... s that is independent the discretization T we leave as learnable parame- ters only the first a entrics of T in each direction and enforce that T have conjugate symmetry.
In particular, parameterize half the corners of the d-dimensional hyperrectangle with 2d-1 hyper- cubes with length size a. That is, is made up of the free- parameter tensors T T..... T2--1 e Cx xaxnxm situated in half of the corners of T. Each corner diag- onally opposite of tensor T, is as- signed the conjugate transpose val- ues of TT. All other values of T are set to zero. This is illustrated the middle-top part of Figure 3.3 for the case d with Ti and Tg. Note that, equivalently, ply FFT-shift to the output of the real FFT so the low-frequency coefficients are located in one block in the middle. that case, we then apply another FFT-shift to the output of the contraction before applying the inverse real FFT. We will use the notation ..). Th for any E [2d-1]. The discrete version of equation then becomes the mapping Ś R X-axm Rexx xeexx--Xaxa defined as
Low high
yuy wo
Kept modes
Redunddart crefficionts
Positive frequencies
Lowe to High
Redundant coeficinss
Kept modes
Negative frequenctes
ytu on Mon
Low high
Kept modes
Real-FF FFT Coefficients
Real-FFT Coefficients after FFT-Shift
Figure 3: Visualization the hyper-corners containing the low- frequency in the real-FFF. On the left, we show the low-frequency modes used by our algorithm. Equivalently, can first apply an FFT-shift (right) and keep the central coefficients.
can ap-
(0) (T (()), X
(2)
where the operation is simply tensor contraction along the last dimension. Specifically, have

---

Published in Transactions on Machine Learning Research (08/2024)
(T (()) (a.5) Σ T(h Ma (4..) 2),
(3)
From equation?? 2 a full FNO layer build by adding a point-wise linear action to 0, bias term, and applying non-linear activation. In particular, from an input E X=xxxx the output I E R Xmxxax given
q(1 ... La. - o (QQ(() .. .a,;) +(6) b)
with 0 R R. a fixed, non-linear activation, and b R" Q Rmxm, T TT.... Tze-1 the learnable parameters of the layer. The full FNO model consists of L E N such layers each with weight
Cov Xx are
tensors TI. Ty that have learnable parameters T Tk....... for any LE [L] and & E [2d In the case m for all layers, we introduce the joint paramcter tensor W E xXxxxxxXXXXxxXX2
so that
W (..., 1) + k: + TT
(4)
A perusal of the above discussion reveals that there (2 + 1)mn + + total parameters in each FNO layer. Note that, since m and n constitute the respective imput and output channels of the layer, the number of parameters can quickly explode due to the exponential scaling factor 24 if many wavenunbers are kept. Preserving a large number of modes could be crucial for applications where the spectral decay of the input or output functions is slow such as image processing or the modeling of multi-scale physics. In the following section, we describe a tensorization method that is able mitigate this growth without sacrificing approximation power.
3.4 Architectural improvements
Our proposed approach uses FNO as a backbone. To improve its performance, we first study various aspects of the Fourier Neural Architecture and perform thorough ablation to validate each aspect. In particular, we propose improvements to the base architecture that improve performance.
Normalization in neural operators While normalization techniques, such as Batch-Normalization Tofle & Szegedy (2015), have proven very successful in training neural networks, additional consideration must be given when applying those to neural operators in order to preserve its properties, notably discrctization invariance. Specifically, it cannot depend on the spatial variables and therefore has to be either a global or function-wise normalization. We investigate several configurations using instance normalization Ulyanov et al. (2016) and layer-normalization Ba et al. (2016), in conjunction with the use-of preactivation He et al. (2016
Channel mixing FNO relies on global convolution realized in the spectral domain. Inspired by previous works, e.g. Guibas et al. 2021), propose adding an ML. in the original space, after each Spectral convolution. In practice, found that two-layer bottleneck MLP works well, c.g. decrease the dimension by half in the first linear layer before restoring it in the second one.
Boundary conditions: domain padding and skip connections Fourier neural operators circumvent the limitation of traditional Fourier methods to inputs with periodic boundaries only, This is achieved through local linear transformation added to the spectral convolution. This can be seen as lincar skip connection. We investigate replacing these with an identity skip-connection and a soft-gated skip- connection Bulat et al. 2020bb
We represent in Figure. the original FNO architocture Li et al. 2021b), subfigure 4(a) the improved version with double (sequential) skip connections (subfigure 4(b) and our best architecture, both with and without preactivation (subfigures [[c] and ((d) respectively).

---

Published in Transactions on Machine Learning Research (08/2024)
Spectral
G+LD
Spectrakcon
Stip
Loyen-orm
LayerNorm
Skip
Lape-Norm
Skip
Spectral Conv
Skip
Additioon
Spectral Conv
Additiion
GeLU
Skep
Addition
34ip
Addition
MIP
GelU
MiP
Sklp
Layer-Norm
Layer-Nomm
Norm
GeUU
Addition
Addition
GeUU
((v)
((m)
G(V)
g(v)
(a) Original FNO.
(b) Double Skip.
(c) Improved backbone pre-
(d) Improved backbone.
activation.
Figure 4: Original FNO and Improved Backbone Architecture. The original FNO architecture Li
et al. 2021b) is composed of simply Spectral Convolution, with a (linear, soft-gating, or identity) skip connection to recover high-frequency information and handle non-periodic inputs [4(a)). We improve the architecture as detailed in section 3.4 In particular, we have a version with a double (sequential) skip
connection 4(b) while our best architecture uses nested skip connections, and can be made both with pre-
or post-activation (subfigures 4(c) and ((d) respectively). The latter, subfigure {(d) is our best architecture.
We also investigate the impact of domain-padding, found by Li (2021b) to improve results, especially for non-periodic inputs, and padding for the multi-grid decomposition. In fact, the multi-grid domain decompasition inherently breaks any periodic boundary conditions (if the function itself is periodic, its patches typically likely are not). The effect of this can be mitigated by padding the inputs, which increases computational cost, resulting in tradeoff between padding and computational efficiency.
3.5 Tensor Fourier Neural Operators
In the previous section, introduced unified formulation of FNO where the whole operator is parametrized by a single parameter tensor W. This enables us to introduce the tensor operator, which parameterizes ef- ficiently W with a low-rank, tensor factorization. We introduce the method for the case of a Tucker decom- position Tucker (1966c), for its flexibility. Other decompositions, such as Canonical Polyadic (CP) Carroll & Chang (1970): : Harshman (1970), can be readily integrated. This joint parametrization has several ad- vantages: it applics low-rank constraint on the entire tensor W, thus regularizing the model. These advantages translate into huge reduction in the number of parameters, ii) better generalization and an operator less prone to overfitting. We show superior performance for cow-ommressioon ratios (up to 200x) and very little performance degradation when largely compressing (> 450x) ) the model, iii) better performance in a low-data regime.
In practice, we express W in a low-rank factorized form, e.g. Tucker or CP. In the case a Tucker factorization with rank .R, R2, R,, R,, Ro), where R, controls the rank across layers, R₁ Ro control the rank across the input and output co-dimension, respectively, and R1.... Ry control the rank across the dimensions of the operator:

---

Published in Transactions on Machine Learning Research (08/2024)
R G (a,t,, (W,)))) ... . (.ra), (1,r)) (U(O)) ((
= Σ
(5)
Here, G is the of size RS x R ..... Ra and U U U U U(() are factor matrices of size (Ry x I.), (Ry x I), (Ro 0), (R₃ a),,.. ,,,...,,,, (Ra x a); respectively.
Note that the mode (dimension) corresponding to the co-dimension can be left uncompressed by setting R, L and U Id. This leads to layerwise compression. Also note that having a rank of 1 along any of the modes would mean that the slices along that mode differ only by a (multiplicative) scaling parameter. Also note that during the forward pass, we can pass directly in factorized form to cach layer by selecting the corresponding rows in U(L) While the contraction in equation 3 can be done using the reconstructed tensor, it can also be done directly by contracting F(0) with the factors of the decomposition. For small, adequately chosen ranks, this can result in computational speedups.
A visualization of the Tucker decomposition of a third-order tensor can be seen in Figure 5). Note that we can rewrite the entire weight parameter for this Tucker case, equivalently, using the more compact n-mode product as:
g
W
W G X1 U x ( Xd+1 U Xd+2 U Xd43 U
We can efficiently perform an iFFT after contraction with the ten- sorized kernel. For any layer the (21.22 coordinate of the matrix- valued convolution function k(z) is as follows,
Figure 5: Illustration of Tucker decomposition. For clarity we show W as a 3rder tensor weight.
Rr
Ri
Ro
Ri
Ra Σ
(n((
G(ri," ra, To,
ml
(i..))) ... U - (ja,To) (l.) exp(2 Σ impix)
This joint factorization along the entire operator allows leverage redundancies both locally and across the entire operator. This leads to large reduction in the memory footprint, with only fraction of the pa- rameter. It also acts as low-rank regularizer on the operator, facilitating training. Finally, through global parametrization, we introduce skip connections that allow gradients to flow through the latent parametriza- tion to all the layers jointly, leading to better optimization.
Importantly, this formulation general and works with any tensor factorization. For instance, we also explore a Canomical-Polyadic decomposition (CP) which can be seen special case of Tucker with a super-diagonal core. Tn that case, we set single rank R and express the weights weighted of R rank-1 tensors. Concretely:
W Σ X.U r=1
(.r)
(6)
U (t.r) ((0) ((L)(.))
where U U@) U(O), U(1), U(d) are factor matrices size (Rx (RxI), (RxO), (Rxa), -... (Rxa), respectively and À E R Note that the CP, contrarily to the Tucker, has a single rank parameter, shared between all the dimensions. This means that maintain the number of parameters the same, R neods be very high, which leads memory issues. This makes CP more suitable for large compression ratios, and indeed, found it leads to better performance at high-compression / very low-rank. In this paper,

---

Published in Transactions on Machine Learning Research (08/2024)
Level
Level
Level
Padded Input Region
Output Region
Level 0
Predicting with padded regions. Local region the input padded used to predict corresponding region in the output.
(b) MG-Domain Decomposition, Progres- sively larger spatial regions are added a local region by subsampling.
Figure Domain decomposition in space (6(a) and our Multi-Grid based approach. (6(b) White squares represent the region of interest while yellow squares the larger embeddings.
also explore the tensor-train decomposition Oseledets 2011 A rank-(1, (1, . Rv, Rr, Ro, R1, 1) TT factorization expresses W as:
W(i
id iaiioiio io. =) G( - Gr Go(io) x (i.)..
Where each of the factors of the decompositions G₁ are third order tensors of size Rs x Iy X R In the experimental section
4.3 we show results TFNO trained with a Tucker, TT and CP factorization.
Separable Fourier Convolution The proposed tensorization approach introduces factorization of the weights in the spectral domain. When a CP Kolda & Bader (2009) is used, this induces separability over the learned kernel. In the case where the number of input channels is equal to the number of output channels (n = m in eq. [ we propose to make this separability explicit by not performing any channel mixing in the spectral domain and relying on the MLP introduced above to do so. Specifically, we can now rewrite the parameters in Eq. as W xx-xxxxxx (notice this has one less dimension that for the non-separable case). The operation in the spectral domain simplifies from a tensor contraction to just an element-wit wise matrix multiplication:
(T (T))( ((..)) 1) = T(4....,.4.77
la,
The separable Spectral convolution can be considered a depthwise convolution performed in the Fourier domain, e.g. without any channel mixing, and with separability along each dimension. The mixing between channels is instead done in the spatial domain. This results in a significant reduction in the number of parameters while having minimal impact performance (we found it necessary to increase the depth of the network, however, to ensure the network retained enough capacity).
3.6 Multi-Grid Domain Decomposition
Having introduced our decomposition in the operator's parameter space, we now introduce our novel multi- grid approach to decompose the problem domain.
Domain decomposition is a method commonly used to parallelize classical solvers for time-dependent PDEs that. based on the principle that the solution for a fixed local region in space depends mostly on the input at the same local region (Chan & Mathew 1994). In particular, since the time-step > 0 of + h), for any point x E and t E R4,, depends most strongly on the points u(y,t) t) for all B(z, r(h)) where B(zrr(h) denotes the ball centered at with radius r(h). This phenomenon is easily seen for the case of the heat. equation where, in one dimension, the
the numerical integrator small, the solution u(x,t+
10

---

Published in Transactions on Machine Learning Research (08/2024)
solution satisfies
((x
(-x--y)) - y)²
u(rt+ x
exp
u(v,t) dy
4h
1444
(- (x- y)2
exp
u(y, t) dy
4h
mh
with the approximation holding since 99.9937% of the kernel's mass is contained within B(x 4h). While some results exist, there is general convergence theory for this approach, however, its empirical success has made it popular for various numerical methods (Albin & Bruno 201T). Domain deomposition ion met are also widely used for boundary value, for example, in solving elliptic PDEs (Dryja & Widlund 1990). Differences in methods vary based on how boundary data is handled between adjecent subdomains. In our proposed appraoch we supplemant global data to each subdomain by using overlapping patches and stacking channels with an increasing receptive between at coarser scales.
To exploit this localization the domain D is split in q E N pairwise-disjoint regions Di, Dr so that D= D=Ug D= D) Each region D, is then embedded into larger one Z, > D2 so that points away from the center of D, have enough information to be well approximated. A model can then be trained so that the approximation G(a 2, lo, x wlD, holds for all j E [g]. This idea is illustrated in Figure 6hhreere where D = [o, 1]² and all Dj. Zj are differently sized squares. This allows the model to be ran fully in parallel hence its time and memory complexities are reduced linearly in q.
Multi-Grid. Domain decomposition works well in classical solvers when the time step h 0 is small because the mapping ((-:() ~ u(:;t++ 1 h) is close to the identity. However, the major advancement made by machine learning-basee baaen operator methods for PDEs is that model can approximate the solution, in one shot, for very large times i.c. h > 1. But, for larger h, the size of Zj relative to Dj must increase to obtain the same approximation accuracy, independently of model capacity. This causes any computational savings made by the decomposition approach to be lost.
To miligate this, we propose a multi-grid based domain decomposition approach where global information is added hierarchieally at different resolutions. While our approach is inspired by the classical multi grid method, it is not based the V-cyele algorithm (McCormick) 1985). For case of presentation, we describe this concept when domain D T is uniformlyy discretized by 2 x 2 points, for some E N, but note that generalizations can readily be made. Given final level L E N, we first sub-divide the domain into 22L. total regions each of size 2- L x 2--1 and denote them D10) ..... Diar. We call this the zeroth level. Then,
around each D for any E 12 consider the square D. of size 2 x 2*- that is equidistant, in every direction, from each boundary of D We then subsample the points in DD")) uniformly by a factor of in each direction, making D have 2=-L 2 points. We call this the first level. We continue this process by considering the squares D(d) of size 2--L+2 2--L+2 around cach DD and subsample them uniformly by a factor of in each direction to again yield squares with 2--L x 2--L points. The process is repeated until the Lth level is reached wherein D/-4) is the entire domain subsampled by a factor of 2-4 in each direction. The input to the model is then the concatenation each of these levels, in particular,
aj 1D J 1x
The process is illustrated for the case L = 2 in Figure 6(b) Since we work with the torus, the region of the previous level is always at the center of the current level.
The intuition behind this method is that since the dependence of points inside a local region diminishes the further we are from that region, is enough to have coarser information, as we go farther. We combine this multi-grid method with the standard domain decomposition approach by building appropriately padded squares 200) of size 2--4 + 2p x 2- + 2p around each ( where E N is the amount of padding to be added in each direction. We then take the evaluations of the input function a at cach level and concatenate them as channels. In particular, we train model so that ((al2))). a ( s n/p000- Since the
model only operates on each padded region separately, we reduce the total number of grid points used from
11

---

Published in Transactions on Machine Learning Research (08/2024)
12.5%
6.2%
L test error
TFNO [Tucker]
3.1%
TFNO [CP] Trim
1.6%
0.8%
100
200 300
400
500
Coapression-Ratio
Figure 7: Tensorization: error in logscale function of the compression ratio. We compare the tensor neural operator with an FNO with the same number of parameters (trimmed). We achieve over 100x compression ratio with better performance that the original FNO
4.50%
Ours [4 regions]
4.00%
Ours [16 regions]
Ours [4 regions, 10x compression]
3.50%
Ours [16 regions, 10x compression]
12 test error
3.00%
Full Field Model
2.50%
2.00%
1.50%
1.00%
2
3
5
Compression-Ratee
Figure MG-Domain Decomposition: error as a function of the domain compression ratio. We compare MG-TFNO with different numbers of multigrid regions both with and without weight tensor compression to full field FNO model. We achieve over 7x input space compression, 10x parameter space compression ratios and better performance than the original FNO.
22s to (2--L+ and define the domain compression ratio as the quotient of these numbers. Furthermore, note that, assuming a is Rd4-valued, model that does not employ our multi-grid domain decompasition uses inputs with d' channels while our approach builds inputs with (L + 1)dA channels. In particular, the the number of regions hence global information is
number of input channels scales only logarithmically added at very little additional cost. Indeed, FNO models are usually trained with internal widths much larger than d hence the extra input channels cause almost additional memory overhead.
12

---

Published in Transactions on Machine Learning Research (08/2024)
4
Experiments
In this section, we first introduce the data, experimental setting and implementation details before empirically validating our approach through thorough experiments and ablations.
4.1 Data.
We experiment on a dataset of 10K training samples and 2K test samples of the two-dimensional Navier- Stokes equation with Reynolds number 500. We also experiment with the one-dimensional viscous Burgers' equation.
Navier-Stokes.
We consider the vorticity form of the two-dimensional Navier-Stokes equation,
Ou + a. w.w
Aw f.
te te (0,7)
Re
(7)
AΦ w/ [ 0= JET LE (0,T)
with initial condition w(0,-) where T 2T) is the torus, f e i2( is a forcing function, and Re 0 is the Reynolds number. Then wtt,)) HT(T;; IR) for any E (0, and s > 0, is the unique weak solution to equation 7 [Temamm 1988). We consider the non-linear operator mapping f + w(,,-) with T= 5 and fix the Reynolds number Re = 500. We define the Gaussian measure μ =V, C) on the forcing functions where we take the covariance C 27(-2 + 97) following the setting in (De Hoop et al. 2022) Input data is obtained by generating i.i.d. samples from by a KL-expansion onto the eigenfunctions of C (Powell et al. 2014). Solutions to equation 7 are then obtained by a pseudo-spectral scheme (Chandler & Kerswelll 2013). Generating a single instance of the data at the 128 128 resilution on a single GPU takes approximately one minute. All machine lcarning methods compared in this work have an exccution time of less than one second, providing, at least, 100x speed-up. The cost of generating the high resolution data scales quadrtically due to the CFL condition and the speed-up due to the machine learning methods becomes much larger.
Burgers' Equation.
We consider the one-dimensional Burgers' equation on the torus,
0.u + utl, = V/llaa
x EET, LE (0,
(8)
t=0 uo:
TET
for initial condition un E LT; and viscosity v O. Then u(t,,) -) (T; for any t R4 and > 0, is the unique weak solution 8 Evans 2010). We consider the non-linear operator uo + uTT, with T = 0.5 or and fix 0.01. We define the Gaussian measure μ NN(0,) where we take the covariance C= 1=/3/2(-00) - + -7 Input data is obtained by generating i.i.d. samples from μ by KL-cxpansion onto the eigenfunctions of C. Solutions to equation are then obtained by pseudt-spectral solver using Heun's 's method. We use 8K samples for training and 2K for testing.
4.2 Implementation details
Implementation We use PyTorrc Paszke al. (2017) for implementing all the models. The tensor operations implemented using TensorLy Kossaifi et al. (2019) and TensorLy-Torch Kossaifi (2021). Our code will be released under the permissive MIT license, as Python package that is well-tested and comes with extensive documentation, to encourage and facilitate downstream scientific applications. Tt will be made available along withe the final version of the manuscript.
We train all models via gradient backpropagation using a mini-batel, size of 16, the Adam optimizer, with learning rate of le-3, weight decay of le-4, for 500 epochs, decreasing the learning rate every 100 epochs by factors of }. The model width is set in all cases to 64 except when specified otherwise (for the Trimmed FNO), meaning that the input was first lifted (with linear layer) from the
13

---

Published in Transactions on Machine Learning Research (08/2024)
Table 2: Comparing the performance of MG-TFNO with previous works, on the Navier-Stokes equations. Our method achieves superior performance with fraction of the parameters while largely compressing the weights (TFNO) and the input-domain (MG-TFNO).
Method
L test error (%) # Params
Model CR Input CR
FNO Li al. (2021b)
1.34
67M
FFNO Tran
20023
1.15
1M
67x
UNO Rahman et al.
2022b)
0.5 %
96M
TUNO (CP)
0.49%
1M
72x
TFNO (CP)
0.29%
890K
75x
TFNO (CP)
0.47%
447K
150x
MG-TENO (CP)
0.49
447K
40x
1.9x
MG-TFNO (Tucker)
0.42
447K
19x
1.9x
number of input channels to that width. The projection layer projects from the width 256 and a prediction linear layer outputs the predictions. 10000 samples were used for training, as well as a separate set of 2000 samples for testing. All experiments are done on a NVIDIA Tesla V100 GPU, The resolution is 128 128 unless specified otherwise.
To disentangle the effect of each of our components, the comparisons between the original FNO, the MG- FNO, TFNO, and the MG-TFNO were conducted in the same setting, with a mini-batch size of 32, modes of 42 and 21 for the height and width, respectively, and an operator width of 64.
For the comparison between our best models, we use all the modes (64 and 32) and mini-batch size of 16, which leads to improved performance for all models but longer training times. For each comparison, the same setting and hyper-paramete parameters were used for all models.
Training the operator. Since MG-TENO predicts local regions which are then stitched together to form a global function without any communication, aliasing effects can occur where one output prediction does not flow smoothly into the next. To prevent this, we train our model using the H Sobolev norm (Czarnecki ct al. 2017 Li et al., 2021a). By matching derivatives, training with this loss prevents any discontinuities from occurring and the output prediction is smooth. We note that this approach is only applicable if the PDE solution is globally H' which may not always be the case, for example, in problems which contains shocks. A new approach for smoothly stiching the patch-wise solution will be needed in such cases and we leave this to future work. We note that all problems considered in our experimental section satisfy the assumption.
4.3 Experimental results
In this section, we empirically verify the performance of our method, by comparing it with previous works, and through thorough ablation studies.
Comparison with previous works In this section, we compare our approach with both the regular FNO Li et al. (2021b), the Factorized-FNV (FFNO) Tran et al. (2023) and the U-Shaped Neural Opera- tor Rahman et al. (2022b).
FFNO separately applied FFT along each mode before combining the results. We used the same settings as described in the original paper TTan et al. (2023). U-NO uses a U-shaped architecture and reduces the spacial resolution while increasing the co-domain size.
In all cases, our approach achieves superior performance with fraction of the parameters, as can be seen in Table 4.4.4
14

---

Published in Transactions on Machine Learning Research (08/2024)
9.0%
TFNO Tucker 2x [test]
25.0%
TFNO [Tucker]
FNO [test]
FNO
8.0%
TFNO Tucker 2x [train]
12.5%
7.0%
FNO [train]
E test error
1? tost error
6.2%
6.0%
5.0%
3.1%
4.0%
1.6%
3.0%
2.0%
0.8%
1.0%
0
2000
4000
6000
8000
10000
100
200
300
400
500
Number of Training Samples
Epochs
Figure 9: Error as a function of the number of training samples (left) and training VS testing loss. We compare TFNO with regular FNO. Note that on the left we show the testing T.2 error while, for training, the = loss is used and that is compared with the H test error on the right. Our approach generalizes better while requiring fewer training samples.
Tensorizing: better compression. In Figure 7 we show the performance of our approach (TNO) com- pared to the original FNO, for varying compression ratios. In the Trimmed-FNO, we adjust the width in order to match the number of parameters in our TNO. We focus the width of the network it was shown to be the most important parameter (De Hoop et al.) 2022). Our method massively outperforms the Trimmed-FNO at every single fixed parameter amount. Furthermore, even for very large compression ratios, our FNO outperforms the full-parameter FNO model. This likely due to the regularizing effect of the tensor factorization on the weight, showing that many of the ones in the original model are redundant.
Tensorizing: better gencralization. Figure (left) shows that our TNO generalizes better with less training samples. Indeed, al every fixed amount of training samples, the TNO massively oulperforms the full- hall-parameter FNO model. Even when only using half the samples, our TNO outperforms the FNO trained the full dataset. Furthermore, Figure g (right) shows that our TNO overfits significantly less than FNO, demonstrating the regularizing effect of the tensor decomposition. This result is invaluable in the PDE setting where very few training samples are typically available due to the high computational cost of traditional PDE solvers.
Multi-Grid Domain Decomposition. In Table 1.4.4 we compare our MG-TENO with the baseline FNO and the TFNO, respectively. MG-TENO enables compressing both the weight tensor but also the input domain. On the other hand, preserving resolution invariance requires padding the patches, which decreases performance, resulting in a tradeoff between input domain compression and prediction accuracy.
We also show the impact of multi-grid domain decomposition on performance in Figure S We find that lower compression ratios (corresponding to a larger amount of padding in the decomposed regions) perform better which is unsurprising since more information is incorporated into the model. More surprisingly, we find that using a larger number of regions (16) performs consistently better than using a smaller number (4) and both can outperform the full-field FNO. This can be due to the fact that: i) the domain decomposition acts as form of data augmentation, exploiting the translational invariance of the PDE and more regions yield larger amounts of data, and ii) the output space of the model is simplified since a function can have high frequencies globally but may only have low frequencies locally. Consistently, we find that the tensor compression in the weights acts as a regularizer and improves performance across the board.
Architectural improvements to the backbone In addition to the ablation performed on our MG- TENO, we also investigate architectural improvements to the FNO backbone, see Sec (3.4 for details. In particular, we find that, while instance normalization docreases performance, layer normalization helps, es- pecially when used in conjunction with a pre-activation. Adding an MLP similarly improves performance, we found that bottleneck (expansion factor of 0.5) works well in practice, resulting in an absolute improvement
15

---

Published in Transactions on Machine Learning Research (08/2024)
Table 3: Impact of our architectural improvements.
Method
Layers
L test error
H' test error
# Params
Model CR
FNO et al.
2021b
1.34%
3.78%
67,142,657
FNO et al.
2021b
0.90%
2.59%
100,705,
0.7x
FNO et al.
2021b
0.73%
2.09%
134,268,161
0.5x
TFNO (CP)
4
0.47%
1.20%
447,100
150x
TFNO (CP)
0.27%
0.74%
662,081
101x
TFNO (CP)
0.22%
0.59%
877,057
77x
Table 4: Resolution invariance of TFNO. Since the model an operator, it is resolution invariant.
particular, here, we trained our model in resolution 128 128 and test on unseen samples in various resolutions and show it generalizes, with virtually no loss of performance higher resolutions unseen during training.
128 128
256 256
512 x 512
1024 1024
Method
L error
H' error
L error
H' error
L error
H' error
L error
H' error
CP TENO 0.3%
0.87% 1.2%
0.3%
0.93%
0.3%
0.93%
0.3%
0.93%
CP MG-TENO
0.49%
0.49%
1.3%
0.49%
1.5%
0.49%
1.6%
of 0.87% in relative L error. We found the ordering of normalization, activation, and weights (inchuding preactivation), did not have significant impact performance. Finally, when not using multi-grid domain dccomposition, the inputs are periodic and padding is not necessary. In that case, not padding the input im- proves performance. We use all these improvements for the backbone of the best version of our MG-TENO, Fig where we show that our improved backbone significantly outperforms the original FNO, while our approach significantly outperforms both, with a small fraction of the parameters, opening the door to the application of MG-TFNO to high-resolution problems.
We experiment with various number of layers and compare the original FNO with improved backbone in Table. 4.3
4.4 Ablation studies
In this section, further study the propertics of our model through ablation studics. We first look at how TFNO suffers less from overfitting thanks to the low-rank constraints before comparing its performance with various tensor decompositions. Finally, we perform ablation studies for our multi-grid domain decomposition on Burger's equation.
4.4.1 Resolution invariance
TFNO is resolution invariant, meaning that it can be trained on one resolution and tested different one. To illustrate this, we show zero-shot super-resolution results: we trained our best model (Table 4.4.4 on images of resolution 128 x 128 and tested it on unseen samples at higher resolutions (256 x 256 and 512 x 512), Table As can be seen, our method does as well on unseen, higher-resolution unseen testing samples as does on the training resolution, confirming the resolution invariance property of our neural operator.
4.4.2 Training higher-resolution with Multi-grid
One important advantage of our multi-grid domain decomposition is that it enables training much larger models on large inputs by distributing over patchos. We demonstrato this, by training on larger resolution (512x512 discretization) and using the largest FNO and TFNO that fits in memory, on V100 GPU. For the original FNO, this corresponds width of 12, first row in table 5 We then compare its performance
16

---

Published in Transactions on Machine Learning Research (08/2024)
Table 5: Training on 512x512. Multi-grid domain decomposition allows us to fit larger models into
memory by distributing patchos in the domain space, thus reaching lower relative error.
Model
Width
Patches
Padding
I error
FNO
12
0
6.1
MG-FNO
42
4
70
2.9
MG-FNO
66
4
2.4
MG-FNO
88
16
40
1.8
Tucker MG-TENO
80
16
46
1.3
with the multigrid approach with a neural operator as large as fits into the same V100 GPUs i.e. each width in the table has been optimized to be as large as memory allows. As we can see, our approach allows to larger model and reaches much lower relative I error.
4.4.3 Overfitting and Low-Rank Constraint
Here, show that lower ranks (higher compressions) lead to reduced overfitting. In Figure 5 we show the training and testing H errors for our TOP with Tucker decomposition at varying compression ratios (2x, 49x and 172x). We can see how, while the test error does not vary much, the gap between training and test errors reduces as we decrease the rank. As we can see, while being the most flexible, Tucker does not perform as well at higher compression ratios. In those extreme cases, CP and Tensor- Train lead to lower errors.
9.0%
TFNO CP 156x [test] FNO [test] TFNO CP 156x [train]
TENO TT 574x [test] FNO [test]
10.0%
8.0%
TFNO TT 574x [train]
Hi test error
7.0%
FNO [train]
test error
8.0%
FNO [train]
6.0%
5.0%
6.0%
4.0%
4.0%
3.0%
2.0%
2.0%
100
200
300
100
300
Epochs
Epochs
(a) Train VS Test error over time for TOP with
(b) Train VS Test error over time for TOP with TT factorization.
CP factorization).
Figure 10: Train/test curve for TOP-CP (10(a) and TOP-TT (10(b))
4.4.4 Tensor- Tensor-Train and TOP
Our approach is independent of the choice of tensor decomposition. We already showed how Tucker is most flexible and works well across all ranks. We also showed that while memory demanding for high rank, a CP decomposition leads to better performance and low rank. Our method can also be used in conjunction with other decompositions, such as tensor-train. To illustrate this, we show the convergence behavior of TNO with a Tensor- Train docomposition for a compression ratio of 178, figure 10(b)
We also compare in Table 4.4.4 our TFNO with different tensor decompositions.
17

---

Published in Transactions on Machine Learning Research (08/2024)
Table 6: Relative I test error of our MG-TENO approach for different tensor decompositions. We empirically found that Tucker works best for small compression ratio, CP excels at large compression ratio (<s 100x) but becomes computationally heavy for smaller ones. TT tends to be unstable at low-compression ratios but preserves a good performance for extreme compression ratio (> 500x).
Method
L2 test error
# Params Model CR
FNO Li et al. (0221b)
1.12%
67 M
0x
TFNO
[Tucker]
0.37%
28 M
2.3x
TFNO
CP
0.46%
808 K
83x
TFNO
TT
1.18%
117 K
574x
Table 7: Ablation comparing the performance on the relative I22 test error of our MG-TENO ap- proach, compared with its parts TFNO and MG-FNO and the regular FNO, Navier-Stokes. CR stands for compression ratio. Tensorization and multi-grid domain decomposition both individually im- prove performance while enabling space savings. The two techniques combined lead to further improvements, enabling large compression for both input and parameter, while outperforming regular FNO.
Method
L test error
# Params
Model CR
Domain CR
FNO (L et al. 2021b)
2.54%
58 M
0x
0x
TFNO [Tucker]
1.39%
41 M
1.5x
0x
TFNO [CP]
2.24%
130 K
482x
0x
MG-FNO
1.43%
58 M
0x
1.4x
MG-TFNO [Tucker] 0.85%
5.5 M
10x
1.78x
MG-TFNO [ucker]
1.89%
5.5 M
10x
7x
4.4.5 Decomposing domain and weights: MG-TFNO.
Tensorization and multi-grid domain decomposition not only improve performance individually but their advantages compound and lead to strictly better algorithm that scales well to higher-resoltion data by decreasing the number of parameters in the model as well as the size of the inputs thereby improving performance as well as memory and computational footprint. Table 7 compares FNO with Tensorization alone, multi-grid domain decomposition alone, and our joint approach combining the two, MG-TENO. In all cascs, for a, we keep 40 Fourier coefficients for height and 24 for the width and use an operator width of 64. Our results imply that, under full parallelization, the memory footprint of the model's inference can be reduced by 7x and the size of its weights by 10x while also improving performance.
Consistently with our other experiments, we find that the tensor compression in the weights acts as a regularizer and improves performance across the board. Our results imply that, under full parallelization, the memory footprint of the model's inference can be reduced by 7x and its weight size by 10x while also improving performance.
4.4.6 Burgers' Equation
We test the efficacy of the standard domain decomposition approach by training on two separate Burgers problems: one with a final time T = 0.5 and one with 7 = 1. As described in Section 3.6 we expect that for T = 1, each region requires more global information thus significantly more padding need to be used in order reach the same error. The results of Figure [T indoed confirm this. The domain compression ratios needed for the approach to reach the performance of the full-field model are higher, indicating the need for incorporating global information. These results motivate our multi-grid domain decomposition approach.
18

---

Published in Transactions on Machine Learning Research (08/2024)
0.20%
Radius 1.0
0.78%
Radius 1.2
Radius 0.8
Radius 1.0
Radius 0.5
Radius 0.8
L test error
0.10%
Full Field Model
L test error
0.39%
Full Field Model
0.05%
0.20%
0.10%
0.02%
1.5
2.0
2.5
3.0
3.5
4.0
5.0
1'5
2.0 Compression-Ratio
2.5
3.0
3.5
Compression-Ratio
Figure 11: Error on Burgers' equation with T 0.5 (left) and T = (right) as a function of domain compression ratio using standard domain decomposition without our multi-grid approach. We evaluate the performance of the standard domain decomposition approach. The radius indicates the size, in physical space, of the padding added to each region.
Model
# Params
GPU RAM (nvml)
Max GPU RAM (PyTorch)
Runtime
FNO
69 M
4.465 Gb
3.400 Gb
0.023
FFNO
M
5.465 Gb
4.493 Gb
0.043 8
UNO
5.048 Gb
3.341 Gb
0.029 s
TUNO [CP-fac]
344
4.952 Gb
3.347 Gb
0.026 s
TUNO [TUCKER
M
4.957 Gb
3.376 Gb
0.022
TUNO TUCKER--Bae
836 K
4.860 Gb
3.234 Gb
0.023
TFNO [CP-fac]
167 K
4.476 Gb
3.416 Gb
0.019
TFNO CC-facc
282 K
4.776 Gb
3.678 Gb
0.025 s
TFNO CC-ffcc
514 K
5.289 Gb
4.209 Gb
0.031 s
TFNO CCrr
514
6.840 Gb
5.792 Gb
0.025 8
TENO [ac
M
4.604 Gb
3.473 Gb
0.022 8
TFNO re-
M
5.254 Gb
4.169 Gb
0.031 s
TENO
TUCKER-
M
4.600 Gb
3.472 Gb
0.022
TENO
TUCKER-faa
724 K
4.355 Gb
3.297 Gb
0.020
Table 8: Training, Network width
4.5 Computational efficiency
In this section, we perform ablation studies to explore the impact of our proposed tensorization on memory consumption and training and inference times. The tensorized formulation enables an efficient implementa- tion obtained by directly contracting the activations with the factors of the decomposition, allowing us further scale the size of the models can fit in the memory. However, this speedup depends heavily on the hardware and parameters. Going forward, plan to further improve computational efficiency and memory savings.
To experimentally measure the gains, we created an FNO keeping all the modes, and we benchmarked two network widths, 64 (Tables 4.5 and 4.5 and 128 (Tables 4.5 and 4.5 for a batch size of 16.
In both cases, we measured the memory usage of both a regular FNO and our proposed Tucker TFNO, on single RTX 3080 GPU, for a single forward pass (inference) and forward-backwe backward pass (training), using [1] pynvml (pynvml. .m..eerieeeetteeeeee/tnoooo and PyTorrhh (torch. cuda. memory aammory_alloceee allocated
19

---

Published in Transactions on Machine Learning Research (08/2024)
Model
# Params
GPU RAM (nvml)
Max GPU RAAM (PyTorch)
Runtime
FNO
69 M
1.967 Gb
0.832 Gb
0.019
FFNO
M
1.522 Gb
0.653 Gb
0.041
UNO TUNO [CP-fac]
79 M
2.133 Gb
1.188 Gb
0.021 s
344
1.840 Gb
0.896 Gb
0.021 s
TUNO (TUCKER-fae fac]
4
1.840 Gb
0.910 Gb
0.020 s
TUNO TUCKER--ia
836
1.784 Gb
0.898 Gb
0.020 8
TFNO [CP-fac]
167
1.709 Gb
0.575 Gb
0.017
TFNO Ca-fac
282 K
1.709 Gb
0.575 Gb
0.018
TFNO CP-fac
514 K
1.709 Gb
0.576 Gb
0.021 8
TFNO
Cre
514 K
1.725 Gb
0.790 Gb
0.027 s
TENO
[T-
M
1.713 Gb
0.579 Gb
0.017 8
TFNO
rec
M
1.528 Gb
0.580 Gb
0.021 8
TENO
(Tac
M
1.709 Gb
0.588 Gb
0.017 8
TENO
[Tac
724 K
1.711 Gb
0.577 Gb
0.016
Table 9: Inference, Network width of 64
Model
# Params
GPU RAM (nvml)
Max GPU RAA (PyTorch)
Runtime
FNO
277 M
2.752 Gb
1.926 Gb
0.039
FFNO
M
2.358 Gb
1.294 Gb
0.080
UNO
96 M
2.198 Gb
1.252 Gb
0.023 8
TUNO [CP-fac]
381
1.717 Gb
0.896 Gb
0.022 s
TUNO [TUCKER fac]
M
1.860 Gb
0.913 Gb
0.021 s
TUNO TUCKER-BRE fac
M
1.719 Gb
0.898 Gb
0.022 s
TENO [CP-fac]
314 K
1.967 Gb
0.896 Gb
0.028
TFNO CC-ffcc
496 K
1.844 Gb
0.897 Gb
0.029
TFNO Cacfac
858 K
1.911 Gb
0.898 Gb
0.033
TENO CC-rec
858 K
2,885 Gb
1.566 Gb
0.062
TENO [f--
M
1.914 Gb
0.902 Gb
0.030 8
TENO re--
M
2.366 Gb
1.102 Gb
0.046 8
TENO
(ac
M
2.022 Gb
0.950 Gb
0.028
TENO
(UCCKER-fac)
M
1.967 Gb
0.905 Gb
0.026
Table 10: Inference, Network width of 128
20

---

Published in Transactions on Machine Learning Research (08/2024)
Model
# Params
GPU RAM (nvml)
Max GPU RAA (PyTorch)
Runtime
FNO
277 M
6.362 Gb
5.435 Gb
0.038
FFNO
5 M
9.805 Gb
8.698 Gb
0.082
UNO TUNO [CP-fac]
96 M
5.054 Gb
3.562 Gb
0.026 8
381 K
4.966 Gb
3.504 Gb
0.023 8
TUNO [TUCKER-EN
M
3.095 Gb
3.600 Gb
0.022 S
TUNO [TUCKER-fac)
M
4.950 Gb
3.480 Gb
0.022 8
TFNO [CP-fac]
314 K
6.033 Gb
4.678 Gb
0.029
TFNO CCP-fac]
496 K
6.485 Gb
4.939 Gb
0.032
TFNO CC-fac]
858 K
6.755 Gb
5.477 Gb
0.038
TFNO Cee
858 K
11.852 Gb
10.721 Gb
0.059 8
TFNO [a-
M
6,513 Gb
4.994 Gb
0.031 8
TFNO rec
M
9.476 Gb
8.503 Gb
0.044 s
TENO
[TUCKER fac]
M
6.889 Gb
5.480 Gb
0.031 8
TENO TUc
M
6.157 Gb
4.858 Gb
0.030 8
Table 11: Training, Network width of 128
4.6 Super-evaluation
In this section, we show some examples of input forcing functions, the corresponding ground-truth target function and the function predicted by our proposed MG-TFNO.
In particular, we visualize the predictions from our proposed method for two different settings: in both cases we train the model on fixed size of 128 x 128 imput functions but we either test with i) the same resolution or ii) with a new, much higher resolution, unseen during training (zero-shot super-evaluation)
Training and evaluating on the same resolution: In this setting, we train our proposed method on the full training set of functions discretized resolution 128 x 128 and test on new unseen samples at the same resolution, Figure. 4.6
Our model correctly learns the target function.
Zero-shot super-evaluation: In this setting, we take the same model as above, trained on the full training set of functions discretized al resolution 128 x 128.
However, this time, we give our model as input new unseen samples at new, wnseen resolution of 1024 1024, much higher than the training one, Figure. 4.6 The model is able to use high-frequency information in the new resolution and produce results good in the training resolution.
A natural question is whether there is an advantage compared to just downscaling the input, feeding it at the training resolution to the model and reupsampling it. In other words, is the model able to leverage high-fre figh-frequency information it has not seen during training.
To answer this question, we generated 1000 samples of the Navicr-Stokes equation, in the same sctting as previously introduced, but with resolution 1024 x 1024. We trained our proposed TFNO the full training set of samples resolution 128 x 128. We then evaluated the model on the high resolution, unseen, test samples of resolution 1024 x 1024 in two settings:
[upscaling] In that case, we downscale the input to the training resolution (128 x 128), predict the solution which we then re-upsample to the target high-resolution (1024 x 1024) for evaluation. In that case, we get a relative 12 error of 4.5%. As expected, the model is able to leverage high-frequency information since that is lost when downsampling.
(zero-shot super-evaluation] In that case, we leverage the fact that our proposed method is a neural operator Li et al. (2021b) and therefore take input functions at any regular discriti-
21

---

Published in Transactions on Machine Learning Research (08/2024)
Figure 12: Visualization of evaluation of our model with samples from the Navier-Stokes equa- tion at resolution 128 x 128. We show the input initial conditions (top-row), the corresponding target ground-truth function (middle-row), and the prediction function from our proposed model (bollom-row)) Our model has learned to correctly predict the target function and can solve the Navier-Stokes equation with very high accuracy for any initial conditions.
zation. We directly feed the samples at higl-resolution (1024 x 1024) and use the prediction (funtions on the same discretization and resolution) for evaluation. In that case, we get relative 12 error over 10 times better of only 0.4%.
This enables us to empirically verify the ability of our neural operator to leverage high-frequency information unscon during training.
5 Conclusion
In this work, we introduced i) a novel tensor operator (TFNO) well as a multi-grid domain decomposition approach which together form MG-TFNO, ii) an operator model that outperforms the FNO with a fraction of the parameters and memory complexity requirements, and iii) architectural improvements to the FNO. Our method scales better, generalizes better, and requires fewer training samples to reach the same performance; while the multi-gria domain decomposition enables parallelism over large input sizes. This is a promising direction for applications high-resolution data.
Acknowledgments
A.Anandkumar is supported by the Bren Named Chair, Schmidt AI 2050 Senior fellow, and ONR (MURI grant N00014-18-12624),
References
Jonas Adler and Ozan Oktem. Solving ill ill-posed inverse problems using iterative deep neural networks. Inverse Problems, 2017.
22

---

Published in Transactions on Machine Learning Research (08/2024)
Figure 13: Visualization of super-evaluation with our proposed models for samples from the Navior-Stokes equation at resolution 1024 x 1024. Here, we trained our proposed method with rank- 128 CP factorization, on data of resolution 128 x 128. However, we evaluate on data at a much higher resolution of 1024 x 1024. This is possible because our approach is a neural operator that can be evaluated and queried on any arbitrary resolution (here on a regular grid as it leverages the Fast Fourier Transform). Even though model has never obtained directly from solver at high-resolution, it is able to generalize. We show the inpul initial conditions (top-row), the corresponding target ground-truth function (middle-row), and the prediction function from our proposed model (bottom-row). Our model has learned to correctly predict the target function and can solve the Navier-Stokes equation with very high accuracy for any initial conditions and new unseen resolution.
the higher-frequency information from the Navier-Stokes equation
23

---

Published in Transactions on Machine Learning Research (08/2024)
Nathan Albin and Oscar P. Bruno. A spectral fe solver for the compressible navier-stokes equations in general domains i: Explicit time-stepping. Journal of Compulational Physics, 230(16):6248-6222 6270, 2011.
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. WrXiv:1607.04550 2016.
arv preprinl
Saakaar Bhatnagar, Yaser Afshar, Shaowu Pan, Karthik Duraisamy, and Shailendra Kaushik. Prediction of aerodynamic flow fields using convolutional neural networks. Computational Mechanics, pp. 1 1-21, 2019.
Kaushik Bhattacharya, Bamdad Hosseini, Nikola B Kovachki, and Andrew M Stuart. Model reduction and neural networks for parametric pdes. arXiv preprint arXiv:205.000 031880, 2020.
Mackenzie Blanusa, Carla J López-Zurita, and Stephan Rasp. The role internal variability in global climate projections of extreme events. arXiv preprint arXiv:2208.0827 2022.
Adrian Bulat, Jean Kossaifi, Georgios Trimiropoulos, and Maja Pantic. Incremental multi-domain learning with network latent tensor factorization. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 34, pp. 10470-10477, 2020a.
Adrian Bulat, Jean Kossaifi, Georgios Teimiropoulos, and Maja Pantic. Toward fast and accurate human pose estimation via soft-gated skip connections. In 2020 15th IEEE International Conference on Aulomatic Face & Gesture Recognition, 2020b.
J. Carroll and J. Chang. Analysis of individual differences in multidimensional scaling via an n-way gener- alization of eckart-young decomposition. Psychometrika, pp. 283-319, 1970.
Tony F. Chan and Tarek P. Mathew. Domain decomposition algorithms. Acta Numerica, 3:61-143, 1994.
Gary J. Chandler and Rich R. Kerswell. Invariant recurrent solutions embedded in turbulent two- dimensional kolmogorov flow. Journal of Fluid Mechanics, 722:554 595, 2013.
Jean- Bana-Baptstee Cordonnier, Andreas Loukas, and Martin Jaggi. Multi-head attention: Collaborate instead of concatenate. arXiv preprint Xiv::20006 16362, 2020.
R. Courant, K. Friedrichs, and H. Lewy. Uber die partiellen differennzengleichungen der mathemat ischen physik. Mathematische annalen, 100(1):32-77 74, 1928.
Wojeiech M Czarnecki, Simon Osindero, Max Jaderberg, Grzegorz Swirszez, and Razvan Pascanu. Sobolev training for neural networks. Advances in Neural Information Processing Systems, 30, 2017.
Tri Dao, Beidi Chen, Nimit S Sohoni, Arjun Desai, Michael Poli, Jessica Grogan, Alexander Liu, Aniruddh Rao, Atri Rudra, and Christopher Ré. Monarch: Expressive structured matrices for efficient and accurate training. In International Conference on Machine Learning, pp. 4690 4721. PMLR, 2022.
Maarten De Hoop, Daniel Zhengyu Huang, Elizabeth Qian, and Andrew M Stuart. The cost-accuracy trade-off in operator learning with neural networks. arXiv preprint ar Xix:22203 13181, 2022.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Un- terthiner, Mostafa. Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, al. An image is worth 16x16 words:
Transformers for image recognition at scale. arXiv preprint arXin:2010. 11929, 2020.
Maksymilian Dryja and Olof B. Widlund. Chapter 16 some domain decomposition algorithms for elliptic problems. In David R. Kincaid and Linda J. Hayes (eds.), Iterative Methods for Large Linear Systems, pp. 273 291. Academic Press, 1990.
Soheil Esmaeilzadeh, Kamyar Azizzadenesheli, Karthik Kashinath, Mustafa Mustafa, Hamdi A Tchelepi, Philip Marcus, Mr Prabhat, Anima Anandkumar, et al. Meshfreeflownet: A constraecd deep continuous space-time super-resolution framework. In SC20: International Conference for High Perfor- mance Computing, Networking, Storage and Analysis, pp. 1-15. IEFF, 2020.
Lawrence C. Evans. Partial differential equations. American Mathematical Society, 2010.
24

---

Published in Transactions on Machine Learning Research (08/2024)
John Guibas, Morteza Mardani, Zongyi Li, Andrew Tao, Anima Anandkumar, and Bryan Catanzaro. Adap- tive fourier neural operators: Efficient token mixers for transformers. arXiv preprinl or orXim:2111 1t:22111.358 2021.
Xiaoxiao Guo, Wei Li, and Francesco Torio. Convolutional neural networks for steady flow approximation. In Proceedings of the 22nd ACM SIGKDD International Conference on Knowledge Discovery and Data Mining, 2016.
Gaurav Gupta, Xiongye Xiao, and Paul Bogdan. Multiwavelet-basen operator learning for differential equa- tions. Advances in Neural Information Processing Systems, 34:24048-2444 240622, 2021.
Julia Gusak, Maksym Kholiavchenko, Evgeny Ponomarev, Larisa Markeeva, Philip Blagoveschensky, Andrzej Cichocki, and Tvan Oseledets. Automated multi-stage compression of neural networks. Oct 2019.
R. A. Harshman. Foundations of the PARAFAC procedure: Models and conditions for an "explanatory" multi-modal factor analysis. UCLA Working Papers in Phonetics, 16:1-84, 1970.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Idontity mappings in deep residual networks. Computer Vision ECCV 2016: 14th European Conference, Amslerdam, The Netherlands, Oclober 11 14, 2016, Proceedings, Part IV 14, pp. 630 645. Springer, 2016.
Sergey Toffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In International conference on machine learning, pp. 448-456. pmlr, 2015.
Majid Janzamin, Rong Ge, Jean Kossaifi, and Anima Anandkumar. Spectral learning matrices and tensors. Foundations and Trends in Machine Learning, 12(5 336-536, 2019a. ISSN 1935-8245. doi: 10.1561/2200000057. URL http:://a.......r//0011611661/// wtt..//0...1111100000007777
doi
Majid Janzamin, Rong Ge, Jean Kossaifi, Anima Anandkumar, ct al. Spectral learning on matrices and lensors. Found. and Trendso in Mach. Learn., 12(5-6):393 536, 2019b.
Yong-Deok Kim, Eunhyeok Park, Sungjoo Yoo, Taelim Choi, Lu Yang, and Dongjun Shin. Compression of deep convolutional neural networks for fast and low power mobile applications. 2016.
Tamara G Kolda and Brett W Bader. Tensor decompositions and applications. SIAM Rev., 51(3):455-55) -00, 2009.
Jean Kossaifi. Tensorly-torch.
2021.
Jean Kossaifi, Yannis Panagakis, Anima Anandkumar, and Maja Pantic. Tensorly: Tensor learning in python. Journal of Machine Learning Research (JMLR), 20(26), 2019.
Jean Kossaifi, Autoine Toisoul, Adrian Bulat, Yannis Panagakis, Timothy Hospedales, and Maja Pantic. Factorized higher-order CNNs with an application 2020.
emotion estimation. pp. 6060-6069,
Nikola Kovachki, Samuel Lanthaler, and Siddhartha Mishra. On universal approximation and error bounds for fourier neural operators. Journal of Machine Learning Research, 22(290):1 76, 2021a.
Nikola Kovachki, Zongyi Li, Burigede Liu, Kamyar Azizzadenesheli, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Neural operator: Learning maps between function spaces, arXio preprint arXiv:2108.08/4 2021b.
Vadim Lebedev, Yaroslav Ganin, Maksim Rakhuba, Ivan V. Oseledets, and Victor S. Lempitsky. Speeding-up convolutional neural networks using fino-tuned CP-dccomposition. 2015.
Martin Leutbecher and Tim N Palmer. Ensemble forecasting. Journal of computational physics, 227(7): 2008.
25

---

Published in Transactions on Machine Learning Research (08/2024)
Zongyi Li, Nikola Kovachki, Kamyar Azizzadnnesheli, Burigede Liu, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Neural operator: Graph kernel network for partial differential equations. arXiv preprint Xri:220033 03485, 2020a.
Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Andrew Stuart, Kaushik Bhattacharya, and Anima Anandkumar. Multipole graph neural operator for parametric partial diffcrential equations. Advances in Newral Information Processing Systems, 33:6755-6768 6766, 2020b.
Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, Andrew Stu- art, and Anima Anandkumar. Markov neural operators for learning chaotic systems. ar preprint arXiv:2106. 06898, 2021a.
Zongyi Li, Nikola Borislavov Kovachki, Kamyar Azizzadenesheli, Burigede liu, Kaushik Bhattacharya, An- drew Stuart, and Anima Anandkumar. Fourier neural operator for parametric partial differential equations. In International Conference on Learning Representations, 2021b.
Zongyi Li, Hongkai Zheng, Nikola Kovachki, David Jin, Haoxuan Chen, Burigede Liu, Kamyar Azizzade- nesheli, and Anima Anandkumar. Physics-informed neural operator for learning partial differential equa- tions. Xix preprint orXiv:2111.09794, 2021c
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 Lu, Pengzhan Jin, and George Em Karniadakis. Deeponet: Learning nonlinear operators for identi- fying differential equations based on the universal approximation theorem of operators. arXiv preprint WrXXvv1910.031000 2019.
S. F. McCormick. Multigrid methods for variational problems: General theory for the v- cycle. SIAM Journal on Numerical Analysis, 22(4):634 643, 1985.
Alexander Novikov, Dmitry Podoprikhin, Anton Osokin, and Dmitry Vetrov. Tensorizing neural networks. pp. 442-450, 2015.
I. V. Oseledets.
Tensor-train decomposition. SIAM J. Sci. Comput., 33(5):2295-22) 2317, September 2011.
Yannis Panagakis, Jean Kossaifi, Grigorios G. Chrysos, James Oldfield, Mihalis A. Nicolaou, Anima Anand- kumar, and Stefanos Zafeiriou. Tensor methods in computer vision and deep learning. Proceedings of the IEEE, 109:863-890, 2021a. URL
.tt../cco..snaantccccccccol.. org/Corpus 233465455
Yannis Panagakis, Jean Kossaifi, Grigorios G. Chrysos, James Oldfield, Mihalis A. Nicolaou, Anima Anand- kunar, and Stefanos Zafeiriou. Tensor methods in computer vision and deep learning. Proceedings of the IEEE, 109(5):863 890, 2021b. doi: 10.1109/ 10.1109/JPROCE 00T.C...02..0073322
Christos Papadopoulos, Yannis Panagakis, Manolis Koubarakis, and Mihalis Nicolaou. Efficient learning of multiple nlp tasks via collective weight factorization on bert. In Findings of the Association for Compu- tational Linguistics: NAACL 2022, pp. 882-890, 2022.
Evangelos E Papalexakis, Christos Faloutsos, and Nicholas D Sidiropoulos. Tensors for data mining and data fusion: Models, applications, and scalable algorithms. ACM Thans. Intell. Syst. and Technol. (TIST), 8 (2):1-44, 2016.
Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang, Zachary DeVito, Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer. Automatic differentiation in PyTorch. 2017.
Jaideep Pathak, Shashank Subramanian, Peter Harrington, Sanjeev Raja, Ashesh Chattopadhyay, Morteza Mardani, Thorsten Kurth, David Hall, Zongyi Li, Kamyar Azizzadenesheli, al. Fourcastnet: A global data-driven high-resolution weather model using adaptive fourier neural operators. arXiv preprint arXiv:2202.1122 11214, 2022.
26

---

Published in Transactions on Machine Learning Research (08/2024)
Catherine E. Powell, Gabriel Lord, and Tony Shardlow. An Introduction to Computational Stochastic PDEs. Texts in Applied Mathematios. Cambridge University Press, United Kingdom, 1 edition, August 2014. ISBN 9780521728522.
Md Ashiqur Rahman, Manuel A Florez, Anima Anandkumar, Zachary E. Ross, and Kamyar Azimadeneshell. Generative adversarial neural operators. arXiv preprint arXit:2205.00017 2022a.
Md Ashiqur Rahman, Zachary Ross, and Kamyar Arizzadenesheli. U-no: U-shaped neural operators.
arXiv preprint 2xxx:2204.11199 2022b.
Nicholas D Sidiropoulos, Lieven De Lathauwer, Xiao Fu, Kejun Huang, Evangelos E Papalexakis, and Christos Faloutsos. Tensor decomposition for signal processing and machine learning. Transactions Signal Processing, 65(13):3551 355-33222,, 2017.
Julia Slingo and Tim Palmer. Uncertainty in weather and climate prediction. Philosoppical Transactions of the Royal Society A: Mathematical, Physical and Engineering Sciences, (491114)))45511740000 2011.
Roger Temam.
dynamical systems in mechanics and physics. Applied mathematical sciences. Springer-Verlag. New York, 1988.
Alasdair Tran, Alexander Mathews, Lexing Xie, and Cheng Soon Ong. Factorized fourier neural operators. The Eleventh International Conference on Learning Representations, 2023. URL http:://openreviies
net/forumiid=tmm11 httddd=tmmiMPP11110 Pa
L. R. Tucker. Some mathematical noles on three-mode factor analysis. Psychomelrika, 31:279 311, 1966c.
Dmitry Ulyanov, Andrea Vedaldi, and Victor Lempitsky. Instance normalization: The missing ingredient for fast stylization. arXiv preprint arXi::1607.0802. 2016.
Gege Wen, Zongyi Li, Kamyar Axizzadenesheli, Anima Anaudkumar, and Sally M Benson. U-fno an enhanced fourier neural pverator-based deep-learning model for multiphase flow. Advances in Water Resowrces 163:104180, 2022.
Yan Yang, Angela F Gao, Jorge C Castellanos, Zachary Ross, Kamyar Axizzadenesheli, and Robert W Clayton. Seismic wave propagation and inversion with neural operators. The Seismic Record, 1(3):126 134, 2021.
Yan Yang, Angela F Gao, Jorge C Castellanos, Zachary E Ross, Kamyar Azizzadenesheli, and Robert W Clayton. Accelerated full seismic waveform modeling and inversion with u-shaped neural operators. arXiv preprint arXiv:2209. 11955, 2022.
Yinhao Zhu and Nicholas Zabaras. Bayesian deep convolutional encoder-decoder networks for surrogate modcling and uncortainty quantification. Journal of Computational Physics, 2018. ISSN 0021-9991.
27


Reviews:
[{"review_id": "SpUgcFrsj0", "reviewer": "Reviewer_PLjK", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n1. MG-TFNO introduces novel tensor factorization and multi-grid domain decomposition techniques that significantly improve the memory efficiency and scalability of Fourier Neural Operators for high-resolution PDEs.\n\n2. The tensor factorization drastically reduces the number of parameters while maintaining or even improving performance, acting as an effective regularizer.\n\n3. The multi-grid domain decomposition enables parallel local predictions with progressively added contextual information, allowing the method to scale to much higher resolution inputs.\n\n4. The strong empirical results on the Navier-Stokes equation demonstrate the superior performance of MG-TFNO compared to regular FNO in terms of accuracy, parameter efficiency, sample efficiency, and generalization.\n\n\nWeaknesses:\n1.  The paper lacks theoretical analysis of the approximation capabilities and convergence properties of MG-TFNO, which would provide a stronger foundation for the empirical results.\n\n2. The impact of tensor factorization rank, multi-grid resolution, and other hyperparameters on the method's performance and tradeoffs could be explored more thoroughly.\n\n3. The H1 smoothness assumption in the multi-grid decomposition may limit the approach to certain classes of problems, and alternative stitching approaches may be needed for PDEs with discontinuities or shocks.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposes a new method called Multi-Grid Tensorized Fourier Neural Operator (MG-TFNO) for learning solution operators of partial differential equations (PDEs) at high resolutions. MG-TFNO addresses the limitations of previous methods in terms of memory complexity and data scarcity by leveraging local and global structures of the problem through a combination of decomposition techniques in both the input domain and the operator's parameter space. The key contributions include:\n\n1. A novel multi-grid based domain decomposition approach that enables parallelization over input samples and reduces the input domain size.\n\n2. A global tensor factorization of the parameters into a high-order latent subspace of the Fourier domain, which significantly reduces the number of parameters.\n\n3. Architectural improvements to the Fourier Neural Operator (FNO) backbone model used in MG-TFNO.\n\nThe method is evaluated on the Navier-Stokes equations, where it achieves better performance than the regular FNO with significantly compressed model size and input domain. The experiments also demonstrate improved generalization, sample efficiency, and reduced overfitting compared to FNO. The approach is general and can be used with different tensor factorizations and backbone architectures. The paper discusses the potential impact of MG-TFNO in enabling the application of neural operators to high-resolution, real-world PDE problems that were previously intractable due to computational and data requirements.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1. Conduct more thorough ablation studies to explore the impact of tensor factorization rank, multi-grid resolution, and other hyperparameters on the method's performance and tradeoffs. \n\n2. Improve the clarity of the writing in some areas, such as the multi-grid domain decomposition description and the analysis of experimental results. Provide clearer intuition for the tensor factorization and locality assumptions to make the work more accessible. \n\n3. Discuss the potential limitations and failure cases of MG-TFNO more explicitly, especially in relation to the H1 smoothness assumption and the handling of discontinuities or shocks in the PDE solutions.", "broader_impact_concerns": "There are no broader impact concerns.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "2KHpEZBlzP", "reviewer": "Reviewer_DbZ6", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths**\nThe high-level idea behind the design is intuitive and well-motivated. Both multi-grid methods and domain decomposition are classic numerical methods for solving high-resolution and large-scale numerical problems, and I think incorporating them into a neural operator makes good sense. The proposed framework seems timely and reasonable to me.\n\n**Weaknesses**\nActually, I don’t have major concerns with the technical methods or the experiments, and I buy the results in the paper. I will list a few small concerns I have, which are, hopefully, not too difficult to address after moderate paper editing.\n\n1. Some claims in the abstract are unclear:\n- “Our approach can be used in any operator learning setting.” This seems to be a very strong claim (“any”). I wonder if it is better to say “can be used in several operator learning settings, e.g., X, Y, and Z” instead.\n- “...on the turbulent Navier-Stokes equations…less than half the error…over 150x compression.” It would be great if this sentence could mention explicitly that the N-S equations are in 2D. Also, it is unclear what method this sentence compares to achieve “less than half the error” and “over 150x compression”.\n- Similarly, it would be great to clarify “150x” and “7x” in the next sentence.\n\n2. I feel certain statements in the intro may be misinterpreted and can really benefit from further clarification:\n- “For example, in weather forecasts, …”, “Conventional PDE solvers..to tackle such problems”, and “For instance, consider a weather forecast…”: While there is nothing technically wrong here, these sentences in the intro, together with the “High-Resolution PDEs” in the title, left me the impression that this paper aims to tackle challenging large-scale PDE problems comparable to those in meteorology, and I was disappointed to find the problems actually solved in this work is 2D 1024x1024 N-S at most.\n- The first two paragraphs state the downside of conventional PDE solvers (paragraph 1) and the advantages of neural PDE solvers (paragraph 2). This is not a balanced statement, as it misses the strengths of conventional solvers (which also hints at the weaknesses of neural network methods), e.g., provable solution convergence, guarantees on error bounds, no bias toward data seen in the training set, and so on. I don’t think it hurts this paper’s value if the intro states all these strengths and weaknesses – this gives readers an unbiased view of both methods.\n- “Since a local prediction depends most strongly on its immediate spatial surroundings” I am not sure I agree with this claim from the technical perspective, and I’d be grateful if it could elaborated in more context. While I agree operators in a PDE are typically local, e.g., Laplacian in Poisson problems, its inverse operator needed in prediction can have a global dependency. That said, I agree that using multigrid is a reasonable idea.\n- Finally, the bullet points summarizing the contribution at the end of the intro are too scattered. I feel merging them into two or three major contributions is perhaps a better way of organizing them.\n\n3. A few questions about certain statements in the Background section:\n- “As we move beyond learning from simple structures to solving increasingly complex problems, the data we manipulate becomes more structured.” I don’t get this claim. Why do complex problems have more structured data?\n- “Of all these prior works, none has been applied to neural operator.” Is this sentence saying tensor decomposition is not explored in neural operators? Do low-rank neural operators and Multipole neural operators [Kovachki 2021b] not fall into this category?\n\n4. A technical comment on “Domain Decomposition” described in Sec. 3.6: I feel what is discussed in this section primarily focused on partitioning the domain, which is only a small part of the full domain decomposition idea. Many efforts in domain decomposition algorithms are devoted to solving the interface data shared by domains, which does not seem to be discussed or needed in this paper. Overall, I feel it is a little bit exaggerated to connect this section to domain decomposition because of this particular reason.\n\n5. A general comment on the experiments and results: While I don’t have major concerns here and I appreciate the results, I feel certain claims can be toned down to better reflect the method’s capability, e.g.,\n- “enables parallelism over huge inputs”: 1024 x 1024 is a respectable resolution in 2D, but it is still premature to call it “huge inputs” if we compare it with what multigrid/domain decomposition can solve in 3D PDE problems (e.g., Poisson).\n- “paves the way to applications on very high-resolution data”. Same as above.\n- “We plan to deploy MG-TFNO to large-scale weather forecasts for which existing deep learning models are prohibitive.” This is ambitious, but I don’t think the gap between what this paper presented and this goal is small. This statement seems to imply the proposed method is more promising than the existing deep learning methods for solving large-scale weather forecasts, and I am not sure this is the right impression to have.\n- I also suggest adding a paragraph discussing the limitations of the proposed approach, e.g., evaluated on 2D problems with grid discretizations only. Extending it to 3D and irregular geometric domains (e.g., organic shapes) may be doable but does not seem trivial.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper presents MG-TFNO, a new variant of neural operators that incorporates multi-grid ideas and low-rank tensor factorization. MG-TFNO learns to solve 2D partial differential equations (PDEs) in varying resolutions (trained on 128 x 128 and tested on at most 1024 x 1024).\n\nIn my opinion, the main contribution is 1) the proposed new neural operator that strategically puts together multigrid/domain decomposition/tensor factorization into a single framework and 2) evaluations and ablation study of this framework on several 2D test problems. While neural operators with multi-scale designs or with tensor factorization alone are not new, I am unaware of an existing work that presents a design similar to MG-TFNO, and I am willing to agree it contains certain novelty.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "See above.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "qOmgBJpa6R", "reviewer": "Reviewer_1Vpu", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths: \n1. The paper is well-written and organized in general. \n2. The proposed methods can address the training and data issues of the current FNO well. The results presented in this paper are promising.\n\nWeaknesses\n1. Some related works still need to be fully discussed. \n2. Some claims are not well supported by the experiments.\n3. The experimental settings are relatively simple.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper proposes a novel neural operator learning regime named multi-grid tensorized neural operator (MG-TFNO). MG-TFNO leverages the following main techniques to improve the existing FNO: 1). the multi-grid approach to handle high-resolution regions; 2). the low-rank tensor factorization technique to improve the generalization and efficiency; 3). an improved FNO architecture with skip connections. The paper mainly conducts experiments on the Navier-Stokes equation and reveals promising results in terms of accuracy and efficiency.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "I think a major revision is required before the manuscript is accepted. The major concerns are listed as follows:\n1. Discussions about the related works. The domain decomposition technique has been widely used in the PINN paper [1-2]. Furthermore, some studies also use local/global information and decomposition ideas for weather forecasting [3]. It would be better to discuss these papers in the background section.\n2. The authors claim the data efficiency advantage as one of the main contributions, while corresponding experiments are not comprehensively conducted. It would be better to compare the performance of each model with the varying numbers of data points.\n3. In the experimental settings, the authors only consider one kind of Navier-Stokes equation with fixed Re = 500. However, it would be questionable to claim such flow as “turbulent flow.” Additional experiments with larger Re are recommended.\n4. What is the number of layers in FFNO during the experiments? One of the main contributions of the FFNO paper is to extend the number of layers for FNO. So, it would be better to clarify such issues.\n\nMinor issues may need to be clarified:\n1. In the Introduction, the authors claim that `This coupling allows us to limit the number of parameters in the model without limiting its expressivity.’ It is worth mentioning that the expressivity and generalizability are not the same. Are there any experimental results (for instance, the comparisons of training error) or existing theoretical results for low-rank networks to support this claim?\n2. In Section 3.3, what is $G$ in $\\alpha<\\frac{1}{G2}…$? Is it a typo?\n3. In Section 3.6, the paper presents the approximation rate (0.999937) in the 1D case. What about 2D and 3D cases? Will the ` curse of dimensionality ’ occur? \n4. No validation set has been introduced in this paper. How did the authors choose the hyper-parameters?\n5. The paper seems to have conducted each experiment once, and no variances were reported. Then, how did the authors choose the random seed?\n\n\n[1] Shukla K, Jagtap A D, Karniadakis G E. Parallel physics-informed neural networks via domain decomposition[J]. Journal of Computational Physics, 2021, 447: 110683.\n\n[2] Kopaničáková A, Kothari H, Karniadakis G E, et al. Enhancing Training of Physics-Informed Neural Networks Using Domain Decomposition–Based Preconditioning Strategies[J]. SIAM Journal on Scientific Computing, 2024: S46-S67.\n\n[3] Gao Z, Shi X, Wang H, et al. Earthformer: Exploring space-time transformers for earth system forecasting[J]. Advances in Neural Information Processing Systems, 2022, 35: 25390-25403.", "broader_impact_concerns": "I think the researchers in the AI4PDE field would be interested in this paper and related topics. Moreover, this paper does not have ethical risks.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The reviewers were very positive about the idea of combining  tensor factorization and multi-grid domain decomposition techniques and believe that this may have impact in the future.
Decision: Accept (poster)