Original: A Closer Look at Learned Optimization: Stability, Robustness, and Inductive Biases
Revised: A Closer Look at Learned Optimization: Stability, Robustness, and Inductive Biases

Original: Anonymous Author(s)
Revised: Anonymous Author(s)

Original: AfﬁliationAddress email
Revised: AfﬁliationAddress email

Original: Abstract
Revised: Abstract

Original: Learned optimizers—neural networks that are trained to act as optimizers—havethe potential to dramatically accelerate training of machine learning models.
Revised: Learned optimizers—neural networks that are trained to act as optimizers—have thepotential to dramatically accelerate training of machine learning models.

Original: However, even when meta-trained across thousands of tasks at huge computationalexpense, blackbox learned optimizers often struggle with stability and generalization when applied to tasks unlike those in their meta-training set.
Revised: However,even when meta-trained across thousands of tasks at huge computational expense,blackbox learned optimizers often struggle with stability and generalization whenapplied to tasks unlike those in their meta-training set.

Original: In this paper, weuse tools from dynamical systems to investigate the inductive biases and stabilityproperties of optimization algorithms, and apply the resulting insights to designinginductive biases for blackbox optimizers.
Revised: In this paper, we use toolsfrom dynamical systems to investigate the inductive biases and stability propertiesof optimization algorithms, and apply the resulting insights to designing inductivebiases for blackbox optimizers.

Original: Our investigation begins with a noisyquadratic model, where we characterize conditions in which optimization is stable, in terms of eigenvalues of the training dynamics.
Revised: Our investigation begins with a noisy quadraticmodel, where we characterize conditions in which optimization is stable, in termsof eigenvalues of the training dynamics.

Original: We then introduce simplemodiﬁcations to a learned optimizer’s architecture and meta-training procedurewhich lead to improved stability, and improve the optimizer’s inductive bias.
Revised: We then introduce simple modiﬁcations toa learned optimizer’s architecture and meta-training procedure which lead to improved stability, and improve the optimizer’s inductive bias.

Original: Weapply the resulting learned optimizer to a variety of neural network training tasks,where it outperforms the current state of the art learned optimizer—at matchedoptimizer computational overhead—with regard to optimization performanceandmeta-training speed, and is capable of generalization to tasks far different fromthose it was meta-trained on.
Revised: We apply the resultinglearned optimizer to a variety of neural network training tasks, where it outperformsthe current state of the art learned optimizer—at matched optimizer computationaloverhead—with regard to optimization performance and meta-training speed, andis capable of generalization to tasks far different from those it was meta-trained on.

Original:  Introduction
Revised:  Introduction

Original: Algorithms for stochastic non-convex optimization are a foundational element of modern neuralnetwork training [1].
Revised: Algorithms for stochastic non-convex optimization are a foundational element of modern neuralnetwork training [1].

Original: Choice of optimization algorithm
Revised: Choice of optimization algorithm

Original: (and the associated hyperparameters) is criticalfor achieving good performance of the underlying model, as well as training stability.
Revised: (and the associated hyperparameters) is criticalfor achieving good performance of the underlying model, as well as training stability.

Original: Practically, thereare few formal rules for choosing optimizers and hyperparameters, with researchers and practitionerstypically defaulting to a small number of optimizers and associated hyperparameters with which theyare familiar.
Revised: Practically, thereare few formal rules for choosing optimizers and hyperparameters, with researchers and practitionerstypically defaulting to a small number of optimizers and associated hyperparameters with which theyare familiar.

Original: Moreover, the algorithms chosen between are usually derived based on analysis in theconvex setting or informal heuristics, and algorithm performance varies substantially across differentreal world tasks [2].
Revised: Moreover, the algorithms chosen between are usually derived based on analysis in theconvex setting or informal heuristics and algorithm performance varies substantially across differentreal world tasks [2].

Original: To resolve these issues, learned optimizers have been proposed [3–8].
Revised: To resolve these issues, learned optimizers have been proposed [3–8].

Original: These optimizers are typicallyeither composed of blackbox function approximators (such as neural networks) or
Revised: These optimizers are typicallyeither composed of blackbox function approximators (such as neural networks) or

Original: hand-designedfunctions containing hyperparameters that are learned.
Revised: hand-designedfunctions containing hyperparameters that are learned.

Original: The distinguishing factor of this class ofoptimizers is the training methodology: because global search over hyperparameters is computationally intractable for optimizers with many parameters, local gradient-based methods are used [7–9].
Revised: The distinguishing factor of this class of optimizers is the training methodology: because global search over hyperparameters is computationallyintractable for optimizers with many parameters, local gradient-based methods are used [7–9].

Original: Thus, the meta-training of these optimizers—in which they are trained to optimize some metric ofperformance for the underlying model such as validation loss—has strong similarities to trainingother neural network models.
Revised: Thus,the meta-training of these optimizers—in which they are trained to optimize some metric of performance for the underlying model such as validation loss—has strong similarities to training otherneural network models.

Original: While learned optimizers are potentially transformative, they suffer from several fundamental ﬂawspreventing their broad uptake.
Revised: While learned optimizers are potentially transformative, they suffer from several fundamental ﬂawspreventing their broad uptake.

Original: The ﬁrst is reduced performance and stability when they are applied in circumstances unlike those in which they are meta-trained—for instance, learned optimizers often diverge when used to train models for more steps than the learned optimizer was applied for duringmeta-training [10–12].
Revised: The ﬁrst is reduced performance and stability when they are applied incircumstances unlike those in which they are meta-trained—for instance, learned optimizers oftendiverge when used to train models for more steps than the learned optimizer was applied for during meta-training [10–12].

Original: The second is instability and inconsistency in the meta-training of the learnedoptimizers themselves—learned optimizer performance is often highly dependent on random seed,and meta-training trajectories can get stuck for many steps and make inconsistent progress [8, 9].
Revised: The second is instability and inconsistency in the meta-training of the learnedoptimizers themselves—learned optimizer performance is often highly dependent on random seed,and meta-training trajectories can get stuck for many steps and make inconsistent progress [8, 9].

Original: In this paper, we:
Revised: In this paper, we:

Original: • Use tools from dynamical systems theory to characterize the stability of the parameterdynamics induced by learned optimizers, via eigenvalue analysis of training in the noisyquadratic setting.•
Revised: • 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 noisyquadratic setting.

Original:  Propose a series of changes to the architecture and training of learned optimizers that improve their stability and inductive bias.
Revised: • Propose a series of changes to the architecture and training of learned optimizers thatimprove their stability and inductive bias.

Original: These changes include incorporating a tunable
Revised: These changes include incorporating a tunable

Original: contribution from a nominal optimizer with descent guarantees, using large magnitude
Revised: contribution from a nominal optimizer with descent guarantees, using large magnitude

Original: weight decay on the optimizer’s parameters, and preconditioning the updates generated by the learned optimizer before applying them to the parameters.
Revised: weight decay on the optimizer’s parameters, and preconditioning the updates generated bythe learned optimizer before applying them to the parameters.

