[Paper short review] Vision Transformer for Small-Size Datasets 논문 리뷰

@ 굵은 글씨는 중요한 내용, 빨간 글씨는 내가 추가한 내용

 

https://arxiv.org/pdf/2112.13492v1.pdf

 

• 정리: SPT(Shifted Patch Tokenization)과 LSA(Localy Self-Attention)를 적용함으로써 기존 vision transformer의 locality inductive bias 부족 문제를 완화함

 

---

 

0. Abstract

Transformer를 vision task에 적용한 vision transformer(ViT)가 높은 성능을 보임

• 하지만, 대규모 dataset을 사용한 pretraining을 기반으로 함 (이는 locality inductive bias 부족 때문)

 

이를 해결하기 위해 본 논문에서는 SPT(Shifted Patch Tokenization)과 LSA(Localy Self-Attention) 제안

• locality inductive bias 부족 문제를 해결하여 small dataset에서도 학습 가능
• 다양한 ViT task에 적용 가능
• 실험 결과 small dataset에서 평균 2.96% 성능 향상 // Swin Transformer에서 4.08% 성능 향상

 

3. PROPOSED METHOD

SPT (그림 (a))를 통해 더 많은 공간 정보를 visual token에 embedding하고 ViT의 locality inductive bias를 증가시킴

1. 입력 이미지를 여러 방향으로 이동 (Shifted Images)
2. 입력 이미지와 연결 (Concatenating)
3. Patch partition 수행 (Patch features)
4. Patch flattening, layer normalization, linear projection 수행 (Visual Tokens)

LSA (그림 (b))를 통해 ViT의 attention을 local로 집중시킴으로써 locality inductive biase를 증가시킴

1. 소프트맥스에 학습가능한 temperature를 적용 (소프트맥스는 temperature scaling을 통해 출력 분포의 smoothness를 제어할 수 있음) = weakly sharp
2. 쿼리와 키로 계산된 유사성 매트릭스의 대각선 성분을 억제하는 대각선 마스킹 적용 (자체 토큰 관계를 제거하고 서로 다른 토큰 간의 attention score를 증가시켜 attention score 분포를 더 선명하게 만듬) = strongly sharp

 

3.1. Preliminary

기본 ViT의 전체 프로세스

 

입력 이미지를 겹치지 않는 patch로 나눈 후 평평하게 만든 벡터를 구한다. 이 프로세스는 다음과 같다.

• = 입력 이미지 (H,W,C = 높이, 너비, 채널)
• = i번째 flattened 벡터
• N = HW/P^2 = patch 수 (P = patch size)

 

각 벡터를 transformer encoder의 hidden 차원의 공간에 linear projection하여 patch embedding을 얻는다 (=T). 이 프로세스는 다음과 같다.

• = 학습가능한 linear projection (d = transformer encoder의 hidden 차원)

visual token의 수용 필드 size는 다음과 같다.

• rtoken = 토큰화의 수용 필드 크기
• rtrans = transformer encoder의 수용 필드 크기 (=1, transformer encoder에서 조정X)
• j, k = convolutional layer의 stride, kernel size (= patch 크기)

224 × 224 이미지를 학습한다고 가정할때, 표준 ViT의 패치 크기는 16이므로 비주얼 토큰의 rtoken도 16인 반면, ResNet50 수용적 필드 크기는 483이다.

• 결과적으로 ViT의 비주얼 토큰은 ResNet50 기능보다 약 30배 작은 수용성 필드 크기를 가진다. 이로 인해, locality inductive bias가 부족하다고 예상된다.
• 따라서 본 논문은 토큰화의 수용적 분야를 늘림으로써 풍부한 공간 정보를 활용하기 위한 SPT를 제안

일반 ViT의 self attention은 다음과 같다.

