Dynamic Long Context Reasoning over Compressed Memory via End-to-End Reinforcement Learning
Zhuoen Chen, Dongfang Li, Meishan Zhang, Baotian Hu, Min Zhang
Abstract
Large Language Models (LLMs) face severe challenges in long-context processing, including quadratic computational costs, information forgetting, and the context fragmentation inherent in Retrieval-Augmented Generation (RAG). We introduce LycheeMemory, a cognitively inspired framework that enables efficient long-context inference via chunk-wise compression and selective memory recall, rather than processing all raw tokens. LycheeMemory segments the input into chunks and encodes each into compressed KV-cache-style representations using a Compressor. A Gate then dynamically selects relevant memory blocks, which a Reasoner iteratively processes with an evolving working memory to solve downstream tasks. The Compressor and Reasoner are jointly optimized via end-to-end reinforcement learning, while the Gate is trained separately as a classifier. Experimental results demonstrate that LycheeMemory achieves competitive accuracy (up to 82% in ablation variants) on multi-hop reasoning benchmarks (e.g., RULER-HQA), successfully extrapolates context length from 7K to 1.75M, and provides a favorable accuracy-efficiency trade-off against strong long-context baselines. Notably, compared to MemAgent, LycheeMemory achieves an average 2× reduction in peak GPU memory usage and a 6× speedup during inference.
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Your agent calls
Luneget_paper_fulltext
Free to start. No credit card required.
Terminal
Install the CLIlune papers fulltext f2064fc2-b922-40c3-a510-ebce3f2dd2a2Builds on18
- Retrieval-Augmented Generation for Knowledge-Intensive NLP TasksPatrick Lewis, Ethan Perez, Aleksandra Piktus, Fabio Petroni et al.NeurIPS 2020 · 19,162 citations
- LoRA: Low-Rank Adaptation of Large Language ModelsEdward J. Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu et al.ICLR 2022 · 18,833 citations
- Transformers are RNNs: Fast Autoregressive Transformers with Linear AttentionAngelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, François FleuretICML 2020 · 2,665 citations
- Efficient Streaming Language Models with Attention SinksGuangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han et al.ICLR 2024 · 1,714 citations
- SnapKV: LLM Knows What You are Looking for Before GenerationYuhong Li, Yingbing Huang, Bowen Yang, Bharat Venkitesh et al.NeurIPS 2024 · 1,019 citations
Related papers
- MemoRAG: Boosting Long Context Processing with Global Memory-Enhanced Retrieval AugmentationHongjin Qian, Zheng Liu, Peitian Zhang, Kelong Mao et al.WWW 2025 · 92 citations
- RetroLM: Retrieval-Augmented KVs for Long-Context ProcessingKun Luo, Zheng Liu, Shitao Xiao, Jiabei Chen et al.AAAI 2026
- Compress, Gather, and Recompute: REFORMing Long-Context Processing in TransformersWoomin Song, Sai Muralidhar Jayanthi, Srikanth Ronanki, Kanthashree Mysore Sathyendra et al.NeurIPS 2025 · 1 citation
- LycheeDecode: Accelerating Long-Context LLM Inference via Hybrid-Head Sparse DecodingGang Lin, Dongfang Li, Zhuoen Chen, Yukun Shi et al.ICLR 2026 · 5 citations
- Which Heads Matter for Reasoning? RL-Guided KV Cache CompressionWenjie Du, Li Jiang, Keda TAO, Xue Liu et al.ICML 2026 · 11 citations
