An Image is Worth 16x16 words

Paper Review 3 : 모델링, 학습, 튜닝, 전이학습, 개선

|Intro

본 포스팅에서 리뷰할 논문은 이전까지 진행한 Transformer 논문 리뷰 및 구현의 포스팅에서 한 번씩 언급되었던 Transformer를 Computer Vision Task에 적용한 것입니다. 엄밀히 말해 Transformer 구조 전체가 적용된 것은 아니고, Transformer의 Encoder만이 활용된 형태입니다.

이를 통해 입력된 사진이 무엇인지 구분하거나, 분류하거나, 객체를 찾는 등의 task를 주로 수행하게 됩니다. 참고로 아직 이 분야에서의 워딩이 명확하게 정의된 것은 아니라, ViT라 칭하면 Vision Transformer의 시발점이 된 본 논문에서 제시한 모델 구조를 의미하기도 하지만, CV 영역에 적용된 Transformer를 칭하기도 합니다. 맥락에 따라 이해하는 것이 나을 듯 싶습니다.

추가로 CV 영역에서 Decoder가 사용되는 경우는 이미지 생성이나 고화질화(Super Resolution) 작업, 때로는 Semantic Segmentation을 위한 task를 수행하는 과정에서 수행되곤 합니다.

|Paper Summary

ViT 구조

ViT(Vision Transformer) 모델의 구조는 그리 복잡하거나 어렵지 않습니다. 기본적인 개념만 이해하면 그 뒤는 오히려 이전에 포스팅한 Transformer보다 간단하고 쉽습니다.

ViT의 발상은 Transformer가 기존에 자연어를 중심으로 대다수의 영역에서 SOTA(State of the Art) 수준의 성능을 보였기 때문에 이를 Computer Vision Task에도 적용해보자는 것이었습니다. 하지만 이 과정에서 발생하는 문제가 있습니다. 각각의 주요단어에 대해 사전을 구축하고, 이를 기반으로 문장을 여러 개의 토큰으로 분할하는 방식으로 모든 문장을 나눌 수 있었던 자연어에 비해 이미지는 어떤 기준으로 토큰으로 나누기 애매했던 것입니다. (특정 객체를 찾아 이를 토큰화한다고 해도, 이 과정에서 색종이 자르듯 핵심객체만 자르고 남은 영역 등에 대한 처리가 애매해지게 됩니다) 때문에 본 논문에서는 그냥 간단히 하나의 이미지를 지정된 사이즈로 나누는 방식으로 토큰화합니다. 

위 그림의 왼쪽 하단을 보면 한 장의 이미지가 분할되어 9장의 이미지가 되고 이를 순서대로 하나씩 모델에 넣는 다는 것을 표현하고 있습니다. 본 논문에서 채택하는 기본 구조는 이처럼 이미지를 지정된 사이즈로 나누는 것입니다. 이 외에 Hybrid Architecture라고 하여 단순히 이미지를 여러 장의 패치로 분할하고 이를 Linear Projection으로 넘기는 것 대신, 사전학습된 CNN 기반 모델로부터 출력한 feature map을 통해 분할하는 방식도 제시하고 있습니다. 다만 본 논문에서 메인 주제로 다루진 않았기에 여기서도 이렇게 간략히 소개만 하고 넘어가도록 하겠습니다.

ViT의 구조는 이미지를 여러 개의 Patch로 분할해 Linear Projection으로 Embedding하고, 클래스 레이블과 Positional Embedding 값을 Transformer Encoder로 넘겨줍니다. n회에 걸친 Encoder 반복 후 MLP Head와 분류할 class 수에 맞는 FC Layer를 통해 최종적인 예측으로 넘어가는 형태입니다. 이전 포스팅에서 구현한 자연어 모델과 다른 점 중 하나는 자연어 모델에 비해 Transformer 구조의 내부 파라미터가 더 크다는 것입니다.

ViT 모델의 파라미터

NLP Transformer의 경우, 기본값으로 6개의 Layer와 Head를 중심으로 구성했던 것에 비해 CV Transformer에서는 그 두 배에 해당하는 12개의 Layer와 Head로 구성되며, d_model의 사이즈도 1.5배가 크게 설정돼있습니다.

본 포스팅에선 이전 Transformer 논문에서 다뤘던 내용과 겹치는 부분이 많기 때문에 Transformer 구조에 대해선 최소 한도로 설명을 진행하며, ViT의 고유특징인 패치와 실제 학습과 파인튜닝 성능 비교 및 기타 실험 결과에 대해 다루도록 하겠습니다. 본 논문에서는 학습에는 224 사이즈 이상의 이미지로 구성된 수백만 장 단위의 이미지 데이터셋을 학습에 사용했습니다. 파인튜닝에는 그보다 작은 데이터셋을 사용하기도 했지만 컴퓨팅 리소스의 한계를 고려해 본 포스팅에선 이미지의 사이즈는 맞추되, 규모는 이보단 작은 데이터셋을 활용할 것입니다. 

본 포스팅에선 224 사이즈의 만 장이 조금 넘는 이미지 데이터셋인 Sports Dataset(Kaggle)을 활용할 것입니다. 본 포스팅을 위해 여러 데이터셋과 방법론을 적용해봤는데 결론적으로 이미지 자체의 사이즈가 최우선적으로 커야 최소한의 성능이 나오기 시작하고, 데이터셋의 규모가 커져야 충분히 성능이 높아진다는 결론을 내렸습니다. 이를 보완할 수 있는 방법은 크게 두 가지가 있었는데요, 첫 번째로는 ConvNeXt와 같이 CNN 계열의 모델을 활용하는 것입니다. 본 논문에서 제시한 Hybrid Structure를 구현하는 셈입니다. 이를 통해 inductive bias를 향상시켜 더 작고&적은 데이터에서도 잘 학습하도록 할 수 있었습니다. 

Vision Transformer for Small-Size Datasets에서 제시하는 방법

