Vision Transformer (ViT)의 Patch Embedding과 Class Token 정리

Vision Transformer (ViT) 구현을 읽을 때 처음 막히기 쉬운 부분은 Transformer Encoder 자체가 아니라 그 앞 단계입니다. (B, C, H, W) 이미지가 언제 (B, N, D) 토큰 리스트로 바뀌고, [CLS]와 위치 정보가 어느 축에 추가되는지 여부입니다. 이 부분을 모호하게 알면, 체크포인트의 해상도 변경이나 태스크 헤드 연결 시 shape이 맞지 않게 됩니다.

ViT의 입력부는 다음 4단계로 나누어 보면 이해하기 쉽습니다.

• 이미지를 겹치지 않는 패치(patch)들로 분할합니다.
• 각 패치를 차원 D로 선형 투영합니다. (sequence axis의 맨 앞에)
• [CLS]를 연결(concat)합니다.
• 같은 shape의 Position Embedding을 element-wise로 더해줍니다.

이 글에서는 ViT 원 논문과 Google Research의 공식 repository를 근거로, 텐서 shape, 최소 PyTorch 구현, Pre-LN Encoder, 해상도 변경 시 주의점 등을 연결하여 설명합니다.

224×224 이미지가 197×768이 되기까지

ViT-B/16에 RGB 이미지를 입력하는 대표적인 예를 사용하겠습니다.

1개의 패치를 flatten한 차원은 (B, 197, 768)이 됩니다.

(B, 3, 224, 224)
-> patch grid (B, 196, 16, 16, 3)
-> flatten (B, 196, 768)
...

16×16×3와

Patch Embedding은 Conv2d로 작성 가능합니다

원 논문의 수식에서는 flatten된 패치를

구현에서는 kernel_size=P, stride=P, out_channels=D의 Conv2d로 한 번에 계산할 수 있습니다. 커널이 1 패치 전체를 덮고 같은 폭으로 이동하기 때문에 패치는 겹치지 않습니다. Conv2d의 각 output channel은 Linear projection의 1 output dimension에 대응합니다.

이는 1회 패치 투영에 대한 등가성입니다. 국소 커널을 여러 층으로 쌓는 CNN 전체와 ViT 전체가 같은 아키텍처인 것은 아닙니다.

입력 embedding만 최소 구현하기

다음 코드는 이미지를 패치 토큰으로 변환하고, [CLS]와 Position Embedding을 더하는 부분만 보여준 교육용 레퍼런스입니다. 임의 해상도 대응이나 체크포인트 변환은 포함하지 않았습니다.

import torch
from torch import Tensor, nn
class VitInputEmbedding(nn.Module):
...

shape 변환의 핵심은 flatten(2).transpose(1, 2)입니다. Conv2d 직후는 (B, D, H/P, W/P)이므로, 공간 축을 모은 후 Transformer가 기대하는 (B, N, D)로 바꿔줍니다.

expand한 Class Token은 batch마다 다른 parameter를 갖지 않습니다. 하나의 학습 가능한 vector를 각 sample의 sequence 맨 앞에 배치하며, backpropagation 시에는 공유 parameter에 gradient가 모입니다.

concat할 것과 add할 것을 구분하기

ViT 입력 수식은 다음 형태입니다.

[CLS] $ ext{patch}_1; ext{patch}_2; ext{patch}_3; ...$ (여기서 세미콜론은 sequence axis에서의 concat입니다.) [CLS]를 추가하면 토큰 수가

Class Token은 정답 레이블이 아닙니다. 모든 입력에서 공통의 초기 parameter로 시작하며, Self-Attention을 통해 패치 정보를 집약합니다. 이미지 분류에서는 마지막 Encoder block을 통과한 0번 토큰을 LayerNorm에 거쳐 classification head로 전달합니다.

Encoder block은 Pre-LN 형태입니다

ViT 원 논문의 Encoder block은 Attention과 MLP 앞에 Layer Normalization을 두는 Pre-LN 구성입니다.

residual connection(입력을 sub-layer 출력에 더하는 경로)은 2개 있습니다. 하나는 Attention을, 다른 하나는 MLP를 우회합니다. 구현에서 Post-LN block으로 대체할 경우 parameter shape은 맞더라도 계산이 같아지지 않으므로, checkpoint architecture의 확인이 필요합니다.

ViT-B에서는 이 block을 12층에 걸쳐 통과시킵니다. 입력과 출력의 shape은 (B, 197, 768)를 유지하지만, 각 token의 값은 모든 token과의 Attention 및 MLP를 통해 업데이트됩니다.

patch size는 Attention cost를 크게 변화시킨다

Self-Attention의 score matrix는 token 수를

Patch size	Grid	Tokens	Attention 요소 수
32	7×7	50	2,500
...
ViTのパッチサイズ32、16、8におけるトークン数とAttention計算量の比較

