Bayesian Optimization of Function Networks

1 Introduction . We consider Bayesian optimization ( BO ) of objective functions defined by a series of time-consumingto-evaluate functions , f1 , . . . , fK , arranged in a directed acyclic network , so that each function takes as input the output of its parent nodes . As we detail below , these problems arise in BO-based policy search in reinforcement learning ( Lizotte et al. , 2007 ) , optimization of complex systems modeled via simulation , and calibration of time-consuming physics-based models . To illustrate , we introduce a running example of vaccine manufacturing ( Sekhon and Saluja , 2011 ) , focusing on the portion of the manufacturing process that uses live cells to produce proteins needed in a vaccine . It begins with a cell culture , in which living cells are grown and used as “ factories ” to produce proteins . This process is controlled by a vector , x1 , containing the temperature , pH , and CO2 content used when growing these cells . The output of this process is the quantity of the desired protein y1 = f1 ( x1 ) , i.e. , the yield of this step , along with other byproducts . This output is passed into a second process , purification , which removes byproducts and is controlled by a vector x2 comprising temperature , pressure , and flow rate . The yield of the desired protein from this second step is y2 = f2 ( x2 , y1 ) . This output enters a third step , formulation , in which we formulate the raw protein into a form that can be distributed as controlled by a third set of parameters . This determines the yield of the overall process y3 = f3 ( x3 , y2 ) . We wish to choose ( x1 , x2 , x3 ) to maximize overall protein yield . This problem is summarized as a function network in Figure 1 . The problem described above and other similar problems can be tackled with Bayesian optimization ( BO ) , which has been shown to perform well compared to other derivative-free global optimization 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . methods for time-consuming-to-evaluate objective functions ( Snoek et al. , 2012 ; Frazier , 2018 ) . A standard BO algorithm would fit a Gaussian process ( GP ) ( Rasmussen and Williams , 2006 ) model on the objective function ( y3 , which depends on ( x1 , x2 , x3 ) ) and use it , along with an acquisition function , to sequentially choose the points to evaluate . Under this standard approach , however , evaluations of the intermediate nodes , f1 , . . . , fK−1 , would be ignored despite being available when computing the objective function . In the example above , this corresponds to looking only at the yield of the overall process , and not of each individual step . In this paper , we introduce a novel BO approach that leverages function network structure for substantially more efficient optimization . This approach models the individual nodes of the network using distinct GPs . This allows incorporating observations of each node ’ s output recursively into a non-Gaussian posterior on the network ’ s overall output . Our approach then chooses the points to evaluate using the expected improvement ( Jones et al. , 1998 ) computed with respect to this implied posterior on the objective function . The non-Gaussian nature of this posterior prevents the expected improvement from having a closed form . However , we show that it can still be efficiently maximized via sample average approximation ( Kleywegt et al. , 2002 ) . Our approach can outperform standard BO by leveraging information from internal nodes unavailable to standard methods . We briefly explain one way this can happen , in the context of the example above . In vaccine manufacturing , each function fk ( xk , yk−1 ) , k = 2 , 3 , is bounded between 0 and yk−1 because new protein can not be created in the purification and formulation stages . Moreover , application experts have a prior on what values for fk ( xk , yk−1 ) /yk−1 should be achievable if xk is set well . Thus , if we see that y3 is unexpectedly poor , information from intermediate nodes can be extremely valuable : if y1 and y2 are as expected , then this suggests the problem is with x3 ; if y1 is as expected but y2 is poor , then the problem is likely with x2 ; and if y1 is poor then the problem is likely with x1 . ( If there is a problem with xk , there may also be a problem with xk′ , k′ > k , but we can focus on fixing xk first . ) Thus , by observing intermediate nodes , we can instantly reduce the effective dimensionality of the input space by a factor of K = 3 . We show that our method is asymptotically consistent , i.e. , that it discovers the global optimum given sufficiently many samples . Remarkably , in contrast with most BO methods , it may do so without measuring densely over the feasible domain , instead leveraging function network structure to exclude regions as unnecessary to explore . This indicates the power of function network structure to improve query efficiency . We demonstrate through numerical experiments that access to additional information available in a problem formulated as a function network can dramatically accelerate optimization . We study four synthetic problems and four real-world problems : a manufacturing problem similar in spirit to the vaccine example above , an active learning problem with a robotic arm , and two problems arising in epidemiology , one calibrating an epidemic model and the other designing a testing strategy to control the spread of COVID-19 . Our method significantly outperforms competing methods that utilize less information , in some cases by ∼5 % and in other cases by several orders of magnitude . 2 Related Work . Our work occurs within BO , a framework for global optimization of expensive-to-evaluate black-box functions that originated with the work of Zhilinskas ( 1975 ) and Močkus ( 1975 ) , and has recently become popular due to its remarkable performance in hyperparameter tuning of machine learning algorithms ( Snoek et al. , 2012 ; Swersky et al. , 2013 ; Wu et al. , 2019 ) . We refer the reader to Shahriari et al . ( 2016 ) and Frazier ( 2018 ) for modern introductions to BO . Our approach can be catalogued as a grey-box BO method since it does not treat the objective function entirely as a black box , and instead exploits known structure to improve sampling efficiency . Other examples of grey-box BO approaches include multi-fidelity BO ( Kandasamy et al. , 2017 ; Wu et al. , 2019 ) , which leverages cheaper approximations of the objective function ; BO of objective functions that are the integral of an expensive-to-evaluate integrand ( Williams et al. , 2000 ; ToscanoPalmerin and Frazier , 2018 ; Cakmak et al. , 2020 ) ; BO of objective functions that are a sum of squared errors ( Uhrenholt and Jensen , 2019 ) ; and , more generally , BO of objective functions that are the composition of an expensive-to-evaluate inner function and a known inexpensive-to-evaluate outer function ( Astudillo and Frazier , 2019 ) , of which our work can be seen as a significant generalization . We refer the reader to Astudillo and Frazier ( 2021 ) for a tutorial on grey-box BO . Our work is also closely related to Marque-Pucheu et al . ( 2019 ) , which proposes a method for efficient sequential experimental design of nested computer codes , also using GPs . This work focuses on the case where there are only two node functions , and one takes as input the output of the other . In contrast with our work , the goal of the proposed method is to learn the output code as accurately as possible within a limited evaluation budget , but optimization is not pursued . Optimization of composite ( a.k.a nested ) functions has also been considered in the gradient-based optimization literature ( Shapiro , 2003 ; Drusvyatskiy and Paquette , 2019 ; Charisopoulos et al. , 2021 ; Balasubramanian et al. , 2020 ) . In contrast with ours , these works assume that evaluations are inexpensive , and often also some form of convexity , along with availability of gradients . However , this literature is similar in spirit to ours in that information from inner functions improves efficiency . Function networks arise in many application areas . While these applications have not , to our knowledge , been previously formulated as specific instances of the general function network model we propose , their literatures are nonetheless relevant to our work . • In engineering and aerospace design , function networks arise in multidisciplinary optimization ( Cramer et al. , 1994 ; Amaral et al. , 2014 ; Benaouali and Kachel , 2019 ) , where simulators focusing on different physical laws are coupled into a function network . For example , a simulation of airflow over a wing may output the forces on the wing to another simulation of mechanical stress on the wing structure . • In drug discovery and materials design , function networks arise from the sequence-structurefunction ( Sadowski and Jones , 2009 ) and composition-structure-property ( Hattrick-Simpers et al. , 2016 ) paradigms . Here , decision variables ( composition , e.g. , what fraction of what raw materials are used ) determine the structure of the material ( how the atoms combine together ) , which in turn determines how the material behaves ( properties ) . • Function networks arise in the design of queuing networks ( Fu and Henderson , 2017 ; Arcelli , 2020 ) . This includes manufacturing systems ( Ghasemi et al. , 2018 ) , where the partiallyprocessed output of one workstation is input to the next workstation , the design of service systems ( Wang et al. , 2020 ) such as hospitals and airport security checkpoints , and choosing traffic signal timings for a city ’ s street network ( Osorio and Bierlaire , 2013 ) . • Finally , function networks arise in reinforcement learning ( Sutton and Barto , 2018 ) and Markov decision processes ( Puterman , 1990 ) , where the transition kernel transforms the state variable at the start of a time range to another state variable at the end of this range . This outputted state variable is the input to the next time range . §5.2 shows an example . 3 Problem Setting . We consider objective functions evaluated via a series of functions , f1 , . . . , fK , arranged in a directed network so that each function in this network takes as input the output of its parent nodes , and assume that evaluating this collection of functions is time-consuming . The network structure is encoded as a directed graph with nodes V = { 1 , . . . , K } and directed edges E = { ( j , k ) : fk take as input the output of fj } . We assume that this graph is acyclic and has a single leaf node whose output is the objective to optimize . Let J ( k ) denote the set of parent nodes of node k1 . We assume , without loss of generality , that the nodes are ordered such that j < k for all j ∈ J ( k ) . In addition to the output of its parent nodes , we 1We allow J ( k ) = ∅ , which corresponds to the root node ( s ) . assume that each function , fk , takes as input a ( potentially empty ) subset , I ( k ) ⊂ { 1 , . . . , D } , of the components of the decision vector x ∈ RD . Let h1 ( x ) , . . . , hK ( x ) denote the values of the K nodes in the function network when it is evaluated at x . These values are defined recursively by hk ( x ) = fk ( xI ( k ) , hJ ( k ) ( x ) ) , k = 1 , . . . , K , i.e. , by evaluating each function in the network as the values of its parent nodes become available . This recursion is well-defined by our assumption that the graph is acyclic . The objective function is then g = hK , and the goal is to solve max x∈X g ( x ) , where X ⊂ RD is a simple compact set , such as a hyper-rectangle . The standard BO approach to this problem models g using a GP prior distribution . This approach iteratively chooses the next point at which to evaluate g as follows . Given n observations of the objective function , g ( x1 ) , . . . , g ( xn ) , it computes the posterior distribution on g , which is itself a GP . It then uses this posterior distribution to compute an acquisition function ( Frazier , 2018 ) that quantifies the value of the information that would result from a function evaluation at any given point . Finally , it chooses the next point to evaluate , xn+1 , as the point that maximizes this acquisition function . Notably , although the outputs of f1 , . . . , fK−1 would be observed as part of this approach , these evaluations would be ignored by standard BO when calculating the posterior distribution and corresponding acquisition function .

Summary: This work investigates the problem of Bayesian optimization of the output of a network of functions. The network is organized as a directed acyclic graph, where each function takes as input the output of its parent nodes as well as (part of ) the input. Contrary to the standard setting that only observes the final output, this work assumes that intermediate nodes can also produce an output (which can be stochastic or deterministic). The authors model the nodes of the network using Gaussian processes, and employ an acquisition function that corresponds to the expected improvement with respect to the implied posterior on the final output under a proper statistical model. The authors optimize the acquisition function using sample average approximation. The authors show that their approach satisfies asymptotic convergence. Various experiments on synthetic and real problems demonstrate that the improved model and acquisition function result in visibly improved solution quality, at the expense of the runtime cost.