비전 트랜스포머(ViT) 구현 및 CIFAR-10에서 훈련 결과 공개

픽셀에서 예측까지: 이미지용 트랜스포머 구축 MIT의 한 학부생이 ViT(비전 트랜스포머), CIFAR-10 데이터셋에서의 훈련 및 결과를 설명합니다. 컨벌루션 신경망(CNN)은 컴퓨터 비전 분야의 대부분의 주요 혁신을 이끌어왔지만, 만약 그것이 우리를 제약해 왔다면 어떨까요? 2020년에 Google의 연구팀은 이런 대膽한 질문을 던졌습니다: 컨볼루션을 완전히 없애고도 세계 최고 수준의 이미지 모델을 만들 수 있을까? 그들의 답변인 비전 트랜스포머(Vision Transformer, ViT)는 딥 러닝의 새로운 시대를 열었습니다. 그림 1: 비전 트랜스포머(ViT) 아키텍처. 출처: Dosovitskiy et al., 2020 (arXiv:2010.11929) 저는 컴퓨터 비전과 생성 모델에 관심이 있는 MIT 학부생입니다. 최근에 비전 트랜스포머를 처음부터 구현하여 그 구조를 더 잘 이해하는 데 집중했습니다. 이 글은 그 과정을 요약한 가이드로, 이론, 시각 자료 및 코드를 결합하여 작성되었습니다. 이 글을 읽으면 PyTorch로 작동하는 ViT와 이를 어떻게 이해할 수 있는지 깊게 이해하게 될 것입니다. 1. 배경과 직관 1.1 순환 모델에서 트랜스포머의 등장 (NLP에서) 2017년 이전, 자연어 처리(NLP) 분야는 RNNs와 LSTMs가 지배했습니다. 이 모델들은 기계 번역부터 언어 모델링까지 다양한 작업을 가능하게 했습니다. 하지만 초기 성공에도 불구하고, 주요한 제약이 있었습니다. 시퀀스를 하나씩 처리해야 하므로 병렬화가 불가능했고, 시퀀스가 길어질수록 이전 토큰의 정보를 유지하기 어려웠습니다. 이러한 병목 현상은 특히 언어의 전반적인 이해가 필요한 작업에서는 확장성을 저해했습니다. 2017년, Google의 연구진은 "Attention Is All You Need"이라는 논문에서 트랜스포머라는 새로운 아키텍처를 제안했습니다. 이 아키텍처는 단순하지만 강력한 아이디어인 자기 주의(self-attention)를 중심으로 구성되었습니다. 각 토큰이 시퀀스 내 모든 다른 토큰을 직접 고려할 수 있도록 하는 것이 특징입니다. 즉, 각 단어는 질문(query), 듣기(key), 정보 수집(value)을 통해 학습합니다. 이 메커니즘은 순환의 필요성을 없애고, 이전 RNNs의 주요 약점을 해결했습니다. 그림 2: RNNs는 입력을 순차적으로 처리하지만, 트랜스포머는 모든 토큰을 동시에 처리합니다. 선의 굵기는 주의의 강도를 나타냅니다. 출처: 저자 불과 2년 만에 트랜스포머 아키텍처는 NLP 분야를 완전히 장악했습니다. 이전 모델들보다 효율적이고, 확장성이 더 좋으며, 장거리 의존성을 더 잘 모델링할 수 있었습니다. 트랜스포머는 BERT, GPT, T5 등의 주요 혁신 모델의 백본이 되었습니다. 그렇다면, 컴퓨터 비전 분야에서는 어떨까요? 당시에는 CNNs가 주도하고 있었지만, 이들도 자체적인 제약이 있었습니다. 컨볼루션은 본질적으로 국소적이어서 장거리 의존성을 포착하기 어렵습니다. 또한 공간적 사전 정보와 세심한 특성 엔지니어링에 크게 의존했습니다. 1.2 주의력이 컨볼루션을 대체할 수 있을까? 비전 분야로의 전환 2020년, Dosovitskiy 등은 "An Image Is Worth 16x16 Words"라는 논문에서 비전 트랜스포머(ViT)를 소개했습니다. 그들은 대담한 아이디어를 제시했습니다: 이미지를 문장처럼 다루면 어떨까? 컨볼루션 필터에 의존하지 않고 이미지를 패치로 나누어 표준 트랜스포머에 입력했습니다. 초기 ViTs는 CNNs와 경쟁하기 위해 거대한 데이터셋이 필요했지만, 이 접근법은 주의 기반 모델이 비전 분야에서도 작동할 수 있다는 것을 입증했습니다. 비전 트랜스포머의 등장 이후, 이는 개선의 물결을 일으켰습니다. 대담한 실험에서 시작된 이 아이디어는 이제 현대 컴퓨터 비전의 핵심 구성 요소가 되었습니다. 2. 비전 트랜스포머의 작동 방식 2.1 패치 임베딩 트랜스포머는 원래 문장의 단어 토큰 같은 시퀀스를 처리하도록 설계되었습니다. 그러나 이미지는 2D 그리드 형태의 픽셀로 이루어져 있습니다. 따라서 이미지를 트랜스포머에 입력하려면 어떻게 해야 할까요? 비전 트랜스포머는 이미지를 겹치지 않는 정사각형 패치(예: 16x16 픽셀)로 나눕니다. 각 패치는 1D 벡터로 평평하게 펴지고, 선형 투사되어 고정 크기의 임베딩으로 변환됩니다. 예를 들어, ViT는 단어 토큰의 시퀀스 대신 이미지 패치 임베딩의 시퀀스를 처리합니다. 그림 3: 패치 임베딩. 출처: Dosovitskiy et al., 2020 (arXiv:2010.11929) 비유: 단어 토크나이저가 문장을 단어 임베딩의 시퀀스로 바꾸는 것처럼, ViT는 이미지를 패치 임베딩의 시퀀스로 바꿉니다. 2.2 클래스 토큰과 위치 임베딩 트랜스포머가 이미지 시퀀스를 올바르게 처리하려면 두 가지 추가 요소가 필요합니다: 모든 패치에서 전역 정보를 집계하는 [CLS] 토큰, 공간 구조를 인코딩하는 위치 임베딩. ViT에서는 특별한 학습 가능한 토큰이 입력 시퀀스의 맨 앞에 삽입됩니다. 자기 주의 과정에서 이 토큰은 모든 패치에 주의를 기울이며, 최종 분류에 사용될 표현으로 변환됩니다. 트랜스포머는 순서에 무관하므로, 모델이 공간적 인식을 갖도록 하기 위해 각 토큰에 고유한 위치 임베딩을 추가합니다. [CLS] 토큰과 위치 임베딩 모두 학습 가능한 매개변수로, 훈련 중 갱신됩니다. 2.3 멀티헤드 자기 주의 (MHSA) 비전 트랜스포머의 핵심은 멀티헤드 자기 주의 메커니즘입니다. 이 부분은 모델이 공간적 거리와 상관없이 이미지 패치가 서로 어떻게 관련되는지 이해할 수 있게 합니다. 단일 주의 함수 대신, MHSA는 입력을 여러 "헤드"로 나눕니다. 각 헤드는 입력의 다른 측면에 집중하도록 학습합니다. 어떤 헤드는 가장자리에, 다른 헤드는 질감에, 또 다른 헤드는 공간 배치에 집중할 수 있습니다. 이들의 출력은 연결되어 원래 임베딩 공간으로 재투영됩니다. 작동 방식, 단계별 설명: 방정식 1: 스케일된 닷 프로덕트 주의 왜 "멀티헤드"인가? 각 헤드는 시퀀스의 다른 부분에 주의를 기울입니다. 이는 모델이 병렬로 복잡한 관계를 이해할 수 있게 합니다. 공간적 인접성뿐만 아니라 의미 구조도 포함됩니다. 그림 4: 왼쪽은 쿼리(Q), 키(K), 값(V)을 사용하여 주의 점수가 계산되는 방식을 보여줍니다. 오른쪽은 여러 주의 "헤드"가 병렬로 다양한 표현을 포착하는 방식을 설명합니다. 출처: Vaswani et al. (arXiv:1706.03762) 2.4 트랜스포머 인코더 자기 주의 메커니즘을 갖춘 후, 이를 더 큰 단위인 트랜스포머 블록 안에 포장합니다. 이 블록은 ViTs와 NLP 트랜스포머의 기본 구성 요소입니다. 각 블록은 전역적으로 주의를 기울이고, 여러 층을 통해 특성을 변환하면서 노멀라이제이션과 잔차 연결을 유지하여 안정성을 보장합니다. ViT 트랜스포머 블록 내부: - 주의 전 노멀라이제이션 (pre-norm) - 노멀라이제이션된 입력에 대한 멀티헤드 자기 주의 적용 - 주의 출력을 다시 더하는 잔차 연결 - 또 다른 노멀라이제이션, followed by a small MLP - MLP 출력을 더하는 또 다른 잔차 연결 그림 5: ViT 인코더 블록의 주의, MLP, 잔차 연결. 출처: 저자 이 구조는 모든 트랜스포머 층에서 반복됩니다 (예: ViT-Base는 12층). 2.5 분류 헤드 다수의 트랜스포머 블록을 통과한 후, 모델이 최종 예측을 생성할 방법이 필요합니다. 비전 트랜스포머에서는 이 작업을 분류 헤드가 수행합니다. 임베딩 단계에서 시퀀스의 맨 앞에 [CLS] 토큰을 추가했습니다. BERT와 마찬가지로, 이 토큰은 자기 주의를 통해 모든 이미지 패치로부터 정보를 집계하도록 설계되었습니다. 모든 트랜스포머 층을 통과한 후, [CLS] 토큰의 최종 임베딩이 전체 이미지를 요약하는 표현으로 사용됩니다. 이 벡터는 간단한 선형 레이어를 통해 클래스 로짓을 출력합니다. 3. 구현 가이드 모든 핵심 모듈 — 패치 임베딩, 멀티헤드 자기 주의, 인코더 블록 —은 처음부터 구현되었습니다. timm의 단축키를 사용하지 않았습니다. 3.1 패치 임베딩 이미지 패치를 임베딩 시퀀스로 변환하기 위해 다음과 같은 트릭을 사용합니다. 이미지에서 패치를 추출하고 평평하게 펴는 수동 for-loop를 작성하는 대신, Conv2d 레이어를 사용하여: 이것이 단일 작업으로 겹치지 않는 패치를 추출하고 학습 가능한 선형 투사를 적용하며, 역전파를 쉽게 합니다. ```python class PatchEmbed(nn.Module): def init(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768): super().init() self.img_size = img_size self.patch_size = patch_size self.num_patches = (img_size // patch_size) ** 2 self.proj = nn.Conv2d( in_chans, embed_dim, kernel_size=patch_size, stride=patch_size ) def forward(self, x): # x shape: [B, 3, 224, 224] x = self.proj(x) # [B, embed_dim, H/patch, W/patch] x = x.flatten(2) # [B, embed_dim, num_patches] x = x.transpose(1, 2) # [B, num_patches, embed_dim] return x ``` 3.2 클래스 토큰과 위치 임베딩 다음은 ViTEmbed 모듈을 정의하는 코드입니다: ```python class ViTEmbed(nn.Module): def init(self, num_patches, embed_dim): super().init() # 학습 가능한 [CLS] 토큰 (모델 당 1개) self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) # [1, 1, D] # 학습 가능한 위치 임베딩 (CLS 토큰을 포함한 각 토큰당 1개) self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim)) # [1, N+1, D] def forward(self, x): batch_size = x.shape[0] # [CLS] 토큰을 배치 크기에 맞춰 확장 cls_tokens = self.cls_token.expand(batch_size, -1, -1) # [B, 1, D] # 패치 임베딩에 CLS 토큰을 추가 x = torch.cat((cls_tokens, x), dim=1) # [B, N+1, D] # 위치 임베딩을 추가 x = x + self.pos_embed # [B, N+1, D] return x ``` 3.3 멀티헤드 자기 주의 비전 트랜스포머의 가장 중요한 부분 중 하나인 멀티헤드 자기 주의를 구현합니다. 각 입력 토큰은 쿼리(Q), 키(K), 값(V) 벡터로 선형 투사됩니다. 주의는 여러 헤드에 걸쳐 병렬로 계산되고, 출력은 연결되어 원래 임베딩 차원으로 재투영됩니다. ```python class MyMultiheadAttention(nn.Module): def init(self, embed_dim, num_heads): super().init() assert embed_dim % num_heads == 0, "embed_dim must be divisible by num_heads" self.embed_dim = embed_dim self.num_heads = num_heads self.head_dim = embed_dim // num_heads # 학습 가능한 Q, K, V 투사 self.q_proj = nn.Linear(embed_dim, embed_dim) self.k_proj = nn.Linear(embed_dim, embed_dim) self.v_proj = nn.Linear(embed_dim, embed_dim) # 최종 출력 투사 self.out_proj = nn.Linear(embed_dim, embed_dim) def forward(self, x): B, T, C = x.shape # [배치, 시퀀스 길이, 임베딩 차원] # 입력을 Q, K, V로 투사 Q = self.q_proj(x) K = self.k_proj(x) V = self.v_proj(x) # 헤드로 재구성: [B, num_heads, T, head_dim] def split_heads(tensor): return tensor.view(B, T, self.num_heads, self.head_dim).transpose(1, 2) Q = split_heads(Q) K = split_heads(K) V = split_heads(V) # 스케일된 닷 프로덕트 주의 scores = torch.matmul(Q, K.transpose(-2, -1)) # [B, heads, T, T] scores /= self.head_dim ** 0.5 attn = torch.softmax(scores, dim=-1) # 값을 주의에 적용 out = torch.matmul(attn, V) # [B, heads, T, head_dim] # 헤드를 재결합 out = out.transpose(1, 2).contiguous().view(B, T, C) # 최종 선형 투사 return self.out_proj(out) ``` 3.4 트랜스포머 인코더 마지막으로, 모든 것을 모듈러 단위인 트랜스포머 블록에 감싸습니다. 이 설계는 모델이 전역적으로 주의를 기울이고, MLP를 통해 특성을 변환하면서 스킵 연결을 통해 안정성을 유지할 수 있게 합니다. ```python class TransformerBlock(nn.Module): def init(self, embed_dim, num_heads, mlp_ratio=4.0): super().init() self.norm1 = nn.LayerNorm(embed_dim) self.attn = nn.MultiheadAttention(embed_dim, num_heads, 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.Linear(int(embed_dim * mlp_ratio), embed_dim) ) def forward(self, x): # 자기 주의와 잔차 연결 x = x + self.attn(self.norm1(x), self.norm1(x), self.norm1(x))[0] # MLP와 잔차 연결 x = x + self.mlp(self.norm2(x)) return x ``` 3.5 모든 것을 함께 조립 비전 트랜스포머의 핵심 구성 요소 — 패치 임베딩, 위치 인코딩, 멀티헤드 자기 주의, 트랜스포머 블록, [CLS] 토큰 —를 모두 조립하여 전체 모델을 만듭니다. ```python class SimpleViT(nn.Module): def init( self, img_size=224, patch_size=16, in_chans=3, embed_dim=768, depth=12, num_heads=12, num_classes=1000 ): super().init() self.patch_embed = PatchEmbed(img_size, patch_size, in_chans, embed_dim) num_patches = (img_size // patch_size) ** 2 self.vit_embed = ViTEmbed(num_patches, embed_dim) # 트랜스포머 블록 쌓기 self.blocks = nn.Sequential(*[ TransformerBlock(embed_dim, num_heads) for _ in range(depth) ]) # 최종 분류 전 노멀라이제이션 self.norm = nn.LayerNorm(embed_dim) # 선형 분류 헤드 (CLS 토큰 사용) self.head = nn.Linear(embed_dim, num_classes) def forward(self, x): # [배치 크기, 채널, 높이, 너비] x = self.patch_embed(x) # -> [B, N, D] x = self.vit_embed(x) # [CLS] 토큰 + 위치 임베딩 추가 x = self.blocks(x) # 트랜스포머 층 통과 x = self.norm(x) # [CLS] 토큰 정규화 return self.head(x[:, 0]) # [CLS] 토큰을 사용한 분류 ``` 4. ViT 훈련 4.1 데이터셋: CIFAR-10 우리는 잘 알려진 벤치마크 데이터셋인 CIFAR-10을 사용하여 비전 트랜스포머를 훈련했습니다. 이 데이터셋은 비행기, 고양이, 배 등 10개 클래스에 걸쳐 60,000개의 이미지를 포함하며, 각 이미지는 32x32 픽셀로 작습니다. 따라서 CIFAR-10은: 4.2 모델 설정: CIFAR-10을 위한 ViT 조정 ViTs는 원래 ImageNet 같은 대규모 데이터셋을 위해 설계되었지만, 제한된 컴퓨팅 자원으로 CIFAR-10에서 훈련을 가능하게 하기 위해 여러 조정을 가했습니다. ```python CIFAR-10을 위한 SimpleViT 구성 재구성 model = SimpleViT( img_size=32, # CIFAR-10 이미지는 32x32 픽셀 patch_size=4, # 4x4 패치 → 64 토큰 in_chans=3, embed_dim=192, # 작은 임베딩 크기 depth=6, # 적은 트랜스포머 블록 num_heads=3, # 192로 나누어 떨어짐 num_classes=10 # CIFAR-10의 클래스 수 ).to(device) ``` 4.3 훈련 설정 모델은 다음과 같이 훈련되었습니다: 훈련은 효율적이었으며, 약 30초 정도의 에폭 시간이 걸렸습니다. 이는: 4.4 결과 우리는 비전 트랜스포머를 30 에폭 동안 훈련하여 총 약 15분이 소요되었습니다. 훈련이 끝난 후, 모델은 CIFAR-10 테스트 세트에서 약 60%의 정확도를 달성했습니다. 이는 모델의 단순함과 비교적 작은 데이터셋 크기를 고려할 때 탄탄한 기준선입니다. 학습 진행 상황: 아래 훈련 플롯에서 확인할 수 있습니다: 그림 6: 30 에폭 동안의 훈련 손실과 테스트 정확도. 출처: 저자 다음은 모델의 예측 예시입니다. 고양이나 개구리와 같은 많은 샘플을 올바르게 식별했지만, 시각적으로 유사한 클래스 (예: 배를 비행기로 잘못 분류)에서는 어려움을 겪었습니다. 그림 7: CIFAR-10 이미지에서의 예측 예시. 출처: 저자 다음 막대 차트는 모든 10개 클래스에서 모델의 성능을 보여줍니다. 특히: 그림 8: 클래스별 정확도. 출처: 저자 5. 업계 전문가의 평가 및 회사이야기 비전 트랜스포머(ViT)의 등장은 딥 러닝 분야에서 혁신적인 변화를 가져왔습니다. 초기 ViTs는 CNNs와 경쟁하기 위해 거대한 데이터셋이 필요했지만, 주의 기반 모델이 비전 분야에서도 성공할 수 있다는 것을 증명했습니다. 이는 모델의 유연성과 병렬 처리 능력을 향상시키며, 장거리 의존성을 더 잘 포착할 수 있는 능력을 부여했습니다. Google의 연구팀은 ViT의 성공을 통해 이미지 처리 분야에서 주의 기반 모델의 가능성을 널리 알렸습니다. 이 모델은 이제 컴퓨터 비전의 핵심 구성 요소로 자리 잡았으며, 다양한 작업에서 우수한 성능을 보여주고 있습니다. ViT는 CNNs의 제약을 극복하고, 복잡한 시각 정보를 효과적으로 처리할 수 있는 새로운 접근 방식을 제공합니다.