Computer Science Mathematics
Computer Science-2021.gz:143501
Computer Science
aef6e2c5722256401bb09976d77dcd49806c1e73
End-to-End Learning of Deep Kernel Acquisition Functions for Bayesian Optimization

For Bayesian optimization (BO) on high-dimensional data with complex structure, neural network-based kernels for Gaussian processes (GPs) have been used to learn flexible surrogate functions by the high representation power of deep learning. However, existing methods train neural networks by maximizing the marginal likelihood, which do not directly improve the BO performance. In this paper, we propose a meta-learning method for BO with neural network-based kernels that minimizes the expected gap between the true optimum value and the best value found by BO. We model a policy, which takes the current evaluated data points as input and outputs the next data point to be evaluated, by a neural network, where neural network-based kernels, GPs, and mutual information-based acquisition functions are used as its layers. With our model, the neural network-based kernel is trained to be appropriate for the acquisition function by backpropagating the gap through the acquisition function and GP. Our model is trained by a reinforcement learning framework from multiple tasks. Since the neural network is shared across different tasks, we can gather knowledge on BO from multiple training tasks, and use the knowledge for unseen test tasks. In experiments using three text document datasets, we demonstrate that the proposed method achieves better BO performance than the existing methods.


Introduction
Bayesian optimization (BO) is an approach for the global optimization of black-box functions that are expensive to evaluate [32,7,39]. BO has been successfully used for a wide variety of applications, such as computer vision [11], recommendations [16,14], chemical design [18], material science [38,46], probabilistic programs [33], and the automatic selection of machine learning algorithms [41,4,44,27,26]. Gaussian processes (GPs) [34] are commonly used as surrogates to model a distribution over target functions. The next data point to be evaluated is selected by an acquisition function calculated using GPs in such for finding optimal points with fewer target function evaluations.
With GPs, it is important to use appropriate kernel functions. Neural networks have been used as components of kernel functions [55,21,8,53,54], which enable us to learn flexible kernel functions for high-dimensional data with complex structure by the high representation power of deep learning. For example, neural network-based kernels are obtained by transforming the input of RBF kernels by neural networks. However, many training data are required for training neural networks. Therefore, the use of neural networks in BO is limited, since obtaining a large amount of training data is expensive.
To learn expressive kernels with a small amount of training data in a test task, transfer learning and meta-learning methods have been proposed [56,6,52,22,20,45,23,31]. With such methods, the knowledge learned from training tasks is transferred to test tasks, where the training tasks are related to but different from test tasks. These existing methods train neural network-based GPs by maximizing the marginal likelihood. Therefore, although they improve the performance of estimating target functions, they do not necessarily improve the BO performance since it is not directly optimized, and the acquisition function to be used is not considered for training.
In this paper, we propose a meta-learning method for training neural network-based kernels by directly maximizing the BO performance with the consideration of the acquisition function to be used. We model a policy, which takes the current evaluated data points as input and outputs the next data point to be evaluated, a neural network, where neural arXiv:2111.00639v1 [stat.ML] 1 Nov 2021 Figure 1: Framework of our proposed method. We are given multiple tasks for training, where feature vectors and their response values are known. Training tasks are used for training our model (deep kernel acquisition function), which consists of neural network-based kernels, Gaussian processes, and mutual information-based acquisition functions. Our model takes a set of evaluated data points as input, and outputs acquisition values. Using the acquisition values, a data point is selected to be evaluated next. For each training epoch, first, a task is randomly selected from the training tasks. Second, we run Bayesian optimization using our model on the selected task and evaluate the gap. Third, the parameters of the neural network-based kernels are updated by minimizing the gap using reinforcement learning, where a loss is backpropagated through the acquisition function and the Gaussian process.
network-based kernels, GPs, and mutual information-based acquisition functions (MI) [9] are used as its layers. By incorporating the MI into the policy network, the neural networks in the kernel function are trained such that the BO performance is improved when MI is used. We call our model the deep kernel acquisition function. Since the predictive mean and variance given the evaluated data points are calculated in a closed form by GP, and the MI is also calculated in a closed form, we can backpropagate the gap through the GP and MI to update the neural network parameters. We can use other acquisition functions, such as expected improvement [29,24] and upper confidence bound [2,42], instead of MI if they are differentiable.
We formulate BO as a Markov decision process, where the current evaluated data points are a state, the next data point to be evaluated is an action, and the BO performance is a reward. For evaluating the BO performance, we use the gap [50,5] which is the difference between the true optimum value and the best value found. The proposed method trains the policy network by minimizing the expected gap using the policy gradient method [43] with a reinforcement learning framework. The expected gap is calculated by running BO procedures on randomly sampled training tasks for each training epoch using an episodic training framework [35,36,40,12,28]. Since the policy network is shared across different tasks, we can gain useful knowledge for BO from a wide variety of training tasks, and use it for unseen test tasks that are different from training tasks. Figure 1 illustrates the training framework of the proposed method.

