MLP-Mixer: An all-MLP Architecture for Vision 논문 리뷰

https://arxiv.org/pdf/2105.01601.pdf

요약

ViT가 좋은 성능을 보이고 있는데 저자는 Convolution과 Attention이 둘 다 필요하지는 않는다고 주장한다. 저자는 MLP-Mixer라는 모델을 제안했는데 image patch와 across patch가 존재하여 각각 채널별, 깊이별 혼합 연산을 수행한다. 큰 데이터셋이나 modern regularization technique이 있는 경우 좋은 성능을 발휘한다.

---

MLP-Mixer의 구조

저자는 MLP-Mixer를 제안하는데 Convolution layer와 self-attention layer 모두 쓰이지 않았다. 오직 행렬곱으로 나타낼 수 있는 Fully connected layer와 비선형함수로만 이뤄져있다.

MLP-Mixer는 3가지 단계로 구성된다.

1. Patch
   원본 입력 이미지를 일정 크기만큼 잘라 Patch를 만들고 Fully connected layer로 연결해 token을 만든다. 이 행렬이 Mixer layer의 입력으로 들어간다.
2. Mixer layer
   Mixer layer는 token-mixing MLP, channel-mixing MLP순으로 수행된다. 각 MLP 연산은 Fully connected layer, GeLU, Fully connected layer로 구성되어있다. 이때 각 MLP 연산 앞에 layer normalization이 적용되고 그 앞에서부터 MLP 연산 뒤까지 skip connection이 연결된 구조이다.
3. Classifier
   Global Average Pooling 한 뒤 Fully connected layer로 연결한다.

좀 더 자세하게는 다음과 같다.

Patch의 크기 $P$와 이미지의 크기는 $(H,W)$는 미리 정의된다. Patch의 개수 $S$는 $HW/P^2$로 정해진다. Patch 하나 당 하나의 Token이 FCN으로 만들어지고 Token의 채널 수를 hidden dimension $C$라 하며 미리 정의된다. Mixer layer의 입력으로 들어가는 행렬 X의 크기는 Token 하나를 행으로 하여 세로로 연결시킨 $(S,C)$ 크기의 행렬이 된다.

Token-mixing MLP는 $X^T$의 각 행을 FCN으로 연결하여 $X^T$의 크기를 $(C, D_S)$로 바꾸고 GELU를 통과시킨 뒤 또 각 행을 FCN으로 연결하여 $(C, S)$로 바꾼다. Channel-mixing MLP는 $X$의 각 행을 FCN으로 연결하여 $(S, D_C)$로 바꿨다가 GELU를 통과시키고 $(S, C)$로 바꾼다. Token-mixing MLP에 쓰인 파라미터의 크기는 $(S, D_S)$, $(D_S, S)$이고 Channel-mixing MLP에 쓰인 파라미터의 크기는 $(C, D_C)$, $(D_C, C)$이다. 이 파라미터는 각 행마다 공유된다. $D_S$와 $D_C$는 MLP layer의 Width라고도 한다.

$(S,C)$크기의 출력을 $(S,1)$으로 각 행마다 평균내 줄인 뒤 FCN을 연결하여 분류작업의 클래스 수로 맞춰 softmax를 통과시킨다.

---

MLP-Mixer의 강점

Mixer layer에서 일어나는 MLP연산은 단순한 행렬곱이므로 Convolution 연산보다 훨씬 간단해서 특화된 최적화 구현을 통해 연산량을 줄일 수 있다.
단순하지만 좋은 성능을 내고, modern regularization technique가 들어가면 state-of-the-art모델 급으로 성능이 좋지만 고도화된 CNN 네트워크보다는 안좋다.

---

Mixer layer의 특징

Mixer layer의 token-mixing MLP와 channel-mixing MLP는 각 행, 열마다 다른 파라미터를 쓰지 않고 같은 파라미터로 학습시킨다.

1. channel-mixing MLP는 같은 파라미터로 학습을 해야 positional invariance(translation invariance)를 보장할 수 있다. 이는 CNN의 가장 두드러지는 특징인데, 대상의 위치가 달라져도 같은 클래스로 분류해야 하기 위해서 보장되어야 하는 특성이다. 따라서 파라미터를 공유해야 하는 것은 자연스러운 것이다.

2. token-mixing MLP에서는 그러나 파라미터를 공유하는 것이 일반적이지 않다. 기존의 Xception같은 경우에서 Seperable Convolution은 depthwise convolution을 할 때 다른 depth끼리 다른 커널으로 학습을 했다. 하지만 MLP-Mixer는 같은 파라미터로 학습하여 각 행과 열의 길이가 증가함에따라 너무 급격하게 파라미터 수가 커지는 것을 막았다. 이러한 결정이 성능에 영향을 크게 미치지 않은 것을 경험적으로 확인하였다.

---

isotorpic vs pyramidal

MLP-Mixer는 isotorpic이다. 가장 처음에 나오는 patch projection을 제외하면 이미지의 해상도가 일정하게 유지된다. 전통적인 기존의 이미지 네트워크는 pyramidal이다. 입력 이후로 채널 수가 많아지고 해상도가 작아지기 때문이다. 참고로 서로 구조를 반대로 조합한 네트워크, isotorpic resnet, pyramidal ViT도 존재한다.

---

실험

아래 3가지 값을 측정하였다.

1. downstream task에서의 정확도
2. pre-training에서의 연산량
3. test할 때의 처리량

