[논문 정리] AIMv2: Multi-modal Autoregressive Pre-training of Large Vision Encoders

https://arxiv.org/abs/2411.14402

---

본 논문은 AIM의 후속 연구이며 아래 정리 게시글 참고
https://velog.io/@bluein/paper-7

---

1. Introduction

• Vision Model의 Pre-training 연구는 시간이 지남에 따라 크게 발전
• 초기에는 특정 Task의 성능을 극대화하기 위해 전문 Model 설계
• 점차적으로 최소한의 적응으로 여러 사전 정의된 Downstream Task에 배포 가능한 Foundation Model 등장
• 그러나 Large Language Models (LLMs)의 놀라운 성공이 Vision Model 활용에 새로운 패러다임 도입
• 기존의 엄격한 사전 설정과 달리, LLMs는 Pre-training된 Model의 능력을 더 효과적으로 활용 가능
• 이러한 변화는 Vision Model의 Pre-training 방법 재고를 요구
• Generative pre-training은 Language Modeling의 지배적인 패러다임으로, 뛰어난 성능과 Scalability을 보임
• Generative pre-training은 Computer Vision에서도 광범위하게 탐구됨, 하지만 여전히 discriminative 방법보다 성능이 낮음
• 예를 들어, El-Nouby 등은 LLMs Pre-training과 유사한 형식을 제안하여 긍정적인 Scaling 특성 시연, 그러나 discriminative 방법과 동일한 성능을 위해서는 훨씬 더 높은 Capacity의 Model 필요
• 반면, contrastive 기법은 Parameter Efficiency가 높지만, Training 및 Scaling이 상당히 어려움
• 이러한 문제를 완화하기 위한 상당한 진전이 있었으나, generative pre-training의 Simplicity와 Scalability, discriminative 접근의 Parameter Efficiency를 결합한 방법 개발에는 여전히 격차가 존재

AIMV2

• AIMV2는 Image Patch와 Text Token을 모두 autoregressive하게 생성하도록 Pre-training된 Open Vision Model 계열 구조
• Pre-training 중, AIMV2는 causal Multi-modal decoder를 사용하여 먼저 Image Patch를 Regression한 후 Text Token을 autoregressive 방식으로 Decoding
• Advantages
  • Implementation 및 Training이 간단하며, 과도하게 큰 Batch Size나 특수한 Inter-batch 통신 방법 필요 없음
  • AIMV2의 아키텍처와 Pre-training 목표는 LLM 기반 Multi-modal Application과 잘 맞아 원활한 통합 가능
  • Image Patch와 Text Token에서 Training 신호 추출, discriminative 목표보다 더 밀집된 supervision 제공
• Performance
  • 다양한 Vision 및 Multi-modal Task에서 뛰어난 성능을 보이는 강력한 일반화 성능
  • Multi-modal Understanding 벤치마크에서 SOTA Vision-Language Pre-training 방법보다 우수한 성능
  • DINOv2을 Open-vocabulary Object Detection 및 Reference Representation Understanding에서 능가, frozen trunk 상태에서 강력한 Recognition 성능
  • Data나 Parameter를 Scaling할 때 일관되게 향상되는 높은 Scalability
  • 여러 현대 도구와의 호환성 시연, 네이티브 Image Resolution 지원 및 zero-shot Recognition 적응 가능

2. Approach

2.1. Pre-training

• Model은 표준 unimodal autoregressive 프레임워크를 Scaling 하여 Image와 Text를 통합한 Multi-modal 설정으로 통합된 Sequence 생성

• 구체적으로, Image $(x)$는 $(I$)개의 겹치지 않는 Patch $(x_i$), $(i$ $in [1, I])$로 분할되어 Token Sequence 형성
• Text Sequence는 subwords $(x_t$), $(t$ $in[I, I+T])$로 분할
• 이 Sequence들은 결합되어 Text Token이 Image Token을 참조할 수 있도록 함
• 항상 Image를 먼저 배치하여 강력한 Visual Feature 생성을 가능하게 함

