PaperlayerView exact revision on arXiv

Source-linked AI summary

MEMENTO: Teaching LLMs to Manage Their Own Context

Vasilis Kontonis, Yuchen Zeng, Shivam Garg, Lingjiao Chen, Hao Tang, Ziyan Wang, Ahmed Awadallah, Eric Horvitz, John Langford, Dimitris Papailiopoulos

arXiv:2604.09852v1cs.AIcs.LG

TL;DR

Long reasoning traces lack mechanisms for organizing and compressing intermediate state. MEMENTO trains models to segment reasoning into blocks, summarize each block, and continue with sparse attention, achieving strong reasoning accuracy with lower KV-cache use while retaining information in memento KV states.

  • Problem

    Reasoning models produce long, unstructured chains of thought without a built-in mechanism to organize intermediate results or discard past tokens.

  • Method

    MEMENTO trains models to segment reasoning into blocks, compress each block into a memento, and mask completed blocks while retaining their memento KV states.

  • Results

    Across three model families, MEMENTO reduces peak KV cache by 2–3× while preserving strong reasoning accuracy; removing implicit KV information lowers AIME’24 accuracy by 15 pp.

  • Takeaways & Limitations

    OPENMEMENTOS provides 228K annotated reasoning traces, and native vLLM block masking supports higher-throughput inference and RL fine-tuning.

Abstract

from arXiv · show

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.

A. Trace Selection

MEMENTO organizes reasoning traces into alternating thinking blocks and mementos, using sparse attention to remove completed blocks from the KV cache.

  • Sparse attention: MEMENTO generates alternating thinking blocks and mementos, then masks each completed thinking block during inference.The model attends to the current block and past mementos rather than retaining all preceding thinking tokens.
  • SFT data generation: The SFT pipeline splits reasoning traces into sentences, scores candidate boundaries, optimizes the cuts, and summarizes each block.
  • Memory reduction: Peak KV cache falls by approximately 2–2.5× across benchmark categories.

1 Introduction

MEMENTO addresses the lack of organization and compression in long reasoning traces by teaching models to create compact mementos and mask completed blocks. Across model families and scales, it preserves strong reasoning accuracy while reducing KV-cache use, and its cached memento states retain additional information from masked blocks.

  • Motivation: Reasoning models lack a built-in mechanism to organize intermediate results or discard past tokens, leaving long chains of thought as flat, costly streams.
  • Method: MEMENTO segments chain-of-thought into coherent blocks, compresses each into a memento, and attends only to past mementos plus the current block.Each block is compressed to approximately 5–20× smaller size on average.
  • Mechanism: A dual information stream combines explicit memento text with implicit information retained in memento KV states after masking.
  • Mechanism: 15 pp accuracy loss on AIME’24 occurs when memento KV states are recomputed without the preceding block context.
  • Results: Across Qwen3, Phi-4-reasoning, and Olmo-3-7B-Think, MEMENTO reduces peak KV cache by 2–3× while maintaining strong reasoning accuracy.Qwen3-32B loses 2.6 pp on AIME’26, and the average accuracy gap across five benchmark groups is 3.5 pp at 32B versus 6.3 pp at 8B.
  • Training data: OPENMEMENTOS provides 228K segmented and summarized reasoning traces for training MEMENTO models.

2 Related Work

Prior work manages context through external memory systems, textual carryovers, learned gist representations, shorter traces, or latent compression. MEMENTO belongs to the line that compresses reasoning while retaining natural-language summaries and using in-place masking.

  • External systems: External context-management systems use summarizers, memory modules, or orchestration logic rather than teaching models to manage context internally.
  • Textual carryovers: InftyThink, Accordion-Thinking, and The Markovian Thinker compress reasoning chunks into carryovers, with several methods adding RL to improve continuation.
  • Learned context management: PENCIL trains models from scratch to erase intermediate reasoning through reduction rules on synthetic tasks.Its reported demonstrations use 25M-parameter models, 2K context length, 3-SAT, and Einstein’s Puzzle.
  • Latent compression: Gist-token approaches evict original tokens while encoding compressed chunk information entirely in hidden representations, losing interpretability.
  • Efficiency methods: Other approaches reduce memory through token skipping, compressed traces, shorter RL-trained reasoning, or latent representations.
  • Naming overlap: Memento is also used by other works for external episodic memory in inference-time agent adaptation.

