Title: Deep Equilibrium Based Neural Operators for Steady-State PDEs
Venue: NeurIPS 2023
Accepted: yes
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

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 blackbox root solver and differentiates analytically through this fixed point resulting in O (1) training memory. Our experiments indicate that FNO-DEQ-based architectures outperform FNO-based baselines with 4 × 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.

1 Introduction

Partial differential equations (PDEs) are used to model a wide range of processes in science and engineering. They define a relationship of (unknown) function and its partial derivatives. Most PDEs do not admit a closed form solution, and are solved using a variety of classical numerical methods such as finite element [LeVeque, 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 a rapidly growing area of research in data-driven approaches to 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 a PDE in some family and output its solution-and are trained by examples of coefficient-solution pairs.

While several architectures for this task have been proposed, the design space-in particular taking into account structural properties of the PDEs the operator is trained on-is still largely unexplored. Most present architectures are based on 'neuralizing' a classical numerical method. For instance, Li et al. [2020a] take inspiration from spectral methods, and introduce FNO: a trained composition of (parametrized) kernels in Fourier space. Brandstetter et al. [2022] instead consider finite-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:

where u : Ω → R d u , a : Ω → R d a and f : Ω → R d f are functions defined over the domain Ω , and L is 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-converging Newton-like iterative schemes [LeVeque, 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 et al., 2021, Marwah et al., 2022] to hypothesize that very deep, but heavily weight-tied architectures would provide a 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, as opposed to making the model deeper-thus increasing its size. To this end, we introduce FNO-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 as directly parameterizing the fixed points of some iterative process; or by explicitly expanding some iterative fixed point solver like Newton's or Broyden's method as an infinitely deep, weight-tied model.

Such an architecture has a distinct computational advantage: implicit layer models effectively backpropagate through the infinite-depth network while using only constant memory (equivalent to a single layer's activations). Empirically, we show that for steady-state PDEs, weight-tied and DEQ based models perform better than baselines with 4 × 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 4 × 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 a 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 a wide variety of steady-state PDE families.

Finally, we create a dataset of pairs of steady-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 is to not only improve upon the existing solvers, but to also replace the more computationally inefficient parts of the solver

with a learned counter part. Second set of approaches are full machine learning based approaches that aim to 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 a neural network to learn a correction term to correct over an existing hand designed solver for a Poisson equation, and also provide convergence guarantees of their method to 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 a low-resolution grid with high accuracy. This work specifically focuses on solving Navier-Stokes equations, their method is more accurate than numerical techniques like Direct Numerical Simulation (DNS) with a low-resolution grid, and is also 80 × more faster. Brandstetter et al. [2022] introduced a message passing based hybrid scheme to train a hybrid solver and also propose a 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 et 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 at 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 a variety of (time-dependent) 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]. Theoretical results pertaining to the neural operators mostly include universal approximation results Kovachki et al. [2021a], Lanthaler et al. [2022] which show that architectures like FNO and DeepONet can indeed approximate the infinite dimension operators.

In this work, we focus on 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 a network, weight-tied FNO architectures can outperform FNO and its variants 4 × its size. We further introduce FNO-DEQ, a deep equilibrium based architecture to simulate an infinitely deep weight-tied network (by solving for a 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 2 (Ω; R d ) ) . For a domain Ω we denote by L 2 (Ω; R d ) the space of square integrable functions g : Ω → R d such that ∥ g ∥ L 2 (Ω) < ∞ , where ∥ g ∥ L 2 (Ω) = (∫ Ω ∥ g ( x ) ∥ 2 ℓ 2 dx ) 1 / 2 .

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 a sequence of kernel integral operator layers followed by non-linear activation functions. Each kernel integral operator layer is a convolution-based kernel function that is instantiated through a linear transformation in Fourier domain, making it less sensitive to the level of

spatial discretization. Specifically, an L -layered FNO G θ : R d u → R d u with learnable parameters θ , is defined as

where P : L 2 (Ω; R d u ) → L 2 ( R d v ; R d v ) and Q : L 2 ( R d v ; R d v ) → L 2 ( R d v ; R d u ) are projection operators, and L l : L 2 ( R d v ; R d v ) → L 2 ( R d v ; R d v ) for l ∈ [ L ] is the l th FNO layer defined as,

Here σ is a non-linear activation function, W l , b l are the l th layer weight matrix and bias terms. Finally K l is the l th integral kernel operator which is calculated using the Fourier transform as introduced in Li et al. [2020a] defined as follows,

where F and F -1 are the Fourier transform and the inverse Fourier transform, with R l representing the learnable weight-matrix in the Fourier domain. Therefore, ultimately, the trainable parameters θ is a collection of all the weight matrices and biases, i.e, θ := { W l , b l , R l , · · · , W 1 , b 1 , R 1 } .

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 a deep feedforward network with L layers :

where x ∈ R n x is the input injection, z [ i ] ∈ R n z is the hidden state of i th layer with z [0] = 0 , and f [ i ] θ : R n x × n z ↦→ R n z is the feature transformation of i th layer, parametrized by θ . Suppose the above model is weight-tied, i.e., f [ i ] θ = f θ , ∀ i , and lim i →∞ f θ ( z [ i ] ; x ) exists and its value is z ⋆ . Further, assume that for this z ⋆ , it holds that f θ ( z ⋆ ; x ) = z ⋆ . Then, equilibrium models can be interpreted as the infinite-depth limit of the above network such that f ∞ θ ( z ⋆ ; x ) = z ⋆

Under certain conditions 2 , and for certain classes of f θ 3 , the output z ⋆ 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 z [ t +1] = f θ ( z [ t ] ; x ) some fixed number of times, and backpropagate through the network to compute gradients. However, this can be computationally expensive. Deep equilibrium (DEQ) models [Bai et al., 2019] explicitly solve for z ⋆ through iterative root finding methods like Broyden's method [Broyden, 1965], Newton's method, Anderson acceleration [Anderson, 1965]. DEQs use implicit function theorem to directly differentiate through the fixed point z ⋆ at equilibrium, thus requiring constant memory to backpropagate through an infinite-depth network:

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., Jacobianfree [Fung et al., 2022] or 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 a unique equilibrium point. However, designing f θ 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:

Definition 2 (Steady-State PDE) . Given a bounded open set Ω ⊂ R d , a steady-state PDE can be written in the following general form:

Here L is a continuous operator, the function u ∈ L 2 ( Ω; R d u ) is the unknown function that we wish to solve for and a ∈ L 2 ( Ω; R d a ) collects all the coefficients describing the PDE, and f ∈ L 2 ( Ω; R d f ) is a function independent of u . We will, for concreteness, assume periodic boundary conditions, i.e. ∀ z ∈ Z d , ∀ x ∈ Ω we have u ( x + z ) = u ( x ) . (Equivalently, Ω := T d = [0 , 2 π ] d can be identified with the torus.) 4 Finally, we will denote u ⋆ : Ω → R as the solution to the PDE.

Steady-state models a 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 iterative methods like Newton updates or conjugate gradient descent, typically with carefully chosen preconditioning to 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 a variational formulation), iterative algorithms can be efficiently 'neuralized', that is, the iterative algorithm can be represented by a compact neural network, so long as the coefficients of the PDE are also representable by a 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 a fixed point of a suitably chosen operator. Namely, we can design an operator G : L 2 (Ω; R d u ) × L 2 (Ω; R d f ) → L 2 (Ω; R d u ) 5 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 P : L 2 (Ω; R d u ) × L 2 (Ω; R d f ) → L 2 ( R d v ; R d v ) × L 2 ( R d v ; R d v ) and Q : L 2 ( R d v ; R d v ) → L 2 ( R d v ; R d u ) , are defined as

