Switch-KD: 시각-언어 모델을 위한 시각 스위치 지식 증류 기법

대형 시각-언어 모델(VLM)을 모바일이나 엣지 기기에 배포하기 위해서는 모델 경량화가 필수적이다. 이 논문은 시각 정보를 언어 모델의 경로로 직접 전달하는 새로운 증류 방식을 통해, 구조 변경 없이도 소형 모델의 멀티모달 이해 능력을 획기적으로 개선했다.

기존의 시각-언어 모델 지식 증류는 시각 정보와 언어 정보를 별도의 경로로 감독하는 경우가 많았다. 하지만 VLM 내부에서 멀티모달 지식은 결국 언어라는 공통의 공간에서 융합된다. 따라서 시각적 특징을 직접 비교하기보다, 시각 정보가 언어 모델을 거쳐 출력되는 '확률 분포' 상에서 교사와 학생을 일치시키는 것이 더 효과적이라는 점에 착안했다.

Switch-KD는 학생 모델이 본 이미지를 교사 모델의 '두뇌(언어 디코더)'에 입력하여 해석하게 만든다. 이를 통해 학생 모델의 시각 인코더가 교사 모델이 이해하기 쉬운 방식으로 정보를 추출하도록 유도한다. 이는 마치 학생의 눈으로 본 세상을 선생님의 지식으로 설명하게 하여, 학생의 시각적 이해도를 선생님 수준으로 끌어올리는 원리와 같다.

또한, 수만 개의 단어 후보 중 실제 정답과 관련 없는 '긴 꼬리(Long-tail)' 영역의 노이즈를 제거하기 위해, 정보가 집중된 상위 로짓 영역을 동적으로 찾아내어 집중 학습한다. 결과적으로 소형 모델은 불필요한 계산을 줄이고 핵심적인 멀티모달 정렬 능력을 빠르게 습득하게 된다.

Switch-KD는 두 가지 주요 경로로 구성된다. 첫 번째는 표준 정렬 경로(Standard Alignment Pathway)로, 학생과 교사의 최종 출력 로짓을 직접 비교한다. 두 번째는 시각 스위치 경로(Visual-Switch Pathway)로, 학생의 시각 인코더 출력을 교사의 프로젝터와 언어 모델에 입력하여 생성된 로짓(z_switch)을 교사의 원래 로짓(z_teacher)과 비교한다. 이 과정은 학생의 시각적 표현이 교사의 언어 체계 내에서 일관되게 해석되도록 강제한다.

핵심 손실 함수인 DBiLD(Dynamic Bi-directional Logits Difference)는 Kneedle 알고리즘을 사용하여 로짓 분포의 곡률이 급격히 변하는 지점(Knee point)을 동적으로 계산한다. [전체 단어 수 N과 정렬된 로짓 값 z_i를 입력으로] → [정규화된 좌표계에서 참조선과의 수직 거리 d_i를 계산하여] → [거리가 최대가 되는 지점 k를 출력하고] → [이 k를 기준으로 정보가 풍부한 상위 영역과 노이즈 영역을 구분한다].

추출된 상위 k개의 로짓들에 대해 내부적인 차이(Pairwise Difference) 벡터를 생성하고, 이를 온도 매개변수 τ가 적용된 Softmax를 통해 확률 분포 p로 변환한다. [차이 벡터 d를 입력으로] → [exp(d/τ) / Σexp(d/τ) 연산을 수행해] → [확률 분포 p를 얻고] → [이 분포 간의 Reverse KL Divergence를 최소화하여 교사와 학생의 상대적 순위 구조를 일치시킨다].

Switch-KD를 적용한 0.5B TinyLLaVA 모델은 10개의 주요 멀티모달 벤치마크에서 기존 SFT 방식 대비 평균 3.6점의 성능 향상을 기록했다. 특히 시각적 세부 이해가 필요한 VizWiz 벤치마크에서는 3.4점의 큰 폭의 상승을 보였으며, MMBench(+4.0), VQAv2(+1.5) 등에서도 우수한 성적을 거두었다.

Ablation Study 결과, 시각 스위치 메커니즘을 제거했을 때 평균 점수가 60.1에서 58.8로 하락하여 해당 모듈의 중요성을 입증했다. 또한 DBiLD 손실 함수는 고정된 k값을 사용하는 방식보다 0.4점 이상의 성능 우위를 보였으며, 이는 모델과 샘플마다 다른 로짓 분포에 유연하게 대응했기 때문으로 분석된다.

교사 모델의 크기에 따른 영향 분석에서는 1.5B 학생 모델의 경우 7B 교사를 사용할 때 가장 높은 성능(65.3점)을 기록했으나, 0.5B 학생 모델의 경우 너무 강력한 7B 교사보다 3B 교사를 사용할 때 오히려 더 나은 학습 효율을 보였다. 이는 학생 모델의 수용 용량에 따른 적절한 교사 선택의 중요성을 시사한다.

Switch-KD 아키텍처는 ViT-Projector-LLM 구조를 따르는 모든 VLM에 범용적으로 적용 가능하다. 핵심은 학생의 시각 인코더 V_s와 교사의 언어 경로(P_t, L_t)를 결합한 하이브리드 추론 경로를 구축하는 것이다. 이를 통해 시각적 특징 공간에서의 직접적인 거리 측정 대신, 텍스트 확률 공간이라는 고차원적인 의미론적 공간에서 모달리티 간 정렬을 수행한다.

DBiLD 손실 함수는 기존의 Forward KL Divergence가 학생의 분포를 과도하게 넓히는 문제를 해결하기 위해 Reverse KL Divergence를 채택했다. 이는 학생 모델이 교사 모델의 높은 확률 영역(High-confidence region)에 집중하게 하여 학습의 안정성을 높인다. 또한 로짓의 절대값이 아닌 상대적 차이(Logits Difference)를 학습함으로써, 교사와 학생 간의 어휘적 편향(Vocabulary bias) 문제를 완화한다.

학습 과정은 Pre-training(PT)과 Distilled Fine-Tuning(DFT)의 2단계로 나뉜다. 실험 결과, 초기 PT 단계보다 실제 하위 태스크를 학습하는 DFT 단계에서 지식 증류를 적용하는 것이 성능 향상에 훨씬 더 결정적인 역할을 하는 것으로 나타났다. 이는 지식 증류가 단순한 정렬을 넘어 복잡한 추론 능력을 전이하는 데 효과적임을 보여준다.

현재 Switch-KD는 교사 모델과 학생 모델 간의 특징 공간(Feature-space) 및 어휘 사전(Vocabulary) 일관성을 전제로 한다. 따라서 아키텍처가 완전히 다른 이종 모델 간의 지식 전이에는 한계가 있을 수 있으며, 향후 아키텍처에 구애받지 않는 매핑 레이어 도입 등의 연구가 필요하다.