Original: • Demonstrate experimentally that the resulting stabilized through ample regularization
Revised: • Demonstrate experimentally that the resulting stabilized through ample regularization

Original: (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-training, and generalizes well to new tasks which are dissimilar from tasks it was meta-trained on.
Revised: (STAR) learned optimizer is more stable and faster to meta-train, is more stable and performant than baseline learned optimizers even when applied for many more training steps thanused during meta-training, and generalizes well to new tasks which are dissimilar from tasksit was meta-trained on.

Original: For instance, STAR generalizes well to a transformer task with 175 ⇥
Revised: For instance, STAR generalizes well to a transformer task with 175 ⇥

Original: more parameters and 5 ⇥ the number of training steps than the MLP it was meta-trained to
Revised: more parameters and 5 ⇥ the number of training steps than the MLP it was meta-trained to

Original: optimize.
Revised: optimize.

Original: Effective learned optimizer generalization after meta-training on a single task has
Revised: Effective learned optimizer generalization after meta-training on a single task has

Original: not previously been demonstrated.
Revised: not previously been demonstrated.

Original:  Related Work
Revised:  Related Work

Original: Learned optimization has seen a recent surge of interest motivated by the success of deep learningmethods in a wide variety of problems [3–8, 13–16].
Revised: Learned optimization has seen a recent surge of interest motivated by the success of deep learningmethods in a wide variety of problems [3–8, 13–16].

Original: These methods ﬁt into a larger class of metalearning methods, which aim to leverage the success of expressive learning algorithms (such as neuralnetworks) to learn learning algorithms [17–19].
Revised: These methods ﬁt into a larger class of metalearning methods, which aim to leverage the success of expressive learning algorithms (such as neuralnetworks) to learn learning algorithms [17–19].

Original: Meta-learning algorithms have been particularlysuccessful and widely investigated in the domain of few-shot learning, in which an agent must learnto make predictions based on a small amount of data.
Revised: Meta-learning algorithms have been particularlysuccessful and widely investigated in the domain of few-shot learning, in which an agent must learnto make predictions based on a small amount of data.

Original: Early approaches to this problem focused onblackbox models such as recurrent networks [20–23] due to their expressive power.
Revised: Early approaches to this problem focused onblackbox models such as recurrent networks [20–23] due to their expressive power.

Original: However, worksthat integrate algorithmic inductive biases—for example by exploiting gradient descent [24, 25],metric learning [26, 27], convex optimization [28, 29], exact or approximate Bayesian inference [30–32], changepoint detection [33, 34], and other algorithmic primitives—have been highly successfulwhen applied to few-shot learning.
Revised: However, worksthat integrate algorithmic inductive biases—for example by exploiting gradient descent [24, 25],metric learning [26, 27], convex optimization [28, 29], exact or approximate Bayesian inference [30–32], changepoint detection [33, 34], and other algorithmic primitives—have been highly successfulwhen applied to few-shot learning.

Original: However, similar investigation of inductive biases for metalearning beyond the few-shot learning setting has been largely absent.
Revised: However, similar investigation of inductive biases for metalearning beyond the few-shot learning setting has been largely absent.

Original: The difﬁculty of stable training of large models led to the development of adaptive optimizationalgorithms (such as Adagrad [35], RMSProp [36], Adam [37], and many other) which are relativelyinvariant to hyperparameter choice.
Revised: The difﬁculty of stable training of large models led to the development of adaptive optimizationalgorithms (such as Adagrad [35], RMSProp [36], Adam [37], and many other) which are relativelyinvariant to hyperparameter choice.

Original: In practice, however, tuning of these hyperparameters is stillnecessary both over the course of training and across optimization problems.
Revised: In practice, however, tuning of these hyperparameters is stillnecessary both over the course of training and across optimization problems.

Original: While further methodshave been developed for automatic online hyperparameter tuning [38–40], a line of work has focusedon meta-learning a neural network that chooses hyperparameters as a function of training history[41–45].
Revised: While further methodshave been developed for automatic online hyperparameter tuning [38–40], a line of work has focusedon meta-learning a neural network that chooses hyperparameters as a function of training history[41–45].

Original: This approach results in inherently better stability of learning algorithms as the output istypically restricted to (in expectation) descent directions.
Revised: This approach results in inherently better stability of learning algorithms as the output istypically restricted to (in expectation) descent directions.

Original: In this work, we ﬁnd such inherent stabilityis crucial for guiding training, especially early in meta-training.
Revised: In this work, we ﬁnd such inherent stabilityis crucial for guiding training, especially early in meta-training.

Original: Shortly after the advent of deep learning, blackbox optimizers were developed which aimed toexploit the expressivity of neural networks in optimizer design.
Revised: Shortly after the advent of deep learning, blackbox optimizers were developed which aimed toexploit the expressivity of neural networks in optimizer design.

Original: Of particular relevance to this workare [46, 47], which use reinforcement learning (combined with guided policy search [48]) and theapproaches taken in [6, 7] which are directly trained via backpropagation through time [49].
Revised: Of particular relevance to this workare [46, 47], which use reinforcement learning (combined with guided policy search [48]) and theapproaches taken in [6, 7] which are directly trained via backpropagation through time [49].

Original: Theguided policy search learning strategy is one of several methods that leverage forms of curriculumlearning to stabilize learning [50].
Revised: Theguided policy search learning strategy is one of several methods that leverage forms of curriculumlearning to stabilize learning [50].

Original: These works all propose networks that ingest (a history of)gradients and return a parameter updates.
Revised: These works all propose networks that ingest (a history of)gradients and return a parameter updates.

Original: As investigated by [8], the long computational graphsnecessitated by meta-training can result in chaotic behavior which makes meta-training extremelynoisy.
Revised: As investigated by [8], the long computational graphsnecessitated by meta-training can result in chaotic behavior which makes meta-training extremelynoisy.

Original: Numerous approaches to address this problem have been proposed.
Revised: Numerous approaches to address this problem have been proposed.

Original: [8, 51] propose truncatedzeroth order optimization which avoids extremely noisy reparameterization gradients combinedwith truncated computation graphs.
Revised: [8, 51] propose truncatedzeroth order optimization which avoids extremely noisy reparameterization gradients combinedwith truncated computation graphs.

Original: [45–47] among others use reinforcement learning, in which thetruncated computational graph is augmented with a value function capturing dependency of future
Revised: [45–47] among others use reinforcement learning, in which thetruncated computational graph is augmented with a value function capturing dependency of futurelosses on the policy.

Revised: [7, 9] leverage input features computed via ﬁxed momentum operators which

Original: [Figure]
Revised: [Figure]

Original: losses on the policy. [7, 9] leverage input features computed via ﬁxed momentum operators whichyield stable-by-design hidden states.
Revised: yield stable-by-design hidden states.

Original: In thispaper we dive into the question of stability of metatraining and ﬁnd that carefully designing stability into blackbox optimizers is both necessary  yields optimizers that both  the prior state of the art (at matched optimizer computationaloverhead) and improve overall stability.
Revised: In this we dive into the question of stability of metatraining and ﬁnd 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 computationaloverhead) and improve overall stability.

