CompoDistill: Attention Distillation for Compositional Reasoning in Multimodal LLMs

1. 연구 배경 및 문제 정의 (Problem)

배경:
멀티모달 대규모 언어 모델 (MLLM) 은 뛰어난 성능을 보이지만, 높은 계산 복잡도로 인해 실제 배포에 어려움이 있습니다. 이를 해결하기 위해 큰 모델 (Teacher) 의 지식을 작은 모델 (Student) 로 전달하는 지식 증류 (Knowledge Distillation, KD) 기법이 주목받고 있습니다. 기존 KD 방법들은 시각적 질문 답변 (VQA) 과 같은 시각 인식 (Visual Recognition) 능력 향상에는 효과적이었습니다.

문제점:
저자들은 기존 KD 방법들이 시각적 지각 (Visual Perception) 능력을 효과적으로 증류하지 못한다는 점을 발견했습니다.

• 시각적 인식 vs 시각적 지각: 인식은 객체 분류 등 단순 특징 식별을 의미하는 반면, 지각은 객체 간의 관계, 속성, 공간적 배치 등을 이해하는 더 복잡한 추론 능력을 의미합니다.
• 관찰: 기존 KD 모델들은 VQA 성능은 향상되지만, 관계 추론이 필요한 구성적 추론 (Compositional Reasoning, CR) 태스크에서는 일반 SFT (Supervised Fine-Tuning) 모델과 비슷한 낮은 성능을 보였습니다.
• 근본 원인: 분석 결과, Teacher 와 Student 모델 간의 시각적 주의 (Visual Attention) 불일치 (Misalignment) 가 주된 원인이었습니다. Student 모델은 Teacher 모델이 중요하게 여기는 이미지 영역 (예: 객체 간의 관계가 있는 부분) 에 집중하지 못하고 관련 없는 영역을 주시하는 경향이 있었습니다.

2. 제안 방법: CompoDistill (Methodology)

저자들은 시각적 지각 능력 향상을 위해 CompoDistill이라는 새로운 증류 프레임워크를 제안했습니다. 이는 Teacher 의 시각적 주의 메커니즘을 Student 에게 명시적으로 정렬 (Align) 하는 것을 핵심으로 합니다.

핵심 구성 요소

1. 시각적 주의 정렬 모듈 (Visual ATtention alignment, VAT):

   • 목적: Student 모델의 시각적 주의 분포를 Teacher 모델의 그것과 정렬합니다.
   • 그룹 레이어 매칭 (Group Layer Matching): Teacher 와 Student 의 레이어 수가 다를 때, 단순 1:1 매칭 대신 Student 의 한 레이어를 Teacher 의 여러 레이어 그룹 (Group) 에 매칭하는 '1 대 다 (One-to-Many)' 전략을 사용합니다. 이는 Teacher 의 풍부한 지각 능력을 Student 가 포괄적으로 학습하도록 돕습니다.
   • 손실 함수: Teacher 그룹의 평균 주의 행렬과 Student 의 주의 행렬 간의 코사인 유사도 (Cosine Similarity) 를 최대화하도록 학습합니다.

2. Teacher 어댑터 페치 모듈 (Teacher Adapter Fetch, TAF):

   • 문제: Teacher 의 주의 메커니즘은 Teacher 고유의 비전 - 언어 공간 (Feature Space) 에 최적화되어 있어, Student 에 직접 적용하면 공간 불일치 (Mismatch) 가 발생합니다.
   • 해결: Student 가 Teacher 와 동일한 시각적 입력을 처리할 수 있도록, Teacher 의 프리트레인된 어댑터 (Adapter) 를 고정 (Freeze) 하고 Student 에게 직접 사용합니다. 여기에 경량화된 MLP 만 추가하여 차원을 맞춥니다. 이를 통해 Teacher 의 주의 메커니즘이 Student 의 특징 공간과 호환되도록 합니다.

