初めに

AdamでLearningRateScheduleをうまく使うと、理論上はRectifiedAdam(RAdam)相当になる。ここでは、TensorFlow/Kerasで実際に実装し、それを適用して学習結果をRectifiedAdam等と比較してみる。

RAdamとは

単純に言えば、Adam+固定Warmupである(以前書いたこちらの記事参照のこと)。
"固定"というのは初期パラメータ(具体的にはbeta_2)であらかじめ決まっているという意味で、TensorFlow/KerasでのLearningRateScheduleで置き換え可能である。
これについてはすでに論文があり、そちらではより簡単な形状でも同等の性能を出せることが示されている。同論文で提案されている方法についても合わせて比較する。
ちなみにその論文では、

RAdam: perform 4 iterations of momentum SGD,then use Adam with fixed warmup

という見出しでRAdamの説明がなされているが、最初の4回だけのmomentumSGDは実質無視してよい、と論文内で結論付けている。

LearningRateSchedule

tf.kerasでは学習率の制御には"tf.keras.callbacks.LearningRateScheduler"を使ってエポック単位で変更することが多いはずだが、RAdamではstep毎(つまりバッチ単位)に学習率を制御しているのでそのままでは同じ処理にできない。
そこで、"tf.keras.optimizers.schedules.LearningRateSchedule"というクラスをOptimizerのlearning_rateに直接渡すとステップ毎に制御できるので、これを使って実現する。
RectifiedAdam相当のほかに、上にあげた論文で提案されている２つのWarmupも実装する。

StepDecay

まずベースとなるクラスを作成しておく。
共通処理として「既定ステップ数を超えたら、学習率を変更する」という単純なStepDecay処理が入っている。
lr=学習率/decay_start=規定ステップ数/decay_ratio=規定ステップ後の学習率の変化となっている。
beta_2はここでは使わないが定義だけしておく。

import tensorflow as tf
class BaseSchedule(tf.keras.optimizers.schedules.LearningRateSchedule):
    def __init__(self, **kwargs):
        super().__init__()
        self.lr = kwargs['lr']
        self.decay_start =  kwargs.get('decay_start', 0.0)
        self.decay_ratio = kwargs.get('decay_ratio', 1.0)
        self.beta_2 = kwargs.get('beta_2', 0.999)
    
    def get_config(self):
        config = {"lr": float(self.lr), 
                  "decay_start": float(self.decay_start),
                  "decay_ratio": float(self.decay_ratio),
                  "beta_2": float(self.beta_2)
                  }
        print('get_config')
        return config

    def __call__(self, step):
        return tf.where(step >= self.decay_start, self.decay_ratio*self.lr, self.lr)

例えば150エポック以降で学習率を1/10したい場合は以下の様にOptimizerに渡す。

steps_per_epoch = 100
lr_schedule = BaseSchedule(lr=0.001, decay_start=steps_per_epoch*150, decay_ratio=0.1)
optimizer = tf.keras.optimizers.Adam( learning_rate = lr_schedule )

RectifedAdam相当

RectifiedAdamの論文から関係のあるところだけ抜粋するとこんな感じになる。
$ \rho_{\infty} \leftarrow 2 / (1-\beta_2) - 1 $
$ \rho_{t} \leftarrow \rho_{\infty} - 2t\beta_2^t/(1-\beta_2^t) $
$ r_t \leftarrow \sqrt{\frac{(\rho_t-4)(\rho_t-2)\rho_{\infty}}{(\rho_{\infty}-4)(\rho_{\infty}-2)\rho_t}} $

MomentumSGDになる区間は学習率0として実装する。
計算のため、beta_2を指定する必要がある。

