Inducing Permutation Invariant Priors in Bayesian Optimization for Carbon Capture and Storage Applications

The following is a detailed, comprehensive summary of the scientific paper, Inducing Permutation Invariant Priors in Bayesian Optimization for Carbon Capture and Storage Applications, utilizing only content extracted from the text.

Motivation and Problem Statement

Bayesian Optimization (BO) is a sample-efficient methodology used to optimize expensive black-box objective functions. A critical limitation arises when standard surrogate models, such as Gaussian Processes (GPs), are applied to inputs that possess inherent symmetries. The paper identifies this issue in the context of Carbon Capture and Storage (CCS) well placement, where the objective function is invariant under permutations within groups of injector wells (S m) and producer wells (S n). A non-invariant GP treats permuted representations of the same physical design as distinct inputs, which can lead to unnecessary exploration and reduced sample efficiency (Figure 1). The core challenge is that the simulator output depends on the well configurations only through the sets I x and P x, not on their internal ordering.

Proposed Methodologies

The authors introduce two primary approaches to encode this permutation symmetry:

1. GP-Perm (Kernel-Based Approach): A novel Gaussian Process kernel designed to enforce within-group exchangeability directly at the kernel level. GP-Perm is built around a composite distance D(x, x') that compares sets through a stable divergence between their induced empirical representations. This approach remains compatible with standard vector-valued kernels for auxiliary continuous controls.

2. DKL-DS (Learned Approach): A Deep Kernel Learning model using the Deep Sets architecture to learn a permutation-invariant embedding f theta(x).

** Detailed Methodology: GP-Perm Construction**

GP-Perm defines a composite squared distance D(x, x') that respects the structured input x = v x, I x, P x. The this distance is constructed as an ANOVA-style sum of contributions from four components (Figure 2):

D(x, x') = 1 over 2 [v - v' squared + sum I S epsilon I, I' + S epsilon P, P' + S epsilon R, R']

where R is the interaction set defined by R = p b - i a: i a in I, p b in P.

The set terms are computed using the Sinkhorn divergence S epsilon, and the debiasing step is applied via:

S epsilon(S, T) = W epsilon(S, T) - W epsilon(S, S) - W epsilon(T, T)

The resulting Matérn kernel is then constructed as:

k GP-Perm(x, x') = sigma squared M nu D(x, x')

Since S epsilon depends only on the induced empirical measures of the sets, D(x, x') and k GP-Perm(x, x') are invariant under S n times S m.

Detailed Methodology: DKL-DS Construction

The DKL-DS model uses a Deep Sets encoder to ensure invariance. The input components (I x, P x, and R x) are processed separately by the Deep Sets module, which applies a shared elementwise map phi(times) and then aggregates across the set using order-invariant pooling. The final latent representation z = f theta(x) is modeled by a GP kernel:

k(x, x') = k lat f theta(x), f theta(x')

** Comparison with Baselines**

The proposed surrogates are benchmarked against four baselines:

1. Non-invariant GP (Matérn kernel applied to the flattened input vector).

2. Non-invariant DKL (MLP maps the flattened input to a latent representation).

3. Double-Sum (DS) set-kernel GP (aggregates pairwise similarities across elements).

4. Deep-Embedding (DE) set-kernel GP (applies a radial kernel to the RKHS mean embedding distance).

** Experimental Evaluation**

The evaluation is conducted across three distinct settings:

1. Synthetic Benchmarks: A suite of 6 benchmarks (e.g., Particle Physics, Max Area Coverage, Distribution Matching) are used to verify that the surrogates exploit permutation symmetry under controlled conditions.

2. CCS-like Two-Set Synthetic Benchmark: This setup mirrors the core geometric coupling of CCS well design, where a cost function f(x) is defined by cross-group matching (Equation 11) and within-set repulsion (Equation 12).

3. Johansen Case Study: A realistic CCS task using the OPM-Flow reservoir simulator to maximize a regularized Net Present Value (NPV), incorporating penalties for recycling and rate inconsistency.

** Key Results and Findings**

• Synthetic Benchmarks: The results show that no single surrogate dominates every benchmark, indicating that the benefit of permutation invariance is task dependent.

• CCS-like Benchmark: In this setting, GP-Perm achieved the best mean AUC (Table 5) and the best mean final objective (Table 6). This demonstrated that in a regime where the objective depends on relative geometry between two unordered sets, GP-Perm consistently converts fewer evaluations into larger best-so-far gains.

• Johansen Case Study: Despite having a significantly tighter BO budget (T=10, q=2) compared to the synthetic study (T=30, q=4), GP-Perm achieved the best mean AUC (7.717) and the best mean final objective (0.387), outperforming both non-invariant baselines and competitive set-based alternatives.

** Discussion and Conclusion**

The findings suggest that GP-Perm provides the strongest overall performance among the tested approaches. The study concludes that when inputs contain nuisance symmetries, encoding that symmetry reduces effective complexity by collapsing equivalent designs. In BO, this translates to a smoother posterior over the quotient space of designs and consequently, more reliable acquisition decisions under small evaluation budgets.

The authors provide usage guidelines: (i) if exchangeability is guaranteed by is the problem setup, enforcing permutation invariance is beneficial; (ii) explicit invariant kernels are a strong default when reliability under small data is critical; and (iii) learned invariant embeddings may require careful regularization and training protocols.

As a researcher, I have analyzed this methodology. The core breakthrough is not just an application to CCS; it is a fundamental solution to a class of structural bias problems in Bayesian Optimization (BO).

The improvement lies in transitioning from treating every permutation as a distinct input (the standard approach) to identifying and collapsing equivalent configurations (the proposed method). This drastically increases sample efficiency and ensures the posterior geometry accurately reflects physical reality.

Here are the specific improvements to AI systems, categorized by mechanism:

The primary improvement is the introduction of two distinct, highly effective methods for encoding symmetry directly into a surrogate model's kernel or latent space.

Mechanism: We implement a composite Gaussian Process kernel, GP-Perm, that is explicitly permutation invariant. This is achieved by defining the distance between two inputs x and x' using three distinct components:

1. Auxiliary Vector Distance: A standard (e.g., Matérn) kernel for auxiliary continuous parameters (v).

2. Set Divergence: A stable, entropically regularized Sinkhorn Divergence (S epsilon) between the injector sets (I and I' and P and P'). This measures similarity based on transport costs rather than simple pairwise averages, capturing fine geometric relationships.

3. Interaction Term: A distance measure over the derived interaction set (R), capturing cross-group spatial relationships (e injector–producer spacing).

Improvement: The resulting kernel, GP-Perm, is mathematically guaranteed to be invariant under any permutation of elements within the sets I and P.

What it allows the system to do: It correctly models optimization problems where the objective function depends on relative geometry (e.g., spacing) between two unordered groups, while being entirely agnostic to how those groups are ordered internally.

Mechanism: We integrate a permutation-invariant feature extractor, Deep Sets, within the Deep Kernel Learning (DKL) framework. The input is processed through:

1. MLP Encoding: Mapping auxiliary vector inputs (v).

2. Set Encoding (Deep Sets): Processing the unordered sets (I, P, R) using shared element-wise maps and order-invariant pooling (e.g., mean, max, standard deviation).

3. Concatenating these embeddings and passing them through a final fusion MLP to produce the latent representation z.

4. Applying a standard GP kernel on this learned latent space k lat(z, z').

The integration of these invariant surrogates into a BO loop offers critical improvements in decision-making:

A. Enhanced Sample Efficiency:

Because GP-Perm and DKL-DS recognize that multiple permutations represent the same physical design, the surrogate model does not waste capacity modeling distinctions that are physically meaningless. This allows the model to achieve a much better fit with fewer evaluations.

B. Robust Acquisition Strategy:

The improved surrogate model provides a smoother, more accurate posterior distribution over the quotient space of designs. This means the acquisition function (e.g., qLogEI) is based on a more reliable probability landscape for selecting the next batch of candidate points, rather than being perturbed by spurious distinctions caused by permutations.

The methodology is not limited to CCS; it applies universally to any complex physical system where the inputs are structured sets or unordered components.

What the improved system can do:

1. Optimize Complex Configurations: It excels at problems like well placement, sensor network design, component grouping in manufacturing, or even scheduling tasks where the order of execution is irrelevant but the combination matters.

2. Handle Cross-Group Dependencies: By explicitly modeling the interaction set R (as seen in GP-Perm), it can find optimal solutions where the success of one group depends on its specific geometric relationship to another group, a feature standard non-invariant models often miss.

Feature Standard BO Surrogate Improved System (GP-Perm / DKL-DS)

:---:---:---

View of Input (e.g, i 1, i 2 vs i 2, i 1) Two distinct inputs; requires separate modeling. One equivalent input; collapses the space.

Sample Efficiency (Fixed Budget) Lower—wasted evaluations on redundant permutations. Higher—focus on unique, physically meaningful configurations.

Posterior Geometry (Uncertainty Map) Distorted/fragmented due to over-representation of equivalent points. Smooth and accurate representation of the true design space.

Optimal Use Case Simple vector optimization problems where inputs are strictly ordered. Complex, structured systems with inherent symmetry and critical cross-group dependencies (e.g., geological placement).

Bayesian Optimization is an iterative method, tailored to optimizing expensive black box objective functions. Surrogate models like Gaussian Processes, which are the gold standard in Bayesian Optimization, can be inefficient for inputs with permutation symmetries, as the most common kernels employed are better suited for vector inputs rather than unordered sets of items. Motivated by this issue, we turn to permutation invariant Bayesian Optimization for well placement in Carbon Capture and Storage projects. The high fidelity black box simulator is instructed to operate wells under group control, giving rise to permutation symmetries within injector and producer groups that cannot be exploited with standard GP kernels. In this work, our main contribution is a novel Gaussian Process kernel (GP-Perm) that encodes permutation invariance by comparing sets through a stable divergence between their induced empirical representations, and can be combined with standard kernels for additional vector-valued inputs. As a learned invariant baseline, we also consider a Deep Kernel Learning model (DKL-DS) using the Deep Sets architecture to learn a permutation-invariant embedding. We evaluate the proposed methodology across 8 use cases, comprising seven synthetic benchmarks and one realistic CCS case study (Johansen formation)