where W (1) P ∈ R d u × d v , W (2) P ∈ R d f × d v , W Q ∈ R d v × d u and b (1) P , b (2) P ∈ R d v , b Q ∈ R d u .

Definition 4 (Input-injected FNO layer) . An input-injected FNO layer L : L 2 ( R d v ; R d v ) × L 2 ( R d v ; R d v ) → L 2 ( R d v ; R d v ) is defined as

where W ∈ R d v × d v , b ∈ R d v and R ( k ) ∈ R d v × d v for all k ∈ [ K ] are learnable parameters.

Note the difference between the FNO layer specified above, and the standard FNO layer Equation 3 is the extra input g to the layer, which in our architecture will correspond to a projection of (some or 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.

Weight-tied architecture I: Weight-tied FNO The first architecture we consider is a weighttied 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-Tied) . We define a M times weight-tied neural operator G θ M as,

such that: P , Q are projection and embedding layers as in Definition 3 an B L : L 2 ( R d v ; R d v ) × L 2 ( R d v ; R d v ) → L 2 ( R d v ; R d v ) is an L -layer FNO block, i.e, B L = L L ◦ L L -1 ◦ L L -2 ◦ L 1 where for all l ∈ [ L ] , L l ( · , P ( f )) 6 is an input-injected FNO block as in Definition 4.

