VIT(Vision Transformer)

요약

• 이전 Vision Task에서 Self-Attention적용의 한계

      Self-Attention을 적용하는 시도는 있었으나, Hardware Accelerators에 비효율적 → ResNet 구조가 SOTA였음

     따라서 기존의 Transformer를 최대한 그대로 적용하고자 함

 

• Attention is All you Need 

      NLP에서 가장 대표적인 구조 "Self-Attention"를 활용한 Transformer

      대표 모델 "BERT"는 Large Dataset(Corpus)를 사전학습(Pre-Train) → 작은 Task에서 미세조정(Fine-Tune) 

 

• Transformer의 장점

      계산 효율성(Efficiency) 및 확장성(Scalability) 

      100B Parameter도 학습 가능!

      데이터셋이 크면 클수록 모델을 키워도 되며, 성능이 포화(Saturate)될 징후 X → 더 키우면 키울수록 더 성능은 높아질 것임

 

• Transformer의 적용 방안

      이미지를 Patch로 분할 후 Sequence로 입력  → NLP에서 단어(Word)가  입력되는 방식과 동일! ( ∵ 논문 제목이 "IMAGE IS WORTH 16X16 WORDS"인 이유)

      Supervised Learning 방식으로 학습

 

• Transformer의 특징

      ImageNet와 같은 Mid-sized 데이터셋으로 학습 시, ResNet보다 낮은 성능을 보임( * ImageNet은 더 이상 큰 데이터셋이 아님)

      JFT-300M 사전 학습 후, Transfer Learning → CNN구조 보다 매우 좋은 성능 달성(SOTA)

      Transformer는 inductive biases가 없음 = Locality와 Translation Equivariance 같은 CNN의 특성이 없음

 

inductive biases란 : 만나지 못한 상황을 해결하게 위해, 추가적인 가정을 활용해서 문제를 해결

VIT 핵심 목적

1.  Large scale pre-training 가능한지

2.  일반적인 Transformer를 바로 적용

3.  Medium-Resolution 이미지에도 적용 가능

VIT Architecture

NLP의 Transformer의 아키텍처를 거의 바로 사용할 수 있다

ViT 구조

• 입력 Embedding

      º Token embedding을 1D Sequence로 입력

         (이미지를 여러개의 패치(16x16 크기)로 자른후에 각 패치별로 1차원 embedding demension(16x16x3 = 768)으로 만든다)

        - 2D Image(xp∈R^(H×W×C) )를  xp∈R^N×(P^2⋅C) 로 변환

        - (P,P) : Image Patch의 해상도

        - N : Patch의 수 = HW/P^2 

        - D : 모든 Layer에서의 동일한 latent Vectore size

          → Flatten 한 Patch를 학습 가능한 Linear Projection( E )을 사용하여, D 차원으로 매핑

      º Patch Embedding이 출력됨

• [CLS] Token

      º BERT의 [class] Token처럼, class token을 concatenate 시키고 각 패치마다 Position Embedding을 더해준다

       - Pre-Training과 Fine-Tuning을 수행하는 Classification Head가 부착 -> cls token과 positional embedding은 모두 학습되는 파라미터

         (class token은 패치가 아닌 이미지 전체의 Embedding을 가지고 있다는 가정하에 최종 classification head에서 사용 / Position Embedding은 각 패치의 순서를 모델에 알려주는 역할을 한다)

• Classification Head

      º Pre-training : 1-hidden Layer인 MLP 

      º Fine-Tuning : 1-linear Layer

• Position embedding

      º Patch Embedding의 Position 정보를 유지하기 위해서, 추가 

      º 2D-Position Embedding을 추가해보았지만 더 좋은 성능 X

        → 이미지이지만, 1D Position Embedding 사용

• Transformer

      º 결과 embedding sequence는 Encoder의 입력으로 들어감 

      º Transformer Encoder

        - Multi-Head로 구성된 self-Attention 메커니즘 적용

      ※ Multi-Head Self Attention: https://jalammar.github.io/illustrated-transformer/

The Illustrated Transformer