class RectifiedAdamSchedule(BaseSchedule):
    def __call__(self, step):
        lr = super().__call__(step)
        sma_inf = 2.0 / (1.0 - self.beta_2) - 1.0
        beta2_power = self.beta_2**step
        sma_t = sma_inf - 2.0 * step * beta2_power / (1.0 - beta2_power)
        r_t =  tf.sqrt( (sma_t - 4.0) / (sma_inf - 4.0) * (sma_t - 2.0) / (sma_inf - 2.0) * sma_inf / sma_t)
        return tf.where(sma_t > 4.0 , r_t*lr, 0.0)

ExponentialWarmup

検証論文で提案されているうちの一つで、以下の数式でRectifiedAdam相当になるという。

$ r_{t}=1-\exp(-(1-\beta_2)\cdot t) $

class ExponentialWarmupSchedule(BaseSchedule):
    def __call__(self, step):
        lr = super().__call__(step)
        return (1.0 - tf.math.exp( -(1.0-self.beta_2)*step ))*lr

LinearWarmup

これも検証論文で提案されているうちの一つだが、単純な線形増加でも同じような結果になるという。
$ r_{t}=\min(1, \frac{1-\beta_2}{2} \cdot t) $

class LinearWarmupSchedule(BaseSchedule):
    def __call__(self, step):
        lr = super().__call__(step)
        return tf.math.minimum( 1.0, (1-self.beta_2)/2*step )*lr

実験

特性確認

以上の仕組みで計算した学習率のグラフはこちら。(Adamのデフォルト値であるlr=0.001、beta_2=0.999とした場合)

ExponentialWarmupがRectifiedAdamとかなり似た形状を持つことがわかる。
またLinearWarmupのグラフを見ると、計算上2000ステップまでの線形増加となったことがわかる。これはエポックごとのステップ数が比較的すくなければ、CallbackのLearningRateSchedulerをつかって簡単に近似できるので、後の実験ではその結果も示す。

同条件で、極簡単な重みの学習曲線を描画したものがこちら。
TensorFlow AddonsのRAdamの結果も加えてあるので、実際のRAdamとの比較もできる。

見てわかる通り、RAdam相当のAdamはRAdamとほぼ同じ特性を持つ。(最初の４ステップでMomentumSGDになっていない分若干のずれが出る)

描画用コード

CIFAR10で実験

CAFAR10を21層のVGGライクなCNNで画像認識させた結果を記す。
学習率は全て0.001、全200エポック中150以降は1/10にする。beta_2は0.999。同じ条件で3回実施。
RAdamはTensorflowAddonsのRectifiedAdamを用いた。
表中の'Linear(epoch)'というのはエポック単位(約100ステップ)で単純増加させた場合の結果を指す。

Optimizer	Shedule	Worst	Median	Best
Adam	RectifiedAdam	95.66	95.78	95.82
Adam	Exponential	95.56	95.62	95.71
Adam	Linear	95.63	95.66	95.72
Adam	Linear(epoch)	95.61	95.73	95.81
RAdam	-	95.64	95.73	95.79
Adam	-	94.80	94.88	95.25

以下のことがわかる。

RAdamを使うとAdamに比べて明白に性能向上している
AdamにWarmupを加えたものは、いずれもRAdamと同等の性能を示した
エポック毎の線形増加でもRAdam同等なので、エポックあたりのステップ数が小さい(この実験では100ステップ)という条件であれば、エポック毎に制御するWarmupで十分である

参考にした論文では、ResNet-50やTransformerでの結果が示されているが、同様にRAdam同等の結果が出たとしている。

まとめ

AdamをRectifiedAdam相当の性能にする方法を示した。
単純にエポックごとに増加させることでも十分性能向上は可能。beta_2=0.999であれば、2000ステップを目安に単調増加させればよい。

参考

ついにAdamを超えた！最新の最適化アルゴリズム「RAdam」解説

RectifiedAdam論文
On the Variance of the Adaptive Learning Rate and Beyond, Liu, L. et al. ICLR 2020

RAdam検証論文(@hasebirokou 氏から情報提供いただいた)
On the adequacy of untuned warmup for adaptive optimization