Accelerating the convergence of Newton's method for nonlinear elliptic PDEs using Fourier neural operators

TL;DR

This work addresses the slow convergence of Newton's method for nonlinear elliptic PDEs discretized by finite differences by learning a mesh-robust initial guess with Fourier neural operators (FNOs). The authors train an FNO to map PDE data $(\Phi,K)$ to an approximate discrete solution, using a discretization-informed loss that couples data fidelity with the PDE residual, enabling mesh-independent predictions. Across 1D and 2D tests of nonlinear and anisotropic diffusion, the FNO-based initialization substantially reduces the number of Newton iterations and CPU time, especially on coarser grids, and remains effective when solving multiple realizations or time-dependent problems. The approach offers a practical offline/online workflow with potential extensions to unstructured geometries and multi-solution PDEs, and it compares favorably to existing line-search and ML-based initializers in both performance and robustness.

Abstract

It is well known that Newton's method can have trouble converging if the initial guess is too far from the solution. Such a problem particularly occurs when this method is used to solve nonlinear elliptic partial differential equations (PDEs) discretized via finite differences. This work focuses on accelerating Newton's method convergence in this context. We seek to construct a mapping from the parameters of the nonlinear PDE to an approximation of its discrete solution, independently of the mesh resolution. This approximation is then used as an initial guess for Newton's method. To achieve these objectives, we elect to use a Fourier neural operator (FNO). The loss function is the sum of a data term (i.e., the comparison between known solutions and outputs of the FNO) and a physical term (i.e., the residual of the PDE discretization). Numerical results, in one and two dimensions, show that the proposed initial guess accelerates the convergence of Newton's method by a large margin compared to a naive initial guess, especially for highly nonlinear and anisotropic problems, with larger gains on coarse grids.

Figures (11)

• Figure 1: Sketch of an FNO, adapted from li2020fourier. The network has $n_c$ input channels ($n_c = 2$ in 1D and $n_c = 5$ in 2D), which are then extrapolated to $n_p$ dimensions by the extrapolation layer $P$; $L$ Fourier layers follow, whose result is projected back to $\mathbb{R}$ to give an approximation of the PDE solution. In the Fourier layers, $\mathbb{F}$ denotes the fast Fourier transform $\mathtt{FFT}$.
• Figure 2: Score with respect to the number of points in the grid, for different values of the batch size $N_b$. From left to right, we display scores $S_\text{data}$, $S_\text{dis}$ and $S_\text{iter}$. Scores $S_\text{data}$ and $S_\text{dis}$ should be as small as possible; score $S_\text{iter}$ should be as large as possible.
• Figure 3: Examples obtained after training the FNO with the loss function $\mathcal{L}_\text{data}^{L^2}$ from \ref{['eq_loss2']} only. From top to bottom, each series of graphs corresponds to a different example. In the first three columns, from left to right, we display $\Phi(x)$, $K(x)$, the solution $U(x)$ (blue line) and its prediction (orange line). The rightmost column compares the residuals of the discretized PDE computed with the exact solution (blue line) and the predicted solution (orange line).
• Figure 4: Examples obtained after training the FNO with the loss function $\mathcal{L}_\text{data}^{L^2}$ from \ref{['eq_loss2']} and the discretization-informed loss function $\mathcal{L}_\text{dis}$ from \ref{['eq_loss1']}. From top to bottom, the four graphs correspond to a different example. In the first three columns, from left to right, we display $\Phi(x)$, $K(x)$, the solution $U(x)$ (blue line) and its prediction (orange line). The rightmost column compares the residuals of the discretized PDE computed with the exact solution (blue line) and the predicted solution (orange line).
• Figure 5: Statistics of the gains $G_\text{iter}$ in number of iterations (left panels) and $G_\text{CPU}$ in CPU time (right panels) for the 1D problem \ref{['1deq']} with $\alpha_0=2$. From top to bottom, we display the results for meshes with $100$, $200$, $400$ and $600$ points. The gains are grouped in bins of varying sizes, and the number of occurrences in each bin is displayed as a histogram.