cs.LGMar 4, 2026

Hierarchical Inference and Closure Learning via Adaptive Surrogates for ODEs and PDEs

Authors: Pengyu Zhang, Arnaud Vadeboncoeur, Alex Glyn-Davies, Mark Girolami

Organizations: Department of Engineering, University of Cambrdige, 7a JJ Thomson Ave, Cambridge, CB3 0FA, UK · The Alan Turing Institute, 96 Euston Rd., London, NW1 2DB, UK

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

Abstract

Inverse problems are the task of calibrating models to match data. They play a pivotal role in diverse engineering applications by allowing practitioners to align models with reality. In many applications, engineers and scientists do not have a complete picture of i) the detailed properties of a system (such as material properties, geometry, initial conditions, etc.); ii) the complete laws describing all dynamics at play (such as friction laws, complicated damping phenomena, and general nonlinear interactions). In this paper, we develop a principled methodology for leveraging data from collections of distinct yet related physical systems to jointly estimate the individual model parameters of each system, and learn the shared unknown dynamics in the form of an ML-based closure model. To robustly infer the unknown parameters for each system, we employ a hierarchical Bayesian framework, which allows for the joint inference of multiple systems and their population-level statistics. To learn the closures, we use a maximum marginal likelihood estimate of a neural network embeded within the ODE/PDE formulation of the problem. To realize this framework we utilize the ensemble Metropolis-Adjusted Langevin Algorithm (MALA) for stable and efficient sampling. To mitigate the computational bottleneck of repetitive forward evaluations in solving inverse problems, we introduce a bilevel optimization strategy to simultaneously train a surrogate forward model alongside the inference. Within this framework, we evaluate and compare distinct surrogate architectures, specifically Fourier Neural Operators (FNO) and parametric Physics-Informed Neural Network (PINNs).

Figures & tables

Figure 1: Hierarchical Bayesian model illustrating how a global hyperparameter ϕ\boldsymbol{\phi} governs a collection of task-specific parameters θ(1:K)\boldsymbol{\theta}^{(1:K)} , each of which generates an associated dataset y(k)\boldsymbol{y}^{(k)} .
Figure 2: Overview of the proposed framework. The algorithm alternates between three main stages: (1) Hierarchical posterior sampling using ensemble MALA to update parameter chains, (2) Lower-level surrogate training optimizing the surrogate model FβF^{\beta} parameterized by β\beta over NN iterations, where the loss function LSurrogate\mathcal{L}_{\mathsf{Surrogate}} could take the forms in ( 27 ), ( 28 ) and ( 29 ), and (3) Upper-level closure learning optimizing the closure model fαf^{\alpha} parameterized by α\alpha to discover the unknown nonlinear dynamics, where the loss function LLML\mathcal{L}_{\mathsf{LML}} follows ( 26 ).
Figure 3: Distribution of ground-truth physical parameters for K=20\mathbf{K=20} mass–damper systems in Experiment 1. Each system’s parameter vector θ={log⁡(k),u0,v0}\boldsymbol{\theta}=\{\log(k),u_{0},v_{0}\} is sampled from the hierarchical prior with hyperparameters μϕ={log⁡(5.0),0.0,2.0}\boldsymbol{\mu}_{\phi}=\{\log(5.0),0.0,2.0\} and τϕ={0.03,2.0,2.0}\boldsymbol{\tau}_{\phi}=\{0.03,2.0,2.0\} . The histograms show the parameter distributions across the 20 systems, while the overlaid curves represent the corresponding ground-truth probability density functions implied by the hierarchical prior. The green dashed lines indicate ground-truth population means. The apparent deviation of the histograms from the ground-truth curves occurs because a limited population of only 20 systems poorly represents the full underlying probability distribution.
Figure 4: Displacement trajectories u(t)\mathbf{u(t)} for 20 nonlinear mass–damper systems. These systems are governed by ( 31 ) and solved using the leapfrog integrator with system-specific parameters in Figure 3 . The right two panels illustrate, for two representative systems, the corresponding sparse and noisy observations generated via observation operator g(k)g^{(k)} .
Figure 5: Comparison of closure model estimation and surrogate performance across different models with K=20\mathbf{K=20} in Experiment 1. The rows correspond to Numerical Solver, Supervised FNO, Physics-Based FNO, and PINNs, respectively. Left Column: Comparison between the learned closure fα(u˙)f^{\alpha}(\dot{u}) (orange) and the ground truth f(u˙)f(\dot{u}) (blue) over the velocity range u˙∈[−6,6]\dot{u}\in[-6,6] . Middle &\& Right Columns: Surrogate predictions of displacement u(t)u(t) for two selected systems (System 3 and System 9). Bottom Panel: Histogram of velocity values in GT dataset.
Figure 6: Trace plots illustrating the evolution of a randomly selected MALA chain for stiffness kk of System 9 across four models with K=20\mathbf{K=20} in Experiment 1. The figures illustrate convergence behaviour relative to the GT parameter value (green dashed line).
ModelNo. of SystemsParameter InferenceClosure MSESurrogate MSE
  Mean MSECoverage  