Original:  Problem Statement
Revised:  Problem Statement

Original: We will consider the problem of training a neural network by optimizing its parameters  2  ✓ R N .
Revised: We will consider the problem of training a neural network by optimizing its parameters  2  ✓ R N .

Original: We aim to minimize the expectation over data of loss function ` :  ⇥ X !
Revised: We aim to minimize the expectation over data of loss function ` :  ⇥ X !

Original: R which acts onparameters  and data (point or minibatch) x 2 X .
Revised: R which acts onparameters  and data (point or minibatch) x 2 X .

Original: We deﬁne our expected loss at iteration t as
Revised: We deﬁne our expected loss at iteration t as

Original: [Equation]
Revised: [Equation]

Original: A learned optimizer is deﬁned by a parameteric update function f ( · ; ✓ ) with meta-parameters ✓ 2 ⇥ ,which acts on a history of training statistics such as parameters, loss values, gradients, and trainingiteration number.
Revised: A learned optimizer is deﬁned by a parameteric update function f ( · ; ✓ ) with meta-parameters ✓ 2 ⇥ ,which acts on a history of training statistics such as parameters, loss values, gradients, and trainingiteration number.

Original: We write these combined input features at time t as z t .
Revised: We write these combined input features at time t as z t .

Original: When used to train a modelwith parameters  , then f ( · ; ✓ ) maps input state z t (as well as optimizer hidden state which we donot explicitly write) to updated parameters  t
Revised: When used to train a modelwith parameters  , then f ( · ; ✓ ) maps input state z t (as well as optimizer hidden state which we donot explicitly write) to updated parameters  t

Original: +1 , with an update of the form
Revised: +1 , with an update of the form

Original: [Equation]
Revised: [Equation]

Original: In this paper we consider only the problem of optimizing the train loss of the inner problem.
Revised: In this paper we consider only the problem of optimizing the train loss of the inner problem.

Original: Aparticular strength of learned optimizers is that they can be trained with respect to validation loss[8, 45].
Revised: Aparticular strength of learned optimizers is that they can be trained with respect to validation loss[8, 45].

Original: However, this capability is independent and complementary to the topics of this paper, sofollowing Metz et al.
Revised: However, this capability is independent and complementary to the topics of this paper, sofollowing Metz et al.

Original: [9] we focus on train loss to remove factors of variation from experimentation.
Revised: [9] we focus on train loss to remove factors of variation from experimentation.

Original: Our goal in meta-learning is thus to ﬁnd the parameters ˆ ✓ that minimize the meta-loss L ( ✓ ; T ) ,
Revised: Our goal in meta-learning is thus to ﬁnd the parameters ˆ ✓ that minimize the meta-loss L ( ✓ ; T ) ,

Original: [Equation]
Revised: [Equation]

Original: under the inner parameter dynamics imposed by (2), where w t is a weighting term (as has been usedpreviously, e.g. [6]), and where T denotes a maximum (possibly inﬁnite) run length.
Revised: under the inner parameter dynamics imposed by (2), where w t is a weighting term (as has been usedpreviously, e.g. [6]), and where T denotes a maximum (possibly inﬁnite) run length.

Original:  Understanding Optimizer Performance: The Noisy Quadratic Setting
Revised:  Understanding Optimizer Performance: The Noisy Quadratic Setting

Original: In this section, we examine performances differences between optimizers, with a particular focus onthe impacts on meta-learning of learned optimizers.
Revised: In this section, we examine performances differences between optimizers, with a particular focus onthe impacts on meta-learning of learned optimizers.

Original: We characterize the behavior of optimizers in the noisy quadratic setting.
Revised: We characterize the behavior of optimizers in the noisy quadratic setting.

Original: This setting—consisting of a quadratic optimization problem with randomlysampled, i.i.d. minima—aims to be representative of optimization with stochastic minibatches[11, 52, 53].
Revised: This setting—consisting of a quadratic optimization problem with randomlysampled, i.i.d. minima—aims to be representative of optimization with stochastic minibatches[11, 52, 53].

Original: Moreover, this setting is a reasonable proxy for wide neural networks [54, 55].
Revised: Moreover, this setting is a reasonable proxy for wide neural networks [54, 55].

Original: The loss 1 at each timestep is
Revised: The loss 1 at each timestep is

Original: [Equation]
Revised: [Equation]

Original:  with i.i.d.
Revised: ) with i.i.d.

Original: ⇠ t ⇠ N (0 , ⌃ ⇠ ) , yielding gradient r t = r  t ` (  t ) =
Revised: ⇠ t ⇠ N (0 , ⌃ ⇠ ) , yielding gradient r t = r  t ` (  t ) =

Original: H (  t  ⇠ t ) .
Revised: H (  t  ⇠ t ) .

Original: Our chosen loss is the 121 average of the loss at each timestep, and thus w t = 1 /T for all t .
Revised: Our chosen loss is the 120 average of the loss at each timestep, and thus w t = 1 /T for all t .

Original: We consider an update of the form
Revised: We consider an update of the form

Original: [Equation]
Revised: [Equation]

Original: [Equation] where g t corresponds to a nominal term, which we do not aim to meta-learn, and P corresponds toour meta-learned term (thus corresponding to ✓ in the notation of the previous section), here takingthe form of a dense preconditioner matrix.
Revised: where g t corresponds to a nominal term, which we do not aim to meta-learn, and P corresponds toour meta-learned term (thus corresponding to ✓ in the notation of the previous section), here takingthe form of a dense preconditioner matrix.

Original: In this section, we will primarily consider nominal termsof the form g t =
Revised: In this section, we will primarily consider nominal termsof the form g t =

Original: ↵ r t , yielding combined autonomous dynamics
Revised: ↵ r t , yielding combined autonomous dynamics

Original: [Equation]
Revised: [Equation]

Original: 4.1 Nominal Terms Shift the Region of Stability
Revised: 4.1 Nominal Terms Shift the Region of Stability

Original: Central to our discussion will be the notion of stability.
Revised: Central to our discussion will be the notion of stability.

Original: In the study of deterministic dynamicalsystems (for example, taking ⌃ ⇠ ! 0 ), asymptotic stability 2 is guaranteed in linear system (6) when⇢ ( A ) = max i |  i ( A ) | < 1 for all eigenvalues  i ( A ) of A .
Revised: In the study of deterministic dynamicalsystems (for example, taking ⌃ ⇠ ! 0 ), asymptotic stability 2 is guaranteed in linear system (6) when⇢ ( A ) = max i |  i ( A ) | < 1 for all eigenvalues  i ( A ) of A .

Original: Moreover, for the linear dynamical systemwe consider in this section, asymptotic stability implies lim T !
Revised: Moreover, for the linear dynamical systemwe consider in this section, asymptotic stability implies lim T !

Original: 1 L ( ✓ ; T ) is ﬁnite, whereas instabilityimplies lim T !
Revised: 1 L ( ✓ ; T ) is ﬁnite, whereas instabilityimplies lim T !

Original: 1 L ( ✓ ; T ) = 1 for both the deterministic and stochastic system.
Revised: 1 L ( ✓ ; T ) = 1 for both the deterministic and stochastic system.

Original: Thus, stability is anecessary condition for achieving optimizers which, in expectation, reduce the loss over long traininghorizons.
Revised: Thus, stability is anecessary condition for achieving optimizers which, in expectation, reduce the loss over long traininghorizons.

Original: However, we emphasize that stability here is a proxy for convergence that yields simpliﬁeddiscussion and analysis in our setting.
Revised: However, we emphasize that stability here is a proxy for convergence that yields simpliﬁeddiscussion and analysis in our setting.

Original: While stability is essential for guaranteeing favorable optimizer performance, the noisy quadraticsetting also gives us insight into the relationship between the stability of underlying dynamicalsystem and the gradient with respect to meta-trained optimizer P .
Revised: While stability is essential for guaranteeing favorable optimizer performance, the noisy quadraticsetting also gives us insight into the relationship between the stability of underlying dynamicalsystem and the gradient with respect to meta-trained optimizer P .

Original: This gradient is essential formeta-training the optimizer.
Revised: This gradient is essential formeta-training the optimizer.

Original: Under the dynamics (6), the parameters at time t > 0 are  t =
Revised: Under the dynamics (6), the parameters at time t > 0 are  t =

Original: A t  0 +
Revised: A t  0 +

Original: [Equation] t  1 k =

Revised: P t  1 k =

Original: A t  k  1 ( I  A )
Revised: A t  k  1 ( I  A )

Original: ⇠ k and thus  t ⇠ N ( A t  0 , ⌃ t ) where
Revised: ⇠ k and thus  t ⇠ N ( A t  0 , ⌃ t ) where

Original: [Equation]
Revised: [Equation]

Original: Thus, our expected loss is
Revised: Thus, our expected loss is

Original: [Equation]
Revised: [Equation]

Original: [Equation] This loss at time t is polynomial in A with degree 2 t (and note that A is constant across time).
Revised: This loss at time t is polynomial in A with degree 2 t (and note that A is constant across time).

Original: Thus,the gradient of this loss term with respect to P is polynomial in the parameters of A with degree2 t  1 .
Revised: Thus,the gradient of this loss term with respect to P is polynomial in the parameters of A with degree2 t  1 .

Original: As a result, instability of the dynamics of  generally implies instability of the gradient of theloss with respect to P , which in turn implies the expected gradient magnitude and gradient varianceboth diverge for long horizons.
Revised: As a result, instability of the dynamics of  generally implies instability of the gradient of theloss with respect to P , which in turn implies the expected gradient magnitude and gradient varianceboth diverge for long horizons.

Original: How can we guarantee stability of the update dynamics?
Revised: How can we guarantee stability of the update dynamics?

Original: We have stated previously that we require 3⇢ ( A ) < 1 , which we can map to eigenvalues of P via the following result.
Revised: We have stated previously that we require 3⇢ ( A ) < 1 , which we can map to eigenvalues of P via the following result.

Original: We make the technicalassumption that P has real eigenvalues and is diagonalizable; this is satisﬁed in the case that P
Revised: We make the technical

Original: [Figure]
Revised: [Figure]

Original:  is symmetric as is common for preconditioner matrices.
Revised: assumption that P has real eigenvalues and is diagonalizable; this is satisﬁed in the case that P is symmetric as is common for preconditioner matrices.

Original: This assumption allows us to strictlyorder eigenvalues of A , and we write  min ( A )
Revised: This assumption allows us to strictlyorder eigenvalues of A , and we write  min ( A )

Original: =
Revised: =

Original: min
Revised: min

Original: i  i ( A ) and  max ( A ) = max i
Revised: i  i ( A ) and  max ( A ) = max i

Original:  i ( A ) .
Revised:  i ( A ) .

Original: Moreimportantly, this assumption clariﬁes the presentation in this primarily pedagogical discussion.
Revised: Moreimportantly, this assumption clariﬁes the presentation in this primarily pedagogical discussion.

Original: Theorem 1.
Revised: Theorem 1.

Original: Let
Revised: Let

Original: A =
Revised: A =

Original: I
Revised: I

Original:  ↵ H  PH with H symmetric positive deﬁnite, P diagonalizable with real 154 eigenvalues, and ↵  0 .
Revised:  ↵ H  PH with H symmetric positive deﬁnite, P diagonalizable with real 153 eigenvalues, and ↵  0 .

Original: Then ⇢ ( A )  1 if
Revised: Then ⇢ ( A )  1 if

Original: [Equation]

Original: [Equation]

Original: [Equation]
Revised: [Equation]

Original: The proof of this result and all other results is provided in the Appendix A.2.
Revised: The proof of this result and all other results is provided in the Appendix A.2.

Original: There are a few keytakeaways from this result.
Revised: There are a few keytakeaways from this result.

Original: The lower bound (9) highlights the stabilizing effect of the nominaloptimization term.
Revised: The lower bound (9) highlights the stabilizing effect of the nominaloptimization term.

Original: If a particular learned optimizer would result in divergence due to moving uphillin a particular direction (corresponding to the eigenvector associated with the minimum eigenvalue)without the nominal term, the addition of the nominal term ↵  0 gives us additional stability margin.
Revised: If a particular learned optimizer would result in divergence due to moving uphillin a particular direction (corresponding to the eigenvector associated with the minimum eigenvalue)without the nominal term, the addition of the nominal term ↵  0 gives us additional stability margin.

Original: The tradeoff is that for the upper bound (10) we lose stability margin.
Revised: The tradeoff is that for the upper bound (10) we lose stability margin.

Original: Thus, we can in general learnonly smaller P (as measured by the induced spectrum of A ).
Revised: Thus, we can in general learnonly smaller P (as measured by the induced spectrum of A ).

Original: In the design of a learned optimizer, the addition of the nominal term can lead to especially largesteps which result in divergence due to violation of (10).
Revised: In the design of a learned optimizer, the addition of the nominal term can lead to especially largesteps which result in divergence due to violation of (10).

Original: As such, we must be cautious to limit themagnitude of ↵ , and this result motivates regularization of P .
Revised: As such, we must be cautious to limit themagnitude of ↵ , and this result motivates regularization of P .

Original: Within the quadratic setting, we areunlikely to practically hit this upper bound as it corresponds to taking steps large enough to oscillateto divergence.
Revised: Within the quadratic setting, we areunlikely to practically hit this upper bound as it corresponds to taking steps large enough to oscillateto divergence.

Original: For neural network optimization, however, optimization tends to predictably (albeitapproximately, due to minibatch dynamics) hit this upper bound [56].
Revised: For neural network optimization, however, optimization tends to predictably (albeitapproximately, due to minibatch dynamics) hit this upper bound [56].

Original: We illustrate the effect on eigenvalues of incorporating a nominal term in Figure 1.
Revised: We illustrate the effect on eigenvalues of incorporating a nominal term in Figure 1.

Original: We investigate thestability properties of the noisy quadratic model in Figure 2, where we ﬁnd the nominal term makeslearned optimizers more stable, and easier to train.
Revised: We investigate thestability properties of the noisy quadratic model in Figure 2, where we ﬁnd the nominal term makeslearned optimizers more stable, and easier to train.

