PaperlayerView exact revision on arXiv

Source-linked AI summary

STAPLE: Static Sparse Upcycling of Pretrained Checkpoints via Embedding Modules

Ranajoy Sadhukhan, Euijin Hong, Sheng Cao, Zechun Liu, Shuowei Jin, Sha Li, Changlong Yu, Xin Liu, Qin Lu, Zhenyu Shi, Priyanka Nigam, Beidi Chen

arXiv:2610.12599v1cs.LG

TL;DR

Sparse upcycling seeks to expand pretrained models without expensive sparse pretraining from scratch, but static memory can perform poorly when integrated naively. STAPLE retains pretrained computation and adds token-indexed embeddings as a residual pathway. On OLMo-2 1B, it improves the ten-task average over continued dense training, with gains across coding, math, and reading-comprehension benchmarks.

  • Problem

    Sparse architectures can add capacity without proportionally increasing active computation, but sparse pretraining is expensive and static memory integration can yield poor returns.

    Source: §1 · pp. 1–2

  • Method

    STAPLE preserves pretrained FFN up-projections and adds zero-initialized token-indexed embeddings as a residual pathway for complementary token-specific information.

    Source: §1 · pp. 2–3

  • Results

    On OLMo-2 1B, STAPLE raises the ten-task average by 1.20 points over continued dense training, with gains of 4.27 points on HumanEval, 2.80 on MBPP, 2.27 on GSM8K, and 1.65 on DROP.

    Source: §1 · p. 3

  • Takeaways & Limitations

    STAPLE suggests static sparse memory can expand pretrained models with practical performance gains as capacity grows, without prohibitive serving cost.

    Source: §1 · p. 2, §6 · p. 10

Abstract

from arXiv · show

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.

1 Introduction

STAPLE expands pretrained dense models with static sparse memory while preserving their existing computation, addressing the cost and sample-efficiency challenges of sparse upcycling. It improves benchmark performance over continued dense training and scales with added memory capacity.

1 Introduction

  • Sparse upcycling aims to expand pretrained checkpoints without costly sparse pretraining from scratch, but existing MoE approaches face expert homogeneity and system overheads.The desired approach should remain sample-efficient, preserve seed-model training and serving efficiency, and scale with sparse capacity.Source: §1 · pp. 1–2
  • Under the same upcycling budget, replacement-style STEM struggles to match continued dense training, whereas STAPLE substantially improves downstream performance by retaining pretrained computation.Figure 1 shows STEM remaining below the dense baseline even after 100B tokens.Source: §1 · p. 2, Figure 1 · p. 2
  • A short warmup with selective freezing particularly improves knowledge-intensive and commonsense tasks, while math and code reasoning are relatively less affected.The Figure 2 comparison uses four embedding layers.Source: §1 · p. 2, Figure 2 · p. 3
  • STAPLE improves average downstream performance as either the number of augmented layers or embedding dimension increases, while added memory can remain largely offloaded to CPU memory.Source: §1 · p. 2
  • STAPLE retains the pretrained FFN up-projection and adds token-indexed embeddings as a residual pathway, letting the new parameters learn complementary token-specific information.Its dense pathway preserves shared pretrained computation; the embedding pathway is zero-initialized.Source: §1 · pp. 2–3
  • On OLMo-2 1B, STAPLE raises the ten-task average by 1.20 points over continued dense training, with gains on HumanEval, MBPP, GSM8K, and DROP.The respective gains are 4.27, 2.80, 2.27, and 1.65 points.Source: §1 · p. 3

2 Method

STAPLE retains the pretrained FFN up-projection and adds zero-initialized token-indexed embeddings. The method studies initialization, capacity scaling, and host-offloaded execution.

2.1 Architecture study: impact of retaining the up-projection

  • Replacement-based STEM fails to match continued dense training even with extra warmup, while retaining the up-projection and adding embeddings residually performs substantially better under the same budget.Sparse token updates compound replacement’s challenge because embeddings must relearn discarded pretrained up-projection features; tested initializations did not improve long-term performance.Source: §2.1 · p. 4

