FNOPE: Simulation-based inference on function spaces with Fourier Neural Operators

TL;DR

FNOPE tackles the challenge of performing Bayesian posterior inference for function-valued parameters in simulators. By integrating Fourier Neural Operators with a flow-matching objective, it delivers discretization-agnostic, scalable posterior estimation and accommodates vector-valued parameters. Across benchmarks and a real glaciology task, FNOPE consistently requires far fewer simulations than baselines while preserving calibration and predictive accuracy, extending SBI to new scientific domains. The approach offers a practical pathway for spatiotemporal inference where high-dimensional, non-uniform data are intrinsic to the problem.

Abstract

Simulation-based inference (SBI) is an established approach for performing Bayesian inference on scientific simulators. SBI so far works best on low-dimensional parametric models. However, it is difficult to infer function-valued parameters, which frequently occur in disciplines that model spatiotemporal processes such as the climate and earth sciences. Here, we introduce an approach for efficient posterior estimation, using a Fourier Neural Operator (FNO) architecture with a flow matching objective. We show that our approach, FNOPE, can perform inference of function-valued parameters at a fraction of the simulation budget of state of the art methods. In addition, FNOPE supports posterior evaluation at arbitrary discretizations of the domain, as well as simultaneous estimation of vector-valued parameters. We demonstrate the effectiveness of our approach on several benchmark tasks and a challenging spatial inference task from glaciology. FNOPE extends the applicability of SBI methods to new scientific domains by enabling the inference of function-valued parameters.

Figures (15)

• Figure 1: Overview. FNOPE approximates the posterior over function-valued parameters of a mechanistic model conditioned on function-valued observations. We use a FNO architecture with a flow matching objective to efficiently represent the function-valued parameters, enabling us to estimate extremely high dimensional posterior distributions at arbitrary discretizations of the domain.
• Figure 2: FNOPE architecture. FNOPE is based on several FNO blocks (gray): A FNO block receives the discretization-dependent spectral features of the function-valued parameters and observations as an input and processes them in a linear layer before transforming it back to the original domain via the approximate inverse transformation. A pointwise linear operation on the input is added, along with embeddings of the flow time, point positions, and vector-valued parameters. We expand this setup to several parallel channels and stack layers. The vector-valued velocities are separately estimated via a MLP.
• Figure 3: Linear Gaussian simulator. Sliced-Wasserstein distance (SWD) to ground truth posterior.
• Figure 4: SIRD model.(a) Posterior conditioned on 40 time points. left: Posterior (mean $\pm$ std.) of the time-varying parameter and ground truth parameters (dashed). right: Two dimensional posterior of vector-valued parameters and ground truth parameters (dot). (b) Posterior predictive (mean $\pm$ std.) of infected, recovered and deceased populations with observations marked. (c)upper: MSE of posterior predictive samples to observations. lower: Simulation-based calibration error of diagonal (SBC EoD). 'Lower bound' refers to the SBC EoD for uniformly sampled posterior ranks (details in Appendix \ref{['app:eval_details']}).
• Figure 5: Darcy flow.(a) Ground truth parameter and posterior samples for a simulation budget of $10^4$ training samples (more posterior samples in Fig. \ref{['app_fig:darcy_samples']}). (b) MSE of posterior predictives to the ground truth observation (FNOPE, FNOPE (fix) and FMPE (raw) visually overlap). (c) Simulation-based calibration Error of Diagonal (EoD) for 50 dimensions. (d) Posterior log-probability of ground truth samples normalized by the number of dimensions (higher is better).