Related work
BO has been formulated as reinforcement learning problems [49,48]. MetaBO [49] models an acquisition function by a neural network that takes the outputs of Gaussian processes, which are the mean and variance of the response value at a query point, as input. The kernel parameters of GPs are estimated by maximizing the marginal likelihood to fit the training data, and then the neural network parameters are estimated by reinforcement learning. On the other hand, the kernel parameters in the proposed method are estimated to improve the BO performance, which enables us to estimate the mean and variance of the response value that are appropriate for BO. Vishunu et al. [48] uses recurrent neural networks to model a policy that takes the sequence of evaluated data points as input. In their model, they do not use GPs, although GPs have been successfully used for a wide variety of BO problems. Other works [49,48] completely depend on neural networks for modeling acquisition functions, which require much training data. In contrast to all of the existing methods, we can exploit the knowledge on acquisition functions in the BO literature with the proposed method using a mutual information-based acquisition function as a layer in the neural network. It would help to improve the performance especially with a small number of training data since we do not need to train the neural network from scratch. Although many meta-learning methods have been proposed [37,3,12,47,40,15] including GP-based methods [20,45,13,23], they are designed for supervised learning, but not for BO.

Proposed method 3.1 Problem formulation
Suppose that we are given S datasets on BO tasks {{(x sn , y sn )} Ns n=1 } S s=1 in the training phase, where x sn ∈ X is the nth feature vector in the sth task, y sn = f s (x sn ) is its scalar response value by task-specific black-box function f s (·), and N s is the number of examples in the sth dataset. We assume that feature spaces X are identical across tasks. The number of data points can be different across tasks. In the test phase, we are given a set of feature vectors for a test task , where x * n ∈ X is the nth feature vector in the test task. Although the test task is related to some of the training tasks, it is different. Our goal is to find point x ∈ D * that has higher response value f * (x) with fewer queries, where f * (·) is the target function of the test task we want to maximize. For simplicity, we assume that a set of feature vectors for each task is given. When no set is given, the proposed method is applicable by generating a set of feature vectors.

Gap minimization
For each timestep in the sth task, BO selects a data point in D s = {x sn } Ns n=1 to evaluate next using policy function a(·): where n st ∈ {1, · · · , N s } is the data point index of the query at the tth timestep in the sth task, N st ⊆ {1, · · · , N s } is the set of evaluated data point indices until the tth timestep, n / ∈ N st represents a data point index that is not included in evaluated data points N st , and Ψ is the parameters of the policy function. When N st is the evaluated data points, we assume that their feature vectors and response values {(x sn , y sn )} n∈Nst are observed, and omit the dependency of {(x sn , y sn )} n∈Nst in the equations.
For the evaluation measurement of BO, the following expected cumulative gap is used, where E s [·] denotes the expectation over tasks, T is the number of data points to be evaluated, p(N st ) is the probability that the evaluated data points are N st at the tth timestep, and is the difference between true maximum value max n∈{1,··· ,Ns} y sn and the maximum value in the evaluated data points until the tth timesteps max n∈Nst y sn in the sth task. Since we are given the set of feature vectors and their response values {(x sn , y sn )} Ns n=1 for the training tasks, we can calculate gap r s (N st ) in Eq. (3) for the training tasks. The probability of the evaluated data points is factorized using a policy function: where using Eq. (1). The data point with the highest policy function value is included in the evaluated data points. We use the initially given evaluated data points for N s0 . This BO formulation can be regarded as a Markov decision process, where evaluated data points N st is the state, data index to be evaluated the next n s,t+1 is the action, and gap r s (N st ) is the negative reward. Therefore, we can train parameters Ψ of the policy function by a reinforcement learning algorithm described in Section 3.4.