Original: 4.2 Preconditioners can Stabilize and Simplify the Design of Update Dynamics
Revised: 4.2 Preconditioners can Stabilize and Simplify the Design of Update Dynamics

Original: Adaptive preconditioners have seen widespread use in large-scale machine learning.
Revised: Adaptive preconditioners have seen widespread use in large-scale machine learning.

Original: Methods suchas Adagrad [35], RMSProp [36] and Adam [37] are some of the most fundamental tools in trainingneural networks.
Revised: Methods suchas Adagrad [35], RMSProp [36] and Adam [37] are some of the most fundamental tools in trainingneural networks.

Original: These approaches typically approximate the inverse square root Hessian, and applya transformation
Revised: These approaches typically approximate the inverse square root Hessian, and apply

Original: [Equation]
Revised: [Equation]

Original: (11) where g t is some nominal optimizer.
Revised: where g t is some nominal optimizer.

Original: Typically this approximation is diagonal which correspondsto maintaining some weighted recursive estimate of the magnitude of each element of g .
Revised: Typically this approximation is diagonal which correspondsto maintaining some weighted recursive estimate of the magnitude of each element of g .

Original: This transformation roughly normalizes the step size, making steps isotropic and thus the effective learningrate may be better controlled.
Revised: Thistransformation roughly normalizes the step size, making steps isotropic and thus the effective learningrate may be better controlled.

Original: If we combine this Hessian preconditioner together with our preconditioner P and nominal step size↵ , the resulting dynamics are
Revised: If we combine this Hessian preconditioner together with our preconditioner P and nominal step size↵ , the resulting dynamics are

Original: [Equation]
Revised: [Equation]

Original: Stability in this preconditioned setting is given by the following result.
Revised: Stability in this preconditioned setting is given by the following result.

Original: Lemma 2.
Revised: Lemma 2.

Original: Consider parameter dynamics given by (12) with assumptions on P, H matching Theorem1.
Revised: Consider parameter dynamics given by (12) with assumptions on P, H matching Theorem1.

Original: Then (12) is stable if (9) holds for (12) and
Revised: Then (12) is stable if (9) holds for (12) and

Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: Note, this results in a looser upper bound than (10) if  max ( H ) > 1 .
Revised: Note, this results in a looser upper bound than (10) if  max ( H ) > 1 .

Original: Why apply the preconditioner to the output of the learned optimizer as opposed to the input gradient?
Revised: Why apply the preconditioner to the output of the learned optimizer as opposed to the input gradient?

Original: Normalization at the output of the learned component results in better robustness than normalizingthe input, as it is (relatively) robust to arbitrary initializations of the blackbox term.
Revised: Normalization at the output of the learned component results in better robustness than normalizingthe input, as it is (relatively) robust to arbitrary initializations of the blackbox term.

Original: We expand onthis in both Section 4.3 (for quadratic problems) and Section 5 (for general learned optimizers).
Revised: We expand onthis in both Section 4.3 (for quadratic problems) and Section 5 (for general learned optimizers).

Original: 4.3 Adaptive Nominal Terms Improve Robust Stability
Revised: 4.3 Adaptive Nominal Terms Improve Robust Stability

Original: We have so far motivated choosing a nominal ↵ > 0 to bias the optimizer dynamics toward descent/stability.
Revised: We have so far motivated choosing a nominal ↵ > 0 to bias the optimizer dynamics toward descent/stability.

Original: However, it is likely that we do not want to leave ↵ ﬁxed over the course of inner orouter (meta) training.
Revised: However, it is likely that we do not want to leave ↵ ﬁxed over the course of inner orouter (meta) training.

Original: In inner training, decreasing the learning rate over the course of training hasbeen shown to improve empirical performance (and is often necessary for guaranteeing convergence[57]), and is almost always done when training neural networks.
Revised: In inner training, decreasing the learning rate over the course of training hasbeen shown to improve empirical performance (and is often necessary for guaranteeing convergence[57]), and is almost always done when training neural networks.

Original: For simple nominal gradient estimators and comparably complex blackbox terms, the blackboxmodel should be able to cancel out the nominal term, and induce the effects of reducing ↵ .
Revised: For simple nominal gradient estimators and comparably complex blackbox terms, the blackboxmodel should be able to cancel out the nominal term, and induce the effects of reducing ↵ .

Original: Setting
Revised: Setting

Original: P = P
Revised: P = P

Original: ⇤  ↵ I with nominal gradient term g t =
Revised: ⇤  ↵ I with nominal gradient term g t =

Original: r t yields closed loop dynamics
Revised: r t yields closed loop dynamics

Original: [Equation]
Revised: [Equation]

Original: (14) which is optimal with respect to meta-loss for optimal P ⇤ .
Revised: 14) which is optimal with respect to meta-loss for optimal P ⇤ .

Original: In this subsection, we argue from the pointof view of robust stability that this strategy is suboptimal relative to direct control of the magnitudeof the nominal and blackbox term.
Revised: In this subsection, we argue from the pointof view of robust stability that this strategy is suboptimal relative to direct control of the magnitudeof the nominal and blackbox term.

Original: We introduce a multiplicative error model 4 which captures the sub-optimality of the learned P duringmeta-training.
Revised: We introduce a multiplicative error model 4 which captures the sub-optimality of the learned P duringmeta-training.

Original: Let P =  ˜ P for diagonal disturbance  2 D .
Revised: Let P =  ˜ P for diagonal disturbance  2 D .

Original: We deﬁne this uncertainty set as
Revised: We deﬁne this uncertainty set as

Original: [Equation]
Revised: [Equation]

Original: Within this error model, we can establish conditions for stability in line with the previous subsections.
Revised: Within this error model, we can establish conditions for stability in line with the previous subsections.

Original: We wish to establish robust stability conditions, which guarantee the stability of the dynamical systemfor all realizations of the disturbance.
Revised: We wish to establish robust stability conditions, which guarantee the stability of the dynamical systemfor all realizations of the disturbance.

Original: Lemma 3.
Revised: Lemma 3.

Original: Let A =
Revised: Let A =

Original: I
Revised: I

Original:  ↵ H  PH with P =  ˜ P .
Revised:  ↵ H  PH with P =  ˜ P .

Original: We will assume ˜ P and  are simultaneously 208 diagonalizable, H symmetric positive deﬁnite, and0 < ↵ < 2 /  max ( H ) .
Revised: We will assume ˜ P and  are simultaneouslydiagonalizable, H symmetric positive deﬁnite, and ↵ > 0 .

Original: Then, ⇢ ( A )  1 for all209  
Revised: Then, ⇢ ( A )  1 for all  2 D if

Original: [Equation]
Revised: [Equation]

Original: [Equation]

Original: [Equation]

Original: [Equation]
Revised: [Equation]

Original: These results generally result in the tightening of the margin of stability.
Revised: These results generally result in the tightening of the margin of stability.

Original: If we choose ˜ P
Revised: If we choose ˜ P

