💬 NLP

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

This paper introduces Byte-Prefix Marginalization (BPM), a novel cross-tokenizer on-policy distillation method that re-expresses teacher distributions over a shared byte space to preserve probability mass and improve student model performance across mathematics and programming benchmarks.

Original authors: Hao Wang, Kun Yuan, Wenlin Zhong, Minglei Zhang, Han Xiao, Ming Sun, Honggang Qi

Published 2026-07-27

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Hao Wang, Kun Yuan, Wenlin Zhong, Minglei Zhang, Han Xiao, Ming Sun, Honggang Qi

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine a world where giant, brilliant AI brains are built by different architects, each using their own unique language to think. One architect might speak in long, flowing sentences, while another chops thoughts into tiny, rapid-fire syllables. This is the reality of modern "Large Language Models" (LLMs). They are incredibly smart, but they often speak different "token" languages. A token is just a chunk of text—like a word, a part of a word, or even a single letter—that the computer uses to process information.

Now, imagine you want to build a smaller, faster, and cheaper AI student that can learn from these giant teachers. This is called "distillation." Usually, you'd just have the student copy the teacher's answers. But here's the catch: if the teacher and student speak different token languages, a direct copy-paste is impossible. It's like trying to teach a student who only speaks English by handing them a textbook written entirely in a code of emojis and symbols. The student sees the symbols, but they don't know which symbol matches which English word. If you try to force a match, you might accidentally teach the student that a picture of a cat means "apple," or worse, you might throw away half the teacher's wisdom because it doesn't fit the student's dictionary.

This paper tackles that exact headache. The researchers wanted to know: Can we teach a student AI to learn from a teacher that speaks a completely different "token" language, without losing any of the teacher's brilliance or accidentally teaching it nonsense?

The Solution: The "Byte-Prefix" Translator

The authors, a team from KwaiKAT and the University of Chinese Academy of Sciences, introduced a clever new method called Byte-Prefix Marginalization (BPM). To understand how it works, let's ditch the complex math for a moment and think about how computers actually see text.

Even though different AIs chop text into different-sized "tokens" (chunks), they all agree on the underlying bytes. Bytes are the tiny, invisible building blocks of text—the raw 0s and 1s that make up every letter, space, and punctuation mark. Whether an AI sees the word "unbelievable" as one giant token, or breaks it into "un," "bel," and "ievable," they all agree that the bytes for "un" come first, followed by the bytes for "bel," and so on.

The authors realized that instead of trying to match the messy, different-sized tokens, they could translate the teacher's wisdom into this shared "byte language" first. Here is how BPM works, step-by-step:

The Byte Map: Imagine the teacher's next prediction is a cloud of probability. The teacher says, "There's a 30% chance the next word is 'unbelievable'."
The Longest Match: The student's tokenizer looks at the bytes for "unbelievable" and asks, "What is the longest chunk of my dictionary that fits at the start of these bytes?" Maybe the student has a token for "un" and another for "belie."
The Transfer: BPM takes the teacher's 30% probability for "unbelievable" and hands it over to the student's token "un" (or "belie," whichever is the longest match). It's like saying, "If the teacher thinks 'unbelievable' is likely, then the student should definitely think 'un' is likely, because 'un' is the first part of that word."
The Safety Net: What if the teacher says something the student's dictionary can't even start with? BPM doesn't just throw that probability away. It puts it into a special "residual" bucket, ensuring that 100% of the teacher's probability mass is preserved. Nothing is lost.

This method creates a perfect, "mass-preserving" target. It guarantees that the student learns exactly what the teacher intended, even if they speak different languages, as long as they agree on the bytes.

The "Whitespace" Trap and the Fix

However, the researchers discovered a funny, tricky problem. When the teacher and student are writing code, sometimes the next "token" is just a bunch of spaces or tabs (indentation). Because different tokenizers chop up spaces differently, the teacher might be saying, "I'm adding three spaces," while the student's tokenizer sees that as "I'm adding one big space token."

