Physics-Informed Inference Time Scaling for Solving High-Dimensional PDE via Defect Correction

TL;DR

SCaSML introduces a novel inference-time framework that upgrades pre-trained PDE surrogates by solving a structurally preserved defect PDE with Monte Carlo methods. The approach preserves the semi-linear structure to enable efficient MLP/MLMC corrections, yielding a final error bounded by the product of surrogate and simulation errors and achieving faster convergence than standalone surrogates or naive solvers. Across problems up to 160 dimensions, SCaSML consistently reduces errors by 20-80% and demonstrates elastic compute, where additional inference-time effort yields meaningful accuracy gains without retraining. This hybrid physics-informed strategy offers practical guidance for reliable, scalable solving of high-dimensional PDEs in scientific and engineering applications.

Abstract

Solving high-dimensional partial differential equations (PDEs) is a critical challenge where modern data-driven solvers often lack reliability and rigorous error guarantees. We introduce Simulation-Calibrated Scientific Machine Learning (SCaSML), a framework that systematically improves pre-trained PDE solvers at inference time without any retraining. Our core idea is to use defect correction method that derive a new PDE, termed Structural-preserving Law of Defect, that precisely describes the error of a given surrogate model. Since it retains the structure of the original problem, we can solve it efficiently with traditional stochastic simulators and correct the initial machine-learned solution. We prove that SCaSML achieves a faster convergence rate, with a final error bounded by the product of the surrogate and simulation errors. On challenging PDEs up to 160 dimensions, SCaSML reduces the error of various surrogate models, including PINNs and Gaussian Processes, by 20-80%. Code of SCaSML is available at https://github.com/Francis-Fan-create/SCaSML.

Figures (29)

• Figure 1: SCaSML framework pipeline.a) SCaSML aims to allocate compute at inference time to further improve the accuracy of a surrogate model. It first fits a surrogate model $\hat{u}$ as an initial estimate of the PDE solution $u$, then leverages stochastic simulation algorithms to approximate the defect$\breve u = u - \hat{u}$ at inference time via formulating $\breve{u}$ as the solution to the structural-preserving law of defect. b) Method for deriving this law. Using an approximate solution of a semi‑linear equation, we derive a new differential equation that characterizes the error, which inherently preserves the semi‑linear structure.
• Figure 2: Flow diagram of SCaSML. We formulate the error $u-\hat{u}$ of surrogate solution $\hat{u}$ as the solution Structural-preserving Law of Defect (\ref{['eq:law_of_defect_semilinear']}), a new semi-linear PDE. At inference time, we approximate $u-\hat{u}$ via solving Structural-preserving Law of Defect using Multilevel Picard (MLP) iteration. The generated estimation of $u-\hat{u}$ helps us to calibrate the surrogate solution $\hat{u}$.
• Figure 3: Empirical verification of the improved scaling law for SCaSML.(a) Conceptual diagram. The final SCaSML error is a product of the surrogate error and the simulation error (Theorem \ref{['thm:l2_error_bound']}). By balancing the computational budget between training the surrogate and performing inference-time simulation, SCaSML achieves a faster overall convergence rate (Corollary \ref{['cor:scaling_law_main']}). (b) Numerical results. We plot the $L^2$ error versus the number of collocation points ($m$) on a log-log scale for a GP surrogate and SCaSML. The slope of the line corresponds to the convergence rate $\gamma$. SCaSML consistently exhibits a steeper slope than the base surrogate, empirically confirming its accelerated convergence.
• Figure 4: Efficiency and performance of the SCaSML methodology.(a) Violin plots showing the distribution of pointwise errors. SCaSML consistently reduces the mean error and tightens the distribution compared to the surrogate (SR) model. (b) Inference-time scaling. As the number of inference-time simulation samples increases, SCaSML's error steadily decreases, demonstrating effective use of additional compute. (c) Summary of performance. The left panel shows that SCaSML (blue stars) consistently achieves lower $L^2$ error than both the surrogate (SR) and naive MLP methods across all problems. The right panel (heatmap) confirms that SCaSML also dominates in $L^\infty$ and $L^1$ error metrics.
• Figure 5: Violin Plot for comparison of the baseline PINN surrogate (black), MLP (gray), applying qudrature SCaSML (teal) to calibrate the PINN surrogate on linear convection-diffusion equation for $d=10,20,30,60$.

Theorems & Definitions (41)

• Definition 2.1: Structural-preserving Law of Defect for Linear PDEs
• Remark 2.2: Regards Training and Inference Separation
• Theorem 4.2: Global $L^2$ Error Bound
• Corollary 4.3: Improved Scaling Law
• Definition B.1: Gauss–Legendre quadrature
• Definition B.2: Quadrature Multilevel Picard Iteration
• Definition B.3: Full-history Multilevel Picard Iteration Hutzenthaler_2020
• Definition D.1: Coordinate System and Vector Norms
• Definition D.2: Measurable Spaces and Functions
• Definition D.3: Probability Space and $L^p$ Norms