기술 요약: X-Token: 투영 기반 교차 토크나이저 지식 증류

문제 제기

지식 증류 (KD) 는 일반적으로 교사 모델과 학생 모델 간에 공유된 토크나이저를 요구하여, 실무자를 특정 모델 계열에 묶어두는 결과를 초래합니다. 이러한 제한은 어휘가 호환되지 않는 더 강력하거나 전문화된 교사 모델을 활용하는 것을 방해하며 (예: Qwen 또는 Phi 모델에서 Llama 학생으로 증류), 다양한 토크나이저 계열 간 다중 교사 증류를 어렵게 만듭니다.

기존의 교차 토크나이저 KD 방법은 표현 기반 (은닉 상태 정렬) 과 로그 분포 기반 (출력 확률 연산) 으로 두 가지 범주로 나뉩니다. 후자는 구조 변경 없이 지속적인 사전 학습에 적합하여 선호됩니다. 그러나 현재 최첨단 로그 기반 방법인 GOLD는 호환되지 않는 토크나이저를 처리할 때 두 가지 치명적인 구조적 한계를 겪습니다:

1. 희귀 토큰 실패: 토크나이저가 텍스트를 다르게 분할할 때 (예: 교사가 다 자리 숫자를 분할하는 반면 학생은 이를 묶음), 중요한 토큰이 "불일치" 하위 집합에 속하게 됩니다. 이러한 토큰은 순위 기반 매칭에서 오는 정체성 무관한 노이즈로 인해 저하되고, 공통 KL 항의 그라디언트에 의해 억제되어 심각한 성능 저하를 초래합니다 (예: GSM8k 정확도가 12.89 에서 2.56 으로 하락).
2. 과도한 보수적 매칭: 엄격한 1 대 1 문자열 동일성은 표면 형태가 다른 유효한 의미 정렬을 배제합니다 (예: 학생 토큰 "Hundreds"가 교사 토큰 "Hund" + "reds"에 해당함), 이로 인해 유용한 지도 신호가 활용되지 못합니다.

방법론: X-Token

X-Token 은 스팬 정렬 (Span Alignment), 투영 행렬 ($W$), 그리고 보완적 손실 공식이라는 세 가지 구성 요소 프레임워크를 통해 이러한 문제를 해결합니다.

1. 시퀀스 정렬

서로 다른 시퀀스 길이와 위치 불일치를 처리하기 위해 X-Token 은 스팬 정렬을 사용합니다. 토큰은 동일한 기본 텍스트로 디코딩되는 청크로 그룹화됩니다. 각 청크 내에서 체인 규칙 병합을 적용하여 토큰별 확률을 단일 청크 수준 분포로 결합하여 증류를 위한 정렬된 단위를 생성합니다.

2. 투영 행렬 ($W$)

희소 투영 행렬 $W \in \mathbb{R}^{|\mathcal{V}_S| \times |\mathcal{V}_T|}$은 학생 토큰 확률을 교사의 어휘 공간으로 매핑합니다. 이는 두 번의 패스를 통해 결정론적으로 구성됩니다:

• 정확 일치 패스: 표준화된 문자열 일치를 가진 토큰에 대해 $W[s, t] = 1$을 설정합니다.
• 다중 토큰 규칙 패스: 일치하지 않는 학생 토큰의 경우, 디코딩된 텍스트를 교사에 의해 재토큰화합니다. 학생 토큰은 감쇠 가중치 (예: $\beta=0.9, \gamma=0.1$) 를 가진 결과적인 교사 토큰 시퀀스로 매핑됩니다.
  이 행렬은 토크나이저 규칙에서 초기화되며, 훈련 중에 선택적으로 미세 조정될 수 있습니다.

3. 보완적 손실 공식

X-Token 은 학생 어휘 내 중요한 토큰 (예: 숫자) 의 커버리지 감사를 기반으로 선택되는 두 가지 손실 모드를 도입합니다:

• P-KL (투영-KL): 중요한 토큰이 "희귀" 집합에 속하는 시나리오 (예: Qwen 교사) 에 설계되었습니다. 분할을 완전히 제거합니다. 학생 분포는 $W$를 통해 교사의 어휘 공간으로 투영되며, 교사 분포와 투영된 학생 분포 간에 직접 KL 발산이 계산됩니다. 이는 오류가 있는 순위 기반 그라디언트와 억제 효과를 제거합니다.
• H-KL (하이브리드-KL): 분할이 구조적으로 타당한 시나리오 (예: Phi-4-mini 교사) 에 설계되었습니다. 하이브리드 구조를 유지하지만 "공통" 집합의 정의를 완화합니다. 엄격한 문자열 동일성 대신, $W$ 하에서 학생 토큰 $s$에 대한 최상위 교사 토큰인 $t^*$인 쌍 $(s, t^*)$을 포함합니다. 이는 정체성 정렬된 쌍에 대한 직접 KL 의 더 날카로운 지도를 유지하면서 근사 동등 정렬을 포착합니다.

