Title: LEARNING TO OPTIMIZE VIA DUAL SPACE PRECONDITIONING

Abstract: Preconditioning an minimization algorithm improve its convergence and can lead to a minimizer in one iteration in some extreme cases. There is currently no analytical way for finding a suitable preconditioner. We present a general methodology for learning the preconditioner and show that it can lead to dramatic speed-ups over standard optimization techniques.

1. INTRODUCTION:

Many problems arising in applied mathematics can be formulated as the minimization of a convex function f : R d → (-∞, +∞] min

x∈R d 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 : R d → (-∞, ∞], the domain of f is the set {x ∈ R d , 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 R d , and is differentiable and strictly convex on D(f ),

• the gradient ∇f is coercive on D(f ): For every sequence

(x i ) i∈N ∈ D(f ) converges to a point on the boundary of the domain of f , ∇f (x i ) → ∞.

Given a Legendre function f , one can define its Fenchel conjugate f (y) = sup x∈R d 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 : R d → (-∞, +∞]. To minimize f , the DPGD algorithm is written

EQUATION

where γ > 0 and p : R d → (-∞, +∞] 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

µ∇ 2 f (3) ∇ 2 p (2) L∇ 2 f ,

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 = arg min 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 R d ) 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 x i ∈ D(f ) according to some distribution and create the dataset {(∇f (x i ), x i ), i ∈ {1, . . . , n}} where the ∇f (x i ) represent the features and x i the labels. The feature space is D(f ) and the label space is D(f ) (Proposition 1). Since ∇f maps ∇f (x i ) to x i (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 µ (θ)

EQUATION

)

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 (x i ) and the labels x i . 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 (x i ) and a diffeomorphism H -1 : D(f ) → [-0.5, 0.5] d to the labels x i . 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.

Figure 1 : General model setup

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 .

Figure 2 : The curve of ∇f is typical of objectives f that we want to minimize

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 cannot 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

x i ∼ µ, then G • ∇f (x i ) 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(J G•∇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.

3.4. ESTIMATION OF THE LOSS:

At this step, µ is chosen and L µ has to be minimized by SGD. As usual, L µ cannot be computed in closed form. The classical idea is to approximate L µ by an empirical mean involving samples from the distribution µ. Sampling from µ turns out to be difficult in our case. Indeed, the use of rejection-based MCMC methods lead to a high rejection rate. Since we can evaluate the density of µ, we suggest to use importance sampling instead. Proposition 4. If x i ∼ π where the density of π is positive, then

E 1 n n i=1 (Model ∇f * (∇f (x i ), θ), x i ) µ(x i ) π(x i ) = L µ (θ),

where µ(x) (resp. π(x)) denotes the density of µ (resp. π) with respect to Lebesgue measure.

We call the map θ → 1 n n i=1 (Model ∇f * (∇f (x i ), θ), x i ) µ(xi) π(xi) the empirical risk. Proposition 4 states that if x i ∼ π, then the empirical risk is an unbiased estimator of the theoretical risk. Since this estimator will be minimized using SGD, we only need to know it up to a constant factor. Therefore we can get rid of the 1/n factor and of constant factors defining the densities of µ and π. The distribution π is called the proposal distribution.

3.5. RELATED WORKS:

Learning an optimization algorithm has already been proposed in the literature. In Li & Malik (2016) ; Andrychowicz et al. (2016) the authors propose to learn a better optimization algorithm by observing its execution over a feature space of objective functions. This outperforms existing hand-engineered algorithms in terms of convergence speed, but require to know the performance of some algorithms over some objectives. Using optimization literature knowledge, our approach allows to reduce the problem to a supervised learning setting in R d , that we can solve with a neural net. In its philosophy, our approach can also be related to deep learning methods for inverse problems used in signal processing, see e.g Lucas et al. (2018) .

4. NUMERICAL EXPERIMENT:

In this section, we run DPGD algorithm after learning ∇f . We first consider a one dimensional toy problem with a ground truth involving power functions as in Maddison et al. (2019) . Then ill-conditioned logistic losses are considered. These functions cannot be minimized (efficiently) by gradient descent. We call LDPGD our approach (for Learned DPGD).

4.1. POWER FUNCTIONS AND LOGISTIC LOSS:

Two class of objective functions are considered. Power functions are defined by

f : x → 1 a Ax -c a , where    a ∈ (2, +∞) c ∈ R d A is a matrix.

(5)

These functions are nonsmooth and non-strongly convex, therefore gradient descent is not guaranteed to minimize them. Power function minimization is the main example of Maddison et al. (2019) .

We also consider the logistic loss

EQUATION

)

