Computationally Efficient High-Dimensional Bayesian Optimization via Variable Selection

We study the problem of globally maximizing a black-box function $f(\mathbf{x})$ with an input domain $\mathcal{X}=[0,1]^{D}$, where the function has some special properties: (1) It is hard to calculate its first and second-order derivatives, therefore gradient-based optimization algorithms are not useful; (2) It is too strong to make additional assumptions on the function such as convexity; (3) It is expensive to evaluate the function, hence some classical global optimization algorithms such as evolutionary algorithms (EA) are not applicable.

Bayesian optimization (BO) is a popular global optimization method to solve the problem described above. It aims to obtain the input $\mathbf{x}^{*}$ that maximizes the function $f$ by sequentially acquiring queries that are likely to achieve the maximum and evaluating the function on these queries. BO has been successfully applied in many scenarios such as hyper-parameter tuning  (Snoek et al., 2012; Klein et al., 2017), automated machine learning (Nickson et al., 2014; Yao et al., 2018), reinforcement learning (Brochu et al., 2010; Marco et al., 2017; Wilson et al., 2014), robotics (Calandra et al., 2016; Berkenkamp et al., 2016), and chemical design (Griffiths and Hernández-Lobato, 2017; Negoescu et al., 2011). However, most problems described above that have been solved by BO successfully have black-box functions with low-dimensional domains, typically with $D\leq 20$ (Frazier, 2018). Scaling BO to high-dimensional black-box functions is challenging because of the following two reasons: (1) Due to the curse of dimensionality, the global optima is harder to find as $D$ increases; and (2) computationally, vanilla BO is extremely time consuming on functions with large $D$. As global optimization for high-dimensional black-box function has become a necessity in several scientific fields such as algorithm configuration (Hutter et al., 2010), computer vision (Bergstra et al., 2013) and biology (Gonzalez et al., 2015), developing new BO algorithms that can effectively optimize black-box functions with high dimensions is very important for practical applications.

A large class of algorithms for high-dimensional BO is based on the assumption that the black-box function has an effective subspace with dimension $d_{e}\ll D$ (Djolonga et al., 2013; Wang et al., 2016; Moriconi et al., 2019; Nayebi et al., 2019; Letham et al., 2020). Therefore, these algorithms first embed the high-dimensional domain $\mathcal{X}$ to a space with the embedding dimension $d$ pre-specified by users, do vanilla BO in the embedding space to obtain the new query, and then project it back and evaluate the function $f$. These algorithms are time efficient since BO is done in a low-dimensional space. Wang et al. (2016) proves that if $d\geq d_{e}$, then theoretically with probability $1$ the embedding space contains the global optimum. However, since $d_{e}$ is usually not known, it is difficult for users to set a suitable $d$. Previous work such as Eriksson and Jankowiak (2021) shows that different settings of $d$ will impact the performance of embedding-based algorithms, and there has been little work on how to choose $d$ heuristically. Letham et al. (2020) also points out that when projecting the optimal point in the embedding space back to the original space, it is not guaranteed that this projection point is inside $\mathcal{X}$, hence algorithms may fail to find an optimum within the input domain.

We develop a new algorithm, called VS-BO (Variable Selection Bayesian Optimization), to solve issues mentioned above. Our method is based on the assumption that all the $D$ variables (elements) of the input $\mathbf{x}$ can be divided into two disjoint sets $\mathbf{x}=\{\mathbf{x}_{ipt},\mathbf{x}_{nipt}\}$: (1) $\mathbf{x}_{ipt}$, called important variables, are variables that have significant effects on the output value of $f$; (2) $\mathbf{x}_{nipt}$, called unimportant variables, are variables that have no or little effect on the output. Previous work such as Hutter et al. (2014) shows that the performance of many machine learning methods is strongly affected by only a small subset of hyperparameters, indicating the rationality of this assumption. We propose a robust strategy to identify $\mathbf{x}_{ipt}$, and do BO on the space of $\mathbf{x}_{ipt}$ to reduce time consumption. In particular, our method is able to learn the dimension of $\mathbf{x}_{ipt}$ automatically, hence there is no need to pre-specify the hyperparameter $d$ as embedding-based algorithms. Since the space of $\mathbf{x}_{ipt}$ is axis-aligned, issues caused by the space projection no longer exist in our method. We theoretically analyze the computational complexity of VS-BO, showing that our method can decrease the computational complexity of both steps of fitting the Gaussian Process (GP) and optimizing the acquisition function. We formalize the assumption that some variables of the input are important while others are unimportant and derive the regret bound of our method. Finally, we empirically show the good performance of VS-BO on several synthetic and real problems.

