STEM: SCALING TRANSFORMERS WITH EMBEDDING MODULES
Authors: Ranajoy Sadhukhan
Sheng Cao
Harry Dong
Changsheng Zhao
Attiano Purpura-Pontoniere
Yuandong Tian
Zechun Liu
Beidi Chen
Venue: ICLR 2026 Poster (2026)
Keywords: Sparse Transformer
Parametric scaling
Embedding Layers
Foundation Models
Pre-training
Model Architecture
Primary area: foundation or frontier models, including LLMs
TL;DR: STEM replaces each FFN up-projection with a per-layer embedding lookup to scale parametric capacity without increasing per-token compute or cross-device communication, yielding FLOP-efficient performance gains.

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 \textbf{STEM} (\emph{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 it knowledge storage capacity. 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 $\sim$3--4\% improvements in average downstream performance, 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 remaining simpler to train and deploy than existing fine-grained sparse models.

Published as a conference paper at ICLR 2026
STEM: SCALING TRANSFORMERS WITH EMBEDDING MODULES
Ranajoy Sadhukhan', Sheng Cao Harry Dong', Changsheng Zhao Attiano Purpura-Pontoniers Yuandong Tian Zechun Liu Beidi Chen
CCannegee Mellon University
Meta Al
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 Em- bedding Modules), 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 asyn- chronous prefetch, and decouples capacity from both per-token FLOPs and cross- device communication. Empirically, STEM trains stably despite extreme spar- sity. It improves downstream performance over dense baselines while reducing per-token FLOPs and parameter accesses (eliminating roughly one-third FFN parameters). STEM learns embedding spaces with large angular spread which enhances its knowledge storage capacity. In addition, STEM strengthens long- context performance: sequence length grows, tivated, yielding practical test-time capacity scaling. Across 350M and IB model scales, STEM delivers up ~3-4% improvements in average downstream perfor- mance, 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 remaining simpler to train and deploy than existing fine-grained sparse models.
distinct parameters
INTRODUCTION
14.50
Medel Architecture
79
$4.25
24.00
60
Salidation
3.375
50
13.00
12.75
Tokens
x20
Context Length
(a) Validation PPI. scores
(b) Context Length Scalability
(c) STEM layer
Figure 1: (a) Validation PPL vs. training tokens for IB STEM dense; (b) N at 8k/16k/32k;; (c) STEM layer: embedding tables offloaded to CPU and token-indexed ones are prefetched to GPU.
Sparse computation is a key mechanism for realizing the benefits predicted by parameter-scaling laws (Kaplan al. 2020 Hoffmann al. 2022) without proportionally increasing per-token com- pute. In particular, Mixture-of oxfeExssss (MoE) (Shazeer et al. 2017 Artetxe et al., 2022 Fedus et al,
2022) models have been adopted in several frontier LLMs (Team 2025b al Dai et al. (2224) because they raise parametric capacity at roughly constant activated FLOPs by sparsely activating a small
subset of experts per token. Recent work (Boix-Adsera & Rigollet 2025
He 2024
Databricks

---

Published as a conference paper at ICLR 2026
2024 Dai et al.] 2024) further advocate for finer-grained sparsity that employs large number of eicro-experts to achieve better expressivity, enhanced knowledge storing capacity, and favorable efficiency metrics.
However, finer granularity introduces nontrivial challenges in both optimization and systems. On the training side, even large fraction of experts can remain under-trained (Huang et al. 2025) due to a highly non- non-uniform routing and result in training instability. While load- balancinnn objectives (Shazeer et al.) 2017] Fedus et 2022; Lepikhin et al.) 2020) can address these issues, they may interfere with the primary objective if not carefully tuned (Dai et al. 2024; Qiu et al. [025 Go & Mahajan 2025). On the systems side, increasing the number of experts typically raises the number of all-to-all messages while shrinking message sizes, degrading bandwidth utilization and ampli- fying communication overhead (Huang et al. 2024; Li et al. 2025b). Finer granularity can also reduce parameter-access locality and degrade kernel efficiency when expert subnetworks become too small for dense linear-algebra kernels to reach high occupancy, yielding suboptimal end-to-end performance. To harness the full potential of fine-grained sparsity, we require: (a) stable optimiza- tion, (b) broad expert utilization (each micro-expert learns useful representations), and (c) negligible expert retrieval latency and communication overhead.
We identify static sparsity as a potential solution to achieve these desired properties. Static sparsity keeps the compute path predictable (no runtime routing latency), enables prefetch and CPU offload- ing (removing the need for inter-node communication). Recently, static sparsity via token-indexed routing has emerged strong performance guarantees. However, such token-based selection strategy lacks context adap- tivity. If applied naively, can reduce the expressivity of the model and degrade quality despite more parameters. Our ablation study in sec. 4.4.3) 4.4.3 highlights the criticality of selecting the suitable module for sparsification.
promising direction (Roller al.. 2021] Google DeepMind (024) with
Based on these observations, we introduce STEM, a static. token-indexed. fine-graines mechanism that replaces only the up-projection in gated FFNs with token-specific vector retrieved from a layer-local embedding table. The gating and down-projection paths are preserved and shared across tokens. We observe that STEM achieves the following:
Better Training Stability: Despite being extremely sparse, STEM does not exhibit any training in- stability issues as usually seen MoE models. Figure 3a shows that unlike MoE models, STEM does not exhibit any loss spikes.
Improved Performance with Larger Knowledge Capacity: STEM learns a representation space for the embeddings that is conducive to better information storage. The learned embeddings exhibit large angular spread (i.e., low pairwise cosine similarity), which reduces representational interfer- ence and improves addressability of the parametric memory. As result, it effectively increases the distinct "slots" available for storing and retrieving information. In our downstream evaluation benchmark, STEM consistently outperforms the dense baseline on snowledge-intensing tasks like, ARC-Challenge (Clark et al. 2018), and OpenBookQA (Mihaylov et al. 2018) by large margins (~9-10%)
Improved Long-context Inference: During long-context inference, STEM activates more distinct pa- rameters as sequence length grows, yielding test-time capacity scaling. As shown in Figure [b] the benefits strengthen with context: on Needle-in-a-Haasttad (NIAH) (Kamradl 2024), the gap over the dense baseline increases from 8.4% to 13%
Training and Inference-ting efficiency: STEM reduces both FLOPs as well as parameter loading cost by eliminating one-third of the parameters in FFN layers. Consequently, it is strictly more efficient during both computation-intensive training and prefilling, as well as in memory-intensive decoding.
We benchmark STEM against the dense baseline with 350M MobileLLM (Liu et al., 2024 and Llama3.2-1B Meta All 2024) model variants. Additionally, we compare with Hash Layer MoEs with the same lotal parameter count. We report results on standard downstream suites across pre- training, mid-training, and context-length extension. STEM improves downstream accuracy by up to ~3-4% while reducing per-token FLOPs and parameter accesses by up to one-third. It also strengthens knowledge retrieval and mathematical reasoning, with gains on GSM8K (Cobbe et al.]
2

---