where r > 0 and for every i, b i ∈ {-1, 1} n and a i ∈ R d . These functions are strongly-convex and smooth. It is easy to control the condition number of f :

Lemma 1. Consider the logistic loss f and the matrix A whose lines are the a i , and denote λ max the largest eigenvalue of the real symmetric matrix A T A. Denote L the smoothness parameter of f1 and κ the condition number of f . Then, L ≤ L = λ max 4n + r, and κ ≤ λ max 4nr + 1.

4.2. EXPERIMENTAL SETUP:

We use a three layers neural network, with leaky ReLU activations and respectively 256 and 128 neurons in the two hidden layers. The square loss is used (y, x) = 1 2 y -x 2 . The proposal distribution π for the importance sampling is taken gaussian. In the case of power functions, the features are rescaled using the map log-rescaling :

x → log(1 + x) if x > 0 -log(1 -x) else.

Justification for this choice is provided in the appendix (section D). Then the rescaled features and the labels are normalized to have zero mean and unit variance. The general architecture is summarized in figure 3 . We see that ∇f is accurately learned the machine learning algorithm. This approximation of ∇f is then used as a preconditioner of DPGD. The value of the objective function and the value of the iterates while running LDPGD is plotted in the next figure. The stepsize is set to 1 during 10 steps and LDPGD is initialized at x 0 = 100. The iterates of gradient descent always diverge (no matter the stepsize choice). On the contrary, LDPGD iterates quickly converge to 357.1 (with ∇f (357.1) = 3.8 × 10 -14 ) which is close to the actual minimizer (356.4 in this example).

4.4. RESULTS FOR LOGISTIC REGRESSION:

Consider f a logistic loss (6) with d = 50, n = 1000, and condition number κ. We train the model during 100 epochs. and represent the evolution of the objective function and the iterates while running LDPGD (after learning ∇f ) and gradient descent (GD).

The algorithms are initialized at (50, 50, 50) T and we use a stepsize for which GD converges linearly2 . This is done for several logistic losses with worsening condition numbers. To vary κ, we vary the parameter r that controls the strong convexity of the objective f (see Lemma 1). Intuitively, if r is high then f is well-conditioned and if r is low f is ill-conditioned. The left figure represents the iterates (each curve corresponds to the evolution of one coordinate of the iterates) and the right figure represents the objective fonction values while running both algorithms. Our careful preconditioning allows LDPGD to outperform GD especially when the objective is ill-conditioned, as predicted by the theory, see section 2.

A PROOF OF PROPOSITION 3:

Recall that G is a diffeomorphism and that f is Legendre. Let U the uniform distribution over [-0.5, 0.5] d and denote q its density with respect to Lebesgue measure. Under the assumption of Proposition 3, ∇f : D(f ) → D(f ) is a diffeomorphism. Therefore, the map M = G • ∇f : D(f ) → [-0.5, 0.5] d is also a diffeomorphism and its inverse map is

M -1 = ∇f • G -1 : [-0.5, 0.5] d → D(f ). Our problem is to find a distribution µ supported by D(f ) such that M (X) = G • ∇f (X) ∼ U if X ∼ µ. Such µ satisfies for every nonnegative measurable function φ, φ(x)dµ(x) = φ(M -1 • M (x))dµ(x) = E X∼µ (φ(M -1 • M (X))) = E U ∼U (φ(M -1 (U ))) = φ(M -1 (u))q(u)du = D(f ) φ(M -1 (M (x)))q(M (x))| det(J M (x))|dx = D(f ) φ(x)q(M (x))| det(J M (x))|dx,

where J M denotes the Jacobian matrix of M and the penultimate equality follows from the change of variable formula (which is allowed because M is a diffeomorphism). Therefore, the only possible solution to our problem is the distribution µ with density proportional to | det(J M (x))| if M (x) ∈ [-0.5, 0.5] d (which is always satisfied by definition of M ) and 0 else. One can check that this density solves the problem.

B PROOF OF PROPOSITION 4:

We apply the importance sampling principle to the measurable nonnegative function (the parameter θ is fixed) φ(x) = (Model ∇f * (∇f (x), θ), x).

We have

E Y ∼µ (φ(Y )) = E X∼π φ(X) µ(X) π(X) = E 1 n n i=1 φ(x i ) µ(x i ) π(x i ) ,

where x i ∼ π.

C PROOF OF LEMMA 1:

