はじめに

最近、AdEMAMix という長期の勾配情報を利用する手法が提案されていました。
色々と違いはありますが、長期の勾配情報を利用する手法は過去にいくつか提案されているため、今後類似の手法が提案された際のために備忘録を作成します。

AdEMAMix

論文

https://arxiv.org/abs/2409.03137

AdamW に長期の勾配情報の利用を追加したものです。
更新式を以下に示します。weight decay と $\alpha, \beta_3$ のスケジューリングについては省略しています。

\begin{align}
m_{1,0}, m_{2,0}, v_0 &\gets 0 \\
m_{1,t} &\gets \beta_1 m_{1,t-1} + (1 - \beta_1) g_t, \hat{m}_{1,t} \gets \frac{m_{1,t}}{1 - \beta_1^t}  \\
m_{2,t} &\gets \beta_3 m_{2,t-1} + (1 - \beta_3) g_t \\
v_t &\gets \beta_2 v_{t-1} + (1 - \beta_2) g_t^2, \hat{v}_t \gets \frac{v_t}{1 - \beta_2^t}  \\
\theta_t &\gets \theta_{t-1}  - \eta \frac{\hat{m}_{1,t} + \alpha m_{2,t}}{\sqrt{\hat{v}_t} + \epsilon}
\end{align}

$m_2$ が長期の勾配情報を保存します。そのため、$\beta3$は$\beta1$よりも1に近い大きな値を利用します。

個人的な疑問点は以下の通りです。

$\beta_3$ をスケジューリングする関係か $m_2$ について Adam で提案された補正を行いません($\hat{m}_2$ を使いません)
長期の勾配情報 $m_2$ は $\alpha$ 倍して $\hat{m}_1$ に加算しています。$\alpha$ は [4, 10] の間で探索しているとあるため、かなり大きいです。そのため、AdamW に対して更新幅がかなり大きくなり、実質学習率を大きくするような効果があります

Lion

論文

https://arxiv.org/abs/2302.06675

更新式を以下に示します。weight decay については省略しています。

\begin{align}
m_0 &\gets 0 \\
c_t &\gets \beta_1 m_{t-1} + (1 - \beta_1) g_t \\
\theta_t &\gets \theta_{t-1} - \eta sign(c_t) \\
m_t &\gets \beta_2 m_{t-1} + (1 - \beta_2) g_t
\end{align}

$m$ に長期の勾配情報を保存します。短期の勾配情報は保存せずに、最新の勾配をそのまま利用します。

AdEMAMix との大まかな違いは以下の通りです。

短期の勾配情報として AdEMAMix では $m_1$ を利用しますが、Lion では勾配そのものを利用します(AdEMAMix の $beta_1 = 0$ と等価)
各パラメータに対する更新量は常に一定で、勾配情報の符号のみを利用します

QHM/QHAdam

論文

https://arxiv.org/abs/1810.06801

QHM の更新式は以下の通りです。
他のアルゴリズムとの関係で、論文中で g で表される変数は m と、$\nu$ で表されるパラメータは $\gamma$ と表記しました。

\begin{align}
m_0 &\gets 0 \\
m_t &\gets \beta m_{t-1} + (1 - \beta) g_t \\
\theta_t &\gets \theta_{t-1} - \eta \{ \gamma m_{t} + (1 - \gamma) g_t\}
\end{align}

QHAdam の更新式は以下の通りです。

\begin{align}
m_0, v_0 &\gets 0 \\
m_t &\gets \beta_1 m_{t-1} + (1 - \beta_1) g_t,   \hat{m}_t \gets \frac{m_t}{1 - \beta_1^t}  \\
v_t &\gets \beta_2 v_{t-1} + (1 - \beta_2) g_t^2, \hat{v}_t \gets \frac{v_t}{1 - \beta_2^t}  \\
\theta_t &\gets \theta_{t-1} - \eta \frac{\gamma_1 \hat{m}_t + (1 - \gamma_1) g_t}{\sqrt{\gamma_2 \hat{v}_t + (1 - \gamma_2)g_t^2} + \epsilon}
\end{align}

$m$ に長期の勾配情報を保存します。短期の勾配情報は保存せずに、最新の勾配をそのまま利用します。

AdEMAMix との大まかな違いは以下の通りです。

QHM は Adam 系ではなく SGD ベースなので勾配の大きさを勾配の移動二乗平均平方根(moving RMS)で補正しません
短期の勾配情報として AdEMAMix では $m_1$ を利用しますが、Lion と同様 QHM/QHAdam では勾配そのものを利用します(AdEMAMix の $beta_1 = 0$ と等価)
QHAdam では更新量を調節するための移動二乗平均平方根(moving RMS) に対しても、$\gamma_2$ を小さくすることで最新の勾配の割合を増やすことができます。ただし、長期の勾配情報を利用するという意味合いは薄いと考えています

AggMo

論文

https://arxiv.org/abs/1804.00325

更新式を以下に示します。変数名は論文のものから変更しています。

\begin{align}
m_{i,0} \gets 0 for all i\\
m_{i,t} \gets \beta_i m_{i,t-1} - g_t, for all i \\
\theta_t \gets \theta_{t-1} + \frac{\eta}{K} \sum_{i = 1}^{K}{m_{i,t}}
\end{align}

任意の数(式中ではKと表記)のモーメンタム係数を利用できます。

AdEMAMix とのおおよその違いは以下の通りです。

Adam 系ではなく SGD ベースなので勾配の大きさを勾配の移動二乗平均平方根(moving RMS)で補正しません
勾配情報を保持するモーメンタム係数を任意の数設定できます(AdEMAMixは $beta_1$ と $beta_2$ の2つ)
複数の勾配情報を混ぜる際に、重みなしの平均を使います(AdEMAMix は古い勾配情報の重みとして $\alpha$ を利用します)

以上