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

## Abstract

Abstract We address the challenge of optimizing meta-parameters (i.e., hyperparameters) in machine learning algorithms, a critical factor influencing training efficiency and model performance. Moving away from the computationally expensive traditional meta-parameter search methods, we introduce MetaOptimize framework that dynamically 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 accounts for long-term effect of step sizes on training, through a discounted sum of future losses. We also introduce low complexity variants of MetaOptimize that, in conjunction with its adaptability to multiple optimization algorithms, demonstrate performance competitive to those of best hand-crafted learning rate schedules across various machine learning applications.

## 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., ; Jin, ).
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., ), its primary domain, where dynamic environments or evolving loss functions necessitate meta-parameters, like step sizes, to adapt to optimal time-varying values rather than settling on a static value as in the stationary case. Nevertheless, this work concentrates on the stationary scenario, demonstrating the competitiveness of meta-parameter optimization even in this case.

In this work, we propose *MetaOptimize* as a framework for optimizing meta-parameters
to minimize a form of regret, specifically accounting for the long-term influence of step sizes on future loss. The framework is applicable to a broad range of meta-parameters, however the primary focus of this paper is on step sizes as a critical meta-parameter that is universally present.

MetaOptimize brings additional benefits beyond simplifying the search process. Firstly, it enables a dynamic step-size adjustment during training, potentially accelerating the learning process. Traditional methods typically require manual customization of learning rate schedules for each problem, often following an optimal pattern of initial increase and subsequent decay (Amid et al., ). As our experiments show, step sizes obtained from MetaOptimize follow similar patterns automatically.

Secondly, varying step sizes across different blocks of a neural network, such as layers or neurons, has been shown to improve performance (Singh et al., ; Howard & Ruder, ). Manually tuning or using grid search for block-wise step-sizes is impractical in networks with numerous blocks.
MetaOptimize framework can automatically manage blockwise step-sizes.

The concept of meta step-size optimization can be traced back to (Kesten, ), Delta-bar-Delta (Sutton, ; Jacobs, ), and its incremental variant, IDBD (Sutton, ). Over the years, numerous methods have been developed to address this challenge, detailed further in Section 8. This research distinguishes itself from prior work through the following key aspects:

- •
We introduce a formalization of step-size optimization as minimizing a specific form of regret, essentially a discounted sum of future losses. We demonstrate how to handle this minimization in a causal manner, by introducing the MetaOptimize framework.
- •
MetaOptimize framework is general in the sense that it can wrap around any first-order optimization algorithm, also called *base update*, (such as SGD, RMSProp (Hinton et al., ), Adam (Kingma & Ba, ), or Lion (Chen et al., ))), for which it optimizes step sizes via an algorithm of desire (such as SGD, Adam, RMSProp, or Lion), called the *meta update*.
- •
We develop approximation methods (Section 6), that when integrated into MetaOptimize, lead to computationally efficient algorithms that outperform state-of-the-art automatic hyperparameter optimization methods on CIFAR10, ImageNet, and language modeling applications (refer to experiments in Section 7).
- •
We show that some existing methods (like IDBD, its extension (Xu et al., ), and hypergradient descent (Baydin et al., )) are specific instances or approximations within the MetaOptimize framework (see Section 5).

## 2 Problem Setting

We introduce a general continual optimization setting that, for a given sequence of loss functions $f_{t}(\cdot):\mathbb{R}^{n}\to\mathbb{R}$, $t=0,1,2,\ldots$, aims to find a sequence of weight vectors $\boldsymbol{w}_{1},\boldsymbol{w}_{2},\boldsymbol{w}_{3},\ldots$ to minimize a discounted sum of future loss:

$$ $F^{\gamma}_{t}\overset{\text{\tiny def}}{=}(1-\gamma)\,\sum_{\tau>t}\gamma^{ \tau-t-1}f_{\tau}(\boldsymbol{w}_{\tau}),$ (1) $$

where $\gamma\in[0,1)$ is a fixed constant, often very close to $1$, called the *discount factor*.
As an important special case, the above setting includes stationary supervised learning if $f_{t}$ are sampled from a static distribution, for all $t$. In this case, minimizing $F^{\gamma}_{t}$ results in rapid minimization of expected loss.

Consider an arbitrary first order optimization algorithm (including but not limited to SGD, RMSProp, Adam, or Lion) for updating $\boldsymbol{w}_{t}$. At each time $t$, this algorithm takes the gradient $\nabla f_{t}(\boldsymbol{w}_{t})$ of the immediate loss function, along with an $m$-dimensional vector $\boldsymbol{\beta}_{t}$ of meta-parameters, and updates $\boldsymbol{w}_{t}$ and possibly some internal variables (e.g., momentum in Adam or trace of gradient squares in RMSProp), based on a fixed update rule $\operatorname{Alg}_{\textrm{base}}$, referred to as the *base-update*,

$$ $\boldsymbol{x}_{t}=\operatorname{Alg}_{\textrm{base}}(\boldsymbol{x}_{t}, \nabla f_{t}(\boldsymbol{w}_{t}),\boldsymbol{\beta}_{t}),$ (2) $$

where $\boldsymbol{x}_{t}\overset{\text{\tiny def}}{=}\operatorname{Stack}(
\boldsymbol{w}_{t},\,\tilde{\boldsymbol{x}}_{t})$ is an $\tilde{n}$-dimensional vector obtained by stacking $\boldsymbol{w}_{t}$ and all internal variables of the algorithm that are being updated (e.g., momentum), denoted by $\tilde{\boldsymbol{x}}_{t}$.
The goal of the MetaOptimize framework is to find a sequence of meta-parameters $\boldsymbol{\beta}_{t}$, for $t=1,2,\ldots$, such that when plugged into the base update, (2), results in relative minimization of $F^{\gamma}_{t}$ defined in (1).

Step-size optimization is a special case of the above framework where at each time $t$, the $m$ dimensional vector $\boldsymbol{\beta}_{t}$ is used to determine
the $n$-dimensional (weight-wise) vector $\boldsymbol{\alpha}_{t}$ of step sizes (typically $m\ll n$), through a fixed function $\sigma:\mathbb{R}^{m}\to\mathbb{R}^{n}$,

$$ $\boldsymbol{\alpha}_{t}=\sigma(\boldsymbol{\beta}_{t}).$ (3) $$

A typical choice is to partition weights of the neural network into $m$ blocks and use step-size
$\exp(\beta)$ within each block for some entry $\beta$ of $\boldsymbol{\beta}$.
Depending on $m$, this can result in a single shared scalar step-size, or layer-wise, node-wise, or weight-wise step sizes.
It is particularly beneficial to consider a function $\sigma$ of the exponential form, mentioned above, because of two reasons (Sutton, ).
First, it ensures that $\boldsymbol{\alpha}_{t}$ will always be positive. Second, a constant change in $\boldsymbol{\beta}_{t}$ would lead to a multiplicative change in $\boldsymbol{\alpha}_{t}$, making it suitable for adapting step sizes with different orders of magnitude.

## 3 Forward and Backward Views

Since the definition of $F^{\gamma}_{t}$ in (1) relies on information forward into the future, minimizing it in a causal way necessitates alternative views; discussed in this section.
In order to motivate our approach, we start by considering a hypothetical meta-parameter optimization algorithm that has oracle access to future information (e.g., future loss), and updates $\boldsymbol{\beta}_{t}$ along the gradient of $F^{\gamma}_{t}$ with respect to $\boldsymbol{\beta}_{t}$; that is for $t=0,1,2,\ldots$,

$$ $\begin{split}\boldsymbol{\beta}_{t+1}&=\boldsymbol{\beta}_{t}-\eta\frac{d}{d\, \boldsymbol{\beta}_{t}}F^{\gamma}_{t}=\boldsymbol{\beta}_{t}-\eta\,(1-\gamma) \,\sum_{\tau>t}\gamma^{\tau-t-1}\frac{d}{d\,\boldsymbol{\beta}_{t}}f_{\tau}( \boldsymbol{w}_{\tau}),\end{split}$ (4) $$

for some fixed *meta step-size*, $\eta>0$.
This *forward-view* update however requires that at time $t$, we have access to $f_{\tau}(\cdot)$ and $\boldsymbol{w}_{\tau}$ for all $\tau>t$, which are typically unavailable.
To circumvent this problem, we adopt an idea similar to *eligibility traces* in reinforcement learning (Sutton, ; Sutton & Barto, ). More specifically, instead of the forward-view update, we introduce an update of the following type, which we call the *backward-view* update. At time $\tau=0,1,2,\ldots$, we let

$$ $\boldsymbol{\beta}_{\tau+1}\leftarrow\boldsymbol{\beta}_{\tau}-\eta\,(1-\gamma )\,\sum_{t<\tau}\gamma^{\tau-t-1}\frac{d}{d\,\boldsymbol{\beta}_{t}}f_{\tau}( \boldsymbol{w}_{\tau}).$ (5) $$

Note that every term $\gamma^{\tau-t-1}\frac{d\,f_{\tau}(\boldsymbol{w}_{\tau})}{d\,\boldsymbol{
\beta}_{t}}$ in the right hand side of (4) also appears in (5), but is applied at time $\tau$ instead of time $t$, which is the earliest time that all required information for computing this term is available. Consequently, the $\widehat{\nabla_{\boldsymbol{\beta}}F}_{\tau}$ defined below serves as a causal proxy for ${d\,F^{\gamma}_{\tau}}/{d\,\boldsymbol{\beta}_{\tau}}$;

$$ $\widehat{\nabla_{\boldsymbol{\beta}}F}_{\tau}\overset{\text{\tiny def}}{=}(1- \gamma)\sum_{t=0}^{\tau-1}\gamma^{\tau-t-1}\frac{d}{d\,\boldsymbol{\beta}_{t}} f_{\tau}(\boldsymbol{w}_{\tau}).$ (6) $$

It follows from chain rule that

$$ $\widehat{\nabla_{\boldsymbol{\beta}}F}_{\tau}=\mathcal{H}_{\tau}^{T}\nabla f_{ \tau}(\boldsymbol{w}_{\tau}),$ (7) $$

where

$$ $\mathcal{H}_{\tau}\overset{\text{\tiny def}}{=}(1-\gamma)\sum_{t=0}^{\tau-1} \gamma^{\tau-t-1}\frac{d\boldsymbol{w}_{\tau}}{d\,\boldsymbol{\beta}_{t}}.$ (8) $$

The ${d\boldsymbol{w}_{\tau}}/{d\,\boldsymbol{\beta}_{t}}$ in (8) denotes the Jacobian matrix of $\boldsymbol{w}_{\tau}$ with respect to $\boldsymbol{\beta}_{t}$.
Therefore, $\mathcal{H}_{\tau}$ is an $n\times m$ matrix such that
$\mathcal{H}_{\tau}\,\boldsymbol{v}$, for any $m\times 1$ vector $\boldsymbol{v}$, equals the change in $\boldsymbol{w}_{\tau}$ if we increment all past $\boldsymbol{\beta}_{t}$ along $\gamma^{\tau-t}\,\boldsymbol{v}$, while taking into account the non-linear dynamics of $\boldsymbol{\beta}$ (i.e., the impact of each $\boldsymbol{\beta}_{t}$ increment on $\boldsymbol{\beta}_{\tau}$ of future times $\tau>t$.)

## 4 MetaOptimize

The general formulation of MetaOptimize framework is given in Algorithm 1.
The idea is to update $\boldsymbol{\beta}_{t}$ via any first order optimization algorithm to minimize $F^{\gamma}_{t}$, while using the surrogate gradient $\widehat{\nabla_{\boldsymbol{\beta}}F}_{t}$ in place of $\nabla_{\boldsymbol{\beta}}F^{\gamma}_{t}$, to preserve causality of the updates.
More specifically, for $t=1,2,\ldots$, let

$$ $\begin{split}\boldsymbol{y}_{t+1}&=\operatorname{Alg}_{\textrm{meta}}\big{(} \boldsymbol{y}_{t},\widehat{\nabla_{\boldsymbol{\beta}}F}_{t}\big{)}= \operatorname{Alg}_{\textrm{meta}}\big{(}\boldsymbol{y}_{t},\mathcal{H}_{t}^{T }\nabla f_{t}(\boldsymbol{w}_{t})\big{)}\end{split}$ (9) $$

be the *meta update*, where $\boldsymbol{y}_{t}\overset{\text{\tiny def}}{=}\operatorname{Stack}(
\boldsymbol{\beta}_{t},\,\tilde{\boldsymbol{y}}_{t})$ is an $\tilde{m}$-dimensional vector obtained from stacking $\boldsymbol{\beta}_{t}$ and all other internal variables $\tilde{\boldsymbol{y}}_{t}$ of the $\operatorname{Alg}_{\textrm{meta}}$ algorithm (e.g., momentum), and the second equality follows from (7). Examples of $\operatorname{Alg}_{\textrm{meta}}$ include SGD, RMSprop, Adam, and Lion algorithms. Note that in all cases, we pass $\widehat{\nabla_{\boldsymbol{\beta}}F}$ to the algorithm as the gradient.

In each iteration, after performing the base update (2), we compute $\mathcal{H}_{t}^{T}\nabla f_{t}(\boldsymbol{w}_{t})$ and plug it into (9) to update $\boldsymbol{y}$ (and in particular $\boldsymbol{\beta}$). In the rest of this section, we present incremental updates for $\mathcal{H}_{t}$.

Let $\boldsymbol{h}_{t}$ be an $nm$-dimensional vector obtained by stacking the columns of the $n\times m$ matrix $\mathcal{H}_{t}$.
It follows from the chain rule that for any times $t$ and $\tau$ with $t\geq\tau$,

$$ $\begin{split}\frac{d\,\boldsymbol{y}_{t+1}}{d\,\boldsymbol{\beta}_{\tau}}&= \frac{d\,\boldsymbol{y}_{t+1}}{d\,\boldsymbol{y}_{t}}\,\frac{d\,\boldsymbol{y} _{t}}{d\,\boldsymbol{\beta}_{\tau}}\,+\frac{d\,\boldsymbol{y}_{t+1}}{d\, \boldsymbol{x}_{t}}\,\frac{d\,\boldsymbol{x}_{t}}{d\,\boldsymbol{\beta}_{\tau} }+\frac{d\,\boldsymbol{y}_{t+1}}{d\,\boldsymbol{h}_{t}}\,\frac{d\,\boldsymbol{ h}_{t}}{d\,\boldsymbol{\beta}_{\tau}},\\ \frac{d\,\boldsymbol{x}_{t+1}}{d\,\boldsymbol{\beta}_{\tau}}&=\frac{d\, \boldsymbol{x}_{t+1}}{d\,\boldsymbol{y}_{t}}\,\frac{d\,\boldsymbol{y}_{t}}{d\, \boldsymbol{\beta}_{\tau}}+\frac{d\,\boldsymbol{x}_{t+1}}{d\,\boldsymbol{x}_{t }}\,\frac{d\,\boldsymbol{x}_{t}}{d\,\boldsymbol{\beta}_{\tau}}+\frac{d\, \boldsymbol{x}_{t+1}}{d\,\boldsymbol{h}_{t}}\,\frac{d\,\boldsymbol{h}_{t}}{d\, \boldsymbol{\beta}_{\tau}},\\ \frac{d\,\boldsymbol{h}_{t+1}}{d\,\boldsymbol{\beta}_{\tau}}&=\frac{d\, \boldsymbol{h}_{t+1}}{d\,\boldsymbol{y}_{t}}\,\frac{d\,\boldsymbol{y}_{t}}{d\, \boldsymbol{\beta}_{\tau}}\,+\frac{d\,\boldsymbol{h}_{t+1}}{d\,\boldsymbol{x}_ {t}}\,\frac{d\,\boldsymbol{x}_{t}}{d\,\boldsymbol{\beta}_{\tau}}+\frac{d\, \boldsymbol{h}_{t+1}}{d\,\boldsymbol{h}_{t}}\,\frac{d\,\boldsymbol{h}_{t}}{d\, \boldsymbol{\beta}_{\tau}}.\end{split}$ $$

Letting

$$ $G_{t}\overset{\text{\tiny def}}{=}{\left[\begin{array}[]{ccc}\begin{array}[]{c }\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{y}_{t+1}}{d\,\boldsymbol{y}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{y}_{t+1}}{d\,\boldsymbol{x}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{y}_{t+1}}{d\,\boldsymbol{h}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\\ \begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{x}_{t+1}}{d\,\boldsymbol{y}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{x}_{t+1}}{d\,\boldsymbol{x}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{x}_{t+1}}{d\,\boldsymbol{h}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\\ \begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{h}_{t+1}}{d\,\boldsymbol{y}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{h}_{t+1}}{d\,\boldsymbol{x}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{h}_{t+1}}{d\,\boldsymbol{h}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\end{array}\right]},$ (10) $$

the above set of equations can be equivalently written as

