MLP-Mixer (2021) 논문 리뷰

MLPMixer의 논문 제목은 MLP-Mixer: An all-MLP Architecture for Vision다. (링크)

 

저자들은 Ilya Tolstikhin, Neil Houlsby, Alexander Kolesnikov, Lucas Beyer, Xiaohua Zhai, Thomas Unterthiner, Jessica Yung, Andreas Steiner, Daniel Keysers, Jakob Uszkoreit, Mario Lucic, Alexey Dosovitskiy다. 

 

Github: 링크

 

이름에서 알 수 있듯이 컨볼루션이나 트랜스포머의 어텐션 구조가 아니라 단순한 MLP 만으로 좋은 성능을 달성함을 보인 논문이다.

 

 

Abstract

CNN은 컴퓨터 비전 분야에서의 핵심 모델이다. 최근에는 Vision Transformer와 같은 어텐션 기반의 모델이 인기를 얻고 있다. 본 논문에서는 컨볼루션과 어텐션이 좋은 성능을 내는데에 충분하지만 반드시 필요한 것은 아님을 보인다. MLP-Mixer라는 multi-layer perceptrons (MLPs)에 기반한 구조를 제시한다. MLP-Mixer는 두 종류의 레이어로 구성된다. 하나는 Image patches에 독립적으로 적용되며 ("mixing" the per-location features) , 다른 하나는 accross patches 패치 전반에 걸쳐서 적용된다 ("mixing" the spatical information). 대규모 데이터셋에 학습하고 최신 regulariziation을 적용하면 MLP-Mixer는 image classification에서 SOTA와 비교해서도 경쟁력있는 점수를 달성한다. 

 

 

1. Introduction

기존에는 CNN이나 트랜스포머 구조의 SOTA 모델들이 많았다. 저자들은 MLP-Mixer (이하 줄여서 Mixer)라는 간단한 대체 구조를 제안한다. 이는 컨볼루션이나 셀프 어텐션을 사용하지 않는다. 대신, MLP 구조를 반복해서 spatial locations와 features channels 모두에 대해서 적용한다. Mixer는 오직 matrix multiplication routines에 의존하며, daya layout (reshapes and transpositions)와 scalar nonlinearities에 따라서 변한다. 

 

 

 

Figure 1은 Mixer의 macro-structure 거시적 구조를 묘사한다. Mixer는 patches x channels의 테이블 형식으로 표시되는 linearly projected image patches (or referred to as tokens)의 시퀀스를 입력으로 받는다. 그리고 이 차원을 계속 유지한다. Mixer는 2가지 종류의 MLP 레이어를 사용하는데 바로 channel-mixing MLPs와 token-miixng MLPs다. Channel-mixing MLPs는 서로 다른 채널들 간의 소통을 담당한다. 토큰 별로 독립적으로 작동하며 테이블의 개별적인 rows 행들을 입력으로 받는다. 

 Token-mixing MLPs는 서로 다른 spatial locations (tokens)의 소통을 담당하는데, 채널 별로 독립적으로 동작하며 테이블의 개별적인 열 columns를 입력으로 받는다. 위 두 가지 레이어들은 각각의 입력의 차원에 끼워넣어져서 각각의 상호작용을 가능하게 한다. 

 

극단적인 경우, Mixer는 CNN의 스페셜 케이스로 볼 수 있다. 1x1 컨볼루션을 channel mixing으로, single-channel depth-wise convolution이면서 full receptive field와 parameter sharing인 것을 token mixing으로 볼 수 있다. 하지만 역으로 전형적인 CNN은 Mixer의 스페셜 케이스라고 할 수 없다. 추가적으로 컨볼루션은 MLPs의 plain matrix multiplication 보다 복잡하기 때문에 추가적인 matrix multipliation의 계산적 감소나 특별한 구현이 필요하다.

 

Mixer는 단순한 구조임에도 불구하고 ImageNet ILSVRC2012 같은 대규모 데이터셋에 대해서 학습하면, 기존의 CNNs와 Transformers에서 제시된 accuracy / cost trade-off에 따라서 SOTA 성능에 근접한다. ImageNet ILSVRC2012에 대해서 Top-1 validation accuracy가 87.94%로, ViT에 비슷하지만 CNN에 비해서 약간 모자란 성능을 보인다. 

 

2. Mixer Architecture

현재 딥 비전 아키텍처는 다음의 2가지 피쳐를 (i) 주어진 spatial location에, (ii) 서로 다른 spatial locations 사이의, 혹은 양쪽을 한번에 혼합한다. CNN에서는 (ii)를 $N \times N$ 컨볼루션과 풀링으로 구현한다. 더 깊은 레이어는 더 큰 receptive field를 가지는데, 1 x 1 컨볼루션은 (i)을 수행하면서 더 큰 커널은 (i)과 (ii)를 모두 수행한다. ViT나 다른 어텐션 기반의 아키텍처의 경우 셀프 어텐션은 (i)과 (ii)를 수행하고, MLP 블록들은 (i)를 수행한다. Mixer 구조의 아이디어는 per-location (channel-mixing) 연산인 (i)와 cross-location (token-mixing) 연산 (ii)를 선명하게 분리하고자 한다. 두 연산 모두 MLPs로 구현가능하다. Figure 1에 아키텍처가 요약하고 있다.

 

