X-Token: 投影引导的跨分词器知识蒸馏

# X-Token: 投影引导的跨分词器知识蒸馏

来源: https://arxiv.org/html/2605.21699

\correspondingauthor sharatht@nvidia\.com

Adithyakrishna Venkatesh Hanasoge, Mingyu Yang, Ali Taghibakhshi, Saurav Muralidharan, Ashwath Aithal, Pavlo Molchanov

###### 摘要

**摘要:** 跨分词器知识蒸馏允许学生模型从词汇表不兼容的教师模型中学习。先前的工作在隐藏状态或 logits 上操作；后者更受欢迎，因为它可以作为无需辅助组件的即插即用替代方案。基于 logit 的方法要么只使用正确 token 的概率，错过教师分布中完整的“暗知识”，要么操作完整的输出分布，依赖于严格的 token 划分和/或缺乏原则的启发式排序。我们识别出基于全分布 logit 方法的两个关键缺陷：(i) **罕见 token 失败**，其中关键 token 落入不匹配的子集（例如，Llama 产生 1100 个多位数字，而 Qwen 以数字拆分方式监督），并在训练期间被抑制，与来自较弱教师的同分词器 KD 相比，GSM8k 从 12.89 降至 2.56；以及 (ii) **过度保守的匹配**，严格的 1 对 1 匹配排除了不同表面形式下的近似等价 token。这些失败需要不同的补救措施：当关键 token 未对齐时消除划分，当对齐可靠时改进划分。我们提出 **X-Token**，一种具有两种互补损失公式的方法，针对这些问题。**P-KL** 移除划分，通过稀疏投影矩阵 \(W\)（从分词器级字符串规则初始化）将学生分布与教师分布对齐，以解决罕见 token 失败。**H-KL** 保留混合形式，同时放宽匹配，在 \(W\) 下将每个学生 token 与其排名最高的教师映射对齐。两个目标共享 \(W\)，并自然扩展到多个教师。实验上，在 Llama-3.2-1B 上，X-Token 在 Qwen3-4B 教师下比当前最先进的方法 GOLD [patiño2025\_unlocking\_on\_policy\_distillation\_for\_any\_model\_family] 平均提高 +3.82 个点，在 Phi-4-Mini 教师下提高 +0.5 个点。此外，双教师设置（Phi-4-mini++Llama-3B）相比单教师蒸馏提高了 +1.3 个点。

## 1 引言

知识蒸馏 (KD) [hinton2015distilling, romero2014fitnets, furlanello2018born, park2019relational] 将教师输出分布中的“暗知识”传递给学生，通常通过每个位置的 Kullback–Leibler (KL) 散度作用于下一个 token 概率分布。这种公式需要共享的分词器，实际上将学生与同族教师绑定。因此，致力于给定分词器的实践者（例如，Llama-3.2-1B [grattafiori2024llama]）无法利用具有不兼容分词器的更强或更专业的教师（例如，Phi-4-mini [abouelenin2025phi], Qwen3-4B [yang2025qwen3]）。这一约束也阻止了跨分词器家族的多个教师蒸馏，限制了将具有互补优势（例如，推理、代码、多语言）的教师组合为统一训练信号的能力。跨分词器蒸馏移除了这一限制，既能从教师分词器锁定中解放出来，又能从多样化来源实现有效的多教师学习。

请参阅标题

图 1: **左:** 多教师蒸馏通过适当的损失路由每个教师——同分词器 Llama-3.2-3B 使用 KL，跨分词器 Qwen3-4B 和 Phi-4-mini 使用 P-KL/H-KL。**右:** X-Token 解决了 GOLD 字符串相等划分的两种失败模式，并在教师之间组合。右上方: 关键 token `201` 在 GOLD 中无匹配，接收到错误信号；P-KL 通过投影 \(W\) 将其连接到 \(V_T\) 中的 `{2,0,1}`。右下方: `Hundreds` 被排除在 GOLD 的 common-KL 项之外；H-KL 通过 \(W\) 的 top-1 映射接纳 (`Hundreds`, `Hund`)。

