Vision Transformer - 이론과 구현

ViT 개요

자연어 처리 분야에서 트랜스포머 모델이 큰 성능 향상을 이뤄내면서 컴퓨터비전 분야에서도 Transformer 구조를 적용하려는 시도가 이루어졌다. Vision Transformer는 2020년 구글의 "An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale" 논문에서 제안된 이미지인식 딥러닝 모델이다.

CNN 과 ViT의 차이

• CNN과 ViT는 이미지 특징을 잘 표현하는 임베딩을 만들고자 하는 목적은 같지만 그 과정에는 큰 차이가 있다. 합성곱 신경망의 임베딩은 이미지 패치 중 일부만 선택하여 학습하며 이를 통해 이미지 전체의 특징을 추출한다.
• 반면 ViT 임베딩은 이미지를 작은 패치들로 나누어 각 패치간의 상관 관계를 학습한다. 이를 위해 Self-Attention 방법을 사용해 모든 이미지 패치가 서로에게 주는 영향을 고려해 이미지의 전체 특징을 추출한다.
• 따라서 좁은 Receptive Field를 가진 CNN은 전체 이미지 정보를 표현하는 데 수많은 계층이 필요하지만,  Transformer 모델은 Attention Distance를 계산하여 오직 한개의 ViT 레이어로 전체 이미지 정보를 쉽게 표현할 수 있다. 
• 또한 ViT는 픽셀 단위로 처리하는 CNN과 달리 패치 단위로 이미지를 처리하기 때문에 더 작은 모델로도 높은 성능을 얻을 수 있다.
• ViT 모델은 입력 이미지의 크기가 고정되어 있어 크기가 다른 이미지를 처리하려면 이미지 크기를 맞추는 전처리가 필요하며 CNN이 이미지의 공간적인 위치정보를 고려하는데 비해 ViT는 패치간의 상대적인 위치정보만 고려하기 때문에 이미지 변환에 취약할 수 있다. 

ViT 모델 구조

ViT 모델은 입력이미지를 트랜스 포머 구조에 맞게 일정한 크기의 패치로 나눈 다음 각 패치를 벡터형태로 변환하는 Patch Embedding과 각 패치와의 관계를 학습하는 Encoder 계층으로 구성된다.

Vision Transformer 구조도, 출처: https://arxiv.org/abs/2010.11929

패치 임베딩

• 패치 임베딩 레이어는 입력 이미지를 작은 패치로 분할하는 과정을 말한다. 작은 패치로 분할하기 위해서는 이미지 크기를 맞추는 전처리가 수행되어야한다. 예를 들어 입력 이미지의 크기가 640x480이었다면 224x224 크기와 같은 정방형 크기로 이미지 크기를 조절한다. 
• 패치 임베딩을 위해서는 Kernel Size와 Stride를 설정해야하며 이 둘은 하이퍼 파라미터로 정의 된다. 커널은 패치의 크기를 의미하며 Stride는 패치가 이동하는 폭을 의미한다.
• 일반적으로 분할된 이미지 패치들을 배열형태로 나열할때 왼쪽에서 오른쪽, 위에서 아래의 순서로 배열한다. 이 배열 가장 왼쪽에 CLS Token(분류 토큰)을 추가한다. 이 분류 토큰은 전체 이미지를 대표하는 벡터로 특정 문제를 예측하는데 사용된다.
• 또한, 위치 임베딩(Position Embedding)을 사용하여 인접한 패치간의 관계를 학습한다. 위치 임베딩은 위치 정보를 임베딩 벡터로 변환하고 기존 이미지 패치 벡터들과 더한다.
• 이후 마지막에 Layer Normalization을 적용하면 패치 임베딩 레이어가 만들어진다.
• 코드로 구현하면 다음과 같다.

import torch
import torch.nn as nn

