Vision Transformer(ViT) - 개념과 원리, 패치부터 분류까지

CNN 없이 이미지를 트랜스포머로 보기

Vision Transformer(ViT)는 이미지 인식 문제를 CNN 대신 Transformer 구조로 직접 풀어낸 모델이다. 이미지를 작은 패치(patch) 로 나눈 뒤, 각 패치를 토큰(token) 처럼 취급하고 Self-Attention으로 패치 간 관계를 학습해 이미지 전체의 전역적 문맥을 잡는다. CNN이 국소 수용 영역에 의존하는 것과 달리, ViT는 멀리 떨어진 영역 간 상관관계를 한 번에 모델링할 수 있고, 데이터가 충분할 때 CNN을 넘어서는 성능을 보인다. 이 글에서는 ViT의 개념과 동작 원리만 다루며, 입력부터 출력까지의 파이프라인과 패치 임베딩·CLS 토큰·트랜스포머 인코더를 PyTorch로 간단히 구현하는 부분까지 정리한다. timm으로 사전 학습된 ViT를 쓰는 실전 예제(멀티 브랜치 분류기) 는 다음 편에서 이어서 다룬다.

(Vision Transformer(ViT)를 바탕으로 재구성했다.)

---

1. ViT가 하는 일 한 줄 요약

ViT는 이미지를 패치 단위 시퀀스로 바꾼 뒤, Transformer Encoder의 Self-Attention으로 전역 관계를 학습하고, 분류 헤드로 예측하는 구조다. 즉 “이미지를 시퀀스로 본다”는 관점 전환으로 컴퓨터 비전과 NLP의 경계를 넘어선 대표 모델이다.

---

2. ViT 파이프라인 – 입력부터 출력까지

2-1. Input Image

모델에 들어가는 원본 이미지다. ViT는 여기에 합성곱을 바로 쓰지 않고, 다음 단계에서 고정 크기 패치로 잘라 쓴다.

2-2. Divide Image into Patches

입력 이미지를 동일한 크기의 패치로 균등 분할한다. 이 단계가 ViT의 핵심 전처리로, 이미지를 패치들의 시퀀스로 재구성하기 위한 준비다.

2-3. Image Patches

각 패치는 이미지의 국소 시각 정보를 담고, 이후 독립적인 처리 단위가 된다. 이 시점에서 이미지는 2차원 행렬이 아니라 여러 패치의 집합으로 표현된다.

2-4. Convert Patches into Tokens

각 패치를 flatten 한 뒤 선형 변환으로 고정 차원 벡터(임베딩)로 만든다. 이 벡터가 Transformer에 들어가는 Patch Token이며, NLP의 토큰과 같은 역할을 한다.

2-5. Transformer Encoder

Patch Token 시퀀스가 Transformer Encoder에 입력된다. Encoder 안에서는:

• Self-Attention: 모든 패치 쌍의 관계 계산
• Feed Forward Network: 비선형 변환
• Residual Connection + Layer Normalization: 학습 안정성·정보 보존

입력을 순차 처리하지 않고 병렬로 계산한다.

2-6. Learning Global Context

Self-Attention 덕분에 ViT는 공간적으로 멀리 떨어진 영역 간 연관성까지 직접 학습한다. CNN의 국소 수용 영역과 구별되는 특성이다.

2-7. Output (Classification Result)

Encoder 출력을 분류 헤드에 넣어 최종 클래스를 예측한다. 보통 CLS 토큰 하나의 표현을 사용한다.

---

3. PyTorch로 ViT 흐름 짚어 보기

작은 이미지(8×8×3)를 4×4 패치로 나누고, 패치 임베딩 → CLS 토큰 → 위치 임베딩 → Transformer Encoder → 분류까지 한 번에 구현해 보자.

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

# 8×8 컬러 이미지 → 4×4 패치 4개
image = np.random.rand(8, 8, 3).astype(np.float32)
image = image.reshape(2, 4, 2, 4, 3)
image = image.transpose(0, 2, 1, 3, 4)  # (2, 2, 4, 4, 3)
patches = image.reshape(-1, 4, 4, 3)    # (4, 4, 4, 3)
patches = torch.tensor(patches)

embedding_dim = 32
num_heads = 4
num_transformer_layers = 2
num_classes = 10
mlp_dim = 256

def patch_embedding(patches, embedding_dim):
    N, Ph, Pw, C = patches.shape
    patch_dim = Ph * Pw * C
    patch_flat = patches.reshape(N, patch_dim).unsqueeze(0)
    proj = nn.Linear(patch_dim, embedding_dim)
    tokens = proj(patch_flat)
    return tokens, proj

tokens, _ = patch_embedding(patches, embedding_dim)
B, N, D = tokens.shape

# CLS 토큰 + 시퀀스
cls_token = nn.Parameter(torch.zeros(1, 1, D))
x = torch.cat([cls_token.expand(B, -1, -1), tokens], dim=1)

# 위치 임베딩
pos_embed = nn.Parameter(torch.randn(1, x.size(1), D) * 0.02)
x = x + pos_embed

# Transformer Encoder
encoder_layer = nn.TransformerEncoderLayer(
    d_model=D, nhead=num_heads,
    dim_feedforward=mlp_dim, batch_first=True
)
transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_transformer_layers)
z = transformer_encoder(x)

# 분류: CLS 위치만 사용
classification_layer = nn.Linear(D, num_classes)
logits = classification_layer(z[:, 0, :])
output = F.log_softmax(logits, dim=-1)
pred_class = output.topk(1, dim=-1).indices.item()

정리하면, ViT는 패치 → flatten → Linear → Patch Token, 그다음 CLS 토큰 추가 + 위치 임베딩 → Transformer Encoder → CLS 출력으로 분류하는 구조다.

---

4. ViT의 특성 정리

• 데이터가 많을수록 Transformer가 CNN보다 성능 향상 폭이 크다.
• 이미지 크기가 바뀌면 패치 크기·개수를 조절해 유연하게 대응할 수 있다.
• 이미지가 커지면 패치 수가 늘어나 시퀀스 길이가 길어지므로, 연산량·메모리를 고려해 패치 크기를 정한다.

---

5. 정리

주제	핵심 포인트
ViT	이미지를 패치 시퀀스로 바꾸고, Transformer Self-Attention으로 전역 문맥을 학습하는 비전 모델이다.
패치 → 토큰	패치를 flatten 후 Linear로 고정 차원 벡터(Patch Token)로 만든다.
CLS + 위치 임베딩	분류용 CLS 토큰을 시퀀스 앞에 붙이고, 위치 임베딩을 더해 Encoder에 넣는다.

---

마치며

• ViT는 “이미지를 시퀀스로” 보는 관점으로, CNN과 다른 방식의 전역 관계 학습을 가능하게 한다.
• 작은 예제로 패치 분할 → 임베딩 → CLS·위치 임베딩 → Transformer Encoder → 분류 흐름을 한 번 구현해 보면, 논문이나 라이브러리 코드 읽기에 도움이 된다.

timm으로 사전 학습된 ViT를 불러와 멀티 브랜치 분류기(4자리 숫자 이미지 분류 등)를 만드는 실전 구현은 다음 편에서 이어서 정리한다.