MEMENTO: Teaching LLMs to Manage Their Own Context

Vasilis Kontonis; Yuchen Zeng; Shivam Garg; Lingjiao Chen; Hao Tang; Ziyan Wang; Ahmed Awadallah; Eric Horvitz

cs.AIApr 10, 2026

MEMENTO: Teaching LLMs to Manage Their Own Context

Authors: Vasilis Kontonis, Yuchen Zeng, Shivam Garg, Lingjiao Chen, Hao Tang, Ziyan Wang, Ahmed Awadallah, Eric Horvitz,

Organizations: Microsoft Research

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.

Figures & tables

Figure 2: OpenMementos dataset distributions by domain (228K samples). (a) Math and code have

\sim

9 blocks/sample; science has

\sim

7. (b) Block sizes range from 2.3K (science) to 3.8K (math) chars. (c) Summary sizes cluster around 509–603 chars across all domains, indicating a stable compression target. (d) Math achieves the tightest compression ratio (median 0.16) due to its larger blocks.

Panel of Figure 3 — Figure 3: Training data scaling. Pass@1 accuracy on AIME24 and AIME25 for Qwen2.5-7B-Instruct fine-tuned on 1K–100K examples. All methods improve monotonically with data size.

			AIME’26		Comp. Math		MATH-500		GPQA-D		LCB v6
			Val	$\Delta$	Val	$\Delta$	Val	$\Delta$	Val	$\Delta$	Val	$\Delta$
Qwen3-8B		Acc	66.8 1.1		54.3 0.3		90.5 0.9		61.4 2.4		73.1 1.0
	Base	Peak KV	2.41		2.71		0.84		1.23		1.76
		AUC KV	25.3		30.9		4.3		6.6		15.6
		Acc	64.7 1.1		49.2 0.3		89.7 1.0		57.8 2.5		70.0 1.0
	Control	Peak KV	2.59		2.82		0.88		1.60		1.89

Table 1: Memento achieves 2–3

\times

peak KV reduction on models with uniform attention layers while maintaining strong reasoning performance; RL on top of our Qwen3-8B further improves accuracy. Olmo-3-7B shows more modest savings (

\sim

0.85–0.93

\times

) due to its hybrid sliding-window architecture ( Section 4 ). The

\Delta

columns show changes of Memento and Mem.+RL relative to Control: accuracy deltas are additive (pp); KV deltas are multiplicative (Method / Control, so 0.39

\times

means 61% KV reduction). Metrics. Accuracy (%): pass@1 accuracy. Peak KV (GB): peak KV cache size; determines the minimum memory required to serve a request. AUC KV (GB

\cdot

ktok): area under the KV-occupancy-vs-token curve, capturing total memory-time cost that penalizes both large footprints and long generations (see Figure 4 for an illustration). Rows. For each model we report Base (unmodified), Control (Base fine-tuned on the same OpenThoughts source traces used to create OpenMementos , but without block/memento annotations, for the same number of training examples), and Memento (SFT on OpenMementos with block masking); for Qwen3-8B we additionally report Memento +RL ( Section A.2.3 ). Olmo-3 Memento competition-math evaluations use 8 generations per problem (vs. 64 for all other models) and are run with HuggingFace Transformers rather than vLLM. See Section A.2.2 for full benchmark and evaluation details.

Panel of Figure 4 — Figure 4: KV cache traces on individual problems (Qwen3-8B, both answers correct). (a) AIME24 P2: typical sawtooth pattern with 6 compactions; peak 0.77 vs 2.17 GB ( $2.8\times$ reduction). (b) AIME24 P26: Memento solves the problem in 12k tokens (vs 23k); frequent compactions keep peak at 0.82 vs 3.41 GB ( $4.2\times$ ). (c) AIME24 P5: Memento generates $3\times$ more tokens (31k vs 10k) with many compactions. Peak is still lower (1.27 vs 1.55 GB), but the total KV area-under-curve is $2.1\times$ higher than the base—a failure mode where block masking induces excessive generation.

Figure 5: CDF of compression ratio (summary chars / block chars). OLMo3-7B and Qwen3-8B achieve the tightest compression on competition math. Phi-4 has the widest compression spread, especially on MATH-500.

Model	Bench.	Base	Memento	Ret.	Jacc.
Qwen3-8B	AIME’24	93.3	90.0	96.4	96.4
	AIME’25	93.3	86.7	92.9	92.9
	AIME’26	86.7	90.0	100.0	96.3
Phi-4-r (14B)	AIME’24	93.3	93.3	100.0	100.0
	AIME’25	93.3	90.0	96.4	96.4
	AIME’26	93.3	90.0	96.4	96.4

Table 2: Problem coverage (pass@64) and solved-set overlap for Base vs. Memento on AIME (

n{=}64

per problem, 30 problems per benchmark).

Figure 6: Majority-vote headroom and Qwen3-8B CISPO ( MiniMax et al., 2025 ) RL trajectory. Left: AIME 2026 maj@

k

for the three Memento SFT models ( Table 1 ); horizontal lines show each Base model’s pass@1. All models match Base accuracy by

k{=}2

–

3

. Middle: per-step RL training accuracy (faint raw trace with 25-step moving average). Right: AIME’25 validation accuracy evaluated every 25 steps, peaking at 66.2% at step 350 (used in Table 1 ). Majority voting uses uniform tie-breaking among tied majority answers.

Panel of Figure 7 — Figure 7: Serving throughput (Qwen3-8B, 1 $\times$ B200 GPU). AIME24 $\times$ 8 repetitions (240 requests, 32K max tokens). Left: Memento sustains $1.75\times$ higher token throughput at full concurrency. Right: $1.58\times$ faster batch completion. Vanilla plateaus as KV cache fills GPU memory.

