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

TL;DR

This paper develops 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.

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 (19)

• Figure 1: Hierarchical Bayesian model illustrating how a global hyperparameter $\boldsymbol{\phi}$ governs a collection of task-specific parameters $\boldsymbol{\theta}^{(1:K)}$, each of which generates an associated dataset $\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^\beta$ parameterized by $\beta$ over $N$ iterations, where the loss function $\mathcal{L}_{\mathsf{Surrogate}}$ could take the forms in \ref{['eq:supervised_loss']}, \ref{['eq:fno_physics_loss']} and \ref{['eq:PINNs_loss']}, and (3) Upper-level closure learning optimizing the closure model $f^\alpha$ parameterized by $\alpha$ to discover the unknown nonlinear dynamics, where the loss function $\mathcal{L}_{\mathsf{LML}}$ follows \ref{['eq:alpha_loss_2']}.
• Figure 3: Distribution of ground-truth physical parameters for $\mathbf{K=20}$ mass–damper systems in Experiment 1. Each system’s parameter vector $\boldsymbol{\theta} = \{\log(k), u_0, v_0\}$ is sampled from the hierarchical prior with hyperparameters $\boldsymbol{\mu}_\phi = \{\log(5.0), 0.0, 2.0\}$ and $\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 $\mathbf{u(t)}$ for 20 nonlinear mass–damper systems. These systems are governed by \ref{['eq:mass_spring_1']} and solved using the leapfrog integrator with system-specific parameters in Figure \ref{['GT_parameters_1']}. The right two panels illustrate, for two representative systems, the corresponding sparse and noisy observations generated via observation operator $g^{(k)}$.
• Figure 5: Comparison of closure model estimation and surrogate performance across different models with $\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^\alpha(\dot{u})$ (orange) and the ground truth $f(\dot{u})$ (blue) over the velocity range $\dot{u} \in [-6, 6]$. Middle $\&$ Right Columns: Surrogate predictions of displacement $u(t)$ for two selected systems (System 3 and System 9). Bottom Panel: Histogram of velocity values in GT dataset.