Published as a conference paper at ICLR 2026
Toyut
Inpun X
Irput xiX
Rouler (lpper)
(LU)(W)))
SELDWWOOD
CPU Endweling Take
W
Expert Ji
Espert
fa
My
W
WHO
Output
Oungur
SWAGLU
MoE
Figure 2: Schematics of (a) SwiGLU FFN, (b) MoE FFN, and (c) STEM with single prefetched token embedding. In MoE FFN, the full FFN module is considered expert.
2021 and MMLU (Hendrycks al. 2021), and shows pronounced improvements Needle-in-a- Haystack Kamradt 2024) at longer contexts.
METHOD
2.1 BACKGROUND
Consider a decoder-only transformer with N layers, vocabulary size V, model width d, and feed- forward width dæ. For given layer the SwiGLU feed-forward block uses a gate projection W Rduxd E [L...., V} denote the vocabulary id the current token, and the corresponding input hidden state of the gth FFN layer is given by Rd. Then the transformation in the FFN layer
up projection W E Rd and down projection W gaxda, Consider,
Ye (SLLU( O (Wxxx)),
where @ denotes elementwise multiplication.
-ixture-of-Expertss (MoE). In MoE, a dense FFN is replaced by K expert FENs {fex and router (xxe) that selects Te(x;) (top-r experts) with mixture weights (xee (Artetxe al. 2022 Fedus et al. 2022). With SwiGLU experts,
fex((x)) = W (SiLU(WZ,,)) O (Wi,x))),
W E Rdxdr
the layer output is
(xx( (x,)
(2)
kx)2(x2)
Token-indexed ceaturee-ffExpeette To eliminate the routing parameters and auxiliary routing loss functions, (Roller et al.) 2021) fixed mapping from input token ids to experts based on random and balanced hash functions. Consequently, the FFN output is computed as,
Ye
Σ kthash(t)
(xxe
2.2 STEM
Unlike MoE alternatives, STEM only replaces the dense up-projection in the SwiGLU FFN with token-indexed vector looked up from per- per-layer table. For layer &, let U2 E RV Xdd be the embedding table. Given input X ERdd the STEM layer computes
Ye (SLLU(W)))) C U.(t)).
(4)
where Uz[t] e 4 is the row of U, corresponding token t and @ denotes elementwise multipli- cation. We provide a simple schematic diagram for dense baseline (SwiGLU FFN), MoE and STEM
in Fig 2
3

---

