Apple study finds stronger gating can slow in-context learning in Mamba

Apple and Harvard researchers found that stronger gating can lead Mamba state-space models to learn a parameter-stored shortcut before retrieving the correct answer from the current prompt. The study of gating in state-space models, published by Apple in November, also found that the same mechanism often helped models handle sequences longer than those used in training.

The finding describes an optimization tradeoff in the paper’s theoretical model and tested Mamba configurations, not a universal limit for every state-space model. Some retrieval failures can arise from training dynamics even when a model can represent the correct retrieval solution.

The researchers separated two behaviors using tasks in which an input had a fixed underlying label but the label shown in each prompt could be randomly changed. In-weights memorization meant storing the usual input-to-label mapping in model parameters. In-context learning meant finding the matching input in the current prompt and returning the label attached there. Memorization could therefore score above chance while still ignoring the prompt-specific answer; perfect performance required retrieval.

In the paper’s simplified Mamba-like model, gating exponentially reduced the influence of earlier tokens. Stronger gating preserved less distant context, weakening the training signal for retrieval while leaving the model able to learn the fixed mapping in its weights.

The authors described the result as a memorization plateau, not necessarily a permanent outcome. Retrieval remained reachable in the reduced theoretical dynamics, but the estimated escape time grew sharply as the initial decay strengthened. Weak gating made retrieval easier to learn, while also allowing noise from more irrelevant tokens to accumulate in long contexts. Stronger decay reduced that noise but imposed locality, so useful information still had to remain within the model’s effective context.

The multi-token retrieval experiments used eight-token keys and values, changed each value token with 10% probability, and trained on contexts containing up to 512 key-value pairs. Memorizing the fixed dictionary produced about 43% sequence accuracy. Across five random seeds, the default Mamba initialization did not leave that baseline during 100,000 training steps, while weaker-gated variants reached retrieval sooner. When evaluation contexts grew to as many as 2,048 pairs, stronger-gated models generalized better.

A more naturalistic test fine-tuned 1.5-billion-parameter Mamba models on a variable-pool version of the Berkeley Function Calling Leaderboard after pretraining on 300 billion tokens. Within the tested range, weaker gating produced moderately better accuracy and fewer syntactically valid calls to tools absent from the prompt. Gating strength did not measurably change length generalization in that experiment, although all tested Mamba variants generalized better than the 2-billion-parameter Transformer baseline as tool pools grew.

The authors said the study does not immediately establish a remedy. Simply weakening, reinitializing or removing gates may sacrifice long-context performance, and the results do not establish the same tradeoff across all architectures, model sizes, datasets or production deployments.