[논문리뷰] CoaT: Co-Scale Conv-Attentional Image Transformers

 

 

참조

 

[논문리뷰] AN IMAGE IS WORTH 16X16 WORDS:TRANSFORMERS FOR IMAGE RECOGNITION AT SCALE(Vision Transformer)

[논문리뷰] CPVT(CONDITIONAL POSITIONAL ENCODINGS FOR VISION TRANSFORMERS)

 

들어가며

 

이 글은 기본적인 CNN, Transformer, Vision Transformer(ViT)에 대한 사전 지식이 있다는 가정하에 작성되었습니다.

 

현재 컴퓨터비전에 CNN은 말할 것도 없고, Transformer 기반의 구조도 둘 다 활발하게 발전되고 있습니다. 생각해 보면 두 메커니즘은 공통점이 있습니다.

 

Transformer는 Query와 Key의 유사도를 구한 행렬과 대응되는 Value 행렬의 곱과 합으로 Attention Score를 구하고, CNN은 커널의 Weight와 대응되는 Feature map과의 곱과 합으로 새로운 Feature를 생성합니다. 이렇게 두 과정은 공통적으로 Weight Sum이 포함되어 있습니다. 

 

하지만, CNN은 학습에서 사용된 커널과 동일한 Weight로 Test를 합니다. 그래서 Feature에 따라 Weight가 Test에서도 유동적으로 변하는 Transformer 보다 일반화 성능이 떨어졌습니다. 하지만, 이런 Transformer에도 Attention Score를 구할 때, H * W 해상도의 Feature map이 곱해져야 해서, 어마무시한 계산량을 요구한다는 단점이 존재하였습니다.

 

그리하여 Transformer, CNN 둘 다 연구가 활발하게 이루어지고 있었고, 최근에는 DETR 이라는 CNN + Transformer의 논문도 나오고 있습니다. 따라서, 본 논문은 이러한 최신 동향에 힘입어 CNN과 Transformer를 결합한 Co-Scale Conv-Attentional Image Transformer(CoaT)를 제안하였고, 이 CoaT는 2가지 새로운 메커니즘을 도입하였습니다.

 

• Conv-Attention Module: 기존 Self-Attention 형태보다 계산적 효율성이 있는 모듈인 Factorized Attention 모듈을 새롭게 제안
  
• Co-Scale 메커니즘: 기존 Multi-Scale과 같이 다른 해상도의 Feature들이 Coarse-to-fine 형태로 서로 상호작용 하는 것을 넘어, Fine-to-Coarse의 형태도 상호작용하는 메커니즘.
  -> 직렬 블록 및 병렬 블럭으로블록으로 총 2가지 블록으로 구성

 

Conv-Attention Module

 

Revisit Self-Attention

들어가기 앞서 기존 Self-Attention의 동작 방식을 보면 아래와 같습니다.

 

 

Feature map \( F \in \mathbb {R}^{H*W*C} \)에 대한 Self-Attention을 진행한다고 했을 때, Query, Key, Value의 입력에 사용되는 토큰의 수는 Class 토큰을 포함하여 \( N = HW + 1 \)입니다. 이 연산은 \( O(N^2) \)의 공간복잡도와 \(O(N^2C) \)로 해상도에 대해서 이산적으로 계산 복잡도가 증가합니다.

 

그래서 ViT는 Feature maps을 16 * 16 Patch로 분할하는 방식을 취했는데, 이렇게 하면 Feature map 각각의 디테일한 부분은 학습하지 못한다는 딜레마에 빠지고 있었습니다.

 

본 논문은 이러한 딜레마를 극복하는 Factorized Attention Mechanism을 제안하였습니다.

 

Factorized Attention

이 Factorized Attention 연산 방식은 아래와 같습니다. 

 

 

저희가 아는 형태랑 많이 다릅니다. 이름 그대로 기존 Attention 연산에서 \( \frac {Q} {\sqrt {C}} \)를 앞으로 인수분해 해버렸습니다. 이렇게 하더라도 Q, K, V가 각각 상호작용하는 메커니즘 자체는 변하지 않는다고 주장합니다.

 

