Bayesian Variational Optimization for Combinatorial Spaces

This paper focuses on Bayesian Optimization in combinatorial spaces. In many applications in the natural science. Broad applications include the study of molecules, proteins, DNA, device structures and quantum circuit designs, a on optimization over combinatorial categorical spaces is needed to find optimal or pareto-optimal solutions. However, only a limited amount of methods have been proposed to tackle this problem. Many of them depend on employing Gaussian Process for combinatorial Bayesian Optimizations. Gaussian Processes suffer from scalability issues for large data sizes as their scaling is cubic with respect to the number of data points. This is often impractical for optimizing large search spaces. Here, we introduce a variational Bayesian optimization method that combines variational optimization and continuous relaxations to the optimization of the acquisition function for Bayesian optimization. Critically, this method allows for gradient-based optimization and has the capability of optimizing problems with large data size and data dimensions. We have shown the performance of our method is comparable to state-of-the-art methods while maintaining its scalability advantages. We also applied our method in molecular optimization.


Introduction
Bayesian optimization (BO) is a powerful framework for tackling global optimization problems involving black-box functions (Jones, Schonlau, and Welch 1998). BO seeks to identify an optimal solution with the minimal possible incurred costs. It has been widely applied, yielding impressive results on many problems in different areas ranging from automatic chemical design (Gómez-Bombarelli et al. 2018) to hyperparameter optimization (Snoek, Larochelle, and Adams 2012) have been reported.
However, this is not true for all types of search spaces, in particular discrete spaces. Consider for example optimizing some black box function on a discrete grid of integers like in Figure 1. In this work we focus on Bayesian optimization of objective functions on combinatorial search spaces consisting of discrete variables where the number of possible configurations quickly explodes. For n categorical variables with k categories the number of possible combinations scales with O(k n ). * Equal contributions. Figure 1: Optimization of a black box function over a discrete, non-differentiable input space. The path from an initial input point highlighted in green is highlighted with arrows ending at the red input with optimal value. Combinatorial BO (Baptista and Poloczek 2018) aims to find the global optima of highly non-linear, black-box objectives for which simple and exact solutions are inaccurate and gradient-based optimizers are not amenable. These objectives typically have expensive and noisy evaluations and thus require optimizers with high sample efficiency. Some simple common examples of typical combinatorial optimization problems include the traveling salesman problem, integer linear programming, boolean satisfiability and scheduling.
The vast majority of the BO literature focuses on continuous search spaces. The reason for this is that BO relies on Gaussian processes and the smoothness from kernel methods used to model functional uncertainty. One first specifies a "belief" over possible explanations of the underlying function f using a probabilistic surrogate model and then combines this with the use of an acquisition function which assesses the expected utility of a set of novel X chosen by solving an inner optimization problem. arXiv:2011.02004v1 [cs.LG] 3 Nov 2020

Contributions
In this work, we develop a simple and efficient Bayesian optimization method for discrete combinatorial spaces by treating discrete variables as random variables and applying continuous relaxations. By employing the reparameterization trick combined with Thompson sampling and variational methods, we are able to apply gradient based optimization to maximize the acquisition function when finding new exploration data. Our method performs better than the current state-of-the-art approaches that we compared to and is highly scalable, since we use a Bayesian neural network surrogate model and gradient-based optimization in the algorithmic inner loop.

Related Work
There has been substantial work done in tackling discrete optimization, here we detail a few different approaches in Bayesian optimization and related methods.

Variational Optimization
There has been interesting relevant work treating discrete and non-differentiable optimization spaces by considering smoothed variational approximations of those spaces like in (Staines and Barber 2012;Wierstra et al. 2008).

Treating discrete spaces as continuous
A basic BO approach to combinatorial inputs is to represent all variables using one-hot encoding, treating all integer-valued variables as values on a real line so that the acquisition function considers the closest integer for the chosen real value. This approach seems better suited for ordinal variables and while it has been used in low dimensional situations (Garrido-Merchán and Hernández-Lobato 2020), it has not proved useful for high dimensional ones.

