arXiv:2602.00800cs.LGcs.AI2026-02被引 1

通过词元索引参数实现模型扩容不增算力,突破传统扩展瓶颈。

JTok: On Token Embedding as another Axis of Scaling Law via Joint Token Self-modulation

  • 引入词元调制向量,以轻量级操作提升模型容量
  • 在相同算力下,性能比原MoE提升35%,如MMLU+4.1
  • 适合追求高效率、低延迟的大型语言模型部署场景

大语言模型通常沿密集维度扩展,导致性能与计算成本近线性增长。尽管MoE架构解耦了容量与计算,但带来巨大内存开销和硬件效率问题。为此,我们提出词元索引参数作为新的、正交的扩展轴,实现容量与FLOPs解耦。具体地,引入联合词元(JTok)与词元混合(JTok-M),通过辅助嵌入表检索调制向量,以轻量元素级操作调制主干网络,算力开销可忽略。在从650M(190M + 460M嵌入)到61B(17B + 44B嵌入)的密集与MoE骨干上进行大量实验,结果表明该方法持续降低验证损失,显著提升下游任务表现(如MMLU+4.1,ARC+8.3,CEval+8.9)。严格的等算力分析证实,JTok-M从根本上改变了质量-算力帕累托前沿,在相同性能下比原始MoE减少35%计算量,并验证了词元索引参数具有可预测的幂律扩展规律。此外,高效实现确保其引入的开销极小。

原文摘要 · Abstract (English)

LLMs have traditionally scaled along dense dimensions, where performance is coupled with near-linear increases in computational cost. While MoE decouples capacity from compute, it introduces large memory overhead and hardware efficiency challenges. To overcome these, we propose token-indexed parameters as a novel, orthogonal scaling axis that decouple model capacity from FLOPs. Specifically, we introduce Joint-Token (JTok) and Mixture of Joint-Token (JTok-M), which augment Transformer layers with modulation vectors retrieved from auxiliary embedding tables. These vectors modulate the backbone via lightweight, element-wise operations, incurring negligible FLOPs overhead. Extensive experiments on both dense and MoE backbones, spanning from 650M (190M + 460M embedding) to 61B (17B + 44B embedding) total parameters, demonstrate that our approach consistently reduces validation loss and significantly improves downstream task performance (e.g., +4.1 on MMLU, +8.3 on ARC, +8.9 on CEval). Rigorous isoFLOPs analysis further confirms that JTok-M fundamentally shifts the quality-compute Pareto frontier, achieving comparable model quality with 35% less compute relative to vanilla MoE architectures, and we validate that token-indexed parameters exhibit a predictable power-law scaling behavior. Moreover, our efficient implementation ensures that the overhead introduced by JTok and JTok-M remains marginal.

模型扩展轻量化MoE词元调制

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。