[Paper Review] CMT 논문 이해하기

『CMT: Convolutional Neural Networks Meet Vision Transformers. CVPR. 2022.』

이번에 소개할 논문은 Convolutional Neural Networks Meet Vision Transformers 이다. ViT의 등장 이후 기존의 CNN에 활용된 기술을 그대로 사장시키는 것이 아니라 어떻게 하면 ViT에 Convolution을 녹일 수 있을까? 하는 호기심에서 비롯된 논문으로 최근까지도 CNN, RNN, Transformer의 여러 조합으로 새로운 모델들이 등장하고 있는 것을 생각하면 충분히 Contribution이 있는 논문이라고 생각한다.

 

어떻게 Vision Transformer와 CNN의 장점을 합쳤는지 그 방법에 집중하면서 읽어보도록 하자.

---

0. Abstract

Transformer의 장점은 Long-range Dependency를 잡아낸다는 것이고, CNN의 장점은 Local Information을 잘 잡아낸다는 것이다. 이에 본 논문에서는 두 가지 모델의 장점을 결합한 새로운 Hybrid Network를 제안한다. 논문에서는 이를 CMT 계열의 모델이라고 지칭하며 적은 연산량으로도 우수한 일반화 성능을 보였다고 말한다.

 

1. Introduction

ViT는 기존 CNN의 Backbone을 Transformer로 대체하여 입력 이미지(224x224x3)을 16x16x3 크기의 196개의 Patch로 나눈 후 Transformer Block에 입력해 글로벌적인 관계를 모델링하고 특징을 추출한다. 이는 최신 Vision Task에서 좋은 성능을 보였으나 아직 비슷한 크기의 CNN 모델보다 낮은 성능을 보이는 경우가 많다.

 

본 논문에서는 그 이유를 다음의 3가지로 분석하였다.

1. Patch 기반 입력 방식의 한계: 이미지는 NLP와 달리 2D 구조와 Local 정보를 가지지만, Transformer에서는 이를 무시함
2. 저해상도 및 다중 스케일 특징 추출의 어려움: Transformer는 고정된 크기의 Patch를 사용하기 때문에 저해상도의 특징이나 다중 스케일 특징을 추출하는 것이 어려움
3. 높은 계산량 및 메모리 사용량: Self-Attention Module은 $O(N^2C)$의 계산량을 요구함. $O(NC^2)$의 계산량을 요구하는 CNN이 상대적으로 효율적

따라서 본 논문에서는 CNN과 Transformer가 결합된 형태의 모델 구조인 CMT를 제안한다. CMT의 주요 특징은 다음과 같다.

• 입력된 이미지는 우선 CNN을 활용한 Convolutional Stem을 거쳐 Fine-grained Feature를 추출한 뒤 CMT Block으로 전달된다.
• CMT Block은 기존 Transformer Block을 개선한 형태로, Depth-wise Convolution을 활용하여 Local 정보를 강화한다. → Stage 1에서 ViT 대비 더 높은 상태의 해상도를 유지할 수 있다(Dense Prediction에 유리).
• CNN과 유사하게 4개의 Convolution Layer를 사용하여 점차 해상도를 줄이고 Feature Dimension을 증가시키는 Stage-wise 구조를 사용하였다. → 다중 스케일 특징을 효과적으로 추출, 고해상도 처리로 인한 연산 부담 감소
• Local Perception Unit(LPN), Inverted Residual Feed-Forward Network(IRFFN)을 도입하여 중간 특징 내에서 Local 및 Global 정보를 동시에 포착하여 Representation 학습 능력을 향상시킴.
• ViT에서 사용한 Class Token 제거 → Average Pooling 사용
• CMT 모델을 다양한 크기로 확장할 수 있도록 단순한 스케일링 전략 제안

ViT에서 사용된 Class Token은 다른 196개의 Patch와 Self-Attention을 통해 정보를 교환하며 전역적인 특징 학습을 진행한다.

 

2. Related Work

대부분 아는 모델이기에 잘 모르는 내용만 서술

Convolutional Neural Network

GENet: Inter-channel 관계를 모델링하여 Feature map의 channel별 중요도를 동적으로 조정하는 매커니즘 도입(SENet도 같은 내용)

NLNet: Self-Attention 기법을 CNN에 통합하여, 모든 공간 위치간의 관계를 학습하도록 설계

 

Vision Transfomers

DeiT: Knowledge Distilation 전략을 도입하여 Transformer 기반 모델이 대규모의 Pretrain 없이도 실용적인 성능을 달성할 수 있음을 입증

T2T-ViT: 인접한 토큰들을 반복적으로 집계하여 하나의 토큰으로 변환

TNT(Transformer in Transformer): Patch-level, Pixel-level Representation을 모두 학습할 수 있도록 내-외부 Transformer Block 구성

PVT(Pyramid Vision Transformer): Pyramid 구조를 도입하여 다중 스케일 Feature map을 생성할 수 있도록 설계 → Dense Prediction 작업에서 강력한 성능

CPVT & CvT: Transformer Block에 Convolutional Projection을 적용

 

3. Approach

3.1. Overall Architecture

Figure 2(b)처럼 기존 DeiT 모델은 입력 이미지를 Non-overlapping Patch로 분할하여 Transformer Block에 입력한다. 때문에 Patch 내부의 정보를 효과적으로 모델링 하는 것이 어렵다. 

→ CNN 기반의 Stem Architecture를 도입하여 지역 정보를 보다 정교하게 추출!

 

Stem Architecture는 3단계 Convolution 연산으로 이루어져 있으며, 디테일한 구성은 논문을 참고.

여러 모델들에서 보이는 이 Stem Architecture를 가장 먼저 사용한 모델은 뭘까...

 

CMT는 최신 CNN 모델처럼 4-stage로 구성되어 있다. 각 Stage마다 Patch Embedding Layer가 존재하며 Feature Map 크기는 1/2로, Dimension은 2배로 확장시킨다. → 4-stage를 거치며 Multi-Scale Feature(4개의 서로 다른 해상도)를 학습!

 

각 Stage 내부에서는 여러 CMT Block을 순차적으로 배치하여 Feature 변환을 수행한다. 이 때 입력 해상도를 유지하면서 지역 및 전역 정보를 모두 학습한다.

 

CMT의 최종 분류는 Global Average Pooling Layer를 통해 이루어진다. 이후 Projection Layer를 통과하고 Softmax 계산을 통해 최종 결과를 출력한다.

 

3.2. CMT Block

제안된 CMT Block은 Figure 2(c)에서 볼 수 있듯이 Local Perception Unit(LPU), Lightweight Multi-head Self-attention(LMHSA), Inverted Residual Feed-forward Network(IRFFN) 3가지의 주요 구성 요소로 이루어져 있다.

 

(1) Local Perception Unit

Vision Task에서 Rotation과 Shift는 흔히 사용되는 Augmentation 기법이다. 일반적으로는 이러한 변환을 적용하여도 Translation-invariance 특징을 가지는데, Transformer 모델은 Positional Encoding을 활용해 Patch간의 순서를 유지하는 방식을 사용하기 때문에 Translation-invariance 특징을 가지지 못한다.

 

또한 Patch 내부의 Local Relation, Structure information을 효과적으로 학습시키기 어려우며, 이러한 문제를 해결하기 위해 논문에서는 Local Perception Unit(LPU)를 도입하였다.

$$LPU(x) = DWConv(x) + X$$

LPU는 지역 정보를 학습하기 위한 연산 방식으로 입력 $X$에 Depth-wise Convolution을 적용한 결과를 추가하여 지역 정보를 강화한다.

 

(2) Lightweight Multi-head Self-attention(LMHSA)

[기존 Self-attention]

$$Attn(Q, K, V) = Softmax(\frac{QK^T}{\sqrt{d_k}})V$$

[LMHSA]

