STEM: Scaling Transformers with Embedding Modules

• URL: http://arxiv.org/abs/2601.10639v1
• Date: Thu, 15 Jan 2026 18:00:27 GMT
• Title: STEM: Scaling Transformers with Embedding Modules
• Authors: Ranajoy Sadhukhan, Sheng Cao, Harry Dong, Changsheng Zhao, Attiano Purpura-Pontoniere, Yuandong Tian, Zechun Liu, Beidi Chen,
• Abstract summary: We introduce STEM, a static, token-indexed approach that replaces the FFN up-projection with a layer-local embedding lookup.<n>This removes runtime routing, enables CPU offload with asynchronous prefetch, and decouples capacity from both per-token FLOPs and cross-device communication.<n>Overall, STEM is an effective way of scaling parametric memory while providing better interpretability, better training stability and improved efficiency.
• Score: 59.26825251273227
• License: http://creativecommons.org/licenses/by/4.0/
• 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.

Related papers

• MSN: A Memory-based Sparse Activation Scaling Framework for Large-scale Industrial Recommendation [19.132874291460936]
  We propose MSN, a memory-based sparse activation scaling framework for recommendation models.<n> MSN retrieves personalized representations from a large parameterized memory and integrates them into downstream feature interaction modules.<n> MSN consistently improves recommendation performance while maintaining high efficiency.
  arXiv  Detail & Related papers  (2026-02-07T12:43:51Z)
• High-Rank Structured Modulation for Parameter-Efficient Fine-Tuning [57.85676271833619]
  Low-rank Adaptation (LoRA) uses a low-rank update method to simulate full parameter fine-tuning.<n>We present textbfSMoA, a high-rank textbfStructured textbfMOdulation textbfAdapter that uses fewer trainable parameters while maintaining a higher rank.
  arXiv  Detail & Related papers  (2026-01-12T13:06:17Z)
• ExplicitLM: Decoupling Knowledge from Parameters via Explicit Memory Banks [4.099810580680816]
  Large language models suffer from knowledge staleness and lack of interpretability due to implicit knowledge storage.<n>We propose ExplicitLM, a novel architecture featuring a million-scale external memory bank storing human-readable knowledge as token sequences.
  arXiv  Detail & Related papers  (2025-11-03T13:53:19Z)
• TopoSizing: An LLM-aided Framework of Topology-based Understanding and Sizing for AMS Circuits [7.615431299673158]
  Traditional black-box optimization achieves sampling efficiency but lacks circuit understanding.<n>We propose TopoSizing, an end-to-end framework that performs robust circuit understanding directly from raw netlists.
  arXiv  Detail & Related papers  (2025-09-17T16:52:46Z)
• Large-Scale Model Enabled Semantic Communication Based on Robust Knowledge Distillation [45.347078403677216]
  Large-scale models (LSMs) can be an effective framework for semantic representation and understanding.<n>However, their direct deployment is often hindered by high computational complexity and resource requirements.<n>This paper proposes a novel knowledge distillation based semantic communication framework.
  arXiv  Detail & Related papers  (2025-08-04T07:47:18Z)
• Pangu Embedded: An Efficient Dual-system LLM Reasoner with Metacognition [95.54406667705999]
  Pangu Embedded is an efficient Large Language Model (LLM) reasoner developed on Ascend Neural Processing Units (NPUs)<n>It addresses the significant computational costs and inference latency challenges prevalent in existing reasoning-optimized LLMs.<n>It delivers rapid responses and state-of-the-art reasoning quality within a single, unified model architecture.
  arXiv  Detail & Related papers  (2025-05-28T14:03:02Z)
• DSMoE: Matrix-Partitioned Experts with Dynamic Routing for Computation-Efficient Dense LLMs [86.76714527437383]
  This paper proposes DSMoE, a novel approach that achieves sparsification by partitioning pre-trained FFN layers into computational blocks.<n>We implement adaptive expert routing using sigmoid activation and straight-through estimators, enabling tokens to flexibly access different aspects of model knowledge.<n>Experiments on LLaMA models demonstrate that under equivalent computational constraints, DSMoE achieves superior performance compared to existing pruning and MoE approaches.
  arXiv  Detail & Related papers  (2025-02-18T02:37:26Z)
• SHERL: Synthesizing High Accuracy and Efficient Memory for Resource-Limited Transfer Learning [63.93193829913252]
  We propose an innovative METL strategy called SHERL for resource-limited scenarios. In the early route, intermediate outputs are consolidated via an anti-redundancy operation. In the late route, utilizing minimal late pre-trained layers could alleviate the peak demand on memory overhead.
  arXiv  Detail & Related papers  (2024-07-10T10:22:35Z)
• Lightweight and Flexible Deep Equilibrium Learning for CSI Feedback in FDD Massive MIMO [13.856867175477042]
  In frequency-division duplexing (FDD) massive multiple-input multiple-output (MIMO) systems, downlink channel state information (CSI) needs to be sent back to the base station (BS) by the users. We propose a lightweight and flexible deep learning-based CSI feedback approach by capitalizing on deep equilibrium models.
  arXiv  Detail & Related papers  (2022-11-28T05:53:09Z)
• HFedMS: Heterogeneous Federated Learning with Memorable Data Semantics in Industrial Metaverse [49.1501082763252]
  This paper presents HFEDMS for incorporating practical FL into the emerging Industrial Metaverse. It reduces data heterogeneity through dynamic grouping and training mode conversion. Then, it compensates for the forgotten knowledge by fusing compressed historical data semantics. Experiments have been conducted on the streamed non-i.i.d. FEMNIST dataset using 368 simulated devices.
  arXiv  Detail & Related papers  (2022-11-07T04:33:24Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.