위와 같은 과정은 \( O(NC + C^2) \)의 공간 복잡도와 \( O(NC^2 \)의 시간복잡도를 가져 N에 대해 선형적인 계산 복잡도를 가지지만, 패치로 나누지 않아 기존의 딜레마를 극복하였습니다. 

 

Convolutional Relative Position Encoding

위 방법으로 Attention 메커니즘에 대한 딜레마를 극복하였지만, 추가로 입력해줘야 할 부분이 남았습니다. 이대로 Attention 연산을 진행하면 각 토큰에 대한 Position Encoding이 되지 않은 채로 진행되기 때문에, 각 토큰 위치에 대한 모호성이 생깁니다.

 

아래는 이 모호성에 대한 예시인데, \( L = softmax(K)^{T} V \in \mathbb {R}^{C*C} \)이고 \( Q \)에서 파생된 2개의 query vectors \( q_1, q_2 \in \mathbb {R}^C \)가 서로 같을 때라 가정한 수식입니다.

 

 

따라서, 위치 정보를 넣어줘야 하는데, 본 논문은 이 과정을 NLP에서 사용된 Relative Position Encoding(RPE) 방식을 선택하였습니다. 그래서 이 Attention을 넣어준 것을 아래와 같이 정의할 수 있습니다.

 

RPE는 자신의 토큰과 다른 토크들과에 대한 위치 정보에 대한 정보를, Embedding Vector으로 학습시켜 상대적인 관계를 모델링하는 것을 말합니다. 아래는 1차원 토큰에 대한 예시를 든 것이고 이에 따라 \( P \) 도 1차원으로 변합니다. 

 

 

저 위 \( w \) 아래 첨자인 -3, -2, -1, 0, 1은 현재 토큰을 기준으로 상대적인 거리에 대한 인덱스를 나타내는 것이고, 원문에서는 우리가 지정한 M 크기의 \( P \) 보다 더 긴 시퀀스가 들어올 수도 있기 때문에 아래와 같은 공식으로 너무 길어지거나 짧아지는 토큰에 대해서는 인덱스를 클리핑 하였습니다.

 

 

예컨대, 현재 토큰의 인덱스 \( i =4 \)이고 다른 토큰의 인덱스 \( j \in [1, 2, 3, 5] \) 일 때, M = 3이면 [-3, -2, -1, 0, 1] -> [-1, -1, -1, 0, +1]의 인덱스 결과가 나옵니다. 이에 따라 Position Encoding \(P \)는 아래와 같이 구할 수 있습니다.

 

 

이후, 이 쿼리 Q에 위에서 도출한 Relative Position Embedding \( P \)을 곱하여 \( E = QP \)를 출력하고, 여기에 벨류 \( V \)를 곱하여 최종적인 Relative attention map \(EV \in \mathbb {R}^{N * C} \) 를 도출합니다. 

 

이렇게 해서 이 \( E\)는 쿼리와 벨류 간, 관계를 모델링하기 위해 학습된다고 주장합니다.

 

접은 글에서 봤다시피 복잡한 수식으로 이루어져 있지만, 실제 논문에서는 이런 이론을 바탕으로 단순 Convolution으로 위 수식을 풀어서 사용했습니다. 위의 \( EV \)는 \(Q(PV) \)로 풀어서 나타낼 수 있고, P와 V 간의 관계를 Convolution을 사용하여 나타내면 아래와 같은 식이 나옵니다. 

 

 

여기서 Depthwise Conv를 사용하는데, 본 논문은 이 방식이 공간, 시간 복잡도에서 기존과 다른 차별점이라고 이야기합니다. 그리고 이제 이 \( EV \)에 대한 위치 정보를 Attention 연산에 삽입하는 과정은 아래와 같습니다.

 

 

먼저 \( EV \)를 구하고, Class 토큰을 포함하는 Attention에 대해 차원을 맞춰주기 위해 0을 concat 해줍니다. 최종적으로, 이 \( EV \)를 Attention 연산에 더해줍니다.

 

Convolutional Position Encoding

본 논문은 위의 Depthwise Conv를 Position 정보에 사용한 방법을 CPVT에서 한 것과 동일하게, Attention 연산을 수행하기 전 input features X 에도 삽입하였다고 합니다. 이리하여, Relative Position Encoding의 효과를 더 극대화시킬 수 있다고 주장합니다.

 

 

위와 같은 흐름들로 전체적인 구조는 아래와 같습니다. 

 

Conv-Attentional Module

 

위 Convolutional Position Encoding에서는 3의 커널 사이즈를 사용했고, Convolutional Relative Position Encoding에서는 각 서로 다른 Attention heads에 대해 커널 사이즈를 3, 5, 7을 가져갔습니다. 

 

최종 출력인 Output Feature Map에서 기존과 동일하게 FFNs를 사용합니다.

 

 

Co-Scale Conv-Attentional Transformer

처음에 기술했듯이, 본 논문은 Multi-Scale과 비슷한 Co-Scale 메커니즘으로 서로  다른 해상도의 Feature들이 Coarse-to-Fine 및 Fine-to-Coarse, 두 가지 형태로 상호작용을 한다고 하였습니다. 이에 따라 직렬 블록과 병렬 블록이 구성됩니다.

 

 

직렬 블록

 

먼저 전체 구조를 나타내면 아래와 같습니다. 

 

 

1. 특정한 비율로 Patch Embed
   -> 이 과정에서 차원이 축소되어 여러 Scale을 가짐(1/32, 1/16, 1/8, 1/4)
   
   
2. Flatten 하여 Sequence 형태로 변경
   
   
3. CLS 토큰 삽입
   
   
4. Conv-Attentional Module 적용
   
   
5. CLS 토큰 불할 및 Image 토큰 2D로 Reshape

 

위와 같은 단계들로, Conv-Attention Module이 적용된, 여러 Scale의 결과를 얻을 수 있습니다. 

 

 

병렬 블록

 

전체 구조는 아래와 같습니다. 

 

 

여기서 병렬 블록은 위의 직렬 블록의 각 Scale에서 출력한 결과들입니다. 각각의 병렬 블록들이 서로 상호작용을 위해 Scale이 Up, Downsampling이 이루어집니다. 이 과정으로 Corase-to-Fine, Fine-to-Corase에 대한 Attention이 가능해지고, 이를 Co-Scale 메커니즘이라고 합니다.

 

이 Co-Scale 메커니즘은 Direct Cross-Layer Attention 방법과 Featuer Interpolation 방법으로 나뉩니다.

 

Direct Cross-Layer Attention에 대한 과정은 아래와 같습니다:

 

• 먼저 각 병렬 블록에 대해 Conv-Attention이 이루어짐.
  
  
• 서로 다른 스케일에 대한 Cross-Attention 연산을 위해 K, V가 적절한 Up, Down sampling이 이루어짐
  
  
• 각 스케일에서 연산된 결과(Conv-Attention, 2 Cross-Attention)를 더하고, FFN으로 통합

 

다음으로, Feature Interpolation에 대한 과정은 아래와 같습니다:

 

• 각 병렬 블록에 대한 Conv-Attention이 이루어짐
  -> Direct Cross-Layer Attention과 같음
  
  
• 각 스케일의 Conv-Attention 결과들이 서로 다른 스케일의 크기에 맞춰 Up, Downsampling이 이루어짐
  
  
• 각 스케일의 피처들(다른 스케일 -> 현재스케일로 Interpolation 된 Up, Down 결과들 + 현재 스케일의 Conv-Attention 결과)을 더하고, FFN으로 통합

 

두 방법 다 최종적으로 출력된 각 스케일의 Class Token을 결합(Concat)하여, Class를 예측합니다. 그리고 경험상 Feature Interpolation 방법이 더 낫다고 하네요.

 

 

또한, 본 논문은 이러한 병렬 블록을 적용하지 않은 버전(CoaT-Lite)도 아래와 같이 나타내었습니다.

왼쪽: CoaT-Lite, 오른쪽: CoaT

 

 

마치며

 

 

위는 COCO에서의 Object Detection 결과인데 성능이 잘 나왔으며, Co-Scale을 적용하였을 때 성능이 월등히 높은 것을 확인할 수 있습니다. 

 

그리고 아래를 보면 Convolution을 적용한 Relative Position Encoding의 효과를 볼 수 있습니다.

 

 

 

오늘은 CoaT 논문을 리뷰해 보았습니다.

긴 글 읽어주셔서 감사합니다.