Random search and evolutionary algorithms
Methods like local search and evolutionary algorithms such as particle search are able to handle black-box functions in discrete spaces. However, these procedures have a variety of problems and are not designed to be sample efficient and hence often prohibitively expensive. Moreover, local search algorithms do not necessarily converge to a global optimum. Other popular techniques such as mathematical programming, e.g., linear, convex, and mixed-integer programming, cannot be applied to black-box functions. Bayesian Optimization with sparse Bayesian linear regression In 2018, BOCS (Baptista and Poloczek 2018) was proposed using sparse Bayesian linear regression instead of GPs. The acquisition function was optimized by a semidefinite programming or simulated annealing that allowed to speed up the procedure of picking new points for next evaluations. However, BOCS has certain limitations which restrict it to problems with low order interactions between variables.

BO with graph kernels
In 2019, COMBO (Oh et al. 2019) was proposed, it quantifies "smoothness" of functions on combinatorial search spaces by utilizing a combinatorial graph with a ARD diffusion kernel to model high-order interactions between variables which thus lead to better performance a) b) Figure 2: Variational approximation of a discrete function b) using a Gaussian q(x|α) making it amenable to gradient based optimization. a) is the contour map of the objective

BO on attributed graphs
Deep Graph Bayesian Optimization on attributed graphs was by proposed by (Cui, Yang, and Hu 2019). They use deep graph neural network to model black-box functions on graph avoiding the cubic complexity of GPs that scales linearly with the number of observations. They test their method on molecular discovery and urban road network design. Maximizing acquisition functions (Wilson, Hutter, and Deisenroth 2018) show that acquisition functions estimated via Monte-Carlo are amenable to gradient-based optimization. They also identify a common family of acquisition functions, including EI and UCB, whose properties lend them to greedy maximization approaches.

Other ML approaches to Combinatorial Optimization
Work has also been done to tackle combinatorial optimization problems using neural networks and reinforcement learning, in (Bello et al. 2016), focusing on the traveling salesman problem, they train a recurrent neural network that, given a set of city coordinates, predicts a distribution over different city permutations using negative tour length as the reward signal.

Background
We give a brief overview of variational and bayesian optimization before introducing our method

Variational optimization
Variational optimization (Staines and Barber 2012;Louppe, Hermans, and Cranmer 2017) is a general optimization technique used to form a differentiable bound on the optima of a non-differentiable function. Given our acquisition function L acq (x) to optimize, we can form a bound: where q(x|α) is a proposed distribution with parameters α over input values x. That is, the minimum of the set of acquisition function values is always less than or equal to any of the average function value. Provided that q(x|α) is flexible enough, the parameters α can be updated to place its mass arbitrarily tight around the optimum. Figure 3: The main steps of the Bayesian variational optimization (BVO) algorithm proposed in this work. This depicts the relaxation of the discrete search space and the direct optimization of the acquisition function by using a variational proposal q(x|α) to find the optimal point in the input space x m .

Bayesian optimization
Bayesian optimization relies on both a surrogate model M θ and an acquisition function L acq (x) to define a strategy for efficiently maximizing a black-box function f . At each "outer-loop" iteration, this strategy is used to choose a set x whose evaluation improves the search procedure.

Surrogate model
The surrogate model M θ provides a probabilistic interpretation of the underlying function f k ∼ p(f |D). Gaussian processes are a natural surrogate model but as mentioned previously, their cost scales as the cube of the number of data. Acquisition functions are expectations over predictive distributions (of new outcomes) p(y|x, D) revealed when evaluating a black box function f (x). This formulation naturally occurs in a Bayesian framework whereby x is determined by accounting for the utility provided by possible outcomes y ∼ p(y|x, D). Denoting the chosen utility function as U (y) we can represent the acquisition functions as

