Static memory expands pretrained models without sparse pretraining
In plain words
STAPLE is a method for expanding a pretrained language model with a large lookup memory. Each input token selects a row from that memory, and the selected information is added to the model's existing feed-forward computation. The new rows start at zero, so the model initially behaves like the original.
Researchers compared this design with continued training of the dense model and with other sparse designs. On OLMo-2 1B, STAPLE raised the average score across ten tasks from 44.39 to 45.59 under the reported training budget.
The research
The practical question is whether extra parameters can add useful capability after pretraining without throwing away the computation already learned or making the model expensive to serve.
That matters because sparse expansion can avoid full retraining, but expert-based designs add routing, memory, communication, and fragmentation costs, while replacement-style static memory may waste limited training on relearning.
STAPLE keeps the pretrained feed-forward path and adds zero-initialized token memory, improving tested scores during limited continued training.
Researchers compared additive memory with replacement in a four-layer ablation, then tested STAPLE-12L against dense continuation and parameter-matched static alternatives on OLMo-2 1B.
Its average score across ten tasks rose from 44.39 to 45.59, with gains in coding, mathematics, and reading comprehension.
The retained pathway and sharp score drop after table removal support a role for learned static memory.
The conclusion is limited to the reported OLMo-2 settings: the 7B test used a reduced training budget, and table removal was performed without retraining on a checkpoint whose remaining weights had co-adapted with memory.
Analysis
Deep
- Authors
- Ranajoy Sadhukhan, Euijin Hong, Sheng Cao, Zechun Liu, Shuowei Jin
- License
- CC BY 4.0