SolverK=5K=51.44×10−21.44\times 10^{-2}100%100\%1.051.05-
K=10K=102.41×10−22.41\times 10^{-2}93.3%93.3\%0.480.48-
K=20K=202.70×10−22.70\times 10^{-2}96.7%96.7\%0.500.50-
K=30K=303.02×10−23.02\times 10^{-2}91.1%91.1\%0.450.45-
FNO (Supervised)K=5K=54.85×10−24.85\times 10^{-2}93.3%93.3\%1.051.051.53×10−31.53\times 10^{-3}
Table 1: Quantitative comparison of the four models across varying numbers of systems ( K∈{5,10,20,30,50,100}\mathbf{K\in\{5,10,20,30,50,100\}} ) in Experiment 1. The table evaluates Parameter Inference performance via the MSE of posterior means relative to GT parameter values and Coverage , defined as the percentage of GT parameters falling within two posterior standard deviations. Additionally, Closure MSE measures the MSE of the learned nonlinear closure fα(u˙)f^{\alpha}(\dot{u}) evaluated on u˙∈[−5,5]\dot{u}\in[-5,5] , while Surrogate MSE quantifies the accuracy of the surrogate model’s trajectory predictions compared to GT trajectories.
Figure 7: Posterior histograms and corresponding approximated probability density functions for three parameters from System 9 across four models with K=20\mathbf{K=20} in Experiment 1. Each panel shows the marginal posterior distribution obtained from the final 1,000 samples of one random MALA chain after convergence. The blue bars represent the posterior samples, the red curve denotes the approximated density, and the green dashed line marks the GT parameter value.
MethodNo. of SystemsParameter InferenceClosure MSESurrogate MSE
  Mean MSECoverage  
