Pretrained models gain sparse capacity without retraining: a zero-initialized memory path lifts ten-task average by 1.20 points
Synopsis
Pretrained dense models can be expanded during continued training instead of pretraining a sparse architecture from scratch. STAPLE keeps the pretrained feed-forward up-projection and adds zero-initialized token-indexed embedding tables as a residual pathway, so the new parameters learn complementary token-specific information. On OLMo-2 1B it raises the ten-task average from 44.39 to 45.59 over continued dense training, with HumanEval up 4.27 points, MBPP 2.80, GSM8K 2.27 and DROP 1.65, and gains persist on OLMo-2 7B.
Interpretation
When expanding a pretrained dense model, the existing computation can be kept and the new static sparse memory added as a complement rather than a replacement. Earlier static sparse approaches, notably STEM, replace the feed-forward up-projection with a token-indexed embedding table; in a four-layer ablation, that replacement fails to match continued dense training even with extra warmup. Under the same training budget, retaining the up-projection and adding the embeddings as a residual pathway performs substantially better; on OLMo-2 1B the ten-task average rises from 44.39 for continued dense training to 45.59.
The new embedding tables start at zero weights, which restores the seed model's computation graph, and a short warmup then teaches them general-purpose memory. The replacement approach lacks a functionality-preserving initialization: random, token-wise mean, principal-component and transformed input-embedding initializations did not improve long-term performance. Warmup trains on a subset of the pretraining corpus with the feed-forward up-projections frozen; it consistently improves MMLU and average commonsense reasoning over direct upcycling, but gives no consistent gains on advanced math and programming tasks.
Static memory capacity scales along two axes, the number of embedding-augmented layers and the embedding dimension, with little throughput cost. Going from 4 to 12 augmented layers raises the average downstream score from 45.10 to 45.59; a low-rank variant shrinks the token-indexed table to a quarter on 1B and reaches 44.97, still above the 44.39 of continued dense training. On OLMo-2 7B, augmenting the first 16 layers raises the average from 57.11 to 58.01 and improves 9 of 10 tasks, with average mathematics up 4.16% and average code up 3.68%.
Keeping the large tables in host memory and moving only the needed rows to the GPU lets billions of new parameters avoid slowing training or generation. The dense baseline keeps all weights resident on the GPU; the MoE upcycling baseline trains 13% slower and needs about 2.9 times the GPU memory for weights and optimizer state. STAPLE-12L holds 91.9 GiB of tables and optimizer moments in host DRAM, trains at 113.9K versus 117.0K tokens/s for dense, stays within 1.5% of dense generation throughput at batch size 64, and uses 58.3% less GPU energy per token than the evaluated MoE implementation.
Perspective
The reported scope is continued-training experiments on two dense backbones, OLMo-2 1B and 7B, using the Dolma 3 Dolmino Mix high-quality mixture, with warmup drawn from roughly 100B tokens of DCLM Baseline 1.0. The main 1B comparison uses about 100B warmup tokens followed by 100B mid-training tokens, while the 7B runs use about 30B warmup plus 70B mid-training because of higher computational cost. System-efficiency results come from a throughput-oriented decoding evaluation and one energy measurement on a single otherwise idle B200 at batch size 64 and context length 256.
The modest MMLU gain on 7B may reflect the shorter warmup, consistent with the 1B warmup ablation, although that explanation is not isolated in the 7B experiments. Removing all embedding tables does not return the model to the original dense seed, because the remaining weights have co-adapted with static memory. Token-row substitution shows the static embeddings causally change outputs, yet consistent context progressively overrides them, consistent with a context-dependent gate controlling how much the memory contributes. The numerical trace links a token-addressable memory intervention to coherent downstream reasoning, but does not establish general edit reliability or an isolated fact representation.
