Multi-Grid Tensorized Fourier Neural  Operator for High Resolution PDEs
Authors: Jean Kossaifi
Nikola Borislavov Kovachki
Kamyar Azizzadenesheli
Anima Anandkumar
Venue: Submitted to ICLR 2023 (2023)
Keywords: Fourier-Neural-Operators
Tensorization
Multi-Grid

Abstract: Memory complexity and data scarcity are two main pressing challenges in learning solution operators of partial differential equations (PDE) at high resolutions. These challenges limited prior neural operator modelsMemory complexity and data scarcity are two main pressing challenges in learning solution operators of partial differential equations (PDE) at high resolutions. These challenges limited prior neural operator modelsMemory complexity and data scarcity are two main pressing challenges in learning solution operators of partial differential equations (PDE) at high resolutions. These challenges limited prior neural operator models to low/mid-resolution problems rather than full scale real-world problems. Yet, these problems possess spatially local structures that is not used by previous approaches. We propose to exploit this natural structure of real-world phenomena to predict solutions locally and unite them into a global solution. Specifically, we introduce a neural operator that scales to large resolutions by leveraging local and global structures through decomposition of both the input domain and the operator's parameter space. It consists of a multi-grid tensorized  neural operator (MG-TFNO), a new data efficient and highly parallelizable operator learning approach with reduced memory requirement and better generalization. MG-TFNO employs a novel multi-grid based domain decomposition approach to exploit the spatially local structure in the data. Using the FNO as a backbone, its parameters are represented 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. In addition, the low-rank regularization it applies to the parameters enables efficient learning in low-data regimes, which is particularly relevant for solving PDEs where obtaining ground-truth predictions is extremely costly and samples, therefore, are limited. We empirically verify the efficiency of our method on the turbulent Navier-Stokes equations where we demonstrate superior performance, with 2.5 times lower error, 10X compression of the model parameters, and 1.8X compression of the input domain size. Our tensorization approach yields up to 400x reduction in the number of parameter without loss in accuracy. Similarly, our domain decomposition method gives a 7x reduction in the domain size while slightly improving accuracy. Furthermore, our method can be trained with much fewer samples than previous approaches, outperforming the FNO when trained with just half the samples.

Under review as a conference paper at ICLR 2023
MULTI-GRID TENSORIZED FOURIER NEURAL OPERA- TOR FOR HIGH RESOLUTION PDEs
Anonymous authors Paper under double-blind review
ABSTRACT
Memory complexity and data scarcity are two main pressing challenges in learn- ing solution operators of partial differential equations (PDE) at high resolutions. These challenges limited prior neural operator models low/mid-resolution prob- lems rather than full scale real-world problems. Yet, these problems possess spa- tially local structures that not used by previous approaches. We propose to exploit this natural structure of real-world phenomena to predict solutions locally and unite them into global solution. Specifically, we introduce a neural operator that scales to large resolutions by leveraging local and global structures through both the input domain and the operator's 's parameter space. It consists of multi-grid tensorized neural operator (MG-TFNO), new data effi- cient and highly parallelizable operator learning approach with reduced memory requirement and better generalization. MG-TENO employs novel multi-grid based domain decomposition approach to exploit the spatially local structure in the data. Using the FNO backbone, its parameters are represented high- order latent subspace of the Fourier domain, through a global tensor factorization, resulting in an extreme reduction in the number parameters and improved gen- eralization. In addition, the low-rank regularization it applies to the parameters enables efficient learning in low-data regimes, which is particularly relevant for solving PDEs where obtaining ground-truth predictions is extremely costly and samples, therefore, are limited. We empirically verify the efficiency of our method the turbulent Navier-Stoker equations where we demonstrate superior perfor- mance, with 2.5 times lower error, 10x compression of the model parameters, and 1.8x compression of the input domain size. Our tensorization approach yields up to 400x reduction in the number of parameter without loss in accuracy. Similarly, our domain decomposition method gives a 7x reduction the domain size while slightly improving accuracy. Furthermore, our method can be trained with much fewer samples than previous approaches, outperforming the FNO when trained with just half the samples.
decomposition
INTRODUCTION
Real-world scientific computing problems often time require repeatedly solving large-scale and high-resolution partial differential equations (PDEs). For instance, weather forecasts, large sys- tems of differential equations are solved to forecast the future state of the weather. Due to internal inherent and aleatoric uncertainties, multiple repeated runs are carried out by weather scientists ev- ery day to quantify prediction uncertainties. Conventional PDE solvers constitute the mainstream approach used tackle such computational problems. However, these methods are known to be slow and memory-intensive. They require immense amount of computing power, are unable to learn and adapt based on observed data, and often times require sophisticated tuning (Slingo & Palmer 2011: Leutbecher & Palmer 2008 Blanusa et al., 2022).
Neural operators are a new class of models that aim at tackling these challenging problems (Li al., 2020b They are mappings between function spaces whose trained models emulate the solution operators of PDEs (Kovachki et al. 2021b). In the context of PDEs, these deep learning models orders of magnitude faster than conventional 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 methodologies (Liu et al. 2022 Li al. 2021b). Among

---

Under review as a conference paper at ICLR 2023
F
F
Tensor Operator Block
Tensor Operator Block
Levelo
Multi-Grid Domain Decomposition
Joint Parameter Decomposition
Overview of our approach. First (left), a multi-grid approach is used to create coarse to fine inputs that capture high-resolution details in a local region while still encoding global context. The resulting regions are fed to tensorized Fourier operator (middle), the parameters of which are jointly represented in single latent space via 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 is ensured via the choice of loss function.
Figure I:
Method
L test error
Params
Model CR
Domain CR
FNO
2.54%
58 M
0x
TENO [Tucker]
1.39%
41 M
1.5x
0x
TENO [CP]
2.24%
130 K
482x
0x
MG-FNO
1.43%
58 M
0x
1.4x
MG-TENO [Tucker]
0.85%
5.5 M
10X
1.78x
MG-TFNO [Tucker]
1.89%
5.5 M
10X
7x
Table 1: Overview of the performance on the relative 2 test error of our MG- TFNO approach, compared with its parts TFNO and MG-FNO and the regular FNO, on Navier-Stokes. CR stands for compression ratio. For each method, we report the relative T,2 error, the number of parameters, and the compression ratios for both the input domain and the number of parameters. Tensorization and multi-grid domain decomposition both individually improve performance while enabling space savings. The two techniques combined lead to further improvements, enabling huge compression for both input and parameter, while largely outperforming all other approaches.
neural operator models, Fourier neural operators (FNOs) in particular, have seen successful appli- cation in scientific computing for the task of learning the solution operator to PDEs well in computer vision for classification, in-painting, and segmentation (Li et al.] 2020a Kovachki et al.) 2021a Guibas et al. 2222]. By leverage spectral theory, FNOs have successfully advanced frontiers in weather forecasts, carbon storage, and seismology (Pathak et al.) 2022 Wen et al.) 2022 Yang al. 202T). While FNOs 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.
This growth may become bottleneck in their application to high-resolution physical simulations such as climate or materials modeling. In general, despite significant speed up and better flexibil- ity, prior works on neural operators suffer from similar memory complexity issues as conventional solvers do on high-resolution problems.
In the worst case, the large memory complexity is required and, fact, unavoidable due to the need for resolving fine scale features globally. However, many real-world problems, possess local structure that is not currently exploited by neural operator methods. For instance, consider a weather forecast where predictions for the next hour heavily dependent on the weather conditions in local regions and minimally on global weather conditions. Incorporating and learning this local structure of the underlying PDEs is the keys overcoming the curse of memory complexity.
2

