PaperlayerView exact revision on arXiv

Source-linked AI summary

The Pensieve Paradigm: Stateful Language Models Mastering Their Own Context

Xiaoyuan Liu, Tian Liang, Dongyang Ma, Deyu Zhou, Haitao Mi, Pinjia He, Yan Wang

arXiv:2602.12108v1cs.AI

TL;DR

Existing LLMs passively operate within externally engineered context and lack agency over their own memory. StateLM trains models to manage context through a learned reasoning loop and memory operations. It reports gains across long-document QA, chat memory, and deep research, while acknowledging context-window limitations from stubs and delayed pruning.

  • Problem

    Current LLMs are stateless passive predictors that rely on externally provided context and cannot actively manage their own memory or reasoning process.

  • Method

    StateLM trains foundation models to self-engineer context using a reasoning loop with memory operations including context deletion, reading, and persistent note-taking.

  • Results

    StateLM consistently outperforms instruct baselines and existing agentic methods across long-context tasks, with 10%–20% chat-memory gains and over 40% deep-research improvement.

  • Takeaways & Limitations

    The results support shifting from passively accumulating context toward models that dynamically refine and control their reasoning states.

  • Takeaways & Limitations

    Deleted content leaves lightweight stubs that can consume context, while delayed or conservative pruning may cause transient context overflow.

Abstract

from arXiv · show

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.

1 Introduction

StateLM addresses the statelessness of conventional LLMs by learning to manage their own context through memory tools. Across long-document QA, multi-turn dialogue, and deep research, it improves performance while maintaining compact internal state.

  • Current LLMs passively complete sequences within externally provided contexts and cannot actively manage their own memory or reasoning process.
  • StateLM combines deleteContext, readChunk, and updateNote to read segments, preserve key information in notes, and remove redundant raw text.
  • 1/4 active context and 5% to 12% average accuracy gains across model scales distinguish StateLM from instruct baselines on long-document QA.
  • 10% to 20% absolute chat-memory gains and up to 52% BrowseComp-Plus accuracy demonstrate StateLM’s performance across diverse scenarios.On BrowseComp-Plus, vanilla LLM counterparts remain around 5%, producing an average improvement of over 40%.
  • StateLM shifts context engineering from human-written workflows to learned model agency through a general-purpose toolkit of memory operations.
  • The study reports significant gains across long-document QA, multi-turn dialogue, and deep research within one unified model.

2 Related Work: Human as the Wizard

Prior context-management systems generally leave the model inside human-designed workflows, even when they add retrieval, memory hierarchies, branching, or summarization. StateLM instead trains the model to become its own context engineer and architect its reasoning loop.

  • Context engineering commonly uses human-written workflows to retrieve and stuff relevant information into a stateless LLM’s prompt.
  • Agentic memory frameworks such as MemGPT and MemOS introduce memory hierarchies that page information into and out of the context window.
  • Context-Folding and ReSum make context management learnable through branching, folding, or periodic summarization of interaction histories.
  • These methods remain constrained by human-designed routines, fixed summarization schedules, or logic hard-coded into training objectives and tool definitions.
  • StateLM trains the model to become its own context engineer rather than merely adapt to a predefined context-management routine.

3 Methodology

