TokAlign: Efficient Vocabulary Adaptation via Token Alignment¶
会议: ACL 2025
arXiv: 2506.03523
领域: LLM 效率 / 词表适配
关键词: 词表替换, Token 对齐, 共现矩阵, 跨语言迁移, 知识蒸馏
一句话总结¶
提出 TokAlign，基于 Token 共现信息学习两个词表之间的一对一映射矩阵，高效替换 LLM 的词表，实现跨语言知识迁移和跨模型 token 级蒸馏。
研究背景与动机¶
核心问题¶
LLM 的 tokenizer 在新领域或新语言上往往效率低下——例如 LLaMA3 的 tokenizer 处理亚美尼亚语时，文本长度是英语的 3.95 倍（相同字节数下）。这直接导致：
- 训练和推理速度下降：低压缩率 = 更长的 token 序列 = 更慢的处理
- 跨模型知识迁移受阻：不同 LLM 使用不同词表，无法进行 token 级蒸馏或集成
- 重新预训练代价巨大：为新 tokenizer 从头训练 LLM 成本极高
现有方法的局限¶
- Focus：依赖两个词表的共有 token 来初始化新 token，当词表重叠少时效果差（初始 PPL 高达 2.9e5）
- WECHSEL：需要双语词典和外部词嵌入，对无词典语言不可用
- ZeTT：训练额外的超网络（hypernetwork）来生成 token 参数，计算成本极高（Pythia-2.8B 需 661 GPU 小时）
- OFA：通过矩阵分解和多语言词嵌入组合，流程复杂
方法详解¶
整体框架¶
TokAlign 从 token-token 共现的角度出发，分三步完成词表对齐：
Step 1: 语料 Tokenization - 使用混合语料（CulturaX 40% + The Stack 30% + Proof-Pile-2 30%） - 分别用源 tokenizer 和目标 tokenizer 对同一语料进行 tokenize - 默认使用 1B tokens 的语料规模
Step 2: Token 表示学习 - 使用 GloVe 从各自的 token 序列中学习 token 表示向量 - 选择 GloVe 而非 CBOW/FastText 的原因：GloVe 利用全局统计共现信息 - 训练开销极低：128 核 CPU 服务器上不到 2 小时
Step 3: Token 对齐 - 基于学到的 token 表示，计算源/目标词表间的逐对 cosine 相似度 - 构建对齐矩阵 \(M_{s \to t}\)，为每个源 token 找到最相似的目标 token - 两个词表共有的 token 直接映射，无需对齐
关键设计¶
1. 对齐质量评估指标
提出两个新指标评估对齐矩阵质量： - BLEU-1：将源 token ID 序列通过对齐矩阵转换后，与目标 token ID 序列计算 BLEU-1（文本匹配视角） - BERTScore：将转换后的 token 序列解码为文本，与原文计算语义相似度
这两个指标与初始预训练 loss 成正比关系，可作为对齐质量的先验判断。
2. 渐进式适配（Progressive Adaptation）
词表替换后的微调分为两个阶段： - 前半程：仅调整 Embedding 层和 lm_head，避免 loss 剧烈波动 - 后半程：全参数一起微调
这种渐进策略显著提升了训练稳定性，避免了 loss spike。
3. 参数初始化
对齐矩阵确定后，目标词表中每个 token 的参数（embedding 和 lm_head）从最相似的源 token 复制初始化，提供了极好的起点。
实验关键数据¶
主实验¶
跨语言迁移（Pythia → Qwen2 词表）：
| 初始化方法 | 初始 PPL (Avg) | + LAT 后 PPL |
|---|---|---|
| Focus | 3.1e5 | 12.4 |
| ZeTT | 3.4e2 | 7.8 |
| TokAlign | 1.2e2 | 6.9 |
→ TokAlign 初始 PPL 比 Focus 低 3 个数量级，比 ZeTT 低 65%
文本压缩率提升：词表替换后，13 种语言的 token 长度平均减少 29.2%
低资源语言惊艳表现：TokAlign 初始化的 Pythia-1B 在 3 种低资源语言上的 PPL 甚至优于同规模 Qwen2
词表适配后性能恢复：
| 方法 | GPU 小时 | 恢复步数 | Avg @6 tasks |
|---|---|---|---|
| Focus | 99.70 | >5k | 45.29 |
| ZeTT | 418.94 | ~5k | 45.48 |
| TokAlign | 99.70 | ~5k | 48.42 |
→ 训练速度比 Focus 快 1.92 倍，仅需 5k 步即可恢复原模型 97.6% 性能
Token 级蒸馏¶
统一词表后可进行 token 级蒸馏，效果远超句子级蒸馏：
| 方法 | ARC-E | BoolQ | HellaSwag | Avg |
|---|---|---|---|---|
| Pythia-1B 原始 | 56.82 | 60.43 | 37.68 | 49.55 |
| + 句子级蒸馏 (Qwen2-7B) | 52.27 | 67.49 | 39.03 | 49.90 |
| + Token 级蒸馏 (Qwen2-7B) | 62.33 | 70.18 | 41.58 | 54.02 |
→ Token 级蒸馏比句子级蒸馏高 4.4%，仅用 235M tokens
→ Pythia-1B 经 token 级蒸馏后性能可比肩原始 Pythia-7B
关键发现¶
- 跨语言 zero-shot ICL：TokAlign 比 Focus 平均高 4.4%，日语 +2.3%，越南语 +2.2%
- 英语能力保留：TokAlign 保留 54.2% 英语能力，远超 Focus 的 40.8%
- 语料鲁棒性：换用 SlimPajama 训练 GloVe 嵌入，结果基本不变
- 初始 loss 大幅降低：Pythia-2.8B 首步 loss 从 17.8（Focus）降至 9.5（TokAlign）
- 相对表示改进：使用 300 个共有 anchor token 将 GloVe 嵌入转为相对表示，对齐质量进一步提升
亮点与洞察¶
- 无监督且无需外部资源：不需要双语词典、外部词嵌入或额外超网络，纯粹从共现统计中学习对齐
- 计算成本极低：GloVe 训练 + 对齐仅需 CPU 上 2 小时，相比 ZeTT 的 661 GPU 小时降低了 3 个数量级
- 打通跨模型知识迁移：统一词表后 token 级蒸馏效果显著，小模型可以高效吸收大模型知识
- 分布式假设的成功应用：借鉴经典的分布式语义假设（co-occurrence → semantics），在 subword token 级别验证了其有效性
- 全词表替换而非扩展：不同于多数方法仅扩展词表，TokAlign 支持完全替换，适用范围更广
局限性¶
- 依赖 GloVe 训练质量：低频 token 的共现统计可能不充分，影响对齐准确性
- 一对一映射限制：源和目标词表 token 粒度不同时，一对一映射可能丢失信息
- 仅在 Pythia 系列验证：主实验基于 Pythia（非 SOTA 模型），在更大规模或更强模型上的效果待验证
- 混合语料比例需调优：CulturaX/Stack/ProofPile 的 40/30/30 比例是经验值，不同场景可能需要调整
- 渐进式微调策略较朴素：简单的两阶段划分可能不是最优，更精细的学习率调度可能有改进空间
相关工作¶
- 词表示学习：Word2Vec、GloVe、FastText 等经典方法，TokAlign 借鉴 GloVe 的全局共现矩阵分解思路
- 词表适配：Focus、WECHSEL、ZeTT、OFA 等方法，各有依赖外部资源或高计算成本的缺点
- 知识蒸馏：token 级 vs 句子级蒸馏的比较，在 NMT 领域已有讨论，本文将其推广到 LLM
- 多语言模型：Qwen2、LLaMA3、Gemma 等多语言 LLM，词表设计对低资源语言的影响
评分¶
| 维度 | 分数 | 说明 |
|---|---|---|
| 新颖性 | ⭐⭐⭐⭐ | GloVe 共现 + token 级对齐的结合很有创意 |
| 理论性 | ⭐⭐⭐ | 方法直觉合理但理论分析偏少 |
| 实验充分性 | ⭐⭐⭐⭐⭐ | 多规模多语言多下游任务，消融全面 |
| 实用价值 | ⭐⭐⭐⭐⭐ | 成本极低效果好，真正实用的词表替换方案 |
| 总体推荐 | ⭐⭐⭐⭐ | 简洁实用的词表适配方法，打通跨模型蒸馏 |