MEMENTO: Teaching LLMs to Manage Their Own Context

• URL: http://arxiv.org/abs/2604.09852v1
• Date: Fri, 10 Apr 2026 19:30:29 GMT
• Title: MEMENTO: Teaching LLMs to Manage Their Own Context
• Abstract summary: Reasoning models think in long, unstructured streams with no mechanism for compressing or organizing their own intermediate state.<n>We introduce MEMENTO: a method that teaches models to segment reasoning into blocks, compress each block into a memento, and reason forward by attending only to mementos.<n>We show that a two-stage SFT recipe on OpenMementos is effective across different model families and scales.
• Score: 50.3558738319336
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• 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.

Related papers

• Language Models Can Control Their Own Attention [42.07765767561011]
  We introduce Declarative Attention (DA), a protocol that elicits the model to declare where it needs to attend within its chain-of-thought.<n>DA significantly reduces total attended tokens during decoding (52.0%, 31.1%) with modest accuracy drops (1.27pp, 2.75pp) that shrink with model scale.
  arXiv  Detail & Related papers  (2026-09-02T15:43:38Z)
• LimiX-2M: Mitigating Low-Rank Collapse and Attention Bottlenecks in Tabular Foundation Models [56.999481798138625]
  LimiX-2M is a 2M- parameter model that outperforms larger TabPFN-v2 and TabICL baselines on widely used benchmarks.<n>We present a unified tokenize-and-route framework for strong Tabular foundation models (TFMs)<n>Results highlight value-aware tokenization and readout-aligned routing as key levers for improving the accuracy--efficiency trade-off in TFMs.
  arXiv  Detail & Related papers  (2026-06-03T06:07:33Z)
• Enabling KV Caching of Shared Prefix for Diffusion Language Models [0.2740273306918099]
  Key-value (KV) caching for shared prefixes is essential for high- throughput large language models (LLM) serving.<n>We propose bicache, the first KV caching technique for shared prefixes in diffusion language models (DLMs)<n>We show that bicache significantly improves serving throughput by 36.3%-98.3% compared to existing techniques without accuracy collapse.
  arXiv  Detail & Related papers  (2026-05-26T05:27:57Z)
• Learning is Forgetting: LLM Training As Lossy Compression [19.24753945415483]
  We argue large language models (LLMs) are best seen as an instance of lossy compression, where over training they learn by retaining only information relevant to their objective(s)<n>We show pre-training results in models that are optimally compressed for next-sequence prediction, approaching the Information Bottleneck bound on compression.
  arXiv  Detail & Related papers  (2026-04-08T20:12:07Z)
• Probing to Refine: Reinforcement Distillation of LLMs via Explanatory Inversion [84.20493238687187]
  We introduce a novel framework that moves beyond simple mimicry to instill a deeper conceptual understanding.<n>underlinetextitFirst, to address pattern memorization, Explanatory Inversion (EI) generates targeted explanatory probes''<n>underlinetextitSecond, to improve generalization, Explanatory GRPO (texttEXGRPO) uses a reinforcement learning algorithm with a novel Dialogue Structure Utility Bonus.
  arXiv  Detail & Related papers  (2026-02-26T23:01:46Z)
• Going Down Memory Lane: Scaling Tokens for Video Stream Understanding with Dynamic KV-Cache Memory [50.30283773196725]
  Existing approaches rely on key-value caching to accumulate frame-level details over time, but use a limited number of tokens per frame.<n>We propose scaling the token budget to enable more granular-temporal understanding and reasoning.
  arXiv  Detail & Related papers  (2026-02-20T18:59:50Z)
• d3LLM: Ultra-Fast Diffusion LLM using Pseudo-Trajectory Distillation [31.922313594074925]
  Diffusion large language models (dLLMs) offer capabilities beyond those of autoregressive (AR) LLMs.<n>Current methods typically focus on only one-side of the coin, targeting either efficiency or performance.<n>We propose d3LLM (Pseudo-Distilled Diffusion Large Language Model), striking a balance between accuracy and parallelism.
  arXiv  Detail & Related papers  (2026-01-12T14:25:36Z)
• Striking the Right Balance between Compute and Copy: Improving LLM Inferencing Under Speculative Decoding [12.302511322703852]
  We propose a new KV cache allocation mechanism called Balancing Memory and Compute (BMC)<n>BMC allocates, once every r iterations, KV tensors with r redundant rows, allowing in-place update without copy overhead for those iterations.<n> BMC achieves a throughput acceleration of up to 1.36x and 2.29x over state-of-the-art inference servers vLLM and DeepSpeed.
  arXiv  Detail & Related papers  (2025-11-15T04:49:23Z)
• Judge Q: Trainable Queries for Optimized Information Retention in KV Cache Eviction [53.83828564664595]
  Large language models (LLMs) utilize key-value ( KV) cache to store historical information during sequence processing.<n>Current methods for KV cache eviction typically utilize the last window from the pre-filling phase as queries to compute the KV importance scores for eviction.<n>We propose Judge Q, a novel training method which incorporates a soft token list.
  arXiv  Detail & Related papers  (2025-09-13T03:34:12Z)
• S2-Attention: Hardware-Aware Context Sharding Among Attention Heads [49.1454481007861]
  Sparse attention selectively attends to a subset of tokens in the context.<n>It remains unclear whether sparse attention can maintain the model's quality at a scale of today's large language models.<n>This paper presents Sparsely-Sharded(S2) Attention, a Triton library that provides kernel optimization for sparse attention customizable at both per-head and per-context-range levels.
  arXiv  Detail & Related papers  (2024-07-25T00:27:07Z)
• EVEREST: Efficient Masked Video Autoencoder by Removing Redundant Spatiotemporal Tokens [57.354304637367555]
  We present EVEREST, a surprisingly efficient MVA approach for video representation learning. It finds tokens containing rich motion features and discards uninformative ones during both pre-training and fine-tuning. Our method significantly reduces the computation and memory requirements of MVA.
  arXiv  Detail & Related papers  (2022-11-19T09:57:01Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.