The Pensieve Paradigm: Stateful Language Models Mastering Their Own Context

TL;DR

This paper introduces StateLM, a new class of language models capable of actively managing their own context and memory tools, leading to significant improvements in long-document understanding and reasoning tasks.

Contribution

We propose StateLM, a model with an internal reasoning loop and memory management tools, enabling dynamic context engineering and surpassing standard LLM performance.

Findings

StateLM outperforms standard LLMs on long-document QA tasks.

StateLM achieves 10-20% accuracy improvements on chat memory tasks.

StateLM reaches up to 52% accuracy on BrowseComp-Plus, compared to 5% for standard models.

Abstract

In the world of Harry Potter, when Dumbledore's mind is overburdened, he extracts memories into a Pensieve to be revisited later. In the world of AI, while we possess the Pensieve-mature databases and retrieval systems, our models inexplicably lack the "wand" to operate it. They remain like a Dumbledore without agency, passively accepting a manually engineered context as their entire memory. This work finally places the wand in the model's hand. We introduce StateLM, a new class of foundation models endowed with an internal reasoning loop to manage their own state. We equip our model with a suite of memory tools, such as context pruning, document indexing, and note-taking, and train it to actively manage these tools. By learning to dynamically engineering its own context, our model breaks free from the architectural prison of a fixed window. Experiments across various model sizes demonstrate StateLM's effectiveness across diverse scenarios. On long-document QA tasks, StateLMs consistently outperform standard LLMs across all model scales; on the chat memory task, they achieve absolute accuracy improvements of 10% to 20% over standard LLMs. On the deep research task BrowseComp-Plus, the performance gap becomes even more pronounced: StateLM achieves up to 52% accuracy, whereas standard LLM counterparts struggle around 5%. Ultimately, our approach shifts LLMs from passive predictors to state-aware agents where reasoning becomes a stateful and manageable process.