tf.kerasでMLP-Mixerの実装

#はじめに

Vision Transformer(ViT)の発表以降、ConvolutionではなくAttentionとMLPを主要な仕組みとするTransformer系の画像認識モデルの研究が活発になっているようだが、最近MLP-MixerというAttentionを使わないモデルが発表されたので、tf.kerasで実装してみた。

#MLP-Mixerとは


発表論文はこちら。

https://arxiv.org/abs/2105.01601

本記事では主に実装面を紹介するものであり、理論的なところは下記の記事に詳しいので詳細はそちらを参照のこと。

https://qiita.com/kanataken/items/0b48b7af0187b7538385

「SOTAを達成した」というわけではなく、ViTよりももっと単純な構成でも結構学習できてしまうことが新発見ということのようだ。

モデルの論文としては珍しく論文にコードが掲載されている。以下、論文内のコードを転載。(公式の実装は[こちら](https://github.com/google-research/vision_transformer/blob/linen/vit_jax/models_mixer.py))

```python

import einops
import flax.linen as nn
import jax.numpy as jnp

class MlpBlock(nn.Module):
  mlp_dim: int
  @nn.compact
  def __call__(self, x):
    y = nn.Dense(self.mlp_dim)(x)
    y = nn.gelu(y)
    return nn.Dense(x.shape[-1])(y)

class MixerBlock(nn.Module):
  tokens_mlp_dim: int
  channels_mlp_dim: int
  @nn.compact
  def __call__(self, x):
    y = nn.LayerNorm()(x)
    y = jnp.swapaxes(y, 1, 2)
    y = MlpBlock(self.tokens_mlp_dim, name='token_mixing')(y)
    y = jnp.swapaxes(y, 1, 2)
    x = x + y
    y = nn.LayerNorm()(x)
    return x + MlpBlock(self.channels_mlp_dim, name='channel_mixing')(y)

class MlpMixer(nn.Module):
  num_classes: int
  num_blocks: int
  patche_size: int
  hidden_dim: int
  tokens_mlp_dim: int
  channels_mlp_dim: int
  @nn.compact
  def __call__(self, inputs, x):
    s = self.patch_size
    x = nn.Conv(self.hidden_dim, (s,s),
                strides=(s,s), name='stem')(inputs)
    x = einops.rearrange(x, 'n h w c -> n (h w) c')
    for _ in range(self.num_blocks):
      x = MixerBlock(self.tokens_mlp_dim, self.channels_mlp_dim)(x)
    x = nn.LayerNorm(name='pre_head_layer_norm')(x)
    x = jnp.mean(x, axis=1)
    return nn.Dense(self.num_classes, name='head',
                    kernel_init=nn.initializers.zeros)(x)
```

jaxで書かれているが、基本的な構成はこのコードだけで完結しており、ViTに比べれば驚くほど単純な構成。
ちなみに筆者によるViT実装の記事はこちら。

https://qiita.com/T-STAR/items/728b5c279bcfb6a6381f

MLP-MixerではViTに存在した"Multi Head Self Attention"と"Position Embedding"の処理が存在しない。
CNNでは途中でstrideを2にしたりしてデータの空間的次元を縮小していくが、ViT同様にそのような処理も一切ない。
処理冒頭でnn.Convがあるが、これはパッチ分割のためで、space_to_depthと全結合の組み合わせでも書けるところ。

複数ブロックを"Token Mixing"と"Channel Mixing"の順次処理で構成するが、"Token Mixing"の部分が"Multi Head Self Attention"に当たるようだ。"Channel Mixing"の部分は新しい名称をつけただけで、ViTと同じ処理になっているようだ。

また、論文中に以下のような記述がある。

> In the extreme case, our architecture can be seen as a very special CNN, which uses 1×1 convolutions for channel mixing, and single-channel depth-wise convolutions of a full receptive field and parameter sharing for token mixing. 

"Token Mixing"がDepth-Wise Convolution、"Channel Mixing"がPoint-Wise Convolution、に其々対応するような特殊なCNNにも見える、ということのようだ。

出力はDenseの直前にmean(Global Average Pooling)が入っているが、これはViTの標準実装では存在しない。(ViTの公式実装ではオプションとして指定できるようになっている)

#実装

tf.kerasでの実装を示す。tensorflowは2.5.0。

```python
import tensorflow as tf
from tensorflow.keras import layers

class SkipConnection(tf.keras.layers.Layer):
    def __init__(self, drop_rate=0.0, **kwargs):
        super().__init__(**kwargs)
        self.drop_rate = drop_rate

    def get_config(self):
        config = {'drop_rate': self.drop_rate}
        base_config = super().get_config()
        return dict(list(base_config.items()) + list(config.items()))

    def call(self, inputs):
        shortcut, x = inputs
        if self.drop_rate != 0.0:
            x = layers.Dropout(rate=self.drop_rate, noise_shape=(None,1,1))(x)
        x = layers.Add()([shortcut,x])
        return x

def mlp_mixer(
    image_size = 224,
    patch_size = 16,
    channels = 3,
    num_classes = 1000,
    num_blocks = 12,
    hidden_dim = 768,
    tokens_mlp_dim = 384,
    channels_mlp_dim = 3072,
    l2_reg=0.0,
    stoch_depth = 0.1,
    dropout = 0.0):

    dense_kwargs = {
        'kernel_initializer':'glorot_normal',
        'bias_initializer': tf.keras.initializers.RandomNormal(stddev=1e-2),
        'kernel_regularizer':tf.keras.regularizers.l2(l2_reg)
        }
    lnorm_kwargs = {'epsilon':1e-6}

    def MLP_Block(x, mlp_dim, dropout, prefix):
        dim = x.shape[-1]
        x = layers.Dense(mlp_dim, name=prefix+'_fc1', **dense_kwargs)(x)
        x = layers.Activation(tf.nn.gelu, name=prefix+'_activation')(x)
        if dropout!=0.0:
            x = layers.Dropout(dropout, name=prefix+'_do')(x)
        x = layers.Dense(dim, name=prefix+'_fc2', **dense_kwargs)(x)
        return x

    def Mixer_Block(x, tokens_mlp_dim, channels_mlp_dim, stoch_depth, dropout, prefix):
        y = layers.LayerNormalization(name = prefix+"_norm")(x)
        y = layers.Permute((2, 1), name = prefix+"_swap_axes")(y)
        y = MLP_Block(y, tokens_mlp_dim, dropout, prefix=prefix+'_Token_Mixing')
        y = layers.Permute((2, 1), name = prefix+"_swap_axes_restore")(y)
        x = SkipConnection( stoch_depth, name=prefix+'_skipconnect1')((x,y))

        y = layers.LayerNormalization(name = prefix+"_norm2")(x)
        y = MLP_Block(y, channels_mlp_dim, dropout,prefix=prefix+'_Channel_Mixing')
        x = SkipConnection( stoch_depth, name=prefix+'_skipconnect2')((x,y))
        return x


    x = inputs = layers.Input(shape=(image_size, image_size, channels), name="Input")
    
    num_patches = (image_size // patch_size) ** 2
    patch_dim = channels * patch_size ** 2
    x = layers.Lambda( lambda y: tf.nn.space_to_depth(y, patch_size), name='Stem_space2depth' )(x)
    x = layers.Dense(hidden_dim, name='Stem_fc', **dense_kwargs)(x)
    x = layers.Reshape((num_patches,hidden_dim), name='Stem_reshape')(x)

    for i in range(num_blocks):
        block_stoch_depth =  (i+1)*stoch_depth/num_blocks
        x = Mixer_Block(x, tokens_mlp_dim, channels_mlp_dim, block_stoch_depth, dropout,prefix=f'Mixer_{i+1}')

    x = layers.LayerNormalization(name="Head_norm", **lnorm_kwargs)(x)
    x = layers.GlobalAveragePooling1D(name='Head_averaging')(x)
    x = layers.Dense(num_classes, kernel_initializer=tf.keras.initializers.Zeros() , name="Head_fc")(x)
    return tf.keras.Model(inputs,x) 
```

基本的に論文内に載っているコードに沿っている。ただし、DropoutとStochasticDepthを使用した旨が論文に書かれているのだが、論文内掲載コードにはその辺が実装されていないようなので、そこは勝手に補足した。Weight Decayに関してはl2_regularizerとして実装した。

#実験

CIFAR-10で実験してみる。ViT系は巨大データセットでの学習が前提のため、小規模のデータセットでの実験にどの程度意義があるかはわからないが、参考程度に。

##設定
学習は以下のような設定で行なった。

|Epoch数 | バッチサイズ |  データ拡張 | 入力画像サイズ| 
|:-:|:-:|:-:|:-:|:-:|:-:|:-:|:-:|:-:|
| 600  | 256  | Mixup |32x32| 

ViT系は学習が遅くエポック数が短いと実力が出ないので、600エポックとかなり長くして実験。
その他オプティマイザ等の設定は以下の通り。

| オプティマイザ | 学習率 | 学習率減衰 |Warm Up | Global <br>Clip Norm | Label <br>Smoothing|
|:-:|:-:|:-:|:-:|:-:|:-:|
| SGD<br>(momentum=0.9)  | 1e-1→1e-3 | cosine | 2000 steps|1.0 | 0.1 |

論文ではAdamを使っているがここではSGDを使った。

比較のため、同じようなパラメータ数でViTモデルを作成して実験した。
ViT-Tiと呼ばれるモデルがあるようなので、それにほぼ同等のパラメータ数になるようにMLP-Mixerの方を調整した。
以下実験でのパラメータ。

| Model | Patch Size | Hidden Dim. | Layers | Heads | Tokens MLP Dim. | Channels MLP Dim. | Stochastic Depth |  Weight Decay |
|:-:|:-:|:-:|:-:|:-:|:-:|:-:|:-:|:-:|
| Mixer-240| 4 | 240  | 12  | - | 120 | 960 | 0.1 | 1e-4 | 
| ViT-Ti | 4 | 192  | 12  | 3 | - | 768 | 0.1 | 1e-4 |

本来は224x224のデータを16x16の領域で分割して14x14データとして処理が進むのだが、ここでは入力画像サイズが32x32なのでパッチ分割のサイズは4x4にし、パッチ分割後は8x8で処理される。
Mixer-240というモデル名は勝手に命名。
ViTは公開済みの自作コードをStochastic Depth対応するように改造して作成。通常のDropoutは無し。

## 実験結果
実験の結果は以下のようになった。差は明らかなので実験は一回のみの結果を掲載。TPUで実施。
(FLOPsは独自計算なので論文等とは多少は差があるはずだが、大きな違いはないはず)

| Model|  パラメータ数 | FLOPs |  認識率 | Time/Step |
|:-:|:-:|:-:|:-:|:-:|
| Mixer-240/4 | 5.8M | 398.1M | 90.11% | 35ms | 
| ViT-Ti/4 | 5.4M | 351.5M | 91.81%  | 45ms  | 

![cifar10_mixer_vs_vit.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/824193/b65ac9d6-0621-b73d-aba3-1c10383d34c6.png)


以下所見。

* ViT-Ti/4の方が最終的な認識率で上回った。
    * 学習曲線を見ると、エポック数をさらに増やせばViTはもう少し成績が伸びそうな気がするが、Mixerはこの設定ではあまり伸びないと思われる。
* 学習初期(200エポック程度まで)はMixerの方が認識率が高い。
    * 十分に学習ステップ数を取らないとViTよりも高性能と解釈してしまう可能性が出てきそう。
* "Time/Step"(一ステップにかかる時間)で比較すると、Mixerの方が短い。
    * FLOPsはほぼ同じはずなので、構造の簡単さによりスループットが上がっていると思われる。
* パラメータ数が同程度のEfficentNet-B0なら200エポック程度でも95%以上は行くはずなので、どちらも特に好成績というわけでは無い。

以下、実験用Google Colabノートへのリンク。

https://colab.research.google.com/gist/T-STAR-LTD/befe2d5687be3d723bb522e65864a719/mlp-mixer-cifar10.ipynb?authuser=1

#まとめ

MLP-Mixerの実装と評価を行なった。
CIFAR10の実験だけで判断するのは早計だが、絶対的な認識性能としてはViTに及ばずとも、処理が軽いという利点は確認できたように思う。
実装コードもシンプルで面白い。現状では実用性はあまり無いと思われるが、これを土台とした改良モデルが色々出るはずなので、今後に期待。

ちなみに、以前ViT-B/16でCIFAR10を学習させた際には90%以上にできなかったのに対し、パラメータ数を大幅に減らしたViT-Tiでは結構あっさり90%以上に到達したことが、個人的には新発見。