• Inductive bias

    º CNN(Locality 가정)이나 RNN(Sequentiality 가정) 경우, Global 한 영역의 처리는 어려움

    º ViT는 일반적인 CNN과 다르게 공간에 대한, "Inductive bias"이 없음

    º 그러므로, ViT는 더 많은 데이터를 통해, 원초적인 관계를 Robust 하게 학습시켜야 함

       - ViT는 MLP Layer에서만 Local 및 Translation Equivariance 함

       - Self-Attention 메커니즘은 Global 함

    º 2-차원 구조 매우 드물게 사용

       - 이미지 Patch를 잘라서 넣는 Input 부분

       - Fine-tune 시, 다른 해상도(Resolution)의 이미지를 위한 위치 임베딩 조정

   ▶ 공간(Spatial) 관계를 처음부터 학습해야 함 (∵ Position Embedding 초기화 시 위치정보 전달 X)

 

• Hybrid Architecture

    º Image Patch의 대안으로, CNN의 Feature Map의 Sequence를 사용할 수 있음

       - CNN의 Feature는 Spatial size가 1x1이 될 수 있음

       - CNN으로 Feature추출 → Flatten → Eq.1의 Embedding Projection ( E ) 적용

Fine-Tuning & Higher Resolution

• 미세조정 시, 해상도 설정

    º 사전 훈련 때보다 더 높은 해상도로 Fine-Tune 하는 것이 성능에 도움이 됨

    º 더 높은 해상도 및 Patch Size 동일 → Sequence 길이 증가

       - Fine-Tuninig 시 : Pre-trained Position Embedding은 더 이상 의미가 없음

       - Pre-Trained Position Embedding을 2D 보간(Interpolation)을 수행

---

코드로 보는 VIT

1. Patch Embedding