$$ $\left[\begin{array}[]{c}\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{y}_{t+1}}{d\,\boldsymbol{\beta}_{\tau}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\\ \begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{x}_{t+1}}{d\,\boldsymbol{\beta}_{\tau}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\\ \begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{h}_{t+1}}{d\,\boldsymbol{\beta}_{\tau}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\end{array}\right]=G_{t}\left[\begin{array}[]{c}\begin{ array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{y}_{t}}{d\,\boldsymbol{\beta}_{\tau}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\\ \begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{x}_{t}}{d\,\boldsymbol{\beta}_{\tau}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\\ \begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{h}_{t}}{d\,\boldsymbol{\beta}_{\tau}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\end{array}\right].$ $$

It follows that

$$ $\sum_{\tau=0}^{t}\gamma^{t-\tau}\left[\begin{array}[]{c}\begin{array}[]{c} \vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{y}_{t+1}}{d\,\boldsymbol{\beta}_{\tau}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\\ \begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{x}_{t+1}}{d\,\boldsymbol{\beta}_{\tau}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\\ \begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{h}_{t+1}}{d\,\boldsymbol{\beta}_{\tau}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\end{array}\right]=G_{t}\left[\begin{array}[]{c}\begin{ array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{y}_{t}}{d\,\boldsymbol{\beta}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\\ \begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{x}_{t}}{d\,\boldsymbol{\beta}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\\ \begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{h}_{t}}{d\,\boldsymbol{\beta}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\end{array}\right]+\,G_{t}\sum_{\tau=0}^{t-1}\gamma^{t- \tau}\left[\begin{array}[]{c}\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{y}_{t}}{d\,\boldsymbol{\beta}_{\tau}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\\ \begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{x}_{t}}{d\,\boldsymbol{\beta}_{\tau}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\\ \begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{h}_{t}}{d\,\boldsymbol{\beta}_{\tau}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\end{array}\right].$ (11) $$

Let

$$ $\displaystyle Y_{t}$ $\displaystyle\overset{\text{\tiny def}}{=}(1-\gamma)\sum_{\tau=0}^{t-1}\gamma^ {t-\tau-1}\frac{d\,\boldsymbol{y}_{t}}{d\,\boldsymbol{\beta}_{\tau}}$ (12) $\displaystyle X_{t}$ $\displaystyle\overset{\text{\tiny def}}{=}(1-\gamma)\sum_{\tau=0}^{t-1}\gamma^ {t-\tau-1}\frac{d\,\boldsymbol{x}_{t}}{d\,\boldsymbol{\beta}_{\tau}},$ (13) $\displaystyle Q_{t}$ $\displaystyle\overset{\text{\tiny def}}{=}(1-\gamma)\sum_{\tau=0}^{t-1}\gamma^ {t-\tau-1}\frac{d\,\boldsymbol{h}_{t}}{d\,\boldsymbol{\beta}_{\tau}}.$ (14) $$

Note also that ${d\,\boldsymbol{x}_{t}}/{d\,\boldsymbol{\beta}_{t}}=0$, ${d\,\boldsymbol{h}_{t}}/{d\,\boldsymbol{\beta}_{t}}=0$, and
${d\,\boldsymbol{y}_{t}}/{d\,\boldsymbol{\beta}_{t}}={d\,\operatorname{Stack}(
\boldsymbol{\beta}_{t},\,\tilde{\boldsymbol{y}}_{t})}/{d\,\boldsymbol{\beta}_{
t}}=\operatorname{Stack}(I,\,0)$.
Plugging these into (11), we obtain

$$ $\left[\begin{array}[]{c}\!Y_{t+1}\!\\ \!X_{t+1}\!\\ \!Q_{t+1}\!\end{array}\right]=G_{t}\left(\gamma\left[\begin{array}[]{c}\!Y_{t} \!\\ \!X_{t}\!\\ \!Q_{t}\!\end{array}\right]+(1-\gamma)\left[\begin{array}[]{c}\begin{array}[]{ c}\vspace{-10pt}\hfil\\ {\left[\begin{array}[]{c}\!I\!\\ \!0\!\end{array}\right]}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\\ 0\\ 0\end{array}\right]\right).$ (15) $$

Matrices $X_{t},Y_{t},Q_{t}$ can be computed iteratively using (15). The matrix $\mathcal{H}_{t}$ in (8) is then obtained from the sub-matrix constituting the first $n$ rows of $X_{t}$, because $\boldsymbol{x}_{t}=\operatorname{Stack}(\boldsymbol{w}_{t},\,\tilde{
\boldsymbol{x}}_{t})$.

To complete Algorithm 1, it only remains to compute the matrix $G_{t}$ in (10). In Appendix A, we calculate $G_{t}$ for common choices of base and meta updates: SGD, AdamW, and Lion. Notably, the first row of $G_{t}$ blocks depends only on $\operatorname{Alg}_{\textrm{meta}}$, and the rest of $G_{t}$ blocks depend only on $\operatorname{Alg}_{\textrm{base}}$. This simplifies the derivation and implementation for various base and meta algorithm combinations.

###### Remark 4.1 .

A key distinction of MetaOptimize from existing meta-parameter optimization methods is that it accounts for the dynamics of the meta-parameters $\boldsymbol{\beta}$—specifically, how changes in the current $\boldsymbol{\beta}$ affect future values of $\boldsymbol{\beta}$. This is captured by the $Y_{t}$ matrix defined in (12), whose influence then propagates into $\mathcal{H}_{t}$ and the meta-update (see (15)). To provide more intuition, lets focus on a simple case with one-dimensional $\beta$ and SGD meta-updates, and consider two case: a)
If $\beta_{t}$ has consistently increased over the recent past trying to track the optimal $\beta$, then $Y_{t}$ will grow large, resulting in significant increments of $H_{t}$. This increases $|H_{t}|$, and improves the tracking of optimal $\beta$.
b) If $\beta_{t}$ has remained nearly constant, suggesting convergence to the optimal value, $Y_{t}$ will shrink, leading to smaller $H_{t}$ increments and smaller updates to $\beta_{t}$. This helps stabilize $\beta$ around its optimal value.

Figure: Algorithm 1 MetaOptimize Framework (for general meta-parameters)

###### Remark 4.1 .

## 5 Reducing Complexity

The matrix $G_{t}$ in (10) is typically large, reducing the algorithm’s practicality. We discuss two approximations of $G_{t}$ for more efficient algorithms.

2$\times$2 approximation: The vector $\boldsymbol{h}_{t}$, formed by stacking $\mathcal{H}_{t}$’s columns, has length $mn$, making $G_{t}$’s last row and column of blocks very large.
Moreover, as shown in Appendix A, the term $d\boldsymbol{h}_{t+1}/d\boldsymbol{x}_{t}$ typically involves third order derivatives of $f_{t}$ with respect to $\boldsymbol{w}_{t}$, which is not practically computable.

In the 2$\times$2 approximation, we resolve the above problems by completely zeroing out all blocks in the last row and also in the last column of blocks of $G_{t}$ in (10). Consequently, we can also remove $Q_{t}$ from the algorithm.
This appears to have minimal impact on the performance, as we empirically observed in simple settings.
Intuitively, the block $d,\boldsymbol{x}_{t+1}/d,\boldsymbol{h}_{t}$ in $G_{t}$ is zero, as $\mathcal{H}_{t}$ doesn’t affect the base update (2). Thus, $Q$ affects $X$, only indirectly, via $Y$.

L-approximation:
Herein, we take a step further, and in addition to the last row and the last column of blocks of $G_{t}$, we also zero out the block in the first row and the second column of $G_{t}$. In other words, we let

$$ $G^{L}_{t}\overset{\text{\tiny def}}{=}{\left[\begin{array}[]{cc}\begin{array}[ ]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{y}_{t+1}}{d\,\boldsymbol{y}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {0}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\\ \begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{x}_{t+1}}{d\,\boldsymbol{y}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{x}_{t+1}}{d\,\boldsymbol{x}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\end{array}\right]},$ (16) $$

and simplify (15) as

$$ $\left[\begin{array}[]{c}\!Y_{t+1}\!\\ \!X_{t+1}\!\end{array}\right]=G^{L}_{t}\left(\gamma\left[\begin{array}[]{c}\!Y _{t}\!\\ \!X_{t}\!\end{array}\right]+(1-\gamma)\left[\begin{array}[]{c}\!I\!\\ \!0\!\\ \hline\cr 0\end{array}\right]\right).$ (17) $$

We have empirically observed that the resulting algorithm typically performs as good as the 2$\times$2 approximation, and even results in improved stability in some cases.

Intuition of MetaOptimize updates:
Algorithm 2 provides a 2$\times$2 approximation of MetaOptimize for the case where both base and meta updates use SGD, and under scalar step-size (detailed derivation in Appendix A). It shows that
$\mathcal{H}_{t}$ traces past gradients, decaying at rate
$\gamma(I-\left[\boldsymbol{\alpha}\right]\nabla^{2}f_{t})$. This decay ensures that if past gradients poorly approximate future ones due to large
$\nabla^{2}f_{t}$ or $\boldsymbol{\alpha}$, their influence fades more rapidly. If the current gradient aligns positively with past gradients (i.e., $-\mathcal{H}_{t}^{T}\nabla f_{t}>0$), the algorithm increases the step-size $\boldsymbol{\alpha}$ for quicker adaptation; if negatively correlated, it reduces the step size to prevent issues like zigzagging.
$Y_{t}$ in (12) reflects the impact of changes in past
$\boldsymbol{\beta}$ on the current value of $\boldsymbol{\beta}$, amplifying the increment in the $\mathcal{H}_{t+1}$ update
if $\boldsymbol{\beta}$ has been consistently rising or falling over the recent past. It is also worth noting that in Algorithm 2, under the L-approximation, $Y_{t}$
remains constant, equal to $I$. A similar phenomenon occurs also when Adam, RMSProp, or Lion algorithms are used instead of SGD.

Figure: Algorithm 2 MetaOptimize with $2\times 2$ approx., $(\operatorname{Alg}_{\textrm{base}},\operatorname{Alg}_{\textrm{meta}})\!=$ (SGD, SGD), and scalar step-size

Containing some existing algorithms as special cases:
Special cases of the above L-approximation method include IDBD algorithm (Sutton, ) and its extension (Xu et al., ), if we limit $\operatorname{Alg}_{\textrm{base}}$ and $\operatorname{Alg}_{\textrm{meta}}$ to SGD algorithm. Refer to Appendix B.1 for more details and proofs.

MetaOptimize also contains the hypergradient-descent algorithm (Baydin et al., ) as a special case, when using SGD for both base and meta updates of MetaOptimze with $\gamma=0$.
Hypergradient-descent updates step size towards minimizing the immediate loss $f_{t}$ rather than discounted sum of future losses, $F^{\gamma}_{t}$, ignoring long-term effects of step size on future loss. See Appendix B.2 for details.

## 6 Hessian-Free MetaOptimize

The step-size optimization algorithms discussed so far typically involve Hessian, $\nabla^{2}f_{t}(\boldsymbol{w}_{t})$, of the loss function. In particular, the Hessian matrix typically appears in the middle column of blocks in the $G_{t}$ matrix; e.g., in the $d\boldsymbol{w}_{t+1}/d\boldsymbol{w}_{t}$ block where $\boldsymbol{w}_{t+1}=\boldsymbol{w}_{t}-\boldsymbol{\alpha}_{t}\nabla f_{t}(
\boldsymbol{w}_{t})$.
Consequently, the update in (15) involves a Hessian-matrix-product of the form $\nabla^{2}f_{t}(\boldsymbol{w}_{t})\mathcal{H}_{t}$, which increases per-step computational complexity of the algorithm.
The added computational overhead would be still manageable if $m$ is small.
In particular for $m=1$ (i.e., the case that a scalar step-size is used for update of all weights), $\mathcal{H}_{t}$ would be a vector; and one can leverage efficient Hessian-vector-product computation techniques that have the same complexity as gradient computation (Pearlmutter, ).

Interestingly, for certain base and meta algorithms, we can eliminate the Hessian without much compromising the performance.
An example of such (base or meta) algorithms is the Lion algorithm (Chen et al., ).
The Lion algorithm, when used as the base algorithm, updates $\boldsymbol{w}_{t}$ as

$$ $\displaystyle\boldsymbol{m}_{t+1}=\rho\,\boldsymbol{m}_{t}+(1-\rho)\,\nabla f_ {t}(\boldsymbol{w}_{t}),$ $\displaystyle\boldsymbol{w}_{t+1}=\boldsymbol{w}_{t}-\boldsymbol{\alpha}_{t} \operatorname{Sign}\big{(}c\,\boldsymbol{m}_{t}+(1-c)\nabla f_{t}\big{)}- \kappa\boldsymbol{\alpha}_{t}\boldsymbol{w}_{t},$ $$

where $\rho,c\in[0,1)$, $\kappa$ is a nonnegative weight-decay parameter, and $\operatorname{Sign}(\cdot)$ is the entry-wise sign function. In the special cases of $c=0$ or $\rho=0$, $\boldsymbol{m}_{t}$ can be eliminated and the above update simplifies to $\boldsymbol{w}_{t+1}=\boldsymbol{w}_{t}-\boldsymbol{\alpha}_{t}\operatorname{
Sign}\big{(}\nabla f_{t}\big{)}-\kappa\boldsymbol{\alpha}_{t}\boldsymbol{w}_{t}$. In this case, it is easy to see that the derivatives of $\boldsymbol{x}_{t}$ in (10) are Hessian-free.
The above argument can be extended to arbitrary values of $c$ and $\rho$.
In Appendix A.1.3 (respectively Appendix A.3.2), we show that if $\operatorname{Alg}_{\textrm{meta}}$ ($\operatorname{Alg}_{\textrm{base}}$) is the Lion algorithm, then the first row (second and third rows) of blocks in $G$ would be Hessian-free. In summary, Algorithm 1 turns Hessian-free, if Lion is used in both base and meta updates. This elimination of Hessian results from flatness of the $\operatorname{Sign}$ function when ignoring the discontinuity at 0 0.

For other algorithms, we may consider their *Hessian-free approximation* by zeroing out any Hessian term in $G_{t}$. The Hessian-free approximation turns out to be a good approximation, especially for base and meta algorithms that involve gradient normalization, like RMSProp and Adam.
Note that, the sign function used in the Lion algorithm is an extreme form of normalization that divides a vector by its absolute value. We could instead use softer forms of normalization, such as normalizing to square root of a trace of squared vector, $\boldsymbol{v}_{t}$, as in RMSProp.
Such normalizations typically result in two opposing Hessian-based terms in $\mathcal{H}_{t}$’s update (stemming from $\frac{d\,\boldsymbol{w}_{t+1}}{d\,\boldsymbol{w}_{t}}$ and $\frac{d\,\boldsymbol{w}_{t+1}}{d\,\boldsymbol{v}_{t}}$ blocks of matrix $G_{t}$), aiming to cancel out, particularly when consecutive gradients are positively correlated.

The main advantage of Hessian-free methods lies in their computational congeniality. For base and meta updates including SGD, RMSProp, AdamW, and Lion, the Hessian-free
2$\times$2 approximation has low computational complexity, requiring only a few vector-products per iteration beyond the computations required for the base and meta updates.
When Hessian terms in
$2\times 2$ approximation of $G_{t}$
are zeroed out; the blocks in $G_{t}$, and therefore the blocks in $X_{t}$ and $Y_{t}$,
become diagonal. Thus, $X_{t}$ and $Y_{t}$ matrices can be simplified to vector forms, eliminating costly matrix multiplications. The same holds for general blockwise step-sizes (e.g., layer-wise and weight-wise step-sizes), leading to computational overheads on par with the scalar case. We note also that for the meta updates mentioned above if we use no weight-decay in the meta update, Hessian-free
2$\times$2 approximation becomes equivalent to Hessian-free
L-approximation. Algorithm 3 presents Hessian-free approximations for some selected base and meta updates: *SGD with momentum (SGDm)*, AdamW, and Lion.

## 7 Experiments

In this section, we evaluate the MetaOptimize framework on image classification and language modeling benchmarks.
Out of several possible combinations of base and meta algorithms and approximations, we report a few Hessian-free combinations from Algorithm 3 that showed better performance. In all experiments, we set the initial step-sizes of MetaOptimize to one or two orders of magnitudes smaller than the range of good fixed step-sizes, with no specific tuning.
We compare MetaOptimize against some popular baselines whose meta-parameters are well-tuned for each task separately. Refer to Appendix C for further experiment details. Codes are available at (Anonymous, ).

Figure: Algorithm 3 Hessian-free MetaOptimize algorithms with 2$\times$2 approximation used in experiments

### 7.1 CIFAR10 dataset

The first set of experiments involve training
ResNet-18 with batch size of 100 on the CIFAR10 (Krizhevsky et al., ) dataset.
Fig. 3 depicts the learning curves of four combinations of (base, meta) algorithms for Hessian-free MetaOptimize, along with the corresponding baselines with well-tuned fixed step sizes.
For MetaOptimize, in addition to scalar step-sizes, we also considered block-wise step-sizes by partitioning layers of the ResNet18 network into six blocks (first and last linear blocks and 4 ResNet blocks).
Fig. 3 demonstrates that each tested base-meta combination of MetaOptimize, whether scalar or blockwise, surpasses the performance of the corresponding fixed step-size baseline.

Figure: Figure 1: Learning curves for selected (base, meta) combinations in CIFAR10.
Refer to caption: x1.png

Interestingly, as demonstrated in Fig. 3, the MetaOptimize algorithms show remarkable robustness to initial step-size choices, even for initial step sizes that are several orders of magnitude smaller than the optimal fixed step-size.

Fig. 3 depicts the blockwise step-sizes for (SGDm, Adam) across different blocks, showing an increasing trend from the first to the last block (output layer), which is generally a desirable pattern. In contrast, in the blockwise versions of (AdamW, Adam), (Lion, Lion), and (RMSProp, Adam) updates, we empirically observed that the first five blocks exhibit similar trends and values, while the last block follows a distinct trend, growing larger and rising at a later time.

### 7.2 ImageNet dataset

We trained ResNet-18 with batch-size 256 on ImageNet (Deng et al., ).
We compared MetaOptimize with scalar step-size against four state-of-the-art hyperparamter optimization algorithms, namely DoG (Ivgi et al., ), gdtuo (Chandra et al., ), Prodigy (Mishchenko & Defazio, ), and mechanic (Cutkosky et al., ), as well as AdamW and Lion baselines with fixed step-sizes, and AdamW with a well-tuned cosine decay learning rate scheduler with a 10k iterations warmup.
Learning curves and complexity overheads are shown respectively in Fig. 5 and Table 1, showcasing the advantage of MetaOptimize algorithms (learning curve of DoG is not depicted due to its relatively poor performance). Unlike CIFAR10, here the blockwise versions of MetaOptimize showed no improvement over the scalar versions. Refer to Appendix D for further details.

Figure: Figure 4: ImageNet learning curves.
Refer to caption: extracted/5900369/Figures/imagenet/Performance_train_acc5_v6.png

**Table 1: Per-iteration wall-clock-time and GPU-space overhead (compared to AdamW).**
|  | ImageNet | TinyStories |  |  |
| --- | --- | --- | --- | --- |
|  | Time | Space | Time | Space |
| AdamW (fixed stepsize)$A^{A}$ | $0\%$ | $0\%$ | $0\%$ | $0\%$ |
| DoG (Ivgi et al., ) | $+45\%$ | $1.4\%$ | $+268\%$ | $0\%$ |
| gdtuo (Chandra et al., ) | $+85\%$ | $64\%$ | $+150\%$ | $21\%$ |
| mechanic (Cutkosky et al., ) | $+42\%$ | $88\%$ | $+9\%$ | $0\%$ |
| Prodigy (Mishchenko & Defazio, ) | $+42\%$ | $13\%$ | $+9\%$ | $0\%$ |
| MetaOptimize (AdamW, Lion) | $+44\%$ | $33\%$ | $+13\%$ | $0\%$ |

### 7.3 Language modeling

For language model experiments, we used the TinyStories dataset (Eldan & Li, ), a synthetic collection of brief stories designed for children aged 3 to 4. This dataset proves effective for training and evaluating language models that are significantly smaller than the current state-of-the-art, and capable of crafting stories that are not only fluent and coherent but also diverse.

