SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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^{\mathrm{th}}$  layer is the output of  $(l - 1)^{\mathrm{th}}$  layer i.e.,  $v_{l} = \mathcal{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  $\mathcal{B}_l$  comprising of three FNO layers  $\mathcal{L}$  i.e.,  $\mathcal{B}_l = \mathcal{L}_{L_1}^l \circ \mathcal{L}_{L_2}^l \circ \mathcal{L}_{L_3}^l$  and three convolutional layers, where  $\mathcal{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 \in \Omega$ .
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  $\mathcal{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,  $\mathrm{FNO}++$  serves as a non-weight-tied counterpart to a weight-tied input-injected network. Like weight-tied networks,  $\mathrm{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 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
<sup>6</sup>We are abusing the notation somewhat and denoting by  $\mathcal{P}(f)$  the second coordinate of  $\mathcal{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  $\{\mathcal{N}(0,(\sigma_k^2))\}_{k = 0}^{M - 1}$ , where  $M$  is the total number of Gaussians we sample from. We set  $\sigma_0^2 = 0$  and  $\sigma_{\mathrm{max}}^2 = \sigma_{M - 1}^2\leq 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^{\star}$  and prediction on test data. The total depth of all networks besides FNO is given by  $6B + 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  $\mathcal{P},\mathcal{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,
$$
\begin{array}{l} - \nabla \cdot (a (x) \nabla u (x)) = f (x), \quad x \in (0, 1) ^ {2} \\ u (x) = 0 \quad x \in \partial (0, 1) ^ {2}. \tag {9} \\ \end{array}
$$
Note that the diffusion coefficient  $a \in L^{\infty}(\Omega)(\Omega; \mathbb{R}_+)$ , i.e., the coefficients are always positive, and  $f \in L^{2}(\Omega; \mathbb{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_{\theta}: L^{2}(\Omega; \mathbb{R}^{d_{u}}) \times L^{2}(\Omega; \mathbb{R}^{d_{a}}) \to L^{2}(\Omega; \mathbb{R}^{d_{u}})$  such that  $u^{\star} = G_{\theta}(u^{\star}, 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  $\mathcal{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\times$  parameters.
We also report results on shallow FNO-DEQ, FNO-WT and  $\mathrm{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  $\mathrm{FNO}++$  with  $7\times$  parameters on clean data and on data with noisy inputs, and perform comparably to deep  $\mathrm{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,
$$
\begin{array}{l} u \cdot \nabla \omega = \nu \Delta \omega + f, \quad x \in \Omega \\ \nabla \cdot u = 0 \quad x \in \Omega \\ \end{array}
$$
where  $\Omega := (0,2\pi)^2$ , and  $u: \Omega \to \mathbb{R}^2$  is the velocity and  $\omega: \Omega \to \mathbb{R}$  where  $\omega = \nabla \times u$ ,  $\nu \in \mathbb{R}_+$  is the viscosity and  $f: \Omega \to \mathbb{R}$  is the external force term. We learn an operator  $G_{\theta}: L^{2}(\Omega; \mathbb{R}^{d_u}) \times L^{2}(\Omega; \mathbb{R}^{d_f}) \to L^{2}(\Omega; \mathbb{R}^{d_u})$ , such that  $u^{\star} = G_{\theta}(u^{\star}, f)$ . We train all the models on data with viscosity values  $\nu = 0.01$  and  $\nu = 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.
Architecture	Parameters	#Blocks	Test error ↓
σ2max=0	(σ2max)i=0.001	(σ2max)t=0.001
FNO	2.37M	1	0.0080 ± 5e-4	0.0079 ± 2e-4	0.0125 ± 4e-5
FNO	4.15M	2	0.0105 ± 6e-4	0.0106 ± 4e-4	0.0136 ± 2e-5
FNO	7.71M	4	0.2550 ± 2e-8	0.2557 ± 8e-9	0.2617 ± 2e-9
FNO++	2.37M	1	0.0075 ± 2e-4	0.0075 ± 2e-4	0.0145 ± 7e-4
FNO++	4.15M	2	0.0065 ± 2e-4	0.0065 ± 9e-5	0.0117 ± 5e-5
FNO++	7.71M	4	0.0064 ± 2e-4	0.0064 ± 2e-4	0.0109 ± 5e-4
S-FNO++	1.78M	0.66	0.0093 ± 5e-4	0.0094 ± 7e-4	0.0402 ± 6e-3
FNO-WT	2.37M	1	0.0055 ± 1e-4	0.0056 ± 5e-5	0.0112 ± 4e-4
FNO-DEQ	2.37M	1	0.0055 ± 1e-4	0.0056 ± 7e-5	0.0112 ± 4e-4
S-FNO-WT	1.19M	0.33	0.0057 ± 3e-5	0.0057 ± 5e-5	0.0112 ± 1e-4
S-FNO-DEQ	1.19M	0.33	0.0056 ± 4e-5	0.0056 ± 5e-5	0.0136 ± 0.011
In general, DEQ-based architectures are slower to train (upto  $\sim 2.5\times$  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 non-weight-tied FNO layers. In general, we observe diminishing returns in  $\mathrm{FNO + + }$  beyond 4 blocks. Additionally, training the original FNO network on more than 4 FNO blocks is challenging, with the network often diverging during training, and therefore we do not include these results in the paper.
Table 1: Results on Darcy flow: clean data (Col 4), noisy inputs (Col 5) and noisy observations (Col 6) with max variance of added noise being  $(\sigma_{\mathrm{max}}^2)^i$  and  $(\sigma_{\mathrm{max}}^2)^t$ , respectively. Reported test error 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.
Architecture	Parameters	#Blocks	Test error ↓
σ2max=0	(σ2max)i=0.001	(σ2max)t=0.001
FNO	2.37M	1	0.184 ± 0.002	0.218 ± 0.003	0.184 ± 0.001
FNO	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	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 2: 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  $(\sigma_{\mathrm{max}}^2)^i$  and  $(\sigma_{\mathrm{max}}^2)^t$ , respectively. Reported test error has been averaged on three different runs with seeds 0, 1, and 2.
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  $\Omega := \mathbb{T}^d$ . We note that all continuous function  $f \in L^{2}(\Omega; \mathbb{R})$  and  $\int_{\Omega} |f(x)| dx < \infty$  can be written as,  $f(x) = \sum_{\omega \in \mathbb{Z}^d} e^{ix^T \omega} \hat{f}_w$ . where  $\{\hat{f}_{\omega}\}_{\omega \in \mathbb{Z}^d}$  are the Fourier coefficients of the function  $f$ . We define as  $L_N^2(\Omega)$  as the space of functions such that for all  $f_N \in L_N^2(\Omega)$  with Fourier coefficients that vanish outside a bounded ball. Finally, we define an
Architecture	Parameters	#Blocks	Test error ↓
σ2max=0	(σ2max)i=0.001	(σ2max)t=0.001
FNO	2.37M	1	0.181 ± 0.005	0.186 ± 0.003	0.178 ± 0.006
FNO	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
FNO++	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
FNO++	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 3: 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  $(\sigma_{\mathrm{max}}^2)^i$  and  $(\sigma_{\mathrm{max}}^2)^t$ , respectively. Reported test error has been averaged on three different runs with seeds 0, 1, and 2.
orthogonal projection operator  $\Pi_N:L^2 (\Omega)\to L_N^2 (\Omega)$ , such that for all  $f\in L^{2}(\Omega)$  we have,
$$
f _ {n} = \Pi_ {N} (f) = \Pi_ {N} \left(\sum_ {\omega \in \mathbb {Z} ^ {d}} f _ {\omega} e ^ {i x ^ {T} \omega}\right) = \sum_ {\| \omega \| _ {\infty} \leq N} \hat {f} _ {\omega} e ^ {i x ^ {T} \omega}. \tag {11}
$$
That is, the projection operator  $\Pi_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^{\star} \in L^{2}(\Omega; \mathbb{R}^{d_u})$  define the solution to a steady-state PDE in Definition 2. Then there exists 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})$  such that,  $u^{\star} = \mathcal{G}(u^{\star}, f)$ . Furthermore, for every  $\epsilon > 0$  there exists an  $N \in \mathbb{N}$  such that for compact sets  $K_{u} \subset L^{2}(\Omega; \mathbb{R}^{d_u})$  and  $K_{f} \subset L^{2}(\Omega; \mathbb{R}^{d_f})$  there exists a neural network  $G_{\theta}: L_N^2(\Omega; \mathbb{R}^{d_u}) \times L_N^2(\Omega; \mathbb{R}^{d_f}) \to L_N^2(\Omega; \mathbb{R}^{d_u})$  with parameters  $\theta$ , such that,
$$
\sup  _ {u \in K _ {u}, f \in K _ {f}} \| u ^ {\star} - G _ {\theta} (\Pi_ {N} u ^ {\star}, \Pi_ {N} f) \| _ {L ^ {2} (\Omega)} \leq \epsilon .
$$
The proof for the above theorem is relatively straightforward and provided in Appendix C. The proof uses the fact that  $u^{\star}$  is a fixed-point of the operator  $G(u, f) = u - (L(u) - f)$ , allowing us to use 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}(\Omega; \mathbb{R}^{d_{u}}) \to L^{2}(\Omega; \mathbb{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^{\prime}(u)$  defines the Frechet derivative of the operator  $L$ . If for all  $u, v \in L^{2}(\Omega; \mathbb{R})$  we have  $\| L^{\prime}(u)v\|_{L^{2}(\Omega)} \geq \lambda \| v\|_{L^{2}(\Omega)}$  and  $\| L^{\prime}(u) - L^{\prime}(v)\|_{L^{2}(\Omega)} \leq \Lambda \| u - v\|_{L^{2}(\Omega)}$  for  $0 < \lambda \leq \Lambda < \infty$ , then for the Newton update,  $u_{t+1} \gets u_t - L^{\prime}(u_t)^{-1} (L(u_t) - f)$ , with  $u_0 \in L^{2}(\Omega; \mathbb{R})$ , there exists an  $\epsilon > 0$ , such that  $\| u_T - u^\star \|_{L^2(\Omega)} \leq \epsilon$  if  $T \geq \log \left( \frac{1}{\epsilon} \right) / \log \left( \frac{2\lambda^2}{\Lambda \| L(u_0) - f \|_{L^2(\Omega)}} \right)$ .
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].
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 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  $\sim 4\times$  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.
TM is supported by CMU Software Engineering Institute via Department of Defense under contract FA8702-15-D-0002. AP is supported by a grant from the Bosch Center for Artificial Intelligence. ZL gratefully acknowledges the NSF (FAI 2040929 and IIS2211955), 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) via Department of Defense contract FA8702-15-D-0002, for their generous support of ACMI Lab's research. JL is supported in part by NSF award DMS-2012286, and AR is supported in part by NSF awards IIS-2211907, CCF-2238523, Amazon Research Award, and the CMU/PwC DT&I Center.
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
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
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:
$$
u ^ {\star} = \tau G _ {\theta} \left(u ^ {\star}, a\right) + (1 - \tau) u ^ {\star} \tag {12}
$$
where  $\tau$  is a tunable damping factor and  $u^{\star}$  is the fixed point computed using Anderson solver in the forward pass. This step can be repeated  $S$  times. We use  $\tau = 0.5$  and  $S = 1$  for Darcy Flow, and  $\tau = 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.