두 번째로는 'Vision Transformer for Small-Size Datasets'에서 제시한 방법을 적용하는 것입니다. 해당 논문에서는 한 장의 이미지를 단순분할해 patch로 만드는 과정에서 공간적 정보가 훼손되고, 질낮은 토큰화가 진행되는 것이 큰 데이터셋이 필요한 주 이유라고 설명하고 있습니다. 이를 해결하기 위해 위 그림에서 볼 수 있듯 한 장의 이미지에 padding과 shift를 적용해 여러 버전의 Patch 그룹을 형성하는 방법을 제안하고 있습니다. 다만 이러한 방법론은 더 작은 데이터셋에서 잘 학습할 수 있지만, 몇 배는 더 많은 VRAM을 사용하게 될텐데요. 또한 이 방법은 이후의 포스팅할 Swin Transformer의 개념과도 유사하기에 본 포스팅에서는 적용하지 않았습니다.

|Training & FineTuning Description

학습과 파인튜닝에 대한 설

본 논문의 주요 내용 중 하나는 다양한 실험을 해봤다는 것입니다. 기본적으로 224 사이즈의 큰 이미지(당시로서는) 데이터셋을 활용해 모델을 학습시킨 후, 이를 더 작거나 큰 데이터셋 등으로 파인튜닝해 학습하는 식입니다. 본 포스팅에서도 이러한 과정을 함께 진행해볼 것입니다. 이를 위해 최대한 본 논문에서 서술된 방식을 구현하되, 어떤 스케줄러를 어떻게 사용했는지, 데이터 증강 등을 어떻게 사용했는지와 같이 다소 애매한 부분은 본 논문의 내용을 저해하지 않는 선에서 직접 구현해 사용했습니다.

다만 수백만 이상의 데이터셋을 직접 학습시키는 것은 한계가 있기 때문에 본 포스팅의 목적을 저해하지 않도록 모델링과 이미지 사이즈는 본 논문의 사항을 그대로 유지하되, 데이터셋의 규모나 그 외 기타 방법론 등을 적용해 효율적으로 학습해보도록 하겠습니다. (본 논문에서 제시하는 방법론들에 대한 설명과 구현은 그대로 진행할 것입니다)

|Modeling

모델링 과정은 우선 torch에 구현된 encoder 등의 클래스를 불러와 전체 구조를 한 번에 살펴보는 것부터 시작해, Patch 분할 등 ViT의 구성요소를 하나씩 구현해보고, 이렇게 구현한 모델을 기반으로 학습과 파인튜닝을 진행해보는 식으로 가보도록 하겠습니다.

|Simple Structure

ViT 전체 코드

놀랍게도 위 코드가 ViT를 구현한 전체 코드입니다! 사실 이미 구현된 모델을 불러오거나 가중치까지도 불러오는데 한 줄이면 충분하지만, 그건 본 포스팅의 목적상 맞지 않기 때문에 이 정도 수준으로만 압축하겠습니다. 위 코드를 살펴보면 크게 네 부분으로 나눠져 있습니다.

[row 4~9]→[row32~37] 이미지가 입력되면 이를 지정된 사이즈의 Patch로 나누고, 채널 사이즈를 고려해 Linear Projection하는 과정을 거칩니다. [row11~15]→[row39~42] 선형변환된 이미지 패치들에 이미지의 클래스와 위치 정보를 추가하기 위해 클래스 cls_token과 pos_embedding 변수를 임의로 생성하고 추가합니다. [row17~21]→[row44~45] 패치 분할 및 선형변환된 이미지에 클래스와 위치 정보가 추가된 값을 Transformer Encoder Block을 반복적으로 통과시킵니다. [row23~29]→[row47~50] 이전까지의 과정을 통해 출력한 결과물 중 클래스 토큰에 대한 정보만을 가져와 마지막 MLP Head를 통과시켜 최종적인 결과물인 클래스 분류 결과를 출력하는 구조입니다.

본 논문에서도 NLP에 사용된 Transformer 구조를 최대한 지켰다고 하는데요, 다만 기존 Transformer와 달리 Dropout을 더 자주 적용했다고 합니다. 예컨데 본래 Transformer 논문에서는 MLP Layer에 대한 Dropout 언급이 없었다면 ViT에서 언급한 사항은 Linear 이후마다 Dropout을 적용하는 식입니다. 실제로 Transformer를 이용해 학습시키다보면 과적합이 문제가 되는 경우가 잦은데, 이를 해결하고자 함으로 비춰집니다.

|Structure Details 

위에서 간략하게 구현한 코드를 이번엔 하나씩 페이퍼에서 제시된 사항을 확인하며 구현해보도록 하겠습니다. 전체 코드는 여기를 확인하시면 됩니다.

|Patch Embeddings

Patch Embedding에 대한 이미지

본 논문의 가장 핵심적인 부분이라고 할 수 있는 Patch Embedding부터 시작해보겠습니다. Patch Embedding은 2차원, n개 채널(흑백은 1, 컬러는 3)을 가진 이미지를 NLP 처리를 위해 고안된 Transformer에 입력으로 넣기 위한 것입니다. 즉, Height x Weight x Channel의 형상(shape)으로 되어있는 이미지를 D(임베딩 차원 수) x N(HW/P^2)의 형상으로 바꿔줘야 한다는 것입니다. 이를 코드로 구현하면 아래와 같습니다.

ViT Patch Embedding Code

[row 18~20] 이미지 사이즈와 패치 사이즈를 입력으로 받습니다. 이미지 사이즈를 패치 사이즈로 나눈 값을 제곱해 전체 패치 개수(N)을 구합니다. 이때 이미지는 가로세로가 같은 정사각형 형태로 가정하며, 가로세로가 다르다면 이에 대해 각기 처리해줘야 합니다. 변환할 임베딩 차원 D는 입력으로 받습니다. 이를 통해 이미지의 사이즈와 채널을 입력받고, 이를 변형할 형상(D, N)까지 구할 수 있습니다.

[row 22~23]→[row27~30] 배치를 제외하면 3차원 데이터를 2차원 데이터로 만들기 위해 kernel size와 stride size가 patch size로 동일한 Convolution Layer를 통과시킨 후 flatten합니다. 이를 통해 (batch, channel, height, width) 형상을 (batch, embed_dim(D), n_patches(N))의 형상으로 변환할 수 있습니다. 여기까지의 코드를 통해 우리는 이미지 입력을 Transformer 입력으로 사용할 수 있는 형상으로 바꿀 준비를 완료했습니다. (참고로 특히 Transformer 구조를 구현할 때에는 데이터의 입출력 형상을 잘 정리해두는 것이 좋습니다.)

Patch Embedding에 대한 수식

