Amortized Proximal Optimization
Anonymous Author(s)
AffiliationAddress email
Abstract
We propose a framework for online meta-optimization of parameters that governoptimization, called Amortized Proximal Optimization (APO).
We first interpretvarious existing neural network optimizers as approximate stochastic proximalpoint methods which trade off the current-batch loss with proximity terms inboth function space and weight space.
The idea behind APO is to amortize theminimization of the proximal point objective by meta-learning the parameters of anupdate rule.
We show how APO can be used to adapt a learning rate or a structuredpreconditioning matrix.
Under appropriate assumptions, APO can recover existingoptimizers such as natural gradient descent and KFAC.
It enjoys low computationaloverhead and avoids expensive and numerically sensitive operations required bysome second-order optimizers, such as matrix inverses.
We empirically test APOfor online adaptation of learning rates and structured preconditioning matricesfor regression, image reconstruction, image classification, and natural languagetranslation tasks.
Empirically, the learning rate schedules found by APO generallyoutperform optimal fixed learning rates and are competitive with manually tuneddecay schedules.
Using APO to adapt a structured preconditioning matrix generallyresults in optimization performance competitive with second-order methods.
Introduction
Many optimization algorithms widely used in machine learning can be seen as approximations toan 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 aproximity term that penalizes the discrepancy from the current iterate: [Equation]
[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 excessivelychanging the parameters, hence preventing aggressive updates.
Moreover, the stochastic PPM hasgood convergence properties [4].
While minimizing Eq.
1 exactly is usually impractical (or at leastuneconomical), solving it approximately (by taking first or second-order Taylor series approximationsto the loss or the proximity term) has motivated important and widely used optimization algorithmssuch as natural gradient descent [1] and mirror descent [6].
Stochastic gradient descent (SGD) [68]itself can be seen as an approximate PPM where the loss term is linearized and the discrepancyfunction is squared Euclidean distance. Minimization of the exact proximal objective in Eq.
1 is un-economical, as it may require tens tohundreds of optimization steps per parameter update. Instead, we propose to learn a parametricupdate rule θ ( t +1) ← u ( θ ( t ) , B ( t ) , ϕ ) described by meta-parameters ϕ , and use a variant of Eq. as a meta-objective to optimize ϕ .
By meta-learning the optimization parameters online, we canamortize the cost of minimizing the PPM objective, by performing one meta-step every K > 1steps of the base optimization.
We consider adapting optimization hyperparameters (such as thelearning rate) for existing optimizers such as SGD and RMSprop [77], as well as learning structuredpreconditioning 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 caseswhere their assumptions (such as the use of linear or quadratic approximations) break down.
Bydoing meta-descent on the optimization parameters, we can amortize the cost of minimizing the PPMobjective, 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 effectiveneural network optimizers can be seen as trading off two different proximity terms: a functionspace 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 singleoptimization hyperparameter to learning a flexible update rule.
We consider two use cases that coverboth ends of this spectrum.
At one end, we consider the problem of adapting learning rates of existingoptimizers, specifically SGD, RMSprop, and Adam [18].
The learning rate is considered one of themost essential hyperparameters to tune [9], and good learning rate schedules are often found by yearsof trial and error.
Empirically, the learning rate schedules found by APO outperformed the best fixedlearning 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 theupdate 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 ofour APO meta-objective with respect to a full preconditioning matrix coincides with damped versionsof natural gradient descent or Gauss-Newton.
While computing and storing a full preconditioningmatrix 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 [23].
APO is more straightforward to implement in current-day deep learning frameworks than EKFACand 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 preconditioningmatrix adapted by APO achieved competitive convergence to other second-order optimizers.
Proximal Optimization and 2 nd -Order Methods: A Unifying Framework
We first motivate the proximal objective that we use as the meta-objective for APO, and relate it toexisting neural network optimization methods.
Our framework is largely based on Grosse [24], towhich readers are referred for a more detailed discussion.
Preliminaries.
Consider a prediction problem with a finite training set D train =
{ ( x ( i ) , t ( i ) ) }
Ni =1 .
For a data point ( x , t ) and parameters θ ∈ R m , let y =
f ( x , θ ) be the prediction of a networkparameterized by θ and L ( y , t ) be the loss.
Our goal is to minimize the cost function J ( θ )
=1 N (cid:80) Ni =
1 L ( f ( x ( i ) , θ ) , t ( i ) ) .
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.
2.1 Proximal Optimization
When we update the parameters on a mini-batch of data, we would like to reduce the loss on thatmini-batch, while not changing the predictions on previously visited examples or moving too far inweight space.
This motivates the following proximal point update:
[Equation]
[Figure]
where D F ( u , θ , x ) =
ρ ( f ( x , u ) , f ( x , θ )) and D W ( u , θ ) = 1 / 2 ∥ u − θ ∥ 22 are discrepancy functions(described below), λ FSD and λ WSD are hyperparameters that control the strength of each discrepancyterm, ˜ x is a sample from the data-generating distribution D , and ρ ( · , · ) is the output-space divergence.
The proximal objective in Eq.
2 consists of three terms.
The first term is the loss on the currentmini-batch.
The second term is the function space discrepancy (FSD) , whose role is to preventthe update from substantially altering the predictions on other data points.
The general idea of the
FSD term has been successful in alleviating catastrophic forgetting [10], fine-tuning pre-trainedmodels [34], and training a student model from a teacher network [29].
The final term is the weight space discrepancy (WSD) , which encourages the update to move theparameters as little as possible.
It can be used to motivate damping in the context of second-orderoptimization [52].
While weight space distance may appear counterproductive from an optimizationstandpoint because it depends on the model parameterization, analyses of neural network optimizationand generalization often rely on network parameters staying close to their initialization in the
Euclidean norm [32, 87, 8].
In fact, Wadia et al.
[82] 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 regressionexample 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 currentexample, 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 currentdata 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 onewhich nonetheless improves performance on nearby examples.
2.2 Connection Between Proximal Optimization and Second-Order Optimization
Method
Loss Approx.
FSD
WSD
Gradient Descent 1 st -order - ✓
Hessian-Free 2 nd -order - ✓
Natural Gradient 1 st -order 2 nd -order ✗
Proximal Point
Exact
Exact ✓
We further motivate our proximal objective byrelating it to existing neural network optimizers.
Ordinary SGD can be viewed as an approximate
PPM update with a first-order approximation tothe loss term and no FSD term.
Hessian-Free optimization [50], a classic second-order optimization method for neural networks, approximatelyminimizes a second-order approximation to theloss on each batch using conjugate gradients.
Itcan be seen as minimizing a quadratic approximation to Eq.
2 with no FSD term.
Amari [1] motivatednatural 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 proximityterm.
Optimizing over a manifold of probability distributions with KL divergence as the proximityterm yields the familiar update involving the Fisher information matrix.
Natural gradient optimizationof neural networks [51, 53] can be interpreted as minimizing Eq.
2 with a linear approximation tothe 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 neuralnetwork optimizers for stability [53].
[Figure]
Algorithm
Amortized Proximal Optimization (APO) — Meta-Learning Optimization Parameters ϕ
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 θ ′ ( ϕ ) :
= u ( θ , ϕ , B ) ▷ Compute the 1-step lookahead parameters
Q ( ϕ ) :
= J B ( θ ′ ( ϕ ))
+ λ FSD / |B ′
| (cid:80) (˜ x , · ) ∈B ′ D F ( θ ′ ( ϕ ) , θ , ˜ x ) + λ WSD / 2 ∥
θ ′ ( ϕ ) − θ ∥▷ Compute meta-objective ϕ ← ϕ − α ∇ ϕ Q ( ϕ )
▷ Update optimizer parameters (e.g. LR or preconditioner) end if θ ← u ( θ , ϕ , B ) ▷ Update model parameters end while
In a more general context, when taking a first-order approximation to the loss and a second-orderapproximation 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 theserelationships 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 onlinemeta-learning of optimization parameters.
Then, we describe two use cases that we explore in thepaper: (1) adapting learning rates of existing base optimizers such as SGD, RMSProp, and Adam,and (2) meta-learning a structured preconditioning matrix.
3.1 Proximal Meta-Optimization
We assume an update rule u parameterized by a vector ϕ which updates the network weights θ on abatch B ( t ) : 1 θ ( t +1) ← u ( θ ( t ) , ϕ , B ( t ) ) .
One use case of APO is to tune the hyperparameters of anexisting optimizer, in which case ϕ denotes the hyperparameters.
For example, when tuning the SGDlearning rate, we have ϕ =
η
and the update is given by: u SGD ( θ , η, B )
= θ − η ∇ θ J B ( θ ) .
Moreambitiously, we could use APO to adapt a full preconditioning matrix P .
In this case, we defineϕ =
P and the update is given by: u Precond ( θ , P , B ) =
θ − P ∇ θ J B ( θ ) .
In Section 2, we introduced a general proximal objective for training neural networks and observedthat many optimization techniques could be seen as an approximation of PPM.
Motivated by thisconnection, we propose to directly minimize the proximal objective with respect to the optimizationparameters.
As shown in Table 1, under linear and quadratic assumptions on the loss term and FSDterm, the proximal objective in Eq. 1 has a closed-form solution; however, when these assumptionsdo not hold, direct minimization of the proximal objective can still be robust.
Adapting a parametricupdate rule allows us to take advantage of the properties of existing optimizers, and also amortizesthe cost of minimizing the proximal objective over the course of training.
We propose to use the following meta-objective, which evaluates the proximal objective at u ( θ , ϕ , B ) : [Equation]
[Equation]
In practice, we estimate the expectations in the meta-objective by sampling two different minibatches, B and B ′ , where B is used to compute the gradient and the loss term, and B ′ is used tocompute the FSD term.
Intuitively, this proximal meta-objective aims to find optimizer parametersϕ that 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 hasbeen made by other mini-batches.
The optimization parameters ϕ are optimized with a stochasticgradient-based algorithm (the meta-optimizer ).
The meta-gradient ∇ ϕ Q ( ϕ ) can be computed viaautomatic differentiation through the one-step unrolled computation graph (Figure 2).
We refer to ourframework as Amortized Proximal Optimization (APO, Algorithm 1).
3.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-objectivewith respect to the learning rate η .
Although adaptive optimizers such as RMSProp and Adam usecoordinate-wise learning rates, they still have a global learning rate which is essential to tune.
APOcan be applied to such global learning rates to find learning rate schedules.
3.3 APO for Adaptive Preconditioning
More ambitiously, we can use the APO framework to adapt the preconditioning matrix, allowing theupdate 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 2.2) hold, then the optimal preconditioning matrix is equivalent to varioussecond-order updates, depending on the choice of the FSD function.
Theorem 1.
Consider an approximation ˆ Q ( P ) to the meta-objective (Eq. 4) where the loss term islinearized 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 secondmoment matrix E B∼D (cid:2) gg ⊤ (cid:3) is non-singular.
Then, the preconditioning matrix which minimizes ˆ Q isgiven 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 derivedfrom 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 themodel’s predictive distribution over y .
Therefore, the optimal preconditioner is the damped naturalgradient preconditioner, P ⋆ =
( F + λ WSD I ) − 1 .
Thus, when APO is used to exactly minimize anapproximate meta-objective, the update it yields coincides with classical second-order optimizationalgorithms, depending on the choice of the FSD function.
3.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 bepositive semidefinite (PSD) so that the transformed gradient is a descent direction [63].
To satisfy these requirements, we adopt a structured preconditioner analogous to that of the EKFACoptimizer [23].
Given a weight matrix W ∈ R m i × m i +1 of a layer, we construct the preconditioningmatrix as a product of smaller matrices:
[Equation]
where A ∈ R m [-i +
 × m i-][Equation]
