기술 요약: 바이트 접두사 한계화(Byte-Prefix Marginalization)를 통한 교차 토크나이저 온폴리시 증류

1. 문제 정의

서로 다른 계열의 오픈 웨이트 언어 모델들은 추론, 코딩, 도구 사용 측면에서 상호 보완적인 강점을 보입니다. 이러한 역량을 하나의 컴팩트한 학생 모델로 통합하는 것은 효율성과 배포 단순성을 위해 바람직합니다. 이를 위해 **온폴리시 증류(On-Policy Distillation, OPD)**가 유용하지만, 표준적인 밀집(dense) OPD는 교사(teacher)와 학생(student) 모델 간의 공유된 토크나이저를 가정합니다. 토크나이저가 서로 다를 경우(이종 토크나이저), 기존의 교차 토크나이저 증류 방식은 전체 어휘에 대한 감독을 위한 세 가지 핵심 요구사항을 충족하지 못합니다:

1. 어휘 완결성 (D1): 대상 분포가 학생의 어휘 내 모든 토큰에 대해 잘 정의되어야 합니다.
2. 바이트 수준 정렬 (D2): 교사의 확률 질량(probability mass)이 단순히 토큰 ID나 확률 순위가 아니라, 그 기저에 있는 바이트 표현의 대응 관계에 따라 학생 토큰으로 매핑되어야 합니다.
3. 질량 보존 (D3): 변환 과정은 특정 학생 토큰으로 매핑될 수 없는 잔여 질량을 포함하여, 교사의 총 확률 질량(합계가 1)을 보존해야 합니다.

기존 방식들은 이러한 속성을 위반합니다: 후보 집합(candidate-set) 접근법은 일치하지 않는 질량을 버리고(D1 위반), 순위 매칭(rank-matching) 방식은 관련 없는 바이트 콘텐츠를 가진 토큰에 질량을 할당하며(D2 위반), 시퀀스 수준 증류는 전체 어휘 분포 매칭 자체를 회피합니다.

2. 방법론: 바이트 접두사 한계화 (Byte-Prefix Marginalization, BPM)

저자들은 교사의 다음 토큰 분포를 공유된 바이트 공간 내에서 학생의 어휘로 재표현하는 **바이트 접두사 한계화(BPM)**를 제안합니다. 서로 다른 토크나이저라도 동일한 바이트 시퀀스로 디코딩된다는 점을 이용하여, BPM은 이 공유된 표현을 활용해 정확한 타겟을 구축합니다.

핵심 메커즘

BPM은 교사 토큰의 바이트 시퀀스가 학생 토크나이저에 의해 탐욕적(greedily)으로 분할될 수 있다는 원리에 기반합니다. 이 방법은 다음과 같이 진행됩니다:

1. 매핑 (Mapping): 각 교사 토큰 $w$에 대해, $w$의 바이트 표현이 $u$의 바이트 표현의 접두사(prefix)가 되는 가장 긴 학생 토큰 $u$를 식별합니다.
2. 집계 (Aggregation): 동일한 학생 토큰으로 매핑되는 모든 교사 토큰으로부터의 확률 질량을 집계합니다.
3. 잔여 처리 (Residual Handling): 완전한 학생 토큰 접두사를 포함하지 않는 교사 토큰의 질량은 명시적인 **잔여 카테고리($\emptyset$)**에 할당됩니다. 이는 질량 보존(D3)과 어휘 완결성(D1)을 보장합니다.
4. 스패닝 수정 (Spanning Correction): 단일 학생 토큰이 여러 교사 토큰을 가로지르는 경우(예: 교사 토큰의 시퀀스를 덮는 학생 토큰), BPM은 교사의 실현된 분할 경로(realized segmentation path)로부터 유도된 체인 인자화 하한(chain-factorized lower bound)을 사용하여 질량을 할당함으로써, 학생이 추론 구간을 닫는 것을 피하도록 학습하는 것을 방지합니다.

훈련 목적 함수

학생은 자신의 생성된 궤적(trajectories) 상에서 학생의 분포와 변환된 타겟 사이의 **일반화된 젠슨-샤논 발산(generalized Jensen–Shannon divergence, $D_\beta$)**을 최소화하도록 훈련됩니다.

• 종료 토큰 처리 (Stop Token Handling): 종료 신호를 처리하기 위해 "스톱 브릿지(stop bridge)" 메커니즘이 도입되었습니다. 이는 실현된 종료 토큰이 선언된 EOS와 다르더라도 학생이 교사의 중단 동작을 학습할 수 있도록 보장합니다.
• 공백 붕괴 완화 (Whitespace Collapse Mitigation): 저자들은 코드 생성 시 흔히 발생하는 모든 공백 행(all-whitespace rows)에 대한 바이트 정밀 감독이 의미적 내용 대신 토크나이저 특유의 분할 노이즈를 전달하여 들여쓰기 오류 및 코드 붕괴를 초래하는 실패 모드를 식별했습니다. 이를 해결하기 위해 BPM은 모든 공백 타겟 위치를 손실 함수에서 제외하는 **사전 계산된 마스크(precomputed mask)**를 사용합니다.