We used the implementation in (Karpathy, ) for training 15M parameter model with a batch size of 128 on the TinyStories dataset. Two combinations of Hessian-free MetaOptimize with scalar step sizes were tested against Lion and AdamW with well-tuned fixed step sizes, AdamW with a well-tuned cosine decay learning rate scheduler with 1k warmup iterations, and the four state-of-the-art step-size adaptation algorithms mentioned in the previous subsection. According to the learning curves, shown in Fig. 5, MetaOptimize outperforms all baselines (with an initial delay due to small initial step-sizes), except for the well-tuned learning rate scheduler within 30k iterations.

### 7.4 Sensitivity analysis

Here, we briefly discuss the sensitivity of MetaOptimize to its meta-parameters.

For the meta-stepsize $\eta$ in MetaOptimize, there is generally no need for tuning, and the default value $\eta=10^{-3}$ works universally well in stationary supervised learning. All experiments in this section used this default value with no sweeping required. The rationale for this choice is that when using Adam, Lion, or RMSProp for meta-updates, the absolute change in $\beta$ per iteration is approximately $\eta\times O(1)\simeq 10^{-3}$. Unless the current stepsize $\alpha$ is already near its optimal value, most $\beta$ updates will consistently move toward the optimal $\beta$. Within 1,000 steps, $\beta$ can change by $O(1)$, nearly doubling or halving $\alpha=\exp(\beta)$. Over 10,000 iterations, $\alpha$ can adjust to stepsizes that are $e^{10}>20,000$ times larger or smaller, allowing $\eta\simeq 10^{-3}$ to efficiently track optimal stepsizes while minimizing unnecessary fluctuations in $\alpha$.

Regarding the discount factor $\gamma$, we used the default value $\gamma=1$ in all experiments and observed minimal sensitivity to $\gamma$ for values $\gamma\geq 0.999$ in a series of preliminary tests. However, performance begins to degrade with smaller values of $\gamma$.

## 8 Related Works

Automatic adaptation of step sizes, has been an important research topic in the literature of stochastic optimization. Several works aimed to remove the manual tuning of learning rates via adaptations of classical line search (Rolinek & Martius, ; Vaswani et al., ; Paquette & Scheinberg, ; Kunstner et al., ) and Polyak step size (Berrada et al., ; Loizou et al., ), stochastic proximal methods (Asi & Duchi, ), stochastic quadratic approximation (Schaul et al., ), hyper-gradient descent (Baydin et al., ), nested hyper-gradient descent (Chandra et al., ), distance to a solution adaptation (Ivgi et al., ; Defazio & Mishchenko, ; Mishchenko & Defazio, ), and online convex learning (Cutkosky et al., ).
A limitation of most of these methods is their potential underperformance when their meta-parameters are not optimally configured for specific problems (Ivgi et al., ). Moreover, the primary focus of most of these methods is on minimizing immediate loss rather than considering the long-term effects of step sizes on future loss.

Normalization techniques
proposed over past few years, such as
AdaGrad (Duchi et al., ), RMSProp, and Adam have significantly enhanced the training process. While these algorithms show promise in the stationary problems, these normalization techniques do not optimize effective step sizes and are prone to have sub-optimal performance especially in the continual learning settings (Degris et al., ).

An early practical step-size optimization method was the Incremental-Delta-Bar-Delta (IDBD) algorithm, introduced in (Sutton, ), which aimed to optimize the step-size vector to minimize a specific form of quadratic loss functions in a continual setting. This algorithm was later extended for neural networks in (Schraudolph, ; Xu et al., ; Donini et al., ), and adapted in (Mahmood et al., ) and (Javed, ) for different meta or base updates beyond SGD.
However, the development of IDBD and its extensions included some implicit assumptions, notably overlooking the impact of step-size dynamics on the formulation of step-size update rules. These extensions are, in essence, special cases of the L-approximation within the MetaOptimize framework.
The current work extends the IDBD research, significantly broadening the framework and establishing a solid basis for the derivations.
IDBD and its extensions have been used in various machine learning tasks including
independent component analysis (Schraudolph & Giannakopoulos, ),
human motion tracking (Kehl & Van Gool, ),
classification (Koop, ; Andrychowicz et al., ),

and reinforcement learning (Xu et al., ; Young et al., ; Javed et al., ).
Refer to (Sutton, ) for a comprehensive history of step-size optimization.

There is also a line of research on the so-called parameter-free optimization that aims to remove the need for step-size tuning with almost no knowledge of the problem properties. Most of these methods are primarily designed for stochastic convex optimization (Luo & Schapire, ; Orabona & Pál, ), while more recent ones (Orabona & Tommasi, ; Ivgi et al., ) were applied to supervised learning tasks with small or moderate sample sizes.

## 9 Limitations and Future Works

Our work represents a step toward unlocking the potential of meta-parameter optimization, with substantial room for further exploration, some of which we outline here:
Hessian:
We confined our experiments to Hessian-free methods for practicality, though Hessian-based algorithms could offer superior performance. These methods, however, face challenges requiring additional research. The Hessian matrix is notably noisy, impacting $\mathcal{H}_{t+1}$ multiplicatively, necessitating smoothing and clipping techniques. Additionally, the Hessian approximates the loss landscape’s curvature but fails to account for non-differentiable curvatures, such as those from ReLU unit breakpoints, significant at training’s end. From a computational perspective, developing low-complexity methods for approximate Hessian matrix products, especially for adjusting step-sizes at the layer and weight levels, is essential.

Blockwise step-sizes:
While step sizes can vary much in granularity, our experiments focused on scalar and blockwise step-sizes. While increasing the number of step sizes is anticipated to enhance performance, our experimental findings in Section 7 reveal that this improvement is not consistent across the MetaOptimize approximations evaluated. Further investigation is needed in future research.

Other approximations:
We explored a limited set of MetaOptimize’s possible approximations, leaving a comprehensive analysis of various approximations for future research.

Other meta-parameters:
Our study was limited to differentiable meta-parameters, not covering discrete ones like batch size or network layer count. We also did not investigate several significant differentiable meta-parameters beyond step-sizes, deferring such exploration to future work.

Automatic Differentiation:
While certain versions of MetaOptimize, such as the L-Approximation, could be implemented using standard automatic differentiation, its applicability to the general case of MetaOptimize remains unclear. Unlike updates for $\boldsymbol{w}$ and $\boldsymbol{\beta}$ (base and meta parameters), the $H$ matrix lacks an explicit incremental formula that can be easily handled by automatic differentiation. For some versions of MetaOptimize, including the Hessian-free approximations used in our experiments, automatic differentiation is unnecessary, as meta updates do not require additional differentiation. Exploring the scope and applicability of automatic differentiation across different MetaOptimize instances is an interesting direction for future research.

Continual learning:
Although continual step-size optimization is primarily aimed at continual learning, this study focused on the stationary case, demonstrating MetaOptimize’s competitiveness in a context that is particularly challenging for it. Investigating the framework within continual learning presents a promising direction for future research.

## Appendix A Step-size Optimization for Different Choices of Base and Meta updates

In this appendix, we derive $G_{t}$ defined in (10) for different choices of algorithms for base and meta updates, and propose corresponding step-size optimization algorithms.

Consider the following partitions of $G_{t}$,

$$ $G^{\textrm{meta}}_{t}\overset{\text{\tiny def}}{=}\left[\begin{array}[]{ccc} \begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{y}_{t+1}}{d\,\boldsymbol{y}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{y}_{t+1}}{d\,\boldsymbol{x}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{y}_{t+1}}{d\,\boldsymbol{h}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\end{array}\right],$ (18) $$

$$ $G^{\textrm{base}}_{t}\overset{\text{\tiny def}}{=}\left[\begin{array}[]{ccc} \begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{x}_{t+1}}{d\,\boldsymbol{y}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{x}_{t+1}}{d\,\boldsymbol{x}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{x}_{t+1}}{d\,\boldsymbol{h}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\\ \begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{h}_{t+1}}{d\,\boldsymbol{y}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{h}_{t+1}}{d\,\boldsymbol{x}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{h}_{t+1}}{d\,\boldsymbol{h}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\end{array}\right].$ (19) $$

Then,

