Error-Conditioned Neural Solvers
By Haina Jiang, Liam Wang, Peng-Chen Chen, Min Seop Kwak, Seungryong Kim, Brian Bell, Jeong Joon Park
"ENS feeds PDE residual as input to network for iterative correction, outperforming residual-minimization methods on ill-conditioned systems and achieving 10x gains on turbulent flow."
Abstract
Neural surrogate models offer fast approximate mappings from PDE parameters to solutions, but they typically treat solving as a purely statistical task: once trained, they struggle to correct their own constraint violations and extrapolate beyond the training distribution. Recent hybrid methods promote physical correctness by targeting the PDE residual via gradient descent or Gauss--Newton steps, but inherit the compute cost and instability of the underlying classical optimizers. We show, theoretically and empirically, that numerically minimizing the PDE residual can be an unreliable proxy for reconstruction accuracy in ill-conditioned systems, explaining why these methods often do not make accurate predictions despite achieving low residuals. We propose error-conditioned Neural Solvers (ENS), built on a different principle: rather than an optimization target, the PDE residual field is passed as a direct input to the network at each iteration, enabling it to read the spatial structure of its own errors and learn an update policy to iteratively correct its predictions. Across four PDE families, ENS attains the highest prediction accuracy in the large majority of settings, with gains reaching $10\times$ on turbulent Kolmogorov flow, while avoiding the expensive compute cost of hybrid methods. ENS's learned correction policy generalizes under distribution shift, including zero-shot parameter changes and cross-equation transfer, where its relative advantage is largest in the ill-conditioned regimes where residual minimization is least reliable. Project website: https://neuralsolver.github.io/.
Technical Analysis & Implementation
Error-Conditioned Neural Solvers (ENS)§
Core Contribution: The paper identifies that minimizing PDE residuals can be an unreliable proxy for solution accuracy in ill-conditioned systems, and proposes ENS: a learned iterative solver that conditions on the residual field as an input feature rather than as an optimization target.
Methodology:
Let the PDE be $\mathcal{N}[u] = f$ with boundary conditions, and denote the discrete approximation as $\mathcal{N}_h(u_h) = f_h$. Given a current prediction $\hat{u}^{(t)}$, the residual $r^{(t)} = f_h - \mathcal{N}_h(\hat{u}^{(t)})$ is computed. Instead of using gradient descent to minimize $\|r\|$, ENS feeds both $\hat{u}^{(t)}$ and $r^{(t)}$ into a neural network $\Phi_\theta$ to predict a correction:
$$\hat{u}^{(t+1)} = \hat{u}^{(t)} + \Phi_\theta(\hat{u}^{(t)}, r^{(t)})$$
This allows the network to learn correction patterns that exploit the spatial structure of the residual, bypassing the ill-conditioning issues of numerical optimization.
Training: ENS is trained in a supervised manner using ground-truth solutions (e.g., from a classical solver). For each training example, an initial prediction is generated (e.g., from a coarse solver or a naive neural network), and the network is trained to predict the true correction $\Delta u = u^* - \hat{u}^{(t)}$. This can be extended to multiple iterations via unfolded training or imitation learning.
Architecture: The network $\Phi_\theta$ is typically a convolutional neural network (or Fourier neural operator) that processes the stacked input $[\hat{u}, r]$ on the spatial grid. For time-dependent PDEs, temporal convolution or recurrent connections may be added.
Code snippet (PyTorch-style):
import torch
import torch.nn as nn
class ENSBlock(nn.Module):
def __init__(self, in_channels=2, hidden_channels=64, out_channels=1):
super().__init__()
self.conv1 = nn.Conv2d(in_channels, hidden_channels, 3, padding=1)
self.conv2 = nn.Conv2d(hidden_channels, hidden_channels, 3, padding=1)
self.conv3 = nn.Conv2d(hidden_channels, out_channels, 3, padding=1)
def forward(self, u, r):
# u: (batch, 1, H, W), r: (batch, 1, H, W)
x = torch.cat([u, r], dim=1) # (batch, 2, H, W)
x = F.relu(self.conv1(x))
x = F.relu(self.conv2(x))
delta = self.conv3(x)
return u + deltaTheoretical Insight: The paper analyzes the condition number of the residual-to-error map. For ill-conditioned PDEs, a small residual does not guarantee a small error. ENS avoids this by directly learning to map residual to correction, effectively inverting the ill-conditioned operator in a data-driven way.
Results: On 4 PDE families (including Poisson, advection-diffusion, Navier-Stokes, and Kolmogorov flow), ENS consistently outperforms both standard neural surrogates and hybrid residual-minimization methods, especially under distribution shift (e.g., zero-shot parameter changes). On turbulent Kolmogorov flow, ENS achieves up to 10x error reduction compared to baselines, with lower compute cost.
Implications: ENS demonstrates that conditioning on the residual as an input, rather than optimizing it, is a more robust and efficient strategy for iterative PDE correction, particularly for ill-conditioned systems.