Cross-Tokenizer On-Policy Distillation via Byte-Prefix Marginalization

2026-07-24

AI Summary

This work addresses the challenge of misaligned output vocabularies in cross-tokenizer knowledge distillation, where the teacher model’s next-token probability distribution cannot be accurately mapped onto the student’s vocabulary, leading to significant information loss. To resolve this, the authors propose Byte-level Prefix Marginalization (BPM), a method that performs a lossless transformation of the teacher’s distribution into the student’s vocabulary space by operating in a shared byte-level representation. BPM achieves precise alignment through byte-prefix matching and explicitly models residual categories for uncovered tokens, recovering the teacher-induced marginal distribution at over 99% of positions. Integrated with online policy distillation and multi-teacher fusion, BPM yields consistent improvements, outperforming the strongest baseline by 3.7–6.6 points in average pass@8 scores across six mathematical and programming benchmarks.

Abstract

Open-weight language models from different families exhibit complementary capabilities, motivating their consolidation into a compact student through on-policy distillation (OPD). However, full-vocabulary OPD typically assumes a shared tokenizer, while existing cross-tokenizer methods may discard teacher probability mass or assign it to student tokens with unrelated content. We introduce Byte-Prefix Marginalization (BPM), which re-expresses the teacher's next-token distribution over the student vocabulary in a shared byte space. Specifically, BPM assigns each teacher token's probability to the longest student token whose byte representation is a prefix of the teacher token's bytes, aggregates mass mapped to the same student token, and places otherwise unmatched mass in an explicit residual category. This produces a vocabulary-complete, byte-aligned, and mass-preserving target for dense OPD. The target exactly recovers the teacher-induced byte-prefix marginal when the relevant prefix does not span multiple teacher tokens (a condition satisfied at more than 99% of training positions) and uses a mass-preserving, chain-factorized lower bound otherwise. Across Qwen3-32B, GLM-Z1-9B-0414, and MiniMax-M2.7 as teachers, BPM consistently outperforms current cross-tokenizer methods on six mathematics and programming benchmarks, improving six-benchmark avg@8 by 3.7-6.6 points over the strongest baselines.

Problem

• cross-tokenizer
• on-policy distillation
• probability mass preservation
• byte-prefix alignment
• language model distillation

Innovation

• Byte-Prefix Marginalization
• On-Policy Distillation
• Cross-Tokenizer Alignment
• Probability Mass Preservation
• Language Model Distillation