Bayesian Modeling and Uncertainty Quantification for Learning to Optimize: What, Why, and How

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 . Instead their solutions ’ usefulness ( sometimes based on their proximity to optima ) , when the optima are unknown , can be very uncertain . Being able to quantify solution uncertainty directly provides calibration with ensured awareness of the solution quality and usefulness ( and another potential benefit is in optimization performance by enhancing the search efficiency ) . For instance , reliable and trustworthy machine learning models demand uncertainty awareness and quantification ( UQ ) during training ( optimizing ) such models , whereas in reality deep neural networks without proper modelling of uncertainty suffer from overconfidence and miscalibration ( Guo et al. , 2017 ) . In another example of 3D prediction of proteinprotein interactions , even though there exists the model uncertainty of the objective function and the data uncertainty of the protein structure data ( Cao & Shen , 2020 ) , state-of-the-art methods only predict several ranked solutions ( Porter et al. , 2019 ) without any associated uncertainty , which is hard for biologists to interpret . Despite progress in optimization with uncertainty-awareness , fundamental 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 ) , whereas inconspicuous attention was paid to the uncertainty arising from the optimizer that is directly responsible for deriving the end solutions . The optimizer is usually prefabricated and fixed . 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 unaltered for the entire process . For Bayesian neural networks training , the update rule of the iterative optimizer is usu- ally chosen off-the-shelf like Adam , SGD , or RMSDrop . In principle , it is important to recognize the existence of an optimization algorithm space where the specific optimizer lies as well as the importance of the optimizer uncertainty , intrinsically defined in the space to the optimization and UQ solutions . In practice , such uncertainty is unwittingly omitted in the current context where an optimizer is treated as a de facto specified sample in the space . A naı̈ve characterization for a hand-built optimizer is to rely on a few of its hyper-parameters ( Wenzel et al. , 2020 ; Lorraine & Duvenaud , 2018 ) , whereas such a parameterization space could be too biased and restricted to span a broader and more complicated optimizer manifold , potentially resulting in inaccurate UQ results ( see Section 4 , confidence estimation for test functions ) . Fortunately , the recent surge of learning to optimize ( L2O ) ( Andrychowicz et al. , 2016a ; Li & Malik , 2016 ; Chen et al. , 2021 ) which is another prominent optimization paradigm , parameterized by a neural network using high-dimensional weights and following the data-driven ethos to learn the update rule adaptively , therefore introducing the less inductive bias and accompanied with the universal approximation ability , indicates a plausible solution to depict the sophisticated and more complete optimizer space . Contributions . The aforementioned introduction reveals and discusses the fundamental question of why modelling the optimizer uncertainty , and our contingent questions would be , what defines the optimizer uncertainty and how to enable UQ during optimization . We give the following answers . ( i ) We define the prior and likelihood of the optimizer ( Neal , 2012 ) , which routinely determine the optimizer posterior with the general product rule and Bayes ’ theorem . The optimizer uncertainty is therefore well-defined . ( ii ) To unseal UQ within optimization , we treat an optimizer as a random sample from an algorithmic space of iterative optimizers . We leverage the surging L2O technique ( Andrychowicz et al. , 2016a ; Li & Malik , 2016 ) with high-dimensional weights for the optimizer space parameterization . ( iii ) We further construct the end-to-end training pipeline via variational inference ( Kingma & Welling , 2013 ; Higgins et al. , 2016 ) to avoid the expensive computational cost of Markov chain Monte Carlo ( MCMC ) and degenerated posterior . The newly trained L2O is referred as uncertainty-aware L2O ( UA-L2O ) . In summary , the core intellectual value of this work is for the first time , we recognize and quantify a novel form of uncertainty that lies in the optimization algorithm space parameterized by L2O , apart from the classical data- or model-based uncertainties ( also known as epistemic and aleatoric uncertainties ( Fox & Ülkümen , 2011 ) ) . Extensive experiments show that UA-L2O had superior capability in uncertainty calibration thanks to the accurate estimation of solution posteriors . It confidence levels well match the chance for the global optimum falling in the corresponding intervals . Intriguingly , although not directly targeted , UA-L2O also outperformed competing methods in optimization performance , as seen in optimizing two out of three test functions in high dimensions , the loss function in data privacy attack , energy functions in four out of five protein-docking cases . 2 PRELIMINARIES . 2.1 LEARNING TO OPTIMIZE . Let us consider an optimization problem : minx f ( x ) where x ∈ Rd . A classic optimizer often iteratively updates x based on a handcrafted rule . For example , the first-order gradient descent algorithm takes an update at iteration t based on the local gradient at the instantaneous point xt : xt+1 = xt − α∇f ( xt ) , where α is the step size . Learning to Optimize ( L2O ) has lots of freedom to use the available information . We define zt as optimization trajectories ’ historical information up to time t , e.g. , the existing iterates x0 , . . . , xt , and/or their gradients ∇f ( x0 ) , . . . , ∇f ( xt ) . L2O models an update rule by a predictor function g of zt : xt+1 = xt − g ( zt ; φ ) , where the mapping of g is parameterized by φ . Finding an optimal update rule can be formulated mathematically as searching for a good φ over the parameter space of g. Practically , g is often a neural network . Since neural networks are universal approximators , L2O has the potential to discover completely new update rules without relying on existing rules . In order to find a desired φ associated with a fast optimizer , ( Andrychowicz et al. , 2016b ) proposed to minimize the weighted sum of the objective function f ( xt ) over a time span T : min φ T∑ t=1 wtf ( xt ) , with xt+1 = xt − g ( zt ; φ ) , t = 0 , . . . , T − 1 , ( 1 ) where w1 , . . . , wT are the weights whose choices vary from case to case and depend on empirical settings . Note that φ determines the objective value through determining the iterates xt . L2O solves the problem ( 1 ) for a desirable φ and correspondingly the update rule g ( zt ; φ ) . A typical L2O workflow is divided into two stages ( Chen et al. , 2021 ) : a meta-training stage that learns the optimizer with a set of similar optimizees from the task distribution ; and a meta-testing stage that applies the learned optimizer to training new unseen optimizees . The meta-training process often occurs offline and is time consuming . However , the online application of the method is ( aimed to be ) time saving . 2.2 THE UNCERTAINTY OF OPTIMIZATION . The choice of optimizers is recognized to remarkably impact the solution quality , especially for non-convex and rugged objectives such as various loss functions for training deep networks . While non-convex optimization is NP-hard and the global optimum is never guaranteed , domain users usually can only blindly take the solutions given by a particular optimizer , or choose from several well-known algorithms based on human experiences or on an ad hoc basis . Such trial-and-error selections are biased and restricted to span a more complete optimizer manifold , and computationally expensive , even intractable when new algorithms are to be discovered by L2O , whose “ fitness ” or confidence on certain problem instances ( especially instances with shifts from the target task distributions ) is never known . Given a specific problem instance , if we can provide a confidence score or quantify the uncertainty of each candidate optimizer , it will certainly improve the selection and calibration of optimization algorithms , in particular the learned optimizers since they are essentially all “ new ” , and hence enhancing the reliability of L2O . As for non-convex objectives , different choices of optimizers , even each algorithm ’ s hyperparameter or initialization variations , can lead to vastly different solutions , implying a new type of uncertainty arising from choosing optimizers . Stochastic optimization methods like random search ( Zhigljavsky , 2012 ) , simulated annealing ( Kirkpatrick et al. , 1983 ) and genetic algorithms ( Goldenberg , 1989 ) inject the randomness into algorithms to reduce solution uncertainty . However , they can not reliably quantify its effect on the final solutions . Several works explored uncertainty quantification during optimization , by Bayesian optimization ( Hennig & Schuler , 2012 ; Wang & Jegelka , 2017 ) or Monte Carlo sampling ( Ortega et al. , 2012 ) . However , those methods are often computationally too expensive , not to mention applied within the already costly L2O process . More importantly , they neither explicitly consider the uncertainty within the optimizer itself , nor concern the automatic learning and discovery of the most suitable optimizer . 3 TECHNICAL APPROACH . To quantify and address such “ fitness ” especially for the L2O parameterization in a principled way , we will introduce this new optimizer uncertainty arising from the choice of optimization algorithms , which is different from the classical data- or model-based uncertainties ( also known as epistemic and aleatoric uncertainties ( Fox & Ülkümen , 2011 ) ) . We will explore : how to quantify optimizer uncertainty ? how to design a learned optimizer so that its uncertainty can be more accurately quantified ? what is the best calibrated optimizer we can choose or learn for a given problem ? The core innovation is to treat an optimizer as a random sample from a continuous algorithmic space , rather than one of a few hand-crafted update rules , so as to model the intrinsic uncertainty within the optimizer . This novel view of algorithmic space is particularly enabled by L2O . Note an L2O update rule is typically parameterized by a neural network g , with its inputs zt and learnable parameters φ . Its learning capacity allows us to “ sample ” optimizers by taking g ( · ; φ ) with different weights φ . That is sufficiently versatile and flexible in practice . 3.1 CONCEPTS FOR OPTIMIZER UNCERTAINTY . Without loss of generality , we define am iterative algorithmic space G , where each point g ∈ G produces an update g ( zt ) dependent on the current/past zero-th order and/or first-order information zt . We treat g as a random vector from G , that leads to defining the optimizer uncertainty as follows : Definition 1 ( Optimizer Uncertainty ) Let G be the algorithmic space , where each point g ∈ G is an optimizer ( omitting φ parameterization ) . We assume that 1. g has a prior distribution p ( g ) ; 2 . Its likelihood can be interpreted as p ( zt|zt0 , g ) = ∏t i=t0+1 p ( xi|zi−1 , g ) , ∀t0 < t. The likelihood factorization is based on causality of first-order optimizers , i.e. , ( xi , zi ) depends solely on { ( xi0 , zi0 ) } i0 < i given g. Thereby , we define the optimizer uncertainty at step t as the posterior of g , that is a conditional distribution on the t steps of the trajectory , using Bayes ’ theorem : p ( g|zt ) ∝ p ( g ) t∏ i=1 p ( xi|zi−1 , g ) . ( 2 ) Intuitively , the prior p ( g ) represents the belief about well-performing optimizers , the likelihood p ( zt|zt0 , g ) represents the probability of observing an optimization trajectory ( data ) under a given optimizer g , and the posterior p ( g|zt ) represents the probability for g being the optimizer generating the observed data . Prior work on hyper-parameter optimization ( HPO ) for classical algorithms ( Feurer & Hutter , 2019 ) could be viewed as a special case of L2O ( e.g. , only a few hyper-parameters are “ learnable ” , which resulting in a much biased and restricted prior p ( g ) ) . The seminal work Bergstra et al . ( 2011 ) optimized hyper-parameters using posterior-based fitness modeling . Recently , Wenzel et al . ( 2020 ) proposed a non-Bayesian approach of hyper-parameter ensembling that can estimate the predictive uncertainty of an optimization algorithm by varying hyper-parameters . Our work generalizes from tuning a few hyper-parameters to learning the entire optimizer space , and extends a Bayesian treatment . Since L2O involves much higher-dimensional parameters compared to HPO , its uncertainty quantification calls for computationally tractable and scalable approaches , which is detailed next .

Summary: This paper proposes a Bayesian Learning to Optimize technique called Uncertainty-Aware Learning to Optimize (UA-L2O). The main contribution of this paper is its introduction of the concept of the space of optimization algorithms and treating an optimizer as a random sample from it, and deriving a proper Bayesian inference algorithm based on that concept. The space of optimization algorithms, to simplify, is expressed as a space of parameters of learned optimizers $\phi$, and a Gaussian prior is placed on those parameters. The posterior inference procedure is approximated with variational inference with reparameterization tricks, and the actual deployment of the learned optimizer is done with Monte-Carlo approximation with learned variational distribution for the optimizer parameters $\phi$. The proposed algorithm is applied to various synthetic and real-world optimization problems and showed some useful calibration behaviors.