こんにちは！AIisFunです！
初めてQiitaに投稿します。よろしくお願いします。
論文概要
Kartik Chandra氏らによる論文「Gradient Descent: The Ultimate Optimizer」がNeurIPS 2022で発表され、話題になっているようです。論文を読んでみましたので、概要をご紹介いたします。
ざっというと、SGD(確率的勾配降下法)などにおける学習率といったハイパーパラメータを誤差逆伝搬法の枠組みの中で、逆伝搬の各ステップで自動で更新、決定してしまうautomatic differentiation (AD)という手法の提案です。SGDだけではなく、複数のハイパーパラメータを有するAdamなどにも適用可能です。さらには、ハイパーパラメータの学習率そのものも、hyper-hyperparameterとして学習可能です。
この仕組みをプログラムに取り入れることは容易で、本論文ではPyTorchでの実装例や注意点、MLP、CNN、RNNに適用した実験例も紹介されています。
では、元論文の順に従って解説してきます。
1. 導入
深層学習の勾配降下法においては学習率α(本論文ではstep sizeと表現されていますが、学習率とほぼ同じ意味で使われていると思われますので、以下、学習率と記します）を適切に決定する必要があります。αを大きくしすぎると学習が収束しませんし、小さくしすぎると学習に時間がかかります。（いきなり脱線しますが、この辺の事情については岡野原大輔著の『ディープラーニングを支える技術２』の第１章に素晴らしい解説があります）。
誤差逆伝搬においてはAdamなどのオプティマイザーを用い、ハイパーパラメータである学習率の初期値$\alpha$、$\beta_1$、$\beta_2$については何となくデフォルト値を用いている、という方も多いかもしれません。しかしながら、Adamは必ずしもSGDより良い結果を与えるというわけでもなく、また、これらのハイパーパラメータの与え方も重要になる場合もあります。本論文では、モデルのパラメータを更新する際に、ハイパーパラメータも勾配降下法の仕組みの中で一緒に更新してしまうという提案を行っています。
Baydinら(2018)による先行する研究があり、SGD(確率的勾配降下法)における学習率の学習を定式化していますが、Chandra氏はその欠点として
- 事前に微分式をマニュアルで導いておく必要があり、SGD以外の他の降下法に応用しようとすると、式が複雑になりすぎて間違いが生じやすい
- モーメンタム法や、ADAMに適用する場合、学習率以外の他のハイパーパラメータには応用できない。
と指摘しています。ハイパーパラメータを含めた式の微分を自動化するautomatic differentiation(AD)の手法により、これらの問題点が解消され、さらには、ハイパーパラメータそのものの学習率となるhype-hyperparameterも効率的に導入可能、さらには、ハイパーパラメータを多段にすることで、ハイパーパラメータの初期値依存性が軽減されると主張しています。
2. ハイパーパラメータ更新の定式化
2.1. ハイパーパラメータの更新式
まずは、Baydinらの先行研究によるハイパーパラメータ公式の定式化に関する説明です。
SGDにおけるモデルパラメータの更新は良く知られている以下の式で表されます。
w_{i+1}=w_i-\alpha \frac{\partial f(w_i)}{\partial w_i} \;\;・・・（１）
ここで、$f$は目的関数、$w$はモデルパラメータ、$i$は学習のステップ、$\alpha$は学習率です。
これを学習率の更新を含める場合に拡張すると、次の2つの式で表されます。
\alpha_{i+1}=\alpha_i-\kappa \frac{\partial f(w_i)}{\partial \alpha_i} \;\;\;\;\;\;\;・・・（２a）\\
w_{i+1}=w_i-\alpha_{i+1} \frac{\partial f(w_i)}{\partial w_i} \;\;\;\;・・・（２b）
ここで、$\kappa$は学習率を更新するための学習率です。
ここで、$\partial f(w_i)/\partial \alpha_i$をBaydinに従い、手作業で微分を行います。
\frac{\partial f(w_i)}{\partial \alpha_i}=\frac{\partial f(w_i)}{\partial w_i}\cdot\frac{\partial w_i}{\partial \alpha_i}=\frac{\partial f(w_i)}{\partial w_i}\cdot\frac{\partial(w_{i-1}-\alpha_i \frac{\partial f(w_{i-1})}{\partial w_{i-1}})}{\partial \alpha_i}\\
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~=\frac{\partial f(w_i)}{\partial w_i}\cdot (-\frac{\partial f(w_{i-1})}{\partial w_{i-1}})
しかし、このような解析的な微分を他の勾配降下法、例えばAdamのパラメータ$\alpha$、$\beta_1$、$\beta_2$などに適用しようとすると、非常に複雑な式になります。
2.2. ハイパーパラメータの自動微分
そこで、本論文では、理論式を導くことなく、ハイパーパラメータの微分を自動で行う方法（AD:Automatic Differentiation)を提案しています。まず、上（１）式に対応するPyTorchの疑似コード以下に示します。
def SGD.__init__(self, alpha):
self.alpha = alpha
def SGD.step(w):
d_w = w.grad.detach()
w = w.detach() - self.alpha.detach() * d_w
ここでは、学習率$\alpha$をテンソルとして表現しており、detach()によって計算グラフからの切り離します。これによって誤差の逆伝搬がその下方、つまり、前の学習ステップに伝わらないようにしています。
次に、上の（２a）、（２b）式に相当する疑似コードを示します。
def HyperSGD.step(w):
# 式(2a)の基づく学習率αの更新
d_alpha = self.alpha.grad.detach()
self.alpha = self.alpha.detach() - kappa.detach() * d_alpha
# 式(2b)に基づくモデルパラメータwの更新
d_w = w.grad.detach()
w = w.detach() - self.alpha * d_w
最後の式で、右辺の.detach()が無いことにご注意ください。
計算グラフと、detachの様子を表したのが下の図です。
左図は$\alpha$を更新しない場合で、モデルパラメータの更新3ステップ分が示されていますが、当然ながら更新ステップを越えて誤差逆伝搬は行われません。右図では$\alpha$も誤差逆伝搬により更新されるため、各更新に$\alpha$による微分が含まれています。code2.pyの最後の式でalphaにdetach()が無い理由がここにあります。一方で、更新のない$\kappa$や、前回ステップの$\alpha$などは切り離されています。
2.3. SGD以外の勾配降下法への適用
以上の手法をAdamなど他のオプティマイザーに適用することも比較的容易です。具体的なコードはgithubに公開されています。
https://github.com/kach/gradient-descent-the-ultimate-optimizer
公開コードにはSGD、Momentum(SGDに含まれている）、AdaGrad、RMSProp、Adamが含まれています。ここで注意すべきことは、例えばAdamのハイパーパラメータ$\beta_i$、$\beta_2$などは0から1の範囲になくてはならないため、ハイパーパラメータ更新時にこの範囲の外にはみ出さぬよう、更新時にシグモイド関数などで変数変換を行うことで値域を0～1に収め、更新後に逆変換で元の範囲に戻す必要があります。また、初期値も0にならないようにセットする必要があります。
2.4. ハイパーパラメータ学習の再帰的なネスティング
これまでの例では$\alpha$の学習率$\kappa$は一定となっていましたが、これ自体を勾配降下法で更新させることができます。その疑似コードは以下の通りです。
def HyperSGD.__init__(self, alpha, opt):
self.alpha = alpha
self.optimizer = opt
def HyperSGD.step(w):
# κの更新
# optimizerはκを学習パラメータとするSGD
self.optimizer.set(self.alpha)
# wとαの更新
# optimizerはαを学習パラメータとするSGD
self.optimizer.step(self.alpha)
d_w = w.grad.detach()
w = w.detach() - self.alpha * d_w
opt = HyperSGD(alpha=0.01, opt=SGD(kappa)
ここで、opt=SGD(kappa)とすることで、ハイパーパラメータの更新にはSGDを適用しています。さらに$\kappa$自体を学習させる場合には、HyperSGD(0.01, HyperSGD(0.01, SGD(0.01)))というように、再帰的にオプティマイザーを積み上げます。このように高階構造にすることで、初期値(ここでは$\alpha=0.01$といったもの）への依存性が低減できます。
3. 実験
論文ではADの効果を確かめる様々な実験結果が示されていますが、その一部を紹介します。
3.1 MNINST
まずは最初の例として、データはMNIST、モデルは全結合層1層(ニューロン数128）の単純なネットワーク、モデルパラメータの更新はSGDで、ハイパーパラメータの更新にはSGD/Adam/AdaGrad/RMSPropを適用しています。バッチサイズは256、エポック数＝30です。
下図で、1つ目の結果が学習率$\alpha$ = 0.01で一定の結果で、これに対し、ハイパーパラメータをADで変化させた場合は、誤答率が大幅に下がっています。なお、SGD(0.0769)のようにピンクの数字で示されたケースでは、すぐ上の方法で得られた最終の$\alpha$を初期値として同じ手法で再度計算を繰り返した結果を示しています。特にSGD＋Adamの組み合わせで得られた結果を初期値として再計算した結果は著しく改善しています。
同様にモデルパラメータの更新をAdamで、ハイパーパラメータの更新をSGDまたはAdam行った結果を下図に示します。モデルパラメータの変化に応じて学習率を変化させていくAdamをベースラインとしているので、改善の度合いは小さめですが、それでもはっきりと改善が見られます。Adam$^\alpha$はAdamのハイパーパラメータのうち、学習率のみを更新の対象としたもので、$\beta_1$、$\beta_2$は更新の対象としていません。Adam$^\alpha$/SGDの組み合わせケースではAdam/SGDのケースに比べわずかに悪化、Adam/Adamの組み合わせでは$\beta_1$、$\beta_2$を更新しなくてもテスト誤答率ほとんど変わりませんでした。
3.2. CIFAR10/ResNet-20
次はCIFAR10画像データセットを用い、10種類の物体のクラスを予測する問題をResNet-20(Heら、2016)で行った実験で、モーメンタム法を適用しています。ベースラインを学習率$\alpha$=0.1、モーメント$\mu$=0.9で一定とした場合としており、訓練時の損失が下図の灰色の線で示されています。青線は、$\alpha$を0.01、0.1、1.0のいずれかで、$\mu$を0.09、0.9、0.99のいずれかでそれぞれ固定した場合の結果を示しています。黒線がADにより$\alpha$と$\mu$を変化させたもので、$\alpha$と$\beta$の初期値に対し、ハイパーパラメータの学習率$\alpha_\alpha$と$\alpha_\mu$を、それぞれ$\alpha^2\cdot 10^{-3}$、$1/(1-\mu)\cdot 10^{-6}$としています。
多くのケースで、ベースラインケースや、学習率とモーメントを一定にしたケースと同程度もしくはより良い結果を出しています。学習率とモーメントの初期値を共に極端な値に設定した右下のケースでは、ベースラインの結果より訓練損失が高くなっているものの、それでも訓練損失はエポックと共に徐々に減少しているのが見て取れます。もっとも、学習率$\alpha$を1.0という極端な値に設定することは通常ないでしょう。
下左図ではエポック数を500まで増やし、①$\alpha$=0.1、$\mu$=0.9と一定にした場合、②上と同じくHeら論文の設定にならい、$\alpha$をエポック100以降、0.01, エポック150以降0.001と変更した場合、③AD法で$\alpha$を変更させていった場合のテストデータに対する誤答率を示しています。下右図はこれらの場合の学習率$\alpha$の変化を示しています。Heらの手法では試行錯誤的に学習率のスケジュールを定めたと思われますが、同様の結果(誤答率)がADにより得られています。
3.3. 高階のハイパーパラメータ最適化
最後にハイパーパラメータの学習率の最適化を高階化(ネスティング)した結果を示します。
3.1.のデータを用い、各階層のハイパーパラメータの初期値を様々に変えた場合のテストデータに対する誤答率を示しています。ただし、全階層におけるハイパーパラメータの初期値は同一の値に設定しています。ハイパーパラメータを変化させない場合（height=0、灰色の線）に対し、ADの結果は階層height=1で$\alpha$の初期値が低い場合と、$\alpha$の初期値が著しく高い場合を除き、結果は良好で、高階の場合、初期値の影響を受けにくいことが見て取れます。
ハイパーパラメータの学習率の初期値は階層ごとに変えたことが良いことを著者は示しており、初段の$\alpha$が$10^{-8}$～$10^{-4}$のときは、下段より順に、$\alpha\cdot 10^2$、$\alpha\cdot 10^0$、$\alpha\cdot 10^{-2}$、初段の$\alpha$が$10^{-3}$以上のときは$\alpha\cdot 10^{-3}$、$\alpha\cdot 10^{-4}$、$\alpha\cdot 10^{-8}$とするのが良いようです。ここは少しトリッキーですね。こうした場合の結果を下図に示します。
一段しかない場合でも、劇的に結果が改善しています。
高階にすると、計算時間の増加が気になるところですが、1段増やしても1-2％しか計算時間が増えなかったとのことです。
4. 制約
2.2.では誤差逆伝搬におけるグラフの切断について説明していますが、たとえばPyTorchにビルトインされているLSTMでは、内部で計算グラフが変更されてしまうために、別途コードする必要があるそうです。なお、今回、説明を省略した実験例では、LSTMによる計算が実装されています。（ビルトインされたモデルにはLSTM以外にも同様の問題がある可能性があり、要注意です）
最後に
学習率といったハイパーパラメータの選択やスケジューリングは面倒で、深く考えずにAdamに任せてしまうことも良くありますが、Adamは便利ながら最適ではないことは3.2の例で示されているとおりです。そういう意味で、本論文は比較的シンプルな方法で学習率を最適化させる、従来のオプティマイザーとは異なる手法を示したと言えるのではないでしょうか。
Qiita初投稿に加え、TEXで数式を書くのも初めてで何かと不備があると思いますが、ご容赦ください。記述方法でアドバイスをいただければ幸いです。