Policy function
We model the policy function in Eq. (1) by incorporating deep kernel learning (DKL) [53], Gaussian process (GP), and a mutual information-based acquisition function (MI) [9] as shown in Figure 2. By using DKL, we can flexibly learn , and kernel parameters α, β, and η in Eqs. (8,9). Finally, acquisition function value a of the query is calculated by mean µ, variance σ 2 , and ξ based on GP-MI in Eq. (6). We omit dependency between ξ and past variance σ 2 . White and gray circles represent observed and hidden variables. Rectangles represent functions.
representations from the given data, which is especially important for such high-dimensional data as text documents. We share neural networks in deep kernels across different tasks, enabling us to learn shared knowledge across tasks, and use it for unseen test tasks. MI has been successfully used for BO tasks. By using MI as a component in our model, we can exploit its useful knowledge designed for BO. Instead of MI, we can incorporate other types of acquisition functions designed for BO if they are differentiable in the proposed method.
With MI, the acquisition function is given by where µ(x; N st , Ψ) and σ 2 (x; N st , Ψ) are the mean and variance at x given N st , and ν > 0 is a hyperparameter. With a GP with a zero mean and kernel function k(·, ·; Ψ), the mean and variance are given by where k st = [k(x, x ; Ψ)] x ∈Nst ∈ R |Nst| is a kernel vector between x and the evaluated data points at timestep t, x,x ∈Nst is the kernel matrix of the evaluated data points with size (|N st | × |N st |), and β > 0 is the observation noise parameter.
We use the RBF kernel with neural network transformation for DKL, where g(·; Θ) is a neural network with parameter Θ, α and η are kernel parameters, and δ is the Kronecker delta function, i.e., δ(x, x ) = 1 if x = x and δ = 0 otherwise. The parameters in the policy function consist of the parameters of the neural networks and the kernel parameters, Ψ = {Θ, α, β, η}. The policy function in Eq. (6) acts as a neural network that takes query point x and set of evaluated data points N st as input, and outputs a scalar value, where neural network g and GP are used as its components. Since the number of evaluated data points N st is different across different timesteps, standard neural networks cannot be used as a policy function when they assume that the input layer's size is fixed. On the other hand, our model can handle different numbers of evaluated data points using the GP framework, where the mean and variance given evaluated data points are calculated in a closed form.

Training procedure
We train parameters Ψ by minimizing the following expected cumulative gap with discount factor γ ∈ (0, 1], Algorithm 1 Training procedure of our model. Randomly sample a task, s ∼ Uniform(1, · · · , S).

4:
Randomly sample initial evaluated data points N s0 .

7:
Sample query data index n st according to Eq. (13). 8: Update evaluated data points N st ← N s,t−1 ∪ n sn .

9:
Calculate negative reward r s (N st ). for t ∈ {1, · · · , T } do 12: Calculate discounted sum of future negative rewards G t = T τ =t γ τ −t r s (N sτ ). Calculate lossL in Eq. (14) and its gradient. 15: Update model parameters Ψ using loss L and its gradient using a stochastic gradient method. 16: end while where evaluated data points N st are selected according to policy a with parameter Ψ. Objective function L(Ψ) is minimized by the policy gradient method [43] using the following gradient, where G t = T τ =t γ τ −t r sτ is the sum of the future gaps with discount factor γ at the tth timestep. The training procedure is shown in Algorithm 1. For each training epoch, we first randomly sample a task in Line 3 and randomly generate initial evaluated data points N s0 ⊂ {1, · · · , N s } in Line 4. Next we generate BO episode n s1 , N s1 , R s1 , · · · , n sT , N sT , R sT using policy a(·; N s,t−1 , Ψ) and transitionp(n st |a(·; N s,t−1 , Ψ)) in Lines 5-10, where we use the following stochastic transition functioñ p(n st = n|a(·; N s,t−1 , Ψ)) = exp(a(xsn;Ns,t−1,Ψ)) instead of the deterministic transition in Eq. (5). Next we calculate the sum of future gaps G t in Lines 11-13. Then, we calculate lossL and update parameters Ψ by minimizing the loss using a stochastic gradient method in Lines 14-15. Since loss L is differentiable, we can backpropagate it through the MI to update the neural network parameters and the kernel parameters. We use the average of the discounted sum of the future rewards as a baseline to reduce the variance [51,17].