Published as a conference paper at ICLR 2026
Table 1: Theoretical efficiency for each decoder FFN layer when replacing the FFN up-projection with token-indenn indexed STEM embedding table. We assume SwiGLU, ignore biases, and count ele- mentwise ops as O(DL).
FEN
STEM
Savings (A)
Prefill training (batch B, sequence length I.)
FLOPs
B(3dgdL + dæ B(2dgdL daL) B(ddaL)
Communication
uniq( BL
Decoding (per step. batch size B) 3dd
Parameter loading cost
2dd Bandd
ddæ
Communication
Notation: model width; D: FEN hidden size; L: context length; Lunig; number of unique tokens in the L-token context; Busi:: number of unique tokens across the batch at decode step (S. B); uniq(BL) number of unique tokens across the BL tokens in training batch.
Notes: Training multiplies both FLOP counts z the usual forward+backward factor, but the saving AFLOPs dDT. remains. Communication doubles during training as gradients the STEM embeddings transferred back to CPU for optimizer update.
2.3 STEM
STEM uses strictly fewer active parameters, and FLOPs for each token. And because of the archi- tectural bias, STEM susceptible to some loss of contextual learning ability. We also introduce hybrid variant of STEM, which retains the up projection matrix in FFN, but complements with an additive token-specific modulation. Concretely, the new variant STEM computes the FFN output as follows.
YE ((LLU((())))) O (WTxe + U2l())
(5)
ANALYSIS
3.1 EFFICIENCY
STEM improves both computation and memory access. During compute-intensive phases (train- ing and prefill), replacing the FFN up-projection with token-indexed embeddings reduces the per- layer FLOPs. During memory-intensive decoding, lowers parameter traffic relative to dense up-projection. Table summarizss the per-layer counts and the resulting savings. Below present a simple theoretical analysis of the training and inference efficiency for single decoder layer.
Training efficiency. Consider a batch of sequences with sequence length L, hidden width d, and FFN hidden size dg. Ignoring elementwise ops and biases, the per-layer training FLOPs (forward backward + weight gradients) can be written as
I p.we
(4 +2 3Ld de),
Estem BALLd + +L 2Ld de
Atta
MLP
The per-layer FLOPs reduction of STEM therefore
A Frain
baaee
I. jastom
BLd dH,
and the corresponding saving fraction is
A Firain
dg
saving fraction
E bassc
4d 2T. + 3d
Plugging in the architecture hyperparameters for each Qwen2.5 model yields saving fractions of 21.7% for Qwen2.5-1.5B, 22.8% for Qwen2.5-3B, 23.9% for Qwen2.5-7B, 19.7% for Qwen2.5- 14B, and 24.8% for Qwen2.5-32B.
4

---

Published as a conference paper at ICLR 2026
Inference efficiency. Prefill efficiency closely matches training efficiency because both compute-bound. In contrast, decoding is primarily memory-bound: the dominant cost loading parameters and KV cache rather than doing FLOPs. For batch size B and context length L, can write the per-layer memory access cost as M base = ( (dd + 2Ld + 3d dit):
Mee
B 2Ld + 4d + 2d dp
KV cache
projection params
The reduction in parameter loading cost
AMee M
dec
-d da.
so the saving fraction is
A
dg
saving fraction
Mass
4d + 2L + 3dg
which matches the FLOPs saving factor during training and prefill. As the batch size grows, the linear layers become increasingly compute-bound bound, and STEM's per-layer FLOPs reduction ensures that this efficiency gain is sustained even in the high-throughput regime.
A key difference from MoE is how cost scales with batch size. STEM, parameter traffic grows mainly with the number of unique tokens seen. In contrast, MoE expert selection expands with batch size and routing diversity; larger batches tend to light up more experts, quickly eroding the sparsity benefit.
3.1.1 VRAM AND COMMUNICATION SAVINGS
MoE models use a lot of VRAM. The expert subnetworks must stay on the GPU, or be fetched repeatedly. Expert parallelism also needs all-to-all communication, even when only few experts are active (Huang et al. 2024 Go & Mahajan 2025). STEM avoids these costs. Its embeddings are token-indexed and local to each layer, so the model can prefetch them without any routing logic. These tables are separate from the matmul weights, so can offload them to CPU memory. In our setups, this frees up roughly one-third the FFN parameter memory. We can also replicate the embedding tables in CPU memory every serving node. This eliminates cross-node expert traffic and the synchronization overhead of expert parallelism.
Prefetching cost. The prefetching cost be greatly reduced by deduplicating the STEM embed- dings of the batched tokens. We can further cut traffic by caching the most frequently used STEM embeddings, using the extra memory save from removing the up-projection matrices. As the model embedding size grows, compute cost increases quadratically, but prefetching cost grows only linearly. This makes CPU-offloaded STEM increasingly attractive and scalable for larger model sizes.
3.2 CONTEXT-LENGTT ADAPTIVE PARAMETER USAGE
Because STEM employs token-indexed, fine- ame-grained sparsity, the number of distinct parameters touched in forward pass grows with the number of unique tokens in the window. Aside from the shared projections
attention (Q/K/V/O) and the gated FFN's gate/down projections, the STEM module draws one vector token ID per layer, repeated tokens reuse the same vector, while novel tokens activate new ones. Let be the context length and Lanig the count of unique token ids in the sequence; with STEM applied at layers S and FFN width dir the STEM-specific parameters activated by a single sequence are
ParamssTea STEM (L) |S] dir Lunig-
In natural text Tunngg typically grows sublinearly (Heaps-like), so longer contexts steadily engage more parameters without increasing per-token FLOPs.
This yields test-time capacity scaling with predictable latency: active parameter count keeps on growing with context length, and does not saturate quickly like in MoEs. The dense gating and down-projection preserve contextual mixing, while the STEM path supplies additional capacity at low overhead, supporting long-context tasks (multi-document RAG, CoT) with near-constant per- token compute. [6] illustutss how STEM outperforms the dense baseline at longer context lengths. Additional long-context evaluation Longhench are provided in Appendix A.2
5

---

Published as a conference paper at ICLR 2026
Table Training hyperparameters by setting. Common: weight decay 0.1, 0.9, 32 0.95,
LR warmup ratio 0.01. Minimum LR is 0.1x peak LR. For 1B pretraining, we follow the OLMO schedule for 5T tokens but stop early at 1T.
Configuration
350M Pretrain
Prefrain
Midtrain
1B Context-Extend
Peak I.R
2e-3
4e-4
3.2e-4
le-5
LR schedule
cosine
cosine
linear
cosine
Batch size
512
512
512
Max sequenee length
2048
4096
4096
32768
Training steps
100,000
500,000
50,000
10,000
Cross-doc masking
No
No
No
Yes
4 EXPERIMENTS
We evaluate STEM against dense and MoE baselines downstream tasks while controlling for (i) training compute (activated FLOPs) and (ii) the number training tokens. MoE variants are configured to match STEM's s total parameter count, and their activated FLOPs are kept comparable to the dense baseline. (Note: STEM uses strictly fewer per-token FLOPs than both baselines.) We study two model scales 350M and IB, performing comprehensive ablations 350M and validating STEM at IB under both pretraining-from-srratt and mid-training insertion. Finally, we assess long-context behavior by further fine-tuning with extended context length. We evaluate the Return on Investment (ROI) defined here as the ratio of model accuracy to training FLOPs determine the training efficiency of each model, as the economic value has become major concern of foundational models. Formally, we define it as:
Model Accuracy (Avg)
Training ROI
Total Training FLOPs
4.1 EXPERIMENTAL SETTING
Datasets. For pretraining, we OLMO-MIX- 1124 (OLM al., 2025 a 3.9T-token cor-
pus built from DCLM (Li et al. 2025a and Dolma 1.7 (Soldaini et al.) 2024); we subsample IT tokens for our runs. For mid- training we mix OLMo-MIX-II2. (65%). NEMOTRON-CC. MATH-V1 1 (5%) (Rabech Karimi Mahabadi, (025), and
PREERONNIIPRRRIIINNGG-G (30%) (NVIDIA et al.] 2025). For context-length extension, we use PROLONG-DATT'
(Gao et al. 2024) (63% long-context 37% short-context) and pack sequences up to 32,768 tokens with cross-document attention masking.
Models. We model architectures from Mob leeileeliae 1eLLM-350M (Liu et al. 2024 and 1,1 a.lama33 3.2 -1B (Meta Al 2024) for evaluations. In both the models, we do not share the input embeddings and the language model head. Unless otherwise noted, one third of FFN layers are replaced at uniform intervals by the sparse alternative. For STEM, the dense up-projection is replaced by an embedding table of size V x dæ in each layer. For Hash layer MoE design, we use top-1 routing and choose the number of experts per layer to match STEM's total parameter count, while keeping activated FLOPs comparable to the dense baseline. We also report ablations that replace one half of FFN layers with STEM, and an extreme setting that replaces all FFN layers except the first
Evaluations. Pretrained checkpoints are evaluated zero-shot on eight common sense reasoning tasks: ARC-Easy, ARC-Challenge (Clark al., 2018), BoolQ (Clark at. 2019), PIQA (Bisk al. 2020 SIQA (Sap et al., 2019), HellaSwag (Zellers et al. 2019), Ooon (Mihaylov et al, 2018 and WinoGrande (ai2 2019). To assess advanced knowledge and mathematical reasoning
for mid-training checkpoints, we report MMLU (Hendrycks et al., (02TT) and GSM8K (Cobbe et al.
2021. J. For long-context behavior after context extension, we use Needle-in a-Hasstack (NIAH)
Kamradi 2024).
Training details. We pretrain the 350M models on 100B tokens and the IB models on 1T tokens. We use the AdamW optimizer with cosine learning rate schedule, 10% warmup, and minimum

---

Published as a conference paper at ICLR 2026
Model Architecture
Model Architecture
4.5
2.9
HashMoE (top-1/26)
STEM (full)
2.8
STEM (one-third)
STEM (half)
Loss Value
3.5
Loss Value
2.7
STEM (one-third)
3.0
2.6
2.5
2.5
2.0
2.4
0.0
0.5
1.0
1.5
2.0 x10
4
Training tokens
FLOPs
x10*
(a) Sparse Architecture Comparison
(b) STEM Layer Count Ablation Study
Figure 3: (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, lower training loss can be achieved at fewer training FLOPs.
learning rate of 0.1x the peak value. After pretraining, we run midtraining stage on 100B tokens, followed by context-extension stage on 20B tokens. The full set of hyperparameters is listed in Table 2
4.2 EXPERIMENTAL RESULTS
STEM demonstrates the benefits of fine- fine-grained sparse scaling by improving downstream perfor- mance with fewer training FLOPs. Interestingly, STEM does not suffer from training instability issues that is often the case for fine-grained MoE models (Databricks) 2024 Dai et al., 2024). In- stead, the geometric properties of the STEM embedding spaces further help improve the training convergence. Figure 3a demonstrates the training stability of STEM compared token-indexed Hash Layers MoE, where HashMoE has more bumpy jumps during the training. Moreover, we see the STEM architecture has larger model capacity (lower training loss tendency) when we scale up the training tokens as the loss curve of STEM crosses over the other two architectures when training tokens increase. Furthermore, even with fewer training FLOPs STEM achieves lower training (bb and validation Ia losses.
4.3 DOWNSTREAM EVALUATION RESULTS
We compare STEM with dense baseline as well as Hash layer MoE at 350M scale. On the other hand, for IB model, compare STEM (with one-third of FFN replacement) with only the dense baseline. In both cases we observe substantial improvement in tasks requiring comparatively more external knowledge such as, Are-Challenge and OpenBookQA, while having modest improvements on the rest of the tasks. Additionally, the improvements on the knowledge-intensin tasks are more significant with increase in FFN replacement with STEM layers. Note all the STEM replacement are replacing the up-projection component of original FFN unless specified in the table.
Upon midtraining the IB STEM model continues outperform the dense baseline on the lan- guage modeling downstream tasks. Additionally, STEM architecture exhibits improvements in rea- soning and knowledge retrieval abilities through GSM8k and MMLU performances.
4.4
ABLATION STUDIES
4.4.1
IMPACT OF STEM LAYER COUNT
To identify the efficacy of STEM layers, vary the number of FFN layers replace with STEM alternative. We place the STEM-based decoder layers at regular intervals, interleaved with regular
ROO is normalized at each hasline for hetter comparison.
STEM defaults to replacing third of FFN layers, also writes STEM-1/3

---

Published as a conference paper at ICLR 2026
Table 3: Downstream accuracy of pretrained models at 350M and IB scales. We report the total number of parameters and the number of active parameters for each model variant. Baseline denotes the dense SwiGLU FFN model. For 350M, in the first few rows, we compare sparse alternatives under similar FLOPs: Hash-MoE (top- 1/16 experts in 1/3 of FFN layers), STEM with 1/3 of FFN layers replaced (including up projection replacement, gate projection replacement, and STEM with an additional up-projection). In the next set of rows, we compare STEM with varying up projection layer replacement ratios (1/3, 1/2, full). For IB, we report the dense baseline and STEM with 1/3 up projection layer replacement.
#Total
#Active
Model
ARC-E
ARC-C
RoolQ
PIQA
SIQA
HSwag
OBQA
SGELOPS RO
350M (Pretraining)
Baseline
0.37
0.37
57.6 66
30.55
58.20
69.42
41.10
49.
34,80
56.35
49.72
0.74
1x
Hash-MoE
1.22
0.37
58.83
36.33
55.44
70.21
43.55
47.56
39.26
53.44
50.58
0.74
1.02x
STEME
1.14
63.01
32.68
60.31
70.18
39.76
52.38
33.00
55.88
50.90
0.70
1.08x
STEM (gmte-proj)
1.14
0.35
$4.56
34.12
59.13
64.92
44.56
43.62
36.91
55.00
49.10
0.70
1.04x
STEM!
STEM-11 1/2
1.85
0.34
62.95
40.00
62.02
70.94
43.70
51.49
46.68
55.78
54,20
0.67
1.20s
70.23
(Pretraining)
Baseline
1,50
1,50
66.98
41.88
73.44
44.09
59.65
55.82
3,00
1x
STEM
6.75
1.41
63.95
42,03
61.66
75.00
44.78
60.37
45.
57.34
$6.63
2,83
1.08A
Table 4: Mid-trained model evaluations (IB).
Model
ARC-E ARC-C
RoolQ
PIQA
SIQA
HellaSwag
OBQA
Winogrande
Avg
GSMRK
MMLU
IB (Mid-training)
Baseline
70.78
42.11
65.84
72.95
47.13
60,39
42.97
57.81
57.50
44,2
29.92
STEM
69.78
44.22
68.54
74.69
45.65
61.90
45.70
57.42
58.49
46.4
32.38
FFN-based decoder blocks. Table]]shows that increasing the number of replacement from one-third to half improves the average downstream performance substantially. However, the improvement slows down beyond that. Note that, with increasing number of replacements, the training FLOPs also decrease, and therefore the overall training ROI still increases. We 1/3) achieves 1.08x training ROI of the baseline, while STEM-1/2 achieves 1.20x and STEM-full achieves 1.33x of the baseline. Figure 3b presents the comparison of the three variants in terms of loss vs training FLOPS.
that the STEM (STEM-
4.4.2 IMPACT OF STEM PLACEMENT
Placement of STEM inside the gated FFN matters. To demonstrate this, compare op- tions: replacing the up-projection vs. the gate-projection. As shown in Table 3 replacing the gate underperforms even the dense baseline, while replacing the up-projection yields consistent gains. In SwiGLU, the gate o (WW*) should depend on the current hidden state to modulate <WW") contextually. Swapping 7or for token-indexcd embedding e makes the gate largely input-independent (d)e)) weakening its context- selection. Moreover, the nonlinearity can be effectively abstracted away by the learned embeddings. and consequently its role weakened. In contrast, applying STEM to the up-projection preserves contextual information in gate computation path and proves to be an optimal fine-grained sparse design.
4.4.3 UP-PROJECTION WITH ADDITIVE EMBEDDING
To further study the optimality of STEM's design, we implement STEM 2.3 that retains up pro- jection and additively modulates its output with the STEM embedding. Although adds more parameters and FL.PP, the downstream performance does not improve.
8

---

Published as a conference paper at ICLR 2026
STEM ID
30
Artivation
Layer Features
25
Love (PS0 end 9
25
Up prpjeciin outputs STEM cobeodinns
20
Decusity
20
6-
10
-1.5
-0.5
0.0
0.5
Cosine Similaritt
Cosine Simiarity
Codine sintlartty
(a) Pair-wise cosine similarity
(b) Layer 10
(c) Layer 10
Figure 4: Geometry of STEM embeddings. (a) Distribution of pairwise cosine similarity of STEM embeddings of sampled layers. (b) Pair-wise cosine similarity distributions of up up-projection output space and STEM embeddings. (c) Cosine similarities are computed between the input hidden states of the down projection matrix. All the plots are provided from the 1B model.
5 STEM CHARACTERISTICS
In this section, we analyze some of the characteristies that STEM embeddings demonstrate. We observe that in each layer the STEM embeddings of different tokens have very low pairwise sine similarity which elicits some desirable properties regarding information storage capacity and training convergence. Additionally, because of the clear mapping between the embeddings and the tokens, STEM models are more interpretable.
5.1
LARGE ANGULAR SPREAD OF STEM EMBEDDINGS
Figure 4a| shows that STEM embeddings exhibit very low pairwise cosine similarity -i... large angular spread. We hypothesize that this property improves the information-retiian
behavior of FFN layers by reducing interference among stored items. Prior work (Geva al. 2021 Meng et al.) 2022) models FFNs as key-value memories: each hidden unit is associated with key given by row of the up-projection W(n) E Rd de doatel and a value given by the corresponding column of the down-projection W'd) E Rednss Xdu; the gate projection provides context-dependent, multiplicative modulation that creates a selective read. In this view, the pre-activation h = ((W(1)z) induces a soft address memory slots (hidden units).
In contrast, STEM replaces the learned affine addressing with a direct, token-indexen address vector, upon which the gate still applies context-dependent modulation. To quantify the geometry of these address vectors, we report the distribution of pairwise cosine similarities between unit-normalized vectors. A distribution concentrated near zero (as in Figure ann Figure 4b) indicates that most gles are close to 90" and thus the angular spread between the vectors is reasonably large. This large angular spread lowers cross-talk between slots and can thereby improve the effective information storage capacity of the FFN memory at fixed width Donoho & Elad (2003); Tropp (2004). Figure 4c demonsttraeees the distribution of pairwise cosine similarities between the address vectors after the modulation applied by the gate projection.
5.2
INTERPRETABILIT OF STEM MODELS
STEM exposes token-indexed, layer-local parameters that act as interpretable FFN addresses, en- abling simple, reversible edits that causally steer factual predictions with high reliability and low collateral change. Because each token has a layer-specific STEM vector 6t,E E Rd we can inter- inference time in transparent way.
For example, Figure 5 shows that we can manually control the top next-token probabilities by per- forming swap layer f,
EGemmany
while leaving all other parameters unchanged. Under the original prompt containing "Spain", the in- tervened model's 's top-k next-token distribution closely matches that the control prompt containing "Germany", illustrating precise, token-indexed knowledge editing.

---

Published as a conference paper at ICLR 2026
0.8
0.8
0.6
0.6
Probabiityy
0.4
0.4
0.4
0.2
02
0.2
0.0
Madrid
the
Barcelona
Berlin
called
the
located
Berlin
called
the
located
(a) Original:
(b) Target:
(c) Intervened:
The capital of Spain is
The capital of Germany is
The capital of Spain is
Figure 5: Knowledge edit. Top-4 next-token probabilities for the original prompt "The capital of Spain is" (left), for the target prompt "The capital of Germany is" (middle), and for the intervened model where we swap the STEM vector éSpain, with EGemany, at every STEM layer keeping the original prompt the same(right). The swap shifts mass from Madrid to Berlin, demonstrating loken-indexen layer-local, and reversible control of factual predictions.
RELATED WORKS
MoE (Shazeer al., 2017 Fedus et al.) 2022) introduced large parametric capacity for LLMs at
near-constant FLOPss through sparse computation. The success of MoE models hinges closely with auxiliary loss function designs (Fedus et al. 2022 Rajjhandari et al. 2022 Qiu al. 2025), and system-level solutions (Huang et al., 2024 Go & Mahajan [025 Wang et al. (024b) that ensure load balance among expert networks, training stability, mitigation of representation collapse (Chi et al. 2022), and tolerable communication overload during training and inference. To avoid the interference of auxiliary routing losses with the training objective, recent works have proposed aux- iliary loss-free approaches (Roller et al. 2021 Wang et al. 2024a) that inject fixed or dynamic routing bias to the MoE model.
Conversely, PKM models (Lample al. 2019) reserve large key-value parametric memory with efficient top-k selection through memory memory-efficient keys arranged in product space. (Lmple al. 2019 He 2024) scales up the parametric memory compared MoE, increases the granularity of sparsity, and avoids the cross-device communication overhead, but at the cost of high memory lookup cost during inference. and under-training issues of the large value memory. These challenges require sophisticated architectural modifications (Huang et al. 2025) and advanced system-level solutions (Berges et al., 2024) to overcome.
Recently, Gemma-3n 3n (Google DeepMind 2024) proposed Per Layer Embeddings (PLE) for small on-device models to complement their limited parametric capacity with token-indexed sparse para- metric memory. However, they do not dispose of original FFN modules, and use a much lower- dimensional PLE only to modulate the FFN output in each layer. These embedding tables are ac- commodated in fast storage, outside GPU HBM memory accommodate larger batch sizes and enable fast prefetching.
CONCLUSION
This work introduced STEM, static, token-indexed design that replaces the FFN up-projection with a layer-local embedding lookup.This decouples parametric capacity from per-token compute and cross-device communication, yielding lower per-token FLOPs and fewer parameter accesses, and enabling CPU offload with asynchronous prefetch. Empirically, STEM trains stably despite extreme sparsity (compared to fine-graiie lne-grained MoE variants), improves accuracy over dense baselines, and exhibits higher effective memory capacity via arge-aagguar-sppre embedding space. It also strengthens tong-context performance by activating more distinct parameters as sequence length grows, providing practical test-time capacity scaling.
10

---

Published as a conference paper at ICLR 2026
8 ETHICS STATEMENT
This work develops and empirically evaluates a novel large language model architecture. All training and evaluation datasets are publicly available and widely used within the research community; no new human-subject data were collected, and no sensitive or proprietary data sources were used. Due to computational resource constraints, experiments were conducted on models with up one billion parameters and evaluated at pre- traiiiiinng and mid-training checkpoints, and the scope of the conclusions should be interpreted accordingly. Future research this direction should continue to assess ethical considerations throughout model development, evaluation, and potential deployment.
REPRODUCIBILITY STATEMENT
This work follows the reproducibility recommendations of ICLR; details necessary replicate sults are referenced rather than repeated here. Section 4.1 documents the training datasets, model architectures, training procedures, and evaluation datasets and protocols referenced throughout the experiments. To facilitate independent verification, code and trained model checkpoints will be released to support full reproducibility.
REFERENCES
Winogrande: An adversarial winograd schema challenge scale. 2019.
Mikel Artetxe, Shruti Bhosale, Naman Goyal, Todor Mihaylov, Myle Ott, Sam Shleifer, Xi Victoria Lin, Jingfei Du, Srinivasan Iyer, Ramakanth Pasunuru, Giri Anantharaman, Xian Li. Shuohui Chen. Halil Akin, Mandeep Baines, Louis Martin, Xing Zhou, Punit Singh Koura, Brian O'Horo, Jeff Wang, Luke Zettlemoyer, Mona Diab, Zoriissa Kozareva, and Ves Stoyanov. Efficient large scale language modeling with mixtures of experts, 2022. URL httt::::/aaxiv.ooo///oo 2112 10684
Yushi Bai, Xin Lv, Jiajie Zhang, Hongchang Lyu, Jiankai Tang, Zhidian Huang, Zhengxiao Du, Xiao Liu, Aohan Zeng, Lei Hou, Yuxiao Dong, Jie Tang, and Juanzi Li. Longbench: A bilingual, mul- titask benchmark for long context understanding, 2024. URL https:://rrxa org/ /
2308. 14508
Vincent-Pierre Berges, Barlas Oğuz, Daniel Haziza, Wen tau Yih, Luke Zettlemoyer, and Gargi Ghosh. Memory layers scale, 2024. URL htttp:://axxiv org/
09764
Yonatan Bisk, Rowan Zellers, Ronan Le Bras, Jianfeng Gao, and Yejin Choi. Piqa: Reasoning about physical commonsense in natural language. In Thirty-Fourth AAAI Conference Artificial Intelligence, 2020.
Enric Bdie-A and Philippe Rigollet. The power of fine-grained experts: Granularity boosts expressivity in mixture of experts, 2025. URL htttps://axxi...o////ab.. wwt.o////aa//...o.//////// 06839
Zewen Chi, Li Dong, Shaohan Huang, Damai Dai, Shuming Ma, Barun Patra, Saksham Singhal, Payal Bajaj, Xia Song, Xian-Ling Mao, Heyan Huang, and Furu Wei. On the representation col- lapse of sparse mixture of experts, 2022. URL /
arxia. org/ abs 2004. 091
Christopher Clark, Kenton Lee, Ming-Wes Wei Chang, Tom Kwiatkowski. Michael Collins, and Kristina Toutanova. Boolq: Exploring the surprising difficulty of natural yes/no questions. NAACL, 2019.
Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge. arXiv:1803.05.3 155.505 2018.
Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jemy Tworek, Jacob Hilton, Reiichiro Nakano, Christopher Hesse, and John Schulman. Training verifiers to solve math word problems. arXiv preprint 1rXiv:2110.14166 14168, 2021.
11

---

Published as a conference paper at ICLR 2026
Damai Dai, Chengqi Deng, Chenggang Zhao, R. X. Xu, Huazuo Gao, Deli Chen, Jiashi Li,
Wangding Zeng, Xingkai Yu, Y. Wu, Zhenda Xie, Y. K. Li, Panpan Huang, Fuli Luo, Chong Ruan, Zhifang Sui, and Wenfeng Liang. Deepseekmoe: Towards ultimate expert specializa- tion in mixture-of-experts language models, 2024. URL
06066
Databricks. Introducing dbrx: A new state-of-the-art open Ilm, 2024. URL https: ///::///wwwwwww www. databri cks. com/b hatmmbllg////e http://ntooouuiii----aaoooo htt.://wwonooooooo.eeeeeeeeeeeee
11
Ac-
cessed: 2025-09-04.
David L. Donoho and Michael Elad. Optimally sparse representation in general (nonorthogonal)
dictionaries via I minimization. Proceedings of the National Academy of Sciences, 100(5):2197- 2202, 2003, doi: 10. (0.1073/pnas04438447700 1073 T073/pnas. 043784711437777000000
wwtp://a.......///oo/o///o////// URL https: /twps//////.ppp......////////CC
William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity, 2022. URL http:://arrxiv.orr//bbb///
03961
Tianyu Gao, Alexander Wettig, Howard Yen, and Danqi Chen. Enabling large language models to generate text with citations. 2024.
Mor Geva, Roei Schuster, Jonathan Berant, and Omer Levy, Transformer feed-forward layers are key-value memories, 2021. URL httpt:ttt/////iiii wwtpoorg/////ooo/o 14913
Scokjin Go and Divya Mahajan. Moetuner: Optimized mixture of expert serving with balanced expert placement and token routing, 2025. URL /MM.upps//aaa2...orr//b////22
Google DeepMind. Gemma 3n documentation, 2024. URL https://el.google...
gemma/ wwt.a//oossssoeeeeeeeeeeee /eemna- 3n Accessed: 2025-09-04.
Xu Owen He. Mixture of million experts, 2024. URL htt.:tttss//aax/..or//ab///////
04153
Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. Measuring massive multitask language understanding. Proceedings of the Interna- tional Conference on Learning Representations (ICLR), 2021.
Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, Tom Hen- nigan, Eric Noland, Katie Millican, George van den Driessche, Bogdan Damoc, Aurelia Guy, Simon Osindero, Karen Simonyan, Erich Elsen, Jack W. Rae, Oriol Vinyals, and Laurent Sifre. Training compute-optimal large language models, 2022. URL https: /ttp:://axxvvvvv.. org/ /bs/
2203.15555
Haiyang Huang, Newsha Ardalani, Anna Sun, Liu Ke, Shruti Bhosale, Hsien-Hsin S, Lee, Carole- Jean Wu, and Benjamin Lee. Toward efficient inference for mixture of experts. In The Thirty- eighth Annual Conference on Neural Information Processing Systems, 2024. URL https:// @gpenreview.net. htttt//orrmmidd=ssccmm
Zihao Huang, Qiyang Min, Hongzhi Huang, Defa Zhu, Yutao Zeng, Ran Guo, and Xun Zhou. Ultra- sparse memory network, 2025.
Greg Kamradt. Limtest needleinahaystack. https::/ httppp///gggaaaa..ooooo////aaaaa LLL.TeTtaaattMMMss........
com wwwpw//tttttt/////../..t/ Needle MitecleenAAae InAHaystack 2024. GitHub repository, accessed accesssse 2025-09-24.
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B. Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models, 2020. URI. htt.t:t:pss// http///avvvv.////// org/ abs 2001. 08361
Guillaume Lample, Alexandre Sablayrolles, Mare' Aurelio Ranzato, Ludovic Denoyer, and Hervé Jégou. Large memory layers with product keys, 2019. URL https::///// //arxv/a.vvvv/via org/ abs/ 9077 05242
12

---

Published as a conference paper at ICLR 2026
Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, and Zhifeng Chen. Gshard: Scaling giant models with condi- tional computation and automatic sharding, 2020. URL
16668
Jeffrey Li, Alex Fang, Georgios Smymiss, Maor Ivgi, Matt Jordan, Samir Gadre, Hritik Bansal, Etash Guha, Sedrick Keh, Kushal Arora, Saurabh Garg, Rui Xin. Niklas Muennighoff, Rein- hard Heckel, Jean Mercat, Mayee Chen, Suchin Gururangan, Mitchell Wortsman, Alon Al- balak, Yonatan Bitton, Marianna Nezhurina, Amro Abbas, Cheng- Yu Hsieh, Dhruba Ghosh, Josh Gardner, Maciej Kilian, Hanlin Zhang. Rulin Shao, Sarah Pratt, Sunny Sanyal, Gabriel II- harco, Giannis Daras, Kalyani Marathe, Aaron Gokaslan, Jieyu Zhang, Khyathi Chandu, Thao Nguyen, Igor Vasiljevic, Sham Kakade, Shuran Song, Sujay Sanghavi, Fartash Faghri, Se- woong Oh, Luke Zettlemoyer, Kyle Lo, Alaaeldin El-Nouby. Hadi Pouransari, Alexander Toshev, Stephanie Wang, Dirk Groeneveld, Luca Soldaini, Pang Wei Koh, Jenia Jitsev, Thomas Kol- lar, Alexandros G. Dimakis, Yair Carmon, Achal Dave, Ludwig Schmidt, and Vaishaal Shankar. Datacomp-Im: In search of the next generation of training sets for language models, 2025a. URL
abspg///b///ee https: pttps://azzee /ttp///aaivvv/v/a otg/ abs 2406 11794
Yan Li, Pengfei Zheng, Shuang Chen, Zewei Xu, Yuanhao Lai, Yunfei Du, and Zhengang Wang. Speculative moe: Communication efficient parallel moe inference with speculative token and expert pre-scheduling, 2025b. URL http::ttt////t /wtpivxxivvvvv////////aa////.vvv www.aiiww......oo...ooooooooooee ab s/ 2503. 04398
Zechun Liu, Changsheng Zhao, Forrest landola, Chen Lai, Yuandong Tian, Igor Fedorov, Yunyang Xiong, Emne Chang, Yangyang Shi, Raghuraman Krishnamoorthi, Liangzhen Lai, and Vikas Chandra. Mobilellm: Optimizing sub-billion parameter language models for on-device use cases, 2024.
2402. /ttp://bbbb222.1445555
Kevin Meng, David Bau, Alex Andonian, and Yonatan Belinkov. Locating and editing factual associ- ations in GPT. Advances in Neural Information Processing Systems, 36, 2022. arXwv:202.00526..
Meta AI. Llama 3.2 1b, 2024. URL/tttss: httttpp:///uggghgffc..c///ree hwtp:///meeeaaaaaaaaaaaaaaaaaaaa Accessed: 2025-09-04.
2-1B 1B
Todor Mihaylov, Peter Clark, Tushar Khot, and Ashish Sabharwal. Can a suit of armor conduct electricity? new dataset for open book question answering. In EMNLP, 2018.
NVIDIA, :, Aarti Basant, Abhijit Khairnar, Abhijit Paithankar, Abhinav Khattar, Adithya Renduch- intala, Aditya Malte, Akhiad Bercovich, Akshay Hazare, Alejandra Rico, Aleksander Ficek, Alex Kondratenko, Alex Shaposhnikov, Alexander Bukharin, Ali Taghibakhshi, Amelia Bar- ton, Ameya Sunil Mahabaleshwarkar, Amy Shen, Andrew Tao, Ann Guan, Anna Shors, Anub- hav Mandarwal, Arham Mehta, Arun Venkatesan, Ashton Sharabiani, Ashwath Aithal, Ashwin Poojary, Ayush Dattagupta, Balaram Buddharaju, Banghua Zhu, Barnaby Simkin, Bilal Kartal, Bita Darvish Rouhani, Bobby Chen, Boris Ginsburg. Brandon Norick, Brian Yu, Bryan Catan- zaro, Charles Wang, Charlie Truong, Chetan Mungekar, Chintan Patel, Chris Alexiuk, Christian Munley, Christopher Parisien, Dan Su, Daniel Afrimi, Daniel Korzekwa, Daniel Rohrer, Daria Gitman, David Mosallanezhad, Deepak Narayanan, Dima Rekesh, Dina Yared, Dmytro Pykhtar, Dong Ahn, Duncan Riach, Eileen Long, Elliott Ning, Eric Chung, Erick Galinkin, Evelina Bakh- turina, Gargi Prasad, Gerald Shen, Haifeng Qian, Haim Elisha, Harsh Sharma, Hayley Ross, Helen Ngo, Herman Sahota, Hexin Wang. Hoo Chang Shin, Hua Huang, lain Cunningham, Igor Gitman, Ivan Moshkov, Jachun Jung, Jan Kautz, Jane Polak Scoweroft, Jared Casper, Jian Zhang. Jiaqi Zeng, Jimmy Zhang, Jinze Xue, Jocelyn Huang, Joey Conway, John Kamalu, Jonathan Cohen, Joseph Jennings. Julien Veron Vialard. Junkeun Yi. Jupinder Parmar, Kari Briski. Kather- ine Cheung, Katherine Luna, Keith Wyss. Keshav Santhanam, Kezhi Kong. Krzysztof Pawelec, Kumar Anik, Kunlun Li, Kushan Ahmadian, Lawrence McAfee, Laya Sleiman, Leon Derezyn- ski, Luis Vega, Maer Rodrigues de Melo, Makesh Narsimhan Sreedhar, Marcin Chochowski, Mark Cai, Markus Kliegl, Marta Soppiiesska-ziibbnsee Matvei Novikov, Mehrzad Samadi, Meredith Price, Meriem Boubdir, Michael Boone. Michael Evans, Michal Bien, Michal Za- walski, Miguel Martinez, Mike Chrzanowski, Mohammad Shoeybi, Mostofa Patwary, Namit Dhameja, Nave Assaf, Negar Habibi, Nidhi Bhatia, Nikki Pope, Nima Tajbakhsh, Nirmal Ku- mar Juluru, Oleg Rybakov, Oleksii Hrinchuk, Oleksii Kuchaiev, Oluwatobi Olabiyi, Pablo Rib- alta, Padmavathy Subramanian, Parth Chadha, Pavlo Molchanov, Peter Dykas, Peter Jin, Piotr
13

---

Published as a conference paper at ICLR 2026
Bialecki, Piotr Januszewski, Pradeep Thalasta, Prashant Gaikwad, Prasoon Varshney, Pritam Gun- decha, Przemek Tredak, Rabeeh Karimi Mahabadi, Rajen Patel, Ran El-Yaniv, Ranjit Rajan. Ria Cheruvu, Rima Shahbazyan, Ritika Borkar, Ritu Gala, Roger Waleffe, Ruoxi Zhang. Russell J. Hewett, Ryan Prenger, Sahil Jain, Samuel Kriman, Sanjeev Satheesh, Saori Kaji, Sarah Yurick, Saurav Muralidharan, Sean Narenthiran, Seonmyeong Bak, Sepehr Sameni, Seungju Han, Shan- mugam Ramasamy, Shaona Ghosh, Sharath Turuvekere Sreenivas, Shelby Thomas, Shizhe Diao, Shreya Gopal, Shrimai Prabhumoye, Shubham Toshniwal, Shuoyang Ding, Siddharth Singh, Sid- dhartha Jain, Somshubra Majumdar, Soumye Singhal, Stefania Alborghetti, Syeda Nahida Akter, Terry Kong, Tim Moon, Tomasz Hliwiak, Tomer Asida, Tony Wang. Tugrul Konuk, Twinkle Vashishth, Tyler Poon, Udi Karpas, Vahid Noroozi, Venkat Srinivasan, Vijay Korthikanti, Vikram Fugro, Vineeth Kalluru, Vitaly Kurin, Vitaly Lavrukhin, Wasi Uddin Ahmad, Wei Du, Wonmin Byeon, Ximing Lu, Xin Dong, Yashaswi Karnati, Yejin Choi, Yian Zhang, Ying Lin, Yong- gan Fu, Yoshi Suhara, Zhen Dong, Zhiyu Li. Zhongbo Zhu, and Zijia Chen. Nvidia nemotron nano 2: An accurate and efficient hybrid haamb--ransformee reasoning model, 2025. URL aattps://axxin http://axxxvv..rgaaoooe www.://gg/...eoooeoo.o/eee 2508. 14444
Team OLMo, Pete Walsh, Luca Soldaini, Dirk Groeneveld, Kyle Lo, Shane Arora, Akshita Bha- gia, Yuling Gu, Shengyi Huang. Matt Jordan, Nathan Lambert, Dustin Schwenk, Oyvind Tafjord, Taira Anderson, David Atkinson, Faeze Brahman, Christopher Clark, Pradeep Dasigi, Nouha Dziri, Michal Guerquin, Hamish Ivison, Pang Wei Koh, Jiacheng Liu, Saumya Malik, William Merrill, Lester James V. Miranda, Jacob Morrison, Tyler Murray, Crystal Nam, Valentina Py- atkin, Aman Rangapur, Michael Schmitz, Sam Skjonsberg, David Wadden, Christopher Wilhelm, Michael Wilson, Luke Zettlemoyer, Ali Farhadi, Noah A. Smith, and Hannaneh Hajishirzi. olmo 2 furious, 2025. URL htt.https/:/aaae www.iii.ooooooooooooooooooooooor abs/ 2501. 60656
Zihan Qiu, Zeyu Huang. Bo Zheng, Kaiyue Wen, Zekun Wang, Rui Men, Ivan Titov, Dayiheng Liu, Jingren Zhou, and Junyang Lin. Demons in the detail: On implementing load balancing loss for training specialized mixture-of-expert models, 2025. URL https::/arxiie http//aaxxi...r/ab.a.aaa.
22501.18777
Shrimai Prabhumoye Mostofa Patwary Mohammad Shoeybi Bryan Catanzaro Rabeeh Karimi Ma- habadi, Sanjeev Satheesh. Nemotron-ce-math: 133 billion-token--call high quality math pre- training dataset. 2025. URL fttppt:tts///axiiio org/ abs /abb//555 2508 15096.
Samyam Rajbhandari, Conglong Li, Zhewei Yao, Minjia Zhang, Reza Yazdani Aminabadi, Am- mar Ahmad Awan, Jeff Rasley, and Yuxiong He. DeepSpeed-MoE: Advancing mixture-of- experts inference and training power next-generation AI scale. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song. Csaba Szepesvari, Gang Niu, and Sivan Sabato (eds.), Proceed- ings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pp. 18332-18346. PMLR, 17-23 Jul 2022. URL https: //proceedings.. www.mm..ppe..111222..m http:s///a..//r..haaaaaaaiaa22 hhwaaaaaaaaaaaaaaaaaaaaaaaaaaaee html
Stephen Roller, Sainbayar Sukhbaatar, Arthur Szlam, and Jason Weston. Hash layers for large sparse models, 2021. URL
wwwp:axwwww.....oaaaaaa/aaaa/aaa wwtp:////.abbb.....o 04426
Maarten Sap, Hannah Rashkin, Derek Chen, Ronan LeBras, and Yejin Choi. Socialiqa: Common- sense reasoning about social interactions, 2019. URL ftttpa://arxiv..oott wttp://wwwr///////////oo// 09728
Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gated mixture-of-exppers layer, 2017. URL https://aaxxxi
Luca Soldaini, Rodney Kinney, Akshita Bhagia, Dustin Schwenk, David Atkinson, Russell Authur, Ben Bogin, Khyathi Chandu, Jennifer Dumas, Yanai Elazar, Valentin Hofmann, Ananya Harsh Jha, Sachin Kumar, Li Lucy, Xinxi Lyu, Nathan Lambert, Ian Magnusson, Jacob Morrison, Niklas Muennighoff, Aakanksha Naik, Crystal Nam, Matthew E. Peters, Abhilasha Ravichander, Kyle Richardson, Zejiang Shen, Emma Strubell, Nishant Subramani, Oyvind Tafjord, Pete Walsh, Luke Zettlemoyer, Noah A. Smith, Hannaneh Hajishirzi, Iz Beltagy, Dirk Groeneveld, Jesse Dodge, and Kyle Lo. Dolma: an open corpus of three trillion tokens for language model pretraining research, 2024. URL htI.:ttss://
org/ abs 020.2 001 59
14

---

Published as a conference paper at ICLR 2026
Zayne Sprague, Xi Ye, Kaj Bostrom, Swarat Chaudhuri, and Greg Durrett. Musr: Testing the limits of thain-of-thouuht with multistep soft reasoning, 2024. URL 2310016044
wttpaaaivvvv...... org/ /bs/
Mirac Suzgun, Nathan Scales, Nathanael Schärli, Sebastian Gehrmann, Tay, Hyung Won Chung. Aakanksha Chowdhery, Quoc V. Le, Ed H. Chi, Denny Zhou, and Jason Wei. Challenging big- bench tasks and whether chain-of- -f-uuughtt can solve them, 2022. httthttss///xxyy
abs / 02202220.... 09261
Qwen Team. Qwen3-max: Just scale it, September 2025a.
Qwen Team. Qwen3 technical report, 2025b. htl.ttgg.//ar11..r3//////2......
J.A. Tropp. Greed is good: algorithmic results for sparse approximation. IEEE Transactions Information Theory, 50(10))2231-22 2242, 2004. doi: 10.1109/11122044884773
Lean Wang, Huazuo Gao, Chenggang Zhao, Xu Sun, and Damai Dai. Auxiliary-loss-free load balancing strategy for mixture-of-experts, 2024a. URL htttttps///ax...oo/r//aaa/////// 15664
Wei Wang, Zhiquan Lai, Shengwei Li, Weijie Liu, Keshi Ge, Ao Shen, Huayou Su, and Dongsheng Li. Pro-prophet: A systematic load balancing method for efficient parallel training of large-scale moe models, 2024b. URL httpttsss///aiiiiiee www.eiiiiioo.aaoooooooooeoeeeeee / 2411. 10003
Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can ma- chine really finish your sentence? In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, 2019.
A APPENDIX
A.1 ADDITIONAL BENCHMARKS
Interestingly, our additional experiments on more contextual reasoning-heavy tasks show that STEM's contextual reasoning skill is better than that of the dense baseline. To directly probe reason- ing beyond parametric knowledge, we evaluate IB-scale baseline and STEM models on -enchhhh Hard (BBH) (Suzgun et al. 2022), MuSRRSpraguu et al., 2024), and the LongBencht Bai al. 2024) multi-hop reasoning and bode-understanding subsets. BBH is a collection of diverse, challenging tasks designed to require multi-step and compositional reasoning. MuSR requires the model to track entities and constraints over a long narrative before answering question. The LongBench multi- hop subset tests reasoning across multiple passages, while the code-understanding subset evaluates comprehension of complex code snippets. As shown in Table STEM consistently outperforms the dense baseline on BBH. MuSR, and on LongBench multi-hop and code-underee ande-understanding tasks across all context-length ranges, indicating that STEM does not impair contextual reasoning and can in fact improve it.
Table 5: Contextual reasoning benchmarks for 1B-scale models. LongBench scores are averaged over tasks within each context-length range.
Model
BBH
MuSR
LongBench
Multi-hop
LongBench Code
< 4k
4 8k
8k
4k
4 8k
8k
Baseline
24.87
35.85
5.72
6.20
6.19
45.37
44.64
41.30 49.60
STEM
27.55
36.38
10.20
8.63
7.82
52.68
52.53
A.2 ADDITIONAL LONG-CONTEXT EVALUATION
Apart from the synthetic task Needle-in-a-hen hayssaaakk further evaluate STEM LongBench, long-context benchmark that spans six task categories, including single- and multi-document ques- tion answering, summarization, few-shot learning, synthetic tasks, and code completion. We group
15

---

Published as a conference paper at ICLR 2026
Table LongBench results (average across tasks) for IB models, grouped by context length.
Model
0-2k
2-4k
4 6k
6-8k
8-10k
10-12k
12k+
Base
24.0
23.8
22.1
22.3
21.9
21.1
23.5
STEM
27.6
27.6
24.4
22.7
23.0
21.7
24.2
test examples by context length and report the average scores in each regime. As shown Ta- ble 6 the IB STEM model consistently matches or outperforms the IB dense baseline across all context-length ranges, indicating that its long-context capabilities extend beyond synthetic tasks.
16


Reviews:
[{"review_id": "gBtoGcolW0", "reviewer": "Reviewer_YuJd", "summary": "The paper proposes a new modification to the Transformer architecture to enhance fine-grained sparsity. The most popular way to achieve this today is via MoE architectures which are known to the tricky to train\nin a stable manner and introduce significant communication overheads during training and inference. In addition MoE training also has to deal with issues such as experts getting under-trained. Load-balancing objectives are commonly used to alleviate this but they come at a certain quality cost typically. \nSTEM proposes instead looking at the standard SwiGLU FFN architecture and replacing the up-projection layer with a token indexed embedding lookup that is specific to each layer. That is a depth L model will have L+1 embedding tables in total. This is a way of introducing static fine-grained sparsity into deeper layers of the model. The benefits of such a change include minimizing communication overhead, enabling CPU offload with asynchronous prefetch while achieving fine-grained sparsity.\n\nThe authors perform pre-training experiments on a 350M parameter model for 100B tokens and a 1B parameter model for 1T tokens. They also perform mid-training and context extension experiments on a 1B model with 100B and 20B token budgets respectively.\n\nThe authors observe better training stability with STEM training. They see consistently improved performance on knowledge-intensive tasks.\nThey also see improved long-context inference performance as STEM activates more distinct parameters as the sequence length grows. STEM also is more efficient in terms of train time and test time FLOPs.\nA key concern with approaches like STEM is the gain in shallow parametric capacity comes with losses in deep contextual reasoning abilities. However, the experiments of the paper show some initial promising results to potentially mitigate this concern.", "questions": "- Can you provide some details on what the wall-clock time or GPU hour impact of STEM is in practice? I believe the proposed benefits can indeed be realized but it might be less than the theoretical maximum of 33% reported.\n\n- Can you provide more details on how you perform the mid-training? How are the STEM embeddings initialized? How long does the loss spike (due to sudden architecture change) take to recover?\n\n- Can you provide more downstream evaluations on more contextual reasoning heavy benchmarks? MMLU at the 1B scale seems barely above random guessing (25%) and moreover MMLU is also quite knowledge intensive.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- The paper proposes an interesting idea to complement the transformer architecture. The authors are aware of the many practical challenges surrounding today’s transformer architecture and take them into account when designing STEM.\n\n- The proposed method can be applied in stages by deciding to only apply it to some portion of the layers of a transformer. This offers for a smoother transition away from the standard architecture and for greater flexibility in architecture design.\n\n- The paper performs an analysis of the pairwise cosine similarity of STEM embeddings vs original FFN up projection’s embeddings and finds that there is less redundancy in STEM embeddings. While this is an interesting finding, it is unclear whether reduced redundancy might negatively affect robustness.", "weaknesses": "- Although the paper correctly accounts for reduced FLOPs, it can often be quite challenging to realize the full benefit of a FLOP reduction in practice. The paper offers a theoretical discussion of how the embedding table reads can be pre-fetched and offloaded onto CPU which makes sense but it would have made a stronger case if wall-clock training, prefill and generation time comaprisons were provided along with HBM usage.\n\n- The concern about the hit to contextual reasoning abilities (since we replace processing of later layers’ highly enriched residual streams with a token indexed embedding) isn’t fully alleviated by the downstream numbers shown. It is nice that we see noticeable improvements in MMLU and GSM8k which are more reasoning heavy tasks but given the scale of the model, the improvement could be coming from the greater parametric knowledge capacity provided by STEM.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "7NjXSJhr0e", "reviewer": "Reviewer_8jLf", "summary": "The paper introduces STEM, a new transformer architecture that replaces the FFN up-projection with a token-indexed, layer-local embedding lookup. This approach offers static sparsity, removing runtime routing and enabling CPU offload, while keeping gate and down-projection layers dense. Experiments on 350M and 1B parameter models show consistent improvements compared with dense and Hash-MoE baselines, with no observed instability. Theoretical analysis, ablations, and interpretability experiments (e.g., token embedding swapping) support the design’s efficiency and explainability advantages.", "questions": "1. How does STEM scale beyond 1B parameters? Are there any observed communication or memory bottlenecks at larger scales (e.g., 7B+)?\n2. What is the exact CPU-GPU data transfer overhead in practical deployments? Are there latency trade-offs during inference?\n3. Could we have some analysis on inference and training efficiency?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "1. Novel Architecture Design. A creative and simple static-sparsity alternative to MoE models that avoids routing overhead and load-balancing complications.\n2. Strong Empirical Validation. Comprehensive experiments across scales (350M, 1B) and tasks demonstrate consistent benefits in both efficiency and accuracy.\n3. Training Stability. Unlike many fine-grained sparse models, STEM avoids loss spikes and under-trained experts.", "weaknesses": "1. Potential Memory Overhead. Each layer’s token-indexed embedding table may become impractical for large vocabularies, despite CPU offload?\n2. Lack of Ablation on Embedding Dimensionality. The effect of embedding size or vocabulary size on performance and stability is unexplored.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "3Af3PQR8xA", "reviewer": "Reviewer_7Z8s", "summary": "The paper introduces STEM, which is a static, token-indexed design that replaces the FFN up-projection with a layer-local embedding lookup.This decouples parametric capacity from per-token compute and cross-device communication, yielding lower per-token FLOPs and fewer parameter accesses, and enabling CPU offload with asynchronous prefetch. STEM is a static, token-indexed, fine-grained mechanism that replaces only the up-projection in gated FFNs with a token-specific vector retrieved from a layer-local embedding table.", "questions": "What are the limitations of the STEM method compared to the competing approaches such as MOE? Presumably something is lost due to the increased sparsity compared to MOE.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "STEM simplifies operation compared to MOE by dense up-projection in the SwiGLU FFN with a token-indexed vector from a per-layer table lookup. \n\nSTEM improves both computation by reducing per-layer FLOPs during training and memory access by lowering parameter traffic relative to a dense up-projection.\n\nExperiments are informative, with sufficient ablations. The paper does a good job at describing the various aspects of the STEM technique and in pointing out the benefits of the approach relative to other methods such as MOE. Design choices are clearly (if tersely) motivated.", "weaknesses": "Paper is dense and focused. It uses a lot of jargon and will not be so accessible to readers not familiar with the ideas and issues specific to this narrow research area. This is not necessarily a weakness, but readers that are not in this area may not appreciate the paper's purpose or contributions. The paper would be better with some diagrams to illustrate the architectural differences between STEM and other approaches, including MOE techniques. Figure 1.c is too small and unclear for this purpose.\n\nIn section 4.4.1 the term \"ROI\" is used in reference to table 3. However, table 3 does not include ROI in the numbers. What does the acronym ROI stand for? \"Return on Investment\"? That doesn't really make sense in this context (what is the \"investment\"? What is the \"return\"?). The text only vaguely defines what is meant by ROI, as \"performance over training flops\". I assume that \"over\" in this context means divided by, and performance means the average of the scores. This should be expressed more clearly and precisely, e.g. as an equation \"ROI = average accuracy/#training flops\". \n\nTable 3 (and other tables in the paper) is lacking an informative caption. There is no definition of the numbers, just the column headings. The column headings just state the problem task (presumably). The caption should say that these numbers represent percent correct on the task test sets.\n\nIt is not clear how figure 3 expresses anything about the \"Geometry\" (which is misspelled here) of the STEM Embeddings. All that it shows are histograms of cosine distances.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "fs7Cq2svky", "reviewer": "Reviewer_eAWT", "summary": "This paper proposes STEM, a sparse transformer architecture that replaces the FFN up-projection with token-indexed embeddings from a layer-local lookup table. The authors claim improved training stability, better downstream performance, and enhanced long-context capabilities compared to dense baselines and MoE alternatives.", "questions": "1. Can you provide results with parameter-matched dense baselines?\n\n2. What are actual wall-clock training/inference times and memory usage?\n\n3. How does this compare to DeepSeekMoE, Mixtral, or other modern MoE methods?", "limitations": "", "rating": 2, "confidence": 2, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. The paper identifies real issues with fine-grained MoE models (training instability, load balancing, communication overhead) and proposes a simpler alternative.\n\n2. The paper provides evaluation across multiple scales (350M, 1B), training regimes (pretraining, midtraining, context extension), and benchmarks, as well as systematic ablations on layer count, placement (gate vs. up-projection), and hybrid variants.\n\n3. Analysis of embedding space properties (angular spread) is valuable.", "weaknesses": "1. There is a parameter mismatch in the experiments, in particular, STEM uses 1.14B parameters vs. 0.37B dense baseline at 350M scale.\n\n2. The results are incocsistent, as table 3 shows STEM often underperforms baseline: 350M: ARC-C (32.68 vs 30.55), SIQA (39.76 vs 41.10), 1B: BoolQ (61.66 vs 64.21), PIQA (75.00 vs 73.44 favors STEM but margin is small). Abstract claims \"up to 3-4% improvement\" but average gains are ~1% or less.\n\n3.  The long-context analysis is limited, as NIAH is a synthetic task.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)