Bayesian experimental design: grouped geometric pooled posterior via ensemble Kalman methods

Abstract: Bayesian experimental design (BED) for complex physical systems is often limited by the nested inference required to estimate the expected information gain (EIG) or its gradients. Each outer sample induces a different posterior, creating a large and heterogeneous set of inference targets. Existing methods have to sacrifice either accuracy or efficiency: they either perform per-outer-sample posterior inference, which yields higher fidelity but at prohibitive computational cost, or amortize the inner inference across all outer samples for computational reuse, at the risk of degraded accuracy under posterior heterogeneity. To improve accuracy and maintain cost at the amortized level, we propose a grouped geometric pooled posterior framework that partitions outer samples into groups and constructs a pooled proposal for each group. While such grouping strategy would normally require generating separate proposal samples for different groups, our tailored ensemble Kalman inversion (EKI) formulation generates these samples without extra forward-model evaluation cost. We also introduce a conservative diagnostic to assess importance-sampling quality to guide grouping. This grouping strategy improves within-group proposal-target alignment, yielding more accurate and stable estimators while keeping the cost comparable to amortized approaches. We evaluate the performance of our method on both Gaussian-linear and high-dimensional network-based model discrepancy calibration problems.

• The paper introduces a grouped geometric pooled posterior (GPP) framework that leverages Ensemble Kalman Inversion (EKI) to significantly cut forward model evaluations in Bayesian experimental design.
• It employs adaptive clustering and a conservative ESS-based diagnostic to mitigate proposal-target mismatches and weight degeneracy, enhancing gradient estimation reliability.
• Experimental results demonstrate lower estimator variance and error reduction in both low- and high-dimensional settings compared to traditional methods.

Grouped Geometric Pooled Posterior for Bayesian Experimental Design via Ensemble Kalman Methods

Context and Motivation

Bayesian Experimental Design (BED) in complex physical systems is dominated by the challenge of estimating expected information gain (EIG) or its gradients, necessitating repeated, computationally intensive posterior inferences across highly heterogeneous outer samples. Traditional approaches—per-outer-sample nested Monte Carlo, SMC, or local Laplace approximations—provide high-fidelity inference but exhibit prohibitive computational burden, scaling with the number of outer samples and the cost of forward model evaluations. Amortized inference strategies attempt computational reuse, e.g., via shared variational surrogates or global proposal distributions, but these often degrade in accuracy due to posterior heterogeneity, especially in the tails.

Grouped Geometric Pooled Posterior: Methodological Advances

This work introduces a grouped geometric pooled posterior (GPP) framework for BED, which navigates the trade-off between accuracy and computation by hierarchically partitioning outer samples into groups with similar posteriors. For each group, the method constructs a group-specific GPP, mitigating proposal-target mismatch and importance weight degeneracy associated with a global, ‘one-size-fits-all’ proposal. The proposal-target fit is quantified and controlled via an adaptive, conservative ESS-based diagnostic, and grouping is performed using a computationally efficient cluster-based strategy.

A major computational bottleneck in previous GPP and nested approaches is the need for forward model evaluations to generate proposal samples for each group. The authors resolve this by tailoring ensemble Kalman inversion (EKI) to generate samples for all grouped GPPs with constant forward model cost—forward evaluations are required only in the prediction step, and subsequent group-specific ensemble updates exploit linear algebra and can reuse model outputs.

EKI-based Sampling Algorithms

In both Gaussian-linear and nonlinear settings, the paper demonstrates that EKI can drive a prior ensemble to match the GPP (or grouped GPP) by leveraging either:

• Stacked-observation formulation: Augmenting the state such that each group’s GPP corresponds to a Gaussian with a block-diagonal covariance, performing a standard EKI update in the high-dimensional stacked space.
• Mean-observation formulation: Deriving the same result via efficient precision-weighted mean and covariance updates, computationally advantageous in homoskedastic BED regimes where all observations share the same noise model.

This EKI-driven group-based proposal yields a robust and numerically stable gradient estimate for BED optimization, and its asymptotic properties are established in the linear-Gaussian regime.

Diagnostic-Guided Grouping with Conservative ESS Bounds