$$ $G_{t}={\left[\begin{array}[]{ccc}\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{y}_{t+1}}{d\,\boldsymbol{y}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{y}_{t+1}}{d\,\boldsymbol{x}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{y}_{t+1}}{d\,\boldsymbol{h}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\\ \begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{x}_{t+1}}{d\,\boldsymbol{y}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{x}_{t+1}}{d\,\boldsymbol{x}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{x}_{t+1}}{d\,\boldsymbol{h}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\\ \begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{h}_{t+1}}{d\,\boldsymbol{y}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{h}_{t+1}}{d\,\boldsymbol{x}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{h}_{t+1}}{d\,\boldsymbol{h}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\end{array}\right]}=\left[\begin{array}[]{c}\begin{array }[]{c}\vspace{-10pt}\hfil\\ {G^{\textrm{meta}}_{t}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\\ \hline\cr\begin{array}[]{c}\vspace{-10pt}\hfil\\ {G^{\textrm{base}}_{t}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\end{array}\right].$ (20) $$

In the sequel, we study base and meta updates separately, because $\operatorname{Alg}_{\textrm{base}}$ and $\operatorname{Alg}_{\textrm{meta}}$ impact disjoint sets of blocks in $G_{t}$. In particular, as we will see, the choice of $\operatorname{Alg}_{\textrm{base}}$ only affects $G^{\textrm{base}}$ while the choice of $\operatorname{Alg}_{\textrm{meta}}$ only affects $G^{\textrm{meta}}$.

Notation conventions in all Appendices: For any vector $\boldsymbol{v}$, we denote by $\left[\boldsymbol{v}\right]$ a diagonal matrix with diagonal entries derived from $\boldsymbol{v}$.
We denote by $\sigma^{\prime}(\boldsymbol{\beta}_{t})$ the Jacobian of $\boldsymbol{\alpha}_{t}$ with respect to $\boldsymbol{\beta}_{t}$.

Before delving into computing $G^{\textrm{base}}_{t}$ and $G^{\textrm{meta}}_{t}$ for different base and meta algorithms, we further simplify these matrices.

### A.1 Derivation of G meta superscript 𝐺 meta G^{\textrm{meta}} italic_G start_POSTSUPERSCRIPT meta end_POSTSUPERSCRIPT for Different Meta Updates

We start by simplifying $G^{\textrm{meta}}$, and introducing some notations.

Note that the meta update has no dependence on internal variables, $\tilde{\boldsymbol{x}}$, of the base algorithm. As a result,

$$ $\frac{d\,\boldsymbol{y}_{t+1}}{d\,\tilde{\boldsymbol{x}}_{t}}=0.$ (21) $$

Then,

$$ $G^{\textrm{meta}}_{t}=\left[\begin{array}[]{ccc}\begin{array}[]{c}\vspace{-10 pt}\hfil\\ {\frac{d\,\boldsymbol{y}_{t+1}}{d\,\boldsymbol{y}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{y}_{t+1}}{d\,\boldsymbol{x}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{y}_{t+1}}{d\,\boldsymbol{h}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\end{array}\right]=\left[\begin{array}[]{cccc}\begin{ array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{y}_{t+1}}{d\,\boldsymbol{y}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{y}_{t+1}}{d\,\boldsymbol{w}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{y}_{t+1}}{d\,\tilde{\boldsymbol{x}}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{y}_{t+1}}{d\,\boldsymbol{h}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\end{array}\right]=\left[\begin{array}[]{cccc}\begin{ array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{y}_{t+1}}{d\,\boldsymbol{y}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{y}_{t+1}}{d\,\boldsymbol{w}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&0&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{y}_{t+1}}{d\,\boldsymbol{h}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\end{array}\right],$ (22) $$

where the third equality is due to (21).
Let

$$ $L_{t}\overset{\text{\tiny def}}{=}\left[\begin{array}[]{c|c|c|c}\quad\nabla f_ {t}(\boldsymbol{w}_{t})^{T}&0&0&0\\ \hline\cr 0&\begin{array}[]{c}\vspace{-9pt}\hfil\\ \quad\nabla f_{t}(\boldsymbol{w}_{t})^{T}\end{array}&0&0\\ \hline\cr 0&0&\begin{array}[]{c}\vspace{-4pt}\hfil\\ \quad\ddots\\ \vspace{-4pt}\hfil\\ \end{array}&0\\ \hline\cr 0&0&0&\begin{array}[]{c}\vspace{-9pt}\hfil\\ \quad\nabla f_{t}(\boldsymbol{w}_{t})^{T}\end{array}\end{array}\right]\,\, \begin{array}[]{l}\begin{array}[]{c}{\color[rgb]{.75,.75,.75}\definecolor[ named]{pgfstrokecolor}{rgb}{.75,.75,.75}\pgfsys@color@gray@stroke{.75} \pgfsys@color@gray@fill{.75}\leftarrow 1}\end{array}\\ \begin{array}[]{c}\vspace{-9pt}\hfil\\ {\color[rgb]{.75,.75,.75}\definecolor[named]{pgfstrokecolor}{rgb}{.75,.75,.75} \pgfsys@color@gray@stroke{.75}\pgfsys@color@gray@fill{.75}\leftarrow 2}\end{ array}\\ \begin{array}[]{c}\vspace{-4pt}\hfil\\ \,\,{\color[rgb]{.75,.75,.75}\definecolor[named]{pgfstrokecolor}{rgb}{ .75,.75,.75}\pgfsys@color@gray@stroke{.75}\pgfsys@color@gray@fill{.75}\vdots} \\ \vspace{-4pt}\hfil\\ \end{array}\\ \begin{array}[]{c}\vspace{-9pt}\hfil\\ {\color[rgb]{.75,.75,.75}\definecolor[named]{pgfstrokecolor}{rgb}{.75,.75,.75} \pgfsys@color@gray@stroke{.75}\pgfsys@color@gray@fill{.75}\leftarrow m}\end{ array}\end{array}$ (23) $$

and recall that $\boldsymbol{h}_{t}$ is a vectorization of $\mathcal{H}_{t}$.
Then,

$$ $\mathcal{H}_{t}\nabla f_{t}(\boldsymbol{w}_{t})=L_{t}\boldsymbol{h}_{t}.$ (24) $$

We now proceed to derivation of $G^{\textrm{meta}}$ for different choices of $\operatorname{Alg}_{\textrm{meta}}$.

#### A.1.1 Meta SGD

Here, we consider SGD for the meta update (9),

$$ $\boldsymbol{\beta}_{t+1}=\boldsymbol{\beta}_{t}-\eta\widehat{\nabla_{ \boldsymbol{\beta}}F}_{t}=\boldsymbol{\beta}_{t}-\eta\,\mathcal{H}_{t}^{T} \nabla f_{t}(\boldsymbol{w}_{t}),$ (25) $$

where $\eta$ is a scalar, called the *meta step size*. In this case, $\boldsymbol{y}_{t}=\boldsymbol{\beta}_{t}$.
It then follows from (25) that

$$ $\frac{d\,\boldsymbol{\beta}_{t+1}}{d\,\boldsymbol{h}_{t}}\,=\,-\eta\frac{d}{d \,\boldsymbol{h}_{t}}\big{(}\mathcal{H}_{t}^{T}\nabla f_{t}(\boldsymbol{w}_{t} )\big{)}\,=\,-\eta\frac{d}{d\,\boldsymbol{h}_{t}}\big{(}L_{t}\boldsymbol{h}_{t }\big{)}\,=\,-\eta L_{t},$ (26) $$

where the second equality is due to (24).
Consequently, from (22), we obtain

$$ $\begin{split}G^{\textrm{meta}}_{t}&=\left[\begin{array}[]{cccc}\begin{array}[] {c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{y}_{t+1}}{d\,\boldsymbol{y}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{y}_{t+1}}{d\,\boldsymbol{w}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&0&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{y}_{t+1}}{d\,\boldsymbol{h}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\end{array}\right]\\ &=\left[\begin{array}[]{cccc}\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{\beta}_{t+1}}{d\,\boldsymbol{\beta}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{\beta}_{t+1}}{d\,\boldsymbol{w}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&0&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{\beta}_{t+1}}{d\,\boldsymbol{h}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\end{array}\right]\\ &=\left[\begin{array}[]{cccc}\begin{array}[]{c}\vspace{-10pt}\hfil\\ {I}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {-\eta\mathcal{H}^{T}_{t}\nabla^{2}f_{t}(\boldsymbol{w}_{t})}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {0}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {-\eta L_{t}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\end{array}\right],\end{split}$ (27) $$

where the last inequality follows from (26) and simple differentiations of (25).
Here, $\nabla^{2}f_{t}(\boldsymbol{w}_{t})$ denotes the Hessian of $f_{t}$ at $\boldsymbol{w}_{t}$.

#### A.1.2 Meta Adam

The meta update based on the Adam algorithm is as follows,

$$ $\begin{split}\bar{\boldsymbol{m}}_{t+1}&=\bar{\rho}\,\bar{\boldsymbol{m}}_{t} \,+\,\mathcal{H}_{t}^{T}\nabla f_{t}(\boldsymbol{w}_{t}),\\ \vphantom{\tilde{\zeta}}\bar{\boldsymbol{v}}_{t+1}&=\bar{\lambda}\,\boldsymbol {v}_{t}\,+\,\big{(}\mathcal{H}_{t}^{T}\nabla f_{t}(\boldsymbol{w}_{t})\big{)}^ {2},\\ \bar{\mu}_{t}&=\left(\frac{1-\bar{\rho}}{{\color[rgb]{.75,.75,.75}\definecolor [named]{pgfstrokecolor}{rgb}{.75,.75,.75}\pgfsys@color@gray@stroke{.75} \pgfsys@color@gray@fill{.75}1-\bar{\rho}^{t}}}\right)\,/\,\sqrt{\frac{1-\bar{ \lambda}}{1-\bar{\lambda}^{t}}},\\ \boldsymbol{\beta}_{t+1}&=\boldsymbol{\beta}_{t}\,-\,\eta\,\bar{\mu}_{t}\frac{ \bar{\boldsymbol{m}}_{t}}{\sqrt{\vphantom{\tilde{\zeta}}\bar{\boldsymbol{v}}_{ t}}}\end{split}$ (28) $$

where $\bar{\boldsymbol{m}}_{t}$ is the momentum vector, $\vphantom{\tilde{\zeta}}\bar{\boldsymbol{v}}_{t}$ is the trace of squared surrogate-meta-gradient.
Since Adam algorithm needs to keep track of $\boldsymbol{\beta}_{t}$, $\bar{\boldsymbol{m}}_{t}$, and $\vphantom{\tilde{\zeta}}\bar{\boldsymbol{v}}_{t}$, we have

$$ $\boldsymbol{y}_{t}=\left[\begin{array}[]{c}\boldsymbol{\beta}_{t}\\ \bar{\boldsymbol{m}}_{t}\\ \vphantom{\tilde{\zeta}}\bar{\boldsymbol{v}}_{t}\end{array}\right].$ (29) $$

Recall the following notation convention at the end of the Introduction section: for any $k\geq 1$, and any $k$-dimensional vector $\boldsymbol{v}=[v_{1},\ldots,v_{k}]$, we denote the the corresponding diagonal matrix by $\left[\boldsymbol{v}\right]$:

$$ $\left[\boldsymbol{v}\right]\overset{\text{\tiny def}}{=}\left[\begin{array}[]{ ccc}v_{1}&\cdots&0\\ \vdots&\ddots&\vdots\\ 0&\cdots&v_{k}\end{array}\right].$ (30) $$

Consequently, from (22), we obtain

$$ $\begin{split}G^{\textrm{meta}}_{t}&=\left[\begin{array}[]{c|cc|c}\begin{array} []{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{y}_{t+1}}{d\,\boldsymbol{y}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{y}_{t+1}}{d\,\boldsymbol{w}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&0&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{y}_{t+1}}{d\,\boldsymbol{h}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\end{array}\right]\\ &=\left[\begin{array}[]{ccc|cc|c}\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{\beta}_{t+1}}{d\,\boldsymbol{\beta}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{\beta}_{t+1}}{d\,\bar{\boldsymbol{m}}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{\beta}_{t+1}}{d\,\vphantom{\tilde{\zeta}}\bar{ \boldsymbol{v}}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{\beta}_{t+1}}{d\,\boldsymbol{w}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&0&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{\beta}_{t+1}}{d\,\boldsymbol{h}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\\ \begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\bar{\boldsymbol{m}}_{t+1}}{d\,\boldsymbol{\beta}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\bar{\boldsymbol{m}}_{t+1}}{d\,\bar{\boldsymbol{m}}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\bar{\boldsymbol{m}}_{t+1}}{d\,\vphantom{\tilde{\zeta}}\bar{ \boldsymbol{v}}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\bar{\boldsymbol{m}}_{t+1}}{d\,\boldsymbol{w}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&0&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\bar{\boldsymbol{m}}_{t+1}}{d\,\boldsymbol{h}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\\ \begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\vphantom{\tilde{\zeta}}\bar{\boldsymbol{v}}{t+1}}{d\,\boldsymbol{ \beta}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\vphantom{\tilde{\zeta}}\bar{\boldsymbol{v}}{t+1}}{d\,\bar{ \boldsymbol{m}}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\vphantom{\tilde{\zeta}}\bar{\boldsymbol{v}}{t+1}}{d\,\vphantom{ \tilde{\zeta}}\bar{\boldsymbol{v}}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\vphantom{\tilde{\zeta}}\bar{\boldsymbol{v}}{t+1}}{d\,\boldsymbol{w} _{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&0&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\vphantom{\tilde{\zeta}}\bar{\boldsymbol{v}}{t+1}}{d\,\boldsymbol{h} _{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\end{array}\right]\\ &=\left[\begin{array}[]{ccc|cc|c}I&-\eta\bar{\mu}_{t}\Big{[}\frac{1}{\sqrt{ \vphantom{\tilde{\zeta}}\bar{\boldsymbol{v}}_{t}}}\Big{]}&\frac{\eta\bar{\mu}_ {t}}{2}\Big{[}\frac{\bar{\boldsymbol{m}}_{t}}{\vphantom{\tilde{\zeta}}\bar{ \boldsymbol{v}}_{t}^{1.5}}\Big{]}&0&0&0\\ 0&\bar{\rho}I&0&\mathcal{H}_{t}^{T}\nabla^{2}f_{t}&0&\frac{d\,\bar{\boldsymbol {m}}_{t+1}}{d\,\boldsymbol{h}_{t}}\\ 0&0&\bar{\lambda}I&2\big{[}\mathcal{H}_{t}^{T}\nabla f_{t}\big{]}\mathcal{H}_{ t}^{T}\nabla^{2}f_{t}&0&\frac{d\,\vphantom{\tilde{\zeta}}\bar{\boldsymbol{v}}_ {t+1}}{d\,\boldsymbol{h}_{t}}\end{array}\right],\end{split}$ (31) $$

where the last equality follows by calculating derivatives of (28).
For the two remaining terms in the last column of $G_{t}$, we have

$$ $\frac{d\,\bar{\boldsymbol{m}}_{t+1}}{d\,\boldsymbol{h}_{t}}\,=\,\frac{d}{d\, \boldsymbol{h}_{t}}\big{(}\mathcal{H}_{t}^{T}\nabla f_{t}(\boldsymbol{w}_{t}) \big{)}\,=\,\eta\frac{d}{d\,\boldsymbol{h}_{t}}\big{(}L_{t}\boldsymbol{h}_{t} \big{)}\,=\eta\,L_{t}.$ (32) $$

where the first equality follows from the update of $\bar{\boldsymbol{m}}_{t+1}$ in (28), and the second equality is due to (24).
In the same vein,

$$ $\frac{d\,\vphantom{\tilde{\zeta}}\bar{\boldsymbol{v}}_{t+1}}{d\,\boldsymbol{h} _{t}}\,=\,\frac{d}{d\,\boldsymbol{h}_{t}}\big{(}\mathcal{H}_{t}^{T}\nabla f_{t }(\boldsymbol{w}_{t})\big{)}^{2}\,=\,\frac{d}{d\,\boldsymbol{h}_{t}}\big{(}L_{ t}\boldsymbol{h}_{t}\big{)}^{2}\,=\,2\big{[}L_{t}\boldsymbol{h}_{t}\big{]} \frac{d}{d\,\boldsymbol{h}_{t}}\big{(}L_{t}\boldsymbol{h}_{t}\big{)}\,=\,2\big {[}L_{t}\boldsymbol{h}_{t}\big{]}\,L_{t}\,=\,2\big{[}\mathcal{H}_{t}^{T}\, \nabla f_{t}(\boldsymbol{w}_{t})\big{]}\,L_{t},$ (33) $$

where the first equality follows from the update of $\vphantom{\tilde{\zeta}}\bar{\boldsymbol{v}}_{t+1}$ in (28), the second equality is due to (24), and the last equality is again from (24).

Plugging (32) and (33) into (31), we obtain

$$ $G^{\textrm{meta}}_{t}=\left[\begin{array}[]{ccc|cc|c}I&-\eta\bar{\mu}_{t}\Big{ [}\frac{1}{\sqrt{\vphantom{\tilde{\zeta}}\bar{\boldsymbol{v}}_{t}}}\Big{]}& \frac{\eta\bar{\mu}_{t}}{2}\Big{[}\frac{\bar{\boldsymbol{m}}_{t}}{\vphantom{ \tilde{\zeta}}\bar{\boldsymbol{v}}_{t}^{1.5}}\Big{]}&0&0&0\\ 0&\bar{\rho}I&0&\mathcal{H}_{t}^{T}\nabla^{2}f_{t}&0&\eta\,L_{t}\\ 0&0&\bar{\lambda}I&2\big{[}\mathcal{H}_{t}^{T}\nabla f_{t}\big{]}\mathcal{H}_{ t}^{T}\nabla^{2}f_{t}&0&2\big{[}\mathcal{H}_{t}^{T}\,\nabla f_{t}\big{]}\,L_{t }\end{array}\right].$ (34) $$

#### A.1.3 Meta Lion

The meta update based on the lion algorithm is as follows

$$ $\displaystyle\bar{\boldsymbol{m}}_{t+1}=\rho\,\bar{\boldsymbol{m}}_{t}+(1-\rho )\,\widehat{\nabla_{\boldsymbol{\beta}}F}_{t},$ (35) $\displaystyle\boldsymbol{\beta}_{t+1}=\boldsymbol{\beta}_{t}-\eta\operatorname {Sign}\big{(}c\,\bar{\boldsymbol{m}}_{t}+(1-c)\widehat{\nabla_{\boldsymbol{ \beta}}F}_{t}\big{)},$ (36) $$

where $\eta$ is a scalar, called the *meta step size*, and $\rho,c\in[0,1)$.
Note that the meta algorithm operates on a low dimensional space. Therefore, we drop the regularizers like weight-decay in the meta updates, as they are primarily aimed to resolve the overfitting problem in high dimensional problems.
Substituting $\widehat{\nabla_{\boldsymbol{\beta}}F}_{t}$ with $\mathcal{H}_{t}^{T}\nabla f_{t}(\boldsymbol{w}_{t})$ we obtain the following meta updates

$$ $\displaystyle\bar{\boldsymbol{m}}_{t+1}=\rho\,\bar{\boldsymbol{m}}_{t}+(1-\rho )\,\mathcal{H}_{t}^{T}\nabla f_{t}(\boldsymbol{w}_{t}),$ (37) $\displaystyle\boldsymbol{\beta}_{t+1}=\boldsymbol{\beta}_{t}-\eta\operatorname {Sign}\big{(}c\,\bar{\boldsymbol{m}}_{t}+(1-c)\mathcal{H}_{t}^{T}\nabla f_{t}( \boldsymbol{w}_{t})\big{)}.$ (38) $$

In this case,

$$ $\boldsymbol{y}_{t}=\left[\begin{array}[]{c}\boldsymbol{\beta}_{t}\\ \bar{\boldsymbol{m}}_{t}\end{array}\right],$ $$

and

$$ $\begin{split}G^{\textrm{meta}}_{t}&=\left[\begin{array}[]{ccccc}\begin{array}[ ]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{y}_{t+1}}{d\,\boldsymbol{y}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{y}_{t+1}}{d\,\boldsymbol{w}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&0&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{y}_{t+1}}{d\,\boldsymbol{h}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\end{array}\right]\\ &=\left[\begin{array}[]{ccccc}\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{\beta}_{t+1}}{d\,\boldsymbol{\beta}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{\beta}_{t+1}}{d\,\bar{\boldsymbol{m}}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{\beta}_{t+1}}{d\,\boldsymbol{w}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&0&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{\beta}_{t+1}}{d\,\boldsymbol{h}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\\ \begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\bar{\boldsymbol{m}}_{t+1}}{d\,\boldsymbol{\beta}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\bar{\boldsymbol{m}}_{t+1}}{d\,\bar{\boldsymbol{m}}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\bar{\boldsymbol{m}}_{t+1}}{d\,\boldsymbol{w}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&0&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\bar{\boldsymbol{m}}_{t+1}}{d\,\boldsymbol{h}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\end{array}\right]\\ &=\left[\begin{array}[]{ccccc}\begin{array}[]{c}\vspace{-10pt}\hfil\\ {I}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {0}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {0}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&0&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {0}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\\ \begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\bar{\boldsymbol{m}}_{t+1}}{d\,\boldsymbol{\beta}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\bar{\boldsymbol{m}}_{t+1}}{d\,\bar{\boldsymbol{m}}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\bar{\boldsymbol{m}}_{t+1}}{d\,\boldsymbol{w}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&0&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\bar{\boldsymbol{m}}_{t+1}}{d\,\boldsymbol{h}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\end{array}\right],\end{split}$ (39) $$

where the last equality follows from (38).
Consider the following block representation of $Y_{t}$:

$$ $Y_{t}=\left[\begin{array}[]{c}B_{t}\\ Y^{\bar{m}}_{t}\end{array}\right].$ (40) $$

Since the base algorithm, does not take $\bar{\boldsymbol{m}}$ as input, as we will see in (42) and (43) of next subsection (Appendix A.2), $\frac{d\,\bar{\boldsymbol{m}}_{t+1}}{d\,\bar{\boldsymbol{m}}_{t}}$ is the only non-zero block of $G_{t}$ in its column of blocks (i.e., $\frac{d\,s_{t+1}}{d\,\bar{\boldsymbol{m}}_{t}}=0$ for every variable $s$ other than $\bar{\boldsymbol{m}}$). Consequently, it follows from (15) that $Y^{\bar{m}}_{t}$ as defined in (40), has no impact on the update of $X_{t+1}$, $B_{t+1}$, and $Q_{t+1}$. Therefore, we can zero-out the rows and columns of $G^{\textrm{meta}}$ that correspond to derivative of $\bar{\boldsymbol{m}}$.
As such we obtain the following equivalent of $G^{\textrm{meta}}$ in (39) from an algorithmic perspective:

$$ $\begin{split}G^{\textrm{meta}}_{t}\equiv\left[\begin{array}[]{cc}\begin{array} []{c}\vspace{-10pt}\hfil\\ {I_{m\times m}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {0}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\\ \begin{array}[]{c}\vspace{-10pt}\hfil\\ {0}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {0}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\end{array}\right].\end{split}$ (41) $$

As a result, we get $B_{t}=I$ for all times $t$.

### A.2 Derivation of G base superscript 𝐺 base G^{\textrm{base}} italic_G start_POSTSUPERSCRIPT base end_POSTSUPERSCRIPT for Different Base Updates

We now turn our focus to computation of $G^{\textrm{base}}$ .
Let us start by simplifying $G^{\textrm{base}}$, and introducing some notations.

Note that the base update has no dependence on internal variables, $\tilde{\boldsymbol{y}}$, of the meta update. As a result,

$$ $\frac{d\,\boldsymbol{x}_{t+1}}{d\,\tilde{\boldsymbol{y}}_{t}}=0.$ (42) $$

Moreover, it follows from the definition of $\mathcal{H}_{t}$ in (8) that

$$ $\frac{d\,\mathcal{H}_{t+1}}{d\,\tilde{\boldsymbol{y}}_{t}}=(1-\gamma)\sum_{t=0 }^{t}\gamma^{t-\tau}\frac{d}{d\tilde{\boldsymbol{y}}_{t}}\,\left(\frac{d \boldsymbol{w}_{t+1}}{d\,\boldsymbol{\beta}_{\tau}}\right)=(1-\gamma)\sum_{t=0 }^{t}\gamma^{t-\tau}\frac{d}{d\boldsymbol{\beta}_{\tau}}\,\left(\frac{d \boldsymbol{w}_{t+1}}{d\,\tilde{\boldsymbol{y}}_{t}}\right)=(1-\gamma)\sum_{t= 0}^{t}\gamma^{t-\tau}\frac{d}{d\boldsymbol{\beta}_{\tau}}\,\left(0\right)=0,$ $$

where the third equality follows from (42). Therefore,

$$ $\frac{d\,\boldsymbol{h}_{t+1}}{d\,\tilde{\boldsymbol{y}}_{t}}=0.$ (43) $$

Note also that $\operatorname{Alg}_{\textrm{base}}$ does not take $\mathcal{H}_{t}$ as input, and therefore,

$$ $\frac{d\,\boldsymbol{x}_{t+1}}{d\,\boldsymbol{h}_{t}}=0.$ (44) $$

Consequently, we can simplify $G^{\textrm{base}}_{t}$ as follows,

$$ $G^{\textrm{base}}_{t}=\left[\begin{array}[]{ccc}\begin{array}[]{c}\vspace{-10 pt}\hfil\\ {\frac{d\,\boldsymbol{x}_{t+1}}{d\,\boldsymbol{y}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{x}_{t+1}}{d\,\boldsymbol{x}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{x}_{t+1}}{d\,\boldsymbol{h}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\\ \begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{h}_{t+1}}{d\,\boldsymbol{y}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{h}_{t+1}}{d\,\boldsymbol{x}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{h}_{t+1}}{d\,\boldsymbol{h}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\end{array}\right]=\left[\begin{array}[]{cccc}\begin{ array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{x}_{t+1}}{d\,\boldsymbol{\beta}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{x}_{t+1}}{d\,\tilde{\boldsymbol{y}}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{x}_{t+1}}{d\,\boldsymbol{x}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{x}_{t+1}}{d\,\boldsymbol{h}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\\ \begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{h}_{t+1}}{d\,\boldsymbol{\beta}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{h}_{t+1}}{d\,\tilde{\boldsymbol{y}}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{h}_{t+1}}{d\,\boldsymbol{x}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{h}_{t+1}}{d\,\boldsymbol{h}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\end{array}\right]=\left[\begin{array}[]{cccc}\begin{ array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{x}_{t+1}}{d\,\boldsymbol{\beta}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&0&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{x}_{t+1}}{d\,\boldsymbol{x}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&0\\ \begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{h}_{t+1}}{d\,\boldsymbol{\beta}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&0&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{h}_{t+1}}{d\,\boldsymbol{x}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{h}_{t+1}}{d\,\boldsymbol{h}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\end{array}\right],$ (45) $$

where the last equality is due to (42), (43), and (44).

On an independent note, consider the following block representation of $Y_{t}$,

$$ $Y_{t}=\left[\begin{array}[]{c}B_{t}-\frac{1-\gamma}{\gamma}I\\ \tilde{Y}_{t}\end{array}\right],$ (46) $$

Therefore,

$$ $\gamma Y_{t}+(1-\gamma)\left[\begin{array}[]{c}I\\ 0\end{array}\right]=\left[\begin{array}[]{c}B_{t}\\ \tilde{Y}_{t}\end{array}\right]$ $$

It then follows from (20) and (15) that

$$ $\left[\begin{array}[]{c}X_{t+1}\\ Q_{t+1}\end{array}\right]=\gamma\,G^{\textrm{base}}_{t}\left[\begin{array}[]{c }\left[\begin{array}[]{c}B_{t}\\ \tilde{Y}_{t}\end{array}\right]\\ X_{t}\\ Q_{t}\end{array}\right].$ (47) $$

Moreover, from the definition of $Y_{t}$ in (12), we have

$$ $\begin{split}\frac{d\,B_{t}}{d\,\boldsymbol{x}_{t}}&=(1-\gamma)\frac{d}{d\, \boldsymbol{x}_{t}}\sum_{\tau=0}^{t}\gamma^{t-\tau}\frac{d\,\boldsymbol{\beta} _{t}}{d\,\boldsymbol{\beta}_{\tau}}=(1-\gamma)\sum_{\tau=0}^{t}\gamma^{t-\tau} \frac{d}{d\,\boldsymbol{\beta}_{\tau}}\left(\frac{d\,\boldsymbol{\beta}_{t}}{d \,\boldsymbol{x}_{t}}\right)=(1-\gamma)\sum_{\tau=0}^{t}\gamma^{t-\tau}\frac{d }{d\,\boldsymbol{\beta}_{\tau}}\big{(}0\big{)}=0,\\ \frac{d\,B_{t}}{d\,\boldsymbol{\beta}_{t}}&=(1-\gamma)\frac{d}{d\,\boldsymbol{ \beta}_{t}}\sum_{\tau=0}^{t}\gamma^{t-\tau}\frac{d\,\boldsymbol{\beta}_{t}}{d \,\boldsymbol{\beta}_{\tau}}=(1-\gamma)\sum_{\tau=0}^{t}\gamma^{t-\tau}\frac{d }{d\,\boldsymbol{\beta}_{\tau}}\left(\frac{d\,\boldsymbol{\beta}_{t}}{d\, \boldsymbol{\beta}_{t}}\right)=(1-\gamma)\sum_{\tau=0}^{t}\gamma^{t-\tau}\frac {d}{d\,\boldsymbol{\beta}_{\tau}}\big{(}I\big{)}=0,\\ \frac{d\,B_{t}}{d\,\boldsymbol{h}_{t}}&=(1-\gamma)\frac{d}{d\,\boldsymbol{h}_{ t}}\sum_{\tau=0}^{t}\gamma^{t-\tau}\frac{d\,\boldsymbol{\beta}_{t}}{d\, \boldsymbol{\beta}_{\tau}}=(1-\gamma)\sum_{\tau=0}^{t}\gamma^{t-\tau}\frac{d}{ d\,\boldsymbol{\beta}_{\tau}}\left(\frac{d\,\boldsymbol{\beta}_{t}}{d\, \boldsymbol{h}_{t}}\right)=(1-\gamma)\sum_{\tau=0}^{t}\gamma^{t-\tau}\frac{d}{ d\,\boldsymbol{\beta}_{\tau}}\big{(}0\big{)}=0.\end{split}$ (48) $$

Finally, recall the definition

$$ $\sigma^{\prime}(\boldsymbol{\beta}_{t})\overset{\text{\tiny def}}{=}\frac{d\, \boldsymbol{\alpha}_{t}}{d\,\boldsymbol{\beta}_{t}}$ (49) $$

as the Jacobian of $\boldsymbol{\alpha}_{t}$ with respect to $\boldsymbol{\beta}_{t}$.

We now proceed to derivation of $G^{\textrm{base}}$ for different choices of $\operatorname{Alg}_{\textrm{base}}$.

### A.3 Base SGD

Base SGD algorithm makes the following base update in each iteration:

$$ $\boldsymbol{w}_{t+1}=\boldsymbol{w}_{t}-\boldsymbol{\alpha}_{t}\nabla f_{t}( \boldsymbol{w}_{t}).$ (50) $$

In this case, $\boldsymbol{x}_{t}=\boldsymbol{w}_{t}$ and $X_{t}=\mathcal{H}_{t}$.
Then, $G^{\textrm{base}}_{t}$ in (45) can be simplified to

$$ $\begin{split}G^{\textrm{base}}_{t}&=\left[\begin{array}[]{cccc}\begin{array}[] {c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{x}_{t+1}}{d\,\boldsymbol{\beta}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&0&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{x}_{t+1}}{d\,\boldsymbol{x}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&0\\ \begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{h}_{t+1}}{d\,\boldsymbol{\beta}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&0&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{h}_{t+1}}{d\,\boldsymbol{x}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{h}_{t+1}}{d\,\boldsymbol{h}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\end{array}\right]\\ &=\left[\begin{array}[]{cccc}\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{w}_{t+1}}{d\,\boldsymbol{\beta}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&0&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{w}_{t+1}}{d\,\boldsymbol{w}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&0\\ \begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{h}_{t+1}}{d\,\boldsymbol{\beta}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&0&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{h}_{t+1}}{d\,\boldsymbol{w}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{h}_{t+1}}{d\,\boldsymbol{h}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\end{array}\right]\\ &=\left[\begin{array}[]{cccc}-\left[\nabla f_{t}(\boldsymbol{w}_{t})\right] \sigma^{\prime}(\boldsymbol{\beta}_{t})&0&I-\left[\boldsymbol{\alpha}_{t} \right]\nabla^{2}f_{t}(\boldsymbol{w}_{t})&0\\ \begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{h}_{t+1}}{d\,\boldsymbol{\beta}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&0&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{h}_{t+1}}{d\,\boldsymbol{w}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{h}_{t+1}}{d\,\boldsymbol{h}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\end{array}\right],\end{split}$ (51) $$

where the last equality follows by computing simple derivatives of $\boldsymbol{w}_{t+1}$ in (50).

We proceed to compute the three remaining entries of $G^{\textrm{base}}_{t}$, i.e.,
${d\,\boldsymbol{h}_{t+1}}/{d\,\boldsymbol{\beta}_{t}}$,
${d\,\boldsymbol{h}_{t+1}}/{d\,\boldsymbol{w}_{t}}$, and
${d\,\boldsymbol{h}_{t+1}}/{d\,\boldsymbol{h}_{t}}$.
Note that
by plugging the first row of $G^{\textrm{base}}_{t}$, given in (51), into (47), and noting that $\mathcal{H}_{t}=X_{t}$, we obtain

$$ $\mathcal{H}_{t+1}=\gamma\big{(}I-\left[\boldsymbol{\alpha}_{t}\right]\nabla^{2 }f_{t}(\boldsymbol{w}_{t})\big{)}\mathcal{H}_{t}\,-\,\gamma\left[\nabla f_{t}( \boldsymbol{w}_{t})\right]\sigma^{\prime}(\boldsymbol{\beta}_{t})\,B_{t},$ (52) $$

for all $t\geq 0$.
By vectorizing both sides of (52) we obtain

$$ $\boldsymbol{h}_{t+1}\,=\,\gamma\left[\begin{array}[]{c}\begin{array}[]{c} \vspace{-10pt}\hfil\\ {\big{(}I-\left[\boldsymbol{\alpha}_{t}\right]\nabla^{2}f_{t}\big{)}\,\mathcal {H}_{t}^{[1]}\,-\,\left[\nabla f_{t}\right]\,\sigma^{\prime}(\boldsymbol{\beta }_{t})\,B_{t}^{[1]}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\\ \hline\cr\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\big{(}I-\left[\boldsymbol{\alpha}_{t}\right]\nabla^{2}f_{t}\big{)}\,\mathcal {H}_{t}^{[2]}\,-\,\left[\nabla f_{t}\right]\,\sigma^{\prime}(\boldsymbol{\beta }_{t})\,B_{t}^{[2]}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\\ \hline\cr\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\vdots}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\\ \hline\cr\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\big{(}I-\left[\boldsymbol{\alpha}_{t}\right]\nabla^{2}f_{t}\big{)}\,\mathcal {H}_{t}^{[m]}\,-\,\left[\nabla f_{t}\right]\,\sigma^{\prime}(\boldsymbol{\beta }_{t})\,B_{t}^{[m]}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\end{array}\right].$ (53) $$

Note that for any pair of same-size vectors $\boldsymbol{a}$ and $\boldsymbol{b}$, we have $\left[\boldsymbol{a}\right]\boldsymbol{b}=\left[\boldsymbol{b}\right]
\boldsymbol{a}$ where $\left[\boldsymbol{a}\right]$ and $\left[\boldsymbol{b}\right]$ are diagonal matrices of $\boldsymbol{a}$ and $\boldsymbol{b}$, respectively. Therefore, (53) can be equivalently written in the following form

$$ $\boldsymbol{h}_{t+1}\,=\,\gamma\left[\begin{array}[]{c}\begin{array}[]{c} \vspace{-10pt}\hfil\\ {\big{(}I-\left[\boldsymbol{\alpha}_{t}\right]\nabla^{2}f_{t}\big{)}\,\mathcal {H}_{t}^{[1]}\,-\,\big{[}\sigma^{\prime}(\boldsymbol{\beta}_{t})\,B_{t}^{[1]} \big{]}\,\nabla f_{t}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\\ \hline\cr\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\vdots}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\\ \hline\cr\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\big{(}I-\left[\boldsymbol{\alpha}_{t}\right]\nabla^{2}f_{t}\big{)}\,\mathcal {H}_{t}^{[m]}\,-\,\big{[}\sigma^{\prime}(\boldsymbol{\beta}_{t})\,B_{t}^{[m]} \big{]}\,\nabla f_{t}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\end{array}\right].$ (54) $$

By taking the derivative of (53) with respect to $\boldsymbol{h}_{t}$, we obtain

$$ $\frac{d\,\boldsymbol{h}_{t+1}}{d\,\boldsymbol{h}_{t}}=\gamma\,\left[\begin{ array}[]{c|c|c|c}\begin{array}[]{c}I-\left[\boldsymbol{\alpha}_{t}\right]\, \nabla^{2}f_{t}(\boldsymbol{w}_{t})\\ \vspace{-6pt}\hfil\\ \end{array}&0&0&0\\ \hline\cr 0&\begin{array}[]{c}\vspace{-6pt}\hfil\\ I-\left[\boldsymbol{\alpha}_{t}\right]\,\nabla^{2}f_{t}(\boldsymbol{w}_{t})\\ \vspace{-6pt}\hfil\\ \end{array}&0&0\\ \hline\cr 0&0&\begin{array}[]{c}\vspace{-4pt}\hfil\\ \ddots\\ \vspace{-4pt}\hfil\\ \end{array}&0\\ \hline\cr 0&0&0&\begin{array}[]{c}\vspace{-6pt}\hfil\\ I-\left[\boldsymbol{\alpha}_{t}\right]\,\nabla^{2}f_{t}(\boldsymbol{w}_{t}) \end{array}\end{array}\right]\,\,\begin{array}[]{l}\begin{array}[]{c}{\color[ rgb]{.75,.75,.75}\definecolor[named]{pgfstrokecolor}{rgb}{.75,.75,.75} \pgfsys@color@gray@stroke{.75}\pgfsys@color@gray@fill{.75}\leftarrow 1\textrm{ st}}\\ \vspace{-6pt}\hfil\\ \end{array}\\ \begin{array}[]{c}\vspace{-6pt}\hfil\\ {\color[rgb]{.75,.75,.75}\definecolor[named]{pgfstrokecolor}{rgb}{.75,.75,.75} \pgfsys@color@gray@stroke{.75}\pgfsys@color@gray@fill{.75}\leftarrow 2\textrm{ nd}}\\ \vspace{-6pt}\hfil\\ \end{array}\\ \begin{array}[]{c}\vspace{-4pt}\hfil\\ \,\,{\color[rgb]{.75,.75,.75}\definecolor[named]{pgfstrokecolor}{rgb}{ .75,.75,.75}\pgfsys@color@gray@stroke{.75}\pgfsys@color@gray@fill{.75}\vdots} \\ \vspace{-4pt}\hfil\\ \end{array}\\ \begin{array}[]{c}\vspace{-6pt}\hfil\\ {\color[rgb]{.75,.75,.75}\definecolor[named]{pgfstrokecolor}{rgb}{.75,.75,.75} \pgfsys@color@gray@stroke{.75}\pgfsys@color@gray@fill{.75}\leftarrow m\textrm{ th}}\end{array}\end{array}.$ (55) $$

In the above equation, note that $d\,B_{t}/d\,\boldsymbol{h}_{t}=0$ due to (48).
Let $\beta_{t}[i]$ and $w_{t}[j]$ denote the $i$th and $j$th entries of $\boldsymbol{\beta}_{t}$ and $\boldsymbol{w}_{t}$, for $i=1,\ldots,m$ and $j=1,\ldots,n$, respectively.
It then follows from (53) and (48) that

$$ $\frac{d\,\boldsymbol{h}_{t+1}}{d\,\boldsymbol{\beta}_{t}}=-\gamma\left[\begin{ array}[]{c|c|c}\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\Big{[}\frac{d\,\boldsymbol{\alpha}_{t}}{d\,\beta_{t}[1]}\Big{]}\nabla^{2}f_{ t}\,\mathcal{H}_{t}^{[1]}\,+\,\left[\nabla f_{t}\right]\,\frac{\partial\, \sigma^{\prime}(\boldsymbol{\beta}_{t})}{\partial\,\beta_{t}[1]}\,B_{t}^{[1]}} \\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\cdots}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\Big{[}\frac{d\,\boldsymbol{\alpha}_{t}}{d\,\beta_{t}[m]}\Big{]}\nabla^{2}f_{ t}\,\mathcal{H}_{t}^{[1]}\,+\,\left[\nabla f_{t}\right]\,\frac{\partial\, \sigma^{\prime}(\boldsymbol{\beta}_{t})}{\partial\,\beta_{t}[{m}]}\,B_{t}^{[1] }}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\\ \hline\cr\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\vdots}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\ddots}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\vdots}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\\ \hline\cr\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\Big{[}\frac{d\,\boldsymbol{\alpha}_{t}}{d\,\beta_{t}[1]}\Big{]}\nabla^{2}f_{ t}\,\mathcal{H}_{t}^{[m]}\,+\,\left[\nabla f_{t}\right]\,\frac{\partial\, \sigma^{\prime}(\boldsymbol{\beta}_{t})}{\partial\,\beta_{t}[1]}\,B_{t}^{[m]}} \\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\cdots}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\Big{[}\frac{d\,\boldsymbol{\alpha}_{t}}{d\,\beta_{t}[m]}\Big{]}\nabla^{2}f_{ t}\,\mathcal{H}_{t}^{[m]}\,+\,\left[\nabla f_{t}\right]\,\frac{\partial\, \sigma^{\prime}(\boldsymbol{\beta}_{t})}{\partial\,\beta_{t}[{m}]}\,B_{t}^{[m] }}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\end{array}\right],$ (56) $$

where $\frac{\partial}{\partial\beta}$ stands for the entry-wise partial derivative of a matrix with respect to a scalar variable $\beta$.
In the same vein, (54) and (48) imply that

$$ $\frac{d\,\boldsymbol{h}_{t+1}}{d\,\boldsymbol{w}_{t}}=-\gamma\left[\begin{ array}[]{c}\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\left[\boldsymbol{\alpha}_{t}\right]\frac{d\,(\nabla^{2}f_{t}(\boldsymbol{w}_ {t})\,\mathcal{H}_{t}^{[1]})}{d\,\boldsymbol{w}_{t}}\,+\,\big{[}\sigma^{\prime }(\boldsymbol{\beta}_{t})\,B_{t}^{[1]}\big{]}\,\nabla^{2}f_{t}(\boldsymbol{w}_ {t})}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\\ \hline\cr\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\vdots}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\\ \hline\cr\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\left[\boldsymbol{\alpha}_{t}\right]\frac{d\,(\nabla^{2}f_{t}(\boldsymbol{w}_ {t})\,\mathcal{H}_{t}^{[{m}]})}{d\,\boldsymbol{w}_{t}}\,+\,\big{[}\sigma^{ \prime}(\boldsymbol{\beta}_{t})\,B_{t}^{[m]}\big{]}\,\nabla^{2}f_{t}( \boldsymbol{w}_{t})}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\end{array}\right].$ (57) $$

Finally, $G^{\textrm{base}}_{t}$ is obtained by plugging (55), (56), and (57) into (51).

In the special case that $\boldsymbol{\beta}$ is a scalar (equivalently $m=1$), and
furthermore $\alpha=\sigma(\beta)=e^{\beta}$, matrix $G^{\textrm{base}}_{t}$ would be simplified to

$$ $G_{t}^{\textrm{base (scalar)}}=\left[\begin{array}[]{cccc}\vspace{6pt}1&-\eta \,\boldsymbol{h}_{t}^{T}\nabla^{2}f_{t}(\boldsymbol{w}_{t})&-\eta\,\nabla f_{t }(\boldsymbol{w}_{t})^{T}\\ \vspace{6pt}-\alpha\nabla f_{t}(\boldsymbol{w}_{t})&I-\alpha\nabla^{2}f_{t}( \boldsymbol{w}_{t})&0\\ -\gamma\alpha\nabla^{2}f_{t}(\boldsymbol{w}_{t})\boldsymbol{h}_{t}-B_{t}\, \alpha\nabla f_{t}(\boldsymbol{w}_{t})&\quad-\gamma\alpha\frac{d\,\left(\nabla ^{2}f_{t}(\boldsymbol{w}_{t})\boldsymbol{h}_{t}\right)}{d\boldsymbol{w}_{t}}-B _{t}\,\alpha\nabla^{2}f_{t}(\boldsymbol{w}_{t})&\gamma\big{(}I-\alpha\nabla^{2 }f_{t}(\boldsymbol{w}_{t})\big{)}\end{array}\right].$ $$

#### A.3.1 Base AdamW

The base update according to the AdamW algorithm (Loizou et al., ) is as follows,

$$ $\begin{split}\boldsymbol{m}_{t+1}&=\rho\,\boldsymbol{m}_{t}\,+\,\nabla f_{t}( \boldsymbol{w}_{t}),\\ \boldsymbol{v}_{t+1}&=\lambda\,\boldsymbol{v}_{t}\,+\,\nabla f_{t}(\boldsymbol {w}_{t})^{2},\\ \mu_{t}&=\left(\frac{1-\rho}{{\color[rgb]{.75,.75,.75}\definecolor[named]{ pgfstrokecolor}{rgb}{.75,.75,.75}\pgfsys@color@gray@stroke{.75} \pgfsys@color@gray@fill{.75}1-\rho^{t}}}\right)\,/\,\sqrt{\frac{1-\lambda}{1- \lambda^{t}}},\\ \boldsymbol{w}_{t+1}&=\boldsymbol{w}_{t}\,-\,\boldsymbol{\alpha}_{t}\mu_{t} \frac{\boldsymbol{m}_{t}}{\sqrt{\boldsymbol{v}_{t}}}-\kappa\boldsymbol{\alpha} _{t}\boldsymbol{w}_{t},\end{split}$ (58) $$

where $\boldsymbol{m}_{t}$ is the momentum vector, $\boldsymbol{v}_{t}$ is the trace of gradient square used for normalization, and $\kappa>0$ is a weight-decay parameter.
Therefore the base algorithm needs to keep track of $\boldsymbol{w}_{t},\boldsymbol{m}_{t},\boldsymbol{v}_{t}$, i.e.,

$$ $\boldsymbol{x}_{t}=\left[\begin{array}[]{c}\boldsymbol{w}_{t}\\ \hfill\boldsymbol{m}_{t}\\ \boldsymbol{v}_{t}\end{array}\right].$ (59) $$

It then follows from (45) that

$$ $\begin{split}G^{\textrm{base}}_{t}&=\left[\begin{array}[]{cccc}\begin{array}[] {c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{x}_{t+1}}{d\,\boldsymbol{\beta}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&0&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{x}_{t+1}}{d\,\boldsymbol{x}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&0\\ \begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{h}_{t+1}}{d\,\boldsymbol{\beta}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&0&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{h}_{t+1}}{d\,\boldsymbol{x}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{h}_{t+1}}{d\,\boldsymbol{h}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\end{array}\right]\\ &=\left[\begin{array}[]{cc|ccc|c}\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{w}_{t+1}}{d\,\boldsymbol{\beta}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&0&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{w}_{t+1}}{d\,\boldsymbol{w}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\frac{d\,\boldsymbol{w}_{t+1}}{d\,\boldsymbol{m}_{t}}& \frac{d\,\boldsymbol{w}_{t+1}}{d\,\boldsymbol{v}_{t}}&0\\ \frac{d\,\boldsymbol{m}_{t+1}}{d\,\boldsymbol{\beta}_{t}}&0&\frac{d\, \boldsymbol{m}_{t+1}}{d\,\boldsymbol{w}_{t}}&\frac{d\,\boldsymbol{m}_{t+1}}{d \,\boldsymbol{m}_{t}}&\frac{d\,\boldsymbol{m}_{t+1}}{d\,\boldsymbol{v}_{t}}&0 \\ \frac{d\,\boldsymbol{v}_{t+1}}{d\,\boldsymbol{\beta}_{t}}&0&\frac{d\, \boldsymbol{v}_{t+1}}{d\,\boldsymbol{w}_{t}}&\frac{d\,\boldsymbol{v}_{t+1}}{d \,\boldsymbol{m}_{t}}&\frac{d\,\boldsymbol{v}_{t+1}}{d\,\boldsymbol{v}_{t}}&0 \\ \hline\cr\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{h}_{t+1}}{d\,\boldsymbol{\beta}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&0&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{h}_{t+1}}{d\,\boldsymbol{w}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{h}_{t+1}}{d\,\boldsymbol{m}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{h}_{t+1}}{d\,\boldsymbol{v}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{h}_{t+1}}{d\,\boldsymbol{h}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\end{array}\right]\\ &=\left[\begin{array}[]{cc|ccc|c}-\mu_{t}\Big{[}\frac{\boldsymbol{m}_{t}}{ \sqrt{\boldsymbol{v}_{t}}}+\kappa\boldsymbol{w}_{t}\Big{]}\sigma^{\prime}( \boldsymbol{\beta}_{t})&0&I-\kappa\left[\boldsymbol{\alpha}_{t}\right]&-\mu_{t }\Big{[}\frac{\boldsymbol{\alpha}_{t}}{\sqrt{\boldsymbol{v}_{t}}}\Big{]}&\frac {\mu_{t}}{2}\Big{[}\frac{\boldsymbol{\alpha}_{t}\boldsymbol{m}_{t}}{ \boldsymbol{v}_{t}^{1.5}}\Big{]}&0\\ 0&0&\nabla^{2}f_{t}&\rho I&0&0\\ 0&0&2\left[\nabla f_{t}\right]\,\nabla^{2}f_{t}&0&\lambda I&0\\ \hline\cr\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{h}_{t+1}}{d\,\boldsymbol{\beta}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&0&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{h}_{t+1}}{d\,\boldsymbol{w}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{h}_{t+1}}{d\,\boldsymbol{m}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{h}_{t+1}}{d\,\boldsymbol{v}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{h}_{t+1}}{d\,\boldsymbol{h}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\end{array}\right]\end{split}$ (60) $$

where the last equality follows from simple derivative computations in (58).

We proceed to compute the terms in the last row of the $G^{\textrm{base}}_{t}$ above.
Consider the following block representation of $X_{t}$,

$$ $X_{t}=\left[\begin{array}[]{c}\mathcal{H}_{t}\\ \hfill X^{m}_{t}\\ X^{v}_{t}\end{array}\right],$ (61) $$

Plugging the first row of $G^{\textrm{base}}_{t}$, given in (60), into (47), implies that

$$ $\mathcal{H}_{t+1}=-\gamma\mu_{t}\Big{[}\frac{\boldsymbol{m}_{t}}{\sqrt{ \boldsymbol{v}_{t}}}+\kappa\boldsymbol{w}_{t}\Big{]}\sigma^{\prime}( \boldsymbol{\beta}_{t})B_{t}\,+\,\gamma\big{(}I-\kappa\left[\boldsymbol{\alpha }_{t}\right]\big{)}\mathcal{H}_{t}\,-\,\gamma\mu_{t}\Big{[}\frac{\boldsymbol{ \alpha}_{t}}{\sqrt{\boldsymbol{v}_{t}}}\Big{]}\,X^{m}_{t}\,+\,\gamma\frac{\mu_ {t}}{2}\Big{[}\frac{\boldsymbol{\alpha}_{t}\boldsymbol{m}_{t}}{\boldsymbol{v}_ {t}^{1.5}}\Big{]}\,X^{v}_{t}.$ (62) $$

for all $t\geq 0$. Note that for any pair of same-size vectors $\boldsymbol{a}$ and $\boldsymbol{b}$, we have $\left[\boldsymbol{a}\right]\boldsymbol{b}=\left[\boldsymbol{b}\right]
\boldsymbol{a}$ where $\left[\boldsymbol{a}\right]$ and $\left[\boldsymbol{b}\right]$ are diagonal matrices of $\boldsymbol{a}$ and $\boldsymbol{b}$, respectively. Therefore, the $i$th column in the matrix equation (62) can be equivalently written as

$$ $\mathcal{H}^{[i]}_{t+1}=-\gamma\mu_{t}\Big{[}\sigma^{\prime}(\boldsymbol{\beta }_{t})B^{[i]}_{t}\Big{]}\frac{\boldsymbol{m}_{t}}{\sqrt{\boldsymbol{v}_{t}}}+ \kappa\boldsymbol{w}_{t}\,+\,\gamma\big{(}I-\kappa\left[\boldsymbol{\alpha}_{t }\right]\big{)}\mathcal{H}^{[i]}_{t}\,-\,\gamma\mu_{t}\Big{[}X^{m\,[i]}_{t} \Big{]}\,\frac{\boldsymbol{\alpha}_{t}}{\sqrt{\boldsymbol{v}_{t}}}\,+\,\gamma \frac{\mu_{t}}{2}\Big{[}X^{v\,[i]}_{t}\Big{]}\,\frac{\boldsymbol{\alpha}_{t} \boldsymbol{m}_{t}}{\boldsymbol{v}_{t}^{1.5}},$ (63) $$

where $B^{[i]}_{t}$, $\mathcal{H}^{[i]}_{t}$, $X^{m\,[i]}_{t}$, and $X^{v\,[i]}_{t}$ stand for the $i$th columns of $B_{t}$, $\mathcal{H}_{t}$, $X^{m}_{t}$, and $X^{v}_{t}$, respectively.
Following similar arguments as in (48),
it is easy to show that

$$ $\begin{split}&\frac{d\,X^{m}_{t}}{d\,\boldsymbol{\beta}_{t}}=\frac{d\,X^{v}_{t }}{d\,\boldsymbol{\beta}_{t}}=0,\\ &\frac{d\,X^{m}_{t}}{d\,\boldsymbol{w}_{t}}=\frac{d\,X^{v}_{t}}{d\,\boldsymbol {w}_{t}}=0,\\ &\frac{d\,X^{m}_{t}}{d\,\boldsymbol{m}_{t}}=\frac{d\,X^{v}_{t}}{d\,\boldsymbol {m}_{t}}=0,\\ &\frac{d\,X^{m}_{t}}{d\,\boldsymbol{v}_{t}}=\frac{d\,X^{v}_{t}}{d\,\boldsymbol {v}_{t}}=0,\\ &\frac{d\,X^{m}_{t}}{d\,\boldsymbol{h}_{t}}=\frac{d\,X^{v}_{t}}{d\,\boldsymbol {h}_{t}}=0.\end{split}$ (64) $$

Note that $\boldsymbol{h}_{t}$ is an $nm$-dimensional vector derived from stacking the columns of $\mathcal{H}_{t}$. Therefore, we consider a block representation of $\boldsymbol{h}_{t}$ consisting of $m$ blocks, each of which corresponds to a column of $\mathcal{H}_{t}$.
By taking the derivative of (62) with respect to $\boldsymbol{h}_{t}$, and using (64), we obtain

$$ $\frac{d\,\boldsymbol{h}_{t+1}}{d\,\boldsymbol{h}_{t}}=\gamma\,\left[\begin{ array}[]{c|c|c|c}\begin{array}[]{c}I-\kappa\left[\boldsymbol{\alpha}_{t}\right ]\\ \vspace{-6pt}\hfil\\ \end{array}&0&0&0\\ \hline\cr 0&\begin{array}[]{c}\vspace{-6pt}\hfil\\ I-\kappa\left[\boldsymbol{\alpha}_{t}\right]\\ \vspace{-6pt}\hfil\\ \end{array}&0&0\\ \hline\cr 0&0&\begin{array}[]{c}\vspace{-4pt}\hfil\\ \ddots\\ \vspace{-4pt}\hfil\\ \end{array}&0\\ \hline\cr 0&0&0&\begin{array}[]{c}\vspace{-6pt}\hfil\\ I-\kappa\left[\boldsymbol{\alpha}_{t}\right]\end{array}\end{array}\right]\,\, \begin{array}[]{l}\begin{array}[]{c}{\color[rgb]{.75,.75,.75}\definecolor[ named]{pgfstrokecolor}{rgb}{.75,.75,.75}\pgfsys@color@gray@stroke{.75} \pgfsys@color@gray@fill{.75}\leftarrow 1\textrm{st}}\\ \vspace{-6pt}\hfil\\ \end{array}\\ \begin{array}[]{c}\vspace{-6pt}\hfil\\ {\color[rgb]{.75,.75,.75}\definecolor[named]{pgfstrokecolor}{rgb}{.75,.75,.75} \pgfsys@color@gray@stroke{.75}\pgfsys@color@gray@fill{.75}\leftarrow 2\textrm{ nd}}\\ \vspace{-6pt}\hfil\\ \end{array}\\ \begin{array}[]{c}\vspace{-4pt}\hfil\\ \,\,{\color[rgb]{.75,.75,.75}\definecolor[named]{pgfstrokecolor}{rgb}{ .75,.75,.75}\pgfsys@color@gray@stroke{.75}\pgfsys@color@gray@fill{.75}\vdots} \\ \vspace{-4pt}\hfil\\ \end{array}\\ \begin{array}[]{c}\vspace{-6pt}\hfil\\ {\color[rgb]{.75,.75,.75}\definecolor[named]{pgfstrokecolor}{rgb}{.75,.75,.75} \pgfsys@color@gray@stroke{.75}\pgfsys@color@gray@fill{.75}\leftarrow m\textrm{ th}}\end{array}\end{array}.$ (65) $$

Let $\beta_{t}[i]$ and $w_{t}[j]$ denote the $i$th and $j$th entries of $\boldsymbol{\beta}_{t}$ and $\boldsymbol{w}_{t}$, for $i=1,\ldots,m$ and $j=1,\ldots,n$, respectively.
Note that ${d\,\boldsymbol{h}_{t+1}}/{d\,\boldsymbol{\beta}_{t}}$ is a block matrix, in the form of an $m\times m$ array of $n\times 1$ blocks, $\frac{d\,\boldsymbol{h}_{t+1}}{d\,\boldsymbol{\beta}_{t}}[i,j]\overset{\text{
\tiny def}}{=}\frac{d\,\mathcal{H}_{t+1}^{[i]}}{d\,\beta_{t}[j]}$, for $i,j=1,\ldots,m$.
It then follows from (62) and (64) that, for $i,j=1,\ldots,m$,

$$ $\begin{split}\frac{d\,\boldsymbol{h}_{t+1}}{d\,\boldsymbol{\beta}_{t}}[i,j]=& \frac{d\,\mathcal{H}_{t+1}^{[i]}}{d\,\beta_{t}[j]}\\ =&-\gamma\mu_{t}\Big{[}\frac{\boldsymbol{m}_{t}}{\sqrt{\boldsymbol{v}_{t}}}+ \kappa\boldsymbol{w}_{t}\Big{]}\left(\frac{\partial\,\sigma^{\prime}( \boldsymbol{\beta}_{t})}{\partial\,\beta_{t}[j]}\right)B^{[i]}_{t}\,+\,\gamma \left(I-\kappa\Big{[}\frac{d\,\boldsymbol{\alpha}_{t}}{d\,\beta_{t}[j]}\Big{]} \right)\mathcal{H}^{[i]}_{t}\\ \,&-\,\gamma\mu_{t}\Big{[}\frac{1}{\sqrt{\boldsymbol{v}_{t}}}\Big{]}\Big{[} \frac{d\,\boldsymbol{\alpha}_{t}}{d\,\beta_{t}[j]}\Big{]}\,X^{m\,[i]}_{t}\,+\, \gamma\frac{\mu_{t}}{2}\Big{[}\frac{\boldsymbol{m}_{t}}{\boldsymbol{v}_{t}^{1. 5}}\Big{]}\,\Big{[}\frac{d\,\boldsymbol{\alpha}_{t}}{d\,\beta_{t}[j]}\Big{]}\, X^{v\,[i]}_{t},\end{split}$ (66) $$

where $\frac{\partial}{\partial\beta}$ stands for the entry-wise partial derivative of a matrix with respect to a scalar variable $\beta$.

In the same vein, it follows from (63) and (64) that

$$ $\frac{d\,\boldsymbol{h}_{t+1}}{d\,\boldsymbol{w}_{t}}=-\gamma\mu_{t}\kappa \left[\begin{array}[]{c}\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\Big{[}\sigma^{\prime}(\boldsymbol{\beta}_{t})B^{[1]}_{t}\Big{]}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\\ \hline\cr\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\vdots}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\\ \hline\cr\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\Big{[}\sigma^{\prime}(\boldsymbol{\beta}_{t})B^{[m]}_{t}\Big{]}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\end{array}\right],$ (67) $$

$$ $\frac{d\,\boldsymbol{h}_{t+1}}{d\,\boldsymbol{m}_{t}}=\gamma\mu_{t}\left[ \begin{array}[]{c}\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\Big{[}\frac{\boldsymbol{\alpha}_{t}X^{v\,[1]}_{t}}{2\,\boldsymbol{v}_{t}^{1. 5}}-\frac{\sigma^{\prime}(\boldsymbol{\beta}_{t})B^{[1]}_{t}}{\sqrt{ \boldsymbol{v}_{t}}}\Big{]}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\\ \hline\cr\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\vdots}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\\ \hline\cr\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\Big{[}\frac{\boldsymbol{\alpha}_{t}X^{v\,[m]}_{t}}{2\,\boldsymbol{v}_{t}^{1. 5}}-\frac{\sigma^{\prime}(\boldsymbol{\beta}_{t})B^{[m]}_{t}}{\sqrt{ \boldsymbol{v}_{t}}}\Big{]}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\end{array}\right],$ (68) $$

$$ $\frac{d\,\boldsymbol{h}_{t+1}}{d\,\boldsymbol{v}_{t}}=\frac{\gamma\mu_{t}}{2} \left[\begin{array}[]{c}\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\Big{[}\frac{1}{\boldsymbol{v}_{t}^{1.5}}\Big{]}\,\Big{[}\big{(}\sigma^{ \prime}(\boldsymbol{\beta}_{t})B^{[1]}_{t}\big{)}\boldsymbol{m}_{t}+ \boldsymbol{\alpha}_{t}\,X^{m\,[1]}_{t}\,-\,\frac{3\boldsymbol{\alpha}_{t} \boldsymbol{m}_{t}X^{v\,[1]}_{t}}{2\,\boldsymbol{v}_{t}}\Big{]}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\par\\ \hline\cr\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\vdots}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\\ \hline\cr\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\Big{[}\frac{1}{\boldsymbol{v}_{t}^{1.5}}\Big{]}\,\Big{[}\big{(}\sigma^{ \prime}(\boldsymbol{\beta}_{t})B^{[m]}_{t}\big{)}\boldsymbol{m}_{t}+ \boldsymbol{\alpha}_{t}\,X^{m\,[m]}_{t}\,-\,\frac{3\boldsymbol{\alpha}_{t} \boldsymbol{m}_{t}X^{v\,[m]}_{t}}{2\,\boldsymbol{v}_{t}}\Big{]}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\end{array}\right].$ (69) $$

Finally, $G^{\textrm{base}}_{t}$ is obtained by plugging (65), (66), (67), (68), and (69) into (60).

#### A.3.2 Base Lion

The lion algorithm, when used for base update, is as follows

$$ $\displaystyle\boldsymbol{m}_{t+1}=\rho\,\boldsymbol{m}_{t}+(1-\rho)\,\nabla f_ {t}(\boldsymbol{w}_{t}),$ (70) $\displaystyle\boldsymbol{w}_{t+1}=\boldsymbol{w}_{t}-\boldsymbol{\alpha}_{t} \operatorname{Sign}\big{(}c\,\boldsymbol{m}_{t}+(1-c)\nabla f_{t}\big{)}- \kappa\boldsymbol{\alpha}_{t}\boldsymbol{w}_{t},$ (71) $$

where $\boldsymbol{m}_{t}$ is called the momentum, $\kappa>0$ is the weight-decay parameter, $\rho,c\in[0,1)$ are constants, and $\operatorname{Sign}(\cdot)$ is a function that computes entry-wise sign of a vector. Let

$$ $\boldsymbol{x}_{t}=\left[\begin{array}[]{c}\boldsymbol{w}_{t}\\ \hfill\boldsymbol{m}_{t}\end{array}\right].$ (72) $$

It then follows from (45) that

$$ $\begin{split}G^{\textrm{base}}_{t}&=\left[\begin{array}[]{cccc}\begin{array}[] {c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{x}_{t+1}}{d\,\boldsymbol{\beta}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&0&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{x}_{t+1}}{d\,\boldsymbol{x}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&0\\ \begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{h}_{t+1}}{d\,\boldsymbol{\beta}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&0&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{h}_{t+1}}{d\,\boldsymbol{x}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{h}_{t+1}}{d\,\boldsymbol{h}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\end{array}\right]\\ &=\left[\begin{array}[]{ccccc}\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{w}_{t+1}}{d\,\boldsymbol{\beta}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&0&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{w}_{t+1}}{d\,\boldsymbol{w}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\frac{d\,\boldsymbol{w}_{t+1}}{d\,\boldsymbol{m}_{t}}&0 \\ \frac{d\,\boldsymbol{m}_{t+1}}{d\,\boldsymbol{\beta}_{t}}&0&\frac{d\, \boldsymbol{m}_{t+1}}{d\,\boldsymbol{w}_{t}}&\frac{d\,\boldsymbol{m}_{t+1}}{d \,\boldsymbol{m}_{t}}&0\\ \begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{h}_{t+1}}{d\,\boldsymbol{\beta}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&0&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{h}_{t+1}}{d\,\boldsymbol{w}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{h}_{t+1}}{d\,\boldsymbol{m}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{h}_{t+1}}{d\,\boldsymbol{h}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\end{array}\right]\\ &=\left[\begin{array}[]{ccccc}-\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\big{[}\operatorname{Sign}\big{(}c\,\boldsymbol{m}_{t}+(1-c)\nabla f_{t}\big{ )}+\kappa\boldsymbol{w}_{t}\big{]}\sigma^{\prime}(\boldsymbol{\beta}_{t})}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\,0&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {I-\kappa\left[\boldsymbol{\alpha}_{t}\right]}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&0&0\\ \frac{d\,\boldsymbol{m}_{t+1}}{d\,\boldsymbol{\beta}_{t}}&0&\frac{d\, \boldsymbol{m}_{t+1}}{d\,\boldsymbol{w}_{t}}&\frac{d\,\boldsymbol{m}_{t+1}}{d \,\boldsymbol{m}_{t}}&0\\ \begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{h}_{t+1}}{d\,\boldsymbol{\beta}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&0&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{h}_{t+1}}{d\,\boldsymbol{w}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{h}_{t+1}}{d\,\boldsymbol{m}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{h}_{t+1}}{d\,\boldsymbol{h}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\end{array}\right]\end{split}$ (73) $$

where the second equality is due to (72) and the last equality follows from (71).
Consider the following block representation of $X_{t}$,

$$ $X_{t}=\left[\begin{array}[]{c}\mathcal{H}_{t}\\ \hfill X^{m}_{t}\end{array}\right].$ (74) $$

Plugging the first row of $G^{\textrm{base}}_{t}$, given in (73), into (47), implies that

$$ $\mathcal{H}_{t+1}=-\gamma\big{[}\operatorname{Sign}\big{(}c\,\boldsymbol{m}_{t }+(1-c)\nabla f_{t}\big{)}+\kappa\boldsymbol{w}_{t}\big{]}\sigma^{\prime}( \boldsymbol{\beta}_{t})\,B_{t}\,+\,\gamma\big{(}I-\kappa\left[\boldsymbol{ \alpha}_{t}\right]\big{)}\mathcal{H}_{t}$ (75) $$

For simplicity of notation, we define the diagonal matrix $S_{t}$ as

$$ $S_{t}\overset{\text{\tiny def}}{=}\big{[}\operatorname{Sign}\big{(}c\, \boldsymbol{m}_{t}+(1-c)\nabla f_{t}\big{)}+\kappa\boldsymbol{w}_{t}\big{]}.$ (76) $$

Then,

$$ $\boldsymbol{h}_{t+1}=\gamma\left[\begin{array}[]{c}\begin{array}[]{c}\vspace{- 10pt}\hfil\\ {-S_{t}\,\sigma^{\prime}(\boldsymbol{\beta}_{t})\,B_{t}^{[1]}\,+\,\gamma\big{( }I-\kappa\left[\boldsymbol{\alpha}_{t}\right]\big{)}\mathcal{H}_{t}^{[1]}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\\ \hline\cr\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\vdots}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\\ \hline\cr\begin{array}[]{c}\vspace{-10pt}\hfil\\ {-S_{t}\,\sigma^{\prime}(\boldsymbol{\beta}_{t})\,B_{t}^{[m]}\,+\,\gamma\big{( }I-\kappa\left[\boldsymbol{\alpha}_{t}\right]\big{)}\mathcal{H}_{t}^{[m]}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}.\end{array}\right]$ (77) $$

It follows that

$$ $\frac{d\,\boldsymbol{h}_{t+1}}{d\,\boldsymbol{m}_{t}}=0,$ (78) $$

and

$$ $\frac{d\,\boldsymbol{h}_{t+1}}{d\,\boldsymbol{w}_{t}}=-\gamma\left[\begin{ array}[]{c|c|c}\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\left[\boldsymbol{e}_{1}\right]\,\sigma^{\prime}(\boldsymbol{\beta}_{t})\,B_{ t}^{[1]}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\cdots&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\left[\boldsymbol{e}_{n}\right]\,\sigma^{\prime}(\boldsymbol{\beta}_{t})\,B_{ t}^{[1]}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\\ \hline\cr\vdots&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\ddots}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\vdots\\ \hline\cr\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\left[\boldsymbol{e}_{1}\right]\,\sigma^{\prime}(\boldsymbol{\beta}_{t})\,B_{ t}^{[m]}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\cdots&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\left[\boldsymbol{e}_{n}\right]\,\sigma^{\prime}(\boldsymbol{\beta}_{t})\,B_{ t}^{[m]}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\end{array}\right],$ (79) $$

where $\boldsymbol{e}_{i}$ is the $i$th unit vector (i.e., an $n$-dimensional vector whose $i$th entry is $1$ and all other entries are zero).
Let $\beta_{t}[i]$ and $\mathcal{H}_{t}^{[i]}$ be the $i$th entry of $\boldsymbol{\beta}_{t}$ and $i$th column of $\mathcal{H}_{t}$, respectively, for $i=1,\ldots,m$. Then,

$$ $\frac{d\,\boldsymbol{h}_{t+1}}{d\,\boldsymbol{\beta}_{t}}=-\gamma\left[\begin{ array}[]{c|c|c}\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\gamma\kappa\Big{[}\frac{d\,\boldsymbol{\alpha}_{t}}{d\,\beta_{t}[1]}\Big{]} \,\mathcal{H}_{t}^{[1]}\,+\,S_{t}\,\frac{\partial\,\sigma^{\prime}(\boldsymbol {\beta}_{t})}{\partial\,\beta_{t}[1]}\,B_{t}^{[1]}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\quad\cdots}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\gamma\kappa\Big{[}\frac{d\,\boldsymbol{\alpha}_{t}}{d\,\beta_{t}[m]}\Big{]} \,\mathcal{H}_{t}^{[1]}\,+\,S_{t}\,\frac{\partial\,\sigma^{\prime}(\boldsymbol {\beta}_{t})}{\partial\,\beta_{t}[{m}]}\,B_{t}^{[1]}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\\ \hline\cr\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\vdots}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\ddots}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\vdots}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\\ \hline\cr\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\gamma\kappa\Big{[}\frac{d\,\boldsymbol{\alpha}_{t}}{d\,\beta_{t}[1]}\Big{]} \,\mathcal{H}_{t}^{[m]}\,+\,S_{t}\,\frac{\partial\,\sigma^{\prime}(\boldsymbol {\beta}_{t})}{\partial\,\beta_{t}[1]}\,B_{t}^{[m]}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\cdots}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\gamma\kappa\Big{[}\frac{d\,\boldsymbol{\alpha}_{t}}{d\,\beta_{t}[m]}\Big{]} \,\mathcal{H}_{t}^{[m]}\,+\,S_{t}\,\frac{\partial\,\sigma^{\prime}(\boldsymbol {\beta}_{t})}{\partial\,\beta_{t}[{m}]}\,B_{t}^{[m]}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\end{array}\right],$ (80) $$

