Deep Equilibrium Based Neural Operators for Steady-State PDEs
Authors: Tanya Marwah
Ashwini Pokle
J Zico Kolter
Zachary Chase Lipton
Jianfeng Lu
Andrej Risteski
Venue: NeurIPS 2023 poster (2023)
Keywords: Deep Equilibrium Models
Partial Differential Equations
Neural Operators

Abstract: Data-driven machine learning approaches are being increasingly used to solve partial differential equations (PDEs). They have shown particularly striking successes when training an operator, which takes as input a PDE in some family, and outputs its solution. However, the architectural design space, especially given structural knowledge of the PDE family of interest, is still poorly understood. We seek to remedy this gap by studying the benefits of weight-tied neural network architectures for steady-state PDEs. To achieve this, we first demonstrate that the solution of most steady-state PDEs can be expressed as a fixed point of a non-linear operator. Motivated by this observation, we propose FNO-DEQ, a deep equilibrium variant of the FNO architecture that directly solves for the solution of a steady-state PDE as the infinite-depth fixed point of an implicit operator layer using a black-box root solver and differentiates analytically through this fixed point resulting in $\mathcal{O}(1)$ training memory. Our experiments indicate that FNO-DEQ-based architectures outperform FNO-based baselines with $4\times$ the number of parameters in predicting the solution to steady-state PDEs such as Darcy Flow and steady-state incompressible Navier-Stokes. Finally, we show FNO-DEQ is more robust when trained with datasets with more noisy observations than the FNO-based baselines, demonstrating the benefits of using appropriate inductive biases in architectural design for different neural network based PDE solvers. Further, we show a universal approximation result that demonstrates that FNO-DEQ can approximate the solution to any steady-state PDE that can be written as a fixed point equation.

Deep Equilibrium Based Neural Operators for Steady-State PDEs
Tanya Marwah
Ashwini Pokle
J. Zico Kolter 1,2 Zachary C. Lipton
Jianfeng Lu 3 Andrej Risteski
CCarnegie Mellon University 2 Bosch Center for Al 3 Duke University
(tmarwah, (tmarwah,apokle apokle, zicokolter zlipton, (rriisteskaaandred Sanfeng@math. duke. edu
edu
Abstract
Data-driven machine learning approaches are being increasingly used solve par- tial differential equations (PDEs). They have shown particularly striking successes when training an operator, which takes as input PDE in some family, and outputs its solution. However, the architectural design space, especially given structural knowledge of the PDE family of interest, is still poorly understood. We seek remedy this gap by studying the benefits of weight-tied neural network architec- tures for steady-state PDEs. To achieve this, we first demonstrate that the solution of most steady-state PDEs can be expressed as fixed point of non-linear oper- ator. Motivated by this observation, we propose FNO-DEQ, a deep equilibrium variant of the FNO architecture that directly solves for the solution of steady-state PDE as the infinite-depth fixed point an implicit operator layer using black- box root solver and differentiates analytically through this fixed point resulting in O(1) training memory. Our experiments indicate that FNO-DEQ- based archi- tectures outperform FNO-based baselines with 4x the number of parameters in predicting the solution to steady-state PDEs such as Darcy Flow and steady- incompressible Navier-Stokes. Finally, we show FNO-DEQ more robust when trained with datasets with more noisy observations than the FNO-based baselines, demonstrating the benefits of using appropriate inductive biases in architectural design for different neural network based PDE solvers. Further, we show a univer- sal approximation result that demonstrates that FNO-DEQ can approximate the solution to any steady-state PDE that can be written as fixed point equation.
Introduction
Partial differential equations (PDEs) are used model a wide range of processes in science and engineering. They define a relationship of (unknown) function and its partial derivatives. Most PDEs do admit closed form solution, and are solved using variety of classical numerical methods such as finite element [Le Veque, 2007], finite volume [Moukalled et al., 2016), and spectral methods [Kopriva, 2009, Boyd, 2001]. These methods are often very computationally expensive, both as the ambient dimension grows, and as the desired accuracy increases.
This has motivated rapidly growing of research in data-driven approaches PDE solving. One promising approach involves learning neural solution operators [Chen and Chen, 1995, Lu et al., 2019, Bhattacharya et al., 2021, Li et al., 2020b], which take in the coefficients of PDE in some family and output its solution-and and are trained by examples of coefficieent-solttinn pairs.
"Equal contribution. Correspondence tmarwah@andrew. cmu. edu and apokle@andrew. cm..
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

While several architectures for this task have been proposed, the design space- in particular taking into account structural properties of the PDEs the operator trained on on-is still largely unexplored. Most present architectures are based on "neuralizing" classical numerical method. For instance, Li et al. [2020a] take inspiration from spectral methods, and introduce FNO: trained composition of (parametrized) kernels in Fourier space. Brandstetter et al. [2022] instead consider fiiie-difference methods and generalize them into (learnable) graph neural networks using message-passing..
Our work focuses on families of PDEs that describe the steady-state of a system (that is, there is no time variable). Namely, we consider equations of the form:
u(x)) = f(x),,
Vo E
where u : Q2 + Rd, a a - Rd and f RR are functions defined over the domain Ωt and a (possibly non-linear) operator. This family includes many natural PDE families like Poisson equations, electrostatic equations, and steady-state Navier-Stokes.
We take inspiration from classical numerical approaches of fast fast-converging Newton-like iterative schemes Veque, 2007, Faragó and Karátson, 2002] to solve steady-state PDEs, as well as recent theoretical works for elliptic (linear and non-linear PDEs) [Marwah et al., 2021, Chen al., 2021, Marwah et al., 2022] to hypothesize that very deep, but heavily weight-tied architectures would provide useful architectural design choice for steady-state PDEs.
In this paper, we show that for steady state equations it is often more beneficial to weight-tie an existing neural operator, opposed to making the model deeper-thus increasing its size. To this end, we introduce ENO-DEQ, a new architecture for solving steady-state PDEs. FNO-DEQ is a deep equilibrium model (DEQ) that utilizes weight-tied FNO layers along with implicit differentiation and root-solvers to approximate the solution of a steady-state PDE. DEQs are a perfect match to the desiderata laid out above: they can be viewed alternately directly parameterizing the fixed points of some iterative process; or by explicitly expanding some iterative fixed point solver like Newton's 's or Broyden's 's method as infinitely deep, weight-fied model.
Such an architecture has a distinct computational advantage: implicit layer models effectively backpropagate through the infinite-dpphh network while using only constant memory (equivalent to single layer's activations). Empirically, we show that for steady-state PDEs, weight-tied and DEQ based models perform better than baselines with 4x the number of parameters, and are robust to training data noise. In summary, we make the following contributions:
We show the benefits of weight-tying as an effective architectural choice for neural operators when applied to steady-state PDEs.
We introduce FNO-DEQ, a FNO based deep equilibrium model (DEQ) that uses implicit layers and root solving to approximate the solution of a steady-state PDE. We further attest to the empirical performance of FNO-DEQ by showing that it performs as well as FNO and its variants with 4x number of parameters.
We show that FNO-DEQ and weight tied architectures are more robust to both input and observation noise, thus showing that weight-tying is useful inductive bias for architectural design for steady-state PDEs.
By leveraging the universal approximation results of FNO [Kovachki et al., 2021a] we show that FNO-DEQ based architectures can universally approximate the solution operator for wide variety of steady-state PDE families.
Finally, we create a dataset of pairs of steady-state state incompressible Navier-Stokes equations with different forcing functions and viscosities, along with their solutions, which we will make public as a community benchmark for steady-state PDE solvers.
2 Related Work
Neural network based approaches for solving PDEs can broadly be divided into two categories. First are hybrid solvers [Bar-Sinai et al., 2019, Kochkov et al., 2021, Hsieh et al., 2019] which use neural networks in conjunction with existing numerical solvers. The main motivation to not only improve upon the existing solvers, but to also replace the more computationally inefficient parts of the solver

---

with learned counter part. Second set of approaches are full machine learning based approaches that aim leverage the approximation capabilities of neural networks [Hornik et al., 1989] to directly learn the dynamics of the physical system from observations.
Hybrid solvers like Hsieh et al. [2019] use neural network learn correction term correct over an existing hand designed solver for a Poisson equation, and also provide convergence guarantees of their method the solution of the PDE. However, the experiments in their paper are limited to linear elliptic PDEs. Further, solvers like Bar-Sinai et al. [2019] use neural networks to derive the discretizations for a given PDE, thus enabling the use of a low-resolution grid in the numerical solver. Furthermore, Kochkov et al. [2021] use neural networks to interpolate differential operators between grid points of low-resolution grid with high accuracy. This work specifically focuses on solving Navier-Stoker equations, their method is more accurate than numerical techniques like Direct Numerical Simulation (DNS) with a low-resolution grid, and is also 80x more faster. Brandstetter et al. [2022] introduced a message passing based hybrid scheme to train hybrid solver and also propose loss term which helps improve the stability of hybrid solvers for time dependent PDEs. However, most of these methods are equation specific, and are not easily transferable to other PDEs from the same family.
The neural network based approach that has recently garnered the most interest by the community is that of the operator learning framework [Chen and Chen, 1995, Kovachki et al., 2021b, Lu al., 2019, Li et al., 2020a, Bhattacharya et al.. 2021], which uses a neural network to approximate and infinite dimensional operator between two Banach spaces, thus learning an entire family of PDEs once. Lu et al. [2019] introduces DeepONet, which uses two deep neural networks, referred to as the branch net and trunk net, which are trained concurrently to learn from data. Another line of operator learning framework is that of neural operators Kovachki et al. [2021b]. The most successful methodology for neural operators being the Fourier neural operators (FNO) [Li et al., 2020a]. FNO uses convolution based integral kernels which are evaluated in the Fourier space. Future works like Tran et al. [2021] introduce architectural improvements that enables one to train deeper FNO networks, thus increasing their size and improving their the performance on variety of (timeddependent) PDEs. Moreover, the success of Transformers in domains like language and vision has also inspired transformer based neural operators in works like Li et al. [2022a], Hao et al. [2023] and Liu et al. [2022]. Thcoretical results pertaining to the neural operators mostly include universal approximation results Kovachki et al. [2021a]. Lanthaler al. [2022] which show that architectures like FNO and DeepONet can indeed approximate the infinite dimension operators.
In this work, we focus steady-state equations and show the benefits of weight-tying in improving the performance of FNO for steady-state equations. We show that instead of making a network deeper and hence increasing the size of network, weight-tied FNO architectures can outperform FNO and its variants 4x its size. We further introduce FNO-DEQ, a deep equilibrium based architecture to simulate an infinitely deep weight-tied network (by solving for fixed point) with O(1) training memory. Our work takes inspiration from recent theoretical works like Marwah et al. [2021], Chen et al. [2021]. Marwah et al. [2022] which derive parametric rates for some-steady state equations, and in fact prove that neural networks can approximate solutions to some families of PDEs with just poly(d) parameters, thus evading the curse of dimensionality.
3 Preliminaries
We now introduce some key concepts and notation.
Definition 1 (L; )).. For a domain Q we denote by (R( the space of square integrable functions g : > Rd such that [1)]) < oc, where (0))))(0) (Jg Il9(2)) dx)'
3.1 Neural Operators
Neural operators [Lu et al., 2019, Li et al., 2020a, Bhattacharya et al., 2021, Patel et al., 2021, Kovachki et al., 2023] are a deep learning approach to learning solution operators which map a PDE to its solution. Fourier Neural Operator (FNO) [Li et al., 2020a] is a particularly successful recent architecture parametrized as sequence of kernel integral operator layers followed by non-linear activation functions. Each kernel integral operator layer is a convolution-based kernel function that instantiated through linear transformation in Fourier domain, making it less sensitive to the level of
3

