Inner-layer Token Self-modulation as Another Scaling Axis for LLMs
Yebin Yang, Huaijin Wu, Jingtao Han, Yu Wang, Xiaohan Qin, Jingzhi Wang, Debing Zhang, Junchi Yan
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.
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Your agent calls
Luneget_paper_fulltext
Free to start. No credit card required.
Terminal
Install the CLIlune papers fulltext 1507c3e3-0285-4725-8105-5b62a0cd3afeBuilds on17
- Measuring Massive Multitask Language UnderstandingDan Hendrycks, Collin Burns, Steven Basart, Andy Zou et al.ICLR 2021 · 7,905 citations
- SGLang: Efficient Execution of Structured Language Model ProgramsLianmin Zheng, Liangsheng Yin, Zhiqiang Xie, Chuyue Sun et al.NeurIPS 2024 · 1,586 citations
- DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language ModelsDamai Dai, Chengqi Deng, Chenggang Zhao, R. X. Xu et al.ACL 2024 · 171 citations
- Scaling Laws for Fine-Grained Mixture of ExpertsJan Ludziejewski, Jakub Krajewski, Kamil Adamczewski, Maciej Pióro et al.ICML 2024 · 149 citations
- Beyond Chinchilla-Optimal: Accounting for Inference in Language Model Scaling LawsNikhil Sardana, Jacob P. Portes, Sasha Doubov, Jonathan FrankleICML 2024 · 144 citations
Related papers
- Mixture of Tokens: Continuous MoE through Cross-Example AggregationSzymon Antoniak, Michal Krutul, Maciej Pióro, Jakub Krajewski et al.NeurIPS 2024 · 6 citations
- Hyperparameter Transfer with Mixture-of-Expert LayersTianze Jiang, Blake Bordelon, Cengiz Pehlevan, Boris HaninICML 2026 · 6 citations
- Optimal Sparsity of Mixture-of-Experts Language Models for Reasoning TasksTaishi Nakamura, Satoki Ishikawa, Masaki Kawamura, Takumi Okamoto et al.ICLR 2026 · 2 citations
- Mixture-of-Recursions: Learning Dynamic Recursive Depths for Adaptive Token-Level ComputationSangmin Bae, Yujin Kim, Reza Bayat, Sungnyun Kim et al.NeurIPS 2025 · 143 citations
- cMoLLM at Scale: Horizontal Scaling Laws for Convolutionally-Gated Mixture-of-LLMsXin Yang, Yemin Wang, Mingda Liu, Letian Li et al.ICML 2026
