[1908.03265] On the Variance of the Adaptive Learning Rate and Beyond [paper-reading]

[yoheikikuta]

論文リンク

https://arxiv.org/abs/1908.03265

公開日（yyyy/mm/dd）

2019/08/08

概要

Adam が抱える学習初期に学習率の分散が発散するという問題に着目し、それを解決する RAdam を提案。
経験的に warm up (学習初期は linear でスケールする小さな学習率で学習し、その後に所望の学習率スケジューリングでの学習につなげる) がうまくいくことを知っているが、その必要性を Adam では学習初期の学習率分散の発散と対応づけて説明。
肝になるのは、エポック数が小さい時に、gradient g が標準正規分布に従う → 1/g が逆カイ分布に従う → これは自由度が小さい時は分散が発散 → Adam の学習率はこれを使うので学習初期に望ましくない振る舞いをしてしまう、という流れである。
分散が大きくなるのが問題なのかという点に関しては、経験的に勾配の各成分の絶対値が比較的小さい平均の周りで大きな分散で分布するということを示して、うまく学習が進まなくなってしまうと主張している。
提案する RAdam では分散が ill-defined な学習初期ではモーメンタム法を使用し、それ以降では補正込みの Adam で学習することでこの問題を回避する。
学習率の初期値によって振る舞いが大きく変わる Adam や SGD とは異なり、RAdam はどのような初期値でも安定して良い性能を発揮することを示した。

著者実装：https://github.com/LiyuanLucasLiu/RAdam

[yoheikikuta]

久しぶりに paper-reading を書こうと思って、GitHub でもかなり注目されている RAdam の論文を読んでみることにした。最近の学習率スケジューリングは catch up してなかったので、新しめのものを読んで知識をつけようという目論み。

既に一通り読んでいて、どうも納得できない点がある。最初にそれを列挙しておこう。

分散が大きくなるのが問題なのかという点に関しては、経験的に勾配の各成分の絶対値が比較的小さい平均の周りで大きな分散で分布するということを示して、うまく学習が進まなくなってしまうと主張している。

これは根拠に乏しい気がするが、OpenReview の著者らのコメントを見ても future work という感じのようだ。こういう場合に問題が発生するというのは simplified case (2-layer CNN とか) では示されているとのこと。

gradient g が標準正規分布に従う

これは {g1, · · · , gt} as i.i.d. random variables drawn from a Normal distribution N(0, σ2) と本文で書かれている。t はステップ数。
重みが mean 0 の正規分布で初期化されているので良いと言っているが、学習のステップ数で（データやモデル構造に依存して）変動していくものなので、同じ N(0, σ2) からサンプルされていると考えるのは成り立たない気がするのだが。
ここは以降の理論的解析の重要な starting point で、逆に言えばこれを認めればあとは簡単な計算をしていけばいいだけなのだが、仮定として強すぎる気がする。
5.3 では zero mean でなくても経験的には同じ議論が成り立つと主張しているが、Figure 2 も同じ確率分布からサンプルされるという主張をサポートしてないように見えるし、うーむという感じ。

誰か詳しい人いたら教えてください。

[yoheikikuta]

学習初期の数ステップで小さい学習率で学習する warmup は経験的にはうまくいくことが知られていたが、ちゃんとした理論的裏付けはない。
これがなぜうまくいくのかということを明らかにして、それを取り込んだ学習率スケジューリング手法を提案する、というのがこの論文の主な取り組みである。

まず、適応的に学習率を変える手法を一般的に表現したものは以下のように書ける。
ここは特に新しいことはないのでガッと本文をコピペ。

$ \psi $ が適応的な学習率で、これの取り扱いがこの論文での主たる関心事となる。

[yoheikikuta]

warmup を入れる場合は最初の数エポックは $ α_t = t α_0 $ のように単純な linear scale で学習をして、それ以降は上記のアルゴリズムで学習していくことになる。

勾配 $ g_t $ はパラメタの次元だけ成分があるので、それぞれの成分の絶対値を取ってヒストグラムを作って分布を見てみたのが以下の図である。なんで絶対値を取ってるのかはよく分からない。
奥から 1 epoch から始まり手前に来ているものほど大きい epoch の結果である (iteration = epoch)。
上は warmup を入れていない場合で、特に学習初期（10 epoch までとか）は分布の歪みが大きい。そして後半の epoch で成分の大きさが warmup あり比べて小さく、これによって学習がうまく進まなかったと主張している。
ちなみに perplexity を比べると warmup なしが 500 程度で warmup ありが 10 程度で大きな差が出ている。

この勾配のヒストグラムだけを見て学習がうまくいったうまくいかないを結論づけるのは難しい気がするが、warmup あるなしで振る舞いが大きく異なり得る、というのは確かにあるらしい。

