SOTA 알고리즘 리뷰 4 - ViT-G/14

안녕하세요! 빅웨이브에이아이의 박정환입니다. 

 

지난 SOTA 알고리즘 포스팅 글은 정형 데이터 분석 분야의 TabNet에 대해서 알아보았는데요,

 

<혹시나 못보신 분들을 위한  링크 https://bigwaveai.tistory.com/7>

 

이번 시간에는 이미지 분류 분야의 SOTA 알고리즘인 ViT(Vision Transformer)-G/14에 대해서 알아보도록 하겠습니다!

 

 

Key Idea

 

트랜스포머(Transformer)란 멀티헤드 어텐션(Multi-Head Attention) 기법이 적용된 인코더 디코더 구조의 모델을 의미합니다.

 

ViT 모델은 트랜스포머 방법론을 이미지 분석 분야에 적용한 것입니다.

 

ViT는 이미지를 고정된 사이즈로 분할하여 선형으로 모델에 입력한 후 포지션 임베딩 및 벡터 시퀀스를 모델의 인코더에 입력합니다.

 

분류 Task를 위해 추가적으로 클래스 토큰을 시퀀스에 추가하는 방식입니다.

 

ViT Architecture

 

ViT 방법론은 구글에서 꾸준하게 연구 중이며, ViT-G/14는 ViT-L/16, ViT-H/14 등과 같이 개발 버전 중 하나입니다.

 

ViT-G/14는 가장 최신의 ViT 알고리즘으로 현재 ImageNet 데이터셋에서 가장 높은 성능을 달성하고 있습니다.

 

ViT-G/14는 한마디로 "ViT 방법론의 응용 버전"이라고 보시면 되겠습니다.

 

모델링

 

ViT-G/14는 총 6가지 특징을 가지고 있습니다.

 

1. "헤드(Head)"를 위한 독립(Decoupled) Weight Decay

 

Loss값이 단순하게 작아지는 경우 특정 가중치의 값들이 커지면서 과적합 문제가 발생할 수 있습니다.

 

WD(Weight Decay)란 과적합 문제를 해결하기 위해  위의 가중치 값을 보정하는 기법입니다.

 

본 논문에서는 딥러닝 네트워크 중 최종 선형 레이어를 헤드(Head)라고 하고, 나머지 가중치를 바디(Body)로 명명했습니다.

 

일반적으로 WD는 전체 모델의 가중치에서 적용이 되지만, 본 논문에서는 헤드의 가중치를 따로 WD를 적용했을 때, Few-shot 테스트에서 큰 성능 향상을 관찰했다고 합니다.

 

Few-shot 테스트란 데이터가 극단적으로 적은 경우에도 모델이 학습을 잘 할 수 있는지 실험하는 것입니다.

 

본 연구에서는 이 현상의 원인을 제대로 파악하지 못했지만, 직관적으로 헤드에 WD를 강화하는 것이 클래스 예측 시 더 큰 마진(Margin)을 가지는 표현성(representations)을 생성한다고 합니다.

 

다음의 그림은 헤드 WD의 효과를 시각화한 것입니다.

Weight Decay 효과

 

 

2. 클래스 토큰 제거를 통한 메모리 절약

 

메모리 오버헤드란 어떤 처리를 하기 위해 들어가는 간접적인 처리 시간 · 메모리 등을 말합니다.

 

기존 ViT 모델의 경우 TPU 하드웨어적 문제로, 최대 50%의 메모리 오버헤드가 발생할 수 있다고 합니다.

 

논문에서는 이를 해결하기 위해, 추가 클래스 토큰을 투입해 비교해보았습니다.

 

그림을 보면,  새로 투입된 GAP(Global Avreage Pooling)와 MAP(Multihead Attention Pooling)가 메모리 효율성이 훨씬 높다는 것을 확인할 수 있었습니다.

 

GAP는 각 Feature Map 상의 노드값들을 모두 더해 평균을 뽑아내는 것을 말합니다.

 

MAP는 멀티헤드 어텐션을 적용해 중요한 값들의 Weight만 뽑아내는 것을 말합니다.

 

 

그 중, 가장 표현력이 뛰어나고 균일한 아키텍처를 생성하는 MAP 헤드를 채택합니다.

헤드 간 비교

 

 

3. Scaling up data

 

ViT/G-14 알고리즘은 대규모 컴퓨터 비전 모델에 많이 사용된 JFT-300M데이터 셋보다 훨씬 대용량인 JFT-3B데이터셋을 사용합니다.

 

딥러닝에서는 데이터 셋이 작을수록, 모델의 깊이 및 너비가 깊어지거나 넓어질수록 과적합이 아주 높은 확률로 발생합니다.

 

일반적으로, 딥러닝 모델은 모델의 규모(Scale)가 커져야 성능이 좋아집니다.

 

그에 반해, 왼쪽 그림을 보시면 ViT모델은 모델의 규모가 커지지 않아도, 더 큰 데이터 셋에서의 성능이 일관되게 상승하는 것을 보여줍니다. 

 

그리고, 데이터 셋의 스케일이 작아도 과적합이 발생하지 않는 것 역시 확인할 수 있었습니다.

 

오른쪽 그림에서는, 모델을 파인 튜닝(Fine-tuning) 했을 때 역시 전반적으로 동일한 양상을 보였습니다.

 

이는, 과적합 및 정규화에 대한 걱정없이 모델의 규모를 더 확장할 수 있다는 것을 말하는 것이기도 합니다.

 

 

 

4. 메모리 효율적 Optimizers

 

딥러닝은 스케일이 큰 모델을 트레이닝 할 때, 모델 파라미터에 필요한 스토리지에 병목 현상(bottleneck)이 발생하게 됩니다.

 

병목 현상이란, 얼마나 많은 데이터를 동시에 처리할 수 있느냐의 문제입니다.

 

즉, 한 번에 처리할 수 있는 데이터의 양보다 처리할 수 있는 능력이 충분하지 않을 경우 발생하는 문제입니다.

 

ViT-G 모델에는 무려 약 20억 개의 파라미터가 있으며 8GB의 메모리를 차지합니다. 

 

게다가, 트랜스포머 훈련에 쓰이는 Adam 옵티마이저는 각 파라미터당 두 개의 추가 부동 소수점 스칼라를 저장하므로 추가 오버헤드가 발생합니다.

 

논문에서는 이를 해결하기 위해, 두 가지 방법을 적용시켜보았습니다.

 

4.1. Half-precision Adam

 

딥러닝 모델의 학습에 일반적으로 단정밀도(Single-Precision)을 사용하는데, 최근 모델 사이즈가 점점 커지고 계산량과 그에 필요한 메모리 크기 등이 기하 급수적으로 커져왔습니다.

 

그래서, 계산량을 줄이기 위해 딥러닝 모델의 학습에 반정밀도(Half-precision)을 사용합니다.

 

반정밀도는 비트(Bit) 수가 절반으로 줄어들어 계산량과 메모리 사용량을 크게 줄일 수 있지만, 숫자를 표현할 수 있는 범위가 크게 줄어듭니다.

 

이는, 양자화(Quantization)의 개념과 유사합니다.

 

하지만 우려와는 달리, 논문에서는 반정밀도로 저장하는 것이 트레이닝 성능 및 결과에 영향을 미치지 않는다는 것을 관찰했고, 옵티마이저 오버헤드를 2배에서 1.5배로 줄일 수 있었습니다.

 

4.2. Adafactor Optimizer

 

하지만, 위와 같은 Optimizer는 여전히 큰 메모리 오버헤드가 발생했습니다. 

 

그래서, Adam 대신 Adafactor라는 옵티마이저를 적용했습니다.

 

Adafactor는 적응성의 경험적 이점을 유지하면서 메모리 사용량을 줄이는 Adam 기반의 확률론적 최적화 방법입니다.

 

메모리 오버헤드는 50%까지 줄어든 반면, 성능은 기존 Adam 옵티마이저와 동등하거나 좀 더 나은 수준을 보여주었습니다. 

 

 

5. Learning-rate schedule

 

대부분 딥러닝, 머신러닝을 사용할 때, 모델이나 옵티마이저에 많은 신경을 쓰지만, Learning-rate schedule에는 상대적으로 그러지 않는 경우가 많습니다.

 

실제 학습을 할 때, 어떤 schedule을 고르는지, 하이퍼 파라미터로 어떤 값을 고르는지에 따라 학습 결과에는 큰 차이가 존재합니다.

 

Learning-rate schedule은 미리 학습을 정해두고, 그 일정에 따라 학습률을 조정하는 방법입니다.

 

Warm-up과 Cooldown 단계 사이에 학습률이 너무 빨리 0으로 줄어들지 않아야 합니다.

 

 

Warm-up은 초기에 학습률을 0으로 설정하고, 이를 일정 기간 동안 선형적으로 높여주는 방식을 말합니다.

 

Cooldown은 증가된 학습률을 점차적으로 줄이는 방식을 말합니다.

 

 

이를 위해, 역제곱근 schedule을 사용합니다.

 

역제곱근 schedule

 

 

역제곱근 schedule은 아주 큰 학습률로 시작해서 작은 학습률로 끝나게 됩니다. 

 

역제곱근 shedule은 학습률이 천천히 0으로 줄어들기 때문에, 적합한 것으로 판단한 것으로 보입니다. 

 

그림(하단)을 보면, 역제곱근(Reciprocal)이 200k, 400k 및 500k 스텝 후 cooldown이 천천히 진행되는 것을 확인 할 수 있습니다.

 

그림(상단)을 보면, 검증 점수(Validation Score,높을수록 좋음)를 볼 수 있는데, 역제곱근이 훌륭한 점수를 얻을 것을 확인 할 수 있었습니다.

 

따라서, 논문에서는 역제곱근 schedule을 채택했습니다.

 

개인적 의문은 역제곱근 schedule은 자주 쓰이지 않는 Learning-rate schedule입니다.

 