The source code to reproduce the results from this study can be found at https://github.com/Kingsford-Group/vsbo. This work has been accepted in AutoML 2023, with camera-ready version https://openreview.net/forum?id=QXKWSM0rFCK1.

The basic framework of BO has two steps for each iteration: First, GP is used as the surrogate to model $f$ based on all the previous query-output pairs $\left(\mathbf{x}^{1\mathrel{\mathop{\mathchar 58\relax}}n},y^{1\mathrel{\mathop{\mathchar 58\relax}}n}\right)$:

Here $y^{1\mathrel{\mathop{\mathchar 58\relax}}n}=[y^{1},\dots,y^{n}]$ is a $n$-dimensional vector, $y^{i}=f(\mathbf{x}^{i})+\epsilon^{i}$ is the output of $f$ with random noise $\epsilon^{i}\sim\mathcal{N}(0,\sigma_{0}^{2})$, and $K(\mathbf{x}^{1\mathrel{\mathop{\mathchar 58\relax}}n},\Theta)$ is a $n\times n$ covariance matrix where its entry $K_{i,j}=k(\mathbf{x}^{i},\mathbf{x}^{j},\Theta)$ is the value of a kernel function $k$ in which $\mathbf{x}^{i}$ and $\mathbf{x}^{j}$ are the i-th and j-th queries respectively. $\Theta$ and $\sigma_{0}$ are parameters of GP that will be optimized each iteration, and $\mathbf{I}$ is the $n\times n$ identity matrix. A detailed description of GP and its applications can be found in Williams and Rasmussen (2006).

Given a new input $\mathbf{x}^{\prime}$, we can compute the posterior distribution of $f(\mathbf{x}^{\prime})$ from GP, which is again a Gaussian distribution with mean $\mu(\mathbf{x}^{\prime}\mid\mathbf{x}^{1\mathrel{\mathop{\mathchar 58\relax}}n},y^{1\mathrel{\mathop{\mathchar 58\relax}}n})$ and variance $\sigma^{2}(\mathbf{x}^{\prime}\mid\mathbf{x}^{1\mathrel{\mathop{\mathchar 58\relax}}n})$ that have the following forms:

Here, $\mathbf{k}(\mathbf{x}^{\prime},\mathbf{x}^{1\mathrel{\mathop{\mathchar 58\relax}}n})=[k(\mathbf{x}^{\prime},\mathbf{x}^{1},\Theta),\dots,k(\mathbf{x}^{\prime},\mathbf{x}^{n},\Theta)]$ is a $n$-dimensional vector.

The second step of BO is to use $\mu$ and $\sigma$ to construct an acquisition function $acq$ and maximize it to get the new query $\mathbf{x}^{new}$, on which the function $f$ is evaluated to obtain the new pair $(\mathbf{x}^{new},y^{new})$:

A wide variety of methods have been proposed that are related to high-dimensional BO, and most of them are based on some extra assumptions on intrinsic structures of the domain $\mathcal{X}$ or the function $f$. As mentioned in the previous section, a considerable body of algorithms is based on the assumption that the black-box function has an effective subspace with a significantly smaller dimension than $\mathcal{X}$. Among them, REMBO (Wang et al., 2016) uses a randomly generated matrix as the projection operator to embed $\mathcal{X}$ to a low-dimensional subspace. SI-BO (Djolonga et al., 2013), DSA (Ulmasov et al., 2016) and MGPC-BO (Moriconi et al., 2019) propose different ways to learn the projection operator from data, of which the major shortcoming is that a large number of data points are required to make the learning process accurate. HeSBO (Nayebi et al., 2019) uses a hashing-based method to do subspace embedding. Finally, ALEBO (Letham et al., 2020) aims to improve the performance of REMBO with several novel refinements.

