Learning to Optimize via Dual space Preconditioning

1 INTRODUCTION . Many problems arising in applied mathematics can be formulated as the minimization of a convex function f : Rd → ( −∞ , +∞ ] min x∈Rd f ( x ) . The resolution of an optimization problem is usually tackled using an optimization algorithm that produces a sequence of iterates converging to some minimizer of f ( Nesterov ( 2018 ) ) . The gradient descent algorithm is a standard optimization algorithm that converges linearly ( i.e exponentially fast ) if f is regular enough . Preconditioned methods ( Nemirovsky & Yudin ( 1983 ) ; Lu et al . ( 2016 ) ; Maddison et al . ( 2019 ) ) are powerful optimization algorithms that converges linearly under weaker assumptions that gradient descent . The performance of these methods relies heavily on the preconditioning – the task of choosing the hyperparameter called the preconditioner . In the case of the Dual space Preconditioned Gradient Descent ( DPGD ) of Maddison et al . ( 2019 ) , an optimal preconditioning can lead to convergence in one iteration . The preconditioner is a function that has to be selected properly w.r.t . f to obtain the desired linear convergence . Although Maddison et al . ( 2019 ) gives some hints to precondition DPGD , there is currently no analytical way for finding a suitable preconditioner . In this paper , we propose to learn the preconditioner of DPGD using a neural network . We make the following contribution : • We propose a supervised learning setting to learn the preconditioner of an optimization algorithm ( DPGD , Maddison et al . ( 2019 ) ) • We present a general methodology that allows to effectively learn the preconditioner while avoiding issues related to this task • We implement this methodology in dimension one and 50 and show that this can lead to dramatic speed-ups . The remainder is organized as follows . The next section provides background knowledge on the DPGD algorithm . Then , in section 3 we present our supervised learning setting . Finally , we apply this methodology in section 4 . Additional developments are provided in the appendix as well as postponed proofs . 2 BACKGROUND ON DUAL PRECONDITIONING . The main message of this section is that a good preconditioner is a preconditioner that is `` close '' to ∇f ? , where f ? is the Fenchel conjugate of f . 2.1 LEGENDRE FUNCTIONS . Before introducing DPGD algorithm , we need to define the class of functions f that we aim to minimize . These functions are called Legendre convex functions and are studied in ( Rockafellar , 1970 , Section 26 ) . We recall their definition and some of their properties . Given a function f : Rd → ( −∞ , ∞ ] , the domain of f is the set { x ∈ Rd , f ( x ) < ∞ } . We denote D ( f ) the interior of the domain of f . Definition 1 . The function f is Legendre if D ( f ) is not empty , and if • the function f is lower semicontinuous on Rd , and is differentiable and strictly convex on D ( f ) , • the gradient∇f is coercive on D ( f ) : For every sequence ( xi ) i∈N ∈ D ( f ) converges to a point on the boundary of the domain of f , ‖∇f ( xi ) ‖ → ∞ . Given a Legendre function f , one can define its Fenchel conjugate f ? ( y ) = supx∈Rd〈x , y〉 − f ( x ) . Proposition 1 ( Rockafellar ( 1970 ) ) . The function f is Legendre if and only if its Fenchel conjugate f ? is Legendre , in which case ∇f is one-to-one between D ( f ) and D ( f ? ) . Moreover , ∇f is continuous in both directions , for every x ∈ D ( f ) , ∇f ? ( ∇f ( x ) ) = x and for every y ∈ D ( f ? ) , ∇f ( ∇f ? ( y ) ) = y . Proposition 2 ( Rockafellar ( 1970 ) ) . If f is Legendre , then it either has no minimizer , or one unique global minimizer x ? ∈ D ( f ) . If f admits a minimizer , then ∇f ( x ? ) = 0 , 0 ∈ D ( f ? ) and x ? = ∇f ? ( 0 ) . 2.2 DPGD ALGORITHM ( MADDISON ET AL . ( 2019 ) ) . Consider a convex Legendre function f : Rd → ( −∞ , +∞ ] . To minimize f , the DPGD algorithm is written x+ = x− γ [ ∇p ( ∇f ( x ) ) −∇p ( 0 ) ] , ( 1 ) where γ > 0 and p : Rd → ( −∞ , +∞ ] are parameters . DPGD is well defined over the set of Legendre convex functions f , see Maddison et al . ( 2019 ) . The map ∇p is called a non-linear preconditioner of ∇f . The preconditioning is the design of∇p . 2.3 CONVERGENCE THEORY OF DPGD . In the sequel , we consider a Legendre function f admitting a minimizer x ? . The preconditioning can dramatically affect the convergence of DPGD . Indeed , if ∇p = ∇f ? , then DPGD algorithm converges in one iteration . To see this , note that in this case x+ = x− [ x−∇p ( 0 ) ] = ∇p ( 0 ) = ∇f ? ( 0 ) , if γ = 1 and recall that∇f ? ( 0 ) is the minimizer of f ( Proposition 2 ) . Of course , this preconditioning is unrealistic since computing∇f ? ( 0 ) is as hard as minimizing f . It is more realistic to aim to find a preconditioner p that satisfies the following conditions . Assumption 1 ( Dual relative smoothness ) . There exists L ≥ 0 such that 〈∇p ( x ) −∇p ( y ) , x− y〉 ≤ L〈∇f ? ( x ) −∇f ? ( y ) , x− y〉 , ∀x , y ∈ D ( p ) . ( 2 ) Assumption 2 ( Dual relative strong convexity ) . There exists µ > 0 such that 〈∇p ( x ) −∇p ( y ) , x− y〉 ≥ µ〈∇f∗ ( x ) −∇f∗ ( y ) , x− y〉 , ∀x , y ∈ D ( p ) . ( 3 ) If p and f ? are twice differentiable , these assumptions can be written as µ∇2f ? ( 3 ) ∇2p ( 2 ) L∇2f ? , for the ordering of nonnegative matrices . These assumptions mean that∇p must be chosen close to∇f ? in some sense . Moreover , they are sufficient for the linear convergence of DPGD . Theorem 1 ( Maddison et al . ( 2019 ) , Informal ) . If assumption 1 holds and if p is Legendre , then the sequence of iterates of DPGD converges to x ? = argmin f . If assumption 2 holds , then the convergence is linear . In comparison , gradient descent ( and many other first order methods , Nesterov ( 2018 ) ) converges for smooth convex functions f ( i.e . functions f such that ∇f is Lipschitz continuous ) . Moreover , the convergence is linear if f is strongly convex . Many functions satisfy assumption 1 ( and assumption 2 ) for some preconditioner p , without being smooth ( or strongly convex ) , see Maddison et al . ( 2019 ) . If f is smooth and strongly convex , the convergence rate of gradient descent depends on the global condition number κ of f ( defined as the Lipschitz constant of ∇f divided by the strong convexity parameter ) . If κ is high , then the linear convergence of gradient descent is slow and f is said ill-conditioned . The function f is said well-conditioned else . When f is ill-conditioned , it might be well-conditioned locally ( i.e. , on small subsets of Rd ) but the rate of convergence of gradient descent only takes into account the global , high condition number . This is not the case for preconditioned algorithms ( Li & Malik ( 2016 ) ; Maddison et al . ( 2019 ) ) . 3 SUPERVISED LEARNING SETTING . In this section , we describe our methodology to learn∇f ? . 3.1 FORMULATION . Our idea is to sample points xi ∈ D ( f ) according to some distribution and create the dataset { ( ∇f ( xi ) , xi ) , i ∈ { 1 , . . . , n } } where the ∇f ( xi ) represent the features and xi the labels . The feature space is D ( f ? ) and the label space is D ( f ) ( Proposition 1 ) . Since∇f ? maps∇f ( xi ) to xi ( Proposition 1 ) , ∇f ? can be seen as the solution of this supervised learning problem . We choose a neural net to modelize∇f ? . To solve our supervised learning problem , we design a machine learning algorithm . The goal of the training of the machine learning algorithm is to minimize the theoretical risk Lµ ( θ ) Lµ ( θ ) = EX∼µ ( ` ( Model∇f∗ ( ∇f ( X ) , θ ) , X ) ) , ( 4 ) where µ is the distribution of the labels X , θ represents the parameters of the neural network , Model∇f ? ( ∇f ( X ) , θ ) is the output of the neural network with input∇f ( X ) and parameter θ , and ` ( Model∇f∗ ( ∇f ( X ) , θ ) , X ) is the loss associated to one sample ( ∇f ( X ) , X ) . To fully specify the machine learning algorithm , we need to fully define Lµ and the way we approximate it . We proceed by • Choosing an architecture to specify the function Model∇f∗ ( this is problem specific , see sections 3.2 , 4.2 ) • Giving a value for µ ( section 3.3 ) • Providing a way to approximate the expectation in ( 4 ) by a finite sum ( section 3.4 ) • Choosing a loss ` ( this is problem specific , see section 4.2 ) . Once this is done , the algorithm is trained using Stochastic Gradient Descent ( SGD ) ( Bottou et al . ( 2018 ) ) . 3.2 GENERAL SETUP . In this section we give a general description of the map Model∇f∗ . According to recommendations in Bengio ( 2012 ) , we first standardize the features∇f ( xi ) and the labels xi . In other words , we apply a diffeomorphism ( i.e , a one-to-one map which is continuously differentiable in both directions ) G : D ( f ? ) → [ −0.5 , 0.5 ] d to the features∇f ( xi ) and a diffeomorphism H−1 : D ( f ) → [ −0.5 , 0.5 ] d to the labels xi . Then a neural network is used to learn a predictor from the standardized features to the standardized labels , see figure 1 . One can obtain a label from a standardized label by applying the inverse map H of H−1 , which can be easily computed ( in practice H−1 is essentially an addition and a multiplication by a positive number ) , see section 4 . 3.3 CHOICE OF µ We found that the choice of the distribution µ that parametrized Lµ is critical for the performance of the machine learning algorithm . The reason is the following . Recall that we want to learn ∇f ? to precondition DPGD algorithm whose goal is to minimize nonsmooth and non-strongly convex ( or ill-conditioned ) objectives f , see section 2 . Let us imagine that we chose an uniform distribution for µ . Then , the situation in dimension one is represented by the figure 2 . The image distribution ν of µ by ∇f over-represents some areas of the feature space while other areas are under-represented . When the distribution of the features is as degenerated as ν , it is obvious that machine learning algorithms can not generalize well ( see Bengio ( 2012 ) ; Mesnil et al . ( 2011 ) ) . Therefore , we have to choose µ carefully . The classical recommendation is to ensure that the input of the neural network ( i.e the standardized features ) are uniformly distributed ( Mesnil et al . ( 2011 ) ) . In other words , we have to find µ such that , if xi ∼ µ , then G ◦ ∇f ( xi ) is uniformly distributed over [ −0.5 , 0.5 ] d. This can be done by using the change of variable formula . Proposition 3 . Assume that ∇f and ∇f ? are differentiable . If µ is the distribution with density proportional to |det ( JG◦∇f ( x ) ) | with respect to Lebesgue measure ( where J denotes the Jacobian matrix ) , then the distribution of G ◦ ∇f ( X ) where X ∼ µ is uniform over [ −0.5 , 0.5 ] d. In the sequel , µ is the distribution whose density is given by Proposition 3 . Inspecting figure 2 , this seems reasonable since we expect X ∼ µ to have a high density when∇f ’ s slope is large and a low density when∇f ’ s slope is small . Remark 1 . It is worth noting that the evaluation of µ doesn ’ t require any inversion as in the inverse transform sampling technique . Besides , second order optimization algorithms usually require ( an approximation of ) the inverse of the Hessian matrix , which we don ’ t need .

Summary: This paper proposes an optimization method with the preconditioning in the framework of supervised learning. The ideal preconditioning is given by the Fenchel conjugate of the optimization function. This paper uses a supervised scenario to find the ideal preconditioning. The authors point out the importance of the sampling distribution then and propose a sampling scheme using the uniform distribution on the space of gradient descents. The samples are used to train neural networks that imitate the mapping of the Fenchel conjugate. The training network is incorporated into the Dual space Preconditioned Gradient Descent (DPGD). Some numerical experiments show the effectiveness of the proposed method comparing to the standard gradient descent method.