2.2 Initialization

  • A roughly 2.5% pretraining-corpus warmup with augmented-layer up-projections frozen improves MMLU and average commonsense, but gives no consistent gains on advanced math or programming.The embeddings remain responsive to the subsequent mid-training distribution and can learn capabilities not emphasized during warmup.Source: §2.2 · pp. 4–5

2.4 System Efficiency

  • 2.04× higher decode throughput at batch size 64 and about 54% less HBM model memory than MoE, with a 1 GB cache replacing 19.73 GB of tables.Caching 5–6% of the vocabulary achieves roughly an 80% hit rate.Source: Figure 4 · p. 5, §2.4 · pp. 5–6
  • STAPLE-12L trains at 113.9K vs. 117.0K tokens/s for dense (within 3%), matches dense GPU weights and optimizer memory, while MoE needs 2.9× the GPU memory.Compact row buffers add about 3% peak HBM.Source: Figure 4 · p. 5, §2.4 · p. 6

2.3 Scalability

  • Increasing augmented layers from 4 to 12 raises average score from 45.10 to 45.59; MBPP improves 1.80 points, HumanEval 1.83, and DROP 0.73, while MMLU is essentially unchanged.The embedding tables are placed in early decoder layers 1–12.Source: §2.3 · p. 5
  • At L = 12, rank-512 STAPLE reaches 44.97 average performance versus 45.59 for full-rank STAPLE and 44.39 for dense continuation, with a 16× smaller token-indexed table.The low-rank module uses V r + dffr parameters per layer and adds 2dffr projection FLOPs per token.Source: §2.3 · p. 5

3 Experimental Setup

Experiments compare STAPLE and sparse upcycling baselines on OLMo-2 backbones, using controlled training budgets and a broad benchmark suite.

Models

  • The OLMo-2 1B comparisons include 4- and 12-table zero-initialized STAPLE, parameter-matched Engram-style memory, and an eight-layer MoE with eight top-1-routed experts.Table 1 also reports OLMo-2 7B with 16-layer STAPLE.Source: §3 · p. 6, Table 1 · p. 7

Datasets

  • Mid-training uses Dolma 3 Dolmino Mix; warmup uses about 100B DCLM tokens, roughly 2.5% of seed pretraining, excluding code and math sources.The warmup-data choice isolates the effect of warmup without additional math- or code-specific supervision.Source: §3 · p. 6

Training

  • The 1B comparisons use about 100B warmup plus 100B mid-training tokens, while 7B uses about 30B plus 70B; dense parameters follow OLMo-2 hyperparameters and embeddings use about 4× their peak learning rate.During STAPLE warmup, augmented-layer up-projections are frozen.Source: §3 · pp. 6–7

Evaluations

  • Evaluation spans commonsense reasoning, general knowledge, reading comprehension, mathematics, and coding, using ARC, HellaSwag, WinoGrande, MMLU, DROP, GSM8K, MATH, MBPP, and HumanEval.The suite uses task-specific zero-, three-, or five-shot settings and reports accuracy, normalized accuracy, F1, exact match, or pass@1 as appropriate.Source: §3 · p. 7

4 Results & Analysis

STAPLE improves over continued dense training on OLMo-2 1B, with gains on reasoning-heavy tasks and further gains as static capacity increases. Removal and embedding-swap interventions show that learned memory contributes to predictions while consistent context can counteract its influence.

4 Results & Analysis

  • STAPLE-12L raises the ten-task average from 44.39 to 45.59 over continued dense training, including gains on HumanEval, MBPP, GSM8K, and DROP.The respective improvements are 4.27, 2.80, 2.27, and 1.65 points.Source: §4 · p. 7
  • On OLMo-2 7B, STAPLE raises the average score from 57.11 to 58.01 and improves 9 of 10 tasks.Average mathematics improves by 4.16% relative to the 7B baseline, and average code performance improves by 3.68%; the experiments use about 30B warmup tokens and 70B mid-training tokens.Source: §4 · pp. 7–8