Another assumption is that the black-box function has an additive structure. Kandasamy et al. (2015) first develops a high-dimensional BO algorithm called Add-GP by adopting this assumption. They derive a simplified acquisition function and prove that the regret bound is linearly dependent on the dimension. Their framework is subsequently generalized by Li et al. (2016); Wang et al. (2017) and Rolland et al. (2018).

As described in the previous section, our method is based on the assumption that some variables are more “important" than others, which is similar to the axis-aligned subspace embedding. Several previous works propose different methods to choose axis-aligned subspaces in high-dimensional BO. Li et al. (2016) uses the idea of dropout, i.e, for each iteration of BO, a subset of variables are randomly chosen and optimized, while our work chooses variables that are important in place of the randomness. Eriksson and Jankowiak (2021) develops a method called SAASBO, which uses the idea of Bayesian inference. SAASBO defines a prior distribution for each parameter in the kernel function $k$, and for each iteration the parameters are sampled from posterior distributions and used in the step of optimizing the acquisition function. Since those priors restrict parameters to concentrate near zero, the method is able to learn a sparse axis-aligned subspaces (SAAS) during BO process. Similar to vanilla BO, the main drawback of SAASBO is that it is very time consuming. While traditionally it is assumed that the function $f$ is very expensive to evaluate so that the runtime of BO itself does not need to be considered, previous work such as Ulmasov et al. (2016) points out that in some application scenarios the runtime of BO cannot be neglected. Spagnol et al. (2019) proposes a similar framework of high-dimensional BO as us; they use Hilbert Schmidt Independence criterion (HSIC) to select variables, and use the chosen variables to do BO. However, they do not provide a comprehensive comparison with other high-dimensional BO methods: their method is only compared with the method in Li et al. (2016) on several synthetic functions. In addition, they do not provide any theoretical analysis.

Given the black-box function $f(\mathbf{x})\mathrel{\mathop{\mathchar 58\relax}}\mathcal{X}\to\mathbb{R}$ in the domain $\mathcal{X}=[0,1]^{D}$ with a large $D$, the goal of high-dimensional BO is to find the maximizer $\mathbf{x}^{*}=\text{argmax}_{\mathbf{x}\in\mathcal{X}}f(\mathbf{x})$ efficiently. As mentioned in the introduction, VS-BO is based on the assumption that all variables in $\mathbf{x}$ can be divided into important variables $\mathbf{x}_{ipt}$ and unimportant variables $\mathbf{x}_{nipt}$, and the algorithm uses different strategies to decide values of the variables from two different sets.

The high-level framework of VS-BO (Algorithm 1) is similar to Spagnol et al. (2019). For every $N_{vs}$ iterations VS-BO will update $\mathbf{x}_{ipt}$ and $\mathbf{x}_{nipt}$ (line 8 in Algorithm 1), and for every BO iteration $t$ only variables in $\mathbf{x}_{ipt}$ are used to fit GP (line 11 in Algorithm 1 ), and the new query of important variables $\mathbf{x}_{ipt}^{t}$ is obtained by maximising the acquisition function (line 12 in Algorithm 1). Unlike Spagnol et al. (2019), VS-BO learns a conditional distribution $p(\mathbf{x}_{nipt}\mid\mathbf{x}_{ipt},\mathcal{D})$ from the existing query-output pairs $\mathcal{D}$ (line 5, 9 in Algorithm 1). This distribution is used for choosing the value of $\mathbf{x}_{nipt}$ to make $f(\mathbf{x})$ large when $\mathbf{x}_{ipt}$ is fixed. Hence, once $\mathbf{x}_{ipt}^{t}$ is obtained, the algorithm samples $\mathbf{x}_{nipt}^{t}$ from $p(\mathbf{x}_{nipt}\mid\mathbf{x}_{ipt}^{t},\mathcal{D})$ (line 13 in Algorithm 1), concatenates it with $\mathbf{x}_{ipt}^{t}$ and evaluates $f(\{\mathbf{x}_{ipt}^{t},\mathbf{x}_{nipt}^{t}\})$.

