Enforcing governing equation constraints in neural PDE solvers via training-free projections

AI summary

Overview

Research area: Scientific machine learning, neural PDE solvers, physics-constrained learning, numerical optimization.

Technical level: Advanced. The paper assumes familiarity with differential equations, discretized operators, Jacobians, and numerical linear algebra.

Scope: A study of two training-free, post hoc projection operators that map the outputs of already-trained neural PDE solvers back onto the manifold of states satisfying the governing equation, evaluated on the 3D Lorenz ODE, the 1D Kuramoto–Shivashinsky equation, and the 2D Navier–Stokes equations.

What This Paper Is About

Neural networks trained to simulate physical systems often produce solutions that look accurate on standard error metrics but violate the underlying physics — for example, failing to satisfy the governing PDE, its initial and boundary conditions, or conservation laws. Enforcing these constraints during training (physics-informed losses, architectural tricks, helper networks) brings its own problems: harder optimization, reduced expressiveness, or extra hyperparameters. This paper asks whether one can instead leave training alone and fix the constraints afterwards, by post-processing each approximate solution so that it lies closer to the feasible set.

Key Contributions

1. Formulates constraint enforcement as a projection problem: given a neural solver output û, find the nearest discretized state u such that h(u) = c, where h concatenates the discretized PDE operator, boundary operator, and initial-condition operator, and c concatenates the corresponding right-hand sides.
2. Evaluates a nonlinear optimization-based projection that minimizes ||u − û|| + λ||h(u) − c|| using LBFGS, requiring no changes to the trained model.
3. Evaluates a local linearization-based projection that linearizes h about û via the Jacobian J_h, yielding the linear system 𝒞u = b, with both a constrained variant (exact projection onto the linearized feasible set) and a relaxed variant (trading constraint satisfaction against proximity to û) — both admitting closed-form solutions.
4. Shows how to scale this to large systems by avoiding matrix inversion: the linear systems are solved with sparse Krylov solvers (CG, GMRES, BiCGSTAB) that only require the Jacobian-vector product (JVP(v) = 𝒞v) and vector-Jacobian product (VJP(v) = 𝒞ᵀv) operators.

Main Findings

• All three projections reduce constraint violation. Across Lorenz, Kuramoto–Shivashinsky, and Navier–Stokes, the constrained, relaxed, and LBFGS projections lower the squared L₂ norm of the residual ||h(u) − c|| relative to the unprocessed network output. For the Lorenz system the paper states all projections cut constraint violation by 70% or more and also lower MSE.

• LBFGS is the strongest method. On Lorenz with the plain MLP baseline, the residual drops from 50.8 to 1.18 (×10⁻⁴ scale) under LBFGS, versus 13.1 for the constrained projection and 13.3 for the relaxed one. On Navier–Stokes at resolution 64, LBFGS reduces the FNO residual from 8.13 to 0.00901 (×10⁻² scale) while MSE falls from 13 to 2.63 (×10⁻¹ scale).

• Physics-informed training does not automatically win. The PINN-MLP for Lorenz has higher constraint violation (51.7) than the standard MLP (50.8); the paper attributes this to the PINN not training as effectively under the same time and resource budget. After projection, PINN + LBFGS reaches 1.23 versus 1.18 for the MLP + LBFGS.

• Projection helps even when the constraint was already partly satisfied. For Kuramoto–Shivashinsky, the PINO baseline has a relatively low residual of 4.27 (×10⁻⁵) at resolution 64 but that rises to 874 at resolution 128 and 896 at resolution 256. PINO + LBFGS brings these to 1.27, 1.11, and 1.18 respectively, and the paper observes that constraint error increases with resolution unless projection is performed.

• Resolution behaviour differs between the two PDEs. For Kuramoto–Shivashinsky the MSE is described as resolution-invariant because the dynamics are sufficiently resolved on all three grids, while constraint error rises with grid size. For Navier–Stokes the baseline MSE stays nearly constant across resolutions because large-scale structures dominate the error, but the missing fine-scale detail shows up as constraint violations.

• Navier–Stokes gains the most in accuracy. The paper reports that once violations are reduced the Navier–Stokes MSE drops significantly, with the LBFGS-projected result recovering fine-scale structures that physics-informed models fail to capture because they never encountered those details in training.

• Linearization degrades away from the linearization point. Figure 2 measures the constraint violation along the LBFGS path u_0:200 from the initial guess u_0 = û to the final solution u* = u_200, with the x-axis given by the path size d(k) = Σ_{i=0}^{k} ||u_{i+1} − u_i|| for k = 0 to 200. First- and second-order Taylor approximations reveal two regimes: linear approximations work well initially but become inaccurate near the solution, while quadratic approximations remain reliable throughout — the paper's explanation for why LBFGS outperforms the linearization-based projections.

