SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Tanya Marwah $^{1*}$  Ashwini Pokle $^{1*}$  J. Zico Kolter $^{1,2}$  Zachary C. Lipton $^{1}$  Jianfeng Lu $^{3}$  Andrej Risteski $^{1}$
$^{1}$ Carnegie Mellon University  $^{2}$  Bosch Center for AI  $^{3}$  Duke University {tmarwah, apokle, zicokolter, zlipton, aristesk} @andrew.cmu.edu jianfeng@math.duke.edu
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.
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.
*Equal contribution. Correspondence to tmarwah@andrew.cmu.edu and apokle@andrew.cmu.edu
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:
$$
L (a (x), u (x)) = f (x), \quad \forall x \in \Omega , \tag {1}
$$
where  $u: \Omega \to \mathbb{R}^{d_u}$ ,  $a: \Omega \to \mathbb{R}^{d_a}$  and  $f: \Omega \to \mathbb{R}^{d_f}$  are functions defined over the domain  $\Omega$ , 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 \times$  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 \times$  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.
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 \times$  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 \times$  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  $\mathcal{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.
We now introduce some key concepts and notation.
Definition 1 ( $L^2(\Omega; \mathbb{R}^d)$ ). For a domain  $\Omega$  we denote by  $L^2(\Omega; \mathbb{R}^d)$  the space of square integrable functions  $g: \Omega \to \mathbb{R}^d$  such that  $\| g \|_{L^2(\Omega)} < \infty$ , where  $\| g \|_{L^2(\Omega)} = \left( \int_{\Omega} \| g(x) \|_{\ell_2}^2 dx \right)^{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_{\theta}:\mathbb{R}^{d_u}\to \mathbb{R}^{d_u}$  with learnable parameters  $\theta$  is defined as
$$
G _ {\theta} := \mathcal {Q} \circ \mathcal {L} _ {L} \circ \mathcal {L} _ {L - 1} \circ \dots \circ \mathcal {L} _ {1} \circ \mathcal {P} \tag {2}
$$
where  $\mathcal{P}:L^2 (\Omega ;\mathbb{R}^{d_u})\to L^2 (\mathbb{R}^{d_v};\mathbb{R}^{d_v})$  and  $\mathcal{Q}:L^2 (\mathbb{R}^{d_v};\mathbb{R}^{d_v})\to L^2 (\mathbb{R}^{d_v};\mathbb{R}^{d_u})$  are projection operators, and  $\mathcal{L}_l:L^2 (\mathbb{R}^{d_v};\mathbb{R}^{d_v})\to L^2 (\mathbb{R}^{d_v};\mathbb{R}^{d_v})$  for  $l\in [L]$  is the  $l^{\mathrm{th}}$  FNO layer defined as,
$$
\mathcal {L} _ {l} \left(v _ {l}\right) = \sigma \left(W _ {l} v _ {l} + b _ {l} + \mathcal {K} _ {l} \left(v _ {l}\right)\right). \tag {3}
$$
Here  $\sigma$  is a non-linear activation function,  $W_{l}, b_{l}$  are the  $l^{th}$  layer weight matrix and bias terms. Finally  $\mathcal{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,
$$
\mathcal {K} _ {l} \left(v _ {l}\right) = \mathcal {F} ^ {- 1} \left(R _ {l} \cdot \left(\mathcal {F} v _ {l}\right)\right) (x) \quad \forall x \in \Omega , \tag {4}
$$
where  $\mathcal{F}$  and  $\mathcal{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  $\theta$  is a collection of all the weight matrices and biases, i.e.,  $\theta \coloneqq \{W_l,b_l,R_l,\dots ,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:
$$
z ^ {[ i + 1 ]} = f _ {\theta} ^ {[ i ]} \left(z ^ {[ i ]}; x\right) \quad \text {f o r} i = 0, \dots , L - 1 \tag {5}
$$
where  $x \in \mathbb{R}^{n_x}$  is the input injection,  $z^{[i]} \in \mathbb{R}^{n_z}$  is the hidden state of  $i^{th}$  layer with  $z^{[0]} = 0$ , and  $f_{\theta}^{[i]}: \mathbb{R}^{n_x \times n_z} \mapsto \mathbb{R}^{n_z}$  is the feature transformation of  $i^{th}$  layer, parametrized by  $\theta$ . Suppose the above model is weight-tied, i.e.,  $f_{\theta}^{[i]} = f_{\theta}, \forall i$ , and  $\lim_{i \to \infty} f_{\theta}(z^{[i]}; x)$  exists and its value is  $z^{\star}$ . Further, assume that for this  $z^{\star}$ , it holds that  $f_{\theta}(z^{\star}; x) = z^{\star}$ . Then, equilibrium models can be interpreted as the infinite-depth limit of the above network such that  $f_{\theta}^{\infty}(z^{\star}; x) = z^{\star}$
Under certain conditions $^2$ , and for certain classes of  $f_{\theta}^3$ , the output  $z^{\star}$  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_{\theta}(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^{\star}$  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^{\star}$  at equilibrium, thus requiring constant memory to backpropagate through an infinite-depth network:
$$
\frac {\partial z ^ {\star}}{\partial \theta} = \left(I - \frac {\partial f _ {\theta} \left(z ^ {\star} ; x\right)}{\partial z ^ {\star}}\right) ^ {- 1} \frac {\partial f _ {\theta} \left(z ^ {\star} ; x\right)}{\partial \theta} \tag {6}
$$
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] 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_{\theta}$  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 of  $f_{\theta}$ , and divergence is possible.
$^3$ Bai et al. [2019] state that  $f_{\theta}$  needs to be stable and constrained. In general, by Banach's fixed point theorem, global convergence is guaranteed if  $f_{\theta}$  is contractive over its input domain.
Definition 2 (Steady-State PDE). Given a bounded open set  $\Omega \subset \mathbb{R}^d$ , a steady-state PDE can be written in the following general form:
$$
L (a (x), u (x)) = f (x), \quad \forall x \in \Omega \tag {7}
$$
Here  $L$  is a continuous operator; the function  $u \in L^2(\Omega; \mathbb{R}^{d_u})$  is the unknown function that we wish to solve for and  $a \in L^2(\Omega; \mathbb{R}^{d_a})$  collects all the coefficients describing the PDE, and  $f \in L^2(\Omega; \mathbb{R}^{d_f})$  is a function independent of  $u$ . We will, for concreteness, assume periodic boundary conditions, i.e.  $\forall z \in \mathbb{Z}^d, \forall x \in \Omega$  we have  $u(x + z) = u(x)$ . (Equivalently,  $\Omega := \mathbb{T}^d = [0, 2\pi]^d$  can be identified with the torus.)<sup>4</sup> Finally, we will denote  $u^\star : \Omega \to \mathbb{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  $\mathcal{G}:L^{2}(\Omega ;\mathbb{R}^{d_u})\times L^2 (\Omega ;\mathbb{R}^{d_f})\to L^2 (\Omega ;\mathbb{R}^{d_u})^5$  such that  $u^{\star} = \mathcal{G}(u^{\star},f)$  and a lot of common (preconditioned) first and second-order methods are natural ways to recover the fixed points  $u^{\star}$ .
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  $\mathcal{P}: L^2(\Omega; \mathbb{R}^{d_u}) \times L^2(\Omega; \mathbb{R}^{d_f}) \to L^2(\mathbb{R}^{d_v}; \mathbb{R}^{d_v}) \times L^2(\mathbb{R}^{d_v}; \mathbb{R}^{d_v})$  and  $\mathcal{Q}: L^2(\mathbb{R}^{d_v}; \mathbb{R}^{d_v}) \to L^2(\mathbb{R}^{d_v}; \mathbb{R}^{d_v})$ , are defined as
$$
\mathcal {P} (v, f) = \left(\sigma \left(W _ {P} ^ {(1)} v + b _ {P} ^ {(1)}\right), \sigma \left(W _ {P} ^ {(2)} f + b _ {P} ^ {(2)}\right)\right),
$$
$$
\mathcal {Q} (v) = \sigma \left(W _ {Q} v + b _ {Q}\right)
$$
where  $W_{P}^{(1)}\in \mathbb{R}^{d_{u}\times d_{v}},W_{P}^{(2)}\in \mathbb{R}^{d_{f}\times d_{v}},W_{Q}\in \mathbb{R}^{d_{v}\times d_{u}}$  and  $b_{P}^{(1)},b_{P}^{(2)}\in \mathbb{R}^{d_{v}},b_{Q}\in \mathbb{R}^{d_{u}}$
Definition 4 (Input-injected FNO layer). An input-injected FNO layer  $\mathcal{L}: L^2(\mathbb{R}^{d_v}; \mathbb{R}^{d_v}) \times L^2(\mathbb{R}^{d_v}; \mathbb{R}^{d_v}) \to L^2(\mathbb{R}^{d_v}; \mathbb{R}^{d_v})$  is defined as
$$
\mathcal {L} (v, g) := g + \sigma \left(W v + b + \mathcal {F} ^ {- 1} \left(R ^ {(k)} \cdot (\mathcal {F} v)\right)\right). \tag {8}
$$
where  $W\in \mathbb{R}^{d_v\times d_v}$ $b\in \mathbb{R}^{d_v}$  and  $R^{(k)}\in \mathbb{R}^{d_v\times d_v}$  for all  $k\in [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 weighted 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_{\theta}^{M}$  as,
$$
G _ {\theta} ^ {M} = \mathcal {Q} \circ \underbrace {\mathcal {B} ^ {L} \circ \mathcal {B} ^ {L} \circ \cdots \circ \mathcal {B} ^ {L}} _ {M t i m e s} \circ \mathcal {P}
$$
<sup>4</sup>This is for convenience of exposition, our methods can readily be extended to other boundary conditions like Dirichet, Neumann etc.
We note that the choice of defining the operator function  $f$  is made for purely expository purposes the operator  $\mathcal{G}$  can be defined as  $\mathcal{G}: L^2(\Omega; \mathbb{R}^{d_u}) \times L^2(\Omega; \mathbb{R}^{d_a}) \to L^2(\Omega; \mathbb{R}^{d_u})$  as well.
such that:  $\mathcal{P},\mathcal{Q}$  are projection and embedding layers as in Definition 3 an  $\mathcal{B}^L:L^2 (\mathbb{R}^{d_v};\mathbb{R}^{d_v})\times$ $L^2 (\mathbb{R}^{d_v};\mathbb{R}^{d_v})\to L^2 (\mathbb{R}^{d_v};\mathbb{R}^{d_v})$  is an  $L$  -layer FNO block, i.e,  $\mathcal{B}^L = \mathcal{L}_L\circ \mathcal{L}_{L - 1}\circ \mathcal{L}_{L - 2}\circ \mathcal{L}_1$  where for all  $l\in [L]$ $\mathcal{L}_l(\cdot ,\mathcal{P}(f))$  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_{\theta}^{M}$  converges to some fixed point as  $M \to \infty$ , unrolling  $G_{\theta}^{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  $\mathcal{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_{\theta} \coloneqq \lim_{M \to \infty} G_{\theta}^{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  $\mathcal{P}, \mathcal{Q}$  and  $\mathcal{B}^L$  in Definition 5, FNO-DEQ is trained to parametrize the fixed point equation  $\mathcal{B}^L(v^\star, \mathcal{P}(f)) = v^\star$  and outputs  $u^\star = \mathcal{Q}(v^\star)$ .
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  $\mathcal{O}(1)$  training memory.