---

Under review as a conference paper at ICLR 2023
In this work, we propose 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), model that exploits locality in physical space by novel multi-grid domain decomposition approach to compress the input domain size by up to 7x while leveraging the global interactions of the model parameters compress them by up to x without any loss of accuracy (see Table
In the input space, to predict the solution in any region of the domain, MG-TFNO decomposes the input domain into small local regions to which hierarchical levels of global information are added in a multi-grid -rid fashion. Since 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-TENO achieves low memory complexity by using high resolution data only locally and coarse resolution data globally. Due to its state-of-the-a performance on PDE problems and efficient FFT-based implementation, we use the FNO as the backbone architecture for our method.
In the parameter space, we exploit the spatio-temporal 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 between all the weights in the Fourier space by jointly parameterizing them with a single tensor, learned in factorized form such as Tucker or Canonical-Polyadic (Kolda & Bader 2009). This coupling allows us to limit the number of parameters in the model without limiting its expressivity. On the contrary, this low-rank regular- ization on the model mitigates wwer-fitting and improves generalization. Intuitively, our method can be thought of as a fully-learned implicit scheme capable of converging in a very few fixed number of iterations, Due to the global nature of the integral kernel transform, the FNO avoids the Courant-FFiedrren Leeyy (CFL) condition plaguing explicit schemes, allowing convergence in 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 reducing the complexity of the optimization landscape.
In summary, we make the following contributions:
We propose MG-TFNO, a novel neural operator parameterized in the spectral domain by a single low-rank factorized tensor, allowing its size to grow linearly with the size of the problem.
Our tensor operator achieve better performance with a fraction of the parameters: we outper- form FNO on solving the turbulent Navier Stokes equations with more than 400x weight compres- sion ratio, Figure 4a
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 5
We introduce 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 no performance degradation.
Combining tensorization with multi-grid domain decomposition leads to our best model. MG- TFNO, which is more efficient in terms of task performance, computation, and memory. MG-TFNO achieves 2.5x lower error with 10x model weight compression, and 1.8x domain compression, T
2 BACKGROUND
Here, we review related works and introduce the background necessary to explain our approach.
Neural Operators. Many physical phenomena are governed by PDEs and wide range of scien- tific and engincering computation problems are based on solving these equations. In recent years, a new perspective to PDEs dictates to formulate these problems 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 jet al. 2019). 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 attempts
3

---

Under review as a conference paper at ICLR 2023
have been made to address this limitation. For example mesh free methods are proposed that locally output mesh-free solution (Esmaeilzadeh et al.) 2020), but they are local and still limited to fixed input gird.
A new deep learning paradigm, neural operators, are proposed as maps between function spaces
et al. 2020b Kovachki al. 2021bb. They are discretization invariants maps. The input functions to neural operators can be presented any discretization, mesh, resolution, or basis. The output functions can be evaluated any point in the domain. Variants of neural operators deploy a variety of Nyström approximation to develop new neural operator architecture. Among these, multi-pole neural operators (Li al., (020c) utilize the multi-pole approach to develop computationally effi- cient neural operator architecture. Inspired by the spectral method, Fourier-basse -bseed neural operators show significant applicability in practical applications (Li et. al. 2020a Yang al.) 2021] Wen al.) 2022 Rahman et al. 2022a). Principle component analysis, wavelet bases, and u-shaped methods
are also considered (Gupta al., 2021 Bhhatacharya et al., 2020; Liu et al., 2022, Rahman et al.) 2022). It is also shown that neural operators solely be trained using PDEs, 2022b Yang et al.
resulting in pwwspcc--ifoorsteee informedd neural operators, opening new venues for hybrid data and equation methods (Li et al. b021b) to tackle problems in scientific computing.
Tensor methods in deep learning. As we move beyond learning from simple structures to solv- ing increasingly complex problems, the data we manipulate becomes more structured. To efficiently manipulate these structures, stucture. For all 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 (three dimensional) tensors, videos are 4th order tensors and so on and so forth. Tensor methods methods generalize linear algebraic methods to these higher-order structures. They have been very successful in various application in computer vision, signal processing, data mining and
need to go beyond matrix algebra and leverage the nempiol
machine learning (Panagakis al. 2021 Janzamin et al. 2019 Sidiropoulos al. 2017 Papalex- akis et al. 2016).
Using tensor decomposition Kolda & Bader 2009). previous works have been able to compress and improve deep networks for vision tasks. Either weight matrix is tensorized and factorized Novikov al. (2015). tensor decomposition is directly to the convolutional kernels before fine-tuning to recover-for lost accuracy, which also allows for an efficient reparametrization of the network (Lebe-
dev et al., 2015 Kim et al. 2016 Gusak et al. 2019). There is a tight link between efficient convo- lutional blocks and tensor factorization and factorized higher-order structures (Kossaifi et al.) 2020). Similar strategies have been applied to multi-task learning (Bulat et al., (020)) and NI.P (Papadopou-
los 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 tensor compression to learning operators and propose Tensor OPerator (TPNO).
METHODOLOGY
Here, we briefly review operator learning well as the Fourier Neural Operator, which we build to introduce our proposed Tensor OPerator (TFNO) as well as the Multi-Grid Domain Decomposi- tion, which together form proposed MG-TENO.
3.1 OPERATOR LEARNING
Let A = {a DA + Rdd) and 24 : {u Dr - Rddu) denote two input and output function spaces respectively. Each function in the input function space A, is a map from a bounded, open set DA C. Rd to the d d-ddimensional Euclidean space. Any function in the output function space l is map from bounded open set Du Id to the de-dimensional Euclidean space. In this work we consider the case D = D = Dus C. I2d.
We aim to learn operator G A 4 which is mapping between the two function spaces. In particular, given a dataset of points {(a,, (y),,))), , y where the pair (a,, uy) are functions satisfying (aa;) = Uy. 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. For
4