Compared to Spagnol et al. (2019), our method is new on the following three aspects: First, we propose a new variable selection method that takes full advantage of the information in the fitted GP model, and there is no hyperparameter that needs to be pre-specified in this method; Second, we develop a new mechanism, called VS-momentum, to improve the robustness of variable selection; Finally, we integrate an evolutionary algorithm into the framework of BO to make the sampling of unimportant variables more precise. The following subsections introduce these three points in detail.

From the theoretical perspective, we prove that running BO by only using those important variables is able to decrease the runtime of both the step of fitting the GP and maximizing the acquisition function. Specifically, we have the following proposition:

The proof is in section B of the appendix. Note that the method for fitting the GP and maximizing the acquisition function under the framework of BoTorch is limited-memory BFGS, which is indeed a QN method. Since the complexity is related to the quadratic of $p$, selecting a small subset of variables (so that $p$ is small) can decrease the runtime of BO. Figure 6 empirically shows that compared to vanilla BO, VS-BO can both reduce the runtime of fitting a GP and optimizing the acquisition function, especially when $n$ is not small.

Let $\mathbf{x}^{*}$ be one of the maximal points of $f(\mathbf{x})$. To quantify the efficacy of the optimization algorithm, we are interested in the cumulative regret $R_{N}$, defined as: $R_{N}=\sum_{t=1}^{N}\left[f(\mathbf{x}^{*})-f(\mathbf{x}^{t})\right]$ where $\mathbf{x}^{t}$ is the query at iteration $t$. Intuitively, the algorithm is better when $R_{N}$ is small, and a desirable property is to have no regret: $\lim_{N\to\infty}R_{N}/N=0$. Here, we provide an upper bound of the cumulative regret for a simplified VS-BO algorithm, called VS-GP-UCB (Algorithm 6 in the appendix). Similar to Srinivas et al. (2009), for proving the regret bound we need the smoothness assumption of the kernel function. In addition, we have the extra assumption that the $D-d$ variables in $\mathbf{x}$ are unimportant (for convenience we index unimportant variables from $d+1$ to $D$ without loss of generality), meaning the absolute values of partial derivatives of $f$ on those $D-d$ variables are in general smaller than those on important variables. Formally, we have the following assumption:

In VS-GP-UCB, the values of unimportant variables are fixed in advance (line 4 in Algorithm 6), denoted as $\mathbf{x}^{0}_{[d+1\mathrel{\mathop{\mathchar 58\relax}}D]}$, and the important variables are queried at each iteration by maximizing the acquisition function upper confidence bound (UCB) (Auer, 2002) with those fixed unimportant variables (line 6 in Algorithm 6). We have the following regret bound theorem of VS-GP-UCB:

Here, $\gamma_{N}\mathrel{\mathop{\mathchar 58\relax}}=\max_{A\subset\mathcal{X}\mathrel{\mathop{\mathchar 58\relax}}|A|=N}\mathbf{I}(\mathbf{y}_{A};\mathbf{f}_{A})$ is the maximum information gain with a finite set of sampling points $A$, $\mathbf{f}_{A}=[f(\mathbf{x})]_{\mathbf{x}\in A}$, $\mathbf{y}_{A}=\mathbf{f}_{A}+\epsilon_{A}$, and $C_{1}=\frac{8}{\log\left(1+\sigma_{0}^{-2}\right)}$.

The proof of Theorem 5.1 is in section C of the appendix. Srinivas et al. (2009) upper bounded the maximum information gain for some commonly used kernel functions, for example they prove that by using SE kernel with the same length scales ($\rho_{i}=\rho_{0}$ for all $i$), $\gamma_{N}=\mathcal{O}\left((\log N)^{D+1}\right)$ so that $\lim_{N\to\infty}(\beta_{N}\gamma_{N})/N=0$. However, every variable is equally important in the SE kernel with the same length scales, which does not obey Assumption 5.1. We hypothesize that by using SE kernel of which the length scales are different such that Assumption 5.1 is satisfied, the statement that $\lim_{N\to\infty}(\beta_{N}\gamma_{N})/N=0$ is also correct, although we do not have proof here.

