YoojLee / paper_review

paper review archive
6 stars 2 forks source link

MetaFormer is Actually What You Need for Vision (2022) #37

Open YoojLee opened 1 year ago

YoojLee commented 1 year ago

image

Abstract

기존의 transformer 성공 뒤에는 attention-based token mixer module이 성능에 가장 많이 기여한다는 common belief가 있었음. 그러나, attention module을 다른 token mixer로 대체한 기존 연구들이 발표되면서 해당 논문은 기존의 common belief에 반하여, transformer 구조가 competent한 건 어떤 token mixer를 사용하느냐가 아니라 "metaformer" 구조때문이라고 주장함. 이를 보이기 위하여 token mixer로 아주 간단한 non-parameteric token mixer인 pooling operation을 사용하여 성능을 평가한 결과, 성능이 기존의 "well-tuned" token mixer based transformer 대비 성능이 매우 좋음을 알 수 있었음. → 이에 해당 논문은 "MetaFormer"의 컨셉을 제안하고 token mixer에 집중하는 게 아니라 이 metaformer의 구조를 발전시키는 것을 제안함.

1. Introduction

image

기존의 transformer encoder는 두 가지 요소로 나눌 수 있음. → attention으로 대표되는 token mixer (논문에서 이르길), 그 외의 channel mlp나 residual connection을 포함한 다른 요소들.

attention을 token mixer의 구체적인 한 형태라고 상정하면, 전체 transformer 구조는 Figure 1(a)에서 볼 수 있듯이 metaformer로 나타낼 수 있음. 기존의 mlp mixer 계열의 구조들이 제안되면서 기존의 token mixer가 attention에 국한되지 않아도 된다는 가능성을 제기함. 이를 바탕으로 저자들은 metaformer 구조 자체가 중요하다는 가설을 제시하고, 이를 poolformer라는 non-parameteric pooling operation을 이용한 token mixer를 통해 이러한 가설이 실제로 valid함을 보여주고자 함.

논문은 실험을 통해 poolformer를 여러 vision task에서 높은 성능을 보임을 입증함.

2. Related work

Transformer는 NLP 분야에서 처음 제안되고 현재는 하나의 표준으로 자리잡음. 이를 Vision 분야에도 적용하여 ViT가 처음 제안되며, 이후에 경량화, shfited window 등의 컨셉을 도입하며 많은 attention 기반의 ViT-like model이 제안됨. 한편, Attention 기반 이외에도 MLP를 이용하여 token mixing을 수행하는 MLP-mixer가 attention 기반의 transformer보다 좋은 성능을 보임. 이후 어떠한 token mixer를 사용하는 것이 좋을지, SOTA 성능을 기록하기 위한 새로운 complicated token mixer를 제안하는 방향으로 많은 연구가 이루어짐. 그러나, 해당 논문에서는 기존의 연구 방향에 대한 의문을 던지고 transformer-like model이 성공한 데에는 어떠한 요소가 실질적인 영향을 미쳤는지에 대한 investigation을 수행함. 저자들은 metaformer라는 transformer-like model의 abstraction이 실질적으로 transformer-like model의 성능을 견인했다고 밝힘. 기존의 연구들 중에도 transformer-like model의 성능을 밝히고자 한 시도는 있었으나 metaformer와 같이 transformer-like model을 general한 concept으로 abstract하고 어떤 general framework의 측면에서 해석한 시도는 없었음.

3. Method

MetaFormer

Figure 1-(a)에서 나와 있듯이, input embedding (ViT의 patch embedding과 같은 작업)을 수행하고 metaformer block에 반복적으로 feed되는 구조임. metaformer block과 같은 경우에는 normalization → token mixer → norm → channel mlp 와 같은 구조를 취하고 있음. 또한, residual connection 역시 추가해주고 있음. token mixer의 경우에는 기존의 attention 혹은 mlp가 될 수 있을 것이며 normalization의 경우에는 layer norm, batch norm 등이 추가될 수 있음. Attention 같은 경우에는 token mixer인 동시에 추가적인 channel mixing 연산도 수행한다는 점.