---

Under review as a conference paper at ICLR 2023
overview of theory and implementation, refer the reader to Kovachki et al. (021b). specifically use the ENO and give details the forthcoming section (Li et al. 2020a)
We
3.2 FOURIER NEURAL OPERATORS
For simplicity, we will work on the d-dimensional unit torus T and first describe single, pre- activation FNO layer mapping R" -valued functions to -a-lul functions. Such layer consti-
tutes the mapping G L2 (3 Rm) ) + (T R2) defined as
g(v) = = FF
Ve LT
where &. E L L (TTTRR [=*"" is function constituting the layer parameters and F- 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 a E N.
To implement T F,F-l are replaced by the discrete fast Fourier transforms F,8-1. Let ê E R51x""Xs x denote the evaluation of the function a uniform grid discretizing T with 8 E N points in each direction. We replace F(x) with a weight tensor T E Cs X-xx Xxmm consisting of the Fourier modes of x which are parameters be learned. To ensure that is parameterized as a I valuud function with fixed, maximum amount of wavenumbers a < min(s1,.... ,r} that is independent of the discretization of TV, leave learnable parameters only the first entries T each direction and enforce that T have conjugate symmetry. In particular, we pa- rameterize half the corners of the d-dimensional hyperrectangle with 2d hypercubes with length size a. That is, T is made up of the free-parameter tensors T., T2 situated in half of the corners of T. Each corner diagonally opposite of tensor Tý is assigned the conjugate transpose values Tj- All other values of T are set to zero. This is illustrated in the middle-top part of Figure for the case d = with Tr and T2. We will use the nota- tion T(k, T(). Tx for any h E [2d-1]. The discrete version of I then becomes the mapping R
Cax--XaxxXm
+ Rs: x=xx x defined as
( - F()), V
where the operation is simply the matrix multiplication contraction along the last dimension. Specifically, we have
T (())
1a (l.
Σ (4.0........)))((((()))) la,
From (2) full FNO layer is build by adding point-wise linear action to û, bias term, and applying a non-linear activation. In particular, from an input E IR x--1x---xa-0 Xeaxm the output q E Roi x.. Kea X given
q(1,--. (4. = (Qu(4, la.:) + +b)
with R R a fixed, non-linear activation, and b RR. Q Raxm T T..... To Cax xaxxxm are the learnable parameters of the layer. The full FNO model consists of E N
such layers each with weight tensors T.. Tz that have learnable parameters T (k..... for any i E [Z.] and h [2d-1]. In the case tensor W E (ax-xxxe Xxxxxxxxxxxxxx so that
for all layers, we introduce the joint parameter
W (..,, 2d - + + (1)
Perusal of the above discussion reveals that there are (2 d )m + total parameters in each FNO layer. Note that, since m and n constitute the respective input and output channels of the layer, the number of parameters can quickly explode due the exponential scaling factor 2dd if many wavenumbers are kept. Preserving large number of modes could be crucial for applications where the spectral decay of the input output functions is slow such as in image processing the modeling of multi-scale physics. In the following section we describe a tensorization method that is able to mitigate this growth without sacrificing approximation power.
5

---

Under review as a conference paper at ICLR 2023
3.3 TENSOR OPERATORS
In the previous section, we introduced unified formulation of FNO where the whole operator is parametrized by single parameter tensor W. This enables us to introduce the tensor operator, which parameterizes efficiently W with a low-rank, tensor factorization. In this paper, we focus mostly the Tucker decomposition, for its flexibility, but other decompositions such as Canonical Polyadic can be readily integrated. This joint parametrization has several advantages: low-rank constraint late into i) a huge reduction in the number of parameters, ii) better generalization and an operator less prone to overfitting. We show superior performance for lowcompression ratios (up to x00x) and very little performance degradation when largely compressing (> 450x) the model, iii) better performance in low-data regime.
applies the entire tensor W, thus regularizing the model. These advantages trans-
In practice, we express W in low-rank factorized form, e.g. Tucker or CP. In the case of a Tucker factorization with rank (R..... Rd, RL, R₁, Ro), where Rg controls the rank across layers, R₃ Ro control the rank across the input and output co-dimension, respectively, and R..... Rı control the rank across the dimensions of the operator:
R
Rd
R1
Ro
R2
W
Σ
Σ
Σ
G(r .. Ta.. "
U (..ra.
r==1
r-
( (U U (
Here, G is the core of size R2 x R₁ x Ro x Rí .... x Ra and U U U U(),... .... Uld) are factor matrices of size (Ry x L), (Ri x I), (Ro (R1 x xa), ..)..... (Ra o)), respectively.
W
Note that the mode (dimension) corresponding to the co- dimension can be left uncompressed, by setting R, T. and UIL) ld. This leads to a 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 T directly in factorized form to each layer by selecting the corresponding rows in U(L) The contraction in equation 3 is then either done using the reconstructed tensor, or by directly contracting F(i) with the factors of the decomposition.
Figure 2: Illustration of Tucker decomposition. For clarity show W as a 3dd-order tensor weight.
This joint factorization along the entire operator allows us to leverage redundancies both locally and across the entire operator. This leads to much reduced memory footprint, with only a fraction of the parameter. It also acts as a low-rank regularizer on the operator, facilitating training. Finally, through the global parametrization, we introduce skip connections that allow gradient to flow through the latent parametrization to all the layers jointly, leading to better optimization.
This formulation is general and works with any tensor factorization. For instance, we also explore Manonical-Polyadde decomposition (CP) which can be diagonal core. In that case, we set a single rank R and express the weights as a weighted sum of rank-1 1 tensors. Concretely:
special case of Tucker with super-
W Σ X (. (.rr (W U(O)(.)) (L))..))
(4)
where ULL), U U()), U ... U(d) are factor matrices size (Rx L), (RxI), (RxO), (Rx a).... (R x a), respectively and À e R Note that the CP. contrarily to the Tucker, has single rank parameter, shared between all the dimensions. This means that to maintain the number of
parameters the same, R needs to be very high, which leads to memory issues. This makes CP more suitable for large compression ratios, and indeed, we found leads to better performance at high-compression / very low-rank.

---

Under review as a conference paper at ICLR 2023
Level
Level
Padded Input Region
Output Region
Level
(a) Predicting with padded regions. Local region in the input is padded and used predict the corre- sponding region in the output.
(b) MG-Domain Decomposition. Progres- sively larger spatial regions are added to lo- cal region by subsampling.
Figure 3: Domain decomposition in space (3a) and our Multi-Grid based approach. (3b) White
squares represent the region of interest while yellow squares the larger embeddings.
3.4 MULTI-GRID DOMAIN DECOMPOSITION
Having introduced our decomposition in the operator's parameter space, we now introduce our novel multi-grid appraoch to decompose the problem domain.
Domain decomposition is method commonly used to parallelize classical solvers for time- dependent PDEs that is based on the principal that the solution for fixed local region in space depends mostly on the input at the same local region (Chan & Mathew 1994).
particular, since
the time-step >> 0 of the numerical integrator is small, the solution u(x. 8 h), for any point x: D and E R,, depends most strongly on the points u(y, t) for all E B (xx, (x)) where B ((Irr(h) denotes the ball centered at r with radius r(h). This phenomenon is easily seen for the case of the heat equation where, in dimension, the solution satisfies
(I - y)² 4h
rr+4h
(-x-)))²
((x - y)² 4h
u(x, L h) x
exp
w(y, dy
exp
u(y, () dy
4h
with the approximation holding since 99.9937% of the kernel's mass is contained within B(x, 4h). While some results exist, there is no general convergence theory for this approach, however its empirical success has made it popular for various numerical methods (Albin & Bruno (011).
To exploit this localization. the domain D is split in q E N pairwise-disjoint regions DI. D, so that = -و= Dy. Each region Dj is then embedded into a larger one Zj C Dj so that points away from the center of Dj have enough information be well approximated. A model can then be trained so that the approximation G z, D, s uID, holds for all j [q]. This idea is illustrated in Figure 3a 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 q.
Multi-Grid. Domain decomposition works well in classical solvers when the time step h > 0 is small because the mapping u(,,)) ~ u(,/++) h)) I) close the identity. However the major advance- ment made by machine learning based operator methods for PDEs is that model can approximate the solution, in one shot, for very large times i.e. h > 1. But, for larger the size of Z, relative to D, 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 mitigate this, we propose multi-grid based domain decomposition approach where global in- formation is added hicrarchically at different resolutions. While our approach is inspired by the classical multi-grid method, it is not based on the V-cycle algorithm (McCormick) 1985). For of presentation, we describe this concept when a domain = T uniformly discretized by 2 x 2 points, for some N, but note that generalizations we first sub-divide the domain into 22L total regions each of size 2--L x 2* and denote them 221. We call this the zeroth level. Then, around each (00) for any j [22*],
readily be made. Given final level L. N,
(
(0))
consider the square D of size 2--4+1 2+11++ that is equidistant, in every direction, from each boundary of (/ We then subsample the points in D) uniformly by factor of 1 in each direc- tion, making D/'" have 2 2 points. We call this the first level. We continue this process by considering the squares ( of size L+2 x 2 around each D2 and subsample them
7

---

Under review as a conference paper at ICLR 2023
4.50%
regions]
4.00%
Ours [16 regions]
(4 regions, compression]
12.5%
Ours [16 regions, 10x comprassion)
Feet ertrr
3.00%
Full Field Model
t7 tode error
6.2%
TFNO (Tucker)
2.50%
Trimmed FNC
2.00%
3.1%
1.50%
1.00%
1.6%
2
6
100
200
303
Compression-Ratie
Compression-Rate
(h) MG-Domain Decomposition: error func- tion the domain compression ratio. We compare MG-TENO with different numbers of multigrid re- gions both with and without weight tensor compres- sion full field FNO model. We achieve over 7x input space compression and 10x parameter space compression ratios with better performance that the original FNO.
(a) Tensorization: error in logscale as a function of the compression ratio. We compare the tensor neu- ral operator with FNO with the same number of parameters (Irimmed). We achieve pression ratio with better performance that the original ENO with Tucker, and over 400x com- pression ratio with CP.
200x com-
Figure 4: Impact of Tensorization and Multi-Grid Domain Decomposition (4b)
uniformly by a factor of + in each direction to again yield squares with 2 2 points. The process is repeated until the Lth level is reached wherein ( the entire domain subsampled by a factor of 2- in each direction. The process is illustrated for the case L = 2 in Figure (3b Since work with the torus, the region of the previous level always at the center of the current level.
The intuition behind this method is that since the dependence of points inside a local region dimin- ishes the further we are from that region, it 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 7,") of size 28-L + 2p 2--L + 2p around each DD where p E N the amount of padding to be added in each direction. We then take the evaluations the input function a at each level and concatenate them as channels. In particular,
train a model so that (all200..... . (l2(())) p( R - Since the model only operates on each padded region sep-
arately, we reduce the total number of grid points used from 22 to (2 + 2p)² and define the domain compression ratio the quotient of these numbers. Furthermore, note that, assuming is ud--vuleed, model that does not employ our multi-grid domain decomposition uses inputs with d channels while our approach builds inputs with (L + L)d, channels. particular, the number of input channels scale only logarithmically in the number of regions hence global information is added very little additional cost. Indeed, FNO models are usually trained with internal widths much larger than A hence the extra input channels cause almost no additional memory overhead.
EXPERIMENTS
Data. We experiment on a dataset 10K training samples and 2K test samples the two- dimensional Navier-Stokes equation with Reynolds number 500. We also experiment with the one- dimensional viscous Burgers' equation. Details about the datasets as well the results on Burgers are given in Appendices A.2 and A.3 respectively.
Training the operator. Since MG-TFNO predicts local regions which are then stitched together to form 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 et al. 2017: Li et al.] 2021a). By matching derivatives, training with this loss prevents any discontinuities from occurring and the output prediction is smooth.
Tensorizing: better compression. In Figure 4 we show the performance of our approach (TNO) compared to the original PNO, for varying compression ratios. In the Trimmed- FNO, adjust the
8

---

Under review as a conference paper at ICLR 2023
16.0%
TENO Tucker 10x [test]
25.0%
TENO [Tucker, 10x compression]
TENO [Tuckor, compression]
14.0%
ENO [tesl]
TFNO Tucker 10x [train]
1'tost crer
12.5%
FNO
[CP. 160x compression]
12.0%
FND [train]
20.0%
0.2%
L'
8.0%
3.18
6.0%
4.0%
1.6%
6
2000
4000
6000
8000
10000
0
100 200
350
400
500
Number Training Samples
Epochs
Figure 5: Error as a function of the number of training samples (left) and training VS testing loss. We compare TPNO with regular FNO. Note that the left we show the testing L error while, for training, the H' loss is used and that is compared with the H' test-error on the right.
width in order to match the number of parameters in our TNO. Our method massively outperforms the Trimmed-FNO at every single fixed parameter amount. Furthermore, even for very large com- pression ratios, our FNO outperforms the full-pan palammeeetter FNO model. This is 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 generalization. Figure 5 (left) shows that our TNO generalizes better with less training samples. Indeed, at every fixed amount of training samples, the TNO massively outper- forms the full-parameter FNO model. Even when only using half the samples, our TNO outperforms the FNO trained on the full dataset. Furthermore, Figure ( (right) shows that our TNO overfits sig- nificantly less than FNO, demonstrating the regularizing effect of the tensor decomposition. This result is invaluable in the PDE setting where we frequently have very little training samples available due to the high computational cost of traditional PDE solvers.
Multi-Grid Domain Decomposition. We show the impact of multi-grid domain decomposition on performance in Figure 4b 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 a form of data augmentation, exploiting the transnational 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 regularizer and improves performance across the board.
Putting it all together: MG-TENO. Tensorization and multi-grid -rid domain decomposition not only improve performance individually, but their advantages compound and lead to strictly better algo- rithm that scales well to higher-resolution data by decreasing the number of parameters in the model as well as the size of the inputs thereby improving performance well memory and computa- tional footprint. Table T shows a comparison of FNO with Tensorization alone, multi-grid domain decomposition alone, and our joint approach combining the two, MG-TFNO. Our results imply that, under aall-parlllliizaionn, the memory footprint of the model's inference can be reduced by 7x and size of its weights by 10x while also improving performance.
CONCLUSION
In this work, we introduced novel tensor operator (TFNO) as well as a multi-grid domain decom- position approach which together form MG-TENO, an operator model that outperforms the FNO with a fraction of the parameters and memory complexity requirements. MG-TENO scales better, generalizes better, and requires fewer training samples to reach the same performance. In our future work, we plan to deploy MG-TENO and tackle very high-resolution large-scale weather forecasts for which existing deep learning models are prohibitive.

---

Under review as a conference paper at ICLR 2023
REFERENCES
Jonas Adler and Ozan Oktem. Solving ill-posed inverse problems using iterative deep neural net- works. Inverse Problems, nov 2017.
Nathan Albin and Oscar P. Bruno. A spectral fo solver for the compressible navier-stokes equations in general domains i: Explicit time-stepping. Journal of Computational Physics, 230(16):6248 6270, 2011.
Saakaar Bhatnagar, Yaser Afshar, Shaowu Pan, Karthik Duraisamy, and Shailendra Kaushik. Predic- tion of aerodynamic flow fields using convolutional neural networks. Computational Mechanics, pp. 1-21, 2019.
Kaushik Bhattacharya, Bamdad Hosseini, Nikola Kovachki, and Andrew M Stuart. Model reduc- tion and neural networks for parametric pdes. arXiv preprint arXiv:2005.03/800 2020.
Mackenzie L Blanusa, Carla J López-Zurita, and Stephan Rasp. The role of internal variability in global climate projections of extreme events. arXiv preprint arXiv:2208.08271 2022.
Adrian Bulat, Jean Kossaifi, Georgios Tzimiropoulos, 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. 2020.
Tony F. Chan and Tarek 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-555 595, 2013.
Jean-Baptiste Cordonnier, Andreas Loukas, and Martin Jaggi. Multi-heal attention: Collaborate instead of concatenate. arXiv preprint 210i:::200616636622 2020.
R. Courant, K. Friedrichs. and H. Lewy. Uber die partiellen differeenzengllichungen der mathema- tischen physik. Mathematische annalen, 100(1):32-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.
Maarten De Hoop, Daniel Zhengyu Huang, Elizabeth Qian, and Andrew Stuart. The cost- accuracy trade-off in operator learning with neural networks. arXiv preprint 1xXx:22033.31811, 2022.
Soheil Esmaeilzadeh, Kamyar Azizzadenesheli, Karthik Kashinath, Mustafa Mustafa, Hamdi A Tchelepi, Philip Marcus, Mr Prabhat, Anima Anandkumar, et al. Meshfreeflownet: A physics- constrained deep continuous space-time super-resolution framework. In SC20: International Conference for High Performance Computing, Networking, Storage and Analysis, pp. 1-15, IEEE, 2020.
Lawrence C. Evans. Partial differential equations. American Mathematical Society, 2010.
John Guibas, Morteza Mardani, Zongyi Li, Andrew Tao, Anima Anandkumar, and Bryan Catan- zaro. Adaptive fourier neural operators: Efficient token mixers for transformers. arXiv preprint arXiv:2111.13587, 2021.
Xiaoxiao Guo, Wei Li, and Francesco Iorio. Convolutional neural networks for steady flow ap- proximation. In Proceedings of the 22nd ACM SIGKDD International Conference Knowledge Discovery and Data Mining, 2016.
Gaurav Gupta, Xiongye Xiao, and Paul Bogdan. Multiwavelet-based operator learning for differen- tial equations. Advances in Neural Information Processing Systems, 34:24048-24000 2021.
Julia Gusak, Maksym Kholiavchenko, Evgeny Ponomarev, Larisa Markeeva, Philip Blagoveschen- sky, Andrzej Cichocki, and Ivan Oseledets. Automated multi-stage compression of neural net- works. Oct 2019.
10

---

Under review as a conference paper at ICLR 2023
Majid Janzamin, Rong Ge, Jean Kossaifi, Anima Anandkumar, et al. Spectral learning on matrices and tensors. Found. and Trends@ in Mach. Learn., 12(5-6):393-536 2019.
Yong-Deok Kim, Eunhyeok Park, Sungjoo Yoo, Taelim Choi, Lu Yang, and Dongjun Shin. Com- pression 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-500, 2009.
Jean Kossaifi. Tensorly-torch. http:/ss//.tttuu.occ////ettts//p 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, Antoine Toisoul, Adrian Bulat, Yannis Panagakis, Timothy M Hospedales, and Maja Pantic. Factorized higher-order CNNs with an application to spatio-temporal emotion estimation. pp. 6060-6069, 2020.
Nikola Kovachki, Samuel Lanthaler, and Siddhartha Mishra. On universal approximation and error bounds for fourier neural operators. Journal of Machine Learning Research, 22((90))1-76,, 2021a.
Nikola Kovachki, Zongyi Li, Burigede Liu, Kamyar Azizzadenesheli, Kaushik Bhattacharya, An- drew Stuart, and Anima Anandkumar. Neural operator: Learning maps between function spaces. arXiv preprint arXiv: 100 100.08818, 2021
Vadim Lebedev, Yaroslav Ganin, Maksim Rakhuba, Ivan V. Oscledets, and Victor S. Lempitsky. Speeding-up convolutional neural networks using fine-tuned CP-decomposition. 2015.
Martin Leutbecher and Tim Palmer. Ensemble forecasting. Journal of computational physics, 227 (7)33515-3539, 2008
Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, An- drew Stuart, and Anima Anandkumar. Fourier neural operator for parametric partial differential equations. arXiv preprint arXiv:2010.08895, 2020a.
Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, An- drew Stuart, and Anima Anandkumar. Neural operator: Graph kernel network for partial differ- ential equations. arXiv preprini arXiv:2003.03448 2020b.
Zongyi Li, Nikola Kovachki, Kamyar Nzizzadenesheli, Burigede Liu, Andrew Stuart, Kaushik Bhat- tacharya, and Anima Anandkumar. Multipole graph neural operator for parametric partial differ- ential equations. Advances in Neural Information Processing Systems, 33:6775-6766 2020c.
Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli. Burigede Liu, Kaushik Bhattacharya, An- drew Stuart, and Anima Anandkumar. Markov neural operators for learning chaotic systems. arXiv preprint
06898, 2021a.
Zongyi Li, Hongkai Zheng, Nikola Kovachki, David Jin, Haoxuan Chen, Burigede Liu, Kamyar Azizzadenesheli, and Anima Anandkumar. Physics- informed neural operator for learning partial differential equations. arXiv preprint arXiv:2111.03794, 2021b.
Burigede Liu, Nikola Kovachki, Zongyi Li, Kamyar Azizzadenesheli, Anima Anandkumar, An- drew M Stuart, and Kaushik Bhattacharya. A learning- based multiscale method and its applica- tion to inelastic impact problems. Journal of the Mechanics and Physics of Solids, 158:104668, 2022.
S. F. McCormick. Multigrid methods for variational problems: General theory for the cycle. SIAM Journal on Numerical Analysis, 22(4):634-643, 1985.
Alexander Novikov, Dmitry Podopríkhin, 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-233 2317, September 2011.
11

---

Under review as a conference paper at ICLR 2023
Yannis Panagakis, Jean Kossaifi, Grigorios G. Chrysos, James Oldfield, Mihalis A. Nicolaou, Anima Anandkumar, and Stefanos Zafeiriou. Tensor methods in computer vision and deep learning. Proceedings of the IEEE, (09(5):833899) 2021. doi: 10. 10.109//R00.22221333745
Christos Papadopoulos, Yannis Panagakis, Manolis Koubarakis, and Mihalis Nicolaou. Efficient learning of multiple nlp tasks via collective weight factorization on bert. Findings of the Association for Computational Linguistics: NAACL 2022, pp. 882-890 092-800,4 2022.
Evangelos E Papalexakis, Christos Faloutsos, and Nicholas D Sidiropoulos. Tensors for data mining and data fusion: Models, applications, and scalable algorithms. ACM Trans. Intell. Syst. and Technol. (TIST), 822):1444, 2016.
Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang. Zachary DeVito, Zeming Lin, Alhan 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, et Four- castnet: A global data-driven high-resolution weather model using adaptive fourier neural opera- tors. arXiv preprint arXiv:2202. 11214, 2022.
Catherine E. Powell, Gabriel Lord, and Tony Shardlow, An Introduction Computational Stochastic PDEs. Texts in Applied Mathematics. Cambridge University Press, United Kingdom, 1 edition, August 2014. ISBN 9780521728522.
Md Ashiqur Rahman, Manuel Florez, Anima Anandkumar, Zachary Ross, and Kamyar Aziz- zadenesheli. Generative adversarial neural operators. arXiv preprint arXiv:2205.33011 2022a.
Md Ashiqur Rahman, Zachary E Ross, and Kamyar Azizzadenesheli. U-no: U-shaped neural oper- ators. arXiv preprint arXiv:2204. 11127, 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. Trans- actions Signal Processing, 65(13):3551-3588 2017.
Julia Slingo and Tim Palmer. Uncertainty in weather and climate prediction. Philosophical Trans- actions of the Royal Society A: Mathematical, Physical and Engineering Sciences, 369(1956): 4751-4767, 2011.
Roger Temam. Infinite- fefinitediimessional dynamical systems in mechanics and physics. Applied mathe- matical sciences. Springer-Verlas Verlag, New York, 1988.
Gege Wen, Zongyi Li, Kamyar Azizzadeneshhli,, Anima Anandkumar, and Sally M Benson. U- fno fno-an enhanced fourier neural operator- deep-learning model for multiphase flow. Ad- vances Water Resources, 163:104188 2022.
Yan Yang, Angela F Gao, Jorge C Castellanos, Zachary E Ross, Kamyar Azizzadenesheli, and Robert W Clayton. Seismic wave propagation and inversion with neural operators. The Seismic Record, 1(3):126-11 134, 2021.
Yan Yang, Angela F Gao, Jorge C Castellanos, Zachary E Ross, Kamyar Arizzadenesheli, and Robert W Clayton. Accelerated full seismic waveform modeling and inversion with u-shaped neural operators. arXiv preprint 22Xiv: 11955 2022.
Yinhao Zhu and Nicholas Zabaras. Bayesian deep convolutional encoder networks for sur- rogate modeling and uncertainty quantification. Journal of Computational Physics, 2018. ISSN 0021-9991.
A APPENDIX
A.1 NOTATION
We summarize the notation used throughout the paper in Table
12

---

Under review as a conference paper at ICLR 2023
Variable
Meaning
Dimensionality
T
Tensor of weights the Fourier domain
Çax-xaxmmX
W
Weight tensor parameterizing the entire operator
xax-xXxxxxxxx"
A
Input function space output function space
Infinite
Infinite Infinite Infinite
a
Input function
Output function
DA
Domain of function
d
Du
Domain of function
d
da
Dimension of the co-domain of the input functions
1
du
Dimension of the co-domain of the output functions Fourier transform
Infinite Infinite
Fourier transform
L
Number of integral operation layers
In N
Layer index
Between 1 and T. Infinite
Point-wise activation operation
b
Bias vector
Function each layer
Infinite
Number of kept frequencies in Fourier space
Between and min(s1, sa)
Table 2: Table of notation
A.2 DATA
In this section, we introduce in detail the datasets used for our experiments.
Navier-Stokes.
We consider the vorticity form of the two-dimensional Navier-Stokes equation,
1
On + V- where T = [0, 2x is the torus, f E i' R) a forcing
Aw + f.
x E te (0,T)
Re
(5)
{ 0=
- Ad
TET te (0,7)
with initial condition -0)
function, and Re > 0 is the Reynolds number. Then wlt.)) E (T(TT; for any t E (0, 7] and 0, is the unique weak solution to 5 Temam 1988). We consider the non-linear operator mapping - (T. with 7 = and fix the Reynolds number Re = 500. We define the Gaussian measure μ = AO, C) the forcing functions where take the covariance C 2 + - following the setting in (De Hoop et al. (022). Input data is obtained by generating i.i.d. samples from μ by KL-expansion onto the eigenfunctions of C (Powell al. 2014). Solutions
then obtained by pseudo-spectral scheme (Chandler & Kerswell 2013
Burgers' Equation.
We consider the one-dimensional Burgers' equation on the torus,
O.u utx Vlllga
T, EE (0,7) T
(6)
21 t=== = to.
x E
for initial condition No E (TT R) and viscosity
O. Then u(t, +) E H*(T, for any 1 E R
and s > 0, the unique weak solution to 6 (Evans) 2010). We consider the non-linear operator
uo -)", with T' = 0.5 or and fix 0.01. We define the Gaussian measure
AO C)
where we take the covariance C 35/2 da + 97) Input data is obtained by generating i.i.d. samples from by a KI--expansion onto the eigenfunctions of C. Solutions to 6 are then obtained by a ppseudo-spectral solver using Heun's method. We use 8K samples for training and 2K testing.
A.3 TENSOR OPERATOR: TENSOR DECOMPOSITION
In this section, we expand further the tensor decomposition composing the tensor operator.
13

---

Under review as a conference paper at ICLR 2023
20.0%
I:NO Jucker 2x [test]
T-NO Tucker 172x [test]
15.0%
[train]
RNO [test]
[test]
THNO lucker 49x [trein]
12.5%
[train]
[train]
FNO [troin]
10.0%
error
7.5%
5.0%
10.0%
100
300
500
Epochs
(a) Train VS Test errors various compression ratios. We observe that the gap hetween training and test error decreases with the rank (e.g inversely pro- portional the compression ratio).
0
100
200
300
400
500
Epachs
(b) Train VS Test error over time for TOP with TT factorization.
Figure 6: Impact of the rank for a TOP Tucker (6a) and train/test curve for a TOP-TT (6b)
We first rewrite the entire weight parameter for the Tucker case, equivalently, using the more com- pact n-mode product as:
U W UO. xa Uld) xa+1 U xd+2 U Xd+3 U
iFFT of the tensorized kernel For any layer I, the (f1,32) coordinate of the matrix-valle valued con- volution function x(x) is follows,
Rr
R1
Ro
R
ns
[<(())])1,
Σ
G(ri... ra, Ts, Tos Tt)
(ii,)) ... Uid) ra))
(i,r)) U (j2, (1,)
exp(22 2m
imsik)
1=1
We also note that other tensor decomposition can be straightforwardly used in our framework, such as the tensor-train decomposition Oseledets (201I). A rank-((1 (1, Ri, RRN, R,, Ro, Rr, 1) TT factorization expresses W as:
W W(i,... ... id, les ins it) G (i1) x x (Ga x ... ...GGi) x ... G
Where each of the factors of the decompositions G3 are third order tensors of size R2 x I3 R We show examples of TENO trained with a TT factorization the coming sections.
A.4 IMPLEMENTATION DETAILS
Implementation We use Py Torch Paszke et al. (2017) for implementing all the models. The tensor operations are implemented using Tensorlyy and Tensorl y-Torch Kossaifi et al. (2019) Kossaifi 202T..
We will release the code and data used upon acceptance of the paper.
Hyper-parameters We train all models via gradient backpropagation using a mini-batch size of 16, the Adam optimizer, with a learning rate of 1e-3, weight decay of 1e for 500 epochs, de- creasing the learning rate every 100 epochs by a factors J. The model width set in all cases to 64 except when specified otherwise (for the Trimmed FNO), meaning that the input was first lifted (with a linear layer) from the number of input channels to that width. The projection layer projects from the width to 256 and prediction linear layer outputs the predictions. 10000 samples were used for training, as well as a separate set of 2000 samples for test. For a. we keep 40 Fourier coefficients for height and 24 for the width. All experiments are done NVIDIA Tesla V 1000 GPU.
14

---

Under review as a conference paper at ICLR 2023
Method
1.2 test error
# Params
Model
FNO
2.54%
58
0x
TOP [Tucker]
2.36%
366 K
172x
TOP [CP]
2.03%
350 K
179x
TOP [TT]
1.83%
353 K
178x
Table 3: Relative L test error of our MG-TENO approach for different tensor decompositions for comparable compression ratio.
A.5 ABLATION STUDIES
In this section, we further study the properties our model through ablation studies. We first look how TFNO suffers less from overfitting thank 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.
A.5.1 OVERFITTING AND LOW-RANK CONSTRAINT
Here, show that lower-ranks (higher compressions), lead to reduced overfitting. In Figure ba we show the training and testing H' errors for our TOP with Tucker decompositions, 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 we decrease the rank. As can see, Tucker, while being most flexible, does not perform as well at higher compression ratios, there CP and Tensor-Traie lead to lower error.
A.5.2 TENSOR-TRAUE AND TOP
Our approach is independent of the choice of tensor decomposition. We already showed how Tucker most flexible and works well across all ranks. We also showed that CP, while memory demanding for high rank, 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 here the convergence behaviour of TNO with Tensor-Train decomposition, for compression ratio of 178, figure 6b
We also compare Table our TENO with different tensor decompositions.
A.5.3 BURGERS' EQUATION
Radius 1.0
0.78%
Radius 1.2
Radius
Radius 0.5
Radius 0.8
Full Field Model
0.39%
Field Model
[3 test
L test
0.20%
0.10%
1.5
200
4'5
5.0
2.5
3.0
3'5
Compression-Ratio
Compression-Ratio
Figure 7: Error on Burgers' equation with T' 0.5 (left) and T (right) 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.
We test the efficacy of the standard domain decomposition approach by training on two separate Burgers problems: we expect that for T = 1, each region requires more global information thus significantly more padding need be used in order reach the same error. The results of F confirms this.
with final time T 0.5 and one with f 1. As described in Section 3.4
15

---

Under review as a conference paper at ICLR 2023
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 multi-grid domain decomposition approach.
16


Reviews:
[{"review_id": "p5EjKSlKFA", "reviewer": "Reviewer_u8oi", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper proposes a neural operator that scales to large resolutions by leveraging local and global structures through decomposition of both the input domain and the operator’s parameter space.\nThis is achieved by writing an FNO operation with a single tensor, and subsequently using low-rank tensor factorization. \n", "summary_of_the_review": "The paper presents a strong idea, which is also practically very relevant. I highly encourage the authors to work on the experiment section and to work out the benefits of the methods as well as relation to other methods in more detail. ", "strength_and_weaknesses": "Strengths:\n-\tThe paper presents a very interesting approach to a well-known problem: how to scale neural PDE surrogates efficiently to large resolutions\n-\tWriting the FNO basic block as single parameter tensor, and then applying matrix factorization on that is really neat. I think the idea has great potential.\n-\tThe multi-grid domain decomposition makes a lot of sense too.  \n\n\nWeakness:\n-\tThe biggest concern are the experiments. I cannot find any information on the grid, and any ablation on the baseline method. I could very well be that FNO with much less modes would perform equally well, which would make the achieved compression rates obsolete. It is really hard to judge the results without knowing the grid sizes and having some visual image of the data. Looking at Figure 3, the problems however don’t really seem to have a very high resolution. What is the temporal resolution? Which solver is used to obtain the data? What are the boundary conditions, stating Reynolds numbers and not the boundary condition is problematic. \n-\tThere is no mentioning of other operator learning methods (e.g. DeepONet from Lu et al.) as well as a discussion of many other neural PDE surrogate methods. There is a huge body of neural PDE solvers/surrogates and omitting 95% of that is not good practice. I however admit that not everything is relevant for this paper. DeepOnet, Adaptive Fourier Neural Operators, multigrid methods, Vision Transformers are.\t    \n-\tConnecting to the previous point: I am also very curious how this relates to Vision Transformers or Adaptive Fourier Neural Operators. Also recently UNets get tremendous performances in high-resolution image generation tasks, how are they doing, or even more interestingly can this approach be applied to e.g. UNets/DeepONets as well?\n-\tWhat do I see in Figure 5 right? What is H1 loss there? Both y axes state L2 loss? That fact that there is such a gap between FNO train vs test can very well be due to the rather basic training scheme (no warmup, lr annealing can surely be improved) or can be a sign that better regularization for the baseline is needed.\n\nLu, Lu, Pengzhan Jin, and George Em Karniadakis. \"Deeponet: Learning nonlinear operators for identifying differential equations based on the universal approximation theorem of operators.\" arXiv preprint arXiv:1910.03193 (2019).\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "c3z0qWpGcS7", "reviewer": "Reviewer_zgdt", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposes to predict solutions of PDE locally and unite them into a global solution. The proposed neural operator scales to large resolutions by leveraging local and global structures by decomposing both the input domain and the operator's parameter space. The proposed method can be trained with much fewer samples than previous approaches, outperforming the FNO when trained with just half the samples. ", "summary_of_the_review": "The paper proposes a new approach but solves only a few simple problems. There are already many similar methods, and it is difficult to say which is better and which is more practical. Studying large-scale examples in three dimensions may be more critical.", "strength_and_weaknesses": "Strength\n\nMG-TFNO outperforms the FNO with a fraction of the parameters and memory complexity requirements.\n\nFigure 1 clearly illustrates the structure of the network.\n\nWeaknesses\n\nHow to handle the Fourier transform of non-periodic problems?\n\nWill the computational cost of the Fourier transform be large?\n\nHow is the differentiation of the Fourier transform implemented?\n\nCan this method make a 3D prediction?\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "BiEGMBm_iuT", "reviewer": "Reviewer_J6V8", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper focuses on the problem of neural surrogates of high-resolution PDEs. To this end, it applies the principles of multi-grid domain decomposition to exploit the local structure in such data as well applies tensor factorization with low-rank regularaization to reduce the number of parameters. The method is evaluated on vorticity form of Navier-Stokes equations and shown to improve performance. Although it's unclear what was the resolution of data used in experiments.", "summary_of_the_review": "There are important details missing from the experiments in the paper to make a clear judgement but the idea of using lower ranked tensor factorization and domain decomposition  makes sense for the PDE domain.", "strength_and_weaknesses": "**Strengths**\n\nGiven GPU memory bottlenecks, different compute-memory tradeoff methods need to be explored as we keep increasing data sizes in deep learning. FNO based methods scale even worse with respect to memory when number of tracked modes increases. Proposed approach based on tensor decomposition methods makes sense for lower rank representation of model weights. \n\n**Weaknesses**\n\nWhether this is _hardware-efficient_ is not explored in this work. Something like Monarch [1] is the appropriate comparison for modern hardware. Matching derivatives during training makes the method even more expensive. Some timing comparisons would be useful.\nUsing % for reporting test error makes it harder to understand whether the underlying errors should be considered good. Similarly it's unclear what is \"high\" resolution in the paper because the experiments don't mention the spatial or temporal resolution. Moreover it would be useful to actually clarify if an external PDE solver was used or if the solver was hand written for datagen. \n\nSome of the claims about \"deep learning models being orders of magnitude faster than conventional solvers\" are based on incorrect comparisons; running on GPUs vs. CPUs, not using state-of-the-art PDE solvers rather comparing against hand written solvers, not taking into account approximation errors etc. \n\nIt's also a bit surprising that older neural operator methods like DeepONets [2, 3] are missing from related works.\n\n[1] https://arxiv.org/abs/2204.00595\n[2] https://arxiv.org/abs/1910.03193\n[3] https://arxiv.org/abs/2111.05512\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject