MetaOptimize: A Framework for Optimizing Step Sizes and Other Meta-parameters — part 1/10

Cleaned-text code points [0, 12223). Bibliography omitted here; source.md includes it.

# MetaOptimize: A Framework for Optimizing Step Sizes and Other Meta-parameters

###### Abstract

We address the challenge of optimizing meta-parameters (hyperparameters) in machine learning, a key factor for efficient training and high model performance. Rather than relying on expensive meta-parameter search methods, we introduce MetaOptimize: a dynamic approach that adjusts meta-parameters, particularly step sizes (also known as learning rates), during training. More specifically, MetaOptimize can wrap around any first-order optimization algorithm, tuning step sizes on the fly to minimize a specific form of regret that considers the long-term impact of step sizes on training, through a discounted sum of future losses. We also introduce lower-complexity variants of MetaOptimize that, in conjunction with its adaptability to various optimization algorithms, achieve performance comparable to those of the best hand-crafted learning rate schedules across diverse machine learning tasks.

###### Keywords:

## 1 Introduction

Optimization algorithms used in machine learning involve meta-parameters (i.e., hyperparameters) that substantially influence their performance. These meta-parameters are typically identified through a search process, such as grid search or other trial-and-error methods, prior to training. However, the computational cost of this meta-parameter search is significantly larger than that of training with optimal meta-parameters ( Dahl et al., 2023 ; Jin, 2022 ) . Meta-parameter optimization seeks to streamline this process by concurrently adjusting meta-parameters during training, moving away from the computationally expensive and often sub-optimal trial and error search methods.

Meta-parameter optimization is particularly important in continual learning ( De Lange et al., 2021 ) , where continually changing environments or evolving loss functions necessitate adaptation of meta-parameters, such as step sizes, to track time-varying optima rather than settling on a static value.

In this work, we propose MetaOptimize , a general framework for optimizing meta-parameters to minimize a form of regret that explicitly accounts for the long-term influence of step sizes on future loss. Although this framework can handle various meta-parameters, we concentrate on step sizes as they are ubiquitous and crucial in practice.

MetaOptimize offers additional advantages beyond reducing search overhead. First, it enables dynamic step-size updates during training, potentially speeding the learning process. Traditional methods typically rely on manually designed learning rate schedules (e.g., initial increase followed by decay ( Amid et al., 2022 ) ), whereas MetaOptimize automatically discovers similar patterns.

Second, adapting step sizes across different network blocks (e.g., layers or neurons) can improve performance ( Singh et al., 2015 ; Howard & Ruder, 2018 ) , yet manually tuning such blockwise step sizes is impractical for large networks. By design, MetaOptimize handles these blockwise adjustments.

The concept of meta step-size optimization dates back to ( Kesten, 1958 ) , Delta-bar-Delta ( Sutton, 1981 ; Jacobs, 1988 ) , and its incremental variant, IDBD ( Sutton, 1992 ) . Numerous methods have emerged over the years (Section 8 ). This work distinguishes itself from prior efforts through the following key aspects: • We introduce a formal approach to step-size optimization by minimizing a specific form of regret, essentially a discounted sum of future losses, and demonstrate how to do this causally via the MetaOptimize framework.

• MetaOptimize is general and can wrap around any first-order optimization algorithm (the base update ), such as SGD, RMSProp ( Hinton, 2012 ) , Adam ( Kingma & Ba, 2014 ) , or Lion ( Chen et al., 2023 ) , while optimizing step sizes via a separate first-order method (the meta update ), such as SGD, Adam, RMSProp, or Lion.

• We develop approximation methods (Section 6 ) that, when incorporated into MetaOptimize, yield computationally efficient algorithms outperforming state-of-the-art automatic hyperparameter optimization methods on various stationary and continual (non-stationary) benchmarks (see Section 7 ).

• We show that some existing methods (like IDBD and its extensions, and hypergradient descent ( Baydin et al., 2017 ) ) are specific instances or approximations within the MetaOptimize framework (Section 5 ).

## 2 Problem Setting

We introduce a general continual optimization setting that, for a given sequence of loss functions f t ​ ( ⋅ ) : ℝ n → ℝ f_{t}(\cdot):\mathbb{R}^{n}\to\mathbb{R} , t = 0 , 1 , 2 , … t=0,1,2,\ldots , aims to find a sequence of weight vectors 𝒘 1 , 𝒘 2 , 𝒘 3 , … \boldsymbol{w}_{1},\boldsymbol{w}_{2},\boldsymbol{w}_{3},\ldots that minimize a discounted sum of future losses: F t γ ​ = def ​ ( 1 − γ ) ​ ∑ τ > t γ τ − t − 1 ​ f τ ​ ( 𝒘 τ ) , F^{\gamma}_{t}\overset{\text{\tiny def}}{=}(1-\gamma)\,\sum_{\tau>t}\gamma^{\tau-t-1}f_{\tau}(\boldsymbol{w}_{\tau}), (1) where γ ∈ [ 0 , 1 ) \gamma\in[0,1) is a fixed discount factor, typically close to 1 1 , called the discount factor . For stationary supervised learning, f t f_{t} are i.i.d. samples from the same distribution, so minimizing F t γ F^{\gamma}_{t} promotes rapid reduction of the expected loss.

Consider an arbitrary first-order optimization algorithm (e.g., SGD, RMSProp, Adam, or Lion) for updating 𝒘 t \boldsymbol{w}_{t} . At time t t , it takes the gradient ∇ f t ​ ( 𝒘 t ) \nabla f_{t}(\boldsymbol{w}_{t}) of the current loss, along with an m m -dimensional meta-parameter vector 𝜷 t \boldsymbol{\beta}_{t} , to update 𝒘 t \boldsymbol{w}_{t} and possibly some internal variables 𝒙 ~ t \tilde{\boldsymbol{x}}_{t} (e.g., momentum in Adam). Denoting 𝒙 t ​ = def ​ Stack ⁡ ( 𝒘 t , 𝒙 ~ t ) \boldsymbol{x}_{t}\overset{\text{\tiny def}}{=}\operatorname{Stack}(\boldsymbol{w}_{t},\,\tilde{\boldsymbol{x}}_{t}) and calling this update rule Alg base \operatorname{Alg}_{\textrm{base}} , we have 𝒙 t + 1 = Alg base ⁡ ( 𝒙 t , ∇ f t ​ ( 𝒘 t ) , 𝜷 t ) . \boldsymbol{x}_{t+1}=\operatorname{Alg}_{\textrm{base}}(\boldsymbol{x}_{t},\nabla f_{t}(\boldsymbol{w}_{t}),\boldsymbol{\beta}_{t}). (2) The goal of MetaOptimize is to determine a sequence 𝜷 t \boldsymbol{\beta}_{t} such that plugging them into the above base update yields a trajectory { 𝒘 t } \{\boldsymbol{w}_{t}\} minimizing F t γ F^{\gamma}_{t} .

Step-size adaptation is a natural special case: at each step t t , the m m -dimensional 𝜷 t \boldsymbol{\beta}_{t} defines an n n -dimensional step-size vector 𝜶 t \boldsymbol{\alpha}_{t} via some fixed function σ : ℝ m → ℝ n \sigma:\mathbb{R}^{m}\to\mathbb{R}^{n} , i.e., 𝜶 t = σ ⁡ ( 𝜷 t ) . \boldsymbol{\alpha}_{t}=\sigma(\boldsymbol{\beta}_{t}). (3) A good choice for σ ⁡ ( ⋅ ) \sigma(\cdot) is exponential, ensuring 𝜶 t \boldsymbol{\alpha}_{t} is always positive and making multiplicative changes in 𝜶 t \boldsymbol{\alpha}_{t} correspond to additive changes in 𝜷 t \boldsymbol{\beta}_{t} ( Sutton, 1992 ) . By partitioning the network weights into blocks, we can learn a shared scalar step-size per block, or even a unique step-size per weight, all handled automatically by MetaOptimize.

## 3 Forward and Backward Views

Because F t γ F^{\gamma}_{t} depends on future losses, minimizing it causally requires an alternative view. Suppose hypothetically we had oracle access to future information (i.e., future loss values and weights). We could update 𝜷 t + 1 = 𝜷 t − η ​ d d ​ 𝜷 t ​ F t γ = 𝜷 t − η ⁡ ( 1 − γ ) ​ ∑ τ > t γ τ − t − 1 ​ d d ​ 𝜷 t ​ f τ ​ ( 𝒘 τ ) , \begin{split}\boldsymbol{\beta}_{t+1}&=\boldsymbol{\beta}_{t}-\eta\frac{\mathrm{d}\,}{\mathrm{d}\,\boldsymbol{\beta}_{t}}F^{\gamma}_{t}\\ &=\boldsymbol{\beta}_{t}-\eta\,(1-\gamma)\,\sum_{\tau>t}\gamma^{\tau-t-1}\frac{\mathrm{d}\,}{\mathrm{d}\,\boldsymbol{\beta}_{t}}f_{\tau}(\boldsymbol{w}_{\tau}),\end{split} (4) where η \eta is a meta step-size. This forward-view update, however, is not causal because we do not have the required future information at time t t .

To address this, we adopt an eligibility-trace-style approach from reinforcement learning ( Sutton, 1988 ; Sutton & Barto, 2018 ) , introducing a backward-view update: 𝜷 τ + 1 ← 𝜷 τ − η ⁡ ( 1 − γ ) ​ ∑ t < τ γ τ − t − 1 ​ d d ​ 𝜷 t ​ f τ ​ ( 𝒘 τ ) , \boldsymbol{\beta}_{\tau+1}\leftarrow\boldsymbol{\beta}_{\tau}-\eta\,(1-\gamma)\,\sum_{t<\tau}\gamma^{\tau-t-1}\frac{\mathrm{d}\,}{\mathrm{d}\,\boldsymbol{\beta}_{t}}f_{\tau}(\boldsymbol{w}_{\tau}), (5) so that terms involving f τ f_{\tau} (and 𝒘 τ \boldsymbol{w}_{\tau} ) appear at time τ \tau (instead of t t ), which is the earliest time that these quantities become available. In the small- η \eta limit, the backward view closely approximates the forward view. 1 1 1 Formally, assuming f t ​ ( ⋅ ) = 0 f_{t}(\cdot)=0 for t < 0 t<0 and for t > T t>T , the final updates under ( 5 ) and ( 4 ) become arbitrarily close as η → 0 \eta\to 0 . More specifically, ( β T ( 5 ) − β T ( 4 ) ) / η → 0 (\beta^{\textrm{\eqref{eq:backward}}}_{T}-\beta^{\textrm{\eqref{eq:forward}}}_{T})/{\eta}\to 0 , where β T ( 5 ) \beta^{\textrm{\eqref{eq:backward}}}_{T} and β T ( 4 ) \beta^{\textrm{\eqref{eq:forward}}}_{T} are the values of β \beta at time T T obtained from updates ( 5 ) and ( 4 ), respectively, starting from the same initial β 0 \beta_{0} . This is because as η → 0 \eta\to 0 , β \beta remains almost constant over [ 0 , T ] [0,T] interval, and the right hand sides of ( 5 ) and ( 4 ) match with accuracy O ⁡ ( η 2 ) O(\eta^{2}) , when summed over [ 0 , T ] [0,T] . See Section 9 for a discussion on more accurate approximations for large η \eta .