본 논문에서 Patch Embedding에 대한 수식을 살펴보면 위 그림의 형광펜이 칠해진 부분입니다. 제일 앞에 클래스 토큰을 넣고, 그 뒤로는 Embedding tokens들이 구현돼 있습니다. 위 코드에서 cls_token을 추가하는 코드를 넣어보도록 하겠습니다.

cls_token이 추가된 patch embedding

Computer Vision Task 중에서는 class token을 활용하지 않는 경우도 많은데요, class token을 통해 정보 집약적인 층을 하나 더 추가하고, 최종적인 mlp layer에 이것만을 넘겨줌으로써 불필요한 연산을 최소화할 수 있지만 이 과정에서 많은 정보를 잃을 수도 있기 때문입니다. 하지만 경험상 단순 클래스 분류 문제의 경우에는 본 논문에서 소개된 바와 같이 class token을 추가해 활용하는 것이 더 성능이 높고 요구되는 컴퓨팅 리소스가 적습니다.

|Positional Embedding

이전 Transformer 논문 리뷰에서 이미지의 위치정보를 고정값으로 주는 방법과 학습가능한 파라미터로 주는 방법에 대해서 소개했었습니다. 고정값으로 주는 방법은 주기 함수인 sin함수와 cos함수를 이용해 이미지의 위상 정보를 전달하는 방식이었으며 학습되는 값이 아니고, 이를 Positional Encoding이라 칭합니다. 반대로 학습가능한 파라미터로 주는 방식은 Positional Embedding이라 칭하며 모델의 학습과정에서 위상 정보를 자체적으로 학습하게 하며 이에 따라 요구되는 컴퓨팅 리소스가 증가할 수 있지만 크게 유의미한 수준은 아니고, 성능 향상의 가능성이 있었습니다.

다양한 pe에 대한 실험

본 논문에서는 이에 더해 positional embedding을 아예 주지 않는 방법과 여러 차원으로 적용하는 방법, 상대적인 pe를 적용하는 방법에 대해 실험했습니다. 위 내용을 요약하면 위치 정보를 주는 방법에 따른 성능 차이는 크게 없지만, 주지 않는 것보다는 유의미한 성능 차이를 보인다는 것입니다. 이는 ViT 자체가 픽셀 단위로 세세하게 위상을 따져서 판단하는 것이 아닌 보다 큰 단위인 Patch로 작업하기 때문에 위상정보를 어떻게 제시하던 고려할 수 있는 여지만 준다면 크게 영향받지는 않는다는 것입니다. 이를 고려해 간단하게 Positional Encoding을 구현하도록 하겠습니다.

PE Code

본 논문에서는 pe의 타입이나 적용 유무에 따른 실험 결과만이 제시되었을 뿐, 명확하게 어떻게 사용했다에 대한 설명이 적습니다. 다만 본 논문에서 Transformer 구조를 최대한 유지했다라고 명시했기 때문에 위 코드에서 scale을 추가하겠습니다. scale을 통해 데이터의 범주가 확장되고, 이를 통해 더 안정적인 Atten tion Score를 얻을 수 있다는 내용 또한 이전의 포스팅을 통해 실험했으니 이를 아래와 같이 수정해 적용하겠습니다. (이러한 Scale은 일반적으로 ViT에서는 잘 쓰이지 않는 방법입니다만, 실험결과 사용하는 쪽이 효과가 미세하게 더 좋거나 비슷한 수준이었으며, NLP Transformer 구조를 최대한 지켰다는 본 논문의 내용에 따릅니다.)

Scale이 적용된 Positional Embedding code

참고로 Positional Embedding을 적용한다고 해서 출력 차원이 변화하거나 하지는 않습니다. 

|Encoder Layer

본 논문에서는 Transformer의 Encoder 구조만을 차용하고 있습니다. Encoder는 복수의 Encoder Layer가 쌓여서 구성됩니다. 이에 대한 수식은 아래와 같습니다.

Encoder Layer에 대한 수식

MLP는 GELU Activation을 통한 비선형성이 강조된 구조이며, 두 개의 선형 레이어로 구축돼 있습니다. 본 논문의 다른 부분에서는 선형 구조 이후마다 Dropout을 적용한다고 밝히기도 했으니 이를 적용토록 하겠습니다. 위 공식을 보면 z0는 class token과 Patch Embedding이 하나로 묶이고, 여기에 Positional Embedding이 합쳐진 결과물입니다. 이를 Layer Normalization을 적용한 뒤 Multi Head Attention 레이어를 거치고, 다시 정규화한 후 MLP 레이어 통과 후 Layer Normalization이 적용됩니다. 각 과정에서 Residual Connection이 고려되어야 합니다. 이를 코드로 살펴보면 아래와 같습니다.

Encoder Layer code

코드의 가독성을 위해 MSA 코드는 구현체를 응용하였습니다. 이 MSA에는 이전 포스팅에서 구현했던 Multi-Head Attention과 사전 정규화가 포함돼 있습니다. Encoder Layer를 보면 레이어마다 Dropout, LN, Residual Connection이 모두 구현돼있는 것을 볼 수 있는데요, ViT 논문에서 언급하는 사항과 공식들에서 과적합에 상당히 신경쓰고 있다는 것을 알 수 있습니다. 참고로 위 코드를 구현하는 과정에서 아래의 지표를 참고하였는데요, MLP의 size가 D의 4배씩으로 설정돼있는 것을 보고 고정적으로 4배를 적용할까 하다가 하이퍼 파라미터로 설정해두었습니다. 

ViT 하이퍼 파라미터 값

[23.12.08 추가] 모델의 전체 파라미터를 측정하기 위해 이전 포스팅에서 사용했던 MHA를 가져왔습니다. 정확히는 Mask가 제외된 버전입니다. 이에 대한 코드는 아래와 같습니다.

직접 구현한 MHA 코드

Transformer를 구현할 때, estimate_params 파라미터를 True로 지정하면 위 클래스를 이용하면 그렇지 않을 경우, torch 구현체를 사용합니다. 기본값은 False입니다. 두 방식 모두 테스트했으나 성능 등의 차이가 존재한다고 보기 어려웠습니다.

|Transformer

Patch Embedding과 Positional Embedding, Encoder Layer가 구현됐다면 남은 것은 이들을 적절히 묶어주는 것입니다. 