4.1 Ablation studies

  • Removing the last 4, 8, or all 12 embedding tables lowers the ten-benchmark average from 45.59 to 41.28, 36.94, or 32.17.Removing all tables especially reduces GSM8K, MBPP, and HumanEval; the remaining weights have co-adapted with memory, so the ablated model is not the original dense seed.Source: §4.1 · p. 8

4.2 Interpretability: Interplay between static memory and contextual reasoning

  • Under a maximum-to-minimum embedding swap, four correct demonstrations shift the order-statistics answer probability to 0.97 for the maximum.Without demonstrations, the swapped embedding gives the minimum 0.93 probability.Source: §4.2 · p. 8, Figure 5 · p. 9
  • For code generation, the swap makes the maximum-finding loop use < with probability 0.40 without demonstrations, falling to 0.11 with four.Across both tasks, static embeddings change outputs while consistent context progressively overrides them, consistent with the context-dependent gate controlling memory contribution.Source: §4.2 · pp. 8–9, Figure 5 · p. 9
  • With the question unchanged, remapping a STEM vocabulary row from 60 to 120 shifts intermediate quantities and the final answer to match a rewritten question.This intervention illustrates how a targeted memory change can propagate through downstream reasoning.Source: §4.2 · p. 9

5 Related Works

Related work covers static token-indexed and n-gram memory, sparse memory adaptation, and sparse upcycling. STAPLE focuses on using deterministically indexed memory to scale downstream performance while retaining pretrained FFN computation.

Static sparse memory

Memory-based adaptation

  • Memory-based adaptation studies include learned key–value layers and sparse memory finetuning, whereas STAPLE examines deterministically indexed memory for broad downstream performance during continued training.The cited work includes Lample et al. (2019), Berges et al. (2024), Lin et al. (2025), and Goyal et al. (2026).Source: §5 · p. 9

Sparse upcycling

  • MoE upcycling reuses pretrained dense weights, with domain-specialized branches, informed router initialization, and partial expert reinitialization among approaches to expert specialization.The cited methods include Komatsuzaki et al. (2023), He et al. (2025), Vavre et al. (2024), Sukhbaatar et al. (2024), Jiang et al. (2025), and Nakamura et al. (2025).Source: §5 · p. 9
  • STAPLE expands token-indexed memory while retaining dense FFN computation, avoiding cloned expert pools and learned routing.Its focus is making sparsely and unevenly updated parameters useful under a limited training budget while largely preserving seed-model throughput.Source: §5 · p. 10

6 Conclusion

STAPLE expands pretrained dense models during mid-training by preserving learned computation and complementing it with scalable static memory. The paper reports consistent gains as memory scales and a favorable score–efficiency comparison with an MoE upcycling baseline, while identifying joint scaling with MoE as future work.

6 Conclusion

  • STAPLE is presented as a simple, sample-efficient way to expand pretrained dense-model capacity during mid-training.Its central idea is to preserve pretrained knowledge and complement it with scalable memory.Source: §6 · p. 10
  • Compared with a serving-efficient MoE upcycling baseline, STAPLE achieves a higher average score, trains faster, and decodes 2× faster with a much smaller GPU memory footprint.The conclusion also reports consistent gains with increasing memory scale, notably on reasoning-heavy tasks.Source: §6 · p. 10
  • Jointly scaling static memory with MoE is identified as a future direction for pushing the performance–efficiency frontier.Source: §6 · p. 10

A Training hyperparameters

Table 3 lists training hyperparameters for warmup and mid-training across the static-memory, Engram-style, and MoE configurations.

A Training hyperparameters

  • During STAPLE warmup, augmented-layer up-projections are frozen; MoE expert weights instead use a separate cosine schedule during warmup and linear decay during mid-training.The MoE peak learning rate is 2 × 10−4 in both stages.Source: Table 3 · p. 14

B Serving latency and throughput

Serving measurements cover batch-size-1 prefill and decode, while a separate GPU energy run finds STAPLE-12L close to dense and below the MoE baseline.

