기술 요약: Hide to See: VLM 증류에서의 시각 기반 추론을 위한 추론 접두사 마스킹

문제 정의

최근 Qwen3-VL-Thinking과 같은 비전 - 언어 모델 (VLM) 의 "생각 - 답변 (think-answer)" 접근법은 최종 답변을 생성하기 전에 중간 추론 단계를 생성함으로써 추론 성능을 크게 향상시켰습니다. 그러나 이러한 긴 추론 궤적은 상당한 계산 비용을 초래하여 자원 제약이 있는 시나리오에서의 배포를 제한합니다. 지식 증류 (KD) 는 대형 교사 모델의 능력을 소형 학생 모델로 이전함으로써 이러한 문제에 대한 해결책을 제시합니다.

think-answer VLM 을 증류하는 데 있어 중요한 과제는 **시각 망각 (visual forgetting)**입니다. 학생 모델이 긴 추론 궤적을 생성할 때, 후속 토큰을 예측하기 위해 누적된 텍스트 추론 접두사(이전에 생성된 토큰) 에 크게 의존하는 경향이 있습니다. 이러한 접두사는 종종 학생이 시각 입력을 다시 검토할 필요 없이 교사의 궤적을 모방하기에 충분한 정보를 포함하고 있습니다. 결과적으로 학생은 시각적 증거를 무시하도록 학습되어 시각 기반 추론 능력이 저하됩니다. 시각 주의 맵이나 토큰 관계를 이전하는 기존 증류 방법들은 텍스트 단서를 통한 이러한 특정 단축 경향성을 해결하지 못합니다.

방법론: Masking-KD

저자들은 의도적으로 가장 영향력 있는 텍스트 단서를 가려 학생 모델이 시각적 증거에 추론을 고정하도록 강제하는 새로운 증류 프레임워크인 Masking-KD를 제안합니다.

핵심 메커니즘

표준 인과적 마스크(학생이 모든 이전 토큰을 볼 수 있게 함) 대신, Masking-KD 는 **주요 추론 접두사 마스킹 (salient reasoning-prefix mask)**을 사용합니다. 이 마스크는 미래 토큰에 대한 접근 (표준 인과성) 과 학생이 교사를 모방하기 위해 의존할 것으로 예상되는 특정 주요 추론 접두사에 대한 접근을 모두 차단합니다. 이러한 고영향 텍스트 단서를 제거함으로써 학생은 누락된 정보를 보완하기 위해 시각적 증거를 활용하도록 강요받습니다.

이 프레임워크는 두 가지 적응형 구성 요소로 구성됩니다:

1. 토큰별 주요 추론 접두사 마스킹 (Token-wise Salient Reasoning-Prefix Masking):

   • 식별: 주요 증류 단계 전에 표준 인과적 마스크 하에서 학생 모델의 **보조 순전달 (auxiliary forward pass)**이 수행됩니다. 이는 현재 토큰을 예측할 때 학생이 가장 중점적으로 주의하는 이전 토큰을 보여주는 **응답 - 응답 주의 맵 (response-to-response attention map, $A^{resp}$)**을 생성합니다.
   • 선택: 각 다음 토큰 예측에 대해, 주의 가중치에 기반하여 상위 $\rho$개의 가장 주요한 접두사를 식별합니다. 이들은 답변으로 가는 가장 강력한 "단축"을 제공하는 토큰들입니다.
   • 마스킹: 식별된 주요 접두사들은 실제 증류 훈련 단계 동안 마스킹됩니다 (주의 로짓에서 $-\infty$로 설정됨). 이를 통해 학생이 이에 주의를 기울이는 것을 방지합니다.

2. 자기 주도적 마스킹 예산 스케줄링 (Self-Paced Masking Budget Scheduling):

   • 각 토큰을 증류하는 난이도는 다릅니다. 쉽게 모방 가능한 토큰 (교사와 학생 간의 편차가 작음) 은 학생이 접두사를 단순히 복사할 수 있다면 약한 학습 신호를 제공합니다.
   • 이 방법은 보조 순전달 동안 교사와 학생 분포 간의 **토큰별 역 KL 발산 (token-wise reverse KL divergence)**을 계산합니다.
   • **누적 비율 임계값 ($\rho_n$)**은 각 토큰의 난이도에 따라 적응적으로 할당됩니다. 쉽게 모방 가능한 토큰 (낮은 KL 발산) 은 학생이 텍스트 단서에 의존하도록 강제하기 위해 높은 마스킹 예산(더 많은 접두사가 마스킹됨) 을 받습니다. 학습이 더 어려운 토큰은 안정적인 훈련을 보장하기 위해 더 적은 마스킹을 받습니다.

