Technical Summary: X-Token: Projection-Guided Cross-Tokenizer Knowledge Distillation

Problem Statement

Knowledge Distillation (KD) typically requires a shared tokenizer between teacher and student models, effectively locking practitioners into specific model families. This restriction prevents leveraging stronger or specialized teachers with incompatible vocabularies (e.g., distilling from a Qwen or Phi model into a Llama student) and hinders multi-teacher distillation across diverse tokenizer families.

Existing cross-tokenizer KD methods fall into two categories: representation-based (aligning hidden states) and logit-distribution-based (operating on output probabilities). The latter is preferred for continual pretraining as it requires no architectural changes. However, the current state-of-the-art logit-based method, GOLD, suffers from two critical structural limitations when handling incompatible tokenizers:

1. Uncommon-Token Failure: When tokenizers fragment text differently (e.g., a teacher splitting multi-digit numerals while the student packs them), critical tokens fall into an "unmatched" subset. These tokens are degraded by identity-agnostic noise from rank-based matching and suppressed by gradients from the common-KL term, leading to severe performance drops (e.g., GSM8k accuracy falling from 12.89 to 2.56).
2. Over-Conservative Matching: Strict 1-to-1 string equality excludes valid semantic alignments where surface forms differ (e.g., a student token "Hundreds" corresponding to teacher tokens "Hund" + "reds"), leaving useful supervision signals unexploited.

Methodology: X-Token

X-Token addresses these issues through a three-component framework: Span Alignment, a Projection Matrix ($W$), and Complementary Loss Formulations.

1. Sequence Alignment

To handle differing sequence lengths and positional misalignment, X-Token employs span alignment. Tokens are grouped into chunks that decode to the same underlying text. A chain-rule merge is applied within each chunk to combine per-token probabilities into a single chunk-level distribution, creating aligned units for distillation.

2. Projection Matrix ($W$)

A sparse projection matrix $W \in \mathbb{R}^{|\mathcal{V}_S| \times |\mathcal{V}_T|}$ maps student token probabilities into the teacher's vocabulary space. It is constructed deterministically in two passes:

• Exact-Match Pass: Sets $W[s, t] = 1$ for tokens with canonicalized string matches.
• Multi-Token-Rule Pass: For unmatched student tokens, the decoded text is re-tokenized by the teacher. The student token is mapped to the resulting sequence of teacher tokens with decaying weights (e.g., $\beta=0.9, \gamma=0.1$).
  The matrix is initialized from tokenizer rules and can be optionally fine-tuned during training.

3. Complementary Loss Formulations

X-Token introduces two loss modes, selected based on a coverage audit of critical tokens (e.g., numerals) in the student vocabulary:

• P-KL (Projection-KL): Designed for scenarios where critical tokens fall into the "uncommon" set (e.g., Qwen teacher). It removes the partition entirely. The student distribution is projected via $W$ into the teacher vocabulary space, and a direct KL divergence is computed between the teacher distribution and the projected student distribution. This eliminates erroneous rank-based gradients and suppressive effects.
• H-KL (Hybrid-KL): Designed for scenarios where the partition is structurally sound (e.g., Phi-4-mini teacher). It retains the hybrid structure but relaxes the definition of the "common" set. Instead of strict string equality, it includes pairs $(s, t^*)$ where $t^*$ is the top-ranked teacher token for student token $s$ under $W$. This captures near-equivalent alignments while preserving the sharper supervision of direct KL on identity-aligned pairs.

4. Multi-Teacher Distillation

X-Token naturally extends to multiple teachers by aggregating per-teacher losses. It employs static weighting (which outperformed adaptive schemes in experiments) to combine signals from heterogeneous teachers, leveraging their complementary strengths (e.g., reasoning vs. commonsense).

5. Dynamic Scaling

To stabilize training, the distillation loss is dynamically rescaled at each step to match the magnitude of the cross-entropy loss, preventing optimization instability caused by varying loss scales.

Key Contributions

1. Deterministic Cross-Token Alignment: Introduction of a sparse projection matrix $W$ constructed via string rules and multi-token decoding, enabling direct vocabulary alignment without auxiliary trainable components.
2. Dual Loss Formulations: The proposal of P-KL and H-KL, which adapt to different tokenizer mismatch regimes. P-KL recovers signal for critical unmatched tokens, while H-KL refines matching for well-aligned tokens.
3. Multi-Teacher Capability: Demonstration that X-Token enables effective distillation from heterogeneous teachers, with results showing that complementary teacher combinations yield superior performance over single-teacher or homogeneous pairings.
4. Robust Sequence Alignment: A scalable, DP-based span alignment method that ensures text-consistent units for distillation.

Experimental Results

Experiments were conducted using Llama-3.2-1B as the student and Qwen3-4B, Phi-4-mini-Instruct, and Llama-3.2-3B as teachers.

• Performance vs. SOTA: X-Token outperforms the current state-of-the-art (GOLD) by +3.82 average points when using a Qwen3-4B teacher (with a 6x improvement on GSM8k, from 2.56 to 15.54) and by +0.5 points with a Phi-4-mini teacher.
• Cross-Family Superiority: In the Qwen regime, X-Token (P-KL) surpassed same-tokenizer KD from a weaker Llama-3.2-3B teacher, demonstrating that cross-tokenizer distillation can exceed same-family limits when the loss formulation is correct.
• Multi-Teacher Gains: A two-teacher setup (Phi-4-mini + Llama-3.2-3B) improved over single-teacher distillation by +1.3 points, highlighting the value of combining diverse teacher strengths.
• Ablation Findings:
  • Loss selection is critical: P-KL fails on Phi-4-mini, and H-KL fails on Qwen3-4B.
  • Jointly learning the projection matrix $W$ yields modest but consistent gains over a frozen, rule-based initialization.
  • Static teacher weighting outperformed adaptive schemes (based on entropy or confidence).

Significance

The paper claims that X-Token demonstrates that careful alignment at both the sequence and vocabulary levels, combined with adaptive loss design, is key to unlocking the full potential of cross-tokenizer knowledge distillation. By addressing the specific failure modes of partition-based methods, X-Token allows practitioners to leverage stronger, specialized, or diverse teachers without being constrained by tokenizer compatibility, effectively breaking the "tokenizer lock-in" that currently limits model scaling and specialization.