Integrated Multi-modal Autoregressive Modeling Process

• Sequence는 다음과 같이 인수분해:
  $P(S) = \prod_{j=1}^{I+T} P(S_j | S_{<j})$
• 여기서 $(S_j$)는 결합된 Image Patch와 Text Token Sequence의 j번째 Token, $(S_{<j}$)는 이전 모든 Token
• 이 통합된 인수분해는 Model이 Sequence에서 다음 Token을 autoregressively 예측할 수 있게 함

Pre-training Setting

• 전용 Vision Encoder가 Raw Image Patch를 처리하고, 이는 Multi-modal Decoder에 전달
• Decoder는 결합된 Sequence에 대해 다음 Token 예측 수행
• Vision Encoder와 Multi-modal Decoder는 각각 prefix와 causal self-attention 연산 사용

Objective Function

• Image와 Text Domain에 대해 별도의 Loss Function 정의:
  $L_{img} = \sum_{i=1}^{I} || \hat{x}_i(x_{<i}; \theta) - x_i ||^2$
  $L_{text} = - \sum_{t=I+1}^{I+T} \log P(x_t | x_{<t}; \theta)$
• 전체 목표는 Model Parameter $theta$ 에 대해 $(L_{text} + alpha * L_{img})$를 최소화
  • $L_{text}$: 각 단계에서 실제 Token의 negative log-likelihood를 측정하는 standard cross-entropy Loss
  • $(L_{img}$
    • $ell2$ Pixel level Regression Loss
    • Model이 예측한 Patch $hat{x}_i(theta)$와 실제 Patch $x_i$ 비교
  • Image Patch는 He 등에 따라 Normalization
  • 실제 Implementation에서는 Multi-modal Decoder의 최종 hidden state를 Image Patch와 Language의 적절한 출력 차원으로 매핑하기 위해 별도의 Linear Layer 사용

2.2. Architecture

Vision Encoder

• AIMV2의 Vision Encoder로 Vision Transformer (ViT) 아키텍처 채택
• 300M에서 3B Parameter 범위의 Vision Encoder 시리즈 Training
• 상세한 Model 사양은 Table 1에 제공

Prefix Attention

• El-Nouby 등을 따름
• Prefix Attention Mask 적용하여 Vision Encoder 내 self-attention 메커니즘 제한
• 이 전략은 추가 Tuning 없이 추론 시 양방향 attention 사용 가능하게 함
• 구체적으로, prefix 길이를 (M \sim U{1,2,...,I-1})로 Random Sampling
• Pixel Loss은 prefix Patch를 제외한 ({ x_i | i > M })에 대해서만 계산

SwiGLU and RMSNorm

• Vision Encoder와 Multi-modal Decoder에 SwiGLU를 Feedforward Network (FFN)으로 통합
• 모든 Normalization Layer를 RMSNorm으로 교체
• 이러한 수정은 Language Modeling에서 SwiGLU와 RMSNorm의 최근 성공을 활용

Multi-modal Decoder

• Image와 Text Modality 모두에 대해 autoregressive 생성 수행하는 통합된 Multi-modal Decoder 채택
• Image Feature와 Raw Text Token을 각각 Linear Projection하여 (R{d{dec}})에 임베딩
• Decoder는 Image와 Text Feature의 결합된 Sequence를 입력으로 받고, self-attention 연산에서 causal attention 사용
• Decoder의 출력은 Image Token과 Text Token 각각을 예측하기 위해 두 개의 별도 Linear Head를 통해 처리
• 모든 AIMV2 Variation에 동일한 Decoder Capacity 사용

Optimization Hyperparameters

• 모든 AIMV2 Model의 Pre-training 중 사용된 최적화 Hyperparameter는 Table A1에 개요

2.3. Data

Dataset Composition

• Open 및 Closed Data셋의 Image-Text Pair를 조합하여 AIMV2 Model Pre-training
• 공개적으로 사용 가능한 DFN-2B와 COYO Data셋 사용
• High Quality Image-Text Pairs (HQITP)라는 독점 Data셋 추가
• Alt-text 외에도 Lai 등의 접근 방식을 따르는 Synthetic Caption 사용
• Data셋의 크기와 각 Data셋에 사용된 샘플링 확률 등 세부 사항은 Table 2에 제공
• 특별히 언급되지 않는 한, 모든 AIMV2 Model은 120억 개의 Image-Text 샘플을 사용하여 Pre-training

2.4. Post-Training

Post-Training Strategies

• 초기 Pre-training 단계에서 높은 성능의 Model 도출, 이후 다양한 post-training 전략을 통해 능력 추가 향상 탐구

High-resolution Adaptation

• 초기 Pre-training 단계에서 고정된 Resolution 224px의 Image Data 사용
• Detection, Segmentation, Multi-modal LLMs와 같은 많은 Downstream Task은 더 높은 Resolution의 Image 처리에 적응된 Model에서 이점
• 따라서, AIMV2 Model을 336px 및 448px Resolution으로 Fine-tuning
• High-resolution adaptation 단계는 Pre-training 단계와 동일한 풀에서 샘플링된 20억 개의 Image-Text Pair 사용, 이 단계에서는 Synthetic Caption 사용하지 않음
• Zhai 등의 관찰과 일치하게, 안정적인 최적화를 유지하기 위해 weight decay를 0으로 설정하는 것이 중요

Native Resolution Fine-tuning

• 전용 Resolution와 종횡비로 Model을 Training 하는 것은 원본 형태의 Image를 처리해야 하는 많은 응용 프로그램에서 유연하지 않음
• 이전 연구인 FlexiViT와 NaViT이 이러한 제한을 해결
• 다양한 종횡비와 Resolution으로 Training하기 위한 다른 접근 방식 채택
• 구체적으로, (B_i)를 Mini-batch 내 Image 수, (A_i)를 Image당 Patch 수, (C)를 Mini-batch 내 총 Image Patch 수로 정의
• Mini-batch (i)에 대해, 영역 (A)를 무작위로 샘플링하고 Image의 종횡비를 유지하면서 이 영역 내에 맞게 크기 조정
• 그 후, (C = A_i \times B_i)가 되도록 Mini-batch 크기 (B_i) 조정
• 이 전략은 Pouransari 등이 가변 컨Text 길이로 LLMs를 Training하는 접근 방식과 유사
• Implementation 시, Sequence 패킹, attention masking, 또는 커스텀 Pooling 연산에 대한 휴리스틱이 불필요
• (A = 2^n) 선택, 여기서 (n)은 잘린 정규 분포 (N(0,1))에서 ([-1,1]) 범위 내에서 샘플링되고 ([7,12])로 Linear 매핑

3. Analysis

• One of the main advantages of AIMV2 is its simplicity
  • Implementation 및 Scaling이 용이
  • 따라서, AIMV2 Model 계열의 Scaling 특성을 조사

3.1. Scaling AIMV2

• Analysis of the Impact of Data Size and Model Capacity Scaling

  • Model 크기를 고정하고 Pre-training 시 본 샘플 수를 변경
  • 이 분석은 Hoffmann 등의 "Approach 1"과 유사

  

• Setup

  • 3억에서 30억 Parameter 범위의 네 가지 Model Capacity Training
  • Pre-training 시 5억에서 64억 Image-Text Pair 샘플 수 변동
  • 모든 Model은 조기 종료 없이 수렴할 때까지 Training
  • 최소한의 계산 비용으로 단일 Model을 각 Capacity에 대해 Training, half-cosine learning rate Schedule 사용
  • 이 실행에서 일곱 개의 중간 체크포인트 선택 후 Linear 쿨다운 적용, 쿨다운 단계 길이는 초기 Pre-training 단계의 20%

• Results

  • Data나 Parameter Scaling 시 일관된 성능 향상 관찰
  • 낮은 Capacity Model의 Data Scaling 시 수익 감소 현상
  • 최적 Model 크기는 Computational Budget에 따라 변동
  • 작은 Computational Budget에서는 큰 Capacity Model이 과소 Training 되어 낮은 Capacity Model보다 성능 저하

3.2. AIMV2 vs. Captioning

• Study on the Role of AIMV2's Image-level Autoregressive Objective

  • Multi-modal autoregressive 목표로 Training된 Model과 Language supervision만으로 Training된 Model 비교

  

• Setup

  • 특별한 언급이 없는 한, ViT-H Backbone과 20억 Image-Text Pair 사용하여 Pre-training
  • 모든 Model은 수렴할 때까지 코사인 learning rate Schedule로 Training
  • Frozen trunk 상태에서 attentive probe 후 IN-1k top-1 정확도 측정

• Results

  • AIMV2의 Image level 목표가 captioning-only baseline 대비 일관되게 성능 향상
  • Model Capacity 및 Pre-training Data 크기 변경 시에도 동일
  • Data나 Model Capacity Scaling 시 두 접근 방식 모두 일관된 향상, 그러나 captioning baseline은 Data Scaling 시 성능 정체 현상 관찰, AIMV2는 이러한 현상 없음

4. Results

• AIMV2는 다양한 Downstream Task에 즉시 활용 가능한 generalist Vision Encoder
  • Recognition, Detection, Caption Generation 및 여러 Multi-modal 벤치마크 포함 다양한 Task에서 성능 평가

4.1. Image Recognition

• Attentive Probing

  • AIMV2 Model의 품질을 Recognition 벤치마크의 off-the-shelf Backbone으로 평가, Table B1에 개요
  • Yu 등이 제안한 attentive probing 설정 채택, Vision Encoder는 frozen 상태 유지, 마지막 Layer Feature 위에 attentive probe classifier만 Training

  

  • 위 테이블에서, probing 실험에 사용된 상세 Hyperparameter는 Table A2에 제공
  • Key Findings
    • AIMV2는 MAE, AIM과 같은 generative unsupervised 방법보다 유의하게 우수, 더 작은 Capacity Model에서도 동일
    • DINOv2 대비 AIMV2-1B 및 AIMV2-H가 여러 벤치마크에서 경쟁력 있는 성능 제공, 특히 IN-1k, Food101, DTD, Cars, Infographic에서 우수
    • 그러나 DINOv2는 iNaturalist 및 fMoW에서 뛰어난 성능 제공
    • Self-supervised Model이 의료 영상 벤치마크(RxRx1, CAM17)에서 강력한 성능 보여, Low-resource Domain에서 self-supervised Training의 중요성 확인
    • 다른 Vision-Language Pre-training된 baseline과 비교 시, AIMV2는 매우 경쟁력 있는 성능
    • 예를 들어, ViT-Large Capacity에서 AIMV2는 대부분의 벤치마크에서 OAI CLIP보다 우수, DFN-CLIP 및 SigLIP보다도 여러 주요 벤치마크(IN-1k, iNaturalist, DTD, Infographic)에서 강력한 성능 달성
    • AIMV2는 DFN-CLIP 및 SigLIP보다 약 1/4의 Data(12B vs. 40B)로 Training되었으며, Training 및 Scaling이 더 용이
    • AIMV2 Model Capacity Scaling은 일관된 성능 향상으로 이어지며, AIMV2-3B가 가장 강력한 결과 도출, 특히 448px Image로 Fine-tuning된 변종은 IN-1k에서 89.5% top-1 정확도 달성
    • Figure 3에서 Model Capacity 및 Image Resolution Scaling 시 IN-1k 성능 명확한 향상 관찰, High-resolution Fine-tuning된 Backbone의 상세 결과는 Appendix B에 제공

• Zero-shot via LiT Tuning

  • AIMV2 Backbone의 LiT와의 호환성 조사, zero-shot 설정으로 Scaling
  • IN-1k zero-shot 성능은 Table 5에 제시
  • Key Findings
    • Multi-modal autoregressive 목표를 가진 AIMV2는 captioning-only baseline 대비 소폭 향상
    • AIMV2-L Model은 CapPa의 결과보다 유리한 성능, 긴 Training 기간 후
    • 전반적으로, AIMV2는 OAI CLIP보다 우수한 zero-shot 성능, 그러나 40B Image-Text Pair로 더 긴 일정에 Training된 SigLIP과 같은 전용 Model보다는 뒤처짐

• Native Resolution

  • 다양한 Resolution 및 종횡비로 Image를 처리하도록 AIMV2 Fine-tuning, 2.4에 상세
  • 이 단계의 품질 평가를 위해 native resolution에 적응된 AIMV2 Encoder와 특정 Resolution으로 조정된 Model 비교, 결과는 Table 6에 제시
  • Key Findings
    • AIMV2-Lnative는 다양한 Resolution에서 강력한 성능 제공, 전용 Model과 비교해 약간의 성능 저하만 발생
    • IN-1k validation set Image의 원본 native resolution 사용 시 87.3% robust accuracy 달성, AIMV2가 exceptional recognition 성능 유지하며 종횡비 및 Resolution에서 높은 유연성 제공

4.2. Object Detection and Grounding

• Demonstration of Additional Features of AIMV2
  • AIMV2의 성능을 Open-Vocabulary Detection (OVD)과 Referring Expression Comprehension (REC)과 같은 Task에서 평가
  • MM-Grounding-DINO에서 소개된 Model 아키텍처를 따르되, ViT-L을 ViTDet 형식으로 Vision Backbone에 적용
  • 결과는 Table 4에 제시
  • OVD Performance Evaluation
    • COCO와 LVIS에서 평가
  • REC Performance Evaluation
    • RefCOCO (RefC), RefCOCO+ (RefC+), RefCOCOg (RefCg)에서 평가
  • Training Datasets
    • Objects365v1, Flickr-30k Entities, GQA, COCO17, RefCOCO 사용
  • DINOv2 Training Setup
    • Window size를 16으로 고정하여 Backbone 간 고정된 계산 비용 보장
  • Results
    • AIMV2는 모든 벤치마크 중 하나를 제외하고 DINOv2 및 다른 Vision-Language Pre-training Model보다 우수한 성능을 보임
    • 특히 LVIS에서 뛰어난 성능 시연
    • 추가로 Closed-vocabulary Detection, Instance Segmentation 및 다양한 Window Sizes에 대한 ablation 결과는 Appendix D에 제공

4.3. Multi-modal Understanding

• Importance of Vision Encoder in the Advancement of Multi-modal Models
  • Large-scale Multi-modal Model의 발전에 Vision Encoder가 중요한 역할 수행
• Quantifying the Performance of AIMV2
  • Liu 등과 유사한 Multi-modal Instruction Tuning 단계 수행
  • McKinzie 등과 유사한 Large-scale Multi-modal Pre-training 후 Few-shot In-Context Learning (ICL) 설정 탐구

4.3.1. Multi-modal Instruction Tuning

• Setup

  • Vision Encoder (예: AIMV2-L)와 LLM (예: Llama 3.0) 사이에 2층 MLP 커넥터 배치
  • 이 단계에서 Vision Encoder의 Parameter는 고정
  • Liu 등과 달리, 커넥터와 LLM을 단일 단계에서 공동 Training
  • 커넥터의 learning rate를 8배 증가시켜 Training
  • 평가 Data셋, Task 프롬프트, Hyperparameter는 Appendix C에 상세히 기술
  • 특별한 언급이 없는 한, Llava SFT mixture와 Llama-3.0 8B LLM Decoder 사용
  • 모든 Model을 단일 epoch 동안 Training

• Evaluation

  • 다양한 질문 응답 벤치마크에서 instruction-tuned Model 평가
    • General Knowledge, Text가 풍부한 Image, Science Domain, Caption Generation 포함
  • AIMV2와 여러 baseline의 결과는 Table 7에 제시
  • Key Findings
    • 가장 작은 Model인 AIMV2-L이 대부분의 벤치마크에서 OAI CLIP, SigLIP, DINOv2보다 우수한 성능 보임
    • Baseline이 더 큰 Capacity이나 높은 입력 Resolution를 사용할 때에도 동일
    • AIMV2 Backbone Capacity Scaling 시 성능 일관되게 향상, AIMV2-3B가 모든 벤치마크에서 최고 성능 도달 (단, 한 가지 제외)

• Varying the LLM and Data Mixture

  • Setup
    • Table 7에 보고된 정규 설정 외에도, AIMV2가 다양한 LLM Decoder 및 Instruction Tuning Data Mix과 함께 유사한 이점을 제공하는지 평가
    • 다음 설정에서 Instruction Tuning 단계 수행:
      1. Cambrian Data mixture와 함께 Llama 3.0 사용
      2. Llava SFT mixture와 함께 Vicuna 1.5 사용
  • Results
    • AIMV2-L과 유사한 크기의 OAI CLIP 및 SigLIP Backbone과 함께 Figure 6에 결과 제시
    • 모든 설정에서 AIMV2는 OAI CLIP과 SigLIP보다 더 강력하거나 적어도 동등한 성능 제공
    • 이러한 결과는 다양한 Multi-modal 파이프라인에서 AIMV2의 Robustness와 호환성을 추가로 입증

• High-Resolution via Tiling

  • Strategy
    • Vision-Language Model의 성능 향상을 위한 인기 있는 전략 중 하나는 Image Resolution 증가
    • Tiling 전략 사용: High-resolution Image를 Pre-training Resolution(예: 224px 또는 336px)에 맞게 동일한 크기의 crop으로 Segmentation
  • Compatibility with AIMV2
    • Crop size를 336px로 설정하고, 672px 및 1008px Image (각각 2×2, 3×3 그리드)에 대해 평가
    • 결과는 Figure 5에 제시
    • 모든 방법의 성능이 높은 Resolution에서 향상됨을 관찰, 특히 TextVQA에서 유의한 향상
    • High-resolution tiling 설정에서도 AIMV2는 baseline보다 우수한 성능 유지

4.3.2. Multi-modal In-Context Learning

• Setup
  • Large-scale Multi-modal Pre-training 설정에서 AIMV2 평가
  • MM1의 Pre-training 전략을 따르되, Vision Encoder를 AIMV2로 교체
  • Interleaved image-text 문서를 사용하여 Pre-training된 Model로 in-context evaluations 가능하게 함
  • ICL 성능은 Table 8에 제시
  • Benchmarks
    • COCO, NoCaps, TextCaps, VQAv2, TextVQA, VizWiz, GQA, OK-VQA
• Results
  • AIMV2는 4-shot 및 8-shot 설정에서 최고의 성능 달성, MM1 시리즈가 채택한 더 높은 Capacity의 DFN-CLIP을 능가
  • 이는 Large-scale Multi-modal 설정에서 AIMV2가 ICL을 효과적으로 활용할 수 있음을 강조

7. Conclusion

• Introduction of AIMV2
  • AIMV2는 Image Patch와 Text Token을 Reconstruction 하는 Multi-modal autoregressive 목표로 Pre-training 된 Vision Encoder
• Advantages of Integrated Objectives
  • Image Recognition, Grounding, Multi-modal Understanding 등 다양한 Task에서 AIMV2의 우수한 성능 발휘 가능
• Reasons for Superior Performance
  • 모든 입력 Token과 Patch의 신호를 활용하여, 다른 방법에 비해 적은 샘플로도 효율적인 Training 가능
• Comparison Results
  • 기존의 self-supervised 및 Vision-Language Pre-training Model보다 일관되게 우수하거나 동등한 성능
  • AIMV2의 강력함을 Vision Encoder로서 입증
• Advantages of Pre-training Process
  • 단순한 Pre-training 과정으로 인해 쉬운 Scalability 보장
  • Vision Model Scaling의 추가 발전을 위한 길을 열어줌