본문 바로가기
Paper

LLaVA-MoD Making LLaVA Tiny via MoE-Knowledge Distillation 논문 읽기

by 남디윤 2024. 12. 26.

LLaVA-MoD: Making LLaVA Tiny via MoE-Knowledge Distillation

날짜: 2024년 12월 26일

 

멀티모달 LLM에 Knowledge Distillation 을 접목시키는거에 관심이 있어서 읽어보게된 논문입니다.

원래 아카이브에서부터 조금 주목받던 논문인데 ICLR 2025에 게재될 것 같네용

 

 

 

1. Introduction

  • 멀티모달 LLM은 시각 인코더와 대형 언어 모델을 통합하여 좋은 성능
    • 큰 크기, 방대한 훈련 데이터 → 챌린징 요소
    • 많은 매개변수의 수는 고급 하드웨어 필요, 느린 추론 속도 → 모바일 장치와 같은 실제 배치 작업 어려움
    • 소형 MLLM 탐색 중요
  • 기존 s-MLLM에 관한 연구는 고품질 데이터 수집과 필터링 프로토콜
    • 효과적일 수 있지만, 본질적으로 모델 용량에 의해 제한됨 (모델의 크기가 작기 때문에 충분한 복잡한 지식을 효과적으로 학습하고 담아내기 어렵다는 한계)
    • MLLM을 위한 증류 distillation 프레임워크를 탐색
      • 문제 1) s-MLLM을 위한 경량 구조를 설계, 동시에 l-MLLM으로부터 복잡한 지식을 효과적으로 흡수, 표현 능력 유지
      • 문제 2) 효과적인 지식 이전 확보
  • LLaVA-MoD: Mixture-of-Expert (MoE)와 Knowledge Distillation을 통해 l-MLLM에서 s-MLLM으로 복잡한 다중 작업 지식을 전수하는 문제를 해결하기 위해 제안된 모델
    • 문제 1 해결)
      • 간단한 방법: l-MLLM을 단순히 축소하여 작은 모델을 만드는 것. 직접적인 크기 감소는 모델의 표현 능력 저하, 멀티모달 처리 성능 저하
      • MoE(Dai et al., 2024; Jiang et al., 2024)의 성공 사례 영감
      • s-MLLM에 MoE 구조를 설계하여 스케일 축소를 균형 있게 수행, 증류 과정에서 희소하게 활성화된 전문가를 통해 복잡한 멀티모달 지식 포착
        • s-MLLM에 여러 개의 피드포워드 네트워크(FFN)와 LLM 블록 내에서 선형 게이트를 장착
        • 각 FFN은 l-MLLM에서 세밀한 지식을 포착하는 전문가 역할
        • 게이트는 상위 k 전문가를 동적으로 선택하여 최적의 지식 전송 경로를 식별하고 훈련 및 추론 효율성을 유지
    • 문제 2 해결)
      • 점진적인 증류 전략 제안
      • 학습 가능한 어댑터를 사용하여 비전 인코더와 LLM을 정렬. dense한 s-MLLM을 초기화하는 것으로 시작
      • 💡Dense
        • 모든 파라미터가 학습에 사용되어 일반적인 정보를 잘 캡처할 수 있도록 구성된 모델
        • 모든 파라미터가 활성화되어 학습에 참여하기 때문에, 이 모델은 "trainable parameters"가 많은 상태
        • Sparse 아키텍처는 전체적인 파라미터 수가 많더라도, 실제 학습이나 추론 과정에서 사용되는 파라미터의 수를 줄여 효율성을 높입니다.
        • LoRA (Low-Rank Adaptation)는 일반적으로 모델의 파라미터를 효율적으로 조정하기 위해 사용하는 기법으로, sparsity와 연관이 있지만, 전통적인 sparse 아키텍처와는 다소 다른 방식. LoRA의 기본 아이디어는 기존 대형 모델의 파라미터를 수정하는 대신, 적은 수의 저랭크(low-rank) 조정 파라미터를 추가하여 모델의 성능을 개선하는 것
      • 연결된 신경망이 전체적으로 활성화된 상태를 의미
    • 이후 두 개의 연속적인 증류 단계를 수행
      • Mimic Distillation (조밀-조밀(D2D)과 조밀-희소(D2S) 두 단계)
        • 첫 번째 단계는 일반 지식을 타겟으로 하여 튼튼한 기반을 마련
          • 초기화된 조밀한 s-MLLM과 l-MLLM 간의 출력 로짓 분포를 정렬하기 위해 일반적인 캡션 및 대화 데이터셋을 사용하여 표준 KD 손실을 적용
        • 두 번째 단계는 다중 작업을 처리하기 위한 전문 지식을 타겟
          • s-MLLM이 조밀한 상태에서 희소한 상태로 변환
      • Preference Distillation
        • l-MLLM은 "좋은" 샘플과 "나쁜" 샘플이 무엇인지에 대한 지식을 제공하여 학생 모델을 위한 기초 참조를 설정
        • s-MLLM은 이 지식을 활용하여 확률 분포를 조정하여 좋은 샘플이 l-MLLM의 샘플보다 높은 확률을 가지도록 하고 나쁜 샘플은 낮은 확률
        • l-MLLM 이상의 판단 능력을 향상시켜 s-MLLM의 환각 감소 능력을 증진
  • LLaVA-MoD 성과
    • 최소한의 활성화된 매개변수, 낮은 계산 자원 유지
    • LLaVA-MoD-2B는 벤치마크에서 Qwen-VL-Chat-7B를 평균 8.8% 초과
    • 단 0.3%의 훈련 데이터와 23%의 학습 가능한 매개변수를 사용
    • 7B 및 13B 매개변수를 가진 RLHF 기반 방법의 성능과 일치

 

 

2. Related Work

  • Multimodal Large Language Models
    • BLIP-2: 시각적 특징을 LLM에 적응시키기 위해 추가적인 중간 구조를 추가
    • Flamingo: 정교하게 섞인 다중 모달 시퀀스를 처리하기 위해 LLM에 추가적인 크로스 어텐션 레이어를 통합
    • LLaVA 및 MiniGPT-4: 시각적 지시 튜닝을 통해 모델의 지침 준수 능력을 향상
    • 일부 연구: 더 강력한 비전 인코더, 보다 강력한 세분화된 시각 이해 능력에 초점
    • 본 논문: 모델 크게 확장 필요 X, 성능을 유지하면서 계산 및 저장 효율성을 개선하기 위해 증류(distillation)와 MoE(Mixture of Experts)를 결합하는 것을 목표
  • Knowledge Distillation.
    • 큰 모델을 교사 모델로 사용하여 그 고급 지식을 더 작은 학생 모델로 전이하는 방법
    • MiniLLM: 학생 모델이 교사 분포의 낮은 확률 영역을 과도하게 추정하는 것을 방지하기 위해 역 KL 발산을 최소화
    • GKD: 일반화된 지식 증류를 도입하고 증류와 RLHF의 통합
    • 일부 연구: 특정 기술을 향상시키기 위해 맥락 증류(Huang 외, 2022) 및 사고 과정 체인(CoT) 증류(Mukherjee 외, 2023; Li 외, 2022; Ho 외, 2022)를 채택
    • 본 연구: MLLM의 지식을 더 작은 희소 MoE 구조로 혁신적으로 증류, 최소한의 비용으로 작은 모델의 멀티 모달 처리 능력을 크게 향상
  • Mixture-of-Experts
    • 전문가 혼합(MoE) 아키텍처(Jacobs 외, 1991): 독립적인 전문가를 활용하여 다양한 샘플을 처리함으로써 성능을 향상
      • 변환기 기반 아키텍처에서 피드포워드 신경망(FFN) 층이 전문가 역할
      • 게이팅 전략(Lepikhin 외, 2020; Fedus 외, 2022)을 통해 희소하게 활성화
      • → 계산 오버헤드를 낮추면서 모델 용량을 효과적으로 증가
    • 잘 훈련된 조밀한 모델의 전문가를 사용하여 전문가를 초기화하는 희소 업사이클링 방법(Komatsuzaki 외, 2022)이 훈련 비용을 줄이기 위해 제안
      • MoE는 언어 모델(Jiang 외, 2024; Dai 외, 2024)뿐만 아니라 비전 모델(Riquelme 외, 2021) 및 비전-언어 모델(Lin 외, 2024a; Shen 외, 2023)에서도 큰 잠재력

 

 

