JTok：通过联合 Token 自调制将 Token 嵌入作为另一维度的扩展规律
机器学习
2026-02-03 v1 人工智能
摘要
大型语言模型（LLM）传统上沿稠密维度扩展，性能与计算成本的近线性增长相伴随。虽然 Mixture of Experts（MoE）通过解耦容量与计算量，但引入大量内存开销和硬件效率挑战。为此，我们提出一种作为新型正交扩展维度的 token 索引参数，以实现模型容量与 FLOPs 的解耦。具体而言，我们引入联合 Token（Joint-Token, JTok）和联合 Token 混合模型（Mixture of Joint-Token, JTok-M），通过检索辅助嵌入表中的调制向量来增强 Transformer 层。这些向量通过轻量级的元素级操作调制主干网络，引入的 FLOPs 开销可忽略不计。对 dense 模型和 MoE 主干进行大规模实验，参数规模从 650M（190M + 460M 嵌入）到 61B（17B + 44B 嵌入）不等，结果表明我们方法在一致降低验证损失的同时，显著提升下游任务性能（例如在 MMLU 上提升 4.1，在 ARC 上提升 8.3，在 CEval 上提升 8.9）。严格的 isoFLOPs 分析进一步确认，JTok-M 本质上重新定义了质量-计算 Pareto 前沿，实现与 vanilla MoE 相当的模型质量时计算量降低 35%。我们还验证了 token 索引参数 exhibits 可预测的幂律扩展规律。此外，我们高效的实现确保 JTok 和 JTok-M 引入的开销保持在边际水平。
引用
@article{arxiv.2602.00800,
title = {JTok: On Token Embedding as another Axis of Scaling Law via Joint Token Self-modulation},
author = {Yebin Yang and Huaijin Wu and Fu Guo and Lin Yao and Xiaohan Qin and Jingzhi Wang and Debing Zhang and Junchi Yan},
journal= {arXiv preprint arXiv:2602.00800},
year = {2026}
}