P=16에서 P=8로 변경하면 patch 수는 약 4배가 되지만, score matrix의 요소 수는 약 15.9배입니다.

작은 patch는 작은 물체나 가는 선을 초기에 압축하기 어려울 수 있지만, activation memory와 latency를 증가시킵니다. patch size는 정확도 parameter뿐만 아니라 token budget을 결정하는 system parameter입니다.

해상도 변경 시 Class Token을 보간하지 않는다

224×224, P=16으로 학습된 Position Embedding은 [CLS]를 포함하여 197 tokens 분량입니다. 같은 patch size로 384×384로 변경하면 24×24 patches가 되어 577 tokens 분량이 필요합니다.

원 논문에서는 patch 위치에 해당하는 196개를 14×14 grid로 되돌린 후, 2D interpolation으로 24×24로 변환합니다. Class Token용 1개는 공간 grid가 아니므로 분리하여 그대로 재결합합니다.

구현에서는 다음 순서를 고정하는 것이 안전합니다.

• position_embedding[:, :1]으로 Class Token 위치를 분리하고 -
  position_embedding[:, 1:]을 원래의 2D grid로 reshape한 후 - bicubic 등의 2D interpolation으로 새 grid에 변환합니다.
• (1, new_num_patches, D)로 되돌린 다음 Class Token 위치와 sequence axis에서 concat 합니다.

interpolation은 학습된 위치를 새로운 grid에 근사 배치하는 처리입니다. 고해상도 data로 fine-tuning을 하지 않는다고 해서 미지의 위치 관계가 자동으로 학습되는 것은 아닙니다. 또한 token 수는 197개에서 577개로 증가하므로, Attention 요소 수는 약 8.58배가 됩니다.

CNN과의 차이는 inductive bias로 정리한다

CNN은 국소 kernel과 weight 공유를 architecture에 내장하고, 층을 쌓아 receptive field(출력에 영향을 주는 입력 범위)를 넓힙니다. original ViT는 patch화된 후 모든 token 간을 Self-Attention으로 연결하여 공간 관계를 Position Embedding과 training data로부터 많이 학습합니다.

Checkpoint를 사용하기 전에 확인해야 할 것

ViT 계열 checkpoint는 model 이름이 비슷하더라도 입력 계약(input contract)이 같다고 할 수 없습니다.

• 패치 크기(patch size)와 학습 해상도(training resolution)
• 이미지 정규화(image normalization)의 평균값(mean) 및 표준편차(standard deviation)
• RGB 채널 순서와 리사이즈/크롭 방법
• Class Token 유무와 풀링(pooling) 방법
• 위치 임베딩(Position Embedding)이 학습된 것인지, 고정된 것인지, 상대적인 위치인 것인지
• Pre-LN/Post-LN, MLP 비율(ratio), 헤드 수(head number)
• 분류 헤드(classification head)의 클래스 수(class number)
• 체크포인트(checkpoint)와 코드의 라이선스(license)

ViT-B/16

이라는 이름만으로는 전처리 과정까지 확정할 수 없습니다. 체크포인트를 제공하는 리포지토리나 모델 카드를 기준으로 삼고, 전처리 또한 같은 버전 경계에서 관리해야 합니다.

또한, 오리지널 ViT의 Class Token + Linear head는 직접적으로 이미지 분류(image classification)를 수행합니다. 객체 탐지(object detection)나 세그멘테이션(segmentation)에서는 패치 단위의 피처(feature)와 태스크별 헤드(task-specific head)가 필요합니다. 최종 Class Token만 추출한다고 해서 바운딩 박스(Bounding Box)나 픽셀 레이블(pixel label)을 얻을 수는 없습니다.

요약

ViT의 입력 구현은 Conv2d를 이용한 패치 프로젝션(patch projection), Class Token의 연결(concat), Position Embedding의 추가(add)라는 세 가지 형태 조작으로 나누어 이해할 수 있습니다. 224×224, P=16의 ViT-B의 경우, 196개의 패치 토큰에 [CLS]를 더한 $(B, 197, 768)$이 인코더 입력입니다.

구현 시 주의할 점은, 패치 내부 차원(dimension)과 히든 차원(hidden dimension)을 혼동하지 않는 것, Pre-LN과 잔차 경로(residual path)를 체크포인트에 맞추는 것, 해상도 변경 시 Class Token을 제외하고 위치 그리드(position grid)만 보간하는 것입니다. 패치 크기(patch size)를 작게 할수록 토큰 수가 늘어나고 어텐션 비용은 제곱으로 증가하므로, 디테일과 계산 자원을 동시에 평가해야 합니다.

ViT 원 논문의 배경, 224×224 이미지로부터의 구체적인 계산, CNN 비교, 분류 모델로서의 한계까지 순서대로 확인하고 싶은 분은 개인 블로그의 ViT 완전판도 참고해 주세요.

참고 자료

토론(Discussion)