If the student tries to copy the teacher's exact probability for these space tokens, it gets confused. It starts learning the teacher's specific "chopping style" for spaces rather than the actual logic of the code. The authors found that in some cases, this confusion caused the student's code to collapse completely—its ability to write working programs dropped from 49% to a disastrous 9%.

To fix this, they added a simple "whitespace mask." It's like a referee that says, "Hey, if the next step is just spaces, don't try to copy the teacher's exact space-token probability. Just ignore that part of the lesson." This tiny rule saved the day, preventing the collapse and letting the student focus on the actual code logic.

The Results: Smarter Students, Faster

The team tested this method by taking three different giant AI teachers (Qwen3-32B, GLM-Z1-9B, and MiniMax-M2.7) and trying to teach them to a single, smaller student model (Qwen3.5-2B). These teachers had very different token languages, making the task quite hard.

They compared their new BPM method against other existing ways of handling different tokenizers. The results were clear:

Better Scores: On six tough benchmarks for math and coding, the BPM-trained students consistently outperformed the other methods. They improved the average score by 3.7 to 6.6 points over the strongest previous attempts.
Closing the Gap: The method managed to close 33.5% to 43.9% of the performance gap between the tiny student and the giant teacher. That's a huge chunk of the teacher's brainpower transferred to the student.
Robustness: The method worked well even when the teacher and student were very different (like the MiniMax teacher), proving that the "byte-prefix" approach is a universal translator.

Why This Matters

This paper suggests that we don't need to force all AI models to use the same dictionary to learn from each other. By using the shared "byte" language as a bridge, we can mix and match the best teachers from different families and teach them to efficient students without losing any knowledge. It's a bit like realizing that while everyone speaks different dialects, they all agree on the alphabet. If you teach the student to read the alphabet first, you can teach them anything, no matter who the teacher is.

The authors showed that this approach is not just a theoretical idea but a practical tool that works today, helping to create smaller, faster, and smarter AI models that can solve complex math and coding problems without needing the massive computing power of their giant teachers.

Technical Summary: Cross-Tokenizer On-Policy Distillation via Byte-Prefix Marginalization

1. Problem Statement

Open-weight language models from different families exhibit complementary strengths in reasoning, coding, and tool use. Consolidating these capabilities into a single compact student model via On-Policy Distillation (OPD) is desirable for efficiency and deployment simplicity. However, standard dense OPD assumes a shared tokenizer between teacher and student. When tokenizers differ (heterogeneous tokenizers), existing cross-tokenizer distillation methods fail to satisfy three critical requirements for full-vocabulary supervision:

Vocabulary Completeness (D1): The target distribution must be well-defined for every token in the student vocabulary.
Byte-level Alignment (D2): Teacher probability mass must map to student tokens based on the correspondence of their underlying byte representations, not just token IDs or probability ranks.
Mass Preservation (D3): The transformation must preserve the teacher's total probability mass (summing to one), including any residual mass that cannot be mapped to a specific student token.

Existing methods violate these properties: candidate-set approaches discard unmatched mass (violating D1), rank-matching methods assign mass to tokens with unrelated byte content (violating D2), and sequence-level distillation avoids full-vocabulary distribution matching entirely.

2. Methodology: Byte-Prefix Marginalization (BPM)

The authors propose Byte-Prefix Marginalization (BPM), which re-expresses the teacher's next-token distribution over the student vocabulary within a shared byte space. Since different tokenizers decode to the same byte sequence, BPM leverages this shared representation to construct an exact target.

Core Mechanism

BPM operates on the principle that a teacher token's byte sequence can be greedily segmented by the student tokenizer. The method proceeds as follows:

Mapping: For each teacher token $w$ , the algorithm identifies the longest student token $u$ such that the byte representation of $u$ is a prefix of the byte representation of $w$ .
Aggregation: Probability mass from all teacher tokens mapping to the same student token is aggregated.
Residual Handling: Mass from teacher tokens that do not contain a complete student token prefix is assigned to an explicit residual category ( $\emptyset$ ). This ensures mass preservation (D3) and vocabulary completeness (D1).
Spanning Correction: In cases where a single student token spans multiple teacher tokens (e.g., a student token covering a sequence of teacher tokens), BPM uses a chain-factorized lower bound derived from the teacher's realized segmentation path to assign mass, preventing the student from learning to avoid closing reasoning spans.

Training Objective

The student is trained on its own generated trajectories by minimizing a generalized Jensen–Shannon divergence ( $D_\beta$ ) between the student's distribution and the transformed target.

Stop Token Handling: A "stop bridge" mechanism is introduced to handle termination signals, ensuring the student learns the teacher's stopping behavior even if the realized stop token differs from the declared EOS.
Whitespace Collapse Mitigation: The authors identify a failure mode where byte-exact supervision on all-whitespace rows (common in code generation) transfers tokenizer-specific segmentation noise rather than semantic content, leading to indentation errors and code collapse. To address this, BPM employs a precomputed mask that excludes all-whitespace target positions from the loss function.

3. Key Contributions

Exact Full-Vocabulary Supervision: BPM is the first method to provide an exact, mass-preserving, and byte-aligned target for full-vocabulary OPD across heterogeneous tokenizers. The authors prove that the target exactly recovers the teacher-induced byte-prefix marginal when the byte prefix does not span multiple teacher tokens—a condition satisfied at >99% of training positions.
Systematic Evaluation: The paper presents a rigorous evaluation across three distinct teacher-student pairs (Qwen3-32B, GLM-Z1-9B, and MiniMax-M2.7 as teachers; Qwen3.5-2B as the student), covering varying degrees of tokenizer distance.
Identification and Fix of Whitespace Collapse: The work uncovers a specific failure mode in byte-exact distillation for code generation (where indentation is corrupted by tokenizer noise) and provides a simple, effective masking solution.

4. Experimental Results

Experiments were conducted on six benchmarks: three mathematics (AIME 2026, HMMT 2026, MATH-500) and three code (HumanEval+, LiveCodeBench, TACO).

Performance Gains: BPM consistently outperformed current cross-tokenizer baselines (SimCT, ULD, GOLD) and the separately trained SeqKD reference.
- Across all three teacher-student pairs, BPM improved the six-benchmark average @8 (avg@8) by 3.7 to 6.6 points over the strongest baselines.
- BPM closed 33.5% to 43.9% of the performance gap between the student and the teacher, compared to 19.9–26.6% for the strongest baselines.
Robustness: The improvements were most pronounced on contest benchmarks (AIME, HMMT) which postdate the teachers' training cutoffs.
Distributional Closeness: While some baselines appeared "closer" to the teacher in terms of predictability (due to collapse into repetition), BPM demonstrated superior fidelity in preserving the student's ability to generate teacher-like text without degenerating.
Code Generation: Without the whitespace mask, the unmasked arm suffered a catastrophic drop in code pass rates (e.g., from 49% to 9% in one pair) due to indentation errors. The masked arm maintained high performance and zero block-level indentation errors.

5. Significance and Claims

The paper claims that BPM establishes a new standard for cross-tokenizer distillation by solving the fundamental misalignment problem through byte-level marginalization.

Theoretical Soundness: The method provides a mathematically exact solution for the vast majority of training positions, preserving the teacher's probability mass without discarding information or misaligning content.
Practical Impact: By enabling the consolidation of diverse, high-performing models into a single efficient student without requiring shared tokenizers, BPM facilitates more flexible and capable model deployment.
Limitations: The authors modestly note that BPM discards gradient signals for positions that do not reach a sync point (less than 1% of positions) and that the whitespace mask, while effective, withholds supervision on ~3.5% of aligned code rows. They also acknowledge that data quality bounds the extent of performance gains achievable.

In summary, the paper argues that fidelity in distillation must be applied carefully: byte-exact alignment is essential for content, but must be withheld for whitespace-only regions to prevent the transfer of tokenizer-specific artifacts.