Universal Cross-Tokenizer Distillation via Approximate Likelihood Matching

本文提出了一种跨分词器蒸馏方法ALM，通过近似似然匹配实现不同分词器间的知识转移，首次在子词到字节级迁移等场景中取得显著效果，并在多个应用案例中优于现有方法。

Large Language Model, Pre-training, Fine-tuning, Efficiency, Representation Learning

Benjamin Minixhofer, Ivan Vulić, Edoardo Maria Ponti

University of Cambridge, University of Edinburgh

Generated by grok-3

Background Problem

大型语言模型（LLMs）的知识蒸馏（Distillation）是一种有效的知识转移方法，但传统方法要求教师模型和学生模型使用相同的分词器（Tokenizer），这限制了可适用的教师-学生模型对，尤其是在分词器类型根本不同（如子词到字节级）的情况下。本文旨在解决这一关键问题，提出了一种跨分词器蒸馏方法，首次实现了在高度异构分词器之间的有效知识转移，同时在同类分词器蒸馏中也显著优于现有方法。

Method

本文提出了**Approximate Likelihood Matching (ALM)**方法，用于跨分词器知识蒸馏，其核心思想和步骤如下：

• 核心思想：通过对齐教师和学生模型在不同分词器下生成的token序列中编码相同文本内容的片段（Chunks），并最小化这些片段的似然差异（Likelihood Difference），实现知识转移。
• 具体实现：
  1. 对输入文本分别使用教师和学生模型的分词器进行分词，计算各自的下一token概率。
  2. 找到对齐的token片段（即编码相同文本部分的片段），使用贪婪对齐策略。
  3. 由于无法直接计算所有可能片段的似然（因数量无限），采用二值化的f-散度（Binarised f-Divergence）近似计算片段似然差异，公式为： $\mathcal{L}_{S,T}^{\text{ALM}}(\boldsymbol{x}) = \sum_{i,j,k,l \in A_c(\boldsymbol{x})} f\big(p_T(\boldsymbol{x}, i:j)^{\frac{1}{\tau}} \mid p_S(\boldsymbol{x}, k:l)^{\frac{1}{\tau}}\big) + f\big(1 - p_T(\boldsymbol{x}, i:j)^{\frac{1}{\tau}} \mid 1 - p_S(\boldsymbol{x}, k:l)^{\frac{1}{\tau}}\big)$ 其中$\tau$为温度参数，$f$为f-散度函数。
  4. 可选地，通过Outcome Chunk Debiasing减少分词偏见对结果的影响，并通过对齐隐藏状态（Hidden States）增加教师信号的丰富性。
  5. 使用GradMag方法平衡多个损失组件的梯度贡献，确保各损失项对训练的贡献均等。
• 关键创新：ALM不仅适用于相似分词器之间的蒸馏，还首次实现了根本不同分词器（如子词到字节级）之间的纯蒸馏（Pure Distillation），无需依赖辅助目标。
• 批判性思考：虽然ALM在理论上提供了跨分词器蒸馏的解决方案，但其对齐策略（贪婪对齐）可能在复杂文本中导致次优匹配，影响蒸馏效果。此外，二值化近似可能丢失部分信息，尤其是在长序列或高复杂性文本中，近似误差可能累积，导致性能下降。作者未充分讨论这些潜在问题及其对实际应用的影响。

Experiment

本文通过三个应用场景（Use Cases）验证了ALM方法的有效性，实验设置和结果如下：

• Use Case 1: Tokenizer Transfer via Self-Distillation：将预训练模型（如Gemma2 2B IT和Llama3.2 3B IT）迁移到不同分词器（如Qwen2和字节级），并测试模型集成（Ensembling）效果。使用Tulu3指令微调数据集训练，评估指标包括PiQA、ARC-C、BoolQ等。结果显示，ALM在子词到子词迁移中优于SFT、DSKD和MinED基线（例如，Gemma2到Qwen2的平均性能从SFT的51.6提升到ALM的55.1）；在子词到字节级迁移中，ALM也显著优于基线（例如，Llama3.2到字节级的平均性能从SFT的52.6提升到ALM的53.5）。集成实验进一步表明，迁移后的模型通过概率平均可以提升整体性能。
• Use Case 2: Large-to-Small Distillation Across Tokenizers：将大型数学专业模型（OpenMath2-Llama3.1-8B）蒸馏到小型通用模型（Gemma2 2B），使用OpenMathInstruct-2数据集训练，评估指标为GSM8K和MATH的零样本准确率。结果显示，ALM和ALM+SFT在平均性能上优于SFT（53.3 vs 51.7），并显著优于DSKD和MinED，达到教师模型性能的约71%。
• Use Case 3: Zero-Shot Tokenizer Transfer Hypernetworks：将ALM应用于零样本分词器迁移超网络训练，评估指标包括PiQA、BoolQ等。结果显示，ALM在多个目标分词器（如GPT2、Mistral v2）上的平均性能优于SFT（例如，Mistral v2从49.6提升到51.9）。
• 实验设计评价：实验设置覆盖了多种实际场景，较为全面，但存在以下问题：1）模型规模较小（2.4B到8B），未验证在大规模模型上的效果；2）字节级迁移的性能差距较大，作者未深入分析原因，仅提出推测性改进方向；3）部分基准测试（如IFEval）结果波动较大，未充分讨论数据或评估方法的潜在影响。
• 批判性思考：虽然ALM在实验中表现优于基线，但其性能提升幅度有限，尤其是在字节级迁移中，与原始模型的差距仍显著，表明方法在处理根本不同的分词粒度时可能存在局限性。此外，计算成本（FLOPs和内存需求）虽然低于部分基线（如DSKD），但仍高于SFT，实际应用中的可行性需进一步验证。

Further Thoughts

ALM方法为跨分词器蒸馏提供了一个全新的视角，尤其是在模型重用和组合性方面的潜力令人印象深刻。未来可以将ALM与其他高效蒸馏技术（如拒绝采样或在线策略蒸馏）结合，进一步提升其在大规模模型和长上下文场景下的适用性。此外，ALM的对齐策略和近似方法可能启发其他领域的知识转移研究，例如跨模态学习（Multimodal Learning），其中不同模态的数据表示也存在类似的分词器异构问题。另一个值得探索的方向是，是否可以通过更复杂的对齐算法（如动态规划而非贪婪对齐）来减少近似误差，从而进一步提升蒸馏效果。总之，ALM虽然在当前实验中表现出色，但其理论边界和实际应用中的可扩展性仍需更多研究和验证，尤其是在面对更复杂的模型架构和数据分布时。