Accordingly, we define a causal gradient estimate ∇ 𝜷 F ^ τ ​ = def ​ ( 1 − γ ) ​ ∑ t = 0 τ − 1 γ τ − t − 1 ​ d d ​ 𝜷 t ​ f τ ​ ( 𝒘 τ ) . \widehat{\nabla_{\boldsymbol{\beta}}F}_{\tau}\overset{\text{\tiny def}}{=}(1-\gamma)\sum_{t=0}^{\tau-1}\gamma^{\tau-t-1}\frac{\mathrm{d}\,}{\mathrm{d}\,\boldsymbol{\beta}_{t}}f_{\tau}(\boldsymbol{w}_{\tau}). It follows from chain rule that ∇ 𝜷 F ^ τ = ℋ τ T ∇ f τ ( 𝒘 τ ) , \widehat{\nabla_{\boldsymbol{\beta}}F}_{\tau}=\mathcal{H}_{\tau}^{T}\nabla f_{\tau}(\boldsymbol{w}_{\tau}), (6) where ℋ τ ​ = def ​ ( 1 − γ ) ​ ∑ t = 0 τ − 1 γ τ − t − 1 ​ d ​ 𝒘 τ d ​ 𝜷 t . \mathcal{H}_{\tau}\overset{\text{\tiny def}}{=}(1-\gamma)\sum_{t=0}^{\tau-1}\gamma^{\tau-t-1}\frac{d\boldsymbol{w}_{\tau}}{\mathrm{d}\,\boldsymbol{\beta}_{t}}. (7) Hence, ℋ τ \mathcal{H}_{\tau} encodes how past 𝜷 t \boldsymbol{\beta}_{t} values cumulatively affect 𝒘 τ \boldsymbol{w}_{\tau} under discounting by γ \gamma .

## 4 MetaOptimize

Algorithm 1 presents the general MetaOptimize framework for learning the meta-parameters 𝜷 t \boldsymbol{\beta}_{t} . At each step t t , we replace the intractable gradient ∇ 𝜷 F t γ \nabla_{\boldsymbol{\beta}}F^{\gamma}_{t} with the causal surrogate ∇ 𝜷 F ^ t \widehat{\nabla_{\boldsymbol{\beta}}F}_{t} to ensure the update is feasible in real time, as discussed in Section 3 . Specifically, we feed ∇ 𝜷 F ^ t = ℋ t T ∇ f t ( 𝒘 t ) \widehat{\nabla_{\boldsymbol{\beta}}F}_{t}=\mathcal{H}_{t}^{T}\nabla f_{t}(\boldsymbol{w}_{t}) (from ( 6 )) into any first-order meta-update rule Alg meta \operatorname{Alg}_{\textrm{meta}} , just like a conventional gradient.

Formally, define 𝒚 t ​ = def ​ Stack ⁡ ( 𝜷 t , 𝒚 ~ t ) \boldsymbol{y}_{t}\overset{\text{\tiny def}}{=}\operatorname{Stack}(\boldsymbol{\beta}_{t},\,\tilde{\boldsymbol{y}}_{t}) as the stack of meta-parameters 𝜷 t \boldsymbol{\beta}_{t} and any internal states 𝒚 ~ t \tilde{\boldsymbol{y}}_{t} of Alg meta \operatorname{Alg}_{\textrm{meta}} (e.g., momentum). The meta-update is then: 𝒚 t + 1 = Alg meta ( 𝒚 t , ℋ t T ∇ f t ( 𝒘 t ) ) . \boldsymbol{y}_{t+1}\;=\;\operatorname{Alg}_{\textrm{meta}}\!\Big(\boldsymbol{y}_{t},\,\mathcal{H}_{t}^{T}\nabla f_{t}(\boldsymbol{w}_{t})\Big). (8) After applying the base update ( 2 ) to produce 𝒙 t + 1 \boldsymbol{x}_{t+1} , we compute ℋ t T ∇ f t ( 𝒘 t ) \mathcal{H}_{t}^{T}\nabla f_{t}(\boldsymbol{w}_{t}) and plug it into ( 8 ) to update 𝒚 t + 1 \boldsymbol{y}_{t+1} (and thus 𝜷 t + 1 \boldsymbol{\beta}_{t+1} ). The remaining question is how to maintain ℋ t \mathcal{H}_{t} defined in ( 7 ), through application of the chain rule.