The core of the grouping algorithm is a conservative, analytic lower bound on the ESS for each outer sample, using a log-normal model of importance weights and the predictive covariance from the EKI forecast ensemble. Problematic outer samples (with estimated ESS below threshold) are adaptively partitioned via unsupervised clustering (typically in the observation space), after which each group’s GPP is generated by reusing the EKI forecast and forming group-specific ensemble updates—crucially, with no extra forward-model cost.

Figure 1: Histogram of all outer‐sample values $y$, annotated by their selection status under the proposed ESS diagnostics. Left orange+red region is Group 1, right orange+red region is Group 2, and blue bins form Group 3.

Experimental Assessment

Low-dimensional Parametric Error Case

In a source inversion PDE testbed with a parametric error in source strength, the authors compare GPP-EKI, AD-EKI, and contrastive diffusion samplers. All approaches recover posteriors and designs consistent with ground truth, with the GPP-EKI ensemble matching the analytic pooled posterior closely.

Figure 2: Posteriors of the physical parameter. From top to bottom, contrastive diffusion, AD-EKI, and GPP-EKI. The beliefs are defined on $[0,1]^2$ and shown zoomed in $[0, 0.6]^2$ for clarity.

Figure 3: Updating trajectory of the error parameter.

GPP-EKI achieves several-fold reduction in PDE solves compared to baselines. The Wasserstein distance between individual posteriors and the GPP-proposal is systematically smaller than that to the prior, but proposal-target mismatches for tail samples are still visible, motivating group-based refinements.

Figure 4: Samples for geometric pooled posterior drawn by EKI

High-dimensional Neural Network Structural Error Case

For structural model discrepancy parameterized by a high-dimensional neural network, conventional ungrouped GPP or AD-EKI becomes computationally infeasible or leads to weight degeneracy. The grouped GPP-EKI approach demonstrates:

• Posterior consistency: The grouped method yields physical and error posteriors as concentrated and unbiased as AD-EKI, outperforming random design selection.

  Figure 5: Posteriors of the physical parameter. From top to bottom, random selection, modified AD-EKI, and GPP-EKI. The beliefs are defined on $[0,1]^2$ and shown zoomed in $[0, 0.6]^2$ for clarity.

• Physical solution accuracy: The grouped method attains significantly lower field errors in reconstructed solutions, particularly near experimental designs.

  Figure 6: Relative error between the solution field and the true field. From top to bottom: contrastive diffusion, AD-EKI, and GPP-EKI.

• Sharp reduction in estimator variance: The standard deviation of the EIG gradient estimator is reduced by more than a factor of three compared to the ungrouped case at identical sample budget—a direct consequence of group-specific pooling reducing importance weight degeneracy.
• Reduced PDE cost: The grouped proposal achieves performance comparable to per-outer-sample AD-EKI and diffusion methods at orders-of-magnitude fewer forward model evaluations.
• Proposal-target fit gain: For tail groups, the grouped GPP is substantially closer to individual posteriors than the global proposal, measured by Wasserstein distances or corner plots on selected NN parameters.

Figure 7: Group 1.

Figure 8: Corner plot of selected coordinates of the parameter vector, comparing individual, global pooled, and grouped pooled posterior samples.

Implications, Limitations, and Prospects

The grouped GPP-EKI framework advances the state of BED in high-dimensional, computationally intensive inverse problems by balancing proposal flexibility and amortization. The group-based proposal matches the heterogeneity of posteriors in the outer-sample ensemble without introducing additional forward-model cost, a property unattainable with MCMC or diffusion-based proposal generation.

The practical impact is immediate in scientific and engineering inverse problems: sequential model discrepancy calibration, optimal sensor placement, or parameterization of nonlinear dynamical systems with high-dimensional error models. Theoretically, the approach provides a scalable recipe for constructing importance-sampling proposals via ensemble methods that adaptively track local posterior geometry. The conservative ESS-diagnostic also enables safe, robust automation of proposal adaptation.

Potential future development directions include hybridization with local refinements, dynamic group resizing, or integration with emulator-based or neural surrogate forward models. For AI fields relying on expensive model-based active learning, design, or uncertainty quantification, this amortized local-pooling methodology offers a template for scaling beyond traditional nested inference.

Conclusion

This work presents a scalable and robust methodology for BED in the presence of heterogeneous and high-dimensional posterior families, built on grouped GPPs generated via EKI. The approach enables substantial computational savings, statistically reliable gradient estimation, and adaptivity in complex design problems, positioning it as a frontier technique for scalable experimental design under computational constraints (2604.18505).