StateLM replaces append-only interaction histories with a learned, stateful loop that actively manages context through reasoning, tools, and context edits. Its methodology combines expert-trajectory supervised learning with filtering and masked next-step training.

  • Problem Setup: StateLM models tool-augmented reasoning as sequential interaction between a policy, an environment, and an evolving textual state.At each round, the model produces a reasoning trace followed by a tool call or termination signal, and the environment returns an observation.
  • Problem Setup: Monotonic accumulation retains all prior actions and observations, eventually exhausting fixed context budgets and degrading long-horizon performance.The standard state grows by ordered concatenation, leaving earlier reasoning and raw tool outputs visible in later prompts.
  • StateLM Framework: StateLM converts the append-only interaction log into a mutable state by letting the policy control which past elements remain visible.Its state-update function can append interactions or modify visible context according to context-management actions.
  • StateLM Framework: The Pensieve combines a persistent external notebook with deletion operations that remove selected interactions while keeping recorded notes accessible.This mechanism supports long-horizon reasoning under fixed context budgets by separating durable notes from transient visible content.
  • StateLM Reasoning Flow: StateLM’s reasoning flow analyzes long inputs, builds searchable indexes, searches and reads relevant chunks, records notes, prunes context, checks budget, and finishes when sufficient information is gathered.The workflow repeatedly applies a search–read–note–delete cycle while the toolkit supplies stateful context-interaction operations.
  • Training: Supervised training uses filtered expert trajectories and progressively evolved states, masking loss to the final assistant turn at each training step.Trajectories are retained only when answers are correct and context-management behavior is appropriate, including timely pruning and necessary reading.

4 Experiment

Experiments across synthetic memory retrieval, long-context reasoning, chat memory, and deep research evaluate StateLM across model scales and compare it with instruct and agentic baselines. StateLM remains robust under long contexts, uses less active context, and improves performance across diverse tasks.

  • Long-Context Reasoning: StateLM outperforms instruct baselines on LongDocQA and Chat Memory while using only 1/4 of the context window, also surpassing other agentic methods at the same budget.Compared with Qwen3-8B, StateLM-8B improves Long Document QA by over 10 percentage points and Chat Memory by 13 points.
  • Long-Context Reasoning: StateLM’s Long Document QA gains are largest when answer evidence appears later, reaching 24 and 25 accuracy points over the instruct model for 4B and 8B variants in the 128–256K range.The evidence-position breakdown reports consistent StateLM improvements across all evidence ranges.
  • Long-Context Reasoning: StateLM-14B-RL achieves up to 52% accuracy on BrowseComp-Plus, compared with around 5% for Qwen3-14B, while the StateLM series averages improvements exceeding 40 points over standard LLMs.The benchmark tests deep research through iterative reasoning over search results.

5 Further Analysis

StateLM’s tool-use behavior adapts to task difficulty and input length while keeping memory operations relatively infrequent. Deliberate training remains important because prompt-only agentic models frequently exceed context budgets and underperform trained StateLMs.

  • 5.1 Tool Use Pattern: StateLM increases reasoning rounds on harder benchmarks, where overall performance is lower.This pattern is reported especially for LongMemEval and BrowseComp-Plus.
  • 5.1 Tool Use Pattern: As input length grows, StateLM primarily increases targeted searches rather than memory updates.The model retains only key information in memory while navigating longer inputs.
  • 5.1 Tool Use Pattern: StateLM invokes memory operations less frequently than fixed-workflow methods that may update memory at every step.The reported pattern indicates relatively cost-efficient memory use across tasks.
  • 5.2 Agentic Baselines: Qwen3-14B agentic models answer up to 30% of questions, but trained StateLMs perform substantially better.Smaller agentic variants perform worse, and frequent context-budget overflow contributes to their errors.
  • 5.3 Error Analysis: BM25-based keyword retrieval often misses evidence for implicit or paraphrased queries because of limited semantic coverage.This is identified as a major StateLM failure mode.
  • 5.3 Error Analysis: Malformed tool calls cause a small fraction of failures, particularly in long-horizon trajectories and smaller models.The issue is classified as a formatting error.
  • 5.3 Error Analysis: Deleted-content stubs can accumulate, while untimely or conservative deletion may cause transient context overflow before pruning takes effect.These limitations can gradually consume the context budget over long trajectories.
  • 5.3 Error Analysis: Future improvements include larger management windows, more advanced search algorithms, and higher-quality long-horizon training trajectories.The authors leave these directions for future work.

6 Conclusion