Original: = P ⇤  ↵ I ,corresponding to our previously discussed cancellation of the nominal term, our dynamics become
Revised: = P ⇤  ↵ I ,corresponding to our previously discussed cancellation of the nominal term, our dynamics become

Original: [Equation]
Revised: [Equation]

Original: 18) which harms the stability margins.
Revised: which harms the stability margins.

Original: This is well understood by taking  =
Revised: This is well understood by taking  =

Original: (1  ✏ )
Revised: (1  ✏ )

Original: I , for any adversarial✏  1 .
Revised: I , for any adversarial✏  1 .

Original: Then, we have update dynamics
Revised: Then, we have update dynamics

Original: [Equation]
Revised: [Equation]

Original:  yielding dynamics matrix A
Revised: (19) yielding dynamics matrix A

Original: = I  ↵✏ H  (1  ✏ ) P ⇤ H .
Revised: = I  ↵✏ H  (1  ✏ ) P ⇤ H .

Original: Here, ↵✏ H is the excess term in the dynamicscompared to if the nominal term had instead been set to 0.
Revised: Here, ↵✏ H is the excess term in the dynamicscompared to if the nominal term had instead been set to 0.

Original: In this expression, ✏ adversarially perturbsthe system toward instability, resulting in a potentially substantial performance drop.
Revised: In this expression, ✏ adversarially perturbsthe system toward instability, resulting in a potentially substantial performance drop.

Original: Instead, we can select ↵ over the course of both inner and outer training via directly controlling themagnitude of ↵ , P .
Revised: Instead, we can select ↵ over the course of both inner and outer training via directly controlling themagnitude of ↵ , P .

Original: This approach corresponds to  a hyperparameter controller, which has both beenshown to enable automatic control of step sizes [45] over the course of inner training.
Revised: This approach corresponds to the a hyperparameter controller, which has bothbeen shown to enable automatic control of step sizes [45] over the course of inner training.

Original: 4.4 Non-Markovian Optimizers Require Joint Stability
Revised: 4.4 Non-Markovian Optimizers Require Joint Stability

Original: We brieﬂy discuss the role of non-Markovian (or hidden state) dynamics in learned optimizers.
Revised: We brieﬂy discuss the role of non-Markovian (or hidden state) dynamics in learned optimizers.

Original: Anextended discussion is in Appendix A.4.
Revised: Anextended discussion is in Appendix A.4.

Original: The role of momentum and stable hidden states has beeninvestigated in detail in both works on optimization [53, 56] and learned optimization [61].
Revised: The role of momentum and stable hidden states has beeninvestigated in detail in both works on optimization [53, 56] and learned optimization [61].

Original: Wesummarize the key points below.
Revised: Wesummarize the key points below.

Original: The core analysis change required is that we must consider the joint stability of the hidden statedynamics and the parameter dynamics together.
Revised: The core analysis change required is that we must consider the joint stability of the hidden statedynamics and the parameter dynamics together.

Original: Such analysis has been done in the case of Polyakmomentum [62], and has been found to yield less restrictive upper stability margins (thus allowing alarger step size).
Revised: Such analysis has been done in the case of Polyakmomentum [62], and has been found to yield less restrictive upper stability margins (thus allowing alarger step size).

Original: In the general case of blackbox optimizer dynamics, the stability properties of thehidden state update have been investigated in [61], and we expand on this in Section 5.
Revised: In the general case of blackbox optimizer dynamics, the stability properties of thehidden state update have been investigated in [61], and we expand on this in Section 5.

Original: Momentum often accelerates convergence in the full-batch optimization setting.
Revised: Momentum often accelerates convergence in the full-batch optimization setting.

Original: However, it has theadditional beneﬁt of ﬁltering stochastic gradients, yielding better expected descent.
Revised: However, it has theadditional beneﬁt of ﬁltering stochastic gradients, yielding better expected descent.

Original: This ﬁlteringbehavior is an important consideration in any learned optimizer operating in the stochastic setting.
Revised: This ﬁlteringbehavior is an important consideration in any learned optimizer operating in the stochastic setting.

Original:  Designing a Better Learned Optimizer
Revised:  Designing a Better Learned Optimizer

Original: In this section, we present a set of regularization strategies and architectural modiﬁcations for learnedoptimizers that leverage the insights gleaned in the previous section.
Revised: In this section, we present a set of regularization strategies and architectural modiﬁcations for learnedoptimizers that leverage the insights gleaned in the previous section.

Original: We apply these modiﬁcations tothe small_fc_lopt optimizer presented and open sourced in [9].
Revised: We apply these modiﬁcations tothe small_fc_lopt optimizer presented and open sourced in [9].

Original: This optimizer is a small (197 weight)
Revised: This optimizer is a small (197 weight)

Original: MLP which is applied elementwise—it takes inputs such as the parameter value, the gradient, andfeatures such as gradient momentum at different timescales, and outputs an update to the parameter.
Revised: MLP which is applied elementwise—it takes inputs such as the parameter value, the gradient, andfeatures such as gradient momentum at different timescales, and outputs an update to the parameter.

Original: This optimizer is applied in parallel to all parameters in the model being trained.
Revised: This optimizer is applied in parallel to all parameters in the model being trained.

Original: We provide a briefdiscussion of each architectural modiﬁcation or regularization strategy below, and provide an exactdescription of the full architecture in Appendix B.
Revised: We provide a briefdiscussion of each architectural modiﬁcation or regularization strategy below, and provide an exactdescription of the full architecture in Appendix B.

Original: We refer to our optimizer as a stabilized throughample regularization (STAR) learned optimizer.
Revised: We refer to our optimizer as a stabilized throughample regularization (STAR) learned optimizer.

Original: This section also presents a limited selection ofexperiments on in-meta-distribution and out-of-meta-distribution performance, with experimentaldetails available in Appendix C and more results available in Appendix D.
Revised: This section also presents a limited selection ofexperiments on in-meta-distribution and out-of-meta-distribution performance, with experimentaldetails available in Appendix C and more results available in Appendix D.

Original: 5.1 New Design Features in the STAR Optimizer
Revised: 5.1 New Design Features in the STAR Optimizer

Original: Bias toward descent.
Revised: Bias toward descent.

Original: We add a nominal term, as described in Section 4.1.
Revised: We add a nominal term, as described in Section 4.1.

Original: We focus on nominalterms based on Adam [37] and AggMo [63], in part due to the input features to our blackbox optimizercontaining momentum at different timescale (as in AggMo) and running exponential moving average(EMA) gradient magnitude estimates (as in Adam).
Revised: We focus on nominalterms based on Adam [37] and AggMo [63], in part due to the input features to our blackbox optimizercontaining momentum at different timescale (as in AggMo) and running exponential moving average(EMA) gradient magnitude estimates (as in Adam).

Original: We experimentally compare different nominalterms in the appendix.
Revised: We experimentally compare different nominalterms in the appendix.

Original: We add a magnitude controller on the nominal descent term, following the discussion in Section 4.3.
Revised: We add a magnitude controller on the nominal descent term, following the discussion in Section 4.3.

