[논문 리뷰/요약] TinyLLaVA: A Framework of Small-scale Large Multimodal Models

TinyLLaVA: A Framework of Small-scale Large Multimodal Models

Journal: Arxiv
Published Date: 2024년 2월 22일

---

• 본 논문은 TinyLLaVA 프레임워크를 소개한다.
• 또한 Small-scale Large Multimodal Model(LMM)을 설계하고 분석한다.
• 실험을 통해 데이터의 품질이 우수하고, 더 나은 훈련 기법을 적용할 경우 작은 모델도 큰 모델과 비슷한 성능을 낼 수 있다는 것을 보여준다.

---

Problem

• 모델의 크기가 커지면 다양한 작업에서 성능이 증가하지만, LM(Large Model)의 경우 훈련/추론 비용이 매우 크기때문에, 접근성이 떨어진다.
• 따라서 Tiny LLM을 사용한 Small-scale Model에 대한 연구가 주목받고 있다.

---

Contributions

• 본 연구에서는 Small-scale LLM을 활용한 폭넓은 LMM 실험 결과를 제시한다.
• 그 결과로, 7B LMM에 비해 성능이 더 우수한 3.1B 모델, TinyLLaVA라는 프레임워크를 제시한다.

---

Methods

TinyLLaVA 모델 아키텍쳐

• 위 이미지는 TinyLLaVA의 모델 아키텍쳐이다.
• TinyLLaVA는 LLaVA의 모델 디자인을 살짝 변형하고 훈련 기법을 다르게 적용한 모델이다.
• TinyLLaVA는 Small-scale LLM($F_\theta$), Vision encoder($V_\varphi$), Connector($P_\phi$)로 구성되며 $\theta, \varphi, \phi$ 는 모두 훈련 파라미터이다.

Small-scale LLM

• Small-scale LLM($F_\theta$) 는 입력 시퀀스 벡터 ${h}^{N-1}{i=0}$ 을 입력으로 받아 예측 ${h}^N{i=1}$ 을 출력한다.
• Tokenizer & Embedding 모듈은 입력 텍스트 시퀀스를 Embedding Space로 매핑하고, 출력될 Embedding을 출력 텍스트로 매핑한다.

VIsion Encoder

• VIsion encoder($V_{\varphi}$)는 이미지를 입력으로 받아 패치 임베딩($\mathbf V={\mathrm v_j \in \mathbb R^{d_x} }^M_{i=j}$)을 출력한다.
• 따라서 $\mathbf V=V_{\varphi}(\mathbf X)$ 이다.
• Vision encoder는 ViT, CNN 둘 다 적용할 수 있는데, CNN의 경우 그리드 피쳐맵을 출력으로 사용한다.

Connector

• Connector($P_{\phi}$)는 패치 시퀀스 ${ \mathbf v_j}^M_{j=1}$ 를 Embedding space ${ h_j }^M_{j=1}$ 로 매핑한다.
• Coonector는 사전 훈련된 LLM과 Vision encoder를 효율적으로 사용할 수 있게 디자인되었다.

Training Pipeline

• TinyLLaVA의 훈련 데이터는 이미지-텍스트 쌍($X, Y$)으로 구성되는데, 텍스트 시퀀스는 Multi-turn conversation으로 구성된다.
• $Y=(Y^1_q, Y^1_a, \cdots, Y^T_q, Y^T_a)$로, $T$는 전체 턴 수, $V^t_q$는 사람의 지시(Instruction) 텍스트, $V^t_a$는 지시에 대한 응답 텍스트다.
• TinyLLaVA는 Pre-training과 Supervised Fine-tuning 두 단계로 나뉘어 훈련된다.

Pre-training for Feature Alignment

• 이 단계에서는 비전과 텍스트 정보를 임베딩 스페이스에서 잘 정렬(Align)하는 것을 목표로 한다.
• Pre-training 단계에서는 Image-caption 스타일 포맷$(X,Y_a)$을 사용한다.
• $Y_a$ 는 응답 데이터로, Multi-turn 형식 데이터에서 응답 텍스트, 즉 이미지에 대한 묘사를 담고있는 텍스트 데이터이다.

$$ p(\mathbf Y_a|\mathbf X) = \prod^{N_a}{i=1}F_\theta(\mathbf y_i|P{\phi}\circ V_{\varphi}(\mathbf X)) $$

• 위 수식은 입력 이미지 $\mathbf X$ 가 주어졌을 때, 응답 $\mathbf Y_a$ 가 발생할 확률이다.
• $i=1$부터 $N_a$ 까지 $\mathbf y_i$ 가 발생할 확률을 모두 곱한 것이다.
• 여기에 로그를 취하여 Autoregression 방식 Training Objective로 사용한다.

$$ \max_{\phi, \theta', \varphi'} \sum^{N_a}_{i=1} \log F_\theta (\mathbf y_i|P_\phi \circ V_\varphi(\mathbf X)) $$

• $\theta', \varphi'$ 는 $\theta, \varphi$ 의 부분집합이다.

Supervised Fine-tuning

