F IRST -O RDER P RECONDITIONING VIA H YPERGRADIENT D ESCENT
Anonymous authors
Paper under double-blind review
A BSTRACT
Standard gradient descent methods are susceptible to a range of issues that can impede training, such as high correlations and different scaling in parameter space.
These difﬁculties can be addressed by second-order approaches that apply a preconditioning matrix to the gradient to improve convergence.
Unfortunately, such algorithms typically struggle to scale to high-dimensional problems, in part because the calculation of speciﬁc preconditioners such as the inverse Hessian or
Fisher information matrix is highly expensive.
We introduce ﬁrst-order preconditioning (FOP), a fast, scalable approach that generalizes previous work on hypergradient descent (Almeida et al., 1998; Maclaurin et al., 2015; Baydin et al., 2017) to learn a preconditioning matrix that only makes use of ﬁrst-order information.
Experiments show that FOP is able to improve the performance of standard deep learning optimizers on several visual classiﬁcation tasks with minimal computational overhead.
We also investigate the properties of the learned preconditioning matrices and perform a preliminary theoretical analysis of the algorithm.
I NTRODUCTION
High-dimensional nonlinear optimization problems often present a number of difﬁculties, such as strongly-correlated parameters and variable scaling along different directions in parameter space (Martens, 2016).
Despite this, deep neural networks and other large-scale machine learning models applied to such problems typically rely on simple variations of gradient descent to train, which is known to be highly sensitive to these difﬁculties.
While this approach often works well in practice, addressing the underlying issues directly could provide stronger theoretical guarantees, accelerate training, and improve generalization.
Adaptive learning rate methods such as Adam (Kingma and Ba, 2014), Adagrad (Duchi et al., 2011), and RMSProp (Tieleman and Hinton, 2012) provide some degree of higher-order approximation to re-scale updates based on per-parameter behavior.
Newton-based methods make use of the curvature of the loss surface to both re-scale and rotate the gradient in order to improve convergence.
Natural gradient methods (Amari, 1998) do the same in order to enforce smoothness in the evolution of the model’s conditional distribution.
In each of the latter two cases, the focus is on computing a speciﬁc linear transformation of the gradient that improves the conditioning of the problem.
This transformation is typically known as a preconditioning , or curvature , matrix.
In the case of quasiNewton methods, the preconditioning matrix takes the form of the inverse Hessian, while for natural gradient methods it’s the inverse Fisher information matrix.
Computing these transformations is typically intractable for high-dimensional problems, and while a number of approximate methods exist for both (e.g., (Byrd et al., 1996; Martens and Grosse, 2015; Grosse and Martens, 2016)), they are often still too expensive for the performance gain they provide.
These approaches also suffer from rigid inductive biases regarding the nature of the problems to which they are applied in that they seek to compute or approximate speciﬁc transformations.
However, in large, non-convex problems, the optimal gradient transformation may be less obvious, or may even change over the course of training.
In this paper, we address these issues through a method we term ﬁrst-order preconditioning (FOP).
Unlike previous approaches, FOP doesn’t attempt to compute a speciﬁc preconditioner, such as the inverse Hessian, but rather uses ﬁrst-order hypergradient descent (Maclaurin et al., 2015) to learn an adaptable transformation online directly from the task objective function.
Our method adds minimal
1: Require: model parameters θ , objective function J , FOP matrix M , learning rate (cid:15) , hypergradient learning rate ρ 2: for t = 1,2,... do 3:
Draw data x ( t ) , y ( t ) ∼ D 4:
Perform forward pass: ˆ y ( t ) =
f θ ( t ) ( x ( t ) ) 5:
Compute loss J ( y ( t ) , ˆ y ( t ) ) 6:
Update inference parameters θ : θ ( t +1) ← θ ( t ) − (cid:15)M ( t ) M ( t ) (cid:62) ∇ θ ( t ) J 7:
Update preconditioning matrices:
M
[Equation]
t
[Equation]
← [Equation] t [Equation]
+
[Equation]
(cid:16)
[Equation]
( t )
[Equation]
[Equation] θ ( t [Equation] 1) J ] (cid:62) + ∇ θ ( t
[Equation]
[Equation] [Equation] θ ( t ) J ] (cid:62) (cid:17) M ( t ) 8: Cache ∇ θ ( t ) J computational and memory cost to standard deep network training and results in improved convergence speed and generalization compared to standard approaches.
FOP can be ﬂexibly applied to any gradient-based optimization problem, and we show that when used in conjunction with standard optimizers, it improves their performance.
To our knowledge, this is also the ﬁrst successful application of any hypergradient method to the ImageNet dataset (Deng et al., 2009).
F IRST O RDER P RECONDITIONING 2. T HE B ASIC A PPROACH
Consider a parameter vector θ and a loss function J .
A traditional gradient update with a preconditioning matrix P can be written as
[Equation]
Our goal is to learn P .
However, while we place no other constraints on our preconditioner, in order to ensure that it is positive semi-deﬁnite, and therefore does not reverse the direction of the gradient, we reparameterize P as the product of a matrix M with itself, changing Equation 1 to:
[Equation]
Under reasonable assumptions, gradient descent is guaranteed to converge with the use of even a random symmetric, positive-deﬁnite preconditioner, as we show in Supplementary Section A.3.
Because θ ( t ) is a function of M ( t − 1) , we can then backpropagate from the loss at iteration t to the previous iteration’s preconditioner via a simple application of the chain rule:
[Equation]
By applying the chain and product rules to Equation 2, the gradient with respect to the preconditioner is then simply
[Equation]
We provide a more detailed derivation in Supplementary Section A.1.
Note that ideally we would compute ∇ M ( t ) J ( t +1) to update M ( t ) , but as we don’t have access to J ( t +1) yet, we follow the example of Almeida et al.
and assume that J is smooth enough that this does not have an adverse effect on training.
The basic approach is summarized in Algorithm 1.
We use supervised learning as an example, but the same method applies to any gradient-based optimization.
The preconditioned gradient can then be passed to any standard optimizer to produce an update for M .
For example, we describe the procedure for using FOP with momentum (Polyak, 1964) in Section 2.4.
For multi-layer networks, in order to make the simplest modiﬁcation to normal backpropagation, we learn a separate M for each layer, not a global curvature matrix.
To get an intuition for the behavior of FOP compared to standard algorithms, we observed its trajectories on a set of low-dimensional optimization problems (Figure 1).
Interestingly, while FOP converged in fewer iterations than SGD and Adam (Kingma and Ba, 2014), it took more jagged paths along the objective function surface, suggesting that while it takes more aggressive steps, it is perhaps also able to change direction more rapidly.
2. L OW -R ANK FOP
If θ is an m × n matrix, and we only apply the preconditioning matrix over input dimensions (and share it across output dimensions), then a full-rank M would necessarily be m × m .
When m is large, preconditioning the gradient becomes expensive.
Instead, we can apply a rank- k M , with
M ∈ R m × k and k < m .
To ensure stable performance at the beginning of training, we initialize the preconditioning matrix as close as possible to the identity matrix, even for a low rank M , so that the algorithm begins as straightforward gradient descent and learns to depart from vanilla SGD over time.
Thus, we set the preconditioner P to be
[Equation]
where I m is the m × m identity matrix and M ij ∼ N (0 , σ 2 ) , where σ 2 is small so that P starts out close to the identity matrix.
This effectively decomposes the update into an isotropic scaling of the gradient (multiplication by I m scaled by the learning rate) and a rotation (multiplication by MM (cid:62) ).
We tested the effect of rank on a simple fully-connected network trained on the MNIST dataset (LeCun and Cortes, 2010).
The results, shown in Figure 2, indicate that FOP is able to accelerate training compared to standard SGD (with momentum 0.9) with all values of k and improve ﬁnal test accuracy even with fairly low values of k .
2. S PATIAL P RECONDITIONING FOR C ONVOLUTIONAL N ETWORKS
In order to further reduce the computational cost of FOP in convolutional networks (CNNs), we implemented layer-wise spatial preconditioners, sharing the matrices across both input and output channels (results shown in Section 4).
More concretely, if a convolutional layer has spatial kernels with shape k × k , we can learn a preconditioner that is k 2 × k 2 .
To implement this, when θ is a 4-tensor of kernels of shape k × k ×
I × O , where I and O are the input and output channels,
[Figure]
The effect of on
[Figure]
respectively, we can reshape it to a matrix of size k 2 × IO , left-multiply it by the learned curvature matrix, and then reshape it back to its original dimensions.
When k is small, as is typically the case in deep CNNs, this preconditioner is small as well, resulting in both computational and memory efﬁciency.
2. FOP FOR M OMENTUM
FOP can be implemented alongside any standard optimizer, such as gradient descent with momentum (Polyak, 1964).
Given a parameter vector θ and a loss function J , a basic momentum update with FOP matrix M is typically expressed as two steps:
[Equation]
where v is the velocity term and α is the momentum parameter.
Combining Equations 6 and 7 allows us to write the full update as
[Equation]
If, as in Maclaurin et al. (2015), we were meta-learning M or only updating M after a certain number of iterations, we would then have to backpropagate through v to calculate the gradient for
M .
As we are updating M online, however, we only need to calculate ∇ M ( t ) J .
Therefore, the update is the same as for standard gradient descent.
The experiments in Section 4 were performed using this modiﬁcation of momentum.
R ELATED W ORK
Almeida et al.
(1998) introduced the idea of using gradients from the objective function to learn optimization parameters such as the learning rate or a curvature matrix.
However, their preconditioning matrix was strictly diagonal, amounting to an approximate Newton algorithm
(Martens, 2016), and they only tested their framework on simple optimization problems with either gradient descent or
SGD.
This is related to common deep learning optimizers such as Adagrad (Duchi et al., 2011) and Adam (Kingma and Ba, 2014), which also learn adaptive per-parameter learning rates, but unlike FOP, do not induce a rotation on the gradient.
As we note in Section 2.4, an advantage of FOP is [Figure] that it may also be used in conjunction with any such optimizer.
More recently, Maclaurin et al.
(2015) applied the backpropagation-across-iterations approach in a neural network context, terming the process hypergradient descent .
Their method backpropagates through multiple iterations of the training process of a relatively shallow network to meta-learn a learning rate.
However, this method can incur signiﬁcant memory and computational cost for large models and long training times.
Baydin et al. (2017) instead proposed an online framework directly inherited from Almeida et al. (1998) that used hypergradient-based optimization in which the learning rate is updated after each iteration.
Our method extends this idea to not only learn existing optimizer parameters (e.g., learning rate, momentum), but to introduce novel ones in the form of a non-diagonal, layer-speciﬁc preconditioning matrix for the gradient.
It’s also important to discuss the relationship of FOP to other, non-hypergradient preconditioning methods for deep networks.
These mostly can be sorted into one of two categories, quasi-Newton algorithms and natural gradient approaches.
Quasi-Newton methods seek to learn an approximation of the inverse Hessian.
L-BFGS, for example, does this through tracking the differences between gradients across iterations (Byrd et al., 1996).
This is signiﬁcantly different from FOP, although the outer product of gradients used in the update for FOP can also be an approximation of the Hessian or Gauss-Newton matrices (Bottou et al., 2016).
Natural gradient methods, such as K-FAC (Martens and Grosse, 2015) and KFC (Grosse and Martens, 2016), which approximate the inverse Fisher information matrix, bear a much stronger resemblance to FOP.
However, there are notable differences.
First, unlike FOP, these methods perform extra computation to ensure the invertibility of their curvature matrices.
Second, the learning process for the preconditioner in these methods is completely different, as they do not backpropagate across iterations.
E XPERIMENTS
We measured the performance of FOP on several visual classiﬁcation tasks.
In order to measure the importance of the rotation induced by the preconditioning matrices in addition to the scaling, we also implemented hypgradient descent (HD) methods to learn a scalar layer-wise learning rate (S-HD) and a per-parameter (PP-HD) learning rate.
The former is the same method implemented by Baydin et al.
(2017), and the latter is equivalent to a strictly diagonal curvature matrix.
We also implement a method we call normalized FOP (FOP-norm), in which we rescale the preconditioning
[Figure]
matrix to avoid any effect on the learning rate and rely solely on standard learning rate settings.
Further details on this method can be found in Supplementary Section A.2.
All experiments were implemented using the TensorFlow library (Abadi et al., 2015) 1 .
4. CIFAR-
For CIFAR-10 (Krizhevsky, 2009), an image dataset consisting of 50,000 training and 10,000 test 32 × 32 RGB images divided into 10 object classes, we implemented a 9-layer convolutional architecture inspired by Springenberg et al.
We trained each model for 150 epochs with a batch size of 128 and initial learning rate 0.05, decaying the learning rate by a factor of 10 after 80 epochs.
For S-HD, PP-HD, and FOP, we use Adam as the hypergradient optimizer with a learning rate of 1e-4.
The results are plotted in Figure 3.
FOP produces a signiﬁcant speed-up in training and improves ﬁnal test accuracy compared to baseline methods, including FOP-norm, indicating that both the rotation and the scaling learned by FOP is useful for learning.
4. I MAGE N ET
The ImageNet dataset consists of 1,281,167 training and 50,000 validation 299 × 299 RGB images divided into 1,000 categories (Deng et al., 2009).
Here, we trained a ResNet-18 (He et al., 2015) model for 60 epochs with a batch size of 256 and an initial learning rate of 0.1, decaying by a factor of 10 at the 25th and 50th epochs.
A summary of our results is displayed in Figure 3.
We can see that the improved convergence speed and test performance observed on CIFAR-10 is maintained on this deeper model and more difﬁcult dataset.
[Figure]
4. H YPERPARAMETER R OBUSTNESS
In addition to measuring peak performance we also tested the effect FOP had on the robustness of standard optimizers to hyperparameter selection.
Baydin et al.
(2017) demonstrated the ability of scalar hypergradients to mitigate the effect of the initial learning rate on performance.
We therefore tested whether this beneﬁt was preserved by FOP, as well as whether it extended to other hyperparameter choices, such as the momentum coefﬁcient.
Our results, summarized in Figure 4, support this idea, showing that FOP can improve performance on a wide array of optimizer settings.
For momentum (Fig. 4a), we see that the performance gap is greater for higher learning rates and momentum values, and in several instances FOP is able to train successfully where pure SGD with momentum fails.
For Adam, the difference is smaller, although in the highest-performance learning rate region adding FOP to Adam outperforms the standard method.
We hypothesize that this smaller difference is in large part due to unanticipated effects that preconditioning the gradient has on the adaptive moment estimation performed by Adam.
We leave further investigation into this interaction for future work.
W HAT IS FOP LEARNING ?
By studying the learned preconditioning matrices, it’s possible to gain an intuition for the effect FOP has on the training process.
Interestingly, we found visual tasks induced similar structures in the preconditioning matrices across initializations and across layers.
Visualizing the matrices (Figure 5a) shows that they develop a decorrelating, or whitening , structure: each of the 9 positions in the 3x3 convolutional ﬁlter sends a strong positive weight to itself, and negative weights to its immediate neighbors, without wrapping over the corners of the ﬁlter.
This is interesting, as images are known to have a high degree of spatial autocorrelation (Barlow, 1961; 1989).
As a mechanism for reducing redundant computation, whitening visual inputs is known to be beneﬁcial for both retinal processing in the brain (Atick and Redlich, 1992; Van Hateren, 1992) and in artiﬁcial networks (Desjardins et al., 2015; Huang et al., 2018).
However, it is more unusual to consider whitening of the learning signal, as observed in FOP, rather than the forward activities.
This learned pattern is accompanied by a shift in the norm of the curvature matrix, as the diagonal elements, initialized to one, shrink in value, and the off-diagonal elements grow.
This shift in distribution is visualized in Figure 5b, and grows stronger deeper in the network.
It is possible that this indicates that standard gradient descent is more ill-conditioned in higher layers, or perhaps equivalently, that a greater degree of decorrelation is helpful for kernels that must disentangle high-level representations.
We can also examine the eigenvalue spectra for the learned matrices across layers (Figure 5c).
We can see that the basic requirement that in general a preconditioning matrix must be positive semideﬁnite, so as not to reverse the direction of the gradient, is met.
However, we also note that the eigenvalues are very small in magnitude, indicating a near-zero determinant.
This results in a matrix that is essentially non-invertible, an interesting property, as quasi-Newton and natural gradient methods seek to compute or approximate the inverse of either the Hessian or Fisher information matrix.
The implications of this non-invertibility are avenues for future study.
Furthermore, the eigenvalues grow smaller, and their distribution more uniform, higher in the network, in accordance with the pattern observed in Figure 5b.
A uniform eigenspectrum is seen as an attribute of an effective preconditioning matrix (Li, 2015), as it indicates an even convergence rate in parameter space.
C ONVERGENCE
Our experiments indicate that the curvature matrices learned by FOP converge to a relatively ﬁxed norm roughly two-thirds of the way through training (Figure 5b).
This is important, as it indicates that the effective learning rate induced by the preconditioners stabilizes.
This allows us to perform a preliminary convergence analysis of the algorithm in a manner analogous to Baydin et al.
Consider a modiﬁcation of FOP in which the symmetric, positive semi-deﬁnite preconditioning matrix P ( t ) =
M ( t ) M ( t ) (cid:62) is rescaled at each iteration to have a certain norm γ ( t ) ,
such that γ ( t ) ≈ || P ( t )
|| 2 when t is small and γ ( t )
≈ p ( ∞ ) when t is large, where p ( ∞ ) is some chosen constant.
Speciﬁcally, as in Baydin et al.
, we set γ ( t ) =
δ ( t ) ||
P ( t ) ||
+ (1 − δ ( t ))
p ( ∞ ) , where δ ( t ) is some function that decays over time and starts training at 1 (e.g., 1 /t 2 ).
This formulation allows us to extend the convergence proof of Baydin et al.
(2017) to FOP, under the same assumptions about the objective function J :
Theorem 1 Suppose that J is convex and L -Lipschitz smooth with (cid:107)∇ θ J (cid:107) <
K for some ﬁxed K and all model parameters θ .
Then θ t → θ ∗ if p ( ∞ ) < 1 /L
and tδ ( t ) → 0 as t → ∞ , where the θ t are generated by (non-stochastic) gradient descent.
Proof.
We can write
[Equation]
where ρ is the hypergradient learning rate.
Our assumptions about the limiting behavior of tδ ( t ) then imply that δ ( t ) (cid:107) P ( t ) (cid:107) → 0 and so γ ( t ) → p ( ∞ ) as t → ∞ .
For sufﬁciently large t , we therefore have γ ( t ) ∈ ( 1 L +1 , 1 L ) .
Note also that as P is symmetric and positive semi-deﬁnite, it will not prevent the convergence of gradient descent (Supplementary Section A.3).
Moreover, preliminary investigation shows that the angle of rotation induced by the FOP matrices is signiﬁcantly below 90 ◦ (Supplementary Figure A.1).
Lillicrap et al.
(2016) showed that such a rotation does not impede the convergence of gradient descent.
Because standard SGD converges under these conditions (Karimi et al., 2016), FOP must as well.
C ONCLUSION
In this paper, we introduced a novel optimization technique, FOP, that learns a preconditioning matrix online to improve convergence and generalization in large-scale machine learning models.
We tested FOP on several problems and architectures, examined the nature of the learned transformations, and provided a preliminary analysis of FOP’s convergence properties (Supplementary Section 6).
There are a number of opportunities for future work, including learning transformations for other optimization parameters (e.g., the momentum parameter in case of the the momentum optimizer), expanding and generalizing our theoretical analysis, and testing FOP on a wider variety of models and data sets. A CKNOWLEDGEMENTS
We would like to thank Kenneth Stanley, Jeff Clune, Rosanne Liu, and members of the Horizons and Deep Collective research groups at Uber AI for useful feedback and discussions.
A S UPPLEMENTAL I NFORMATION A. D ERIVATION OF U PDATE FOR M ( t )
Our update rule for the model parameter θ , given by Equation 2, is
[Equation]
Then by applying the the chain and product rules the update is given by
[Equation]
In order to control for the scaling induced by FOP and measure the effect of its rotation only [Equation] we introduce normalized FOP
[Equation]
which performs the following parameter update:
[Equation]
where n is the ﬁrst dimension of M .
This update has the effect of normalizing the preconditioner
M ( t ) M ( t ) (cid:62) , then re-scaling the update by √ n to match the norm of gradient descent, as (cid:107) I n (cid:107) = √ n , where I n is the identity matrix of size n .
A. C ONVERGENCE OF P RECONDITIONED G RADIENT D ESCENT
We demonstrate that given a random, symmetric positive semi-deﬁnite preconditioning matrix P , gradient descent will still converge at a linear rate, modeling our proof after that of Karimi et al.
Theorem 2 Consider a convex, L -Lipschitz objective function f ( θ ) with global minimum f ∗ which obeys the Polyak-Łojasiewicz (PL) Inequality (Polyak, 1963), [Equation]
[Equation]
Then applying the gradient update method given by
[Equation]
[Equation]
where P is a real, symmetric positive semi-deﬁnite matrix and ρ = 2 λ min − λ 2max L is the step size, where λ min and λ max are the minimum and maximum eigenvalues of P , respectively, results in a global linear convergence rate given by
[Equation]
Proof.
First, assume a step-size of 1 /L .
Given that f is L -Lipschitz continuous, we can write
[Equation]
where ∇ f denotes ∇ θ k f .
Plugging in the gradient update equation gives
[Equation]
Let Q (cid:62) Λ Q =
P be the eigendecomposition of P , such that the columns of Q are the orthonormal eigenvectors and Λ is a diagonal matrix whose entries are the eigenvalues of P , which are all nonnegative due to symmetric positive semi-deﬁniteness of P .
We can then rewrite the ﬁrst term of
Equation 13 as
[Equation]
We now change our basis, letting G = Q ∇ f and deﬁne λ min = min
i { λ i } and λ max = max i { λ i } , where λ i are the eigenvalues.
Then we have
[Equation]
Examining the second term of Equation
[Equation]
we have
[Equation] [Equation]
Combining the results of Equations 14 and 15 gives
[Equation]
We can then revert to the original basis:
[Equation]
giving us
[Equation]
By the PL inequality we have that (cid:107)∇ f (cid:107) 2 ≥ 2 µ ( f ( θ k ) − f ∗ ) for some µ > 0 .
Plugging this in gives
[Equation]
Then let ρ = 1 L (cid:0) 2 λ min − λ 2max (cid:1) .
We have
[Equation]
Rearranging and subtracting f ∗ from both sides gives
[Equation]
Applying Equation 22 recursively gives the desired convergence:
[Equation]
Thus this preconditioned gradient descent converges with a step size 2 λ min − λ 2max L .
Standard gradient descent converges with a step size 1 /L under these assumptions.
We also note that even if P changes over the course of training, the required step-size will vary, but convergence will still occur.
[Figure]