---

spatial discretization. Specifically, an L-layered FNO Gg R - Rd with learnable parameters 6, defined as
Go Q Lt. L .... o
where P: (9, (I2 Rd.)) and L (IR;; ( + L (IRd; Bdd) are projection
operators, and L (RR,; Rd,) + I2 (R Rd:) for I [L] is the 1th FNO layer defined as,
(v,) (Wrer + bt + K(()))))
Here a is a non-linear activation function, Wi, by are the 1 layer weight matrix and bias terms. Finally Kr is the 1th integral kernel operator which is calculated using the Fourier transform as introduced in et al. [2020a] defined as follows,
(()))=F' (Ra (Fur)) (x)
Vo E Ω
where F and JF1 are the Fourier transform and the inverse Fourier transform, with RG representing the learnable weight-matrix in the Fourier domain. Therefore, ultimately, the trainable parameters a collection of all the weight matrices and biases, i.e, A (Wi, by, R,.... ,W, R,}.
3.2 Equilibrium Models
Equilibrium models [Liao et al., 2018, Bai et al., 2019, Revay et al., 2020, Winston and Kolter, 2020] compute internal representations by solving for a fixed point in their forward pass. Specifically, consider deep feedforward network with L layers:
2 [++1] - (2()) for i O,..,, L-1
where x E IR is the input injection, ا E R is the hidden state of ith layer with 20 = 0. and So IR. xn. + Th is the feature transformation of 7th layer, parametrized by 0. Suppose the above model is weight-tied, i.e., f! fe. Vi, and lim, fo (21); exists and its value is 2*. Further, assume that for this z", it holds that fo (=*; =" Then, equilibrium models can be interpreted as the infinite-eppth limit of the above network such that fø (=";=:)=
Under certain conditions and for certain classes of fo², the output 2" of the above weight- tied network is a fixed point. A simple way to solve for this fixed point is to use fixed point iterations, i.e., repeatedly apply the update 2|++1] (2(1))); some fixed number of times, and backpropagate through the network to compute gradients. However, this can be computationally expensive. Deep cquilibrium (DEQ) models [Bai et al., 2019] explicitly solve for 2 through iterative root finding methods like Broyden's method [Broyden, 1965], Newton's method, Anderson acceleration [Ander- son, 1965]. DEQs use implicit function theorem to directly differentiate through the fixed point 2* at equilibrium, thus requiring constant memory to backpropagate through an infinite-depth network:
0z*
0fa(2=;z) Oz* -1-
- -11-
0
(6)
00
00
Computing the inverse of Jacobian can quickly become intractable as we deal with high-dimensional feature maps. One can replace the inverse-Jacobian term with an identity matrix i.e., Jacobian- free [Fung et al., 2022] an approximate inverse-Jacobian [Geng et al., 2021] without affecting the final performance. There are alternate formulations of DEQs [Winston and Kolter, 2020] that guarantee existence of unique equilibrium point. However, designing fe for these formulations can be challenging, and in this work we use the formulation by Bai et al. [2019].
4 Problem setting
We first formally define the system of steady-state PDEs that we will solve for:
"The fixed point can be reached if the dynamical system is globally contractive. This is usually not true in practice for most choices fe. and divergence possible.
"Bai et al. [2019] state that fa needs to be stable and constrained. In general, by Banach's fixed point theorem. global convergence guaranteed if fo is contractive its input domain.

---

Definition 2
(Steady- Steady-Statteaa PDE). Given bounded open set 9 C R steady-state PDE can be written in the following general form:
M((x)) (()))) = f(x), Yx EΩ
(7)
Here is a continuous operator, the function u E L (92; (4.) is the unknown function functioon that we wish to
solve for and E L2 (S); IR)) collects all the coefficients describing the PDE, and f E L2 (S2; Rdr) a function independent of u. We will, for concreteness, assume periodic boundary conditions, i.e.
dz E Zd, Vr E Q2 we have u(x + 2) = u(æ). (Equivalently, St := T [0, 2x) can be identified with the torus.) Finally, we will denote u² + R as the solution the PDE.
Steady-state models system at stationarity, i.e, when some quantity of interest like temperature or velocity no longer changes over time. Classical numerical solvers for these PDEs include iter- ative methods like Newton updates or conjugate gradient descent, typically with carefully chosen preconditioning ensure benign conditioning and fast convergence. Furthermore, recent theoretical works [Marwah et al., 2021, Chen et al., 2021, Marwah et al., 2022] show that for many families of PDEs (e.g., steady-state elliptic PDEs that admit variational formulation), iterative algorithms can be efficiently "neuralized", that is, the iterative algorithm can be represented by compact neural network, so long as the coefficients of the PDE are also representable by compact neural network. Moreover, the architectures constructed in these works are heavily weight-tied.
We will operationalize these developments through the additional observation that all these iterative schemes can be viewed as algorithms to find fixed point of a suitably chosen operator. Namely, we can design an operator G (D); Rdx) x L2 (22; [ + (2); ) such that u² G(u f) and a lot of common (preconditioned) first and second-order methods are natural ways to recover the fixed points u*.
Before describing our architectures, we introduce two components that we will repeatedly use. Definition 3 (Projection and embedding layers). A projection and embedding layer, respectively laye
P (2 (dd x (22 - 1.- [2R Rddd) x (2R Rdd) and Q (RRd Rd) 7
L3(RRd Rdd), are defined as
P(v.S) = (o (w/P) 0 +b)) a (wf2) {bb (p)))
2(v) = (Wqu+ bQ)
where W E Rddxxd, xd. w xd, Wq E Rd. xd. and bp 11 Rdd ba Rdd..
Definition 4 (Input-injected FNO layer). An input-injected FNO layer & T:((RRe;RR) x L' Rdd) + L² (RR,, Rd)) is defined as
[(v,g) = g (Wu+b++ (1t
(8)
where W Rd.xde b dRd. and R(K) E Rd,xde for all E [K] are learnable parameters.
Note the difference between the FNO layer specified above, and the standard FNO layer Equation the extra input g to the layer, which in our architecture will correspond all) of the PDE coefficients. We also note that this change to the FNO layer also enables us to learn deeper FNO architectures, as shown in Section 5, With this in mind, we can discuss the architectures we propose.
projection of (some or
Weight-tied architecture I: Weight-tied FNO The first architecture we consider is weight- tied version of FNO (introduced in Section 3), in which we repeatedly apply (M times) the same transformation in each layer. More precisely, we have:
Definition 5 (FNO Weight-Tiedh We define a M times weight-tied neural operator GDM as,
Gb Q B² o
B² op
M times
'This for convenience of exposition, methods can readily be extended other boundary conditions like Dirichet, Neumann etc. 'We note that the choice of defining the operator with the forcing function f is made for purely expository purposes the operator G can be defined G L (L) 1dd)) x L 22( 12 + LL2 (d) well.
5

---

such that: P. Q are projection and embedding layers as in Definition 3 an B² (BR :R)) x (TRd,, RRd)) + 1.2 (IRd,, R24)) is an I-layer FNO block, ie, B² = for all I E [L] Ci(,P(())) is an input-injected FNO block as in Definition 4.
Lu o L o L o b where
Weight-tied architecture II: FNO-DEQ In cases where we believe a weight-tied G converges to some fixed point as M + oc, unrolling G. for a large M requires a lot of hardware memory for training: training the model requires one to store intermediate hidden units for each weight-tied layer for backpropagation, resulting in a (M) increase in the amount of memory required.
To this end, we use Deep Equilibrium models (DEQs) which enables us to implicitly train G9 lim G. by directly solving for the fixed point by leveraging black-box root finding algorithms like quasi-Newton methods, [Broyden, 1965, Anderson, 1965]. Therefore approach as an infinite-depth (or infinitely unrolled) weight-tied network. We refer to this architecture as FNO-DEQ.
think of this
Definition 6 (FNO-DEQ). Given P. Q and B in Definition 5, FNO-DEQ is trained to parametrize the fixed point equation B² (v",P()) = V² and oulputs u² = Q(=*).
Usually, it is non-trivial to differentiate through these black-box root solvers. DEQs enable to implicitly differentiate through the equilibrium fixed point efficiently without any need to backpropa- gate through these root solvers, therefore resulting in O(1) training memory.
5 Experiments
Network architectures. We consider the following network architectures in our experiments.
FNO: We closely follow the architecture proposed by Li et al. [2020a] and construct this network by stacking four FNO layers and four convolutional layers, separated by GELU activation [Hendrycks and Gimpel, 2016). Note that in our current set up, recover the original FNO architecture if the input to the 1hh layer is the output of (I - 1)th layer ie., t = (
Improved ENO (FNO++ ): The original FNO architecture suffers from vanishing gradients, which prohibits from being made deeper [TTan et al., 2021]. We overcome this limitation by introducing residual connections within each block of FNO, with each FNO block B, comprising of three FNO layers C i.e., B, = Ct, o C/,, Ct, and three convolutional layers, where L is defined in Eq. (8).
Weight- tied network (FNO-WT): This is the weight-tied architecture introduced in Definition 5, where we initialize (z) 0 for all x E Ω.
FNO-DEQ: As introduced in Definition 6, FNO-DEQ is a weight-tied network where we explicitly solve for the fixed point in the forward pass with root finding algorithm. We use Anderson acceleration [Anderson, 1965] as the root solver. For the backward pass, we use approximate implicit gradients [Geng al., 2021] which are light-weight and more stable in practice, compared implicit gradients computed by inverting Jacobian.
Note that both weight-tied networks and FNO-DEQs leverage weight-tying but the two models differ in the ultimate goal of the forward pass: DEQs explicitly solve for the fixed point during the forward pass, while weight-tied networks trained with backpropagation may may-not reach a fixed point [Anil et al., 2022]. Furthermore, DEQs require O(1) memory, as they differentiate through the fixed point implicitly, whereas weight-tied networks need explicitly create the entire computation graph for backpropagationn, which can become very large as the network depth increases. Additionally, FNO++ serves as a non weight-tied counterpart to weight-tied input-injected network. Like weight- tied networks, FNO++ does not aim to solve for a fixed point in the forward pass.
Experimental setup. We test the aforementioned network architectures on two families of steady- state PDEs: Darcy Flow equation and stcady-state Navier-Stole equation for incompressible fluids. For experiments with Darcy Flow, we the dataset provided by Li et al. [2020a], and generate our own dataset for steady-state Navier-Stokes. For more details on the datasets and the data generation processes we refer to Sections B. and B.2 of the Appendix. For each family of PDE, we train networks under 3 different training setups: clean data, noisy inputs and noisy observations. For
'We are abusing the notation somewhat and denoting by P(f) the second coordinate of P, which only depends on f.

---

experiments with noisy data, both input and observations, we add noise sampled from a sequence of standard Gaussians with increasing values of variance {0
(o M-1 where M is the total number of Gaussians we sample from. We set of = 0 and o resolution of the grid. Thus, the training data includes equal number of PDEs with different levels of Gaussian noise added to their input or observations. We add noise to training data, and always test on clean data. We follow prior work [Li et al., 2020b] and report relative L2 norm between ground truth u" and prediction on test data. The total depth of all networks besides FNO is given by 6B + 4, where B the number of FNO blocks. Each FNO block has FNO layers and convolutional layers. In addition, we include the depth due to P. Q. and an additional final FNO layer and convolutional layer. We further elaborate upon network architectures and loss functions in in Appendix A.
- S 1/r, where is the
5.1 Darey Flow
For our first set of experiments we consider stationary Darey Flow equations, a form of linear elliptic PDE with in two dimensions. The PDE is defined follows,
-V (a(x) Va(x)) = fz),
E (0. 1)
u(x) 0
E 0(00
Note that the diffusion coefficient a E _ (2) RR), i.e., the coefficients are always positive, and f (9( Rd4) ) is the forcing term. These PDEs are used to model the steady-state pressure of fluids flowing through porous media. They can also be used to model the stationary state of the diffusive process with u(x) modeling the temperature distribution through the space with defining the thermal conductivity of the medium. The task is to learn an operator Go L2; Bd)) ) x (22 (4) L (22 Rd)) such that u² Gg(u",
We report the results of our experiments on Darey Flow in Table 1. The original FNO architecture does not improve its performance with increased number of FNO blocks B. FNO++ with residual connections scales better but saturates at around FNO blocks. contrast, FNO-WT and FNO-DEQ with just a single FNO block outperform deeper non-we non-weight-tied architectures on clean data and on data with noisy inputs. When observations are noisy, FNO-WT WT and FNO-DEQ outperform FNO+ with similar number of parameters, and perform comparably to FNO++ with 4x parameters.
We also report results on shallow FNO-DEQ, FNO-WT and FNO++ architectures. An FNO block in these shallow networks has single FNO layer instead of three layers. In our experiments, shallow weight-tied networks outperform non-weight-tied architectures including FNO++ with 7x parameters on clean data and on data with noisy inputs, and perform comparably to deep FNO++ on noisy observations. In case of noisy observations, we encounter training instabilitty issues in FNO-DEQ. We believe that this shallow network lacks sufficient representation power and cannot accurately solve for the fixed point during the forward pass. These errors in fixed point estimation accumulate over time, leading to incorrect values of implicit gradients, which in turn result in training instability issues.
5.2 Steady-state Navier-Stokes Equations for Incompressible Flow
We consider the steady-state Navier-Stoke Stokes equation for an incompressible viscous fluid in the vorticity form defined on a torus, i.c., with periodic boundary condition,
u. Vw = vdw + f. x E
(10)
V uu= 0
E Ω
where Ω = (0,2 22, and u : Ω - R² is the velocity and w Ω + where w = x u,
E R. is the viscosity and f f Ω + RR is the external force term. We learn an operator Go L2( Rd x L',, Rd)) + (2( Rdd). such that u* = GG(u", f). We train all the models on data with viscosity values 0.01 and 0.001, and create dataset for steady-state incompressible Navier-Stokes, which we will make public as a community benchmark for steady-state PDE solvers.
Results for Navier-Stokes equation have been reported in Table 2 and Table 3. For both values of viscosity, FNO-DEQ outperforms other architectures for all three cases: clean data, noisy inputs and noisy observations. FNO-DEQ is more robust to noisy inputs compared to non-weight-tiee architectures. For noisy inputs, FNO-DEQ matches the test-error of noiseless case in case of viscosity 0.01 and almost matches the test-error of noiseless case in case of viscosity 0.001. We provide additional results for noise level 0.004 Appendix E. FNO-DEQ and FNO-WT consistently outperform non- woi-weight-tied architectures for higher levels of noise well.
7

---

Test error
Architecture
Parameters
#Blocks
o max
(o 0.001
(o 0.001
FNO
2.37M
0.0080 + 5e-4
0.0079 + 2c-4
0.0125 + 4e-5
FNO
4.15M
0.0105 6e-4
0.0106 = 4e-4
0.0136 2e-5
FNO
7.71M
0.2550 + 2c-8
0.2557 + 8c-9
0.2617 + 2c-9
FNO+
2.37M
0.0075 + 2e-4
0.0075 = 2e-4
0.0145 7e-4
FNO++
4.15M
2
0.0065 + 2c-4
0.0065 = 9
0.01 + 5e-5
FNO+
7.71M
0,0064 2
0.0064 2e-4
0.0109 5e-4
S-FNO+
1.78M
0.66
0.0093 + 5c-4
0.0094 = 7c-4
0.0402 + 6c-3
FNO-WT
2.37M
0.0055 + 1e-4
0.0056 = 5e-5
0.0112 1 4e-4
FNO-DEQ
2.37M
0.0055 + 1e-4
0.0056 = 7e-5
0.01 + 4c-4
S-FNO-WT
1.19M
0.33
0.0057 + 3e-5
0.0057 = 5e-5
0.0112 + le-4
S-FNO-DEQ 1.19M
0.33
0.0056 4c-5
0.0056 5c-5
0.0136 0.011
Table 1: Results on Darcy flow: clean data (Col 4).noisy inputs (Col 5) and noisy observations (Col 6) with variance of added noise being max)) and (o ) respectively. Reported test has been averaged on three different runs with seeds 0, 1, and 2. Here, S-FNO++, S-FNO-WT and S-FNO-DEQ are shallow versions of FNO++, FNO-WT and FNO-DEQ respectively.
In general, DEQ-based architectures slower train (upto ~2.5% compared to feedforward networks of similar size) as we solve for the fixed point in the forward pass. However, their inductive bias provides performance gains that cannot be achieved by simply stacking con-weight-tied FNO layers. In general, we observe diminishing returns in FNO++ beyond 4 blocks. Additionally, training the original FNO network on more than 4 FNO blocks challenging, with the network often diverging during training, and therefore we do not include these results in the paper.
Test error 4
Architecture
Parameters
#Blocks
0
(o2 (axax)* ) 0.001
max)) 0.001
FNO
2.37M
1
0.184 + 0.002
0.218 + 0.003
0.184 0.001
ENO
4.15M
2
0.162 + 0.024
0.176 = 0.004
0.152 = 0.005
FNO
7.71M
4
0.157 + 0.012
0.187 + 0.004
0.166 = 0.013
FNO+
2.37M
1
0.199 + 0.001
0.230 = 0.001
0.197 = 0.001
FNO+
4.15M
2
0.154 + 0.005
0.173 - 0.003
0.154 0.006
FNO+
7.71M
4
0.151 = 0.003
0.165 = 0.004
0.149 = 0.003
FNO-WT WT
2.37M
1
0.151 + 0.007
0.173 = 0.017
0.126 = 0.027
FNO-DEQ
2.37M
1
0.128 + 0.004
0.144 + 0.007
0.136 = 0.003
Table Results on incompressible steady-state Navier-Stokes (viscosit=00.001): clean data (Col 4),
noisy inputs (Col 5) and noisy observations (Col 6) with max variance of added noise being ? and (o 2 respectively. Reported test error has been averaged on three different runs with seeds 1, and 2.
Universal Approximation and Fast Convergence of FNO-DEQ
Though the primary contribution of our paper is empirical, we show (by fairly standard techniques) that FNO-DEQ is universal approximator, under mild conditions the steady-state PDEs. More- over, we also show that in some cases, we can hope the fixed-point solver can converge rapidly.
As noted in Definition 2, have Ω T We note that all continuous function f L2 R) and Jo (xx) dr < oc can be written as, fxx Eweza coefficients of the function f. We define as 122 (S2) as the space of functions such that for fN L2x(0) (2) with Fourier coefficients that vanish outside bounded ball. Finally, define an
fu.. where (fw) are the Fourier

---

Test error 1
Architecture
Parameters
#Blocks
max 0
(o ) 0.001
maxx)
0.001
FNO
2.37M
1
0.181 = 0.005
0.186 = 0.003
0.178 0.006
ENO
4.15M
2
0.138 + 0.007
0.150 + 0.006
0.137 = 0.012
FNO
7.71M
4
0.152 + 0.006
0.163 + 0.002
0.151 - 0.008
ENO+
2.37M
1
0.188 + 0.002
0.207 = 0.004
0.187 = 0.003
FNO+
4.15M
2
0.139 + 0.004
0.153 + 0.002
0.140 0.005
ENO+
7.71M
4
0.130 + 0.005
0.151 = 0.004
0.128 = 0.009
FNO-WT
2.37M
1
0.099 + 0.007
0.101 = 0.007
0.130 = 0.044
FNO-DEQ
2.37M
1
0.088 + 0.006
0.099 + 0.007
0.116 - 0.011
Table Results on incompressible steady-state Navier-Stokes (viscosity=0.00) clean data (Col 4),
noisy inputs (Col 5) and noisy observations (Col 6) with max variance of added noise being (o max ) and (a 2 max)" respectively. Reported test error has been averaged on three different runs with seeds
1, and 2.
orthogonal projection operator 1Ix : (S2) - (2)), such that for all f E (22) we have,
fn (N) II
Σ WE2A
Sweet"u
E
facir"
(11)
WWlloOEN
That is, the projection operator N takes an infinite dimensional function and projects it to a finite dimensional space. We prove the following universal approximation result:
Theorem 1. Let u E (S2;; Rd.) define the solution to a steady-state PDE in Definition 2, Then
there exists an operator G (2;; Rd)) x (S2; Rdr) + (22; Rd=) such that, u² G(u", f)) Furthermore, for every c > 0 there exists an N E N such that for compact sets Kn C (S2; Rd,) and K C (2; R24) there exists a neural network GB [3; Rdd) x L3 (St; Rdd)) - L3 (Ω; Rd.) with parameters 0, such thut,
|ut - GT Ny NyJ) Sc FEK,
sup
The proof for the above theorem is relatively straightforward and provided in Appendix C. The proof uses the fact that u² is a fixed-point of the operator G(u, f) the the results in Kovachki al. [2021a] that show continuous operator can be approximated by a network as defined in Equation 2. Note that the choice of G' is by no means unique: one can "universally approximate" any operator G(u, = u - A(L(u) - f), for a continuous operator A. Such G can be thought of as form of "preconditioned" gradient descent, for preconditioner A. For example, a Newton update has the form G(u, f) = u L'(u)= (L(u) f) where L' (.) + L2( Rdd) the Frechet derivative the operator L.
(I.u)= - f), allowing us to use
The reason this is relevant is that the DEQ can choose to universally approximate a fixed-point equation for which the fixed-point solver it is trained with also converges rapidly. As an example, the following classical result shows that under Lax-Miilramm-ii conditions (a kind of strong convexity condition), Newton's 's method converges doubly exponentially fast:
Lemma 1 (Faragó and Karátson [2002], Chapter 5). Consider the PDE defined Definition 2, such that such that I.' I'(u) defines the Frechet derivative of the operator for all u, e
d, = d, d
L (22 R) we have [| (1)(u))ll 2 A||p|| (2))) and [['(u) L' L'())lrr(()) S A e|l for <
S < co, then for the Newton update, U4+1 <- ut - L' T'(u,) -1 (I(w) f), with ug E T,2 RR)
there exists an O, such that ||u n² (1EP(0) S if > log (log (1) / log [ - (1122(0)
2A
For completeness, we include the proof of the above lemma in the Appendix (Section D). We note that the conditions of the above lemma are satisfied for elliptic PDES like Darcy Flow, as well as many variational non-linear elliptic PDEs (e.g., those considered in Marwah et al. [2022]). Hence, we can expect FNO-DEQS to quickly converge to the fixed point, since they employ quasi-Newton methods like Broyden and Anderson methods [Broyden, 1965, Anderson, 1965].

---

7 Conclusion
In this work, we demonstrate that the inductive bias of deep equilibrium models and reedht-tied networks in general-make. them ideal architectures for approximating neural operators for steady- state PDEs. Our experiments on steady-state Navier-Stokes equation and Darcy flow equations show that weight-tied models and FNO-DEQ perform outperform FNO models with 4x the number of parameters and depth. Our findings indicate that FNO-DEQ and weight-tied architectures are, in general, more robust to both input and observation noise compared to non non-weight-tin neightttieddde architectures, including FNO. We believe that our results complement any future progress in the design and development of PDE solvers [Tran et al., 2021, Li et al., 2022b] for steady-state PDEs, and hope that our work motivates the study of relevant inductive biases that could be used to improve them.
8 Acknowledgements
TM is supported by CMU Software Engineering Institute via Department of Defense under contract 1A8702-15-D-000 AP is supported by grant from the Bosch Center for Artificial Intelligence. ZL. gratefully acknowledges the NSF (FAI 2040929 and IIS221 1955), UPMC, Highmark Health, Abridge, Ford Research, Mozilla, the PwC Center, Amazon AI, JP Morgan Chase, the Block Center, the Center for Machine Learning and Health, and the CMU Software Engineering Institute (SEI) for their generous support of ACMI Lab's
via Department of Defense contract FA8702-15-D-D00 research. JL. is supported part by NSF award DMS-2012286, and AR is supported in part by NSF awards IIS-221 1907, CCF-2238523, Amazon Research Award, and the CMU/PwC DT&& Center.
References
Donald Anderson. Iterative procedures for nonlinear integral equations. Journal of the ACM (JACM), 1965. 4, 6, 9, 13
Cem Anil, Ashwini Pokle, Kaiqu Liang, Johannes Treutlein, Yuhuai Wu, Shaojie Bai, J Zico Kolter, and Roger B Grosse. Path independent equilibrium models can better exploit test-time computation. Advances in Neural Information Processing Systems, 35:7796-7809, 2022. 6
Shaojie Bai, J Zico Kolter, and Vladlen Koltun. Deep equilibrium models. Advances in Neural
Information Processing Systems, 32, 2019.
Yohai Bar-Sinai, Stephan Hoyer, Jason Hickey, and Michael P Brenner. Learning data-driven discretizations for partial differential equations. Proceedings of the National Academy of Sciences, 116(3))11334411559, 2019. 2,
Cx K Batchelor and George Keith Batchelor. An introduction to fluid dynamics. Cambridge university press, 1967. 14
Kaushik Bhattacharya, Bamdad Hosseini, Nikola B Kovachki, and Andrew M Stuart. Model reduction and neural networks for parametric PDEs. The SMAI journal of computational mathematics, 7: 121-157, 2021. 3.3
John P Boyd. Chebyshev and Fourier spectral methods. Courier Corporation, 2001.
Johannes Brandstetter, Daniel Worrall, and Max Welling, Message passing neural pde solvers. arXiv preprint arXiv:2202.03376 2022. 2,3
Charles G Broyden. A class of methods for solving nonlinear simultancous equations. Mathematics of computation, 1965. 4, 6, 9, 13
Tianping Chen and Hong Chen. Universal approximation to nonlinear operators by neural networks with arbitrary activation functions and its application to dynamical systems. IEEE Transactions on Neural Networks, 6(4):911-917 1995. 3,
Ziang Chen, Jianfeng Lu, and Yulong Lu. On the representation of solutions to elliptic pdes in barron spaces. Advances in neural information processing systems, 34:6454-6466 2021. 2, 3,
10

---

Gideon Dresdner, Dmitrii Kochkov, Peter Norgaard, Leonardo Zepeda-Núñez, Jamie A. Smith, Michael Brenner, and Stephan Hoyer. Learning correct spectral methods for simulating turbulent flows. 2022. doi: 10.48850/ARXII22007... 2207. 14
URL htt.htt/s///artt //tp:i/////i.////aeeaaeeeeeaa org/ abs/
István Faragó and János Karátson. Numerical solution of nonlinear elliptic problems via precondi- tioning operators: Theory and applications, volume 11. Nova Publishers, 2002. 2, 9, 19
Samy Wu Fung, Howard Heaton, Qiuwei Li, Daniel McKenzie, Stanley Osher, and Wotao Yin. Jfb: Jacobian-fres backpropagation for implicit networks. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 36, pages 6648-6633
2022. 4
Zhengyang Geng, Xin- Yu Zhang, Shaojie Bai, Yisen Wang, and Zhouchen Lin. On training implicit models. Advances in Neural Information Processing Systems, 34:24247-24260, 2021. 4, 6, 13
Zhongkai Hao, Zhengyi Wang, Hang Su, Chengyang Ying, Yinpeng Dong, Songming Liu, Ze Cheng, Jian Song. and Jun Zhu. Gnot: A general neural operator transformer for operator learning. In International Conference Machine Learning, pages 12556-12569, PMLR, 2023. 3
Dan Hendrycks and Kevin Gimpel. Gaussian error linear units (gelus). arXiv preprint
/ 2016. 6
Kurt Hornik, Maxwell Stinchcombe, and Halbert White. Multilayer feedforward networks are universal approximators. Neural networks, 2(5):359-366, 1989. 3
Jun-Ting Hsieh, Shengjia Zhao, Stephan Eismann, Lucia Mirabella, and Stefano Ermon, Learning neural pde solvers with convergence guarantees. arXiv preprint arXiv: 1906.01200, 2019. 2, 3
Dmitrii Kochkov, Jamie A Smith, Ayya Alieva, Qing Wang. Michael P Brenner, and Stephan Hoyer. Machine saamingg-aceelene computational fluid dynamics. Proceedings of the National Academy of Sciences, 118(11):22011884118 18, 2021. 2, 3, 14
David A Kopriva. Implementing spectral methods for partial differential equations: Algorithms for scientists and engineers. Springer Science & Business Media, 2009.
Nikola Kovachki, Samuel Lanthaler, and Siddhartha Mishra. On universal approximation and error bounds for Fourier neural operators. The Journal of Machine Learning Research, 22(1):13237- 13312, 2021a. 2, 3, 9, 15, 18
Nikola Kovachki, Zongyi Li, Burigede Liu, Kamyar Azizzadenesheli, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Neural operator: Learning maps between function spaces. arXiv preprint ar arXiv:2108. 1784412008.0848,, 2021b. 3
Nikola Kovachki, Zongyi Li. Burigede Liu, Kamyar Azizzadenesheli, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Neural operator: Learning maps between function spaces with applications PDEs. Journal of Machine Learning Research, 24(89):1-97, 2023.
Samuel Lanthaler, Siddhartha Mishra, and George E Karniadakis. Error estimates for deeponets: deep learning framework in infinite dimensions. Transactions of Mathematics and Its Applications, 6(1):tnac001, 2022. 3
Randall LeVeque. Finite difference methods for ordinary and partial differential equations: steady- state and time-dependent problems. SIAM, 2007. 1.2
Zijie Li, Kazem Meidani, and Amir Barati Farimani. Transformer for partial differential equations' operator learning. arXiv preprint arXiv:2205. 13671, 2022a.
Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Fourier neural operator for parametric partial differential equations. arXiv preprint arXiv:2010.08893 2020a. 2, 3, 4, 6, 13
Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Neural operator: Graph kernel network for partial differential equations. arxiv preprint arXiv:2003.03483, 2020b. 1,

---

Zongyi Li, Daniel Zhengyu Huang, Burigede Liu, and Anima Anandkumar. Fourier neural operator with learned deformations for pdes on general geometries. arXiv preprint wrxiv:200.005209 2022b. 10
Renjie Liao, Yuwen Xiong, Ethan Fetaya, Lisa Zhang, KiJung Yoon, Xaq Pitkow, Raquel Urtasun, and Richard Zemel. Reviving and improving recurrett back prop-aappgattion In International Conference on Machine Learning, pages 1082-3091. PMLR, 2018. 4
Xinliang Liu, Bo Xu, and Lei Zhang. Mitigating spectral bias for the multiscale operator learning with hierarchical attention. October 2022. URL http: /ttp:///aavvvvvaaeeaeeeeeaia berexiv.or//aee aa/aabb/222211 10890. 3
Lu Lu, Pengzhan Jin, and George Em Karniadakis. Deeponet: Learning nonlinear operators for identifying differential equations based on the universal approximation theorem of operators. arXiv preprint 1eXX:::9100031999 2019. 1,3
Tanya Marwah, Zachary Lipton, and Andrej Risteski. Parametric complexity bounds for approx- imating PDEs with neural networks. Advances in Neural Information Processing Systems, 34: 15044-15055, 2021. 2, 3.5
Tanya Marwah, Zachary C Lipton, Jianfeng Lu, and Andrej Risteski. Neural network approximations of PDEs beyond linearity: Representational perspective. arXiv preprint arXiv:2210.12101, 2022.
2, 9,9
Fadl Moukalled, Luca Mangani, Marwan Darwish, F Moukalled, Mangani, and M Darwish. The finite volume method. Springer, 2016.
Suhas V Patankar and D Brian Spalding. A calculation procedure for heat, mass and momentum transfer in dhree-dimensional parabolic flows. In Numerical prediction of flow, heat transfer. furbulence and combustion, pages 54-73. Elsevier, 1983. 13
Ravi G Patel, Nathaniel A Trask, Mitchell A Wood, and Eric C Cyr. A physics-informed operator regression framework for extracting data-driven continuum models. Computer Methods Applied Mechanics and Engineering. 373:1 13500, 2021.
Max Revay, Ruigang Wang, and Ian R Manchester. Lipschitz bounded equilibrium networks. arXiv preprint arXiv:2010.017332, 2020.
Alasdair Tran, Alexander Mathews, Lexing Xie, and Cheng Soon Ong. Factorized Fourier neural operators. arXiv preprint arXiv:2111.13802 2021. 10
Ezra Winston and J Zico Kolter. Monotone operator equilibrium networks. Advances neural information processing systems, 13:10718-1072 2020. 4
12

---

Appendix
A Implementation Details
Training details. We train all the networks for 500 epochs with Adam optimizer. The learning rate set to 0.001 for Darcy flow and 0.005 for Navier-Stokes. We use learning rate weight decay of lc-4 for both Navier-Ster Stokes and Darey flow. The batch size is set to 32. In case of Darcy flow, we also use cosine annealing for learning rate scheduling. We run all our experiments NVIDIA RTX A6000, NVIDIA GeForce RTX 2080 Tř and 3080 Ti. All networks can easily fit on a single NVIDIA RTX A6000, but training time varies between the networks.
a combination of
For FNO-DEQ, we use Anderson solver [Anderson, 1965] to solve for the fixed point the forward pass. The maximum number of Anderson solver steps is kept fixed at 32 for Dary Flow, and 16 for Navier Stokes. For the backward pass, we use phantom gradients [Geng al., 2021] which are computed as:
u* rGo("", a) + (1 -)
(12)
where is tunable damping factor and u² is the fixed point computed using Anderson solver in the forward pass. This step can be repeated times. We use T = 0.5 and S for Darey Flow, and 0.8 and S 3 for Navier-Stokes.
For the S-FNO-DEQ used in Table 1, we use Broyden's method [Broyden, 1965] to solve for the fixed point in the forward pass and use exact implicit gradients, computed through implicit function theorem as shown in Eq. (6), for the backward pass through DEQ. The maximum number of solver steps is fixed at 32.
For weight-tied networks, we repeatedly apply the FNO block to the input 12 times for Darcy flow, and 6 times for Navier-Stokes.
Network architecture details. The width of an FNO layer set to 32 across all the networks. Additionally, we retain only 12 Fourier modes FNO layer, and truncate higher Fourier modes. We use the code provided by Li et al. [2020a] to replicate the results for FNO, and construct rest of the networks on top of this as described in Sec.
B Datasets
B.1 Darcy Flow
As mentioned in Sec. 5 we use the dataset provided by Li et al. [2020a] for our experiments with steady-state Darcy-Flow.
All the models are trained on 1024 data samples and tested on 500 samples. The resolution of original images is 421 x 421 which downsample to 85 x 85 for our experiments. For experiments with noisy inputs/observations, the variance of Gaussian noise that we add to PDEs are [0, 1e-9, 1e-8, 1e-7, le-6, le-5, le-4, 1e-3].
B.2 Steady-State Incompressible Fluid Navier-Stoke
Vw vAw
s E Ω
V u
x E Ω
To generate the dataset for steady-state Navier-Stokes, instead of solving the steady state PDE using steady-state solvers like the SIMPLE algorithm Patankar and Spalding [1983], we first choose the solution -n = V x u² of the PDE and then generate the corresponding equation, i.e. calculate the corresponding force term f = V ='w".
To generate the solutions w* we forward propagate relatively simple initial distribution of wo (sampled from a Gaussian random field) through a time-dependent lavier-Stokes equation in the vorticity form for short period of time. This ensures our dataset contains solutions w" that are rich
13

---

a
u*
Figure Samples from Darcy Flow
and complex. Precisely, recall the Navier-Stokes equations in their vorticity form:
Ow((z,) t) + u(x, t) Vw(z, t) nAw (z,t) g(x)
(0,2 t [0,
V T.(wx, f)
(0, 2n 2 [0,
(13)
(xx, 0) = wa( x E (0, 2nr)2
where g(x) V (z)) and ğ(z) sin(5x1) Jal is divergence free foreing term and æ (I x2)
are the two coordinates of the input vector. We forward propagate the equations Equation 13 using a pseudo-spectral method using the functions provided in JAX-CFD [Kochkov al., 202 1,
Dresdner et al., 2022] package. The initial vorticity wa is sampled from a Gaussian random field (0, (5 + 25A) 25.))-25)),, which is then made divergence free. We forward propagate the Navier- Stokes equation in Equation 13 for time T 0.5 with dt 0.002 to get w (11 æ), which we choose as the solution to the steady-state PDE in Equation 10, i.c. w for Equation 10.
the stream function V [Batchelor and Batchelor, 1967] to calculate (ON /0x,, 0 /0r/)r))) by solving the Poisson equation A = w in the Fourier domain. Furthermore, since f u² Vw" vAw" use the stream function calculate (f1, f2), i.e., the different components of the force term.
Subsequently,
We use 4500 training samples and 500 testing samples. The input to the network is the vector field f (fi, and we learn map that outputs the vorticity wt. The resolution of grid used to
14

---

generate the dataset is 256 x 256 which we downsample to 128 x 128 while training the models. For experiments with noisy inputs/obserratiion,, we consider values of maximum variance of Gaussian noise: 1e-3 and 4e-3. The variances of the Gaussian noise that we add to the PDEs for the latter case [0, 1e-9, le-8, le-7, le-6, le-5, le-4, le-3, 2e-3, 4e-3]. However, when conducting experiments with variance of le-3, we exclude the last two values of variance from this list.
f1
f2
w
Figure Samples from Steady-state Navier-Stokes dataset with viscosity 0.001. Each triplet visualizes the inputs f1, fa and the ground truth output i.e. w".
C Proof of Universal Approximation
The proof of the universal approximation essentially follows from the result on the universal approxi- mation capabilities of FNO layers in Kovachki et al. [2021a], applied to ((v,)) = v - (Lu f). For the sake of completeness, we reitarate the key steps.
For simplicity, we will assume that d, = d, = dg = 1. (The results straightforwardly generalize.) We will first establish some key technical lemmas and introduce some notation and definitions useful for the proof for Theorem 1.
Definition 7. An operator T R)) + L2 R is continuous at u L2( R. if for every 0, there exists [[[(x) L(v))l<P(O) S.
O, such that for all E (2( with ||u =ll 6, we have
15

---

f,
f2
w
Figure Samples from Steady-state Navier-Stokes dataset with viscosity 0.01. Each triplet visualizes the inputs fi, f2 and the ground truth output i.e. w".
First, we approximate the infinite dimensional operator g (.2 x T22 - (2) by projecting the functions in (2( to miite-dimensional approximation L3 (St), and considering the action of the operator on this subspace. The linear projection we use is the one introduced in Equation 11. More precisely show the following result,
Lemma 2. Given continuous operator L (2) + L2(Ω) as defined in Equation 1, let us (S)) as g(v u- (L(u) f). Then, for every > 0 there exists an N E N such that for all v. f in any compact set K C L?()), the operator
define an operator (2) L2)
IN = GIT fT f) is an e-approximation of G(v, f), Le,, we have,
sup [|g(e, f) GN(u, f)l.f Sc. v.feK
Proof. Note that for an é there exists an N N(c, d) such that for all v E K we have
sup ||0 TxellzA() Sc.
16

---

Therefore, using the definition of IN we can bound the (92 norm of the difference between G and IN as follows,
[ f) - [ f2)))
Hy x (g( - (v Ml + [I N( Nx (l)
$ [9(e, f) - 1Ivg(v, (aN((o,f)lzagg 212 + [[,Ng((, [1 [Iv
We first bound the term / as follows:
[g(u, f) 11 flle
||e (I(u)
II (v (I(o) M
= 10y-1lye -1 +
+ [L(v) - (x)
=c+c+ [L()) - - I())llz))))))
(14)
Since L. continuous, for compact sets K C (.2)),, I(K) compact as well. This is because: (1) for any E K, (())lll))))) is finite; (2) for any K, [I [|[())] + Cllu )())) Therefore, for every e > O, there exists an N E N such that
sup REK
[L()) - nI M 6.
Substituting the above result in Equation 14, have
[| f) [v ( S 3c.
(15)
Similarly, for all E K where K is compact,
bound Term IT following,
[g f) [I L2(0)
< ||0 (L(v) f - 1I (L(IINU) - -x/)) [[(v) ([Ixv)))L(())))
< - IIvll
|lʃ /1
Setct -(Ix)))1100)
(16)
Now, since K and L ( - (12) is continuous operator, there exists modulus of continuity (an increasing real valued function) a E [0, oo), such that for all E K, we have
[ L(IIN)) $ a (Ilv - lxe]l<>(0))
Hence for every >> 0 there exists N E N such that,
alI l[velle(()) Se.
Plugging these bounds in Equation 16, we get,
[G(v.)) - G(INv,, 11 IN S 3e.
(17)
Therefore, combining Equation 15 and Equation 17 then for e > 0, there exists an E N, such that for all f E K we have
|6(v, (on $ 6e. e.fEK
sup
(18)
Taking é' = 6e proves the claim.
Proof of Theorem I. For Lemma we know that there exists finite dimensional projection for the operator G, defined as (u, f) such that for all v, E ( we have
[1 f) GN f)l Se.
Now using the definition of GN (v, f) we have
(v( INg IN.))
HI LII Ix
17

---

From Kovachki et al. [2021a], Theorem 2.4 we know that there exists an FNO network Gg. of the form defined in Equation 2 such that for all E K, where is compact set, there exists e" we have
sup vek
|[ L(Ixv) Gyz (le2 st
(19)
Finally, note that from Lemma D.1 in Kovachki al. [2021a], have that for any E K, there exists an FNO layers Gøs E (S2) and Gg- E (S2) defined in Equation such that
[ Go (1)) c vek
sup
(20)
and
sup fek
[ - Gør Gorll12(2)) se
(21)
for e > 0 and e >
Therefore there exists an é > such that there is an FNO network Ge L2((2) x (2() - L?
where 0 == { 8", of) such that
sup (E
[] Glly,,l S
(22)
Now, since we know that w" is the fixed point of the operator have from Lemma and Equation 22,
[Q(u",f] - Allaa I - (llaa)))) + 11 f)(u',
s8+c
D Fast Convergence for Newton Method
Definition 8 (Frechet Derivative in (2))).). For a continuous operator F -.(( + (2)), the
Frechet derivative E (22 is linear operator ('))) (.2(2) 1,2 such that for all w E (22 we have
|[ u)) F(u) F'u)(w))) 00
lim Moll 2200
|
Lemma Given the operator L. 12((2) + 1,2(2) with Frechet derivative I/, such that for all u (Q)), we have ([((u)(v))l<>())) Al[v|l22(0), then L'(u)- exists and we have, for all
U1, 02 E (Q):
1. [L((w) (1 & |
2. |0 v2 v2llz & ((m,) 1(v2)l
Proof. Note that for all u, v E L2( we have,
[I'(w)D'Ilzz<a) 2 Al|e'llz(())
Taking (v') have
[Lu(() (L'(w)- || (w(e)l))) ||ll 2 (u())()())))
For part 2, note that there exists a c E [0, 1] such that
[[v)) ((22))lca(a) 2 inf [L'(evi (1 c)o2)|la|1r v2ll<z(0) Alli 02ll<> cE(0.1)
18

---

We now show the proof for Lemma 4. The proof is standard and can be found in Faragó and Karátson [2002], however we include the complete proof here for the sake of completeness.
We restate the Lemma here for the convenience of the reader.
Lemma 4 (Faragó and Karátson [2002], Chapter 5). Consider the PDE defined Definition 2, such that du = de dj = 1. such that (u) defines the Frechet derivative of the operator L. ff.
all u, V E (2;; RR) we have (u)(l 2 X| 7 and [[T - (1) (()))) S Al A|1-
(|())10)) for 0 < X SA < oc, then for the Newton update, Ur41 4- U2 (wv) - (L(u) - -)),
with uo RS2; there exists an > 0, such that ||ur - u² =²1l)=(0) S if T z
21 log (log (1) / log (IL
Proof of Lemma 4. Re-writing the updates in Lemma 4 as,
Up+1 = Wt + Pt L'(ue)pt = (L((ue) f)
(23)
(24)
Now, upper bounding ((u++1) f for all x E we have, [(u+1(())) f(x)
+ (((+)())) =I(u(<x)) (xx) + L/ ((((a())+(((+(()))) (x)))))))+++(( uc(x))
-(ur(x)) - f(x) + L'(u(()))p + S (L'(u(()) + - - L' ('(u(())))
Pe
= -f, (I'(u(()) + (me+1 (x)x==()))))))))))))))))) (((x(())))
=
where we use Equation 24 in the final step.
Taking (2())) norm on both sides and using the fact that [I L'(v)ll (3(0) S a =|| we have
[[(u++1) fll $ L' Aelluc++ u.
Noting that for all Ω we have u4+1 ur pt, and using the fact that for all u, [[ L2(0)) & (0))) we have, || ([(((w)))l s
[[((u+1) f1) ~ /
$ 1/22|pp [1 (12:200
S A/21| L' L'w)) (L(ue) - f) 122
A
$ 2x [[ f)1
where we use the result from Lemma 3 in the last step.
Therefore we have
2'-1 (T(ua) - -1²²
[[(( 01(22(0) v
2X²
21-1
[[((u=1) 1f0=ffl1e f((c2(0) v
(L(wo) - L(u²))²²
2X
2'
A
|lu41 n" (12(0) <
[[L(uo) L( ((u'))(2
2Xa
Wee that this condition different from the condition the inner-product in the submitted version of the paper. which had. (L'(u), o) 2 X|j
We note that this rate is different from the one in the submitted version of the paper.
19

---

Therefore, if
A 21a
[[(uo) l(u")ll $
then have
|1e+1 10041-001((000 Se,
for
T 2 log log (') /log (AlL((uo) fllm2())) )
2X2
E
Additional experimental results
We provide additional results for Navier-Stokes equation for noisy inputs and observations in Table and Table 5. For these experiments, the maximum variance of Gaussian noise added to inputs and observations is 0.004. We observe that weight-tied FNO and FNO-DEQ outperform weni-wett-ttedd architectures.
Test error
Architecture
Parameters
#Blocks
max
(o 0.004
(o 0.001
FNO
2.37M
1
0.184 = 0.002
0.238 = 0.008
0.179 = 0.004
FNO
4.15M
2
0.162 0.024
0.196 = 0.011
0.151 0.010
FNO
7.71M
4
0.157 + 0.012
0.216 = 0.002
0.158 + 0.009
FNO+
2.37M
1
0.199 + 0.001
0.255 0.002
0.197 0.004
FNO++
4.15M
2
0.154 1 0.005
0.188 0.006
0.157 0.006
FNO++
7.71M
4
0.151 1 0.003
0.184 1 0.008
0.147 1 0.004
ENO-WT
2.37M
1
0.151 0.007
0.183 0.026
0.129 0.018
ENO-DEQ
2.37M
1
0.128 = 0.004
0.159 = 0.005
0.121 = 0.015
Table 4: Results on incompressible Steady-State Navier-Stokes (viscosity=0.001): clean data
(Col 4), noisy inputs (Col 5) and noisy observations (Col 6) with max variance of added noise being (a )' and ( max respectively. Reported test error has been averaged on three different runs with seeds 0, 1, and 2. $ indicates that the network diverges during training for one of the seeds.
Test error
Architecture
Parameters
#Blocks
0 max
(o2 (o
0.004
(o 0.004
FNO
2.37M
1
0.181
0.005
0.207 0.003
0.178
0.008
FNO
4.15M
2
0.138 0.007
0.163 0.003
0.137 0.006
FNO
7.71M
4
0.152 + 0.006
0.203 + 0.055
0.151 + 0.008
FNO+
2.37M
1
0.188
0.002
0.217 0.001
0.187 0.005
FNO+
4.15M
2
0.139 + 0.004
0.170 + 0.005
0.138 + 0.005
FNO+
7.71M
4
0.130 + 0.005
0.168 + 0.007
0.126 + 0.007
FNO-WT
2.37M
1
0.099
0.007
0.159
0.029
0.123 0.023
FNO-DEQ
2.37M
1
0.088 1 0.006
0.104 1 0.001
0.116 + 0.005
Table 5: Results on incompressible Steady-State Navier-Stokes (viscosity=0.01): clean data (Col
4), noisy inputs (Col 5) and noisy observations (Col 6) with max variance of added noise being (o and ( respectively. Reported test error has been averaged three different runs with seeds 0, 1, and 2. $ indicates that the network diverges during training for one of the seeds.
20

---

Convergence analysis of fixed point. We report variations in test error, absolute residual (Ga(z,) Zell2. and relative residual [[Ga (G,()))-zi) with an increase in the number of solver steps while solving for the fixed point FNO-DEQ, for both Darcy Flow (See Table 6) and Steady-State Navier Stokes (See Table 7). We observe that all these values decrease with increase in the number of fixed point solver iterations and eventually saturate once we have a reasonable estimate of the fixed point. We observe that increasing the number of fixed point solver iterations results in better estimation of the fixed point. For steady state PDEs, we expect the test error to reduce as the estimation of the fixed point improves. Furthermore, at inference time we observe that the test error improves (i.e. reduces) with increase in the number fixed point solver iterations even though the FNO-DEQ is trained with fewer solver steps. For Navier-Stoke with viscosity 0.01, inference time we get test MSE loss of 0.0744 with 48 solver steps from 0.0847 when used with 24 solver steps.
2
This further bolsters the benefits of DEQs (and weight-tied architectures in general) for training neural operators for steady-state PDEs. Moreover, performance saturates after certain point once we have a reasonable estimate of the fixed point, hence showing that more solver steps stabilize to the same solution.
Solver steps
Absolute residual +
Relative residual +
Test Error
212.86
0.8533
0,0777
18.
0.0878
0.0269
0.3530
0.00166
0.00567
16
0.00239
1.13e-5
0.00566
32
0.000234
1.1e-6
0,00566
Table 6: Convergence analysis of fixed point for noiseless Darcy Flow: The test error, absolute residual ||Go z and relative residual [[Ca(z,))zz decrease with increase in the number of fixed point [x
solver iterations. The performance saturates after a certain point once we have a reasonable estimate of the fixed point. We consider the noiscless case, where do not add any noise to inputs targets.
Solver steps
Absolute residual +
Relative residual +.
Test Error
544.16
0.542
0.926
397.75
0,408
0,515
16
150.33
0.157
0.147
37.671
0.0396
0.0847
48
5.625
0.0059
0.0744
64
3.3
0.0034
0.0746
Table 7: Convergence analysis of fixed point for noiseless incompressible Steady-State Navier-
Stokes with viscosity=0.01: The test error, absolute residual |GGg (zc) Zell2 and relative residual
[G((x:)=:ll decrease with increase in the number of fixed point sol ver iterations. The performance 2
saturates after certain point once have reasonable estimate of the fixed point. We consider the noiseless case, where we do not add any noise to inputs or targets.
21

---

10a
FNO-DEQ #tlocks
FNO-WT #Blocks
FNO+
PNO #Blecks
FND eBlocks
MSE
PNO #Blocks
10-1
Epochs
(a) Training Loss Curve
10a
FNO+ #Blocks
PNO #Blocks 4
FND eBlocks
PND #Blocks
FNO dBlocks
Test
10-1
100
200
300
500
Epochs
(b) Test Loss Curve
Figure 4: Training and Test Loss Curves for Steady-State Navier- Stokes with viscosity 0.01. The axis the number of epochs and y axis is the MSE loss in log scale. Note that while all the models converge to approximately the same MSE loss value while training, DEQs and weight-tied networks get a better test loss in fewer epochs.
22


Reviews:
[{"review_id": "H5zUCxPqfc", "reviewer": "Reviewer_HGbF", "summary": "The paper tackles the problem of solving steady-state PDEs with weight-tying FNOs. The authors argue that instead of stack multiple FNO layers with different parameters, repeatedly applying one FNO layer computation is a better choice. This hypothesis is motivated by the fact that steady-state PDEs are solved for fixed-point solutions, where evolving PDE further will not change the solution. Moreover, instead of directly stack the same FNO layer and hope the fix point can be reached, the authors propose to use more advanced fixed-point solving method (FNO-DEQ). By using the fix-point properties, DEQ can have smaller training memory usage but is slower. Experiments are conducted on Darcy flow and naiver-stokes equations.", "questions": "- How is the speed of computing the implicit gradient?\n- In Definition 3, shouldn't all functions map from $\\Omega$?", "limitations": "None", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper is well structured,\n- The proposed method leverages well the fact that the solutions are fixed points.\n- Experimental results look good. ", "weaknesses": "- There are very few discussions on the convergence of the fixed point solution. For example, if we apply the FNO layers more times than during training, would the solution remain the same?\n- There are two major differences between FNO-WT and FNO-DEQ. The first is in forward pass FNO-WT directly applies the network multiple times while FNO-DEQ use fixed point solver. The second is in the backward pass FNO-WT directly propagates through the computation graph while FNO-DEQ use implicit gradients. These two components seem independent to each other. For example, since FNO-WT in some sense also solve for the fixed point solution, can we use the implicit gradient to train FNO-WT? ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their positive feedback and response! We are glad that the reviewer finds that our paper is well structured and likes the empirical results. We address some of the concerns raised by the reviewer below: \n\n**Re: A discussion on the convergence of the fixed point.** \nIncreasing the number of fixed point solver iterations results in a better estimation of the fixed point. For steady state PDEs, we expect the test error to reduce as the estimation of the fixed point improves. Furthermore, at inference time we observe that the test error improves (i.e. reduces) with increase in the number of fixed point solver iterations even though the DEQ is trained with fewer solver steps.\n\nWe report these empirical results for Darcy flow and steady state Navier Stokes (viscosity 0.01) in Tables 1 and 2 in the attached PDF. For example, for Navier-Stokes with viscosity 0.01, at inference time we get a test MSE loss of 0.0744 with 48 solver steps from 0.0847 when used with 24 solver steps. \n\nThis further bolsters the benefits of DEQs (and weight-tied architectures in general) for training neural operators for steady-state PDEs. Moreover, performance saturates after a certain point once we have a reasonable estimate of the fixed point, hence showing that more solver steps stabilize to the same solution. \n\nWhile we only show convergence graphs for the noiseless experiments due to space constraints, we will include similar results for all the PDEs as well as experiments with noise in the final version of the paper. \n\n**Re: Difference between FNO-WT and FNO-DEQ.**\nBoth FNO-WT and FNO-DEQ leverage weight-tying (i.e., applying the same transformation at each layer) as the fundamental architectural choice. Where they differ is at the final aim of the forward pass: while FNO-WT will have a *fixed depth* computation (that may or may not approach a fixed point), FNO-DEQ explicitly is trained to find/tends to a fixed point. Therefore, since FNO-WT might fail to converge to a fixed point, implicit gradients cannot be used for FNO-WT. \n\nWe refer the reviewer to Bai et al. (2019) for further details on the difference between equilibrium models and traditional weight-tied models that repeat a transformation for a fixed depth.\n\n**Re: Speed of computing implicit gradients**\nWe use phantom gradients [2] to compute approximate implicit gradients. These gradients are very light-weight, and fast to compute as they require backward pass through a single FNO block only a couple of times (1 time for Darcy Flow, and 3 times for Navier Stokes). More details of phantom gradients are included in Line 455-459 in the Supplementary Material.\n\n**Re: Clarification on the domains of functions.**\nWe note that since the projection operator $\\mathcal{P}$ is a function from $\\Omega$ to $\\mathbb{R}^d$, the domain for all the other functions in the composition would be $\\mathbb{R}^d$. We will ensure that this is more clear in the final version of the paper.  \n\n\n\n[1] Bai, Shaojie, J. Zico Kolter, and Vladlen Koltun. \"Deep equilibrium models.\" Advances in Neural Information Processing Systems 32 (2019).\n\n[2] Geng, Zhengyang, et al. \"On training implicit models.\" Advances in Neural Information Processing Systems 34 (2021): 24247-24260.\n\n\n"}, {"review_id": "aI7y3JotIU", "reviewer": "Reviewer_wRdk", "summary": "This research examined the solution of steady-state Partial Differential Equations (PDEs) using Fourier Neural Operator (FNO) based architecture. The author introduced a fix-point iteration mechanism into the FNO framework, leading to the proposal of weight-tied FNO and FNO Deep Equilibrium (FNO-DEQ) models. Comparative analyses revealed that these newly proposed architectures outperformed the traditional FNO when solving standard Darcy Flow and Navier-Stokes Equations.", "questions": "The paper doesn't provide a clear explanation regarding the practical implementation of the Fourier Neural Operator Deep Equilibrium (FNO-DEQ) model. The inclusion of a brief paragraph or section detailing its implementation would greatly enhance the reader's understanding and potentially facilitate the model's broader adoption.", "limitations": "The author focused only on the standard Darcy flow and Navier Stokes problem, which the classical FNO already works well. An immediete question is that if the FNO-DEQ can be applied to more challenging PDEs (for example 3D Navier Stokes), to show its effectiveness, where classical FNO struggles. ", "rating": 9, "confidence": 5, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "This paper is well written, providing clear and rigorous mathematical definitions of the problem, its underlying theory, and the proposed solution.\n\nThe innovation of incorporating the fixed-point iteration mechanism (via the contraction mapping theorem) into the Fourier Neural Operator (FNO) is novel and captivating. The results clearly demonstrate a significant improvement when applying this technique to steady-state Partial Differential Equations (PDEs).\n\nFurther enhancing the strength of the paper, the author proves a universal approximation theorem for the FNO Deep Equilibrium (FNO-DEQ) model, which assures the boundedness of the approximation. This theoretical validation lends additional credibility and rigor to the findings.\n", "weaknesses": "The paper doesn't provide any loss versus training epochs data, which would allow us to assess whether the training had indeed reached convergence (as well as to compare the speed of convergence). While it's assumed that convergence must have been achieved for the results shown in Tables 2 and 3, the absence of this specific data inhibits a more comprehensive understanding of the model's training process.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for the encouraging review and feedback. We are glad that the reviewer finds our paper as novel and captivating!!  Please find our replies to some of your comments below. We promise to make the corresponding changes to the camera-ready version of our draft to incorporate your suggestions.\n\n**Re: Loss vs training epochs data and convergence.**\nThank you for your feedback! We agree that adding convergence plots will help towards a more comprehensive understanding of our methodology. We have added the convergence plots for Naiver Stokes with viscosity 0.01 in the attached PDF (Figure 1). The key observation is that while all the models converge to approximately the same train MSE value, the convergence of the test loss differs, i.e., DEQs and weight tied networks get a better test loss in fewer epochs. \nWe also note that the convergence plots for Darcy Flow and Navier-Stokes with viscosity 0.001 follow similar trends. We will add the convergence plots for all the PDEs and models in the final version of the paper. \n\n**Re: Practical implementation and clarification.**\nWe provide a detailed description of the implementation and architectural details in the Appendix (Section A) of the supplementary material. For example, we note that all the networks for 500 epochs with Adam optimizer with L2 weight penalty coefficient of 1e-4. The learning rate is set to 0.001 for Darcy flow and 0.005 for Navier-Stokes, with batch size 32. Further information about the training details of the DEQ architecture are provided in the Appendix (Section A) as well. We will restructure the draft to include the key implementation details earlier in the paper. \n\nFinally, we will release our code (submitted as supplementary material) detailing the implementation details and appropriate documentation on how to reproduce our results along with the camera-ready version of the paper.\n\n**Re: Standard Darcy Flow and Navier-Stokes problem where FNO already works well.**\nWe agree that applying FNO-DEQ and FNO-WT to more challenging PDEs is an important step towards showcasing the benefits of the proposed architectures, and we believe scaling our methodologies to more complex PDEs is a fertile ground for future work. However, we also would like to note that ours is the first paper to benchmark the performance of FNO (and FNO++) along with weight-tied models on steady-state 2D Navier-Stokes PDEs! \n"}, {"review_id": "YeG8oJCZVN", "reviewer": "Reviewer_Kt2x", "summary": "The paper proposes two weight-tied neural network architectures for solving steady-state partial differential equations (PDEs) using the universal approximation capabilities of neural networks. The first architecture is a weight-tied version of Fourier Neural Operators (FNO), while the second architecture is a Deep Equilibrium Model (DEQ) that uses black-box root finding algorithms to implicitly train the model. The paper shows that both architectures outperform existing methods on benchmark problems and can be used to learn efficient solvers for PDEs. The contributions of the paper include the introduction of the weight-tied FNO and FNO-DEQ architectures, as well as the demonstration of their effectiveness in solving steady-state PDEs.", "questions": "Refer to the weaknesses section.", "limitations": "Refer to the weaknesses section.", "rating": 5, "confidence": 1, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "The main strengths of the paper are:\n1. The paper proposes a new architecture called FNO-DEQ that uses weight-tied neural network layers to solve steady-state partial differential equations (PDEs). The architecture is based on the observation that the solution of most steady-state PDEs can be expressed as a fixed point of a non-linear operator.\n2. The paper shows that FNO-DEQ outperforms other non-weight-tied architectures with 4x the number of parameters in predicting the solution to steady-state PDEs such as Darcy Flow and steady-state incompressible Navier-Stokes.\n3. The paper demonstrates that FNO-DEQ and weight-tied architectures are more robust to both input and observation noise compared to non-weight-tied architectures, including FNO.\n4. The paper leverages the universal approximation results of FNO to show that FNO-DEQ can universally approximate the solution operator for a wide variety of steady-state PDE families.", "weaknesses": "Some potential limitations of the paper are:\n1. The proposed FNO-DEQ architecture may not be applicable to all types of PDEs, as it is specifically designed for steady-state PDEs. Further research is needed to explore the effectiveness of weight-tied architectures for other types of PDEs.\n2. The paper focuses on the empirical performance of the proposed approach and does not provide a detailed theoretical analysis of why weight-tying is effective for steady-state PDEs. A more rigorous theoretical analysis could provide deeper insights into the underlying mechanisms of the proposed approach.\n3. The paper does not compare the proposed approach to other state-of-the-art methods for solving steady-state PDEs, such as finite element methods or spectral methods. A more comprehensive comparison could provide a better understanding of the relative strengths and weaknesses of different approaches.\n4. The proposed approach may require more computational resources than other methods, as it involves solving for the fixed point of an implicit operator layer. This could limit its scalability to larger or more complex PDEs.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for your positive review! Please find our responses to the concerns you raised below.\n\n**Re: Proposed architectures not applicable to all PDEs.**\nThere are a multitude of different forms of PDEs, each with their unique characteristics that may or may not be efficiently modeled by a single neural network architecture. In our paper we show that for a *broad class of PDEs*, namely *steady state PDEs*, there are specific architectural choices that can outperform the baselines with much fewer parameters. We fully agree that charting the space of effective architectures for various PDE families is a major outstanding question in using machine learning methods for PDE solvers. \n\n**Re: Lack of theoretical benefits of why weight-tying is beneficial for steady-state PDEs.**\nWe agree that more theoretical understanding of the benefits of different kinds of architectures is an exciting avenue for further research, and we hope that there is a concerted effort towards studying and designing architectures for PDEs by the machine learning community! \nHowever, we would like to note that in general showing separations theoretically between the performance (statistical or algorithmic) of different architectures is incredibly challenging to show theoretically—e.g. sample complexity separations between convnets and fully connected networks are very poorly understood, and most of the results are under very strong assumptions [4]. \n\nThat being said, as we mention in our paper, there indeed are some theoretical motivations towards using weight-tied architectures for steady-state PDEs as established in the previous works of [1,2,3] which we use as a motivation for our work.\n\n**Re: Comparison with finite element method and spectral methods.**\nOur paper builds on neural operators (in particular Fourier Neural Operators (FNO) [5]), which have many benefits over fixed algorithms like finite element methods or spectral methods. Some of the benefits are: \n- **Computational efficiency**: Neural operators learn a solution operator for an entire family of PDE, which has an inference time of just 0.005s compared to the 2.2s of spectral methods when used to solve Navier-Stokes equations [5]. Similar gains can be observed for many families of PDEs. \n- **Robustness to discretization**: Neural operators (especially FNO) are robust to changes in discretization at inference time, that is they tend to perform well even if some of the training data is very coarsely discretized.\n\nWe refer the reviewer to [5] and related works for a more comprehensive overview of the benefits of neural operators.\n\n**Re: Scaling to more complex PDEs.**\nIt is true that the computational overhead would increase with more complex PDEs, however that is true for *all* neural network based PDE solvers and is not limited to FNO-DEQ or FNO-WT. In fact, we believe weight-tied and DEQ based architectures are better suited for more complex steady-state PDEs since the memory used by DEQs for the backward pass is going to be constant, i.e., O(1),  but for deeper architectures in FNO and FNO++,  the memory required would increase with the number of layers in the architecture. \n\nHowever, as mentioned in our paper, training DEQs can be slow as we solve for a fixed point in the forward pass. Through use of approximate implicit gradients and with careful selection of hyperparameters, the training can be made faster. Further, when compared to a non-weight-tied network of a similar depth, the overhead due to fixed point solving is marginal, both in terms of compute and memory requirements.\n\n[1] Marwah, Tanya, et al. \"Neural Network Approximations of PDEs Beyond Linearity: A Representational Perspective.\" International Conference on Machine Learning. PMLR, 2023.\n\n[2] Marwah, Tanya, Zachary Lipton, and Andrej Risteski. \"Parametric complexity bounds for approximating PDEs with neural networks.\" Advances in Neural Information Processing Systems 34 (2021): 15044-15055.\n\n[3] Chen, Ziang, Jianfeng Lu, and Yulong Lu. \"On the representation of solutions to elliptic pdes in barron spaces.\" Advances in neural information processing systems 34 (2021): 6454-6465.\n\n[4] Li, Zhiyuan, Yi Zhang, and Sanjeev Arora. \"Why are convolutional nets more sample-efficient than fully-connected nets?.\" arXiv preprint arXiv:2010.08515 (2020).\n\n[5] Li, Zongyi, et al. \"Fourier neural operator for parametric partial differential equations.\" arXiv preprint arXiv:2010.08895 (2020)\n\n\n"}, {"review_id": "LSLRx0V1VF", "reviewer": "Reviewer_SM2z", "summary": "The paper proposes a method combining Fourier Neural Operator (FNO) with the Deep Equilibrium Model (DEQ) for more efficient operator learning under noisy conditions. Overall, the paper is well-written with clear methods and theory, and some effective experimental results. It's a decent piece of work, and I recommend a weak acceptance. However, I believe there are still areas in the paper that could be improved.", "questions": "None", "limitations": "None.", "rating": 6, "confidence": 5, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "1. The paper is well-written, with a clear presentation of the method and theoretical underpinnings.\n\n2. The experiments show some positive results, demonstrating the potential utility of the proposed approach.\n\n ", "weaknesses": "1. The paper lacks a thorough literature review. It hardly mentions operator learning methods based on attention, which have demonstrated significant advantages in many areas, such as Navier-Stokes equations and problems on irregular geometric regions. Therefore, I believe the authors should add some references [1,2,3] to cover these works.\n\n2. Each experiment in the paper compares the effects of adding noise to either the input or output. However, the improvements observed in noisy settings do not seem to be as significant as those without noise. This makes me wonder why the authors considered these experiments necessary.\n\n3. Additionally, as a new model structure, DEQ should be able to be combined with many other approaches, such as DeepONet or other neural operators. However, the paper seems not to mention this possibility and only tries to combine it with FNO. It would be interesting to see an exploration of the potential of DEQ in combination with other models.\n\nReferences\n1. Transformer for Partial Differential Equations' Operator Learning (https://arxiv.org/abs/2205.13671)\n\n2. GNOT: A General Neural Operator Transformer for Operator Learning (https://arxiv.org/abs/2302.14376) \n\n3. HT-Net: Hierarchical Transformer based Operator Learning Model for Multiscale PDEs (https://arxiv.org/abs/2210.10890)\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their feedback and are encouraged to see that the reviewer finds the paper to be well-written with positive experimental results. \n\nPlease find our replies to some of your comments and concerns:\n\n**Re: Lacking references to attention based operator learning.**\nWe apologize for not including the mentioned papers in our literature review and will be sure to include them in the Related Work section in the revised manuscript. \n\n**Re: Importance of experiments with noise.**\nThe primary motivation behind showing results with added noise was to emulate real-world scenarios where the different physical quantities in a PDE are measured using sensors—since they are likely to incur noise in the observations. \nWe further note that the relative decrease in the MSE loss for weight-tied architectures (FNO-DEQ and FNO-WT) is smaller than when compared to FNO and FNO++ especially when trained with noisy inputs. For example, for Navier-Stokes with viscosity 0.01 we only see a 7% decrease in the performance of FNO-DEQ (best performing weight-tied model) versus 23% decrease in best performing FNO++ model. \nTherefore, weight tied architectures are indeed an effective inductive bias for solving steady-state PDEs with neural operator frameworks.\n\n\n**Re: Weight-tied DeepONets**\nIn principle, weight-tieing and DEQs can be combined with other neural operators like DeepONets. We chose FNO in part due its performance on various benchmark tasks in PDEbench [1]. We certainly hope that experiments with other operator architectures will be performed, by us and the community, and think this is fertile ground for future work! \n\n\n[1] Takamoto, Makoto, et al. \"PDEBench: An extensive benchmark for scientific machine learning.\" Advances in Neural Information Processing Systems 35 (2022): 1596-1611.\n"}]

Author rebuttal:
We thank the reviewers for their feedback and the detailed comments, and are encouraged to see the overall positive response for our paper! We hope that we have sufficiently addressed all the questions and comments posed by the reviewers in their individual responses. 

Furthermore, we have also attached a PDF with relevant figures and numbers that accompany our individual replies to all the reviewers. 


Meta-review: All reviewers agree that the paper brings new insights and significance to the problem by investigating solutions of steady-state PDEs using Fourier Neural Operator (FNO). Also during the discussion phase, the authors were able to provide useful feedback to the comments given by the reviewers, which was helpful in addressing their technical concerns. As a whole, I recommend acceptance of the paper for this submission.
Decision: Accept (poster)