Bayesian Optimization with High-Dimensional Outputs

1 Introduction . Many problems in science and engineering involve reasoning about multiple , correlated outputs . For example , cell towers broadcast signal across an area , and thus signal strength is spatially correlated . In randomized experiments , treatment effects on multiple outcomes are naturally correlated due to shared causal mechanisms . Without further knowledge of the internal mechanisms ( i.e. , in a “ black-box ” setting ) , Multi-task Gaussian processes ( MTGPs ) are a natural model for these types of problems as they model the relationship between each output ( or “ task ” ) while maintaining the gold standard predictive capability and uncertainty quantification of Gaussian processes ( GPs ) . Many downstream analyses require more of the model than just prediction ; they also involve sampling from the posterior distribution to estimate quantities of interest . For instance , we may be interested in the performance of a complex stock trading strategy that requires modeling different stock prices jointly , and want to characterize its conditional value at risk ( CVaR ) [ 49 ] , which generally requires Monte Carlo ( MC ) estimation strategies [ 10 ] . Or , we want to use MTGPs in Bayesian Optimization ( BO ) , a method for sample-efficient optimization of black-box functions . Many state of the art BO approaches use MC acquisition functions [ 60 , 3 , 4 ] , which require sampling from the posterior distribution over new candidate data points . Drawing posterior samples from MTGPs means sampling over all tasks and all new data points , which typically scales multiplicatively in the number of tasks ( t ) and test data points ( n ) , e.g . like O ( n3t3 ) [ 52 , 6 ] . For problems with more than a few tasks , posterior sampling thus quickly becomes 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . intractable due to the size of the posterior covariance matrix . This is especially problematic in the case of many real-world problems that can have hundreds or thousands of correlated outputs that should be modelled jointly in order to achieve the best performance . For instance , the cell tower signal maps in Figure 1 each contain 2,500 outputs ( pixels ) . In this problem , we aim to jointly tune the downtilt angle and transmission power of the antennas on each cell tower ( locations shown in red ) to optimize a global coverage quality metric , which is a known function of power and interference at each location [ 21 ] . Since simulating the power and interference maps given a parameterization is computationally costly , traditionally one might apply BO to optimize the aggregate metric . At its core , this problem is a composite BO problem [ 3 , 4 ] , so we expect an approach that models the constituent outcomes at each pixel individually to achieve higher sample efficiency . However , modelling each pixel using existing approaches used for BO is completely intractable in this setting , as we would have to train and sample from a MTGP with 5,000 tasks . To remedy the poor computational scaling with the number of tasks , we exploit Matheron ’ s rule for sampling from GP posterior distributions [ 13 , 59 ] . We derive an efficient method for MTGP sampling that exploits Kronecker structure inherent to the posterior covariance matrices , thereby reducing the complexity of sampling from the posterior to become effectively additive in the combination of tasks of data points , i.e . O ( n3 + t3 ) , as compared to O ( n3t3 ) . Our implementation of Matheron ’ s rule draws from the exact posterior distribution and does not require random features or inducing points , unlike decoupled sampling [ 59 ] . More specifically , our contributions are as follows : • We propose an exact sampling method for multi-task Gaussian processes that has additive time costs in the combination of tasks and data points , rather than multiplicative ( Section 3 ) . • We demonstrate empirically how large-scale sampling from MTGPs can aid in challenging multi-objective , constrained , and contextual Bayesian Optimization problems ( Section 4 ) . • We introduce a method for efficient posterior sampling for the High-Order Gaussian Process ( HOGP ) model [ 64 ] , allowing it to be used for Bayesian Optimization ( Section 3.2 ) . This advance allows us to more efficiently perform BO on high-dimensional outputs such as images — including optimizing PDEs , optimizing placements of cell towers for cell coverage , and tuning the mirrors of an optical interferometer which optimizes over 65,000 tasks jointly ( Section 4.3 ) . The rest of the paper is organized as follows : First , in Section 2 we review GPs , MTGPs , and sampling procedures from the posterior in both GPs and MTGPs . In Section 3 , we review Matheron ’ s rule for sampling from GP posteriors and explain how to employ it for efficient sampling from MTGP models including the HOGP model . In Section 4 , we illustrate the utility of our method on a wide suite of problems ranging from constrained BO to the first demonstration of large scale composite BO with the HOGP . Please see Appendix A for discussion of the limitations and broader impacts of our work . Our code is fully integrated into BoTorch , see https : //botorch.org/tutorials/composite_ bo_with_hogp and https : //botorch.org/tutorials/composite_mtbo for tutorials . 2 Background . 2.1 Bayesian Optimization . In Bayesian Optimization ( BO ) , the goal is to minimize an expensive-to-evaluate black-box function , i.e. , finding minx2X f ( x ) , by constructing a surrogate model to emulate that function . Gaussian processes ( GPs ) are often used as surrogates due to their flexibility and well-calibrated uncertainty estimates . BO optimizes an acquisition function defined on the predictive distribution of the surrogate model to select the next point ( s ) to evaluate on the true function . These acquisition functions are often written as intractable integrals that are typically evaluated using Monte Carlo ( MC ) integration [ 60 , 3 ] . MC acquisition functions rely on posterior samples from the surrogate model , which should support fast sampling capabilities for efficient optimization [ 4 ] . BO has been applied throughout machine learning , engineering , and the sciences , and many extensions to the setting described above exist . We focus on multi-task BO ( MTBO ) , where f ( x ) is composed of several correlated tasks [ 55 , 15 ] . There are many sub-classes of MTBO problems : constrained BO uses surrogate models to optimize an objective subject to black-box constraints [ 29 , 25 , 30 ] , contextual BO models a single function that varies across different contexts or environments [ 38 , 11 , 26 ] , multi-objective BO aims to explore a Pareto frontier across several objectives [ 36 , 37 , 23 , 24 , 17 ] , and composite BO considers the setting of a differentiable objective function defined on the outputs of a vector-valued black-box function [ 56 , 3 , 4 ] . In all of these problems , the setting is similar : several outputs are modelled by the surrogate , whether the output is a constraint , another objective , or a separate context . As the outputs are often correlated , multi-task Gaussian processes , which model the relationships between the outputs in a data-efficient manner , are a natural and common modeling choice . 2.2 Gaussian Processes . Single Output Gaussian Processes : We briefly review single output GPs , see Rasmussen and Williams [ 48 ] for a more detailed introduction . We assume that y = f ( x ) + `` , f ⇠ GP ( 0 , k✓ ( x , x0 ) ) , and `` ⇠ N ( 0 , 2 ) , where f is the noiseless latent function and y are noisy observations of f with standard deviation . k✓ ( x , x0 ) is the kernel with hyperparameters ✓ ( we will drop the dependence on ✓ for simplicity ) ; we use KAB : = k✓ ( A , B ) to refer to the evaluated kernel function on data points A and B , a matrix which has size |A|⇥ |B| . The predictive distributions over new data points , xtest , is given by the conditional distribution of the Gaussian distribution . That is , p ( f ( xtest ) |D , ✓ ) = N ( µ⇤f |D , ⌃ ⇤ f |D ) , where D : = { X , y } is the training dataset of size n = |X| and µ ⇤ f |D = KxtestX ( Ktrain + 2 I ) 1y , ( 1 ) ⌃⇤f |D = Kxtestxtest KxtestX ( Ktrain + 2 I ) 1KXxtest , ( 2 ) with Ktrain : = KXX . For simplicity , we will drop the subscripts f |D in all future statements . Computing the predictive mean µ⇤ and variance ⌃⇤ requires O ( n3 ) time and O ( n2 ) space when using Cholesky decompositions for the linear solves [ 48 ] . Sampling is usually performed by f ( xtest ) | ( Y = y ) = µ ⇤ + ( ⌃⇤ ) 1/2z , ( 3 ) where z ⇠ N ( 0 , I ) . Computing s samples at ntest test points from the predictive distribution costs O ( n3 + sn2test + n 2 ntest + n3test ) , computed by adding up the cost of all of the matrix vector multiplications ( MVMs ) and matrix solves . For fixed xtest we can incur the cubic terms only once by re-using Cholesky factorizations of Ktrain + 2I and ⌃⇤ for each sample . To reduce the time complexity , we can replace all matrix solves with r < n steps of conjugate gradients ( CG ) and the Cholesky decompositions with rank r < n Lanczos decompositions ( an approach called LOVE [ 45 ] ) . These change the major scaling from n3 down to rn2 and the overall time complexity to O ( rn2 + srntest + rnntest + rn2test ) [ 45 , 28 ] . In general , r ⌧ n is used and is usually accurate to nearly numerical precision [ 28 ] . We provide additional details in Appendix B. Multi-Output Gaussian Processes : One straightforward way of modelling multiple outputs is to consider each output as an independent GP , modelled in batch together , either with shared or independent hyperparameters [ 48 , 28 , 25 ] . However , there are two major drawbacks to this approach : ( i ) the model is not able to model correlations between the outputs , and ( ii ) if there are many outputs then maintaining a separate model for each can result in high memory usage and slow inference times . To remedy these issues , Higdon et al . [ 33 ] propose the PCA-GP , using principal component analysis ( PCA ) to project the outputs to a low-dimensional subspace and then use batch GPs to model the lower-dimensional outputs . We consider multi-task Gaussian processes ( MTGPs ) with the intrinsic co-regionalization model ( ICM ) , which considers the relationship between responses as the product of the data and task features [ 32 , 6 , 1 ] . We focus on this model due to its popularity and simplicity , leaving a similar derivation of the linear model of co-regionalization to Appendix C.1.1 . Given inputs x and x0 belonging to tasks i and j , respectively , the covariance under the ICM is k ( [ x , i ] , [ x0 , j ] ) = kD ( x , x0 ) kt ( i , j ) . Given n data points X with the n associated task indices I , the covariance is a Hadamard product of the data kernel and the task kernel , Ktrain = KXX KII , and the response y is still of size n. We term this implementation of multi-task GPs “ Hadamard MTGPs ” in our experiments . In general , there is no easily exploitable structure in this model . If we observe each task at each data point ( the so-called block design case ) , the covariance matrix becomes Kronecker structured , e.g . Ktrain = KXX ⌦KT , where KXX is the data covariance matrix and KT is the t⇥ t task covariance matrix between tasks [ 6 ] , and we now have nt scalar responses . To simplify our exposition , we assume that KT is full-rank ( this is not required as we can use pseudo-inverses in place of inverses ) . Thus , the GP prior is vec ( y ) ⇠ N ( 0 , KXX ⌦KT ) , where y is a matrix of size n⇥ t , and vec ( y ) is a vector of shape nt . The GP predictive distribution is given by p ( f⇤|xtest , D ) = N ( µ⇤ , ⌃⇤ ) , where µ ⇤ = ( Kxtest , X ⌦KT ) ( KXX ⌦KT + 2 InT ) 1vec ( y ) , ⌃⇤ = ( Kxtest , xtest ⌦KT ) ( Kxtest , X ⌦KT ) ( KXX ⌦KT + 2 InT ) 1 ( K > xtest , X ⌦KT ) . ( 4 ) The kernel matrix on the training data , KXX ⌦ KT , is of size nt ⇥ nt , which under standard ( Cholesky-based ) approaches yields inference cubic and multiplicative in n and t , that is O ( n3t3 ) . However , the Kronecker structure can be exploited to compute the posterior mean and variance in O ( nt ( n+ t ) + n3 + t3 ) , which is dominated by the individual cubic terms [ 50 , 54 ] . Sampling from the posterior distribution in ( 4 ) produces additional computational challenges as we must compute a root ( e.g . Cholesky ) decomposition of ⌃⇤ , which naively costs O ( ( ntestt ) 3 ) plus an additional Cholesky decomposition of ( KXX ⌦KT + 2I ) 1 , which similarly costs O ( ( nt ) 3 ) time [ 6 ] . Thus , the time complexity of drawing s samples is multiplicative in n and t , O ( ( nt ) 3 + ( ntestt ) 3 + s ( ( nt ) 2 + ( ntestt ) 2 ) ) . Using CG and LOVE reduces the complexity ; see Appendix B.4 . High-Order Gaussian Processes : Recently , Zhe et al . [ 64 ] proposed the high-order Gaussian process ( HOGP ) model , a MTGP designed for matrix- and tensor-valued outputs . Given outputs y 2 Rd1⇥···⇥dk ( e.g . a matrix or tensor ) , the covariance is the product of the data dimension and each output index ( il and i0l respectively ) k ( [ x , i1 , · · · , ik ] , [ x 0 , i 0 1 , · · · , i 0 k ] ) = k ( x , x0 ) k ( vi1 , v 0 i01 ) · · · k ( vik , v 0 i0k ) , where i1 , · · · , ik are the indices for the output tensor and v1 , · · · , vk are latent parameters that are optimized along with the kernel hyper-parameters . Thus , KT in the MTGP framework is replaced by a chain of Kronecker products , so that the GP prior is vec ( y ) ⇠ N ( 0 , KXX ⌦ K2 ⌦ · · · ⌦ Kk ) . Exploiting the Kronecker structure , computation of posterior mean , posterior variance and hyper-parameter learning takes O ( n3+ Pk i=2 d 3 i +n Qk i=1 di ) time as d1 = n. Zhe et al . [ 64 ] demonstrate that the HOGP outperforms PCA-GPs [ 33 ] , among other high-dimensional output GP methods , with respect to prediction accuracy . However , their experiments measure only the error of the predictive mean , rather than quantities that use the predictive variance ( such as the negative log likelihood or calibration ) , and they do not provide a way to sample from the posterior distribution .

Summary: This paper presents the idea of scalable multi-task Gaussian processes using a trick known as Matheron's rule. Efficient sampling from multi-task GPs is essential in solving many black box optimization problems with high-dimensional outputs including multi-objective constrained BO and composite functions. Matheron's rule is used to reduce the dependency of the number of tasks from a multiplicative factor to an additive factor. Experimental results are presented validating the benefits of the approach wrt. the computational resource requirements and also sample complexity.