• 이 단계에서는 Multi-turn conversation 데이터 $(X,Y)$ 를 사용한다.
• 이 때, 응답 토큰의 집합을 $\mathbb A = { \mathbf y| \mathbf y \in \mathbf Y^t_a, for ; any ; t=1, \dots,T }$ 라 한다.
• Fine-tuning에서도 마찬가지로 Autoregressive로 Objective를 정의하면 아래와 같다.

$$ \max_{\phi, \theta', \varphi'} \sum^N_{i=1} \mathbb I(\mathbf y_i \in \mathbb A) \log F_\theta(\mathbf y_i | P_\phi \circ V_{\varphi}(\mathbf X)) $$

• 여기서 $\mathbb I(\mathbf y_i, \in \mathbb A)$ 항은, 예측된 $\mathbf y_i$ 가 $\mathbb A$ 안에 속하면 1, 아니면 0을 반환한다.

---

Experiments

Small-scale LLM, Vision encoder 환경을 보여주는 테이블

• 위 테이블은 TinyLLaVA에서 사용하는 Small-scale LLM과 Vision encoder를 보여준다.
• Abb는 약칭, HF path는 사전학습 소스를 나타낸다.

Small-scale LLMs

• 대표적인 Small-scale LLM 3개를 선정했다.
  • TinyLlama (1.1B)
  • StableLM-2-1.6B (1.6B)
  • Phi-2 (2.7B)

VIsion Encoders

• 비전 인코더로는 CLIP-Large (0.3B)를 사용한다.
• 실험을 통해 SigLIP이 더 높은 성능을 보일때도 있어서, 프레임워크에 추가.

Connector

• LLaVA-1.5를 따라, GELU 활성화 함수를 포함한 2개의 레이어로된 MLP를 커넥터로 사용한다.
• 또한 Resamplers를 커넥터로 활용하는 방법도 실험.

💡

Resamplers란? 이미지와 같은 비전 데이터의 출력 형식을 조정하여 LLM이 처리할 수 있는 형태로 변환하는 모듈. 쉽게 말해 이미지 임베딩 모듈이다.

Training Data

학습에 사용된 두 가지 데이터셋

• 데이터셋의 품질이 LMM 성능에 미치는 영향을 연구하기 위해 2개의 데이터셋을 사용한다.
• LLaVA-1.5와 ShaerGPT4V에서 제안된 두 가지 데이터셋을 사용했다.
• SFT(Supervised Fine-Tuning)은 665k개로 같고, Pre-training 단계 데이터 개수만 다르다.

Training Recipes

• LLaVA-1.5의 훈련 방식과 ShaerGPT4V의 훈련 방식 두 가지 방법을 적용하여 비교한다.

Base Recipe(LLaVA-1.5)

• PT stage
  • Connector만 학습, 나머지 모델은 Frozen
  • lr : 1e-3, Batch : 256
  • 1 epoch
• SFT stage
  • Vision encoder 고정, Connector, LLM 학습
  • lr : 2e-5, Batch : 128
  • 1 epoch

Share Recipe(ShareGPT4V)

• PT stage
  • Connector는 Base recipe의 방법으로 초기화
  • Vision encoder의 처음 12개 층은 동결, 나머지 부분만 업데이트
  • lr : 2e-5, Batch : 256
  • 1 epoch
• SFT stage
  • Base recipe와 동일

 

VQA 데이터셋에 대한 VLM 성능 벤치마크

• LLM의 L, V, Q, LM은 각각 LlaMA, Vicuna, Qwen, MobileVLM의 약자이다.
• *은 Test dataset이 Training dataset에 포함되어있다는 의미이다.
• 최고 성능은 굵게, 두번째 성능은 밑줄로 표시되어있다.
• 약 30억개의 파라미터를 가진 TinyLLaVA가 7B, 8.2B 모델에 비해 비교적 성능이 좋은 것을 확인할 수 있다.

 

Ablation Study

Training Recipes

Training Recipe에 따라 모델 크기가 작아도 성능이 좋을 수 있음을 보여줌

• 빨간색은 Share recipe로 훈련된 모델을, 파란색은 Base recipe로 훈련된 모델을 나타낸다.
• Training recipe에 따라, 더 작은 모델이 더 좋은 성능을 낼 수 있다는 것을 보여준다.

 

Resampler vs MLP

Resampler, MLP 성능 비교

• Connector로 Resampler와 MLP를 사용했을 때를 비교해서 나타낸 것이다.
• 대부분 MLP를 사용했을 때 성능이 더 높다.

 

Small-scale LLM

LLM Ablation study

• Vision encoder가 각각 CLIP, SigLIP일 때, 3개의 Small-scale LLM(TinyLlama, StableLM-2, Phi-2)에 대한 실험 결과이다.
• Small-scale LLM이 TinyLlama → StableLM-2 → Phi-2 순서로 성능이 좋다.

 

Vision encoder

Vision encoder Ablation study

• 위 실험과 유사하게, 이번에는 Vision encoder의 비교이다.
• SigLIP이 CLIP보다 성능이 전반적으로 좋다.

Reviews

• 작은 크기의 LLM을 사용하여 모델의 크기를 줄이긴했으나 여전히 Pre-training은 단일 GPU로는 불가능할 크기이다..