# STEM: Scaling Transformers with Embedding Modules

arXiv: https://arxiv.org/abs/2601.10639 · HTML: https://mlsanity.com/paper/2601.10639/

## Key facts

- Cited by 12 papers in the map (data through 14 Oct 2026).
- Builds on 30 of its 41 references that are in the map.
- First posted on arXiv 15 Jan 2026.

## Builds on (most cited first, 30 of 30)

- [Training Verifiers to Solve Math Word Problems](https://mlsanity.com/paper/2110.14168/) (2021)
- [Measuring Massive Multitask Language Understanding](https://mlsanity.com/paper/2009.03300/) (2020)
- [Scaling Laws for Neural Language Models](https://mlsanity.com/paper/2001.08361/) (2020)
- [Think you have Solved Question Answering? Try ARC, the AI2 Reasoning Challenge](https://mlsanity.com/paper/1803.05457/) (2018)
- [HellaSwag: Can a Machine Really Finish Your Sentence?](https://mlsanity.com/paper/1905.07830/) (2019)
- [Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer](https://mlsanity.com/paper/1701.06538/) (2017)
- [Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity](https://mlsanity.com/paper/2101.03961/) (2021)
- [PIQA: Reasoning about Physical Commonsense in Natural Language](https://mlsanity.com/paper/1911.11641/) (2019)
- [Locating and Editing Factual Associations in GPT](https://mlsanity.com/paper/2202.05262/) (2022)
- [BoolQ: Exploring the Surprising Difficulty of Natural Yes/No Questions](https://mlsanity.com/paper/1905.10044/) (2019)
- [Can a Suit of Armor Conduct Electricity? A New Dataset for Open Book Question Answering](https://mlsanity.com/paper/1809.02789/) (2018)
- [Challenging BIG-Bench Tasks and Whether Chain-of-Thought Can Solve Them](https://mlsanity.com/paper/2210.09261/) (2022)
- [GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding](https://mlsanity.com/paper/2006.16668/) (2020)
- [Sparse Autoencoders Find Highly Interpretable Features in Language Models](https://mlsanity.com/paper/2309.08600/) (2023)
- [Transformer Feed-Forward Layers Are Key-Value Memories](https://mlsanity.com/paper/2012.14913/) (2020)
- [Enabling Large Language Models to Generate Text with Citations](https://mlsanity.com/paper/2305.14627/) (2023)
- [DeepSpeed-MoE: Advancing Mixture-of-Experts Inference and Training to Power Next-Generation AI Scale](https://mlsanity.com/paper/2201.05596/) (2022)
- [Hash Layers For Large Sparse Models](https://mlsanity.com/paper/2106.04426/) (2021)
- [MobileLLM: Optimizing Sub-billion Parameter Language Models for On-Device Use Cases](https://mlsanity.com/paper/2402.14905/) (2024)
- [Auxiliary-Loss-Free Load Balancing Strategy for Mixture-of-Experts](https://mlsanity.com/paper/2408.15664/) (2024)
- [On the Representation Collapse of Sparse Mixture of Experts](https://mlsanity.com/paper/2204.09179/) (2022)
- [Large Memory Layers with Product Keys](https://mlsanity.com/paper/1907.05242/) (2019)
- [Mixture of A Million Experts](https://mlsanity.com/paper/2407.04153/) (2024)
- [Demons in the Detail: On Implementing Load Balancing Loss for Training Specialized Mixture-of-Expert Models](https://mlsanity.com/paper/2501.11873/) (2025)
- [Nemotron-CC-Math: A 133 Billion-Token-Scale High Quality Math Pretraining Dataset](https://mlsanity.com/paper/2508.15096/) (2025)
- [MoETuner: Optimized Mixture of Expert Serving with Balanced Expert Placement and Token Routing](https://mlsanity.com/paper/2502.06643/) (2025)
- [Memory Augmented Language Models through Mixture of Word Experts](https://mlsanity.com/paper/2311.10768/) (2023)
- [Semantic Parallelism: Redefining Efficient MoE Inference via Model-Data Co-Scheduling](https://mlsanity.com/paper/2503.04398/) (2025)
- [The power of fine-grained experts: Granularity boosts expressivity in Mixture of Experts](https://mlsanity.com/paper/2505.06839/) (2025)
- [Pro-Prophet: A Systematic Load Balancing Method for Efficient Parallel Training of Large-scale MoE Models](https://mlsanity.com/paper/2411.10003/) (2024)

## Built on by (most cited first, 12 of 12)

- [On the Design of Qwen3.8-Next Architecture: Evaluation, Efficiency, and Training Stability](https://mlsanity.com/paper/2608.30320/) (2026)
- [Scaling Embeddings Outperforms Scaling Experts in Language Models](https://mlsanity.com/paper/2601.21204/) (2026)
- [MeKi: Memory-based Expert Knowledge Injection for Efficient LLM Scaling](https://mlsanity.com/paper/2602.03359/) (2026)
- [Memory Grafting: Scaling Language Model Pre-training via Offline Conditional Memory](https://mlsanity.com/paper/2605.20948/) (2026)
- [When Data Is Scarce: Scaling Sparse Language Models with Repeated Training](https://mlsanity.com/paper/2606.01155/) (2026)
- [Beyond N-gram: Data-Aware X-GRAM Extraction for Efficient Embedding Parameter Scaling](https://mlsanity.com/paper/2604.21724/) (2026)
- [Efficient foundation decoders for fault-tolerant quantum computing](https://mlsanity.com/paper/2606.27119/) (2026)
- [TIDE: Every Layer Knows the Token Beneath the Context](https://mlsanity.com/paper/2605.06216/) (2026)
- [More Than Memory: Task-Conditioned Signed FFN Writes in Long-Context Retrieval](https://mlsanity.com/paper/2607.16254/) (2026)
- [MoME: Mixture-of-Memory Embeddings for Context-Aware Sparse Lookup](https://mlsanity.com/paper/2609.15126/) (2026)
- [FactorEngram: Factorized N-gram Memory with Basis-Level Gating for Language Models](https://mlsanity.com/paper/2609.35578/) (2026)
- [STAPLE: Static Sparse Upcycling of Pretrained Checkpoints via Embedding Modules](https://mlsanity.com/paper/2610.12599/) (2026)

## Abstract (arXiv)

Fine-grained sparsity promises higher parametric capacity without proportional per-token compute, but often suffers from training instability, load balancing, and communication overhead. We introduce STEM (Scaling Transformers with Embedding Modules), a static, token-indexed approach that replaces the FFN up-projection with a layer-local embedding lookup while keeping the gate and down-projection dense. This removes runtime routing, enables CPU offload with asynchronous prefetch, and decouples capacity from both per-token FLOPs and cross-device communication. Empirically, STEM trains stably despite extreme sparsity. It improves downstream performance over dense baselines while reducing per-token FLOPs and parameter accesses (eliminating roughly one-third of FFN parameters). STEM learns embedding spaces with large angular spread which enhances its knowledge storage capacity. More interestingly, this enhanced knowledge capacity comes with better interpretability. The token-indexed nature of STEM embeddings allows simple ways to perform knowledge editing and knowledge injection in an interpretable manner without any intervention in the input text or additional computation. In addition, STEM strengthens long-context performance: as sequence length grows, more distinct parameters are activated, yielding practical test-time capacity scaling. Across 350M and 1B model scales, STEM delivers up to ~3--4% accuracy improvements overall, with notable gains on knowledge and reasoning-heavy benchmarks (ARC-Challenge, OpenBookQA, GSM8K, MMLU). Overall, STEM is an effective way of scaling parametric memory while providing better interpretability, better training stability and improved efficiency.