MEMENTO: Teaching LLMs to Manage Their Own Context

TL;DR

MEMENTO introduces a method for teaching large language models to segment reasoning, compress intermediate states into summaries, and attend only to these summaries, reducing context size and computational load.

Contribution

The paper presents a novel training approach and dataset for enabling models to manage their reasoning context via segmentation and summarization, improving efficiency and accuracy.

Findings

Models trained with MEMENTO maintain strong accuracy on benchmarks.

Achieves approximately 2.5x reduction in KV cache usage.

Extends vLLM to improve throughput by about 1.75x.

Abstract

Reasoning models think in long, unstructured streams with no mechanism for compressing or organizing their own intermediate state. We introduce MEMENTO: a method that teaches models to segment reasoning into blocks, compress each block into a memento, i.e., a dense state summary, and reason forward by attending only to mementos, reducing context, KV cache, and compute. To train MEMENTO models, we release OpenMementos, a public dataset of 228K reasoning traces derived from OpenThoughts-v3, segmented and annotated with intermediate summaries. We show that a two-stage SFT recipe on OpenMementos is effective across different model families (Qwen3, Phi-4, Olmo 3) and scales (8B--32B parameters). Trained models maintain strong accuracy on math, science, and coding benchmarks while achieving ${\sim}2.5\times$ peak KV cache reduction. We extend vLLM to support our inference method, achieving ${\sim}1.75\times$ throughput improvement while also enabling us to perform RL and further improve accuracy. Finally, we identify a dual information stream: information from each reasoning block is carried both by the memento text and by the corresponding KV states, which retain implicit information from the original block. Removing this channel drops accuracy by 15\,pp on AIME24.