3. 3 단계 학습 전략:

   • Stage 1 (Distilled Pre-Training, DPT): TAF 모듈을 사용하여 어댑터를 초기화하고, Teacher 의 어댑터와 언어 모델링 손실 (LLM loss) 및 KL 발산 손실을 통해 시각 - 언어 공간 정렬.
   • Stage 2 (Distilled Fine-Tuning, DFT): VAT 모듈을 활성화하여 Teacher 의 시각적 주의 패턴을 Student 에게 증류 (Attention Distillation).
   • Stage 3 (Supervised Fine-Tuning, SFT): 최종적으로 Student 만을 사용하여 지시 따르기 (Instruction Following) 능력을 강화.

3. 주요 기여 (Key Contributions)

1. 문제 발견 및 분석: MLLM 의 지식 증류에서 기존 연구가 간과했던 '시각적 지각 능력 증류의 실패' 를 최초로 규명하고, 그 원인이 시각적 주의 불일치임을 주의 맵 (Attention Map) 분석을 통해 입증했습니다.
2. 새로운 프레임워크 제안: 시각적 인식과 지각 능력을 모두 증류할 수 있도록 VAT와 TAF 모듈을 통합한 CompoDistill 을 제안했습니다.
3. 성능 달성: 기존 KD 방법들의 한계를 극복하여, 구성적 추론 (CR) 태스크에서 4B 모델 수준의 성능을 내면서도 2B 모델 크기를 유지하며 VQA 성능도 유지했습니다.

4. 실험 결과 (Results)

• 성능 향상: CompoDistill(2B) 은 기존 KD 방법들 (LLaVA-KD, LLaVA-MoD 등) 보다 구성적 추론 (CR) 태스크 (SugarCrepe, SADE, BiVLC, Winoground) 에서 압도적으로 높은 성능을 보였습니다. 특히 4B Teacher 모델의 지식을 효과적으로 전달받아 2B Student 가 4B 모델에 버금가는 CR 성능을 달성했습니다.
• VQA 성능 유지: CR 성능이 크게 향상되었음에도 불구하고, 시각적 인식 능력을 평가하는 VQA 태스크 (GQA, TextVQA 등) 에서도 기존 KD 방법들과 유사하거나 더 나은 성능을 유지했습니다.
• 할루시네이션 감소: 객체 간의 관계에 대한 이해도가 높아져 관계 할루시네이션 (Relational Hallucination) 이 감소하는 것을 R-Bench 와 Reefknot 벤치마크를 통해 확인했습니다.
• 확장성 및 일반화:
  • 더 큰 Teacher 모델 (7B) 을 사용할 경우 Student 의 성능이 추가 향상됨을 확인했습니다.
  • 다른 아키텍처 (MobileLLaMA) 에 적용해도 효과가 입증되어 방법론의 일반화 가능성을 보였습니다.
  • 데이터 양을 늘릴 때 CR 성능이 특히 민감하게 반응하여 데이터 스케일링의 중요성을 확인했습니다.

5. 의의 및 결론 (Significance)

이 논문은 멀티모달 모델의 효율화 (Knowledge Distillation) 연구에 중요한 통찰을 제공합니다. 단순히 출력 (Logit) 만을 맞추는 기존 증류 방식이 아니라, 모델 내부의 시각적 주의 메커니즘 (Attention Mechanism) 을 정렬하는 것이 복잡한 추론 능력을 전달하는 핵심임을 증명했습니다.

CompoDistill 은 적은 계산 비용으로 작은 모델이 복잡한 시각적 관계 추론 능력을 갖추도록 하여, 실제 환경에서 배포 가능한 고성능 경량 MLLM 개발의 새로운 방향성을 제시합니다. 이는 특히 시각적 지각 (Visual Perception) 이 필요한 실제 애플리케이션 (로봇 제어, 의료 영상 분석, 복잡한 장면 이해 등) 에 큰 기여를 할 것으로 기대됩니다.