Deep Equilibrium Based Neural Operators for Steady-State PDEs

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.

Paper

Similar papers

Peer review

Reviewer SM2z6/10 · confidence 5/52023-06-16

Summary

The paper proposes a method combining Fourier Neural Operator (FNO) with the Deep Equilibrium Model (DEQ) for more efficient operator learning under noisy conditions. Overall, the paper is well-written with clear methods and theory, and some effective experimental results. It's a decent piece of work, and I recommend a weak acceptance. However, I believe there are still areas in the paper that could be improved.

Strengths

1. The paper is well-written, with a clear presentation of the method and theoretical underpinnings. 2. The experiments show some positive results, demonstrating the potential utility of the proposed approach.

Weaknesses

1. The paper lacks a thorough literature review. It hardly mentions operator learning methods based on attention, which have demonstrated significant advantages in many areas, such as Navier-Stokes equations and problems on irregular geometric regions. Therefore, I believe the authors should add some references [1,2,3] to cover these works. 2. 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. 3. 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. References 1. Transformer for Partial Differential Equations' Operator Learning (https://arxiv.org/abs/2205.13671) 2. GNOT: A General Neural Operator Transformer for Operator Learning (https://arxiv.org/abs/2302.14376) 3. HT-Net: Hierarchical Transformer based Operator Learning Model for Multiscale PDEs (https://arxiv.org/abs/2210.10890)

Questions

None

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

5: You are absolutely certain about your assessment. You are very familiar with the related work and checked the math/other details carefully.

Soundness

4 excellent

Presentation

4 excellent

Contribution

4 excellent

Limitations

None.

Reviewer Kt2x5/10 · confidence 1/52023-07-06

Summary

The paper proposes two weight-tied neural network architectures for solving steady-state partial differential equations (PDEs) using the universal approximation capabilities of neural networks. The first architecture is a weight-tied version of Fourier Neural Operators (FNO), while the second architecture is a Deep Equilibrium Model (DEQ) that uses black-box root finding algorithms to implicitly train the model. The paper shows that both architectures outperform existing methods on benchmark problems and can be used to learn efficient solvers for PDEs. The contributions of the paper include the introduction of the weight-tied FNO and FNO-DEQ architectures, as well as the demonstration of their effectiveness in solving steady-state PDEs.

Strengths

The main strengths of the paper are: 1. 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. 2. 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. 3. 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. 4. The paper leverages the universal approximation results of FNO to show that FNO-DEQ can universally approximate the solution operator for a wide variety of steady-state PDE families.

Weaknesses

Some potential limitations of the paper are: 1. 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. 2. 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. 3. 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. 4. 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.

Questions

Refer to the weaknesses section.

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

1: Your assessment is an educated guess. The submission is not in your area or the submission was difficult to understand. Math/other details were not carefully checked.

Soundness

2 fair

Presentation

3 good

Contribution

3 good

Limitations

Refer to the weaknesses section.

Reviewer wRdk9/10 · confidence 5/52023-07-06

Summary

This research examined the solution of steady-state Partial Differential Equations (PDEs) using Fourier Neural Operator (FNO) based architecture. The author introduced a fix-point iteration mechanism into the FNO framework, leading to the proposal of weight-tied FNO and FNO Deep Equilibrium (FNO-DEQ) models. Comparative analyses revealed that these newly proposed architectures outperformed the traditional FNO when solving standard Darcy Flow and Navier-Stokes Equations.

Strengths

This paper is well written, providing clear and rigorous mathematical definitions of the problem, its underlying theory, and the proposed solution. The 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). Further 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.

Weaknesses

The paper doesn't provide any loss versus training epochs data, which would allow us to assess whether the training had indeed reached convergence (as well as to compare the speed of convergence). While it's assumed that convergence must have been achieved for the results shown in Tables 2 and 3, the absence of this specific data inhibits a more comprehensive understanding of the model's training process.

Questions

The paper doesn't provide a clear explanation regarding the practical implementation of the Fourier Neural Operator Deep Equilibrium (FNO-DEQ) model. The inclusion of a brief paragraph or section detailing its implementation would greatly enhance the reader's understanding and potentially facilitate the model's broader adoption.

Rating

9: Very Strong Accept: Technically flawless paper with groundbreaking impact on at least one area of AI/ML and excellent impact on multiple areas of AI/ML, with flawless evaluation, resources, and reproducibility, and no unaddressed ethical considerations.

Confidence

5: You are absolutely certain about your assessment. You are very familiar with the related work and checked the math/other details carefully.

Soundness

4 excellent

Presentation

3 good

Contribution

4 excellent

Limitations

The author focused only on the standard Darcy flow and Navier Stokes problem, which the classical FNO already works well. An immediete question is that if the FNO-DEQ can be applied to more challenging PDEs (for example 3D Navier Stokes), to show its effectiveness, where classical FNO struggles.

Reviewer HGbF7/10 · confidence 3/52023-07-13

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$?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

None

Reviewer HGbF2023-08-15

Thanks for the authors' response. While deep equilibrium model is new to me, I believe this is an interesting paper. I feel some theoretical analysis on the convergence would improve the presentation but I think this should be optional. I will increase my score.

Reviewer wRdk2023-08-20

Thanks for the rebuttal. All my questions are well answered.

Authorsrebuttal2023-08-21

Dear reviewer, thank you for your review and feedback! Since today is the last day of the author-reviewer discussion period, please let us know if there are any outstanding questions that we can help answer or clarify.

Authorsrebuttal2023-08-21

Dear reviewer, thank you for your review and feedback! Since today is the last day of the author-reviewer discussion period, please let us know if there are any outstanding questions that we can help answer or clarify.

Program Chairsdecision2023-09-21

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC