STEM: Scaling Transformers with Embedding Modules

TL;DR

STEM tackles the challenge of scaling transformer capacity without proportional compute by introducing static, token-indexed embedding modules that replace only the FFN up-projection while preserving the gate and down-projection. This yields lower per-token FLOPs, CPU offload via asynchronous prefetch, and a memory-augmented, interpretable parameter space with embeddings tied to specific tokens. The approach delivers stable training under high sparsity, improves knowledge-dependent benchmarks, strengthens long-context performance, and enables simple knowledge editing by swapping per-layer embeddings. Taken together, STEM provides a practical, interpretable route to larger parametric memory with improved efficiency and stability for large-scale language models.

Abstract

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.

Figures (7)

• Figure 1: (a) Validation PPL vs. training tokens for 1B STEM vs. dense; (b) Needle-in-a-Haystack at 8k/16k/32k; (c) STEM layer: embedding tables offloaded to CPU and token-indexed ones are prefetched to GPU.
• Figure 2: Schematics of (a) SwiGLU FFN, (b) MoE FFN, and (c) STEM with a single prefetched token embedding. In MoE FFN, the full FFN module is considered as one expert.
• Figure 3: Knowledge injection/edit demonstration. Input text remains the same (Country: Spain), but internally the PLE used for the token is swapped from Spain to Germany, flipping the generated capital/paragraph from Madrid to Berlin.
• Figure 4: STEM-based knowledge editing schemes for length-mismatched source ($n_s$) and target ($n_t$) entity tokenizations.
• Figure 5: (a) Training Stability. Unlike Hash layer MoE, the 350M STEM model does not show any training loss spikes. (b) Performance scaling with more STEM layers. With more STEM layers, a lower training loss can be achieved at fewer training FLOPs.