B Serving latency and throughput

  • Table 4 reports batch-size-1 prefill and decode latency and throughput, with all serving measurements at context length 256 on one NVIDIA B200.STAPLE-12L offloads 12 embedding tables totaling 19.73 GB to CPU DRAM and uses a 1 GB GPU cache.Source: B Serving latency and throughput · p. 14, Table 4 · p. 14
  • The MoE serving baseline uses eight interleaved layers with eight full-width experts per layer, top-1 routing, and all experts resident in HBM.The table caption notes routing and token-reordering overhead but gives no latency or throughput values in the supplied passage.Source: Table 4 · p. 14
  • 18.92 mJ per generated token for STAPLE-12L is approximately 1.2% above dense and 58.3% below the MoE baseline.The measurements used one otherwise idle B200 GPU at batch size 64 and context length 256.Source: B Serving latency and throughput · p. 14

C Implementation Details

The implementation keeps large embedding tables in CPU DRAM and retrieves only workload-required rows; inference uses distinct prefill and decode paths with retrieval overlapped with backbone execution.

C Implementation Details

  • Large STEM tables remain in CPU DRAM, while inference uses a small GPU-resident LFU cache and training materializes a compact buffer over the distinct token set.The buffer is described as spanning the distinct token set from Section 2.4.Source: C Implementation Details · pp. 14–15

C.1 Inference-time embedding offload

  • Prefill favors bulk deduplicated transfers, whereas decode uses GPU-driven fine-grained lookup to minimize exposed latency.Both retrieval paths overlap embedding retrieval with backbone execution.Source: Table 5 · p. 15, §C.1 · p. 15
  • For very short prefills, the implementation uses the GPU-driven mapped-host decode path to avoid CPU gathers and explicit H2D copies.This path is used when CPU-gather setup dominates.Source: §C.1 · p. 15

C.2 Training-time embedding offload

Training offloads full STEM tables and optimizer states to host memory, materializing only active rows on the GPU and staging compact gradients asynchronously.

C.2 Training-time embedding offload

  • The token union and inverse map are computed once per step and reused by every STEM layer, while full tables and fp32 AdamW moments stay in host memory.The training pipeline is identified as Algorithm 1.Source: §C.2 · p. 15
  • When compact gradients become ready, the pipeline enqueues asynchronous D2H copies into an alternate staging buffer.The next-step gathers follow forward consumption order to bound exposed pipeline-fill latency for the first augmented layer.Source: §C.2 · p. 15
  • The compact U × dff row buffer is the autograd leaf, so backward avoids allocating a dense V × dff gradient.Gradient staging is double buffered, and H2D and D2H transfers use separate streams.Source: §C.2 · p. 15
  • Table allocations are interleaved across NUMA memory controllers, with optional huge pages, while a passive OpenMP idle policy avoids competition with the GPU-launch thread.A dedicated host worker polls latency-critical CUDA-event waits.Source: §C.2 · p. 15

D Tracing an implicit numerical update through downstream reasoning

A vocabulary-row remapping tests whether static memory can alter a numerical premise during multi-step GSM8K reasoning. The intervention changes downstream quantities consistently with a rewritten-question control, but the example does not establish general edit reliability.

D Tracing an implicit numerical update through downstream reasoning

  • The intervention remaps the vocabulary row for 60 to the row for 120 wherever the source id occurs, including generated continuations, while leaving input text and dense weights unchanged.This uses four-shot chain-of-thought context on examples the unmodified model solves correctly, separate from the five-shot benchmark evaluation.Source: D Tracing an implicit numerical update… · p. 16
  • 60→120 in the STEM row remapping changes the monthly quantities from 60, 180, and 126 to 120, 360, and 252, and the total from 366 to 732.Figure 6 shows the original output, row update, and rewritten-question control; the row update matches the control's intermediate quantities and final answer.Source: Figure 6 · p. 16, D Tracing an implicit numerical update… · p. 16
  • Agreement with the rewritten-question control supports interpreting the row update as changing the effective numerical premise used by subsequent computation.The example does not establish general edit reliability or an isolated fact representation, and it cannot separate prompt-processing effects from generated-token feedback.Source: D Tracing an implicit numerical update… · p. 16

Related papers

Loading 2610.12599v1…