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: Abstract
Revised: Abstract

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

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

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

Original: 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.
Revised: 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.

Original:  Introduction
Revised:  Introduction

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

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

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

Original: Practically, there are 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 are familiar.
Revised: Practically, there are 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 are familiar.

Original: Moreover, the algorithms chosen between are usually derived based on analysis in the convex setting or informal heuristics, and algorithm performance varies substantially across different real world tasks [2].
Revised: Moreover, the algorithms chosen between are usually derived based on analysis in the convex setting or informal heuristics, and algorithm performance varies substantially across different real 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 typically either composed of blackbox function approximators (such as neural networks) or hand-designed functions containing hyperparameters that are learned.
Revised: These optimizers are typically either composed of blackbox function approximators (such as neural networks) or hand-designed functions containing hyperparameters that are learned.

Original: The distinguishing factor of this class of optimizers 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 computationally intractable 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 of performance for the underlying model such as validation loss—has strong similarities to training other 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 other neural network models.

Original: While learned optimizers are potentially transformative, they suffer from several fundamental ﬂaws preventing their broad uptake.
Revised: While learned optimizers are potentially transformative, they suffer from several fundamental ﬂaws preventing 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 during meta-training [10–12].
Revised: 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 36th Conference on Neural Information Processing Systems (NeurIPS 2022). diverge 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 learned optimizers 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 learned optimizers 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 parameter dynamics induced by learned optimizers, via eigenvalue analysis of training in the noisy quadratic 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 noisy quadratic 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 that improve their stability and inductive bias.

Original: These changes include incorporating a tunable contribution from a nominal optimizer with descent guarantees, using large magnitude weight decay on the optimizer’s parameters, and preconditioning the updates generated by the learned optimizer before applying them to the parameters.
Revised: These changes include incorporating a tunable contribution from a nominal optimizer with descent guarantees, using large magnitude weight decay on the optimizer’s parameters, and preconditioning the updates generated by the learned optimizer before applying them to the parameters.

Original: Demonstrate experimentally that the resulting stabilized through ample regularization (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 than used during meta-training, and generalizes well to new tasks which are dissimilar from tasks it was meta-trained on.
Revised: Demonstrate experimentally that the resulting stabilized through ample regularization (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 than used during meta-training, and generalizes well to new tasks which are dissimilar from tasks it was meta-trained on.

Original: For instance, STAR generalizes well to a transformer task with 175 × more parameters and 5 × the number of training steps than the MLP it was meta-trained to optimize.
Revised: For instance, STAR generalizes well to a transformer task with 175 ⇥ more parameters and 5 ⇥ the number of training steps than the MLP it was meta-trained to optimize.

Original: Effective learned optimizer generalization after meta-training on a single task has not previously been demonstrated.
Revised: Effective learned optimizer generalization after meta-training on a single task has 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 learning methods 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 learning methods 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 neural networks) 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 neural networks) to learn learning algorithms [17–19].

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

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

Original: However, works that 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 successful when applied to few-shot learning.
Revised: However, works that 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 successful when 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 optimization algorithms (such as Adagrad [35], RMSProp [36], Adam [37], and many other) which are relatively invariant to hyperparameter choice.
Revised: The difﬁculty of stable training of large models led to the development of adaptive optimization algorithms (such as Adagrad [35], RMSProp [36], Adam [37], and many other) which are relatively invariant to hyperparameter choice.

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

Original: While further methods have been developed for automatic online hyperparameter tuning [38–40], a line of work has focused on meta-learning a neural network that chooses hyperparameters as a function of training history [41–45].
Revised: While further methods have been developed for automatic online hyperparameter tuning [38–40], a line of work has focused on 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 is typically restricted to (in expectation) descent directions.
Revised: This approach results in inherently better stability of learning algorithms as the output is typically restricted to (in expectation) descent directions.

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

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

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

Original: The guided policy search learning strategy is one of several methods that leverage forms of curriculum learning to stabilize learning [50].
Revised: The guided policy search learning strategy is one of several methods that leverage forms of curriculum learning 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 graphs necessitated by meta-training can result in chaotic behavior which makes meta-training extremely noisy.
Revised: As investigated by [8], the long computational graphs necessitated by meta-training can result in chaotic behavior which makes meta-training extremely noisy.

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

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

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

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

Original: In this paper 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 computational overhead) and improve overall stability.
Revised: In this paper 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 computational overhead) 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 φ ∈ Φ ⊆ 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 (cid:96) : 
Revised: We aim to minimize the expectation over data of loss function ` :  ⇥ X !

Original: Φ × X → R which acts on parameters φ and data (point or minibatch) x ∈ X .
Revised: R which acts on parameters  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 θ ∈ Θ , which acts on a history of training statistics such as parameters, loss values, gradients, and training iteration 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 training iteration 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 model with parameters φ , then f ( · ; θ ) maps input state z t (as well as optimizer hidden state which we do not explicitly write) to updated parameters φ t
Revised: When used to train a model with parameters  , then f ( · ; ✓ ) maps input state z t (as well as optimizer hidden state which we do not explicitly write) to updated parameters  t

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

Original: [Equation] [Equation][Equation]
Revised: [Equation] [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: A particular strength of learned optimizers is that they can be trained with respect to validation loss [8, 45].
Revised: A particular 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, so following Metz et al.
Revised: However, this capability is independent and complementary to the topics of this paper, so following 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: t = w t L t ( φ t ) [Equation] (3) under the inner parameter dynamics imposed by (2), where w t is a weighting term (as has been used previously, 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 used previously, 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 on the impacts on meta-learning of learned optimizers.
Revised: In this section, we examine performances differences between optimizers, with a particular focus on the 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 randomly sampled, 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 randomly sampled, 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 ∇ t = ∇ φ t (cid:96) ( φ 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 average of the loss at each timestep, and thus w t = 1 /T for all t .
Revised: Our chosen loss is the 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][Equation]
Revised: [Equation]

Original: where g t corresponds to a nominal term, which we do not aim to meta-learn, and P corresponds to our meta-learned term (thus corresponding to θ in the notation of the previous section), here taking the 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 to our meta-learned term (thus corresponding to ✓ in the notation of the previous section), here taking the form of a dense preconditioner matrix.

Original: In this section, we will primarily consider nominal terms of the form g t = α ∇ t , yielding combined autonomous dynamics
Revised: In this section, we will primarily consider nominal terms of the form g t = ↵ r t , yielding combined autonomous dynamics

Original: [Equation]
Revised: [Equation]

Revised: [Table]

Revised: [Table]

Original: [Figure]
Revised: [Figure]

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 dynamical systems (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 dynamical systems (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 system we consider in this section, asymptotic stability implies lim T →∞ L ( θ ; T ) is ﬁnite, whereas instability implies lim T →∞ L ( θ ; T )
Revised: Moreover, for the linear dynamical system we consider in this section, asymptotic stability implies lim T ! 1 L ( ✓ ; T ) is ﬁnite, whereas instability implies lim T !

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

Original: Thus, stability is a necessary condition for achieving optimizers which, in expectation, reduce the loss over long training horizons.
Revised: Thus, stability is a necessary condition for achieving optimizers which, in expectation, reduce the loss over long training horizons.

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

Original: 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 .
Revised: 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 .

Original: This gradient is essential for meta-training the optimizer.
Revised: This gradient is essential for meta-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 + (cid:80) t − 1 k =

Revised: A t  0 + P t  1 k =

Original: A t − k − 1 ( I − A ) ξ k and thus φ t ∼ N ( A t φ 0 , Σ t ) where
Revised: A t  k  1 ( I  A ) ⇠ 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][Equation]
Revised: [Equation]

Original: 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 degree 2 t − 1 .
Revised: Thus, the gradient of this loss term with respect to P is polynomial in the parameters of A with degree 2 t  1 .

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

Original: [Figure]
Revised: [Figure]

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ρ ( A ) < 1 , which we can map to eigenvalues of P via the following result.
Revised: We have stated previously that we require⇢ ( A ) < 1 , which we can map to eigenvalues of P via the following result.

Original: We make the technical 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.
Revised: We make the technical 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 strictly order eigenvalues of A , and we write λ min ( A )
Revised: This assumption allows us to strictly order 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: More importantly, this assumption clariﬁes the presentation in this primarily pedagogical discussion.
Revised: More importantly, this assumption clariﬁes the presentation in this primarily pedagogical discussion.

Original: Theorem 1.
Revised: Theorem 1.

Original: Let A =
Revised: Let A =

Original: I − αH − PH with H symmetric positive deﬁnite, P diagonalizable with real eigenvalues, and α ≥ 0 .
Revised: I  ↵ H  PH with H symmetric positive deﬁnite, P diagonalizable with real eigenvalues, and ↵  0 .

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

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 key takeaways from this result.
Revised: There are a few key takeaways from this result.

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

Original: If a particular learned optimizer would result in divergence due to moving uphill in 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 uphill in 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 learn only smaller P (as measured by the induced spectrum of A ).
Revised: Thus, we can in general learn only 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 large steps 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 large steps which result in divergence due to violation of (10).

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

Original: 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.
Revised: 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.

Original: For neural network optimization, however, optimization tends to predictably (albeit approximately, due to minibatch dynamics) hit this upper bound [56].
Revised: For neural network optimization, however, optimization tends to predictably (albeit approximately, 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 the stability properties of the noisy quadratic model in Figure 2, where we ﬁnd the nominal term makes learned optimizers more stable, and easier to train.
Revised: We investigate the stability properties of the noisy quadratic model in Figure 2, where we ﬁnd the nominal term makes learned 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 such as Adagrad [35], RMSProp [36] and Adam [37] are some of the most fundamental tools in training neural networks.
Revised: Methods such as Adagrad [35], RMSProp [36] and Adam [37] are some of the most fundamental tools in training neural networks.

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

Original: [Equation]
Revised: [Equation]

Original: where g t is some nominal optimizer.
Revised: where g t is some nominal optimizer.

Original: Typically this approximation is diagonal which corresponds to maintaining some weighted recursive estimate of the magnitude of each element of g .
Revised: Typically this approximation is diagonal which corresponds to 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 learning rate may be better controlled.
Revised: This transformation roughly normalizes the step size, making steps isotropic and thus the effective learning rate 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 Theorem 1.
Revised: Consider parameter dynamics given by (12) with assumptions on P, H matching Theorem 1.

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 normalizing the 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 normalizing the input, as it is (relatively) robust to arbitrary initializations of the blackbox term.

Original: We expand on this in both Section 4.3 (for quadratic problems) and Section 5 (for general learned optimizers).
Revised: We expand on this 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 or outer (meta) training.
Revised: However, it is likely that we do not want to leave ↵ ﬁxed over the course of inner or outer (meta) training.

Original: In inner training, decreasing the learning rate over the course of training has been 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 has been 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 blackbox model 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 blackbox model should be able to cancel out the nominal term, and induce the effects of reducing ↵ .

Original: Setting
Revised: Setting

Original: P = P ∗ − αI with nominal gradient term g t = ∇ t yields closed loop dynamics
Revised: P = P ⇤  ↵ I with nominal gradient term g t = r t yields closed loop dynamics

Original: [Equation]
Revised: [Equation]

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

Original: 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.
Revised: 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.

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

Original: Let P = ∆ ˜ P for diagonal disturbance ∆ ∈ 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 system for all realizations of the disturbance.
Revised: We wish to establish robust stability conditions, which guarantee the stability of the dynamical system for all realizations of the disturbance.

Original: Lemma 3.
Revised: Lemma 3.

Original: Let A =
Revised: Let A =

Original: I − αH − PH with P = ∆ ˜ P .
Revised: I  ↵ H  PH with P =  ˜ P .

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

Original: Then, ρ ( A ) ≤ 1 for all ∆ ∈ D if
Revised: Then, ⇢ ( A )  1 for all  2 D if

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 [Equation]

Original: [Equation]
Revised: [Equation]

Original: 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 − (cid:15) )
Revised: (1  ✏ )

Original: I , for any adversarial (cid:15) ≤ 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: yielding dynamics matrix A

Original: = I − α(cid:15)H − (1 − (cid:15) ) P ∗ H .
Revised: = I  ↵✏ H  (1  ✏ ) P ⇤ H .

Original: Here, α(cid:15)H is the excess term in the dynamics compared to if the nominal term had instead been set to 0.
Revised: Here, ↵✏ H is the excess term in the dynamics compared to if the nominal term had instead been set to 0.

Original: In this expression, (cid:15) adversarially perturbs the system toward instability, resulting in a potentially substantial performance drop.
Revised: In this expression, ✏ adversarially perturbs the 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 the magnitude of α, P .
Revised: Instead, we can select ↵ over the course of both inner and outer training via directly controlling the magnitude of ↵ , P .

Original: This approach corresponds to a hyperparameter controller, which has both been shown to enable automatic control of step sizes [45] over the course of inner training.
Revised: This approach corresponds to a hyperparameter controller, which has both been 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: An extended discussion is in Appendix A.4.
Revised: An extended discussion is in Appendix A.4.

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

Original: We summarize the key points below.
Revised: We summarize the key points below.

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

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

Original: 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 5.
Revised: 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 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 the additional beneﬁt of ﬁltering stochastic gradients, yielding better expected descent.
Revised: However, it has the additional beneﬁt of ﬁltering stochastic gradients, yielding better expected descent.

Original: This ﬁltering behavior is an important consideration in any learned optimizer operating in the stochastic setting.
Revised: This ﬁltering behavior 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 learned optimizers 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 learned optimizers that leverage the insights gleaned in the previous section.

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

Original: This section also presents a limited selection of experiments on in-meta-distribution and out-of-meta-distribution performance, with experimental details available in Appendix B and more results available in Appendix C.
Revised: This section also presents a limited selection of experiments on in-meta-distribution and out-of-meta-distribution performance, with experimental details available in Appendix B and more results available in Appendix C.

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 nominal terms based on Adam [37] and AggMo [63], in part due to the input features to our blackbox optimizer containing momentum at different timescale (as in AggMo) and running exponential moving average (EMA) gradient magnitude estimates (as in Adam).
Revised: We focus on nominal terms based on Adam [37] and AggMo [63], in part due to the input features to our blackbox optimizer containing momentum at different timescale (as in AggMo) and running exponential moving average (EMA) gradient magnitude estimates (as in Adam).

Original: We experimentally compare different nominal terms in the appendix.
Revised: We experimentally compare different nominal terms 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 on the small MLP (requiring ﬁve additional weights).
Revised: This magnitude controller consists of one additional output head with an exponential nonlinearity on the small MLP (requiring ﬁve additional weights).

Original: The combination of magnitude control applied to a 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 with the blackbox term.
Revised: The combination of magnitude control applied to a 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 with 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 on stable eigenvalues in Section 4.1, we apply relatively heavy weight decay ( L 2 regularization) in outer training to the parameters of the blackbox term.
Revised: In order to discourage violations of the upper bound on stable eigenvalues in Section 4.1, we apply relatively heavy weight decay ( L 2 regularization) in outer 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 magnitude regularization 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 magnitude regularization for arbitrary inputs (for reasonably-sized inputs) and thus achieve better generalization.

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

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

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

Original: Because these EMA terms are maintained as inputs to the network already, the expense of this transformation is only the cost of dividing the blackbox output by the preconditioner.
Revised: Because these EMA terms are maintained as inputs to the network already, the expense of this transformation is 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 combined parameter/hidden state dynamics in Section 4.4.
Revised: We discussed the importance of considering the stability of the combined parameter/hidden state dynamics in Section 4.4.

Original: 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].
Revised: 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].

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

Original: 5.2 Overview of the STAR Optimizer
Revised: 5.2 Overview of the STAR Optimizer

Original: We apply these modiﬁcations to the small_fc_lopt optimizer presented and open sourced in [9].
Revised: We apply these modiﬁcations to the 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, and features 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, and features 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, and the only interaction between parameter updates is through tensor-level input features.
Revised: This optimizer is applied in parallel to all parameters in the model being trained, and the only interaction between parameter updates is through tensor-level input features.

Original: We refer to Appendix A of [9] for a complete explanation of the optimizer, but we review the basic details here.
Revised: We refer to Appendix A of [9] for a complete explanation of the optimizer, but we review the basic details here.

Original: The optimizer is parameterized as
Revised: The optimizer is parameterized as

Original: [Equation]
Revised: [Equation]

Original: In this expression, β 1 , β 2 are small constants which in [9] and here are set to 0 .
Revised: In this expression,  1 ,  2 are small constants which in [9] and here are set to 0 .

Original: The terms d ( z t ) and m ( z t ) , 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.
Revised: The terms d ( z t ) and m ( z t ) , 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.

Original: The architecture of the optimizer is an MLP with two hidden layers, each with a width of four.
Revised: The architecture of the optimizer is an MLP with two hidden layers, each with a width of four.

Original: This limits the number of total parameters to 197, yielding high computational efﬁciency versus most learned optimizers.
Revised: This limits the number of total parameters to 197, yielding high computational efﬁciency versus most learned optimizers.

Original: The input features of the network are the parameter value, various parameterwise momentum terms, EMA gradient norms estimates, several adafactor-based [67] features which aggregate tensor-level information, and a parameterization of the training step.
Revised: The input features of the network are the parameter value, various parameterwise momentum terms, EMA gradient norms estimates, several adafactor-based [67] features which aggregate tensor-level information, and a parameterization of the training step.

Original: Our modiﬁed update takes the form f ( z t ) =
Revised: Our modiﬁed update takes the form f ( z t ) =

Original: f b ( z t ) + f g ( z t ) where f g ( · ) is the nominal term and f b ( · ) is the blackbox term.
Revised: f b ( z t ) + f g ( z t ) where f g ( · ) is the nominal term and f b ( · ) is the blackbox term.

Original: The nominal term is structured as
Revised: The nominal term is structured as

Original: [Equation]
Revised: [Equation]

Original: where m g ( · ) is a magnitude term for the nominal term and g ( z t ) is the nominal term; in our experiments we use a combined AggMo [63] and Adam [37] for this term.
Revised: where m g ( · ) is a magnitude term for the nominal term and g ( z t ) is the nominal term; in our experiments we use a combined AggMo [63] and Adam [37] for this term.

Original: Critically, this term corresponds directly to a descent direction without being passed through a network, biasing the update toward descent.
Revised: Critically, this term corresponds directly to a descent direction without being passed through a network, biasing the update toward descent.

Original: The blackbox term is structured as
Revised: The blackbox term is structured as

Original: [Equation]
Revised: [Equation]

Original: (22) where v ( z t ) is our speciﬁed preconditioner term and m b ( · ) is the blackbox magnitude term.
Revised: (22) where v ( z t ) is our speciﬁed preconditioner term and m b ( · ) is the blackbox magnitude term.

Original: The difference between this parameterization and (20) is the normalization term.
Revised: The difference between this parameterization and (20) is the normalization term.

Original: In these terms, β 1 , β 2 , β 3 , βare hand-speciﬁed constants, and the neural network has three output heads.
Revised: In these terms,  1 ,  2 ,  3 , are hand-speciﬁed constants, and the neural network has three output heads.

Original: The addition of the extra head as well as the multiplicative factor on that term adds an extra ﬁve parameters to the original 197-parameter optimizer, as well as one more for β 1 scaling the entire nominal term.
Revised: The addition of the extra head as well as the multiplicative factor on that term adds an extra ﬁve parameters to the original 197-parameter optimizer, as well as one more for  1 scaling the entire nominal term.

Original: 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 β 3 to account for this change in magnitude.
Revised: 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  3 to account for this change in magnitude.

Original: [Figure]
Revised: [Figure]

Original: [Figure]
Revised: [Figure]

Original: 5.3 The STAR Optimizer Improves Performance
Revised: 5.3 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 the performance vs. optimizer overhead Pareto frontier.
Revised: [9], which is on the performance vs. optimizer overhead Pareto frontier.

Original: We refer to small_fc_lopt as “Blackbox” in our experiments, and emphasize that the STAR optimizer is architecturally identical to the blackbox optimizer other than the modiﬁcations documented in the previous subsection.
Revised: We refer to small_fc_lopt as “Blackbox” in our experiments and emphasize that the STAR optimizer is architecturally identical to the blackbox optimizer other than the modiﬁcations documented in the previous subsection.

Original: We also compare to a hyperparameter-controller learned optimizer (“Hyperparam”).
Revised: We also compare to a hyperparameter-controller learned optimizer (“Hyperparam”).

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

Original: Experimental details are provided in Appendix B and in our open source code.
Revised: Experimental details are provided in Appendix B and in our open source code.

Original: 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.
Revised: 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.

Original: 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.
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-training loss.
Revised: It generalizes in 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 blackbox optimizer 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 optimizer in 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 STAR optimizer 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 the meta-train horizon and the underlying task) [72], 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 critical to achieving broadly applicable optimizers.
Revised: While meta-generalization is helped by training across varying settings (including varying the meta-train horizon and the underlying task) [72], 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 critical to achieving broadly applicable optimizers.

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

Original: In contrast, we achieve dramatically 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 further baseline results (including heavily-tuned non-learned optimizers) on the tasks evaluated in Figure 3.
Revised: [9] for further baseline results (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: We apply the optimizer trained on the ﬁrst task (a small MLP applied to Fashion MNIST) to a wide variety of learning tasks including large models such as a transformer [73].
Revised: We apply the optimizer trained on the ﬁrst task (a small MLP applied to Fashion MNIST) to a wide variety of learning tasks including large models such as a transformer [73].

Original: Remarkably, the STAR
Revised: Remarkably, the STAR

Original: [Figure]
Revised: [Figure]

Original: optimizer—which was trained in an extremely limited setting—generalizes well to different network sizes, nonlinearities, and datasets (including to language tasks).
Revised: optimizer—which was trained in an extremely limited setting—generalizes well to different network sizes, nonlinearities, and datasets (including to language tasks).

Original: STAR nearly uniformly ourperforms the baseline blackbox model.
Revised: STAR nearly uniformly ourperforms the baseline blackbox model.

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

Original: Adam model.
Revised: Adam model.

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

Original: See the Appendix for learning 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 have shown that incorporating stabilizing inductive biases results in strong performance both in-distribution and out-of-distribution.
Revised: This paper has addressed the role of stability and inductive biases in learned optimizers, and we have shown that incorporating stabilizing inductive biases results in strong performance both in-distribution and out-of-distribution.

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

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

Original: Informally, stabilization of the computation graph (as we have deﬁned it) is a sufﬁcient condition to 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 condition to 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 (with appropriate technical conditions) will lead to a global minimum.
Revised: Our motivation from this comes from the convex setting, in which always moving downhill (with appropriate technical conditions) will lead to a global minimum.

Original: 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, 80].
Revised: 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, 80].

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

Original: Moreover, while we 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 ﬂat minima [81, 82].
Revised: Moreover, while we 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 ﬂat minima [81, 82].

Original: Acknowledgments and Disclosure of Funding
Revised: Acknowledgments and Disclosure of Funding

Original: We thank Daniel Freeman and Rohan Sinha for helpful conversations and comments in the development of this work.A The Noisy Quadratic Setting: Additional Details
Revised: We thank Daniel Freeman and Rohan Sinha for helpful conversations and comments in the development of this work.

Original: In this section we extend our discussion of the noisy quadratic model (NQM).

Original: We ﬁrst discuss stability in the NQM.

Original: We then provide proofs for the results in Section 4.

Original: We also extend our discussion of robust stability and the stability of models with hidden states.

Original: A.1 Stability

Original: In this subsection we expand on the short discussion of key stability results in the body of the paper.

Original: We will primarily discuss stability of the nominal system.

Original: This is motivated by a few results.

Original: Lemma 4.

Original: Consider the dynamical system deﬁned by (6) with all eigenvalues of H ﬁnite.

Original: We write the expected total loss for horizon T as L ( θ ; T ) .

Original: Consider also the deterministic system given by (6) , in which Σ ξ → 0 .

Original: We will write the total loss under these dynamics as L ( θ ; T ) .

Original: Then,

Original: [Equation]

Original: for all T > 0 .

Original: Proof.

Original: Note that

Original: [Equation]

Original: Note that H and Σ ξ are PSD, as is Σ t .

Original: Thus, the trace term is strictly non-negative for all time.

Original: Lemma 5.

Original: Let ρ ( A ) < 1 with the total loss deﬁned as previously (again with all eigenvalues of H ﬁnite), and let A be diagonalizable.

Original: Then,

Original: L ( θ ; T ) < ∞ (A.3) for all T > 0 .

Original: Proof.

Original: First, note that L ( θ ; T ) ≤ L ( θ ; T + 1) for all T as both terms in the sum are non-negative.

Original: Thus, it sufﬁces to show lim T →∞ L ( θ ; T ) < ∞ .

Original: We will address the two terms in the sum in (8) in order.

Original: First, note that

Original: [Equation]

Original: where V is a matrix with columns corresponding to eigenvectors of A and Λ is a diagonal matrix containing the eigenvalues of A .

Original: We will write the stacked eigenvalues of A as λ ∈ R N .

Original: Then, we can rewrite the ﬁrst loss term as

Original: [Equation]

Original: where z 0 =

Original: V − 1 φ 0 , Z = diag ( z 0 ) , and where the second line uses the fact that diag ( z ) λ = diag ( λ ) z .

Original: Thus, we have

Original: [Equation]

Original: Note that | λ ti | = | λ i | t .

Original: We will exchange the sums over i and t and note that, because | λ i | < 1 for all i , ∞ (cid:88)

Original: [Equation]

Original: |

Original: [Equation]

Original: (A

Original: [Equation]

Original: 8) for each i , and thus the sum over i is ﬁnite.

Original: Thus, we have shown the ﬁrst term in the loss is ﬁnite—not including the factor of 1 /T .

Original: As T → ∞ , considering the factor of 1 /T , this ﬁrst term vanishes.

Original: We now have to bound the second term in the sum over time in (8).

Original: Note that

Original: [Equation]

Original: and 0 ≤ Σ t , as it is PSD and symmetric by construction and thus this term in the loss is non-negative.

Original: We now need to construct an upper bound.

Original: Note that showing the boundedness of the two-norm of Σ t implies the trace is also bounded.

Original: Moreover, we have

Original: [Equation]

Original: and thus it is sufﬁcient to show the boundedness of the

Original: A t − k − 1 (cid:107)

Original: [Equation]

Original: Because ρ ( A ) < 1 , we have ρ ( A k ) < 1 ; this yields boundedness of the inﬁnite sum by following the same arguments as the ﬁrst part of the proof.

Original: We have thus shown that stability of the nominal model implies non-divergence of the stochastic model, and that the deterministic model provides a lower bound on the performance of the stochastic model.

Original: So, based on the above results, instability of the deterministic model (in which the total cost goes to inﬁnity [83] in the inﬁnite horizon) implies inﬁnite cost for the stochastic model and stability of the deterministic model implies convergent (non-inﬁnite) cost for the stochastic model.

Original: Thus, we discuss the behavior of the deterministic system as a proxy objective for the behavior of the stochastic system.

Original: A.2 Proofs from Section

Original: Proof of Theorem 1.

Original: First, note that (9) corresponds to positive semi-deﬁniteness of αI + P .

Original: Due to this, as well as H > 0 , we have λ

Original: i (( αI + P ) H ) real and non-negative for all i .

Original: Thus,

Original: [Equation]

Original: 15) if λ i (( αI + P ) H ) ≤ 2 for all i .

Original: Due to the positive (semi-)deﬁniteness of αI + P and H , we have

Original: [Equation]

Original: max

Original: [Equation]

Original: [Equation]

Original: 16) and thus

Original: [Equation]

Original: [Equation]

Original: 17) ensures (10).

Original: Due to the positivity of α , we have

Original: [Equation]

Original: ) which gives our desired result.

Original: Proof of Lemma 2.

Original: When we combine the Hessian preconditioner H − 12 together with our preconditioner P and nominal step size α , the resulting dynamics are

Original: [Equation]

Original: (A

Original: [Equation]

Original: 19) which yields a nominal term − αH − 12 ∇ t .

Original: The second term is

Original: [Equation]

Original: where, note, H − 12 PH 12 is a similarity transformation of P due to the positive deﬁniteness of H , and thus

Original: [Equation]

Original: with λ i ( αI + H − 12 PH 12 ) =

Original: α + λ i ( P ) .

Original: Plugging in for ∇ t , we have H − 12 ∇ t =

Original: H 12 ( φ t − ξ t ) which yields λ

Original: i ( H 12 ) =

Original: (cid:112) λ i ( H ) , given these results and following the bound on the maximum eigenvalue in Theorem 1 proves the result.

Original: Proof of Lemma 3.

Original: Substituting for P in Theorem 1, we have

Original: [Equation]

Original: We have assumed ∆ and ˜ P are simultaneously diagonalizable, and thus

Original: [Equation]

Original: for eigenvector matrices V , and thus λ i (∆ ˜ P ) =

Original: λ j (∆) λ k ( ˜ P ) for all i and some j, k .

Original: Note that because ∆ is diagonal, the eigenvalues correspond to the diagonal entries.

Original: We will consider the minimum eigenvalue bound ﬁrst.

Original: Consider two cases.

Original: First, we will consider the λ min ( ˜ P ) ≥ 0 case.

Original: As we know α > 0 and λ j (∆) > 0 for all j , this inequality is always satisﬁed.

Original: In the second case, λ min ( ˜ P ) < 0 , the right hand side has

Original: [Equation]

Original: Thus, to summarize, we have shown the above lower bounds the minimum eigenvalue of ∆ ˜ P and thus satisfying it results in α

Original: + λ i (∆ ˜ P ) ≥ 0 for all i .

Original: We will now plug in to (10) with P = ∆ ˜ P .

Original: Note that, based on the simultaneous diagonalizability, λ i ( αI + ∆ ˜ P )

Original: =

Original: α + λ j (∆) λ k

Original: ( ˜ P ) .

Original: Again, we consider two cases.

Original: If λ min ( ˜ P ) ≤ 0 , then the bound is immediately satisﬁed.

Original: Thus, the largest bound is provided by

Original: [Equation]

Original: and plugging this into (10) yields the desired result.

Original: A.3 Robust Stability

Original: Our robust stability analysis in Section 4.3 relies on a diagonal multiplicative error model, with close connections to D-stability [ 58 ] .

Original: There are many possible robustness conditions and forms of stability analysis within neural network models, most of which originate in the control theory and nonlinear dynamics communities [84, 85].

Original: We highlight that a purely additive model (as opposed to the multiplicative model we use) would not necessarily result in the multiplicative factor in (16) and (17), and instead would result in an additive factor.

Original: Moreover, a term of the form (cid:15)αH (as appears in our dynamics with cancellation) would not necessarily be present for an additive error model.

Original: Fundamentally, however, we believe our multiplicative error model is a reasonable one, with a combined multiplicative/additive model also being reasonable and general, and resulting in similar stability bounds.

Original: A.4 Non-Markovian Optimizers

Original: Analysis of the stability of optimizers with momentum shows that the addition of momentum improves stability margins [56] and helps counteract noise [86].

Original: We refer the reader to [56], Appendix B for an analysis of stability margins for Polyak and Nesterov momentum, and [87] for a comprehensive exploration of the dynamics of Polyak-style momentum.

Original: This improvement in the stability margin is intuitive: if one considers the marginal stability case in which exact oscillation occurs on the sharpest eigenmode, momentum will act in the downhill direction at each timestep (for reasonably chosen hyperparameters) and move the system into a strictly stable regime.

Original: However, we emphasize that stability of the combined system requires analysis of the second order (double integrator) system induced by the interaction of the momentum hyperparameters with the stepsize parameters.

Original: This analysis is critical in the case of more expressive learned optimizers which may have recurrence.

Original: We believe this case—which is computationally and analytically challenging—is a primary direction of future work.

Original: B Experimental Details

Original: In this section we provide details of the experiments featured in the body of the paper.

Original: Our experiments use JAX [88].

Original: For all experiments we use PES [51] with a truncation length of 50 , with a standard deviation of 0 .

Original: We always target the mean training loss (across all timesteps of the inner optimization).

Original: In all experiments we use AdamW [89] as our meta-optimizer; this is identical to

Original: Adam other than how the weight decay is handled.

Original: We use zero weight decay except when explicitly stated otherwise, and apply gradient clipping of 1 . 0 .

Original: B.1 Noisy Quadratic Model

Original: In our NQM experiments, we considered a two-dimensional quadratic program in which

Original: [Equation]

Original: [Equation]

Original: with the initial state φ 0 ∼ N (0 , 10 I ) and the noise ξ t ∼ N (0 , I ) for all t .

Original: The curves in Figure 2a correspond to values

Original: T ∈ { 1 , 2 , 3 , 4 , 5 , 10 , 25 , 50 , 100 , 150 , 200 , 250 , 300 , 400 , 500 , 600 , 700 , 800 , 900 , 1000 } .

Original: (A.31)

Original: The plot visualizes the trace of the empirical gradient variance matrix

Original: [Equation]

Original: where ∇ i is sample i (with respect to the noise instantiation) of the gradient of the loss with respect to the parameters and

Original: [Equation]

Original: The ﬁgure is plotted for 50 values of α (with P = 0 ) at 500 seeds per sample.

Original: In Figures 2 b,c we consider a total episode length of T = 50 .

Original: In this case, PES reduces to simple ES [8].

Original: Plots show the mean over ﬁve seeds and the one standard error conﬁdence intervals.

Original: In Figure 2 b, the dotted line corresponds to the larger eigenvalues and the solid line corresponds to the smaller one.

Original: In Figures 2 b,c we apply EMA smoothing with coefﬁcient 0 .

Original: 95 to each curve (including the conﬁdence intervals).

Original: B.2 In-Distribution Experiments

Original: The two tasks considered for our in-distribution experiments are an MLP applied to Fashion MNIST and a ConvNet applied to CIFAR10.

Original: The MLP model consists of two hidden layers with 128 hidden units each and ReLU activations.

Original: The

Original: ConvNet consists of three hidden layers with 3 × 3 kernels and ReLU activations.

Original: The layers of 32, 64, and 64 units respectively and strides of 2, 1, and 1 respectively.

Original: The conv layers are then followed by a fully connected layer.

Original: For both tasks, batch sizes of 128 are used.

Original: For more details on the tasks, we refer the reader to [ 9 ] .

Original: In our experiments, the “Hyperparam” optimizer corresponds to our architecture in which the blackbox output head is removed, and thus the optimizer corresponds to the nominal term with the magnitude control head.

Original: The “Blackbox” optimizer corresponds to removing the nominal gradient term, which also corresponds to the small_fc_lopt optimizer of [9].

Original: In our optimizers we use weight decay with AdamW, where the weight decay is speciﬁed as a multiplier on the outer learning rate.

Original: We use an outer learning rate of 10 − 4 for all of our experiments, for both tasks.

Original: The training horizon for the tasks are 2000 steps.

Original: In all of our experiments, we set β i = 10 − 3 for all i .

Original: In all plots, we apply EMA smoothing with a coefﬁcient of 0 . 98 .

Original: B.3 Out-of-Distribution Experiments

Original: We apply our optimizer, trained on the MLP/Fashion MNIST task, to a wide variety of other ML tasks.

Original: In particular, we apply it to 16 different tasks that span a wide range of architectures and datasets, from relatively similar to the training task—larger MLPs on CIFAR10, for example—to radically different tasks such as training models with ES gradients, recurrent models such as LSTMs, and transformers.

Original: These tasks are taken from the learned_optimization package [9], which is open source 6 .

Original: We add two baselines on these tasks: an Adam model and an NAdamW (Nesterov-accelerated AdamW) model.

Original: These baselines were also included in [9], and the tasks are all available in the learned optimization package.

Original: We note that the learning rate of the Adam model has been tuned for each task via a coarse search (half orders of magnitude steps), while the NAdamW model has been tuned for each task via random search.

Original: In particular, 1000 random hyperparameter values were evaluated for each task.

Original: Thus, the Adam baseline represents a reasonable baseline, comparable to standard training practice, while the NAdamW baseline represents an extremely strong baseline that is unachievable in typical hyperparameter tuning.

Original: C Further Experimental Results

Original: In this section we present further experimental results which extend and complement the results from the body of the paper.

Original: The details of these experiments are the same as previous unless otherwise stated.

Original: [Figure]

Original: [Figure]

Original: C.1 Alternative Views of Main Results

Original: Figure App.1 shows individual seeds underlying Figure 3.

Original: There are two things to note about these ﬁgures.

Original: First, seeds which perform poorly often perform very poorly, diverging or oscillating around losses much worse than other seeds.

Original: This is to be expected from dynamical systems run for very long horizons, in which small amounts of instability are ampliﬁed over long timeframes.

Original: In spite of this, the results in Figure App.1 validate our results even if outliers are ignored.

Original: C.2 Ablations

Original: Figures App.2, App.

Original: [Figure]

Original: and App.4 show ablations of the various components of the STAR optimizer.

Original: Figure App.2 shows other values of weight decay in outer training.

Original: For values between 0.1 and 0.

Original: [-(times the learning rate) as we use in our experiments, a good tradeoff between stabilizing behavior outside of the training horizon and avoid over-damping the dynamics (as in the case of 1 . 
-]
Revised: [-(times the learning rate) as we use in our experiments, a good tradeoff between stabilizing behavior outside of the training horizon and avoid over-damping the dynamics (as in the case of 1 . 
-]

Original: weight decay) is achieved.

Original: We note that the poor training performance of 0 .

Original: weight decay in Figure App.2b is due to an outlier that did not converge well in meta-training and low weight decay typically does not substantially harm meta-training.

Original: Figure App.3 shows the impact of the Adam-style normalizing preconditioner.

Original: Generally, better meta-losses are achieved with the preconditioning compared to without; it is less clear what the beneﬁt for generalization is, with performance of with/without the preconditioner being mixed.

Original: Generally, due to the strong improvements to in-task performance and limited degredation in out-of-distribution performance, we advocate for this term to be included.

Original: Figure App.4 shows the impact of removing the nominal term.

Original: Generally, without the nominal term, meta-training is slower and converges to worse solutions.

Original: Interestingly, for the ﬁnal trained optimizer, performance on the Fashion MNIST MLP seems comparable between all models; this is very much not the case for the CIFAR model.

Original: Generally, we have found performance at all points in meta-training and generalization to be uniformly improved by the addition on nominal terms, and strongly endorse the inclusion of these terms.

Original: C.3 Full Generalization Experiments

Original: We evaluate the learned optimizers trained on Fashion MNIST on a wide variety of tasks taken from the learned_optimization package, with results visualized in Figures App.5 and App.6.

Original: Note that both of these ﬁgures are the same other than the speciﬁc tasks evaluated, and they are split for space reasons.

Original: For a description of each task, see the source.

Original: Again, we note that the NAdamW baseline is an extremely strong baseline tuned to each task, whereas our model was trained only on the Fashion

Original: MNIST task.

Original: [Figure]

Original: MLP

Original: [Figure]

Original: [Figure]

Original: [Figure]

Original: [Figure]

Original: [Figure]