3. Method

  • Mixture-of-Experts(MoE)와 지식 증류를 사용하여 효율적인 s-MLLM을 구축하기 위한 새로운 프레임워크
  • 두 가지 주요 구성 요소
    1. s-MLLM의 아키텍처 설계: MoE가 있는 희소 s-MLLM을 설계하여 훈련 및 추론 효율성을 유지하면서 전문화된 전문가 지식을 습득할 수 있는 능력을 강화
    2. 증류 메커니즘: l-MLLM에서 희소 s-MLLM으로 지식을 전이하기 위한 점진적 증류 메커니즘을 설계. 모방 증류(mimic distillation)와 선호도 증류(preference distillation)의 두 단계가 포함

 

3.1 Architecture Design of Sparse s-MLLM

  • s-MLLM 정의
    • s-MLLM의 기본 아키텍처는 비전 인코더, 대규모 언어 모델(LLM), 및 비전-언어(VL) 어댑터의 세 가지 주요 구성 요소로 구성
  • Sparsify s-MLLM
    • 본 연구의 s-MLLM 구축 원리: LLM의 크기를 축소, 비전 인코더와 비전-언어 어댑터는 수정하지 않는 것
    • 축소 목표를 달성하기 위해
      • MoE 아키텍처를 통합하여 조밀한 s-MLLM을 희소화
      •  
        • N 개의 피드포워드 네트워크(FFN)를 전문가 모듈로 복제하는 희소 업사이클링 기술을 적용
        • 적절한 전문가를 동적으로 활성화하기 위해 전문가 할당 확률을 예측하는 라우터로 선형 레이어를 도입
        • 각 시퀀스의 토큰 x 에 대해, 먼저 N 전문가의 라우팅 값을 계산
          • $r = \text{Softmax}(x \cdot W_r)$
          • W_r은 라우터의 가중치 행렬
          • r의 각 요소 r_i 는 i 번째 전문가를 활성화할 확률
          • Top-k 전략을 적용하여 가장 높은 k 라우팅 값을 가진 활성화 전문가를 결정
        • 비활성화된 전문가의 라우팅 값은 0으로 설정
          • 최종 출력에 기여하지 않도록 효과적으로 제외
          • 출력 y 는 활성화된 전문가의 기여를 해당 라우팅 값으로 가중합하여 계산
         
    • LLM 내부의 특정 구조를 희소화하는 방식
      • 내부의 특정 부분(FFN 등)을 희소 구조로 변환하여 효율성을 극대화하는 방식
      • LLM 자체의 핵심 아키텍처는 유지되며, 변경되는 부분은 전문가 모듈과 라우팅 메커니즘에 국한

 

3.2 Progressive Distillation

  • 두 단계
    • 모방 증류: s-MLLM π_S는 일반 및 특정 지식 l-MLLM πT로부터.
    • 선호 증류: 선호 증류 단계에서 πS는 πT의 선호 지식을 습득하여 출력을 더욱 세분화하고 환각 감소
      • π_S, π_T 는 각 MLLM 모델을 의미
      • image 1
  • 초기화
    • 증류 전에 먼저 학습 가능한 어댑터를 통해 비전 인코더를 LLM과 정렬하여 잘 초기화된 밀집 버전 πS를 얻는 것을 목표
      • LLMϕ 및 ViTχ는 이미 풍부한 시각 및 언어 지식을 캡처했기 때문에 고정 상태로 유지
      • 비전과 언어 영역 간의 격차를 연결하기 위해서만 Projω를 최적화
  • Mimic Distillation
    • 일반에서 전문화된 모방 증류를 수행하며, 이는 두 단계로 구성.
      • 밀집에서 밀집으로(D2D)와 밀집에서 희소로(D2S) 지식을 πS로 전송
      • D2D 단계에서 πS의 밀집 구조를 활용하여 일반 지식을 습득
      • D2S 단계에서는 그것을 희소 구조로 변환하여 복잡한 전문 지식을 습득
    • a) Dense-to-Dense Distillation
      • l-MLLM의 일반 지식을 복제하는 것을 목표
      • ViTχ는 고정 상태로 유지하고, LLMϕ와 Projω를 공동 최적화. 학습 가능한 매개변수 θ = {ω, ϕ}
      • 일반적인 이미지-캡션 쌍 및 대화 데이터셋을 활용
      • 훈련 목표는 s-MLLM과 l-MLLM의 출력 로짓 간의 쿨백-라이블러 발산(KLD)을 최소화하는 것
    • b) Dense-to-Sparse Distillation
      • l-MLLM의 전문 지식을 s-MLLM으로 이전하여 복잡한 작업에서 우수한 성능을 달성하는 데 중점
      • s-MLLM의 밀접한 형식에서 이 지식을 직접 학습하는 것은 다양한 작업의 지식 구조때문에 불충분 가능성
        • LLMϕ 내에 N개의 피딩포워드 네트워크(FFN)를 복제
        • 라우터로서 MLP 레이어를 추가
        • 매개변수 ϕe를 갖는 전문가를 형성
        • 희소 아키텍처는 s-MLLM이 서로 다른 입력에 따라 가장 관련성이 높은 전문가를 선택적으로 활성화할 수 있게 하여 l-MLLM의 전문 지식을 모방하는 데 큰 장점을 제공
      • 훈련을 위해 우리는 다중 작업 데이터를 활용하여 전문가와 어댑터만 업데이트
        • 전문가를 선택하기 위해 Top-k 라우팅 전략을 사용
        • 학습 가능한 매개변수는 θ = {ω, ϕe}
        • 이전 단계와 유사하게, 우리는 훈련 목표로 KLD를 채택
        • 교차 엔트로피 손실을 최소화
  • Preference Distillation
    • 목표: l-MLLM에서의 선호 지식을 증류하여 s-MLLM이 정확할 뿐만 아니라 합리적인 응답을 생성하도록 안내하는 것
      • 환각을 줄이는 데 중요
      • 훈련 과정에서는 정밀하게 쌍을 이룬 긍정적 응답 $y^+$와 부정적 응답 $y^-$로 구성된 선호 데이터를 효과적으로 사용
      • s-MLLM이 l-MLLM에 비해 긍정적인 응답에는 더 높은 확률을 부여하고 부정적인 응답에는 더 낮은 확률을 부여하도록 최적화하는 것
    • 두 가지 주요 최적화 측면을 포함
      • s-MLLM은 긍정적인 응답과 부정적인 응답을 구별하는 데 있어 교사 모델과 정렬하는 것
      • s-MLLM은 긍정적인 응답에 더 높은 확률을, 부정적인 응답에 더 낮은 확률을 부여하여 l-MLLM을 초월하는 것
    • s-MLLM에서 전문가와 적응기만 훈련하며 Top-k 라우팅 전략을 사용하여 전문가를 선택
    • 훈련 가능한 매개변수는 θ = {ω, ϕe}

 

 

4. Experiments

4.1 Experimental Settings

  • Implementation Details
    • ViT-MLP-LLM 아키텍처 사용
    • l-MLLM 및 s-MLLM에 대해 서로 다른 크기의 Qwen-1.5/2를 LLM으로 사용
      • l-MLLM은 7B 매개변수로 구성
      • s-MLLM은 1.8B 및 0.5B 매개변수로 구성
      • Qwen-1.5 7B를 사용하여 Qwen-1.5 1.8B 및 Qwen-1.5 0.5B을 증류
    • 멀티모달 벤치마크
  • Training Datasets
    • 오픈 소스 데이터셋에서 5M 샘플로 구성. 각 훈련 단계는 서로 다른 데이터셋 사용
    • 초기화 단계: 0.6M의 일반 캡셔닝 샘플을 사용하여 시각적 모달리티와 언어 모달리티 간의 격차 해소
    • 모방 증류:
      • 2.4M의 일반 캡셔닝 및 대화 샘플을 사용하여 l-MLLM에서 일반적인 지식을 증류
      • 1.4M의 다양한 작업 데이터, 포함 VQA, 문서, 과학, OCR을 사용하여 l-MLLM에서 전문 지식을 증류
    • 선호 증류: l-MLLM에서 선호 지식을 전달하기 위해 80K의 선호 데이터 샘플이 사용
  • Evaluation Benchmarks
    • 다양한 하위 작업을 포함하여 멀티모달 이해 및 추론 능력을 포괄적으로 평가
    • 일반 VQA, 텍스트 지향 VQA 및 과학 VQA를 포함하여 다양한 VQA 작업에서 실험을 수행
      • 일반 VQA: 일반적인 시각적 이해 및 관계적 추론을 테스트
      • TextVQA: 텍스트 지향 VQA 작업에 사용되어 이미지 내 텍스트의 세밀한 시각적 인식 및 이해에 중점
      • ScienceQA: 과학적 지식을 측정하는 데 활용
      • POPE, Object HalBench, MMHal-Bench: 여러 환각 벤치마크

 

