[논문 요약/리뷰] MoVE-KD: Knowledge Distillation for VLMs with Mixture of Visual Encoders

 

MoVE-KD: Knowledge Distillation for VLMs with Mixture of Visual Encoders

Author : Cao, Jiajun, et al.
Journal : Arxiv
Keyword : Knowledge distillation
Published Date : 2025년 1월 3일

---

• 다양한 VLM(CLIP, EVA, ConvNeXt)의 Vision encoder들은 Vision-language 하위 분야에서 좋은 성능을 낸다고 알려져있고, 이러한 다양성은 Vision encoder 관련 연구가 중요한 주제로 부상하고있다.
• 기존 연구들은 Vision encoder들의 다양성을 반영하기 위해, Feature concatenation, Attention mechanism등을 사용해 하나의 VLM 내에서 다중 인코더를 사용하는 방식을 채택했다.
• 그러나 이 방법은 하나의 Vision encoder만을 사용했을 때와 비교해서 Computational cost와 Model Complexity을 증가시켜 모델 효율성과 확장성을 감소시킨다.
• 따라서 본 연구는 “다양한 인코더들의 특징을 하나의 비전 인코더에 증류할 수 있을까?” 라는 질문에 대해 탐구한다.
• 기존 KD 기법은 주로 일대일(One-to-one) 증류에 초점을 맞추고 있기 때문에, 다양한 데이터셋, Objective로 사전 학습된 모델을 동시에 증류하는 문제는 충분히 탐구되지 않았다.

---

Contributions

• LVLMs로부터 다양한 Vision encoder를 융합합할 수 있는 MoVE-KD 프레임워크를 제안한다.
• Attention-guided KD 정규화를 도입한다.
• LLaVA 및 LLaVA-NeXT 모델에 적용하여 SoTA 성능을 달성한다.

---

Method

Overview

MoVE-KD 모델의 Pipeline

• 본 논문에서는 MoVE-KD를 제안한다. 이는 여러 Vision encoder를 활용한 새로운 지식 증류 방법이다.
• 먼저 Encoder Adpater를 통해 VLM들의 출력을 같은 통합된 특징 공간에 매핑한다.
• 사전학습된 CLIP 모델의 [CLS] 토큰의 Attention을 바탕으로, 교사 인코더들과 비주얼 토큰들의 Weight를 계산한다.
• 교사 가중치와 토큰 가중치로부터 Weighted sum으로 KD Loss를 계산한다.
• 또한 여러 출처의 지식을 학습하면서 발생할 수 있는 충돌을 완화하기 위해, 학생 인코더 내에 Misture-of-LoRA-Experts(MoLE) 구조를 통합한다.
• 최종 Training Objective는 KD Loss와 LLM으로부터 생성된 Text Loss를 최소화하는 것이다.

Encoder Apader

• 교사 VLM에서 출력된 Visual token들은 표현 공간이 일치하지 않기때문에 학생 모델의 Visual token과 바로 정렬될 수 없다. 따라서 두 층의 MLP layer로 구현된 Encoder Adpater를 도입한다.

Mixture-of-LoRA-experts (MoLE)

MoLE 모듈

• 학생 인코더는 사전학습된 CLIP의 Vision encoder로부터 초기화되어 교사 토큰에 대해 Fine-tuning한다.
• 하지만 학생 인코더를 직접 Fine-tuning하는 과정에서 문제가 발생했다.
  • 목표 데이터셋에 직접 Fine-tuning하면 과적합이나 치명적인 망각 문제(Catastrophic forgetting)가 발생하여 모델의 정확도와 일반화 능력에 영향을 미친다.
  • 또한, 다양한 교사 인코더들의 표현을 하나의 모델로 통합하는 것이 어렵다.
• 위 문제를 해결하기 위해 MoLE 아키텍쳐를 도입한다.

---

Mixture-of-Experts(MoE) + LoRA(Low-Rank Adaptation)

• 위 이미지를 보면 Layer Norm과 FFN 사이 Router가 존재한다.
• 이 Router는 세 개의 전문가(Experts) 중 입력에 맞는 전문가를 활성화한다. 위 사진에서는 3번 전문가가 활성화 된 것으로 볼 수 있다.
• 이 때 MoE의 Output은 아래와 같이 수식화할 수 있다.

$$ F^*(x)=F(x)+E_i(x) \ \mathrm{with} ; i=\textrm{argmax}(\mathrm{Softmax}(f(x))) $$

• $F(x)$는 원래의 FFN 출력을 의미하고, $f(x)$는 Router로 Linear layer이다. $E_i$는 $i$번째 전문가 네트워크이다.
• 즉 전문가 네트워크 중 Router의 Softmax값이 가장 큰 전문가 네트워크가 선택되는 것이다.
• 기존 MoE에서는 FFN 레이어를 복제하여 개별 전문가로 사용했는데, 이는 파라미터 수와 훈련 시간이 크게 증가하기 때문에 LoRA를 사용하여 효율적으로 수행한다.

---

Attention-guided KD Regularization

• 학생 모델은, 여러 개의 교사 모델로부터 필요한 정보와 필요하지 않은 정보를 구분해서 학습해야 한다. 따라서 적절한 제약 조건으로 Distillation을 규제하는 방법이 필요하다.
• 즉 학생 모델의 토큰이 교사 모델의 토큰으로부터 잘 정제되어야 하는 것이고, 이상적인 Distillation loss는 아래와 같다.