Data
We experimentally evaluated the effectiveness of our proposed method using text document datasets with a BO task to find a document that resembles an unknown target document with fewer evaluations. This BO task simulated a recommendation system, where a user seeks a document that matches her preferences by evaluating a limited number of documents. The user's evaluation for each document is the target function, where evaluation is expensive since the documents must be read. We represented a document as bag-of-words x sn = (x sn1 , · · · , x snJ ), or a word frequency vector, where x snj is the number of occurrences of the jth vocabulary term in the nth document for the sth user, and J is the vocabulary size. Each user was a task, where preferences were different across users, but the target functions of different users were related since the vocabulary terms were shared across users, and some users have similar preferences.
We used the following three text document datasets: Reuter, NeurIPS, and WebKB. The Reuter data were constructed from documents that appeared on the Reuters newswire categorized into eight classes 1 [1]. The NeurIPS data were constructed from conference papers in Neural Information Processing Systems 1-12 categorized into 13 research topics 2 . The WebKB data were constructed from web pages categorized into four universities 3 [10]. We omit vocabulary terms that occurred in less than 50 documents as well as documents that contained fewer than 50 words. The number of documents was 1,909 and the number of vocabulary terms was 799 in the 20News data, 592 and 1,655 in the NeurIPS data, and 3,264 and 1,365 in the WebKB data.
We calculated the similarity by the negative Euclidean distance between the representations obtained by neural networks. For the neural network, we used a four-layered feed-forward neural network with 32 hidden units and a rectified linear unit, ReLU(x) = max(0, x), for the activation. The neural network was trained by minimizing the cross-entropy loss using category labels as in the text classification. We used the last hidden layer as the representations of documents. By transforming the word frequency vectors to the representations by the neural network, the similarity reflected the topics of documents that were related to the category labels. For each BO task, we randomly selected a target document and 500 candidate documents, and calculated their similarity, where the target document was considered the user's most preferred document. We generated {20, 40, · · · , 100} training, 20 validation, and 50 test tasks for each dataset. For each validation and test task, an initial evaluated data point was randomly selected. We ran ten experiments for each dataset, and evaluated the performance based on the average of the cumulative gaps,R = 1 ST S s=1 T t=1 r s (N st ), in the test BO tasks.

Proposed method setting
We used a four-layered feed-forward neural network with 32 hidden units and ReLU activation for neural network g in the kernel functions. The discount factor was γ = 0.99. The number of queries was 10 for each task, and the number of initial evaluated data points was |N s0 | = 1. We used three types of acquisition functions: MI, expected improvement (EI) [29,24], and upper confidence bound (UCB) [2,42]. With EI, the acquisition function is given by where φ and Φ are the probability and cumulative density functions of the standard normal distribution, and With UCB, the acquisition function is given by where ν > 0 is the hyperparameter. For the hyperparameter of MI and UCB, we used ν = log 2 10 −6 like in a previous work [9]. We optimized the model parameters using Adam [25] with learning rate 10 −3 and batch size 16. The validation data were used for early stopping, and the maximum number of epochs was 1,000. We pretrained the model parameters by maximizing the marginal likelihood. We implemented the proposed method with PyTorch [30].

