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: The authors devise a novel method for exact multi-task Gaussian process (GP) sampling that dramatically improves time costs, from multiplicative to additive in the combination of tasks and data points. Their method works by exploiting Kronecker structures in covariance matrices and then applying Matheron's identity. This allows them to perform Bayesian optimisation (BO) for tens of thousands of correlated outputs, which was hardly possible previously.