Sparse computation can scale the parametric capacity of LLMs without proportionally increasing active computation, but pretraining sparse architectures remains expensive and data-hungry. Existing sparse upcycling methods mainly target Mixture-of-Experts models, where expert homogeneity, sample inefficiency, and system overheads can limit practical scalability. Static sparse memory offers predictable access and token-specific updates that can help address these challenges, but naively integrating it into pretrained computation yields poor returns. In this work, we introduce STAPLE, a recipe built upon STEM for integrating static sparse memory into pretrained models. Whereas STEM replaces FFN up-projections with token-indexed embeddings, STAPLE retains them and adds the embeddings as a residual pathway, preserving pretrained computation while learning complementary token-specific information in a scalable manner. On OLMo-2 1B, STAPLE improves over continued dense training by 4.27 points on HumanEval, 2.80 on MBPP, 2.27 on GSM8K, and 1.65 on DROP. Average performance increases with static capacity, and improvements persist on the larger OLMo-2 7B backbone. STAPLE also outperforms static sparse baselines such as STEM and an Engram-style model, and achieves a higher average score than an MoE upcycling baseline while decoding 2x faster with less than half its GPU model memory and measured GPU energy per token in our throughput-oriented decoding evaluation. Importantly, despite introducing billions of parameters late in training, our host-offloaded implementation keeps generation throughput within 1.5% and training throughput within 3% of the dense baseline, while caching only the ~5% most frequent embeddings on GPU. These results show that static sparse memory can provide a sample-efficient, scalable, and sustainable path for expanding pretrained models without sparse pretraining from scratch.
Figures & tables
Figure 1: Overview of STAPLE . (a) STAPLE improves over the dense baseline early in mid-training, while replacement-style STEM upcycling remains below it. The horizontal axis counts mid-training tokens, excluding the preceding warmup or dense continuation stage. (b) Downstream performance improves as static sparse capacity is increased through embedding rank or the number of augmented layers. (c) Domain averages normalized to continued dense training; the largest gain is in code, with smaller gains in mathematics, knowledge, and reading comprehension.
Figure 2: Domain-specific performance comparison among STEM upcycling, STAPLE and STAPLE with no warmup. Each model uses four embedding layers. The horizontal axes count mid-training tokens, excluding the separate warmup stage where used. Panels (a–c) illustrate the importance of retaining the FFN up-projections in STAPLE as it consistently outperforms STEM in every aspect. Panels (a,c) show the importance of warmup in improving general knowledge and commonsense task performances. Conversely, panel (b) suggests that math and code reasoning performances are relatively less affected.
Figure 3: Pair-wise cosine similarity analysis between mean token-wise up-projection output and STEM embeddings in different layers of STAPLE with 12 STEM tables. Overall, the STEM embeddings are quite orthogonal to the respective mean up-projection output vectors.
Figure 4: Training and serving efficiency of host-offloaded STAPLE . (a) Training throughput and GPU-resident weights plus optimizer state; parameter counts include the STEM tables. STAPLE keeps its bf16 tables and fp32 AdamW moments in host DRAM (hatched), matching the dense GPU footprint, whereas the MoE baseline (Sec. 3 ) keeps all experts and their optimizer states on the GPU. (b) Decode throughput on one NVIDIA B200 at batch size 64 and context length 256, labeled with GPU memory for weights plus, for STAPLE , a 1 GB cache of frequently accessed rows. Latency and prefill results are in Table 4 .
Commonsense
Knowledge
RC
Math
Code
Method
ARC-E
ARC-C
HSwag
WinoG
MMLU
DROP
GSM8K
MATH
MBPP
HEval
Avg.
OLMo-2 1B: baselines and architecture/initialization comparisons
Baseline
75.42
47.44
65.18
66.22
50.30
41.91
37.76
7.68
28.80
23.17
44.39
MoE upcycling (8 experts, top-1)
75.04
45.73
66.82
64.72
49.96
40.74
39.80
7.64
28.00
28.05
44.65
Engram-style
77.40
47.35
66.05
64.01
50.98
41.68
38.21
8.44
29.20
24.39
44.77
STEM ( L=12 )
74.37
42.06
60.79
61.56
44.30
37.00
26.00
6.94
22.60
23.78
39.94
Table 1: Upcycling comparison and design ablations. All models in the first block use the OLMo-2 1B seed; the Engram-style baseline’s added embedding memory is parameter-matched to STAPLE -12L. The replacement-style STEM and additive STAPLE models both use 12 embedding layers. The second block reports the OLMo-2 7B backbone with L=16 for STAPLE . Best values are bolded within each backbone block.
Commonsense
Knowledge
RC
Math
Code
Method
ARC-E
ARC-C
HSwag
WinoG
MMLU
DROP
GSM8K
MATH
MBPP
HEval
Avg.
Baseline
75.42
47.44
65.18
66.22
50.30
41.91
37.76
7.68
28.80
23.17
44.39
STAPLE ( L=4 )
76.98
46.25
66.00
65.82
51.27
42.83
39.35
7.08
29.80
25.61
45.10
STAPLE ( L=12 )
76.60
45.99
68.27
63.77
51.21
43.56
40.03
7.44
31.60
27.44
45.59
STAPLE ( L=12 , r=512 )
76.18
46.67
66.99
65.59
50.49
41.91
39.80
6.28
30.20
25.61
44.97
Table 2: STAPLE scalability on OLMo-2 1B. Increasing the number of augmented layers L and the embedding dimension r are separate capacity axes; r=512 denotes the low-rank variant, while the other STAPLE rows use full-width embeddings. All rows use the same continued-training budget. Best values are bolded within this table.
Figure 5: Static memory ablation and controlled embedding interventions. (a) Benchmark-score retention after removing the last 4, 8, or all 12 embedding tables from trained STAPLE -12L. (b) Order statistics: with the STEM embedding for ‘maximum’ swapped for that of ‘minimum’, probability over the five sorted candidates (heatmap) and expected answer rank (line) as correct demonstrations are added. (c) Code generation: probability that the generated maximum-finding loop compares with < under the same swap, alongside the unmodified model and the instruction rewritten to ‘minimum’.
Parameter group
Training setup
Optimization schedule
Stage
Parameters
Seq. Len.
Batch Size
Train Steps
Warmup Steps
Resume Step
End Step
LR Scheduler
Peak LR
Min LR Ratio
Weight Decay
Warmup
STEM/STAPLE tables
4096
512
50,000
500
0
50,000
cosine
4×10−4
0.6
1×10−4
Engram n-gram tables
4096
512
50,000
500
0
50,000
cosine
4×10−4
0.6
1×10−4
Backbone
–
–
50,000
0
1,907,359
2,384,186
cosine
4×10−4
0.1
1×10−1
MoE expert weights
–
–
50,000
500
0
50,000
cosine
2×10−4
0.91
1×10−1
Mid-training
STEM/STAPLE tables
4096
512
50,000
0
0
50,000
linear
3×10−4
0
1×10−4
Table 3: Training hyperparameters for warmup and mid-training. STEM/STAPLE tables are the token-indexed static sparse embeddings; Engram n-gram tables are the hashed n-gram embeddings of the Engram-style baseline. During STAPLE warmup, pretrained up-projections in augmented layers are frozen. The backbone warmup schedule resumes the base model’s pretraining cosine schedule at its final checkpoint. MoE expert weights use a separate cosine schedule during warmup and linear decay during mid-training, with a peak learning rate of 2×10−4 in both stages.
Architecture
HBM model memory
Linear FLOPs / token
Decode B1 ms / tok/s
Decode B64 ms / tok/s
Prefill B1 ms / tok/s
Prefill B64 ms / tok/s
Dense
2.97 GB GPU weights
∼ 2.56 G
1.224 / 817
1.817 / 35.2K
1.419 / 180.5K
34.184 / 479.3K
STAPLE -12L
2.97 GB GPU weights + 1 GB cache
∼ 2.56 G
1.237 / 809
1.845 / 34.7K
1.524 / 168.0K
34.678 / 472.5K
STEM -12L
2.57 GB GPU weights + 1 GB cache
∼ 2.16 G
1.151 / 869
1.671 / 38.3K
1.413 / 181.1K
27.580 / 594.1K
MoE (8 experts, top-1)
8.61 GB GPU weights, all 8 experts
∼ 2.56 G + routing
1.951 / 513
3.761 / 17.0K
2.975 / 86.1K
36.179 / 452.9K
Table 4: Serving efficiency on one NVIDIA B200. All measurements use context length 256; latency and throughput are reported in ms and tok/s. HBM memory includes resident weights and GPU caches. STAPLE -12L offloads 12 embedding tables (19.73 GB in BF16 precision) to CPU DRAM with a 1 GB GPU cache. MoE denotes the upcycled MoE baseline of Sec. 3 (eight interleaved MoE layers, each with eight full-width experts and top-1 routing), which keeps all experts in HBM and incurs routing and token-reordering overhead.
Phase
Retrieval path
Overlap / execution
Prefill
• Deduplicate prompt tokens once. • Probe the GPU LFU cache and gather only cache-miss rows from CPU. • Pack misses into pinned host buffers and issue bulk asynchronous H2D copies.
• Reuse the same inverse token map across all STAPLE layers. • Prefetch later-layer rows while earlier transformer layers execute. • Pipeline successive requests/chunks and fuse compact-row reconstruction with the FFN using torch.compile .
Decode
• Store every four adjacent STAPLE tables in a token-major CPU layout. • Launch a GPU lookup kernel that probes the LFU cache. • On a miss, read the row directly from CUDA-registered host memory; no CPU gather or explicit H2D copy is required.
• Capture grouped lookups and model execution in a multi-stream CUDA graph. • Launch later lookup groups on a side stream while earlier layers compute. • Synchronize only when a group is first consumed (layers 1, 5, and 9 for STAPLE -12L).
Figure 6: A token-indexed memory intervention propagates through a GSM8K solution. Full output text for the original model, the STEM row update, and the rewritten-question control, with line wrapping adjusted for readability. Highlights mark the critical numerical changes. With the question unchanged, the row update produces the same intermediate quantities and final answer as explicitly replacing 60 with 120 in the question.
Memory-augmented language models combine persistent or temporary memory with neural computation to retain, retrieve, and use information beyond the active context.
Dynamic Sparse Training (DST) offers a promising paradigm for improving the training and inference efficiency of deep neural networks; however, we find that in large language model training, DST can suffer from optimization instability, manifested as loss spikes after topology updates. In this work, we show that the naive use of standard Adam-based optimizers leads to a cold-start issue for newly regrown parameters, resulting in excessively large updates and disrupted training dynamics. To address this issue, we propose Sparse Memory-Efficient Training (SMET), which stabilizes DST with optimizer warm-up and improves training progress through density-aware learning-rate scaling. SMET further reduces memory consumption by storing gradients and optimizer states only for active parameters. We provide a theoretical analysis of the update behaviors under SMET, showing improved optimization stability. Extensive experiments demonstrate that SMET enables stable, scalable, and memory-efficient sparse pre-training of LLMs, paving the way for sparse training as a practical alternative to dense training. Our code is publicly available at: https://github.com/QiaoXiao7282/SMET.
Qiao Xiao, Boqian Wu, Patrik Okanovic +6
Eindhoven University of Technology · University of Luxembourg · University of Twente +1
The quadratic complexity of attention imposes severe memory and computational bottlenecks on Large Language Model (LLM) inference. This challenge is particularly acute for emerging agentic applications that require processing multi-million token sequences. We propose STS, a sparse attention mechanism that requires no model retraining. STS leverages the key insight that tokens identified as important by a smaller draft model are highly predictive of important tokens for a larger target model. By integrating into speculative decoding frameworks, STS repurposes the draft model's attention scores to dynamically construct a token-and-head-wise sparsity mask. This mask effectively prunes the expensive attention computation in the target LLM. Our evaluation shows that STS achieves a 2.67x speedup operating at approximately 90% sparsity on representative benchmark NarrativeQA, maintaining negligible accuracy degradation compared to dense attention. STS establishes a new state-of-the-art on the sparsity-accuracy trade-off, outperforming prior techniques by enabling higher sparsity levels for a given accuracy budget.
Ceyu Xu, Jiangnan Yu, Yongji Wu +1
The Hong Kong University of Science and Technology · UC Berkeley
Scaling autoregressive large language models (LLMs) has driven unprecedented progress but comes with vast computational costs. In this work, we tackle these costs by leveraging unstructured sparsity within an LLM's feedforward layers, the components accounting for most of the model parameters and execution FLOPs. To achieve this, we introduce a new sparse packing format and a set of CUDA kernels designed to seamlessly integrate with the optimized execution pipelines of modern GPUs, enabling efficient sparse computation during LLM inference and training. To substantiate our gains, we provide a quantitative study of LLM sparsity, demonstrating that simple L1 regularization can induce over 99% sparsity with negligible impact on downstream performance. When paired with our kernels, we show that these sparsity levels translate into substantial throughput, energy efficiency, and memory usage benefits that increase with model scale. We will release all code and kernels under an open-source license to promote adoption and accelerate research toward establishing sparsity as a practical axis for improving the efficiency and scalability of modern foundation models.