4.2 Main Results

  • LLaVA-MoD의 두 가지 측면에서의 장점을 강조하기 위해 실험 수행
    • 성능과 효율성
    • 성능: 이해 중심의 벤치마크(Tab. 2)와 환각 중심의 벤치마크(Tab. 3)를 평가
    • 효율성: 데이터 샘플과 모델 크기를 기준으로 비교를 제시
  • 이해 중심의 벤치마크 Comprehension-Oriented Benchmarks
    • LLaVA-MoD는 이해 중심 벤치마크에서 1B 및 2B 모델 크기 중 SOTA 평균 결과를 달성
    • 본 연구 접근 방식: 대규모 MLLM에서 희소 MoE 아키텍처를 증류하여 소규모 MLLM을 효과적으로 학습시킨다는 점을 강조
    • Samples 는 각 모델이 학습한 데이터 수
    •  
    • image 2
  • 환각 중심의 벤치마크 Hallucination-Oriented Benchmarks
    • LLaVA-MoD는 환각을 완화하는 데 있어 주목할 만한 성능을 보여 주며, 심지어 그 교사 모델을 초월하기도 함
    • 두 가지 측면의 선호 증류에 기인
      • 긍정적인 응답에 대해 더 높은 확률을 할당함으로써, 선호 증류는 LLaVA-MoD가 올바르고 관련 있는 정보를 제공하는 데 집중하도록 유도
      • 부정적인 응답에 대해 더 낮은 확률을 할당함으로써, 선호 증류는 부정확하거나 근거 없는 정보를 억제
    • 최근 RLHF 기반 모델도 초과
    • image 3
  • 효율성 비교 Efficiency Comparison
    • 표2
    • 2B 크기의 LLaVA-MoD는 단지 5M 샘플을 필요로 하며 훈련 및 추론 중 총 2.2B 매개변수의 극히 일부만 활성화
    • 데이터, 매개변수 및 계산 자원의 보다 효율적인 활용을 가능
    • 효율적인 훈련과 추론을 초래할 뿐만 아니라 자원이 제한된 환경에서도 고성능 모델을 배포하는 실용적인 솔루션을 제공

 

4.3 Ablation Study

  • 이해력과 환각 벤치마크에서 선호 증류의 영향을 조사
  • 선호 증류의 영향
    • 선호 증류는 환각을 완화
    • s-MLLM의 환각을 현저하게 줄이는 반면, 이해 능력이 일관되게 향상되지 않는다는 점에서 이러한 관찰은 LLM의 경우와 일치
    • 선호 최적화는 환각을 줄이는 데 우선 순위를 두며, 종종 이해 능력의 감소를 초래

