[ViT 시리즈] Vision Transformer 논문 리뷰 보단 메모.

---

본격적인 Vision Transformer 관련 논문들을 리뷰를 하기 전에, 공부한 내용을 정리할 겸해서 간단하게 메모장을 끄적이는 느낌으로 정리를 해볼까 합니다.

 

---

 

1. 사전 지식

---

■ Attention mechanism

• Encoder ↔️ Decoder 사이의 correlation을 바탕으로 특징을 추출해 나가는 과정.
• Decoder로 부터 query가 나옴.
• Encoder로 부터 key, value가 나옴.

 

📝 참고 그림

출처: https://blog.promedius.ai/transformer/

 

 

■ Self-Attention mechanism

• 입력 데이터로 부터 query, key, value가 계산된다. 그 이후에는 Attention mechanism과 동일한 과정으로 진행.
• 데이터 내의 상관 관계를 바탕으로 특징을 추출해 나가는 과정.
  • 입력 데이터로 부터 query, key, value를 구함.
  • query, key의 similarity를 통해 weight를 결정. 
• Transformer Encoder는 Self-Attention mechanism 이용.

■ Inductive bias

네트워크 목적

새로운 데이터에 대한 좋은 성능을 내는 것을 목표로 한다.

 

Inductiva bias란?

새로운 데이터에 대해 좋은 성능을 내기 위해 모델에 사전적으로 주어지는 가정들을 inductive bias라고 부른다.

• 예를 들면, convolution filter를 활용해 local feature patterns을 추출한 수 있는 것.

■ CNN vs Transformer

CNN

멀리 떨어진 두 정보들을 통합할 때 여러개의 layer를 거쳐야 한다.

2차원의 local feature patterns을 유지하면서 layer를 통과한다. 이는, convolution filter를 활용해 local feature patterns을 추출 하는 것이며 이것이 바로 inductive bias이다.

 

Transformer

• 순서
  • 5x5 image가 주어졌을 때, 1D 1D vector로 불어냄.
  • 1D vector를 Linear Projection 거친 후 patch embedding에 self-attention을 적용하면 Query, Key, Value를 계산하게 됨
  • Query, Key의 Similarity를 통해 weight를 결정.
  • Query, Key의 Similarity를 통해 weight를 결정.
• self-attention layer 하나만 거쳐도 멀리 떨어진 정보들을 교환할 수 있다.

• 데이터를 1D로 만든 후 self-attention을 통해 layer를 통과함. 이는 2차원의 local feature patterns이 유지되지 않아 CNN에 비해 inductive bias가 적은 모델이라고 말할 수 있다.
• inductive bias가 적은 만큼 모델의 자유도가 높아 더 많은 데이터를 학습할 수 있다. 이는, ViT의 단점으로 작용하는데 이 문제를 극복한 방법이 DeiT이다.

 

 

2. ViT(Vision Transformer)

---

■ 특징

• NLP 분야에서 사용된 Transformer를 응용하여 Vision task에서 사용할 수 있도록 고안된 첫 논문 ViT
• Transformer만 사용해서 image classification task에 적용
• Architecture의 hyper-parameter에 따라 여러가지 모델들은 제안하고 있음.

 

■ 학습 과정

학습 순서

1. 이미지 입력
2. 16x16 patch로 자름
3. 각각의 patch들을 1D-vector로 풀어냄
4. 1D vector로 만들어진 patch들을 Linear Projection을 통해 768 차원의 각 patch의 embedding vector로 표현.

   • class PatchEmbedding(nn.Module):
         def __init__(self, in_channels: int = 3, patch_size: int = 16, emb_size: int = 768):
             self.patch_size = patch_size
             super().__init__()
             '''
     				self.projection = nn.Sequential(
                 # break-down the image in s1 x s2 patches and flat them
                 Rearrange('b c (h s1) (w s2) -> b (h w) (s1 s2 c)', s1=patch_size, s2=patch_size),
                 nn.Linear(patch_size * patch_size * in_channels, emb_size)
             )
     				'''
     				self.projection = nn.Sequential(
                 # using a conv layer instead of a linear one -> performance gains
                 nn.Conv2d(in_channels, emb_size, kernel_size=patch_size, stride=patch_size),
                 Rearrange('b e (h) (w) -> b (h w) e'),
             )
                     
         def forward(self, x: Tensor) -> Tensor:
             x = self.projection(x)
             return x

5. 각 patch embedding에 classification token, position embedding 추가.
6. Transformer Encoder
7. Transformer Output
8. classification

 