Weight-tied architecture II: FNO-DEQ In cases where we believe a weight-tied G θ M converges to some fixed point as M →∞ , unrolling G θ M 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 O ( M ) increase in the amount of memory required.

To this end, we use Deep Equilibrium models (DEQs) which enables us to implicitly train G θ := lim M →∞ G θ M by directly solving for the fixed point by leveraging black-box root finding algorithms like quasi-Newton methods, [Broyden, 1965, Anderson, 1965]. Therefore we can think of this approach as an infinite-depth (or infinitely unrolled) weight-tied network. We refer to this architecture as FNO-DEQ .

Definition 6 (FNO-DEQ) . Given P , Q and B L in Definition 5, FNO-DEQ is trained to parametrize the fixed point equation B L ( v ⋆ , P ( f )) = v ⋆ and outputs u ⋆ = Q ( v ⋆ ) .

Usually, it is non-trivial to differentiate through these black-box root solvers. DEQs enable us to implicitly differentiate through the equilibrium fixed point efficiently without any need to backpropagate 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, we recover the original FNO architecture if the input to the l th layer is the output of ( l -1) th layer i.e., v l = B l -1 ( v l -1 ) .

Improved FNO (FNO++ ) : The original FNO architecture suffers from vanishing gradients, which prohibits it from being made deeper [Tran et al., 2021]. We overcome this limitation by introducing residual connections within each block of FNO, with each FNO block B l comprising of three FNO layers L i.e., B l = L l L 1 ◦ L l L 2 ◦ L l L 3 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 v 0 ( x ) = 0 for all x ∈ Ω .

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 a root finding algorithm. We use Anderson acceleration [Anderson, 1965] as the root solver. For the backward pass, we use approximate implicit gradients [Geng et al., 2021] which are light-weight and more stable in practice, compared to 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 or 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 to explicitly create the entire computation graph for backpropagation, which can become very large as the network depth increases. Additionally, FNO++ serves as a non weight-tied counterpart to a weight-tied input-injected network. Like weighttied 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 steadystate PDEs: Darcy Flow equation and steady-state Navier-Stokes equation for incompressible fluids. For experiments with Darcy Flow, we use 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.1 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

experiments with noisy data, both input and observations, we add noise sampled from a sequence of standard Gaussians with increasing values of variance {N (0 , ( σ 2 k )) } M -1 k =0 , where M is the total number of Gaussians we sample from. We set σ 2 0 = 0 and σ 2 max = σ 2 M -1 ≤ 1 /r , where r is the 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 L 2 norm between ground truth u ⋆ and prediction on test data. The total depth of all networks besides FNO is given by 6 B +4 , where B is the number of FNO blocks. Each FNO block has 3 FNO layers and convolutional layers. In addition, we include the depth due to P , Q , and an additional final FNO layer and a convolutional layer. We further elaborate upon network architectures and loss functions in in Appendix A.

5.1 Darcy Flow

For our first set of experiments we consider stationary Darcy Flow equations, a form of linear elliptic PDE with in two dimensions. The PDE is defined as follows,

Note that the diffusion coefficient a ∈ L ∞ (Ω)(Ω; R + ) , i.e., the coefficients are always positive, and f ∈ L 2 (Ω; R d f ) is the forcing term. These PDEs are used to model the steady-state pressure of fluids flowing through a 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 a defining the thermal conductivity of the medium. The task is to learn an operator G θ : L 2 (Ω; R d u ) × L 2 (Ω; R d a ) → L 2 (Ω; R d u ) such that u ⋆ = G θ ( u ⋆ , a ) .

