[WIP] "MetaFormer Baselines for Vision" を読む

## 目次

- [目次](#目次)
- [この記事の簡単な説明](#この記事の簡単な説明)
- [アブストラクト](#アブストラクト)
- [1. イントロダクション(追記予定)](#1-イントロダクション追記予定)
- [2. 手法](#2-手法)
  - [2.1. MetaFormer の概念の再確認](#21-metaformer-の概念の再確認)
  - [2.2. IdentitiyFormer と RandFormer](#22-identitiyformer-と-randformer)
  - [2.3. ConvFormer と CAFormer](#23-convformer-と-caformer)
  - [2.4. MetaFormer を改善するテクニック](#24-metaformer-を改善するテクニック)
    - [2.4.1. StarReLU(追記予定)](#241-starrelu追記予定)
    - [2.4.2. その他の修正(追記予定)](#242-その他の修正追記予定)
- [3. 実験結果(追記予定)](#3-実験結果追記予定)
- [4. 参考文献](#4-参考文献)

## この記事の簡単な説明

- [MetaFormer Baselines for Vision](https://arxiv.org/abs/2210.13452) の論文を読む。  
- MetaFormer [2] のトークンミキサー(<font color="Blue">[2] の論文では、トークンミキサーとして、アテンション、MLP、プーリングを採用して比較している</font>)をいろいろ変えても精度良好(ImageNet-1Kに対して、80% より大な精度)である。
- <font color="Blue">青文字</font>は私の考察やコメントを表す。
- <font color="Red"> この記事は製作途中です。順次追記していきます。また、今後更新するときはなるべく簡潔に読みやすくしていこうと思います。</font>

## アブストラクト

[MetaFormer [2]](https://arxiv.org/abs/2111.11418) のトークンミキサーを変えて、性能を比較する。論文の要約は、次の通りである。
1. <b>MetaFormer は、性能の下限を保証する。</b>  
   IdentityFormer(トークンミキサーに恒等写像を採用(<font color="Blue">一番単純なトークンミキサー</font>))だけでも ImageNet-1K に対して、80% より大な精度を達成する。
2. <b>MetaFormer は、任意のトークンミキサーでうまく機能する。</b>  
    RandFormer(トークンをランダムに混ぜる行列を使用)でも、81% より大な精度(IdentityFormer よりも精度良好)であるため、他のトークンミキサーを採用しても性能は保証される。
3. <b>MetaFormer は SOTA である。</b>
   5年前に遡るトークンミキサーだけでも、MetaFormerからインスタンス化されたモデルは、既に SOTA である。   
   - <b>ConvFormer は、ConvNeXt より高性能</b>  
        一般的な Depthwise Separable Convolution をトークンミキサーを採用する(純粋な CNN とみなすことができる) ConvFormer は、強力な CNN モデル ConvNeXt を上回る。
    - <b>CAFormer は、ImageNet-1K に対して新記録を達成</b>  
         Depthwise Separable Convolution を下段に、Vanilla Self-Attention を上段に適用するトークンミキサーを採用する CAFormer は、ImageNet-1K に対して新記録を達成した。：224×224 の解像度で、外部データや蒸留を用いない通常の教師あり学習で、85.5% の精度を達成した。

また、GELU よりも 71% FLOPs を削減する、新しいアクティベーション StarReLU を発見した。具体的には、StarReLU は、分配シフトの緩和に特化した、Squared ReLU の変形である。StarReLU は、他のニューラルネットワークと並ぶ MetaFormer のようなモデルで可能性を秘めている。  
コード：　https://github.com/sail-sg/metaformer

## 1. イントロダクション(追記予定)
<b>追記予定。</b>

## 2. 手法

### 2.1. MetaFormer の概念の再確認 


$I$ を入力として、次のように特徴量(トークンと呼ばれる)抽出をする。  
$$
X = InputEmbedding(I) \tag{1}
$$
トークン $X\in \mathbb{R}^{N\times C}$($N$: 長さ、$C$: チャンネル数)、MetaFormer ブロックに入力される。
$$
X' = X + \text{TokenMixer}(\text{Norm}_1(X)) \tag{2}
$$

$$
X = X' + \sigma(\text{Norm}_2(X')W_1)W_2 \tag{3}
$$
ただし、$Norm_1(・), Norm_2(・)$ は正規化、$TokenMixer(・)$ は、主にトークン間の情報伝達を行うミキサーを表す。また、$\sigma(・)$ は活性化関数、$W_1, W_2$ は MLP の学習可能なパラメータを表す。  
<font color="Blue">元の情報 X に、正規化 + トークミンキサー、正規化 + MLP の重み積算 + アクティベーション を足しているため、勾配消失を防ぐ効果があると考えられる。</font>  
トークンミキサーを具体的なモジュールとして指定することで、MetaFormer は特定のモデルにインスタンス化される。  

### 2.2. IdentitiyFormer と RandFormer

<figure>
  <img alt="image" src="https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3514428/9aefc0e0-6739-8423-e177-44843e0c2c55.png" width=100%>
  <figcaption>図 1. IdentityFormer, RandFormer, ConvFormer, CAFormer</figcaption>
</figure>

まずは、トークンミキサーの基本形を採用して、MetaFormer の能力を証明する。

- <b>IdentitiyFormer</b>

  次の恒等写像を考える。
  $$
  IdentityMapping(X) = X \tag{4}
  $$
  恒等写像は、トークンを混ぜず、実際、トークンミキサーとはみなされない。便宜上、トークンミキサーとみなして、他のトークンミキサーと比較していく。  
  図 1 (e) のように、IdentityFormer ブロックを作成し、図 1 (a) のように、ダウンサンプリング + IdentityFormer ブロックを4回繰り返す構造を IdentityFormer フレームワークとする。

- <b>RandFormer</b>  

  もう1つの基本のトークンミキサーは、global random mixiing である。
  $$
  RandomMixing(X) = XW_R \tag{5}
  $$
  ただし、$X\in \mathbb{R}^{N\times C}$($N$ を長さ、$C$ チャンネル数)、$W_R \in \mathbb{R}^{N\times N}$をランダムに初期化後固定された行列とする。固定のための余分なパラメータとトークン数の2乗の計算量(<font color="Blue">$N\times N$</font>)が必要なため、トークン数が大きい場合には適さない。コードは次の通りである。
  ```(python)
  # RandomMixing 
  import torch
  import torch.nn as nn
  # Identity mapping
  from torch.nn import Identity
  # Random mixing
  class RandomMixing(nn.Module):
      def __init__(self, num_tokens=196):
          super().__init__()
          self.random_matrix = nn.parameter.Parameter(
              data=torch.softmax(
                  torch.rand(num_tokens,num_tokens),
                  dim=-1),
              requires_grad=False)
      
      def forward(self, x):
          B, H, W, C = x.shape
          x = x.reshape(B, H * W, C)
          x = torch.einsum(
              'mn, bnc -> bmc',
              self.random_matrix, x)
          x = x.reshape(B, H, W, C)
          return x
  ```

  図 1 (f) のように、RandFormer ブロックを作成し、図 1 (b) のように、ダウンサンプリング + RandFormer ブロックを2回、ダウンサンプリング + IdentityFormer ブロックを2回繰り返す構造を RandFormer フレームワークとする(計算量を考慮して、RandFormer ブロックは2回までとしている)。

  MetaFormer [2] に出てきた PoolFormer と公正に比較するために、PoolFoermerV2 も作成する。 

- <b>モデル構成について</b>

  <figure>
    <figcaption>表 1. IdentityFormer, RandFormer, PoolFormerV2 のモデル構成</figcaption>
    <img alt="image" src="https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3514428/3c4c1f37-63ed-27ff-9ae2-b238ad7b2937.png" width=100%>
  </figure>

  - 表記等
    - C: チャンネル数
    - L: ブロック数
    - T: トークンミキサーの種類
      1. Id: Identity Mapping
      2. Rand: Random Mixing
      3. Pool: Pooling  
    - <font color="Blue">Size 「S*」について、「S」はチャンネル数 $C = (64, 128, 320, 512)$ を表し、「*」はブロック数 L の合計値を表す。</font>
    -  <font color="Blue">Size 「M*」について、「M」はチャンネル数 $C = (96, 192, 384, 768)$ を表し、「*」はブロック数 L の合計値を表す。</font>
    - <font color="Blue"> 各トークンミキサーブロックにおける、MLPのチャンネル数が、C の値に対応している。</font>
    - <font color="Blue"> 分類器(Classifier Head)では、Global Average Pooling + Norm + FC。</font>
    - <font color="Blue"> 正規化(Norm)は、レイヤー正規化。</font>

### 2.3. ConvFormer と CAFormer

トークンミキサーを変更して、SOTA となるモデルを探る。

- <b> ConvFormer </b>
  まず、トークンミキサーとして、 MobileNetV2 [3] に採用されている、Inverted Separable Convolution を採用する。

  $$
  Convolutions(X) = Conv_{pw2}(Conv_{dw}(\sigma (Conv_{pw1}(X)))) \tag{6}
  $$

  ただし、$Conv_{pw1}(・), Conv_{pw2}(・)$ は、Pointwise Convolution、$Conv_{dw}(・)$ は Depthwise Convolution、$\sigma (・)$ は非線形活性化関数を表す。  
  カーネルサイズは 7、拡張率は 2 (チャンネル数が 2 倍になる)とする。  
  トークンミキサーに ConvFormer ブロック(図 1 (g))を採用して Metaformer を ConvFormer としてインスタンス化する(図 1 (c)、ダウンサンプリングと ConvFormer の 4 回繰り返し)。

- <b> CAFormer </b>
  トークンミキサーとして、Transformer の Vanilla Self-Attention [4] を適用する Transformer ブロック(図 1 (h)) を採用する。ただし、RandFormer と同様に計算量を考慮して、最初の 2 回は ConvFormer ブロックを繰り返し、次の 2 回 Transformer ブロックを繰り返すことによって、CAFormer をインスタンス化する(図 1 (d))。

- <b>モデル構成について</b>

  <figure>
    <figcaption>表 2. ConvFormer, CAFormer のモデル構成</figcaption>
    <img alt="image" src="https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3514428/9df1d184-2738-7fad-1234-d64c5fc0a4e7.png" width=100%>
  </figure>

  - 表記等(表 1 と異なる点のみ記載)
    - T: トークンミキサーの種類
      1. Conv: Separable Convolution
      2. Attn: Vanilla Self-Attention
    - <font color="Blue"> 分類器(Classifier Head)では、Global Average Pooling + Norm + MLP。</font>

### 2.4. MetaFormer を改善するテクニック

#### 2.4.1. StarReLU(追記予定)
<b>追記予定。</b>

#### 2.4.2. その他の修正(追記予定)
<b>追記予定。</b>


## 3. 実験結果(追記予定)
<b>追記予定。</b>


## 4. 参考文献

[1] Yu, Weihao, Chenyang, Si, Pan, Zhou, Mi, Luo, Yichen, Zhou, Jiashi, Feng, Shuicheng, Yan, Xinchao, Wang. "MetaFormer Baselines for Vision". IEEE Transactions on Pattern Analysis and Machine Intelligence 46. 2(2024): 896–912. https://arxiv.org/abs/2210.13452

[2] Weihao Yu, , Mi Luo, Pan Zhou, Chenyang Si, Yichen Zhou, Xinchao Wang, Jiashi Feng, Shuicheng Yan. "MetaFormer Is Actually What You Need for Vision." (2022). https://arxiv.org/abs/2111.11418

[3] Mark Sandler, , Andrew Howard, Menglong Zhu, Andrey Zhmoginov, Liang-Chieh Chen. "MobileNetV2: Inverted Residuals and Linear Bottlenecks." (2019). https://arxiv.org/abs/1801.04381

[4] Ashish Vaswani, , Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, Illia Polosukhin. "Attention Is All You Need." (2023). https://arxiv.org/abs/1706.03762