1. 학습 가능한 linear projection이 각 토큰에 적용되어 Query, Key 및 Value를 얻는다.
2. Query와 Key의 dot-product 연산을 통해 토큰 간의 의미 관계를 나타내는 유사성 매트릭스, 즉 R ∈ R^((N+1)×(N+1))를 계산
3. R의 대각선 성분은 self 토큰 관계를 나타내고, 비대각 성분은 다른 토큰과의 관계를 나타낸다.

• , = 각각 query, key에 대한 학습가능한 linear projection (dq, dk = 각각 query, key의 차원)

 

  4. R을 value 차원의 제곱근으로 나눈 후 소프트맥스 함수를 적용하여 attention score matrix를 구한다.

  5. 식 (5)와 같이 attention score matrix와 value의 dot-product 연산을 통해 self attention을 수행한다.

• =  value에 대한 학습가능한 linear projection (dv = value의 차원)

식 (5)는 관계성이 높은 token들끼리의 attention score가 커지도록 설계되어 있지만 두 가지 원인으로 인해 ViT의 attention은 관계성에 상관없이 유사한 경향을 보인다.

1. Query(xEq)와 key(xEk)는 동일한 입력 토큰에서 선형으로 투영되어 self 토큰 관계는 보통 서로 다른 토큰 간 관계보다 크다. 따라서 식 (5)의 소프트맥스는 self 토큰 관계에 상대적으로 높은 점수를 주고, 서로 다른 토큰 간 관계에 작은 점수를 준다.
2. 식 (5)에서 R을 √dk로 나눈 이유는 Softmax가 small gradient를 갖는 것을 방지하기 위함이다. 그러나 이는 오히려 소프트맥스 함수의 높은 temperature로 작용하여 attention score 분포의 평활화를 야기한다. 아래 표에서 확인할 수 있듯이, 소프트맥스의 temperature가 √dk보다 작을 때 오히려 높은 성능을 보인다. 

 

3.2. Shifted Patch Tokenization

3.2.1 Formulation

1. 입력 이미지에 대해 4개의 대각선 방향으로 patch 크기의 절만만큼 공간적 이동 (=S)

• S외에 다양한 시프트 전략을 사용할 수 있지만, 본 연구 내 SPT의 모든 실험은 S를 따름

 

2. 이동된 image는 입력 이미지와 동일한 크기로 자른 후 입력 이미지와 concat 수행

 

3. concat된 feature를 겹치지 않는 patch로 나누고 식(1)과 같이 패치를 flatten한다. 그 후, LN(Layer Normalization) 및 linear projection을 통해 visual token을 얻는다. 

 

전체 프로세스는 다음 식과 같다.

• s^i ∈ R^(H×W×) = S에 따라 i번째 이동된 이미지
• Es ∈ R^(P^2·C·(Ns+1)×ds) = 학습 가능한 linear projection
• ds = transformer encoder의 hidden 차원
• Ns = S에 의해 이동된 이미지의 수

 

3.2.2 Patch Embedding Layer

SPT를 patch embedding layer로 사용하는 방법

• class token을 visual token에 concat한 후 positional embedding을 추가한다 (class token은 전체 이미지의 표현 정보를 가진 token, positional embedding은 visual token에 위치 정보 제공). 

• Xcls ∈ R^(dS) = class token
• Epos ∈ R^((N+1)×dS) = 학습 가능한 positional embedding
• N = 식(6)에서 embeddde token의 수

 

3.2.3 Pooling Layer

토큰화는 3D를 2D에 embedding하는 과정이다. 이 과정을 통해 공간 크기를 감소시킬 수 있고, 다음과 같이 SPT를 풀링 레이어로 사용하여 visual token을 줄이는 것을 제안한다.

1. class token과 visual token을 분리
2. 2D형태의 visual token은 R : R^(N×d) → R^((H/P)×(W/P)×d)로 변형된다.
3. 식(6)의 SPT를 적용하여 토큰 수가 줄어든 새로운 visual token이 embedded된다.
4. linear projection class token을 concat한다.

