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

TL;DR

The paper addresses inefficiencies in dense and Mixture-of-Experts scaling by introducing token-indexed parameters as a new scaling axis. It presents Joint-Token (JTok) and Mixture of Joint-Token (JTok-M), which inject token-specific modulation vectors from embedding tables into Transformer layers using lightweight Hadamard gating, decoupling capacity from FLOPs. The authors establish a scaling framework with isoFLOPs and derive a scale-invariant compute saving of about $35\%$ for JTok-M, along with a predictable power-law scaling of the token-indexed parameters. Empirically, token-indexed methods yield consistent gains across dense and MoE backbones up to $17$B, with substantial downstream improvements and only modest system overhead, illustrating token-indexed parameters as a practical and robust scaling direction for future LLMs.

Abstract

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.

Figures (8)

• Figure 1: Architecture of JTok/JTok-M. JTok (left) augments each Transformer layer with a token-indexed table. Each token retrieves a modulation vector, applies norm and a learnable per-dimension scaling, and forms a lightweight multiplicative gate to modulate the FFN update via element-wise products. JTok-M (right) generalizes JTok by maintaining a pool of token-indexed modulators and using a router conditioned on the hidden state to select a sparse Top-$K$ mixture per token; the mixed modulator is normalized+scaled and injected as an additional residual alongside the backbone update. Both are plug-in bypass modules implemented with table lookups and element-wise operations, allowing retrieval to be overlapped with backbone and adding negligible compute.
• Figure 2: Training loss for dense and MoE backbones and corresponding JTok and JTok-M variants. The top row shows JTok's performance on Dense-S(190M), M(505M), and L(1B) backbones. The bottom row shows results on Dense-XL (1.5B), 1.5B-A250M MoE, and 3.2B-A0.5B MoE backbones. In all settings, JTok (and JTok-M for MoE backbone) achieves a consistently and significantly lower training loss.
• Figure 3: Downstream accuracy trajectories on the 17B-A2B MoE backbone. We track the zero/few-shot performance on six representative benchmarks (MMLU, ARC-C, HellaSwag, CMMLU, C-Eval, and Xiezhi) throughout the 570B-token pretraining course. JTok-M (blue) consistently surpasses the backbone baseline (orange) from the early stages, yielding substantial gains in both knowledge-intensive and reasoning tasks without saturation.
• Figure 4: JTok-M (fixed $(\eta,\rho)$) improvements remain stable and consistent across different backbone model sizes.
• Figure 5: IsoFlOPs profile of vanilla MoE models. Each curve shows test loss vs. per-token gFLOPs at fixed FLOPs, with optimal points defining the efficient frontier.