现有的跨分词器 KD 方法分为两大类别：**基于表示**的方法，在嵌入或隐藏状态级别对齐教师和学生（例如，DSKD [zhang2024dual]），以及**基于 logit 分布**的方法，直接操作输出分布（例如，ULD [boizard2024towards], GOLD [patiño2025\_unlocking\_on\_policy\_distillation\_for\_any\_model\_family]）。后一类方法在持续预训练规模下特别有吸引力，因为它们不需要额外的可训练组件，并且可以作为标准 KD 损失的即插即用替代方案集成，无需修改模型架构或引入额外的前向传播。GOLD [patiño2025\_unlocking\_on\_policy\_distillation\_for\_any\_model\_family] 方法应用了一种混合损失，将 token 划分为一个 1 对 1 字符串匹配的 **common** 子集（使用 KL 散度训练）和一个 **uncommon** 剩余部分（遵循 ULD [boizard2024towards] 通过基于排名的 L1 匹配）。然而，这种混合设计表现出两个结构性的局限性。首先，**罕见 token 失败**：当分词器以不同方式分割时（例如，Qwen3 拆分多位数字，而 Llama-3 将它们打包为单个 token），**关键 token**——其正确预测直接决定任务准确率的 token（例如，GSM8k 中的多位数字）——被强制进入不匹配的子集。这些 token 随后受到 (i) 来自基于排名匹配的与身份无关的噪声，以及 (ii) 来自 common-KL 项通过全词汇表 softmax 作用的抑制性梯度而退化。其次，**过度保守的匹配**：严格的字符串相等排除了跨分词器的等价 token 对，包括表面形式和教师多 token 分解——即使划分在其他方面良好形成（例如，学生 token `Hundreds` 对应教师 token `Hund` 后跟 `reds`，但严格的匹配因缺乏精确字符串相等而丢弃了这一对应关系）。