3. 주요 기여

1. 정확한 전체 어휘 감독: BPM은 이종 토크나이저 간의 OPD를 위해 정확하고 질량이 보존되며 바이트가 정렬된 타겟을 제공하는 최초의 방법입니다. 저자들은 바이트 접두사가 여러 교사 토큰을 가로지르지 않을 때 타겟이 교사에 의해 유도된 바이트 접두사 한계(byte-prefix marginal)를 정확히 복구함을 증명하였으며, 이 조건은 99% 이상의 훈련 위치에서 충족됩니다.
2. 체계적 평가: 본 논문은 세 가지 서로 다른 교사-학생 쌍(교사: Qwen3-32B, GLM-Z1-9B, MiniMax-M2.7 / 학생: Qwen3.5-2B)에 대해 엄격한 평가를 수행하였으며, 이는 토크나이저 간의 거리 차이를 포괄합니다.
3. 공백 붕괴의 식별 및 해결: 본 연구는 바이트 정밀 증류가 코드 생성 시(들여쓰기가 토크나이저 노이즈에 의해 손상되는 경우) 일으키는 특정 실패 모드를 밝혀내고, 간단하면서도 효과적인 마스킹 솔루션을 제공합니다.

4. 실험 결과

실험은 세 가지 수학 벤치마크(AIME 2026, HMMT 2026, MATH-500)와 세 가지 코딩 벤치마크(HumanEval+, LiveCodeBench, TACO)를 포함한 총 6개의 벤치마크에서 수행되었습니다.

• 성능 향상: BPM은 현재의 교차 토크나이저 베이스라인(SimCT, ULD, GOLD) 및 별도로 훈련된 SeqKD 참조 모델보다 일관되게 우수한 성능을 보였습니다.
  • 세 가지 교사-학생 쌍 모두에서, BPM은 가장 강력한 베이스라인 대비 **six-benchmark average @8 (avg@8)**을 3.7 ~ 6.6 포인트 개선했습니다.
  • BPM은 가장 강력한 베이스라인들이 보여준 19.9~26.6%의 격차와 비교하여, 학생과 교사 사이의 성능 격차를 **33.5% ~ 43.9%**까지 좁혔습니다.
• 강건성: 성능 향상은 교사의 훈련 컷오프 이후의 컨테스트 벤치마크(AIME, HMMT)에서 가장 두드러졌습니다.
• 분포적 근접성: 일부 베이스라인은 (반복으로 인한 붕괴 때문에) 예측 가능성 측면에서 교사와 더 "가까워" 보였으나, BPM은 퇴보 없이 교사 같은 텍스트를 생성하는 학생의 능력을 보존하는 데 있어 더 우수한 충실도(fidelity)를 입증했습니다.
• 코드 생성: 공백 마스크가 없는 경우, 마스크가 없는 버전은 들여쓰기 오류로 인해 코드 통과율이 급격히 하락했습니다(예: 한 쌍에서 49%에서 9%로 하락). 마스크를 적용한 버전은 높은 성능과 제로(0)의 블록 레벨 들여쓰기 오류를 유지했습니다.

5. 의의 및 주장

본 논문은 BPM이 바이트 수준의 한계화를 통해 근본적인 불일치 문제를 해결함으로써 교차 토크나이저 증류의 새로운 표준을 세운다고 주장합니다.

• 이론적 타당성: 이 방법은 대부분의 훈련 위치에서 정보를 버리거나 내용을 잘못 정렬하지 않고 교사의 확률 질량을 보존하는 수학적으로 정확한 솔루션을 제공합니다.
• 실질적 영향: 서로 다른 고성능 모델을 공유된 토크나이저 없이도 하나의 효율적인 학생 모델로 통합할 수 있게 함으로써, BPM은 더 유연하고 유능한 모델 배포를 용이하게 합니다.
• 한계점: 저자들은 BPM이 동기화 지점(sync point)에 도달하지 못하는 위치(1% 미만)에 대한 그래디언트 신호를 버린다는 점과, 공백 마스크가 효과적이긴 하지만 정렬된 코드 행의 약 3.5%에 대한 감독을 제한한다는 점을 겸허히 언급했습니다. 또한 데이터 품질이 달성 가능한 성능 향상의 범위를 제한한다는 점도 인정했습니다.

요약하자면, 본 논문은 **충실도(fidelity)**를 적용할 때 주의가 필요하다고 주장합니다: 바이트 정밀 정렬은 콘텐츠를 위해 필수적이지만, 토크나이저 특유의 아티팩트 전파를 방지하기 위해 공백 전용 영역에서는 감독을 유보해야 합니다.