위와 같은 이유라면, 코사인 schedule을 사용하면 어떨까 라는 생각을 하였습니다.

 

코사인 schedule 역시 학습률이 천천히 0으로 줄어들기 때문입니다.

 

큰 학습률에서 머무는 시간이 비교적 짧은 역제곱근과 비교했을 때, 코사인 schedule은 더 오래 머물기 때문에, 성능이 더 좋지 않을까 라는 생각을 해보았습니다.

 

Learning-rate 및 Validation Score schedule별 비교

 

 

6. 모델 차원 선택

 

앞서 언급했듯이, ViT 모델에는 많은 파라미터가 있습니다.

 

그래서, ViT 모델은 XLA(Accelerated Linear Algebra) 컴파일러를 통해 런타임 속도와 메모리 공간을 위한 최적화를 진행합니다.

 

XLA은 소스코드를 변경하지 않고, TensorFlow 모델을 가속화할 수 있는 컴파일러입니다.

 

XLA는 모델을 컴파일하기 위해 복잡한 방법을 사용하는데, 결과적으로 어떤 모델이 최적인지 예측하기 어렵습니다. 

 

XLA에 최적의 조합을 낼 수 있는 모델을 찾기 위해 논문에서는 다양한 형태의 ViT를 대량으로 인스턴스화하고 훈련을 시도하는 시뮬레이션을 진행했습니다.

 

깊이,너비,헤드 및 MLP너비를 변경하지만 패치 크기는 14픽셀로 유지합니다.

 

이러한 방식으로 모델이 장치의 메모리에 맞는지 여부를 측정합니다.

 

 

아래 그림은 시뮬레이션의 결과를 요약한 것입니다.

 

 

 

주황색 블록 : 수정하지 않은 ViT

 

녹색 블록 : Half-precision을 결합한 Adam + ViT (4.1에서 언급했던 모델)

 

파란색 블록 : Adafactor + ViT (4.2에서 언급했던 모델, ViT-G)

 

흰색 블록 : 어떤 설정에서도 메모리에 맞을 수 없음

 

 

각 블록은 하나의 모델 구성에 해당합니다.

 

블록색의 농도는 상대적 훈련 속도에 해당합니다.

 

 

시물레이션 결과, 파란색 블록이 압도적으로 많고, 농도가 짙은 것을 확인할 수 있었습니다.

 

이는, XLA에 가장 적합한 최적의 모델은 파란색 블록(ViT-G)이란 것이겠죠?

 

모델 아키텍쳐 세부사항

표를 보시면, G/14, g/14 모델 아키텍쳐 구성이 가장 범용적인 것을 확인할 수 있습니다.

 

GFLOPs는 파라미터와 함께 모델이 얼마나 큰지, 효율적인지를 나타내는 지표를 말합니다.

 

그림과 표에 요약된 대로, 메모리에 적합한 모델은 G/14인 것으로 나타났습니다.

 

 

 

기존 SOTA와의 비교

 

 

기존 SOTA와 비교를 해보았을 때, ImageNet, INet V2, VTAB(light)에서는 우위를 점했습니다.

 

INet ReaL, ObjectNet에서는 우위를 점하진 못했지만, 근소한 차이였습니다.

 

 

마무리

 

ViT-G/14는 ViT 모델 중 가장 최신의 알고리즘으로써, 기존 SOTA와 비교했을 때 더 뛰어난 성능을 보여주고 있습니다.

 

뛰어난 성능과 더불어, 계산량 및 메모리 부문에서도 아주 효율적인 움직임을 보여주었습니다.

 

앞으로, ViT 모델의 개발자들은 이 모델 아키텍쳐를 통해 더 적은 컴퓨팅 리소스로 훈련할 수 있는 모델을 설계할 수 있을 것입니다. 

 

 

 

직접 번역한 것이다 보니 틀린 부분이나 어색한 부분이 있을지도 모르겠습니다.

 

지적 및 피드백은 감사한 마음으로 언제든 환영이니 댓글로 달아주시면 더 감사하겠습니다!

 

감사합니다 '_' !!

 

인용 자료

An image is worth 16x16 words: Transformers for image recognition at scale A. Dosovitskiy, L. Beyer, A. Kolesnikov, D. Weissenborn, X. Zhai, T. Unterthiner, et al. arXiv preprint arXiv:2010.11929 2020

Scaling vision transformers X. Zhai, A. Kolesnikov, N. Houlsby and L. Beyer arXiv preprint arXiv:2106.04560 2021

Attention is all you need A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, et al. Advances in neural information processing systems 2017 Pages: 5998-6008

Set transformer:A framework for attention-based permutation-invariant neural networks.    Juho Lee, Yoonho Lee, Jungtaek Kim, Adam Kosiorek, Seungjin Choi, and Yee Whye Teh. In ICML, 2019.

https://deepinsight.tistory.com/106 [Steve-Lee's Deep Insight]

https://paperswithcode.com/method/adafactor

https://hoya012.github.io/blog/Bag-of-Tricks-for-Image-Classification-with-Convolutional-Neural-Networks-Review/