Bayesian Learning to Optimize: Quantifying the Optimizer Uncertainty

1 INTRODUCTION . Computational models of many real-world applications involve optimizing non-convex objective functions . As the non-convex optimization problem is NP-hard , no optimization algorithm ( or optimizer ) could guarantee the global optima in general , and instead , their solutions ’ usefulness ( sometimes based on their proximity to the optima ) , when the optima are unknown , can be very uncertain . Being able to quantify such uncertainty is important to not only assessing the solution uncertainty after optimization but also enhancing the search efficiency during optimization . For instance , reliable and trustworthy machine learning models demand uncertainty awareness and quantification during training ( optimizing ) such models , whereas in reality deep neural networks without proper modeling of uncertainty suffer from overconfidence and miscalibration ( Guo et al. , 2017 ) . In another application example of protein docking , although there exists epistemic uncertainty of the objective function and the aleatoric uncertainty of the protein structure data ( Cao & Shen , 2020 ) , state-ofthe-art methods only predict several single solutions ( Porter et al. , 2019 ) without any associated uncertainty , which makes those predictions hard to be interpreted by the end users . Various optimization methods have been proposed in response to the need of uncertainty awareness . Stochastic optimization methods like random search ( Zhigljavsky , 2012 ) , simulated annealing ( Kirkpatrick et al. , 1983 ) , genetic algorithms ( Goldenberg , 1989 ) and particle swarm optimization ( Kennedy & Eberhart , 1995 ) injected the randomness into the algorithms in order to reduce uncertainties . However , these methods do not provide the uncertainty quantification ( UQ ) of solutions . Recently , there have been growing interests in applying inference-based methods to optimization problems ( Brochu et al. , 2010 ; Shapiro , 2000 ; Pelikan et al. , 1999 ) . Generally , they transfer the uncertainties within the data and model into the final solution by modelling the posterior distribution over the global optima . For instance , Bijl et al . ( 2016 ) uses sequential Monte Carlo to approximate the distribution over the optima with Thompson sampling as the search strategy . Hernández-Lobato et al . ( 2014 ) uses kernel approximation for modelling the posterior over the optimum under Gaussian process . Ortega et al . ( 2012 ) ; Cao & Shen ( 2020 ) directly model the posterior over the optimum as a Boltzmann distribution . They not only surpass the previous methods in accuracy and efficiency , but also provide easy-to-interpret uncertainty quantification . Despite progress in optimization with uncertainty-awareness , significant open questions remain . Existing methods consider uncertainty either within the data or the model ( including objective functions ) ( Kendall & Gal , 2017 ; Ortega et al. , 2012 ; Cao & Shen , 2020 ) . However , no attention was ever paid to the uncertainty arising from the optimizer that is directly responsible for deriving the end solutions with given data and models . The optimizer is usually pre-defined and fixed in the optimization algorithm space . For instance , there are several popular update rules in Bayesian optimization , such as expected improvement Vazquez & Bect ( 2010 ) or upper confidence bound Srinivas et al . ( 2009 ) , that are chosen and fixed for the entire process . For Bayesian neural networks training , the update rule is usually chosen off-the-shelf , such as Adam , SGD , or RMSDrop . The uncertainty in the optimizer is intrinsically defined over the optimizer space and important to the optimization and UQ solutions . However , such uncertainty is unwittingly ignored when the optimizer is treated as a fixed sample in the space . To fill the aforementioned gap , the core intellectual value of this work is to recognize and quantify a new form of uncertainty , that lies in the optimization algorithm ( optimizer ) , besides the classical data- or model- based uncertainties ( also known as epistemic and aleatoric uncertainties ) . The underlying innovation is to treat an optimizer as a random sample from the algorithmic space , rather than one of a few hand-crafted update rules . The key enabling technique is to consider the algorithmic space being parameterized by a neural network . We then leverage a Boltzmann-shaped posterior over the optimizers , and approximate the posterior locally as Gaussian distributions through variational inference . Our approach , Bayesian learning to optimize ( BL2O ) , for the first time addresses the modeling of the optimizer-based uncertainty . Extensive experiments on optimizing test functions , energy functions in a bioinformatics application , and loss functions in the image classification and data privacy attack demonstrate that compared to the start-of-art methods , BL2O substantially improves the performance of optimization and uncertainty quantification , as well as calibration and out-of-domain detection in classification . In the following sections , we first review related methods in details and reveal the remaining gap . We then formally define the problem of optimization with uncertainty quantification and point out the optimizer as a source of uncertainty . After formally defining the optimizer space , the optimal optimizer as a random vector in the space , and the optimizer uncertainty , we propose our novel model , BL2O . And lastly , we compare our BL2O with both Bayesian and non-Bayesian competing methods on extensive test functions and real-world applications . 2 RELATED WORK . Many works ( Wang & Jegelka , 2017 ; Hennig & Schuler , 2012 ) studied optimization with uncertainty quantification under the framework of Bayesian optimization ( Shahriari et al. , 2016 ; Brochu et al. , 2010 ) . In these studies , multiple objectives are sampled from the posterior over the objectives ( p ( f |D ) ) , where D is the observed data . Each sampled objective is optimized for obtaining samples of the global optima : w∗ so that the empirical distribution over w∗ can be built . Approximation is much needed since those approaches need optimization for every sample . For instance , Henrández-Lobato et al . ( 2014 ) uses kernel approximation to approximate the posterior distirbution . Another line of work uses various sampling schemes for estimating the density of posterior distributions . For instance , Bijl et al . ( 2016 ) uses sequential Monte Carlo sampling . De Bonet et al . ( 1997 ) designs a randomized optimization algorithm that directly samples global optima . These methods are much more efficient , but their performance heavily depends on the objective landscapes . Moreover , a few studies ( Ahmed et al. , 2016 ; Lizotte , 2008 ; Osborne et al. , 2009 ; Wu et al. , 2017 ) in Bayesian optimization utilize first-order information to boost the performance of optimization . For instance , Osborne et al . ( 2009 ) uses gradient information to improve the covariance matrix in Gaussian process . Wu et al . ( 2017 ) embeds the derivative knowledge into the acquisition function which is optimized in every iteration . Finally , there are approaches ( Ortega et al. , 2012 ; Cao & Shen , 2020 ) that directly model the shape of posterior as the Boltzmann distributions : p ( w∗|D ) ∝ exp ( −αf ( w∗ ) ) , where α is the scheduled temperature constant . They automatically adjust α during the search in order to balance the exploration-exploitation tradeoff . They beat previous work in terms of both efficiency and accuracy . However , as revealed earlier in the Introduction , none of the methods above consider the uncertainty within the optimizer . 3 METHODS . Notation . We use a bold-faced uppercase letter to denote a matrix ( e.g . W ) , a bold-faced lowercase letter to denote a vector ( e.g . w ) , and a normal lowercase letter to denote a scalar ( e.g . w ) . 3.1 PROBLEM STATEMENT . The goal of optimization is to find the global optimum for an objective function f ( w ) w.r.t . w : w∗ = arg min w f ( w ) . ( 1 ) w∗ is assumed unknown and treated as a random vector in this study . Once a optimizer obtains ŵ , its estimate of w∗ , it is important to assess the quality and the uncertainty of the solution . Considering that many real-world objective functions are nonconvex and noisy in ŵ , solution quality is often measured by ||ŵ−w∗|| , the proximity to the global optimum rather than that to the optimal function value . Examples include energy functions as the objective and RMSDs as the proximity measure in protein docking ( Lensink et al. , 2007 ) . Therefore , the goal of uncertainty quantification ( UQ ) here is the following : P ( ||ŵ −w∗|| 6 rσ|D ) = σ ( 2 ) where rσ is the upper bound of ||ŵ − w∗|| at σ confidence level , and D denotes samples during optimization . Such UQ results additionally provide confidence in the solution ŵ and improve model reliability for end users . To calculate the probability defined in Eq 2 and perform UQ , a direct albeit challenging way is to model the posterior over w∗ ( p ( w∗|D ) ) and then sample from the posterior . When the optimizer g is regarded fixed in existing literature , the posterior is actually p ( w∗|D , g ) . A central contribution of ours is to further consider the optimizer as a source of uncertainty , model it as a random vector in an optimizer space , and perform posterior estimation of p ( w∗|D ) . 3.2 OPTIMIZER UNCERTAINTY : A FRAMEWORK . An optimizer is directly responsible for optimization and thus naturally a source of solution uncertainty . To address this often-neglected uncertainty source , we first define the space of optimizers and then model an optimizer as a point in this space . Considering that many widely-used optimizers are iterative and using first-order derivatives , we restrict the optimizer space as follows : Definition 3.1 ( ( First-order Iterative ) Optimizer Space ) We define a first-order , iterative algorithmic space G , where each point g ∈ G is an iterative optimizer , that has the following mapping : g ( { ∇f ( wτ ) } tτ=1 ) = δwt , where ∇f ( w ) τ and δwt are the gradient and the update vector at τ th and tth iteration , respectively . Here we use g ( · ) to denote a pre-defined update rule and the resulting optimizer . For instance , in gradient descent , g ( { ∇f ( wτ ) } tτ=1 ) = −α∇f ( wt ) , where α is the step size . Now that the optimizer space is defined , we next define the ( unknown ) optimal optimizer and its uncertainty . Definition 3.2 ( Optimal Optimizer ) We define the optimal optimizer g∗ ∈ G as the optimizer that can obtain the lowest function value with a fixed budget T : g∗ = arg min g∈G ( T∑ t=1 f ( wtg ) ) ( 3 ) where wtg = w t−1 g + g ( { ∇f ( wτg ) } t−1τ=1 ) is the parameter value at tth iteration updated through the optimizer g. In practice the optimal optimizer g∗ is unknown so we treat g∗ as a random vector and formally define the optimizer uncertainty as follows : Definition 3.3 ( Optimizer Uncertainty ) Let G be the algorithmic space , where each point g ∈ G is an optimizer . We assume there is a prior distribution over the optimal optimizer g∗ as p ( g∗ ) . We also assume a likelihood distribution as p ( D|g∗ ) , whereD are the observed data ( sample trajectory ) given g∗ . Then we define the optimizer uncertainty through p ( g∗|D ) ∝ p ( D|g∗ ) p ( g∗ ) . To inject the optimizer uncertainty into p ( w∗|D ) , it is straightforward to have the following integration for posterior estimation : p ( w∗|D ) = ∫ p ( g∗|D ) p ( w∗|D , g∗ ) dg ( 4 )

Summary: The paper considers the question of quantifying the uncertainty that arises from the optimiser used to perform inference in a given model. Taking a Bayesian approach, the aim is to deduce the posterior over the space of optimisers. The form for the posterior is chosen to be a Boltzmann distribution which is then approximated with a multivariate Gaussian using a KL divergence. The parameterisation for the posterior is defined using an LSTM neural network.