$$LightAttn(Q, K, V)=Softmax(\frac{QK'^T}{\sqrt{d_k}}+B)V'$$

기존 Self-attention은 $QK^T$의 크기가 Patch 수 제곱이 되므로 입력 해상도가 커질수록 연산량이 $O(n^2)$로 증가한다. 본 논문에서는 이를 해결하기 위해 Lightweight Multi-head Self-attention(LMHSA)을 도입하였다.

 

이 모듈의 핵심은 Key와 Value의 해상도를 줄이는 것이다. Key와 Value의 해상도를 줄임으로써 연산량을 크게 줄이고, Multi-head Self-attention 방식으로 확장하여 여러개의 Head를 연결한 n x d 크기의 Sequence를 만든다.

1. 우선 K, V에 Depth-wise Convolution을 적용한다. → K', V'
2. 학습 가능한 Relative Position Bias $B$를 추가하여 학습된 위치 정보가 다른 해상도에서도 쉽게 조정될 수 있도록 구성

기존 Transformer의 Positional Encoding은 해상도가 변할 경우 문제가 발생하였지만 Relative Position Bias는 해상도 변화에 따라 동적으로 조정될 수 있음(Bicubic Interpolation 활용). → 다양한 해상도 입력에 대해 강건한 성능

 

(3) Inverted Residual Feed-forward Network(IRFFN)

[기존 ViT의 FFN]

$$FFN(x) = GELU(XW_1+b_1)W_2+b_2$$

[IRFFN]

$$IRFFN(X)=Conv(\mathcal{F}(Conv(X)))$$

$$\mathcal{F}=DWConv(X)+X$$

CMT는 기존 ViT의 FFN 구성을 개선하였다. MobileNetv2의 Inverted Residual Block을 기반으로 설계하였으며, Figure 2(c)를 통해 그 구성을 확인할 수 있다.

 

Convolution에 더해 Residual Connection을 도입한 이유는 깊은 네트워크에서도 Gradient가 효과적으로 전달되게 만들기 위함이다.

 

최종 CMT 블록을 정리하면 다음과 같다.

$$Y_i = LPU(X_{i-1})$$

$$Z_i=LMHSA(LN(Y_i))+Y_i$$

$$X_i=IRFFN(LN(Z_i))+Z_i$$

각 Stage에서 CMT Block을 쌓아 Feature Transformation, Feature Aggregation을 수행한다.

 

3.3. Complexity Analysis

(1) ViT 연산량

• $O(MHSA) = 2nd(d_k+d_v)+n^2(d_k+d_v)$
• $O(FFN)=2nd^2r$
• $r$: FFN의 확장 비율로 ViT에서는 4

→ $O(Transformer Block) = 12nd^2+2n^2d$

 

(2) CMT 연산량

• $O(LPU)=9nd$
• $O(LMHSA)=2nd^2(1+\frac{1}{k^2})+\frac{2n^2d}{k^2}$
• $O(IRFFN)=8nd^2+36nd$

→ $O(CMT Block) = 10nd^2(1+\frac{0.2}{k^2} + \frac{2n^2d}{k^2}+45nd)$

 

CMT는 ViT보다 연산량이 낮다!!!

 

3.4. Scaling Strategy

본 논문에서는 Transformer 기반 Network에 적합한 새로운 Compound Scaling 전략을 제안한다.

$$Depth : \alpha^{\phi},\quad Dimension : \beta^{\phi}, \quad Resolution : /gamma^{\phi}$$

이 때 $\phi$는 Compound Coefficient로 모델 크기를 조정하는 역할을 수행한다. 본 논문에서는 실험적으로 $\alpha = 1.2, \beta = 1.3, \gamma = 1.15$로 하이퍼파라미터를 설정하였다.

→ 계산량 증가와 성능 향상 사이의 균형

Compound Scaling은 EfficientNet에서 제안된 방법으로 Depth, Width, Resolution을 동시에 최적의 비율로 증가시켜 모델 성능을 높이는 것과 동시에 계산 비용을 최소한으로 증가시키는 방법이다.

기존에는 성능 향상을 위해 모델 크기를 확장시키기 위해 Depth 증가(깊은 네트워크), Width 증가(채널 수 증가), Resolution 증가(입력 이미지 확대) 중 한가지 방식만을 사용하였다. 하지만 이 경우 모델의 연산량과 성능이 불균형해지는 문제가 발생하여 이를 동시에 조정하는 Compound Scaling이 제안되었다.

 

본 논문에서는 이 Scaling Strategy를 활용하여 다양한 크기의 CMT 모델을 설계하였다(CMT-Ti, CMT-XS, CMT-S, CMT-B).

 

4. Experiments

본 논문에서는 CMT 아키텍처의 효과를 검증하기 위해 Image Classification, Object Detection, Instance Segmentation 등의 Task에서 실험을 진행하였다.

4.1. ImageNet Classification

ImageNet 데이터 셋 기반 Image Classification 실험에서 CMT 모델은 기존 Transformer 및 CNN 기반 모델들보다 더 적은 연산량으로 더 높은 정확도를 달성하였다. 특히 최근 Transformer 기반 모델들이 EfficientNet보다 낮은 성능을 보인 것과 달리 CMT는 그보다 높은 성능을 낼 수 있었다.

→ CMT Block의 Hybrid 구조가 효율적이며 Local + Global 정보를 효과적으로 학습한다는 것을 확인!

 

4.2. Ablation Study

Stage-wise Architecture

ViT, DeiT와 같은 기존 Transformer 기반 모델들은 Single-scale Feature map만을 생성할 수 있었다. 하지만 Dense Prediction Task에서는 Multi-scale 정보가 중요하고, 때문에 주로 CNN에서 활용하는 Stage 기반의 설계가 필요하다.

 

본 논문에서는 DeiT를 Stage기반의 아키텍처로 변경하여 기존 모델 대비 성능이 향상하는 것을 확인해, Stage-wise Architecture의 효과를 입증하였다(Table 3).

 

CMT Block

CMT 모델에서 LPU, IRFFFN과 같은 구성 요소가 성능에 미치는 영향을 분석하였다(Table 5). 그 결과 Stem, LPU, IRFFN, Shortcut 들이 성능 향상에 필수적인 것을 확인하였다.

 

Normalization Function

Transformer와 CNN은 일반적으로 다른 정규화 기법을 사용한다. Transformer는 주로 Layer Normalization을, CNN은 Batch Normalization을 사용하는데, CMT에서는 실험을 통해 Layer Normalization과 Batch Normalization을 적절히 조합하는 것이 좋은 성능을 내는 것을 확인하였다.

 

Scaling Strategy

기존 Scaling Strategy와 비교하여 본 논문에서 제안한 Compound Scaling Strategy의 효과를 분석하였다(Table 4). Compound Scaling을 사용했을 때의 성능이 더 좋은 것을 확인.

 

4.3. Transfer Learning

COCO Dataset을 사용해 Object Detection, Instance Segmentation Task에서도 모델 성능을 확인하였다. Object Detection에서는 RetinaNet을, Instance Segmentation에서는 Mask R-CNN 모델을 사용하였으며 Backbone을 CMT-S로 변경하여 그 성능을 확인하였다.

 

그 결과, Table 6, Table 7에서 볼 수 있듯이 두 Task에서 모두 기존 모델들 대비 좋은 성능을 보였으며, Inference Speed 또한 효율적인 속도를 달성하였다.

 

CIFAR 10, CIFAR 100 등 다양한 Transfer Learning Dataset에 대해서도 Fine-Tuning하여 평가를 진행하였으며, 그 결과는 Table 8에서 확인할 수 있다.

 

5. Conclusion

본 논문에서는 CMT(CNNs Meet Transformers)라는 새로운 Hybrid 아키텍처를 제안하였다. 이 아키텍처는 기존 Transformer 모델의 한계(지역 정보 부족, 고해상도 처리시 높은 연산비용)를 극복하고 CNN, Transformer 각각의 장점을 살려 효율적인 모델 아키텍처를 구성하였다. 또한 Compound Scaling Strategy를 도입하여 다양한 모델 변형을 생성하였다.

CMT의 아키텍처를 살펴보면, 초반 부분에는 CNN 관련 기법으로 구성되어 있으며, 후반 부분은 Transformer의 Attention + FFN 구조를 차용하였다. 더불어 CNN의 4-stage Multi-scale, Compound Scaling 기법을 활용해 성능을 고도화 한 것으로 보여진다. 공간 정보가 없는 Patch 대신 Stem과 CNN Block을 Transformer 이전에 통과시켜 그 정보를 유지시킨 것이 성능 향상에 포인트이지 않나 싶다.

---

세 줄 요약

1. CNN의 Local Information 학습 능력 + Transformer의 Long-term Dependency 학습 능력을 결합한 Hybrid 모델인 CMT(CNNs Meet Transformers) 모델 제안
2.  Stem, LPU(Local Perception Unit), LMSHA(Lightweight Multi-head Self-attention), IRFFN(Inverted Residual Feed-forward Network)와 같은 구성 요소를 도입하여 연산량을 최소화 하면서 높은 정확도를 유지함
3. Compound Scaling을 활용해 Depth, Width, Resolution을 균형 있게 확장하여 다양한 크기의 모델을 설계함