HierarchicalK=5K=54.19×10−24.19\times 10^{-2}86.7%86.7\%1.181.183.80×10−33.80\times 10^{-3}
K=10K=102.44×10−22.44\times 10^{-2}93.3%93.3\%1.151.151.66×10−31.66\times 10^{-3}
K=20K=206.37×10−26.37\times 10^{-2}90.0%90.0\%1.271.273.97×10−33.97\times 10^{-3}
K=30K=305.59×10−25.59\times 10^{-2}92.2%92.2\%1.151.152.64×10−32.64\times 10^{-3}
K=50K=505.06×10−25.06\times 10^{-2}81.3%81.3\%0.790.792.29×10−32.29\times 10^{-3}
Table 2: Comparison of hierarchical and non-hierarchical frameworks trained using PINNs as surrogate model in Experiment 1. Reported metrics are the same as in Table 1 .
Figure 8: Comparison of posterior parameter distributions between Hierarchical and Non-Hierarchical Bayesian frameworks using a PINN-based surrogate with K=20\mathbf{K=20} in Experiment 1. The figure compares the uncertainty quantification performance for three parameters from System 3. The top row (Hierarchical) and bottom row (Non-Hierarchical) show distinct posterior densities. GT values are indicated by green dashed lines, overlaid on the posterior samples (blue histograms) and their PDF approximations (red curves).
Figure 9: MALA chains of posterior samples for stiffness kk of System 9 for Hierarchical and Non-Hierarchical Bayesian frameworks in Experiment 1. The figures compare the convergence rate between hierarchical and non-hierarchical models.
Figure 10: MALA chains of posterior samples for hyperparameters inferred using the PINN-based hierarchical Bayesian framework in Experiment 1. The panels display the MALA chains for μϕ\boldsymbol{\mu}_{\phi} and τϕ\boldsymbol{\tau}_{\phi} . The top row corresponds hyperprior means ( μk,μu0,μv0\mu_{k},\mu_{u_{0}},\mu_{v_{0}} ), while the bottom row shows the variances ( τk,τu0,τv0\tau_{k},\tau_{u_{0}},\tau_{v_{0}} ). Green dashed lines represent ground-truth values.
SamplesSolverFNO (Supervised)FNO (Physics-based)PINNsPINNs (Non-Hierarchical)
K=5K=50.140s0.549s0.099s0.038s0.040s
K=10K=100.171s0.615s0.104s0.038s0.042s
K=20K=200.201s0.712s0.124s0.040s0.045s
K=30K=300.238s0.763s0.157s0.041s0.046s
K=50K=50-0.842s0.232s0.048s0.058s
K=100K=100-1.158s0.489s0.091s0.091s
Table 3: Computational efficiency comparison across considered models for varying number of systems ( K∈{5,10,20,30,50,100}\mathbf{K\in\{5,10,20,30,50,100\}} ) in Experiment 1. The reported values represent the average time (in seconds) required to complete one full training epoch within the bilevel framework.
Figure 11: Distribution of ground-truth physical parameters for 20 Darcy flow systems in Experiment 2. Each system’s parameter vector θ={z1,z2,z3}\boldsymbol{\theta}=\{z_{1},z_{2},z_{3}\} is sampled from the hierarchical prior with hyperparameters μϕ={0.2,−0.1,0.2}\boldsymbol{\mu}_{\phi}=\{0.2,-0.1,0.2\} and τϕ={0.5,0.5,0.5}\boldsymbol{\tau}_{\phi}=\{0.5,0.5,0.5\} .
Figure 12: Representative ground truth Darcy flow fields generated using fixed-point numerical solver. The plots display the fields for three selected systems (System 8, 9, and 19) from the dataset of K=20K=20 systems. The red markers indicate spatial locations of observation data.
Figure 13: Comparison of closure model estimation and surrogate performance across different models with K=20\mathbf{K=20} in Experiment 2. The rows correspond to GT field generated by fixed-point solver, and results from Supervised FNO, Physics-Based FNO, and PINNs, respectively. Left Column: Comparison between the learned closure model fα(u)f^{\alpha}(u) (orange) and the ground truth f(u)f(u) (blue). Middle &\& Right Columns: Surrogate predictions for two systems (System 8 and System 19). Bottom Panel: Histogram of uu in GT dataset.
Figure 14: Posterior histograms and corresponding approximated posterior density functions for three parameters from System 9 across three models with K=20\mathbf{K=20} in Experiment 2. Each panel shows the marginal posterior distribution obtained from the final 1,000 samples of one random MALA chain after convergence.
ModelNo. of SystemsParameter InferenceClosure MSESurrogate MSE
  Mean MSECoverage  
FNO (Supervised)K=10K=103.39×10−23.39\times 10^{-2}93.3%93.3\%5.01×10−35.01\times 10^{-3}8.92×10−48.92\times 10^{-4}
K=20K=201.43×10−21.43\times 10^{-2}93.3%93.3\%3.95×10−33.95\times 10^{-3}3.48×10−53.48\times 10^{-5}
K=30K=302.37×10−22.37\times 10^{-2}90.0%90.0\%3.28×10−33.28\times 10^{-3}3.57×10−53.57\times 10^{-5}
FNO (Physics-based)K=10K=101.17×10−11.17\times 10^{-1}56.7%56.7\%4.28×10−24.28\times 10^{-2}7.20×10−37.20\times 10^{-3}
K=20K=202.47×10−22.47\times 10^{-2}91.7%91.7\%1.49×10−21.49\times 10^{-2}6.35×10−56.35\times 10^{-5}
Table 4: Quantitative comparison of three models across varying numbers of systems ( K∈{10,20,30}\mathbf{K\in\{10,20,30\}} ) in Experiment 2. The table evaluates Parameter Inference , Closure MSE and Surrogate MSE with metrics the same as in Table 1 . Closure MSE is evaluated over range u∈[0,2.5]u\in[0,2.5] .
Figure 15: Comparison of posterior parameter distributions between Hierarchical and Non-Hierarchical Bayesian frameworks using a PINN-based surrogate with K=20\mathbf{K=20} in Experiment 2. The figure compares the uncertainty quantification performance for three parameters from System 3 for both hierarchical and non-hierarchical setups.
MethodNo. of SystemsParameter InferenceClosure MSESurrogate MSE
  Mean MSECoverage  