We report the results of our experiments on Darcy 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 4 FNO blocks. In contrast, FNO-WT and FNO-DEQ with just a single FNO block outperform deeper non-weight-tied architectures on clean data and on data with noisy inputs. When observations are noisy, FNO-WT and FNO-DEQ outperform FNO++ with a similar number of parameters, and perform comparably to FNO++ with 4 × parameters.

We also report results on shallow FNO-DEQ, FNO-WT and FNO++ architectures. An FNO block in these shallow networks has a single FNO layer instead of three layers. In our experiments, shallow weight-tied networks outperform non-weight-tied architectures including FNO++ with 7 × 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 instability 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-Stokes equation for an incompressible viscous fluid in the vorticity form defined on a torus, i.e., with periodic boundary condition,

where Ω := (0 , 2 π ) 2 , and u : Ω → R 2 is the velocity and ω : Ω → R where ω = ∇ × u , ν ∈ R + is the viscosity and f : Ω → R is the external force term. We learn an operator G θ : L 2 (Ω; R d u ) × L 2 (Ω; R d f ) → L 2 (Ω; R d u ) , such that u ⋆ = G θ ( u ⋆ , f ) . Wetrain all the models on data with viscosity values ν = 0 . 01 and ν = 0 . 001 , and create a 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-tied 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 in Appendix E. FNO-DEQ and FNO-WT consistently outperform non-weight-tied architectures for higher levels of noise as well.

6 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 a universal approximator, under mild conditions on the steady-state PDEs. Moreover, we also show that in some cases, we can hope the fixed-point solver can converge rapidly.

As noted in Definition 2, we have Ω := T d . We note that all continuous function f ∈ L 2 (Ω; R ) and ∫ Ω | f ( x ) | dx < ∞ can be written as, f ( x ) = ∑ ω ∈ Z d e ix T ω ˆ f w . where { ˆ f ω } ω ∈ Z d are the Fourier coefficients of the function f . We define as L 2 N (Ω) as the space of functions such that for all f N ∈ L 2 N (Ω) with Fourier coefficients that vanish outside a bounded ball. Finally, we define an

orthogonal projection operator Π N : L 2 (Ω) → L 2 N (Ω) , such that for all f ∈ L 2 (Ω) we have,

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 ⋆ ∈ L 2 (Ω; R d u ) define the solution to a steady-state PDE in Definition 2, Then there exists an operator G : L 2 (Ω; R d u ) × L 2 (Ω; R d f ) → L 2 (Ω; R d u ) such that, u ⋆ = G ( u ⋆ , f ) . Furthermore, for every ϵ > 0 there exists an N ∈ N such that for compact sets K u ⊂ L 2 (Ω; R d u ) and K f ⊂ L 2 (Ω; R d f ) there exists a neural network G θ : L 2 N (Ω; R d u ) × L 2 N (Ω; R d f ) → L 2 N (Ω; R d u ) with parameters θ , such that,

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 ) = u -( L ( u ) -f ) , allowing us to use the the results in Kovachki et al. [2021a] that show a 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, f ) = u -A ( L ( u ) -f ) , for a continuous operator A . Such a G can be thought of as a form of 'preconditioned' gradient descent, for a preconditioner A . For example, a Newton update has the form G ( u, f ) = u -L ′ ( u ) -1 ( L ( u ) -f ) , where L ′ : L 2 (Ω; R d u ) → L 2 (Ω; R d u ) is the Frechet derivative of the operator L .

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-Milgram-like conditions (a kind of strong convexity condition), Newton's method converges doubly exponentially fast:

Lemma 1 (Faragó and Karátson [2002], Chapter 5) . Consider the PDE defined Definition 2, such that d u = d v = d f = 1 . such that L ′ ( u ) defines the Frechet derivative of the operator L . If for all u, v ∈ L 2 (Ω; R ) we have ∥ L ′ ( u ) v ∥ L 2 (Ω) ≥ λ ∥ v ∥ L 2 (Ω) and ∥ L ′ ( u ) -L ′ ( v ) ∥ L 2 (Ω) ≤ Λ ∥ u -v ∥ L 2 (Ω) for 0 < λ ≤ Λ < ∞ , then for the Newton update, u t +1 ← u t -L ′ ( u t ) -1 ( L ( u t ) -f ) , with u 0 ∈ L 2 (Ω; R ) , there exists an ϵ > 0 , such that ∥ u T -u ⋆ ∥ L 2 (Ω) ≤ ϵ if T ≥ log ( log ( 1 ϵ ) / log ( 2 λ 2 Λ ∥ L ( u 0 ) -f ∥ L 2 (Ω) )) .

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 weight-tied networks in general-makes them ideal architectures for approximating neural operators for steadystate 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 ∼ 4 × 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-weight-tied 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.

References

Donald G 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. 4

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(31):15344-15349, 2019. 2, 3

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. 1, 3

John P Boyd. Chebyshev and Fourier spectral methods . Courier Corporation, 2001. 1

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 simultaneous 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. 1, 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-6465, 2021. 2, 3, 5

Gideon Dresdner, Dmitrii Kochkov, Peter Norgaard, Leonardo Zepeda-Núñez, Jamie A. Smith, Michael P. Brenner, and Stephan Hoyer. Learning to correct spectral methods for simulating turbulent flows. 2022. doi: 10.48550/ARXIV.2207.00556. URL https://arxiv.org/abs/ 2207.00556 . 14

István Faragó and János Karátson. Numerical solution of nonlinear elliptic problems via preconditioning 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-free backpropagation for implicit networks. In Proceedings of the AAAI Conference on Artificial Intelligence , volume 36, pages 6648-6656, 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 on Machine Learning , pages 12556-12569. PMLR, 2023. 3

Dan Hendrycks and Kevin Gimpel. Gaussian error linear units (gelus). arXiv preprint arXiv:1606.08415 , 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 learning-accelerated computational fluid dynamics. Proceedings of the National Academy of Sciences , 118(21):e2101784118, 2021. 2, 3, 14

David A Kopriva. Implementing spectral methods for partial differential equations: Algorithms for scientists and engineers . Springer Science & Business Media, 2009. 1

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):1323713312, 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 arXiv:2108.08481 , 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 to PDEs. Journal of Machine Learning Research , 24(89):1-97, 2023. 3

Samuel Lanthaler, Siddhartha Mishra, and George E Karniadakis. Error estimates for deeponets: A deep learning framework in infinite dimensions. Transactions of Mathematics and Its Applications , 6(1):tnac001, 2022. 3

Randall J LeVeque. Finite difference methods for ordinary and partial differential equations: steadystate 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. 3

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.08895 , 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.03485 , 2020b. 1, 7

Zongyi Li, Daniel Zhengyu Huang, Burigede Liu, and Anima Anandkumar. Fourier neural operator with learned deformations for pdes on general geometries. arXiv preprint arXiv:2207.05209 , 2022b. 10

Renjie Liao, Yuwen Xiong, Ethan Fetaya, Lisa Zhang, KiJung Yoon, Xaq Pitkow, Raquel Urtasun, and Richard Zemel. Reviving and improving recurrent back-propagation. In International Conference on Machine Learning , pages 3082-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://arxiv.org/abs/2210.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 arXiv:1910.03193 , 2019. 1, 3

Tanya Marwah, Zachary Lipton, and Andrej Risteski. Parametric complexity bounds for approximating 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, 3, 5, 9

Fadl Moukalled, Luca Mangani, Marwan Darwish, F Moukalled, L Mangani, and M Darwish. The finite volume method . Springer, 2016. 1

Suhas V Patankar and D Brian Spalding. A calculation procedure for heat, mass and momentum transfer in three-dimensional parabolic flows. In Numerical prediction of flow, heat transfer, turbulence 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 in Applied Mechanics and Engineering , 373:113500, 2021. 3

Max Revay, Ruigang Wang, and Ian R Manchester. Lipschitz bounded equilibrium networks. arXiv preprint arXiv:2010.01732 , 2020. 4

Alasdair Tran, Alexander Mathews, Lexing Xie, and Cheng Soon Ong. Factorized Fourier neural operators. arXiv preprint arXiv:2111.13802 , 2021. 3, 6, 10

