A Closer Look at Learned Optimization: Stability, Robustness, and Inductive Biases
Authors: James Harrison
Luke Metz
Jascha Sohl-Dickstein
Venue: NeurIPS 2022 Accept (2022)
Keywords: meta-learning
learned optimization

Abstract: Learned optimizers---neural networks that are trained to act as optimizers---have the potential to dramatically accelerate training of machine learning models. However, even when meta-trained across thousands of tasks at huge computational expense, blackbox learned optimizers often struggle with stability and generalization when applied to tasks unlike those in their meta-training set. In this paper, we use tools from dynamical systems to investigate the inductive biases and stability properties of optimization algorithms, and apply the resulting insights to designing inductive biases for blackbox optimizers. Our investigation begins with a noisy quadratic model, where we characterize conditions in which optimization is stable, in terms of eigenvalues of the training dynamics. We then introduce simple modifications to a learned optimizer's architecture and meta-training procedure which lead to improved stability, and improve the optimizer's inductive bias. We apply the resulting learned optimizer to a variety of neural network training tasks, where it outperforms the current state of the art learned optimizer---at matched optimizer computational overhead---with regard to optimization performance and meta-training speed, and is capable of generalization to tasks far different from those it was meta-trained on. 

A Closer Look at Learned Optimization: Stability, Robustness, and Inductive Biases
James Harrison, Luke Metz, Jascha Sohl-Dickstein Google Research, Brain Team
(jamesharrison, lmetz, jaschasd)@google com
Abstract
Learned optimizers- neerall networks that are trained to act optimizers have the potential to dramatically accelerate training of machine learning models. How- ever, even when meta-trained across thousands of tasks at huge computational expense, blackbox learned optimizers often struggle with stability and generaliza- tion when applied to tasks unlike those in their meta-training set. In this paper, we use tools from dynamical systems to investigate the inductive biases and stability properties of optimization algorithms, and apply the resulting insights designing inductive biases for blackbox optimizers. Our investigation begins with a noisy quadratic model, where characterize conditions in which optimization is sta- ble, in terms of eigenvalues the training dynamics. We then introduce simple modifications to learned optimizer's 's architecture and meta- veeta-trainng procedure which lead to improved stability, and improve the optimizer's ' inductive bias. We apply the resulting learned optimizer where it outperforms the current state of the art learned optimizer-at at matched optimizer computational overhead-with regard to optimization performance and meta- vecta-training speed, and is capable of generalization to tasks far different from those was meta-trained on.
variety of neural network training tasks,
Introduction
Algorithms for stochastic non-convex optimization are a foundational element of modern neural network training [ Choice of optimization algorithm (and the associated hyperparameters) is critical for achieving good performance the underlying model, as well as training stability. Practically, there few formal rules for choosing optimizers and hyperparameters, with researchers and practitioners typically defaulting to a small number of optimizers and associated hyperparameters with which they familiar. Moreover, the algorithms chosen between usually derived based analysis in the convex setting or informal heuristics, and algorithm performance varies substantially across different real world tasks
These optimizers are typically either composed of blackbox function approximators (such neural networks) or hand-designed functions containing hyperparameters that are learned. The distinguishing factor of this class of optimizers is the training methodology: because global search over hyperparameters is computation- ally intractable for optimizers with many parameters, local gradient- baseed methods are used 77- p.. Thus, the meta-training of these optimizers in which they trained optimize some metric of performance for the underlying model such as validation loss-has strong similarities to training other neural network models.
To resolve these issues, learned optimizers have been proposed B
While learned optimizers are potentially transformative, they suffer from several fundamental flaws preventing their broad uptake. The first is reduced performance and stability when they are applied in circumstances unlike those in which they are meta-trained--ie for instance, learned optimizers often
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

diverge when used to train models for more steps than the learned optimizer was applied for during meta-training 10-12. The second instability and inconsistency in the meta-training of the learned optimizers themselves leeaaeee optimizer performance is often highly dependent on random seed, and meta-training trajectories can get stuck for many steps and make inconsistent progress
In this paper, we:
Use tools from dynamical systems theory to characterize the stability of the parameter dynamics induced by learned optimizers, via eigenvalue analysis of training in the noisy quadratic setting.
Propose a series of changes to the architecture and training of learned optimizers that improve their stability and inductive bias. These changes include incorporating tunable contribution from a nominal optimizer with descent guarantees, using large magnitude weight decay the optimizer's parameters, and preconditioning the updates generated by the learned optimizer before applying them to the parameters.
Demonstrate experimentally that the resulting stabilized through ample regularization (STAR) learned optimizer is more stable and faster to meta-train, is more stable and perfor- mant than baseline learned optimizers even when applied for many more training steps than used during meta-traine meta-training, and generalizes well to new tasks which are dissimilar from tasks it was meta-trained on. For instance, STAR generalizes well a transformer task with 175x x more parameters and 5x the number of training steps than the MLP it was meta-trained to optimize. Effective learned optimizer generalization after meta-traie teattttaining on single task has not previously been demonstrated.
2 Related Work
Learned optimization has seen a recent surge of interest motivated by the success of deep learning methods in a wide variety of problems B 13 These methods fit into a larger class of meta- learning methods, which aim to leverage the success of expressive learning algorithms (such neural networks) to learn learning algorithms [1711]. Meta-learning algorithms have been particularly successful and widely investigated in the domain of few-shot learning, which an agent must learn to make predictions based on a small amount of data, Early approaches to this problem focused on blackbox models such as recurrent networks (2--23) due their expressive power. However, works that integrate algorithmic inductive biases for example by exploiting gradient descent [24 25
metric learning 126 27,, convex optimization 28 29) exact or approximate Bayesian inference 30
32). changepoint detection 33 [d] and other algorithmic primitives have been highly successful when applied to few-shot learning. However, similar investigation of inductive biases for meta- learning beyond the few-shot learning setting has been largely absent.
The difficulty of stable training of large models led to the development of adaptive optimization algorithms (such as Adagrad [35]. RMSProp [6). Adam (37], and many other) which are relatively invariant to hyperparameter choice. In practice, however, tuning of these hyperparameters necessary both over the course of training and across optimization problems. While further methods have been developed for automatic online hyperparameter tuning B8 40, a line of work has focused on meta-learning neural network that chooses hyperparameters as a function of training history This approach results in inherently better stability of learning algorithms as the output is typically restricted to (in expectation) descent directions. In this work, we find such inherent stability crucial for guiding training, especially early in meta-training.
still
Shortly after the advent of deep learning, blackbox optimizers were developed which aimed to exploit the expressivity of neural networks in optimizer design. Of particular relevance this work are [46] 47), which use reinforcement learning (combined with guided policy search 4X and the approaches taken in 6 / which are directly trained via backpropagation through time [4 The guided policy search learning strategy is one of several methods that leverage forms of curriculum learning stabilize learning (50). These works all propose networks that ingest (a history of) gradients and return a parameter updates. As investigated by [8]. the long computational graphs necessitated by meta-training can result chaotic behavior which makes meta-training extremely noisy. Numerous approaches to address this problem have been proposed. [. 51 propose truncated zeroth order optimization which avoids extremely noisy reparameterizztioon gradients combined with truncated computation graphs. 45 47
among others reinforcement learning, in which the

---

truncated computational graph is augmented with value function capturing dependency of future losses on the policy. 7 leverage input features computed via fixed momentum operators which yield stable-by-design hidden states. In this paper we dive into the question of stability of meta- training, and find that carefully designing stability into blackbox optimizers is both necessary, and yields optimizers that both outperform the prior state of the art (at matched optimizer computational overhead) and improve overall stability.
Problem Statement
We will consider the problem of training a neural network by optimizing its parameters & E + C R-V We aim to minimize the expectation over data of loss function f + x X' R which acts parameters & and data (point or minibatch) x E X. We define our expected loss at iteration t as
( Fez, [f(xre Q)].
A learned optimizer is defined by parameteric update function ((; 0) with meta-parameters 0 O which acts on a history of training statistics such as parameters, loss values, gradients, and training iteration number. We write these combined input features at time f as 2,. When used to train model with parameters o, then f(; maps input state 24 (as well as optimizer hidden state which we do not explicitly write) to updated parameters 01+1, with an update of the form
Φ1-1 = 02 f(zx;)) 0)-
(2)
In this paper consider only the problem of optimizing the train loss of the inner problem. A particular strength of learned optimizers is that they can be trained with respect to validation loss
8 45) However, this capability is independent and complementary to the topics of this paper, so following Metz et al. Y we focus on train loss to remove factors of variation from experimentation. Our goal in meta-learning is thus to find the parameters that minimize the meta-loss L(0;T),
0 = arg min L(O;T), Bee
L(:: T) = Σ we Co(s), E=1
(3)
under the inner parameter dynamics imposed by 2 where wy is weighting term (as has been used previously, c.g. [6], and where T denotes a maximum (possibly infinite) run length.
4 Understanding Optimizer Performance: The Noisy Quadratic Setting
In this section, we examine performances differences between optimizers, with a particular focus on the impacts on meta-learning of learned optimizers. We characterize the behavior of optimizers in the noisy quadratic setting. This setting-consisting of quadratic optimization problem with randomly sampled, i.i.d. minima-ain aims to be representative of optimization with stochastic minibatches
11 52 53 Moreover, this setting is a reasonable proxy for wide neural networks 54 53
The loss at each timestep is
{():) - H(( -xa
(4)
with i...d & (0, Eg), yielding gradient V. V ((rt
H( - {.). Our chosen loss is the
average of the loss at each timestep, and thus wi = 1/7 for all t.
We consider update of the form
Φ1+1= 6, - (g. + PV.)
(5)
where St. corresponds a nominal term, which do not aim to meta-learn, and P corresponds to our meta-learned term (thus corresponding to 0 in the notation of the previous section), here taking the form of dense preconditioner matrix. In this section, will primarily consider nominal terms of the form 9t a, yielding combined autonomous dynamics
Φ (I (al + P) o, + (al PP &
I-A
Note that is also common to consider a deterministic loss of the form of H with a noisy gradient V. = H + &.. These settings are roughly equivalent, with our setting accumulating an extra tr(HEg)) at each timestep in expectation.
3

---

Nominal term
Weight decay
Preconditioning
Combined
Imag
-1-
(a)
Real (A)
(b)
Real (A)
(c)
Real (A)
(d)
Real (A)
Figure| 1: Nominal terms, preconditioning, and weight decay improve the inductive bias of
learned optimizers. We visualize eigenvalues of the induced parameter update map A (Equation 6 when a simple linear learned optimizer is applied to train a quadratic model. Eigenvalues within the unit circle correspond to stable optimization, while those outside the unit circle will lead diverging training trajectories. Red circles and blue Xes correspond to eigenvalues before and after making each intervention, respectively. Arrows show how an intervention will tend to shift eigenvalues as the intervention magnitude is increased. (a) Incorporating a nominal term causes eigenvalues to have more negative real part, and decay toward the real line. This can cause training dynamics to remain stable even when the learned optimizer alone would cause gradient ascent in some direction. (b) Increasing weight decay pulls the entries in P 'towrrd O, turning off the learned optimizer and thus pulls the eigenvalues of A toward 1. (c) Pre-conditioning reduces the impact of the problem's 's Hessian H. This leads to eigenvalues that depend more on properties of the optimizer, and less on properties of the problem being optimized. (d) Through a weighted combination of the three interventions of (a-c) all eigenvalues are mapped into the unit circle, and the learned optimizer becomes stable.
4.1 Nominal Terms Shift the Region of Stability
Central to our discussion will be the notion of stability. In the study of deterministic dynamical Fl is guaranteed in linear system (6) when p(A) max; [Ax(A)] < for all eigenvalues A;(AA of A. Moreover, for the linear dynamical system we consider in this section, asymptotic stability implies lim L(0; T) is finite, whereas instability implies lim L(0: T) = oo for both the deterministic and stochastic system. Thus, stability is a necessary condition for achieving optimizers which, in expectation, reduce the loss over long training horizons. However, we emphasize that stability here a proxy for convergence that yields simplified discussion and analysis in our setting.
systems (for example, taking Σg - 0), asymptotic stability
While stability is essential for guaranteeing favorable optimizer performance, the noisy quadratic setting also gives us insight into the relationship between the stability of underlying dynamical system and the gradient with respect to meta-trained optimizer P. This gradient is essential for meta-training the optimizer. Under the dynamics O the parameters at time t > 0 are d,
A d' + EE= 1 [1 A) and thus oi (N do. Me) where
A - 11 - A) -
(7)
0=
Thus, our expected loss is
T 000 (TT(T') (A') HA'd + +H((HEE++ +E))).
L(O T) =
(8)
This loss at time f is polynomial in A with degree 2t (and note that A is constant across time). Thus, the gradient of this loss term with respect to P is polynomial in the parameters of A with degree 2t - 1. As a result, instability of the dynamics of á generally implies instability of the gradient the loss with respect to P. which in turn implies the expected gradient magnitude and gradient variance both diverge for long horizons.
Asymptotic stability is defined by limiting behavior converging to a constant, in our case chosen to be O, lim d, = 0. Since our noisy quadratic model has global minimum at O, this corresponds to reaching minimum loss. See Appendix A.I for more discussion of stability and related minor technical results.

---

10
1.0
10"
0.8
[AIA]]
0.6
Average
10%
Gradient
10a
0.4
0.0
0.3
10%
0.1
0.4
0.2
10-1
0.00
0.25
0.50
0.75
1.00
2500
5000
7500
10000
0
2500 5000
7500 10000
(a)
IM
Meta-iterations
(c)
Meta-iterations
Nominal terms improve the stability and trainability of learned optimizers. (a) We adjust the coefficient of the nominal term for the optimizer described in Equation from 0 to 2/ (Amxx(H), The plot shows the trace of the variance of the meta-parameter gradient as a function of a, with different colors correspond trajectory lengths from T (green) T 1000 (blue). For a within our stability bounds, we achieve lower variance gradient estimates for all problem lengths T, with particularly dramatic reduction for larger T. (b) The magnitude of the eigenvalues of the dynamics matrix. A, as a function of meta-iteration during meta-traintn veea-training. These are shown for 2a 2a/Amae (max(())) E {0, 0.1, 0.2, 0.3, 0.4], corresponding to the legend items. Lower variance gradients in pane (a) correspond to more stable, smaller magnitude, eigenvalues. (c) The meta-loss (here corresponding to the mean loss up to = 50) for varying a. Colors correspond to the legend of the center figure. Improving the inductive bias of the learned optimizer, by adjusting leads to near-optimal behavior in under 1000 meta-training iterations, while poorly chosen biases may require 100x as many iterations, or fail to train entirely. See Appendix B. for experimental details.
Figure 2:
How can guarantee stability of the update dynamics? We have stated previously that we require p(A) < 1, which we can map to eigenvalues of via the following result. We make the technical assumption that P has real eigenvalues and is diagonalizable; this is satisfied in the case that P is symmetric as is common for preconditioner matrices. This assumption allows us to strictly order cigenvalues of A, and we write Amin (A) min; A (A) and Amus (A) max; (AA) More importantly, this assumption clarifies the presentation in this primarily pedagogical discussion.
Theorem 1.
Les A = I aH PH with H symmetric positive definite, P diagonalizable with real eigenvalues, and a 2 0. Then ((A) if
a S Amin (P)
(9)
2
Amaxx P)
(10)
Amax
(
The proof of this result and all other results is provided in the Appendix A.2 There are few key takeaways from this result. The lower bound (9) highlights the stabilizing effect of the nominal optimization term. If particular learned optimizer would result in divergence due to moving uphill in particular direction (corresponding to the eigenvector associated with the minimum eigenvalue) without the nominal term, the addition of the nominal term a gives us additional stability margin. The tradeoff is that for the upper bound (TO) we lose stability margin. Thus, we can in general learn only smaller P (as measured by the induced spectrum of A).
In the design of a learned optimizer, the addition of the nominal term can lead to especially large steps which result in divergence due violation of (10). As such, we must be cautious limit the magnitude a. and this result motivates regularization of P. Within the quadratic setting, we are unlikely to practically hit this upper bound as it corresponds to taking steps large enough to oscillate to divergence. For neural network optimization, however, optimization tends to predictably (albeit approximately, due to minibatch dynamics) hit this upper bound (56]
We illustrate the effect on eigenvalues of incorporating a nominal term in Figure | We investigate the stability properties of the noisy quadratic model in Figure learned optimizers more stable, and easier to train.
2 where we find the nominal term makes
Note that our analysis will include the case where p(A) = 1. implying marginal stability in which the system neither converges nor diverges.
5

---

4.2 Preconditioners can Stabilize and Simplify the Design of Update Dynamies
Adaptive preconditioners have seen widespread use in large-scale machine learning. Methods such as Adagrad [35] RMSProp [6] and Adam Kr are some of the most fundamental tools in training neural networks. These approaches typically approximate the inverse square root Hessian, and apply transformation
91 =
H 9
(11)
where St is some nominal optimizer. Typically this approximation diagonal which corresponds to maintaining some weighted recursive estimate of the magnitude of each element of g. This transformation roughly normalizes the step size, making steps isotropic and thus the effective learning rate may be better controlled.
If we combine this Hessian preconditioner together with our preconditioner P and nominal step size a, the resulting dynamics are
- P)VV.
(12)
Stability in this preconditioned setting is given by the following result.
Lemma 2. Consider parameter dynamics given by with assumptions P, H matching Theorem Then is stable if holds for and
maax (P) Vimace H)
(13)
Note, this results in a looser upper bound than (10 if Amax (H) > 1.
Why apply the preconditioner to the output of the learned optimizer as opposed to the input gradient? Normalization at the output of the learned component results in better robustness than normalizing the input, as it is (relatively) robust to arbitrary initializations of the blackbox term. We expand on this in both Section (for quadratic problems) and Section S (for general learned optimizers).
4.3 Adaptive Nominal Terms Improve Robust Stability
We have so far motivated choosing a nominal > 0 to bias the optimizer dynamics toward de- scent/stability. However, is likely that do not want leave a fixed over the course of inner or outer (meta) training. In inner training. decreasing the learning rate over the course of training has been shown improve empirical performance (and is often necessary for guaranteeing convergence 57), and is almost always done when training neural networks.
For simple nominal gradient estimators and comparably complex blackbox terms, the blackbox model should be able to cancel out the nominal term, and induce the effects of reducing a. Setting P p* al with nominal gradient term gt V. yields closed loop dynamics
Φc+1 = Φi V. (P- aT)V: = di P'V,
(14)
which is optimal with respect to meta-loss for optimal P* In this subsection, we argue from the point of view of robust stability that this strategy is suboptimal relative to direct control of the magnitude of the nominal and blackbox term.
We introduce multiplicative error mode which captures the sub-optimality of the learned P during meta- veta-training. Let P = AP for diagonal disturbance A E D. We define this uncertainty set as
D AA E NNN, A = diag(d), < d, S de, d E R
(15)
Within this error model, we can establish conditions for stability in line with the previous subsections. We wish to establish robust stability conditions, which guarantee the stability of the dynamical system for all realizations of the disturbance.
*Our diagonal multiplicative error model is related to a standard formulation within the analysis of the robust stability of linear dynamical systems, known D D-stability (58 60). although we emphasize that D D-stability analysis is usually in continuous time.

---

Lemma 3. Let A aH - PH with P= AP. We will assume P and A are simultaneously diagonalizable, H symmetric positive definite, and 0 < )mao Then, p(A) S for all
S Amina (P)
(16)
max; di
2 iax;; d; Amaax H))
mmxx ())
a
(17)
These results generally result in the tightening of the margin of stability. If we choose P p al, corresponding to previously discussed cancellation of the nominal term, our dynamics become
Φ1+1 = 6: - (a(I - A) + Ve.
(18)
which harms the stability margins. This is well understood by taking A (1- (1-c)d, for any adversarial 6 & 1. Then, we have update dynamics
Φt+1 01 (ac ())c)
(19)
yielding dynamics matrix A = I -1 (1- compared to if the nominal term had instead been set O. In this expression, adversarially perturbs the system toward instability, resulting in a potentially substantial performance drop.
PPHH.. H. Here, H is the excess term in the dynamics
Instead, can select over the course of both inner and outer training via directly controlling the magnitude of a, P. This approach corresponds to hyperparameter controller, which has both been shown to enable automatic control of step sizes [5 over the course of inner training.
4.4
Non-Markovian Optimizers Require Joint Stability
We briefly discuss the role of Man-Maakovian (or hidden state) dynamics in learned optimizers. An extended discussion is in Appendix AA44 The role of momentum and stable hidden states has been investigated detail in both works on optimization
53 56 and learned optimization 61 We
summarize the key points below.
The core analysis change required is that we must consider the joint stability of the hidden state dynamics and the parameter dynamics together. Such analysis has been done in the case of Polyak momentum [22] and has been found yield less restrictive upper stability margins (thus allowing larger step size). In the general case of blackbox optimizer dynamics, the stability properties of the hidden state update have been investigated in [61), and we expand on this in Section [
Momentum often accelerates convergence in the full-batch optimization setting. However, it has the additional benefit of filtering stochastic gradients, yielding better expected descent. This filtering behavior an important consideration in any learned optimizer operating in the stochastic setting.
Designing a Better Learned Optimizer
In this section, we present a set of regularization strategies and architectural modifications for learned optimizers that leverage the insights gleaned in the previous section. We refer to our optimizer as stabilized through ample regularization (STAR) learned optimizen This section also presents limited selection of experiments on in-meta-distribution and out-of-meta-distribe with experimental details available in Appendix B and more results available in Appendix C
performance,
5.1
New Design Features in the STAR Optimizer
Bias toward descent. We add a nominal term, as described in Section We focus on nominal terms based Adam 37 and Ao 63) in part due to the input features blackbox optimizer containing momentum at different timescale (as in AggMo) and running exponential moving average (EMA) gradient magnitude estimates (as in Adam). We experimentally compare different nominal terms in the appendix.
'The code for our optimizer is available here: http:///p./gggg.ooooooo////////p o.to:///t////rrrr...ootttaaaa/a/
-py

---

We add magnitude controller on the nominal descent term, following the discussion in Section This magnitude controller consists one additional output head with an exponential nonlinearity on the small MIP (requiring five additional weights). The combination of magnitude control applied to a nominal optimizer (such as Adam) makes our full nominal term equivalent to a hyperparameter- controller learned optimizer, albeit with a controller that leverages substantial computation reuse with the blackbox term.
Magnitude control via weight decay. In order to discourage violations of the upper bound on stable eigenvalues in Section 4.1 we apply relatively heavy weight decay (L2 regularization) in outer training to the parameters of the blackbox term. By directly controlling weight magnitudes (as opposed controlling magnitudes via regularizing network outputs) we achieve magnitude regularization for arbitrary imputs (for reasonably-sized inputs) and thus achieve better generalization.
Preconditioner-style normalization. As discussed, we an adaptive inverse EMA preconditioner in our nominal term to better bias our model toward descent. As suggested by Section 4.2 we apply the same preconditioner to the output of the blackbox term as well as the network inputs. Because these EMA terms are maintained as inputs to the network already, the expense of this transformation only the cost of dividing the blackbox output by the preconditioner.
Stable hidden states. We discussed the importance of considering the stability of the combined parameter/hidden state dynamics in Section 4.4 It is critical to design the hidden state update dynamics to bias toward stability; indeed, this is a well-known fact in the study of recurrent networks in general 64 (66 and has been discussed in 7 61.. In this paper, in previous works [ 9 use exponential moving average, momentum-styll hidden states which are stable by design.
5.2 Overview of the STAR Optimizer
We apply these modifications to the small_fc_lpt optimizer presented and open sourced in (9] This optimizer is small (197 weight) MLP which is applied elementwise-it takes inputs such as the parameter value, the gradient, and features such as gradient momentum at different timescales, and outputs an update to the parameter. This optimizer applied in parallel to all parameters in the model being trained, and the only interaction between parameter updates is through tensor-level input features. We refer to Appendix A of g for a complete explanation of the optimizer, but we review the basic details here. The optimizer is parameterized
f(x)) 3, (2z) exp(32mp (2)))
(20)
In this expression, B1, B2 are small constants which in [9] and here are set to 0.001. The terms d(x;) and m(²)), corresponding to direction and magnitude terms respectively (so named because the magnitude term only positively scales the direction output), are heads of the neural network. The architecture of the optimizer is an MLP with two hidden layers, each with a width of four. This limits the number of total parameters to 197, yielding high computational efficiency versus most learned optimizers. The input features of the network are the parameter value, various parameter- wise momentum terms, EMA gradient norms estimates, several adafactor- [67 features which aggregate tensor-level information, and parameterization of the training step.
Our modified update takes the form f f=+) (zx) + f(())) where fg(-) is the nominal term and fo(-) is the blackbox term. The nominal term is structured as
(21) where m. is a magnitude term for the nominal term and g(2;) is the nominal term; in our experiments we use combined AggMo [x and Adam (37 for this term. Critically, this term corresponds directly to a descent direction without being passed through a network, biasing the update toward descent. The blackbox term is structured as
So (z:) = 31 exp(32m (z:))g(z)))
fo(z)) = P3 d(2))
v(z)) cxp(3amp(z1)
(22)
where v((;) is our specified preconditioner term and m²(.) 1) is the blackbox magnitude term. The differ- ence between this parameterization and 20) is the normalization term. In these terms, B1, 32, 3s. 84 are band-specified constants, and the neural network has three output heads. The addition of the extra head as well as the multiplicative factor on that term adds an extra five parameters to the original 197-paraeeter optimizer, as well as one more for 31 scaling the entire nominal term. Note that the addition of the v(-) term in the denominator of the blackbox term has magnitude comparable to the mean gradient magnitude, which requires choosing 33 to account for this change in magnitude.
8

---

MLP on Fashion MNIST
CNN on CIFARIO
0.50
1.8
0.45
Meta-loss
0.40
Meta l005
1.6
1.4
0.35
1.2
0.30
0.25
20000
40000
60000
80000
100000
20000
40000
60000
80000
100000
(a)
Meta-iteration
(b)
Meta-iteration
0.5
1.0
Hyperparam
STAR (WD-0.1)
Blackbox
STAR (WD=0.5)
0.3
0.6
0.4
0.0
0.0
(c)
2000
4000
6000
8000 10000 12000 14000
(d)
2000
4000
6000
8000 10000 12000 14000 Inner Iteration
Inner Iteration
Figure 3: The STAR learned optimizer trains faster, achieves better fully-trained optimizers, and has better stability than prior learned optimizers. We visualize meta-training and inner training results for our STAR optimizer, at different values of weight decay, well as for purely blackbox optimizer (the small_cc_lop optimizer from Metz et al. () and a hyperparameter controller model. The upper row shows meta-training curves for (a) two hidden layer MLP on Fashion MNIST, and (b) a three layer CNN on CIFARIO. The lower row shows inner training curves for (c) Fashion MNIST, and (d) CIFARIO. The meta-training objective consists of the average loss over the first 2000 inner iterations, with this horizon indicated by the black dashed line. The blackbox models diverge outside of their meta-training horizon. The STAR model, on the other hand, remains stable when applied for more steps than used during meta veta-training.
5.3 The STAR Optimizer Improves Performance
As described and shown in Figure S the STAR optimizer shows improved stability and faster meta- training than baselines, including the small_fcllop optimizer from Metz et al. M which is on the performance vs. optimizer overhead Pareto frontier. We refer to small_fe_ lopt as "Blackbox" in our experiments, and emphasize that the STAR optimizer architecturally identical to the blackbox optimizer other than the modifications documented in the previous subsection. We also compare to learned optimizer ("Hyperparam"). This optimizer consists of our nominal term with the magnitude control head, but removing the blackbox term. Experimental details are provided in Appendix B and in our open source code. Experiments on additional tasks, with other values of weight decay, ablations of the primary components of the STAR optimizer, and visualization of different random seeds as opposed to aggregated statistics, are provided in Appendix C
hyperparameter-conten
Most interestingly, the STAR optimizer continues to perform strongly even when optimizing for almost two orders of magnitude more steps than it was applied for during meta-training titt-training.. It gen- eralizes in this way to longer training runs without sacrificing performance on the meta-training loss. Divergence or poor performance outside of the meta-training setting (such as for the blackbox optimizer in Figure 31t) is primary limitation to the application of learned optimizers. In contrast, the STAR optimizer has controllable and reliable stability behavior outside of the meta-train set- ting. While meta- meta-generallzztion is helped by training across varying settings (including varying the meta-train horizon and the underlying task) 722, this inherent robustness the ability of a model to generalize well to tasks despite not being meta vecia-trained with the explicit goal of generalization is critical to achieving broadly applicable optimizers. We emphasize that the usual approach to meta wiwnggeenaiiaaionarrreee
meta-training across many large tasks-is extremely expensive, typically requiring weeks of computation on millions dollars of hardware. In contrast, we achieve dramatically improved stability fast and for free. We refer the reader to Metz et al. [ for further baseline results (including heavily- tunee nea-ledrned optimizers) on the tasks evaluated in Figure 3
We further explore the inherent generalization abilities of the STAR optimizer in Figure We apply the optimizer trained on the first task (a small MI. applied to Fashion MNIST) to a wide variety of learning tasks including large models such as transformer A Remarkably, the STAR

---

2.5
Hyperparam
STAR (WD=0.5)
7
Blackbox
2.0
STAR (WD=0.11
NAdamW
Loos
1.0
(8)
2000
4000
iteration
6000
8000
10000
(b)
0
2000
4000
6000
8000
10000
7.0
7.0
6.5
6.5
6.0
6.0
5.5
Lo55
5.5
5.0
5.0
4.5
4.5
4.0
4.0
2000
4000
6000
8000
10000
2000
4000 6000
8000
10000
(c)
Inner iteration
(d)
Inner Iteration
Figure 4: After only being meta-trained to optimize MLP Fashion MNIST for 2000 inner- iterations, the STAR learned optimizer is able to generalize to never before seen problems. We show performance
the following tasks: (a) A 3 hidden layer MLP, with layer norm, trained
Cifar10. (b) A shallow Resnet-like 668 model trained on 32x32 ImageNet [6% (c) A 256 unit LSTM 70 language model trained LMIB ZII. (d) A 5 layer, 256 hidden size decoder-only transformer (also on LMIB). In all cases, the blackbox optimizer diverges, while the STAR models, with appropriately chosen weight decay, continue descend on the loss.
optimizer-which was trained in extremely limited setting-generalize well to different network sizes, nonlinearities, and datasets (including to language tasks). STAR nearly uniformly ourperforms the baseline blackbox model. The performance of STAR (for the correct choice of weight decay) is comparable to (and occasionally substantially better than, as in Figure 4h) a hyperparameter-tunee Adam model. These tasks were selected from subset of evaluation tasks. See the Appendix for learning curves for all generalization experiments.
6 Discussion
This paper has addressed the role of stability and inductive biases learned optimizers, and we have shown that incorporating stabilizing inductive biases results in strong performance both in-distribution and out-of-distribution. Much further work on injecting stability into learned optimizers remains to be done, and designing inductive biases for learned optimizers is potentially work that straddles traditional optimization theory and the study of learned optimizers. There is a deep literature the stability of general computation graphs that and should, be exploited to allow more grained control of stability properties of optimizers (7477))
exciting new line of
Informally, stabilization of the computation graph (as we have defined it) is sufficient condition to guarantee that an optimizer moves downhill on the training loss landscape (in expectation). Our motivation from this comes from the convex setting, in which always moving downhill (with appropriate technical conditions) will lead to a global minimum. However, it is not clear if our induced biases are optimal for the training of neural networks, and in general, it is unclear the extent to which neural network training resembles convex optimization problems [55] 79 SC Further study of which inductive biases desirable for optimizing networks is necessary. Moreover, while have shown strong generalization properties arise from our stability properties, there may be further desirable inductive biases that could be induced to improve generalization, such as for example biasing networks toward flat minima 81 82
Acknowledgments and Disclosure of Funding
We thank Daniel Freeman and Rohan Sinha for helpful conversations and comments in the develop- ment of this work.
10

---

References
[1] Léon Bottou. Stochastic gradient learning in neural networks. Proceedings of Neuro-Numes, 1991.
[2] Robin M Schmidt, Frank Schneider, and Philipp Hennig. Descending through crowded valley- benchmarking deep learning optimizers. International Conference on Machine Learning (ICML), 2021. 2021.
[3] Yoshua Bengio, Samy Bengio, and Jocelyn Cloutier. Learning a synaptic learning rule. Université de Montréal, Département d' d'informatique de recherche opérationnelle, 1990.
[4] Samy Bengio, Yoshua Bengio, Jocelyn Cloutier, and Jan Geesei. On the optimization of synaptic learning rule. In Conference on Optimality in Artificial and Biological Neural Networks, 1992.
[5] Thomas Philip Runarsson and Magnus Thor Jonsson. Evolution and design of distributed learning rules. In IEEE Symposium on Combinations of Evolutionary Computation and Neural Networks, 2000.
[6] Marein Andrychowicz, Misha Denil, Sergio Gomez, Matthew W Hoffman, David Pfau, Tom Schaul, Brendan Shillingford, and Nando De Freitas. Learning to learn by gradient descent by gradient descent. Neural Information Processing Systems (NeurIPS). 2016.
[7] Olga Wichrowska, Niru Maheswaranathan, Matthew W Hoffmann, Sergio Gomez Colmenarejo, Misha Denil, Nando Freitas, and Jascha Sohl-Dickstein. Learned optimizers that scale and generalize. In International Conference on Machine Learning (ICML), 2017.
[8] Luke Metz, Niru Maheswaranathan, Jeremy Nixon, Daniel Freeman, and Jascha Sohl-Dickstein. Understanding and correcting pathologies in the training of learned optimizers. In International Conference Machine Learning (ICML), 2019.
[9] Luke Metz, C. Daniel Freeman, James Harrison, Niru Maheswaranathan, and Jascha Sohl- Dickstein. Practical tradeoffs between memory, compute, and performance in learned optimizers. In Conference on Lifelong Learning Agents (CoLLAs). 2022.
[10] Kaifeng Lv, Shunhua Jiang, and Jian Li. Learning gradient descent: Better generalization and longer horizons. arXiv: /70.::/7700033333 2017.
[11] Yuhuai Wu, Mengye Ren, Renjie Liao, and Roger Grosse. Understanding short-horizon bias in stochastic meta-optimization. In International Conference Learning Representations (ICLR), 2018.
[12] Luke Metz, Niru Maheswaranathan, C Daniel Freeman, Ben Poole, and Jascha Sohl-Dickstein. Tasks, stability, architecture, and compute: Training more effective learned optimizers, and using them to train themselves. arXiv:2009. 2020.
[13] Luca Franceschi, Paolo Frasconi, Saverio Salzo, Riccardo Grazzi, and Massimiliano Pontil. Bilevel programming for hyperparameter optimization and meta-learning. In International Conference on Machine Learning (ICML), 2018
[14] Luke Metz, Niru Maheswaranathan, Jonathon Shlens, Jascha Sohl-Dickstein, and Ekin D Cubuk. Using learned optimizers to make models robust to input noise. arXiv preprint arXiv:1906.033 2019.
[15] Jiayi Shen, Xiaohan Chen, Howard Heaton, Tianlong Chen, Jialin Liu, Wotao Yin, and Zhangyang Wang. Learning minimax optimizer: A pilot study, In International Conference on Learning Representations, 2020.
[16] Wenqing Zheng, Tianlong Chen, Ting-Kuei Hu, and Zhangyang Wang. Symbolic learning to optimize: Towards interpretability and scalability. In International Conference on Learning Representations (ICLR), 2022.
[17] Jürgen Schmidhuber. Evolutionary principles in self-referential learning, or on learning how to learn: the meta-meta-.. hook. PhD thesis, Technische Universität München, 1987.
11

---

[18] Timothy Hospedales, Antreas Antoniou, Paul Micaelli, and Amos Storkey. Meta-learning in neural networks: A survey. arXiv:2004.05339 2020.
[19] Joaquin Vanschoren. Meta-learning: A survey, arXiv:1870.03545 2018.
[20] Adam Santoro, Sergey Bartunov, Matthew Botvinick, Daan Wierstra, and Timothy Lillicrap. Meta-learning with Learning (ICML), 2016.
neural networks. International Conference on Machine
[21] Sepp Hochreiter, A Steven Younger, and Peter Conwell. Learning to learn using gradient descent. In International Conference Artificial Neural Networks, 2001.
[22] Jane X Wang. Zeb Kurth-Nelson, Dhruva Tirumala, Hubert Soyer, Joel Leibo, Remi Munos, Charles Blundell, Dharshan Kumaran, and Matt Botvinick. Learning to reinforcement learn. 1rXiv:1611.05700 2016.
[23] Yan Duan, John Schulman, Xi Chen, Peter L Bartlett, Iya Sutskever, and Pieter Abbeel. RL2: Fast reinforcement learning via slow reinforcement learning. arXiv:1611.02777 2016.
[24] Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adapta- tion of deep networks. In International Conference on Machine Learning (ICML), 2017.
[25] Alex Nichol, Joshua Achiam, and John Schulman. On first-order meta-learning algorithms. arXiv:1803.00299, 2018.
[26] Jake Snell, Kevin Swersky, and Richard Zemel. Prototypical networks for few-shot learning. In Neural Information Processing Systems (NeurIPS), 2017.
[27] Mengye Ren, Eleni Triantafillou, Sachin Ravi, Jake Snell, Kevin Swersky, Joshua B Tenen- baum, Hugo Larochelle, and Richard S Zemel. Meta-learning for semi-supervised few-shot classification. arXiv:1803.0067. 2018.
[28] Luca Bertinetto, João F Henriques, Philip HS Torr, and Andrea Vedaldi. Meta-learning with differentiable closed-form solvers. In International Conference on Learning Representations (ICLR), 2019.
[29] Kwonjoon Lee, Subhransu Maji, Avinash Ravichandran, and Stefano Soatto. Meta-learning with differentiable convex optimization. In IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2019.
[30] James Harrison, Apoorva Sharma, and Marco Pavone. Meta-learning priors for efficient online Bayesian regression. In Workshop the Algorithmic Foundations of Robotics (WAFR), 2018.
[31] Erin Grant, Chelsea Finn, Sergey Levine, Trevor Darrell, and Thomas Griffiths. Recasting gradient-based meta-learning hierarchical Bayes. In International Conference on Learning Representations (ICLR), 2018.
[32] John Willes, James Harrison, Ali Harakeh, Chelsea Finn, Marco Pavone, and Steven Waslander. Bayesian embeddings for few-shot open world recognition. IEEE Transactions on Pattern Analysis & Machine Intelligence, 2022.
[33] Anusha Nagabandi, Chelsea Finn, and Sergey Levine. Deep online learning via meta-learning: Continual adaptation for model-based RL. In International Conference on Learning Represen- tations (ICLR), 2019.
[34] James Harrison, Apoorva Sharma, Chelsea Finn, and Marco Pavone. Continuous meta-learning without tasks. In Neural Information Processing Systems (NeurIPS), 2020.
[35] John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of Machine Learming Research, 12(7), 2011.
[36] Tijmen Tieleman and Geoffrey Hinton. Lecture 6.5-RMSprop: Divide the gradient by a running average of its recent magnitude. COURSERA: Neural networks for machine learning, 2012.
12

---

[37] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. International Conference on Learning Representations (ICLR). 2015.
[38] Atilim Gunes Baydin, Robert Cornish, David Martinez Rubio, Mark Schmidt, and Frank Wood. Online learning rate adaptation with hypergradient descent. vrXiv:1703.04782, 2017.
[39] Ted Moskovitz, Rui Wang, Janice Lan, Sanyam Kapoor, Thomas Miconi, Jason Yosinski, and Aditya Rawal. First-order preconditioning via hypergradient descent. arXiv:/910.08461, 2019.
[40] Juhan Bae, Paul Vicol, Jeff Z HaoChen, and Roger Grosse. Amortized proximal optimization. arXiv:2203.00089, 2022.
[41] Christian Daniel, Jonathan Taylor, and Sebastian Nowozin. Learning step size controllers for robust neural network training. In AAAI Conference Artificial Intelligence, 2016.
[42] Chang Xu, Tao Qin, Gang Wang, and Tie-Yan Liu. Reinforcement learning for learning rate control. arXiv: 1.Kiv:705.11100 11159, 2017.
[43] Zhen Xu, Andrew M Dai, Jonas Kemp, and Luke Metz. Learning an adaptive learning rate schedule. 109.99190, 2019.
[44] Jeffrey Ichnowski. Paras Jain, Bartolomeo Stellato, Goran Banjac, Michael Luo, Francesco Borrelli, Joseph E Gonzalez, Ion Stoica, and Ken Goldberg. Accelerating quadratic optimization with reinforcement learning. Neural Information Processing Systems (NeurIPS), 2021.
[45] Diogo Almeida, Clemens Winter, Jie Tang, and Wojciech Zaremba. A generalizable approach learning optimizers. arXiv:2106.000 00958, 2021.
[46] Ke Li and Jitendra Malik. Learning to optimize.
2016.
[47] Ke Li and Jitendra Malik. Learning to optimize neural nets. arXiv:1703.00441.200 2017.
[48] Sergey Levine and Vladlen Koltun. Guided policy search. International Conference on Machine Learning (ICML), 2013.
[49] Dougal Maclaurin, David Duvenaud, and Ryan Adams. Gradient-based hyperparameter opti- mization through reversible learning. In International Conference on Machine Learning (ICML), 2015.
[50] Tianlong Chen, Weiyi Zhang, Zhou Jingyang, Shiyu Chang. Sijia Liu, Lisa Amini, and Zhangyang Wang. Training stronger baselines for learning optimize. In Neural Information Processing Systems (NeurlPS), 2020.
[51] Paul Vicol, Luke Metz, and Jascha Sohl- Dohk-Dickstein. Unbiased gradient estimation in unrolled computation graphs with persistent evolution strategies. In International Conference on Machine Learning (ICML), 2021.
[52] Tom Schaul, Sixin Zhang, and Yann LeCun. No more pesky learning rates. In International Conference on Machine Learning (ICML), 2013.
[53) Guodong Zhang, Lala Li, Zachary Nado, James Martens, Sushant Sachdeva, George Dahl, Chris Shallue, and Roger B Grosse. Which algorithmic choices matter which batch sizes? insights from a noisy quadratic model. In Neural Information Processing Systems (NeurIPS), 2019.
[54] Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kemnl:: Convergence and generalization neural networks. In Neural Information Processing Systems (NeurIPS), 2018.
[55] Jachoon Lee, Lechao Xiao, Samuel Schoenholz, Yasaman Bahri, Roman Novak, Jascha Sohl- Dickstein, and Jeffrey Pennington. Wide neural networks of any depth evolve as linear models under gradient descent. In Neural Information Processing Systems (NeurIPS), 2019.
[56) Jeremy M Cohen, Simran Kaur, Yuanzhi Li, Zico Kolter, and Ameet Talwalkar. Gradient descent neural networks typically occurs at the edge of stability. arXiv:2103.00063 2021.
13

---

[57] Herbert Robbins and Sutton Monro. A stochastic approximation method. The annals of mathematical statistics, 1951.
[58] Charles R Johnson. Sufficient conditions for d-stability. Journal of Economic Theory, 1974.
[59] Daniel Hershkowitz. Recent directions in matrix stability. Linear Algebra and Applications, 1992.
[60] MC De Oliveira, JC Geromel, and Liu Hsu. LMI characterization of structural and robust stability: the discrete-time case. Linear Algebra and its applications, 1999.
[61] Niru Maheswaranathan, David Sussillo, Luke Metz, Ruoxi Sun, and Jascha Sohl-Dickstein. Reverse engineering learned optimizers reveals known and novel mechanisms. In Neural Information Processing Systems (NeurIPS), 2021.
[62] Boris T Polyak. Some methods of speeding up the convergence of iteration methods. USSR computational mathematics and mathematical physics, 1964.
[63] James Lucas, Shengyang Sun, Richard Zemel, and Roger Grosse. Aggregated momentum: Stability through passive damping. arXiv: 1804.00323 2018.
[64] John Miller and Moritz Hardt. Stable recurrent models. arXiv: 10369, 2018.
[65] Razvan Pascanu, Tomas Mikolov, and Yoshua Bengio. On the difficulty of training recurrent neural networks. In International Conference on Machine Learning (ICML), 2013.
[66] Liang Jin, Peter N Nikiforuk, and Madan M Gupta. Absolute stability conditions for discrete- time recurrent neural networks. IEEE Transactions on Neural Networks, 1994.
[67] Noam Shazeer and Mitchell Stern. Adafactor: Adaptive learning rates with sublinear memory cost. In International Conference on Machine Learning (ICML), 2018.
[68] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2016.
[69] Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2009.
[70] Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 1997.
[71] Ciprian Chelba, Tomas Mikolov, Mike Schuster, Qi Ge, Thorsten Brants, Phillipp Koehn, and Tony Robinson. One billion word benchmark for measuring progress in statistical language modeling. arXiv:1312.3005, 2013.
[72] Luke Metz, Niru Maheswaranathan, Ruoxi Sun, C Daniel Freeman, Ben Poole, and Jascha Sohl-Dickstein. Using thousand optimization tasks to learn hyperparameter search strategies. arXiv:2002. / /88,, 2020.
[73] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Eukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Neural Information Processing Systems (NeurIPS), 2017.
[74] Samuel S Schoenholz, Justin Gilmer, Surya Ganguli, and Jascha Sohl-Dickstein. Deep in- formation propagation. In International Conference on Learning Representations (ICLR), 2017.
[75] Ge Yang and Samuel Schoenholz. Mean field residual networks: On the edge of chaos. In Neural Information Processing Systems (NeurIPS). 2017.
[76] James Martens, Andy Ballard, Guillaume Desjardins, Grzegorz Swirszez, Valentin Dalibard, Jascha Sohl-Dickstein. and Samuel S Schoenholz. Rapid training of deep neural networks without skip connections or normalization layers using deep kernel shaping. Journal of Machine Learning Research, 2022.
14

---

[77] Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces. arXiv:2111.00396, 2021.
[78] Albert Gu, Tri Dao, Stefano Ermon, Atri Rudra, and Christopher Ré. Hippo: Recurrent memory with optimal polynomial projections. In Neural Information Processing Systems (NeurIPS), 2020.
[79] Stanislav Fort and Stanislaw Jastrzebski. Large scale structure of neural network loss landscapes. In Neural Information Processing Systems (NeurIPS), 2019.
[80] Stanislav Fort, Gintare Karolina Dziugaite, Mansheej Paul, Sepideh Kharaghani. Daniel M Roy. and Surya Ganguli. Deep learning versus kernel learning: an empirical study of loss landscape geometry and the time evolution the neural tangent kernel. In Neural Information Processing Systems (NeurIPS), 2020.
[81] Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang. On large-batch training for deep learning: Generalization gap and sharp minima. arXiv:1609.0483 2016.
[82] Pratik Chaudhari, Anna Choromanska, Stefano Soatto, Yann LeCun, Carlo Baldassi, Christian Borgs, Jennifer Chayes, Levent Sagun, and Riccardo Zecchina. Entropy-sgd: Biasing gradient descent into wide valleys. Journal of Statistical Mechanics: Theory and Experiment, 2019.
[83] Brian Anderson and John Moore. Optimal control: linear quadratic methods. 2007.
[84] Jinde Cao, De-Shuang Huang, and Yuzhong Qu. Global robust stability of delayed recurrent neural networks. Chaos, Solitons & Fractals, 2005.
[85] Huaguang Zhang, Zhanshan Wang, and Derong Liu. A comprehensive review of stability analysis of continuous-time recurrent neural networks. IEEE Transactions on Neural Networks and Learning Systems, 2014.
[86] Aaron Defazio. Momentum via primal averaging: Theoretical insights and learning rate schedules for non-convex optimization. arXiv:2010.004066 2020.
[87] Gabriel Goh. Why momentum really works. Distill, 2(4);e6, 2017.
[88] James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Ganderman-Milner and Qiao Zhang. JAX: composable transformations of Python+NummPy programs, 2018.
[89] Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv:1711.05101 2017.
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] We discuss limitations in Section
(c) Did you discuss any potential negative societal impacts of your work? [N/A] Our work aims accelerate network training, and so there are no major negative societal impacts beyond accelerating the inherent risks of machine learning.
(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 of assumptions of all theoretical results? [Yes]
(b) Did you include complete proofs of all theoretical results? [Yes] All proofs are available in the appendix.
15

---

3. If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either in the supplemental material
URL)? [Yes]
(b) Did you specify all the training details (e.g.. data splits, hyperparameters, how they were chosen)? [Yes] These details are available in the appendix.
(c) Did you report error bars (e.g., with respect to the random seed after running exper- iments multiple times)? [Yes] Details on the error statistics are also available in the appendix.
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)? [Yes] Also available in the appendix.
4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) your work uses existing assets, did you cite the creators? [Yes] We primarily build upon the learned optimization package, and cite [ We also cite the relevant other tools and datasets throughout the text.
(b) Did you mention the license of the assets? [Yes] The learned optimization package is open source, and specific license details available in the Github repository.
(c) Did you include any new assets either in the supplemental material or as a URL? [Yes] (d) Did you discuss whether and how consent was obtained from people whose data you're 're using/curating? [N/A] We use standard publicly available datasets. Information this either available at each reference for each dataset, or is not available to us.
(e) Did you discuss whether the data you are using/curating contains personally identifiable information offensive content? [No] We only commonly available datasets; information on this for these datasets is widely available if it is available at all.
5. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, applicable? [N/A]
(b) Did you deseribe any potential participant risks, with links to Institutional Review Board (IRB) approvals, applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation? [N/A]
16


Reviews:
[{"review_id": "yUTD9DNxcW9", "reviewer": "Reviewer_sccL", "summary": "This paper proposes a method to improve learned optimizer generalization by heavily regularizing both training and the architecture of the learned optimizer. These regularizations are derived from a theoretical analysis of the noisy quadratic problem from the perspective of dynamical systems. Using these insights, the paper presents experiments demonstrating the ability to generalize to much larger and very different problems, as well as much longer optimization depths.\n", "questions": "- What is the computational/memory overhead of the learned optimizer used? How does the training curve compare when plotted against training time instead of iteration/epoch?\n- Is the blackbox magnitude term also an output head of the MLP, so the 3 heads are m_g, m_b, d?\n", "limitations": "- Evaluations are relatively limited in scale, with small architectures and datasets.\n- The efficacy of this method in practice as an actual optimizer is questionable, though it seems positioned more as a path to eventual efficacy.\n", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 2, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths\n- The theoretical analysis provides good insight that leads to regularization techniques that will be useful for the area of learned optimization\n- Experiments show reasonable improvements over a tuned Adam optimizer.\n\nWeaknesses\n- Evaluations are all on models much smaller than used in practice. The largest model used appears to be a “Resnet Like” ConvNet on CIFAR and a small transformer.\n- The actual, exact regularizations that are constructed based on the theoretical analysis are far too obfuscated, given that this is the primary contribution of the paper. Tracking down a good summary of the actual method made takes quite a while and ends in the appendix, for example, section B.1 should probably be in the main text, since this is very important for readers who are considering applying this method.\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": {}}, {"review_id": "edqwpxq9cL5", "reviewer": "Reviewer_b5Ki", "summary": "Learned optimizers suffer from being unable to extrapolate for a greater number of steps than they were trained on, and further the learning of these optimizers themselves are unstable. The authors propose that these learned optimizers can be studied under the lens of linear dynamical system, and by constraining the eigenvalues of the update map $A$ via introducing a nominal term, weight decay and preconditioning, the training of such learned optimizers become much more stable and are better able to generalize across greater time steps and even other domains.", "questions": "I was curious if learned optimizers did indeed diverge (or are more likely to diverge) when $\\rho (A) > 1$. Do we find that this is true empirically for existing learned optimizers? \n\nGeneralization of the learned optimizer is interesting and I would expect the preconditioner matrix P and nominal term $\\alpha$ provided by the learned optimizer to be domain dependent, and would be expecting more of the divergence in Fig. 4c, STAR (WD=0.1). Do you have an explanation why your method generalizes rather than diverge?\n\nThe Hyperparam experimental setting doesn't seem to be properly described in the paper, and it would be good if you could provide some details on that.\n\n**Post Rebuttal** The authors clarified in detail the questions above, and I have changed my rating to reflect that.", "limitations": "N/A", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths\n- The perspective of viewing learned optimizers through the lens of linear dynamical systems is refreshing and enlightening, and could potentially explain why they diverge, and the authors provide theoretical substantiation as well\n- Strong empirical results in both generalizing across time (in terms of additional time steps optimized) and across different domains\n\nWeaknesses\n- The flow of the paper is rough on the edges and could use a bit more intuition in some areas, particularly on the section on applying the preconditioning to the output vs to the input gradient\n- Minor typos in the paper (e.g. \"Incorporating *a* nominal terms\" in pg. 3 Fig 1a, \"to the *a* hyperparameter\" in line 217, \"inner *trainin*\" in pg 8)\n\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": {}}, {"review_id": "__enIgB2nx_", "reviewer": "Reviewer_nHvn", "summary": "This paper aims to find the inductive biases and properties for a good optimizer and hence improves upon the existing learnt optimizers.\n", "questions": "Can the authors present a reason why the optimizers fail to stabilize when the input is of high resolution as shown in Appendix", "limitations": "None", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strenths\n– An analysis of stability of the learnt optimizers.\n– Designs new optimizer based on the properties that effect the stability.\n\nWeakness\n– Based on the definition of stability, convergence is not guaranteed. As the limit T->\\inf L(\\theta;T) can be finite even when the loss is oscillating.\n– The results do improve on the existing optimizer, the results are shown on very limited setting.\n– The results fail on high resolution imagenet and on bigger architectures.\n– Due to the above reason, this is not useful in the practical settings.\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": {}}, {"review_id": "SjBG62-z5BT", "reviewer": "Reviewer_Bc3c", "summary": "The paper studies the stability of the learning process from the lens of dynamical systems and characterizes the stability of the optimization process in terms of the eigenvalues of the training dynamics. They use the resulting insights to propose modifications to the learned optimizer’s architecture and the meta-learning procedure to improve its stability and generalization properties. They demonstrate the effectiveness of their modifications by testing the learnt optimizer on a variety of tasks. Importantly, they show that the resulting learned optimizer generalizes much better to tasks very different from the tasks it was meta-trained on.", "questions": "I think both the points in the Weaknesses sections are fixable and would appreciate the authors making amends along those directions.\n\n**Minor Comments**\n1. In line 299 of the paper, the authors say “STAR’s performance is comparable (and occasionally substantially better, 300 as in Figure 4a) to heavily hyperparameter-tuned Adam-based model”. From looking at the plots though, this sounds a bit disingenuous. STAR seems to be doing better on the CIFAR plots but worse on the other datasets (which is reasonable btw, I don’t expect it to work better, but I think it should be made clear in the text). So I would prefer if the paper talked about the differences in performance between tasks and the plausible causes.\n2. Typos - line 217 “the a” -> “a”", "limitations": "The paper discusses certain limitations of the current approach w.r.t the assumptions regarding the convexity of the optimization problem and the need for baking in more inductive biases in order to generalize across tasks. However, I would like to see a deeper discussion on why the proposed learned optimizer still fails to outperform standard optimizers like Adam/NAdamW on a variety of tasks. ", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths**\n1. I liked the first principles approach of analyzing the stability of the existing learned optimizers presented by the paper. Moreover, the paper used the analysis to propose well motivated improvements to the architecture which resulted in significant performance gains. \n2. The paper demonstrates impressive results on various tasks. The improvements on tasks that the optimizer wasn’t trained on were particularly impressive!\n3. The paper was also clearly written and easy to follow.\n\n**Weaknesses**\n1. Ablation/analysis of the stability of the optimizer and validate if it was indeed the stability causing the issues.\n2. The experiments show that the optimizer does reasonably well when tested on tasks very different from the tasks it was trained on. However, I would also like to see how well the optimizer performs when trained on a very wide range of tasks. Does it perform better than the standard optimizers (like adam) across all possible tasks? How much better/worse does it perform compared to an optimizer trained for the specific task? etc.\n\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