Command Palette
Search for a command to run...
크로스 토크나이저 온폴리시 증류에 대한 재고찰: 정렬 커버리지에서 감독 신뢰성으로
크로스 토크나이저 온폴리시 증류에 대한 재고찰: 정렬 커버리지에서 감독 신뢰성으로
Bingxi Hou Guochao Jiang Guofeng Quan Weiqing Li Wenfeng Feng Guohua Liu Yuewei Zhang
초록
온폴리시 증류(OPD)는 교사 피드백을 활용해 학생이 스스로 생성한 결과에 대해 학습하도록 만든다. 서로 다른 토크나이저를 사용할 경우 교사와 학생의 예측을 비교하려면 시퀀스 수준과 어휘 수준 모두에서 정렬이 필요하다. 본 논문에서는 이러한 정렬 커버리지를 확장하는 것이 학습 향상으로 이어지는지 조사한다. 수학적 추론 및 코드 생성에 관한 세 개의 이종 교사-학생 쌍에서, 상당한 어휘 불일치에도 불구하고 엄격한 1:1 그룹이 이미 대부분의 학생 생성 토큰을 포괄한다. 증류 전에 학생 모델들로부터 샘플링된 응답들에서 공유 어휘는 엄격하게 정렬된 위치에서 평균적으로 교사와 학생 확률 질량의 거의 전부를 보존한다. 각 엄격 위치에서 공유 어휘 중 학생이 선택한 상위 16개 부분집합으로 역 KL(reverse KL)을 제한하면 전체 공유 어휘 OPD에 필적하는 정확도를 달성하고 평가된 크로스 토크나이저 기준선들을 능가한다. 불일치 그룹의 스팬(span) 로그 확률에 평균 제곱 오차 감독을 추가하면 완전한 감독 커버리지를 얻지만 정확도는 오히려 감소한다. 엄격 손실만으로 학습한 체크포인트에서 스팬 그래디언트는 엄격 그래디언트와 약하거나 음의 방향 일치를 보이며 이에 비해 크기가 커진다. 이러한 진단은 스팬 감독 추가로 인한 정확도 하락을 설명하는 데 도움이 될 수 있다. 우리의 발견은 정렬 커버리지 극대화에서 감독 신뢰성 우선으로의 전환을 뒷받침한다. 엄격 위치에서의 압축된 감독은 약하게 정렬되거나 충돌하는 학습 신호를 도입하는 더 넓은 커버리지보다 더 효과적일 수 있다.
One-sentence Summary
Researchers at Alibaba Cloud Computing show that in cross-tokenizer on-policy distillation, strict 1:1 aligned groups already cover most student-generated tokens, and restricting reverse KL to a student-selected top-16 shared-vocabulary subset at strict positions matches full shared-vocabulary OPD while outperforming baselines, whereas adding span MSE supervision degrades accuracy and motivates a shift from maximizing alignment coverage to prioritizing supervision reliability.
Key Contributions
- The paper analyzes alignment coverage in cross-tokenizer on-policy distillation across mathematical reasoning and code generation tasks, showing that strict 1:1 aligned groups cover most student-generated tokens and retain nearly all teacher and student probability mass at strict positions despite vocabulary mismatch.
- It demonstrates that restricting reverse KL distillation to a student-selected top-16 subset of the shared vocabulary at strict positions achieves accuracy comparable to full shared-vocabulary distillation while outperforming the evaluated cross-tokenizer baselines.
- The work shows that adding span log-probability MSE supervision on mismatch groups reduces downstream accuracy, and its gradient diagnostics reveal weak or negative agreement with strict gradients along with growing relative magnitude, motivating a shift from maximizing alignment coverage to prioritizing supervision reliability.
Introduction
On-policy distillation trains a language model on its own generated responses using teacher feedback, which matters because it aligns the student with teacher preferences at the states the student actually visits. When the student and teacher use different tokenizers, the same response can have different token boundaries and next-token distributions over different vocabularies, so cross-tokenizer distillation must handle both sequence-level and vocabulary-level alignment. Prior methods such as rank matching, learned mappings, likelihood matching, byte-level outputs, and multi-token grouping have aimed to recover more supervision, but they generally emphasize alignment coverage without fully testing whether the added supervision is useful for learning. The authors examine strict cross-tokenizer distillation and find that strict alignment already covers most student-generated tokens despite large static vocabulary gaps, and that a compact student-selected top-k subset of the shared vocabulary retains nearly all distillation gains. By contrast, adding span log-probability MSE for unmatched spans achieves complete coverage but reduces accuracy, motivating a shift from alignment coverage to supervision reliability.
Method
On-Policy Distillation (OPD) trains a student language model by sampling trajectories from the student itself and aligning its distribution to that of a teacher model on the prefixes the student actually visits. This process can be viewed as a dense KL-constrained reinforcement learning setup where the teacher distribution induces a token-level reward. With a shared tokenizer, OPD minimizes the following objective:
LOPD(θ)=Ex∼D,y∼πθ(⋅∣x)[i=1∑LKL(πθ(⋅∣x,y<i)∥πT(⋅∣x,y<i))].When extending this approach across model families with different tokenizers, the same text can be assigned different token boundaries and vocabulary entries. Cross-Tokenizer OPD must therefore account for alignment at both the sequence and vocabulary levels. The authors leverage token-group alignment to score and align the decoded student response using both models' tokenizers. The resulting student and teacher response-token sequences are denoted by y=(y1,…,yL) and v=(v1,…,vn). By retaining the token offsets shared by both sequences, the response is partitioned into pairs of aligned token groups:
S(y)={(Srθ,SrT)}r=1R.These groups are divided into strict 1:1 groups, where a single student token and a single teacher token span the same interval, and mismatch groups, which require multiple tokens on at least one side to construct the same span.
To align the predictions, the authors identify a shared vocabulary V∩=Vθ∩VT by matching underlying tokens. For strict 1:1 groups, the student and teacher distributions are restricted and renormalized over this shared vocabulary. The strict Cross-Tokenizer objective then sums the reverse KL divergence over strictly aligned positions:
L1:1(θ)=Ex∼D,y∼πθ(⋅∣x)r∈A1:1(y)∑KL(πˉθ(⋅∣x,y<ir)∥πˉT(⋅∣x,v<jr)).As shown in the figure below, diagnostic evaluations of this cross-tokenizer approach reveal that strict alignment remains high during training, shared vocabulary captures nearly all probability mass in both models, and excluding mismatch groups yields the best performance.
To evaluate the learning value of the remaining mismatch groups, the authors introduce span supervision. For each mismatch group, the probabilities of the observed token paths are computed for both the student and the teacher. These probabilities are matched using a mean squared error loss over the mismatch groups in a log-probability formulation:
Lspan(θ)=Ex∼D,y∼πθ(⋅∣x)r∈Amis(y)∑(logqθ(r)−logqT(r))2.The total training loss combines the strict objective and the span supervision with a weighting factor λ:
Lλ(θ)=L1:1(θ)+λLspan(θ).As shown in the figure below, the impact of mismatch weights on full-average accuracy across different teacher-student pairs demonstrates that strict supervision alone performs best.
By varying λ, the authors control the influence of the span loss while maintaining complete supervision coverage. However, empirical results indicate that adding span supervision for mismatch groups tends to reduce downstream accuracy across tested weights, confirming that strict matching retains the most useful supervision.
Experiment
The experiments study cross-tokenizer knowledge distillation on three teacher-student pairs using a shared pool of mathematics and code prompts, with downstream evaluation on math and code benchmarks. They validate that strict token alignment covers most positions on student-generated trajectories despite large static vocabulary mismatch, while adding span supervision for remaining mismatch groups consistently hurts accuracy. The shared vocabulary retains nearly all predictive probability mass, and distillation on a small student-selected top-k subset keeps most of the strict-supervision gains while outperforming four baselines. Gradient diagnostics show that span-loss gradients align poorly with strict-loss gradients and grow in relative magnitude during training, helping explain the negative effect of complete coverage.
Strict token coverage on student trajectories stayed high across the evaluated model pairs even when static vocabulary overlap differed substantially. The pair with the lowest vocabulary Jaccard overlap also had the highest strict coverage under both student and teacher tokenizations. Coverage remained stable across training windows, indicating that large static vocabulary mismatch can coexist with strict alignment at most positions. Strict student and teacher token coverage stayed high across all model pairs despite static vocabulary overlap ranging from about 39% to about 65%. Granite to Phi showed the lowest static vocabulary Jaccard overlap but the highest strict coverage under both tokenizations. Within each pair, strict coverage varied little across training windows, changing by at most 1.43 percentage points for student tokens and 3.75 for teacher tokens.
Strict full-vocabulary distillation consistently improves math and full-average accuracy over baseline alternatives and the base model across all three cross-tokenizer pairs, with the largest full-average gain for Granite-to-Qwen. Code accuracy is more mixed, but strict full remains at or near the best. A compact student-selected subset preserves nearly all of the strict full learning benefit. Strict full supervision leads over all baseline alternatives in math and full-average accuracy across every teacher-student pair. Student-selected top-k subsets retain almost all of the strict full distillation gain, with no consistent further benefit from a larger subset.
The evaluation examines cross-tokenizer distillation across several teacher-student pairs, measuring token coverage and downstream accuracy. Strict token coverage remains high across pairs even when static vocabulary overlap differs substantially, indicating that vocabulary mismatch does not prevent close sequence alignment. Strict full-vocabulary distillation consistently improves math and full-average accuracy over baselines, with the largest full-average gain for Granite-to-Qwen, while code accuracy is more mixed. A compact student-selected subset preserves nearly all of this strict distillation benefit, with no consistent advantage from using a larger subset.