Ezra Winston and J Zico Kolter. Monotone operator equilibrium networks. Advances in neural information processing systems , 33:10718-10728, 2020. 4

Appendix

A Implementation Details

Training details. We train all the networks for 500 epochs with Adam optimizer. The learning rate is set to 0.001 for Darcy flow and 0.005 for Navier-Stokes. We use learning rate weight decay of 1e-4 for both Navier-Stokes and Darcy 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 on a combination of NVIDIA RTX A6000, NVIDIA GeForce RTX 2080 Ti and 3080 Ti. All networks can easily fit on a single NVIDIA RTX A6000, but training time varies between the networks.

For FNO-DEQ, we use Anderson solver [Anderson, 1965] to solve for the fixed point in 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 et al., 2021] which are computed as:

where τ is a tunable damping factor and u ⋆ is the fixed point computed using Anderson solver in the forward pass. This step can be repeated S times. We use τ = 0 . 5 and S = 1 for Darcy 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 in 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. 5.

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 × 421 which we downsample to 85 × 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, 1e-6, 1e-5, 1e-4, 1e-3].

B.2 Steady-State Incompressible Fluid Navier-Stoke

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 ω ⋆ := ∇× u ⋆ of the PDE and then generate the corresponding equation, i.e. calculate the corresponding force term f = u ⋆ · ∇ ω ⋆ -ν ∆ ω ⋆ .

To generate the solutions ω ⋆ , we forward propagate a relatively simple initial distribution of ω 0 (sampled from a Gaussian random field) through a time-dependent Navier-Stokes equation in the vorticity form for a short period of time. This ensures our dataset contains solutions ω ∗ that are rich

a

u*

and complex. Precisely, recall the Navier-Stokes equations in their vorticity form:

where g ( x ) = ∇× ˜ g ( x ) and ˜ g ( x ) = sin(5 x 1 ) ˆ x 2 is a divergence free forcing term and x = ( x 1 , x 2 ) 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 et al., 2021, Dresdner et al., 2022] package. The initial vorticity ω 0 is sampled from a Gaussian random field N (0 , (5 3 / 2 ( I +25∆) -2 . 5 )) , which is then made divergence free. We forward propagate the NavierStokes equation in Equation 13 for time T = 0 . 5 with dt = 0 . 002 to get ω (1 , x ) , which we choose as the solution to the steady-state PDE in Equation 10, i.e, ω ⋆ for Equation 10.

Subsequently, we use the stream function Ψ [Batchelor and Batchelor, 1967] to calculate u = ( ∂ Ψ /∂x 1 , ∂ Ψ /∂x 2 ) by solving the Poisson equation ∆Ψ = ω in the Fourier domain. Furthermore, since f = u ⋆ · ∇ ω ⋆ -ν ∆ ω ⋆ we use the stream function to calculate ( f 1 , f 2 ) , i.e., the different components of the force term.

We use 4500 training samples and 500 testing samples. The input to the network is the vector field ˜ f = ( f 1 , f 2 ) and we learn a map that outputs the vorticity ω ⋆ . The resolution of grid used to

generate the dataset is 256 × 256 which we downsample to 128 × 128 while training the models. For experiments with noisy inputs/observations, we consider two 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 are [0, 1e-9, 1e-8, 1e-7, 1e-6, 1e-5, 1e-4, 1e-3, 2e-3, 4e-3]. However, when conducting experiments with a variance of 1e-3, we exclude the last two values of variance from this list.

f1

f2

*3

C Proof of Universal Approximation

The proof of the universal approximation essentially follows from the result on the universal approximation capabilities of FNO layers in Kovachki et al. [2021a], applied to G ( v, f ) = v -( Lv -f ) . For the sake of completeness, we reitarate the key steps.

For simplicity, we will assume that d u = d v = d f = 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 : L 2 (Ω; R ) → L 2 (Ω; R ) is continuous at u ∈ L 2 (Ω; R ) if for every ϵ > 0 , there exists a δ > 0 , such that for all v ∈ L 2 (Ω) with ∥ u -v ∥ L 2 (Ω) ≤ δ , we have ∥ L ( u ) -L ( v ) ∥ L 2 (Ω) ≤ ϵ .

f

f2

*3

