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: The authors looked at uncertainty quantification (UQ) of optimization algorithms. Working under the L2O framework and adopting the Bayesian approach, they used a neural network to model the updating rule, and endowed the network weights with a prior. By treating the solution iterates and their gradients as data, the authors defined a conditional likelihood on them. The resulting posterior on the network weights is then approximated using a variational algorithm. To conduct inference, the authors computed the predictive density of the solution through Monte Carlo sampling. This end-to-end pipeline is called uncertainty aware L2O (UA-L2O) and it is applied to quantify uncertainty in various learning tasks.