Original: This magnitude controller consists of one additional output head with an exponential nonlinearity onthe small MLP (requiring ﬁve additional weights).
Revised: This magnitude controller consists of one additional output head with an exponential nonlinearity onthe small MLP (requiring ﬁve additional weights).

Original: The combination of magnitude control applied toa nominal optimizer (such as Adam) makes our full nominal term equivalent to a hyperparametercontroller learned optimizer, albeit with a controller that leverages substantial computation reuse withthe blackbox term.
Revised: The combination of magnitude control applied toa nominal optimizer (such as Adam) makes our full nominal term equivalent to a hyperparametercontrolled learned optimizer, albeit with a controller that leverages substantial computation reusewith the blackbox term.

Original: Magnitude control via weight decay.
Revised: Magnitude control via weight decay.

Original: In order to discourage violations of the upper bound onstable eigenvalues in Section 4.1, we apply relatively heavy weight decay ( L 2 regularization) inouter training to the parameters of the blackbox term.
Revised: In order to discourage violations of the upper bound onstable eigenvalues in Section 4.1, we apply relatively heavy weight decay ( L 2 regularization) inouter training to the parameters of the blackbox term.

Original: By directly controlling weight magnitudes(as opposed to controlling magnitudes via regularizing network outputs) we achieve magnituderegularization for arbitrary inputs (for reasonably-sized inputs) and thus achieve better generalization. 
Revised: By directly controlling weight magnitudes(as opposed to controlling magnitudes via regularizing network outputs) we achieve magnituderegularization for arbitrary inputs (for reasonably-sized inputs) and thus achieve better generalization. [Figure] [Figure]

Original: Preconditioner-style normalization.
Revised: Preconditioner-style normalization.

Original: As discussed, we use an adaptive inverse EMA preconditionerin our nominal term to better bias our model toward descent.
Revised: As discussed, we use an adaptive inverse EMA preconditionerin our nominal term to better bias our model toward descent.

Original: As suggested by Section 4.2, we applythe same preconditioner to the output of the blackbox term as well as the network inputs.Because
Revised: As suggested by Section 4.2, we applythe same preconditioner to the output of the blackbox term as well as the network inputs.

Original: [Figure] [Figure] these EMA terms are maintained as inputs to the network already, the expense of this transformationis only the cost of dividing the blackbox output by the preconditioner.
Revised: Becausethese EMA terms are maintained as inputs to the network already, the expense of this transformationis only the cost of dividing the blackbox output by the preconditioner.

Original: Stable hidden states.
Revised: Stable hidden states.

Original: We discussed the importance of considering the stability of the combinedparameter/hidden state dynamics in Section 4.4.
Revised: We discussed the importance of considering the stability of the combinedparameter/hidden state dynamics in Section 4.4.

Original: It is critical to design the hidden state updatedynamics to bias toward stability; indeed, this is a well-known fact in the study of recurrent networksin general [64–66] and has been discussed in [7, 61].
Revised: It is critical to design the hidden state updatedynamics to bias toward stability; indeed, this is a well-known fact in the study of recurrent networksin general [64–66] and has been discussed in [7, 61].

Original: In this paper, as in previous works [7, 9], weuse exponential moving average, momentum-style hidden states which are stable by design.
Revised: In this paper, as in previous works [7, 9], weuse exponential moving average, momentum-style hidden states which are stable by design.

Original: 5.2 The STAR Optimizer Improves Performance
Revised: 5.2 The STAR Optimizer Improves Performance

Original: As described and shown in Figure 3, the STAR optimizer shows improved stability and faster metatraining than baselines, including the small_fc_lopt optimizer from Metz et al.
Revised: As described and shown in Figure 3, the STAR optimizer shows improved stability and faster metatraining than baselines, including the small_fc_lopt optimizer from Metz et al.

Original: [9], which is on theperformance vs. optimizer overhead Pareto frontier.
Revised: [9], which is on theperformance vs. optimizer overhead Pareto frontier.

Original: We refer to small_fc_lopt as “Blackbox” in ourexperiments, and emphasize that the STAR optimizer is architecturally identical to the blackboxoptimizer other than the modiﬁcations documented in the previous subsection.We also compare to ahyperparameter-controller learned optimizer (“Hyperparam”).
Revised: We refer to small_fc_lopt as “Blackbox” inour experiments, and emphasize that the STAR optimizer is architecturally identical to the blackboxoptimizer other than the modiﬁcations documented in the previous subsection.

Original: This optimizer consists of our nominalterm with the magnitude control head, but removing the blackbox term.

Original: Experimental details areprovided in Appendix C and in our open source code.
Revised: Experimental detailsare provided in Appendix C and in our open source code.

Original: Experiments on additional tasks, with othervalues of weight decay, ablations of the primary components of the STAR optimizer, and visualizationof different random seeds as opposed to aggregated statistics, are provided in Appendix D.
Revised: Experiments on additional tasks, with othervalues of weight decay, ablations of the primary components of the STAR optimizer, and visualizationof different random seeds as opposed to aggregated statistics, are provided in Appendix D.

Original: Most interestingly, the STAR optimizer continues to perform strongly even when optimizing foralmost two orders of magnitude more steps than it was applied for during meta-training.
Revised: 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.

Original: It generalizes in this way to longer training runs without sacriﬁcing performance on the meta-trainingloss.
Revised: It generalizesin this way to longer training runs without sacriﬁcing performance on the meta-training loss.

Original: Divergence or poor performance outside of the meta-training setting (such as for the blackboxoptimizer in Figure 3d) is a primary limitation to the application of learned optimizers.
Revised: Divergence or poor performance outside of the meta-training setting (such as for the blackbox optimizerin Figure 3d) is a primary limitation to the application of learned optimizers.

Original: In contrast,the STAR optimizer has controllable and reliable stability behavior outside of the meta-train setting.
Revised: In contrast, the STARoptimizer has controllable and reliable stability behavior outside of the meta-train setting.

Original: While meta-generalization is helped by training across varying settings (including varying themeta-train horizon and the underlying task) [67], this inherent robustness—the ability of a model togeneralize well to new tasks despite not being meta-trained with the explicit goal of generalization—iscritical to achieving broadly applicable optimizers. We emphasize that the usual approach tometa-generalization—large-scale meta-training across many large tasks—is extremely expensive,
Revised: Whilemeta-generalization is helped by training across varying settings (including varying the meta-trainhorizon and the underlying task) [67], this inherent robustness—the ability of a model to generalize well to new tasks despite not being meta-trained with the explicit goal of generalization—is

Original: [Figure]
Revised: [Figure]

Original: typically requiring weeks of computation on millions of dollars of hardware.
Revised: critical to achieving broadly applicable optimizers. We emphasize that the usual approach to metageneralization—large-scale meta-training across many large tasks—is extremely expensive, typicallyrequiring weeks of computation on millions of dollars of hardware.

Original: In contrast, we achievedramatically improved stability fast and for free.
Revised: In contrast, we achieve dramatically improved stability fast and for free.

Original: We refer the reader to Metz et al.
Revised: We refer the reader to Metz et al.