Denote L the smoothness constant of f (i.e the Lipschitz constant of its gradient) and λ its strong convexity parameter. For every

x ∈ R d , ∇f (x) = 1 n n i=1 -b i a i 1 1 + exp (b i a i , x )) + rx ∇ 2 f (x) = 1 n n i=1 b 2 i a i a T i exp (b i a i , x )) (1 + exp (b i a i , x ))) 2 + rI d ,

where I d is the identity matrix. Note that for all a = 1, a (1+a) 2 ≤ 1 4 , and that b 2 i ≤ 1. It follows that

∇ 2 f A T A 4n + rI d .

Therefore, L ≤ λmax 4n +r. Since κ = L/λ, the result follows from the fact that r ≤ λ (see Equation 6).

D LOG-RESCALING:

For power functions, ∇f (x) can range in -10 100 , 10 100 d if x is sampled from a gaussian. Assume d = 1 that we use Doubles precision in the implementation. This means that all the values we manipulate are stored with a precision of approximately 15 decimal places.

If we naively standardize (∇f (x i )) i∈ [1,n] to have zero mean and unit variance, we rescale numbers in the range [-10 100 , 10 100 ] d , to the range [-0.5, 0.5] d . This means that a precision of 10 -15 in the standardized scale represents a precision of 10 85 in the original scale.

In practice, this loss of precision means that inputs that are different before standardization, might be indistinguishable after the standardization. As a result, the neural network might receive many different outputs for a single input value.

As it tries to minimize the empirical risk, the model will tend to associate to input values an average of the corresponding outputs. As a consequence, the regressed model will be a piecewise constant function, and will be constant on subsets of diameter 10 85 . This is especially problematic for inputs close to 0. Let us denote N 0 the neighborhood of 0 on which the regressed model is constant equal to Model ∇f * (0, θ). Any x ∈ D(f ) for which ∇f (x) ∈ N 0 is a fixed point of the DPGD alogorithm:

x + = x - 1 L p [Model ∇f * (∇f (x), θ) -Model ∇f * (0, θ)] = x - 1 L p [0] = x

The consequence is very undesirable: the solution offered by the DPGD alogorithm might be x even though ∇f (x) = O(10 85 ) is far from zero. 

:

i.e, the Lipschitz constant of ∇f

GD converges linearly if the stepsize is smaller than 1/L, where L is the Lipschitz constant of ∇f . Therefore, it is enough to take a stepsize smaller than 1/L where L is defined in Lemma 1

Figures and Tables:

Figure 3: Model with log-rescaling

Figure 4: Learning ∇f * . Left: Learning curves. Right: predictions.

Figure 5: Minimizing a one dimensional power function. Left: Objective function values. Right: iterates values.

Figure 6: Logistic regression with r = 0.5

Figure 9: Regressed model on the imprecise dataset

Figure 10: Log-rescaling.

Peer review (iclr2020):
This paper attempts to learn a preconditioner for optimization, specifically for the Dual space preconditioned descent (DPGD). - The techniques used to learn the preconditioner are heuristic, not scalable and without justification or ablation studies. - It does not compare against *standard* optimization techniques that construct data-driven preconditioners such as Adam or Adagrad or even to more Newton, natural gradient methods that use the Hessian or the Fisher information matrix as preconditioners. It shows ad-hoc synthetic experiments in dimensions 1 and 50. This is clearly not enough. Detailed review below: - Section 2: Please explain why Legendre functions are useful in ML. For assumption 1, 2; it needs to be explained why these hold for a given f*. What constraints do you need on f? What functions satisfy these? Please explain this explicitly. - Section 3: What is the number of points x_i needed in high dimensions to learn? Is it even possible to scale up this method to high dimensions? - Constructing mu requires computing the determinant of the Jacobian. What is the computational complexity? Moreover, it seems that we need access to the 
abla f(x) for all x in D(f)? - Please state all the assumptions in the beginning rather than introducing one at a time in the propositions. - Remark 1: It is unclear that the cost of an inverse Hessian matrix is more than the procedure proposed in this paper. - Section 3.5: Please explain what is the advantage of this learned optimizer compared to other methods? Note that there is literature on non-smooth optimization and methods like sub-gradient descent can be used in this case. - What is the justification for the selection of the loss function and log-rescaling? - The result of Lemma 1 is standard. Please acknowledge this. - Section 4: *The step-size is set to 1*. It seems that the optimizer has been overfit and engineered to work on this specific problem. Either these decisions need to be justified, there needs to be an ablation study or there needs to be a larger set of experiments.