B AYESIAN M ODELING AND U NCERTAINTY Q UANTIFICATION FOR L EARNING TO O PTIMIZE : W HAT , W HY , AND H OW
A BSTRACT
Optimizing an objective function with uncertainty awareness is well-known to improve the accuracy and conﬁdence of optimization solutions.
Meanwhile another relevant but very different question remains yet open : how to model and quantify the uncertainty of an optimization algorithm (a k a., optimizer) itself?
To close such a gap, the prerequisite is to consider the optimizers as sampled from a distribution, rather than a few prefabricated and ﬁxed update rules.
We ﬁrst take the novel angle to consider the algorithmic space of optimizers and provide deﬁnitions for the optimizer prior and likelihood, that intrinsically determine the posterior and therefore uncertainty.
We then leverage the recent advance of learning to optimize (L2O) for the space parameterization, with the end-to-end training pipeline built via variational inference, referred to as uncertainty-aware L2O ( UA-L2O ).
Our study represents the ﬁrst effort to recognize and quantify the uncertainty of the optimization algorithm.
The extensive numerical results show that, UA-L2O achieves superior uncertainty calibration with accurate conﬁdence estimation and tight conﬁdence intervals, suggesting the improved posterior estimation thanks to considering optimizer uncertainty.
Intriguingly, UA-L2O even improves optimization performances for two out of three test functions, the loss function in data privacy attack and four of ﬁve cases of the energy function in protein docking.
Our codes are released at https://github.com/Shen-Lab/Bayesian-L2O .
I NTRODUCTION
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 the 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 beneﬁt is in optimization performance by enhancing the search efﬁciency).
For instance, reliable and trustworthy machine learning models demand uncertainty awareness and quantiﬁcation (UQ) during training (optimizing) such models, whereas in reality deep neural networks without proper modelling of uncertainty suffer from overconﬁdence 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; Lyu et al., 2021; Cao & Shen, 2020), whereas inconspicuous attention was paid to the uncertainty arising from the optimizer that is directly re- sponsible for deriving the end solutions.
The optimizer is usually prefabricated and ﬁxed.
For instance, there are several popular update rules in Bayesian optimization, such as expected improvement (Vazquez & Bect, 2010) and upper conﬁdence 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 usually chosen off-the-shelf, such as Adam, SGD, and RMSDrop.
In principle, it is important to recognize the existence of an optimization algorithm space where a speciﬁc optimizer lies as well as the importance of the optimizer uncertainty, intrinsically deﬁned 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 speciﬁed 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, conﬁdence estimation for test functions).
Fortunately, the recent surge of learning to optimize (L2O) (Andrychowicz et al., 2016a; Li & Malik, 2016; Cao et al., 2019a; You et al., 2020; Li et al., 2020; Chen et al., 2020b;a; Shen et al., 2021; 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 deﬁnes the optimizer uncertainty and how to enable UQ during optimization.
We give the following answers.
(i)
We deﬁne 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-deﬁned.
(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 to as uncertainty-aware L2O ( UA-L2O ).
In summary, the core intellectual value of this work is, for the ﬁrst 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 & ¨ Ulk¨umen, 2011)).
Extensive experiments show that UA-L2O had superior capability in uncertainty calibration thanks to the accurate estimation of solution posteriors.
It conﬁdence levels well matched 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 ﬁve protein-docking cases.
P RELIMINARIES 2. L EARNING TO O PTIMIZE
Let us consider an optimization problem: min x f ( x ) where x ∈ R d .
A classic optimizer often iteratively updates x based on a handcrafted rule.
For example, the ﬁrst-order gradient descent algorithm takes an update at iteration t based on the local gradient at the instantaneous point x t : x t +1 = x
t − α ∇ f ( x t ) , where α is the step size.
Learning to Optimize (L2O) has lots of freedom to use the available information.
We deﬁne z t as optimization trajectories’ historical information up to time t , e.g., the existing iterates x 0 , . . .
, x t , and/or their gradients ∇ f ( x 0 ) , . .
. , ∇ f ( x t ) .
L2O models an update rule by a predictor function g of z t : x t +1 = x
t − g ( z t ; φ ) , 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 ﬁnd a desired φ associated with a fast optimizer, (Andrychowicz et al., 2016b) proposed to minimize the weighted sum of the objective function f ( x t ) over a time span T : min φ
T (cid:88) t = w t f ( x t ) , with x t +1 = x
t − g ( z t ; φ ) , t = 0 , . . .
, T − 1 , (1) where w 1 , . . .
, w T 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 x t .
L2O solves the problem (1) for a desirable φ and correspondingly the update rule g ( z t ; φ ) .
A typical L2O workﬂow 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 new unseen optimizees.
The meta-training process often occurs ofﬂine and is time consuming.
However, the online application of the method, meta testing, is (aimed to be) time saving.
2. T HE U NCERTAINTY OF O PTIMIZATION
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 “ﬁtness” or conﬁdence on certain problem instances (especially instances with shifts from the target task distributions) is never known.
Given a speciﬁc problem instance, if we can provide a conﬁdence 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 cannot reliably quantify its effect on the ﬁnal solutions.
Several works explored uncertainty quantiﬁcation during optimization, by Bayesian optimization (Hennig & Schuler, 2012; Wang & Jegelka, 2017) or Monte Carlo sampling (Ortega et al., 2012).
Several works in Bayesian optimization explored uncertainty quantiﬁcation during optimization, by modeling the distributions over the functional space (Hennig & Schuler, 2012; Wang & Jegelka, 2017) or directly over the optimum (Ortega et al., 2012; Cao & Shen, 2020).
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.
T ECHNICAL A PPROACH
To quantify and address such “ﬁtness” 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 & ¨Ulk¨umen, 2011)).
We will explore: how to quantify optimizer uncertainty?
how to design a learned optimizer so that its uncertainty can be more accurately quantiﬁed?
what is the best calibrated optimizer one 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 that an L2O update rule is typically parameterized by a neural network g , with its inputs z t and learnable parameters φ .
Its learning capacity allows us to “sample” optimizers by taking g ( · ; φ ) with different weights φ .
That is sufﬁciently versatile and ﬂexible in practice.
3. C ONCEPTS FOR O PTIMIZER U NCERTAINTY
Without loss of generality, we deﬁne an iterative algorithmic space G , where each point g ∈ G maps to an update g ( z t ) dependent on the current/past zero-th order and/or ﬁrst-order information z t .
We treat g as a random vector from G , that leads to deﬁning the optimizer uncertainty as follows:
Deﬁnition 1 (Optimizer Uncertainty) Let G be the algorithmic space, where each point g ∈ G is an optimizer (omitting φ parameterization).
We assume that
g has a prior distribution p ( g ) ; 2.
Its likelihood can be interpreted as p ( z t | z t 0 , g ) =
(cid:81)
[-= t 
-][Equation]
t 0 +
[Equation]
( x i
[Equation]
z [Equation] − 1 , g ) , ∀ t 0 < t .
The likelihood factorization is based on the causality of iterative optimizers, i.e., ( x i , z i ) depends solely on { ( x i 0 , z i 0 ) } i 0 <
i given g .
Thereby, we deﬁne 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:
[Equation]
Intuitively, the prior p ( g ) represents the belief about well-performing optimizers, the likelihood p ( z t | z t 0 , g ) represents the probability of observing an optimization trajectory (data) under a given optimizer g , and the posterior p ( g | z t ) 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 results in a much biased and restricted prior p ( g ) ).
The seminal work Bergstra et al.
(2011) optimized hyper-parameters using posterior-based ﬁtness 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 quantiﬁcation calls for computationally tractable and scalable approaches, which is detailed next.
3. L2O WITH B AYESIAN U NCERTAINTY Q UANTIFICATION
Now that we have the optimizer space G parameterized in φ , an optimizer g is modeled as a random vector in the space.
That allows for reducing the modeling of optimizer uncertainty p ( g | z t ) to modeling the posterior p ( φ | z t ) .
For simplicity, we use a Gaussian prior of φ : p ( φ ) ∝ exp( − 12 λ (cid:107) φ (cid:107) 22 ) where λ is a constant.
We also follow the widely-adopted idea in (Gal & Ghahramani, 2016; Bishop & Tipping, 2003; Ortega et al., 2012; Cao & Shen, 2020), to view the loss function value as proportional to the negative logarithmic likelihood, i.e., p ( x t | z t − 1 , g )
=
p ( x t | z t − 1 , φ ) ∝ exp( − f ( x φ t )) .
We can then calculate the posterior of the optimizer based upon equation (2):
[Equation] [Equation]
Training.
Maximum a posteriori (MAP) estimation of the optimizer parameters leads to a degenerate distribution and point estimation without uncertainty.
Moreover, directly maximizing the optimizer posterior in Eq.
3 via Markov chain Monte Carlo (MCMC) would encounter computational intractability considering the high-dimensional space of the optimizer space parameterized in
L2O by neural network weights.
Instead, we introduce a variational distribution q ( φ ; θ ) (Kingma & Welling, 2013) with the variational parameter θ
and accordingly we introduce the following objective function whose minimization is equivalent to maximizing the evidence lower bound (ELBO) (see Appendix D for the derivation):
[Equation]
We refer to the L2O under the new meta-training loss (4) w.r.t θ as uncertainty-aware L2O
( UAL2O ), which can be trained end-to-end with the Gaussian parametrization and the reparametrization [Figure] trick (Kingma & Welling, 2013; Higgins et al., 2016).
The feed-forward computational chart for training UA-L2O is depicted in Figure 1.
Deployment .
With the well-trained UA-L2O and the found θ ∗ , we ideally assume the function values of the trajectory are submartingales that E
[Equation] ∞ .
Then based upon martingale convergence theorem (Doob, 1953) lim
[Equation]
exists almost surely, i.e. it converges with inﬁnite optimization iterations.
Within the ﬁnite steps,
UA-L2O has a well-deﬁned approximated posterior q ( φ ; θ ∗ ) , and hence the baked-in ability to quantify optimizer uncertainty at the solution.
Therefore, the posterior of solutions is calculated by the integral with the variational distribution q ( φ ; θ ∗ ) as:
[Equation]
which can be estimated by Monte Carlo sampling.
Speciﬁcally, for estimating the solution posterior, the learned model is run 10,000 times with random initialization and different trajectories z t accordingly.
The optimizer parameter θ is sampled from q ( θ | φ ∗ ) once per iteration, following variational inference as in Kingma & Welling (2013) and Higgins et al.
E XPERIMENTS
We examine our proposed UA-L2O on optimizing: (i) non-convex test functions, (ii) loss functions in data privacy attack, and (iii) energy functions in predicting 3D protein-protein interactions (protein docking).
We compare UA-L2O with three non-Bayesian methods: manually-designed Adam (Kingma & Ba, 2014) and particle swarm optimization (PSO) (Kennedy & Eberhart, 1995) as well as DM LSTM (Andrychowicz et al., 2016a), an L2O method.
We also compare to Adam and PSO with hyper-parameter Bayesian optimization (Adam-BO and PSO-BO), Adam with learning rate ensembles (Adam lr Ensemble) and Adam with stochastic gradient MCMC (Adam Noisy Gradient).
Lastly we compare to Bayesian active learning (BAL) (Cao & Shen, 2020).
For Adam and PSO, we make their hyper-parameters probabilistic to calculate the posterior as shown in Table 6 in Appendix C. All algorithms are run for 10,000 times with random initialization to obtain the empirical posterior distributions p ( x ∗ | z T ) , following Eq.
5 for UA-L2O and its counterparts for others.
We access the optimization uncertainty with the following expression (Ortega et al., 2012;
Cao & Shen, 2020):
[Equation] [Figure]
where lb σ and ub σ are the lower and upper bounds at the conﬁdence level of σ , and we choose the one with the lowest function value (we also study the mean value in Appendix E) to be the estimated solution ˆ x .
Ablation on different runs and comparison with competitors in wall-clock time are presented in Appendix E. For the choice of assessing the proximity of solutions rather than function values | f ( x ∗ ) − f
(ˆ x ) | , the reason is that in many optimization applications (e.g. privacy attack and protein docking), the objective function actually plays the role of the “tool” instead of the “goal”, thus the function value is not the key metric to assess the quality of solutions in particular when the objective function is non-convex and noisy.
Evaluation metrics.
UA-L2O is designed for uncertainty awareness in the optimizer.
To assess the calibration of uncertainty quantiﬁcation (UQ), we use the accuracy of conﬁdence, following (Ortega et al., 2012; Cao & Shen, 2020).
In other words, we compare expected conﬁdence levels σ as in Eq.
6 (0.8 and 0.9 in this paper) and the empirical probability ε σ of “success” when the corresponding conﬁdence intervals indeed contain the global optima.
When there is a tie in the major assessment metric σ − ε σ , the tightness of the conﬁdence interval for the solution’s proximity to the global optimum (Eq. 6), r σ = ub σ − lb σ , serves as a tie-breaker.
A schematic illustration of the uncertainty assessment is shown in Figure 2.
Although UA-L2O is designed for the uncertainty awareness rather than the quality of optimization solutions, we still assess potential beneﬁt in optimization performance, using the expected distance between the optimized solutions and the global optimum E x ∗ ∼ p ( x ∗ | z T ) (cid:107) x ∗ − x true (cid:107) 2 .
The lower distances indicate the better solution quality (optimization performance).
Two-stage training.
Due to the extremely high-dimensional optimizer space and the rugged posterior landscape, it is over-challenging to directly train the model through the equivalency of ELBO optimization in Equation (4).
We therefore perform two-stage training where we ﬁrst train our model in a non-Bayesian way and then use the resulting θ as the warm start for ﬁne-tuning the mean variational parameters in the second stage.
We examine the performance of the warm-start model in
Appendix E to show the beneﬁt of the second-stage training.
Implementations.
The model is implemented in Tensorﬂow 1.13 (Abadi et al., 2016) and optimized by Adam (Kingma & Ba, 2014).
For the L2O architecture, we use the coordinate-wise LSTM from (Andrychowicz et al., 2016a) containing 10,282 free parameters.
The meta-training and test data are described in Appendix B. For all experiments, the length of LSTM is set to be 20, with 5,000 training epochs for both training stages.
We tune the λ values in equations (3) and (4) from { 0.1, 1, 10 } , validated by validation performance in test functions and by the tightness of conﬁdence intervals in privacy attack and protein docking.
The ablation study on λ is presented in Appendix C showing that UA-L2O is not sensitive to the value of λ .
4. N ON -C ONVEX T EST F UNCTIONS
Setup.
We ﬁrst evaluate the performance on the test functions in the global optimization benchmark (Jamil & Yang, 2013).
We choose three extremely rugged, non-convex functions, Rastrigin,
Ackley and Griewank with analytical forms shown in Appendix B, paired with dimensionalities in { 3 , 6 , ..., 30 } .
Results.
We compare UA-L2O with competitors in optimization and UQ performances shown in Figure 3.
In the most relevant assessment metric, the accuracy of conﬁdence levels (bottom two rows), UA-L2O performs the best in all but few cases for all the three functions, outperforming hand-engineered optimizers, non-Bayesian or Bayesian, and L2O without uncertainty awareness (DM LSTM).
The few exceptions, such as the Griewank function in 27 dimensions, could not be attributed to the higher dimensionality (as that was not the case for UA-L2O for Rastrigin and Ackley functions).
DM LSTM also showed overconﬁdence in such exceptions but the few exceptions cannot be attributed to L2O either.
Given that UA-L2O consistently outperforms competitors in conﬁdence accuracy, comparing the tightness of conﬁdence intervals alone (Figure 9) is not meaningful.
Note that other methods, when having tighter conﬁdence intervals, were often too overconﬁdent to realize that the true global optima were inside the intervals much less than their conﬁdence levels suggested.
We also examined potential but not intended beneﬁts of UA-L2O in optimization performance (top row of Figure 3).
Interestingly, UA-L2O even improved optimization performances for Ackley and Griewank functions when the dimensionality is above 12, while being slightly worse in Rastrigin.
[Figure]
These results echo our conjecture that, although our central goal is not for better optimization solutions but for their better uncertainty awareness, better uncertainty calibration could sometimes improve optimization performances
[Figure] We further plot the conﬁdence for solution regions segmented in (cid:107) x ∗ − ˆ x (cid:107) 2 versus the corresponding optimization performance (cid:107) x ∗ − x true (cid:107) 2 (see Appendix A for detailed procedure) in Figure 10.
We ﬁnd that when UA-L2O achieves better optimization performance (Ackley and Griewank as opposed to Rastrigin), it generates solutions with higher conﬁdence in the regions closer to the global optima and with narrower distributions (more certain).
In contrast the competitors do not show such trends.
4. D ATA P RIVACY A TTACK
Setup.
We next apply our model to an application that critically needs UQ.
As many machine learning models are deployed publicly, it is important to avoid leaking private and sensitive information, such as ﬁnancial data and health data.
Data privacy attack (Nasr et al., 2018) studies this problem by playing the role of hackers and attacking machine learning models to quantify the risk of privacy leakage.
Better attacks would help models to be better prepared for privacy defense.
We use the model and dataset in (Cao et al., 2019b), where each input has 9 features involving patient genetic information shown in Appendix B, and the output is the probability of the clinical signiﬁcance (pathogenicity) for genetic variants in patients.
We study the following model inversion attack (Fredrikson et al., 2015): given 5 features x (cid:48) ∈ [0 , 1] 5 out of 9 and the label y of each patient, the privacy attack aims to recover the rest 4 features x true ∈ [0 , 1] 4 (potentially sensitive patient information).
Therefore, the optimization is min x ∈ [0 , 1] 4 (cid:80) ( g ( x (cid:48) , x ) − y ) 2 for all patients where g is a trained predictor.
The closeness between the optimized and real input features can quantify the risk of information leakage and the quality of the attack.
We evaluate our method for all test cases in (Cao et al., 2019b).
Results.
We report the UQ and optimization performances of UA-L2O in Table 1.
In UQ, UA-L2O’s performance (the accuracy of estimated conﬁdence) stands out against competitors, although not as dominant as it does to competitors in test functions. Even in optimization, UA-L2O outperforms all competitors, which again demonstrate the potential beneﬁt of uncertainty-awareness to optimization.
[Table]
Note that PSO had slightly more accurate conﬁdence but much looser intervals, suggesting ﬂat posteriors and leading to worse optimization.
We argue that, unlike pure optimization problems in test functions
[Table]
the optimization for privacy attack involves the generalization issue in the machine learning paradigm: less generalizable samples pose challenges to not only optimization but also UQ.
To verify this argument
[Table]
we plot optimization performance for individual samples versus their estimated posterior precision in Figure 4 showing a positive correlation: samples easier to generalize are closer to the optima (more leftward in the x -axis)
[Table]
and more precise in estimated conﬁdence (more downward in the y -axis).
The tightness of conﬁdence intervals is also well correlated with solution quality
[Table]
[Figure]
Similar to that in Section 4.
[Figure]
we also plot conﬁdence versus optimization performance in Figure
[Figure]
to further demonstrate the calibration capability of UA-L2O. To demonstrate the practical use of developed UQ [Figure] we conduct an experiment on “failure” case detection [Figure]
that we use assessed uncertainties to probe the optimization performance of individual runs. Results in Appendix F show UA-L2O provides the best detection. 4. E NERGY F UNCTIONS FOR P ROTEIN D OCKING
[Figure]
Setup.
We lastly examine UA-L2O using a bioinformatics application: predicting the 3D structures of proteincomplexes (Smith & Sternberg, 2002) or protein docking.
Ab initio protein docking can be recast as optimizing a noisy and expensive energy function in a highdimensional conformational space (Cao & Shen, 2020).
While solving such optimization problems still remains difﬁcult, quantifying the uncertainty of resulting solutions is even more challenging.
The compared state-of-the-art method is BAL (Cao & Shen, 2020).
We calculate the energy function (objective function f ( x ) ) in a CHARMM19 force ﬁeld as in (Moal & Bates, 2010).
25 protein-protein complexes are chosen from the protein docking benchmark set 4.0 (Hwang et al., 2010) as the training set, which is shown in Appendix B. For each complex, we choose 5 starting points (top-5 structure predictions from ZDOCK (Pierce et al., 2014)).
In total, our training set includes 125 samples.
Moreover, we parameterize the search space for ﬂexible docking as R 12 as in BAL.
The resulting f ( x ) is fully differentiable in the search space.
We only consider 100 interface atoms due to the computational concern.
The number of iterations for one training epoch is 600 and in total we have 5,000 training epochs.
Both BAL and UA-L2O have 600 iterations during the testing stage. [Figure] [Table]
[Figure] [Table]
[Figure]
[Table]
Results. We compare UA-L2O and BAL in Table 2. For UQ, UA-L2O not only achieves a more accurate conﬁdence as shown in Figure 6 and Appendix C (three of ﬁve cases had ground truths located in 90% conﬁdence intervals), but also has a clearer trend that the better solutions correspond to tighter intervals (more certain) as in Figure 5 (with Pearson correlation coefﬁcient of 0.59 versus -0.8 for BAL).
Even for optimization, UA-L2O outperforms other methods in four out of ﬁve cases.
The results show signiﬁcant advantages of UA-L2O compared with the state-of-the-art BAL in both
UQ and optimization for protein docking.
4. R ESULT S UMMARY
We brieﬂy summarize the results as follows
• In general, UA-L2O achieves better optimization performance with a few exceptions.
This echos the hypothesis that the assessment of uncertainty is able to enhance the search efﬁciency and effectiveness during optimization.
• UA-L2O provides the most accurate posterior estimation most of the time.
This empirical evidence veriﬁes the effectiveness of our uncertainty-aware L2O inspired from theory.
• Accurate posterior estimation is a key for UA-L2O’s superior calibration capability , such that UA-L2O generates solutions closer to optima with higher conﬁdence and tighter conﬁdence intervals, whereas competitors suffer from overconﬁdence or miscalibration.
Such calibration capability is critical to detect the failure cases in real-world applications.
C ONCLUSIONS
Current optimization algorithms, even with uncertainty-awareness, do not address the uncertainty arising within the optimizer itself.
To close the gap, in this paper we attempt to ask and address three fundamental questions, why modelling the optimizer uncertainty, what deﬁnes the optimizer uncertainty, and how to enabling UQ during optimization.
We ﬁrst emphasize the uncertainty arising from the optimizer is a crucial source that directly responses for the end solutions deriving, except for data- and model-uncertainty.
Next, we deﬁne the prior and likelihood of the optimizer which determines the optimizer posterior and the optimizer uncertainty.
We further leverage learning to optimize (L2O) for the optimizer parameterization, treating an optimizer as a random sample from an algorithmic space of iterative optimizers, with the end-to-end training pipeline built via variational inference.
Extensive experiments show UA-L2O achieves the optimization performance, superior in two out of three test functions with the high variable dimension, best in the loss function in data privacy attack, and exceeded in the energy function for protein docking in four out of ﬁve cases.
Besides, it delivers the most accurate posterior estimation and calibration capability, with the solutions closer to the ground truth of a larger population and tighter conﬁdence intervals.
Our study represents the ﬁrst effort to recognize and quantify the uncertainty of the optimization algorithm, with extensive support from numerical results and analysis. A CKNOWLEDGMENT
This project was in part supported by NSF (CCF-1943008 to YS, CCSS-2113904 to ZW).
Portions of this research were conducted with the advanced computing resources provided by Texas A&M
High Performance Research Computing.
R EPRODUCIBILITY S TATEMENT
To ensure reproducibility, we use the Machine Learning Reproducibility Checklist v2.0, Apr.2020 (Pineau et al., 2021).
• For all models and algorithms presented, – A clear description of the mathematical settings, algorithm, and/or model.
We clearly describe all of the settings, formulations, and algorithms in Section 3.
– A clear explanation of any assumptions.
We state our assumptions clearly in Section 3.
– An analysis of the complexity (time, space, sample size) of any algorithm.
We do not make the analysis.
• For any theoretical claim , – A clear statement of the claim.
We provide a clear statement in Section 3.
– A complete proof of the claim.
We do not make theoretical proofs.
• For all datasets used, check if you include: – The relevant statistics, such as number of examples.
We provide all the related statistics in Section 4 and Appendix B. – The details of train/validation/test splits.
We give this information in our repository.
– An explanation of any data that were excluded, and all pre-processing step.
We give this information in our repository.
– A link to a downloadable version of the dataset or simulation environment.
Our repository contains all of the instructions to run experiments on the datasets in this work.
We put it in supplementary materials.
– For new data collected, a complete description of the data collection process, such as instructions to annotators and methods for quality control.
We do not collect or release new datasets.
• For all shared code related to this work, check if you include: – Speciﬁcation of dependencies.
We give installation instructions in the README of our repository.
– Training code.
The training code is available in our repository.
– Evaluation code.
The evaluation code is available in our repository.
– (Pre-)trained model(s).
We do not release trained models.
– README ﬁle includes table of results accompanied by precise command to run to produce those results.
We include a README with detailed instructions to reproduce all of our experimental results.
• For all reported experimental results , check if you include: – The range of hyper-parameters considered, method to select the best hyperparameter conﬁguration, and speciﬁcation of all hyper-parameters used to generate results.
We provide all details of hyper-parameter tuning in Section 4.
– The exact number of training and evaluation runs.
Ten thousand independent repetitions are conducted for each experiments.
– A clear deﬁnition of the speciﬁc measure or statistics used to report results.
We provide all details of evaluation metrics in Section 4.
– A description of results with central tendency (e.g. mean) & variation (e.g. error bars).
We report conﬁdence intervals for all experiments in Section 4.
– The average runtime for each result, or estimated energy cost.
We do not report the running time or energy cost.
– A description of the computing infrastructure used.
A clear description is presented in Section 3.
A PPENDIX A C ONFIDENCE E STIMATION P ROCEDURES
In evaluation, for the conﬁdence versus optimization performance scatter plot, we ﬁrst perform
Monte Carlo sampling on the solution posterior, and then bin the solutions to estimate their fraction and corresponding average distance to the optimum.
Detailed procedures can be found in Algorithms 1.
B D ATA D ESCRIPTION
The analytical forms of three test functions are shown in Table 3.
The meta-training data are a broad family of similar functions (Cao et al., 2019a).
For Rastrigin, it is f D ( x )
= (cid:107)
A x − b (cid:107) 22 − c (cid:80) Di =1 10 cos(2 πx i )
+ 10 D where A, b, c are sampled from i.i.d normal distributions; for Ackley, it is f D ( x ) = − 20 exp( − 0 .
2 (cid:112) 0 .
5 (cid:107)
A x − b (cid:107) 22 ) − c (cid:80) Di
=1 exp(cos(2 πx i ) /D )
+ exp(1)
+ 20 ; for
Griewank, it is f D ( x ) =
1 + (cid:80) Di =1 (cid:107)
A x − b (cid:107) 22 / 4000 − c (cid:81) Di =1 cos( x i ) .
For privacy attack, we use 9-dimension features of privacy attack data generated from MutPred (Pejaver et al., 2017), which captures mutational impacts on protein structure, dynamics, and function, and grouped hierarchically into a custom oncology based on their inherent relationships.
The list of features can be found in Table 4.
The meta-training is performed on 1,593 patients and we evaluate in 318 patients (Cao et al., 2019b).
For protein docking, the PDB IDs of meta-training samples are shown in Table 5 and we test in proteins with PDB in Table 2.
[Table]
C M ORE E VALUATION R ESULTS
We provide more evaluation results as follows.
(i)
The ablation of λ values on three test functions is shown in Figure 7.
(ii)
The UQ results for test cases 1AHW 3 and 1JMO 4 in protein docking are plotted in Figure 8.
(iii) Conﬁdence intervals corresponding to results in Figure 3.
[Table] [Table]
[Figure]
[Table]
[Figure] [Figure]
D C OMPUTATION OF UA-L2O O BJECTIVE F UNCTION
We ﬁrst rewrite the evidence as:
[Equation]
[Equation]
where − KL[ q ( φ ; θ ) || p ( z T , φ | z t 0 )] is ELBO and KL[ q ( φ ; θ )
|| p ( φ | z T )] is our objective.
Therefore, minimizing KL[ q ( φ ; θ ) || p ( φ | z T )] is equivalent to maximizing ELBO (Yin & Zhou, 2018). [Figure]
[Figure]
We rewrite our objective as:
[Equation]
E A BLATION S TUDY
We perform ablation studies on sample size in uncertainty quantiﬁcation of UA-L2O in Table 7, mean versus minimum of ˆ x computation and wall-clock time comparison in Table 8, on privacy attack experiment on gene BRCA1/2 in Table 9.
Results show UA-L2O is not sensitive to sample size or mean/minimum calculation of ˆ x , and it is comparably efﬁcient with competitor optimizers.
Moreover, we include a study of the wall-time versus dimensionality ranging from 10 to 2,000, using the optimization of Rastrigin function in Figure 10. We report UA-L2O’s wall-time for both metatraining and meta-testing. The results above show that the meta-training time grows almost linearly with the dimensionality when dimensionality > 200 (which is partly attributed to the fact that the number of samples is now adopted as a value correlated with the dimension). So meta-training even in millions of dimensions for deep-learning model parameters is expensive yet still manageable, especially using multiple GPU cores.
Most importantly, meta-testing was much faster and the time remained ﬂat with regards to the dimensions, that once UA-L2O is trained, its deployment remains fast and scales well to the dimensionality. In fact, L2O usually takes the setting of ofﬂine training (which could be time-consuming) and online deployment (which is fast enough and scales well) Chen et al.
We further compare with baselines as learning rate ensemble and stochastic gradient MCMC of
Adam, shown in Table 11.
We also make comparison with warm-start UA-L2O to show the improvement of UA-L2O training
[Table]
[Table]
[Table] [Table] [Table] [Table]
F U TILIZATION OF O PTIMIZER U NCERTAINTY
We demonstrate the “usefulness” of the proposed optimizer uncertainty via performing the experiment of “failure detection” as a basic downstream task
[Table]
Speciﬁcally, an optimization solution is deﬁned as “failure” if its distance to the global optimum is beyond a threshold 0 .
5 √ d (where d is the dimension and the threshold was chosen for approximately balanced success/failure cases). Instead of ﬁxing the conﬁdence level σ and estimating the conﬁdence interval r (as observed in Figure 4 that better optimization correlated better with tighter conﬁdence intervals), here we let each method estimate the conﬁdence level σ at ﬁxed r < 0 . 5 √ d (“success”) for the optimization solution in each trajectory and use σ to classify the quality of such solutions over 1000 trajectories. The classiﬁcation performance was assessed using the area under the receiver-operator characteristic curve (AUROC) and the area under the precision-recall curve (AUPRC).
The results on privacy attack are shown in Table 12, indicating the best posterior estimation of UA-L2O.
[Table]