4. 다중 교사 증류

X-Token 은 교사별 손실을 집계함으로써 자연스럽게 여러 교사로 확장됩니다. 이질적인 교사들의 신호를 결합하기 위해 정적 가중치 (실험에서 적응형 방식보다 우수한 성능을 보임) 를 사용하여 (예: 추론 대 상식) 그들의 보완적 강점을 활용합니다.

5. 동적 스케일링

훈련을 안정화하기 위해 증류 손실은 각 단계에서 교차 엔트로피 손실의 크기와 일치하도록 동적으로 재조정되어, 변하는 손실 규모로 인한 최적화 불안정성을 방지합니다.

주요 기여

1. 결정론적 교차 토큰 정렬: 문자 규칙과 다중 토큰 디코딩을 통해 구성된 희소 투영 행렬 $W$의 도입으로, 보조 학습 가능 구성 요소 없이 직접적인 어휘 정렬을 가능하게 합니다.
2. 이중 손실 공식: 서로 다른 토크나이저 불일치 체제에 적응하는 P-KL과 H-KL 제안. P-KL 은 중요한 불일치 토큰에 대한 신호를 복구하는 반면, H-KL 은 잘 정렬된 토큰에 대한 매칭을 정제합니다.
3. 다중 교사 기능: X-Token 이 이질적인 교사로부터 효과적인 증류를 가능하게 함을 입증하며, 보완적인 교사 조합이 단일 교사 또는 동질적 페어링보다 우수한 성능을 보인다는 결과를 제시합니다.
4. 강건한 시퀀스 정렬: 증류를 위한 텍스트 일관성 단위를 보장하는 확장 가능한 DP 기반 스팬 정렬 방법.

실험 결과

실험은 학생으로 Llama-3.2-1B를, 교사로는 Qwen3-4B, Phi-4-mini-Instruct, Llama-3.2-3B를 사용하여 수행되었습니다.

• SOTA 대비 성능: Qwen3-4B 교사를 사용할 때 X-Token 은 현재 최첨단 (GOLD) 보다 +3.82 평균 점수만큼 우수하며 (GSM8k 에서 2.56 에서 15.54 로 6 배 개선), Phi-4-mini 교사와 함께 사용할 때 +0.5 점만큼 우수합니다.
• 교차 계열 우위: Qwen 체제에서 X-Token(P-KL) 은 더 약한 Llama-3.2-3B 교사로부터의 동일 토크나이저 KD 를 능가하여, 손실 공식이 정확할 때 교차 토크나이저 증류가 동일 계열의 한계를 초과할 수 있음을 입증했습니다.
• 다중 교사 이득: 두 교사 설정 (Phi-4-mini + Llama-3.2-3B) 은 단일 교사 증류보다 +1.3 점 향상되어, 다양한 교사 강점의 결합 가치를 강조합니다.
• 절대 실험 결과:
  • 손실 선택이 중요합니다: P-KL 은 Phi-4-mini 에서 실패하고, H-KL 은 Qwen3-4B 에서 실패합니다.
  • 투영 행렬 $W$를 공동 학습하는 것은 고정된 규칙 기반 초기화보다 소폭이지만 일관된 이득을 제공합니다.
  • 정적 교사 가중치는 적응형 방식 (엔트로피 또는 신뢰도 기반) 보다 우수한 성능을 보였습니다.

의의

이 논문은 X-Token 이 시퀀스 및 어휘 수준 모두에서의 세심한 정렬과 적응형 손실 설계의 결합이 교차 토크나이저 지식 증류의 잠재력을 완전히 unlocking 하는 열쇠임을 보여준다고 주장합니다. 분할 기반 방법의 특정 실패 모드를 해결함으로써, X-Token 은 실무자가 토크나이저 호환성에 구애받지 않고 더 강력하거나 전문화되거나 다양한 교사를 활용할 수 있게 하여, 현재 모델 확장 및 전문화를 제한하는 "토크나이저 락인"을 효과적으로 깨뜨립니다.