and

$$ $\frac{d\,\boldsymbol{h}_{t+1}}{d\,\boldsymbol{h}_{t}}=\gamma\,\left[\begin{ array}[]{c|c|c|c}\begin{array}[]{c}I-\kappa\left[\boldsymbol{\alpha}_{t}\right ]\\ \vspace{-6pt}\hfil\\ \end{array}&0&0&0\\ \hline\cr 0&\begin{array}[]{c}\vspace{-6pt}\hfil\\ I-\kappa\left[\boldsymbol{\alpha}_{t}\right]\\ \vspace{-6pt}\hfil\\ \end{array}&0&0\\ \hline\cr 0&0&\begin{array}[]{c}\vspace{-4pt}\hfil\\ \ddots\\ \vspace{-4pt}\hfil\\ \end{array}&0\\ \hline\cr 0&0&0&\begin{array}[]{c}\vspace{-6pt}\hfil\\ I-\kappa\left[\boldsymbol{\alpha}_{t}\right]\end{array}\end{array}\right]\,\, \begin{array}[]{l}\begin{array}[]{c}{\color[rgb]{.75,.75,.75}\definecolor[ named]{pgfstrokecolor}{rgb}{.75,.75,.75}\pgfsys@color@gray@stroke{.75} \pgfsys@color@gray@fill{.75}\leftarrow 1\textrm{st}}\\ \vspace{-6pt}\hfil\\ \end{array}\\ \begin{array}[]{c}\vspace{-6pt}\hfil\\ {\color[rgb]{.75,.75,.75}\definecolor[named]{pgfstrokecolor}{rgb}{.75,.75,.75} \pgfsys@color@gray@stroke{.75}\pgfsys@color@gray@fill{.75}\leftarrow 2\textrm{ nd}}\\ \vspace{-6pt}\hfil\\ \end{array}\\ \begin{array}[]{c}\vspace{-4pt}\hfil\\ \,\,{\color[rgb]{.75,.75,.75}\definecolor[named]{pgfstrokecolor}{rgb}{ .75,.75,.75}\pgfsys@color@gray@stroke{.75}\pgfsys@color@gray@fill{.75}\vdots} \\ \vspace{-4pt}\hfil\\ \end{array}\\ \begin{array}[]{c}\vspace{-6pt}\hfil\\ {\color[rgb]{.75,.75,.75}\definecolor[named]{pgfstrokecolor}{rgb}{.75,.75,.75} \pgfsys@color@gray@stroke{.75}\pgfsys@color@gray@fill{.75}\leftarrow m\textrm{ th}}\end{array},\end{array}.$ (81) $$

