Learning Relevant Contextual Variables Within Bayesian Optimization

TL;DR

Contextual factors often influence black-box objectives in Bayesian optimization, but their relevance is rarely known and some contexts can be controlled at a cost. The paper introduces Sensitivity-Analysis-Driven CBO (SADCBO), which uses a sensitivity-analysis-based variable selection (via Feature Collapsing) to identify a subset of contextual variables to optimize, and employs an adaptive phase switch (based on an early-stopping bound) to alternate between observational and interventional querying. The method builds GP surrogates on the selected context subset and uses standard acquisition (UCB) on the joint space, with a cost-aware adjustment to compare relevance against optimization cost through a per-variable cost term. Empirical results on 4 real-world and 4 synthetic tasks show that SADCBO delivers consistent improvements over baselines, and can be combined with sparsity-promoting surrogates like SAASBO to further boost performance. This framework offers a practical and explainable approach to cost-efficiently learn which contextual variables matter and how to best leverage them in optimization workflows.

Abstract

Contextual Bayesian Optimization (CBO) efficiently optimizes black-box functions with respect to design variables, while simultaneously integrating contextual information regarding the environment, such as experimental conditions. However, the relevance of contextual variables is not necessarily known beforehand. Moreover, contextual variables can sometimes be optimized themselves at an additional cost, a setting overlooked by current CBO algorithms. Cost-sensitive CBO would simply include optimizable contextual variables as part of the design variables based on their cost. Instead, we adaptively select a subset of contextual variables to include in the optimization, based on the trade-off between their relevance and the additional cost incurred by optimizing them compared to leaving them to be determined by the environment. We learn the relevance of contextual variables by sensitivity analysis of the posterior surrogate model while minimizing the cost of optimization by leveraging recent developments on early stopping for BO. We empirically evaluate our proposed Sensitivity-Analysis-Driven Contextual BO (SADCBO) method against alternatives on both synthetic and real-world experiments, together with extensive ablation studies, and demonstrate a consistent improvement across examples.

Figures (10)

• Figure 1: Sensitivity analysis on $\mathcal{D}_t^{\textnormal{BO}}$ characterizes variable importance at the optimum faster than $\mathcal{D}_t$.Top left: 2D black-box objective together with the queries produced along a BO trajectory. Initial samples are represented by empty dark-colored triangles, newly obtained samples as dots with an increasingly lighter color. Top right: Best value found during the optimization trial. Bottom left: Sensitivity indices for $z^{(1)}$ and $z^{(2)}$ averaged over $\mathcal{D}_t^{\text{BO}}$. As we converge to the optimum, $\mathcal{D}_t^{\text{BO}}$ mainly involves samples close to the optimum, leading to a different variable relevance ranking (iteration 30 to the end; $z^{(1)}$ is more relevant) compared to the early iterations (10 to 30; $z^{(2)}$ is more relevant). Bottom right: Sensitivity indices computed on the whole dataset $\mathcal{D}_t$ do not converge as quickly and do not capture the shift in relevance close to the optimum.
• Figure 2: Benchmark of the different methods. (a) On real-world datasets, SADCBO (red curve with white markers) performs on par with other baselines and is the top performer for the Robot Pushing task. (b) On synthetic functions, SADCBO outperforms other baselines in three cases out of four. (c) Histograms of phase switching criteriong time for SADCBO computed for the Hartmann6D (c.1) and Hartmann4D problems (c.2). (d) Inclusion probability of each contextual variable for SADCBO computed for the Hartmann6D (d.1) and Hartmann4D problems (d.2). Each panel shows the mean $\pm 2$ standard error across $N=100$ trials.
• Figure 3: Combining SADCBO with sparsity-enforcing surrogate SAASBO. For any variable, the associated query cost is 1. $p(\mathbf{z}) = \mathcal{U}([0,1]^c)$. The combination is fruitful and improves the performances of SAASBO.
• Figure 4: Assessing SADCBO’s phase switching criterion on the Hartmann6D function. The iteration selected by the adaptive stopping criterion implemented in SADCBO yields one of the best BO trials. Each curve is computed as an average of 10 different random seeds.
• Figure S1: Flowchart of the proposed method SADCBO.