Patch Embedding, Positional Embedding, Encoder Layer를 이용해 구성한 모델 구조

제일 처음 봤던 것과 같이 모델 구조는 입력된 이미지를 패치로 나누며 클래스 토큰을 추가하고, 위치 임베딩을 더한 뒤, Encoder Layer를 반복적으로 순환하고, 이를 정리(norm+dropout)하고, 클래스 집약적인 정보가 담겨있을 cls_token에 해당하는 첫 번째 레이어를 마지막 MLP Head로 넘겨주는 것입니다. 이렇게 모델을 구축했다면 테스트 데이터를  넣어 출력되는 형태가 정상적으로 나타나는지를 살펴보아야 합니다.

모델의 출력 테스트

우리는 이후의 학습 과정에서 224x224 사이즈의 sports 데이터셋에 학습한 뒤, 이를 32x32 사이즈의 cifar10 데이터셋에 파인튜닝하는 과정을 거칠 것입니다. 이때 사전학습된 공식 모델을 불러와 학습 과정과 성능 등을 테스트해보고, Attention Map을 그리며 본 포스팅을 마무리 짓도록 하겠습니다.

참고로 이렇게 구현한 모델 구현체를 torchsummary를 통해 요약해보면 다음과 같습니다.

구현한 모델의 파라미터 수(좌)와 모델에서 제시된 모델 변주별 파라미터 수

이는 torch summary가 직접 구현한 클래스에서만 정상적으로 파라미터를 측정하고, 사전 구현체(nn.Multiheadattention)를 사용할 경우  그 내부에서 사용되는 파라미터를 측정하지 못하기 때문입니다. 만일 이전에 구현했던 사전 정규화(LN), 멀티헤드어텐션(MHA)을 직접 구현할 경우 다음과 같이 전체 파라미터가 측정되게 됩니다.

직접 구현된 MHA를 적용했을 경우

총 8,570만 개의 파라미터를 가지고 있어 본 논문에서 반올림해 제시한 8,600만개와 같은 수준임을 알 수 있습니다.

|Scheduler

본 논문에서 제시하는 스케줄러는 초기 웜업 단계 이후 학습율을 조금씩 떨어뜨리는 형상을 띄고 있습니다. 이는 이전 포스팅에서 제시한 Noam Scheduler와도 유사한 모습입니다.

학습과 관련된 하이퍼 파라미터 등

위 그림에서 제시한 LR decay method들을 살펴보면 linear이거나 cosine입니다. cosine 함수가 적용된 스케줄러는 일정한 주기를 기준으로 학습율이 최초 상태의 일정 비율로 돌아오도록 합니다. 이를 그림으로 살펴보면 아래와 같습니다.

cosine learning rate graph

위 그림은 주기를 전체 스탭의 1/4로 설정한 것입니다. 이로 인해 최초 학습율 0.001에서 0.0005로 떨어지고 올라오는 과정을 4번 거치게 됩니다. 주기를 1로 설정한 cosine scheduler와 선형으로 감소하는 scheduler의 learing rate를 비교하면 아래와 같습니다.

cosine decay vs linear decay

본 논문에서는 cosine scheduler의 주기를 명확하게 언급하지 않았고, 만약 전체 step을 하나의 주기로 잡으면 위 그림과 같이 linear 스케줄러에 비해 다소 완만한 곡선을 그릴 뿐입니다. 때문에 여기서는 전체 스탭을 하나의 주기로 잡고, warmup과 cosine이 결합된 스케줄러를 적용하도록 하고, 이를 정의해보도록 하겠습니다.

scheduler code

위 코드는 warmup step까지는 선형적으로 학습율을 올리다가 warmup step 이후부터는 cosine 함수에 기초해 학습율을 떨어뜨립니다. 주기는 전체 스탭으로 했기 때문에 다시 최초 학습율로 돌아오지는 않습니다. 이에 대한 학습율 그래프를 그려보면 아래와 같습니다.

lr with scheduler graph

참고로 본 논문에서 제시한 warmup step은 10,000 step으로 ImageNet(1,200,000)과 ImageNet-21k(14,000,000)의 데이터 사이즈를 고려하면, ImageNet의 경우 전체 스탭이 1,200,000(Data Size)÷4096(Batch Size)×300(Epoch) = 87,891입니다. 이 중 warmup step의 비율을 구하면, 

10,000(Warmup) / 87,891(Total Step) = 약 11.4%입니다. ImageNet-21k의 전체 스탭은 14,000,000(Data Size) ÷ 4096(Batch Size) × 30or90(Epoch) = 102,540or307,617입니다. 이중 warmup step의 비율은 10,000(warmup) / 102,540or307,617 = 약 10%or3.3%입니다. 즉, 이러한 warmup step은 학습에 있어서 모델의 성능이 충분히 올라오기 전에 데이터셋에 대해 과적합하는 것을 방지하는 효과가 있기에 채택되지만, 데이터셋과 모델의 특징을 고려해야 한다는 것으로 보여집니다.

본 포스팅에서는 약 27,000 Step의 학습을 진행하였습니다. 다만 데이터셋이 작아 Warmup 및 cosine scheduler를 사용할 경우 충분한 성능이 나오지 않았습니다. 여기서 충분한 성능의 기준은 제가 임의로 잡았으며, Train/Valid/Test 데이터셋으로 구분하는 Holdout을 적용 후 학습했을 때 Test 데이터셋에 대해 정확도와 F1-Score가 모두 0.4 이상은 되어야 한다고 보았습니다. 이러한 기준을 잡기 위해 HuggingFace의 Timm 라이브러리에 구축된 ViT 구현체를 불러와 테스트하며 평가하였습니다.

|Optimizer & Loss

학습을 위한 하이퍼파리미터 테이블

본 논문에선 과적합 방지를 위해 label smoothing, dropout, normalization 등의 기법을 적극적으로 활용했다는 언급만 있을뿐 Loss에 대한 별다른 언급은 없습니다. 그렇다면 이 또한 Transformer 논문의 내용을 준수해 label smoothing 0.1이 적용된 CrossEntropyLoss를 사용하도록 하겠습니다. 또한 Optimizer는 Adam을 사용했습니다. 본 논문에서는 learning rate와 weight decay는 위 그림과 같이 모델과 데이터 사이즈마다 다르게 적용되었고, 일반적으로 학습 데이터셋의 규모와 이미지 사이즈 등을 고려해 위 그림에서 색칠된 부분의 옵션(다소 높은 lr과 decay)을 사용했습니다. 하지만 저희가 사용할 데이터셋은 만 개가 조금 넘는 아주 작은 데이터셋으로 Transformer 구조를 학습시키기엔 무척이나 적은 양입니다. 