훈련 목적 함수

학생은 마스킹된 컨텍스트 하의 예측 분포와 전체 컨텍스트 하의 교사 분포 간의 역 KL 발산을 최소화하도록 훈련됩니다. 손실 함수는 다음과 같이 정의됩니다:
$$\mathcal{L}_{Distill} = \frac{1}{N} \sum_{n=1}^{N} \sum_{y \in \mathcal{V}} p_s(y \mid \mathbf{x}_v, \mathbf{x}_q, \mathbf{y}_{<n}, \tilde{\mathbf{M}}) \log \frac{p_s(y \mid \mathbf{x}_v, \mathbf{x}_q, \mathbf{y}_{<n}, \tilde{\mathbf{M}})}{p_t(y \mid \mathbf{x}_v, \mathbf{x}_q, \mathbf{y}_{<n}, \mathbf{M})}$$
여기서 $\tilde{\mathbf{M}}$은 주요 추론 접두사 마스크이고 $\mathbf{M}$은 표준 인과적 마스크입니다.

주요 기여

1. Masking-KD 프레임워크: 주요 추론 접두사를 마스킹하여 학생이 텍스트 단축에 의존하는 대신 시각적 증거를 활용하도록 강제함으로써 시각 기반 사고를 강화하는 새로운 증류 패러다임.
2. 토큰별 적응형 마스킹: 디코딩 단계마다 중요한 컨텍스트가 다르다는 점을 인정하여, 각 특정 토큰 예측에 대해 가장 영향력 있는 추론 접두사를 동적으로 식별하고 마스킹하는 메커니즘.
3. 자기 주도적 예산 스케줄링: 각 토큰의 증류 난이도 (KL 발산) 에 따라 마스킹 강도를 조정하여 쉽게 모방된 토큰에 대한 학습 신호를 증폭시키는 적응형 전략.
4. 실증적 검증: 기존 VLM 증류 방법 및 자기 증류 베이스라인보다 우수한 성능을 입증한 광범위한 실험.

실험 결과

저자들은 Geometry-3K, MathVista, We-Math, MMK12, MathVerse, LogicVista, MMMU-Pro를 포함한 다중 모달 추론 벤치마크에서 Masking-KD 를 평가했습니다.

• 성능 향상: Masking-KD 는 모든 모델 크기 (2B, 4B, 8B) 에서 최첨단 (state-of-the-art) 결과를 달성했습니다. 특히, 증류된 4B 모델은 증류되지 않은 8B 교사 모델을 능가했으며, 2B 모델은 증류되지 않은 4B 교사 모델을 능가했습니다.
• 베이스라인 비교: 이 방법은 표준 응답 수준 증류뿐만 아니라 시각 패턴을 명시적으로 이전하는 방법 (예: LLaVA-KD, CompoDistill, Align-TI) 보다 현저히 우수한 성능을 보였습니다.
• 자기 증류: 이 프레임워크는 학생이 자신의 교사로 역할을 하는 자기 증류 설정에서도 효과적이었으며, OPSD 와 같은 다른 자기 증류 방법보다 우수한 성능을 발휘했습니다.
• 시각 주의 분석: 정성적 및 정량적 분석은 Masking-KD 모델이 베이스라인에 비해 생성 전반에 걸쳐 더 높은 시각 주의 비율을 유지하여 "시각 망각" 현상을 효과적으로 완화했음을 확인했습니다. 텍스트 단서가 마스킹되었을 때 모델들이 관련 이미지 영역을 더 강하게 활성화하는 것이 관찰되었습니다.

중요성 및 주장

이 논문은 think-answer VLM 을 증류하는 데 있어 주요 병목 현상은 누적된 추론 접두사가 제공하는 텍스트 단축을 학생이 활용하여 시각 입력을 소홀히 한다는 점이라고 주장합니다.

Masking-KD 의 중요성은 교사 모델의 시각적 지식을 (주의 맵 등을 통해) 명시적으로 이전하는 것에서 학생이 자신의 시각 기반 추론 과정을 개발하도록 유도하는 것으로의 전환에 있습니다. 가장 명백한 텍스트 단서를 "숨김"으로써, 이 방법은 학생이 문제를 해결하기 위해 이미지를 "보게" 만듭니다. 저자들은 이 접근법이 내부 시각 패턴 모방에만 초점을 맞춘 이전 방법들보다 VLM 추론 증류에 더 효과적이라고 결론지으며, 이는 긴 체인 추론 작업에서 시각 망각의 근본 원인을 직접적으로 해결하기 때문입니다.