Error-Conditioned Neural Solvers Haina Jiang*1
Liam Wang*1
Seungryong Kim2 1
University of Michigan
2
Peng-Chen Chen1
Brian Bell3 KAIST AI
* Equal contribution.
3
Min Seop Kwak1,2
Jeong Joon Park1,†
Los Alamos National Laboratory
† Corresponding author.
arXiv:2606.27354v1 [cs.LG] 25 Jun 2026
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× 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: neuralsolver.github.io
1
Introduction
Neural operators dramatically accelerate PDE solving by replacing expensive numerical solvers with fast, learned mappings from parameters to solution fields [Lu et al., 2021, Li et al., 2020]. Yet they typically treat solving as a purely statistical regression problem, which introduces two fundamental limitations: they typically have no mechanism to assess or correct constraint violations in their own outputs and struggle to extrapolate reliably beyond the training distribution. Traditional numerical methods [Aliabadi, 2020, Ŝolín, 2005] avoid these by construction, but their cost scales rapidly with resolution and dimensionality, making them the primary bottleneck in large-scale simulation applications. Recent hybrid methods incorporate the PDE residual at test time to impose physics-based corrections. These methods share a common structure: the residual is used as an optimization target via gradient descent in PINO [Li et al., 2024], or constraint-manifold projection in DiffusionPDE [Huang et al., 2024] and PCFM [Utkarsh et al., 2025], relying on expensive numerical procedures to translate residual information into solution corrections. Moreover, we show theoretically and empirically that this numerical optimization can expose these hybrid methods to a residual-reconstruction gap: minimizing the PDE residual is often an unreliable proxy for solution accuracy in ill-conditioned systems (Prop. 1), explaining why they tend to achieve low prediction accuracy despite achieving low residuals. We present error-conditioned Neural Solvers (ENS), a new framework built on a different principle: the PDE residual field should be a direct input to the network, instead of an optimization target. At each iteration, ENS receives the current solution prediction alongside its PDE residual field, and learns to produce a correction by directly reading the spatial structure of the error. ENS is trained under reconstruction supervision alone, where the residual is an input signal rather than an objective to minimize, directly sidestepping the residual-reconstruction gap (Prop. 1) that undermines hybrid methods. Applied recurrently, ENS progressively refines the solution without explicit numerical optimization, with computation cost cheaper than existing hybrid methods, and with robustness to initialization that enables generalization to unseen equations.
1
Reconstruction loss PDE residual POSEIDON
Reconstruction loss PDE residual PINO
Reconstruction loss PDE residual FNO
ENS
Reconstruction loss PDE residual
Figure 1: Prediction and Physical accuracy of PDE solvers. For each method we show the reconstruction loss (L1 error) and PDE residual field r(u) = F(u; f ) computed on the same test instance. Feed-forward solvers (FNO, POSEIDON) exhibit large residual fields, indicating systematic violation of the governing equations. While the test-time PDE optimization of PINO reduces the PDE error, it does not translate to accurate solution fields (following Prop. 1). Our ENS produces both near-perfect reconstruction and physical accuracy. Corresponding fields are visualized on the same color scale. We evaluate ENS across four PDE families and increasingly challenging regimes—in-distribution prediction, super-resolution, coefficient extrapolation, and cross-equation transfer. ENS attains the most accurate predictions in most of the tested settings, by up to an order of magnitude on the most illconditioned problems, and achieves low PDE residuals together with low reconstruction error (Fig. 1), despite never directly minimizing the residual. Under distribution shift, ENS’s advantage is largest in the ill-conditioned regimes where minimizing the residual is least reliable: data-driven operators degrade sharply, while residual-minimizing methods remain competitive in benign regimes but are slow and lose their advantage as conditioning worsens. ENS further exhibits initialization-robust convergence: trajectories spanning seven orders of magnitude in initial residual converge to the same residual floor (Fig. 4), a reliability property consistent with its strong extrapolation in ill-conditioned settings. Our contributions are: (i) ENS, a recurrent residual-conditioned Neural Solver achieving up to an order-of-magnitude improvement in reconstruction accuracy across diverse extrapolation settings; (ii) a characterization of ENS, including initialization robustness, ablation experiments, and its design choices; and (iii) a theoretical and empirical analysis of the residual-reconstruction gap (Prop. 1), explaining why existing hybrid methods often struggle to recover accurate solutions despite achieving low residuals.
2
Related Work
Neural operators learn mappings between function spaces for fast prediction across PDE families [Li et al., 2020, Lu et al., 2021, Kovachki et al., 2023], with foundation models extending this to large-scale multi-PDE pretraining [Herde et al., 2024]. These feed-forward models are physics-blind at inference: they receive no feedback about their own prediction error and can incur large constraint violations even when outputs appear plausible. PINNs [Raissi et al., 2019] instead impose physics on a single instance but require per-instance retraining (Fig. 4). To add inference-time physics feedback, hybrid methods use the PDE residual as a correction signal. Optimization-based methods reach a low residual through a separate inference-time procedure—parameter fine-tuning [Li et al., 2024], residual-guided diffusion [Huang et al., 2024, Jacobsen et al., 2025], constraint projection [Christopher et al., 2024, Utkarsh et al., 2025], or zero-shot constrained sampling [Cheng et al., 2025]—while analytical methods apply a closed-form accelerated linearized correction at the prediction [Cao et al., 2023, Jha, 2024, Huang and Perdikaris, 2026]. In every case the residual is an optimization target, not an input to the network; we analyze the resulting cost and reliability in Sec. 3 and Tab. 1. Iterative learned refinement has also been studied without residual conditioning, using noise levels [Lippe et al., 2023] or corrections embedded in a coarse solver [Wei et al., 2026]. Concurrently to our work, PRISMA [Sawhney et al., 2025] moves residual information inside a diffusion neural operator via spectral attention, sharing our motivation. It differs in both mechanism and problem: PRISMA’s residual is computed from noisy observations and conditions a generative denoiser for inverse problems, whereas ENS computes r(k) from the current prediction and feeds it to a corrector that produces an explicit update, closing the feedback loop of Def. 1 for forward solving under distribution shift. Appendix Sec. E provides detailed descriptions and covers generative and unrolled-network methods.
2
3
The Unified View of Hybrid Residual Correction Methods
The hybrid methods in Sec. 2 use the PDE residual as an error signal at inference time. Although they differ in detail, they share a common structure: the residual is used either as a target for gradient-based updates or as the input to a linear solve. Tab. 1 groups them by the classical numerical scheme each one corresponds to. We show that this shared structure can lead to two failure modes: minimizing the residual can be an unreliable proxy for reconstruction accuracy, and the second-order updates they rely on are stable only close to the solution. Table 1: Taxonomy of PDE solving methods by how the residual enters the network. Existing methods use the residual as a training signal, external gradient, external projection, or external linear solve. ENS is the first Neural Solver to receive the residual as a direct input. Inference costs are measured on 20-frame Navier-Stokes scenes. Init. sensitivity indicates whether convergence requires initialization close to u∗ . Method
Baselines
Residual Usage
Inference Cost (s.)
Init. Sensitive
Feed-forward GD-PINNs GD optim. GN projection
FNO, POSEIDON PINO DiffusionPDE PCFM
None (training only) External gradient External gradient External projection
0.0086, 3.5 44 191 251
✗ ✓ ✓ ✓
ENS (Ours)
—
Network input
0.19
✗
First-order. PINO [Li et al., 2024] performs PINNs-style gradient descent through backpropagation on the residual loss: θ̂ ← θ̂ − α∇θ ∥r(û)∥2 . (1) DiffusionPDE [Huang et al., 2024] applies the first-order update as diffusion guidance: x̂t−1 = denoise(xt ) − γ∇xt ∥r(û0 )∥2 .
(2)
In both cases the correction is applied externally, with the residual used as an optimization target. Second-order (Gauss-Newton). PCFM [Utkarsh et al., 2025] applies a penalized Gauss-Newton correction at each flow step, balancing proximity to the current estimate ûτ ′ with soft constraint satisfaction evaluated at the flow-extrapolated point u + γvθ (u, τ ′ ), where γ = 1 − τ ′ : 2
uτ ′ = arg min ∥u − ûτ ′ ∥2 + λ ∥h(u + γvθ (u, τ ′ ))∥ . u
(3)
PCFM’s scheme inherits Gauss-Newton’s initialization sensitivity and, for full PDE residuals, per-step cost grows to O(n2 )–O(n3 )—equivalent to or worse than classical solvers; PCFM’s claimed O(n) cost holds only for low-dimensional constraints (m ≪ n) and does not extend to full PDE residuals. Moreover, the regularization term ∥u − ûτ ′ ∥2 is symptomatic of optimization pitfalls below shared by the hybrid methods. Proposition 1 (Residual-Reconstruction Gap). For a PDE with true solution u⋆ satisfying r(u⋆ ) ≈ 0, minimizing ∥r(u)∥2 does not reliably reduce reconstruction error. Specifically, let Jr (u⋆ ) denote the Jacobian of r at the true solution u⋆ , with smallest singular value σmin (Jr (u⋆ )). For any ε > 0 there exists ũ with ∥r(ũ)∥2 ≤ ε yet: ε ∥ũ − u⋆ ∥2 ≥ . (4) σmin (Jr (u⋆ )) When σmin is small—as in high-wavenumber Helmholtz or low-viscosity Navier-Stokes—residual minimization becomes an arbitrarily unreliable proxy for reconstruction accuracy. This follows from classical conditioning theory; we apply it to explain why existing hybrid methods achieve low residuals yet poor reconstruction for ill-conditioned systems. Geometrically, the zero-residual manifold is nearly flat along the smallest singular direction of Jr (u⋆ ), permitting large deviations from u⋆ at negligible residual cost. This analysis explains why the low PDE-residuals of hybrid methods do not always lead to superior reconstruction accuracy (Tab. 3). 3
k=2
k=4
k = 21
k=7
k = 59
Error-Conditioned Neural Solver
Reconstruction L1 Loss
k=1
Input: f, g
Predictor P
PDE Residual
û(0)
×K
Compute residual r(k) = F(û(k); f) r(k)
û(k)
Corrector C(f, r, û)
δu
δû(k) û(k+1) = û(k) + β·δû(k)
Solution û(K)
Figure 3: ENS inferenceFigure 2: Evolution of reconstruction loss, PDE residual r(û(k) ), and learned correction δ û(k) across iterations. Note that the correction time schematic diaδ û(k) = Cϕ (r(û(k) ), ·) reflects the spatial structure of r(û(k) ) and drives down gram. the PDE residual as well as the reconstruction prediction error. Remark 1 (Initialization Sensitivity of Linearization Methods). Let M bound the curvature ∥D2 r∥op near u⋆ . Newton and Gauss–Newton corrections are contractive only within a basin of radius O(σmin (Jr (u⋆ ))/M ) around the true solution u⋆ [Deuflhard, 2011]. Linearization-based hybrid methods can therefore grow unreliable outside the basin (see Fig. 4a, Tab. 9), e.g., when distribution shift pushes the initialization far from u⋆ . ENS sidesteps this: its learned corrector Cϕ replaces local linearization with a nonlinear map over the full residual field, and training on reconstruction loss avoids the objective-level failure of Prop. 1.
4
Error-Conditioned Neural Solvers
The methods surveyed above share a common structure: whether via gradient descent or analytical linearization, hybrid methods apply residual as an optimization target rather than read directly as a dynamic error signal, while feed-forward neural operators do not see the residual at all. We propose Neural Solvers, a new class of methods with a simple principle: the network should read and correct its own errors, akin to the iterative corrections of classical solvers. Definition 1 (Neural Solver). A Neural Solver consists of an initialization û(0) and a learned corrector Cϕ that iteratively refines the solution û(k) by reading its explicit error signal e(k) : û(k+1) = û(k) + β · Cϕ û(k) , e(k) ; f, g , for k = 0, . . . , K − 1 (5) where e(k) encodes the current prediction’s failure to satisfy a constraint, and β > 0 is a step size. Neural Operators correspond to the special case where K = 0 in Def. 1—they produce a solution estimate û(0) = Nθ (f, g) in a feed-forward pass, with no corrector applied. Without error feedback, the network cannot assess or correct its own prediction at inference time, regardless of its depth or capacity. The choice of error signal e(k) defines a family of Neural Solvers—for instance, the PDE residual field for physics constraint violation, observation mismatch for inverse problems, or symmetry residuals for equivariance constraints. In this work, we instantiate e(k) as the PDE residual field r(k) = F(û(k) ; f ), the most direct available measure of physical constraint violation. We also refer this instantiation as ENS (Error-Conditioned Neural Solver). A predictor network Pθ produces an initial estimate, and a learned corrector Cϕ iteratively refines it by reading the residual field: û(0) = Pθ (f, g), δ û(k) = Cϕ û(k) , r(k) ; f, g , û(k+1) = û(k) + β δ û(k) (6) where r(k) = F(û(k) ; f ) is recomputed at each step. The residual field enters Cϕ as a spatial input channel—the network reads where and how the solution currently violates the governing equations and learns a nonlinear correction policy over it (Fig. 2). The characterizations of ENS are as follow. 4
Figure 4: PDE residual vs. compute time on the nonlinear Helmholtz equation. Top right: predictions at t = 2.8 s. Bottom (a): residual trajectories; all methods except PINNs share the noisy initialization (top-left). Pure Newton and Gauss-Newton diverge (> 1020 ) from Jacobian unreliability far from the solution; PINNs converge very slowly; gradient descent stagnates near its initialization. Recent hybrid methods (DiffusionPDE, PINO-TTOP, PCFM) inherit these limitations, motivating ENS’s learned correction scheme. ENS reduces the residual by over an order of magnitude within the same budget and converges to a common floor from initial residuals spanning 7 orders of magnitude (Bottom (b)). Learned correction over local linearization. Classical Newton iteration computes corrections by solving Jr (û(k) )δu(k) = −r(k) , inverting the local Jacobian at each step. This is unreliable far from u∗ (Remark 1) and prohibitively expensive for large systems. ENS replaces Jacobian inversion with Cϕ , which learns a nonlinear correction policy across training that generalizes across initializations and equation instances without requiring Jacobian computation at inference. Unlike existing methods that treat the residual as an optimization target, ENS treats it as a reference input, sidestepping the residual reconstruction gap (Prop. 1). Training. We run the full correction loop and supervise every intermediate prediction against the ground truth (Fig. 3): K 2 1 X 1 (k) L= û − ugt , (7) K +1 n 2 k=0
(0)
where û = Pθ (f, g) is the predictor output and û(k) for k ≥ 1 are successive corrections; we use K = 5 correction steps in training, giving six supervised predictions. All trainable components, the predictor Pθ and the corrector Cϕ , are optimized jointly under L by backpropagation; we detach the input to each correction step so that gradients do not propagate across iterations. This training scheme exposes Cϕ to the distribution of residuals it actually encounters at inference, including the compounding errors of imperfect intermediate corrections. The step size β is a fixed hyperparameter during training. At inference, Cϕ runs stably for more steps than the five used in training: we iterate until the residual norm ∥r(û(k) )∥2 converges (Fig. 4b). The step size β is reused unchanged across all settings except super-resolution, where the finer evaluation grid alters the scale of the residual and we adopt a smaller β (selected by line search on a small held–out set). Apart from this, only the number of correction steps varies across settings. We instantiate ENS with a modified-FNO backbone for the static equations and a transformer-based backbone [Li et al., 2025] for the turbulent flows. Consistent with Prop. 1, adding a PDE-residual loss term during training did not improve reconstruction accuracy in our settings, so we train on L alone. Remark 2 (Initialization Robustness). ENS drives ∥r(û(k) )∥2 to a consistent floor independently of initialization, as verified empirically in Fig. 4b across initializations spanning seven orders of magnitude in initial residual. Remarkably, although Cϕ is strictly initialized with Pθ during training, it generalizes to diverse corrupted initializations, signaling that ENS learns a robust correction policy rather than a simple 5
statistical mapping. Note that this independence holds for the PDE residual but not the reconstruction loss: by Prop. 1, different initializations may converge to different points on the low-residual manifold. We observe empirically that this initialization robustness is essential for ENS’s strong extrapolation performance, which enables ENS to recover accurate solutions even when the predictor’s initial estimate is far from the true solution under distribution shift.
5
Experiments
Setup. We evaluate ENS on four PDE families—linear and nonlinear Helmholtz (∆u + κ2 u + λu3 = f ), Darcy flow (−∇ · (a(r)∇u) = f ), Poisson (∆u = af ), and Navier–Stokes in vorticity form (∂τ w + v · ∇w = ν∆w + f , ∇ · v = 0) with diagonal-wave forcing for Navier–Stokes and monochromatic forcing for Kolmogorov flow. We train one model per equation on an in-distribution dataset and testing on heldout in-distribution data plus four out-of-distribution regimes: super-resolution (128 → 256), parameter extrapolation (wavenumber, forcing, viscosity), and cross-equation transfer (details in Tab. 7). We report relative L2 reconstruction error and PDE-residual MSE—by Prop. 1 these are not interchangeable in ill-conditioned regimes, where the residual can be driven low while the solution stays inaccurate. Models. ENS pairs a predictor Pθ , which produces the initial estimate, with a recurrent corrector Cϕ , which reads the current prediction together with its PDE residual and outputs a refinement applied over K steps (Sec. 4). Both are instantiated as FNOs with CNN lifting and projection layers for the static PDEs, and as a transformer-based VideoPDE backbone [Li et al., 2025] for the turbulent Navier–Stokes and Kolmogorov flows. Baselines—FNO [Li et al., 2020], PINO [Li et al., 2024] with and without testtime optimization (TTOP), POSEIDON [Herde et al., 2024], DiffusionPDE [Huang et al., 2024], and PCFM [Utkarsh et al., 2025]—use the authors’ official implementations and defaults settings (PINOTTOP to residual convergence in 500/1,000 steps; DiffusionPDE’s two-frame representation; PCFM’s Gauss–Newton projection over 50 flow steps); we omit methods these already compare. The static FNO and PINO baselines keep their original lifting and projection rather than ENS’s CNN variant: adding the CNN does not clearly improve their prediction accuracy with mixed signals (Tab. 12), so they are run in their original configurations. Finally, we follow PINO’s implementations for finite-difference and Fourier-space residual computations. We refer to supplementary document for details. Accuracy. ENS attains the lowest reconstruction error in the large majority of settings while keeping the residual low (Tabs. 2–3; qualitative results in Figs. 5–6). On Helmholtz it is best on both metrics across all four regimes. On Navier-Stokes it is best in-distribution and under viscosity- and forcing-shift, regimes where feed-forward operators collapse (POSEIDON’s error exceeds 0.5 under forcing-shift), and on turbulent Kolmogorov flow, the most ill-conditioned family, its margin is largest across all settings, as predicted in Prop. 1. The baselines display the residual–reconstruction gaps that motivates reporting both metrics: residual-minimizing hybrids (PINO-TTOP, PCFM) reach low residuals but leave reconstruction errors high in ill-conditioned regimes (e.g., wavenumber and viscosity increase). Fig. 8 shows PINOTTOP’s L2 error rising for wavenumber κ = 3 even as its residual falls, while its L2 error slightly decreases for wavenumber κ = 1, when trained on κ = 2. ENS converges for all cases. The one regime where ENS is not the most accurate is NS super-resolution, where it is outperformed by PINO-TTOP, whose spectral (Fourier) parameterization is resolution-invariant by construction, and by POSEIDON, a multiscale transformer that generalizes well across scales, neither property of which ENS’s backbone has. Efficiency. ENS runs in 0.10 s/sample (static) and 0.19 s/sample (NS) in-distribution—roughly 100× and 230× faster than PINO-TTOP and orders of magnitude faster than DiffusionPDE and PCFM—while remaining the most accurate learned method. Out-of-distribution settings need more correction steps and thus more runtime, but ENS stays far faster than the hybrid methods and, on the expensive NS family, below the numerical solver (0.19 s in-distribution vs. 23.3 s/sample). This places it on the accuracy– compute Pareto front against the learned baselines (Fig. 13, 14); the TTOP- and diffusion-based hybrids sit off the front, spending large compute on residuals that do not translate into accuracy. Ablations. Controlled studies isolate the source of ENS’s gains; full details are in App. F. Conditioning, not computation. With architecture and step count fixed, replacing the residual input with a zero field stalls both losses, whereas ENS reduces both (Fig. 7): the gains come from reading the residual, not added iterations. Not the backbone. On ENS’s own backbone, the operator baselines and their PINO and TTOP 6
variants still trail ENS in reconstruction error in every regime (Tabs. 12, 13, 14); TTOP instead reproduces the gap of Prop. 1, driving the residual low while reconstruction stays high, mild on Navier–Stokes but severe on Kolmogorov flow, where its L2 error is about 50× ENS’s. Input, not objective. Conditioning on the physics-loss gradient lowers the residual but stalls reconstruction (Fig. 15), and a residual training term did not help. Backbone needs expressivity. The residual input contains high-frequency information that must be representable by the backbone; on Helmholtz a pure FNO diverges while its CNN-augmented variants converge (Tab. 11). Transformer backbones are also expressive enough to converge. A learned policy. Although trained with five correction steps, ENS keeps improving past 100 iteration at inference and stops when the residual plateaus (Fig. 8, 11), indicating a learned correction rather than a fixed scheduled mapping. Table 2: Forward prediction results for the Helmholtz problem. Time refers to the mean per-sample inference latency (In-distribution). Bold indicates the best result, and underlining indicates the second-best. Cross-equation refers to inferencing on non-linear Helmholtz while trained on Poisson. PINO(TTOP) and PCFM attain low PDE residuals but high L2 error, showing the residual–reconstruction gap (Prop. 1). In-distribution
Wavenumber-shift L2 Rel
PDE Res
Cross-equation L2 Rel
PDE Res
Super-resolution
Method
Time(s)
L2 Rel
PDE Res
L2 Rel
PDE Res
ENS (ours) FNO PINO PINO(TTOP) POSEIDON DiffusionPDE PCFM
0.1 0.0065 0.0065 10 0.18 191 251
3.04e-03 1.20e-02 4.19e-02 6.44e-02 2.70e-02 9.85e-02 1.08e-01
3.90e-01 2.59e-02 5.19e-01 8.73e-02 6.54e+01 1.82e-02 4.52e-01 4.21e+02 2.39e-01 3.80e+02 5.21e-01 1.33e+03 9.06e-02 1.28e+04 4.98e+01 2.43e-01 6.02e+01 5.27e-01 6.42e+03 1.29e-01 2.24e+03 9.65e+00 1.61e-01 1.28e+01 4.08e-01 3.31e+01 1.68e-01 4.17e+01 2.57e+04 2.40e-01 2.35e+04 4.80e-01 3.27e+04 5.17e-01 6.02e+04 2.19e+04 2.69e-01 2.35e+04 3.90e-01 4.21e+05 7.72e-01 2.29e+04 OOM OOM 3.91e+00 1.78e-01 7.02e+00 6.30e-01 1.62e+01
Table 3: Forward prediction results for the Navier-Stokes (20 frames). Time refers to the mean per-sample inference latency (In-distribution). For ENS, distribution-shift requires more iterations and thus longer inference latency. DiffusionPDE only predicts the first and last frames, thus unable to compute PDE residuals. PCFM goes out of memory at 2562 resolution. The gap is mild here: PINO(TTOP) reaches the lowest residuals yet trails ENS in L2 (Prop. 1) except super-resolution. In-distribution PDE Res
Viscosity-shift L2 Rel
PDE Res
Forcing-shift L2 Rel
PDE Res
Super-resolution
Method
Time(s)
L2 Rel
L2 Rel
PDE Res
ENS (ours) FNO PINO PINO(TTOP) POSEIDON DiffusionPDE PCFM
0.19 0.0086 0.0086 43.7 3.5 229.22 192.01
2.22e-03 4.69e-05 3.05e-02 1.12e-02 5.31e-03 7.40e-04 3.98e-02 9.11e-02 3.54e-02 8.24e-04 1.36e-01 1.18e-02 5.57e-01 3.09e-02 7.51e-02 8.33e-03 3.57e-02 8.45e-04 1.37e-01 1.20e-02 5.57e-01 3.10e-02 7.59e-02 8.48e-03 1.17e-02 2.33e-05 5.42e-02 3.79e-04 1.45e-02 4.39e-05 1.73e-02 7.69e-05 4.73e-03 6.69e-05 7.45e-02 1.17e-03 5.66e-01 2.07e-02 8.83e-03 4.45e-04 5.59e-01 N/A 6.14e-01 N/A 7.52e-01 N/A 6.83e-01 N/A 1.34e+00 4.79e+02 7.18e-01 8.33e-01 9.32e-01 1.00e+04 OOM OOM
Table 4: Forward prediction results for the Kolmogorov flow (40 frames). DiffusionPDE only predicts the first and last frames, thus unable to compute PDE residuals. PCFM goes out of memory in our 40-frames settings. Note the severe residual-reconstruction gap (Prop. 1) of PINO(TTOP). In-distribution
Viscosity-shift
Forcing-shift
Super-resolution
Method
Time(s)
L2 Rel
PDE Res
L2 Rel
PDE Res
L2 Rel
PDE Res
L2 Rel
PDE Res
ENS(Ours) FNO PINO PINO(TTOP) POSEIDON DiffusionPDE PCFM
0.43 0.0093 0.0093 72.7 6.97 229.27 N/A
8.08e-03 5.41e-02 8.05e-02 2.11e-01 5.62e-01 8.19e-01 OOM
2.50e-01 1.22e+00 3.49e+00 4.83e-02 7.35e+00 N/A OOM
3.42e-02 3.54e-01 3.13e-01 4.66e-02 6.16e-01 7.37e-01 OOM
8.05e-01 2.08e+00 3.51e+00 3.18e-03 1.19e+01 N/A OOM
2.06e-02 6.66e-01 7.04e-01 2.06e-01 1.04e+00 8.19e-01 OOM
3.72e-01 1.15e+01 1.22e+01 2.33e-02 1.97e+01 N/A OOM
1.98e-02 1.60e-01 1.32e-01 4.45e+00 6.65e-01 9.45e-01 OOM
5.55e+00 1.49e+00 4.85e+00 7.74e-02 1.53e+01 N/A OOM
7
Input
Ground truth
ENS
FNO
PINO
PINO(TTOP)
POSEIDON
DiffusionPDE
PCFM
(a) Helmholtz à Poisson
(b) Poisson à Nonlinear Helmholtz
Out-of-Memory x
(c) 𝟏𝟐𝟖× 𝟏𝟐𝟖 à 𝟐𝟓𝟔 × 𝟐𝟓𝟔 (Helmholtz)
Out-of-Memory x
Diverged x
(d) Darcy (f = 100) à Darcy (f = 200)
Diverged x
Figure 5: Qualitative comparison on static PDEs. Reconstructed fields (top) and point-wise L1 errors (bottom) across three OOD regimes: cross-equation (HZ → PS, PS → nonlinear HZ); resolution shift(HZ, 128×128 → 256×256); and forcing shift (Darcy, f = 100 → f = 200).
6
Conclusion and Limitations
We introduced Error-Conditioned Neural Solvers (ENS), a recurrent framework that treats the PDE residual field as a direct input rather than an external optimization objective, and learns a nonlinear correction policy over the residual field. Compared to existing hybrid methods, which apply first- or second-order numerical schemes to the residual, ENS attains the highest prediction accuracy in the large majority of settings across diverse PDE families. We further show that minimizing the PDE residual is an unreliable proxy for reconstruction accuracy in ill-conditioned systems, which helps explain the gap between low residuals and suboptimal reconstructions in existing hybrid methods. ENS’s advantage is pronounced on distribution-shift tasks, including changed coefficients and cross-equation transfer, suggesting a path toward a general neural solver for unseen systems. Our experiments are limited to relatively simple 2D systems and assume the governing equations are known at inference. Because ENS reads the residual as an input rather than minimizing it, we hypothesize that it is more forgiving of imperfect residuals from discretization error or partially known equations, which would mislead an optimization-based corrector. Extending ENS to three-dimensional problems and to settings with real, noisy observations is an exciting direction for future work.
8
Input
Ground truth
ENS
FNO
PINO
PINO(TTOP)
POSEIDON
DiffusionPDE
PCFM
(a) Viscosity-shift (Navier-Stokes)
(b) Forcing-shift (Navier-Stokes)
Out-of-Memory x
(c) Forcing-shift (Kolmogorov Flow)
Out-of-Memory x
Out-of-Memory x
(d) !"#× !"# à "%& × "%& (Kolmogorov Flow)
Out-of-Memory x
Figure 6: Qualitative comparison on dynamic PDEs. Reconstructed fields (top) and point-wise L1 errors (bottom) across three OOD regimes: viscosity shift (NS, ν=10−4 → 10−5 ); forcing shift (NS, 2π → 4π; KF, 8π → 10π); and resolution shift (KF, 128×128 → 256×256).
Kolmogorov Flow (Rec. loss) Kolmogorov Flow (PDE residual)
Helmholtz (Rec. loss) Helmholtz (PDE residual)
PDE residual (k = 3) PDE residual (k = 1)
ENS (Res. feedback) Ablation (Zero-field)
101
Loss
Mean Squared Error (MSE)
103
10 1 10 3 10 5 0
2
4
6
8
Iteration Step k
10
12
ENS TTOP
105 104 103 102 101 100 10 1 10 2 100
14
L2 loss (k = 3) L2 loss (k = 1)
101
102 Iteration Step k
103
Figure 7: Effectiveness of residual conditioning. Figure 8: PDE residual and relative L2 error We compare the loss trajectories for ENS (solid lines) during test-time refinement on Helmholtz against the ablation baseline (dashed lines) where (κ = 1, 3). ENS (solid) reduces both within tens zero-field is conditioned instead of the PDE resid- of steps; PINO-TTOP (dashed) needs far more itual field. The results show that ENS with residual erations and, for κ = 3, its L2 error rises even as conditioning is able to reduce both losses, whereas its residual falls—the residual–reconstruction gap: the ablation baseline fails to do so. when σmin is small, lowering the residual may not improve the solution.
9
References F. M. Aliabadi. Boundary element methods. In Encyclopedia of continuum mechanics, pages 182–193. Springer, 2020. M. Andrychowicz, M. Denil, S. Gomez, M. W. Hoffman, D. Pfau, T. Schaul, B. Shillingford, and N. De Freitas. Learning to learn by gradient descent by gradient descent. Advances in neural information processing systems, 29, 2016. S. Bai, J. Z. Kolter, and V. Koltun. Deep equilibrium models. Advances in neural information processing systems, 32, 2019. L. Cao, T. O’Leary-Roseberry, P. K. Jha, J. T. Oden, and O. Ghattas. Residual-based error correction for neural operator accelerated infinite-dimensional bayesian inverse problems. Journal of Computational Physics, 486: 112104, 2023. C. Cheng, B. Han, D. Maddix, A. F. Ansari, A. Stuart, M. W. Mahoney, and B. Wang. Gradient-free generation for hard-constrained systems. In International Conference on Learning Representations, volume 2025, pages 100510–100539, 2025. J. K. Christopher, S. Baek, and F. Fioretto. Constrained synthesis with projected diffusion models. Advances in Neural Information Processing Systems, 37:89307–89333, 2024. P. Deuflhard. Newton Methods for Nonlinear Problems: Affine Invariance and Adaptive Algorithms, volume 35 of Springer Series in Computational Mathematics. Springer, Berlin, Heidelberg, 2011. M. Herde, B. Raonić, T. Rohner, R. Käppeli, R. Molinaro, E. De Bezenac, and S. Mishra. Poseidon: Efficient foundation models for pdes. Advances in Neural Information Processing Systems, 37:72525–72624, 2024. J. Ho, A. Jain, and P. Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840–6851, 2020. E. Hoogeboom, V. G. Satorras, C. Vignac, and M. Welling. Equivariant diffusion for molecule generation in 3d. In International conference on machine learning, pages 8867–8887. PMLR, 2022. J. Huang, G. Yang, Z. Wang, and J. J. Park. Diffusionpde: Generative pde-solving under partial observation. Advances in Neural Information Processing Systems, 37:130291–130323, 2024. X. Huang and P. Perdikaris. Physicscorrect: A training-free approach for stable neural pde simulations. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 40, pages 22057–22065, 2026. C. Jacobsen, Y. Zhuang, and K. Duraisamy. Cocogen: Physically consistent and conditioned score-based generative models for forward and inverse problems. SIAM Journal on Scientific Computing, 47(2):C399–C425, 2025. P. K. Jha. Residual-based error corrector operator to enhance accuracy and reliability of neural operator surrogates of nonlinear variational boundary-value problems. Computer Methods in Applied Mechanics and Engineering, 419:116595, 2024. T. Karras, M. Aittala, T. Aila, and S. Laine. Elucidating the design space of diffusion-based generative models. Advances in neural information processing systems, 35:26565–26577, 2022. N. B. Kovachki, Z. Li, B. Liu, K. Azizzadenesheli, K. Bhattacharya, A. M. Stuart, and A. Anandkumar. Neural operator: Learning maps between function spaces with applications to pdes. J. Mach. Learn. Res., 24:89:1–89:97, 2023. URL https://jmlr.org/papers/v24/21-1524.html. E. Li, Z. Wang, J. Huang, and J. J. Park. Videopde: Unified generative pde solving via video inpainting diffusion models. arXiv preprint arXiv:2506.13754, 2025. Z. Li, N. Kovachki, K. Azizzadenesheli, B. Liu, K. Bhattacharya, A. Stuart, and A. Anandkumar. Fourier neural operator for parametric partial differential equations. arXiv preprint arXiv:2010.08895, 2020. Z. Li, H. Zheng, N. Kovachki, D. Jin, H. Chen, B. Liu, K. Azizzadenesheli, and A. Anandkumar. Physics-informed neural operator for learning partial differential equations. ACM/IMS Journal of Data Science, 1(3):1–27, 2024. P. Lippe, B. Veeling, P. Perdikaris, R. Turner, and J. Brandstetter. Pde-refiner: Achieving accurate long rollouts with neural pde solvers. Advances in Neural Information Processing Systems, 36:67398–67433, 2023. L. Lu, P. Jin, G. Pang, Z. Zhang, and G. E. Karniadakis. Learning nonlinear operators via deeponet based on the universal approximation theorem of operators. Nature machine intelligence, 3(3):218–229, 2021.
10
I. Price, A. Sanchez-Gonzalez, F. Alet, T. R. Andersson, A. El-Kadi, D. Masters, T. Ewalds, J. Stott, S. Mohamed, P. Battaglia, et al. Gencast: Diffusion-based ensemble forecasting for medium-range weather. arXiv preprint arXiv:2312.15796, 2023. M. Raissi, P. Perdikaris, and G. E. Karniadakis. Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations. Journal of Computational physics, 378:686–707, 2019. Y. Saad. Iterative methods for sparse linear systems. SIAM, 2003. C. Saharia, W. Chan, H. Chang, C. A. Lee, J. Ho, T. Salimans, D. J. Fleet, and M. Norouzi. Palette: Image-to-image diffusion models, 2022. URL https://arxiv.org/abs/2111.05826. M. Sawhney, A. Neog, M. Khurana, and A. Karpatne. Beyond loss guidance: Using pde residuals as spectral attention in diffusion neural operators. arXiv preprint arXiv:2512.01370, 2025. A. Schneuing, C. Harris, Y. Du, K. Didi, A. Jamasb, I. Igashov, W. Du, C. Gomes, T. L. Blundell, P. Lio, et al. Structure-based drug design with equivariant diffusion models. Nature Computational Science, 4(12):899–909, 2024. P. Ŝolín. Partial differential equations and the finite element method. John Wiley & Sons, 2005. Y. Song, J. Sohl-Dickstein, D. P. Kingma, A. Kumar, S. Ermon, and B. Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.13456, 2020. L. N. Trefethen, A. E. Trefethen, S. C. Reddy, and T. A. Driscoll. Hydrodynamic stability without eigenvalues. Science, 261(5121):578–584, 1993. U. Utkarsh, P. Cai, A. Edelman, R. Gomez-Bombarelli, and C. V. Rackauckas. Physics-constrained flow matching: Sampling generative models with hard constraints. arXiv preprint arXiv:2506.04171, 2025. J. L. Watson, D. Juergens, N. R. Bennett, B. L. Trippe, J. Yim, H. E. Eisenach, W. Ahern, A. J. Borst, R. J. Ragotte, L. F. Milles, et al. De novo design of protein structure and function with rfdiffusion. Nature, 620 (7976):1089–1100, 2023. H. Wei, A. Franz, B. List, and N. Thuerey. Inc: An indirect neural corrector for auto-regressive hybrid pde solvers. Advances in Neural Information Processing Systems, 38:110182–110216, 2026. T. Xie, X. Fu, O.-E. Ganea, R. Barzilay, and T. Jaakkola. Crystal diffusion variational autoencoder for periodic material generation. arXiv preprint arXiv:2110.06197, 2021.
11
A
Theoretical Analysis: Proofs and Extended Discussion
This appendix provides a proof for the proposition in Section 4, plus extended discussion of the relationship between ENS and classical numerical methods. Throughout, n denotes the total degrees of freedom of the discretized system (n = N 2 for a 2D grid of resolution N × N , with analogous bounds for 3D).
A.1
Proof of Prop. 1: Residual-Reconstruction Gap
The solution manifold {u : r(u) ≈ 0} changes slowly along the direction of the smallest singular vector of Jr (ugt ) — one can move far from ugt while the residual barely changes. The proof constructs an explicit perturbation exploiting this direction to create a lower bound on the worst-case reconstruction error. Linear case. Let F(u) = Au − f with Augt = f . Let v be the right singular vector of A corresponding to σmin , with ∥v∥2 = 1, and define ũ = ugt + αv. Since A = U ΣV ∗ : ∥Av∥2 = ∥Σemin ∥2 = σmin
(8)
using V ∗ v = emin and ∥U · ∥2 = ∥ · ∥2 . The residual and reconstruction error are then exact: ∥r(ũ)∥2 = ασmin ,
∥ũ − ugt ∥2 = α.
(9)
Setting α = ε/σmin gives ∥r(ũ)∥2 = ε and ∥ũ − ugt ∥2 = ε/σmin , a lower bound on the worst-case reconstruction error. Nonlinear case. Let F be twice Fréchet-differentiable near ugt with r(ugt ) = 0 and ∥D2 F∥op ≤ M on Bδ (ugt ). Let v be the right singular vector of Jr (ugt ) corresponding to σmin , ∥v∥2 = 1, and define ũ = ugt + αv. By Taylor expansion: ∥r(ũ)∥2 ≤ ασmin + 12 M α2 .
(10)
2 2 Set α = ε/(2σmin ) for any ε ∈ (0, 4σmin /M ). Then 12 M α2 = ε2 M/(8σmin ) < ε/2, giving:
∥r(ũ)∥2 < ε
∥ũ − ugt ∥2 =
but
ε , 2σmin (Jr (ugt ))
(11)
a lower bound on the worst-case reconstruction error. Worst-case conditioning lower bound. Let r : RN → RN denote a discretized residual map with r(u⋆ ) = 0, and write J⋆ = Jr (u⋆ ). Assume J⋆ is full rank with singular values σmax (J⋆ ) = L,
σmin (J⋆ ) = µ,
χ := cond(J⋆ ) = L/µ.
After rescaling the residual, take L = 1, so that µ = χ−1 . For an iterate u = u⋆ + e in the local linearized regime, r(u) = J⋆ e + o(∥e∥). Suppose the finite-precision linear algebra induces perturbations at effective relative scale εN := cN umach , where cN captures dimension-dependent accumulation of roundoff errors. The lower bounds follow from the least observable singular direction. Let vmin , wmin be right and left singular vectors satisfying J⋆ vmin = χ−1 wmin , J⋆−1 wmin = χvmin . For a Newton correction, the exact linearized step solves J⋆ p = −r(u), and hence gives p = −e. However, perturbing the right-hand side by ζ = ηwmin yields J⋆ pcomp = −r(u) − ζ,
pcomp = −e − ηχvmin .
Thus the post-correction reconstruction error satisfies ∥u + pcomp − u⋆ ∥2 = ηχ. 12
Taking η = εN ∥e∥2 gives the worst-case lower bound ∥u + pcomp − u⋆ ∥2 ≥ εN χ. ∥e∥2 If the perturbation instead acts on the Jacobian, the same extremal two-dimensional model A = diag(1, χ−1 ) with perturbation in the weak coordinate gives ∥u + pcomp − u⋆ ∥2 εN χ ≥ , ∥e∥2 1 − εN χ
εN χ < 1.
For Gauss–Newton implemented through the normal equations, J⋆⊤ J⋆ p = −J⋆⊤ r(u), the coefficient matrix H = J⋆⊤ J⋆ has condition number cond(H) = cond(J⋆ )2 . Applying the same extremal construction to H = diag(1, χ−2 ) and perturbing the weak coordinate by relative size εN gives εN χ2 ∥u + pGN,comp − u⋆ ∥2 , ≥ ∥e∥2 1 − εN χ2
εN χ2 < 1.
When εN χ2 ≥ 1, the perturbed normal-equation matrix may become singular or nearly singular in the least observable direction, so no nontrivial uniform reconstruction guarantee follows from the normal-equation solve alone. For a 128 × 128 image, N = 16384; using double precision and the conservative dense-solve scaling cN = N gives εN ≈ 16384 · 2.2 × 10−16 ≈ 3.6 × 10−12 , so the worst-case amplification scales as 3.6 × 10−12 χ for Newton and 3.6 × 10−12 χ2 for Gauss–Newton normal equations, with the latter becoming unreliable at the scale χ ≳ (3.6 × 10−12 )−1/2 ≈ 5.3 × 105 . This is easily encountered in practice. Remark 3. When a PDE is conditioned such that σmin → 0, this reconstruction error can be arbitrarily large. Remark 4 (Instantiation for Helmholtz and Navier-Stokes). For F(u; f ) = ∆u + κ2 u + λu3 − f , the Jacobian Jr (ugt ) = ∆ + κ2 I + 3λu2gt I is self-adjoint on L2 (Ω), so its singular values coincide with the absolute values of its eigenvalues. As κ grows, near-resonances between ∆ + κ2 I and the spatially varying term 3λu2gt I drive σmin (Jr ) → 0 without requiring exact cancellation. The second derivative bound is M = 6λ∥ugt ∥∞ . Together these place large-κ Helmholtz in the ill-conditioned regime where the gap 1/(2σmin ) is large, consistent with the pronounced residual-reconstruction discrepancy in Tab. 10. For Navier-Stokes, Jr (ugt ) = νJ∆ − Jadv (ugt ). As ν → 0 the viscous term vanishes and σmin (Jr ) is controlled by the advection operator alone, which is known to be ill-conditioned for smooth complex flows [Trefethen et al., 1993]. In both cases, methods minimizing ∥r(u)∥2 at test time navigate along these flat directions, achieving low residual without commensurate reconstruction accuracy. ENS, trained on reconstruction loss, optimizes the correct objective directly.
13
A.2
Discussion of Remark 2: Why ENS is Initialization-Robust
Remark 2 is an empirical observation rather than a formal theorem. Residual convergence is initialization-independent. To our surprise, although our ENS network Cϕ is trained only on the FNO-based (from Pθ ) initialization, our experiment (Fig. 4) shows convergence to the almost identical error floor from diverse initializations. We hypothesize that, instead of learning the pure statistical mapping from input fields to output fields, our error-conditioning forces Cϕ to focus on reducing the residuals until the discretization limit, provided in diverse patterns during training, allowing robustness over various residual patterns and magnitude. This capability of ENS leads to its exceptional extrapolation ability and contrasts with GD and GN, whose correction directions depend on Jr (û(k) ), which could get arbitrarily unreliable far from ugt . Reconstruction loss is not initialization-independent. All initializations converge to the same residual floor but not necessarily to the same point on the low-residual manifold {u : r(u) ≈ 0}. By Prop. 1, points on this manifold can lie at distance up to ε/σmin from ugt , so better initializations yield lower reconstruction error despite reaching the same residual floor.
A.3
Per-Step Complexity
Gradient Descent. Computing ∇u ∥r(u)∥2 = 2Jr (u)⊤ r(u) via reverse-mode automatic differentiation requires one forward pass (O(n) for sparse F) and one backward pass (O(n)). Total: O(n), and O(n log n) for a transformer architecture. Gauss-Newton. The update (Jr⊤ Jr )δu = −Jr⊤ r requires: 1. Forming Jr (u) ∈ Rn×n : O(n2 ) via n backward passes 2. Forming Jr⊤ Jr ∈ Rn×n : O(n3 ) naively, O(n2 ) exploiting sparsity 3. Solving the n × n linear system: O(n3 ) directly, O(n2 ) via Krylov methods with sparse Jr Total: O(n2 ) with Krylov solvers, O(n3 ) naively. Classical solvers (multigrid, CG) exploit the hierarchical sparsity of F to achieve O(n) to O(n log n) for the same system. Gauss-Newton is therefore O(n) to O(n2 ) times more expensive per step than classical solvers, while providing weaker convergence guarantees. PCFM’s Schur complement. PCFM reduces cost by solving only the m × m Gram matrix JJ ⊤ , costing O(m3 + mn). For m ≪ n (e.g., m = 1 for scalar mass conservation), this is O(n). For full PDE residual enforcement where m = n, this recovers O(n3 ). PCFM’s claimed O(n) efficiency therefore does not extend to general PDE residuals. ENS. One residual evaluation F(û(k) ; f ) costs O(n) for sparse F. One forward pass of Cϕ costs O(n) for a U-Net or FNO backbone and O(n log n) for a transformer-based neural operator. Total: O(n) to O(n log n) per refinement step — comparable to a single linear Krylov solve, and O(n) to O(n2 ) times cheaper per step than Gauss-Newton for the same residual dimension, and comparable to black-box feedforward neural operators like FNO.
B
Implementation Details
Architecture. The predictor Pθ maps the PDE inputs (f, g) to an initial estimate û(0) in a single forward pass. At each step k, the corrector Cϕ takes (f, g, û(k) , r(k) ), concatenated along the channel dimension, and outputs a correction field δ û(k) at the same spatial resolution. The predictor and corrector use the same backbone, which depends on the problem family. For the static equations (Helmholtz, Darcy, Poisson) we use a Fourier Neural Operator (FNO) with four Fourier layers, 64 hidden channels, and 20 Fourier modes per spatial dimension; the standard FNO lifting and projection modules are replaced with 3-layer CNNs to better preserve local spatial structure (5-layer CNNs for Darcy). For the turbulent flows (Navier–Stokes and Kolmogorov), both networks instead use a transformer-based VideoPDE backbone [Li et al., 2025], which better captures the spatio-temporal structure of these systems. Training. Both networks are trained jointly end-to-end by unrolling the full K-step correction loop. At each training step: (i) Pθ produces û(0) , (ii) residuals r(k) are recomputed at each step as described in
14
Sec. C.2, (iii) all K intermediate predictions are supervised: Ldata =
K−1 B 2 1 X 1 X (k) ûi − ugt,i , K Bn i=1 2
(12)
k=0
where B is the batch size. Supervising all intermediate steps rather than only the final output provides dense gradient signal to both networks and exposes Cϕ to the full distribution of residual fields it encounters at inference time, including compounding errors from imperfect early corrections. We use K = 5 unrolled steps during training. Both networks are optimised jointly with AdamW (lr = 1 × 10−4 , weight decay 10−4 ). All experiments are conducted on a single RTX 6000 Ada with a batch size B = 32. We set the correction step size to β = 0.05 for static PDEs and Kolmogorov flow, and β = 0.1 for Navier-Stokes. Inference. At inference time, the loop runs for Kinf steps determined by the computational budget and the amount of distribution shift, with r(k) recomputed from the updated solution at every step. For out-of-distribution experiments we increase k to address the slower convergence. Specifically, for the Helmholtz equation, we increase k from 15 in-distribution to 80 for wave-number extrapolation and to 120 for super-resolution. For Navier–Stokes, we increase k from 5 to 15 for viscosity shift and to 10 for forcing shift. For Kolmogorov flow, we increase k from 12 to 15 for both viscosity and forcing shifts, and to 90 for super-resolution. Performance continues to improve beyond k = 5 until the residual floor εC is reached. The corrector seems robust to imperfect initialization quality since it was trained with simulated residuals across the full range of correction stages.
B.1
Diffusion ENS (Extension)
As an exploratory extension of the ENS principle to generative solvers, we augment a base EDM model Dbase [Karras et al., 2022], pretrained from scratch with the standard EDM denoising objective and then freeze, with a trainable error-conditioned model DENS of the same architecture with an added meta-conditioning input. At each denoising step, Dbase produces a base estimate ûbase , whose residual is 0 concatenated with the estimate to form the meta-feature: mt = F(ûbase ; f ), ûbase (13) 0 0 DENS then produces a residual-corrected prediction: ûENS = DENS (ut , σ, [y, mt ]) 0
(14)
replacing the standard x̂0 estimate throughout the EDM sampling loop. Here y is a sparse observation of the input field, comprising 500 random sensor measurements of the source term f . mt is recomputed at each sampling step, providing dynamically adapted residual corrections throughout the trajectory. Architecture and Training. Both Dbase and DENS share the same EDM SongUNet backbone (∼55M parameters), conditioned on σ via Fourier feature embeddings. The conditioning fields [y, mt ] are concatenated as additional input channels and incorporated additively through a conditioning encoder at each block. Dbase is pretrained with the standard EDM denoising objective and then frozen; only DENS is trained—a single forward pass through the frozen Dbase per step suffices to obtain ûbase and the resulting 0 residual, with no simulation loop required. We train DENS with AdamW (lr = 2 × 10−4 , weight decay 10−4 ) for 300k steps. Evaluation Tasks and Protocols. We train and evaluate Diffusion ENS on three PDEs: Helmholtz (κ = 1), Darcy flow (binary permeability), and Poisson, all discretized on 128 × 128 grids, with 50,000 training and 100 test samples per PDE. We report relative L2 error and PDE residual MSE, computed on the predicted field via interior second-order central differences (h = 1/127). Inference. Sampling uses 50 Heun ODE steps (EDM schedule, σmin = 0.002, σmax = 80, ρ = 7) with no gradient computation at inference time. This contrasts with DiffusionPDE [Huang et al., 2024], which requires 2,000 steps with per-step gradient guidance, leading to highly expensive compute cost. Baselines. We compare Diffusion ENS against representative baselines: the original DiffusionPDE model [Huang et al., 2024], PCFM [Utkarsh et al., 2025], and a Palette-style [Saharia et al., 2022] observation-conditioned EDM model. All methods are conditioned on the same 500 random sensor observations of the input field and evaluated on the same 100 test instances per PDE. DiffENS+GN augments DiffENS with a single Gauss– Newton projection of the final sample onto the PDE residual: one least-norm step λ = (Jr Jr⊤ + εI)−1 r(û), 15
Table 5: Diffusion results for the sparse observation scenarios. Instead of providing the full source term f , we subsample and provide partial observation, on which generative models excel Huang et al. [2024]. In addition to our diffusion extension of ENS (DiffENS), we apply one step of Gauss Newton projection on the sampled solution (DiffENS+GN). Time indicates the mean sampling duration in seconds on an NVIDIA RTX 6000 Ada GPU. Helmholtz Method
Time (s)
DiffENS DiffENS+GN DiffusionPDE PCFM Palette
5.8 6.5 215 191 2.8
Darcy
Poisson
L2
PDE p-in (C)
L2
PDE p-in (C)
L2
PDE p-in (C)
0.0214 0.0214 0.1483 0.1557 0.0237
1.68e+02 1.54e-04 3.17e+05 3.13e-02 6.88e+03
0.0298 0.0312 0.0597 0.0833 0.0318
4.56e+00 5.23e+01 2.29e+03 1.65e+02 1.24e+03
0.0212 0.0213 0.0713 0.1619 0.0237
1.57e+02 1.22e-03 5.93e+04 5.74e-03 5.44e+03
û ← û − Jr⊤ λ with ε = 10−6 . DiffusionPDE [Huang et al., 2024] uses the original authors’ pretrained checkpoints with a 2,000-step Heun sampler and per-step DPS guidance on the observation and PDE residual. PCFM [Utkarsh et al., 2025] uses a U-Net backbone trained from scratch, sampled with 50 flow-matching steps and a per-step Gauss–Newton projection to minimize the PDE residual. Palette denotes a Palette-style observation-conditioned EDM model on the same backbone, sampled with 50 Heun steps without residual conditioning. Results are reported in Tab. 5.
C
Experiment Details
C.1
Dataset Preparation
All static PDE datasets (Helmholtz, Poisson, Darcy) are generated by sampling input fields from Gaussian Random Fields (GRFs) and solving the resulting sparse linear systems using standard sparse direct solvers; we use Newton’s method for non-linear Helmholtz. Navier-Stokes initial vorticity fields are sampled from GRFs; spatial derivatives are computed in Fourier space with dealiasing, and time integration uses a Crank-Nicolson scheme for the viscous term with explicit treatment of the nonlinear advection. For feed-forward ENS, each static PDE and Navier-Stokes dataset contains 1,000 training samples, while Kolmogorov flow dataset contains 5,000 training samples. For Diffusion ENS, each dataset contains 50,000 training samples. All datasets contain 100 test samples. Data generation times are reported in Tab. 6. The detailed data generation protocols for both training and in-distribution test data are described below, while the configurations for extrapolation test data are summarized in Tab. 7. Table 6: Data generation time per sample across PDE families. Static PDE datasets are generated on CPU, while the time-dependent Navier-Stokes and Kolmogorov flow datasets are generated on GPU to accelerate numerical simulation. Equation
Linear Helmholtz
Poisson
Nonlinear Helmholtz
Darcy
Navier-Stokes
Kolmogorov Flow
Time (s)
0.04
0.04
1.00
0.40
23.30
15.5
C.1.1
Linear/Nonlinear Helmholtz ∆u(r) + κ2 u(r) + λ u(r)3 = f (r), u(r) = 0,
r∈Ω
r ∈ ∂Ω
(15)
The wavenumber κ controls oscillation frequency and λ introduces cubic nonlinearity. For the linear case (λ = 0) we train with κ = 1; for the nonlinear case we train with κ = 2, λ = 1. Input forcing f is sampled from a GRF with smoothness parameter α = 2.0 and length scale τ = 3.0 . C.1.2
Poisson ∆u(r) = a f (r), u(r) = 0,
16
r∈Ω
r ∈ ∂Ω
(16)
The scaling factor a = 2 is used for training. Poisson and Helmholtz also serve as the mutual crossequation transfer target for each other — models trained on Poisson are evaluated on Helmholtz without fine-tuning, and vice versa, exploiting the shared Laplacian structure that a residual-conditioned corrector can generalize across. C.1.3
Darcy Flow −∇ · a(r) ∇u(r) = f (r), u(r) = 0,
r∈Ω
r ∈ ∂Ω
(17)
Following the setup of FNO, the coefficient field a(r) is generated by first sampling a GRF g ∼ N 0, (−∆ + τ 2 I)−α , and then applying a pointwise thresholding operator ( 12, g(r) ≥ 0, a(r) = 3, g(r) < 0. resulting in a binary-valued permeability field. The forcing is kept fixed f (r) = 100. C.1.4
Navier-Stokes/Kolmogorov flow ∂t w(r, t) + v(r, t) · ∇w(r, t) = ν ∆w(r, t) + f (r),
r ∈ Ω, t ∈ (t0 , T ]
∇ · v(r, t) = 0
(18)
where w = ∇ × v is the vorticity and ν = 10−4 is the viscosity. During training we use sinusoidal forcing: f (r) = 0.1 sin(2π(r1 + r2 )) + cos(2π(r1 + r2 )) (19) We consider the time interval t ∈ [2, 6], where the initial vorticity w0 at t = 2 is used to predict the next 20 frames {w1 , . . . , w20 } with a temporal resolution of ∆t = 0.2. We also study the Kolmogorov flow setting, where the forcing term is given by f (r) = −4 cos(8πr2 ), the viscosity is set to ν = 5 × 10−4 , and solutions are collected over the interval t ∈ [2, 4] with temporal resolution ∆t = 0.05. Note that these time steps correspond to the dataset sampling intervals. The numerical simulations are performed using a much finer time step of ∆t = 10−4 .
C.2
PDE Residual Computation
Unless otherwise specified, all PDE residuals reported in the tables are computed as the mean squared error (MSE) of the residual field. Helmholtz, Poisson, Darcy. Second-order central finite differences on a uniform n × n grid with spacing h = 1/(n − 1), differentiable with respect to u. For nonlinear Helmholtz at grid point (i, j): rij =
ui+1,j − 2ui,j + ui−1,j ui,j+1 − 2ui,j + ui,j−1 + h2 h2 2 3 + κ uij + λuij − fij
(20)
Navier-Stokes and Kolmogorov Flow. Given a predicted vorticity trajectory w ∈ RT ×N ×N , we compute the PDE residual using spectral differentiation in space and finite differences in time. Specifically, the temporal derivative ∂t w is approximated by forward and backward differences at the first and last frames, respectively, and by centered differences for all intermediate frames. Spatial derivatives are computed in Fourier space. We recover the stream function by ψ̂(k) =
ŵ(k) 4π 2 |k|2
(21)
The velocity is then obtained as v = −∂x ψ
u = ∂y ψ,
(22)
and the vorticity derivatives ∂x w, ∂y w, and ∆w are also computed spectrally. The residual is defined as r = ∂t w + u∂x w + v∂y w − ν∆w − f 17
(23)
Table 7: Training vs. evaluation parameters for each extrapolation regime.
C.3
PDE
Regime
Training
Evaluation
linear Helmholtz linear Helmholtz linear Helmholtz
Extrapolation Super-resolution Cross-equation
κ=1 128 × 128 Helmholtz
κ=3 256 × 256 Poisson
nonlinear Helmholtz nonlinear Helmholtz
Extrapolation Super-resolution
κ = 2, λ = 1 128 × 128
κ = 3, λ = 4 256 × 256
Poisson Poisson Poisson
Scaling shift Super-resolution Cross-equation
a=2 128 × 128 Poisson
a=1 256 × 256 nonlinear Helmholtz
Navier-Stokes Navier-Stokes Navier-Stokes
Viscosity shift Forcing shift Super-resolution
ν = 10−4 2π freq. 128 × 128
ν = 10−5 4π freq. 256 × 256
Darcy flow Darcy flow
Forcing shift Super-resolution
f = 100 128 × 128
f = 200 256 × 256
Kolmogorov flow Kolmogorov flow Kolmogorov flow
Forcing shift Viscosity shift Super-resolution
8π freq. ν = 5 × 10−4 128 × 128
10π freq. ν = 10−3 256 × 256
Extrapolation Protocols
All evaluations are zero-shot — no fine-tuning on target distributions. We summarize the extrapolation settings in Tab. 7.
C.4
Baseline Details
For all baselines, we follow their official implementations whenever possible. For static PDEs, FNO and PINO employ a 4-layer FNO-2D backbone with 64 hidden channels and 20 Fourier modes, while for dynamic PDEs they use a 4-layer FNO-3D backbone with 64 hidden channels and 8 Fourier modes. Following the official configurations, FNO is trained using only the reconstruction loss (relative H 1 loss). PINO is trained with reconstruction loss and PDE residual for static PDEs, and additionally incorporates a first-frame regularization term for dynamic PDEs. PINO-TTOP initializes from the pretrained PINO model and performs test-time optimization using only the PDE residual for static PDEs, and the PDE residual together with first-frame regularization for dynamic PDEs, all following their official implementation. We use 500 optimization steps for static PDEs and 1,000 steps for dynamic PDEs as in codebase. For each extrapolation setting , PINO-TTOP is re-run from the pretrained PINO model and optimized independently. For POSEIDON, we initialize from the pretrained POSEIDON-B model and fine-tune it to be converged for 200 epochs on each downstream task. Our DiffusionPDE baselines for time-dependent PDEs (NS, KF) uses a two-frame representation (modeling only initial and final states as a joint distribution) in contrast to the full spatiotemporal fields used by PCFM and ENS. We reuse author-provided checkpoints for Darcy, Poisson, and linear Helmholtz; we train custom DiffusionPDE models for other PDEs, fine-tuning the NS model from the author’s non-bounded NS weights. Following the original implementation, we use a Heun EDM sampler with 2,000 steps and DPS-style soft guidance on the PDE residual and input observation at each denoising step. Our PCFM baseline uses the original Functional Flow Matching (FFM) framework with an FNO backbone (including 3D FNO for time-dependent PDEs). We train all PCFM models from scratch following the original training hyperparameters. During inference (50 ODE steps), PDE constraints are enforced via a hard projection at every integration step using a Gauss-Newton solver; we use the authors’ default sampler without guided interpolation. For time-dependent PDEs (NS, KF), the PCFM applies IC and mass conservation constraints, following the original authors’ implementation.
18
D
Additional Ablation Analysis
We expand on the ablations summarized in the main text, covering (i) residual conditioning versus computational depth, (ii) matched-backbone comparisons that isolate the method from the architecture, (iii) backbone expressivity, (iv) the dynamics of the learned correction, and (v) residual versus gradient conditioning. Residual conditioning versus computational depth. To confirm that ENS’s improvement comes from reading the residual rather than from additional iterations, we replace the residual input with a zero field while holding the architecture and step count fixed. Figs. 9 and 10 show the reconstruction error and PDE residual over correction steps: the ablated model reduces neither quantity, whereas ENS drives both down, so the effect comes from residual conditioning rather than the recurrent computation itself. The same figures show that out-of-distribution inputs reach the same floors as in-distribution inputs but require more correction steps, which the residual-plateau stopping rule accommodates automatically. The architecture is not the source. ENS uses a CNN-augmented FNO backbone for the static PDEs and a transformer-based VideoPDE backbone for the turbulent flows. To rule out the backbone as the source of the gains, we place the strongest baselines—the plain operator, its physics-informed (PINO) variant, and its test-time-optimized (TTOP) variant—on ENS’s own backbone for each family (Tabs. 12, 13, 14). Across Helmholtz, Navier–Stokes, and Kolmogorov flow, the baselines methods do not show clear improvement against their original FNO architectures with mixed signals. ENS retains the lowest reconstruction error in every regime against the baselines that share its architecture. The TTOP variant reproduces the residual–reconstruction gap of Prop. 1 on the matched backbone: it reaches comparable or lower PDE residuals while leaving reconstruction error high, and the gap widens with ill-conditioning—mild on Navier–Stokes, where VideoPDE(TTOP) attains a competitive L2 of 8.69e−3, and severe on Kolmogorov flow, where VideoPDE(TTOP) attains the lowest residual yet an L2 of 4.34e−1, roughly 50× worse than ENS (in-distribution). The ENS’s advantage is therefore attributable to residual conditioning, not the backbone. Backbone expressivity. Residual conditioning nonetheless requires a backbone able to represent the high-frequency residual field. Tab 11 sweeps backbones on the linear Helmholtz equation: a pure FNO diverges, its CNN-augmented variants converge (CNN-FNO-CNN best, and in the fewest steps), and a U-Net converges but to a substantially worse solution. Spectral truncation prevents a pure FNO from representing the residual it is asked to read, while the local feature extraction added by CNN lifting and projection restores this capacity. ENS thus works with any sufficiently expressive backbone rather than a single specific one, including the transformer-based architecture applied to the turbulent flows. Dynamics of the learned correction. Figs. 11 and 12 visualize the reconstruction loss, PDE residual, and the learned correction δ û(k) across iterations for extrapolated and Gaussian-noised source fields. The correction magnitude is largest in early iterations and decays as the prediction approaches the solution manifold; under distribution shift, convergence requires roughly 60–90 iterations, well beyond the in-distribution case. That ENS recovers accurate solutions even from noised inputs—initializations far from any seen in training—indicates that the corrector learns a robust correction policy rather than a fixed mapping from the training distribution. Residual versus gradient conditioning. Finally, we ask whether the corrector benefits from the residual field itself or from its physics-loss gradient ∇uk Lres , with Lres = 12 ∥R(uk )∥22 . Fig. 15 compares the two: gradient conditioning accelerates residual minimization but stalls the reconstruction loss, whereas residual conditioning reduces both together. Conditioning on the gradient biases the corrector toward driving the residual to zero—the very objective Prop. 1 shows to be an unreliable proxy—confirming that the residual is most useful as an input signal, not as an optimization target.
E
Related Work (Continued)
Classical numerical solvers. The dominant classical approaches to solving PDEs — finite element methods, finite difference methods, and spectral methods — rely on discretizing the governing equations into large linear or nonlinear systems [Aliabadi, 2020, Ŝolín, 2005]. Although these methods come with 19
strong convergence guarantees, their cost scales unfavorably with resolution and dimensionality, and iterative solvers such as multigrid and Krylov methods [Saad, 2003] still require significant per-instance computation. ENS is motivated precisely by this bottleneck: it amortizes the correction process across a family of PDE instances, replacing expensive per-instance numerical optimization with a single learned recurrent policy. Feed-forward Neural Operators We discuss related work on feed-forward and neural-numerical hybrid methods. In Appendix, we cover more comprehensive related work, including generative models and unrolled networks. Neural operators learn mappings between function spaces, enabling fast prediction across entire PDE families. FNO [Li et al., 2020] parameterizes the integral kernel in Fourier space; DeepONet [Lu et al., 2021] provides universal approximation guarantees via a branch-trunk decomposition; and a unified theoretical treatment is given in Kovachki et al. [2023]. Foundation models such as POSEIDON [Herde et al., 2024] extend operator learning to large-scale pretraining over diverse PDE families. All of these methods are physics-blind at inference time: the network receives no feedback about its own prediction error, which can lead to large constraint violations even when outputs appear visually plausible. While PINNs [Raissi et al., 2019] instead parameterize the solution of a single PDE instance, it requires expensive retraining per instance (Fig. 4). External residual correction. To improve the physical accuracy of feed-forward models, a growing body of work incorporates the PDE residual as a correction signal at inference time, applied externally to learned network. Optimization-based methods include PINO [Li et al., 2024], which performs instance-wise fine-tuning of network parameters by minimizing ∥r(û)∥2 ; DiffusionPDE [Huang et al., 2024] and CoCoGen [Jacobsen et al., 2025], which steer the reverse diffusion process with residual gradients; PDM [Christopher et al., 2024], which projects each diffusion step onto the feasible constraint set; PCFM [Utkarsh et al., 2025], which applies penalized Gauss-Newton corrections at each flow-matching step; and ECI sampling [Cheng et al., 2025], which adapts pre-trained flow-matching models to satisfy hard constraints in zero-shot manner. Analytical approaches instead construct a closed-form correction from the residual directly: Cao et al. [2023] and Jha [2024] solve a linear variational problem at the neural operator’s prediction, yielding a one-shot Newton-like correction; PhysicsCorrect [Huang and Perdikaris, 2026] similarly applies a linearized least-squares solve using a precomputed Jacobian pseudoinverse. Despite their differences, none of these methods feeds the residual field as a direct input channel to a learned corrector network. Learned iterative solvers and algorithm unrolling. Iterative correction has also been explored without explicit residual feedback. PDE-Refiner [Lippe et al., 2023] applies diffusion-inspired multi-step refinement using noise as the correction signal across refinement levels, enabling accurate modeling of all frequency components in long-horizon rollouts. INC [Wei et al., 2026] embeds learned corrections within the governing equations of a coarse numerical solver, provably reducing autoregressive error amplification in chaotic regimes. Both demonstrate the value of iterative learned correction, but neither uses the PDE residual field as an input conditioning. Learning to optimize [Andrychowicz et al., 2016] and Deep Equilibrium Models [Bai et al., 2019] each demonstrate that iterative correction procedures can be learned rather than hand-designed. ENS shares this iterative spirit but differs in a key respect: the corrector is explicitly conditioned on the spatial structure of the PDE residual field, giving the network direct access to where and how its prediction violates the governing equations rather than learning a general-purpose update rule. Diffusion and generative models for scientific problems. Diffusion models [Ho et al., 2020, Song et al., 2020, Karras et al., 2022] have driven a wave of applications across the natural sciences: RFdiffusion [Watson et al., 2023] and DiffSBDD [Schneuing et al., 2024] enable de novo protein and drug design; [Hoogeboom et al., 2022] generate 3D molecular conformations with equivariant architectures; CDVAE [Xie et al., 2021] generates stable crystal structures via score-matching; and GenCast [Price et al., 2023] produces probabilistic weather forecasts that surpass operational numerical prediction systems. DiffusionPDE [Huang et al., 2024] applies generative priors to PDE solving under partial observation, using physical residuals as external guidance, the approach our Diffusion ENS extension replaces with a learned, residual-conditioned corrector internal to the denoising loop.
20
Table 8: Forward prediction results for the Poisson problem. Bold indicates the best result, and underlining indicates the second-best. In-distribution
Super-resolution
Extrapolation
Cross-equation
Method
L2 Rel
PDE Res
L2 Rel
PDE Res
L2 Rel
PDE Res
L2 Rel
PDE Res
ENS(Ours) FNO PINO PINO(TTOP) POSEIDON DiffusionPDE PCFM
3.63e-03 1.79e-02 7.01e-02 1.37e-01 3.28e-02 6.24e-02 8.20e-01
3.61e+00 6.98e+02 5.04e+02 8.63e+01 8.90e+04 5.61e+03 2.15e+01
1.77e-02 2.00e-02 6.52e-02 2.15e-01 3.75e-02 7.83e-01 OOM
1.93e+02 1.51e+04 1.24e+04 4.63e+02 2.32e+05 4.15e+04 OOM
3.57e-02 1.00e+00 9.99e-01 3.58e-01 9.94e-01 9.95e-01 2.65e+00
4.40e+00 1.59e+03 1.13e+03 2.41e+01 9.55e+04 6.92e+03 5.42e+03
4.19e-02 5.18e-01 5.32e-01 4.64e-01 4.80e-01 5.17e-01 1.48e+00
9.40e-01 3.70e+03 1.39e+03 5.27e+01 3.27e+04 1.76e+04 2.01e+02
Table 9: Forward prediction results for the Darcy flow problem. Bold indicates the best result, and underlining indicates the second-best. The PCFM projection for the Darcy constraint was numerically unstable and caused the intermediate latent to increase geometrically and overflow, indicated by "diverged". In-distribution
F
Extrapolation
Super-resolution
Method
L2 Rel
PDE Res
L2 Rel
PDE Res
L2 Rel
PDE Res
ENS(Ours) FNO PINO PINO(TTOP) POSEIDON DiffusionPDE PCFM
1.08e-03 9.45e-03 7.06e-02 2.35e-01 1.94e-02 8.28e-01 diverged
3.72e+03 1.42e+05 2.95e+04 1.51e+03 2.54e+06 1.96e+03 diverged
8.13e-03 5.01e-01 5.19e-01 3.87e-01 4.96e-01 9.00e-01 diverged
1.15e+04 1.53e+05 4.24e+04 5.80e+03 2.63e+06 1.82e+03 diverged
7.01e-03 1.29e-02 6.79e-02 5.74e-01 1.83e-02 5.40e-01 OOM
3.84e+04 1.28e+06 2.37e+05 6.24e+03 6.44e+06 7.48e+03 OOM
Additional Results
We report remaining quantitative results that we omitted in the main paper. These include results for non-linear Helmholtz in Tab. 10, Darcy Flow in Tab. 9, and Poisson equation in Tab. 8.
21
Table 10: Forward prediction results for the nonlinear Helmholtz problem. Bold indicates the best result, and underlining indicates the second-best. In-distribution
Extrapolation
Super-resolution
Method
L2 Rel
PDE Res
L2 Rel
PDE Res
L2 Rel
PDE Res
ENS(Ours) FNO PINO PINO(TTOP) POSEIDON DiffusionPDE PCFM
5.06e-03 2.48e-02 4.51e-02 7.80e-02 7.02e-02 1.36e-01 7.89e-02
6.42e+00 3.39e+02 6.14e+01 1.45e+01 4.57e+04 2.60e+05 2.53e+00
6.31e-02 2.47e-01 2.46e-01 1.98e-01 2.77e-01 4.64e-01 1.75e-01
1.76e+01 2.59e+02 5.66e+01 1.29e+01 3.48e+04 5.27e+01 5.53e+00
1.09e-02 1.33e-01 5.94e-02 1.83e-01 4.16e-01 7.34e-01 OOM
3.55e+01 1.25e+04 3.41e+03 1.22e+02 9.72e+05 6.04e+05 OOM
Table 11: Comparative analysis of ENS across various architectural backbones for the linear Helmholtz equation (In-Distribution). Bold indicates the best result, and underlining indicates the second-best. Helmholtz L2 Rel
PDE Res
k
Convergence
CNN-FNO-CNN 3.04e-03 3.90e-01 FNO 3.16e-02 9.03e+03 CNN-FNO 4.66e-03 5.19e+01 FNO-CNN 1.44e-02 9.19e+00 U-net 1.17e-01 7.20e+03
15 20 45 20
✓ ✗ ✓ ✓ ✓
Model
Table 12: Ablation study on architecture for Helmholtz. Bold indicates the best result, and underlining indicates the second-best. In-distribution Method
L2 Rel
Wavenumber-shift
PDE Res
L2 Rel
PDE Res
Cross-equation L2 Rel
PDE Res
Super-resolution L2 Rel
PDE Res
ENS(Ours) 3.04e-03 3.90e-01 2.59e-02 5.19e-01 8.73e-02 6.54e+01 1.82e-02 4.52e-01 CNN-FNO-CNN 1.80e-02 5.29e+02 2.37e-01 5.26e+02 5.08e-01 2.51e+04 1.37e-01 3.54e+04 CNN-FNO-CNN(PINO) 2.33e-02 3.92e-01 2.35e-01 1.22e+01 9.21e-01 9.48e+02 1.87e-01 3.69e+03 CNN-FNO-CNN(TTOP) 2.31e-02 6.48e-02 3.59e-02 9.55e-02 6.07e-01 1.34e+00 1.93e-01 1.60e+00
Table 13: Ablation study on architecture for Navier-Stokes. Bold indicates the best result, and underlining indicates the second-best. In-distribution Method
L2 Rel
PDE Res
Viscosity-shift L2 Rel
PDE Res
Forcing-shift L2 Rel
PDE Res
Super-resolution L2 Rel
PDE Res
ENS(Ours) 2.22e-03 4.69e-05 3.05e-02 1.12e-02 5.31e-03 7.40e-04 3.98e-02 9.11e-02 VideoPDE 3.02e-02 2.38e-03 1.22e-01 2.42e-02 5.30e-01 4.53e-02 3.86e-01 7.96e-01 VideoPDE(PINO) 2.99e-02 3.91e-04 1.31e-01 2.56e-03 5.24e-01 2.02e-02 4.09e-01 9.91e-01 VideoPDE(TTOP) 8.69e-03 1.47e-05 7.09e-02 4.97e-04 2.42e-02 1.56e-04 7.81e-02 9.60e-04
Table 14: Ablation study on architecture for Kolmogorov flow. Bold indicates the best result, and underlining indicates the second-best. In-distribution Method
L2 Rel
PDE Res
Viscosity-shift L2 Rel
PDE Res
Forcing-shift L2 Rel
PDE Res
Super-resolution L2 Rel
PDE Res
ENS(Ours) 8.08e-03 2.50e-01 3.42e-02 8.05e-01 2.06e-02 3.80e-01 1.98e-02 5.55e+00 VideoPDE 4.12e-02 3.42e+00 5.41e-01 2.35e+01 5.04e-01 3.40e+01 5.56e-01 2.57e+03 VideoPDE(PINO) 4.08e-02 7.73e-01 5.33e-01 8.36e+00 4.81e-01 1.53e+01 5.36e-01 2.36e+03 VideoPDE(TTOP) 4.34e-01 1.09e-01 2.54e-01 6.54e-02 2.14e-01 2.77e-02 7.15e-01 1.42e+00
22
Without PDE Residual With PDE Residual With PDE Residual (Extrapolation) With PDE Residual (Extrapolation)
With PDE Residual
Without PDE Residual
Figure 9: Reconstruction loss (L1 difference) over iterations. From top to bottom, we visualize the L1 error trajectory for the ablation baseline without PDE residual conditioning, ENS model on in-distribution data, ENS model on out-of-distribution data. Clearly, the effectiveness of ENS does not come from the increased computation only, evidenced by the failure of the ablation baseline. Note that out-of-distribution cases take more iterations to converge.
Figure 10: PDE residual over iterations. From top to bottom, we visualize the PDE residual trajectory for the ablation baseline without PDE residual conditioning, ENS model on in-distribution data, ENS model on out-of-distribution data. Clearly, without conditioning on the residual field, the PDE residual even increases during the correction process.
23
Reconstruction Loss PDE Residual u u
PDE Residual (noisy input)
Reconstruction Loss
Figure 11: Visualization of the Reconstruction loss, PDE residual, and the corresponding learned correction term δ û(k) for extrapolated source input fields. The learned correction closely aligns with the PDE residual and progressively removes high-residual regions, leading to a monotonic reduction in both errors. Due to distribution shift, it takes 60~90 iterations for ENS to converge.
Figure 12: Visualization of the Reconstruction loss, PDE residual, and the corresponding learned correction term δ û(k) for out-of-distribution, Gaussian-noised source input fields. Despite the corrupted input field, the learned correction exhibits strong spatial alignment with the PDE residual, resulting in a monotonic decrease in both the PDE residual and reconstruction error.
24
In-distribution
Super-resolution
Viscosity-shift
Forcing-shift
Inference time [s]
102 101 100 10 1 10 2 10 1
10 2
Relative L2 error
100
10 2
10 1
ENS (Ours)
FNO
PINO
Relative L2 error
10 1
100
10 2
PINO (TTOP)
POSEIDON
Relative L2 error
100
10 1
10 2
Relative L2 error
DiffusionPDE
100
PCFM
Figure 13: Accuracy vs. inference cost on Navier–Stokes, per regime. Lower-left is better; the gray curve in each panel traces that regime’s Pareto frontier and the dashed line marks the classical solver (≈23s/sample). ENS is the most accurate learned method on the frontier for in-distribution, viscosity-shift, and forcing-shift, while POSEIDON is more accurate on super-resolution at higher cost. ENS’s inference time is 0.19s in-distribution, 0.5s under viscosity-shift, 0.33s under forcing-shift, and 2.6s for super-resolution, as the number of correction steps adapts to difficulty. The correction and diffusion hybrids (PINO-TTOP, DiffusionPDE, PCFM) lie above the solver line, i.e. slower than solving the PDE directly. FNO and PINO are offset horizontally by ±12% for visibility, as their accuracy is nearly identical.
In-distribution
Super-resolution
Viscosity-shift
Forcing-shift
Inference time [s]
102 101 100 10 1 10 2 10 2
10 1
100
Relative L2 error
101
10 2
ENS (Ours)
10 1
100
101
FNO
PINO
PINO (TTOP)
Relative L2 error
10 2
10 1
100
Relative L2 error POSEIDON
101
10 2
10 1
100
Relative L2 error
101
DiffusionPDE
Figure 14: Accuracy vs. inference cost on Kolmogorov flow, per regime. Lower-left is better; the gray curve in each panel traces that regime’s Pareto frontier and the dashed line marks the classical solver (≈15.5 s/sample). ENS is the most accurate learned method on the frontier in all four regimes, with inference time 0.43 s in-distribution, 0.50 s under viscosity- and forcing-shift, and 3.05 s for super-resolution, as the number of correction steps adapts to difficulty. PINO-TTOP reaches low PDE residuals but high reconstruction error (Prop. 1), seen as its rightward position despite test-time optimization and most extreme under super-resolution (L2 ≈ 4.5). PCFM runs out of memory in the 40-frame setting and is omitted.
25
Mean Squared Error (MSE)
ENS (Res. Conditioning): Rec. Loss ENS (Res. Conditioning): PDE Residual
ENS (Grad. Conditioning): Rec. Loss ENS (Grad. Conditioning): PDE Residual
103 101 10 1 10 3 10 5 0
2
4
6
8
Iteration Step k
10
12
14
Figure 15: Comparison of conditioning strategies for the ENS corrector. We evaluate the optimization trajectories when conditioning on the PDE residual R(uk ) versus the physics-gradient ∇uk Lres . While gradient-conditioning accelerates the minimization of the PDE residual, it results in a stagnation of reconstruction loss. In contrast, residual-conditioning facilitates the simultaneous convergence of both PDE residual and reconstruction loss.
26