PoolFormer

image 논문에서는 metaformer framework에 average pooling을 token mixer로 활용하여 poolformer라는 구조를 만들어냄. poolformer가 좋음을 보여주는 것이 아니라 token mixing을 위해 별도의 설계 없이 아주 간단한 non-parametric pooling을 통해서도 준수한 성능이 나옴을 보여주면서 metaformer 구조 자체의 중요성을 보여주기 위함임.

pooling이 simple하다는 이유 이외에도 기존의 popular token mixer인 self-attention과 spatial MLP의 경우에는 토큰 시퀀스 길이에 complexity가 quadratic하게 증가하기 때문에 길이가 긴 시퀀스는 다루기 힘들지만 (특히, spatial mlp 같은 경우에는 token 개수가 증가할 수록 parameter도 증가하는 구조), pooling의 경우에는 시퀀스 길이에 따라 추가되는 parameter 없이 linear하게 complexity가 증가한다는 장점이 있음. 따라서, 이러한 pooling operation을 기반으로 hierarchical한 구조를 취하여 4개의 stage로 구성된 poolformer를 설계함.

자세한 구조는 아래와 같음.

image

discussion 간단한 non-parametric pooling 구조라고는 하지만 실질적으로 블록 사이에 patch embedding을 반복적으로 끼워넣으면서 token mixing을 암묵적으로 수행하는 것으로 볼 수 있을 듯함. 좀 더 자세히 얘기하자면, poolformer에서 사용한 patch embedding은 convolution 기반으로, stride가 patch size보다 작음. 이렇게 될 경우에는, 기존의 transformer-like model에서 주로 가정하던 non-overlapping patch에 대한 가정이 사라지게 됨. 이렇게 되면, 결국에 각 패치는 실질적으로 특정 location에 대한 정보를 공유하는 구조임. Token mixing이 암묵적으로 수행된다고 볼 수 있을 것 같음. 이는 추후 ablation에서 전혀 의미 있는 token mixing이 이루어지지 않는 경우에도 어느 정도의 성능이 기록된다는 점에서 암묵적인 token mixing이 이루어진다고 볼 수 있을 것 같음. 또한, 실질적으로 convolution 필터 개수도 달라지면서 channel 축에 대한 projection도 이루어지는 것으로 보임 (물론, ablation에 따르면 이는 patch embedding이 가져오는 token mixing의 효과보다 더 적긴 함). 이러한 점에서, 정확히 여기서 제안하는 metaformer의 컨셉이 정말 Figure1과 같다면, poolformer는 metaformer의 적합한 demonstration이 아니라는 생각이 듦. 또한, 기존의 transformer-like model의 대부분이 non-overlapping patches를 가정하고 있고 구조 자체의 flexible함이 장점이라고 생각이 드는데 이미지 데이터에 맞게 너무 많은 prior가 들어가 있지 않나라는 아쉬움이 듦. 그럼에도 불구하고 정말 metaformer의 컨셉이 norm → token mixing → (residual connection) → norm → channel mixing → (residual connection)의 블록 형태가 이어져있는 구조만을 의미하는 거라면, 위의 의문은 그냥 minor한 이슈로 볼 수 있을 것 같음.

image

논문에서는 metaformer 구조 자체가 token mixer choice보다 더 중요하다는 점을 밝히기 위해 위와 같은 non-parametric pooling 연산으로 token mixing을 수행함. Pooling 연산 수식을 보면 자기 자신을 빼주는 걸 볼 수 있는데, 여기서 자기 자신을 빼주는 이유는 MetaFormer Block이 residual connection이 있기 때문이라고 주장함.

이에 대해 조금 더 생각해보면, 우리가 residual connection을 왜 하는지에 대해 조금 더 생각해보면 됨.

4. Experiments

크게 3가지 태스크(Image classification, Object detection & Instance segmentation, Semantic segmentation)에 대해서 성능 평가를 진행함. 성능 평가의 목적은 SOTA보다는 정말 poolformer 자체가 여러 task에서 competitive한 성능을 보여주는지에 초점을 맞춤.

