BudgetMem: Learning Selective Memory Policies for Cost-Efficient Long-Context Processing in Language Models

2025-11-07

AI Summary

To address the prohibitively high computational and memory overhead of large language models (LLMs) in long-context processing, this paper proposes a budget-constrained Selective Memory architecture. The method innovatively integrates a learnable memory gating mechanism with a multidimensional importance scoring function—incorporating BM25, entity density, TF-IDF, discourse markers, and positional bias—to dynamically select and retain salient information under strict memory constraints, thereby departing from the conventional full-document retrieval-augmented generation (RAG) paradigm. Evaluated on an enhanced Llama-3.2-3B-Instruct model, the approach achieves only a 1.0% F1 drop on long-document tasks while reducing memory consumption by 72.4%. Crucially, its performance advantage over baseline RAG methods grows substantially with increasing document length, demonstrating superior scalability and efficiency in resource-constrained long-context settings.

Technology Category

• Machine Learning: Large Multimodal Models (LMMs)
• Natural Language Processing: (Large) Language Models
• Data Mining & Knowledge Management: Conversational Systems for Recommendation & Retrieval

Application Category

• Search and Retrieval-Augmented AI: Retrieval-Augmented Generation (RAG) and multi-modal RAG
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation

Abstract

Large Language Models (LLMs) face significant computational and memory constraints when processing long contexts, despite growing demand for applications requiring reasoning over extensive documents, multi-session dialogues, and book length texts. While recent advances have extended context windows to 100K-1M tokens, such approaches incur prohibitive costs for resource constrained deployments. We propose BudgetMem, a novel memory augmented architecture that learns what to remember rather than remembering everything. Our system combines selective memory policies with feature based salience scoring (entity density, TF-IDF, discourse markers, position bias) to decide which information merits storage under strict budget constraints. Unlike existing retrieval augmented generation (RAG) systems that store all chunks, BudgetMem employs learned gating mechanisms coupled with BM25 sparse retrieval for efficient information access. Through comprehensive experiments on 700 question answer pairs across short (237 tokens) and long (5K-10K tokens) documents with Llama-3.2-3B-Instruct, we demonstrate that BudgetMem achieves remarkable results on long documents: only 1.0% F1 score degradation while saving 72.4% memory compared to baseline RAG. We validate our approach through budget sensitivity analysis (testing 7 budget ratios), naive baseline comparisons, and document length analysis, showing that BudgetMem's benefits increase with document length. Our work provides a practical pathway for deploying capable long context systems on modest hardware, democratizing access to advanced language understanding capabilities.

Problem

• Reducing computational and memory costs in long-context language model processing
• Selectively storing information under strict memory budget constraints
• Improving efficiency over traditional retrieval-augmented generation systems

Innovation

• Selective memory policies with feature-based salience scoring
• Learned gating mechanisms with BM25 sparse retrieval
• Memory-augmented architecture operating under strict budget constraints