3 OPENMEMENTOS Dataset

OPENMEMENTOS converts unstructured reasoning traces into coherent blocks paired with dense mementos through structure-aware splitting, scored boundaries, algorithmic segmentation, and iterative refinement. The resulting 228K-sample dataset provides stable summaries and substantial block-level compression across math, code, and science.

  • Dataset motivation: OPENMEMENTOS contains 228K traces annotated with semantically coherent blocks and mementos because ordinary reasoning traces lack natural boundaries.
  • Design rationale: Direct LLM segmentation failed because it had to jointly optimize block coherence, size balance, and semantic boundaries across many partitions.
  • Pipeline: The pipeline factors segmentation into local boundary scoring, algorithmic global optimization, and LLM-based summary judging.
  • Trace preparation: The dataset starts from 228K OpenThoughts-v3 reasoning traces and protects code and multi-line mathematics during sentence splitting.
  • Sentence splitting: Structure-aware splitting reduces candidate boundaries by approximately 2×, from 397 to 187 per trace on average.
  • Segmentation: Boundary scores range from 0 for mid-thought locations to 3 for major transitions, and segmentation favors strong boundaries while penalizing uneven block sizes.
  • Memento generation: Mementos target terse, logically complete state representations containing information needed by subsequent blocks.The target length is approximately 15–25% of the original tokens.
  • Quality control: Two judge-refinement iterations raise the memento pass rate from 28% to 92% at the ≥8/10 quality threshold.

4 Training the MEMENTO Models

MEMENTO models are trained with a two-stage SFT recipe that first teaches the block–memento format and then trains under masked-context attention. Across model families, the method generally preserves reasoning accuracy while reducing KV-cache usage, with savings constrained by hybrid attention architectures and occasional excessive generation.

  • Training procedure: Two-stage SFT separates learning the block–memento format from operating without access to masked content.Stage 1 uses full causal attention; later training introduces the harder masked-context constraint.
  • Training data scaling: All three data-scaling methods improve monotonically from 1K to 100K examples, while vanilla OpenThoughts achieves the highest accuracy at every budget.The comparison uses Qwen2.5-7B-Instruct on AIME24 and AIME25.
  • Accuracy and memory: The accuracy gap within Qwen3 shrinks from −6.3 pp at 8B to −3.5 pp at 32B across five benchmark groups.This pattern suggests larger models manage compressed context more effectively.
  • KV-cache behavior: Peak KV cache and KV AUC fall by 2–3× and 2–3.5×, respectively, on competition math, although individual problems can exhibit higher total KV cost from excessive generation.Figure 4 includes a case where MEMENTO’s KV area-under-curve is 2.1× higher than the base despite a lower peak.
  • Architecture effects: Olmo-3-7B transfers without architecture-specific changes, but its hybrid sliding-window attention limits peak savings to ∼0.85–0.93× and can worsen AUC.Only its eight full-attention layers benefit from eviction because 75% of layers already cap their cache at 4096 tokens.
  • Accuracy and memory: MEMENTO maintains strong reasoning performance across Qwen3, Phi-4-reasoning, and Olmo-3-7B-Think while reducing peak KV cache by 2–3× on uniform-attention models.On AIME’26, Qwen3-32B reaches 72.6% versus 75.2% for the base model, a 2.6 pp gap.
  • Compression behavior: Mementos usually remain compact: across four model families and four benchmarks, they measure ∼260–615 characters, with strongest compression on competition math at 9–27×.The bulk of blocks achieve 5–20× compression, while shorter-block benchmarks show 6–9× compression.