Inference mode	Pass@1
Memento att. (normal)	66.1%
Memento att. + restart	50.8%
$\Delta$	$-$ 15.3 pp

Table 3: KV ablation on full AIME24 (30 problems, Qwen3-8B memento attention checkpoint, 32K generation). Block masking accuracy is from the 64-repetition evaluation in Table 1 ; the restart experiment uses 8 repetitions.

Panel of Figure 8 — Figure 8: Restart ablation. Step 1 : $M_{i}$ is generated with full attention to $T_{i}$ (same as normal memento attention). Step 2 : KV cache is discarded and recomputed via prefill over prompt $+$ $M_{1..i}$ only— $T_{i}$ is masked, so $M_{i}$ ’s KV states no longer encode block information. The 15 pp accuracy drop ( Table 3 ) shows the KV channel carries significant reasoning capacity.

Appendix figures & tables4 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 10 — Figure 10: Iterative memento refinement on a reasoning block about NBA playoff probabilities. The initial memento (pass 1, score 5/10) describes the approach but omits critical formulas. After judge feedback requesting the specific recurrence and transition probabilities, the refined memento (pass 2, score 8/10) captures the full computational state: function definition, base case, target, and recurrence relation.

Benchmark	Source	# Problems	Answer format
AIME 2024	HuggingFaceH4/aime_2024	30	Integer (0–999), \boxed{}
AIME 2025	yentinglin/aime_2025	30	Integer (0–999), \boxed{}
AIME 2026	MathArena/aime_2026	30	Integer (0–999), \boxed{}
HMMT Feb 2023	MathArena/hmmt_feb_2023	30	Math expression
HMMT Feb 2024	MathArena/hmmt_feb_2024	30	Math expression
HMMT Feb 2025	MathArena/hmmt_feb_2025	30	Math expression

Appendix

Table 5: Benchmark details. Competition-math benchmarks use 64 generations per problem; MATH-500, GPQA, and LiveCodeBench use 2 generations. All use temperature 0.6 and 32K max output tokens.

Panel of Figure 11 — Figure 11: Multi-stage SFT ablation on AIME 2024, AIME 2025, and GPQA-Diamond (Pass@1, $n{=}8$ , Qwen2.5-7B). OT = OpenThoughts only; OM/Full = OpenMementos Full Attention; OM/Mem = OpenMementos Memento Attention; 2-Stg = OT $\to$ OM/Full; 3-Stg = OT $\to$ OM/Full $\to$ OM/Mem (Ours). Training directly on OpenMementos from the base model (OM variants) substantially underperforms vanilla reasoning SFT (OT). Our three-stage pipeline enables block masking while retaining strong performance.

Figure 12: Toy transformer KV probing ( keep0 ). Left: Probe accuracy by layer depth—deeper layers carry more signal, especially for the masked condition. Center: Masked probe accuracy decays with distance from the target block but remains above chance even 7 hops away. Right: Leakage stays constant (17–19%) across training checkpoints even as task accuracy climbs from 77% to 95%, confirming the channel is architectural.

LLM Context Management

Large language model context management studies how to select, compress, retrieve, and organize information within bounded model inputs.

See all 244 papers

Structured Thoughts For Improved Reasoning And Context Pruning
Jul 11, 2026Zain Sarwar, Supriyo Chakraborty, Berkcan Kapusuzoglu +5Efficient Language Model Reasoning Structured CoT Reasoning
University of Chicago · Capital One
MemoSight: Unifying Context Compression and Multi Token Prediction for Reasoning Acceleration
Apr 16, 2026Xinyu Liu, Xin Liu, Bo Jin +8Multi-Token Prediction KV Caching
School of Computer Science and Engineering, Northeastern University, China · Meituan Inc. · NiuTrans Research, Shenyang, China
ZoomR: Memory Efficient Reasoning through Multi-Granularity Key Value Retrieval
Apr 13, 2026David H. Yang, Yuxuan Zhu, Mohammad Mohammadi Amiri +4Memory-Efficient Inference Efficient Language Model Reasoning
Rensselaer Polytechnic Institute · IBM Research

Jul 11, 2026J/K move · Enter open · S save

Structured Thoughts For Improved Reasoning And Context Pruning

Zain Sarwar, Supriyo Chakraborty, Berkcan Kapusuzoglu, Chia-Hsuan Lee+4

University of Chicago · Capital One

Efficient Language Model Reasoning Structured CoT Reasoning CoT Reasoning LLM Reasoning Structured Reasoning LLM Context Management+1

Open paper AskPDF arXiv

Large language models (LLMs) excel at generating long chains of thought, but long reasoning traces are often verbose and memory-inefficient. In this work, we introduce Structured Thoughts, a framework that organizes reasoning into alternating <try> and <outcome> blocks: <try> captures exploratory scratch work, while <outcome> contains the distilled conclusion of that step. We construct a dataset of structured thoughts by segmenting reasoning traces into <try> blocks and prompting an LLM to summarize each step into its corresponding <outcome>. Fine-tuning pretrained foundation models on this reformatted data produces models that adopt the structured reasoning style, leading to performance gains of up to 8.08% on reasoning benchmarks compared to standard SFT. The explicit structure also enables context pruning: after each <try>/<outcome> pair, the <try> can be pruned, allowing the model to retain conclusions without keeping the full scratch work in the context. A proof-of-concept pruning implementation achieves an average of 85% memory / context savings with an 8.67% performance drop across mathematical tasks.

Abstract

LLM Context Management

Explore similar work

Structured Thoughts For Improved Reasoning And Context Pruning

MemoSight: Unifying Context Compression and Multi Token Prediction for Reasoning Acceleration

ZoomR: Memory Efficient Reasoning through Multi-Granularity Key Value Retrieval