[논문리뷰] LocalViT: Bringing Locality to Vision Transformers

참조

 

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

 

들어가며

 

본 논문은 ViT를 기반으로 작성된 논문입니다.

 

수용필드가 Global 한 Transformer 구조를 적용한 ViT는 컴퓨터 비전에서 CNN과 대등한 성능을 가지고 있었습니다. 하지만, 본 논문에서는 Global 수용필드를 가지는 Transformer의 구조상, 객체의 라인이나 엣지 모양등의 Locality는 부족하다고 주장합니다. 

 

여기서 Locality는 커널을 Sliding window 방식으로 움직이는 CNN에 내장된 것입니다. 그래서 이 Global 수용필드를 가지는 Transformer 구조에 Locality를 내장하면 어떨까 하고 이 연구가 시작되었습니다.

 

본론부터 이야기하자면, 이 Locality는 ViT의 Feed-Forward Network(FFN) 부분에 삽입됩니다. FFN에 삽입된 Locality의 검증은 FNN과 Inverted Residual을 비교함으로써 발견하였습니다. 이 Inverted Residual은 MobileNetv2에 적용된 구조이므로 추후에 다루도록 하겠습니다. 

 

여하튼, FFN에 Locality를 삽입한 메커니즘으로 ViT를 기반으로 한 여러 모델에 적용했을 때 우수한 성능을 기록했다고 합니다. 이 Locality 메커니즘의 자세한 내용은 아래에서 살펴보도록 하겠습니다. 

 

Methodology

 

본 논문이 ViT 구조를 따라갑니다. ViT는 Self-Attention과 FFN으로 구성되는데 이 FFN에 Locality를 삽입한다고 앞서 말씀드렸습니다. 먼저, 기존 ViT은 입력과 출력이 같은 Squence-to-Sequence 방식이며 아래와 같이 동작합니다.

 

Sequence 관점

입력 이미지 \( X \in \mathbb {R}^{C*H*W} \)가 패치로 나뉘며, 토큰들의 시퀀스로 변환됩니다{\( \hat {X}_i \in \mathbb {R}^d | i=1,..., N \)}. 여기서 \( d = C * p^2 \)이고, \( N = \frac {HW} {p^2} \)입니다. 이러한 토큰들을 Linear layer를 거쳐 Patch Embedding \( \hat {X} \in \mathbb {R}^{N*d} \)가 형성됩니다.

 

이 \( \hat {X} \)는 Self-Attention 연산 후, FFN로 입력됩니다. 먼저 Self-Attention 연산 아래와 같이 이루어지고, 

 

 

FFN의 수식을 나타내면 아래와 같습니다. 

 

 

여기서 \( W_1 \in \mathbb {R}^{d * \gamma d} \)이고, \( W_2 \in \mathbb {R}^{\gamma d * d} \) 이며, \(f() \)는 Activation 함수입니다. Self-Attention 연산으로 학습된 feature representations을 더 풍부하게 하기 위해, \( \gamma \) 만큼 확장하는 방식을 채택합니다(보통 확장 비율은 4). 

 

 

이상 기존 방식인 Sequence 관점으로 나타내었는데, 전술했듯이 본 논문은 FFN과 Inverted Residual을 비교하였습니다. Inverted Residual은 CNN이 적용되므로, 격자 관점도 살펴보겠습니다. 

 

격자 관점

어렵게 생각할 거 없이 MLP를 적용한 FFN에 CNN을 적용하면 아래와 같은 그림이 나오고, 이 구조가 MobileNet V2에서 사용한 Inverted Residual과 똑같이 생겼습니다.

 

이 과정은 먼저 Sequence에 Convolution을 적용해야 하기 때문에, Sequence를 Image로 Reshape하는 것으로 시작합니다. 

 

 

Seq2Img 안의 \( Z \)는 Sequence 관점에서 사용한 Self-Attention의 출력과 동일한 것입니다. 여기서 \(h = H/p \)이고 \( w W/p \)이고, 결과적으로 2D feature map이 출력됩니다. 

 

다음으로, MLP 부분을 1*1 Convolution 부분으로 바꿔주기만 하면 아래와 같은 식이 나옵니다. 

 

 

 

기존의 1D 형식인 Squence가 2D feature map으로 Reshape 되어 각 픽셀 간의 근접성이 확보되며, 여기서 Convolution 연산만 해주면 Locality를 주입할 수 있다고 주장합니다.

 

Locality

 

위 방식으로 각 픽셀간의 근접성이 확보되어 Convolution 연산을 하였습니다, 각 픽셀이 독립적으로 연산되는 1*1 Convolution의 경우 인접한 픽셀들과의 상호작용이 부족합니다. 이러한 과정은 기존 MLP를 사용하는 FFN과 별반 다를 게 없습니다.

 

따라서, 본 논문 2개의 1*1 Convolution 사이에, 인전 픽셀들과의 상호작용을 위해 3*3 Convolution을 삽입하는데 Depth-wise Convolution(DWConv)을 삽입합니다. 이 3*3 DWConv가 Locality의 핵심입니다. 

 

 

다시 식을 재정의하면 아래와 같이 나타낼 수 있습니다. 

 

 

여기서 \( W^r_1 \)은 1*1 Conv(\(d, \gamma d \)), \( W^r_2 \)은 1*1 Conv(\( \gamma d, d \)), \( W^r_2 \)은 3*3 Conv(\( \gamma d, \gamma d \)), \( f \)는 Activation 함수입니다.

 

 

Class token

기존 ViT는 최종적으로 이미지 분류를 위해 Class token을 삽입하여 Self-Attention, FFN을 통과하였습니다. 

 

 

\( X_{cls} \in \mathbb {R}^{1*d} \)가 Class token이며, 기존 \( N \)개의 이미지 토큰들로 이루어진 시퀀스에 결합하여 총 \( N + 1 \) 개가 만들어집니다.

 

Class Token은 Self-Attention에서 이미지 토큰들과 정보 공유를 하지만, FFN에서는 다시 분리되어 통과합니다.

 

 

위 \(Z \) 만 FFN를 통과 후 \( Y \)를 출력하고, 다시 Class token과 결합하여 최종 이미지를 분류합니다. 

 

 

 

이 Class token과 이미지 토큰을 결합하여 FFN을 통과시켰는데, 결합하지 않고 FFN에 통과시킨 결과와 똑같았다고 합니다. 이에 따라, 위 FFN은 이미지 토큰에 대해서만 변환을 수행한다고 해석할 수 있습니다.

 

 

마치며

 

 

이 구조를 ViT기반 여러 모델에 적용하였을 때, 상당한 성능 향상을 이끌 수 있었습니다. 

 

 

또한, 앞단(Low)에 이 구조를 넣었을 때가, 뒷단(High, Mid)에 넣었을 때 보다 우월한 성능을 기록하고 있습니다. 이에 따라, 앞단에 Locality를 주입해 줘야 뒤의 레이어까지 이 Locality가 퍼진다고 해석할 수 있습니다. 

 

 

\( \gamma \)의 값이 커질수록 self-Attention 연산으로 학습된 feature representations을 더 풍부할 수 있음도 증명되었고요.

 

이 밖에도 Activation 함수를 바꿔가며 큰 성능 향상을 관찰하였지만, 이 부분은 관심 있으시면 원문을 참고하시길 바랍니다!

 

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

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