MEMENTO: Teaching LLMs to Manage Their Context
Vasilis Kontonis · Yuchen Zeng · Shivam Garg · Lingjiao Chen · Hao Tang · Ziyan Wang · Ahmed Hassan Awadallah · Eric Horvitz · John Langford · Dimitris Papailiopoulos
Memento teaches reasoning LLMs to segment chain-of-thought into blocks, compress each into a summary, and mask completed blocks, while maintaining strong reasoning capabilities. We release OpenMementos (228K annotated traces).
摘要
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 \emph{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, OLMo3) 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 up to $2\times$ throughput improvement while also enabling us to perform RL and further improve accuracy. Finally, we identify a \emph{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 15pp on AIME24.