[리뷰] MLP-Mixer: An all-MLP Architecture for Vision

11 min readMay 16, 2021

튜닝의 끝은 순정이라고 했던가…??

지난해 11월 Google Research, Brain 팀은 Vision Transformer을 발표하고 Image Classification에서 SOTA를 달성하며 Transformer가 Vision에 유효하다라는 것을 증명하였다. Google Research, Brain 팀은 2021년 5월 CNN도 아니고 self-attention (Transformer)도 아닌 DNN의 가장 기본적인 알고리즘인 MLP을 이용하여 Vision에 도전한 “MLP-Mixier” 논문을 발표하였다.

논문의 저자를 살펴보면 Ilya Tolstikhin*를 제외하고 Vision Transformer의 저자인 Neil Houlsby∗ , Alexander Kolesnikov∗ , Lucas Beyer∗ , Xiaohua Zhai, Thomas Unterthiner, Alexey Dosovitskiy 등이 모두 공저자로 포함되었다. (*은 동일한 contribution)

MLP-Mixer의 아키텍처

MLP-Mixer는 1) per-patch linear embeddings 2) Mixer layers 3) Classifier head 등 크게 3가지 레이어로 구성되어 있다.

Per-patch linear embeddings

가장 먼저 MLP-Mixer의 입력 이미지는 patches x channels 형태로 “linearly projected image patches”의 시퀀스이다. 이미지를 분할하여 patch로 만드는 방법은 이전 ViT의 그것과 동일하다. 하지만 ViT와 달리 positional encoding이 필요하지 않고 각 patch는 hidden dimension(C)로 변환시킨다.

Mixer layers

현재 딥러닝 비전 구조는 두 가지 feature를 혼합하여 사용하고 있다.

(i) at a given spatial location (e.g., 특정 pixel에 대한 연산)
(ii) between different spatial locations (e.g., 주변 pixel 간 연산)

예를들어 CNN의 경우, 1x1 Conv 동작은 (i), NxN Conv(N>1)과 pooling 동작은 (ii)에 해당되며 ViT를 비롯한 self-attention기반 알고리즘은 MLP-blocks 동작은 (i), self-attention 레이어는 (i)와 (ii)를 동시에 수행한다.

MLP-Mixer 논문의 특징은 두가지 features, 즉 per-location (channel-mixing) 동작과 cross-location (token-mixing) 동작을 분리하고 이 두 동작을 MLPs로 구현하고 있다는 점이다.

Mixer는 두가지 형태의 MLP 레이어(Channel-Mixing MLPs & Token-Mixing MLPs)를 포함한다. (아래 그림은 Mixer 레이어의 구조 및 연산식을 설명한다.)

Channel-Mixing MLPs: 이 MLP은 각 token을 독립적으로 처리하며, 입력으로서 임베딩 테이블의 각 행(row)을 가져옴
Token-Mixing MLPs: 각 채널을 독립적으로 처리하며, 입력으로서 임베딩 테이블의 각 열(column)을 가져옴

CNN과의 유사성

Mixer는 “An all-MLP architecture for vision”이라는 부제를 붙였지만, MLP는 현재 CNN 알고리즘으로 구현이 가능하다.

Per-patch fully-connected layer: 입력 이미지로부터 patch를 만들 때 P x P conv (with Stride P, No padding)을 수행하는 것과 동일함
Channel-Mixing MLPs: 1x1 conv
Token-Mixing MLPs: weight-shared depth-wise convolution

MLP-Mixer의 실험 조건

MLP-Mixer 논문은 SOTA를 달성하기 보다는 단순한 MLP기반 모델이 오늘날 최고의 conv기반 모델과 어텐션기반 모델과 비교하였을 때 경쟁력이 있다는 것을 보이는 것이다. 이를 위해 1) downstream task의 정확도, 2) pre-training의 전체 계산 비용 3) inference 시 throughput을 여타 Image Classification 알고리즘과 비교하고 있다. 실험 조건은 다음과 같다.

Downstream tasks

ILSVRC2012 “ImageNet” (1.3M examples, 1k classes, Original validation labels(“ImNet”)과 cleaned-up “ReaL” labels)
CIFAR-10/100 (50k examples, 10/100 classes)
Oxford-IIIT Pets (3.7k examples, 36 classes)
Oxford Flowers-102 (2k examples, 102 classes)
Visual Task Adaptation Benchmark (1k examples, 119 diverse datasets)

Pre-training data

ILSVRC2021 ImageNet
ImageNet-21k (14M examples, 21k classes)
JFT-300M (300M examples, 18k classes)

Metrics

모델의 계산 비용과 품질간 trade-off를 평가하기 위해 2가지 metric 사용
(1) TPU-v3 가속기 상 전체 pre-training 시간
(2) TPU-v3 상에서 images/sec/core throughput

Models

