[논문 리뷰] CMT: Convolutional Neural Networks Meet Vision Transformers

---

CMT: Convolutional Neural Networks Meet Vision Transformers

CNN의 Locality 개념을 단순히 ViT에 주입하고자했던 Swin Transformer와 같은 연구들과 달리 정말 Vision Transformer에 CNN을 결합한 방식으로 효과적인 결과를 얻어낸 CMT 논문을 리뷰해보자.

 

Abstract

---

• 학계의 다양한 연구를 통해 Vision Transformer는 이미지를 전반적으로 확인하는 Long Range Dependency에 강한 모습을 보이며, 이미지 인식 분야에 성공적으로 적용되고 있다. 
• 하지만 여전히 CNN과 ViT는 그 연산량과 성능에 있어서 꽤나 큰 간극이 존재한다.
• 본 논문에서는 단순히 Transformer만으로 Vision Task를 해결하지 않고, 고성능의 CNN을 함께 사용하는 방식을 취함으로써 기존 ViT 뿐만 아니라 CNN Legacy 모델들 또한 능가하는 모델을 제안한다.
• Long Range에 강력한 Transformer와 Feature Map을 통한 Local Feature에 강력한 CNN의 장점을 결합하여 기존 방식보다 더 성능이 높으면서도 비교적 연산량이 적은 CMT 모델을 제안한다.
• CMT-S는 적은 연산량으로 ImageNet에서 83.5%의 Top-1 Accuracy를 달성하였고, CIFAR10에서는 99.2%, CIFAR100에서는 91.7%의 높은 정확도를 달성했다.

 

Introduction

---

• 최근 몇년간 CNN의 압도적인 활약으로 인해 Computer Vision 분야는 특징점을 잘 추출해내어 이미지를 분류하는등의 Task에서 크나큰 발전을 이룩하였다.
• 반면, Self Attention 기반의 Transformer 모델은 Token 간의 거리가 멀때도 그 연관관계를 효과적으로 찾아낼 수 있는 강점이 있기 때문에 NLP분야에서 높은 성능을 보이고 있다.
• 최근, ViT와 같이 이러한 Transformer 구조를 Vision 영역에 도입하고자 하는 시도가 증가하고 있는데, Object Detection, Segmentation, Classification등에서 꽤나 유의미한 성능을 보이고, 때로는 CNN을 능가하기도 하면서 아주 유망하다고 평가되고 있다.
• 그럼에도 불구하고, 동일한 크기의 Small CNN 모델에 비해서 ViT는 확연히 성능이 떨어지는 모습을 보이는데, 본 논문의 저자들은 그 이유를 3가지로 정의했다.
• 첫째로, ViT의 경우 단순히 이미지를 패치형태로 나누고 1D Sequence의 형태로 입력을 진행한다. 비록 Transformer의 이점덕분에 Patch들 간의 Long Range Dependency를 효과적으로 찾아내긴 하지만, 어찌됐던 이는 이미지의 2D Locality와 같은 근본적인 차이점을 고려하지 않게 된다.
• 둘째로, ViT의 경우 항상 fix된 patch size를 사용하여 Multi Scale Feature Map을 추출할 수 없으며, 특히 Low Resolution Feature Map에 취약한 모습을 보인다.
• 셋째로, ViT는 이미지를 패치로 잘라 입력하는 방식을 취하게 되고, 이미지 크기에 대해 Quadratic한 연산량이 필요로 한다. ViT는 O(N^2*C)인 반면, CNN은 O(N*C^2) 형태이므로, COCO와 같이 1333x800 해상도, Cityscape와 같은 2048x1024등을 학습시키는데에 무리가 있다.
• 본 논문에서는 이러한 CNN과 ViT를 적절히 결합한 Novel Architecture인 CMT를 제안한다.

 

CMT Architecture

---

• 저자들은 CNN과 Transformer의 장점을 모두 활용하기 위해 위와 같은 Hybrid 방식의 Architecture를 제안한다.
• DeiT(ViT) 방식의 경우 이미지를 패치로 잘라서 입력을 하고 Linear Projection을 통해 이를 합치다보니 기존의 위치 정보가 많이 손실되고, Poor Modeling이 되는 한계점이 존재한다.
• 이를 해결하기 위해서 CMT에서는 CMT Stem이라는, 마치 ResNet의 도입부와 같이 Input을 처리하여 Input의 사이즈를 줄이고 핵심 정보를 추출해낸다.
• 그리고 기존 ResNet, EfficientNet과 같은 CNN구조들과 같이 4 Stage 구조를 사용하였는데, 이는 마치 ResNet50, ResNet101과 같은 CNN의 표준 모형을 따르는 것처럼 보인다.
• CMT 블록 내부에서는 Local Feature와, Long Range Dependency를 모두 고려할 수 있도록 제작되었으며 최종 Output은 Global Average Pooling Layer와 함께 FC를 거쳐 Classification Layer로 연결된다. CNN의 기본 구조와 거의 일치하는 구조이기 때문에, 다양한 영역에서 Backbone으로 CNN을 대체하여 사용되기 용이하다.
• 위와 같이 4 Stage 중간중간에 2x2 Convolution Layer를 Stride 2와 함께 사용하여 마치 Modern CNN구조와 같이 Feature Map의 Scale을 점점 바꿔가면서 학습을 진행하는 효과를 가지며, 이를 통해 입력 이미지에 대한 Multi Scale Feature Map을 얻어낼 수 있고, Detection 혹은 Segmentation과 같은 downstream task에 효과적인 모습을 보인다.

 

CMT Block

---

LPU(Local Perception Unit)

• Rotation과 Shift는 Vision Task에서 흔히 사용되는 Augmentation 기법이다. 이러한 변형은 모델의 최종 예측에 영향을 미치면 안된다는 Translation-Invariance를 만족시키기 위함이다.
• 하지만 기존 ViT 구조에서 사용하던 Absolute Position Encoding은 각각의 Patch에 대해 Unique한 값을 할당하기 때문에, 이러한 Invariance의 원칙을 손상시킨다. 또한 이러한 방식은 이미지의 Locality 정보를 무시할 수 있으며, Patch의 구조적인 특성 또한 간과될 수 있다.
• 이러한 한계점을 개선하기 위해서 본 논문에서 제안하는 LPU 방식은 Depth Wise Convolution을 3x3로 진행하여 Locality 정보를 추출해내고 Residual을 더해준다.

 

Lightweight Multi-head Self-attention

기존 Attention

LightweightAttention과 Relative Position Bias

• 기존 Self Attention 방식의 경우, 입력 X가 nxd의 형태를 가질 때, 이를 Linear하게 Transform하여 Query와 Key, Value를 각각 쪼개서 처리하였지만, 이는 매우 많은 연산량을 요구로하게 된다.
• 본 논문의 저자들은 보다 k x k 크기의 depth-wise convolution을 stride k와 함께 진행하여 Query와 Key의 크기를 감소시켜 연산량을 줄였다.
• 또한 Swin Transformer에서 사용되었던 Relative Position Bias 개념을 채택하였다. Swin Transformer에서는 그냥 단순히 좌표의 상대거리를 통해 Relative Position을 구했던 반면, 해당 논문에서는 이를 임의로 초기화 한 후 학습 가능한 방식으로 학습을 진행하였다.

 

Inverted Residual Feed-forward Network

• 기존 ViT의 경우, FFN으로 2개의 Linear Layer를 사용하였다. 또한 첫번째 Layer는 expansion인 4만큼 차원을 증가시켜주고, 두번째 Layer를 거치면서 다시 원래의 차원으로 돌아오는 방식이다. (ResNet의 BottleNeck구조)

• 본 논문에서 저자들은 기존 Bottleneck 구조인 wide-narrow-wide 방식이 아닌 MobileNetV2에서 사용됐던 narrow-wide-narrow 방식을 사용한다. 이는 이미 bottleneck에 충분히 필요한 데이터가 모두 있다고 가정하고, input과 output의 차원을 줄여서 메모리 사용량을 개선시킨다. 해당 방식은 1x1 Conv를 거쳐 채널을 확장시킨 뒤, 3x3 depth wise convolution을 수행하고 다음 1x1 Conv에서 다시 채널을 축소시키는 방식이라서 Inverted Residual Feed-Forward Network로 명명되었다.
• 또한 깊은 Layer 구조에서의 효과적인 학습을 위해서 ResNet 방식으로 Shortcut을 더해주어 정보를 보충하는 방식을 적용하였다.

 

Scaling Strategy

• EfficientNet 구조의 영감을 받아 Compound Scaling 전략을 도입하였다.
• 이를 통해 Depth, Dimension, Resolution에 대해서 가장 최적의 값을 Grid하게 Search하게 된다.

 

 

Result

---

ImageNet Classification

 

 

Transfer Learning

 

 

Object Detection, Instance Segmentation

 

Conclusion

---

• 본 논문에서는 Vision Transformer들의 장점인 Long Range Dependency 추론 능력과, CNN의 강점인 Locality를 고려한 추론 능력의 이점을 모두 얻기 위해서 두가지 모델을 함께 섞은 Hybrid Architecture인 CMT를 제안하였다.
• Inverted Residual FFN, Lightweight MHSA등을 적용하여 연산량을 효과적으로 감소시켰으며, 이로 인해 적은 FLOPS를 가지면서도 CNN을 포함한 기존 모델들에 비해서 높은 정확도를 가진다.
• Classification 뿐만 아니라, 기본적으로 마치 CNN과 같은 Architecture 구조를 가지고 있으므로 기존 CNN을 사용하던 Backbone에 손쉽게 대체될 수 있으며, 실제로 downstream task에서도 개선된 성능을 보여주었다.