downstream task는 다음과 같다.

• ILSVRC2012 (#1.3M, 1k classes)
• CIFAR10/100 (#50k, 10/100 classes)
• Oxford IIIT Pets (#3.7k, 36 classes)
• Oxford Flowers 102 (#2k, 102 classes)
• VTAB (Visual Task Adaptaion Benchmark, 19개의 다양한 데이터셋, 각 #1k)

pre-training 과정은 다음과 같다.

1. ILSVRC2012와 ILSVRC2021 두 데이터셋에 대하여 사전학습.
2. 큰 스케일의 데이터셋에서도 실험하기 위해 JFT-300M으로도 사전학습. (#300M, 18k classes)
3. Dosovitskiy et al., Kolesnikov et al.의 방식대로 데이터셋의 중복을 제거.

pre-training 세팅은 다음과 같다.

• 입력 데이터의 resolution은 224
• optimizer는 Adam, $\beta_1=0.9, \beta_2=0.999$
• learning rate는 linear warmup, linear decay
• batch size 4096
• weight decay, gradient clipping at global norm 1
• JFT-300M에 대하여 Szegedy et al.의 cropping technique 적용, random horizontal flipping도 추가로 적용
• ImageNet에 대하여 RandAugment, mixup, dropout, stochastic depth를 통해 data augmentation과 regularization 적용
• 이러한 테그닉은 timm library에서 영감을 받았다고 한다.
  

fine-tuning 과정은 다음과 같다.

• optimizer는 모멘텀 SGD
• learning rate는 linear warmup, cosine decay
• batch size 512
• weight decay는 사용하지 않음

pre-training에서 사용한 이미지의 크기는 224이다. Patch의 크기는 고정되어있어서 이미지의 크기가 바뀌면 Token의 개수가 바뀌기 때문에 모델 내부의 MLP layer의 weight의 크기가 모두 바뀌어야 한다. 입력 이미지의 크기에 따라서 pre-training을 하기엔 번거로우므로 224크기의 이미지로 pre-training을 한 모델을 통해 fine-tuning에서는 $224*K$ ($K$는 자연수) 크기의 고해상도 이미지로도 처리할 수 있도록 하기 위해 다음과 같은 방법이 쓰였다.

• 입력 이미지가 $K$배가 되면 Token의 수는 $S$개에서 $K^2 * S$개가 된다.
• Token mixing MLP에서의 weight 행렬의 크기도 $(D_S,S)$에서 (행렬 연산을 수행하기 위해) 세로 길이가 $K^2 * S$가 되어야 한다.
• 이때 pre-training의 파라미터를 활용하기 위해 block-diagonal matrix를 통해 초기화를 하였고 이를 위해서 weight 행렬의 크기는 $(K^2 * D_S,K^2 * S)$가 된다.
• 자세히는 $K^2$개의 기존 파라미터의 복사본이 main diagonal에 배치된 형태이다.

측정한 metric은 다음과 같다.

• TPU-v3로 돌린 총 pre-training 시간
• TPU-v3에서 코어 수와 시간(초)에 따른 이미지 처리량
• batch size에 따라 처리량이 다를 것이므로 batch size를 다양하게 실험해보고 가장 높은 케이스를 기록
• 정확도는 Top1 downstream accuracy로 측정하였다.

실험결과

Mixer의 성능을 비교하기 위해 ViT와 BiT도 함께 같은 조건에서 실험하였다. 결과는 다음과 같다.

• 비슷한 크기의 ViT나 BiT과 비교하여 성능은 비슷하지만 속도는 2배정도 빨랐다.
• pre-training 데이터셋이 커짐에 따라 정확도가 증가하는 추세는 ViT나 BiT보다 더 가파르다.
• 작은 pre-training 데이터셋에서 너무 Overfit이 잘되는 경향을 보였다. Regularization이 중요하다.

inductive bias가 ResNet과 차이가 있음을 보였다. 원본 이미지에서 $(16,16)$ 크기로 나누고 임의의 순서대로 모든 이미지에 대하여 동일하게 섞은 뒤 학습을 시킨 결과는 다음과 같다.

• Mixer는 섞지 않았을 때와 섞었을 때의 결과가 완전히 일치하였다.
• ResNet은 성능이 큰 폭으로 깎였다.
• Patch 단위가 아닌 pixel 단위로 섞었을 때는 둘 다 성능이 하락했지만 ResNet이 훨씬 더 많이 하락했다.

---

관련 연구

• ViT는 Convolutional based model보다 inductive bias(locality)가 적다. 따라서 작은 데이터에서는 학습하기가 힘들고 Overfitting이 되기 쉽다. 하지만 작은 데이터에서도 Regularization을 비중있게 가하여 학습이 잘 되도록 한 연구가 있다.
  Training data-efficient image transformers & distillation through attention.
• Mixer의 MLP block의 디자인은 아래 Transformer를 도입한 연구에서 따온 것이다.
  Attention is all you need.
• Mixer의 Patch를 나눠 Token으로 만들어 다루는 것은 아래 ViT를 도입한 연구에서 따온 것이다.
  An image is worth 16x16 words: Transformers for image recognition at scale.
• ResNet의 3x3 convolution을 self-attention layer로 바꾸거나 또다른 attention-like 매커니즘을 개발하는 연구도 활발히 진행 중이다.
• Image Classification의 최고 모델은 data augmentation이나 large scale pre-training등 다양한 방법으로 발전하고 있는 추세이다.