HaloNets: Resnet-like 구조 (3x3 Conv 대신 self-attention 레이어 사용)
Big Transfer(BiT): ResNets (pre-trained on ImageNet-21k or JFT-300M)
NFNets: Normalizer-free ResNets (w/ several optimizations for ImageNet)
MPL: EfficientNet 구조. (Pretrained on JFT-300M, Meta-pseudo labeling from ImageNet)
ALIGN: EfficientNet 구조. (Pretrain image encoder & language encoder on noisy web image text pairs)

Mixer 구조 (Layers/Patch resolution/Parameters)

Small: S/32(8/32x32/10M), S/16(8/16x16/10M)
Base: B/32(12/32x32/46M), B/16(12/16x16/46M)
Large: L/32(24/32x32/188), L/16(24/16x16/189)
Huge: H/14(32/14x14/409)

MLP-Mixer의 실험 결과

Transfer performance, inference throughput & training cost

Mixer-L/16을 ImageNet-21k로 pre-training한 경우, ViT-L에 비해 ImNet top-1 기준 1.15%가량 낮은 84.15%이지만 throughput(img/sec/core)는 105로 VIT-L에 비해 3배 이상을 달성하였다.

Upstream dataset을 JFT-300M으로 변경하여 Mixer-H/14을 pre-training한 경우, BiT-R보다 더 높고 ViT-H에 근접한 87.94%까지 정확도를 높이면서 ViT-H에 비해 2.5배 높은 throughput(img/sec/core)을 달성하였다.

Inference throughput(5th column)에 의해 정렬. MLP-Mixer (pink), Attention기반 (blue), CNN기반(yellow). Avg 5는 5개 downstream tasks의 평균 성능

Accuracy-Compute Trade-off

위 테이블을 accuracy와 compute에 걸리는 Total pre-training TPUv3 core-days로 그리면 아래와 같은 그래프를 얻을 수 있다. Mixer(pink)가 BiT와 같은 CNN 계열(yellow)보다 더 적은 compute 리소스를 요구하는지 좀더 확실히 알 수 있다.

The role of the model scale

Mixer 논문은 두가지 독립적인 방법으로 모델을 scale 하여 Mixer/ViT/BiT를 비교하였다.

(1) 모델 사이즈 증가(layer의 개수, hidden dimension, MLP width): pre-training compute와 test-time throughput에 동시에 영향
(2) Fine-tuning 시 입력 이미지 resolution을 증가: test-time throughput에만 영향을 줌

아래 그래프는 ImageNet validation top-1 accuracy를 total pre-training(왼쪽)과 throughput(오른쪽)을 비교하였다. (pre-training: JFT-300M & fine-tuning: 224x224 resolution ) 전체적으로 Mixer와 ViT가 BiT보다 더 우수한 정확도 성능을 보이며 Mixer가 ViT보다 전체적인 throughput이 더 우수한 것을 알 수 있다.

The role of the pre-training dataset Size

아래 그림은 JFT300M으로 pre-training 시 training size에 따른 정확도를 나타낸 것이다. JFT300M의 가장 작은 subset(3%)으로 pre-training 하였을 때 모든 Mixer 모델은 overfit 되었지만 JFT300M의 모든 dataset(100%)으로 pre-training하였을 때는 BiT는 물론 ViT보다 더 좋은 정확도를 보인다. 또한 JFT300M의 모든 dataset을 이용하는 경우, epoch을 좀더 늘렸을 때 ViT에 비해 Mixer의 정확도가 증가할 여지가 더 클 것으로 예상된다.

Training size 증가에따른 ImageNet Top-1 정확도의 증가 추세: JFT300의 subset — 3%(233 epoch), 10%(70 epoch), 30%(23 epoch), 100%(7 epochs)를 사용함

Insight

Mixer를 사용하였을 때 주요 장점은 아키텍처의 단순성과 계산 속도에 있다. ViT의 복잡도는 제곱에 비례하지만, Mixer의 channel은 패치 사이즈에 독립적이므로 일반적인 CNN 과 마찬가지로 Mixer의 복잡도는 이미지의 픽셀 수에 비례하는 장점이 있다. 또한 각 레이어는 항상 동일한 사이즈의 입력을 받고 각 레이어는 2개의 MLP 블록과 GELU 활성화 함수만 있는 간단한 구조이므로 가속기 구조를 간단히 하는데 유리할 것으로 생각된다.

Mixer는 ViT에 비해 정확도는 낮지만 throughput이 높아 실제 inference 시 더 적은 Compute 리소스를 요구하므로 실제 서비스에 적용하여도 충분히 경쟁력이 있을 것으로 생각된다.

하지만 MLP는 OI(Operational Intensity. flops/byte)가 낮은 연산으로 대표적인 Memory-bound 연산으로 분류되므로 Mixer 네트워크의 규모에 따라 더 높은 Memory 대역폭이 필요할 것으로 예상할 수 있다.

레퍼런스

[1] MLP-Mixer: An all-MLP Architecture for Vision: https://arxiv.org/pdf/2105.01601.pdf

[2] An Image is worth 16x16 words: Transformers for Image Recognition at Scale: https://arxiv.org/pdf/2010.11929.pdf