Introduction
Transformer는 다양한 NLP 태스크에서 이제 지배적으로 사용되고 있습니다. 비전 태스크에 처음으로 Transformer를 성공적으로 적용한 것으로 알려진 ViT는 이미지를 서로 겹치지 않는 16x16 패치로 잘라 이 패치를 토큰으로 사용합니다. Transformer의 성공에도 불구하고 아직까지 그 성능은 적은 양의 데이터에 학습되었을 때 비슷한 크기의 CNN에 미치지 못하고 있는데요, 이는 아마도 ViT가 비전 문제를 해결하는데 특화된 CNN의 구조를 갖지 않기 때문일 것입니다. 본 논문에서는 convolution을 ViT 구조에 적용함으로써 성능과 robustness를 높이고자 합니다.
위 표는 대표적인 Vision Transformer 구조와 본 논문에서 제안하는 CvT를 비교한 것입니다. CvT는 convolution 구조를 차용한 덕분에 PE가 필요하지 않다는 장점을 가집니다.
Proposed Method
CvT의 전체적인 구조는 아래와 같이 CNN과 같이 multi-stage(3-stage)로 구성된 convolution-based transformer 연산으로 이루어져있습니다. 각각의 스테이지를 지나면서 마치 stride를 가진 convolution과 같이 token map (CNN의 feature map에 해당)의 spatial resolution은 점차 줄어들게 됩니다.
- Convolutional Token Embedding
각 스테이지의 첫 번째 모듈인 convolutional token embedding 연산은 local spatial contexts를 모델링하기 위해 만들어졌습니다. 이 연산의 입력은 2D 이미지 혹은 이전 레이어의 2D output token map 입니다.
Stage 에서는 에서 새로운 토큰인 로의 매핑을 학습합니다. 이 함수는 kernel size = s x s, stride = s - o, 그리고 padding = p인 convolution 입니다.
새로운 token map 은 다음과 같은 크기를 가집니다.
은 크기로 flatten된 뒤 layer normalization으로 정규화됩니다.
- Convolutional Projection for Attention
이 연산은 Transformer에 존재하는 Multi-Head Self-Attention (MHSA)의 original position-wise linear projection을 depth-wise separable convolutions으로 치환한 것입니다.
먼저, 토큰들은 2D token map으로 변환됩니다. 그 다음, 커널 사이즈 s를 가지는 depth-wise separable convolution이 적용됩니다. 마지막으로 토큰은 1D로 flatten되어 다음 연산에 적용됩니다.
연산을 도식화하면 다음과 같습니다.
커널 사이즈가 1이라면 이 Convolutional Projection은 원래 Transformer의 projection layer와 같게 되므로, 이는 원래 연산의 generalization이라고 볼 수 있습니다.
Experiment
- Setup
평가를 위해 이미지넷-1k, 22k를 사용했습니다.
이미지넷-22k에 pre-training된 모델은 CIFAR-10/100, Oxford-IIIT-Pet, Oxford-IIIT-Flower 데이터셋에 각각 fine-tuning 되었습니다.
- Model variants
다양한 크기의 모델로 실험하기 위해 Transformer block 갯수를 바꿔가며 아래와 같이 실험했습니다. CvT-X는 Transformer block이 X개 사용되었음을 의미합니다. "W"는 채널 갯수를 늘린 wide 모델을 의미합니다.
- Comparison to SOTA
결과는 아래와 같습니다. ImageNet, ImageNet-Real, ImageNet-V2 데이터셋에 각각 비교하였습니다. 기존 Transformer 모델에 비해 더 적은 파라미터와 FLOPs를 가지고도 더 높은 성능을 얻었습니다.
- Downstream task transfer
파라미터 수가 훨씬 많은 BiT 모델에 비해서 downstream task에서 더 높은 transfer 성능을 보여주었습니다.
- Ablation study - Position Embedding
CvT에는 convolution을 통해 local context가 자연스럽게 embedding되었기 때문에 Position Embedding이 필요치 않습니다. 아래는 PE 유무에 따른 비교 결과입니다. DEiT와는 다르게 CvT는 PE가 없어도 동일한 성능을 보여줍니다.
- Ablation study - Convolution Token Embedding
Convolution Token Embedding을 non-overlapping patch embedding으로 변경하였을 때 성능 감소 효과와 그에 따른 PE의 효과를 분석한 결과입니다.
- Ablation study - Convolutional Projection
아래 실험 결과는 Convolutiona Projection에서 key, value projection에 사용하는 convolution의 stride를 1 또는 2로 하였을 때의 연산량고 성능을 비교한 것입니다. Stride=2로 하였을 때 30%의 연산량 감소 대비 작은 성능 감소가 있어 stride=2를 기본값으로 사용하였습니다.
다음 결과는 Convolution Projection의 Linear projection 대비 효과를 측정한 것입니다. 모든 stage에 적용하였을 때 가장 높은 성능을 보였습니다.
Conclusions
본 논문에서는 Vision Transformer에 convolution을 적용하여 Transformer와 convolution의 장점을 모두 취할 수 있는 방법을 연구하였습니다. 실험을 통해 convolution token embedding과 multi-stage 디자인에 적용된 convolutional projection이 연산량을 적게 가져가면서도 높은 성능을 얻을 수 있다는 것을 보였습니다. 또한, convolution을 적용함으로 인해 내재된 local context structure 덕분에 positional embedding은 더이상 필요하지 않게 되었으므로 다양한 입력 해상도를 요구하는 여러 비전 태스크에서 활용성이 높아졌습니다.
Share article