Source-linked AI summary
STEM: Scaling Transformers with Embedding Modules
Ranajoy Sadhukhan, Sheng Cao, Harry Dong, Changsheng Zhao, Attiano Purpura-Pontoniere, Yuandong Tian, Zechun Liu, Beidi Chen
TL;DR
Fine-grained sparsity seeks greater parametric capacity without proportional per-token computation, but faces instability, load-balancing, and communication challenges. STEM replaces the FFN up-projection with static token-indexed layer-local embeddings, improving efficiency, stability, accuracy, interpretability, and long-context capacity. Across 350M and 1B models, it improves accuracy by up to ∼3–4% while reducing per-token FLOPs and parameter accesses by up to one-third.
Problem
Fine-grained sparsity offers higher parametric capacity without proportional per-token compute but faces training instability, load imbalance, and communication overhead.
Method
STEM replaces the FFN up-projection with a token-specific vector from a layer-local embedding table while preserving the gate and down-projection paths.
Results
STEM improves downstream accuracy by up to ∼3–4% while reducing per-token FLOPs and parameter accesses by up to one-third, with stable training and stronger long-context performance.
Takeaways & Limitations
STEM scales parametric memory with improved efficiency, interpretability, and practical test-time capacity scaling as sequence length grows.
Abstract
from arXiv · showhide
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.
1 Introduction
STEM addresses instability, load-balancing, and communication challenges in fine-grained sparsity with a static token-indexed design that replaces only the FFN up-projection. It improves efficiency, stability, accuracy, interpretability, and long-context capacity through token-specific embeddings.
- Motivation: Fine-grained sparsity can leave experts under-trained, destabilize optimization, and increase communication overhead through fragmented expert exchanges.Load balancing can help but may interfere with the primary objective if not carefully tuned.
- Method: STEM replaces only the gated FFN up-projection with a token-specific vector from a layer-local embedding table while preserving shared gate and down-projection paths.The static token-indexed design avoids runtime routing and supports predictable computation.
- Training stability: STEM remains stable despite extreme sparsity and does not exhibit the loss spikes observed in MoE models.The comparison is reported in Figure 5a.
- Knowledge capacity: STEM embeddings have large angular spread, reducing representational interference and improving addressability for parametric memory.The resulting storage capacity is associated with approximately 9–10% gains over dense baselines on ARC-Challenge and OpenBookQA.
- Interpretability: Token-indexed embeddings make individual micro-experts interpretable and allow output distributions to be steered by swapping table indices without changing the input text.These interventions expose localized, modular, editable, and attributable factual knowledge.
- Long-context inference: As context length grows, STEM activates more distinct parameters, strengthening long-context performance and enabling practical test-time capacity scaling.On Needle-in-a-Haystack, the gap over the dense baseline increases from 8.4% to 13%.
- Efficiency and accuracy: STEM improves downstream accuracy by up to ∼3–4% while reducing per-token FLOPs and parameter accesses by up to one-third.Reported evaluations use 350M and 1B model variants and include knowledge and reasoning benchmarks such as GSM8K and MMLU.
2 Background
The background contrasts dense FFNs, routed MoEs, hash-layer MoEs, and embedding-based alternatives. It emphasizes that finer expert granularity increases communication latency and load imbalance, while static or token-indexed approaches trade routing flexibility for predictable execution and interpretability.
- SwiGLU FFN: A decoder-only Transformer’s SwiGLU FFN uses gate and up projections followed by elementwise multiplication and a down projection.The formulation is defined for layer ℓ with model width d and feed-forward width dff.
- Mixture-of-Experts: MoE replaces a dense FFN with K expert FFNs and a router that selects a small set of top-r experts using mixture weights.The selected experts and weights depend on the layer’s hidden state.
- Hash-layer MoE: Hash-layer MoE removes trainable routing and auxiliary losses by fixing a balanced mapping from token IDs to experts.This preserves sparse selection while making the assignment static.
- Communication overhead: Increasing expert granularity fragments peer-to-peer exchanges into many small payloads, making communication latency-dominated and reducing overlap with computation.This raises end-to-end layer latency under expert parallelism.
- Scaling expert count: Word-level expert granularity raises load skew because Zipfian token frequencies concentrate traffic in a few experts and leave many rarely activated.The resulting imbalance harms statistical and systems efficiency through stragglers, padding or drops, and synchronization costs.
- Embedding-based alternatives: Per Layer Embedding shares the gate and down projections across expert subnetworks, stores tables in node-local CPU memory, and prefetches them as needed.Unlike MoWE, its tables are not sharded across multiple devices, avoiding high all-to-all communication.
3 Method
STEM replaces the FFN up-projection with layer-local token embeddings while retaining dense gating and down-projection, improving efficiency, interpretability, and capacity scaling. Its token-indexed design supports knowledge editing and long-context parameter growth without runtime routing.
- STEM design: STEM replaces only the FFN up-projection with a layer-local token embedding table, while preserving shared gate and down-projection paths.The design builds on layer-wise embeddings but differs from PLE by dispensing with the additional regular FFN block and using full FFN-dimensional embeddings.
- STEM design: Replacing the gate projection hurts downstream performance, whereas replacing the up-projection improves it because gating provides context-dependent modulation.The down-projection is retained because replacing it would break the model’s forward path.
- Information storage and interpretability: STEM embeddings exhibit larger angular spread than standard FFN address vectors, reducing redundancy and potentially improving knowledge attribution and information storage.The embeddings serve as token-specific address vectors modulated by the context-dependent gate.
- Information storage and interpretability: Minimal modifications to STEM embeddings can steer generation without changing the input text, enabling token-level knowledge attribution and editing without additional interpretability computation.The demonstrated mechanism replaces embeddings associated with source tokens to produce output consistent with target entities.
- Efficiency and context scaling: Token-indexed embeddings enable CPU offloading and prefetching without routing logic, while longer contexts activate more distinct parameters and support predictable test-time capacity scaling.Repeated tokens reuse vectors, whereas novel tokens activate new vectors; active parameters therefore grow with context length without increasing per-token FLOPs.
- Knowledge editing: STEM supports length-mismatched knowledge editing through replacement schemes for source and target entities with different tokenization lengths.When the source span is longer, the paper considers padding and related replacement strategies; the approach also transfers across geographic entity types.
4 Experiments
STEM is evaluated against dense and Hash-MoE baselines across model scales, training settings, downstream tasks, and sparsity levels. It improves efficiency, training stability, downstream performance, and training ROI, with gains depending on STEM placement and replacement ratio.
- Experimental Setup: STEM is evaluated at 350M and 1B scales against dense and Hash-MoE baselines under controlled training conditions.The evaluation covers pretraining, mid-training, context-length extension, and downstream reasoning and knowledge tasks.
- Training Stability and Scaling: STEM achieves lower training losses with fewer training FLOPs and avoids the loss spikes observed in Hash layer MoE.The 350M STEM model remains stable, and increasing STEM layers can reduce loss at fewer training FLOPs.
- Downstream Evaluation: STEM improves knowledge-intensive performance, with larger gains on ARC-Challenge and OpenBookQA and more modest gains on other tasks.The knowledge-task improvements become more pronounced as the number of replaced FFN layers increases.
- Downstream Evaluation: The 1B STEM model continues to outperform the dense baseline after midtraining, including on GSM8K and MMLU reasoning and knowledge retrieval evaluations.These results extend the observed downstream improvements beyond pretrained checkpoints.
- Ablations: Training ROI rises from 1.08x for STEM-1/3 to 1.20x for STEM-1/2 and 1.33x for STEM-full relative to the baseline.Increasing replacement reduces training FLOPs, while performance gains slow beyond half replacement.
- Ablations: Replacing the up-projection improves performance, whereas replacing the gate-projection underperforms the dense baseline.The gate must preserve context-dependent modulation, while the up-projection supplies the token-indexed address.
5 STEM Characteristics
STEM embeddings form widely separated, token-indexed addresses that support memory storage and interpretable factual editing. Swapping layer-local token vectors can reversibly shift factual predictions while preserving other parameters.
- Embedding Geometry: STEM embeddings exhibit low pairwise cosine similarity, indicating large angular spread and potentially reduced interference among stored items.The geometry is measured using pairwise cosine similarities of unit-normalized embedding vectors.
- Embedding Geometry: The embedding geometry is compared with up-projection output spaces and down-projection input hidden states in the 1B model.Figure 6 reports these cosine-similarity distributions across sampled layers and representation spaces.
- Knowledge Editing: Swapping eSpain,ℓ with eGermany,ℓ shifts probability mass from Madrid to Berlin while leaving the original prompt unchanged.The intervention is reversible and changes the top next-token distribution without modifying other model parameters.
- Interpretability: STEM exposes token-indexed, layer-local parameters that function as interpretable FFN addresses for factual knowledge editing.The design supports direct intervention on the vector associated with a token at a particular layer.
6 Related Works
Related approaches use sparse computation or token-indexed memory to increase capacity, but they retain important optimization, communication, lookup, or expressivity trade-offs. STEM is positioned against these alternatives through its static token-indexed design.
- Mixture-of-Experts: MoE models increase parametric capacity at near-constant FLOPs but depend on load balancing, auxiliary objectives, and systems solutions for stable operation.Fine-grained expert routing can leave experts under-trained and introduce instability or communication overhead.
- Key-Value Memory: PKM scales parametric memory with efficient top-k selection and avoids cross-device communication, but incurs high inference lookup costs and under-training of value memory.Its product-space key design differs from STEM’s direct token-indexed addressing.
- Per-Layer Embeddings: Gemma-3n’s per-layer embeddings complement rather than replace FFN modules and use lower-dimensional tables for output modulation.Those tables are stored outside GPU HBM to support larger batches and fast prefetching.
7 Conclusion
STEM replaces FFN up-projections with layer-local, token-indexed embedding lookups, decoupling parametric capacity from per-token compute and cross-device communication. The resulting system combines efficiency, stable sparse training, improved accuracy, interpretable editing, and long-context capacity scaling.
- Design and Efficiency: STEM replaces the FFN up-projection with a layer-local embedding lookup while retaining dense gate and down-projection paths.This static, token-indexed design supports CPU offload with asynchronous prefetch.
- Results: STEM lowers per-token FLOPs and parameter accesses while improving accuracy over dense baselines and remaining stable under extreme sparsity.Its embedding space also has higher effective memory capacity through large angular spread.
- Long-Context Scaling: Longer contexts activate more distinct parameters, providing practical test-time capacity scaling.The conclusion connects this behavior to STEM’s token-indexed parametric memory.
A.1 Additional Benchmarks
STEM improves contextual reasoning relative to dense baselines across challenging reasoning benchmarks and context-length ranges. The evaluation covers multi-step, long-narrative, multi-passage, and code-understanding tasks.
- Contextual reasoning: STEM consistently outperforms the dense baseline on BBH, MuSR, and LongBench multi-hop and code-understanding tasks across all context-length ranges.BBH tests multi-step and compositional reasoning; MuSR requires tracking entities and constraints over long narratives; LongBench evaluates multi-passage and code reasoning.
- Evaluation setup: The evaluation uses 1B-scale models and averages LongBench scores over tasks within each context-length range.
- Contextual reasoning: The results indicate that STEM does not impair contextual reasoning and can improve it relative to the dense baseline.
A.2 Additional Long-context Evaluation
Additional LongBench evaluation shows that STEM maintains or improves performance relative to a dense baseline across context-length ranges. The evaluation averages six task categories by context-length regime.
- Long-context performance: The 1B STEM model consistently matches or outperforms the 1B dense baseline across all LongBench context-length ranges.
- Benchmark coverage: LongBench covers single- and multi-document question answering, summarization, few-shot learning, synthetic tasks, and code completion.
- Evaluation setup: LongBench scores are averaged across tasks within each context-length range.
A.3 Additional Architecture Ablation Study
The architecture ablation study examines alternatives to the core STEM design. One hybrid variant retains the FFN up-projection and adds token-specific modulation.
- Architecture alternatives: The study evaluates additional architecture alternatives to assess the optimality of STEM’s design principles.
- Design trade-offs: STEM uses fewer active parameters and FLOPs per token, but its architectural bias can reduce contextual learning ability.
- Hybrid variant: STEM† retains the FFN up-projection and adds an additive token-specific modulation.