Stateful memory traversal improves learning speed in long context models
Retrieval as Traversal: Value-Indexed Memory for Stateful Long-Context Agents
Artificial Intelligence
Summary
Long-context models try to remember lots of information over time, but searching through everything each step can be slow. The authors suggest a new way to keep track of memories by moving step-by-step through nearby related memories, which can be faster to learn. They tested a simple version of this on a small task where the model needed to remember a few things, and it learned quicker than other methods. Their idea is that remembering in a stateful, connected way is easier than treating each retrieval independently. More experiments are needed to test the full method.
What this means in practice
- •For ml engineers: Use stateful memory traversal to speed up training of models handling long sequences with overlapping information.
- •For robotics developers: Enable robots with limited memory tokens to recall related past states more efficiently in tasks like navigation.
Tested on one dataset.
Authors
Daniel Strauß, Zhenshan Bing, Alois Knoll
Abstract
Long-context models repeatedly solve closely related retrieval problems. Attention or global retrieval can retain large stores, but each forward pass again searches for what matters. Recurrent compression avoids this search only by deciding earlier what information to preserve. Neither mechanism directly exploits that useful memories at neighbouring retrieval steps may substantially overlap. We propose an alternative in which retrieval is a persistent local traversal through an address space. We separate a memory's content from its address. Utility is defined by the information revealed by retrieving the content, while geometry determines only how cheaply that information can be reached. This leads to a cost-regularised indexing objective that preserves information value inside local neighbourhoods while charging for every memory made locally accessible. Its marginal-value form gives both attractive and exclusion pressures. Thus, complementary memories should be co-retrievable when the additional value exceeds retrieval cost, whereas redundant memories should not occupy the same retrieval set solely because they are individually useful. As a proof of principle, we test only the simplest fixed-geometry instance on MiniGrid Memory. A persistent local traversal head reaches $0.95$ mean return after $0.59$M training steps with a five-token context, compared with $0.94$M for Gaussian absolute addressing and more than $2$M for categorical absolute addressing. The experiment does not test learned geometry or repeated retrieval. Instead, it motivates the hypothesis that stateful retrieval may be easier to learn than stateless. To test the full proposed method further experimentation is necessary.