The paper introduces StateLM as a state-aware agent that actively manages context through tool-based operations. Across memory retrieval and long-context reasoning tasks, it reports consistent gains over instruct baselines and existing agentic methods.

  • 6 Conclusion: StateLM actively manages its own context through tool-based operations across synthetic memory retrieval and realistic long-context reasoning tasks.The conclusion characterizes this as a state-aware agentic system.
  • 6 Conclusion: 10%–20% gains on the chat memory task accompany over 40% improvements on the agentic deep research task.These are the principal empirical outcomes reported in the conclusion.
  • 6 Conclusion: StateLM shifts reasoning from passively accumulated context toward dynamically refined and controlled reasoning states.The conclusion presents this as a direction for future foundation models and agentic systems.

A Training Dataset Construction

The training data construction section reports dataset statistics and a filtering-based procedure for preparing reinforcement-learning data from LongBench v2.

  • A Training Dataset Construction: The section provides SFT dataset statistics by filtering stage and RL dataset statistics.These are identified in the dataset-statistics tables.
  • A Training Dataset Construction: LongBench v2 multiple-choice questions are converted into open-ended questions by removing their original option lists.The conversion diversifies the problem type.
  • A Training Dataset Construction: An LLM-based judge determines transferability by checking whether the correct answer remains unique, verifiable, and unambiguous.The judge is instantiated as gpt-oss-120b.

B Training Configurations

The training configurations specify separate supervised fine-tuning and reinforcement-learning setups, including hardware, rollout, KL, sampling, and training-step settings.

  • B Training Configurations: SFT trains each model on two machines, each equipped with 8 H20 GPUs.The detailed SFT configurations are provided in the training table.
  • B Training Configurations: RL uses GRPO with rollout batch size 32, rollout number 8, and KL coefficient 0.001.KL loss is enabled during reinforcement-learning training.
  • B Training Configurations: RL samples 8 trajectories per rollout for the 8B model and 2 for the 14B model, with training lasting 32 steps.These settings produce the StateLM-RL variants.

C.1 Generation

The experiments use Qwen3 non-thinking mode and recommended sampling settings for StateLMs and instruct baselines, with benchmark-specific adjustments and distinct grading procedures.

  • Generation: Qwen3 non-thinking mode is used for StateLMs and instruct baselines because thinking mode requires a 32K output space per step.The authors state that thinking mode is computationally expensive and reduces the available context window.
  • Generation: Recommended Qwen3 sampling parameters are adopted for StateLMs and instruct baselines, with detailed settings reported in Table 8.
  • Generation: BrowseComp-Plus uses the official Qwen3 configuration, augmenting Qwen3 instruct models with a BM25 search tool matching StateLM’s search algorithm.
  • Evaluation: Long-document QA multiple-choice answers are graded with a rule-based script that also accepts correct option content without its corresponding letter.LongMemEval and BrowseComp-Plus follow official evaluation guidelines and use an LLM-based judge.
  • Evaluation: The RL reward module uses rule-based grading for multiple-choice questions and an LLM-based judge for open-ended questions.The judge model is gpt-oss-120b, and its grading prompt appears in Figure 8.

D Tool Use Analysis

This section reports StateLM tool-use statistics across benchmarks for the 4B and 8B models, alongside the prompts used for tool-assisted long-context processing and judging.

  • Tool-use statistics: Tool-use patterns are reported for StateLM-4B and StateLM-8B across benchmarks, with each table providing the benchmark’s mean input length.
  • Prompts: Figure 7 presents the system prompt used to train and infer StateLMs.
  • Evaluation prompt: Figure 8 shows the LLM Judge prompt, which evaluates student answers against correct answers and treats incomplete or ambiguous responses as incorrect.
  • Prompts: The long-context assistant prompt emphasizes factual answers grounded in provided context while minimizing context consumption.
  • Prompts: Figure 9 provides the system prompt for the Qwen3-Agentic models.
Loading 2602.12108v1…