전체 프로세스는 식 (8)과 같다.

• = 학습 가능한 linear projection
• d's = 다음 stage의 hidden 차원

 

3.3. Locality Self-Attention Mechanism

3.3.1 Diagonal Masking
소프트맥스에서 자체 토큰 관계를 배제함으로써 서로 다른 토큰 간 관계에 더 큰 점수를 부여하는 방법

수식은 다음과 같다.

• = 마스킹된 유사성 행렬의 각 성분

 

3.3.2 Learnable Temperature Scaling

학습 과정 동안 ViT가 스스로 소프트맥스 temperature를 결정

• 아래는 Softmax temperature를 식(5)에서 학습 가능한 파라미터로 사용했을 때 depth에 따른 평균 학습 temperature

  

• 일반적으로 소프트맥스의 낮은 temperature는 분포가 날카로워진다. 즉, 학습 가능한 temperature scaling은 attention score 분포를 날카롭게 한다.
• 식(5)를 기준으로 대각선 마스킹과 학습 가능한 temperature scaling이 모두 적용된 LSA는 다음과 같다.

• τ = 학습 가능한 temperature

 

이를 통해 attention score 분포의 평활화 문제를 해결

• 아래는 모든 head에 대한 총 Kullback–Leibler Divergence (KLD)의 depth 별 평균

• T = 학습가능한 temperature 적용
• M = 대각선 마스킹 적용
• L = 전체 LSA 적용 (T+M)

• L-ViT가 ViT보다 약 0.5 크기때문에 LSA가 attention score 분포의 평활화 문제를 완화한다는 것을 알 수 있음 (KLD 의 평균이 낮을수록 attention score 분포가 평평하다는 의미)

 

4. EXPERIMENT

4.1. SETTING

4.1.1 Environment and Dataset

4.1.2 Model Configurations

4.1.3 Training Regime

 

4.2. QUANTITATIVE RESULT

4.2.1 Image Classification

 

4.2.2 Ablation Study

 

4.3. QUALITATIVE RESULT

• SPT와 LSA를 ViT에 적용했을 때, 최종 class token의 attention score 시각화 결과 (모델명 앞에 SL이 붙은 것이 SPT와 LSA를 적용한 것)
• attention이 target class에 집중되는 것을 확인할 수 있음
• locality inductive bias를 증가시키고 ViT의 attention을 개선하도록 유도하는 것을 발견

 

5. CONCLUSION

small dataset에서 ViT를 학습하기 위해, ViT의 locality inductive bias을 높이기 위한 SPT와 LSA를 제안한다. SPT는 특정 변환을 통해 풍부한 공간 정보를 시각적 토큰에 embedding하고, LSA는 학습 가능한 매개 변수를 가진 소프트맥스를 통해 ViT가 local에 attend하도록 유도한다.

SPT와 LSA는 독립적으로 상당한 성능 향상을 달성할 수 있으며, 모든 ViT task에 적용된다. 따라서 본 연구를 통해 ViT가 small dataset를 처음부터 학습하고 더 발전할 수 있는 기회를 제공한다.

 

My Discussion

Why this paper?

vision transformer에 관련된 논문을 읽다보니 큰 dataset에 대해서 학습을 해야되고, pretrain된 모델이 있어야된다?는 것과 같은 제약들이 존재했는데 이를 완화한 논문이라 생각하고 읽음

 

--> input이미지에 대해 augmentation을 수행한 이미지를 동시에 활용한다는 것, 대각 토큰은 사용하지 않는 것, temperature scaling을 사용하는 것이 주요 특징

--> augmentation을 따로 사용하지 않고, 동시에 embedding을 한다는 점이 재밌고, 대각선 대각토큰을 삭제하는게 단순한데 효과적인 아이디어라 신기함

--> temperature scaling을 제대로 이해함

--> 추후 논문 쓸 때, 적용해봐도 좋을것같음 (성능 향상에 많은 도움이 될듯)