■ 학습을 통해 결정되는 Parameter

Classification token

classification을 위해 사용되는 token

class PatchEmbedding(nn.Module):
    def __init__(self, in_channels: int = 3, patch_size: int = 16, emb_size: int = 768):
        self.patch_size = patch_size
        super().__init__()
        self.projection = nn.Sequential(
            # using a conv layer instead of a linear one -> performance gains
            nn.Conv2d(in_channels, emb_size, kernel_size=patch_size, stride=patch_size),
            Rearrange('b e (h) (w) -> b (h w) e'),
        )
        
        self.cls_token = nn.Parameter(torch.randn(1,1, emb_size))
        
    def forward(self, x: Tensor) -> Tensor:
        b, _, _, _ = x.shape
        x = self.projection(x)
        cls_tokens = repeat(self.cls_token, '() n e -> b n e', b=b)
        # prepend the cls token to the input
        x = torch.cat([cls_tokens, x], dim=1)
        return x

Position embedding

• patch의 위치 정보를 가지고 있는 embedding
• position embedding을 사용 안하면?
  • 이미지 입력 → 일전 크기의 patch로 자름 → patch를 sequence로 생각하여 transformer Encoder로 입력 → 각 patch가 어떤 위치에서 왔는지 위치 정보가 손실됨. → 학습 안됨
• (classification token, patch embedding) + position embedding 이 Transformer Encoder로 입력됨.
• classification token, position token은 학습에 의해 결정되는 parameter입니다.

