はじめに

RectifiedAdamの実装を調べていたら、いくつか疑問点があったので記事にする。

RectifiedAdamとは

細かいところは参考に記載した記事等をみてもらうとして、開発までの流れをものすごく簡単に言うと以下の様になるだろう。

とりあえず、Warmupを入れると成績が上がる。
しかし、Warmupの設計には経験則しかなかった。
そこで、数学的に根拠のあるWarmupをAdamに組み込んだ。

RectifiedAdamでは下記の主張がされている。

Adamよりも性能向上
学習率を大きく変えても同じような性能になる

これらは本当だろうか？

RectifiedAdam内のWarmupとはどんなものか

論文にはいろいろ根拠が書いてあるが、結局のところ初期パラメータのbeta_2の値で決まる係数を学習率に掛ける、ということになっている。

元の式はこんな感じに書き換え可能。($\rho_t>4$の場合のみ抜粋)

$ \rho_{\infty} \leftarrow 2 / (1-\beta_2) - 1 $
$ \rho_{t} \leftarrow \rho_{\infty} - 2t\beta_2^t/(1-\beta_2^t) $
$ l_t \leftarrow \sqrt{(1-\beta_2^t)/v_t}$
$ r_t \leftarrow \sqrt{\frac{(\rho_t-4)(\rho_t-2)\rho_{\infty}}{(\rho_{\infty}-4)(\rho_{\infty}-2)\rho_t}} $
$ \hat{\alpha_t} \leftarrow \alpha_tr_{t}l_t$
$ \theta_t \leftarrow \theta_{t-1} - \hat{\alpha_t}\hat{m_t}$

$r_t$が1.0の場合は普通のAdamだから、$r_t$だけがRectifiedAdamの寄与する所で、ここを調節することでWarmupの代わりとする。
$r_t$は$t$と$\beta_2$だけできまるので、結局Warmupの形はbeta_2で決まるということだ。

$r_t$を実際にグラフにしたものがこちら。

Adamのデフォルトはbeta_2=0.999だが、この場合6000step辺りまでいかないとWarmupがおわらない。Warmup終了後が本番と考えると、少なくとも10000stepぐらいは実施しないと十分な学習ができないと思われる。仮に100step/epochとすると100epochということになる。
この辺を理解しないでAdamと同じ感覚で短いStep数で学習を終わらせてしまうと、Adamより性能が低下することになるだろう。
また、beta_2が0.99と0.999では大きく違うので、この辺も調整が必要になる可能性がある。つまりハイパーパラメータの調節要素が増えるということも意味する。

ちなみに、$\rho_t<4$では$r_t$が計算できないため、最初の短い期間のみMomentum付きSGDとなる。この期間はWarmupなしと考えて近似すると$r_t=1$としてよいだろう。その仮定で、最初のほうのステップのみクローズアップすると以下の様な図になる。

おおむね5ステップ分はWarmupなしで、その後急激に0になる。Warmupの目的からすると最初は学習率0付近であるべきところのはずで、なぜこのような設計にするのか理解できない。厳密に計算できなくても(例えば直線で近似するなどして)0からつながるようにすべきではなかろうか。

初期ステップでのAdamとの動作の違いをグラフにしてみたものがこちら。
(TensorFlow Addonsに実装されているRectifiedAdamと、TensorFlowのAdamで比較)

グラフ描画コード

初期値1.0から0.0へ重みが変化する様子で、傾きが大きいほど速く学習していることを示す。(詳細は別記事など参照のこと)
RAdamは5step目からはWarmup期間に入るので一気に傾きが平らになるが、その前の区間では逆にAdamよりも急速に学習してしまっている。条件にもよると思われるが、本来は学習を押さえるべきところで、逆の効果を発揮しているように見える。

性能向上？

十分なステップ数が確保できるならば、Warmupを入れれば性能は上がる場合は多いので、RectifiedAdamの性能が上なのは当然の結果で別に驚くこともないように思える。Warmupを使えば同様の性能向上は可能なので(これは論文にも書いてある)、RectifiedAdamの特性というほどのことではないのではないか。

上述したように、ステップ数が少ない場合は性能低下する可能性がある。実際にこのページの実験では性能低下を報告している。

学習率の違いの影響は？

元論文では実験結果とともに、下記のようにある。

by rectifying the variance of the adaptive learning rate, RAdam improves the robustness of model training and can adapt to different learning rates of a broader range.

幅広いレンジの学習率に適応できるとしか書いてないが、これを拡大解釈すると「学習率の調節不要」となってしまう。
実際、筆者は論文の実験結果のグラフをみて一瞬そう思ったのだが、そもそも実質Warmupを加えただけなのにそんなことが可能なのだろうか。
これは実験で確認する。

実験

TensorFlow AddonsのRectifiedAdamとTensorFlowのAdamでCIFAR10の認識率で比較した。beta_2はデフォルトの0.999。全200エポックで150エポック以降は学習率1/10にする。1エポックあたり約100ステップなので、合計約20000ステップ。2回学習させてテストデータの正解率を記載。

アルゴリズム	learning rate	Low	High
RAdam	1e-1	95.05	95.10
Adam	1e-1	93.56	94.10
RAdam	1e-2	95.70	95.74
Adam	1e-2	94.21	94.32
RAdam	1e-3	95.55	95.79
Adam	1e-3	94.80	95.25
RAdam	1e-4	92.46	92.80
Adam	1e-4	93.05	93.17

２回のうち好成績のほうだけでグラフにしたものがこちら。

1e-1～1e-3までは、RAdamでは認識率が向上し、その差も減っているので主張通りではある。ただし差が減ったとはいえ違いは明確にあるので、学習率の設定が不要になったわけではない。
また、1e-4に関してはAdamよりも性能が低下した。学習率が小さいとRAdamではWarmupが入る分学習が遅れてしまうためと思われる。つまり、性能が必ずしも向上するわけではない。

まとめ

というわけで、筆者の意見をまとめると以下の通り。

RectifiedAdamはやや複雑な式のWarmupを内部に実装しただけで革新的なところは無い
Warmupがあることが隠蔽されているため、十分なステップ数が必要なことを理解していないと十分に学習しないまま終了してしまう可能性がある
ハイパーパラメータの調節が不要というわけではない
最初の数ステップの学習率は高すぎて、Warmupの目的からずれているのではないか

十分ステップ数があれば性能向上することは確かだが、LearningRateSchedulerをつかって従来のAdamと組み合わせれば同等の結果はでるはずなので、あえて使用する必要はないように思う。

参考

ついにAdamを超えた！最新の最適化アルゴリズム「RAdam」解説
 On the Variance of the Adaptive Learning Rate and Beyond, Liu, L. et al. ICLR 2020