Original: [9] for furtherbaseline results (including heavily-tuned non-learned optimizers) on the tasks evaluated in Figure 3.
Revised: [9] for further baselineresults (including heavily-tuned non-learned optimizers) on the tasks evaluated in Figure 3.

Original: We further explore the inherent generalization abilities of the STAR optimizer in Figure 4.
Revised: We further explore the inherent generalization abilities of the STAR optimizer in Figure 4.

Original: Weapply the optimizer trained on the ﬁrst task (a small MLP applied to Fashion MNIST) to a widevariety of learning tasks including large models such as a transformer [68].
Revised: Weapply the optimizer trained on the ﬁrst task (a small MLP applied to Fashion MNIST) to a widevariety of learning tasks including large models such as a transformer [68].

Original: Remarkably, the STARoptimizer—which was trained in an extremely limited setting—generalizes well to different networksizes, nonlinearities, and datasets (including to language tasks).
Revised: Remarkably, the STARoptimizer—which was trained in an extremely limited setting—generalizes well to different networksizes, nonlinearities, and datasets (including to language tasks).

Original: STARnearly uniformly ourperformsthe baseline blackbox model.
Revised: STAR uniformly ourperforms thebaseline blackbox model.

Original: The performance of STAR (for the correct choice of weight decay) iscomparable to (and occasionally substantially better than, as in Figure 4a) a hyperparameter-tuned Adam model.
Revised: STAR’s performance is comparable (and occasionally substantially better,as in Figure 4a) to heavily hyperparameter-tuned Adam-based model.

Original: These tasks were selected from a subset of evaluation tasks.
Revised: These tasks were selected froma subset of evaluation tasks.

Original: See the Appendix forlearning curves for all generalization experiments.
Revised: See the Appendix for learning curves for all generalization experiments.

Original:  Discussion
Revised:  Discussion

Original: This paper has addressed the role of stability and inductive biases in learned optimizers, and we haveshown that incorporating stabilizing inductive biases results in strong performance both in-distributionand out-of-distribution.
Revised: This paper has addressed the role of stability and inductive biases in learned optimizers, and wehave shown that incorporating stabilizing inductive biases results in strong performance both indistribution and out-of-distribution.

Original: Much further work on injecting stability into learned optimizers remains tobe done, and designing inductive biases for learned optimizers is potentially an exciting new line ofwork that straddles traditional optimization theory and the study of learned optimizers.
Revised: Much further work on injecting stability into learned optimizersremains to be done, and designing inductive biases for learned optimizers is potentially an excitingnew line of work that straddles traditional optimization theory and the study of learned optimizers.

Original: There is adeep literature on the stability of general computation graphs that can, and should, be exploited toallow more ﬁne-grained control of stability properties of optimizers [73–77].
Revised: There is a deep literature on the stability of general computation graphs that can, and should, beexploited to allow more ﬁne-grained control of stability properties of optimizers [73–77].

Original: Informally, stabilization of the computation graph (as we have deﬁned it) is a sufﬁcient conditionto guarantee that an optimizer moves downhill on the training loss landscape (in expectation).
Revised: Informally, stabilization of the computation graph (as we have deﬁned it) is a sufﬁcient conditionto guarantee that an optimizer moves downhill on the training loss landscape (in expectation).

Original: Our motivation from this comes from the convex setting, in which always moving downhill (withappropriate technical conditions) will lead to a global minimum.
Revised: Our motivation from this comes from the convex setting, in which always moving downhill (withappropriate technical conditions) will lead to a global minimum.

Original: However, it is not clear if ourinduced biases are optimal for the training of neural networks, and in general, it is unclear the extentto which neural network training resembles convex optimization problems [55, 78, 79].
Revised: However, it is not clear if ourinduced biases are optimal for the training of neural networks, and in general, it is unclear the extentto which neural network training resembles convex optimization problems [55, 78, 79].

Original: Further studyof which inductive biases are desirable for optimizing networks is necessary.
Revised: Further studyof which inductive biases are desirable for optimizing networks is necessary.

Original: Moreover, while wehave shown strong generalization properties arise from our stability properties, there may be furtherdesirable inductive biases that could be induced to improve generalization, such as for examplebiasing networks toward ﬂat minima [80, 81].Checklist For all authors...( a)
Revised: Moreover, while wehave shown strong generalization properties arise from our stability properties, there may be furtherdesirable inductive biases that could be induced to improve generalization, such as for examplebiasing networks toward ﬂat minima [80, 81].

Original: Do the main claims made in the abstract and introduction accurately reﬂect

Original: the

Original: paper’scontributions and

Original: scope?

Original: [Yes

Original: (b) Did you describe the limitations of your work?

Original: [Yes] We discuss

Original: limitations in

Original: Section

Original: (c) Did you discuss any potential negative societal impacts of your work?

Original: [N/A]

Original: Our work

Original: aims to accelerate network training, and so there are no major negative societal impacts

Original: beyond accelerating the inherent risks of machine learning.

Original: (d) Have you read the ethics review guidelines and ensured that your paper conforms to

Original: them?

Original: [Yes]

Original: If you are including theoretical results...

Original: (a) Did you state the full set of assumptions of all theoretical results?

Original: [Yes]

Original: (b) Did you include complete proofs of all theoretical results?

Original: [Yes] All proofs are available

Original: in the appendix.

Original: If you ran experiments...

Original: (a) Did you include the code, data, and instructions needed to reproduce the main experi-

Original: mental results (either in the supplemental material or as a URL)?

Original: [Yes]

Original: (b) Did you specify all the training details (e.g., data splits, hyperparameters,

Original: how they

Original: were chosen)?

Original: [Yes] These details are available in the appendix.

Original: (c) Did you report error bars (e.g., with respect to the random seed after running exper-

Original: iments multiple times)?

Original: [Yes] Details on the error statistics are also available in the

Original: appendix.

Original: (d) Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)?

Original: [Yes] Also available in the appendix.

Original: If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...

Original: (a) If your work uses existing assets, did you cite the creators?

Original: [Yes] We primarily build

Original: upon the learned optimization package, and cite [9].

Original: We also cite the relevant other

Original: tools and datasets throughout the text.

Original: (b) Did you mention the license of the assets?

Original: [Yes] The learned optimization package is

Original: open source, and speciﬁc license details are available in the Github repository.

Original: (c) Did you include any new assets either in the supplemental material or as a URL?

Original: [No]

Original: (d) Did you discuss whether and how consent was obtained from people whose data you’reusing/curating?

Original: [N/A] We use standard publicly available datasets.

Original: Information on this

Original: is either available at each reference for each dataset, or is not available to us.

Original: (e) Did you discuss whether the data you are using/curating contains personally identiﬁable

Original: information or offensive content?

Original: [No] We use only commonly available datasets;

Original: information on this for these datasets is widely available if it is available at all.

Original: If you used crowdsourcing or conducted research with human subjects...

Original: (a) Did you include the full text of instructions given to participants and screenshots, if

Original: applicable?

Original: [N/A]

Original: (b) Did you describe any potential participant risks, with links to Institutional Review

Original: Board (IRB) approvals, if applicable?

Original: [N/A]