Beyond Overwriting: Selective Memory Integration for Long-Context Reasoning

Efficiently processing long contexts remains a key challenge for large language models (LLMs). The memorize-while-reading paradigm addresses this challenge by processing documents chunk by chunk and dynamically updating a textual memory, with subsequent work introducing retrieval to recover lost evidence and gating mechanisms for controlling when to update memory and stop reading. Despite these advances, two challenges remain: (i) repeated overwrite-based updates lack explicit protection for previously stored information, potentially causing information loss that is difficult to recover through retrieval; and (ii) optimization based on rule-based rewards provides only indirect guidance for learning when to update memory and stop reading, which may lead to unstable optimization. To address these challenges, we propose GraftMem, which separates evidence acquisition from incremental memory integration through acquisition–integration (AcIn) updates. Combined with selective memory updating, this design effectively mitigates information loss from repeated overwriting. GraftMem also incorporates an early-exit mechanism to improve efficiency. To strengthen training, we combine multi-level reward shaping with decision supervision: alongside trajectory-level and state-level rewards, we directly supervise memory-update and early-exit decisions. Together, these signals alleviate supervision sparsity and support more stable policy optimization. Extensive experiments demonstrate that GraftMem achieves competitive accuracy on long-context reasoning tasks while substantially accelerating inference, achieving a favorable accuracy–efficiency trade-off.