Search papers, labs, and topics across Lattice.
The paper introduces StateLM, a novel language model architecture that incorporates an internal reasoning loop and memory management tools (context pruning, document indexing, note-taking) to actively manage its own context. This allows the model to dynamically engineer its context, overcoming the limitations of fixed-window architectures in standard LLMs. Experiments demonstrate StateLM's superior performance on long-document QA, chat memory, and complex research tasks, achieving significant accuracy gains compared to standard LLMs across various model sizes.
LLMs can now actively manage their own context using internal memory tools, leading to accuracy gains of up to 52% on complex reasoning tasks compared to standard LLMs.
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.