5 Improving Accuracy via RL

RL is applied after MEMENTO SFT to improve single-sample accuracy while retaining compressed-context inference. The experiments show that majority voting can recover baseline accuracy and that RL further improves several benchmarks while preserving much of the KV-cache advantage.

  • Coverage under compression: With 64 completions per problem, MEMENTO’s pass@64 coverage is nearly identical to the base, with an average gap of 2.6 pp and 96.4% solved-set Jaccard similarity.The comparison covers AIME 2024, 2025, and 2026 across three model families.
  • RL setup: RL rollouts use sparse block-masked attention through the custom vLLM engine, matching the inference-time masking pattern.CISPO is used as the RL optimizer, with a KL penalty to limit drift from the SFT checkpoint.
  • RL results: CISPO fine-tuning raises Qwen3-8B MEMENTO accuracy from 57.3 to 64.9 on AIME’26 and from 45.1 to 49.4 on Comp. Math.GPQA-D also rises from 55.8 to 62.9, above the 61.4 vanilla baseline.
  • RL trade-off: After RL, peak KV cache increases from 1.08 to 1.48 GB but remains below the 2.71 GB vanilla footprint.The result preserves much of MEMENTO’s memory advantage while improving single-sample accuracy.
  • Recovering accuracy: Majority voting at k=3 recovers base-model accuracy without additional training, indicating that SFT mainly reduces consistency rather than removing knowledge.Figure 6 reports majority-vote headroom for the three MEMENTO SFT models.

6 Inference and the Implicit KV Channel

MEMENTO combines in-place block masking with retained memento KV states, reducing serving memory pressure while preserving an implicit channel from masked reasoning blocks. Ablations and probing show that this channel contributes materially to downstream accuracy and carries recoverable, directional information.

  • Serving throughput: 1.75× higher token throughput and 1.58× faster batch completion are achieved by MEMENTO on Qwen3-8B at full concurrency.The comparison uses AIME24 × 8 repetitions, 240 requests, 32K maximum tokens, and one B200 GPU.
  • Memento attention: Memento attention retains KV entries computed with block context after masking the original block, so future tokens can attend to implicitly encoded block information.Normal memento attention generates each memento while attending to its full block, then masks the block while retaining the memento KV cache.
  • KV ablation: A 15 pp accuracy drop after KV recomputation shows that memento KV states carry significant information from masked blocks beyond standalone memento text.The restart condition discards the original cache and recomputes KVs with past blocks masked, removing their implicit block information.
  • Implicit KV probing: 26.7% and 23.0% masked passcode accuracy for Qwen3-8B and Qwen3-32B exceed the 10% chance baseline, while causal controls remain at chance.The probe injects passcodes into a target block and predicts digits from downstream memento KV states that cannot directly attend to that block.
  • Implicit KV probing: Deeper layers carry more leaked signal, with masked accuracy rising from 10.8% to 26.5% in Qwen3-8B and from 12.8% to 22.4% in Qwen3-32B.The same deeper-layer concentration appears under the direct probing condition.

7 Conclusion

The conclusion presents MEMENTO as a learned approach for compressing reasoning context while retaining strong accuracy across models and benchmarks. Its central finding is that explicit memento text and implicit KV representations jointly preserve information from masked blocks.

  • Method: MEMENTO teaches models to segment reasoning, compress blocks into mementos, and mask completed blocks with sparse attention.The method targets internal context management during reasoning.
  • Results: 2–3× peak KV cache reduction is demonstrated across Qwen3, Phi-4-reasoning, and Olmo-3-7B-Think while preserving strong reasoning accuracy.The conclusion summarizes results across three model families and multiple scales.
  • Implicit KV channel: Removing the implicit KV channel degrades accuracy by 15 pp, distinguishing MEMENTO from methods that discard context after summarization.Mementos transmit information through both explicit summary text and implicit KV representations.
  • Resources: OPENMEMENTOS provides 228K annotated reasoning traces, while a vLLM fork supports native block masking for further inference and RL research.The released infrastructure is intended to facilitate additional research on MEMENTO.