[Equation]
+1 , B ∈ R m i × m i , and S ∈ R m i × m i [-+
-]+1 are small block matrices.
Here, ⊗denotes the Kronecker product, diag ( · ) denotes the diagonalization operator, and vec ( · ) denotes thevectorization operator.
This parameterization is memory efficient: it requires m 2 i + m
i +
+ m i m i
+1parameters to store, as opposed to m
i m
i +1 parameters for a full preconditioning matrix.
It isstraightforward to show that the structured preconditioner in Eq.
5 is PSD, as it takes the form
CDC ⊤ , where D is PSD.
The preconditioned gradient can be computed efficiently by using theproperties of the Kronecker product: P S vec ( ∇ W J B ( θ ))
= vec ( B ( S 2 ⊙ B ⊤ ∇ W J B ( θ ) A )
A ⊤ ) ,where ⊙ denotes elementwise multiplication.
This is tractable as it only requires four additionalmatrix multiplications and elementwise multiplication of small block matrices in each layer whenupdating the parameters.
While EKFAC uses complicated covariance estimation and eigenvaluedecomposition 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 performingeigendecompositions of) the block matrices, our structured representation incurs less computationper iteration than EKFAC.
While we defined an optimizer with the same functional form as EKFAC, it is not immediatelyobvious 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 assumptionsunderlying KFAC [53], then the structured preconditioner minimizing Eq.
4 coincides with KFAC:
Corollary 2.
Suppose that (1) the assumptions for Theorem 1 are satisfied, (2) the FSD term measuresthe 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 approximatemeta-objective recovers the KFAC update, which can be represented using Eq. 5.
The proof is in Appendix I. If the KFAC assumptions are not satisfied, then APO will generally learna different preconditioner.
This may be desirable, especially if differing probabilistic assumptionslead to different update rules, as is the case for KFAC applied to convolutional networks [25, 37].
3.5 Computation and Memory Cost
Computation Cost.
Computing the FSD term requires sampling an additional mini-batch fromthe 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 steprequires three forward passes to compute the proximal objective.
It additionally requires a backwardpass through the 1-step unrolled computation graph (Figure 2) to compute the gradient of the proximalmeta-objective Q ( ϕ ) with respect to ϕ .
This overhead can be reduced by performing a meta-updateonly once every K iterations: the overhead will consist of 3 /K additional forward passes and 1 /Kadditional 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 computingthe proximal meta-objective.
In the context of structured preconditioner adaptation, we further needto store block matrices A , B , and S (in Eq. 5) for each layer, as in KFAC and EKFAC.
Related Work
Gradient-Based Learning Rate Adaptation.
Maclaurin et al.
[49] backpropagate through the fullunrolled training procedure to meta-optimize learning rate schedules offline .
This is expensive, as itrequires completing a full training run to make a single hyperparameter update.
A related approach isto unroll the optimization for a small number of steps and perform truncated backpropagation [15, 21].
Hypergradient descent [5] adapts the learning rate to minimize the expected loss in the next iteration.
Second-Order Optimization.
Although preconditioned methods have better convergence ratesthan first-order methods [7, 63], storing and computing the inverse of preconditioning matrices isimpractical for high-dimensional problems.
To mitigate these computational issues, Hessian-freeoptimization [50, 54] 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 [53, 25, 55, 23, 26, 76], a diagonal matrix [18, 35], or a low-rank matrix [38, 59].
However, these approximate second-order methods may not be easy to implement in deep learningframeworks, and can still be expensive as they often require matrix inversion or eigendecomposition.
Gradient-Based Preconditioner Adaptation.
Moskovitz et al.
[60] learn the preconditioningmatrix with hypergradient descent.
Meta-curvature [65] and warped gradient descent [40, 20] adaptthe preconditioning matrix that yields effective parameter updates across diverse tasks in the contextof few-shot learning.
In Appendix B, we provide extended related work, and we provide a conceptualcomparison of meta-optimization methods in Table 6.
Experiments
Our experiments investigate the following questions: (1) How does the structured preconditioningmatrix 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 andmanual decay schedules commonly used in the literature?
We used APO to meta-learn the preconditioning matrices for a broad range of tasks, including severalregression datasets, autoencoder training, image classification on CIFAR-10 and CIFAR-100 usingseveral network architectures, neural machine translation using transformers, and low-precision(16-bit) training.
Several of these tasks are particularly challenging for first-order optimizers.
Inaddition, we used APO to tune the global learning rate for multiple base optimizers – SGD, SGD withmomentum (denoted SGDm), RMSprop, and Adam – on CIFAR-10/100 classification with severalnetwork architectures.
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.g. SGDm-APO).
Experiment details and additional experiments, including ablationstudies, are given in Appendix C and J, respectively.
5.1 Meta-Learning Preconditioners
Poorly-Conditioned Regression.
We fist considered a regression task traditionally used to illustratefailures of neural network optimization [67].
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 xand 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 showtraining curves for SGDm, Adam, Shampoo [26], KFAC, and APO-Precond.
As the problem isill-conditioned, 2 nd -order optimizers such as Shampoo and KFAC converge faster than 1 st -ordermethods.
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 datasetsfrom the UCI regression collection [17].
We trained a 2-layer MLP with 100 hidden units perlayer and ReLU activations for 500 epochs.
The training curves for each optimizer are shown in
Figure 3.
We found that APO-Precond consistently achieved competitive convergence compared toother second-order optimizers and reached lower training loss than all baselines.
[Figure]
[Figure]
Image Reconstruction.
We trained an 8-layer autoencoder on MNIST [39]; this is known to be achallenging optimization task for first-order optimizers [30, 54, 53].
We followed the experimentalset-up from Martens & Grosse [53], where the encoder and decoder consist of 4 fully-connectedlayers with sigmoid activation.
The decoder structure is symmetric to that of the encoder, and they donot have tied weights.
The logistic activation function and the presence of a bottleneck layer makethis 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 lossesfor each optimizer in Figure 4 (middle).
APO-Precond converges faster than first-order methodsand achieves competitive training loss to other second-order methods (although there remains aperformance gap compared with KFAC).
[Figure]
[Table]
Image Classification.
To investigate whetheradapting the preconditioner with APO improves generalization, we conducted image classification experiments on CIFAR-10 and CIFAR-100.
We trained
LeNet [39], AlexNet [36], VGG-16 [70] (w/o batchnorm [31]), ResNet-18, and ResNet-32 [28] architectures for 200 epochs on batches of 128 images.
The test accuracies for SGDm, Adam, KFAC, and
APO-Precond are shown in Table 2.
We found that APO-Precond achieved competitive generalizationperformance to SGDm and KFAC.
In particular, forarchitectures without batch normalization (LeNet, AlexNet, and VGG-16), APO-Precond improved thetest accuracy substantially.
Neural Machine Translation.
We also applied APO-Precond to the IWSLT’14 German-to-Englishtranslation task [14].
We used a Transformer [78] composed of 6 encoder and decoder layers, withword embedding and hidden vector dimensionality 512.
We compared APO-Precond to SGDm and
AdamW [47].
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 trainingpresents a challenge for second-order optimizers suchas KFAC which rely on matrix inverses that may besensitive to quantization noise.
We trained LeNetand ResNet-18 with 16-bit floating-point arithmeticto examine if APO-Precond is applicable in training the networks in lower precision.
We usedthe experimental setup from Section 5.1 but stored parameters, activations, and gradients in 16-bitprecision.
We found that KFAC required a large damping factor to maintain stability, and thisprevented it from fully utilizing curvature information.
In contrast, as APO-Precond does not requirematrix inversion, it remained stable with the same choice of FSD and WSD weights we used in the fullprecision experiments.
The final test accuracies on ResNet-18 for SGDm, KFAC, and APO-Precondare shown in Table 3. [Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table] [Figure] [Table]
[Figure]
[Table]
5.2 Meta-Learning Learning Rates
Image Classification on MNIST.
First, we compared SGDm and RMSprop to their APO-tunedvariants 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) showsthe training loss achieved by each approach; we found that for both base optimizers, APO improvedconvergence 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 [28], ResNet34, and WideResNet(WRN-28-10) [85].
For ResNet32, we trained for 400 epochs, and the decayed baseline used a stepschedule with 10 × decay at epochs 150 and 250, following [48].
For ResNet34 and WRN-28-10,we trained for 200 epochs, and the decayed baseline used a step schedule with 5 × decay at epochs60, 120, and 160, following [85].
For CIFAR-100, we used WRN-28-10 with the same schedule asfor CIFAR-10.
For each of the base optimizers, we compared APO to (1) the optimal fixed learningrate and (2) a manual step learning rate decay schedule.
The test accuracies for each base optimizerand 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 baseoptimizer.
Using APO to tune the global learning rate yields higher test accuracy than the best fixedlearning 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 updaterule.
As the meta-parameters are updated only once per K steps of optimization, APO incurs minimalcomputational overhead.
We applied APO to two settings: (1) meta-learning the global learning ratefor existing base optimizers (e.g., SGD, RMSprop, and Adam); and (2) meta-learning structuredpreconditioning matrices, which provides a new approach to 2 nd -order optimization.
Compared tomethods such as KFAC, APO eliminates the need to compute matrix inverses, yielding improvedefficiency and numerical stability.
On a range of tasks, we showed that APO is competitive with2 nd -order methods, and improves generalization compared to baseline 1 st - and 2 nd -order optimizers.
[-of Theorem 
-]
Appendix
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 tothat studied in Wu et al.
[84], which suffers from short-horizon bias.
• In Section F, we provide a derivation of Eq. 3.
• In Section G, we derive classic first- and second-order optimization algorithms (gradientdescent, Newton’s method, and natural gradient) starting from the proximal objective for thePPM (Eq.
2).• In Section H, we provide the proof for Theorem 1, which shows that exactly optimizing theproximal meta-objective with respect to the preconditioner P can recover various first- andsecond-order optimization methods.•
In Section I, we list the KFAC assumptions and prove Corollary 2, which shows that underthe assumptions of Theorem 1 and the KFAC assumptions, exactly optimizing the proximalmeta-objective yields the KFAC update.• In Section J, we show an ablation over λ WSD and λ FSD , and evaluate how well APO performswith different meta-update intervals.
A Table of Notation
[Table]
B Extended Related Work
Short-Horizon Bias.
Short-horizon bias (SHB) has been identified as a core challenge in themeta-optimization of optimizer parameters such as learning rates.
Wu et al.
[84] analyzed the settingwhere 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 approximatedusing a single mini-batch.
This meta-objective yields undesirable behaviour where the learning ratedecreases rapidly to reduce fluctuations of the loss caused by stochastic mini-batch evaluation.
Whileour proximal meta-objective is greedy (e.g. it relies on a single lookahead step), APO empirically doesnot suffer from SHB.
We attribute this to the fact that the first term in our proximal meta-objectiveevaluates the loss on the same mini-batch used to compute the gradient rather than a randomlyselected mini-batch.
We refer readers to Appendix E for an ablation demonstrating the differencebetween 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 [75], KF-RTRL [61], OK [11]) can perform online optimization without sufferingfrom SHB in principle (but still suffering from hysteresis).
However, they do not scale well to manyparameters, and thus are be used to tune the preconditioner.
Other approaches like PES [81] usea finite-differences approximation to the gradient, which does not scale to the dimensionality ofthe preconditioner.
RTRL has been applied to hyperparameter adaptation by Franceschi et al.
[21],and a method for LR adaptation that interpolates between hypergradient descent and RTRL, called
MARTHE, was introduced in [16].
Black-Box Hyperparameter Optimization.
Finding good optimization hyperparameters is alongstanding problem [9].
Black-box methods for hyperparameter optimization, such as grid search,random search [12], and Bayesian optimization [71, 74, 72], are expensive, as they require performingmany complete training runs, and can only find fixed hyperparameter values (e.g., a constant learningrate).
Hyperband [43] can reduce the cost by terminating poorly-performing runs early, but is stilllimited to finding fixed hyperparameters.
Population Based Training (PBT) [33] trains a population ofnetworks simultaneously, and throughout training it terminates poorly-performing networks, replacestheir 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, butbecause 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 Polyakstep size methods [27, 79, 80, 45] for automatic learning rate adaptation.
In a similar vein, Rolinek &
Martius [69] proposed a method that tracks an estimated minimal loss and tunes a global learningrate to drive the loss on each mini-batch to the minimal loss.
Learned Optimization.
Some authors have proposed learning entire optimization algorithms [41,42, 2, 83, 56, 83].
Li & Malik [41] view this problem from a reinforcement learning perspective,where the state consists of the objective function J and the sequence of prior iterates { θ ( t ) } andgradients {∇ θ J ( θ ( t ) ) } , and the action is the step ∆ θ .
In this setting, the update rule α is a policythat can be found via policy gradient methods [73].
Approaches that learn optimizers must be trainedon a set of objective functions { f 1 , . . .
, f n } drawn from a distribution F .
This setup can be restrictiveif we only have one instance of an objective function.
In addition, the initial phase of training theoptimizer 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.
APOrequires only the objective function of interest as it does not have trainable parameters.
Few-Shot Meta-Learning Methods.
Meta-SGD [44] extends MAML [19] by learning not onlya shared initialization which can be fine-tuned in a few steps for many different tasks but also adiagonal 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 diagonalpreconditioner 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 isadapted online during training. Meta-Curvature (MC) [65] meta-learns a preconditioning matrix to improve generalization after asmall number of (preconditioned) gradient steps in MAML-style meta-learning. Meta-Curvaturerequires a task distribution for training, and relies on the assumption that related tasks have similarcurvature; the method learns a single, fixed preconditioning matrix that is applied across tasks (and tonew tasks at meta-test time).
In contrast, APO does not require a task distribution, and dynamicallyadapts the meta-learned preconditioning matrix during a single training run, rather than keeping itfixed as done in MC.
Lee and Choi [40] 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 ofweights to update for each task.
Similarly to Meta-SGD and Meta-Curvature, MT-nets require a taskdistribution for meta-training, and only learn fixed meta-parameters which are not adapted duringtraining of new tasks.
APO is a generic framework for online meta-learning of optimization parameters.
We demonstrateits effectiveness for two specific types of meta-parameters: the preconditioning matrix and theglobal learning rate.
Our main contribution is the use of a principled meta-objective based on theproximal point method, which we prove can recover existing first- and second-order optimizationalgorithms given appropriate approximations to the loss term and FSD term.
We focused on tuningthe KFAC-style preconditioner as this is a popular block-diagonal second-order method, and allowsus to draw the connection to KFAC in Corollary 1.
However, the APO meta-objective can be appliedto tune different parameterizations as well; in particular, one could use interleaved warp-layers toparameterize a preconditioner similarly to [40, 20].
C Experimental Details
C.1 Computing Environment
All experiments were implemented using the PyTorch [66] and JAX [13] libraries, and we ran allexperiments 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 = 10in 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 andpreconditioner adaptation.
Learning Rate Adaptation Setup.
For our learning rate adaptation experiments, we used RMSPropas the meta-optimizer with a meta-learning rate of 0.1.
For SGD-APO and SGDm-APO, we set theinitial learning rate to 0.1, while for RMSProp-APO and Adam-APO, we set the initial learning rateto 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 themeta-optimizer with a meta-learning rate chosen from { 1 e- 4 , 3 e- 5 } .
We initialized the block matricessuch 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 updatesthe parameters with SGDm (or Adam for Transformer) but still meta-learns the preconditioningmatrix (e.g. performing meta-gradient steps on P ).
When updating the parameters, we scaled thepreconditioned 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 theinput features and targets to have a zero mean and unit variance.
We used batch size of 128 andtrained the network for 500 epochs without weight decay.
We conducted hyperparameter searches over learning rates for all baseline models, making choicesbased on the final training loss.
With the chosen set of hyperparameters, we repeated the experiments3 times with different random seeds.
For SGDm, we set the momentum to 0.9 and swept over thelearning rates {1, 0.3, 0.1, 0.03, 0.01, 0.003, 0.001, 0.0003, 0.0001}.
For Adam, we performed agrid search on learning rate of {1e-2, 3e-3, 1e-3, 3e-4, 3e-4, 1e-4}.
For Shampoo, we swept overthe 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} anddamping values of {3e-2, 1e-2, 3e-3, 1e-3, 3e-4, 1e-4}.
For APO-Precond, we set λ FSD = 0 and gridsearched over λ WSD = { 3 , 1 , 0 . 1 , 0 . 01 } .
C.4 Image Reconstruction
We used the same experimental set-up from Martens & Grosse [53] for the deep autoencoder experiment.
The loss function was defined to be the binary entropy and we further added a regularizationterm λ WD 2 ∥ θ ∥ 2 to the loss function, where λ WD = 10 − 5 .
The layer widths for the autoencoder wereset 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 onthe final training loss.
With the chosen set of hyperparameters, we repeated the experiments 3 timeswith different random seeds.
For SGDm, we set the momentum to 0.9 and performed a grid searchover the learning rates {1, 0.3, 0.1, 0.03, 0.01, 0.003, 0.001, 0.0003, 0.0001}.
For Adam, we sweptover learning rate of {3e-2, 1e-2, 3e-3, 1e-3, 3e-4, 3e-4, 1e-4, 3e-5, 1e-5}.
For Shampoo, we triedsetting learning rates in range of {10, 5, 1, 0.3, 0.1, 0.03, 0.01, 0.003, 0.001}.
For KFAC, we did agrid search on learning rates of {0.03, 0.01, 0.003, 0.001, 0.0003, 0.0001, 0.00003, 0.00001} anddamping values of {3e-1, 1e-1, 3e-2, 1e-2, 3e-3, 1e-3, 3e-4, 1e-4}.
For APO-Precond, we performeda 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 [85] fortraining 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 of0.5.
We further normalize the inputs with per-channel mean and standard deviations.
We performedthe 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 [85] andselected 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 thenetwork for 200 epochs.
With the chosen set of hyperparameters, we repeated the experiments 3 timeswith different random seeds and reported the mean accuracy on the test dataset.
For SGDm, we setthe 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 architectureswithout batch normalization.
As batch normalization makes the parameters scale-invariant [3], weimposed 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 measurethe KL divergence.
We also searched over the weight decay in range of {5e-4, 1e-4, 5e-5} for alloptimizers.
Learning Rate Adaptation.
For ResNet32, we trained for 400 epochs, and used a manual schedulethat 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 thelearning rate by a factor of 5 at epochs 60, 120, and 160, following Zagoruyko & Komodakis [85].
For the baseline optimizers, we performed grid searches over the fixed learning rateor initial learning rate for a fixed step schedule, as well as the weight decay.
Forall 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 ratesin { 1 . 0 , 0 . 3 , 0 . 1 , 0 .
For RMSprop and Adam, we searchedover learning rates in { 0 . 1 , 0 . 03 , 0 . 01 , 0 . 003 , 0 .
For eachof 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 [78] composed of 6 encoder and decoder layers, with a word embeddingand hidden vector dimension of 512.
The architecture has feed-forward size of 1024, 4 attentionheads, dropout value 0.3, and weight decay value 0.0001.
For the AdamW baseline, we used awarmup-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 thepractice from Zhang et al.
[86], we used a diagonal block matrix in the structured preconditionerof the embedding weight matrix to reduce memory overhead.
We used the Fairseq toolkit [64] toconduct 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 termto 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 basedon 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
[Figure]
D.1 Rosenbrock Function
[Figure]
We validated APO on the two-dimensional Rosenbrock func-
[Figure]
tion defined as:
[Figure]
[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 [Figure] preconditioning matrix.
Both APO-tuned methods outper- [Figure] formed vanilla SGD with the optimal fixed learning rate (chosen through a careful grid search).
Because Rosenbrock has ill-
[Figure] conditioned curvature, second-order optimization with APO- [Figure] Precond dramatically speeds up the convergence and achievesa lower loss.
[Figure]
D.2 SVHN
We also evaluated APO on the Google Street View House Numbers dataset (SVHN) [62], using a
WideResNet 16-8.
We followed the experimental setup of Zagoruyko & Komodakis [85], and usedthe same decay schedule for the baseline, which decays the learning rate by 10 × at epochs 80 and120.
The optimal fixed learning rate achieves test accuracy fluctuating around 97.20%, while themanual schedule achieves 98.17%.
APO rapidly converges to test accuracy 97.96%, outperformingthe baseline while being slightly worse than the manual schedule (Figure 7).
[Figure]
D.3 CIFAR-10
APO Preconditioning Adaptation.
We show the plots for for experiments listed in Table 2 in
Figure 8 and Figure 9.
[Figure]
[Figure] [Figure]
Furthermore, we show the test accuracy as a function of wall-clock time with various values of Kin Figure 10.
We observed that APO is robust to the values of K and setting K = 10 , as done inour main experiments, show a significant improvement in the test accuracy over SGDm while onlyintroducing 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 comparableto the manual schedule.
The learning rate schedules discovered by APO are nontrivial: most exhibit alarge increase in the learning rate at the start of training to make rapid progress, followed by a gradualdecrease 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 decayschedules.
[Table]
[Table] [Table] [Figure]
[Figure] [Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
values over 4 random restarts.
[Figure] [Figure]
D.4 CIFAR-
[Figure]
APO Learning Rate Schedules.
We show training curve plots for the Adam base-optimizer below. [Figure]
We find that APO achieves better performance than the fixed LR, and is comparable with the manualschedule
[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]
[Table]
D.6 Comparison to HD, L4, and FOP
In this section, we compare APO to Hypergradient Descent (HD) [5] and L4 [69] for the learning rateadaptation problem, and First-Order Preconditioning (FOP)
[60] for the preconditioning adaptationproblem.
D.6.1 Hypergradient Descent
Here, we compare to Hypergradient Descent (HD) [69], a method for online learning rate adaptation.
The meta-objective used in HD is the expected value of the training loss after updating the learningrate, E [ L ( u ( θ, ϕ )] .
This is the same objective that was analyzed theoretically by Wu et al.
[84], thatwas shown to lead to sub-optimal behavior when the objective is stochastic.
In particular, [84] showedthat this meta-objective encourages the learning rate to drop rapidly, which reduces the expected lossin the short-term by eliminating fluctuations in the loss caused by random mini-batch sampling, atthe cost of worse long-term behavior (as discussed in Appendix B).
In our experiments, we foundthat HD suffers from the short horizon bias issue; the only way to prevent the LR from reducing tooquickly is to use a very small meta-learning rate, which becomes more heuristic.
Empirically (andas 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.
Becausethe 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 initiallearning 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, sowe left out this optimizer.
For fair comparison to APO, we used the same exponential learning rateparameterization for both APO and HD.
For SGD-HD and SGDm-HD, we performed a grid searchover initial learning rates {1e-1, 1e-2, 1e-3}, meta-learning rates {3.0, 1.0, 1e-1, 1e-2, 1e-3}, andweight decay coefficients {1e-2, 1e-3, 5e-4, 1e-4, 0.0}.
For Adam-HD, we performed a grid searchover initial learning rates {1e-2, 1e-3, 1e-4}, meta-learning rates {3.0, 1.0, 1e-1, 1e-2, 1e-3}, andweight 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 decentlywhen run with a large initial learning rate and a small meta-learning rate, which has the effect ofreducing the LR slowly rather than dramatically (as is expected from the analysis in [84]).
However,using a larger meta-learning rate, which optimizes the meta-objective more rapidly, exhibits theshort horizon bias issue.
In Figures 23,25, we plot the learning rates adapted by HD with a smallmeta-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 affectstraining and leads to substantially worse test accuracy.
In summary, these results provide evidencethat the HD meta-objective is not a very good target for meta-optimization.
[Table]
[Table]
[Figure]
[Figure]
[Figure]
[Figure]
D.6.2 Comparison to L4
We also compared APO to the L4 optimizer [69], which adapts a global learning rate in an onlinefashion based on Polyak’s update rule.
We applied L4 to tune the global learning rate for four baseoptimizers—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 tunefor 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 accuracyand learning rate adaptation obtained with
[Figure]
[Figure]
[Figure]
[Figure]
[Figure]
[Table]
Image Reconstruction.
Following Section 5, we trained an 8-layer autoencoder on MNIST andcompared APO with FOP using the same experimental set-up from Appendix C.4.
For FOP, we set themomentum to 0.9 and performed grid search over the learning rates and hypergradient learning rates inrange of { 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 show the training losses in Figure 31.
APO-Precond converges faster than FOP and achievescompetitive 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 withthe same set-up from Appendix C.5.
For FOP, we fixed the momentum to 0.9 andagain 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-Precondachieved competitive generalization performance to SGDm and KFAC on all tasks, FOP consistentlyunderperformed compared to KFAC.
E Importance of Using the Same Mini-batch
In this section, we verify empirically the importance of computing the first term of the meta-objectiveon the same mini-batch that is used to compute the gradient of the base optimizer.
Recall that ourmeta-objective is:
[Equation]
Note that the loss term is computed using the same mini-batch B used for the update u ( θ , ϕ , B ) .
Analternative is to evaluate the loss term on a randomly-sampled mini-batch B ′ ∼ D , yielding:
[Equation]
If we set λ FSD = λ
WSD = 0 , then Eq.
7 becomes the meta-objective analyzed by Wu et al.
[84], whichwas used to illustrate the short-horizon bias issue in stochastic meta-optimization.
Figures 32,33 show the result of using Eq.
6 vs. Eq.
7 for meta-optimizing the learning rate of a ResNet32 model on CIFAR-10.
We observe that when using a random mini-batch B ′ to computethe loss term, the learning rate decays rapidly, preventing long-term progress; in contrast, using thesame mini-batch B for the loss term yields a reasonable LR schedule which does not suffer from theshort-horizon bias issue empirically.
[Figure]
[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. [Figure]
[Equation]
Recall that D F ( u ( θ , ϕ , B ) , θ , ˜ x ) =
ρ ( f (˜ x , u ( θ , ϕ , B )) , f (˜ x , θ )) .
The function evaluationf (˜ 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 theloss and FSD computations.
However, with the interpretation of the FSD term in Eq. 8 as a Monte
Carlo estimate of the expectation E ˜ x
∼D [ D F ( u ( θ , ϕ , B ) , θ , ˜ x )] , using the same mini-batch for the lossand dissimilarity would yield a biased estimate.
The effect of using a different vs the same mini-batchto compute the FSD term is shown in Figure 34.
[Figure]
F Approximate Closed-Form Solution for the PPM
In Section 2.1, we introduced a general update rule for the stochastic PPM, which is defined asfollows: [Equation]
[Equation]
We first take the infinitesimal limit by letting λ FSD → ∞ .
Then, the optimal update u ⋆ will stay closeto the current parameters θ ( t ) and we can approximate the loss function with the first-order Taylorapproximation.
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 FSDterm would be:
[Equation]
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 theweight space discrepancy function to be the squared Euclidean distance D W ( u , θ ( t ) )
= 1 / 2 ∥ u − θ ( t ) ∥ 22 .
Combining these insights, we approximate Eqn. 9 by linearizing the loss and taking a quadraticapproximation to the FSD term:
[Equation]
Taking the gradient with respect to u and setting it equal to 0 , we have:
[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 theproximal objective introduced in Section 2.2.
Given a mini-batch B , recall that we consider thefollowing proximal objective for optimizing the model parameters θ :
[Equation]
[Table]
J [Table]
( u )
[Table]
F ( [Table] , θ , x )
[Table]
12 || u − θ ||
[Table] 1 st -order - - ✓
[Table] 2 nd -order - - ✗ [Table]
2 nd [Table] order - - ✓
[Table] 1 st -order 2 nd -order
KL ✗
[Table] 1 st -order 2 nd [Table]
KL ✓
[Table] 1 st - [Table] 2 nd
[Table] ✗
[Table] ✓
[Table]
✓
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 aninfinitely large step.
Thus, each step of gradient descent minimizes an approximation to the proximalobjective which we denote by P GD θ ( t ) ( θ ) :
[Equation]
(15) where:
[Equation]
[Equation]
Now, we set the gradient ∇ θ P GD θ ( t ) ( θ ) = 0 and solve for θ : [Equation]
[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 Newtonupdate, which incorporates the WSD term; setting λ
WSD
= 0 leads to the undamped Newton update asa special case.
Each step of Newton’s method minimizes an approximation to the proximal objectivewhich we denote by P Newton θ ( t )
[Equation] [Equation]
[Equation]
where: P Newton
[Equation]
In Eq.
21, H = ∇ 2 θ J B ( θ ) .
Now, we set the gradient ∇ θ P Newton θ ( t ) ( θ ) = 0 and solve for θ :
[Equation]
θ P Newton
[Equation]
( t
[Equation]
[Equation] [Equation] [Equation] [Equation]
[Equation]
G.3 Generalized Gauss-Newton Method
Next, we consider taking the first-order Taylor series approximation to the loss term, and an arbitraryfunction-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, G = E ˜ x
∼D [ ∇ 2 θ D F ( θ , θ ( t ) , ˜ x )] ,we have:
[Equation]
If we set the gradient ∇ θ P θ ( t ) ( θ ) = 0 and solve for θ , we obtain:
[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 toclassic 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 shownin Table 1.
These results parallel those for directly optimizing the parameters θ using the proximalobjective, but optimizing for the meta-parameters ϕ = P .
Throughout this exposition, we use the notation:
[Equation]
31) where g = ∇ θ J B ( θ ) .
Note that g implicitly depends on the data B ; we use this shorthand to simplifythe exposition.
Theorem.
Consider an approximation ˆ Q ( P ) to the meta-objective (Eq. 4) where the loss term is linearized around the current weights θ
and the FSD term is replaced by its second-order approximationaround θ :
[Equation]
Denote the gradient on a mini-batch as g = ∇ θ J B ( θ ) , and assume that the second moment matrix 982 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 θ .
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 areimplicit.
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 minimumof d , so the first two terms are 0.
Denote the network output on example ˜ x by z = f (˜ x , θ ) .
Tocompute ∇ 2 θ d ( θ ) , we have:
[Equation]
To simplify notation, we let G denote the expectation of the Hessian of the FSD function, G ≡ 993 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]
[Equation] Assuming that the second moment matrix E B∼D [ gg ⊤ ] is non-singular, we have:
[Equation]
[Equation] [Equation] Thus, P ⋆ = [Equation]
− [Equation] Below, we provide detailed derivations for several special cases (e.g. specific assumptions on the lossterm and FSD term) which yield gradient descent, damped Newton’s method, and natural gradientdescent.
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, we have: [Equation]
[Equation]
Thus, the optimal preconditioner is P =
λ WSD I , which corresponds to using a global learning rate
of1 λ
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 weobtain a damped Newton update, while if λ WSD = 0 we obtain the undamped update.
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, if we take the second-order Taylor series approximation to the loss term and solve for theoptimal 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-orderapproximation to the FSD term, where the output-space discrepancy measure ρ is taken to be the KLdivergence.
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]
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 ( 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 astructured preconditioner, where P is block-diagonal with the ℓ th block (corresponding to the ℓ thlayer in the network) given by A ℓ ⊗ B ℓ .
I.1 KFAC Assumptions
This exposition of the KFAC assumptions is based on [24].
First we briefly introduce some notationthat will be used in the assumptions.
Suppose a layer of a neural network has weights W ℓ and that itsinput is the activation vector from the previous layer, a ℓ − 1 .
Then the output of the layer is a lineartransformation of the input, followed by a nonlinear activation function σ :
[Equation] [Equation]
(66) where s ℓ are the pre-activations.
In the backward pass, we have the following activation and weightderivatives:
[Equation]
KFAC [53] 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 uncor- related when w i and w j belong to different layers.
If this assumption is satisfied, then theFisher 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 measuresthe KL divergence, and (3) λ
WSD
=
and λ FSD = 1 .
Moreover, suppose that the parameters θ satisfythe 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. 5.
Proof.
If the KFAC assumptions are satisfied, then F − 1 is block-diagonal with the blockcorresponding 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 structuredparameterization (Eq. 5).
Thus, APO recovers the KFAC update.
[Figure]
Ablations
[Figure] J.1 Ablation Over λ [Figure] and λ
[Figure] Figure 35 provides an ablation over λ FSD for training an MLP on MNIST using APO to tune theglobal learning rate for RMSprop (e.g. RMSprop-APO).
[Figure]
layer MLP with 1000 hidden units per layer on
MNIST, using RMSprop-APO. Figure 36 provides an ablation over λ WSD and λ FSD for training AlexNet on CIFAR-10 using APO toadapt the preconditioning matrix.
We kept the other proximity weight fixed when performing theexperiments.
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 37(a,b) that when using APO to tune the global learning rate, APO is robust to theinitial learning rate of the base optimizer, achieving nearly identical training loss and test accuracyusing initial learning rates that span 6 orders of magnitude, from 1e-7 to 1e-1.
Next, we performed an ablation to evaluate how well APO performs with different meta-updateintervals (i.e. when making more or less frequent updates to the learning rate and preconditioningmatrix during training).
We used APO to tune the learning rate, and experimented with performingmeta-updates once every 10, 20, 50, and 100 base optimization iterations.
We trained a ResNet32model on CIFAR-10, and used SGDm as the base optimizer.
The results are shown in Figure 38.
Wefound that APO is robust to the meta-update interval, performing almost identically with respect totraining loss, test accuracy, and the adaptive learning rate schedule, for each meta-update interval. [Figure]
[Figure]
To further examine how meta-update intervals affect the preconditioning adaptation [Figure] we trained
AlexNet model on CIFAR-10.
We used APO to learn the preconditioning matrix and experimentedwith performing meta-updates once every 10, 20, 50, and 100 base optimization iterations.
The resultsare shown in Figure 39.
As in the learning rate adaptation problem, we found that APO performssimilarly using intervals from 10 to 100, achieving higher test accuracy than the baseline method.
[Figure]