TensorFlow/KerasでAdaBeliefを実装してみた

# はじめに

AdaBeliefという最適化アルゴリズムが発表されたらしいので、実際に実装して評価してみた。
以前、TensorFlow/KerasでのOptimizerの実装方法を[記事](https://qiita.com/T-STAR/items/6a489ee8f3b6b213bede)にしたのだが、こちらはその続編でもある。この辺に興味のない方は実験結果だけみてもらえればよい。

## 環境
 * TensorFlow 2.3.0
 * Google Colab

# まずはMomentumSGDから

ベースとなるMomentumSGDのコードを載せる。
前回の記事で作成したものと基本的には同じだが改造にむけて若干の修正が入っている。複数のアルゴリズムが実装されることになるので、MultiOptimizerV1とした。この時点ではVanillaSGDとMomentumSGDが実装ずみ。

```Python:MultiOptimizerV1
from tensorflow.python.keras.optimizer_v2.optimizer_v2 import OptimizerV2
class MultiOptimizerV1(OptimizerV2):
    def __init__(self, learning_rate=0.01, momentum=0.0, centered=True, name='MultiOptimizerV1', **kwargs):
        super().__init__(name, **kwargs)
        self._set_hyper('learning_rate', kwargs.get('lr', learning_rate))
        self._set_hyper('decay', self._initial_decay)
        self.momentum = momentum
        self.centered = centered if momentum!=0.0 else False

    def get_config(self):
        config = super(MultiOptimizerV1, self).get_config()
        config.update({
            'learning_rate': self._serialize_hyperparameter('learning_rate'),
            'decay': self._serialize_hyperparameter('decay'),
            'momentum': self.momentum,
            'centered': self.centered,
        })
        return config
    def _create_slots(self, var_list):
        for var in var_list:
            if self.momentum!=0.0:
                self.add_slot(var, 'm')

    def _resource_apply_dense(self, grad, var, apply_state=None):
        var_device, var_dtype = var.device, var.dtype.base_dtype
        lr_t = self._decayed_lr(var_dtype)
        local_step = tf.cast(self.iterations+1, var_dtype)
        updates = []

        if self.momentum!=0.0:
            # MomentumSGD
            m = self.get_slot(var, 'm')
            m_t = m.assign( self.momentum*m + (1.0-self.momentum)*grad )
            if self.centered:
                denominator = 1.0-(self.momentum**local_step)
                m_t_hat = m_t / denominator
            else:
                m_t_hat = m_t
            updates.append(m_t)
        else:
            m_t_hat = grad

        # Update var
        var_update = var.assign(var - lr_t*m_t_hat)
        updates.append(var_update)
        return tf.group(*updates)

    def _resource_apply_sparse(self, grad, var, indices):
        raise NotImplementedError("not implemented")
```

# RMSpropを経由してAdamへ

MultiOptimizerV1でmomentum=0とするとVanillaSGDなのだが、その学習率を適応的に更新するのがRMSpropなので、コードは共通で機能追加できる。
以下、追加したものがMultiOptimizerV2。

```Python:MultiOptimizerV2
from tensorflow.python.keras.optimizer_v2.optimizer_v2 import OptimizerV2
from tensorflow.python.keras import backend_config

class MultiOptimizerV2(OptimizerV2):
    def __init__(self, learning_rate=0.01, momentum=0.0, centered=True, rho=0.0, centered_v=True, epsilon=1e-7, name='MultiOptimizerV2', **kwargs):
        super().__init__(name, **kwargs)
        self._set_hyper('learning_rate', kwargs.get('lr', learning_rate))
        self._set_hyper('decay', self._initial_decay)
        self.momentum = momentum
        self.centered = centered if momentum!=0.0 else False
        self.rho = rho
        self.centered_v = centered_v if rho!=0.0 else False
        self.epsilon = epsilon or backend_config.epsilon()

    def get_config(self):
        config = super(MultiOptimizerV2, self).get_config()
        config.update({
            'learning_rate': self._serialize_hyperparameter('learning_rate'),
            'decay': self._serialize_hyperparameter('decay'),
            'momentum': self.momentum,
            'centered': self.centered,
            'rho': self.rho,
            'centered_v': self.centered_v,
            'epsilon': self.epsilon,
        })
        return config
    def _create_slots(self, var_list):
        for var in var_list:
            if self.momentum!=0.0:
                self.add_slot(var, 'm')
            if self.rho!=0.0:
                self.add_slot(var, 'v')

    def _resource_apply_dense(self, grad, var, apply_state=None):
        var_device, var_dtype = var.device, var.dtype.base_dtype
        lr_t = self._decayed_lr(var_dtype)
        local_step = tf.cast(self.iterations+1, var_dtype)

        updates = []
        if self.momentum!=0.0:
            # MomentumSGD
            m = self.get_slot(var, 'm')
            m_t = m.assign( self.momentum*m + (1.0-self.momentum)*grad )
            if self.centered:
                denominator = 1.0-(self.momentum**local_step)
                m_t_hat = m_t / denominator
            else:
                m_t_hat = m_t
            updates.append(m_t)
        else:
            m_t_hat = grad

        if self.rho!=0.0:
            #RMSprop
            v = self.get_slot(var, 'v')
            v_t = v.assign( self.rho*v + (1.0-self.rho)*(grad**2) )
            if self.centered_v:
                denominator = 1.0-(self.rho**local_step)
                lr_t = lr_t / (((v_t/denominator)**0.5)+self.epsilon)
            else:
                lr_t = lr_t / ((v_t**0.5)+self.epsilon)
            updates.append(v_t)

        # Update var
        var_update = var.assign(var - lr_t*m_t_hat)
        updates.append(var_update)
        return tf.group(*updates)

    def _resource_apply_sparse(self, grad, var, indices):
        raise NotImplementedError("not implemented")
```

パラメータにRMSProp同様にrhoとepslionが追加された。追加でcentered_vというのも追加されているが、これはバイアス補正をするフラグ。
新たに内部の計算で使用するためにadd_slotで'v'を追加している。

これでRMSpropが実装されたわけだが、ここでmomentumを設定するとAdamになる。beta_1がmomentum、beta_2がrhoにあたる。
アルゴリズムとパラメータは下記のような関係になっている。

|Algorithm|momentum|centered|rho|centered_v|
|---|---|---|---|---|
|VanillaSGD | =0.0 | False| =0.0 | False |
|MomenutumSGD | >0.0 | False| =0.0 | False |
|RMSprop | =0.0| False| >0.0 | False |
|Adam | >0.0 | True| >0.0 | True |

TensorFlow/KerasのSGDの実装ではmomentumを設定すると他のアルゴリズムと学習率の対応が一致しなくなるが、この実装ではそういうことは起きない。

# そしてAdaBelief

詳細はこちらの[記事](https://qiita.com/omiita/items/5012afa3cba4d73a7aed)で。ここでは[論文](https://arxiv.org/abs/2010.07468)に記載されている数式をもとに実装。

```Python:MultiOptimizerV3
from tensorflow.python.keras.optimizer_v2.optimizer_v2 import OptimizerV2
from tensorflow.python.keras import backend_config
class MultiOptimizerV3(OptimizerV2):
    def __init__(self, learning_rate=0.01, momentum=0.0, centered=True, rho=0.0, centered_v=True, epsilon=1e-7, belief=False, name='MultiOptimizerV3', **kwargs):
        super().__init__(name, **kwargs)
        self._set_hyper('learning_rate', kwargs.get('lr', learning_rate))
        self._set_hyper('decay', self._initial_decay)
        self.momentum = momentum
        self.centered = centered if momentum!=0.0 else False
        self.rho = rho
        self.centered_v = centered_v if rho!=0.0 else False
        self.epsilon = epsilon or backend_config.epsilon()
        self.belief = belief if rho!=0.0 and momentum!=0.0 else False

    def get_config(self):
        config = super(MultiOptimizerV3, self).get_config()
        config.update({
            'learning_rate': self._serialize_hyperparameter('learning_rate'),
            'decay': self._serialize_hyperparameter('decay'),
            'momentum': self.momentum,
            'centered': self.centered,
            'rho': self.rho,
            'centered_v': self.centered_v,
            'epsilon': self.epsilon,
            'belief': self.belief,
        })
        return config
    def _create_slots(self, var_list):
        for var in var_list:
            if self.momentum!=0.0:
                self.add_slot(var, 'm')
            if self.rho!=0.0:
                self.add_slot(var, 'v')

    def _resource_apply_dense(self, grad, var, apply_state=None):
        var_device, var_dtype = var.device, var.dtype.base_dtype
        lr_t = self._decayed_lr(var_dtype)
        local_step = tf.cast(self.iterations+1, var_dtype)

        updates = []
        
        if self.momentum!=0.0:
            # MomentumSGD
            m = self.get_slot(var, 'm')
            m_t = m.assign( self.momentum*m + (1.0-self.momentum)*grad )
            if self.centered:
                denominator = 1.0-(self.momentum**local_step)
                m_t_hat = m_t / denominator
            else:
                m_t_hat = m_t
            updates.append(m_t)
        else:
            m_t_hat = grad

        
        if self.rho!=0.0:
            #RMSprop
            if self.belief: # AdaBelief
                g2 = (grad-m_t)**2
                epsilon_b = self.epsilon
            else:
                g2 = grad**2
                epsilon_b = 0.0

            v = self.get_slot(var, 'v')
            v_t = v.assign( self.rho*v + (1.0-self.rho)*g2 )

            if self.centered_v:
                denominator = 1.0-(self.rho**local_step)
                lr_t = lr_t / ((((v_t+epsilon_b)/denominator)**0.5)+self.epsilon)
            else:
                lr_t = lr_t / ((v_t**0.5)+self.epsilon)
            updates.append(v_t)

        # Update var
        var_update = var.assign(var - lr_t*m_t_hat)
        updates.append(var_update)
        return tf.group(*updates)

    def _resource_apply_sparse(self, grad, var, indices):
        raise NotImplementedError("not implemented")

```
belief=TrueとすればAdaBeliefになる。momentumとrhoの設定は必須。
実装からわかる通り、メモリ消費はAdamと同じ。計算量もほぼ同じと考えてよいだろう。

# 実験

## 特性確認

Optimizerを直接実行して、違いをみてみる。具体的な内容は[過去記事](https://qiita.com/T-STAR/items/c77f5c318d3f3141aad4)と同じなので、そちらを参照のこと。

![adabelief_curb.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/824193/bf1829de-ecfb-d1cc-4697-6e8a69768875.png)

![Bumpy1_MultiOpt (learning_rate=0.02, decay=0.0).gif](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/824193/cbedfb73-22aa-0cfd-58ac-3a9d281ca014.gif)

赤がAdaBeliefになるが、明らかに他のアルゴリズムとは違う動きになる。最適値に到達するまで最も早い。到達後の振動はAdamより若干大きいが、ほぼ同じくらいのようだ。

## 画像認識

CIFAR10の画像認識率をAdamとAdaBeliefで比較する。

* 条件
 * VGGライクの21層のCNN
 * Horizontal Flip,Shift,Cutoutでデータ拡張
 * バッチサイズ512、全200エポックで150エポック以降は学習率を1/10に
 * learning_rate=0.001,momentum=0.9,rho=0.999で実施
 * AdamはTensorFlowの実装を使いパラメータは同じ。
 * Google Colab のTPUで実施

各３回合計６回実施してValidationでのAccuracyをまとめた結果がこちら。

|Algorithm|Worst|Median|Best|
|---|---|---|---|
|Adam|94.80|94.88|95.25|
|AdaBelief|95.45|95.63|95.64|

明白にAdaBeliefの結果がAdamを上回った。ネットワークによって性能向上の増減はあるはずなので、興味のある方は自分で確認してもらいたい。

# まとめ

軽い気持ちでAdaBeliefを実装してみたら本当に性能が上がっていて驚いた。
WeightDecayやWarmupなどと組み合わせたら、さらに性能向上すると思われる。
Adamを超えるという宣伝文句のアルゴリズムはいくつもあるが、ハイパーパラメータが増えたりして結局使いづらいものが多い。こちらはそのような欠点がないので、Adamにかわる定番になる可能性も十分あるように思われる。

#参考

[出きたてホヤホヤ！最新オプティマイザー「AdaBelief」を解説！](https://qiita.com/omiita/items/5012afa3cba4d73a7aed)

[今度こそAdamを超えるのか⁈NeurIPS 2020で注目された「AdaBelief」とは](https://qiita.com/z-lai/items/b4d02f1050e99c59ad37)

#関連記事

[TensorFlow/Kerasでカスタム最適化アルゴリズムを実装する](https://qiita.com/T-STAR/items/6a489ee8f3b6b213bede)
[最適化アルゴリズムを単独実行で比較する(総合編)](https://qiita.com/T-STAR/items/c77f5c318d3f3141aad4)
[TensorFlow+KerasでCutoutを実装/評価する](https://qiita.com/T-STAR/items/0aa81d95ceaa1cdc634d)