Lookahead Optimizerについて

はじめに

TensorFlow Addonsでも提供されている、Lookahead Optimizerについて簡単に記事にする。

Lookahead Optimizerとは

論文はこちら。
Lookahead Optimizer: k steps forward, 1 step back(19 Jul 2019)

日本語による解説記事はこちら。
論文解説「Lookahead Optimizer: k steps forward, 1 step back」

以下の図が動作を示したもの。

簡単に動作を説明すると次のとおり。

1. 重みを記憶
2. 所定のステップ数だけOptimizeを実行して重みを更新(これはSGDでもAdamでもなんでもよい)
3. 所定のステップに達したら、最初に覚えた重みと現在の重みの間のどこかに重みを更新(つまり戻る)
4. 以下繰り返し

ハイパーパラメータとしては、上記のステップ数(sync_period)と、戻る際の重みの割合(slow_step_size)となる。
TensorFlow Addonsの実装では、sync_period=6、slow_step_size=0.5になっており、この場合は戻る際にちょうど中間の値になる。

戻るので速度が遅くなるような気がするが、論文では収束が速くなったとしている。「いそがば回れ」、ということらしいが、他のOptimizerと組みわせることも含めて、最近(2020/12)発表されたSAMと似ている気がする。
その他にも、汎化性能向上や、学習率の変化にたいしてもロバストになったと主張されている。

tf.kerasでの利用法

これはTensorFlowAddonsをつかえば割と簡単で、こんな感じで最初にOptimizerを渡せばよい。その後は学習率の設定も内部で受け渡ししてくれるので、Schedulerでも普通のOptimizerとして扱える。

import tensorflow as tf
import tensorflow_addons as tfa
optimizer = tfa.optimizers.Lookahead(tf.keras.optimizers.Adam(0.001))

実際の動作

単純な最適化を実施した場合の比較はこちら。

Lookaheadでは'step back'していることがわかる。

一方、ResNet50でCIFAR10を学習させた結果がこちら。CosineDecayで学習率は0.2から0.002になるように設定。

あまり明確な違いは出ないが、Lookaheadを使用したほうが認識率が低く収束した(SGDでは96.46、Lookaheadでは95.91)。何度か設定を変えたりしたが一貫してLookaheadを使用したほうが性能が悪化した、という残念な結果だった。
特に学習が速いわけでもないが、遅くなってもいないようだ。

以下の論文では効果があったというので、条件次第では効果があるようだ。
Lookahead optimizer improves the performance of Convolutional Autoencoders for reconstruction of natural images