[딥러닝 논문리뷰] AIM: Scalable Pre-training of Large Autoregressive Image Models (Apple, 2024)

Apple에서 2024년 1월 large pretrained image model인 AIM(Autoregressive Image Models)을 발표했다. 코드와 model weight이 Github에 공개되어 있다.

 

논문 링크: https://arxiv.org/pdf/2401.08541.pdf

GitHub: https://github.com/apple/ml-aim/tree/main

 

 

 

AIM은 LLM에 영감을 받아 만들어진 대규모 vision 모델이다. BEiT (2021), Masked autoencoder(MAE) (2021) 등이 masked language modeling (MLM)을 통해 사전학습 시킨 것과 다르게, 주어진 패치로 다음 패치를 예측하는 autoregressive object를 이용하여 사전학습을 진행했다.

 

AIM의 주요 contribution은 vision 모델도 LLM과 유사한 scaling property를 보일 수 있다는 것을 증명했다는 점이다. DINOv2 (2023) 에서는 142M장의 이미지로 460M 모델을 학습시켰지만, vision 모델은 LLM에서의 scaling law를 따르지 않는다고 주장했고, MAE에서도 비슷하게 얘기했다. 그러나 AIM은 2B장 이미지에 대해 7B 모델을 autoregressive objective로 성공적으로 학습시켰으며, 이 정도의 규모에서도 saturation이 일어나지 않는다는 점을 토대로 large-scale vision model의 새로운 지평을 열 가능성을 확인했다.

 

 

 

 

Related works

본 연구에서는 iGPT (2020) 에서 사용한 autoregressive objective를 사전학습에 사용했다. 또다른 pretrained vision model인 BEiT, MAE 등은 BERT에 영감을 받은 MLM방식을 사용했다. Contrastive method들도 라벨 없이 사전학습한다는 점에서 유사하다고 볼 수 있다. 이들은 작은 모델 사이즈에서는 좋은 성능을 내지만 scaling에는 어려움이 있다. 

 

 

Method

Dataset

 

DFN dataset에서 2B장 이미지를 추출하고 여기에 ImageNet을 섞어서 사용했다.

 

 

Objective

 

 

이미지를 K개의 패치로 overlap 없이 자르고, next patch prediction을 한다.

 

$\sum_x \sum_k -\log P(x_k|x_{<k})$

 

MAE와 유사하게, normalized pixel-level regression loss를 사용한다.

 

$\min_\theta \frac1K\sum_{k=1}^K||\hat x _k (\theta)-x_k||_2^2$

 

 

Architecture

 

기본적으로 ViT 구조를 사용했다. 상세한 파라미터는 다음과 같다.

 

 

이전 sequence의 패치만 이용해 attention을 수행하도록 하는 causal mask를 적용했다. 그러나 downstream task에서는 bidirectional self-attention을 수행해야 하기 때문에 이러한 방식은 성능을 떨어트린다.

 

이를 해결하고자 “prefix transformer”를 도입했다. 처음 S개의 패치가 "prefix"에 해당한다. 이들은 나머지 패치를 예측하기 위한 context로 사용되고, 이들은 autoregression prediction에서 제외된다.

 

 

Downstream tasks

 

Downstream task에 대해 평가할때는 backbone은 모두 freeze하고 classification head만 train했다. Large model 이니만큼 전체를 다시 finetuning 해서 쓰는것은 너무 낭비이다.

 

그러나 pretrain task에서는 패치단위 prediction만 했기 때문에 이미지 단위 token이 없었다. Classification 등의 이미지 단위 prediction을 위해서는 패치 feature들에 global average pooling을 할 수도 있지만, attention pooling operation을 통해 global descriptor 계산하는 로직을 추가했다.

 

 

 

Results

Impact of scaling

 

 

Pretrain loss와 classification accuracy간의 상관관계를 확인할 수 있다. Pretrain loss가 줄어들수록 downstream task의 성능이 높아지는 것을 통해 pretrain objective를 잘 설정했다고 볼 수 있다. 또한 모델 사이즈를 키울수록 성능이 좋아지는 것을 확인할 수 있다. 이는 LLM에서의 양상과 유사하다.

 

 

Ablations

 

 

Method에서 설명한 다양한 구조에 대한 ablation. 참고로 autoregression pattern은 패치를 어떤 순서로 넣어줄 것인가에 대한 내용인데, 일반적으로 생각하는 가로-세로 순서가 가장 좋았다고 한다.

 

 

 

Pretrain objective

 

MLM 방식과도 비교를 진행했다. MLM보다 autoregressive가 좋다고 한다.

 

 

 

Comparison with other pretrained models

 

 

다른 pretrained 모델들과의 비교이다. DINOv2를 제외하고 AIM이 모두 이겼는데, DINOv2는 더 높은 해상도 이미지를 이용했다고 한다. 그리고 DINOv2는 여러가지 자질구레한 학습 트릭들에 크게 의존하고 있는데, AIM은 학습 방법이 매우 간단하다고 한다 (...)

 

근데 AIM을 제외하고 2등인 iBOT과 비교해봐도 iBOT의 파라미터 수는 300M개정도이다. AIM-0.6B의 절반 정도인데도 훨씬 더 좋은 성능을 보인다. 아마 표에 파라미터 갯수를 안 적어놓은 것은 이런 불리함 때문 아니었을까..

 

 

후기

이전 pretrain vision 모델들은 MLM을 주로 사용하곤 했는데 본 논문에서는 2020년 iGPT에서 사용한 autoregressive task를 이용해 모델을 학습시켰다. vision 모델에서도 LLM처럼 모델 사이즈가 커질수록 성능이 증가하는 scaling law가 작용할 수 있음을 보였다.

 

그러나 몇가지 의문점은 ..

1. 마지막 result에서 보듯이 파라미터 수가 엄청나게 큰것 치고 성능이 압도적으로 좋은건 아님
2. 학습 장비와 학습에 소요된 시간을 안적어놓음. 그리고 깃헙에 학습 코드와 loss 코드가 없다.
3. 파라미터 수를 2B까지 키워도 성능이 saturation 되지 않았다고 하는데, 그러면 DINOv2를 이길때까지 한번더 scaling을 안한 이유는..?