Problem Definition
Given a black box function f that is defined over a discrete structured domain X of feasible points, our goal is to find a global optimizer We focus on two main scenarios with search spaces consisting of • binary variables with x ∈ X = {0, 1} d , where x i equals one if a certain element i is present or zero otherwise, and • categorical variables x ∈ X = {0, 1, 2, . . . , k} d where x i is a selected category from 1 − k or zero if not present.
For example, we can generally associate a binary variable with an edge in a graph-like structure.

Bayesian variational optimization
The proposed Bayesian variational optimization (BVO) algorithm is a simple combination of three main parts: • A Bayesian neural network surrogate model with Thompson sampling to approximate the predictive distribution. This allows us to form an unbiased estimate of the acquisition function.
• Variational optimization (Staines and Barber 2012) of the acquisition function.
• A continuous relaxation of the search space for gradientbased optimization of the acquisition function through a categorical reparameterization (Maddison, Mnih, and Teh 2016).
The BVO algorithm are detailed in the next two sections and summarized graphically in figure 3 as well as in algorithm 1 .

Bayesian Neural Network surrogate
Flexible function approximation with reasonable uncertainty quantification can be done with Bayesian neural networks (Blundell et al. 2015). Recent work (Hernández-Lobato et al. 2017), in bayesian optimization has made successful use of BNNs as a surrogate model and we further apply these models in this work. In BNNs, we suppose that observing x provides independent, conditional, normal observations with mean M θ (x) and finite variance σ 2 , such that we have total likelihood Figure 4: Relaxing a discrete space from figure 1 with Gaussian noise to directly optimize some function: a) is the nondifferentiable discrete space, and b) is the differentiable input space of an the function using Gaussians for q(x|α).
We approximate the true posterior p(θ|D) with a diagonal multivariate Gaussian By making use of the reparameterization trick (Kingma and Welling 2014) θ = µ + σ , we can stochastically maximize a lower bound on the marginal log-likelihood We use Thompson sampling to approximate the acquisition function with a single sample from the approximate posterior θ ∼ q(θ) so that p(y|x, D) = p(y|x, θ)p(θ|D)dθ ≈ p(y|x, θ).

Variational optimization of the acquisition function
Since the input to the acquisition function L acq is discrete it is non-differentiable and we cannot optimize it with gradients. However, we can use variational optimization (Staines and Barber 2012;Louppe, Hermans, and Cranmer 2017) to form a differentiable bound on min x L acq (x) ≤ E x∼q(x|α) [L acq (x)] and optimize that in place using a single sample Monte-Carlo estimator of p(y|x, D) and placing a proposal distribution q(x|α) on the input space x with parameters α. Now we can optimize this input space by proxy through optimization of the proposal q(x|α). Now we must consider what is an appropriate proposal distribution q(x|α) for example in the discrete space of figure  1 we could use a Gaussian distribution as depicted in figure  4. Most search spaces we consider will consist of binary or categorical variables, hence it makes sense to model q(x|α) as a product of Bernoulli's or Multinomial distributions.

Algorithm 1: Bayesian Variational Optimization
Given target function f , model M, acquisition L acq , relaxation G, and initial data D for i ← 1 to T do Fit model M with D Randomly initialize a set of variables α for j ← 1 to N do Sample relaxed discrete variables X ← G(α)

Continuous Relaxation of the input space through Categorical Reparameterization
To optimize the acquisition function directly with Monte-Carlo gradients, we are required to choose the proposal q(x|α) so that we can form a differentiable reparameterization such that x = η(α, ) (Kingma and Welling 2013) and hence the path-wise gradients from L acq (x(α)) to α can be computed as where p( ) is a base noise distribution. Motivated by this requirement, we use the concrete distribution (Maddison, Mnih, and Teh 2016) as our proposal distribution, which allows us to form a continuous approximation of categorical random variables by using the the Gumbel-Softmax trick. We can sample from the concrete distribution via g is drawn from the Gumbel distribution with u drawn from uniform distribution between (0, 1). For the binary case, we can sample from the binary concrete distribution via g is drawn from the logistic distribution with u drawn from uniform distribution between (0, 1). In both cases, λ is the temperature. The advantage for using the Gumbel-Softmax instead of Softmax for discrete/categorical variables is that, with certain temperature, the sampled distribution can be closer to one-hot representation.

