Exploiting Gradients in Bayesian Inference of Expensive Simulators

AI summary

Overview

Research area: Simulation-based inference / Bayesian inverse problems, combining Bayesian optimization, Gaussian process surrogate modeling, and automatic differentiation.

Technical level: Advanced. The reader needs familiarity with Gaussian processes, posterior inference, acquisition functions, and forward- versus reverse-mode automatic differentiation.

Scope: The paper shows how simulator gradients can be folded into the BOSIP (Bayesian Optimization for Simulator Inverse Problems) framework via gradient-enhanced Gaussian processes, and measures when that pays off under a fixed simulation budget.

What This Paper Is About

Many scientific simulators are built from differential equations, and a common task is to infer the posterior distribution over their input parameters from noisy real-world observations. When a single simulation takes hours, one cannot explore the posterior exhaustively, so the BOSIP framework uses a Gaussian process surrogate plus Bayesian optimization to pick the most informative parameters to simulate. The authors' question is whether gradient information that modern simulators can now supply — but which is rarely used for inference — can make this process converge faster within the same simulation budget.

Key Contributions

1. A description of how gradients can be incorporated into the GP-based surrogate models inside the BOSIP inference loop, showing that this substantially accelerates posterior convergence and improves approximation quality within a limited simulation budget.
2. Validation across seven benchmarks (analytical toy problems plus PDE-governed simulators), reporting significant and consistent improvements over gradient-free baselines.
3. A cost accounting showing that in problems where parameter dimensionality exceeds output (proxy) dimensionality, reverse-mode differentiation retains its performance benefit after charging for the added computational cost, whereas for forward-mode differentiation the inference speed-up does not outweigh the added costs.
4. A demonstration that the benefits of reverse-mode gradients grow in higher-dimensional parameter spaces, presented as a promising avenue for future work.

Main Findings

• Gradient information accelerates inference: Across the seven benchmarks, gradient-enhanced BOSIP consistently outperformed the gradient-free baseline in posterior approximation quality, and the authors report significant improvement in convergence speed.
• Reverse-mode survives the cost adjustment: For problems where proxy dimensionality is lower than the number of parameters (δ_dim < x_dim), the performance benefit is maintained even after multiplying the simulation count by (1 + δ_dim) to account for gradient cost.
• Forward-mode does not survive the cost adjustment: Where proxy dimensionality exceeds parameter dimensionality (δ_dim ≥ x_dim), the cost-adjusted performance of gradient-enhanced BOSIP is worse than the gradient-free approach; the adjustment multiplies the simulation count by (1 + x_dim).
• Data efficiency is always lower with gradients: Normalizing by total scalar values obtained from the simulator (outputs plus gradient components) — each simulation additionally provides x_dim · δ_dim gradient values, giving an adjustment of (1 + x_dim) — the gradient-enhanced approach is always less data-efficient. In reverse mode this is offset because the added cost is only (1 + δ_dim), where δ_dim < x_dim, so the data obtained is cheaper.
• Higher dimensions favor reverse mode: On the 6D benchmarks, convergence speed increased substantially (the 6D runs at a budget of 500 simulations did not fully converge). For the same cost of a reverse-mode gradient pass, more gradient information is obtained across all input dimensions.
• Proxy design matters numerically: On the SIR benchmark, whose response surface has sharp ridges where gradients attain huge values, the gradient-enhanced surrogate suffered numerical issues; using a proxy variable that smooths the response surface greatly improved performance. The gradient-enhanced runs on this benchmark were truncated to 50 simulations because numerical issues impeded further convergence.
• The core distinction from classical BO: Figure 1 shows, on a toy problem with uniform prior, linear simulator y = x, Gaussian likelihood N(y, 0.2), and observation z_o = 0 (MAP solution x̂ = 0, true posterior N(0, 0.2)), that classical BO concentrates data at the posterior mode, while BOSIP spreads data across high-probability regions to learn the posterior shape.

Methodology in Plain English

The target is the parameter posterior p(x|z_o), which is proportional to the prior times the likelihood, where the likelihood depends on the simulator output f(x). Because the simulator f is expensive, the authors replace it with a Gaussian process surrogate that predicts a "proxy variable" δ = (φ ∘ f)(x) — a transformation of the simulator output chosen to make the surrogate's job easier. The surrogate's predictive distribution is pushed through the likelihood mapping ψ to give a cheap estimate of the posterior and, importantly, a cheap estimate of the uncertainty in that posterior estimate.

That uncertainty drives the active learning loop: a Bayesian optimization procedure (MaxVar acquisition) repeatedly picks the next parameter x that maximizes the posterior variance, evaluates the simulator there, adds the point to the dataset, refits the surrogate, and repeats until the simulation budget is exhausted.