weight decay에 따른 모델의 가중치 변화

위 그림에서 보다시피 weight decay가 높아질수록 모델의 가중치 업데이트가 빠르게 규제화가 이뤄지게 됩니다. 이는 큰 데이터셋에서는 긍정적으로 작용할 수 있지만, 저희가 사용할 작은 데이터셋에서는 부정적인 영향을 보이기 쉽습니다. 또한 우리가 사용하는 Adam Optimizer는 Weight Decay가 정상적으로 이뤄지지 않는 단점이 있고, 이를 보완하기 위한 새로운 optimizer인 AdamW가 나와있습니다. 다만 이에 대해선 본 논문의 취지와 무관하니, 이후의 포스팅에서 다뤄보도록 하겠습니다.

따라서 저희는 dropout은 적용하되, weight decay는 적용치 않도록 하겠습니다. 

Optimizer Betas

참고로 사용한 Optimizer 또한 본 논문의 위 그림과 같이 Adam을 사용하였으며, betas는 아래 그림의 b1과 b2를 의미합니다. 이는 Adam Optimizer의 기본값이라 별도로 지정하지 않아도 괜찮습니다.

optimizer, loss, scheduler code

위 코드 이미지에서 주석처리한 optimizer와 scheduler가 본 논문에서 사용된 방식입니다. 하지만 이는 작은 데이터셋에 대해 불리하게 작용하는 것을 실험적으로 알아냈기 때문에 아래의 방식으로 학습하였습니다. 또한 적은 데이터셋의 한계를 극복하기 위해 다양한 증강기법을 적용했기 때문에 선형으로 하강하는 scheduler 대신 일정 기준(STEP, 성능)마다 성능이 떨어지는 scheduler인 StepLR과 ReduceLROnPlateau를 사용하는 게 더 성능이 좋았습니다.

|Train etc

학습 파트의 전체 코드는 여기를 참조하시면 됩니다.

다음은 모델 학습입니다. 우선 본 포스팅을 위해 다양한 시도를 해봤는데요, 시도한 것들은 아래와 같습니다.

---

1. cifar10/100 데이터셋을 64,128,224 사이즈로 늘려(bilinear, lanczos) 학습

2. l2l의 Mini ImageNet, Standford Univ의 Tiny ImageNet을 본래 사이즈 및 사이즈 키워(bilinear, lanczos) 학습

결론적으로 최소한 기본적인 ViT 구조로는 충분한 학습이 이뤄지지 않았습니다. 이를 검증하기 위해 자체 구현한 모델링 외에도 huggingface의 timm 라이브러리의 구현체와 비교하며 학습을 진행하였습니다. 구현체의 가중치를 가져와 파인튜닝할 때에는 성능이 높았지만, 구현체라 하더라도 처음부터 학습할 때에는 직접 구현한 모델과 성능차이가 없었습니다. 모든 학습은 10,000의 warmup step을 유지하기 위해 50,000 step 이상 진행할 수 있도록 진행하였습니다.

하지만 현실적으로 컴퓨팅 리소스를 다소 비효율적으로 사용하는 ViT를 이용해 본 논문에서 제시한 바대로 ImageNet을 학습시키기에는 컴퓨팅 리소스와 시간이 여의치 않으므로, Kaggle에서 가져온 sports 데이터셋을 이용해 학습해보도록 하겠습니다. 해당 데이터셋은 224x224 해상도의 이미지 데이터셋으로 100종에 달하는 스포츠를 분류하는 데이터셋입니다. Transformer를 훈련하기에는 학습 데이터셋 규모가 13,493장으로 적지만 데이터 증강을 통해 이 부분을 해소하는 식으로 진행해보도록 하겠습니다.

|Model Update

작은 데이터셋에 Transformer Model을 학습시키려다보니 성능을 위한 추가적인 방법론의 적용이 필요했습니다. 최초에는 ViT의 depth/head/dim을 줄임으로서 더 작은 데이터셋에 최적화되도록 유도해봤고 어느정도의 성과는 보였지만 만족할만한 수준은 아니었습니다(수치로 나타내면 정확도 기준 0.2~0.3 정도에 불과했습니다). 

때문에 모델의 성능을 위해선 ViT의 모델이 깊어야 했고(depth), 다양한 관점을 바라봐야 했으며(head), 넓어야(dim) 했습니다. 하지만 이러한 문제는 모델의 학습 속도를 느리게 만들뿐 아니라, 성능이 충분히 높아지기도 전에 포화돼버리고 말았습니다. 이를 해결하고자 속도와 성능적 측면에서 본 논문에서 제시되지 않은 기법들을 적용했습니다.

우선 성능 측면을 위해 아래의 두 가지 방법론을 적용했습니다.

|Performance Side #1: LayerScale

LayerScale은 이전 Transformer 논문 리뷰에서도 사용했었던 Positional Embedding을 더하기 전에 원본 값에 scale을 더해주던 개념을 확장한 것이라고 보면 됩니다. 이러한 Scale을 각 transformer block의 출력에 곱함으로써 이전 포스팅에서 봤던 것과 같이 transformer 모델의 안정성과 성능을 증가시키는 효과가 있습니다. 이 기법은 Facebook에서 낸 'Going deeper with Image Transformers'[5] 논문에서 처음 언급되었습니다. 해당 논문에서는 LayerScale 외에도 Class Token을 나중에 추가해 Class Token 위한 별도의 Attention 프로세스를 적용하는 등의 방법도 제시하기도 하였지만, 본 포스팅의 주제에 벗어나니 LayerScale만 간략히 보고 넘어가도록 하겠습니다.

LayerScale에 대한 설명[5]

