[논문] MLP-Mixer: An all-MLP Architecture for Vision

CNN은 computer vision 분야에서 기본이 되는 모델이다. 최근에 attention-based networks들인 Vision Transformer도 제시된다. 하지만 convolution이나 attention이 좋은 성능을 내는데에 충분하지만 필수적이지는 않다.

본 논문에서는 muti-layer perceptron (MLP)에만 기반하는 MLP-Mixer를 제시한다. 큰 데이터셋들과 최근 정규화 기법들을 사용했을 때 MLP-Mixer는 이미지 분류 benchmark들에서 경쟁력있는 성능을 보여준다.

 

 

 

1 Introduction

기존 computer vision 분야에서는 CNN이 de-facto standard이다. 최근에는 self-attention layer들에 기반한 Vision Transformer (ViT)가 소개되어 SOTA 성능을 보여준다. ViT는 hand-crafted 시각 요소들과 inductive biases를 제거하고 raw data에서 직접적으로 학습하는 트렌드를 따른다.

 

본 논문에서는 MLP-Mixer 구조를 제시한다. MLP-Mixer는 MLP만을 사용해 convolution이나 self-attention을 대체한다.

전체적인 구조는 Figure1에서 묘사된다.

• Linearly projected 이미지 patch(또는 token)들을 input으로 받는다 (patches x channels)
• Mixer는 두가지 MLP layer를 사용한다:
  • Channel-mixing MLP
    • Channel 간 소통을 가능하게 해준다
    • 각 token에 독립적으로 적용되고 각 row를 input으로 받는다
  • Token-mixing MLP
    • 서로 다른 위치적 장소(token) 간 소통을 가능하게 한다
    • 각 channel에 독립적으로 적용되고 각 column을 input으로 받는다

제시된 구조는 극단적으로 특별한 CNN으로 보여질 수 있다. 하지만 convolution은 간단한 matrix multiplication보다 더 복잡하다.

 

간단한 구조에도 불구하고 Mixer는 경쟁력있는 결과를 보여준다.

Pre-training을 사용하면 기존 CNN과 Transformer의 SOTA와 거의 비슷한 성능을 보여준다.

하지만 ViT처럼 일반적인 vision task에 모두 적합하지는 않다.

 

 

 

2 Mixer Architecture

최근 deep vision 구조들은 feature들을 섞는 layer들로 구성되어 있다.

• i) 각 공간적인 위치에 있는 feature들
• ii) 서로 다른 공간적 위치에 있는 feature들

CNN의 경우

• (ii)를 NxN convolution과 pooling을 통해 구현한다
• (i)는 1x1 convolution을 통해 구현된다
• 큰 kernel들은 (i)과 (ii)를 모두 수행한다

Vision Transformer나 다른 attention 기반 구조들의 경우

• (i)와 (ii)는 self-attention으로 수행되고
• (i)는 MLP-block으로 수행된다

Mixer의 구조는 이 per-location 작업인 (i)와 cross-location 작업인 (ii)를 완전히 분리하자는 아이디어에 기반한다.

 

Figure 1이 전체적인 Mixer 구조를 묘사한다.

• Input
  • S개의 non-overlapping 이미지 patch
    • $ S = HW / P^2 $
  • 각 patch는 hidden dim C로 project된다
    • 이때 모든 patch들은 같은 projection matrix를 사용한다
  • $X \in R^{S \times C} $
• Layer
  • Mixer는 같은 크기의 layer들로 이루어진다
  • 각 layer는 두 개의 MLP block으로 구성된다:
    • Token-mixing MLP
      • X의 각 column에 적용된다 ($X^T$)
      • $R^S -> R^S$
    • Channel-mixing MLP
      • X의 각 row에 적용된다
      • $R^C -> R^C$
  • 각 MLP block은 input의 각 row에 독립적으로 적용되는 두 개의 fully-connected layer와 nonlinearity로 구성된다
    • $ \mathbf{U}{*,i} = \mathbf{X}{,i} + \mathbf{W}_2 \sigma \left( \mathbf{W}1 \text{LayerNorm} (\mathbf{X}){,i} \right), \quad \text{for } i = 1 \dots C, $
    • $ \mathbf{Y}{j,*} = \mathbf{U}{j,} + \mathbf{W}_4 \sigma \left( \mathbf{W}3 \text{LayerNorm} (\mathbf{U}){j,} \right), \quad \text{for } j = 1 \dots S. $
      • $\sigma$: element-wise nonlinearity (GELU)
      • $D_S, D_C$: token-mixing과 channel-mixing MLP들의 튜닝 가능한 hidden width이다
        • $D_S, D_C$는 input patch 수와 독립적으로 선택된다
        • 따라서 input patch 수에 대해 선형 계산 복잡도를 가진다

 