The new ingredient is gradients. Each proxy dimension is modeled as a multidimensional GP that conditions jointly on function values and on their partial derivatives with respect to x, using the augmented mean and kernel built from the standard kernel plus its Jacobian and Hessian (following the gradient-enhanced GP formulation of Solak et al., cited as [24]). The gradient of the proxy is obtained by the chain rule from the analytically known Jacobian of φ and the simulator Jacobian obtained through forward- or reverse-mode automatic differentiation. No new kernel hyperparameters are introduced; only d-dimensional gradient noise deviations are added per proxy dimension. The MaxVar acquisition and the final posterior estimate are then computed by extracting the marginal predictive distribution of δ from the augmented model.

Performance is scored by the total variation distance between the normalized BOSIP posterior approximation and the true benchmark posterior, approximated by Monte Carlo integration with samples drawn from the prior. Each benchmark is repeated 20 times with 3 initial data points each, and medians are reported. Budgets were 100 simulations for analytical problems, 200 for PDE simulators, and 500 for the 6D extensions (the 6D experiments used no kernel hyperparameter optimization). The 6D problems were built by averaging three copies of a 2D simulator, f_6D(x_1:6) = (1/3)(f_2D(x_1:2) + f_2D(x_3:4) + f_2D(x_5:6)), keeping the same likelihood and observation as the 2D case.

Why This Matters

Impact on research: The paper argues that gradient information — increasingly available from automatic differentiation in differential-equation simulators — is an untapped signal for simulation-based inference, and it provides a concrete rule for when it is worth using: when parameter dimensionality exceeds proxy/output dimensionality, so that the cheap reverse mode applies.

Real-world applications:

• Calibrating expensive physical simulators (fluid dynamics, computational biology, ecology, epidemiology are cited as domains) against observed data.
• Estimating physical quantities that cannot be directly measured.
• Plasma physics and fusion-related modeling, supported in part by the EUROfusion Consortium (Grant Agreement No 101052200).
• PDE-governed 2D and 3D simulators, of which the paper benchmarks three.

Industry relevance: Any group that pays per simulator run — engineering design, climate and geoscience modeling, medical device or process simulation — gains a direct sample-efficiency argument, but only under the condition the paper identifies. The result that forward-mode differentiation is not worth its cost is equally actionable: it tells practitioners not to spend budget on gradients that scale with input dimensionality.

Future Directions

• Adapting acquisition functions beyond MaxVar to the augmented predictive distribution; the authors note this requires marginalizing over the gradients and is left for future work.
• Improving Gaussian process scalability, since GP cost grows with the number of data points and the biggest gradient benefits appear in high-dimensional problems that need larger simulation budgets; alternative surrogate models that exploit gradients are suggested.
• Better handling of numerical instability from sharp response-surface ridges, via more careful proxy variable design, gradient clipping, or manually increasing gradient noise deviation hyperparameters (the latter two possibly reducing predictive accuracy).
• Further exploration of higher-dimensional parameter spaces, where the reverse-mode advantage is shown to increase.

Target Audience

Researchers and practitioners working on simulation-based inference, Bayesian calibration of expensive computer models, and Bayesian optimization with Gaussian process surrogates. It is most useful to those who already have a differentiable simulator — particularly one where the number of input parameters exceeds the number of outputs — and who need to decide whether paying for gradients is justified under a fixed computational budget. Readers without a background in Gaussian processes or automatic differentiation will need to consult the cited BOSIP paper [18] for the underlying framework.

Authors’ abstract

Simulators based on differential equations are ubiquitous in science and engineering. They are often used in simulation-based inference to evaluate the posterior distribution of the input parameters based on real-world observations of the simulator outputs. However, inference becomes challenging when individual simulator evaluations are computationally expensive. In such cases, a Bayesian optimization-based active learning approach with Gaussian process surrogate models has been used to maximize the information obtained from a limited simulation budget. Recently, gradients of simulator outputs with respect to input parameters have become increasingly available, yet they are rarely exploited for inference. Even though we only need to learn the simulator input-output relationship, gradient information can provide an additional valuable signal to guide the active learning procedure. This is of particular interest in the case of expensive simulators, when sample efficiency is crucial. In this paper, we demonstrate how incorporating gradient information into the Gaussian process surrogate accelerates Bayesian optimization-based inference under a limited simulation budget. Our results show significant improvement in convergence speed from using gradient information. For reverse-mode differentiation, the inference efficiency gains are maintained when accounting for the additional computational cost. In contrast, for forward-mode differentiation, the inference speed-up does not outweigh the computational costs. These results indicate that gradient-enhanced surrogates are beneficial primarily in problems where the number of parameters exceeds the output dimensionality, where reverse-mode differentiation is efficient.