Amortized Proximal Optimization
Authors: Juhan Bae
Paul Vicol
Jeff Z. HaoChen
Roger Baker Grosse
Venue: NeurIPS 2022 Accept (2022)
Keywords: second-order optimization
learning rate adaptation
proximal point method
meta-learning

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.

Amortized Proximal Optimization
Juhan Bae" 1.2, Paul Vicol 1.2 Jeff Z. HaoChen', Roger Grosse
University of Toronto, Vector Institute, Stanffrr University {jbae, pvicol, rgrosse)@cs. jiaochen@stanffrd.eed
toronto. toroono edu
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. 'he idea behind APO is to amortize the mini- mization of the proximal point objective by meta-learning the parameters an update rule. We show how APO can be used to adapt learning rate or 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 second-orde optimizers, such 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 competitive with manually tuned decay schedules. Using APO to adapt a structured preconditioning matrix gener- ally 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 Js: IR -> IR on a mini-batch B, plus proximity term that penalizes the discrepancy from the current iterate:
((++1) + arg Jges (u) + (D 0()), wERm
(I)
where :)::)) measures the discrepancy between two vectors and is hyperparameter that con- trols 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. [ exactly is usually impractical (or at least uneconomical), solving it approximately (by taking first second-order Taylor series approximations to the loss or the proximity term) has motivated important and widely used optimization algorithms such natural gradient descent [I] 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.
Equal Contribution
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

Inspired by the idea that the PPM is useful algorithm to approximate, we propose to amortize the minimization of Eq. by defining 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 778), 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 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. I leaves 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 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 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 [9]. 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 '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 -ewstn. While computing and storing a full preconditioning matrix for 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 avoids the need to compute eigen- decompositions. Empirically, we evaluate APO for learning structured preconditioners image reconstruction, image classification, and neural machine translation tasks. The preconditioning matrix adapted by APO achieved competitive convergence to other second-order optimizers.
regression,
2 Preliminaries
Consider a prediction problem from some input space X' to an output space T. We are given finite training set Dtrain = (x(),, ) For a data point (x, t) and parameters 0 E RT, let y = fx, 0) be the prediction of a network parameterized by O and L(y, t) be the loss. Our goal is to minimize the cost function:
1 N = N
[(0)
C(f(x() ())
(2)
We use [o( to denote the mean loss mini-batch of examples B (())), We summarize our notation in Appendix A
3 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.

---

WSD Only
FSD Only
WSD FSD
Initialization
Data
Current Example
Figure 1: ID illustration of the exact proximal update on regression problem with batch size of inspired
by Grosse 25 The weight the discrepancy term(s) (Awsp and ArsD) is decreased from red to blue.
3.1 Proximal Optimization
When we update the parameters on 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:
( t arg min JBCe (u) + > o x)] AwsD O(!)), u
where Dr(u, 0, x) p([(x u); [(x, 0)) and 1 012 are discrepancy functions (described below). Here, ÅESD and AwsD are hyperparameters that control the strength of each discrepancy term, x is a random data point sampled from the gentratnnn distribution D, and p(.))) is the output-space divergence.
The proximal objective in Eq. consists of three terms. The first term is the loss on the current mini- batch. 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 [II, fine-tuning pre- treierr 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 possible. It can be used 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 In fact, Wadia et 83 showed that pure second order optimizers (i.e. ones without WSD regularization) are unable to generalize in the overparameterized setting.
Euclidean norm 33 88
Figure iillutratss the effects of the WSD and FSD terms the exact PPM update for 1D regression example with a batch size of 1. If the proximal objective includes only the loss and WSD term (i.e. AESD = 0), the PPM update makes the minimal change 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. AwSD 0), the update carves 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 local adjustment to the predictions, but one which nonetheless improves performance on nearby examples.
3.2 Connection Between Proximal Optimization and Second-Order Optimization
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- -Fee optimization [5I). a classic second-order opti- mization method for neural networks, approxi- mately minimizes a second-order approximation to the loss on each batch using conjugate gradi- ents. It can be seen as minimizing quadratic approximation to Eq. 3 with no FSD term.
Method
Loss Approx.
FSD
WSD
Gradient Descent
1"-order
Hessian-Free
2nd order
Natural Gradient
order
2dd orddrr
Proximal Point
Exact
Exact
Table 1: Classical 1s and nd optimization algorithms interpreted as minimizing approximations of the prox- imal objective in Eq. 3 using 1 or 2nd order Taylor expansions of the loss or FSD terms.
3

---

Amari motivated natural gradient descent (NGD) as steepest descent method with an infinitesimal step size; this justifies a first-order approximation the loss term and a second-order approximation to the proximity term. Optimizing over a manifold of probability distributions with KI. divergence as the proximity term yields the familiar update involving the Fisher information matrix. Natural gradient optimization of neural networks [52, [4] can be interpreted minimizing Eq. 3 with 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 practical neural network optimizers for stability (54).
In more general context, when taking first-order approximation to the loss and second-order approximation the FSD, the update rule given closed form as:
0 (+)))) R 0
(AFSD + AwsD VoJJ((
(4)
where G is the Hessian of the FSD term. The derivation is shown in Appendix All of these rela- tionships are summarized in Table
1 and derivations of all of these claims are given in Appendix G
4 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 parameterized by vector df which updates the network weights @ batch B(t
( + (() & B)).
(5)
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 & r) and the update is given by:
USGD (0, B) 0 (V)
(6)
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:
UPrecond O, O,P.B) = 0 PVeeo)((
(7)
In Section 3 introduced 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 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 parametric update rule is that we can amortize the cost of minimizing the proximal objective throughout training.
We propose to use the following meta-objectie objetivbe, which evaluates the proximal objective at u(0, Φ, B):
2(()) EB-D [s( (enEEg.)-DDDr(ud 6,B), (.B)), 0, x)]
AWsD (1,(0, -
013
In practice, we estimate the expectations in the meta-objective by sampling two different mini-batches, B and 13', where B is used to compute the gradient and the loss term, and 13' is used to compute the FSD term. Intuitively, this proximal meta-objective aims to find optimizer parameters Φ that
'The update rule may also depend on state maintained by the optimizer, such as the second moments in RMSprop (78). This state is treated fixed by APO, so we suppress it to avoid clutter.

---

Meta-Objective
Q
Loss on Current Minibatch
J& (+))
1-Step Lookahead
Current Parameters
Step w/ Optimization Parameters <
Weight Weight-Space Discrepancy (WSD)
0
0 - 0 PVOJa(
Anya
|0 (
B
Mini-batch
Function-Space Discrepancy (FSD)
Apsp
Ex.-ees Dp ( 0, x)
Mini-batch Gradient
B
Figure 2: Amortized Proximal Optimization (APO). In each meta-optimization step, perform one-step lookahead from the current parameters 0 to obtain updated parameters 0' (Φ). where & denotes optimization parameters (e.g. learning rate 1) or preconditioner P). The objectvbe ((Φ) then evaluates the proximal point objective at (.)). Note that the loss term in ② (() is computed on the same data that was used to compute the gradient for the lookahead step, B, while the FSD term computed using different datapoint (x, t) Dunin- The optimization parameters & are updated via the meta-gradient V (()).
Algorithm 1 Amortized Proximal Optimization (APO) Meta-Leamin Optimization Parameters ø
Require: f (initial model parameters), & (initial optimization parameters), K (meta-update
interval), (meta-LR)
Require: AwSD (weight-space discrepancy term weighting), ÄpSD (function-space discrepancy term weighting)
while not converged, iteration do
B ~ Duain
> Sample mini-batch to compute the gradient and loss term > Perform meta-update every K iterations Sample additional mini-batch to compute the FSD term Compute the 1-step lookahead parameters
if t mod K 0 then
B' ~ Durin
0' (Φ)) = u(0, ②b, B)
((0) = Js (0)) + /B80/1111 E (R,) (6), + Awsp/2 (0) - 0
> Compute meta-objective
Φ Φ- a 2(Φ)
Update optimizer parameters (e.g. LR preconditioner)
end
w(0 5, B)
Update model parameters
end while
minimize the loss on the current mini-batch, while constraining the size the step with the FSD and WSD terms so that it does not overfit the current other mini-batches.
and undo progress that has been made by
The optimization parameters & are optimized with a stochastic gradient- algorithm (the meta- optimizer). The gradieg V ((Φ) 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
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 (O,
(0,,, n, B3) be the 1 1-step lookahead of parameters and minimize the proximal meta-objective
with respect to the learning rate 77. Although adaptive optimizers such as R.SPrrpp and Adam use coordinate-wiid learning rates, they still have global learning rate which is essential tune. APO can be applied to such global learning rates to find learning rate schedules (that depend on ÄesD or AwSD).
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 uprecoma ( P. B)) denote the parameters after preconditioned gradient step and adapt the preconditioning matrix P according to framework.
5

---

If the assumptions made when deriving the second-ordes methods (detailed in Section 3.2) hold, then the optimal preconditioning matrix equivalent to various second-order updates, depending on the choice of the FSD function.
Theorem 1. Consider an approximation ((P) to the meta-objective ( (Eq. S where the loss term is linearized around the current weights O and the FSD term is replaced by its second-order approxima- tion around O. Denote the gradient on mini-batch as g = VoJB O), and assume that the second gg is non-singular. Then, the preconditioning matrix which minimizes Q is where G denotes the Hessian of the FSD evaluated at O.
moment matrix Eg-D
given by P (Arsp G I
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 informa- tion matrix F = Py-D.y-Pyiill Vo log p(y xx V) o log p(y) x, 0) T) where pyix (0) denotes the model's predictive distribution over Therefore, the optimal preconditioner is the damped natural gradient preconditioner, P (F + AwsDI) when AESD = 1. Thus, when APO used exactly minimize an approximate meta- teea-oobecctie,, the update it yields coincides with classical second-org -odeer 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, would like to enforee the preconditioner to be positive semidefinite (PSD) so that the transformed gradient a descent direction [64].
To satisfy these requirements, we adopt structured preconditioner analogous to that of the EKFAC optimizer (24. Given a weight matrix W e m,m Xmi41 of a layer, we construct the preconditioning matrix as a product of smaller matrices:
Ps- (A & (e) B)T,
where A E IR IR X B E m,, and S E Rm, 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 + m + m.m++1 parameters to store, opposed to m parameters for a full preconditioning matrix. It is straightforward to show that the structured preconditioner in Eq. 2 is PSD, as it takes the form CDCT, where D is PSD. The preconditioned gradient can be computed efficiently by using the properties of the Kronecker product:
Psvec( w (s))) = vec + B BwVJs((0)))A)
(10)
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 (e)( vec(B) vec(S) _T. As APO does not require inverting (or performing cigendecomposstions of) the block matrices, our structured representation incurs less computation per iteration than EKFAC.
While defined 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 shows that if certain conditions are satisfied, including the assumptions underlying KFAC (54), then the structured preconditioner minimizing Eq, coincidees with KFAC:
Corollary 2. Suppose that (1) the assumptions for Theorem Jre satisfied, (2) the FSD term measures the KL divergence, and (3) AwsD O and XESD
Moreover, suppose that the parameters satisfy the KFAC assumptions listed in Appendix ! 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 [ 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, 0) and f(x, u(O, Φ 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 backward pass through the 1-step unrolled computation graph (Figure(2) to compute the gradient the proximal meta-objective Q (()) with respect to Φ. This overhead can be reduced by performing 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, further need to store block matrices A, B, and S (in Eq. 9) for each layer, as KFAC and EKFAC.
5 Related Work
We provide extended related work and 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 offfine. This is expensive, it requires completing a full training run to make a single hyperparameter update. A related approach is unroll the optimization for small number of steps and perform truncated backpropagation [16 22 Micaelli et al. (58) perform offline hyperparameter optimization using forward-mode (FDS) rather than eoverse-mode gradient accumulation. FDS performs update to the hyperparameters after each full training run, and thus requires multiple training runs, in contrast to APO, which operates within single training run. Hypergradient descent (6) adapts the learning rate minimize the expected loss in the next iteration.
Second-Order Optimization. Although preconditioned methods have better convergence rates than first-order methods B (44, storing and computing the inverse of preconditioning matrices is impractical for -igesiimessional problems. To mitigate these computational issues, Hessian- free optimization (51)55 55] approximates Newton's '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,
diagonal matrix [19 36), or a low-rank matrix [39] 60)
However, these approximate second-order methods may not easy to implement in deep learning frameworks, and can still be expensive as they often require matrix inversion eigendecompositiin.
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.
6 Experiments
Our experiments investigate the following questions: (1) How does the structured preconditioning matrix adapted by APO perform in comparison to cxisting first- and second-order optimizers? (2) 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 broad range of tasks, including several regression datasets, autoencoder training, image classification on CIFAR-10 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- classification with several network architectures.
7

---

10
Data: Slice
Data: Protein
Data: Parkinsons
5c0m
KFAC APO-Precond
10th
Train lass
10
Adam
Thain loss
Train Loss
30-2
10-3
10-1
10
10-1
100
200
500
400
500
100
200
500
400
500
100
200
500
400
500
Epoch
Epoch
Epoch
Figure A comparison SGDm, Adam, KFAC, Shampoo, and APO-Precoon UCI regression tasks. Across all tasks, APO- Precond achieves lower loss with competitive convergence compared to second-ord order optimizers.
Data: Synthetic
Data: MNIST
10a
8x100
SGD
5GDm
16-1
AMSppsp--rO
loss
107
Shamoce
Adam
Training Loss
Adam
Shampoo
S00m
5GDmArO
20
KFAC
6
10"
APO
16-4
6.10
2000
4000
0000
8000
10000
200
1000
Epoch
Epoch
Epoch
Figure 4: Left: Synthetic data for poorly-conditoondd regression; Middle: Deep autoencoder on MNIST;
Right: Tuning the global learning rate with APO-we we show the training loss for two-layer MLP trained on MNIST, using SGDm and RMSprop (solid lines), and their APO-tuned variants (dashed lines).
We denote our method that adapts the structured preconditioning matrix as -PP--Prcoond.. The method that tunes the global learning rate of base optimizer is denoted as "Base-APO" (c.g. SGDm-APO). Experiment details and additional experiments, including ablation studies, are given in Appendix and respectively.
6.1 Meta-Learning Preconditioners
Poorly-Conditionce Regression. We fist considered a regression task traditionally used illustrate failures of neural network optimization (68). The targets are given by t Ax, where A is ill- conditioned matrix with xAA) = 10t0 We trained two-layer linear network f'xx, 0) = W x In Figure (left), we show training curves for SGDm, Adam, Shampoo (27], KFAC, and APO-Precond. As the problem is ill-conditioned, 2nd_order optimizers such as Shampoo and KFAC converge faster than -rder methods. APO-Precond performs comparably to 2nd -od-or optimizers with lower loss than KFAC.
and minimized the objective (0) Ex []] [|AAx W₂ Wyxx x |
UCI Regression. Next, we validated APO -PO-ood on the Slice, Protein, and Parkinsons datasets from the UCI regression collection 18 We trained 2-layer ML. 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 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 [3] 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 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 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 (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).
We used AdamW optimizer 48 for training Transformer model.
8

---

Task
Model
SGDm
Adam
KFAC
APO-Precond
CIFAR-10
LeNet
75.73
73.41
76.63
77.42
CIFAR-10
AlexNet
76.27
76.09
78.33
81.14
CIFAR-10
VGG16
91.82
90.19
92.05
92.13
CIFAR-10
ResNet-18
93.69
93.27
94.60
94.75
CIFAR-10
ResNet-32
94.40
93.30
94.49
94.83
CIFAR-100
AlexNet
43,95
41.82
46,24
52.35
CIFAR-100
VGG16
65.98
60.61
61.84
67.95
CIFAR-100
ResNet-18
76.85
70.87
76.48
76.88
CIFAR-100
ResNet-32
77.47
56.58
75.70
77.41
SVHN
ResNet-18
96.19
95.59
96.08
96.89
IWSLT14
Transformer
31.43
34.60
34.62
Table Test accuracy on CIFAR-10 and CIFAR-100, and BLEU score on the IWSLT' 14 German-to- English translation dataset for various optimizers.
Image Classification. To investigate whether adapting the preconditioner with APO improves generalization, we conducted image classification experiments on CIFAR-10 and CIFAR-100 1000 We trained LeNet AO), AlexNet [37. VGG-16 [71] (w/o batch norm (32), ResNet-18, and ResNet- 32 [29] 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 generalization performance to SGDm and KFAC. In particular, for architectures without batch normalization (LeNet, AlexNet, and VGG-16), APO-Precond improved the test accuracy substantially.
Neural Machine Translation. To verify the effectiveness of APO various tasks, we applied APO-Precond on the IWSLT' 14 German-to-Bee Ceemmisto-Engllihh translation task [[5). We used Transformer [79] composed of encoder and decoder layers, with word embedding and hidden vector dimensionality 512. We compared APO-Precond to SGDm and Adam [48]. For AdamW, we used a warmup- then-decay learning rate schedule widely used in practice, and for SGD and APO-Precond, we kept
the learning rate fixed after the warmup. In T 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.
Low Precision Training. Low precision train- ing presents a challenge for second-order op-
Task
Model
SGDm
KFAC
APO-P
CIFAR-10
LeNet
75.65
74.95
77.25
timizers such KFAC which rely on matrix inverses that may be sensitive to quantization noise. We trained LeNet and ResNet-18 with 16-bit floating-pots point arithmetic to examine if APO-Precond is applicable training the net- works 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 IPO-Precond does not require matrix inversion, 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
CIFAR-10
ResNet-18
94.15
92.72
94.79
CIFAR-100
ResNet-18
73.53
73.12
75.47
Table 3: Test accuracy of 16-hit networks on CIFAR-10 and CIFAR-100.
6.2 Meta-Learning Learning Rates
Image Classification on MNIST. First, we compared SGDm and RMSprop to their APO-tuned variants train MLP on MNIST. We used two-layer MLP with 1000 hidden units per layer and ReLU nonlinearities, and trained on mini-batches of size 100 for 100 epochs. Figure F (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, experimented with three network architectures: ResNet32 [2), ResNet34, and Net

---

18
500m Li ScDm Deckyed LR
3.9
Decayed 500m
Training Loss
SCDm-APD
Test: Accuracy.
Learning Rate
12-4
ka-1
0.99
3.00
x0-2
13
Fane
2.03
50Dm Decayed L.
SCDm-APD
15 20c 424
100 225 Epoch
230
1.00
300 125 Epach
130
62c
Iceration
Figure 5: WideResNet 28-10 on CIFAR-I0, using SGD with momentum (SGDm). We compare the training loss (left), test accuracy (middle), and learning rate schedules (right) of baselines with fixed and decayed learning rates, and APO. The training loss plot uses hyperparameters chosen based on training loss, while the middle and right plots use hyperparameters chosen based on validation accuracy. The shaded regions show the min/max values over 4 random restarts.
C-10, ResNet-32
C-10,
ResNet-34
C-10,
WRN
28-10
C-100, WRN
28-10
Fixed
Decay
APO
Fixed
Decay
APO
Fixed
Decay
APO
Fixed
Decay
APO
SGD
90.07
93.30
92.71
93.00
93.54
94.27
93.38
94.86
94.85
76,29
77.92
76.87
SGDm
89.40
93.34
92.75
92.99
95.08
94.47
93.46
95.98
95.50
74.81
81.01
79.33
RMSprop
89.84
91.94
91.28
92.87
93.87
93.97
92.91
93.60
94.22
72.06
76.06
74.17
Adam
90.45
92.26
91.81
93.23
94.12
93.80
92.81
94.04
93.83
72.01
75.53
76.33
Table 4: Tuning the global LR for CIFAR-10 10 ("C--0") and CIFAR- 100
("C--10")): We compare the test accuracies achieved by the optimal fixed learning rate, the manual step decay schedule, and the APO-adapted schedule, using ResNet-32 [29) ResNet-34, and WideResNet 28-10 [6)] Results are the mean of four random restarts. APO outperforms optimal fixed learning rates, and often competitive with manual schedules. APO generally achieves test accuracy comparable to manual schedules in fewer training iterations (App. D
(WRN-28-10) 10) 86. For ResNet32, we trained for 400 epochs, and the decayed baseline used a step schedule with 10x decay at epochs 150 and 250, following [49]. For ResNet34 and WRN-28- 10, we trained for 200 epochs, and the decayed baseline used step schedule with 5x decay at epochs [86]. For CIFAR-100, we used WRN-28- 10 with the same schedule as for CIFAR-10 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-tumed variants are shown in Table 4 In addition, Figure SS shows the training loss, test accuracy and learning rate adaptation for WRN-28- 10 on CIFAR-10, using SGDm 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.
60, 120, and 160, following
7 Conclusion
We introduced Amortized Proximal Optimization (APO), a framework for online meta-learning of op- timization parameters which approximates the proximal point method by learning 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: (I) meta-learning the global learning rate for existing base optimizers (c.g., SGD, RMSprop. and Adam) and (2) meta-learning structured preconditioning matrices, which provides new approach to -oddr 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 2nd-order methods, and improves generalization compared to baseline _st and 2nd -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.orecoringg www.eerinniituuteee..e.........e ai/ ../artnners).
10

---

References
[1] Amari, S.-I. Natural gradient works efficiently in learning. Neural Computation, 10(2):251-276 1998.
[2] Andrychowicz, M., Denil, M., Gomez, S., Hoffman, M. W. Pfau, D., Schaul, T., Shillingford, B., and De Freitas, N. Learning to learn by gradient descent by gradient descent. In Advances in Neural Information Processing Systems, pp- 3981-3989, 2016.
[3] Arora, S., Li, Z.,, and Lyu, K. Theoretical analysis of auto rate-tuning by batch normalization. arXiv preprint arXiv: 181.00391, 2018.
[4] Asi. H. and Duchi, J. C. Stochastic (approximate) proximal point methods: Convergence, optimality, and adaptivity. SIAM Journal on Optimization, 29(3):2257-2290 2019.
[5] Bae, J., Ng. N., Lo, A., Ghassemi, M., and Grosse, R. If influence functions are the answer, then what is the question? arXiv preprint arXiv:2209005364 2022.
[6] Baydin, A. G., Cornish, R., Rubio, D. M., Schmidt, M., and Wood, F. Online learning rate adaptation with hypergradient descent. arXiv preprint FrXXiv:7703.44205 2017.
[7] Beck, A. and Teboulle, M. Mirror descent and nonlinear projected subgradient methods for convex optimization. Operations Research Letters, 31(3):167-175,
2003. 2003.
[8] Becker, S. and Le Cun, Y. Improving the convergence of back propppropngttionn learning with second order methods. In Proceedings of the 1988 Connecticut Summer School, 1988.
[9] Belkin, M., Hsu, D., Ma, S., and Mandal, S. Reconciling modern machine-learning practice and the classical virinnari trade-off. Proceedings of the National Academy of Sciences, 15849)215849-15854 2019.
[10] Bengio, Y. Practical recommendations for gradient-based training of deep architectures. In Neural Networks: Tricks of the Trade, pp. 437-478. Springer, 2012.
[11] Benjamin, A. S., Rolnick, D., and Kording, K. Measuring and regularizing networks in function space. arXiv preprint arXiv: 08289, 2018
[12] Benzing, F., Gauy, M. M., Mujika, A., Martinsson, A., and Steger, A. Optimal Kronecker- sum approximation of real time recurrent learning. In International Conference Machine Learning, pp. 604-613, 2019.
[13] Bergstra, and Bengio. Y. Random search for hyper-paramest optimization. Journal of Machine Learning Research, 13(Feb):281-34) 2012.
[14] Bradbury, J., Frostig, R., Hawkins, P., Johnson, M. J., Leary, C., Maclaurin, D., Necula,
G., Paszke, A., VanderPlas, J., Wanderman-Miihe Milne, S., and Zhang, Q. JAX: Composable transformations of Python+NumPy programs, 2018. URL/http://
thub. com/ hoon/gooole//aa
[15] Cettolo, M., Niehues, J., Stüker, S., Bentivogli, L., and Federico, M. Report on the 11th IWSLT evaluation campaign. In Proceedings of the International Workshop on Spoken Language Translation, Hanoi, Vietnam, volume 57, 2014.
[16] Domke, J. Generic methods for optimization-ban modeling, In Artificial Intelligence and Statistics, pp. 318-326, 2012.
[17] Donini, M., Franceschi, L., Pontil, M., Majumder, O., and Frasconi, P. MARTHE: Scheduling the learning rate via online hypergradients. arXiv preprint arXiv:1910.08525, 2019.
[18] Dua, D. and Graff, C. UCI machine learning repository, 2017. URL http:t////a//riee.... ics uci edu/ml
[19] Duchi, J., Hazan, E.,, and Singer, Y. Adaptive subgradient methods for online learning and stochastic optimization. Journal of Machine Learning Research, 12(7), 2011.
[20] Finn, C., Abbeel, P., and Levine, S. aonettag meta-learning for fast adaptation of deep networks. In International conference machine learning, pp. 1126-11335 PMLR, 2017.

---

[21] Flennerhag, S., Rusu, A. A., Pascanu, R., Visin, F., Yin, H., and Hadsell, R. Meta-learning with warped gradient descent. arXiv preprint arXiv:1909. 00025, 2019.
[22] Franceschi, L., Donini, M., Frasconi, P., and Pontil, M. Forward and reverse gradient-base hyperparameter optimization. In International Conference on Machine Learning, pp. 1165- 1173, 2017.
[23] Gao, B., Gouk, H., Lee, H. B., and Hospedales, T. M. Meta mirror descent: Optimiser learning for fast convergence. arXiv preprint arXiv:2203.02711. 2022.
[24] George, T., Laurent, C., Bouthillier, X., Ballas, N., and Vincent, P. Fast approximate natural gradient descent in a Kronecker-factored cigenbasis. arXiv preprint arXiv: 03884, 2018.
[25] Grosse, R.
University of Toronto CSC2541, Topics in Machine Learning: Neu-
ral Net Training Dynamics, Chapter 4: Second- Optimization. Lecture
Notes, 2021. httth://////tr cs. ta.tooroto cecentt.edu///ee redu//"grosse////e "trgroosee/ourrss///ccss/ coourses/csc2544 / readings/ _eaoood__eond____ee aerea__ddrrrppee pdf
[26] Grosse, R. and Martens, J. A Kronecker-factored approximate Fisher matrix for convolution layers. In International Conference on Machine Learning, pp. 573-582, 2016.
[27] Gupta, V., Koren, T., and Singer, Y. Shampoo: Preconditioned stochastic tensor optimization. In International Conference on Machine Learning, pp. 1842-1850, 2018.
[28] Hazan, E. and Kakade, S. Revisiting the polyak step size. arXiv preprint arXiv:1905.00311 2019.
[29] He, K., Zhang, X., Ren, S., and Sun, J. Deep residual learning for image recognition. In Conference on Computer Vision and Pattern Recognition, pp. 770-778, 2016
[30] Hinton, G., Vinyals, O., and Dean, J. Distilling the knowledge in a neural network. arXiv preprint irXiv:1503.02533 2015.
[31] Hinton, G. E. and Salakhutdinov, R. R. Reducing the dimensionality of data with neural networks. Science, 31335786):504-507, 2006.
[32] loffe, S. and Szegedy, C. Batch normalization: Accelerating deep network training by reducing internal covariate shift. arXiv preprint arXiv.1502.03167, 2015.
[33] Jacot, A., Gabriel, F., and Hongler, C. Neural tangent kernel: Convergence and generalization in neural networks. arXiv preprint 106.:0800.. 2018
[34] Jaderberg, M., Dalibard, V., Osindero, S., Czarnecki, W. M., Donahue, J., Razavi, A., Vinyals, O., Green, T., Dunning, I., Simonyan, K., et al. Population-basse training of neural networks. arXiv preprint 17X1v:7111.094::: 2017.
[35] Jiang, H., He, P., Chen, W., Liu, X., Gao, J., and Zhao, T. Smart: Robust and efficient fine-tuning for pre-trained natural language models through principled regularized optimization. arXiv preprint arxiv:191 10911.000 03430033770 2019.
[36] Kingma, D. P. and Ba, J. Adam: A method for stochastic optimization. In International Conference on Learning Representations, 2015.
[37] Krizhevsky, A., Sutskever, I., and Hinton, G. ImageNet classification with deep convolutional neural networks, In Advances in Neural Information Processing Systems, pp. 1097-1105, 2012.
[38] Laurent, C., George, T., Bouthillier, X., Ballas, N., and Vincent, P. An evaluation of Fisher approximations beyond Kronecker factorization, 2018. URL httpp:://openrevvee net/ forumiid
[39] Le Roux, N., Manzagol, P-A.,, and Bengio, Y. Topmoumoute online natural gradient algorithm.
In Advances Neural Information Processing Systems, pp. 449-856, 2007.
[40] LeCun, The MNIST database of handwritten digits. http://yann. lecun. comfexdb/mnisl/, 1988.
12

---

[41] Lee, Y. and Choi, S. Gradient-based meta-learning with learned layerwise metric and subspace. In International Conference on Machine Learning, pp. 2927-2936, 2018.
[42] Li, K. and Malik, J. Learning to optimize. arXiv preprini arXiv: 1.XXv:1606.018855 2016.
[43] Li, K. and Malik, J. Learning optimize neural nets. arXiv preprint arXiv: (XXXv:1703.0044::: 2017.
[44] Li, L.. Jamieson, K., DeSalvo, G., Rostamizadeh, A., and Talwalkar, A. Hyperband: Bandit-based configuration evaluation for hyperparameter optimization. arXiv preprint 060600, 2016.
[45] Li, Z., Zhou, F., Chen, F., and Li, H. Meta-sgd: Learning to learn quickly for few-shot learning. arXiv preprint arXiv: 1eei110077..333 2017.
[46] Loizou, N., Vaswani, S., Laradji, I. H., and Lacoste-Julien, S. Stochastic polyak step-size for sgd: An adaptive learning rate for fast convergence. In International Conference on Artificial Intelligence and Statistics, pp. 1506-1314, PMLR, 2021.
[47] Lorraine, J., Vicol, P., and Duvenaud, Optimizing millions of hyperparameters by implicit differentiation. In International Conference on Artificial Intelligence and Statistics, pp. 1540- 1552. PMLR, 2020.
[48] Loshchilov, I. and Hutter, F. Decoupled weight decay regularization. arXiv preprint arXiv:1711.05101. 2017.
[49] Lucas, J., Sun, S., Zemel, R., and Grosse, R. Aggregated momentum: Stability through passive damping, In International Conference on Learning Representations, 2019.
[50] Maclaurin, D., Duvenaud, D., and Adams, R. Gradient- based hyperparameter optimization through reversible learning. In International Conference Machine Learning, pp. 2113-2122, 2015.
[51] Martens, J. Deep learning via Hessian-free optimization. In International Conference Machine Learning, 2010.
[52] Martens, J. New insights and perspectives on the natural gradient method. arXiv preprint arXiv:1412. 2014.
[53] Martens, J. Second-order optimization for neural networks. PhD) Thesis. University of Toronto, 2016.
[54] Martens, J. and Grosse, R. Optimizing neural networks with Kronecker-factored approximate curvature. In International Conference on Machine Learning, pp. 2408-2411 2015.
[55] Martens, J. and Sutskever. I. Learning recurrent neural networks with Hessian-free optimization. In International Conference on Machine Learning. 2011.
[56] Martens, J., Ba, J., and Johnson, M. Kronecker-factored curvature approximations for recurrent neural networks. In International Conference on Learning Representations, 2018.
[57] Metz, I.., Maheswaranathan, N., Nixon, J., Freeman, C. D., and Sohl-Dickstein, J. Learned opti- mizers that outperform SGD wall-clock and validation loss, arXiv preprint arXiv: /8/0. 1810.0080. 2018.
[58] Micaelli, P. and Storkey, A. Non-greedy gradient-basee based hyperparameter optimization over long horizons. arXiv preprint arXiv:2007.07869 2020.
[59] Michal, R. and Georg, M. L4: Practical loss-based stepsize adaptation for deep learning. arXiv preprint arXiv:1802.00077 2018
[60] Mishkin, A., Kunstner, F., Nielsen, D., Schmidt, M., and Khan, M. E. Slang: Fast structured covariance approximations for Bayesian deep learning with natural gradient. arXiv preprint arXiv:1811.04504, 2018.
13

---

[61] Moskovitz, T., Wang, R., Lan, J., Kapoor, S., Miconi, T., Yosinski, J., and Rawal, A. First-order preconditioning via hypergradient descent. arXiv preprint arXiv:1910.08461, 2019.
[62] Mujika, A., Meier, F,, and Steger, A. Approximating real-time recurrent learning with random Kronecker factors. arXiv preprint arXiv: 1005. 2018.
[63] Netzer, Y. Wang. T., Coates, A., Bissacco, A., Wu, B., and Ng. A. Y. Reading digits in natural images with unsupervised feature learning. In NIPS Workshop Deep Learning and Unsupervised Feature Learning, 2011.
[64] Nocedal, J. and Wright, S. Numerical Optimization. Springer, 1999.
[65] Ott, M., Edunov, S., Baevski, A., Fan, A., Gross, S., Ng. N., Grangier. D., and Auli, M. fairseq: A fast, extensible toolkit for sequence modeling. In Proceedings of NAACL-HLT 2019: Demonstrations, 2019.
[66] Park, E. and Oliva, J. B. Meta-curvature. arXiv preprint arXiv: 2019.
[67] Paszke, A., Gross, S., Massa, F., Lever, A., Bradbury. J., Chanan, G., Killeen, T., Lin, Z., Gimelshein, N., Antiga, L., Desmaison, A., Kopf, A., Yang. E.,, Vito, Z., Raison, M., Tejani, A., Chilamkurthy, S., Steiner, B., Fang, L., Bai, J., and Chintala, Pytorch: An imperative style, high-performance deep learning library. In Advances in Neural Information Processing Systems 32, pp. 8024-8035. 2019.
[68] Rahimi, A. and Recht, B. Reflections on random kitchen sinks, 2017. URL http://www argmin 2eet/2017/12/05/kitcheee
[69] Robbins, H. and Monro, S. A stochastic approximation method. The Annals of Mathematical Statistics, pp. 400-407, 1951.
[70] Rolinek, M. and Martius, G. L4: Practical loss-based stepsize adaptation for deep learning. In Advances in Neural Information Processing Systems, pp. 6433 6443, 2018.
[71] Simonyan, K. and Zisserman, A. Very deep convolutional networks for large-scale image recognition. arXiv preprint arXiv:1409, 1556, 2014.
[72] Snoek, J., Larochelle, H., and Adams, R. Practical Bayesian optimization of machine learning algorithms. Advances in Neural Information Processing Systems, pp. 2951-2959, 2012.
[73] Snock, J., Rippel, O.. Swersky, K.,, Kiros, R., Satish, N., Sundaram, N., Patwary, M., Prabhat, P., and Adams, R. Scalable Bayesian optimization using deep neural networks. In International Conference Machine Learning, pp. 2171-2180, 2015.
[74] Sutton, S., McAllester, D. A., Singh, S. P., and Mansour, Y. Policy gradient methods for reinforcement learning with function approximation. In Advances in Neural Information Processing Systems, pp. 1057-1063, 2000.
[75] Swersky, K., Snoek, J., and Adams, R. P. Freeze-thaw Bayesian optimization. arXiv preprint arXiv:/406. 1406.3896, 2014.
[76] Tallec, C. and Ollivier, Y. Unbiased online recurrent optimization. arXiv preprint 17X1v:170..05043, 2017.
[77] Tang, 7., Jiang, F., Gong, M., Li, H., Wu, Y., Yu, F., Wang, Z., and Wang, M. SKFAC: Training neural networks with faster Kronecker-factorce approximate curvature. In Conference Computer Vision and Pattern Recognition, pp. 13479-13487, 2021.
[78] Ticleman, T. and Hinton, G. Lecture R--MSMopppp Divide the gradient by running average of its recent magnitude. COURSERA: Neural Networks for Machine Learning, 2012.
[79] Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, t.., and Polosukhin, I. Attention all you need. In Advances in Neural Information Processing Systems, pp. 1898860088 2017.
14

---

[80] Vaswani, S., Mishkin, A., Laradji, I., Schmidt, M.. Gidel, G., and Lacoste-Julien, S. Painless stochastic gradient: Interpolation, line-search, and convergence rates. Advances in Neural Information Processing Systems, pp. 3727-3740, 2019.
[81] Vaswani, S., Laradji, Kunstner, F., Meng, Y., Schmidt, M., and Lacoste-Julien, Adaptive gradient methods converge faster with paramcterization (but you should do line-search). arXiv preprint arXiv:2006.06833 2020.
[82] Vicol, P., Metz, L., and Sohl- Sohl-Dickstein, J. Unbiased gradient estimation in unrolled computation graphs with persistent evolution strategies. International Conference on Machine Learning. pp. 10553-10563, 2021.
[83] Wadia, N., Duckworth, D., Schoenholz. S. S., Dyer, E., and Sohl-Dickstein, J. Whitening and second order optimization both make information in the dataset unusable during training, and can reduce or prevent generalization. International Conference on Machine Learning, pp. 10617-10629, 2021.
[84] Wichrowska, O., Maheswaranathan, N., Hoffmann, M. W., Colmenarejo, G., Denil, M., de Fre- itas, N., and Sohl-Dickstein. J. Learned optimizers that scale and generalize. In International Conference on Machine Learning, pp. 3751-3760) 2017.
[85] Wu, Y., Ren, M., Liao, R., and Grosse, R. Understanding short-horizon bias in stochastic meta-optimization. In International Conference on Learning Representations, 2018.
[86] Zagoruyko, S. and Komodakis, N. Wide residual networks. arXiv preprint arXiv: 2016.
[87] Zhang, G., Li, I.,, Nado, Z., Martens, J., Sachdeva, S., Dahl, G., Shallue, C., and Grosse, R. B. Which algorithmic choices matter which batch sizes? insights from noisy quadratic model. Advances in neural information processing systems, 32:8 19--897-,,7 2019.
(88) Zhang, G., Martens, J., and Grosse, R. Fast convergence of natural gradient descent for overparameterized neural networks. arXiv preprint
2019.
Checklist
1. 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 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
(b) Did you include complete proofs of all theoretical results? [Yes] We provide the proof of Theorem in Section H In Section we list the KFAC assumptions and prove
Corollary 2. Additionally,
Section [F] we provide 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
3. If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main exper- imental results (either in the supplemental material or as a URL.)? [Yes] We include code in the supplementary material.
15

---

(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 the random seed after running experi- ments 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 (c.g.. type of GPUs, internal cluster, or cloud provider)? [Yes] We provide these details in Appendix C.I
4. 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 as a URL? [N/A]
(d) Did you discuss whether and how consent obtained from people whose data you'l 'e 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 crowdsoureing 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 participants and the total amount spent on participant compensation? [N/A]
16


Reviews:
[{"review_id": "CDbDEE9bYxU", "reviewer": "Reviewer_NLDo", "summary": "This paper proposes Amortized Proximal Optimization (APO) which is an online meta-optimization for optimization parameters. The modern neural network update rule can be interpreted by proximal point update where it find minimization of three different function w.r.t update rule: 1) the current mini-batch loss, 2) function discrepancy, and 3) weight discrepancy. The authors also reveal that this interpretation is related to other update rule such as Gradient Descent, Hessian-Free, and Natural Gradient. Based on this interpretation, they propose a new meta-objective by sampling two different mini-batches when updating meta-parameters. This paper uses the proposed meta-optimization for optimizing learning rate adaptation and preconditioning matrix for the update rule, and the extensive experimental results shows the effectiveness of APO.", "questions": "- This paper use an efficient approximation for optimizing the meta-objective, by considering only 1-step horizon.  If one update the meta-parameter for every $K$ step, it is natural for me to unroll all the previous K steps (if it is possible). For example, either FMD (RMD is possible for preconditioning matrix) [3], or implicit differentiation [2] are feasible to unroll them. Can you elaborate the reason why one should use 1-step approximation?\n\n[3]  Luca Franceschi, Michele Donini, Paolo Frasconi, and Massimiliano Pontil. Forward and reverse gradient-based hyperparameter optimization. In International Conference on Machine Learning, pages 1165–1173. PMLR, 2017.", "limitations": "I did not find any potential negative societal impact of this paper.", "rating": 7, "confidence": 5, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "- Originality\n\nMany other approaches that meta-learns optimization parameters focus on minimizing training (or validation) loss w.r.t to the parameters. The main originality of this paper is the proposed meta-objective which includes two different discrepancy term, and it is novel. \n\n---\n\n- Quality\n\nThe intuition of this paper is well-motivated throughout the paper, by relating modern neural networks update rule to the proximal point method. Based on the claim on the relatedness, the derivation of the meta-objective seems technically sound to me. The experiments are sufficient for me to understand the effectiveness of the propose method.\n\n---\n\n- Clarity\n\nThis paper is well-written (especially the Figure 2 and Algorithm 1 are very nice to understand the paper).\n\n---\n\n- Significance\n\nThis paper is really nice to be adapted to online optimization of any continuous hyperparameters. Further, it is easy to be combined  with other hyperparameter optimization techniques, since the main modification is the meta-objective; for example, one can use implicit differentiation [1] to obtain $\\nabla Q(\\phi)$ instead of using automatic differentiation of 1-step lookahead.\n\n---\n\n- Weakness\n\nI think this paper is really good, but one thing to improve this paper (I am not going to degrade my score if authors would not address it) is that it is worthwhile to compare it with other hyperparameter optimization methods such as Hypergradient Descent [2]. I think the comparison would provide a clear insight to understand the effectiveness of the proposed meta-objective, since the algorithm is quite similar to [2] except that the discrepancy terms.\n\n[1] Jonathan Lorraine, Paul Vicol, and David Duvenaud. Optimizing millions of hyperparameters by implicit differentiation. In International Conference on Artificial Intelligence and Statistics, 464 pages 1540–1552. PMLR, 2020.\n\n[2] Atilim Gunes Baydin, Robert Cornish, David Martinez Rubio, Mark Schmidt, and Frank Wood. Online learning rate adaptation with hypergradient descent. arXiv preprint arXiv:1703.04782, 2017.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ["I don’t know"]}, "extra_text": {}}, {"review_id": "AocYwfamExA", "reviewer": "Reviewer_4GgD", "summary": "This paper proposes Amortized Proximal Optimization method, which parameterizes proximal point method and performs online meta-learning the parameters during training. Specifically, they emphasize the importance of considering both function space discrepancy and weight space discrepancy at the same time. They consider two meta-learning scenarios: 1) learning rate scheduling and 2) Kronecker-factored preconditioning matrix. Meta-learning is done with usual one-step lookahead method, but w.r.t. the same training minibatch. The experimental results demonstrate the efficacy of the proposed method over diverse learning scenarios, in terms of training convergence and generalization performance.", "questions": "- How the hyperparameters $\\lambda_{FSD}$ and $\\lambda_{WSD}$ have been selected?\n- Why in the outer-loop you used the same minibatch to compute the loss, which seems a little bit unusual?\n- Typo: L279, \"Figure 4 (right)\" --> \"Figure 4 (middle)\"", "limitations": "They did not particularly mentioned the limitations of this work.", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "# Strengths\n- paper is very clearly written and easy to follow\n- Nice introduction motivating why we need both FSD and WSD at the same time in Figure 1, which seems to be the main contribution of this paper.\n- Connection between the proposed method and the existing second-order optimization methods is well explained.\n- Computationally efficient parameterization for the preconditioning matrix, as well as reasonable one-step lookahead meta-learning procedure\n- Experiments are very extensive, and overall achieve competitive performance compared to the existing optimizers without learnable parameters.\n\n# Weaknesses\n- The biggest and significant weakness of this paper is that they did not compare with any of the existing learnable optimizers. Although the existing meta-learned optimizers mostly consider offline setting (Li et al, Lee et al, Park et al, Flennerhag et al.), Some recent work consider online setting such as Micaelli et al., and Gao et al. Especially, Micaelli et al. proposed to perform online optimization for learning rate, weight decay parameter, and momentum parameter at the same time by using forward mode differentiation, which seems very relevant to this submission. I strongly encourage the authors to compare their model to Micaelli et al. at least for the learning rate scheduling scenario. Gao et al. also seems very relevant to this work as they parameterize Bregman divergence of mirror descent. I believe there should be much more literatures that I'm not aware of, so I will defer it to the discussions with other reviewers. Nonetheless, I believe that the authors should compare with at least one or two existing learnable optimizers. I will adjust my score (either decrease or increase) based on this.\n\n\n# References\n- Li et al, Meta-SGD: Learning to Learn Quickly for Few-Shot Learning, 2017\n- Lee et al, Gradient-based Meta-learning with Learned Layerwise Metric and Subspace, ICML 2018\n- Park et al, Meta-Curvature, NeurIPS 2019\n- Flennerhag et al, Meta-Learning with Warped Gradient Descent, ICLR 2020\n- Micaelli et al, Gradient-based Hyperparameter Optimization Over Long Horizons, NeurIPS 2021\n- Gao et al., Meta Mirror Descent: Optimiser Learning for Fast Convergence, ICLR 2022 workshop", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "5LlaBXnToaR", "reviewer": "Reviewer_nHHL", "summary": "This paper introduces an online meta-optimization method called Amortized Proximal Optimization (APO) that can dynamically adjusts parameters like learning rate and preconditioning matrices during the training process.  Before introducing APO, background information is provided to motivate the necessity for both  a function space distance (FSD) term and a weight space distance (WSD) term for the proximal objective. Then, prior 1st and 2nd order optimization methods are interpreted as proximal point methods with various approximations for the loss, FSD, and WSD.  APO is proposed as a meta-optimizer that can be combined with existing optimization strategies to dynamically adjust meta-parameters.  Experiments show APO adaptively learns good learning rate schedules that outperform fixed learning rates and are competitive with manually defined learning rate schedules.  When applied to preconditioning matrices, APO achieves optimization performance competitive with more computationally expensive second order approaches.  ", "questions": "- Is there any intuition for the empirical results in Section E of the Appendix where the same minibatch needs to be used to compute the the gradient of the base optimizer and the loss term of the meta-objective while a different minibatch should be used for the FSD term?\n- How should APO be extended to adapt multiple meta-parameters? e.g. both learning rate and preconditioning matrix?  Are there meta-parameters that are not compatible with APO? can momentum be adapted with APO?  What about generalization/regularization parameters like weight-decay?\n- It seems like manual learning rate schedules outperfrom APO-lr due to using smaller learning rates towards the end of training.  Why does APO not learn to decay LR even more/faster?  Should the \\lambda parameters of APO be adjusted during training as well to promote faster convergence?  \n- I suggest including a legend for Figure 4 left.\n- See additional suggestions/questions in Quality section of the box above.\n\n__Typos__\n- Footnote 2 is missing.\n- ln 289, extra \"a\"", "limitations": "The authors do not discuss the consistent underperformance of APO relative to using a manual decay schedule for many vision tasks in Table 4.  I think it is worth understanding why there's this gap and see if it's possible to bridge this gap with modifications to APO.  I also think a discussion of potential additional compute cost associated with tuning \\lambda for APO would be useful.  Finally, the authors make it sound like APO is applicable to arbitrary meta-parameters but the questions above regarding optimizing multiple meta-parameters simultaneously and regularization related meta-parameters are unanswered.", "rating": 7, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n- Well written paper that provides useful insights into connections between proximal optimization and popular optimization methods and the role of FSD and WSD in improving performance.  \n- Extensive empirical results across multiple types of problems and architectures.  \n- Strong results for APO-Precond that show improved performance for poorly conditioned regression problems and FP16 training.  \n\nWeaknesses:\n- Presentation of experimental results can be improved.\n- APO underperforms SGDm with manual decay schedule in nearly all tasks in table 4. \n- APO in some sense trades off selection a learning rate and decay schedule with tuning \\lambda for the proximal penalty terms.  Perhaps this is not surprising since we shouldn't expect a free lunch and a more powerful optimization should require more tuning.\n- The decision to use a separate batch for the FSD term of the meta-objective feels heuristic and is only supported by empirical evidence in Appendix E.  \n\n\n## Originality \nI am not too familiar with work on developing new optimization methods but I found the presentation and methodology to be interesting and original.\n\n## Quality\nThe overall quality of the paper is high with adequate detail provided in the supplementary material for all experiments.  I do think the experimental results need to be improved however in the following ways:\n- Error bars/std need to be provided for Table 2 and are also highly recommended for all other tables.  For Table 2 in particular, the differences between KFAC and APO-P on architectures with batch-norm layers seem very small and likely insignificant.  \n- Table 4 can use bolding/underlines for users to quickly identify the best optimizer in each setting.  I think this will show that manual decay schedule does the best in nearly every setting.  While this result detracts from the benefit of APO, I think it's worth investigating how to close this gap.  \n- What is the computational cost for KFAC and APO-P in Section 5.1?\n\n## Clarity\nThe paper is well written in most places.  I do have a few suggestion/questions:\n- VGG-16 is stated to not have batch-norm in ln 286-287 but then in ln 296, only LeNet and AlexNet are considered to not have batch-norm.\n- ln 32-34 can be worded better.\n- ln 142-145: This part was confusing for me and I could have used a concrete example.\n\n## Significance\nI think this work can potentially have high significance if the authors are able to close the gap between manual decay and learning lr with APO and show how APO can be applied to multiple meta-parameters.  It still has impact in the current form\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept