brunch

라이킷 5 댓글

You can make anything
by writing

C.S.Lewis

계정을 잊어버리셨나요?

by Qscar Feb 01. 2024

Convolutional ViT(CvT)

Paper Review 6 : Modeling, Training + a

|Intro

이번에 리뷰할 논문은 Convolution의 개념을 ViT에 주입한 'CvT: Introducing Convolutions to Vision Transformers'입니다.

ViT Paper에서 제시하는 Hybrid Architecture에 대한 설명

이는 ViT 논문에서도 제시됐었던 resnet 혹은 간단한 CNN 기반 레이어를 이용해 패치 임베딩을 실시하는 Hybrid 구조에 대한 개념에서 더욱 진화한듯한 개념입니다. 이전에 리뷰한 swin과 같은 해에 공개된 논문이지만, swin보다 먼저 제시된 개념이고 계층적 구조와 같은 개념들에 대해선 앞선 swin 포스팅들에서 충분히 리뷰했기에 패스하고 넘어가려던 논문이지만, 혹시나 하고 논문을 읽어보니 제법 재미있는 내용들이 포함돼 간단하게라도 리뷰하고자 합니다. 추가적으로 직접 테스트 결과 작은 사이즈(224 이하)에 대한 학습 및 성능에서는 간단한 대체 및 추가 코드만으로 swin 이상의 성능을 보이기도 했습니다.

ImageNet Benchmark - SwinV2-G #16

ImageNet Benchmark - CvT-21 #261

큰 사이즈의 이미지에서는 테스트해보지 않았지만, 아무래도 swin이 우세할 듯 싶어 Benchmark를 찾아보니 위 그림과 같았습니다. ImageNet 데이터셋에 대한 분류 문제에서 swin 모델의 경우 90.17%로 16위에 랭크돼있는 반면, CvT 모델의 경우 84.9%로 261위에 랭크돼있습니다. 다만 파라미터는 swin이 3,000M, CvT는 32M으로 약 100배 정도의 파라미터 수 차이가 존재했습니다.

본 포스팅의 주제인 CvT는 이전에 포스팅한 ViT 및 Swin에서의 내용과 많이 유사합니다. 다만 비효율적으로 생각되는 부분도 존재하는데요, 때문에 본 포스팅에선 논문에서 제시한 바를 있는 그대로 재현하는 것보다는 이전 포스팅에서 다뤘던 개념들을 응용해 보다 효율적으로 작동할 수 있는 구조로 구현하는 것에 초점을 맞출 것입니다.

|Paper Summary

CvT의 전체 구조

본 논문에 적용된 방법론들은 크게 두 가지입니다. 하나는 Convolutional Token Embedding이고, 다른 하나는 Convolutional Projection입니다. 간단히 설명하면, Convolutional Token Embedding은 기존의 ViT에서 이미지를 단순하게 지정된 사이즈의 패치로 분할하는 것 대신, Convolution Layer를 통해 패치화하는 것입니다. 또한 Convolutional Projection은 이렇게 패치화된 데이터에 대해 linear projection하기 전 Convolution Layer를 이용해 Projection한 후 넘겨주는 것입니다(정확히는 Separable Depthwise Convolution을 적용하는 것인데, 코드상으로는 기존 Linear Projection 전에 Depthwise Convolution을 추가한 것뿐처럼 보입니다).

이러한 구조를 응용해 CvT에서는 크게 3단계로 구성된 3 Stage Hierachical Structure를 구현했고, 가로세로 2D로 구성된 이미지 데이터에 대해 보다 적절한 처리를 통해 파라미터 효율적인 모델 구조를 구현했다고 제시하고 있습니다.

|Modeling

CvT의 구조는 ViT의 구조에서 Patch Embedding을 convolutional token embedding으로 대체하고, 기존에 선형으로만 이뤄지던 qkv projection의 앞에 convolutional projection을 추가하는 방식입니다(기존의 Linear Projection을 Convolution Projection으로 대체하는 것인데, 이때 사용되는 Separable Depthwise Convolution이 Depthwise Convolution+Linear의 결합으로 구현될 수 있기에 이렇게 표현했습니다, 이에 대해선 후에 다시 살펴보도록 하겠습니다). 이때 Convolution Layer의 padding, kernel size, stride를 조정해 스테이지가 반복될수록 계층적인 구조(가로세로는 줄어들지만, 채널이 증가)가 되며 추상화 수준이 높아지는 식입니다. 때문에 전체적인 모델링을 진행하기에 앞서 핵심 요소들을 먼저 구현하고, 이에 대해 데이터의 형상(shape)이 적절하게 변화하는지 확인해보겠습니다.

참고로 직접 구현하는 과정에서 본 논문과 다른 부분이 존재하는데요, 첫 번째는 클래스 토큰입니다. 우리는 이전 swin 포스팅 과정에서 굳이 class token을 사용하는 것보다 마지막 레이어를 global average pooling한 값을 이용하는 것이 더 성능적으로 우수함을 확인했기 때문에, 본 논문과 달리 여기서는 class token을 활용하지 않습니다.

두 번째로 다른 부분은 stage별 head의 개수입니다. 이전에 transformer 리뷰 당시 transformer 구조의 비선형성을 대표하는 부분이 바로 multi-head attention이라고 리뷰한 적이 있었습니다. 다만 본 논문에서 제시한 구조에서는 스테이지별 헤드의 개수가 [1,3,6]으로 첫 번째 스테이지에서 사용되는 head가 하나뿐입니다. 이는 충분한 비선형성 내재했다고 보기 어렵다 판단해, 하나 늘려 스테이지별 헤드 개수를 [2,3,6]로 적용함으로써 Multi-head Attention 구조에서 오는 비선형성을 보완하였습니다.

|Main Components

여기서 구현할 핵심 요소 두 가지는 위에서도 언급한 Convolutional Token Embedding과 Convolutional Projection입니다. 이 두 가지는 각 스테이지의 제일 앞에 위치하며, 그 이후의 여러 과정을 거치더라도 형상이 바뀌지 않습니다. 때문에 이 두 가지 요소들을 구현하고, 이 두 가지 요소들로만 구현된 스테이지를 구현해보도록 하겠습니다. 이 부분에 대한 코드는 여기를 참조하시면 됩니다.

|Convolutional Token Embedding

Convolution Token Embedding에 대한 설명

Convolution Token Embedding의 대상은 2D Image 혹은 2D 형태로 reshape된 Token Map입니다. 2D Image는 우리가 최초의 입력으로 사용하는 데이터이고, token map은 각 스테이지의 결과물입니다. 위 그림에서 새로운 가로세로 길이에 대한 공식은 convolution 공식인데요, CvT에서는 이를 이용해 적절한 stride, padding, kernel size 등을 조합해 hierachical한 구조(width x height = sequence length는 줄이면서 channel은 늘리는 구조)를 구현합니다.

이를 코드로 구현하면 다음과 같습니다.

Convolutional Token Embedding

위 코드는 입력된 2D reshape Image/Token Map에 대해서 Convolution을 적용한 뒤, normalization 및 형상 변환을 취하고 있습니다. 이때 stride를 적용해 각 패치(kernel)가 겹치도록 함으로써 패치 간의 관계와 지역성, 상대적 위치를 고려하게 됩니다. 이로 인해 CvT에서는 별도의 Positional Embedding 혹은 Positional Encoding을 적용하지 않습니다.

|Convolutional Projection

CvT에서는 Squeezed Convolution Projection을 사용합니다. 이는 Q, K, V를 projection할 때, K와 V의 크기를 Q보다 작은 사이즈로 projection해 Attention Score 등을 계산하는 방식을 위해 적용되는 방법입니다. 이는 K와 V의 Sequence Length만 같으면, Attention Score의 결과로 출력되는 Attention Map 크기는 동일하기 때문에 가능한 일입니다. 예컨데 Query가 4x4 feature map으로 나와 길이 16의 sequence가 되고, Key와 Value가 2x2 feature map으로 나와 길이 4의 sequence가 되더라도 이를 연산하면 (16x1)x(1x4)=(16x4)가 되었다가, (16x4)x(4x1)=(16x1)→[reshape]→(4x4)가 되게 됩니다. 이는 Convolution을 이용한 Projection의 결과를 이용한 것으로, K와 V의 사이즈를 줄이더라도 영역성과 상대적 위치를 잘 반영할 수 있다는 특징으로 인해 보다 효율적인 연산이 가능해지는 근거가 됩니다.

본 논문에서 제시하는 Convolutional Projection의 구조

본 논문에서는 Convolutional Projection을 구현하기 위해 흔히 MobileNet 등에 사용되는 Depth-wise seperable convolution을 사용한다고 언급하고 있습니다. 이는 Depth-wise Convolution 이후에 Point-wise Convolution을 적용하는 것으로, 이를 통해 더 적은 파라미터를 통해 더 효율적으로 연산이 가능해지게 됩니다.

구체적으로 Depth-wise Convolution은 input dim과 group의 수가 동일한 Convolution 연산 방식으로, 채널 별로 독립적인 컨볼루션 필터가 적용돼 공간적 특징을 추출하는데 초점을 맞춥니다. 이후에 진행되는 Point-wise Convolution은 1x1 컨볼루션을 이용해 각 위치에서 모든 채널의 출력을 결합해, 채널 간의 관계를 학습하고 출력 feature map의 채널 수를 조절(여기선 동일한 차원을 유지합니다)하게 됩니다. 이러한 두 개의 컨볼루션 레이어를 결합함으로써 입력된 2D token map의 공간적 정보와 채널 정보 간의 복잡한 관계를 효과적으로 학습해, 파라미터 효율적인 모델을 구현해낼 수 있는 것입니다.

1x1 Convolution의 경우, 단일 선형계층으로 대체되기도 하는데요 이는 두 코드의 성질이 거의 유사하기 때문입니다. 실제로 초기 가중치를 동일한 값으로 초기화하고, 두 레이어의 출력을 비교해보면 다음 그림과 같이 동일한 값을 출력하게 됩니다.

동일한 값으로 초기화된 1x1 conv와 linear의 결과 비교

물론 완전히 같지는 않습니다. 다만 실행속도에서는 다소 차이가 있는데요, 다음 그림에서 볼 수 있다시피 직접 실험한 결과에선 1x1 Conv가 Linear보다 실행속도가 더 빨랐습니다.

1x1 Conv와 Linear의 실행속도 비교

Convolution Layer의 경우, 최초 실행 시에는 속도가 조금 느리지만 반복적으로 실행될수록 속도가 더욱 빨라지며 실질적으로 두 방법론에 따른 파라미터 수 차이나 vram 메모리 차이는 없습니다. 이에 대한 코드를 살펴보고 싶다면, 여기를 눌러 확인하실 수 있습니다.

본 논문에서는 이러한 Convolution Projection을 하는 과정에서 squeeze 개념을 추가로 적용합니다. 이는 key와 value의 크기를 줄여 압축적으로 공간 정보를 표현하고, 압축된 매트릭스 상에서 attention score를 계산하는 방법입니다. 이에 대한 개념을 본 논문에서는 아래 그림과 같이 표현하고 있습니다.

(b) Convolution Projection과 (c)Squeezed Convolution Projection

위 그림을 보면 stride=1로 진행하던 query와 달리, key와 value를 projection할 때는 stride=2를 적용함으로써 훨씬 작은 차원의 매트릭스로 연산하게 됩니다. 이를 통해 key와 value에 대한 토큰 개수는 1/4로 줄어들게 되며, 이는 곧 연산의 감소로 이어지게 됩니다. 물론 이러한 과정에서 성능 하락 문제가 존재할 수 있지만, 이미지 데이터는 근접한 픽셀과 유사한 값을 가지는 특성이 있기 때문에 이러한 성능 하락은 최소화(minimal performace penalty)되며, Convolutional Projection을 통해 보완될 수 있습니다.

이를 코드로 구현하면 아래와 같습니다.

Squeezed convolutional attention code

위 코드는 2차원 token map으로 형상 변환을 한 뒤, Squeezed Convolutional Projection을 적용하고, attention score를 계산하는 코드입니다. 실제 공식 코드에서는 pointwise의 적용을 Linear Layer를 이용하였으나, 위에서 테스트한 바와 같이 1x1 conv를 그대로 적용하는 것이 성능측면에선 문제가 없으면서도 속도 측면에서 유리하며 논문 구현이라는 관점에서 봤을 땐 위와 같이 작성하는 것이 직관적이기 때문에 위와 같이 구현했습니다. 성능의 차이는 없습니다.

|Stage별 입출력 데이터 형상 확인

Convolutional Embedding과 Squeezed Convolutional Projection이 모두 정의됐다면 이를 불러와 적용할 수 있는 간단한 Block을 구현해보도록 하겠습니다. 이에 대한 코드는 다음과 같습니다.

attention block code

CvT의 전체 구현을 위해선 이러한 Block을 이용한 ViT Stage를 구축하고, 이를 포괄적으로 사용하기 위한 CvT 클래스를 정의해야 합니다. 이 과정을 위해 MLP와 같은 레이어 등이 추가되어야 하구요. 하지만 이러한 과정을 통해 최종적으로 출력되는 데이터의 형상은 바뀌지 않으며, 여기선 이해를 돕기 위해 아주 간단한 구조로 구현했습니다.

위 코드를 통해 스테이지 별로 입/출력되는 데이터의 형상을 살펴보면 다음과 같습니다.

Stage 1의 입출력 데이터 형상

첫 번째 스테이지의 결과 채널은 3에서 64로 늘어났고, 공간차원은 224x224에서 56x56으로 4배씩 줄어들었습니다.

Stage 2의 입출력 데이터 형상

두 번째 스테이지의 결과 채널은 유지되었고, 공간차원은 56x56에서 28x28로 두 배씩 줄어들었습니다.

Stage 3의 입출력 데이터 형상

마지막 세 번째 스테이지의 결과 채널은 유지되고, 공간차원은 28x28에서 14x14로 다시 두 배씩 줄어들었습니다. 즉, CvT는 채널을 늘린 상태에서 공간차원을 줄여가며 hierachical structure로 진행된다는 것을 확인할 수 있습니다.

|Full Modeling

본 논문에서 제시하는 CvT의 유형은 총 세 가지가 있으며, 각 스테이지별 반복수(정확히는 Transformer Block이 반복되는 횟수)를 합친 값을 뒤에 붙이는 식으로 간단하게 표현하고 있습니다.

Architecture of CvT Variants

위의 사항을 구현하기 위해 추가로 정의할 사항은 Transformer Block을 포함하는 Transformer 클래스와 MLP, 그리고 이들을 가져와 예측을 위한 Head를 붙이는 Model 클래스가 필요합니다. 또한 ViT 포스팅 과정에서 확인한 ViT의 성능과 비교하기 위해 ViT에 추가적으로 적용했던 LayerScale을 추가적으로 적용하고, 모델의 빠른 연산을 위해 GeLU 활성화 함수를 빠르게 구현한 QuickGeLU 함수를 적용토록 하겠습니다.

|Layer Scale & MLP in Transformer Block

우선 LayerScale입니다. LayerScale은 ViT 포스팅에서 확인했다시피 트랜스포머 블록에 아주 작은 인자를 곱함으로써 트랜스포머 블록의 성능을 다소간 향상시키는 기법입니다. 이에 대한 자세한 설명은 ViT 포스팅 중 여기를 참조하시면 됩니다. 이에 대한 코드는 다음과 같습니다.

LayerScale Code

MLP에 대한 코드는 본 논문에서 별도의 언급이 없었고, 그저 ViT를 최대한 활용한다는 식의 언급만 있었기에 ViT에서 사용한 MLP 구조를 그대로 구현합니다. MLP는 굳이 별도로 구현하지 않고, Transformer Block 내부에 한 번에 구현하였습니다. 코드는 다음과 같습니다.