Methodology in Plain English

The researchers treat a trained neural solver's output as a starting guess and ask: what is the closest state that actually obeys the physics? The physics is written as a system of equations — the discretized PDE, the boundary conditions, and the initial conditions — collected into one function h that must equal a known vector c. The difference h(u) − c is the residual, or constraint violation.

They test two ways of driving that residual to zero. The first is a straightforward numerical optimization: minimize a combination of how far you moved from the network's guess and how large the residual still is, weighted by λ, solved with LBFGS. The second exploits calculus: approximate the constraint as linear around the network's guess using its Jacobian, then solve a linear system. That linear system has an exact closed-form solution when you want the constraints met exactly, and a relaxed closed-form solution when you allow a tradeoff controlled by λ. Because the Jacobian can be enormous, they do not build the matrix; instead they use iterative solvers that only need to multiply by it or its transpose, which modern autodiff frameworks provide as Jacobian-vector and vector-Jacobian products.

Experiments compare, for each of three dynamical systems, a purely data-driven model (MLP for Lorenz, FNO for KS and NS) and a physics-informed counterpart (PINN for Lorenz, PINO for KS and NS). Each is measured by test-set MSE and by the squared L₂ residual, before and after each projection. The hyperparameter λ differs by problem: 1000 for the Lorenz relaxed projection, 10 for Kuramoto–Shivashinsky, and ||û|| for Navier–Stokes.

Why This Matters

The work reframes physics-consistency as a post-processing step rather than a training burden. That matters because it separates two concerns that physics-informed training entangles: fitting the data, and satisfying the equations. It also shows that constraints from genuinely dynamical, nonlinear PDEs — where satisfaction at one timestep does not imply consistency across time — are tractable, whereas prior constrained-sampling work focused mostly on simpler per-step linear conditions such as divergence-free fields.

Real-world applications include:

• Fluid simulation and engineering design — aerodynamics, turbulence modelling, and any setting where incompressibility and Navier–Stokes residuals must hold for a simulation to be trusted.
• Weather and climate modelling — long-horizon rollouts of atmospheric models where small per-step errors accumulate into physically implausible trajectories.
• Medical imaging and physiological modelling — enforcing physiologically plausible parameter ranges and dynamics, which the paper cites as a motivating example.
• Particle and molecular systems — where mass and momentum conservation are the constraints of interest, and the paper notes such violations are common.

Industry relevance: the projection is training-free, so it can be bolted onto an existing surrogate model without retraining or access to the training pipeline. The paper is explicit that this comes at a cost: the projections cannot currently be inserted as differentiable layers inside a network because backpropagating through the optimization is too expensive, and reducing projection overhead is described as essential for scaling to larger systems.

Future Directions

• Differentiable projection operators. The paper identifies integration of the projection as a differentiable layer as an open goal, which would enable end-to-end training with guaranteed-consistent outputs; the stated blocker is the computational cost of backpropagation through the optimization.
• Iterative relinearization. Because the analysis shows linear approximations are locally effective but degrade rapidly, the authors suggest that repeatedly relinearizing — in the style of Sequential Quadratic Programming — could improve on single-shot linear projection if it can be implemented efficiently.
• Scaling to larger systems. Reducing projection overhead is called out as necessary both for larger problem sizes and for end-to-end training.
• Generative models. The authors propose extending the technique to generative architectures, where projecting sampled trajectories onto physically consistent manifolds could improve temporal coherence — a setting they note is complicated because PDE determinism can yield degenerate posteriors with essentially unique solutions.

Target Audience

Researchers and practitioners in scientific machine learning and computational physics who train neural surrogates for dynamical systems and need those surrogates to respect governing equations. It is most useful to readers already comfortable with PDE discretization, Jacobians, and Krylov solvers, and to engineers who want to improve an existing trained model without retraining it. Readers looking for an introduction to physics-informed learning will find the framing useful but the experimental detail demanding.

Authors’ abstract

Neural PDE solvers used for scientific simulation often violate governing equation constraints. While linear constraints can be projected cheaply, many constraints are nonlinear, complicating projection onto the feasible set. Dynamical PDEs are especially difficult because constraints induce long-range dependencies in time. In this work, we evaluate two training-free, post hoc projections of approximate solutions: a nonlinear optimization-based projection, and a local linearization-based projection using Jacobian-vector and vector-Jacobian products. We analyze constraints across representative PDEs and find that both projections substantially reduce violations and improve accuracy over physics-informed baselines.