위 그림 중 오른쪽은 Layer Normalization 등의 적용차이에 따른 LayerScale의 적용 방식을 소개하고 있고, 결론적으로 가장 오른쪽의 (d)가 그들이 실험결과 평가한 가장 좋은 성능을 내는 구조였습니다. Layer Normalization 이후 FFN이나 Self Attention이 진행되고, 그 출력값에 대해 Layer Scale을 하는 것입니다. 이러한 Layer Scale은 왼쪽 그림의 강조된 부분과 같이 학습가능한 파라미터여야 합니다. 이를 코드로 구현하면 아래와 같습니다.

Layer Scale의 코드

|Performance Side #2: DropPath

두 번째로 추가한 것은 DropPath입니다. dropout이 레이어 내부의 특정 뉴런을 무작위로 비활성화하는 것이라면 DropPath는 여러 레이어로 구현된 블록 혹은 네트워크에서 특정 레이어를 무작위로 비활성화하는 것입니다. 이는 dropout보다 큰 단위의 생략을 가능하게 하고, 이를 통해 깊은 레이어로 구성된 모델 구조에서 일반화 성능을 높이는데 기여합니다. (dropout과 동일하게 학습 단계에서만 작동하며, 추론, 평가 시에는 dropout이 레이어 내에서 일정 확률로 적용되는 방식과 달리 dropPath는 아예 비활성화된다는 차이가 있습니다.)

DropPath의 설명[6]과 그림[7]

DropPath는 2016년 'Deep Networks with Stochastic Depth'[6]라는 논문에서 제시되었습니다. 위의 왼쪽 그림이 해당 논문에서 소개하는 DropPath에 대한 설명이고, 이에 대한 그림은 FracktalNet이라는 모델 구조를 제시한 후속 논문에서 가져왔습니다. 해당 논문의 이름은 'FractalNet: Ultra-Deep Neural Networks without Residuals'[7]입니다. 이를 구현한 코드는 아래와 같습니다.

dropPath 코드

코드는 그리 복잡하지 않습니다. 복수의 레이어가 병렬로 처리되는 모델 구조에서 drop 확률을 이용해 입력 텐서와 동일한 형태의 binary matrix를 생성하고, 이를 출력값에 곱해주는 방식입니다. 이를 통해 일부 레이어 출력을 0으로 만드는 식이며, 8~9번째 줄의 코드를 통해 drop 확률이 0이거나 학습 시가 아닌 경우에는 원본 출력을 그대로 출력하게 됩니다.

참고로 이러한 DropPath를 적용했을 경우, 학습 과정에서만 일부 레이어 출력이 생략되고 제한된 정보로 추론하게 됩니다. 이러한 Drop 계열의 기능들은 학습에서만 온전하게 기능하고, 그 외 평가나 추론 단계에서는 제대로 작동하지 않는다는 특징을 가지고 있으며, 이는 엄밀히 말해 이전에 모델 최적화 포스팅에서 사용한 가지치기(Pruning)와는 다르게 제한된 정보/환경 하에서 온전히 기능하는 뉴런/레이어의 성능을 높이는 방식입니다. 이로 인해 train loss가 val loss보다 지속적으로 높게 나타나거나 성능이 극대화되는 시점에 가서야 수렴하거나 더 낮아지는 경향을 보이기도 하는 것 또한 DropPath의 특징입니다.

|Performance Side #3: Data Augmentation

데이터 수가 부족하기 때문에 데이터 증강은 필수적이었습니다. 최종적으로 사용한 증강기법에 대한 코드는 아래와 같습니다.

데이터 증강 코드

sports 데이터셋은 ImageNet과 유사한 데이터 분포를 보이고 있기 때문에, ImageNet과 동일한 Normalize를 적용하였습니다. 그 외에는 이미지의 임의의 영역을 자르고 크롭한 뒤 본래의 이미지 사이즈에 맞추는 RandomResizedCrop, 좌우반전을 위한 RandomHorizontalFlip, 그리고 무작위로 이미지의 일부 영역을 지우는 RandomErasing을 적용하였습니다. 이러한 증강 기법의 적용은 데이터를 변환하는데 추가적인 시간을 소요되도록 한다는 단점이 있지만, 그만큼 적은 데이터로도 모델의 성능을 끌어올릴 수 있다는 장점이 있습니다.(이 이외에도 다양한 증강을 적용해봤지만, 큰 의미가 없거나 성능이 오히려 하락하기도 하였습니다)

|Training Speed Side #1: Fused Attention

이번에는 모델의 성능보다는 학습 속도를 개선하기 위한 과정들입니다. 첫 번째로 적용한 것은 Fused Attention인데요, 기존의 Attention보다 빠르고 메모리 사용량이 적습니다. 이는 하드웨어에 따라 최적화되는 성능이 다른데요, 이를 직접 구현하는 것보다는 torch에 구현된 attention을 적용하면 자동으로 본인의 하드웨어에 맞는 최적화가 실행되게 됩니다. 이를 코드로 구현하면 아래와 같습니다.

fused attention을 적용하기 위한 코드

위 코드는 기존에 구현한 Multi Head Attention에 구현된 코드 일부입니다. 위의 if문으로 넘어가는 식으로 진행되면 torch의 scaled dot product attention이 적용되고, 이를 통해 자동으로 적용가능한 Fused Attention 등이 탐색되고 적용됩니다. 이와 관련된 정보는 아래 이미지 혹은 그 출처[8]를 살펴보시기 바랍니다.

scaled dot product attention 적용 시의 note[8]

|Training Speed Side #2: Automatic Mixed Precision(AMP)

AMP는 기존에 float32 데이터 타입으로 연산하던 데이터, 모델 구조 및 가중치 등에 대해 float16으로 변환가능한 것은 변환 후 진행하는 방식입니다. 이를 통해 메모리 사용량을 줄이고, gpu 학습 속도를 높일 수 있습니다.  이를 적용하기 위해서는 아래와 같이 학습 코드를 변경해주어야 합니다.

train_better.py 코드의 일부

메모리 효율성과 속도를 향상시키기 위한 두 가지 방법을 적용한 결과에 대해 정리하면 아래와 같습니다.

speedup을 위한 method 적용 결과

위 지표는 nvidia a100 80gb 환경에서 sports dataset을 이용해 테스트되었습니다. 배치 사이즈의 경우, 간단히 256/400/512로만 테스트했고, 3에포크를 돌려 평균을 냈습니다. 조금 더 정밀하게 테스트할 경우에는 위 지표보다도 더 높은 수준의 성능 향상이 가능하기도 합니다만, 여기선 간단하게만 테스트하고 정리했습니다. 이 실험에 대한 코드는 여기를 참조하시면 됩니다.