Experiments
To demonstrate the BVO algorithm, we preform several experiments to validate our approach in a variety of discrete spaces of variable size and structure, including: 1. binary variables: Ising sparsification and contamination control, 2. categorical variables: Pest control and molecular optimization with SELFIES, 3. computational complexity on variable dimensions and data size.

Optimization setup
Due to the model-dependent nature of the BVO method, we carried out a hyperparameter search for the different objectives. This includes trying different BNN models, such as various activation functions (Tanh, ReLU), number of layers (2-6), and layer sizes (50-200). We also explored different acquisition functions (EI, SR, PI listed in (Wilson, Hutter, and Deisenroth 2018)), CONCRETE relaxation temperatures (0.1-1.0), optimization batch size (16-512), and the scaling factors (10 1 − 10 7 ) on the loss function. The training of the models and optimization were computed on Nvidia V100SXM2 GPUs from cloud servers with 8 core CPUs and 32 GB of RAM.

Binary variables
Binary variable optimization is a ubiquitous problem in computer science. Therefore, the benchmarking of BVO with respect to this problem gives us a handle to compare to several methods from distinct areas of computer science. For the binary variable optimization tasks, we benchmark the BVO method for Ising sparsification and contamination control tasks. Figure 5 shows our treatment of binary variables, and how we construct q(x|α) as a product of Bernoulli distributions that are relaxed for gradient based optimization.
x ∼ q(x|α) The Ising sparsification problem entails the deletion of interactions in a zero-field Ising model while maintaining a similar probability distribution to that of the original model. The zero-field Ising model can be expressed as where z ∈ {−1, 1} n , the interaction matrix J p ∈ R n×n , and Z p is the normalization factor. For this problem, our goal is to find an approximate model q(z) with an interaction matrix that are partially dropped out J q ij = x ij J p ij , where x ij ∈ {0, 1}. For the model q(z) to approximate to p(z), the objective function for optimization is where the first term is a Kullback-Leibler divergence for distribution similarity and the second term is for maximizing dropouts. In our benchmark, we use the same objective function in the COMBO code base, which is a 4 × 4 Ising grid. Since the interactions are only the nearest neighbors, there consiste only 24 possible interaction dropouts. All the interactions J p ij are sampled randomly from interval [0.05, 0.5]. The optimization includes 20 randomly initialized points followed by 150 points of iterative selection for evaluation. Contamination control. The contamination control problem is a binary optimization problem for minimizing contamination in a simulated food supply chain. At each stage i, z i ∈ [0, 1] represents the portion of contaminated food, where as the binary control variables x i ∈ {0, 1} represents whether to control and decontaminate the food. Contamination at each stage may have different costs c i . When no control x i = 0, the contaminated portion of the food in next stage will increase; where as x i = 1, certain portion of the contaminated food will be decontaminated. The overall target of the contamination control is to minimize The first term of the target is the cost of decontamination, and the last term is for x regularization. The second term is the penalty for every stages that have contaminated food over a certain threshold u, and it is calculated for the mean of multiple starting food contamination portion z k 1 . In this benchmark, we follow the COMBO baselines by using T = 100, u = 0.1, with 100 food contamination stages and d=21 binary variables. The optimization were 20 random initializing points with 250 optimizations. Results. The results for Bayesian variational optimization on binary variables are shown in table 1. In the binary optimization problem, we can see that BVO is similar or better in some cases to state-of-the-art discrete optimization methods. However, in both Ising sparsification and contamination control problems, BVO has a higher standard deviation over the runs. This could be due to variational sampling for optimization. The optimal values by optimization iterations are plotted in figure 6 (ab) for Ising sparsification and Contamination control, which also includes the baseline comparisons.