class PatchEmbedding(nn.Module):
    def __init__(self, img_size, patch_size, in_channels, embed_dim):
        super(PatchEmbedding, self).__init__()
        self.img_size = img_size
        self.patch_size = patch_size
        self.num_patches = (img_size // patch_size) ** 2
        self.proj = nn.Conv2d(in_channels, embed_dim, kernel_size=patch_size, stride=patch_size)

    def forward(self, x):
        x = self.proj(x)  # Shape: [batch_size, embed_dim, num_patches^(1/2), num_patches^(1/2)]
        x = x.flatten(2)  # Shape: [batch_size, embed_dim, num_patches]
        x = x.transpose(1, 2)  # Shape: [batch_size, num_patches, embed_dim]
        return x

• img_size: 이미지 사이즈 (예: 224이면 224x224 이미지)
• patch_size: 패치의 크기 (예: 16이면 16x16 패치)
• in_channels: 입력 이미지의 채널 수 (예: RGB → 3)
• embed_dim: 각 패치를 임베딩한 벡터의 차원 수

임베딩 레이어

class ViTEmbedding(nn.Module):
    def __init__(self,img_size, patch_size, in_channels, embed_dim,dropout_rate=0.1):
        super(ViTEmbedding,self).__init__()
        self.patch_embedding = PatchEmbedding(img_size, patch_size, in_channels, embed_dim)
        num_patches = self.patch_embedding.num_patches

        # CLS 토큰
        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
        # Leanable Positional Embedding
        # (패치수 + 1) 만큼의 포지셔널 인코딩. 학습가능한 파라미터로 설정
        self.pos_embedding = nn.Parameter(torch.zeros(1, num_patches +1 , embed_dim))

        # Dropout Layer (Transformer 입력 직전에 적용)
        # 입력 값에 regulaization 적용
        self.dropout = nn.Dropout(dropout_rate)

        # 초기화
        nn.init.trunc_normal_(self.cls_token, std=0.02)
        nn.init.trunc_normal_(self.pos_embedding, std=0.02)

    def forward(self, x):
        x = self.patch_embedding(x) # [B,N,D]
        B, N, D = x.shape

        # CLS 토큰을 배치 수 만큼 복제해서 맨 앞에 추가
        cls_tokens = self.cls_token.expand(B,-1,-1)
        x = torch.cat((cls_tokens, x), dim=1)

        # Positional Embedding 추가
        x = x + self.pos_embedding

        return self.dropout(x) # [B,N+1,D] (N+1 = N + CLS token)

인코더 계층

ViT 모델은 이미지를 일정한 크기의 패치로 나눠 각 패치를 벡터로 변환한 후 인코더 계층에 입력으로 전달해 다양한 쿼리, 키, 값  임베딩의 관계를 학습한다. Transformer의 멀티 헤드 어텐션의 수식은 아래 포스팅에 정리했다.

2025.04.14 - [Artificial intelligence/Deep Learning] - Transformer - Multi Head Attention

Transformer - Multi Head Attention

N개의 인코더 레이어를 반복적으로 적용한 후, 마지막 레이어에서는 분류 토큰의 특징 벡터를 추출한다. 이 분류 토큰 벡터는 이미지 데이터를 잘 표현하는 특징 벡터로 간주된다. 이러한 특징 벡터는 이후 다양한 이미지 분류 및 검색문제를 해결하는데 사용된다.

마찬가지로 코드로 구현해보면 다음과 같다

class TransformerEncoderBlock(nn.Module):
    def __init__(self,embed_dim, num_heads, mlp_ratio=4.0,drop_out=0.1):
        super(TransformerEncoderBlock, self).__init__()
        self.norm1 = nn.LayerNorm(embed_dim)
        # Multi-head Self Attention
        self.attn = nn.MultiheadAttention(embed_dim,num_heads,dropout=drop_out,batch_first=True)
        self.norm2 = nn.LayerNorm(embed_dim)
        self.mlp = nn.Sequential(
            nn.Linear(embed_dim, int(embed_dim * mlp_ratio)),
            nn.GELU(),
            nn.Dropout(drop_out),
            nn.Linear(int(embed_dim * mlp_ratio), embed_dim),
            nn.Dropout(drop_out)
        )

    def forward(self, x):
        # Multi-head Self Attention + residual
        x_res = x
        x = self.norm1(x)
        attn_output, _ = self.attn(x,x,x) # query, key, value
        x = x_res + attn_output # Residual Connection

        # MLP + residual
        x_res = x
        x = self.norm2(x)
        x = x_res + self.mlp(x)

        return x

Layer Normalization은 안정적인 학습을 위한 정규화 방식이며, Residual Connection은 정보를 보존하고 gradient의 흐름을 유지한다. MLP는 두개의 Linear + GELU Activation으로 구성된 FeedForward Layer이다. 이제 이 구성품들을 가지고 classfication Head와 결합하면 위의 그림에서 봤던 아키텍처가 완성된다.