Transformer Block의 구조

이 Block의 구조는 ViT에서 구현한 것과 같습니다. 그저 Attention을 위한 코드만 바뀌었다고 보시면 됩니다. Normalize 방식 또한 swin v2에서 제시한 post-res norm 대신 ViT와 Swin V1에서 제시된 구조인 pre-norm을 그대로 적용하였습니다.

|Quick GELU

다음으로는 Quick GELU입니다. Transformer를 이용한 모델 구조를 구현할 때, 가장 많이 활용되는 활성화 함수 중 하나가 바로 GELU인데요, 이는 GELU가 ReLU나 Swish와 같은 다른 활성화 함수에 비해 복잡한 문제를 푸는데 더욱 특화돼 있기 때문입니다.

ReLU, Swish, GeLU의 그래프 비교

예컨데 ReLU는 아주 간단한 연산으로 빠른 속도를 자랑하지만, 음수 입력을 완전히 배제해 정보 손실과 그레디언트 소실 문제가 발생했습니다. Swish의 경우 시그모이드 함수를 이용해 음수 입력에 대해서도 작은 값이 활성화됨으로써 ReLU의 단점을 보완하였습니다. 그리고 GELU의 경우, 입력값을 정규분포의 누적분포함수(CDF)를 사용해 변환함으로써 더욱 복잡한 문제를 풀 수 있게 되었지만, 그만큼 연산량이 늘어난다는 단점을 가지고 있습니다. 이러한 단점을 보완하기 위해 GELU의 그래프를 모방해 빠르게 연산하면서도 GELU의 장점을 그대로 가지게 됩니다.

GELU와 Quick GeLU 그래프 비교

실제로 Quick GELU와 일반 GELU의 그래프는 위 그림과 같이 약간의 차이는 있지만 거의 동일합니다. 그렇다면 성능 차이는 어떨까요? 성능 차이를 비교하기 위해 큰 차원의 텐서를 이용해 천 번씩 연산한 시간들을 비교해보도록 하겠습니다.

GELU와 Quick GELU의 속도 비교

두 번에 걸쳐 비교한 결과 약 12배 정도의 속도 차이가 난다는 것을 확인할 수 있습니다. 이러한 이유로 여기서는 Quick GELU를 도입해 적용토록 하겠습니다. Quick GELU 구현 코드는 아래와 같습니다.

Quick GELU code

이렇게 구현한 함수를 활성화 함수 대신 적용하면 됩니다.

|Vision Transformer

이제 이렇게 구현한 클래스들을 한데 묶어 Transformer 구조를 완성시킵니다.

Transformer

위 코드 또한 ViT 코드의 구조와 동일하며 ViT에서 Patch Embedding으로 패치분할하던 구조를 Convolutional Embedding으로 교체하고, Convolutional Projection이 적용된 복수의 transformer block으로 구성돼 있습니다.

|CvT Model

마지막으로 위에서 정의한 Transformer를 이용해 Stage별로 반복하면서 최종 출력값을 입력받아 평균을 취하고, 클래스 숫자로 분류하는 헤드를 추가한 모델 전체입니다.

CvT Model Code

이렇게 구현한 모델에 CvT-13의 파라미터를 입력하면 총 파라미터 수는 다음과 같습니다.

CvT-13 구현체의 파라미터

본 논문에서 제시하는 CvT-13의 파라미터 수는 총 19.98M입니다. 위 그림에서 볼 수 있듯 직접 구현한 모델의 파라미터 수는 이보다 적은데요, 본 논문과 달리 class token을 적용하지 않은 결과로 보입니다. 이는 약 85M의 파라미터를 가졌던 ViT에 비해 각각 1/4 정도의 파라미터 수입니다.

|Training