First, we approximate the infinite dimensional operator G : L 2 (Ω) × L 2 (Ω) → L 2 (Ω) by projecting the functions in L 2 (Ω) to a finite-dimensional approximation L 2 N (Ω) , and considering the action of the operator on this subspace. The linear projection we use is the one introduced in Equation 11. More precisely we show the following result,

Lemma 2. Given a continuous operator L : L 2 (Ω) → L 2 (Ω) as defined in Equation 1, let us define an operator G : L 2 (Ω) × L 2 (Ω) → L 2 (Ω) as G ( v, f ) := v -( L ( v ) -f ) . Then, for every ϵ > 0 there exists an N ∈ N such that for all v, f in any compact set K ⊂ L 2 (Ω) , the operator G N = Π N G (Π N v, Π N f ) is an ϵ -approximation of G ( v, f ) , i.e., we have,

Proof. Note that for an ϵ > 0 there exists an N = N ( ϵ, d ) such that for all v ∈ K we have

Therefore, using the definition of G N we can bound the L 2 (Ω) norm of the difference between G and G N as follows,

We first bound the term I as follows:

Since L is continuous, for all compact sets K ⊂ L 2 (Ω) , L ( K ) is compact as well. This is because: (1) for any u ∈ K , ∥ L ( u ) ∥ L 2 (Ω) is finite; (2) for any v ∈ K , ∥ L ( v ) ∥ L 2 (Ω) ≤ ∥ L ( u ) ∥ L 2 (Ω) + C ∥ u -v ∥ L 2 (Ω) . Therefore, for every ϵ > 0 , there exists an N ∈ N such that

Substituting the above result in Equation 14, we have

Similarly, for all v ∈ K where K is compact, we can bound Term II as following,

Now, since v ∈ K and L : L 2 (Ω) → L 2 (Ω) is a continuous operator, there exists a modulus of continuity (an increasing real valued function) α ∈ [0 , ∞ ) , such that for all v ∈ K , we have

Hence for every ϵ > 0 there exists an N ∈ N such that,

Plugging these bounds in Equation 16, we get,

Therefore, combining Equation 15 and Equation 17 then for ϵ > 0 , there exists an N ∈ N , such that for all v, f ∈ K we have

Taking ϵ ′ = 6 ϵ proves the claim.

Proof of Theorem 1. For Lemma 2 we know that there exists a finite dimensional projection for the operator G , defined as G N ( v, f ) such that for all v, f ∈ L 2 (Ω) we have

Now using the definition of G N ( v, f ) we have

From Kovachki et al. [2021a], Theorem 2.4 we know that there exists an FNO network G θ L of the form defined in Equation 2 such that for all v ∈ K , where K is a compact set, there exists an ϵ L we have

Finally, note that from Lemma D.1 in Kovachki et al. [2021a], we have that for any v ∈ K , there exists an FNO layers G θ f ∈ L 2 (Ω) and G θ v ∈ L 2 (Ω) defined in Equation 3 such that

and for ϵ v > 0 and ϵ f > 0 .

Therefore there exists an ˜ ϵ > such that there is an FNO network G θ : L 2 (Ω) × L 2 (Ω) → L 2 (Ω) where θ := { θ L , θ v , θ f } such that

Now, since we know that u ⋆ is the fixed point of the operator G we have from Lemma 2 and Equation 22,

D Fast Convergence for Newton Method

Definition 8 (Frechet Derivative in L 2 (Ω) ) . For a continuous operator F : L 2 (Ω) → L 2 (Ω) , the Frechet derivative at u ∈ L 2 (Ω) is a linear operator F ′ ( u ) : L 2 (Ω) → L 2 (Ω) such that for all v ∈ L 2 (Ω) we have

Lemma 3. Given the operator L : L 2 (Ω) → L 2 (Ω) with Frechet derivative L ′ , such that for all u, v ∈ L 2 (Ω) , we have ∥ L ′ ( u )( v ) ∥ L 2 (Ω) ≥ λ ∥ v ∥ L 2 (Ω) , then L ′ ( u ) -1 exists and we have, for all v 1 , v 2 ∈ L 2 (Ω) :

Proof. Note that for all u, v ′ ∈ L 2 (Ω) we have,