$$ \mathcal L_{kd} = \sum^m_{i=1} W^{(tea)}i \sum^N{j=1}(W^{(tok)}j+\frac{1}{N})\mathrm{MSE}(V^{(t)}{i,j}V^{(s)}_j) $$

• $m$은 교사 모델의 개수, $n$은 토큰의 길이이고 $W^{(tea)}, W^{(tok)}$는 교사 모델과 토큰의 가중치 벡터이다. 또한 $V^{(t)}, V^{(s)}$는 각각 교사, 학생의 토큰이다.
• $W^{(tea)}, W^{(tok)}$는 위에서 설명한 것처럼 사전학습된 CLIP의 [CLS] Attention으로부터 만들어진다.

 

• CLIP의 [CLS] Attention을 사용한 이유는, 위 그림에서처럼 이미지의 핵심 영역을 잘 포착하기 때문이다.

Token weight

• 따라서 CLIP의 [CLS] 토큰 $V^{(cls)}$와 나머지 토큰 $V^{(res)}$간 Attention을 계산하고 이를 각 토큰의 가중치로 사용한다.

$$ W(\mathrm {tok}) = \mathrm{Softmax}(\frac{(V^{(cls)}W^{(Q)}) \cdot (V^{(res)}W^{(V)})^T}{\sqrt d}) $$

• $W^{(tok)}$는 토큰의 가중치이고, $W^{(Q)}, W^{(V)}$는 CLIP의 출력 레이어의 Query, Key를 의미한다.

Teacher weight

• 교사 수준에서의 규제는 [CLS] 토큰과 $i$번째 교사의 토큰 $V^{(t)}_i$ 사이 Attention 평균값에 Softmax를 취하여 가중치를 계산한다.

$$ W^{(tea)} = \mathrm{softmax}(mean(\frac{V^{(cls)} \cdot V^{(t)^T}_i}{\sqrt d})) $$

• $W^{(tea)}$는 교사 가중치를 의미한다.

Overall loss

• 전체적인 손실은 Text loss와 KD loss의 합으로 계산한다.

$$ \mathcal L_{total} = \mathcal L_{text} + \lambda_{kd} \cdot \mathcal L_{kd} $$

• $\lambda$는 KD loss 항의 가중치이다.

---

Experiment

• Model : 증류 모델은 LLaVA와 LLaVA-NeXT 사용, 교사 인코더는 CLIP, EVA-02, ConvNeXT 사용.
• Dataset : Pre-training 단계에서 LLaVA Visual Instruct Pretrain LCS-558K 사용, Fine-tuning 단계에서 LLaVA-1.5, LLaVA-NeXT 데이터셋 사용
• Pre-training : 학생 인코더의 MoLE, 인코더 어댑터, 프로젝션 매개변수 가중치 조정
• Fine-tuning : 교사 인코더를 제외한 모든 매개변수 업데이트
• Hyper parameters : CLIP 교사 인코더에 0.8 가중치 지정, MoLE의 Expert 수는 3, LoRA의 rank r은 32로 설정, Distillation loss의 가중치는 0.5
• Setup : 16 * A800 GPU 사용

 

• MoVE-KD는 VLM에서 Vision encoder를 증류하는 최초의 사례이기 때문에, 같인 기준에서 비교 가능 사례가 존재하지 않는다고 언급.
• 위 테이블은 1.7B, 7B, 13B 모델에 대한 비교 벤치마크 기재, 여기서 비교되고 있는 모델은 Distillation을 거친 모델이 아니다.

 

• 위 그래프는 8개 벤치마크에 대한 MoVE-KD의 Abliation study이다.
• 각 구성 모듈을 점진적으로 추가하며 실험한 결과이다.

 

• MoLE에 대한 w/o 실험 결과 비교 테이블이다.
• 기존 모델과 비교했을 때, MoLE를 도입함으로써 Trainable parameter 늘어나게 된다. 이는 전체 파라미터의 0.3% 정도이지만, 성능 향상이 단순히 파라미터 수의 증가때문일 수도 있기 때문에, 이 비교 실험을 추가한 것이다.
• 단순히 MoLE만을 도입하는 것만으로는 성능이 감소하는 경우도 생기는 것으로 이를 증명한다.
• MoLE는 증류와 함께 사용되면 교사 인코더 모델의 증류 과정에서, 여러개의 인코더로부터 정보를 잘 혼합해야하는데, 이 과정에서 정보의 선택을 잘 도와주는 역할을 한다는 것이다.

---

Conclustion

• LVLM을 위해 서로 다른 인코더를 통합하는 최초의 프레임워크 Mixture-of-Visual-Encoder Knowledge Distillation (MoVE-KD)를 소개합니다.
• LoRA와 MoE를 활용하여 교사 인코더들의 충돌을 완화하고 특성을 보존했다.

---

Reivew

• 구조가 복잡한데, Teacher weight나 Token weight를 계산하는 부분이 수식만 있고 설명이 명확하지 않은게 아쉽다.
• MoLE만의 효과가 아니라고 하기엔 전체 성능 증가량과 비교했을 때, 반을 넘게 차지하는 경우가 많다.