class PatchEmbed(nn.Module):
    def __init__(self, img_size, patch_size, in_chans=3, embed_dim=768):
        super(PatchEmbed, self).__init__()
        self.img_size = img_size
        self.patch_size = patch_size
        self.n_patches = (img_size // patch_size) ** 2

        self.proj = nn.Conv2d(
            in_channels=in_chans,
            out_channels=embed_dim,
            kernel_size=patch_size,
            stride=patch_size,
        )  # Embedding dim으로 변환하며 패치크기의 커널로 패치크기만큼 이동하여 이미지를 패치로 분할 할 수 있음.

    def forward(self, x):
        x = self.proj(x)  # (batch_size, embed_dim, n_patches ** 0.5, n_patches ** 0.5)
        x = x.flatten(2)  # 세번째 차원부터 끝까지 flatten (batch_size, embed_dim, n_patches)
        x = x.transpose(1, 2)  # (batch_size, n_patches, embed_dim)
        return x

VIsion Transformer는 전혀 CNN을 사용하지 않는다고 하였다. 그런데 중간에 nn.Conv2d() 가 떡하니 있어 의아할 수 있다. 하지만 자세히 보면 kernerl_size와 stride가 패치 사이즈(16)로 되어 있기 때문에 서로 겹치지 않은 상태로 16x16의 패치로 나눈다는 의미로 해석할 수 있다.

입력 이미지 사이즈가 384x384 라고 했을때 Convolution을 수행하게 되면 차원이 (n, 768, 24, 24) 가 될 것이고 여기서  flatten과 transpose를 사용해서 (n, 576, 768)의 각 패치별(576개) 1차원 벡터(768 embed dim)로 표현 가능하다.

 

2. Multi Head Attention

class Attention(nn.Module):
    def __init__(self, dim, n_heads=12, qkv_bias=True, attn_p=0., proj_p=0.):
        super(Attention, self).__init__()
        self.n_heads = n_heads
        self.dim = dim
        self.head_dim = dim // n_heads
        self.scale = self.head_dim ** -0.5  # 1 / root(self.head_dim)
        '''
        # 나중에 query와 key를 곱하고 softmax를 취하기전에 scale factor로 나눠주는데 이 scale factor의 역할은 
        query @ key 의 값이 커지게 되면 softmax 함수에서의 기울기 변화가 거의 없는 부분으로 가기때문에 gradient vanishing
        문제를 해결하려면 scaling을 해주어야 한다고 Attention is all you need 논문에서 주장
         
        '''
        self.qkv = nn.Linear(dim, dim*3, bias=qkv_bias)
        self.attn_drop = nn.Dropout(attn_p)
        self.proj = nn.Linear(dim, dim)
        self.proj_drop = nn.Dropout(proj_p)

    def forward(self, x):
        n_samples, n_tokens, dim = x.shape
        if dim != self.dim:
            raise ValueError

        qkv = self.qkv(x)  # (n_samples, n_patches+1, dim*3)
        qkv = qkv.reshape(
            n_samples, n_tokens, 3, self.n_heads, self.head_dim
        )  # (n_samples, n_patches+1, 3, n_heads, head_dim)
        qkv = qkv.permute(2, 0, 3, 1, 4)  # (3, n_samples, n_heads, n_patches+1, head_dim)
        q, k, v = qkv[0], qkv[1], qkv[2]  # 각각의 n_heads끼리 query, key, value로 나눔
        k_t = k.transpose(-2, -1)  # (n_samples, n_heads, head_dim, n_patches+1)  dot product를 위한 transpose
        # dot product를 통해 query와 key사이의 유사도를 구함
        dp = (q @ k_t) * self.scale  # (n_samples, n_heads, n_patches+1, n_patches+1)  @: dot product (2x1)@(1x2)=(2x2)
        attn = dp.softmax(dim=-1)  # attention (n_samples, n_heads, n_patches+1, n_patches+1)
        attn = self.attn_drop(attn)

        weighted_avg = attn @ v  # (n_samples, n_heads, n_patches+1, head_dim)
        # 원래 차원으로 되돌림.
        weighted_avg = weighted_avg.transpose(1, 2)  # (n_samples, n_patches+1, n_heads, head_dim)
        weighted_avg = weighted_avg.flatten(2)  # concat (n_samples, n_patches+1, dim)

        x = self.proj(weighted_avg)  # linear projection (n_samples, n_patches+1, dim)
        x = self.proj_drop(x)
        return x

• self.qkv 에서 dim을 3배로 키우는 이유는 query, key, value를 분할 하기 위함
• query와 key를 dot product를 하고 softmax를 취함으로써 둘의 연관성을 구한다.
• 그다음 softmax를 취하기 전에 이 attention score를 scale로 나눠주게 되는데 attention score값이 커지게 되면 softmax함수에서 기울기변화가 없는 부분으로 가기 때문에 gradient vanishing을 막기 위함이다.
• softmax를 취한후 value를 곱해 최종 attention을 구하게 된다.
• value를 곱하는 이유는 관련이 있는 단어들은 그대로 남겨두고 관련이 없는 단어들은 작은 숫자(점수)를 곱해 없애버리기 위함.

3. MLP(Multi Layer Perceptron)

class MLP(nn.Module):
    def __init__(self, in_features, hidden_features, out_features, p=0.):
        super(MLP, self).__init__()
        self.fc1 = nn.Linear(in_features, hidden_features)
        self.act = nn.GELU()
        self.fc2 = nn.Linear(hidden_features, out_features)
        self.drop = nn.Dropout(p)

    def forward(self, x):
        x = self.fc1(x)
        x = self.act(x)
        x = self.drop(x)
        x = self.fc2(x)
        x = self.drop(x)
        return x

• MLP는 아주 간단하게 hidden dimension으로 한번 갔다가 돌아오도록 되어있고 hidden dimension은 base model에서 3072로 하고있다.
• 여기서 activation으로 GELU라는 것을 사용하는데 GELU는 Gaussian Error Linear Unit의 약자이며 다른 알고리즘보다 수렴속도가 빠르다는 특징을 가지고 있다.

AF 비교

4. Transformer Encoder Block

class Block(nn.Module):
    def __init__(self, dim, n_heads, mlp_ratio=4.0, qkv_bias=True, p=0., attn_p=0.):
        super(Block, self).__init__()
        self.norm1 = nn.LayerNorm(dim, eps=1e-6)
        self.attn = Attention(
            dim,
            n_heads=n_heads,
            qkv_bias=qkv_bias,
            attn_p=attn_p,
            proj_p=p
        )
        self.norm2 = nn.LayerNorm(dim, eps=1e-6)
        hidden_features = int(dim * mlp_ratio)  # 3072(MLP size)
        self.mlp = MLP(
            in_features=dim,
            hidden_features= hidden_features,
            out_features=dim,
        )

    def forward(self, x):
        x = x + self.attn(self.norm1(x))
        x = x + self.mlp(self.norm2(x))
        return x

• Vision Transformer의 Encoder를 반복하는 block이다.

 

 

참고

https://visionhong.tistory.com/25

[논문리뷰] Vision Transformer(ViT)

https://kmhana.tistory.com/27

[논문요약] Vision분야에서 드디어 Transformer가 등장 - ViT : Vision Transformer(2020)