It follows from (22), (73), and (78) that in the $G_{t}$ matrix, $\frac{d\,\boldsymbol{m}_{t+1}}{d\,\boldsymbol{m}_{t}}$ is the only non-zero block in its corresponding column of blocks. Consequently, it follows from (15) that $X^{m}_{t}$, as defined in (74), has no impact on the update of $\mathcal{H}_{t+1}$, $Y_{t+1}$, and $Q_{t+1}$. Therefore, the rows and columns of $G^{\textrm{base}}$ that correspond to derivative of $\boldsymbol{m}$ can be completely removed from $G^{\textrm{base}}$.
By removing these rows and columns from $G^{t}$, the matrix update (15) simplifies to

$$ $\left[\begin{array}[]{c}Y_{t+1}\\ \mathcal{H}_{t+1}\\ Q_{t+1}\end{array}\right]=\gamma\,\left[\begin{array}[]{cc}\begin{array}[]{c} \begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{y}_{t+1}}{d\,\boldsymbol{y}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\\ \begin{array}[]{c}\vspace{-10pt}\hfil\\ {\left[\begin{array}[]{cc}\begin{array}[]{c}\vspace{-10pt}\hfil\\ {-\big{[}\operatorname{Sign}\big{(}c\,\boldsymbol{m}_{t}+(1-c)\nabla f_{t}\big {)}\big{]}\sigma^{\prime}(\boldsymbol{\beta}_{t})}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\,\,\,\,\begin{array}[]{c}\vspace{-10pt}\hfil\\ {0}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\\ \begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{h}_{t+1}}{d\,\boldsymbol{\beta}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\,\,\,\,\begin{array}[]{c}\vspace{-10pt}\hfil\\ {0}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\end{array}\right]}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\end{array}&\begin{array}[]{cc}\begin{array}[]{c}\vspace {-10pt}\hfil\\ {\frac{d\,\boldsymbol{y}_{t+1}}{d\,\boldsymbol{w}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{y}_{t+1}}{d\,\boldsymbol{h}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\\ \begin{array}[]{c}\vspace{-10pt}\hfil\\ {I-\kappa\left[\boldsymbol{\alpha}_{t}\right]}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&0\\ \begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{h}_{t+1}}{d\,\boldsymbol{w}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}&\begin{array}[]{c}\vspace{-10pt}\hfil\\ {\frac{d\,\boldsymbol{h}_{t+1}}{d\,\boldsymbol{h}_{t}}}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\end{array}\end{array}\right]\,\left(\left[\begin{array} []{c}Y_{t}\\ \mathcal{H}_{t}\\ Q_{t}\end{array}\right]+(1-\gamma)\left[\begin{array}[]{c}\begin{array}[]{c} \vspace{-10pt}\hfil\\ {\left[\begin{array}[]{c}I\\ 0\end{array}\right]}\\ \vspace{-10pt}\hfil\\ \end{array}\phantom{a}\\ 0\\ 0\end{array}\right]\right),$ (82) $$

where ${d\,\boldsymbol{h}_{t+1}}/{d\,\boldsymbol{\beta}_{t}}$, ${d\,\boldsymbol{h}_{t+1}}/{d\,\boldsymbol{w}_{t}}$, and ${d\,\boldsymbol{h}_{t+1}}/{d\,\boldsymbol{h}_{t}}$ are given in (80), (79), and (81), respectively; and the blocks in the first row depend on the meta update.

## Appendix B Exiting Step-size Optimization Algorithms as Special Cases of MetaOptimize

In this appendix we show that some of the existing step-size optimization algorithms are special cases of the MetaOptimize framework. In particular, we first consider the IDBD algorithm (Sutton, ) and its extension (Xu et al., ), and then discuss about the HyperGradient algorithm (Baydin et al., ).

### B.1 IDBD and Its Extensions

Sutton () proposed the IDBD algorithm for step-size optimization of a class of quadratic loss functions. In particular, it considers loss functions of the form

$$ $f_{t}(\boldsymbol{w}_{t})=\frac{1}{2}\big{(}\boldsymbol{a}_{t}^{T}\boldsymbol{ w}_{t}-b_{t}\big{)}^{2},$ (83) $$

for a given sequence of feature vectors $\boldsymbol{a}_{t}$ and target values $b_{t}$, for $t=1,2,\ldots$.
Moreover, Sutton () assumes weight-wise step sizes, in which case $\boldsymbol{\beta}_{t}$ has the same dimension as $\boldsymbol{w}_{t}$.
The update rule of IDBD is as follows:

$$ $\displaystyle\boldsymbol{g}_{t}$ $\displaystyle\leftarrow(\boldsymbol{a}_{t}^{T}\boldsymbol{w}_{t}-b_{t})\, \boldsymbol{a}_{t},$ (84) $\displaystyle\boldsymbol{\beta}_{t+1}$ $\displaystyle\leftarrow\boldsymbol{\beta}_{t}-\eta\,\boldsymbol{h}_{t}\, \boldsymbol{g}_{t},$ (85) $\displaystyle\boldsymbol{\alpha}_{t+1}$ $\displaystyle\leftarrow\exp\big{(}\boldsymbol{\beta}_{t+1}\big{)},$ (86) $\displaystyle\boldsymbol{w}_{t+1}$ $\displaystyle\leftarrow\boldsymbol{w}_{t}-\boldsymbol{\alpha}_{t+1}\, \boldsymbol{g}_{t},$ (87) $\displaystyle\boldsymbol{h}_{t+1}$ $\displaystyle\leftarrow\left(1-\boldsymbol{\alpha}_{t+1}\boldsymbol{a}_{t}^{2} \right)^{+}\boldsymbol{h}_{t}-\boldsymbol{\alpha}_{t+1}\boldsymbol{g}_{t},$ (88) $$

where $(\cdot)^{+}$ clips the entries at zero to make them non-negative, aimed to improve stability. Here, $\boldsymbol{g}_{t}$ is the gradient of $f_{t}(\boldsymbol{w}_{t})$ and $\boldsymbol{a}_{t}^{2}$ in the last line is a vector that contains diagonal entries of the Hessian of $f_{t}$. The updated values of $\boldsymbol{\beta}$ and $\boldsymbol{w}$ would remain unchanged, if instead of the vector $\boldsymbol{h}_{t}$, we use a diagonal matrix $\mathcal{H}_{t}$ and replace (85) and (88) by

$$ $\begin{split}\boldsymbol{\beta}_{t+1}&\leftarrow\boldsymbol{\beta}_{t}-\eta\, \mathcal{H}_{t}\,\boldsymbol{g}_{t},\\ \mathcal{H}_{t+1}&\leftarrow\left(1-\left[{\boldsymbol{\alpha}}_{t+1} \boldsymbol{a}^{2}\right]\right)^{+}\mathcal{H}_{t}-\left[\boldsymbol{\alpha}_ {t+1}\boldsymbol{g}_{t}\right].\end{split}$ (89) $$

Note that $\left[\boldsymbol{a}^{2}\right]$ is a matrix that is obtained from zeroing-out all non-diagonal entries of the Hessian matrix of $f_{t}$.
It is easy to see that the above formulation of IDBD, equals the L-approximation of MetaOptimize framework when we use SGD for both base and meta updates, and further use a diagonal approximation of the Hessian matrix along with a rectifier in the update of $\mathcal{H}_{t}$.

An extension of IDBD beyond quadratic case has been derived in (Xu et al., ). Similar to IDBD, they also consider weight-wise step sizes, i.e., $m=n$. The update of step sizes in this method is as follows:

$$ $\begin{split}\boldsymbol{\beta}_{t+1}&\leftarrow\boldsymbol{\beta}_{t}-\eta\, \mathcal{H}_{t}^{\mathsf{T}}\,\nabla f_{t}(\boldsymbol{w}_{t})\\ \boldsymbol{\alpha}_{t+1}&\leftarrow\exp(\boldsymbol{\beta}_{t+1}),\\ \boldsymbol{w}_{t+1}&\leftarrow\boldsymbol{w}_{t}-\boldsymbol{\alpha}_{t+1}\, \nabla f_{t}(\boldsymbol{w}_{t}),\\ \mathcal{H}_{t+1}&\leftarrow\big{(}I-\left[\boldsymbol{\alpha}_{t+1}\right]\, \nabla^{2}f_{t}(\boldsymbol{w}_{t})\,\big{)}\,\mathcal{H}_{t}-\big{[} \boldsymbol{\alpha}_{t+1}\,\nabla f_{t}(\boldsymbol{w}_{t})\big{]}.\end{split}$ $$

Similar to IDBD, it is straightforward to check that the above set of updates is equivalent to the L-approximation of MetaOptimize framework that uses SGD for both base and meta updates, except for the fact that the above algorithm uses $\boldsymbol{\alpha}_{t+1}$ in $\boldsymbol{w}_{t+1}$ and $\mathcal{H}_{t+1}$ updates whereas MetaOptimize uses $\boldsymbol{\alpha}_{t}$. This however has no considerable impact since $\boldsymbol{\alpha}_{t}$ varies slowly.

### B.2 Hyper-gradient Descent

HyperGradient descent was proposed in (Baydin et al., ) as a step-size optimization method. It considers scalar step size with straightforward extensions to weight-wise step sizes, and at each time $t$, updates the step size in a direction to minimize the immediate next loss function. In particular, they propose the following additive update for step sizes, that can wrap around an arbitrary base update:

$$ $\begin{split}\boldsymbol{\alpha}_{t}&=\beta_{t}\,\boldsymbol{1}_{n\times 1},\\ \beta_{t+1}&=\beta_{t}-\eta\,\frac{d\,f_{t}(\boldsymbol{w}_{t})}{d\,\beta_{t-1 }}=\beta_{t}-\eta\,\nabla f_{t}(\boldsymbol{w}_{t})^{T}\,\frac{d\,\boldsymbol{ w}_{t}}{d\,\beta_{t-1}}.\end{split}$ (90) $$

The last update can be equivalently written as

$$ $\begin{split}\beta_{t+1}&=\beta_{t}-\eta\,\mathcal{H}_{t}^{T}\,\nabla f_{t}( \boldsymbol{w}_{t}),\\ \mathcal{H}_{t+1}&=0\times\mathcal{H}_{t}+\frac{d\,\boldsymbol{w}_{t+1}}{d\, \beta_{t}}.\end{split}$ (91) $$

The step-size update in (91) can be perceived as a special case of MetaOptimize in two different ways. First, as a MetaOptimize algorithm that uses SGD as its meta update and approximate the $G_{t}$ matrix in (10) by zeroing out all of its blocks except for the top two blocks in the first column.
From another perspective, the additive HyperGradient descent in (91) is also equivalent to a MetaOptimize algorithm that uses SGD as its meta update and sets $\gamma=0$. Note that setting $\gamma$ equal to zero would eliminate the dependence of $\mathcal{H}_{t+1}$ on $X_{t}$ and $Q_{t}$, as can be verified from (15). This would also render the $\beta$ updates ignorant about the long-term impact of step size on future losses.

## Appendix C Experiment Details

In the appendix, we describe the details of experiments performed throughout the paper. In our experiments on CIFAR10 and ImageNet dataset, we used a machine with four Intel Xeon Gold 5120 Skylake @ 2.2GHz CPUs and a single NVIDIA V100 Volta (16GB HBM2 memory) GPU. For TinyStories dataset, we used a machine with four AMD Milan 7413 @ 2.65 GHz 128M cache L3 CPUs and a single NVIDIA A100SXM4 (40 GB memory) GPU. In all experiments, the meta step size $\eta$ is set to $10^{-3}$. The meta-parameters used in the considered optimization algorithm for CIFAR10, ImageNet, and TinyStories are given in Table 2, Table 3, and Table 4, respectively. In the experiments, we performed a grid search for $\rho,\bar{\rho}\in\{0.9,0.99,0.999\}$, $\lambda,\bar{\lambda}\in\{0.99,0.999\}$, and $c,\bar{c}\in\{0.9,0.99\}$. Regarding baselines with fixed step sizes, we did a grid search for the learning rate in the set $\{10^{-5},10^{-4},10^{-3},10^{-2},10^{-1}\}$. We set $\gamma$ equal to one in all experiments. Moreover, in ImageNet (respectively TinyStories) dataset, for AdamW with the learning rate scheduler, we considered a cosine decay with 10k (respectively 1k) steps warmup (according to extensive experimental studies in (Chen et al., ) (respectively (Karpathy, ))) and did a grid search for the maximum learning rate in the set $\{10^{-5},10^{-4},10^{-3}\}$.

Regarding other baseline algorithms, for DoG, although it is a parameter-free algorithm, its performance is still sensitive to the initial step movement. We did a grid search for the initial step movement in the set $\{10^{-9},10^{-8},10^{-7},10^{-6}\}$ and reported the performance for the best value. In all experiments of DoG, we considered the polynomial decay averaging. For Prodigy, we used the default values of parameters as suggested by the authors in github repository.
For gdtuo, we considered the following (base, meta) combinations: (RMSprop, Adam), (Adam, Adam), and (SGD with momentum, Adam) and chose the best combination. For mechanic, we did experiments for the base updates of SGDm, Lion, and Adam and considered the best update.
In order to have a fair comparison, in mechanic and gdtuo, we used the same initial step size as MetaOptimize.

Regarding the complexity overheads reported in Table 1, for AdamW with fixed step-size we used the Pytroch implementation of AdamW. For all other baselines, we used the implementation from the Github repository provided along with (and cited in) the corresponding paper. For MetaOptimize, we used the implementation in (Anonymous, ). Note that the implementation of MetaOptimize in (Anonymous, ) is not optimized for time or space efficiency, and smaller complexity overheads might be achieved with more efficient codes. For each algorithm, the wall-clock time overhead and GPU space overhead are computed by $(T_{\mathrm{Alg}}-T_{\mathrm{AdamW}})/T_{\mathrm{AdamW}}$ and $\big{(}B^{\max}_{\mathrm{AdamW}}/B^{\max}_{\mathrm{Alg}}\big{)}-1$, respectively; where $T_{\mathrm{Alg}}$ and $T_{\mathrm{AdamW}}$ are per-iteration runtimes of the algorithm and AdamW, and $B^{\max}_{\mathrm{Alg}}$ and $B^{\max}_{\mathrm{AdamW}}$ are the maximum batch-sizes that did not cause GPU-memory outage for the algorithm and AdamW.

**Table 2: The values of meta-parameters used in CIFAR10 dataset.**
| Base Update | Meta Update (if any) | $\rho$ | $\lambda$ | $\kappa$ | $c$ | $\bar{\rho}$ | $\bar{\lambda}\vphantom{A^{A^{A}}}$ | $\bar{c}$ | $\alpha_{0}$ | $\eta$ | $\gamma$ |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| AdamW | Fixed step size | 0.9 | 0.999 | 0.1 | - | - | - | - | $10^{-5}$ | - | 1 |
| Adam, Scalar | 0.9 | 0.999 | 0.1 | - | 0.9 | 0.999 | - | $10^{-6}$ | $10^{-3}$ | 1 |  |
| Adam, Blockwise | 0.9 | 0.999 | 0.1 | - | 0.9 | 0.999 | - | $10^{-6}$ | $10^{-3}$ | 1 |  |
| Lion | Fixed step size | 0.99 | - | 0.1 | 0.9 | - | - | - | $10^{-4}$ | - | 1 |
| Lion, Scalar | 0.99 | - | 0.1 | 0.9 | 0.99 | - | 0.9 | $10^{-6}$ | $10^{-3}$ | 1 |  |
| Lion, Blockwise | 0.99 | - | 0.1 | 0.9 | 0.99 | - | 0.9 | $10^{-6}$ | $10^{-3}$ | 1 |  |
| RMSprop | Fixed step size | - | 0.999 | 0.1 | - | - | - | - | $10^{-5}$ | - | 1 |
| Adam, Scalar | - | 0.999 | 0.1 | - | 0.9 | 0.999 | - | $10^{-6}$ | $10^{-3}$ | 1 |  |
| Adam, Blockwise | - | 0.999 | 0.1 | - | 0.9 | 0.999 | - | $10^{-6}$ | $10^{-3}$ | 1 |  |
| SGDm | Fixed step size | 0.9 | - | 0.1 | - | - | - | - | $10^{-3}$ | - | 1 |
| Adam, Scalar | 0.9 | - | 0.1 | - | - | - | - | $10^{-6}$ | $10^{-3}$ | 1 |  |
| Adam, Blockwise | 0.9 | - | 0.1 | - | - | - | - | $10^{-6}$ | $10^{-3}$ | 1 |  |

**Table 3: The values of meta-parameters used in ImageNet dataset.**
| Base Update | Meta Update | $\rho$ | $\lambda$ | $\kappa$ | $c$ | $\bar{\rho}$ | $\bar{\lambda}\vphantom{A^{A^{A}}}$ | $\bar{c}$ | $\alpha_{0}$ | $\eta$ | $\gamma$ |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| AdamW | Fixed step size | 0.9 | 0.999 | 0.1 | - | - | - | - | $10^{-5}$ | - | 1 |
| Lion, Scalar | 0.9 | 0.999 | 0.1 | - | 0.99 | - | 0.9 | $10^{-6}$ | $10^{-3}$ | 1 |  |
| Lion, Blockwise | 0.9 | 0.999 | 0.1 | - | 0.99 | - | 0.9 | $10^{-6}$ | $10^{-3}$ | 1 |  |
| Lion | Fixed step size | 0.99 | - | 0.1 | 0.9 | - | - | - | $10^{-5}$ | - | 1 |
| Lion, Scalar | 0.99 | - | 0.1 | 0.9 | 0.99 | - | 0.9 | $10^{-6}$ | $10^{-3}$ | 1 |  |
| Lion, Blockwise | 0.99 | - | 0.1 | 0.9 | 0.99 | - | 0.9 | $10^{-6}$ | $10^{-3}$ | 1 |  |
| SGDm | Lion, Scalar | 0.9 | - | 0.1 | 0.9 | - | - | - | $10^{-5}$ | $10^{-3}$ | 1 |

**Table 4: The values of meta-parameters used in TinyStories dataset.**
| Base Update | Meta Update (if any) | $\rho$ | $\lambda$ | $\kappa$ | $c$ | $\bar{\rho}$ | $\bar{\lambda}\vphantom{A^{A^{A}}}$ | $\bar{c}$ | $\alpha_{0}$ | $\eta$ | $\gamma$ |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| AdamW | Fixed stepsize | 0.9 | 0.999 | 0.1 | - | - | - | - | $10^{-5}$ | - | 1 |
| Adam, Scalar | 0.9 | 0.999 | 0.1 | - | 0.9 | 0.999 | - | $10^{-6}$ | $10^{-3}$ | 1 |  |
| Lion | Fixed stepsize | 0.99 | - | 0.1 | 0.9 | - | - | - | $10^{-4}$ | - | 1 |
| Lion, Scalar | 0.99 | - | 0.1 | 0.9 | 0.99 | - | 0.9 | $10^{-6}$ | $10^{-3}$ | 1 |  |

## Appendix D Further Experiments

In Figure , we provide the results for the blockwise version of MetaOptimize for two combinations of (AdamW, Lion) and (Lion, Lion). As can be seen, they showed no improvement over the scalar version.

Figure: Figure 6: Comparison of blockwise version of MetaOptimize with the scalar version in ImageNet dataset.
Refer to caption: extracted/5900369/Figures/imagenet/Performance_train_blockwise.png

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 5
presentation: 2
questions: Please discuss Equation (5) and the connection between the proposed method to bilevel optimization.
soundness: 2
strengths: The proposed framework contains some existing algorithms as specials cases.

Their explanations are clear and facilitates easy understanding.
summary: The paper provides a framework to dynamically update hyperparameters such as learning rate during training.
weaknesses: While the explanations provided are clear and easy to understand, some aspects are not thoroughly covered, e.g. Equation (5) would benefit from additional intuition or theoretical support.

There are established lines of work on gradient-based hyperparameter optimization using bilevel optimization. Discussing the connection between the proposed method and these existing approaches would strengthen the paper.

---

confidence: 3
contribution: 2
initial_rating: 3
presentation: 1
questions: 1. In terms of time-space complexity comparison, MetaOptimize does not show a significant advantage on ImageNet but on TinyStoires. Can the authorise explain why? 

Also, see the weaknesses section.
soundness: 2
strengths: 1. The motivation that learning learning rate is clear and practical. 
2. The operation details are discussed and given in the appendix for readers to have a broad insight.
summary: The submission aims to adapt the learning rate while model training. By minimising the regret of the historical model training, the learning rate is updated by the meta gradient computed by a proposed MetaOptimise framework wrapped with the proposed approximation method to reduce the computational cost, especially that caused by the Hessian matrix. The method is empirically texted on both image and language datasets with learning curve and computational cost shown as demonstration for efficiency.
weaknesses: 1. The paper is not presented clearly. 
2. The experiment results do not cover the generalisation performance of the trained models by different optimisations. For example, there is no direct comparison of the accuracy of the test set. The top 5 accuracy applied in the paper can reflect the performance of the optimisers but usually top 1 accuracy is more important. 
3. As learning rate learning is usually a special case of optimiser learning, some previous works are missing for comparison [1,2,3]. 
4. The proposed backward-view update is similar to the forward mode differentiation (FMD) in [1] and the follow-up study on learning rate learning paper [2]. In addition, Equation 1 only contains the training information, the motivation for optimising the learning rate against this objective function is not clear and a follow-up question can be what benefits are achieved by optimising this. 
5. I believe approximate Hessian with either diagonal matrix or block diagonal matrix is common in the precondition optimisation field, L-approximation is another special ly format which is parallel with the others. Stating it as containing existing algorithms as special cases may be overclaiming. Also, it could be interesting to know that with diagonal matrix approximation, the Meta-optimize works and it will definitely reduce computational cost. 
6. Suggestions: The reference format is a little mass, which does not affect my score for this paper.  For example, DoG is published in ICML 2023 while in the reference it is the arXiv version. 

[1] Khodak M, Balcan MF, Talwalkar AS. Adaptive gradient-based meta-learning methods. Advances in Neural Information Processing Systems. 2019;32.
[2] Micaelli P, Storkey AJ. Gradient-based hyperparameter optimization over long horizons. Advances in Neural Information Processing Systems. 2021 Dec 6;34:10798-809.
[3] Gao B, Gouk H, Lee HB, Hospedales TM. Meta mirror descent: Optimiser learning for fast convergence. arXiv preprint arXiv:2203.02711. 2022 Mar 5.

---

confidence: 2
contribution: 2
initial_rating: 5
presentation: 2
questions: (1)	Can you give more theoretical explanations for why gradient descent in Eq.(5) can minimize the future loss defined in Eq.(1) ? Or can you give more theoretical analysis for the benefits of introduction of eligibility traces into hyperparameter optimization?
soundness: 3
strengths: (1)	This paper provides a novel MetaOptimize framework to meta-learn hyperparameters of base machine learning algorithms. In particular, the introduction of eligibility traces (from reinforcement learning) into the field of hyperparameter optimization is new.

(2)	The further proposed approximation version and Hessian-free version make great efforts to reduce the complexity of the MetaOptimize framework.

(3)	The proposed MetaOptimize framework in Algorithm 1 is general and really wraps around any first-order optimization algorithm.

(4)	Experiments on image classification and language modeling benchmarks are extensive, demonstrating the effectiveness of the proposed algorithm when compared with hand-crafted learning rate schedules.
summary: This paper proposes a MetaOptimize framework that dynamically adjusts hyperparameters (like step size) of base algorithms (like SGD and Adam), to minimize a discounted sum of future losses. Specifically, this work adopts an idea similar to eligibility traces in reinforcement learning to construct the future loss, uses base algorithm to minimize this future loss to output the weight vectors, leverages the weight vectors to compute the Jacobian matrix, and finally takes the Jacobian matrix as input of a meta algorithm to update the hyperparameters across different tasks/times. Complexity-reduced variants of MetaOptimize (include approximation version and Hessian-free version) are also proposed. Experiments on classification and language tasks are further conducted to validate the effectiveness of the proposed MetaOptimize framework. I did not check the proof line by line, because all the proofs in the appendix are about derivative computation, and such computation is basic.
weaknesses: **Major**

(1)	I admire that the introduction of eligibility traces into hyperparameter optimization is novel. However, when introducing an existing technique (from reinforcement learning) into your research problem (hyperparameter optimization), some theoretical or insightful analysis should be provided. For example, what is the insightful relationship between the objective in Eq.(5) and the future loss in Eq.(1)? Minimizing the objective in Eq.(5) can provably minimize the future loss in Eq.(1)? Such theoretical analysis is necessary, otherwise the current version in my opinion looks more like a direct application of eligibility traces (i.e. Eq.(5)) into hyperparameter optimization, hence resulting in limited insights.

(2)	The introduction of the Jacobian matrix $\mathcal{H} _{t}$ in line 222 is novel, but the bilevel-optimization (i.e. base-meta algorithm optimization) framework in Algorithm 1 is not new for me. Besides, it will be better to utilize existing theoretical analysis for bilevel programming (e.g. [1] and references therein) to derive specific theoretical results (e.g. the stability or generalization guarantees) for the proposed MetaOptimize framework in Algorithm 1. 


**Minor** 

(3)	There exists too many blank spaces in the paper. For example, between Eq.(10) and Eq.(11), I would suggest use $(\frac{\mathrm{d}y _{t+1}}{\mathrm{d} \beta _{\tau}}, \frac{\mathrm{d}x _{t+1}}{\mathrm{d} \beta _{\tau }}, \frac{\mathrm{d}h _{t+1}}{\mathrm{d} \beta _{\tau}})^{T} = G _{t}(\frac{\mathrm{d}y _{t}}{\mathrm{d} \beta _{\tau}}, \frac{\mathrm{d}x _{t}}{\mathrm{d} \beta _{\tau }}, \frac{\mathrm{d}h _{t}}{\mathrm{d} \beta _{\tau}})^{T}$ to save space, and give more theoretical explanations. Similar problem also exists in Eqs.(12)-(16).

(4)	Line 075, sum of future loss -> sum of future losses

(5)	Eq.(2) has a misleading typo: $x _{t} = Alg _{base} (x _{t}, \nabla f _{t}(w _{t}), \beta _{t})$?

(6)	Throughout the whole paper, the differential operator is $\mathrm{d}$ (mathrm{d}), not $d$.

(7)	Line 138, times $\tau > t$.) ->  times $\tau > t$).

(8)	Line 206, two case

(9)	Line 208, what does $|H _{t}|$ mean?

(10)	Line 234, $d,x _{t+1}/d, h _{t}$.

(11)	Line 318, zeroed out; -> zeroed out, 


Reference

[1] Stability and Generalization of Bilevel Programming in Hyperparameter Optimization. NeurIPS2021.

Meta-review: The paper introduces a method to dynamically optimize the learning rate of a machine learning algorithm. The method is based on eligibility trace and can wrap around any first-order optimization algorithm.

The reviewers in general agree that the paper is easy to follow, the introduction of eligibility trace to hyperparameter optimization is novel and interesting, and the method is general since it can wrap around any first-order optimization algorithm.

However, the reviewers also expressed concerns about some important weaknesses: the paper does not report generalization performance (which is important for machine learning algorithms), and the added generalization results are in general not satisfactory; the theoretical contribution of the paper may be limited, and hence the paper may be an application of eligibility trace to hyperparameter optimization which may limit its technical novelty.

Therefore, rejection is recommended.

Decision: Reject