A.1.1 Full Prompts

The full prompts specify how MEMENTO segments reasoning and compresses each block into a judged, information-dense state summary. They emphasize preserving mathematical continuity and critical reasoning state.

  • Boundary scoring: Boundary scores from 0 to 3 assess whether a reasoning trace can be split without disrupting semantic coherence, topic continuity, or logical flow.Scores of 2–3 indicate clear transitions or natural chapter boundaries.
  • Boundary scoring: Mathematical derivations must remain intact, with low scores for breaks after colons, equations, calculation introductions, or continuing derivation markers.High scores are reserved for completed derivations followed by topic shifts.
  • State compression: The state compressor produces a minimal, information-dense summary preserving definitions, variables, assumptions, constraints, and key intermediate results.Its objective is to minimize summary tokens while retaining logically relevant information.
  • Summary judging: A separate judge scores mementos on formulas, numerical values, methods, validation, hallucinations, and result-first structure.The rubric evaluates whether the compressed state can replace the original reasoning block.

A.1.2 Worked Examples

The worked example shows iterative refinement of a memento for an NBA playoff probability calculation. Judge feedback converts an incomplete procedural summary into a representation containing the computational state needed for continuation.

  • Initial memento: An initial memento for an NBA playoff probability block scores 5/10 because it describes the approach but omits critical formulas.The underlying block defines f(n,a,b), gives f(0,0,0)=1, targets f(6,3,3), and specifies alternating home-team probabilities.
  • Iterative refinement: After feedback requesting the recurrence and transition probabilities, the refined memento scores 8/10 and captures the function, base case, target, and recurrence relation.The refinement restores the specific computational details needed to represent the block’s full state.

A.2 Training and Evaluation Details

MEMENTO training uses staged SFT to teach block-memento formatting before sparse attention, with custom block masking carried through training and inference. The implementation tracks block state in the KV cache, evicts completed reasoning blocks after summaries, and supports evaluation, RL improvements, and architectural probing.

  • Training procedure: Two-stage SFT first trains with full attention, then adapts the model to memento attention using identical loss and hyperparameters but a sparse block mask.Stage 1 selects the best AIME24 checkpoint; Stage 2 masks completed block content from subsequent queries.
  • Sparse attention implementation: The block cache identifies BLOCK, SUMMARY, and OTHER tokens, then masks completed reasoning tokens after each <|summary end|> is generated.Its state persists across autoregressive steps through the KV cache, avoiding full-sequence rescanning.
  • Data format: 32,768-token ChatML sequences encode each assistant response as repeated reasoning-block and summary spans followed by a final answer.Training data is pre-tokenized in HuggingFace Arrow format, with dynamic padding handled by the batch collator.
  • RL and efficiency: MEMENTO+RL improves every reported benchmark over MEMENTO SFT while increasing peak KV from 1.08 to 1.48 GB, still 45% below vanilla’s 2.71 GB.KV AUC rises from 10.7 to 16.4 GB·ktok, roughly half of vanilla’s 30.9.
  • Training ablation: The full three-stage pipeline enables block masking while retaining strong performance, whereas direct OPENMEMENTOS training substantially underperforms vanilla reasoning SFT.On Qwen2.5-7B, the pipeline reaches 32.7% on AIME 2024, 28.7% on AIME 2025, and 45.5% on GPQA-Diamond, the highest GPQA-Diamond score among configurations.
  • KV-channel analysis: Toy-transformer probes reproduce production-model leakage, with masked accuracy reaching 26.2% in the last layer versus 13.1% in the first and remaining above chance at block index=+7.Masked leakage stays roughly constant at 17–19% across checkpoints while task accuracy rises from 77% to 95%.
Loading 2604.09852v1…