Comparing methods
We compared the proposed method with the following five methods: DKL, GP, RL, MetaBO, and Random. DKL is deep kernel learning [53], where GPs with neural network-based kernels were trained by maximizing the marginal likelihood, where y s = (y s1 , · · · , y sNs ) is the vector of response values in the sth task, K s ∈ R Ns×Ns is the kernel matrix between feature vectors in the sth task, and Normal(·|µ, Σ) is the multivariate normal distribution with mean µ and covariance Σ. We used the same neural network architecture with the proposed method. GP is a Gaussian process with an RBF kernel, where kernel parameters α, β, and η were trained by maximizing the marginal likelihood in Eq. (18). In DKL and GP, we used three types of acquisition functions: MI, EI, and UCB.
RL is a reinforcement learning-based method with a policy modeled by neural networks as in a previous work [48]. For the neural networks, we used deep sets [57] for taking a set of evaluated data points as input, where z st is a vector representation of the evaluated data points at the tth timestep, and f RL and g RL are three-layered feed-forward neural networks with 32 hidden units. The policy was modeled by concatenating the representation and a query data point, where [·, ·] represents the concatenation, and h RL is a four-layered feed-forward neural network with 32 hidden units.
MetaBO is a meta-learning method for BO based on reinforcement learning [49], where the acquisition functions were modeled by a neural network that took the GP's outputs and query data point, where h M is a four-layered feed-forward neural network with 32 hidden units, and µ(x sn ) and σ 2 (x sn ) are the mean and variance estimated by the GP with RBF kernels. The kernel parameters were trained by maximizing the marginal likelihood. With RL and MetaBO, we trained neural networks by the policy gradient method, as in our proposed method. Figure 3 shows the average test cumulative gap with different numbers of training tasks. The proposed method achieved the lowest gap in all cases. As the number of training tasks increased, the gap decreased with the proposed method, RL, and MetaBO, indicating that the reinforcement learning-based methods learned how to select queries in BO using a wide variety of training tasks. With DKL, the decrease based on the number of training tasks was smaller than the reinforcement learning-based methods because it was trained by fitting to the training data by maximizing the marginal likelihood. Thus it did not directly improve the BO performance. RL's gap was high, especially when the number of training tasks was small. Since RL needs to train the policy from scratch using training tasks, it requires many training tasks. On the other hand, since MetaBO used GP's result for modeling the policy, it outperformed RL. Furthermore, for modeling the policy, the proposed method uses acquisition functions, i.e., MI, EI, or UCB, which were developed for BO by experts. Therefore, the proposed method outperformed MetaBO and RL. GP's gap was high, indicating that learning kernels with neural networks is important for modeling high-dimensional data. Figure 4 shows the average test gap with different amounts of evaluated data points using MI acquisition functions. The gaps of the proposed method and DKL were almost the same when the number of the evaluated data points was large. However, the gap of the proposed method was better than DKL when the number of the evaluated data points was small. This result demonstrates the effectiveness of training neural networks by minimizing the expected cumulative gap for finding higher response values with fewer evaluations. Although MetaBO's gap was low when the number of evaluated data points was relatively small, the gap's decrease based o the number of evaluated data points was slower than the proposed method. Table 1 shows the computational time in seconds for training with computers with GeForce GTX 1080 Ti GPUs. The training times with DKL and GP, which were trained with the supervised learning framework, were shorter than the  other methods, which were trained with a reinforcement learning framework. It is because reinforcement learning-based methods run BO for each training epoch. The proposed method and MetaBO required more training time than RL since they calculated GPs in their methods.

Conclusion
We proposed a meta-learning method for learning neural network-based kernels for BO to minimize the expected cumulative gap between the true optimum value and the best found value. The proposed method learns knowledge in various training tasks, and can use it for unseen test tasks by sharing neural networks. Since the proposed method incorporates acquisition functions, such as mutual information and expected improvement, for modeling a policy that selects the next data point to be evaluated, we can employ the knowledge in the existing acquisition functions, which enables us to learn policies with fewer training tasks. Although our results are encouraging, we must extend our approach in several future directions. First, we will apply it to other types of high-dimensional structured data, such as images and time-series, using convolutional and recurrent neural networks in kernel functions. Second, we plan to improve our proposed method using reinforcement learning techniques, which includes actor-critic methods [19].