JTok: On Token Embedding as another Axis of Scaling Law via Joint Token Self-modulation
On the map: Mixture-of-experts → Expert routing strategies
Abstract (arXiv)
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.
Key facts. Cited by 2 papers in the map (data through 7 Oct 2026). Placed in Expert routing strategies, one of 12 directions in Mixture-of-experts. Builds on 50 of its 59 references that are in the map. First posted on arXiv 31 Jan 2026.
Builds on
50 of 59 references.
Closest ancestors · 3
Same direction, newest first
- 2024Ultra-Sparse Memory Network19 Nov
- 2024Mixture of A Million Experts4 Jul
- 2021Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity11 Jan
Other references · 47
- 2017Attention Is All You Needcited by 59,932
- 2015Deep Residual Learning for Image Recognitioncited by 51,809
- 2014Adam: A Method for Stochastic Optimizationcited by 36,960
- 2021Training Verifiers to Solve Math Word Problemscited by 9,743
- 2025Qwen3 Technical Reportcited by 8,195
- 2021Evaluating Large Language Models Trained on Codecited by 8,053
- 2020Measuring Massive Multitask Language Understandingcited by 7,768
- 2020Scaling Laws for Neural Language Modelscited by 6,411
- 2021Measuring Mathematical Problem Solving With the MATH Datasetcited by 5,203
- 2018Think you have Solved Question Answering? Try ARC, the AI2 Reasoning Challengecited by 4,710
- 2019HellaSwag: Can a Machine Really Finish Your Sentence?cited by 4,018
- 2021Program Synthesis with Large Language Modelscited by 3,509
35 more
- 2017Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layercited by 3,413
- 2022Training Compute-Optimal Large Language Modelscited by 2,927
- 2017TriviaQA: A Large Scale Distantly Supervised Challenge Dataset for Reading Comprehensioncited by 2,895
- 2023GPQA: A Graduate-Level Google-Proof Q&A Benchmarkcited by 2,748
- 2024LiveCodeBench: Holistic and Contamination Free Evaluation of Large Language Models for Codecited by 1,906
- 2019Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelismcited by 1,829
- 2022Challenging BIG-Bench Tasks and Whether Chain-of-Thought Can Solve Themcited by 1,767
- 2023SGLang: Efficient Execution of Structured Language Model Programscited by 1,107
- 2024DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Modelcited by 1,093
- 2019DROP: A Reading Comprehension Benchmark Requiring Discrete Reasoning Over Paragraphscited by 883
- 2024DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Modelscited by 834
- 2023C-Eval: A Multi-Level Multi-Discipline Chinese Evaluation Suite for Foundation Modelscited by 694
- 2023CMMLU: Measuring massive multitask language understanding in Chinesecited by 436
- 2019Learning Deep Transformer Models for Machine Translationcited by 309
- 2024Scaling Laws for Fine-Grained Mixture of Expertscited by 166
- 2023Beyond Chinchilla-Optimal: Accounting for Inference in Language Model Scaling Lawscited by 155
- 2024Byte Latent Transformer: Patches Scale Better Than Tokenscited by 137
- 2019Large Memory Layers with Product Keyscited by 120
- 2024Scaling Laws with Vocabulary: Larger Models Deserve Larger Vocabulariescited by 113
- 2024Hyper-Connectionscited by 81
- 2019Investigating Prior Knowledge for Challenging Chinese Machine Reading Comprehensioncited by 80
- 2024Understanding Warmup-Stable-Decay Learning Rates: A River Valley Loss Landscape Perspectivecited by 75
- 2024nGPT: Normalized Transformer with Representation Learning on the Hyperspherecited by 68
- 2025SuperBPE: Space Travel for Language Modelscited by 58
- 2023Xiezhi: An Ever-Updating Benchmark for Holistic Domain Knowledge Evaluationcited by 52
- 2025Towards Greater Leverage: Scaling Laws for Efficient Mixture-of-Experts Language Modelscited by 44
- 2025Over-Tokenized Transformer: Vocabulary is Generally Worth Scalingcited by 42
- 2024Memory Layers at Scalecited by 39
- 2024Reconciling Kaplan and Chinchilla Scaling Lawscited by 36
- 2025Pre-training under infinite computecited by 24
- 2024Sigmoid Gating is More Sample Efficient than Softmax Gating in Mixture of Expertscited by 23
- 2025Scaling Embedding Layers in Language Modelscited by 19
- 2024Large Vocabulary Size Improves Large Language Modelscited by 11
- 2025Mixture of Lookup Expertscited by 11
- 2025Scaling Fine-Grained MoE Beyond 50B Parameters: Empirical Evaluation and Practical Insightscited by 2
Built on by
Citing papers per year: 2026: 2.
2 papers.
- 2026xHC: Expanded Hyper-Connectionscited by 3
- 2026MoME: Mixture-of-Memory Embeddings for Context-Aware Sparse Lookupcited by 0
- 2026MoME: Mixture-of-Memory Embeddings for Context-Aware Sparse Lookup14 Sep 2026
- 2026xHC: Expanded Hyper-Connections16 Jul 2026
Similar, not cited
Close in wording, no citation either way
- 2026FastE: Readout-Triggered Token Compression for LLM Embedding Inference
- 2025Helix Parallelism: Rethinking Sharding Strategies for Interactive Multi-Million-Token LLM Decoding
- 2026SMELT: Scaling Laws for Compute-Matched MoE Looped Transformers
- 2026Beyond Token Eviction: Mixed-Dimension Budget Allocation for Efficient KV Cache Compression
- 2026Budgeting Bytes: A Windowed Storage Roofline and Dual-Budget Architecture Ablations for Storage-Bound LLM Decoding ↗
- 2024Efficient LLM Inference using Dynamic Input Pruning and Cache-Aware Masking
- 2025TRACE: Unlocking Effective CXL Bandwidth via Lossless Compression and Precision Scaling
- 2026SUN: Shared Use of Next-token Prediction for Efficient Multi-LLM Disaggregated Serving