流れとしては、学習初期で分布の歪みが大きい場合に学習がうまくいかないと仮定する → これを引き起こす一つの要因は $ v_t $ の各要素の分散が大きい時（ちなみに演算は全て element-wise である）→ この分散が大きいというのは $ \psi $ の解析をする事で正当化される → warmup はそれを防ぐ役割を担っている → これを受けて新しい手法を提案、という感じである。

[yoheikikuta]

まずは parameter β を忘れて ψ$ (g) = \sqrt{ 1/g_1^2 } $ という場合を考える。
そして We view {g1, · · · , gt} as i.i.d. random variables drawn from a Normal distribution N(0, σ2) である。これは冒頭で言ったように強すぎる気がするが、ここでは認めておく。

• {g1, · · · , gt} を N(0, σ2) からサンプルしたものと仮定
• さらに t=1 の場合のみを考える
• g1^2 は自由度 1 のスケールされたカイ二乗分布に従う（カイ二乗分布は標準正規分布に従う確率変数の二乗和が従う分布である。スケールされたの意味は分散を吸収するように変数変換したの意味）
• 1/g1^2 は自由度 1 のスケールされた逆カイ二乗分布に従う（逆カイ二乗分布の定義そのもの）
• √1/g1^2 は自由度 1 のスケールされた逆カイ分布に従う
• これは（平均や）分散が発散するため、上で議論した状況を作ってしまうので悪さをしてしまう。
• それゆえに Var[α x] = α^2 Var[x] を使って小さい α で分散を抑えるのが効くはずで、それが warmup で実現されている

カイ二乗分布とは何だったのか、使うのがあまりに久しぶりだったので復習がてら確率密度関数 (x > 0) を思い出しておく。

自由度 1 の逆カイ分布の分散が発散するという計算は以下。
そもそも平均も発散してしまうが、強引にそれを無視。論文の発散度合いと違う（変数変換すると $ y^{-3/2} e^{-y} $ なので）ので誰か間違っているが、発散するというのは良いのでこれで進む。
誰か真面目に計算した人がいたら教えてください。

[yoheikikuta]

warmup が variance reduction の働きをしていることを踏まえ、簡単な手法として

• Adam-2k: 最初の 2000 epochs はパラメタやモメンタムを更新せずに learning rate だけ更新
• Adam-eps: ε = 10^-4 という大きめの値にする

を提案。
これの結果は以下で、Adam-2k は warmup と同様の性能（2000 epoch 目を他の手法の 1 epoch 目と合わせている）で、Adam-eps は warmup なしよりはいいけど少し悪いというものになっている。

これで論文の主張が正しそうなのではないかという雰囲気が出てくる。
以降では理論解析をもう少し押し進めて雰囲気を強める。

[yoheikikuta]

(1) 式に戻って、epoch 数が小さいときの $ \psi $ を考えてみる。
ルートの中身は exponential moving average の逆数になっているが、epoch 数が小さい、すなわち $ t $ が小さい場合に注目する。
Adam では $ β_2 = 0.999 $ などがよく使われる値で、$ t $ が小さければ $ β_2^t $ は $ (1 - ε)^t \sim 1 - ε t $ という一次の Taylor 展開が良い近似となっていて、leading order は単なる平均になる。

この形にしてしまえば、$ g_i $ が N(0, σ2) に従うというものだったので、$ t / \sum g_i^2 $ は自由度 $ t $ のスケールされた逆カイ二乗分布に従うことになる。
これにより、元の $ \psi $ のルートの中身は、対応する自由度を $ ρ $ としたときに、自由度 $ ρ $ のスケールされた逆カイ二乗分布に従うということが正当化される。

ここまで来ると、以降はあまり真面目に読まずとも論文の言わんとすることは分かる。
これまでの議論により、自由度が小さいうちは逆カイ（二乗）分布の分散は発散してしまうものであった。
しかしある程度自由度が大きくなると有限の値になる。
ということで自由度を求めて、それが分散を ill-defined にするうちは適応的な学習率は諦めて単純なモメンタム法を使う、というのが以降でやりたいことである。

よって、少なくとも $ t $ が小さいうちは {g1, · · · , gt} を N(0, σ2) から i.i.d でサンプルしたものとする という仮定が解析の根幹であることが分かる。
（前述のようにこれは強すぎる気はしているが...）

[yoheikikuta]

Theorem 1 で自由度が大きくなる毎にスケールされた逆カイ二乗分布の分散が減少することが示されているが、これは定理とすべきほどのものでもない単純な計算である。
Γ関数の定義を思い出せば follow するのは難しくないのでここでは割愛。
ちなみに自由度が 4 より大きければ分散が有限値になる。これは具体的な計算は全く必要なくて x の次数を眺めれば分かる。

自由度 ρ を β の言葉で書き直そう。
諸々の定義の導入が面倒なので、論文の議論をそのままコピペしよう。
ちなみに引用されているのは Duke 大のビジネススクールの教授のノート（ https://people.duke.edu/~rnau/Notes_on_forecasting_with_moving_averages--Robert_Nau.pdf ) である。