class PatchEmbedding(nn.Module):
    def __init__(self, in_channels: int = 3, patch_size: int = 16, emb_size: int = 768, img_size: int = 224):
        self.patch_size = patch_size
        super().__init__()
        self.projection = nn.Sequential(
            # using a conv layer instead of a linear one -> performance gains
            nn.Conv2d(in_channels, emb_size, kernel_size=patch_size, stride=patch_size),
            Rearrange('b e (h) (w) -> b (h w) e'),
        )
        self.cls_token = nn.Parameter(torch.randn(1,1, emb_size))
        self.positions = nn.Parameter(torch.randn((img_size // patch_size) **2 + 1, emb_size))

        
    def forward(self, x: Tensor) -> Tensor:
        b, _, _, _ = x.shape
        x = self.projection(x)
        cls_tokens = repeat(self.cls_token, '() n e -> b n e', b=b)
        # prepend the cls token to the input
        x = torch.cat([cls_tokens, x], dim=1)
        # add position embedding
        x += self.positions
        return x

시각화

cos = torch.nn.CosineSimilarity(dim=1, eps=1e-6)
fig = plt.figure(figsize=(10,10))
fig.suptitle("Visualization of position embedding similarities", fontsize=24)

for i in range(1, pos_embed.shape[1]):
    sim = F.cosine_similarity(pos_embed[0, i:i+1], pos_embed[0, 1:], dim=1)
    sim = sim.reshape((14, 14)).detach().cpu().numpy()
    ax = fig.add_subplot(14, 14, i)
    ax.axes.get_xaxis().set_visible(False)
    ax.axes.get_yaxis().set_visible(False)
    ax.imshow(sim)

 

■ Architecture

Transformer Encoder

• ViT는 Vanilla Transformer(Attent is all you need)구조와는 다른 구조로 구성되어 있다.
• Transformer는 layer를 깊게 쌓게되면 학습이 어렵다는 단점이 있다.
  • 이런 단점을 극복하여 학습이 되게 하기 위해서는 layer normalization의 위치가 중요하다는 것이 후속 연구들을 통해 증명이 되었다.
• ViT는 NLP 에서 보여진 layer를 깊게 쌓으면 학습이 어렵다는 문제와 이를 극복하기 위해 layer normalization 위치가 중요하다는 연구를 받아드린 Vision Transformer이다.
  • 기존 Transformer: Multi-head attention → normalization
  • ViT: normalization → Multi-head attention

## Transformer Encoder에 입력으로 들어가는 tensor
transformer_input = torch.cat((model.cls_token, patches), dim=1) + pos_embed

■ Attention in Transformer Encoder

Self-Attention

순서

• layer normalization 거친 후 얻어진 $z$
• weight metrix $W_Q$, $W_K$, $W_V$에 곱해져서 Query, Key, VAlue 값들이 계산됨(⭐️ weight metrix의 학습을 통해 attention이 학습)
• Query, Key 값들의 dot product[torch.einsum('bhqd, bhkd -> bhqk', queries, keys)]를 통해 Similarity 계산
• Softmax를 취해 0 ~ 1의 attention score계산
• 계산된 attention score에 Value를 곱해 최종적인 output을 얻음.

Multi-Head Self-Attention

• ViT base model에서는 self-attention 12번 수행.
• 768차원 → Multi-head self-attention 통과 → 64차원으로 감소 → self-attention을 12번 수행하니 768차원으로 증가 → 결국 차원수가 동일하게 유지됨

class MultiHeadAttention(nn.Module):
    def __init__(self, emb_size: int = 768, num_heads: int = 8, dropout: float = 0):
        super().__init__()
        self.emb_size = emb_size
        self.num_heads = num_heads
        self.keys = nn.Linear(emb_size, emb_size)
        self.queries = nn.Linear(emb_size, emb_size)
        self.values = nn.Linear(emb_size, emb_size)
        self.att_drop = nn.Dropout(dropout)
        self.projection = nn.Linear(emb_size, emb_size)
        self.scaling = (self.emb_size // num_heads) ** -0.5

    def forward(self, x : Tensor, mask: Tensor = None) -> Tensor:
        # split keys, queries and values in num_heads
        queries = rearrange(self.queries(x), "b n (h d) -> b h n d", h=self.num_heads)
        keys = rearrange(self.keys(x), "b n (h d) -> b h n d", h=self.num_heads)
        values  = rearrange(self.values(x), "b n (h d) -> b h n d", h=self.num_heads)
        # sum up over the last axis
        energy = torch.einsum('bhqd, bhkd -> bhqk', queries, keys) # batch, num_heads, query_len, key_len
        if mask is not None:
            fill_value = torch.finfo(torch.float32).min
            energy.mask_fill(~mask, fill_value)
            
        att = F.softmax(energy, dim=-1) * self.scaling
        att = self.att_drop(att)
        # sum up over the third axis
        out = torch.einsum('bhal, bhlv -> bhav ', att, values)
        out = rearrange(out, "b h n d -> b n (h d)")
        out = self.projection(out)
        return out

MLP (1 hidden layer fully-connected layer)

class FeedForwardBlock(nn.Sequential):
    def __init__(self, emb_size: int, expansion: int = 4, drop_p: float = 0.):
        super().__init__(
            nn.Linear(emb_size, expansion * emb_size),
            nn.GELU(),
            nn.Dropout(drop_p),
            nn.Linear(expansion * emb_size, emb_size),
        )

Transformer output

■ 실험적 특징

학습 팁

• pretraining dataset의 해상도는 줄이고, fine tunning dataset의 해상도를 증가시키면 모델 성능 ↑
  • Emerald project(Covid-19 Kaist model)에로 위와 같은 방법으로 적용한거로 알고 있음.
• Transformer을 활용하면 기존 CNN 모델보다 효율적인 학습이 가능함. (하지만 ViT 특징만 이용한다면 오히려 단점으로 적용될 수 있음)

데이터가 많이 필요하다!

• CNN에 비해 inductive bias가 적음
• → 입력 이미지를 패치로 구성한 후 데이터를 1D vector로 만들어 self-attention을 통과함 (이 때, sequence로 들어가니 position embedding은 필수!) ⇒ 따라서 2차원의 지역적인 정보가 유지되지 않음.
• inductive bias가 적으면 그만큼 모델의 자유도는 ↑ ⇒ 이는 많은 데이터로 학습이 가능하다. (결국은 데이터가 많이 필요하다)

데이터가 많이 필요한 ViT의 단점을 극복한 논문이 DeiT(facebook)

• Knowlege Distiliation, Data Augmentation을 활용한 논문.
• 순서
  • 입력 데이터를 Teacher model, Student model에 입력으로 넣어줌
  • Teacher model의 output이 Teacher model의 지식이다 라고 가정
  • Teacher과 Student의 output을 KL-divergence를 이용하여 계산
    • 이는 T와 S의 분포 사이의 거리를 최대한 줄여주기 위함.
  • Teachcer과 Student의 분포의 거리를 줄여 Teacher 모델의 지식을 Student 모델로 전이해 주겠다 라는 방법.
• 핵심 내용
  • CNN의 inductive bias를 Student model(DeiT) 모델에 넣어줄 수 있기 때문에 Teacher model를 CNN을 사용하여 활용하는게 더 유리함.

 

 

Reference

• https://www.youtube.com/watch?v=bgsYOGhpxDc&t=580s
• https://blog.promedius.ai/transformer/
• https://github.com/xmu-xiaoma666/External-Attention-pytorch