위에서 언급됐듯이 각 row (column)에 같은 channel-mixing MLP (token-mixing MLP)가 적용된다.

• Channel-mixing MLP에 parameter tying을 적용하는 것은 널리 쓰이는 방식이다.
  • Positional invariance와 convolution의 주요한 요소를 제공한다
• 하지만 channel들을 교차하는 parameter tying을 거의 쓰이지 않는다.
  • 이는 기존 convolution과 다르게 모든 channel에 같은 kernel을 사용하는 token-mixing MLP block의 특징 때문이다.

 

MLP layer 이외에 Mixer는 다른 표준적인 구성 요소들을 사용한다.

• Skip-connection, layer normalization
• Global average pooliing layer 이후에 linear classifier를 사용한다

 

 

 

3 Experiments

MLP-Mixer 모델의 정확도를 평가하기 위해서 중간 크기에서 큰 크기의 dataset으로 pre-train을 진행하고 작거나 중간 크기의 데이터셋으로 downstream 분류 작업을 수행한다.

크게 세개의 평가 요소에 집중한다:

• Downstream task의 정확도
• Pre-training에 사용된 전체 계산 비용
• Test-time throughput

 

Downstream tasks

Downstream task에는 다음과 같은 데이터셋을 사용한다:

• ILSVRC2012 ImageNet
  • 기존 validation label, cleaned-up ReaL label
• CIFAR-10/100
• Oxford-IIIT Pets
• Oxford Flowers-102
• Visual Task Adaptation Benchmark

 

Pre-training

Pre-training에는 두 데이터셋을 사용한다:

• ILSVRC2021 ImageNet
• ImageNet-21k
• JFT-300M

모든 pre-training dataset에서 downstream task에 포함된 test set은 de-duplicated 된다.

Setup:

• 224x224 화질
• Adam (beta1 = 0.9, beta2=0.999)
• Linear learning rate warmup 10k, linear decay
• Batch size 4096
• Weight decay
• Global norm 1에서 gradient clipping

JFT-300M의 경우 "Going deeper with convolutions."에서 제시된 cropping 기술과 random horizontal flipping을 적용한다.

ImageNet과 ImageNet-21k의 경우 추가적인 증강과 정규과 기술을 사용한다.

• RandAugment, mixup, dropout, stochastic depth

 

Fine-tunning

Setup

• Momentum SGD
• Batch size 512
• Gradient clipping at global norm 1
• Cosine learning rate scheduler with linear warmup

추가적으로 pre-training에 사용된 이미지보다 더 좋은 화질로 fine-tunning을 진행한다.