exponential moving average は simple moving average として近似することがよくあるらしい。
全然知らないが、単純な移動平均とするというのはまあありそうな話ではある。exponential moving average の寄与を期間を適切に定めた simple moving average で表現する、というだけだろう。
そのような期間 $ f(t, β_2) $ を定め、それを求める、というのみである。

これまでの議論を思い出せば、この $ f(t,β_2) $ が自由度として働いているわけなので、(5) 式は自由度 $ f(t,β_2) $ のスケールされた逆カイ二乗分布に従うということになる。
以降ではこの $ f(t,β_2) $ を $ ρ_t $ と書く。

また、$ t $ を無限大に飛ばしたものを $ ρ_∞ $ と定義している。

[yoheikikuta]

問題の分散であるが、やりたいことは分散を大きくしないように保つことであった。
スケールされた逆カイ二乗分布の分散は自由度が大きいほど小さくなる性質があったので、$ ρ_∞ $ の値 $ C_{var} $ を取るように factor $ r_t $ を掛けることにする。

ψ の解析的な表式も求まっているのでこれで終了かと思いきや、数値的に不安定なので工夫を施す。
ちなみに何で不安定なのかはちゃんと追ってないけど、ベータ関数が入っててそこの積分ということなんだろう。
期待値の周りで Taylor 展開して一次の項を取る ( $ \psi^2 = E[\psi^2] - (E[\psi^2] - \psi^2) $ として第二項を微少量と考えて一次まで残す)と、単純な計算で以下が得られる。

$ \psi^2 $ が自由度 $ ρ_t $ のスケールされた逆カイ二乗分布に従うことから、この近似で得た $ Var[\psi^2] / 4 E[\psi^2] $ を使えば、$ Var[ψ] $ は以下のように計算できる。スケールされた逆カイ二乗分布のモーメントはガンマ関数を考えれば計算は難しくないし、面倒であれば wikipedia を見ればよい（$ τ^2 = 1/σ^2 $ と読み替えればよい)。

最終的に求めたかった $ r_t $ はこれを分母にして分子は $ ρ_∞ $ にしたもので全体の√を取ったものであったので、以下が必要な rectification term である。
これらは全て $ β_2 $ で表現できているので、あとは各ステップで単純な演算によってこの term を求めればよいということになる。

[yoheikikuta]

以上までで Rectified Adam (RAdam) アルゴリズムの準備が整った。
結論だけ見れば難しい操作はなく、分散が ill-defined なうちはモメンタム法を使い、well-defined になれば rectification term を計算して乗じることで一定の分散をキープした Adam で学習していくというものになっている。

[yoheikikuta]

実験も色々とやっているが、代表的なものを紹介しておく。
Adam よりは常に良いが、うまく tuning した SGD には劣る、という感じ。
学習率の初期値依存性が少ないので、あれこれ試行錯誤するよりは RAdam を使っておけば一件落着！というようには見える。
自分はまだ試してないけど、これが広く成り立つのであれば確かにかなり有用であろう。

warmup とも比較していて、RAdam と同程度の性能を出すには warmup をどこまでやるかを tune しないといけないという結論を出している。

RAdam は細かい tuning いらずでかなり良い性能を発揮する、というのが利点である。
これは色んな人が色んなデータで試してどれくらい有効なのかが検証されていくのが楽しみだ。

[yoheikikuta]

最後に、論文で仮定したものがどれくらい正そうかを検証している。

分散の計算で期待値周りで Taylor 展開した正当性は以下の図で示されているとのこと。
自由度が大きいところではいいけど、小さいところでは比較的ズレは大きい。この手法では自由度が小さいところが肝なのでやや気になるが、十分小さいという主張。

gi が平均 0 の正規分布に従うという仮定は、0 じゃないかもしれないので $ μ $ として $ g_i $ をサンプルしてオリジナルの適応的学習率と提案手法を比較している。結果が下の図で iteration が小さいうちはオリジナルの値が大きくて提案手法が小さいので、確かに提案手法で分散を抑えられているという主張。
これがどういう実験かイマイチ掴みかねているが、自分の理解では $ g_i $ が N(μ,σ2) に従うというのは受け入れて、その上で { $ g_1, ..., g_t $ } 系列を作り、これらを用いてそれぞれの $ \psi $ を計算して比較したということなのだろう。
自分としては $ g_i $ が N(μ,σ2) に従うというのがやはり納得いってないのだが...

[yoheikikuta]

ということで一通り読んだ。
できあがった手法を見てみると複雑ではないし、結果も良さそうに見えるので色々なところでテストされて真に良い手法かどうかが明らかになるのは楽しみである。

解析としても難しくはないのだが、仮定の部分で自分は納得がいっていない。
誰か詳しい人がいたら教えてください。