Amortized Proximal Optimization
Abstract
We propose a framework for online meta-optimization of parameters that govern optimization, called Amortized Proximal Optimization (APO).
We first interpret various existing neural network optimizers as approximate stochastic proximal point methods which trade off the current-batch loss with proximity terms in both function space and weight space.
The idea behind APO is to amortize the minimization of the proximal point objective by meta-learning the parameters of an update rule.
We show how APO can be used to adapt a learning rate or a structured preconditioning matrix.
Under appropriate assumptions, APO can recover existing optimizers such as natural gradient descent and KFAC.
It enjoys low computational overhead and avoids expensive and numerically sensitive operations required by some second-order optimizers, such as matrix inverses.
We empirically test APO for online adaptation of learning rates and structured preconditioning matrices for regression, image reconstruction, image classification, and natural language translation tasks.
Empirically, the learning rate schedules found by APO generally outperform optimal fixed learning rates and are competitive with manually tuned decay schedules.
Using APO to adapt a structured preconditioning matrix generally results in optimization performance competitive with second-order methods.
Moreover, the absence of matrix inversion provides numerical stability, making it effective for low-precision training.
Introduction
Many optimization algorithms widely used in machine learning can be seen as approximations to an idealized algorithm called the proximal point method (PPM).
When training neural networks, the stochastic PPM iteratively minimizes a loss function J B : R m → R on a mini-batch B , plus a proximity term that penalizes the discrepancy from the current iterate:
[Equation]
where D ( · , · ) measures the discrepancy between two vectors and λ > 0 is a hyperparameter that controls the strength of the proximity term.
The proximity term discourages the update from excessively changing the parameters, hence preventing aggressive updates.
Moreover, the stochastic PPM has good convergence properties [4].
While minimizing Eq.
1 exactly is usually impractical (or at least uneconomical), solving it approximately (by taking first or second-order Taylor series approximations to the loss or the proximity term) has motivated important and widely used optimization algorithms such as natural gradient descent [1] and mirror descent [7].
Stochastic gradient descent (SGD) [69] itself can be seen as an approximate PPM where the loss term is linearized and the discrepancy function is squared Euclidean distance.
Inspired by the idea that the PPM is a useful algorithm to approximate, we propose to amortize the minimization of Eq. 1 by defining a parametric form for an update rule which is likely to be good at minimizing it and adapting its parameters with gradient-based optimization.
We consider adapting optimization hyperparameters (such as the learning rate) for existing optimizers such as SGD and RMSprop [78], as well as learning structured preconditioning matrices.
By choosing a structure for the update rule inspired by existing optimizers, we can take advantage of the insights that went into those optimizers while still being robust to cases where their assumptions (such as the use of linear or quadratic approximations) break down.
By doing meta-descent on the optimization parameters, we can amortize the cost of minimizing the PPM objective, which would otherwise take many steps per parameter update.
Hence, we call our approach Amortized Proximal Optimization (APO) .
Eq.
1 leaves a lot of freedom for the proximity term.
We argue that many of the most effective neural network optimizers can be seen as trading off two different proximity terms: a function space discrepancy (FSD) term which penalizes the average change to the network’s predictions, and a weight space discrepancy (WSD) term which prevents the weights from moving too far, encouraging smoothness to the update and maintaining the accuracy of second-order approximations.
Our meta-objective includes both terms.
Our formulation of APO is general, and can be applied to various settings, from optimizing a single optimization hyperparameter to learning a flexible update rule.
We consider two use cases that cover both ends of this spectrum.
At one end, we consider the problem of adapting learning rates of existing optimizers, specifically SGD, RMSprop, and Adam [19].
The learning rate is considered one of the most essential hyperparameters to tune [10], and good learning rate schedules are often found by years of trial and error.
Empirically, the learning rate schedules found by APO outperformed the best fixed learning rates and were competitive with manual step decay schedules.
Our second use case is more ambitious.
We use APO to learn a preconditioning matrix, giving the update rule the flexibility to represent second-order optimization updates such as Newton’s method,
Gauss-Newton, or natural gradient descent.
We show that, under certain conditions, the optimum of our APO meta-objective with respect to a full preconditioning matrix coincides with damped versions of natural gradient descent or Gauss-Newton.
While computing and storing a full preconditioning matrix for a large neural network is impractical, various practical approximations have been developed.
We use APO to meta-learn a structured preconditioning matrix based on the EKFAC optimizer [24].
APO is more straightforward to implement in current-day deep learning frameworks than EKFAC and is also more computationally efficient per iteration because it avoids the need to compute eigendecompositions.
Empirically, we evaluate APO for learning structured preconditioners on regression, image reconstruction, image classification, and neural machine translation tasks.
The preconditioning matrix adapted by APO achieved competitive convergence to other second-order optimizers.
Preliminaries
Consider a prediction problem from some input space X to an output space T . We are given a finite training set D train =
{ ( x ( i ) , t ( i ) ) } N i =1 .
For a data point ( x , t ) and parameters θ ∈ R m , let y =
f ( x , θ ) be the prediction of a network parameterized by θ and L ( y , t ) be the loss.
Our goal is to minimize the cost function:
[Equation]
We use J B ( θ ) to denote the mean loss on a mini-batch of examples B = { ( x ( i ) , t ( i ) )
} Bi
=1 .
We summarize our notation in Appendix A.
Proximal Optimization and Second-Order Methods: A Unifying Framework
We first motivate the proximal objective that we use as the meta-objective for APO, and relate it to existing neural network optimization methods.
Our framework is largely based on Grosse [25], to which readers are referred for a more detailed discussion.
[Figure]
3.1 Proximal Optimization
When we update the parameters on a mini-batch of data, we would like to reduce the loss on that mini-batch, while not changing the predictions on previously visited examples or moving too far in weight space.
This motivates the following proximal point update:
[Equation]
where D F ( u , θ , x ) =
ρ ( f ( x , u ) , f ( x , θ )) and D W ( u , θ ) = 1 / 2 ∥ u − θ ∥ 22 are discrepancy functions (described below). Here, λ FSD and λ WSD are hyperparameters that control the strength of each discrepancy term, ˜ x is a random data point sampled from the data-generating distribution D , and ρ ( · , · ) is the output-space divergence.
The proximal objective in Eq.
3 consists of three terms.
The first term is the loss on the current minibatch.
The second term is the function space discrepancy (FSD) , whose role is to prevent the update from substantially altering the predictions on other data points.
The general idea of the FSD term has been successful in alleviating catastrophic forgetting [11], fine-tuning pre-trained models [35], computing influence functions [5], and training a student model from a teacher network [30].
The final term is the weight space discrepancy (WSD) , which encourages the update to move the parameters as little as possible.
It can be used to motivate damping in the context of second-order optimization [53].
While weight space distance may appear counterproductive from an optimization standpoint because it depends on the model parameterization, analyses of neural network optimization and generalization often rely on network parameters staying close to their initialization in the
Euclidean norm [33, 88, 9].
In fact, Wadia et al.
[83] showed that pure second-order optimizers (i.e. ones without WSD regularization) are unable to generalize in the overparameterized setting.
Figure 1 illustrates the effects of the WSD and FSD terms on the exact PPM update for a 1D regression example with a batch size of 1.
If the proximal objective includes only the loss and WSD term (i.e. λ FSD = 0 ), the PPM update makes the minimal change to the weights which fits the current example, resulting in a global adjustment to the function which overwrites all the other predictions.
If only the loss and FSD terms are used (i.e. λ WSD = 0 ), the update carves a spike around the current data point, failing to improve predictions on nearby examples and running the risk of memorization.
When both WSD and FSD are penalized, it makes a local adjustment to the predictions, but one which nonetheless improves performance on nearby examples.
3.2 Connection Between Proximal Optimization and Second-Order Optimization
[Table]
We further motivate our proximal objective by relating it to existing neural network optimizers.
Ordinary SGD can be viewed as an approximate
PPM update with a first-order approximation to the loss term and no FSD term.
Hessian-Free optimization [51], a classic second-order optimization method for neural networks, approximately minimizes a second-order approximation to the loss on each batch using conjugate gradients.
It can be seen as minimizing a quadratic approximation to Eq.
3 with no FSD term.
Amari [1] motivated natural gradient descent (NGD) as a steepest descent method with an infinitesimal step size; this justifies a first-order approximation to the loss term and a second-order approximation to the proximity term.
Optimizing over a manifold of probability distributions with KL divergence as the proximity term yields the familiar update involving the Fisher information matrix.
Natural gradient optimization of neural networks [52, 54] can be interpreted as minimizing Eq.
3 with a linear approximation to the loss term and a quadratic approximation to the FSD term.
While NGD traditionally omits the WSD term in order to achieve parameterization invariance, it is typically included in practical neural network optimizers for stability [-[53]. [Figure] Algorithm 
-][54].
[-B ∼ D train ▷ Sample mini-batch to compute the gradient and loss term if t mod K = 
-]
In a more general context, when taking a first-order approximation to the loss and a second-order approximation to the FSD, the update rule is given in closed form as:
[Equation]
where G is the Hessian of the FSD term.
The derivation is shown in Appendix F. All of these relationships are summarized in Table 1, and derivations of all of these claims are given in Appendix G.
Amortized Proximal Optimization
In this section, we introduce Amortized Proximal Optimization (APO), an approach for online meta-learning of optimization parameters.
Then, we describe two use cases that we explore in the paper: (1) adapting learning rates of existing base optimizers such as SGD, RMSProp, and Adam, and (2) meta-learning a structured preconditioning matrix.
4.1 Proximal Meta-Optimization
We assume an update rule u parameterized by a vector ϕ which updates the network weights θ on a batch B ( t ) : [Equation]
One use case of APO is to tune the hyperparameters of an existing optimizer, in which case ϕ denotes the hyperparameters.
For example, when tuning the SGD learning rate, we have ϕ =
η
and the update is given by:
[Equation]
More ambitiously, we could use APO to adapt a full preconditioning matrix P .
In this case, we define ϕ =
P and the update is given by:
[Equation]
In Section 3, we introduced a general proximal objective for training neural networks and observed that many optimization techniques could be seen as an approximation of PPM.
Motivated by this connection, we propose to directly minimize the proximal objective with respect to the optimization parameters.
While still being able to take advantage of valuable properties of existing optimizers, direct minimization can be robust to cases when the assumptions (such as linear and quadratic approximation of the cost) do not hold.
Another advantage of adapting a parametric update rule is that we can amortize the cost of minimizing the proximal objective throughout training.
We propose to use the following meta-objective, which evaluates the proximal objective at u ( θ , ϕ , B ) :
[Equation]
In practice, we estimate the expectations in the meta-objective by sampling two different mini-batches, B and B ′ , where B is used to compute the gradient and the loss term, and B ′ is used to compute the FSD term.
Intuitively, this proximal meta-objective aims to find optimizer parameters ϕ that [Figure] Algorithm Proximal ϕ
Require: θ (initial model parameters), ϕ (initial optimization parameters), K (meta-update interval), α (meta-LR)
Require: λ WSD (weight-space discrepancy term weighting), λ FSD (function-space discrepancy term weighting) while not converged, iteration t do B ∼ D train ▷ Sample mini-batch to compute the gradient and {+loss term if t mod K = 
 then ▷ Perform meta-update every K iterations B ′ ∼ D train ▷ Sample additional mini-batch to compute the FSD term θ ′ ( ϕ ) :+}
