1. 项目概述
在大型语言模型（LLM）的发展历程中，扩展定律（Scaling Law）揭示了模型性能与参数规模、计算资源之间的幂律关系。传统方法主要通过增加密集参数或采用混合专家（MoE）架构来提升模型能力，但这些方法往往面临计算成本线性增长或内存开销大的挑战。JTok与JTok-M提出了一种全新的扩展维度——令牌索引参数（token-indexed parameters），通过轻量级的元素级操作调制主干网络，在几乎不增加计算开销的前提下显著提升模型性能。
这项工作的核心创新在于：将传统沿着密集维度的扩展转变为沿着令牌维度的扩展。具体来说，每个Transformer层都配备一个可学习的令牌嵌入表，通过检索特定令牌的调制向量来对主干网络进行轻量级调制。这种方法突破了传统扩展方式的瓶颈，为模型能力提升开辟了新的路径。
2. 核心设计思路
2.1 传统扩展方式的局限性
当前LLM的扩展主要沿着三个维度：
- 密集参数扩展：增加模型宽度和深度，但计算成本（FLOPs）随之线性增长
- 数据扩展：增加训练token数量，但高质量文本数据日益稀缺
- MoE架构：通过稀疏激活专家网络解耦容量与计算，但面临内存开销和硬件效率挑战
这些方法都面临边际效益递减的问题。例如，当模型参数超过一定规模后，单纯增加参数反而可能导致性能下降；MoE架构虽然理论上可以无限扩展专家数量，但实际上路由效率和专家利用率会迅速饱和。
2.2 令牌索引参数的设计理念
JTok的核心思想是将传统静态的词嵌入扩展为动态的层间调制机制。具体实现包含三个关键设计：
- 分层调制：每个Transformer层维护独立的令牌嵌入表，允许不同层对同一令牌进行差异化处理
- 轻量操作：采用Hadamard积等元素级操作进行调制，计算开销可忽略不计
- 容量解耦：调制参数规模与主干计算完全解耦，可以独立扩展
这种设计使得模型能够在不增加FLOPs的情况下，通过扩大令牌嵌入表来提升容量。实验表明，当令牌索引参数与主干参数比例达到50:1时，仍能保持稳定的性能提升。
2.3 JTok-M的混合专家扩展
JTok-M在JTok基础上引入了混合专家思想，进一步提升了模型的表达能力：
- 多专家池：每个令牌在每个层配备多个调制专家，形成丰富的调制空间
- 轻量路由：基于隐藏状态动态选择Top-K专家进行混合
- 负载均衡：引入辅助损失函数确保各专家得到均衡训练
这种设计既保留了JTok的计算效率优势，又通过动态路由机制增强了模型的上下文感知能力。在17B MoE主干上的实验表明，JTok-M可以带来平均8%的下游任务性能提升。
3. 技术实现细节
3.1 JTok的调制机制
JTok的数学表达简洁而有效。对于第ℓ层的令牌x：
从嵌入表Eℓ中检索对应的调制向量：
e = Eℓ[x] # 维度d的向量应用层归一化和可学习的逐维度缩放：
p = 1 + sℓ ⊙ Normε(e) # sℓ是可学习缩放参数调制MLP的输出：
Δm̂ = Δm ⊙ p # 元素级乘法
这种设计有几点关键考虑：
- 加1操作保留了原始MLP输出的基准值
- 层归一化确保调制向量的尺度稳定
- 逐维度缩放允许模型自适应调整各维度的调制强度
3.2 JTok-M的路由机制
JTok-M的路由过程包含以下步骤：
计算专家logits：
g = (RMSNorm(h))ᵀR # R∈ℝ^(d×ne)是路由矩阵选择Top-K专家并计算混合权重：
G = TopK(g, K) w_i = σ(g_i)/∑σ(g_j) # 使用sigmoid而非softmax混合专家输出：
e = ∑(w_i E_i[x]) # 加权求和
这种路由设计相比传统MoE有几个优势：
- sigmoid激活避免专家间的竞争关系
- 更细粒度的专家利用率控制
- 与主干计算完全解耦，可异步执行
3.3 系统优化策略
为实现高效实现，团队开发了多项系统级优化：
- 异步预取：令牌嵌入的检索与主干计算重叠执行
- 令牌去重：利用令牌频率的Zipf分布特性，对重复令牌只检索一次
- 嵌入并行：将大型嵌入表分片到多GPU，减少单卡内存压力
- CPU卸载：推理时将不活跃的嵌入表保留在主机内存
这些优化使得JTok-M在17B模型上的训练吞吐损失小于7%，推理延迟增加控制在7.3%以内。
4. 实验验证与分析
4.1 性能提升验证
在多样化基准测试中，JTok系列展现出显著优势：
| 模型类型 | 参数量 | MMLU(Δ) | ARC(Δ) | CEval(Δ) |
|---|---|---|---|---|
| Dense-XL | 1.5B | +4.55 | +5.84 | +6.21 |
| MoE 1.5B-A250M | 1.5B | +6.16 | +5.51 | +7.89 |
| MoE 17B-A2B | 17B+44B | +4.11 | +8.28 | +8.93 |
特别值得注意的是，这些提升是在FLOPs基本不变的情况下实现的。例如，17B+44B的JTok-M变体相比纯17B MoE，计算量仅增加约5%，但性能提升相当于35%的计算节省。
4.2 扩展定律验证
通过系统的isoFLOPs分析，研究验证了令牌索引参数也遵循幂律扩展：
- 在固定计算预算下，测试损失与令牌参数数量呈现稳定的log-linear关系
- 扩展曲线与传统密集参数平行，验证了其作为独立扩展维度的有效性
- 最优扩展比例η=Nn/Nc约在30-50之间，超过此范围后收益递减
这种可预测的扩展行为为未来模型设计提供了重要指导。开发者可以根据计算预算，在密集参数、稀疏专家和令牌参数之间进行最优分配。
5. 实际应用建议
基于项目经验，我们总结出以下实践建议：
参数比例选择：
- 对于密集主干，建议η=20-30
- 对于MoE主干，建议η=40-50
- 过大的η会导致收益递减
路由配置技巧：
- 专家数量ne=4-6为宜
- Top-K选择K=2最佳
- 使用sigmoid而非softmax计算权重
训练优化：
- 嵌入表学习率设为主干的1/5
- 逐步增加η的课程学习策略
- 定期检查专家利用率，避免退化
部署考量：
- 优先考虑CPU卸载方案
- 对高频令牌实现缓存机制
- 使用8bit量化压缩嵌入表
这些经验来自对数十种配置的实验验证，可以帮助开发者快速获得最佳实践。
6. 未来发展方向
令牌索引参数为LLM扩展开辟了新的可能性，几个有潜力的方向包括：
- 动态容量分配：根据任务难度自动调整η值
- 跨层参数共享：在深层和浅层间共享部分嵌入表
- 多模态扩展：将概念应用于视觉token的调制
- 条件计算：基于调制强度预测跳过某些层计算
这种方法也可能影响硬件设计，推动面向大规模嵌入检索的专用加速器发展。
从个人实践角度看，JTok系列最令人振奋的不只是性能提升本身，而是它展示了一种新的模型扩展范式。当业界普遍认为Transformer的扩展路径已经趋于成熟时，这项工作提醒我们：通过重新思考参数组织的基本方式，仍然可能发现意想不到的高效路径。这种基础性创新正是推动领域前进的关键动力。