4.1. Image classification

Setup

ImageNet-1K. data augmentation 적용, AdamW를 이용하여 300 에폭 동안 학습하고 weight decay는 0.05로 기존 ViT보다는 적게 줌. learning rate scheduling (warm up for 5 epochs and later cosine schedule was applied) + label smoothing (0.1 of factor) + stochastic depth + layer scale 등의 regularization technique 을 활용함. dropout은 적용되지 않음. 또한, Layer Norm 같은 경우 기존의 layer norm은 channel dimension에 대해서만 통계량을 계산하지만 여기에서는 token과 channel dimension에 대한 통계량을 계산했다는 점 (즉, 여기에서도 token에 대한 정보를 활용한다는 점임 → 간접적인 token mixing이라고 볼 수 있을 듯)

Results

image image

크게 ConvNet과의 비교를 수행하고, Attention과 spatial mlp token mixer를 사용한 metaformer 구조와의 비교를 수행하고 있음. poolformer의 경우 전반적으로 좋은 성능을 보일 뿐 아니라 같은 MACs (Multiply-accumulate operations)와 model size 하에서 baseline보다 더 높은 성능을 보여주고 있음. 이는 반대로 생각하면 같은 성능을 내기 위해 보다 적은 parameter가 필요하다는 점. naive한 token mixer를 사용하여도 기존에 sophiscated token mixer를 사용한 구조들보다 더 좋은 성능을 보였다는 점에서 중요한 것은 metaformer 구조 자체라는 것을 보여주고자 함. 한편, ResNet 구조와의 비교를 통해서 ResNet은 neural convolution이라는 더 높은 locality spatial aggregation 성능을 갖는 token mixer를 사용함에도 poolformer보다 좋지 못한 성능을 보였다는 점에서 metaformer concept 자체가 가져다주는 이점을 어느 정도 설명하고 있음.

4.2. Object detection and instance segmentation

Setup

COCO benchmark. RetinaNet, Mask R-CNN 구조를 활용하여 backbone을 갈아끼우는 형태로 비교 실험을 진행함. baseline으로는 파라미터 수가 비슷한 ResNet 구조를 사용함. ImageNet-1K로 학습된 가중치와 함께 새로 추가된 부분은 Xavier init을 활용하여 가중치 초기화를 진행함.

Results

image

모든 경우에 baseline을 outperform하는 것을 확인할 수 있었음.

4.3. Semantic segmentation

Setup

ADE20K에 대한 실험. Semantic FPN 구조에 backbone을 갈아끼우는 형태로 비교. 가중치 초기화는 OD의 경우와 동일함. 베이스라인은 역시 비슷한 파라미터 수를 갖는 베이스라인으로 구성함. ResNet, ResNext, PVT (transformer 기반)을 베이스라인으로 활용함.

Results

image semantic segmentation의 경우에도, token mixing을 위해서 pooling에만 의존하는 구조임에도 불구하고 poolformer를 backbone으로 사용한 경우가 competitive한 성능을 보여줌. 이는 역시 MetaFormer 구조 자체가 great potential을 갖고 있고 성능에 가장 많이 기여한다는 주장을 뒷받침해줌.

4.4. Ablation studies

classification에 대해 ablation을 수행함. image

Token mixers

Table 5에 따르면, identity mapping과 global random matrix (fixed random weight → softmax 적용)를 사용함에도 불구하고 Top-1 accuracy의 손실이 baseline 대비 거의 없음을 알 수 있음. identity mapping과 global random matrix를 활용한다는 것은 결국에는 의미 있는 token mixing은 거의 이루어지지 않는다는 점인데 그럼에도 불구하고 매우 준수한 성능을 보여주고 있음. 이는 MetaFormer 컨셉 자체가 성능에 많은 기여를 한다는 주장에 힘을 실어줌. → 이 부분에 대해서는 조금 물음표인 게 과연 의미 있는 token mixing이 거의 이루어지지 않는데도 저렇게 성능이 나오는 게 가능한 걸까? patch embedding 과정에서 strided convolution이 이루어졌다는 점에서 token mixing의 효과를 얻었기 때문이 아닐까?