Mixer는 $S$개의 non-overlapping image patches 이미지 패치의 시퀀스를 입력으로 받는다. 각각은 원하는 hidden dimension 인 $C$로 project 된다. 이를 통해서 Two-dimensional real-valued input table $X \in \mathbb{R}^{S \times C}$를 얻는다. 이때 원본 이미지의 해상도는 $(H, W)$, 개별 패치의 해상도는 $(P, P)$, 패치의 숫자 $S = HW / P^2$다. 동일한 projection matrix에 의해서 모든 패치들이 linearly projected 된다. Mixer는 동일한 크기의 복수의 레이어로 구성되며, 개별 레이어를 2개의 MLP 블록들로 구성된다. 첫 번째 블록은 token-mixing MLP로 $X$의 컬럼에 대해서 동작한다. (Transposed input table인 $X^\top$에 동작한다. ) Token-mixing MLP는 $\mathbb{R}^S \rightarrow \mathbb{R}^S$로 매핑하며 모든 컬럼들에 대해서 shared다. 두 번째는 channel-mixing MLP로 $X$의 로우에 대해서 동작하며, $\mathbb{R}^C \rightarrow \mathbb{R}^C$로 매핑한다. 이 역시 모든 로우에 대해서 shared다. 둘 다 fully-connected layers와 nonlinearity로 구성되며 입력 데이터 텐서의 개별 row에 대해 독립적으로 적용된다.

 

Mixer 레이어들은 다음처럼 쓸 수 있다 (레이어 인덱스는 생략한다). 

 

$U_{*, i} = X_{*, i} + W_{2} \sigma ( W_1 \text{LayerNorm} (X)_{*, i} ), $ for $i = 1, ..., C$  

 

$Y_{*, i} = U_{*, i} + W_{4} \sigma ( W_3 \text{LayerNorm} (U)_{*, i} ), $ for $i = 1, ..., S$  

 

$\sigma$는 element-wise nonlinearity로 GELU다. $D_S$와 $D_C$는 tunable hidden widths로 각각 token-mixing MLPs에, 그리고 channel-mixing MLPs에 존재한다. 따라서 quadratic인 ViT와 다르게 네트워크의 계산 복잡도는 linear다. $D_S$는 input patches의 수에 독립적으로 선택되며, $D_C$는 patch size에 독립적으로 선택된다. $D_C$가 patch size에 독립적이기 때문에 CNN과 다르게 전반적인 복잡도는 선형적이다. 

 

Mixer는 skip-connections와 layer normalization도 적용한다. ViT와 다르게 position embeddings는 적용하지 않는데 그 이유는 input tokens의 순서에 token-mixing MLPs가 민감하기 떄문이다. 마지막으로 Mixer는 global average pooling 레이러를 활용해서 표준적인 classification head로 사용하여, 뒤 이은 linear classifier로 예측한다. JAX/Flax로 구조를 작성했다. 코드는 Supplementary E에 제시되어있다. 

 

3. Experiments

Downstream tasks

ILSVRC2012의 ImageNEt의 original validation labels와 cleaned-up ReaL labesl를 사용한다. CIFAR-10과 100, Oxford-IIIT Pets, Oxford Flowers-102, Visual Task Adaption Benchmark (VTAB-1k)에 대해서 적용했다.

 

Pre-training

ILSVRC2021 ImageNet을 사용하는데 이는 ILSVRC2012를 포함하는 데이터다. 또한 JFT-300M도 활용해서 사전학습했다. ViT 논문과 BiT (Big Transfer)를 따라서 dwonstream tasks의 test sets에 대하여 pre-training datasets를 de-duplicate한다. 모든 이미지의 해상도는 224이며, $\beta_1$ = 0.9, $\beta_2$ = 0.999을 가지는 Adam, 그리고 linear learning rate와 warmup은 10k steps, linear decay, 그리고 배치 사이즈는 4096, weight decay를 적용하고, gradient clipping은 global norm이 1로 적용한다. JFT-300M에 대해서는 cropping과 random horizontal flipping으로 증강한다. ImageNet에 대해서는 추가적인 증강과 regularization을 적용하는데, RandAugment, mixup, dropout, stochastic depth를 사용한다. 이는 timm 라이브러리에 영감을 받았다. 하이퍼파라미터의 자세한 내용은 Supplementary B에 소개되어있다.

 

Fine-tuning 

파인튜닝에서는 momentum SGD, batch size 512, gradient clipping at global norm 1, and a cosine learning rate schedule with a linear warmup를 사용한다.

 

Models

Base B, Large L, Huge H로 다양한 Mixer 구조가 있다. 자세한 사항은 아래 Table 1에 나와있다. 

 

 

 

 

 

 

 

 

Mixer가 ViT나 BiT에 근사한 성능임을 Table 2와 3에서 확인할 수 있다.