|Training

모델의 학습 에포크는 1,000, 총 학습스탭은 27,000입니다. 위에서 언급했던 바와 같이 warmup을 적용하지 않은 StepLR과 ReduceLROnPlateau 스케줄러를 활용하였습니다. 또한 성능 비교를 위해 huggingface의 timm을 통해 불러온 구현체와 같이 학습시켜보았습니다. 

직접 구현한 vit의 성능

timm의 vit 구현체 성능

위의 모델 성능은 다소 과적합돼있습니다. 그 이유는 위 평가지표는 1,000 에포크 째에 학습된 모델로 평가한 것이기 때문이며, 실제 valid/test에 좋은 성능을 보이는 모델은 이보다 이전 epoch에서 저장돼었기 때문입니다. 실제로 보다 이전의 모델로 테스트했을 경우엔 성능이 약간 더 높습니다.

|Label Smoothing 조정

본 학습을 위해서 다양한 실험을 진행했었는데요, 성능을 최대한으로 높이는 과정에서 어떤 방법을 쓰더라도 모델의 성능이 일정 수준(1)에 수렴하는 경향을 발견했고, 이는 label smoothing으로 인한 레이블 노이즈라고 분석하였습니다. 이를 분석하기 위해 학습과 테스트에 동일한 데이터셋을 이용했고, 모든 문제와 정답을 알고 있는 상황에서조차 모델의 성능은 아래와 같았습니다.

label smoothing 0.1일 때의 성능

이는 레이블의 개수가 적을 때 발생하는 문제로, 만약 우리가 풀 문제가 실제 ImageNet과 같이 수만 개 이상의 레이블이었다면 label smoothing은 일반화 성능을 높이는 역할을 할 수 있었겠지만, sports 데이터셋과 같이 레이블의 개수가 수천 개조차도 되지 않는 경우에는 오히려 모델의 성능을 하락시키는 결과를 낳았습니다. 이를 다시 검증하고자 label smoothing을 0.0으로 두고 진행했을 때에는 아래와 같은 지표가 관찰되었습니다.

label smoothing 0.0일 때의 성능

위 그림을 살펴보면 훨씬 적은 epoch임에도 loss가 훨씬 빠르게 떨어지는 것을 확인할 수 있습니다. 하지만 confusion matrix상 성능은 동일한데요, 이는 label smoothing의 적용여부와 무관하게 해당 모델이 이미 문제와 답을 모두 아는 상황에서 성능을 최대한으로 학습했기 때문입니다. 이러한 label smoothing issue는 기존과 같은 방식으로 데이터셋을 나눠 학습하는 holdout과 같은 방법론을 적용했을 때 문제가 더 커질 수 있습니다. 그 이유는 loss가 적절히 줄어들지 않아, 이로 인해 역전파가 온전히 일어나지 않기 때문에 발생하는 것으로 유사한 성능을 낼 수도 있지만, label smoothing이 존재할 경우 더 많은 epoch의 학습이 필요했습니다. (cifar10, cifar100, sports 데이터셋을 이용해 테스트해본 결과를 바탕으로 한 분석입니다)

결론적으로 Label Smoothing을 적용해도 적절한 scheduler와 조금 더 많은 epoch가 제시된다면 모델 성능이 충분히 나왔으며, 일반화 성능은 더 높기도 하였습니다. 하지만 컴퓨팅 리소스가 제한적이라면 Label Smoothing을 사용하지 않는 것도 하나의 방법으로 보입니다.

|Attention Map

학습시킨 모델의 attention map을 출력하면 아래와 같습니다. 이와 관련된 코드는 여기를 참조하시면 됩니다. 아래 이미지는 모든 encoder layer들의 attention map을 합쳐서 표현했습니다. 만일 코드가 궁금하시거나 각 레이어별 attention map까지 보고 싶으시다면 github의 원본코드를 참조하시면 됩니다.

attention map visualization 1

위 이미지는 수구 이미지인데요, 학습시킨 모델이 왼쪽의 그림 중 어디에 주목하고 있는지를 시각화한 것입니다. 오른쪽 그림에 보면 사람 혹은 사람들의 복장과 공에 주목하고 있는 것을 볼 수 있습니다. 다른 이미지로 확인해보겠습니다.

attention map visualization 2

여성 피겨스케이팅 사진입니다. 본 데이터셋에서 가장 어려운 클래스 중 하나인데요, 그 이유는 본 데이터셋에는 남성 피겨스케이팅, 여성 피겨스케이팅, 듀오 피겨스케이팅이 별개의 클래스로 존재하기 때문입니다. 따라서 학습된 모델은 위 이미지의 종목이 피겨스케이팅임을 구분하는 것을 넘어 솔로인지 듀오인지, 성별은 어떤지를 구분해야 합니다. 이를 구분하기 위해서인지 위 그림에서 보다시피 피겨스케이팅을 하는 사람을 주목하고 있는 것을 확인할 수 있습니다.

|FineTuning

파인튜닝이라는 것은 결국 잘 학습시킨 모델의 백본을 이용해 비슷한 도메인의 데이터셋에 학습시킬 때, 학습에 필요한 리소스를 줄이는 행위입니다. 우리는 224 사이즈의 sports 데이터셋에 모델을 학습시켰고, 이를 통해 학습시킨 모델은 이미지를 분류하기 위해 이미지로부터 주요 특징을 추출하는 능력을 갖추게 되었습니다. 이를 가시적으로 확인한 것이 위의 attention map입니다. 그렇다면 이 모델을 파인튜닝해본 뒤 성능을 확인해보도록 하겠습니다.

|Fine Tuning Method #1: Img Size - 224x224

전체 코드는 여기를 참조하시면 됩니다.

첫 번째 파인튜닝 방법입니다. 이는 기존에 학습한 모델의 구조에 따라 데이터를 변형해 사용하는 것입니다. 이는 모델의 구조를 변경하지 않아도 되는 편리함이 있지만, 원본 데이터를 변형시켜야 하는 부담이 있습니다. 일반적으로 낮은 해상도의 이미지를 높은 해상도로 높여서 작업할 경우 성능이 높아지는 경향이 있지만, 데이터셋의 특징에 따라 달라질 수도 있으니 주의해야 합니다.