다른 token mixer에 대해서는 depth-wise convolution 적용(이 경우에는 depth-wise conv가 더 powerful한 token mixer이므로 조금 더 성능이 좋게 나옴), pooling size 변경 등의 시도를 추가했으나 성능에 큰 영향은 없었음. Hyperparameter 등에 큰 영향 없이 token mixer를 뭘 쓰든 성능이 어느 정도 준수하게 나온다는 점을 미루어 보았을 때, MetaFormer 컨셉 자체의 중요성이 역시 demonstrate되고 있음.

Normalization

poolformer의 경우, token 간의 정보를 더 활용하는 MLN을 Normalization 방법으로 활용함. 이는 기존의 Layer Norm이나 Batch Norm보다 더 좋은 성능을 보임. 아예 Normalization을 제거하는 것은 학습 안정성을 저해하며 converge가 제대로 되지 않음.

Activation

ReLU나 SiLU로 대체하여 실험한 결과, 성능의 차이는 거의 없었음.

Other components

Residual connection과 channel mlp를 제거하는 실험을 수행함. 그 결과, 성능이 매우 떨어짐을 알 수 있었음. Residual connection 같은 경우에는 제거한 결과 성능이 0.1%로 random pick의 결과가 나옴. 즉, model이 유의미한 쪽으로 전혀 수렴되지 못함을 보여줌. channel mlp 역시 비슷한 결과가 나옴. → 여기서 또 궁금한 점은 channel mlp의 경우에는 제거했을 경우 성능이 5.7로 급감했으나 token mixer의 경우에는 identity나 random weight을 사용한다고 한들 성능 손실이 거의 없었다는 점에서 보았을 때 그럼 channel mixing이 미치는 영향이 token mixer보다 훨씬 크다는 점인가? 라는 의문이 듦.

Hybrid stages

앞서 밝힌 것과 같이 pooling은 token 시퀀스 길이에 크게 구애받지 않고 핸들링이 다른 token mixer 대비 쉬움. 반면, 다른 token mixer들은 global information의 aggregation이 쉬움. 앞 부분에 pooling을 배치하고, 뒷 부분에 attention과 spatial MLP를 배치한다면 효율적이면서도 효과적인 정보 aggregation이 이루어질 수 있을 것임.

이러한 intuition에 따라 stage에 따라 token mixer의 구성을 앞단에 pool, 뒷단에 attention or spatial MLP를 구성하였을 때 baseline에서 추가적인 성능 향상이 이루어졌음. 이는 pooling과 token mixer를 적절히 combine하는 방향으로 metaformer 구조의 모델들이 성능 개선을 이루어내는 방향으로 발전될 수 있음을 시사함.

5. Conclusion and future work

MetaFormer라는 transformer-like model의 abstraction을 제안하고 이를 poolformer라는 specific한 구조를 통해 보여줌. 이때, poolformer는 token mixer로서 매우 간단한 pooling operator를 활용하면서도 기존의 transformer-like model과 견줄 만한 성능을 보여주면서 결국 transformer-like model의 성능에 가장 큰 기여를 하는 것은 MetaFormer 구조 그 자체임을 입증함. 추후의 연구 방향으로는 poolformer를 self-sup 혹은 transfer learning 등 다른 학습 세팅에서도 활용 가능한지 탐구해볼 수 있으며 이러한 metaformer 컨셉이 nlp에도 valid한지 확인하는 과정으로도 이루어질 수 있을 것임.

Wrap-ups

Strengths

Weaknesses

주장 자체는 동의하는 바이고 매우 신선하다고 생각하지만, 그걸 보여주기 위한 poolformer 구조 자체가 좋은 demonstration tool인지에 대해 조금 고민을 많이 해보게 되었던 것 같음.

YoojLee commented 1 year ago

https://github.com/sail-sg/poolformer/issues/51