[Equation]
u ( θ , ϕ , B )
▷ Compute the 1-step lookahead parameters
[Equation]
▷ Compute meta-objective ϕ
[Equation]
▷ Update optimizer parameters (e.g. LR or preconditioner) end if θ ← u ( θ ,
, B ) ▷ Update model parameters end while minimize the loss on the current mini-batch, while constraining the size of the step with the FSD and
WSD terms so that it does not overfit the current mini-batch and undo progress that has been made by other mini-batches.
The optimization parameters ϕ are optimized with a stochastic gradient-based algorithm (the metaoptimizer ).
The meta-gradient ∇ ϕ Q ( ϕ ) can be computed via automatic differentiation through the one-step unrolled computation graph (Figure 2).
We refer to our framework as Amortized Proximal Optimization (APO, Algorithm 1).
4.2 APO for Learning Rate Adaptation
One use case of APO is to adapt the learning rate of an existing base optimizer such as SGD.
To do so, we let u SGD ( θ , η, B ) be the 1-step lookahead of parameters and minimize the proximal meta-objective with respect to the learning rate η .
Although adaptive optimizers such as RMSProp and Adam use coordinate-wise learning rates, they still have a global learning rate which is essential to tune.
APO can be applied to such global learning rates to find learning rate schedules (that depend on λ FSD or λ WSD ).
4.3 APO for Adaptive Preconditioning
More ambitiously, we can use the APO framework to adapt the preconditioning matrix, allowing the update rule to flexibly represent various second-order optimization updates.
We let u Precond ( θ , P , B ) denote the parameters after 1 preconditioned gradient step and adapt the preconditioning matrix P according to our framework.
If the assumptions made when deriving the second-order methods (detailed in Section 3.2) hold, then the optimal preconditioning matrix is equivalent to various second-order updates, depending on the choice of the FSD function.
Theorem 1.
Consider an approximation ˆ Q ( P ) to the meta-objective (Eq. 8) where the loss term is linearized around the current weights θ and the FSD term is replaced by its second-order approximation around θ .
Denote the gradient on a mini-batch as g = ∇ θ J B ( θ ) , and assume that the second moment matrix E B∼D (cid:2) gg ⊤ (cid:3) is non-singular.
Then, the preconditioning matrix which minimizes ˆ Q is given by P ⋆ = ( λ FSD G + λ WSD I ) − 1 , where G denotes the Hessian of the FSD evaluated at θ .
The proof is provided in Appendix H. As an important special case, when the FSD term is derived from the KL divergence between distributions in output space, G coincides with the Fisher information matrix F = E x ∼D , y ∼ P y | x ( θ ) (cid:2) ∇ θ log p ( y | x , θ ) ∇ θ log p ( y | x , θ ) ⊤ (cid:3) , where P y | x ( θ ) denotes the model’s predictive distribution over y .
Therefore, the optimal preconditioner is the damped natural gradient preconditioner, P ⋆ =
( F + λ WSD I ) − 1 when λ FSD = 1 .
Thus, when APO is used to exactly minimize an approximate meta-objective, the update it yields coincides with classical second-order optimization algorithms, depending on the choice of the FSD function.
4.4 Structured Preconditioner Adaptation
In the previous sections, the discussion assumed a full preconditioning matrix for simplicity.
However, a full preconditioner is impractical to represent for modern neural networks.
Moreover, for practical stability of the learned preconditioned update, we would like to enforce the preconditioner to be positive semidefinite (PSD) so that the transformed gradient is a descent direction [64].
To satisfy these requirements, we adopt a structured preconditioner analogous to that of the EKFAC optimizer [24].
Given a weight matrix W ∈ R m i × m i +1 of a layer, we construct the preconditioning matrix as a product of smaller matrices:
[Equation]
where A ∈ R m {+i +
+}
× m
i +1 , B ∈ R
m i × m i , and S ∈ R m i × m i {++
+} are small block matrices.
Here, ⊗ denotes the Kronecker product, diag ( · ) denotes the diagonalization operator, and vec ( · ) denotes the vectorization operator.
This parameterization is memory efficient: it requires m 2 i + m
i +
+ m i m
i + parameters to store, as opposed to m
i m
i +1 parameters for a full preconditioning matrix.
It is straightforward to show that the structured preconditioner in Eq.
9 is PSD, as it takes the form
CDC ⊤ , where D is PSD.
The preconditioned gradient can be computed efficiently by using the properties of the Kronecker product:
[Equation]
where ⊙ denotes elementwise multiplication.
This is tractable to compute as it only requires four additional matrix multiplications and elementwise multiplication of small block matrices in each layer when updating the parameters.
While EKFAC uses complicated covariance estimation and eigenvalue decomposition to construct the block matrices, in APO, we meta-learn these block matrices directly, where ϕ
= [ vec ( A ) ⊤ , vec ( B ) ⊤ , vec ( S ) ⊤ ] ⊤ .
As APO does not require inverting (or performing eigendecompositions of) the block matrices, our structured representation incurs less computation per iteration than EKFAC.
While we defined an optimizer with the same functional form as EKFAC, it is not immediately obvious whether the preconditioner which is actually learned by APO will be at all similar.
A
Corollary of Theorem 1 shows that if certain conditions are satisfied, including the assumptions underlying KFAC [54], then the structured preconditioner minimizing Eq.
8 coincides with KFAC:
Corollary 2.
Suppose that (1) the assumptions for Theorem 1 are satisfied, (2) the FSD term measures the KL divergence, and (3) λ
WSD
=
and λ FSD = 1 .
Moreover, suppose that the parameters θ satisfy the KFAC assumptions listed in Appendix I.
Then, the optimal solution to the approximate meta-objective recovers the KFAC update, which can be represented using Eq. 9.
The proof is in Appendix I. If the KFAC assumptions are not satisfied, then APO will generally learn a different preconditioner.
This may be desirable, especially if differing probabilistic assumptions lead to different update rules, as is the case for KFAC applied to convolutional networks [26, 38].
4.5 Computation and Memory Cost
Computation Cost.
Computing the FSD term requires sampling an additional mini-batch from the training set and performing two additional forward passes for f (˜ x , θ ) and f (˜ x , u ( θ , ϕ , B )) .
Combined with the loss term evaluated on the original mini-batch, one meta-optimization step requires three forward passes to compute the proximal objective.
It additionally requires a backward pass through the 1-step unrolled computation graph (Figure 2) to compute the gradient of the proximal meta-objective Q ( ϕ ) with respect to ϕ .
This overhead can be reduced by performing a meta-update only once every K iterations: the overhead will consist of 3 /K additional forward passes and 1 /K additional backward passes per iteration, which is small for modest values of K (e.g., K = 10 ).
Memory Cost.
APO requires twice the model memory for the 1-step unrolling when computing the proximal meta-objective.
In the context of structured preconditioner adaptation, we further need to store block matrices A , B , and S (in Eq. 9) for each layer, as in KFAC and EKFAC.
Related Work
We provide extended related work and a conceptual comparison of meta-optimization methods in
Appendix B and Table 6, respectively.
Gradient-Based Learning Rate Adaptation.
Maclaurin et al.
[50] backpropagate through the full unrolled training procedure to meta-optimize learning rate schedules offline .
This is expensive, as it requires completing a full training run to make a single hyperparameter update.
A related approach is to unroll the optimization for a small number of steps and perform truncated backpropagation [16, 22]. Micaelli et al.
[58] perform offline hyperparameter optimization using forward-mode (FDS) rather than reverse-mode gradient accumulation.
FDS performs one update to the hyperparameters after each full training run, and thus requires multiple training runs, in contrast to APO, which operates within a single training run.
Hypergradient descent [6] adapts the learning rate to minimize the expected loss in the next iteration.
Second-Order Optimization.
Although preconditioned methods have better convergence rates than first-order methods [8, 64], storing and computing the inverse of preconditioning matrices is impractical for high-dimensional problems.
To mitigate these computational issues, Hessian-free optimization [51, 55] approximates Newton’s update by only accessing the curvature matrix through
Hessian-vector products.
Other works impose a structure on the preconditioner by representing it as a
Kronecker product [54, 26, 56, 24, 27, 77], a diagonal matrix [19, 36], or a low-rank matrix [39, 60].
However, these approximate second-order methods may not be easy to implement in deep learning frameworks, and can still be expensive as they often require matrix inversion or eigendecomposition.
Gradient-Based Preconditioner Adaptation.
There has been some prior work on meta-learning preconditioners.
Moskovitz et al.
[61] learn the preconditioning matrix with hypergradient descent.
Meta-curvature [66] and warped gradient descent [41, 21] adapt the preconditioning matrix that yields effective parameter updates across diverse tasks in the context of few-shot learning.
Experiments
Our experiments investigate the following questions: (1) How does the structured preconditioning matrix adapted by APO perform in comparison to existing first- and second-order optimizers?
How does the learning rate adapted by APO perform compared to optimal fixed learning rates and manual decay schedules commonly used in the literature?
We used APO to meta-learn the preconditioning matrices for a broad range of tasks, including several regression datasets, autoencoder training, image classification on CIFAR-10 and CIFAR-100 using several network architectures, neural machine translation using transformers, and low-precision (16-bit) training.
Several of these tasks are particularly challenging for first-order optimizers.
In addition, we used APO to tune the global learning rate for multiple base optimizers – SGD, SGD with momentum (denoted SGDm), RMSprop, and Adam – on CIFAR-10 and CIFAR-100 classification with several network architectures. [Figure]
[Figure] We denote our method that adapts the structured preconditioning matrix as “APO-Precond”.
The method that tunes the global learning rate of a base optimizer is denoted as “Base-APO” (e [Figure] g. SGDm-APO).
Experiment details and additional experiments, including ablation studies, are given in Appendix C and J, respectively.
6.1 Meta-Learning Preconditioners
Poorly-Conditioned Regression.
We fist considered a regression task traditionally used to illustrate failures of neural network optimization [68].
The targets are given by t =
Ax , where A is an illconditioned matrix with κ ( A ) = 10 10 .
We trained a two-layer linear network f ( x , θ ) =
W 2 W 1 x and minimized the objective J ( θ ) =
E x ∼N ( 0 , I ) (cid:2) ∥
Ax − W 2 W 1 x ∥ 2 (cid:3) .
In Figure 4 (left), we show training curves for SGDm, Adam, Shampoo [27], KFAC, and APO-Precond.
As the problem is ill-conditioned, 2 nd -order optimizers such as Shampoo and KFAC converge faster than 1 st -order methods.
APO-Precond performs comparably to 2 nd -order optimizers with lower loss than KFAC.
UCI Regression.
Next, we validated APO-Precond on the Slice, Protein, and Parkinsons datasets from the UCI regression collection [18].
We trained a 2-layer MLP with 100 hidden units per layer and ReLU activations for 500 epochs.
The training curves for each optimizer are shown in
Figure 3.
By tuning the preconditioning matrix during training, APO-Precond consistently achieved competitive convergence compared to other second-order optimizers and reached lower training loss than all baselines.
Image Reconstruction.
We trained an 8-layer autoencoder on MNIST [40]; this is known to be a challenging optimization task for first-order optimizers [31, 55, 54].
We followed the experimental set-up from Martens & Grosse [54], where the encoder and decoder consist of 4 fully-connected layers with sigmoid activation.
The decoder structure is symmetric to that of the encoder, and they do not have tied weights.
The logistic activation function and the presence of a bottleneck layer make this a challenging optimization problem compared with most current-day architectures.
We compare
APO-Precond with SGDm, Adam, Shampoo, and KFAC optimizers and show the training losses for each optimizer in Figure 4 (middle).
APO-Precond converges faster than first-order methods and achieves competitive training loss to other second-order methods (although there remains a performance gap compared with KFAC).
[Table]
Image Classification.
To investigate whether adapting the preconditioner with APO improves generalization [Table] we conducted image classification experiments on CIFAR-10 and CIFAR-100.
We trained LeNet [40], AlexNet [37], VGG-16 [71] (w/o batch norm [32]), ResNet- [Table] and ResNet32 [29] architectures for 200 epochs on batches of 128 images.
The test accuracies for SGDm, Adam, KFAC
[Table]
and APO-Precond are shown in Table 2.
We found that APO-Precond achieved competitive generalization performance to SGDm and KFAC.
In particular, for architectures without batch normalization (LeNet, AlexNet
[Table] and VGG-16), APO-Precond improved the test accuracy substantially.
Neural Machine Translation.
To verify the effectiveness of APO on various tasks, we applied APO-Precond on the IWSLT’14 German-to-English translation task [15].
We used a Transformer [79] composed of 6 encoder and decoder layers, with word embedding and hidden vector dimensionality 512.
We compared APO-Precond to SGDm and AdamW [48].
For AdamW, we used a warmupthen-decay learning rate schedule widely used in practice, and for SGD and APO-Precond, we kept the learning rate fixed after the warmup.
In Table 2, we show the final test BLEU score for SGDm, AdamW, and APO-Precond.
While keeping the learning rate fixed, we achieved a BLEU score competitive with AdamW.
[Table]
Low Precision Training.
Low precision training presents a challenge for second-order optimizers such as KFAC which rely on matrix inverses that may be sensitive to quantization noise.
We trained LeNet and ResNet-18 with 16-bit floating-point arithmetic to examine if APO-Precond is applicable in training the networks in lower precision.
We used the experimental setup from Section 6.1 but stored parameters, activations, and gradients in 16-bit precision.
We found that KFAC required a large damping factor to maintain stability, and this prevented it from fully utilizing curvature information.
In contrast, as APO-Precond does not require matrix inversion, it remained stable with the same choice of FSD and WSD weights we used in the full precision experiments.
The final test accuracies on ResNet-18 for
SGDm, KFAC, and APO-Precond are shown in Table 3.
6.2 Meta-Learning Learning Rates
Image Classification on MNIST.
First, we compared SGDm and RMSprop to their APO-tuned variants to train an MLP on MNIST.
We used a two-layer MLP with 1000 hidden units per layer and
ReLU nonlinearities, and trained on mini-batches of size 100 for 100 epochs.
Figure 4 (Right) shows the training loss achieved by each approach; we found that for both base optimizers, APO improved convergence speed and obtained substantially lower loss than the baselines.
Image Classification on CIFAR-10 & CIFAR-100.
For learning rate adaptation on CIFAR-10, we experimented with three network architectures: ResNet32 [29], ResNet34, and WideResNet [Figure] [Table] [Figure]
[Table] (WRN-28-10) [86]. For ResNet32, we trained for 400 epochs, and the decayed baseline used a step schedule with 10 × decay at epochs 150 and 250, following [49].
For ResNet34 and WRN-28-10, we trained for 200 epochs, and the decayed baseline used a step schedule with 5 × decay at epochs 60, 120, and 160, following [86].
For CIFAR-100, we used WRN-28-10 with the same schedule as for CIFAR-10.
For each of the base optimizers, we compared APO to (1) the optimal fixed learning rate and (2) a manual step learning rate decay schedule.
The test accuracies for each base optimizer and their APO-tuned variants are shown in Table 4.
In addition, Figure 5 shows the training loss, test accuracy and learning rate adaptation for WRN-28-10 on CIFAR-10, using SGDm as the base optimizer.
Using APO to tune the global learning rate yields higher test accuracy than the best fixed learning rate, and is competitive with the manual schedule.
Conclusion
We introduced Amortized Proximal Optimization (APO), a framework for online meta-learning of optimization parameters which approximates the proximal point method by learning a parametric update rule.
As the meta-parameters are updated only once per K steps of optimization, APO incurs minimal computational overhead.
We applied APO to two settings: (1) meta-learning the global learning rate for existing base optimizers (e.g., SGD, RMSprop, and Adam) and (2) meta-learning structured preconditioning matrices, which provides a new approach to 2 nd -order optimization.
Compared to methods such as KFAC, APO eliminates the need to compute matrix inverses, yielding improved efficiency and numerical stability.
On a range of tasks, we showed that APO is competitive with 2 nd -order methods, and improves generalization compared to baseline 1 st - and 2 nd -order optimizers. Acknowledgements
We thank Michael Zhang for valuable feedback on this paper and thank Alston Lo for helping setting up the experiments.
We would also like to thank Saminul Haque, Jonathan Lorraine, Denny Wu, and
Guodong Zhang, and our many other colleagues for their helpful discussions throughout this research.
Resources used in this research were provided, in part, by the Province of Ontario, the Government of Canada through CIFAR, and companies sponsoring the Vector Institute ( www.vectorinstitute. ai/partners ). Checklist For all authors... (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope?
[Yes] (b) Did you describe the limitations of your work?
[Yes] (c) Did you discuss any potential negative societal impacts of your work?
[N/A]
Our work is fundamental in nature and does not enable immediate misuse.
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them?
[Yes] 2.
If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?
[Yes] We list the KFAC assumptions in Section I. (b) Did you include complete proofs of all theoretical results?
[Yes] We provide the proof of Theorem 1 in Section H.
In Section I, we list the KFAC assumptions and prove Corollary 2.
Additionally, in Section F, we provide a derivation of Eq. 4, and in Section G, we derive classic first- and second-order optimization algorithms (gradient descent, Newton’s method, and natural gradient) starting from the proximal objective for the PPM (Eq. 3).
If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
[Yes] We include code in the supplementary material.
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
[Yes] We provide the training details for all of our experiments in Appendix C. (c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[Yes] For several of our experiments, we plot shaded regions which indicate the min/max values obtained after running the experiments 4 random seeds.
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?
[Yes] We provide these details in Appendix C.1.
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?
[Yes] We cited all datasets and model architectures we used in our experiments (Section 6).
(b) Did you mention the license of the assets?
[N/A] (c) Did you include any new assets either in the supplemental material or as a URL?
[N/A] (d) Did you discuss whether
and how consent was obtained from people whose data you’re using/curating?
[N/A] (e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content?
[N/A] 5.
If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if applicable?
[N/A] (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?
[N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
[N/A]
This appendix is structured as follows
• In Section A, we provide an overview of the notation we use throughout the paper. • In Section B, we provide a discussion of additional related work.
• In Section C, we provide experiment details.
• In Section D, we provide additional experimental results.
• In Section E, we provide an ablation investigating the difference between evaluating the loss term on the same minibatch used to compute the base optimizer step vs a separate randomlysampled minibatch.
Using a separate random minibatch yields a similar meta-objective to that studied in Wu et al.
[85], which suffers from short-horizon bias.
• In Section F, we provide a derivation of Eq. 4.
• In Section G, we derive classic first- and second-order optimization algorithms (gradient descent, Newton’s method, and natural gradient) starting from the proximal objective for the PPM (Eq. 3).
• In Section H, we provide the proof for Theorem 1, which shows that exactly optimizing the proximal meta-objective with respect to the preconditioner P can recover various first- and second-order optimization methods.
In Section I, we list the KFAC assumptions and prove Corollary 2, which shows that under the assumptions of Theorem 1 and the KFAC assumptions, exactly optimizing the proximal meta-objective yields the KFAC update. • In Section J, we show an ablation over λ WSD and λ FSD , and evaluate how well APO performs with different meta-update intervals.
A Table of Notation
[Table]
optimizer
[Table]
e.g. stochastic gradient descent)
[Table]
u (
[Table] , ϕ , B ) [Table] the data B is implicit
[Table]
B Extended Related Work
Short-Horizon Bias.
Short-horizon bias (SHB) has been identified as a core challenge in the meta-optimization of optimizer parameters such as learning rates.
Wu et al.
[85] analyzed the setting where the meta-objective is the expected loss after an optimization step E B∼D train [ J B ( u ( θ , ϕ , B ))] .
Note that the expectation over the entire dataset is intractable and can be typically approximated using a single mini-batch.
This meta-objective yields undesirable behaviour where the learning rate decreases rapidly to reduce fluctuations of the loss caused by stochastic mini-batch evaluation.
While our proximal meta-objective is greedy (e.g. it relies on a single lookahead step), APO empirically does not suffer from SHB.
We attribute this to the fact that the first term in our proximal meta-objective evaluates the loss on the same mini-batch used to compute the gradient rather than a randomly-selected mini-batch.
We refer readers to Appendix E for an ablation demonstrating the difference between using the same and a different mini-batch in computing the first term in our meta-objective.
[Table]
Learned Optimizers
[Table]
training runs for each meta-parameter update); and 5) which meta-parameters it can tune.
Forward-Mode Differentiation.
In principle, short horizon meta-objectives cannot model longterm behavior.
Forward-mode autodiff like Real-Time Recurrent Learning (RTRL) and its approximations (UORO [76], KF-RTRL [62], OK [12]) can perform online optimization without suffering from SHB in principle (but still suffering from hysteresis).
However, they do not scale well to many parameters, and thus are be used to tune the preconditioner.
Other approaches like PES [82] use a finite-differences approximation to the gradient, which does not scale to the dimensionality of the preconditioner.
RTRL has been applied to hyperparameter adaptation by Franceschi et al.
[22], and a method for LR adaptation that interpolates between hypergradient descent and RTRL, called
MARTHE, was introduced in [17].
Black-Box Hyperparameter Optimization.
Finding good optimization hyperparameters is a longstanding problem [10].
Black-box methods for hyperparameter optimization, such as grid search, random search [13], and Bayesian optimization [72, 75, 73], are expensive, as they require performing many complete training runs, and can only find fixed hyperparameter values (e.g., a constant learning rate).
Hyperband [44] can reduce the cost by terminating poorly-performing runs early, but is still limited to finding fixed hyperparameters.
Population Based Training (PBT) [34] trains a population of networks simultaneously, and throughout training it terminates poorly-performing networks, replaces their weights with a copy of the weights of a better-performing network, perturbs the hyperparameters, and continues training from that point.
PBT can find a coarse-grained learning rate schedule, but because it relies on random search, it is far less efficient than gradient-based meta-optimization.
Loss-Based Learning Rate Adaptation.
Several works have proposed modern variants of Polyak step size methods [28, 80, 81, 46] for automatic learning rate adaptation.
In a similar vein, Rolinek &
Martius [70] proposed a method that tracks an estimated minimal loss and tunes a global learning rate to drive the loss on each mini-batch to the minimal loss.
Learned Optimization.
Some authors have proposed learning entire optimization algorithms [42, 43, 2, 84, 57, 84].
Li & Malik [42] view this problem from a reinforcement learning perspective, where the state consists of the objective function J and the sequence of prior iterates { θ ( t ) } and gradients {∇ θ J ( θ ( t ) ) } , and the action is the step ∆ θ .
In this setting, the update rule α is a policy that can be found via policy gradient methods [74].
Approaches that learn optimizers must be trained on a set of objective functions { f 1 , . . .
, f n } drawn from a distribution F .
This setup can be restrictive if we only have one instance of an objective function.
In addition, the initial phase of training the optimizer on a distribution of functions can be expensive.
Learned optimizers require offline training, and typically must use long unrolls to ameliorate short horizon bias, making training expensive.
APO requires only the objective function of interest as it does not have trainable parameters.
Few-Shot Meta-Learning Methods.
Meta-SGD [45] extends MAML [20] by learning not only a shared initialization which can be fine-tuned in a few steps for many different tasks but also a diagonal preconditioning matrix used for the fine-tuning gradient steps (e.g., the inner-loop of MAML, adapting to a specific task).
Meta-SGD requires a task distribution, and only learns a fixed diagonal preconditioner which is not adapted during the task-adaptive optimization procedure.
In contrast,
APO learns a more expressive, block-diagonal preconditioner from a single training run, which is adapted online during training.
Meta-Curvature (MC) [66] meta-learns a preconditioning matrix to improve generalization after a small number of (preconditioned) gradient steps in MAML-style meta-learning.
Meta-Curvature requires a task distribution for training, and relies on the assumption that related tasks have similar curvature; the method learns a single, fixed preconditioning matrix that is applied across tasks (and to new tasks at meta-test time).
In contrast, APO does not require a task distribution, and dynamically adapts the meta-learned preconditioning matrix during a single training run, rather than keeping it fixed as done in MC.
Lee and Choi [41] proposed Transformation Networks (T-nets), which learn an activation-spacemetric that modifies the update direction and step size taken to adapt to each task in a few-shot setting, and Mask Transformation Networks (MT-nets), which extend T-nets by learning which subset of weights to update for each task.
Similarly to Meta-SGD and Meta-Curvature, MT-nets require a task distribution for meta-training, and only learn fixed meta-parameters which are not adapted during training of new tasks.
APO is a generic framework for online meta-learning of optimization parameters.
We demonstrate its effectiveness for two specific types of meta-parameters: the preconditioning matrix and the global learning rate.
Our main contribution is the use of a principled meta-objective based on the proximal point method, which we prove can recover existing first- and second-order optimization algorithms given appropriate approximations to the loss term and FSD term.
We focused on tuning the KFAC-style preconditioner as this is a popular block-diagonal second-order method, and allows us to draw the connection to KFAC in Corollary 1.
However, the APO meta-objective can be applied to tune different parameterizations as well; in particular, one could use interleaved warp-layers to parameterize a preconditioner similarly to [41, 21].
C Experimental Details C.1 Computing Environment
All experiments were implemented using the PyTorch [67] and JAX [14] libraries, and we ran all experiments on NVIDIA P100 GPUs.
C.2 Meta-Optimization Setup
In all experiments, we perform 1 meta-update every 10 updates of the base optimizer (e.g. K =in Algorithm 1).
Hence, the computational overhead of APO is just a small fraction of the original training procedure.
We perform grid searches over λ FSD and λ WSD for both learning rate and preconditioner adaptation.
Learning Rate Adaptation Setup.
For our learning rate adaptation experiments, we used RMSProp as the meta-optimizer with a meta-learning rate of 0.1.
For SGD-APO and SGDm-APO, we set the initial learning rate to 0.1, while for RMSProp-APO and Adam-APO, we set the initial learning rate to 3 e- 4 .
These specific values are not important for good performance; we show in Appendix J.2 that
APO is robust to the choice of initial learning rate.
Preconditioner Adaptation Setup.
For preconditioner adaption experiments, we used Adam as the meta-optimizer with a meta-learning rate chosen from { 1 e- 4 , 3 e- 5 } .
We initialized the block matrices such that the preconditioner is the identity, so that our update is equivalent to SGD at initialization.
Furthermore, we used a warm-up phase of 3000 iterations at the beginning of training that updates the parameters with SGDm (or Adam for Transformer) but still meta-learns the preconditioning matrix (e.g. performing meta-gradient steps on P ).
When updating the parameters, we scaled the preconditioned gradient by a fixed constant of 0.9 for stability.
C.3 Regression on UCI Collection
We used Slice, Protein, and Parkinson data from the UCI collection.
In training, we normalized the input features and targets to have a zero mean and unit variance.
We used batch size of 128 and trained the network for 500 epochs without weight decay.
We conducted hyperparameter searches over learning rates for all baseline models, making choices based on the final training loss.
With the chosen set of hyperparameters, we repeated the experiments 3 times with different random seeds.
For SGDm, we set the momentum to 0.9 and swept over the learning rates {1, 0.3, 0.1, 0.03, 0.01, 0.003, 0.001, 0.0003, 0.0001}.
For Adam, we performed a grid search on learning rate of {1e-2, 3e-3, 1e-3, 3e-4, 3e-4, 1e-4}.
For Shampoo, we swept over the learning rates {10, 5, 1, 0.3, 0.1, 0.03, 0.01, 0.003, 0.001} as suggested by Gupta et al.
For KFAC, we did a grid search on learning rates of {0.03, 0.01, 0.003, 0.001, 0.0003, 0.0001} and damping values of {3e-2, 1e-2, 3e-3, 1e-3, 3e-4, 1e-4}.
For APO-Precond, we set λ FSD = 0 and grid searched over λ WSD = { 3 , 1 , 0 . 1 , 0 . 01 } .
C.4 Image Reconstruction
We used the same experimental set-up from Martens & Grosse [54] for the deep autoencoder experiment.
The loss function was defined to be the binary entropy and we further added a regularization term λ WD 2 ∥ θ ∥ 2 to the loss function, where λ WD = 10 − 5 .
The layer widths for the autoencoder were set to be [784, 1000, 500, 250, 30, 250, 500, 1000, 784] and we used the sigmoid activation function.
We trained the network for 1000 epochs with the batch size of 512.
We conducted extensive hyperparameter searches for all baseline models, making choices based on the final training loss.
With the chosen set of hyperparameters, we repeated the experiments 3 times with different random seeds.
For SGDm, we set the momentum to 0.9 and performed a grid search over the learning rates {1, 0.3, 0.1, 0.03, 0.01, 0.003, 0.001, 0.0003, 0.0001}.
For Adam, we swept over learning rate of {3e-2, 1e-2, 3e-3, 1e-3, 3e-4, 3e-4, 1e-4, 3e-5, 1e-5}.
For Shampoo, we tried setting learning rates in range of {10, 5, 1, 0.3, 0.1, 0.03, 0.01, 0.003, 0.001}.
For KFAC, we did a grid search on learning rates of {0.03, 0.01, 0.003, 0.001, 0.0003, 0.0001, 0.00003, 0.00001} and damping values of {3e-1, 1e-1, 3e-2, 1e-2, 3e-3, 1e-3, 3e-4, 1e-4}.
For APO-Precond, we performed a grid search over λ FSD = { 0 . 3 , 0 . 1 } and λ WSD = { 3 , 1 , 0 . 3 , 0 . 1 } .
We set the FSD to measure the
KL divergence.
C.5 Image Classification
CIFAR-10 & CIFAR-100.
We used the standard procedure from Zagoruyko & Komodakis [86] for training convolutional neural networks.
The images are zero-padded with 4 pixels on each side and then a random 32 × 32 crop is extracted from the image, horizontally flipped with the probability of 0.5.
We further normalize the inputs with per-channel mean and standard deviations.
We performed the extensive search over the hyperparameters for all networks.
We held out 5k examples from the
CIFAR-10 and CIFAR-100 to form the validation set following Zagoruyko & Komodakis [86] and selected the hyperparameters with the highest validation accuracy.
With the chosen hyperparameters, we re-trained the network with the full training dataset and reported the final test accuracy.
Preconditioning Adaptation.
Across all experiments, we used the batch size of 128 and trained the network for 200 epochs.
With the chosen set of hyperparameters, we repeated the experiments 3 times with different random seeds and reported the mean accuracy on the test dataset.
For SGDm, we set the momentum to 0.9 and performed a grid search over the learning rates {0.3, 0.1, 0.03, 0.01, 0.003, 0.001, 0.0003, 0.0001}.
For Adam, we swept over learning rate of {3e-2, 1e-2, 3e-3, 1e-3, 3e-4, 3e-4, 1e-4}.
For KFAC, we did a grid search on learning rates of {0.03, 0.01, 0.003, 0.001, 0.0003, 0.0001, 0.00003, 0.00001} and damping values of {3e-2, 1e-2, 3e-3, 1e-3, 3e-4, 1e-4}.
For APO-Precond, we performed a grid search over λ FSD = { 0 . 3 , 0 . 1 } and λ WSD = { 3 , 1 , 0 . 3 , 0 . 1 } for architectures without batch normalization.
As batch normalization makes the parameters scale-invariant [3], we imposed a higher regularization in the function space and lower regularization in the weight space, searching over λ FSD = { 3 , 1 } and λ WSD = { 0 . 3 , 0 . 1 , 0 .
We set the FSD term to measure the KL divergence.
We also searched over the weight decay in range of {5e-4, 1e-4, 5e-5} for all optimizers.
Learning Rate Adaptation.
For ResNet32, we trained for 400 epochs, and used a manual schedule that decays the learning rate by a factor of 10 at epochs 150 and 250, following Lucas et al.
For
ResNet34 and WideResNet 28-10, we trained for 200 epochs, with a manual schedule that decays the learning rate by a factor of 5 at epochs 60, 120, and 160, following Zagoruyko & Komodakis [86].
For the baseline optimizers, we performed grid searches over the fixed learning rate or initial learning rate for a fixed step schedule, as well as the weight decay.
For all base optimizers and APO-tuned variants, we searched over weight decay values in { 0 . 01 , 0 . 003 , 0 .
For SGD and SGDm, we searched over learning rates in { 1 . 0 , 0 . 3 , 0 . 1 , 0 . 03 , 0 . 01 , 0 . 003 , 0 . 001 , 0 . 0003 , 0 . 0001 } .
For RMSprop and Adam, we searched over learning rates in { 0 . 1 , 0 . 03 , 0 . 01 , 0 . 003 , 0 . 001 , 0 . 0003 , 0 . 0001 , 0 . 00003 , 0 . 00001 } .
For each of the APO-tuned variants (e.g. SGDm-APO), we kept λ WSD = 0 and searched over λ FSD ∈ { 0 . 3 , 0 . 1 , 0 .
C.6 Neural Machine Translation
We trained a Transformer [79] composed of 6 encoder and decoder layers, with a word embedding and hidden vector dimension of 512.
The architecture has feed-forward size of 1024, 4 attention heads, dropout value 0.3, and weight decay value 0.0001.
For the AdamW baseline, we used a warmup-then-decay learning rate schedule widely used in practice, and for the SGD baseline and
APO-Precond, we kept the learning rate fixed after the warmup.
For APO-Precond, following the practice from Zhang et al.
[87], we used a diagonal block matrix in the structured preconditioner of the embedding weight matrix to reduce memory overhead.
We used the Fairseq toolkit [65] to conduct all experiments.
For SGDm, we grid searched over the learning rates in {10, 3, 1, 0.3, 0.1, 0.03, 0.01, 0.003, 0.001} and tried both the fixed learning rate and inverse sqrt learning rate schedules.
For AdamW, we set β 1 = 0 .
9 , β 2 = 0 .
98 , and ϵ = 10 −
and searched over the learning rates in range {0.05, 0.001, 0.0005, 0.0001}.
AdamW also used the inverse sqrt learning rate schedule with 4000 warmup steps.
For APO-Precond, we searched over λ FSD = { 0 . 3 , 0 . 1 } and λ WSD = { 3 , 1 , 0 . 3 , 0 . 1 } and let FSD term to measure the KL divergence.
As we found a fixed learning rate schedule to work best for SGDm, we also used a fixed learning rate schedule for APO-Precond.
We selected the hyperparameters based on the BLEU score on the validation set and reported on the final test BLEU score on the checkpoint, which achieved the highest validation BLEU score.
D Additional Results D.1 Rosenbrock Function
[Figure]
Both the learning rate schedules and the preconditioner adapted by APO outperform the optimal fixed learning rate for
SGD. We validated APO on the two-dimensional Rosenbrock function defined as:
[Equation]
with initialization ( x, y ) =
(1 , − 1 . 5) .
In Figure 6, we used
APO to tune the learning rate for SGD, as well as the full preconditioning matrix.
Both APO-tuned methods outperformed vanilla SGD with the optimal fixed learning rate (chosen through a careful grid search).
Because Rosenbrock has illconditioned curvature, second-order optimization with APOPrecond dramatically speeds up the convergence and achieves a lower loss.
D.2 SVHN
We also evaluated APO on the Google Street View House Numbers dataset (SVHN) [63], using a
WideResNet 16-8.
We followed the experimental setup of Zagoruyko & Komodakis [86], and used the same decay schedule for the baseline, which decays the learning rate by 10 × at epochs 80 and 120.
The optimal fixed learning rate achieves test accuracy fluctuating around 97.20%, while the manual schedule achieves 98.17%.
APO rapidly converges to test accuracy 97.96%, outperforming the baseline while being slightly worse than the manual schedule (Figure 7).
[Figure]
D.3 CIFAR-
APO Preconditioning Adaptation.
We show the plots for for experiments listed in Table 2 in
Figure 8 and Figure
[Figure]
[Figure]
Furthermore, we show the test accuracy as a function of wall-clock time with various values of in Figure 10.
We observed that APO is robust to the values of K and setting K = 10 , as done in our main experiments, show a significant improvement in the test accuracy over SGDm while only introducing a small computational overhead.
[Figure]
APO Learning Rate Schedules.
We show the plots for each base-optimizer and network architecture as below.
We find that APO achieves better performance than the best fixed LR, and is comparable to the manual schedule.
The learning rate schedules discovered by APO are nontrivial: most exhibit a large increase in the learning rate at the start of training to make rapid progress, followed by a gradual decrease in the learning rate to fine-tune the solution as optimization approaches a local optimum.
The APO learning rate schedules often span two orders of magnitude, similarly to manual decay schedules.
[Table]
[Figure]
[Table] [Table]
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
[Table] [Figure]
D.4 CIFAR-
APO Learning Rate Schedules.
We show training curve plots for the Adam base-optimizer below.
We find that APO achieves better performance than the fixed LR, and is comparable with the manual schedule.
[Figure]
D.5 16-bit Neural Network Training
We show the training curve plots for training 16-bit ResNet-18 on CIFAR-10 and CIFAR-100 in
Figure 22 .
[Figure]
D.6 Comparison to HD, L4, and FOP
In this section [Figure] we compare APO to Hypergradient Descent (HD) [6] and L4 [70] for the learning rate adaptation problem, and First-Order Preconditioning (FOP)
[61] for the preconditioning adaptation problem.
D.6.1 Hypergradient Descent
Here, we compare to Hypergradient Descent (HD) [70], a method for online learning rate adaptation.
The meta-objective used in HD is the expected value of the training loss after updating the learning rate, E [ L ( u ( θ, ϕ )] .
This is the same objective that was analyzed theoretically by Wu et al.
[85], that was shown to lead to sub-optimal behavior when the objective is stochastic.
In particular, [85] showed that this meta-objective encourages the learning rate to drop rapidly, which reduces the expected loss in the short-term by eliminating fluctuations in the loss caused by random mini-batch sampling, at the cost of worse long-term behavior (as discussed in Appendix B).
In our experiments, we found that HD suffers from the short horizon bias issue; the only way to prevent the LR from reducing too quickly is to use a very small meta-learning rate, which becomes more heuristic.
Empirically (and as discussed in more detail in Appendix E), the APO meta-objective does not suffer from this issue, allowing us to use larger meta-learning rates without having the LR decrease too quickly.
Because the meta-learning rate for HD must be small, HD is not robust to the initial learning rate; the initial
LR must be tuned in order to obtain decent performance.
In contrast, APO is very robust to the initial learning rate, behaving nearly identically for initial learning rates spanning 6 orders of magnitude, from 1e-1 to 1e-7 (shown in Appendix J.2).
Comparison.
We ran hypergradient descent versions of SGD, SGDm, and Adam, denoted SGD-HD,
SGDm-HD, and Adam-HD, respectively.
There was no existing implementation of RMSprop-HD, so we left out this optimizer.
For fair comparison to APO, we used the same exponential learning rate parameterization for both APO and HD.
For SGD-HD and SGDm-HD, we performed a grid search over initial learning rates {1e-1, 1e-2, 1e-3}, meta-learning rates {3.0, 1.0, 1e-1, 1e-2, 1e-3}, and weight decay coefficients {1e-2, 1e-3, 5e-4, 1e-4, 0.0}.
For Adam-HD, we performed a grid search over initial learning rates {1e-2, 1e-3, 1e-4}, meta-learning rates {3.0, 1.0, 1e-1, 1e-2, 1e-3}, and weight decay coefficients {1e-3, 5e-4, 1e-4, 0.0}.
On the Meta-Learning Rate for HD.
In our experiments, we found that HD only works decently when run with a large initial learning rate and a small meta-learning rate, which has the effect of reducing the LR slowly rather than dramatically (as is expected from the analysis in [85]).
However, using a larger meta-learning rate, which optimizes the meta-objective more rapidly, exhibits the short horizon bias issue.
In Figures 23,25, we plot the learning rates adapted by HD with a small meta-learning rate (0.001) and a large meta-learning rate (1.0), for each of SGD-HD, SGDm-HD, and Adam-HD.
We also plot the corresponding test accuracies for each of these experiments in
Figures 24,26, to show that the LR adapted using a large meta-learning rate catastrophically affects training and leads to substantially worse test accuracy.
In summary, these results provide evidence that the HD meta-objective is not a very good target for meta-optimization.
[Table] [Table]
[Figure]
[Table] [Figure] [Figure]
[Figure]
D.6.2 Comparison to L
We also compared APO to the L4 optimizer [70], which adapts a global learning rate in an online fashion based on Polyak’s update rule.
We applied L4 to tune the global learning rate for four base optimizers—vanilla SGD without momentum, SGD with momentum coefficient 0.9 (denoted SGDm),
RMSprop, and Adam—used to train ResNet32 and ResNet34 on CIFAR-10.
For each experiment, we performed a grid search over weight decay coefficients {1e-3, 5e-4, 1e-4, 0.0} and α values {0.1, 0.15, 0.3}, where α controls the scaling of the step size, and is the primary hyperparameter to tune for L4.
For the other L4 hyperparameters, we used the defaults γ = 0 .
9 and τ = 0 .
The results are shown in Table 11 and Table 12.
We found that L4 obtained poor generalization, underperforming the other baselines and APO.
Figures 27,28 and Figures 29,30 show the test accuracy and learning rate adaptation obtained with each L4-tuned optimizer, on CIFAR-10 classification with ResNet32 and ResNet34.
[Figure] [Figure]
[Figure] [Figure]
D.6.3 First-Order Preconditioning
[Figure]
[Table]
Image Reconstruction.
Following Section 6, we trained an 8-layer autoencoder on MNIST and compared APO with FOP using the same experimental set-up from Appendix C.4.
For FOP, we set the momentum to 0.9 and performed grid search over the learning rates and hypergradient learning rates in range of { 1 , 0 . 3 , 0 . 1 , 0 .
and { 10 − 2 , 10 − 3 , 10 − 4 , 10 − 5 , 10 − 6 } , respectively.
We show the training losses in Figure 31.
APO-Precond converges faster than FOP and achieves competitive training loss to KFAC (although there remains some performance gap with KFAC).
[Figure]
Image Classification.
We also compared APO with FOP on image classification task with the same set-up from Appendix C.5.
For FOP, we fixed the momentum to 0.9 and again performed grid search over the learning rates and hypergradient learning rates in { 1 , 0 . 3 , 0 . 1 , 0 . 03 , 0 . 01 , 0 . 003 , 0 . 001 , 0 . 0003 , 0 . 0001 } and { 10 − 2 , 10 − 3 , 10 − 4 , 10 − 5 , 10 − 6 } , respectively.
We also searched over the weight decay in range of {5e-4, 1e-4, 5e-5}.
The test accuracies of APO-Precond and FOP are summarized in Table 13.
While APO-Precond achieved competitive generalization performance to SGDm and KFAC on all tasks, FOP consistently underperformed compared to KFAC.
D.7 Tuning Other Optimization Parameters
Here, we show proof-of-concept results for tuning different optimizer parameters, other than the learning rate or preconditioner.
In particular, we focus on two parameters of the RMSprop optimizer that are rarely tuned by hand.
The RMSprop update rule is:
[Equation]
We consider adapting ϵ and ρ (the power to which s t is raised in the denominator, which is fixed to ρ = 12 in standard RMSprop).
Both ϵ and ρ can be interpreted as having a damping effect on the update.
We investigated the effect of tuning { ϵ, ρ } using APO on two tasks: 1) training a two-hidden-layer MLP with 1000 hidden units per layer on FashionMNIST; and 2) training a
ResNet32 on CIFAR-10.
We set the learning rate to the best fixed value for the baseline: η = 0 .
001 for CIFAR-10 and η = 0 .
0001 for FashionMNIST.
We used RMSprop with learning rate 0.01 for meta-optimization, and we performed a grid search over λ FSD ∈ { 1e-3 , 1e-4 , 1e-5 } and λ WSD ∈ { 1e-1 , 1e-2 , 1e-3 , 1e-4 , 1e-5 , 1e-6 } .
We compared to baseline RMSprop with its optimal fixed learning rates ( η = 0 .
001 for CIFAR-10 and η = 0 .
0001 for FashionMNIST) and with the default values ϵ = 1e-8 and ρ = 0 .
The results for FashionMNIST and CIFAR-10 are shown in
Figures 32 and 33, respectively.
We found that on each of these tasks, tuning { ϵ, ρ } resulted in faster training and lower final training loss.
[Figure]
[Figure]
E Importance of Using the Same Mini-batch
In this section
[Figure]
we verify empirically the importance of computing the first term of the meta-objective on the same mini-batch that is used to compute the gradient of the base optimizer.
Recall that our meta-objective is:
[Equation]
[Equation] Note that the loss term is computed using the same mini-batch B used for the update u ( θ , ϕ , B ) .
An alternative is to evaluate the loss term on a randomly-sampled mini-batch B ′ ∼ D , yielding:
[Equation]
If we set λ FSD = λ
WSD = 0 , then Eq.
14 becomes the meta-objective analyzed by Wu et al.
[85], which was used to illustrate the short-horizon bias issue in stochastic meta-optimization.
Figures 34,35 show the result of using Eq. 13 vs. Eq.
14 for meta-optimizing the learning rate of a
ResNet32 model on CIFAR-10.
We observe that when using a random mini-batch B ′ to compute the loss term, the learning rate decays rapidly, preventing long-term progress; in contrast, using the same mini-batch B for the loss term yields a reasonable LR schedule which does not suffer from the short-horizon bias issue empirically.
[Figure]
Same Mini-Batch for the FSD Term.
Another alternative to the meta-objective is to compute the
FSD term using the same mini-batch used to compute the loss term, as shown in Eq.
[Equation]
[Figure] Recall that D F ( u ( θ , ϕ , B ) , θ , ˜ x ) =
ρ ( f (˜ x , u ( θ , ϕ , B )) , f (˜ x , θ )) .
The function evaluation f (˜ x , u ( θ , ϕ , B )) must be performed to compute the loss term; thus, using the same mini-batch
B for the FSD term can allow us to reduce computation by re-using this function output in both the loss and FSD computations.
However, with the interpretation of the FSD term in Eq. 15 as a
Monte Carlo estimate of the expectation E ˜ x
∼D [ D F ( u ( θ , ϕ , B ) , θ , ˜ x )] , using the same mini-batch for the loss and dissimilarity would yield a biased estimate.
The effect of using a different vs the same mini-batch to compute the FSD term is shown in Figure [Figure]
F Approximate Closed-Form Solution for the PPM
In Section 3.1, we introduced a general update rule for the stochastic PPM, which is defined as follows:
[Equation]
We first take the infinitesimal limit by letting λ FSD → ∞ .
Then, the optimal update u ⋆ will stay close to the current parameters θ ( t ) and we can approximate the loss function with the first-order Taylor approximation.
Moreover, since the FSD term E ˜ x ∼D [ D F ( u , θ ( t ) , x )] is minimized when u = θ ( t ) , we have ∇ u E ˜ x
∼D [ D F ( u , θ ( t ) , x )] | u
= θ ( t ) = 0 .
Hence, the second-order Taylor expansion of the FSD term would be:
[Equation]
(17) where G = ∇ 2 u E ˜ x
∼D [ D F ( u , θ ( t ) , x )] | u
= θ ( t ) is the local Hessian of the FSD term.
We further let the weight space discrepancy function to be the squared Euclidean distance D W ( u , θ ( t ) )
= 1 / 2 ∥ u − θ ( t ) ∥ 22 .
Combining these insights, we approximate Eqn. 16 by linearizing the loss and taking a quadratic approximation to the FSD term:
[Equation]
Taking the gradient with respect to u and setting it equal to 0 , we have: [Equation]
[Equation]
Rearranging the terms, we arrive at the approximate closed-form solution:
[Equation]
G Optimization Algorithms Derived from the Proximal Objective
In this section, we show how several classical optimization algorithms can be derived from the proximal objective introduced in Section 3.2.
Given a mini-batch B , recall that we consider the following proximal objective for optimizing the model parameters θ :
[Equation]
Method
Loss Term
[Equation]
FSD Term Approx. D
[Equation]
FSD Term Choice of D F
WSD Term 12 || u − θ ||
Gradient Descent 1 st -order - - ✓
Newton’s Method 2 nd -order - - ✗
Damped Newton’s Method 2 nd -order - - ✓
Natural Gradient 1 st -order 2 nd -order
KL ✗
Damped Natural Gradient 1 st -order 2 nd -order
KL ✓
Generalized Gauss-Newton 1 st -order 2 nd -order
Bregman ✗
Damped Generalized Gauss-Newton 1 st -order 2 nd -order
Bregman ✓
Direct Proximal Optimization
Exact
Exact
Any ✓
G.1 Gradient Descent
To derive vanilla gradient descent, we take the first-order Taylor series approximation to the loss term
J B ( θ ) about θ ( t ) , and set λ FSD = 0 .
We need to have λ
WSD > 0 to ensure that we do not take an infinitely large step.
Thus, each step of gradient descent minimizes an approximation to the proximal objective which we denote by P GD θ ( t ) ( θ ) :
[Equation]
where:
[Equation]
Now, we set the gradient ∇ θ P GD θ ( t ) ( θ ) = 0 and solve for θ :
[Equation]
G.2 Newton’s Method
To derive Newton’s method, we take the second-order Taylor series approximation to the loss term
J B ( θ ) about θ ( t ) , and set λ FSD = 0 .
Below we derive the general form for the damped Newton update, which incorporates the WSD term; setting λ
WSD
= 0 leads to the undamped Newton update as a special case.
Each step of Newton’s method minimizes an approximation to the proximal objective which we denote by P Newton
[Equation] [Equation]
where:
[Equation]
In Eq.
[Equation] H = ∇ 2 θ J B ( θ ) .
Now [Equation] we set the gradient ∇ θ P Newton θ ( t )
[Equation]
0 and solve for θ :
[Equation]
G.3 Generalized Gauss-Newton Method
Next
[Equation] we consider taking the first-order Taylor series approximation to the loss term [Equation] and an arbitrary function-space discrepancy D F approximated by its second-order Taylor series expansion.
We will use the notation z
= f ( x , θ ) to denote the outputs of a neural network with parameters θ on inputs x .
We have the following proximal objective:
[Equation]
Taking the second-order approximation to D F (with respect to θ ) [Equation] we have:
[Equation]
Using the chain rule [Equation] we can derive ∇ 2 θ D F ( θ , θ ( t ) )
| θ = θ ( t ) as follows:
[Equation]
Letting G denote the expectation of the Hessian of the FSD function
[Equation] G = E ˜ x ∼D [ ∇ 2 θ D F ( θ , θ ( t ) , ˜ x )] , we have: [Equation]
[Equation]
If we set the gradient ∇ θ P θ ( t ) ( θ ) = 0 and solve for θ
[Equation]
we obtain:
[Equation]
[Equation]
H Optimizing the 1-Step Meta-Objective Recovers Classic Methods
In this section, we show that when we approximate the loss term and FSD term of our meta-objective
Q ( ϕ ) and solve for the analytic arg min ϕ Q ( ϕ ) , we recover the preconditioners corresponding to classic first- and second-order algorithms, including gradient descent, Gauss-Newton, Generalized
Gauss-Newton, and natural gradient.
The preconditioners used by each of these methods are shown in Table 1.
These results parallel those for directly optimizing the parameters θ using the proximal objective, but optimizing for the meta-parameters ϕ = P .
Throughout this exposition, we use the notation:
[Equation]
where g = ∇ θ J B ( θ ) .
Note that g implicitly depends on the data B ; we use this shorthand to simplify the exposition.
Theorem.
Consider an approximation ˆ Q ( P ) to the meta-objective (Eq. 8) where the loss term is linearized around the current weights θ and the FSD term is replaced by its second-order approximation around θ :
[Equation]
Denote the gradient on a mini-batch as g = ∇ θ J B ( θ ) , and assume that the second moment matrix
E B∼D (cid:2) gg ⊤ (cid:3) is non-singular.
Then, the preconditioning matrix which minimizes ˆ Q is given by
P [Equation]
− 1 , where G denotes the Hessian of the FSD evaluated at θ .
Proof.
We have the following meta-objective (where the loss term is expanded using its first-order
Taylor series approximation):
[Equation]
Let us first derive the second-order Taylor series approximation to the discrepancy term
D F ( u ( θ , ϕ , B ) , θ , ˜ x ) .
To simplify notation, let d ( θ ′ )
≡ D F ( θ ′ , θ , ˜ x ) where θ and the data ˜ x are implicit.
Then, the second-order expansion of d about θ is:
[Equation]
We have d ( θ ) =
D F ( θ , θ , ˜ x ) =
ρ ( f (˜ x , θ ) , f (˜ x , θ ))
= 0 , and ∇ d ( θ ) = 0 because θ is a minimum of d , so the first two terms are 0.
Denote the network output on example ˜ x by z = f (˜ x , θ ) .
To compute ∇ 2 θ d ( θ ) , we have:
[Equation]
To simplify notation [Equation] we let G denote the expectation of the Hessian of the FSD function [Equation]
G ≡ E ˜ x ∼D [ ∇ 2 θ d ( θ ) ].
Plugging this into the meta-objective, we have:
[Equation]
Next, we take the gradient with respect to P and set it to 0 to solve for P :
[Equation]
Assuming that the second moment matrix E B∼D [ gg ⊤ ] is non-singular, we have:
[Equation]
Thus, P ⋆ = ( λ FSD G + λ WSD I ) − 1 .
Below, we provide detailed derivations for several special cases (e.g. specific assumptions on the loss term and FSD term) which yield gradient descent, damped Newton’s method, and natural gradient descent.
H.1 Gradient Descent
For gradient descent, we consider the first-order Taylor series approximation of the loss term about θ and set λ FSD = 0 .
We have:
[Equation]
Next, we take the gradient with respect to P and set it to 0 to solve for P ⋆ :
[Equation]
Assuming that the second moment matrix E B∼D [ gg ⊤ ] is non-singular [Equation] we have:
[Equation]
Thus [Equation] the optimal preconditioner is P =
λ WSD I , which corresponds to using a global learning rate of 1 λ WSD , as we would expect from steepest descent with the WSD term.
H.2 (Damped) Newton’s Method
For Newton’s method, we consider the second-order Taylor series expansion of the loss term about θ , and set λ FSD = 0 .
Here, we consider the generic case where we have λ WSD ≥ 0 : if λ WSD > 0 we obtain a damped Newton update, while if λ WSD = 0 we obtain the undamped update.
We have:
[Equation]
Next
[Equation]
we take the gradient with respect to P and set it to 0 to solve for P ⋆ :
[Equation]
[Equation]
Assuming that the second moment matrix E B∼D [ gg ⊤ ] is non-singular [Equation] we have:
[Equation]
[Equation] Thus, if we take the second-order Taylor series approximation to the loss term and solve for the optimal P ⋆ through the 1-step lookahead θ ′ ( P ) , we obtain the (damped) inverse Hessian, P ⋆ = ( H + λ WSD I ) − 1 .
H.3 Natural Gradient
For natural gradient, we consider a first-order approximation to the loss term and a second-order approximation to the FSD term, where the output-space discrepancy measure ρ is taken to be the KL divergence.
The second derivative of the KL divergence yields the Fisher information matrix F ; here, we write
F = J ⊤ z θ H
ρ J z θ .
Plugging this into the meta-objective, we have:
[Equation]
[Equation]
Next, we take the gradient with respect to P and set it to 0 to solve for P ⋆ :
[Equation]
[Equation] Assuming that the second moment matrix E ( x , t ) ∼D [ gg ⊤ ] is non-singular, we have:
[Equation]
I Optimizing the 1-Step Meta-Objective Yields the KFAC Update
In the previous section, we derived the optimal solutions to various approximate proximal objectives, optimizing over an unconstrained preconditioner P .
Here, we consider an analogous setup for a structured preconditioner, where P is block-diagonal with the ℓ th block (corresponding to the ℓ th layer in the network) given by A ℓ ⊗ B ℓ .
I.1 KFAC Assumptions
This exposition of the KFAC assumptions is based on [25].
First we briefly introduce some notation that will be used in the assumptions.
Suppose a layer of a neural network has weights W ℓ and that its input is the activation vector from the previous layer, a ℓ − 1 .
Then the output of the layer is a linear transformation of the input, followed by a nonlinear activation function σ :
[Equation]
73) where s ℓ are the pre-activations.
In the backward pass, we have the following activation and weight derivatives:
[Equation]
KFAC [54] makes the following assumptions on the neural network being optimized:
• The layers are independent, such that that the pseudo-derivatives dw i and dw j are uncorrelated when w i and w j belong to different layers.
If this assumption is satisfied, then the Fisher information matrix F will be block-diagonal.
• The activations { a ℓ } are independent of the pre-activation pseudo-gradients {D s ℓ } .
If a ℓ −is independent of D s ℓ , then the ℓ th block of the Fisher is:
[Equation]
where:
[Equation]
I.2 Proof of Corollary
Corollary.
Suppose that (1) the assumptions for Theorem 1 are satisfied, (2) the FSD term measures the KL divergence, and (3) λ
WSD
=
and λ FSD = 1 .
Moreover, suppose that the parameters θ satisfy the KFAC assumptions listed in Appendix I.
Then, the optimal solution to the approximate metaobjective recovers the KFAC update, which can be represented using the structured preconditioner in
Eq. 9.
Proof.
If the KFAC assumptions are satisfied, then F − 1 is block-diagonal with the block corresponding to the ℓ th layer expressed as G ℓ =
A − 1 ℓ ⊗ B − 1 ℓ where A ℓ =
E [¯ a ℓ − 1 ¯ a ⊤ ℓ − 1 ] and
B ℓ = E
[ D s ℓ D s ⊤ ℓ ] .
By Theorem 1, the optimal solution to the approximate meta-objective is
P ⋆ = F − 1 .
Hence P ⋆ is this block-diagonal matrix, which can be expressed using our structured parameterization (Eq. 9).
Thus, APO recovers the KFAC update.
J Ablations J.1 Ablation Over λ WSD and λ FSD
Figure 37 provides an ablation over λ FSD for training an MLP on MNIST using APO to tune the global learning rate for RMSprop (e.g. RMSprop-APO).
[Figure]
Figure 38 provides an ablation over λ WSD and λ FSD for training AlexNet on CIFAR-10 using APO to adapt the preconditioning matrix.
We kept the other proximity weight fixed when performing the experiments.
We found that APO-Precond is robust in various ranges of proximity weights λ
[Figure]
J.2 Robustness to Initial Learning Rate and Meta-Update Interval
We show in Figure 39(a,b) that when using APO to tune the global learning rate, APO is robust to the initial learning rate of the base optimizer, achieving nearly identical training loss and test accuracy using initial learning rates that span 6 orders of magnitude, from 1e-7 to 1e- [Figure]
[Figure]
Next
[Figure]
we performed an ablation to evaluate how well APO performs with different meta-update intervals (i
[Figure]
e. when making more or less frequent updates to the learning rate and preconditioning matrix during training).
We used APO to tune the learning rate
[Figure]
and experimented with performing meta-updates once every 10, 20, 50, and 100 base optimization iterations. We trained a ResNet model on CIFAR-
[Figure]
and used SGDm as the base optimizer.
The results are shown in Figure 40.
We found that APO is robust to the meta-update interval, performing almost identically with respect to training loss, test accuracy
[Figure]
and the adaptive learning rate schedule, for each meta-update interval
[Figure]
To further examine how meta-update intervals affect the preconditioning adaptation we trained
AlexNet model on CIFAR-10.
We used APO to learn the preconditioning matrix and experimented with performing meta-updates once every 10, 20, 50, and 100 base optimization iterations.
The results are shown in Figure 41.
As in the learning rate adaptation problem, we found that APO performs similarly using intervals from 10 to 100, achieving higher test accuracy than the baseline method.
[Figure]