我们提出 **X-Token**（图 1 (https://arxiv.org/html/2605.21699#S1.F1)），它解决了当前方法的局限性并做出了以下贡献：
- **确定性跨分词器对齐**。我们引入一个稀疏投影矩阵 \(W\)，通过规范化的字符串匹配和多 token 解码规则构建，实现跨分词器的直接对齐；\(W\) 可以在 KD 期间可选地微调以获得额外收益。
- **互补的损失公式 (P-KL 和 H-KL) 和损失选择标准**。P-KL 移除划分并对齐完整分布，而 H-KL 通过 \(W\) 下的 top-排名映射放宽匹配。对 token 类别（例如，数字）的简单覆盖审计指导选择：当关键 token 落在 common 集合之外时使用 P-KL，否则使用 H-KL。P-KL 在 Qwen3-4B 下比 GOLD 平均提高 +3.82 个点（包括 GSM8k 上 **6倍** 提升，2.56→**15.54**）；H-KL 在 Phi-4-mini 下增加了一致的 +0.5 个点增益。
- **跨分词器家族的多教师 KD**。X-Token 能够从异构教师进行蒸馏。我们展示 **互补性** 是关键（Phi-4-mini++Llama-3.2-3B 比单教师 KD 提高 +1.3），并且简单的 **静态** 加权优于自适应方案。
- **用于 KD 的鲁棒序列对齐**。我们提供确定性的、可扩展的基于 DP 的学生和教师输入序列对齐方法，用于 KD。

## 2 方法

X-Token 由三个组件组成：(i) 跨度对齐以生成文本一致的单位，(ii) 投影矩阵 \(W\) 以桥接词汇表，以及 (iii) 两个互补的损失公式 P-KL 和 H-KL，以及可选的多教师扩展。它们共同实现了不匹配分词器之间的蒸馏。所有损失公式都作用于通过跨度对齐和链式规则合并获得的块级分布，该合并通过自回归乘积将每个对齐跨度内的每个 token 概率组合成单个块级分布。

### 2.1 序列对齐

当教师和学生使用不同的分词器 \(T_S, T_T\) 时，token 序列长度不同且缺乏位置对应关系，使得逐位置 KD 变得不明确。我们通过 **跨度对齐** 来解决这个问题，将 token 分组为块 \(\{ (A_k^S, A_k^T) \}_{k=1}^K\)，这些块解码为相同的底层文本。然后我们在每个块上应用链式规则合并，获得块级分布 \(\hat{p}_S^{(k)}\) 和 \(\hat{p}_T^{(k)}\)，这些作为蒸馏的对齐单位。这种方法受到 [minixhofer2025universal] 的启发。

### 2.2 X-Token 投影矩阵 \(W\)

即使在对齐之后，教师和学生的分布也是定义在不同词汇表上的。我们引入一个投影矩阵 \(W \in \mathbb{R}^{|V_S| \times |V_T|}\)，将学生 token 概率映射到教师词汇表空间，其中 \(V_S\) 和 \(V_T\) 分别表示学生和教师的词汇表。

请参阅标题

图 2: Llama-3.2 学生和 Qwen-3 教师的投影矩阵 \(W\) 子集。精确匹配包括 `_the`, `_cat`, 和 `_run`。对于没有精确匹配的 token，应用多 token 规则：例如，`201` → (2,0,1)，希腊前缀 `_περι` 映射到五个教师子 token，较低权重的条目（带阴影）被截断（top-K=4）。

如图 2 (https://arxiv.org/html/2605.21699#S2.F2) 所示，投影矩阵 \(W \in \mathbb{R}^{|V_S| \times |V_T|}\) 将每个学生 token 映射到教师 token 的加权组合。我们确定性地分两遍构建 \(W\)。

**(1) 精确匹配遍**：对于每个 \((s, t) \in V_S \times V_T\)，其解码后的字符串在规范化后（统一空格前缀如 Ġ/_ 和新行标记）匹配，设置 \(W[s, t] = 1\)。我们在索引数组或矩阵时，互换地使用 \(s\) 和 \(t\) 及其在规范词汇表顺序下的整数索引。

**(2) 多 token 规则遍**：对于每个剩余的学生 token \(s\)，将其解码文本在教师分词器下重新分词，得到序列 \((\tau_0, ..., \tau_\ell) \in V_T\)（这里 \(\tau_i \in V_T\) 是此重新分词中的第 \(i\) 个索引），并设置 \(W[s, \tau_i] = \beta \gamma^i\)，其中 \((\beta, \gamma) = (0.9, 0.1)\)。每行截断为前 4 个条目并进行归一化。矩阵 \(W\) 在训练前构建一次，并且可以在 KD 期间可选地进行微调；完整伪代码在附录 8 (https://arxiv.org/html/2605.21699#S8) 中提供。

### 2.3 知识蒸馏

我们采用标准的 KD 目标 [hinton2015distilling]，但将其应用于对齐后的块。给定块级分布 \(\{ \hat{p}_S^{(k)}, \hat{p}_T^{(k)} \}_{k=1}^K\)，我们在 top-K 教师 logits 上计算 KL（\(K=8192\)）：

\[
\mathcal{L}_{\mathrm{KD}} = \frac{1}{K} \sum_{k=1}^K \mathrm{KL} \bigl( \hat{p}_T^{(k)} \;\|\; \hat{p}_S^{(k)} \bigr).
\]

### 2.4 混合损失公式

我们首先形式化 GOLD [patiño2025\_unlocking\_on\_policy\_distillation\_for\_any\_model\_family] 中使用的基于划分的混合损失，这作为我们损失变体的基线。该公式将词汇表划分为一个 1 对 1 字符串匹配的 **common** 子集 \(\mathcal{C}\) 和 uncommon 剩余部分 \(\mathcal{U}_S, \mathcal{U}_T\)。它在 \(\mathcal{C}\) 上应用直接 KL，在 \(\mathcal{U} = \mathcal{U}_S \cup \mathcal{U}_T\) 上应用排名排序的 \(L_1\) 匹配：

\[
\begin{aligned}
\mathcal{L}_{\mathrm{common}}^{(k)} &= \sum_{(s, t) \in \mathcal{C}} \hat{p}_T^{(k)}[t] \bigl( \log \hat{p}_T^{(k)}[t] - \log \hat{p}_S^{(k)}[s] \bigr), \\
\mathcal{L}_{\mathrm{ULD}}^{(k)} &= \Bigl\| \mathrm{sort}_\downarrow \bigl( \hat{p}_S^{(k)} |_{\mathcal{U}_S} \bigr) - \mathrm{sort}_\downarrow \bigl( \hat{p}_T^{(k)} |_{\mathcal{U}_T} \bigr) \Bigr\|_1, \\
\mathcal{L}_{\mathrm{GOLD}}^{(k)} &= \lambda_{\mathrm{KL}} \mathcal{L}_{\mathrm{common}}^{(k)} + \lambda_{\mathrm{ULD}} \mathcal{L}_{\mathrm{ULD}}^{(k)}.
\end{aligned}
\]

虽然这种混合公式使得跨分词器 KD 成为可能，但它在 uncommon 集合中的 token 上引入了不良的梯度行为，我们接下来进行分析。

### 2.5 P-KL：解决混合损失中的错误和抑制性梯度

GOLD 的混合损失在 uncommon 学生 logits 上引出了两种不良的梯度行为（图 1 (https://arxiv.org/html/2605.21699#S1.F1)）。

**来自基于排名匹配的错误梯度：** ULD 项 \(\mathcal{L}_{\mathrm{ULD}}\) 通过排名匹配 uncommon 集合中的 token，将每个学生 token 与相似排名的教师 token 配对，而不是语义对应。这产生了与身份无关的梯度，使关键 token（例如，数字）与不相关的教师 token（例如，特殊字符）不对齐，降低了监督质量。

**来自 common-KL 项的抑制性梯度：** common-KL 项 \(\mathcal{L}_{\mathrm{common}}\)（公式 1 (https://arxiv.org/html/2605.21699#S2.E1)）使用全词汇表 softmax 计算。虽然 uncommon token 在损失中不明显出现，但归一化对所有 logits 产生梯度，降低了 \(\mathcal{U}\) 中 token 的相对概率。详细证明见附录。

这些效应共同导致对 uncommon token 的监督弱或误导，特别是当关键 token 落入 \(\mathcal{U}\) 时（例如，Llama 的 `1,100` 多位数字在数字拆分 Qwen 分词器下），导致性能下降（例如，GSM8k 降至 2.56，而来自较弱教师的同分词器 KD 为 12.89）。

为了解决这一局限性，我们利用 **P-KL**，它将学生分布 \(\hat{p}_S^{(k)}\) 投影到教师词汇表空间 \(\tilde{p}_S^{(k)}\)，实现与教师分布的直接对齐（图 1 (https://arxiv.org/html/2605.21699#S1.F1)）。这里，\(i\) 索引学生词汇表 \(V_S\)，\(j\) 索引教师词汇表 \(V_T\)：

\[
\tilde{p}_S^{(k)}[t] = \sum_{s \in V_S} W[s, t] \cdot \hat{p}_S^{(k)}[s], \quad \mathcal{L}_P^{(k)} = \mathrm{KL} \bigl( \hat{p}_T^{(k)} \,\|\, \tilde{p}_S^{(k)} \bigr).
\]

这个公式用教师感知的监督替代了两种错误源，覆盖所有 token，包括那些在 \(\mathcal{U}\) 中的（例如，`201` 映射到 `[2,0,1]`），通过直接将学生分布与教师分布对齐，恢复了划分丢弃的引导。

### 2.6 H-KL：放宽 1 对 1 匹配

当没有关键 token 被路由到 \(\mathcal{U}\) 时——我们在学生词汇表上按类别进行审计（表 8 (https://arxiv.org/html/2605.21699#S8)）——...（原文未结束，但我们翻译到这里）