모델의 학습 과정 또한 ViT와 동일한 옵션을 활용합니다. ViT와의 차이라면 옵티마이저를 AdamW로 하였다는 것 정도입니다. 때문에 ViT와 같은 옵션이라고 보기보다는 Swin과 같은 옵션이라고 보는 것이 더 적절합니다. 때문에 학습과 관련된 데이터 증강, 세팅은 이전에 포스팅한 Swin 게시물을 보시거나, 제 github에서 코드로 보시는 것을 추천드립니다. 스케줄러는 cosine scheduler with warmup을 그대로 사용하였고, 200epoch, 500epoch 학습하였습니다.

200epoch 학습결과(좌)와 500epoch 학습결과(우)

이를 과거 swin의 성능과 비교해보기 위해 swin 포스팅에서 사용했던 지표를 그대로 가져와 확인해보겠습니다. 아래 그림들은 swin v1과 v2를 각각 1,000에포크 학습시킨 결과입니다. 참고로 학습에 사용한 swin v1 Tiny의 파라미터수는 약 2,827만 개, v2 Tiny는 약 2,764만 개로 CvT보다 40% 정도의 파라미터가 많은 모델들입니다.

swin v1 1,000epoch 학습결과(좌), swin v2 1,000 epoch 학습결과(우)

모델을 학습시키며 현재 다음에 리뷰할 ConvNext V1 논문을 읽고 있는데요, 컴퓨터 비전 영역에서 꾸준히 사랑받고 있는 ConvNet을 현대화(modernize)시키기 위해 ViT와 Swin 등에 사용된 학습 방법이나 트랜드를 ResNet에 적용함으로써 Swin 이상의 성능을 내는 것을 확인할 수 있었습니다. 모델의 구조 자체가 복잡하거나 어렵지 않음에도 성능향상이 쉽게 이뤄져 다음 포스팅을 통해 ConvNext를 단순 구현하는 것뿐 아니라, 본 포스팅에서 구현한 CvT에도 이를 적용해 더욱 현대화시켜 성능을 테스트해보도록 하겠습니다.

|Modernizing CvT

ConvNext V1 논문을 읽으며 Convolution의 현대화라는 개념을 적용함으로써 모델 성능을 높일 수 있다해 저도 이를 반영해 학습시켜보았습니다. 구체적으로 적용한 개념은 다음과 같습니다.

1. swin과 같이 hierachical한 총 네 개의 스테이지로 구성하고, 각 스테이지의 비율을 [1:1:3:1]로 맞춤

2. 더 큰 kernel size(7)을 사용 : Convolutional Projection 과정에 사용

3. inverted bottleneck : 기존 MLP에만 사용했던 구조를 Convolutional Projection 과정에 적용

4. Batch Normalization 대신 Layer Normalization 사용

이 외의 내용에 대해서는 이미 개념적으로 적용돼 있거나, 적용시 효율이 좋지 않았습니다. 이에 대한 결과는 다음과 같습니다. 우선 파라미터 수입니다.

CvT Next의 파라미

해당 모델의 파라미터 수는 13.61M으로 CvT 대비 2/3 수준이며, Swin과 비교해서는 절반도 되지 않는 수준입니다. 이에 대한 결과입니다. 좌측은 CvT를 1,000 에포크 학습시켰을 때 가장 높았던 성능, 우측은 CvT Next를 적용해 1,000에포크 학습시켰을 때 가장 높았던 성능입니다.

CvT 1,000 에포크 성능(좌)와 CvT Next 1,000에포크 성능(우)

CvT는 아쉽게도 swin v1에 비해선 손색이 있는데요, 파라미터 수를 조금 늘려 v1 수준에 맞추면 성능이 높아지지 않을까 예측해봅니다. 반면 CvT Next는 파라미터수가 줄어들었음에도 모델 성능이 꽤 높아졌습니다. 여태까지 학습한 모델 중 가장 성능이 높았던 swin v1과 같은 수준입니다. 이에 대한 코드는 여기를 살펴보시면 됩니다.

|Reference

[1] Haiping Wu, Bin Xiao, Noel Codella, Mengchen Liu, Xiyang Dai, Lu Yuan, Lei Zhang. CvT: Introducing Convolutions to Vision Transformers. https://arxiv.org/abs/2103.15808

[2] Microsoft. CvT. https://github.com/microsoft/CvT