Categorical variables
For the categorical variable (figure 7) optimization tasks, we benchmark our method to other baseline methods on the Pest control task. To demonstrate BVO in applications, we aim to directly optimize molecules with SELFIES representation. Pest control. The Pest control problem is similar to contamination control, but with more control choices (4 different pesticide l). The pest situation is also updated with a similar dynamics to the contamination control problem but with different pesticide efficacy. However, there are also more changes in dynamic control rates, spread rates, and an reduce cost changes when a pesticide are used more often. For this benchmark, there are 5 different categories (no action + 4 pesticides) and 21 pest control stages. The optimization target is similar to equation 4, without the regularization term. For optimization, we use 20 random initial points with 300 optimizations. The results of pest control optimization are shown in table 2 along with the baseline comparisons. Our method is on par to the COMBO optimization results, while better than the other methods. The minimum values by optimization iterations are plotted in figure 6 (c) along random search, TPE and COMBO, where we see that BVO performs very well in terms of final optimal values and variances.

Method
Pest   ating correct SMILES representation of a valid molecules are difficult, due to both physical meaning and grammar constraints. This could lead to low generative validity, or complicated reinforcement learning algorithms for generating the molecules. SELFIES (Krenn et al. 2019) is an improved version of SMILES representation, which encodes the graph in Chomsky type-2 context-free grammar that generates 100% valid molecules. Figure 8 shows an example on how a molecule is represented in SELFIES. In this optimization, our search space is 200 SELFIES tokens, each token with 25 categories. The tokens include construction of rings, branches, atoms and its bonds. For the molecular target, we optimize for the highest penalized LogP score (solubility). The search space for molecule optimization in this case is on the order of 10 279 . This scale is very difficult for GP based optimization methods.
Due to the large search space and complicated sequence dependencies in SELFIES, we optimize the molecule by initializing 2048 SELFIES strings, and parallel evaluating 512 samples for each optimization for 1200 steps. The model is built with 1D convolution layers and pooling layers for long sequence learning. The optimization were repeated over 10 times. The maximum score achieved for BVO is shown in table 3, with a maximum score of 4.94. As a comparison, GVAE have a maximum score of 2.94 and CVAE with 1.98. This advantage for BVO with SELFIES could be due to the fact that our molecular search space can access outside of the ZINC dataset. Nevertheless, the BVO on SELFIES has a better optimization LogP score compared to a random search using SELFIES.

Scalability of BVO
Bayesian optimization often relies on Gaussian processes (GP) for the surrugate model. However, GP has the limitation that computation costs scale cubically to data size. This may lead to difficulty on categorical optimization tasks when they have high variable dimensions and requires large data size for training. In figure 9, we measured the computation time for BVO and GP (COMBO) on the pest control experiment. For comparison, the mean of first 3 points of each curve are normalized to 1. The computation scaling regarding to variable dimension for BVO is linear scaling (order = 1.09), in contrast to GP which is quadratically O(n 2 ) (order = 1.93). Regarding with data size (optimization iterations), GP scales between O(n 2 ) − O(n 3 ) (order = 2.49), while BVO can scale by constant (fixed number of batches per epoch) or by linear (fixed epochs of data size). The complexity analysis here shows that BVO is more practical then GP for optimizing large problems, which we demonstrated on SELFIES optimization.

Conclusion
In this paper, we developed a simple and efficient Bayesian variational optimization method for discrete combinatorial spaces. By employing reparameterization and Thompson sampling, we are able to apply gradient based optimization to maximize the acquisition function when finding new exploration data. While our method performs better than the current state-of-the-art approaches, the algorithm also allows high dimension and large data size optimization without suffering computation limitations. As a demonstration, we optimized the molecules with BVO on SELFIES representation, which consist of large search spaces.