Li et al. (2016) also derives a regret bound for its dropout algorithm (Lemma 5 in Li et al. (2016)). Compared to the regret bound in Srinivas et al. (2009) (Theorem 2 in Srinivas et al. (2009)), both Li et al. (2016) and our work have an additional residual in the bound, while ours contains a small coefficient $\alpha$. In the case when $\alpha\to 0$, the bound in Theorem 5.1 is the same as that in theorem 2 of Srinivas et al. (2009) and there is no regret. These results show the necessity of the variable selection since it can help decrease the value of $\alpha$. In addition, compared to fixing unimportant variables in VS-GP-UCB, sampling from the CMA-ES posterior may further decrease the residual value.

We compare VS-BO to a broad selection of existing methods: vanilla BO, REMBO and its variant REMBO Interleave, Dragonfly, HeSBO and ALEBO. The details of implementations of these methods as well as hyperparameter settings are described in section D of the appendix.

We propose a new method, VS-BO, for high-dimensional BO that is based on the assumption that variables of the input can be divided to two categories: important and unimportant. Our method can assign variables into these two categories with no need for pre-specifying any crucial hyper-parameter and use different strategies to decide values of the variables in different categories to reduce runtime. The good performance of our method on synthetic and real-world problems further verify the rationality of the assumption. We show the computational efficiency of our method both theoretically and empirically. In addition, information from the variable selection improves the interpretability of BO model: VS-BO can find important variables so that it can help increase our understanding of the black-box function. We also notice that in practice vanilla BO under the framework of BoTorch usually has a good performance if the runtime of BO does not need to be considered, especially when the dimension is not too large ($D<100$). However, this method is usually not considered as a baseline to compare with in previous high-dimensional BO work.

We also find some limitations of our method when running experiments. First, when the dimension of the input increases, it becomes harder to do variable selection accurately. Therefore, embedding-based methods are still the first choice when the input of a function has thousands of dimensions. It might be interesting to develop new algorithms that can do variable selection robustly even when the dimension is extremely large. Further, Grad-IS might be invalid when variables are discrete or categorical, therefore new methods for calculating the importance score of these kinds of variables are needed. These are several directions for future improvements of VS-BO. It is also interesting to do further theoretical study such as investigating the bounds on the maximum information gain of kernels that satisfy Assumption 5.1.

1. 1.

   For all authors…

   1. (a)

      Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope? [Yes] The last paragraph of the introduction section describes the paper’s contributions.

   2. (b)

      Did you describe the limitations of your work? [Yes] The last paragraph of the conclusion section describes some limitations of our work.

   3. (c)

      Did you discuss any potential negative societal impacts of your work? [N/A] There will be no potential negative societal impacts of our work.

   4. (d)

      Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes] I read it carefully.

2. 2.

   If you are including theoretical results…

   1. (a)

      Did you state the full set of assumptions of all theoretical results? [Yes] It is in section 5.

   2. (b)

      Did you include complete proofs of all theoretical results? [Yes] It is in section B and C.

3. 3.

   If you ran experiments…

   1. (a)

      Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)? [Yes] Codes will be attached in the supplementary material.

   2. (b)

      Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Yes] They are described in section D.

   3. (c)

      Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)? [Yes] For each case we repeat $20$ independent runs, and we plot mean and standard deviation in our figures.

   4. (d)

      Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)? [Yes] They are described in section D.

4. 4.

   If you are using existing assets (e.g., code, data, models) or curating/releasing new assets…

   1. (a)

      If your work uses existing assets, did you cite the creators? [Yes]

   2. (b)

      Did you mention the license of the assets? [N/A]

   3. (c)

      Did you include any new assets either in the supplemental material or as a URL? [Yes] They are described in section D.

   4. (d)

      Did you discuss whether and how consent was obtained from people whose data you’re using/curating? [N/A]

   5. (e)

      Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [N/A]

5. 5.

   If you used crowdsourcing or conducted research with human subjects…

   1. (a)

      Did you include the full text of instructions given to participants and screenshots, if applicable? [N/A]

   2. (b)

      Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable? [N/A]

   3. (c)

      Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation? [N/A]