[Paper review] MLP-Mixer An all-MLP Architecture for Vision

이번에 리뷰하고자 하는 논문에서 소개하는 MLP-Mixer는 Vision Transfomer가 등장한 이후 소개된 모델이다. Vision Transfomer은 Convolution 대신 'Self Attention' 이라는 개념을 중심적으로 적용함으로써, Convolution을 뛰어넘는 SOTA를 달성할 수 있었다. Vision Transfomer의 영향으로, Computer Vision에서는 Model Architecture부분에서 Convolution을 사용하지 않는 흐름이 나타났다. 오늘 소개할 논문 역시, 이러한 흐름 속에서 탄생한 모델로써, Convolution을 사용하지 않는 모델이다. 하지만 여기에 더불어, Attention의 개념조차 모델에 적용하지 않았다. MLP-Mixer는 오직 'MLP', 즉 우리가 딥러닝을 처음 배울때 사용하는 개념인 '다중 퍼셉트론' 구조만을 모델에 적용했다. 흥미로운 것은 이렇게 간단한 원리로 만들어진 모델이 기존에 SOTA를 달성하는 모델 성능들과 거의 뒤쳐지지 않는다는 점이다. 이제 본격적으로 모델의 구조를 분석해보면서, 모델을 설명해보겠다.

논문 사진 참조

 

Mixer는 어떤 구조를 가지고 있기에 MLP만으로 좋은 성능을 낼 수 있었을까?

Vision Transfomer의 논문을 읽어본 사람은 알겠지만, 큰 구조는 비슷해 보인다. 하지만 자세히 보면 확실히 그 차이점을 알 수 있다. Vision Transfomer와 완전히 같다고 할 수 있는 부분은 모델 앞부분, 즉 이미지를 patch 단위로 분할하여, 병렬적으로 input으로 넣어주고, patch embedding을 거치는 것이다. 이렇게 patch embedding까지 거친 이미지는 차원이 SxC로 표현된다. (S는 분할한 patch의 개수, C는 embedding 거친 후 channel의 dimension) 이후 MLP-Mixer의 핵심 구조라 할 수 있는 Mixer Layer를 통과하게 된다. Mixer Layer는 크게 두가지의 기능을 하는 MLP block으로 이루어져 있다. 첫번째는 Token-Mixing MLP block이다. 이 block은 말그대로 patch(=token)의 정보들을 섞는 block이다. 이를 통해 모델은 이미지의 spatial location을 global하게 학습할 수 있다고 한다. 두번째 Channel-Mixing MLP block으로, channel의 정보만을 섞는 block이다. 만약 첫번째 block만 거치게 된다면, 모델은 각 channel마다 뽑아낼 수 있는 정보들을 학습하기 힘들다. 따라서 두번째 block을 둠으로써, image의 spatial location와는 독립적으로, channel에 해당하는 정보를 모델이 학습하는 것이다. 각 MLP block은 FullyConnected-GELU-FullyConnected으로 이루어져있다. 또한 ResNet의 skip connection원리 역시 적용했고, Transfomer와 마찬가지로 LayerNorm을 적용했다. 참고로 Mixed Layer의 개수는 하이퍼 파라미터이고, 논문에서는 최대 32 층 까지 쌓은 모델을 실험했다.

이제 모델을 알아봤으니, 성능을 한번 보자. 논문에서는 결과를 두가지 측면에서 바라본 것을 확인할 수 있다. 첫번째는 '모델의 Accuracy'이고 두번째는 '계산 효율성'이다. 밑에 결과를 보면 쉽게 확인할 수 있다. (Throughput이 크고, TPUv3 core-day가 작을수록, 계산이 효율적(=모델 연산 속도가 빠름)이라고 할 수 있다.)

 

결론적으로, "Mixer는 성능면에서는 기존 SOTA에 버금가고(SOTA는 달성하지 X), 계산적인 측면에서는 다른 모델들보다 더 빠르다. 그렇기 때문에 충분히 경쟁력있는 모델이다."

논문 사진 참조

위 결과 외에도 논문을 읽어보면서 Mixer의 또다른 중요한 특징을 알 수 있었다. 그것은 'Transfomer'와 마찬가지로 큰 데이터에서 pretrain할 수록, 그 성능이 점점 더 오르는 것이다. 즉 많은 데이터의 장점을 그대로 살릴 수 있다는 말이다. 이게 가능한 이유는 'Inductive Bias'가 약하기 때문이다. Mixer의 구조를 다시 올라가서 보면, input을 global하게 Mixing하는 구조이다. 이러한 구조에서는 결국 Inductive Bias가 적을 수 밖에 없고, 이는 데이터가 어떤 형태가 들어와도 제약이 적다는 것을 의미한다.

참고로, Inductive Bias는 모델을 설계할때 초기에 세우는 assumption이라고 생각하면 된다. 가령 CNN에서는 locality라는 가정(=image의 local 영역에서 중요한 특징들을 뽑아낼 수 있다.)이 대표적인 Inductive Bias이다. 이러한 Inductive Bias가 있으면, 데이터 양이 적을때는 모델이 좀 더 잘 예측을 할 수 있게 도와주는 역할을 한다. 하지만, 데이터 양이 충분히 많아지면 성능향상에 오히려 방해가 된다.

 

이번에 MLP-Mixer를 리뷰하면서, 요즘 Computer Vision의 추세는 'CNN없이도 좋은 성능을 내는 모델'을 만들어 낼 수 있는 쪽으로 간다는 점을 느꼈고, 이는 결국 모델 자체가 가지는 Inductive Bias를 줄이는 방향으로 연구가 진행된다는 것을 알 수 있었다. Inductive Bias가 적은 모델은 큰 데이터에서 훈련할때, 기존의 CNN 기반의 모델들과 다르게 성능이 더욱 향상될 수 있다는 점에서, 데이터가 많아질 미래사회를 생각할때 더 적합한 모델이지 않을까라는 생각이 든다.

 

아래는 논문 링크이다

https://arxiv.org/abs/2105.01601

MLP-Mixer: An all-MLP Architecture for Vision

업데이트> MLP-Mixer를 파이토치로 직접 구현해보았다.

https://github.com/ChoiDae1/MLP_Mixer-PyTorch