다만 이 경우에는 우리가 분류할 클래스의 수에 따라 최종 출력층(head)를 수정해줘야 합니다. 이에 대한 코드는 아래와 같습니다.

모델을 불러와 head를 수정하는 코드

우선 구현체의 성능을 보기 전에 공식 구현체의 성능을 확인해보겠습니다. 크게 두 가지 버전이 있는데요, 하나는 사전에 학습되지 않은 버전을 최초 학습하는 것이고, 다른 하나는 sports 데이터셋에 학습시킨 후 cifar10 데이터셋에 파인튜닝하는 것입니다.

공식 구현체 학습 성능

왼쪽 그림이 사전구현된 모델이지만 가중치가 초기화된 모델이고, 오른쪽 그림이 사전학습된 가중치까지 로드된 모델의 파인튜닝 성능입니다. 즉, 왼쪽 그림은 사전구현체가 cifar10 데이터셋에 최초로 학습된 것이고, 오른쪽 그림만 파인튜닝한 결과입니다. 둘 다 10번의 epoch 동안 학습을 진행했고, 사전학습된 오른쪽 파인튜닝한 경우가 빠르게 성능이 최적화되었습니다.

다음으로는 직접 구현한 모델의 경우입니다.

직접 구현한 모델의 성능

직접 구현한 모델 또한 위와 같이 왼쪽이 최초 학습된 모델이고, 오른쪽이 사전학습된 가중치를 불러와 파인튜닝하는 경우입니다. 이 또한 파인튜닝한 쪽이 빠르게 성능이 최적화됩니다.

마지막으로 사전학습된 모델 간 confusion matrix 비교입니다.

224x224 사이즈로 파인튜닝한 결과. 공식구현체(좌), 자체구현체(우)

본래 구현된 모델 구조로 파인튜닝했을 경우, 본 포스팅을 통해 자체적으로 구현하고 학습시킨 모델보다 공식 구현체의 성능이 더 높게 나왔습니다. 이는 사전학습된 데이터의 규모 등의 차이로 인한 것으로 보여집니다.

|Fine Tuning Method #1: Img Size - 32x32

전체 코드는 여기를 참조하시면 됩니다.

파인튜닝의 두 번째 방법은 이미지의 사이즈에 따라 모델의 구조를 변경하는 경우인데요, 이 경우에는 단순히 클래스 수에 따라 최종 출력층(head)만 수정하는 것뿐 아니라 입력층 및 그와 관계된 레이어들도 수정해줘야 합니다. ViT의 경우에는 입력층 및 그와 관계된 레이어가 두 개가 있는데요, 하나는 Patch Embedding이고 다른 하나는 Positional Embedding입니다. 이에 대한 코드는 아래와 같습니다.

구현체의 입출력층을 수정하는 코드

이 경우 또한 공식 구현체와 직접 구현체의 성능을 비교해보고 본 포스팅 마치겠습니다.

32x32 사이즈로 파인튜닝한 결과. 공식구현체(좌), 자체구현체(우)

새로운 사이즈에 대한 파인튜닝을 위해 입출력층을 모두 바꿨을 때에는 근소하지만 자체구현한 모델의 성능이 더 우수했습니다. 이는 사전구현체와 직접 구현한 모델의 백본인 Transformer Encoder의 성능이 비슷하기 때문으로 해석됩니다. 이미지 사이즈가 본래와 같았던 224 사이즈로 키웠을 경우보다 둘 다 성능이 낮으며, 사이즈를 키웠을 경우엔 최종 출력층인 head와 Patch Embedding, 그리고 Positional Embedding에 존재하던 사전학습된 가중치가 이미지를 올바르게 해석하는데 더 긍정적인 역할을 할 수 있었던 것으로 분석해볼 수 있을 것 같습니다.(위의 224 사이즈 파인튜닝에서 공식 구현체보다 성능이 떨어지는 이유도 여기에 있다 볼 수 있을 것 같습니다.)

|Reference

[1] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, Neil Houlsby. An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale. https://arxiv.org/pdf/2010.11929v2.pdf

[2] ViT-pytorch. visualize_attention_map. https://github.com/jeonsworld/ViT-pytorch/blob/main/visualize_attention_map.ipynb

[3] Seung Hoon Lee, Seunghyun Lee, Byung Cheol Song. Vision Transformer for Small-Size Datasets. https://arxiv.org/pdf/2112.13492.pdf

[4] HuggingFace. Vision Transformer. https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/vision_transformer.py

[5] Hugo Touvron, Matthieu Cord, Alexandre Sablayrolles, Gabriel Synnaeve, Hervé Jégou. Going deeper with Image Transformers. https://arxiv.org/abs/2103.17239

[6] Gao Huang, Yu Sun, Zhuang Liu, Daniel Sedra, Kilian Weinberger. Deep Networks with Stochastic Depth. https://arxiv.org/abs/1603.09382

[7] Gustav Larsson, Michael Maire, Gregory Shakhnarovich. FractalNet: Ultra-Deep Neural Networks without Residuals. https://arxiv.org/pdf/1605.07648v4.pdf

[8] torch documentation. scaled_dot_product_attention. https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html

|Log

2023.12.05 | Paper Reading, Test code 작성, 컨셉 기획

2023.12.06 | 초안 작성, Simple ViT 모델링

2023.12.07 | 전체 모델 구조 구현, 구현체와 성능 비교 -> 모델 학습

2023.12.08 | 모델 학습 결과 확인 및 수정보완 -> 파인튜닝

2023.12.09 | 파인튜닝 결과 확인 및 비교, Attention Map 구현

2023.12.10 | 새로운 데이터셋 모색(cifar100→tiny Imagenet) 및 적용, 학습

2023.12.11 | 새로운 데이터셋에 맞게 포스팅 내용 수정

2023.12.12 | 모델 학습 성능 평가 → Bad, 대안 모색

2023.12.13 | Sports Dataset에 대해 학습, 작은 데이터셋에 최적화하기 위해 모델 구조 개선 등

2023.12.14 | 적용한 기법 등 정리, 비교

2023.12.15 | Attention Map 새로 그리고, FineTuning 결과 비교, 포스팅 정리

2023.12.18 | 최종 검토