• Input resolution이 바뀌면 모델 구조도 바뀌어야한다
• Hidden layer의 넓이를 $D_S$에서 $D_{S'}$로 바꾼다.
• 그 후 weight matrix $W'_2$를 기존 $W_2$의 matrix의 복사본을 대각선에 가지는 block-diagonal matrix 로 초기화한다
  • 이 기법은 $S' = K^2S$일때만 가능하다 (K는 자연수)

VTAB-1k의 경우 224 화질과 448 화질로 각각 fine-tuning을 진행한다.

 

Metrics

모델의 계산 비용과 정확도 trade-off를 평가한다.

• 계산 비용
  • TPU-v3에서의 pre-traning time
    • 세 가지 요소가 결합된 평가 요소이다 
      • 이론적 FLOPs, 계산 효율성, 데이터 효율성
  • TPU-v3에서의 throughput
    • 다양한 batch size를 시도하고 가장 높은 throughput을 보고한다
• 정확도
  • Fine-tunning이후 top-1 downstream 정확도
    • 이 지표 비용이 너무 높은 경우 few-shot accuracy로 대체한다

 

Models

MLP-Mixer의 다양한 변종과 CNN, attention-based 모델들을 비교한다.

위 테이블은 MLP-Mixer의 다양한 setting들을 비교한다.

• Attention-basd
  • Vision Transformer (ViT) 
  • HaloNet-H4
• Convolution-based
  • Big Transfer (BiT)
  • NFNet-F4
  • MPL
  • ALIGN

 

3.1 Main results

위 테이블은 실험 결과를 요약한다.

• ImNet은 기존 ImageNet의 validation을 사용한 결과이다
• ReaL은 cleaned-up label을 사용한 결과이다
• Avg 5는 모든 downstream task (ImageNet, CIFAR-10, CIFAR-100, Pets, Flowers)의 평균 성능이다

ImageNet-21k로 pre-train된 결과 Mixer는 비교적 비슷한 정확도를 보여준다.

• Regularization이 없으면 Mixer는 overfit 하는 경향을 보여준다

Pre-train data가 더 커지면 크게 향상된 정확도를 보여준다.

• BiT보다 0.5% 높고, ViT보다 0.5% 낮은 정확도이다
• 하지만 Mixer는 ViT보다 2.5배 빠르게 작동하고 BiT보다도 2배 빠르게 작동한다

위 Figure 2 왼쪽은 정확도와 계산 비용 trade-off를 나타내는 그래프이다.

 

 

3.2 The role of the model scale

이 섹션에서는 더 작은 Mixer model을 비교해본다.

모델은 크게 두 가지로 확장될 수 있다.

• Pre-training에 model 크기를 증가 (layer 수, hidden dim, MLP 넓이)
• Fine-tuning에 이미지 화질 증가

실험에서는 다양한 Mixer 설정을 비슷한 크기를 가진 ViT와 BiT 모델들과 비교한다.

위 Table 3과 Figure 3은 실험 결과를 요약해 보여준다.

• Mixer-B/16은 ViT-B/16 모델보다 훨씬 더 낮은 정확도를 보여준다
  • Training loss는 비슷했기 때문에 Mixer가 overfitting에 취약하다는 것을 보여준다
• Pre-training dataset 크기가 커질수록 Mixer의 성능은 지속적으로 향상된다.
  • JFT-300M으로 pre-train 되었을 때 ViT-H/14에 비해 0.3% 낮은 성능과 2.2배 빠른 학습 속도를 보여준다.

 

 

3.3 The role of the pre-training dataset size

위 결과에 따르면 pre-training dataset이 커질수록 Mixer의 성능은 상당히 향상된다. 이 섹션에서는 그 이유를 더 자세히 연구한다.

 

Mixer-B/32, Mixer-L/32, Mixer-L/16 모델을 각각 JFT-300M의 3%, 10%, 30%, 100% subset으로 pre-training을 진행한다.

추가적으로 Mixer-L/16 모델을 JFT-3B 데이터셋으로 pre-training을 진행한다.

 

Figure 2 오른쪽 그래프에 결과를 보고한다.

• 작은 subset으로 pre-train 되었을 때는 모든 Mixer 모델이 크게 overfit 한다.
• 데이터셋이 커질수록 Mixer-L/32와 Mixer-L/16의 성능은 BiT보다 빠르게 향상된다.
  • ViT의 경우 비슷한 향상을 보이지만 Mixer 모델에 비해서는 향상폭이 크지 않다.
    • 이는 ViT의 self-attention layer들이 Mixer 구조가 발견하는 true underlying distribution보다 덜 적합하기 때문이라 추론할 수 있다.

 

 

3.4 Invariance to input permutations

이 섹션에는 Mixer와 CNN구조의 inductive bias 차이를 연구한다.

Mixer-B/16과 ResNet50x1 모델을 JFT-300M으로 pre-training한 후:

• 16x16 patch들 순서를 바꾸고 각 patch의 pixel 순서를 공유된 permutation으로 변경
• 모든 이미지의 pixel을 랜덤하게 섞기 (모든 이미지에 같은 permutation 사용)

위 Figure 4는 실험 결과와 섞인 이미지 샘플을 몇 개 보여준다.

• Mixer의 경우 예상했듯이 같은 patch안에 있는 pixel 순서가 중요하지 않다
• ResNet의 경우 pixel의 위치에 크게 의존하고 pixel 순서가 바뀌면 성능도 크게 하락한다
• Global permutation의 경우 마찬가지로 Mixer는 적은 성능 하락을 보인다.

 

 

3.5 Visualization

CNN의 첫 layer들은 지역적인 영역에 있는 pixel들에 대해 Gabor-like detector를 학습한다.

MLP가 이와 비슷한 방법으로 정보를 처리하는지 의문이 생긴다.

Figure 5는 Mixer의 첫 세개 token-mixing MLP의 hidden unit들을 보여준다.

• 몇몇 feature들은 전체 이미지에 적용되고 몇몇은 작은 영역에만 적용된다
• 깊은 layer들도 확인 가능한 구조를 가지고 있지 않다
• Hyperparameter에 따라 학습되는 unit의 구조다 달라진다

 

 

 

4 Related work

---

 

 

 

5 Conclusion

본 논문에서는 vision task를 위한 아주 간단한 구조를 제시한다.

실험을 통해 이런 간단한 구조도 SOTA 기법들과 유사한 accuracy-compuataional resource trade-off를 보이는 것을 확인했다.

 

실용적인 측면에서 제시된 모델로 학습된 feature들과 CNN과 Transformer로 학습된 feature들의 차이점 연구가 쓸모있을 수 있다.

이론적인 측면에서 다양한 feature들에 숨겨진 inductive bias를 이해하고 일반화에 주는 영향을 연구할 수 있다.