mlsanityData to 7 Oct 2026

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

arXiv 2602.00800 ↗First posted 31 Jan 2026cs.LG cs.AIcited by 2

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

  1. 2024Ultra-Sparse Memory Network19 Nov
  2. 2024Mixture of A Million Experts4 Jul
  3. 2021Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity11 Jan

Other references · 47

  1. 2017Attention Is All You Needcited by 59,932
  2. 2015Deep Residual Learning for Image Recognitioncited by 51,809
  3. 2014Adam: A Method for Stochastic Optimizationcited by 36,960
  4. 2021Training Verifiers to Solve Math Word Problemscited by 9,743
  5. 2025Qwen3 Technical Reportcited by 8,195
  6. 2021Evaluating Large Language Models Trained on Codecited by 8,053
  7. 2020Measuring Massive Multitask Language Understandingcited by 7,768
  8. 2020Scaling Laws for Neural Language Modelscited by 6,411
  9. 2021Measuring Mathematical Problem Solving With the MATH Datasetcited by 5,203
  10. 2018Think you have Solved Question Answering? Try ARC, the AI2 Reasoning Challengecited by 4,710
  11. 2019HellaSwag: Can a Machine Really Finish Your Sentence?cited by 4,018
  12. 2021Program Synthesis with Large Language Modelscited by 3,509
35 more
  1. 2017Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layercited by 3,413
  2. 2022Training Compute-Optimal Large Language Modelscited by 2,927
  3. 2017TriviaQA: A Large Scale Distantly Supervised Challenge Dataset for Reading Comprehensioncited by 2,895
  4. 2023GPQA: A Graduate-Level Google-Proof Q&A Benchmarkcited by 2,748
  5. 2024LiveCodeBench: Holistic and Contamination Free Evaluation of Large Language Models for Codecited by 1,906
  6. 2019Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelismcited by 1,829
  7. 2022Challenging BIG-Bench Tasks and Whether Chain-of-Thought Can Solve Themcited by 1,767
  8. 2023SGLang: Efficient Execution of Structured Language Model Programscited by 1,107
  9. 2024DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Modelcited by 1,093
  10. 2019DROP: A Reading Comprehension Benchmark Requiring Discrete Reasoning Over Paragraphscited by 883
  11. 2024DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Modelscited by 834
  12. 2023C-Eval: A Multi-Level Multi-Discipline Chinese Evaluation Suite for Foundation Modelscited by 694
  13. 2023CMMLU: Measuring massive multitask language understanding in Chinesecited by 436
  14. 2019Learning Deep Transformer Models for Machine Translationcited by 309
  15. 2024Scaling Laws for Fine-Grained Mixture of Expertscited by 166
  16. 2023Beyond Chinchilla-Optimal: Accounting for Inference in Language Model Scaling Lawscited by 155
  17. 2024Byte Latent Transformer: Patches Scale Better Than Tokenscited by 137
  18. 2019Large Memory Layers with Product Keyscited by 120
  19. 2024Scaling Laws with Vocabulary: Larger Models Deserve Larger Vocabulariescited by 113
  20. 2024Hyper-Connectionscited by 81
  21. 2019Investigating Prior Knowledge for Challenging Chinese Machine Reading Comprehensioncited by 80
  22. 2024Understanding Warmup-Stable-Decay Learning Rates: A River Valley Loss Landscape Perspectivecited by 75
  23. 2024nGPT: Normalized Transformer with Representation Learning on the Hyperspherecited by 68
  24. 2025SuperBPE: Space Travel for Language Modelscited by 58
  25. 2023Xiezhi: An Ever-Updating Benchmark for Holistic Domain Knowledge Evaluationcited by 52
  26. 2025Towards Greater Leverage: Scaling Laws for Efficient Mixture-of-Experts Language Modelscited by 44
  27. 2025Over-Tokenized Transformer: Vocabulary is Generally Worth Scalingcited by 42
  28. 2024Memory Layers at Scalecited by 39
  29. 2024Reconciling Kaplan and Chinchilla Scaling Lawscited by 36
  30. 2025Pre-training under infinite computecited by 24
  31. 2024Sigmoid Gating is More Sample Efficient than Softmax Gating in Mixture of Expertscited by 23
  32. 2025Scaling Embedding Layers in Language Modelscited by 19
  33. 2024Large Vocabulary Size Improves Large Language Modelscited by 11
  34. 2025Mixture of Lookup Expertscited by 11
  35. 2025Scaling Fine-Grained MoE Beyond 50B Parameters: Empirical Evaluation and Practical Insightscited by 2

Built on by

Cited by year (2026 to 7 Oct)

Citing papers per year: 2026: 2.

2 papers.

Similar, not cited

Close in wording, no citation either way

  1. 2026FastE: Readout-Triggered Token Compression for LLM Embedding Inference
  2. 2025Helix Parallelism: Rethinking Sharding Strategies for Interactive Multi-Million-Token LLM Decoding
  3. 2026SMELT: Scaling Laws for Compute-Matched MoE Looped Transformers
  4. 2026Beyond Token Eviction: Mixed-Dimension Budget Allocation for Efficient KV Cache Compression
  5. 2026Budgeting Bytes: A Windowed Storage Roofline and Dual-Budget Architecture Ablations for Storage-Bound LLM Decoding ↗
  6. 2024Efficient LLM Inference using Dynamic Input Pruning and Cache-Aware Masking
  7. 2025TRACE: Unlocking Effective CXL Bandwidth via Lossless Compression and Precision Scaling
  8. 2026SUN: Shared Use of Next-token Prediction for Efficient Multi-LLM Disaggregated Serving