Inner-layer Token Self-modulation as Another Scaling Axis for 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 ReToken and MoRT, 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 190M to 9.8B parameters, demonstrate that our approach consistently reduces validation loss and significantly improves downstream task performance (e.g., +7.3 on ARC-C, +6.3 on GSM8K). Rigorous isoFLOPs analysis further confirms that MoRT 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 ReToken and MoRT remains marginal.
Lay Summary
Large language models usually become better by making the whole model bigger, but this also makes them much more expensive to train and run. A popular alternative, called mixture-of-experts, adds extra model capacity without using all of it at once, but it can still be difficult to deploy efficiently because it requires moving many parameters around in memory. In this work, we explore a different way to scale language models. Instead of adding larger neural network layers, we give each text token access to small learned vectors that can adjust how the model processes that token inside each layer. Our two methods, ReToken and MoRT, act like lightweight token-specific controls: they add useful model capacity through simple element-wise operations, while adding very little extra computation. Across dense and mixture-of-experts language models, this approach consistently improves training loss and downstream performance on knowledge, reasoning, coding, and math benchmarks. Our scaling experiments suggest that MoRT can reach the same model quality with about 35% less training compute. This points to token-indexed parameters as a practical new direction for building stronger and more efficient language models.