Vision Transformer for Small-Size Datasets 논문리뷰

그림5

2021년 12월, 최근 발표된 Vision Transformer 관련 논문으로 대량의 데이터가 필요한 기존의 ViT 문제를 해결한 논문입니다.

소규모 데이터셋에서도 scratch에서(처음부터) 학습할 수 있는 SPT(Shifted Patch Tokenization) 및 LSA(Locality Self-Attention) 모듈을 제안하여, ViT에 모두 적용하여 실험한 결과 대표적인 소형 데이터셋인 TinyImageNet(클래스별 500개 데이터, 총 200개 클래스)에서 평균 2.96%의 성능 향상, Swin Transformer는 4.08% 성능 향상

 

1. Introduction

• 기존 ViT의 문제점
  1) 열악한 토큰화 : non-overlapping patches를 사용함으로 인해 visual token들이 상대적으로 더 작은 receptive field를 가지게 된다. 그래서 인접 픽셀과의 공간적 관계가 visual token에 충분히 포함되지 않는다.
  2) 열악한 attention mechanism : 이미지 데이터의 특징 차원은 자연어나 오디오 신호의 특징보다 훨씬 크기때문에 임베디드 토큰의 수는 필연적으로 많다. 그래서 토큰의 attention score의 분포가 smooth해진다. 즉, ViT가 중요한 visual token에 locally(가까이) attend(주의를 기울이다)할 수 없다.
  * 위의 두가지 주요 문제는 target class에 집중하기 어려운, 매우 불필요한 attention을 유발하게 된다. 그래서 ViT가 일반적으로 background에 집중하고, 대상 클래스의 모양을 잘 포착하지 못하게 한다. (그림5 참조)
• 위의 문제 해결을 위해 제시한 모듈
  1) SPT(Shifted Patch Tokenization) : TSM(Temporal Shift Module)에서 파생된 것, TSM은 feature의 일부 시간 채널을 이동시키는 효과적인 시간 모델링을 뜻함, 이에 영감을 받아 공간적으로 이동된 이미지를 입력 이미지와 함께 토큰화하는 효과적인 공간 모델링을 제안함 -> 표준 토큰화보다 ViT에 더 넓은 receptive field를 제공
  2) LSA(Locality Self-Attention) : ViT가 locally attend할 수 있도록 하는 모듈, self-token을 배제하고 softmax 함수에 학습가능한 temperature를 적용하여 attention score 분포의 smoothing(평활화) 현상을 완화한다. LSA는 각 토큰이 자신과 관련성이 큰 토큰에 더 집중하도록 하여, local에서 작동하도록 attention을 유도한다.
• 제안된 SPT와 LSA는 add-on 모듈 형태로 다양한 ViT에 구조적 변경 없이 쉽게 적용할 수 있고, 성능을 효과적으로 향상시킬 수 있다. (그림 1을 보면, 정확도도 향상되고, 시간당 처리 가능한 수도 많아진 것, 즉 모델 속도 개선도 된 것을 볼 수 있다.)

그림1, 정확도 비교, 별표는 모듈을 적용한 경우의 정확도

 

그림2

2. Related Works

• small dataset으로 ViT쓰면 잘 안돼서 이 논문에서 그걸 제안함

 

3. Method

• 위의 그림2의 (a) : SPT, 입력 이미지를 여러 방향(그림2에서는 4개)으로 공간적으로 이동하고 입력 이미지와 concat해서 처리함. 그런 다음 patch partitioning이 표준 ViT처럼 적용되고, visual token에 삽입하기 위해 patch flattening(평탄화) > layer 정규화 > 선형 투영의 3단계가 수행됨
• 그림2의 (b) : LSA, 일반적으로 softmax 함수는 temperature scaling을 통해서 출력 분포의 평활도를 제어할 수 있다. (참고 : https://3months.tistory.com/491), LSA는 주로 softmax 함수의 temperature 매개변수를 학습해서 attention score 분포를 sharp하게 한다. (그림2-b의 상단부를 보면 Attention Score가 smooth한 거에서 weakly sharp, strongly sharp한 이미지 확인 가능), 그리고 Query와 Key에 의해 계산된 유사도 행렬의 대각 성분을 강제로 억제하는 대각 마스킹을 적용해서 자체 토큰 관계(self-token relation)을 제거한다. (본인 자신이랑은 관계가 100%일 것이니 제거하는 게 맞다고 생각합니다.) 이 마스킹을 함으로써, 다른 토큰들 간의 attention score를 상대적으로 증가시킬 수 있고, attention score 분포를 더 선명하게 만든다.

3.1. 표준 ViT의 self-attention 메커니즘의 토큰화 및 공식화 방법

x : input Image (H x W x C), 높이, 너비, 채널

1) ViT는 input 이미지를 non-overlapping(겹치지 않는) 패치들로 나누고, vector의 sequence를 얻기 위해 패치들을 평탄화한다. 

2) tokenization : transformer encoder의 hidden dimmension 영역으로 선형 투영(linearly projecting)을 통해 patch embedding을 얻는다. 각각의 patch embedding은 transformer encoder의 visual token input이 된다. 

이하생략

 

3.2. SPT (Shifted Patch Tokenization)

 

3.2.1. Formulation

- 각 input image를 4개의 대각선 방향(좌상, 우상, 좌하, 우하)로 패치 크기의 절반만큼 공간적으로 이동시킨다.

- 다음으로 이동된 특징을 input image와 동일한 크기로 자른 다음 input과 concat한다.

- 연결된 feature를 non-overlapping patches로 나누고 패치를 평탄화한다. (기존방식처럼)

- layer normalization(LN) 및 선형 투영(linear projection)을 통해 visual token을 얻는다.

 

3.2.2. Patch Embedding Layer

- SPT를 패치 임베딩 레이어로 사용하는 방법을 설명

- 클래스 토큰을 visual token에 연결한 다음, positional embedding을 추가한다. 여기서 클래스 토큰은 전체 이미지의 representation 정보가 있는 토큰이며, positional embedding은 visual token에 위치 정보만을 제공한다. 클래스 토큰을 사용하지 않으면 SPT의 출력에 위치 임베딩만 추가된다. 

 

3.2.3. Pooling Layer

- Tokenization(토큰화) : 3D-tensor feature를 2D-matrix feature로 처리하는 프로세스, HxWxC 3차원인 x를 Nxd의 2차원인 T(x)로 만드는 과정

- 그래서 tokenization을 풀링 계층으로 사용하면 visual token의 수를 줄일 수 있다. -> SPT를 풀링 레이어로 사용하는 걸 제안함

- 먼저 class token과 visual token을 분리하고 2D matrix 형태의 visual token을 공간 구조의 3D tensor로 재구성한다. 그리고 3.2.1의 식이 적용되면 토큰 수가 감소한 새로운 visual token이 포함된다. 마지막으로 선형 투영된 class token은 앞에서 구한 새로운 visual token과 연결된다.

 

 

3.3. LSA (Locality Self-Attention Mechanism)

 

3.3.1 Diagonal Masking (대각 마스킹)

- 대각 마스킹을 이용해서 self-token보다 다른 토큰에 더 집중되도록 처리

 

3.3.2 Learnable Temperature Scaling

- ViT가 학습 과정에서 스스로 softmax temperature를 결정할 수 있도록 하여, attention score 분포의 평활화 문제를 해결함

 

참고자료

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

https://github.com/aanna0701/SPT_LSA_ViT

GitHub - aanna0701/SPT_LSA_ViT: Implementation for Visual Transformer for Small-size Datasets