image 4

  • KTO Stabilizes Preference Distillation
    • 다양한 선호 최적화 손실의 훈련 안정성을 조사
    • KTO(Ethayarajh et al., 2024)와 DPO(Rafailov et al., 2024)를 사용
    • 두 방법 모두 환각 완화를 증가
    • 샘플이 좋은지 나쁜지를 나타내는 KTO가 s-MLLM이 l-MLLM을 초과하는 데 더 나은 신호를 제공함을 의미
    • 💡
    • KTO와 DPO*주로 다중 모달 학습 및 지식 증류 과정에서 사용KTO는 2024년에 Ethayarajh 등에서 제안된 방법으로, 샘플이 "좋은"지 "나쁜"지를 판별하여 s-MLLM(작은 다중 모달 언어 모델)이 l-MLLM(큰 다중 모달 언어 모델)을 초월할 수 있도록 더 나은 신호를 제공합니다. KTO는 각 샘플의 질을 직접 판단할 수 있는 기능이 있어 모델의 성능 개선에 도움을 줍니다.
      연구 결과에 따르면, KTO는 DPO보다 우수하여 환각(hallucination)을 감소시키는 데 더 효과적임을 보였습니다.DPO는 2024년에 Rafailov 등에서 제안된 방법으로, 두 샘플을 비교하여 어느 것이 더 나은지를 결정합니다. 즉, DPO는 샘플 간의 상대적인 질을 평가하는 방식입니다.
      그러나 연구에서는 DPO가 KTO에 비해 덜 효과적이라고 여겨졌습니다. 이는 DPO가 상대적인 비교에 의존하기 때문에 발생하는 결과로, KTO가 더 직접적으로 샘플의 질을 반영할 수 있다는 장점이 있습니다.
    • DPO (Decision Preference Optimization):
    • KTO (Known Truth Optimization):
    • 모두 모델의 학습 과정에서 사용되는 손실 함수
  • 훈련 전략의 영향
    • KD Facilitates MoE Training
      • KD가 MoE 훈련에 미치는 영향을 감독된 미세 조정(SFT)과 비교하여 조사
      • E4T2의 동일한 희소 구성으로 수행 (4개의 전문가 중 2개가 활성화되는 구성.)
      • KD가 MoE 아키텍처를 효과적으로 훈련하기 위한 더 나은 최적화 신호를 제공
      • image 5
    • General-to-Specialized KD Boosts Performance
      • 제안한 D2D+D2S를 D2S와 비교하여 모방 증류에서 일반화에서 전문화로의 접근 방식의 영향을 조사
      • D2D+D2S가 D2S를 초과하여 평균 4.0%의 향상
      • D2D+D2S가 일반 및 전문 역량 사이의 균형을 이루는 효과적인 지식 전이
      • D2D+D2S는 D2S가 필요로 하는 GPU 시간의 62.5%만을 소모함으로써 계산적으로도 더 효율적
      • image 6
    • Focus on Response Distillation Improves Generalization
      • 응답 증류와 응답+지침 증류를 비교하여 증류 목표의 영향을 조사
      • 응답 증류가 응답+지침 증류를 초과하여 평균 3.1% 향상
      • 응답에 집중하는 것이 자기 회귀 모델링에서의 지식 증류에 더 효과적
      • image 7
  • 모델 아키텍처의 영향
    • Sparse Architecture Facilitates Knowledge Transfer.
      • 희소 아키텍처와 밀집 아키텍처를 비교하여 증류에서의 효과를 탐구
      • 구성은 E4T1이며, 네 개의 전문가가 초기화되고 최상위 전문가가 활성화
      • 결과는 희소 아키텍처가 MoE를 활용하여 l-MLLM으로부터 다양한 지식을 효과적
      • image 8
    • Teacher Capacity Matters in Knowledge Distillation.
      • 교사 용량은 지식 증류에서 중요
      • 강한 교사로 Qwen-1.5-7B를, 약한 교사로 Qwen-1.5-4B를 사용
      • Qwen-1.5-1.8B의 경우, 7B 크기의 교사를 사용할 때 4B 크기의 교사에 비해 평균 2.2%의 이득을 기록
      • 학생과 교사 간의 용량 차이가 클 경우 효과적인 지식 전이를 방해 가능성
      • 중간 용량의 "중간 교사"를 활용하면 격차를 좁혀 부드러운 지식 전이가 가능
      • image 9

 

 

5. Conclusion

  • LLaVA-MoD를 소개
  • 대규모 모델로부터의 지식 증류를 통해 소규모 다중 모달 언어 모델의 효율적인 훈련을 위한 새로운 프레임워크
  • MLLM 증류에서 두 가지 주요 문제
    • 효율성과 표현성 균형을 위한 MoE 설계를 통한 s-MLLM 아키텍처 개선 및 점진적인 지식 전이 전략 구현