[Github]
[arXiv](2024/01/16 version v1)
Abstract
Autoregressive objective를 통해 ViT를 크게 확장하여 downstream task에서 강력한 성능을 보여주는 대규모 비전 모델인 AIM(Autoregressive Image Model) 구축
Pre-training Dataset
Common Crawl에서 Data Filtering Nework로 필터링된 12.8B text-image pair가 있는 DFN dataset에서 alignment score가 상위 15%인 DFN-2B dataset이 있다.
LLM 사전 훈련의 일반적인 관행에서 착안하여 p = 0.8로 DFN-2B에서 샘플링하고 p = 0.2로 ImageNet-1K에서 이미지를 샘플링한다. 이러한 데이터셋을 DFN-2B+라고 한다.
Approach
Training Objective
일반적인 자기회귀 목표를 따르며,
손실은 negative log-likelihood로 정의된다.
Prediction loss
피라미터 θ에 대한 normalized pixel-level regression loss:
Architecture
Backbone은 ViT이다.
깊이보다는 너비 확장을 우선시한다.
Prefix Transformer
자기회귀 모델에서 현재 생성 중인 패치 이후의 패치를 보지 못하도록 하는 causal mask는 그 작동 방식이 ViT와 달라 다운스트림 작업에서 성능이 저하된다.
이 문제를 해결하기 위해 PrefixLM과 같이 초기 패치 일부를 prefix로 지정하여 모든 패치에서 볼 수 있도록 mask를 제거한다. 이를 통해 causal mask가 없는 다운스트림 작업에 적응할 수 있다.
 |
 |
MLP prediction heads
사전 훈련 목표에 너무 특화되는 것을 막기 위해 훈련에만 특정한 head를 사용하는 것은 일반적인 관행이다.
본문에서는 최종 레이어 위에 N개의 MLP를 추가하여 각 패치를 독립적으로 처리하는 단순한 설계를 채택했다.
Straightforward implementation
- Transformer와 MLP의 입력에 sinusoidal positional embedding
- MLP 확장 비율 x4
- Bias term과 [CLS] token 제거
- MLP head block 12개
- 출력 패치는 패치별로 정규화
- 기타 등등...
Downstream adaptation
대규모 모델은 fine-tuning 하기 까다롭기 때문에 다운스트림 작업에서는 classification head만 훈련한다.
손실은 패치마다 각각 계산되어 image-level의 토큰이 없기 때문에 global avg pooling 대신 선형 분류기에 입력하기 전 multihead attention pooling을 사용한다.
Multihead attention pooling은 각 attention head 별로 학습 가능한 쿼리 q와 각 패치에 대한 cross-attention을 수행하여 global descriptor p̂을 얻는 것이다.
각 패치 pi, 학습 가능한 쿼리 벡터 q, attention head h, projection weight W에 대해:
Multihead attention pooling의 장점은 패치 수준의 손실을 사용할 때, 다운스트림 작업에서 적은 수의 피라미터를 추가하면서도 global feature를 집계할 수 있다는 것이다.
Results
Impact of scaling
LLM과 같이 모델 크기, 반복 횟수, 데이터셋 크기가 증가함에 따라 성능이 포화되지 않고 꾸준히 증가함을 볼 수 있다.
Downstream evaluation
Ablation