Taking v = L ′ ( u ) -1 ( v ′ ) , we have

For part 2 , note that there exists a c ∈ [0 , 1] such that

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 d u = d v = d f = 1 . such that L ′ ( u ) defines the Frechet derivative of the operator L . If for all u, v ∈ L 2 (Ω; R ) we have ∥ L ′ ( u ) v ∥ L 2 (Ω) ≥ λ ∥ v ∥ L 2 (Ω) 7 and ∥ L ′ ( u ) -L ′ ( v ) ∥ L 2 (Ω) ≤ Λ ∥ u -v ∥ L 2 (Ω) for 0 < λ ≤ Λ < ∞ , then for the Newton update, u t +1 ← u t -L ′ ( u t ) -1 ( L ( u t ) -f ) , with u 0 ∈ L 2 (Ω; R ) , there exists an ϵ > 0 , such that ∥ u T -u ⋆ ∥ L 2 (Ω) ≤ ϵ if 8 T ≥ log ( log ( 1 ϵ ) / log ( 2 λ 2 Λ ∥ L ( u 0 ) -f ∥ L 2 (Ω) )) .

Proof of Lemma 4. Re-writing the updates in Lemma 4 as,

Now, upper bounding L ( u t +1 ) -f for all x ∈ Ω we have,

where we use Equation 24 in the final step.

Taking L 2 (Ω) norm on both sides and using the fact that ∥ L ′ ( u ) -L ′ ( v ) ∥ L 2 (Ω) ≤ Λ ∥ u -v ∥ L 2 (Ω) , we have

Noting that for all x ∈ Ω , we have u t +1 -u t = p t , and using the fact that for all u, v ∥ L ′ ( u ) -1 v ∥ L 2 (Ω) ≤ 1 λ ∥ v ∥ L 2 (Ω) we have, ∥ L ′ ( u t ) p t ∥ L 2 (Ω) ≤ 1 λ ∥ p t ∥ L 2 (Ω)

where we use the result from Lemma 3 in the last step.

Therefore we have

Therefore, if then we have

for

E Additional experimental results

We provide additional results for Navier-Stokes equation for noisy inputs and observations in Table 4 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 non-weight-tied architectures.

‡ indicates that the network diverges during training for one of the seeds.

Convergence analysis of fixed point. Wereport variations in test error, absolute residual ∥ G θ ( z t ) -z t ∥ 2 , and relative residual ∥ G θ ( z t ) -z t ∥ 2 ∥ z t ∥ 2 with an increase in the number of solver steps while solving for the fixed point in 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 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 FNO-DEQ is trained with fewer solver steps. 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.

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 a certain point once we have a reasonable estimate of the fixed point, hence showing that more solver steps stabilize to the same solution.

(a) Training Loss Curve

(b) Test Loss Curve

100

FNO-DEQ #Blocks 1

FNO-WT #Blocks 1

FNO++ #Blocks 1

FNO++ #Blocks 2

FNO++ #Blocks 4

Train MSE Loss

FNO #Blocks 1

FNO #Blocks 2

FNO #Blocks 4

10-1

0

100

200

300

400

500

Epochs

100

FNO-DEQ #Blocks 1

FNO-WT #Blocks 1

FNO++ #Blocks 1

FNO++#Blocks2

FNO++ #Blocks 4

Test MSE Loss

FNO #Blocks 1

FNO #Blocks 2

FNO #Blocks 4

10-1

0

100

200

300

400

500

Epochs

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n1. 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\nWEAKNESSES:\n1. 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\nQUESTIONS:\nNone\n\nLIMITATIONS:\nNone.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-06-16"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nSome 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.\n\nQUESTIONS:\nRefer to the weaknesses section.\n\nLIMITATIONS:\nRefer to the weaknesses section.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 1, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThis 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\nWEAKNESSES:\nThe 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.\n\nQUESTIONS:\nThe 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.\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 9, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-06"}

comments: 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.

STRENGTHS:
- The paper is well structured,
- The proposed method leverages well the fact that the solutions are fixed points.
- 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?
- 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?

QUESTIONS:
- How is the speed of computing the implicit gradient?
- In Definition 3, shouldn't all functions map from $\Omega$?

LIMITATIONS:
None
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-13