HierarchicalK=10K=104.89×10−24.89\times 10^{-2}83.3%83.3\%2.62×10−12.62\times 10^{-1}1.09×10−31.09\times 10^{-3}
K=20K=201.53×10−21.53\times 10^{-2}93.3%93.3\%8.31×10−38.31\times 10^{-3}1.56×10−41.56\times 10^{-4}
K=30K=302.56×10−22.56\times 10^{-2}90.0%90.0\%1.11×10−21.11\times 10^{-2}9.33×10−49.33\times 10^{-4}
Non- HierarchicalK=10K=106.21×10−26.21\times 10^{-2}80.0%80.0\%5.33×10−15.33\times 10^{-1}1.69×10−31.69\times 10^{-3}
K=20K=201.46×10−21.46\times 10^{-2}93.3%93.3\%9.62×10−39.62\times 10^{-3}1.79×10−41.79\times 10^{-4}
Table 5: Comparison of hierarchical and non-hierarchical approaches trained using PINNs as surrogate model in Experiment 2. Reported metrics are the same as in Table 1 .
SamplesSolverFNO (Supervised)FNO (Physics-based)PINNsPINNs (Non-Hierarchical)
K=10K=104.158s4.158s3.121s3.121s0.686s0.686s0.132s0.132s0.134s0.134s
K=20K=2010.374s10.374s3.837s3.837s1.119s1.119s0.222s0.222s0.223s0.223s
K=30K=3017.962s17.962s4.962s4.962s1.934s1.934s0.378s0.378s0.381s0.381s
      
Table 6: Computational efficiency comparison across considered models for varying number of systems ( K∈{10,20,30}\mathbf{K\in\{10,20,30\}} ) in Experiment 2.
Figure 16: Distribution of ground-truth physical parameters for 20 Burgers’ systems in Experiment 3. Each system’s parameter vector θ={log⁡(ν),z}\boldsymbol{\theta}=\{\log(\nu),z\} is sampled from the hierarchical prior with hyperparameters μϕ={log⁡(0.05),1.4}\boldsymbol{\mu}_{\phi}=\{\log(0.05),1.4\} and τϕ={0.07,0.2}\boldsymbol{\tau}_{\phi}=\{0.07,0.2\} .
Figure 17: Representative ground-truth Burgers’ fields generated using numerical solver. The plots display the fields for two selected systems (System 1 and 14) from the dataset of K=20K=20 systems. The red markers indicate spatial locations of observation data.
Figure 18: Closure model estimation and surrogate performance of supervised FNO with K=20\mathbf{K=20} in Experiment 3. Left: Closure model estimation with histograms of ground-truth velocity data provided. Middle &\& Right: Ground truth and surrogate prediction of solution field of System 14.
Figure 19: MALA traces and posterior histograms of parameters from System 14 in Experiment 3. Top row: MALA traces of parameters ν\nu and zz . Bottom row: posterior histograms of last 1,000 samples of MALA traces and their corresponding PDF approximations.
More in this topic

PDE Inverse Problems

Partial differential equation (PDE) inverse problems recover hidden fields, parameters, or sources from observations governed by physical laws.

See all 106 papers

Explore similar work

CardsList

  1. Meta-Inverse Physics-Informed Neural Networks for High-Dimensional Ordinary Differential Equations

    May 5, 2026Zhao Wei, Kenneth Hor Cheng Koh, Sheng Yuan Chin +3Meta-LearningPhysics-Informed ML

    Centre for Frontier AI Research (CFAR) and the Institute of High Performance Computing (IHPC), Agency for Science, Technology and Research (A*STAR), 1 Fusionopolis Way, #16-16 Connexis, Singapore 138632, Singapore · Singapore Institute of Food and Biotechnology Innovation (SIFBI), Agency for Science, Technology and Research (A*STAR), 31 Biopolis Way, Nanos, Singapore 138669, Singapore · A*STAR Skin Research Labs (A*SRL), Agency for Science, Technology and Research (A*STAR), 8A Biomedical Grove, Immunos, Singapore 138648, Singapore +1

  2. Neural network surrogates with uncertainty quantification for inverse problems in partial differential equations

    Jun 18, 2026Christian Jimenez-Beltran, Aretha L. Teckentrup, Antonio Vergari +1Neural Surrogate ModelingUncertainty Quantification

    School of Mathematics and Maxwell Institute for Mathematical Sciences University of Edinburgh · School of Informatics University of Edinburgh

  3. SOLIS: Physics-Informed Learning of Interpretable Neural Surrogates for Nonlinear Systems

    Apr 16, 2026Murat Furkan Mansur, Tufan KumbasarNeural Surrogate ModelingParameter Estimation

    AI and Intelligent Systems Laboratory, Istanbul Technical University, 34469, Istanbul, Türkiye