Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Chandra Vamsi Krishna Alla, Harish Naidu Gaddam, Manohar Kommi

TL;DR

BudgetMem tackles the challenge of cost-efficient long-context processing by learning selective memory policies that decide what to store under strict budgets. It combines a dual-tier memory (episodic for recent content and semantic for older content) with a hybrid retrieval module (dense + BM25) and a budget-aware compression pipeline to keep information within a small external memory while preserving answerability. In experiments on 700 QA cases across short and long documents with Llama-3.2-3B-Instruct, BudgetMem achieves only a ~1% F1 degradation on long documents while saving up to ~72% memory, with benefits increasing as document length grows. The approach is practical for deployment on modest hardware, provides citation-grounded generation, and offers a reproducible framework for developing memory-augmented LLMs aimed at real-world, resource-constrained applications.

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.

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

TL;DR

BudgetMem tackles the challenge of cost-efficient long-context processing by learning selective memory policies that decide what to store under strict budgets. It combines a dual-tier memory (episodic for recent content and semantic for older content) with a hybrid retrieval module (dense + BM25) and a budget-aware compression pipeline to keep information within a small external memory while preserving answerability. In experiments on 700 QA cases across short and long documents with Llama-3.2-3B-Instruct, BudgetMem achieves only a ~1% F1 degradation on long documents while saving up to ~72% memory, with benefits increasing as document length grows. The approach is practical for deployment on modest hardware, provides citation-grounded generation, and offers a reproducible framework for developing memory-augmented LLMs aimed at real-world, resource-constrained applications.

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.

Paper Structure

This paper contains 61 sections, 20 equations, 3 figures, 5 tables.

Table of Contents

  1. Introduction
  2. Motivation and Problem Statement
  3. Our Approach: BudgetMem
  4. Contributions
  5. Paper Organization
  6. Related Work
  7. Long-Context Transformers
  8. Retrieval-Augmented Generation
  9. Memory-Augmented Neural Networks
  10. Efficient LLM Inference
  11. Methodology
  12. Problem Formulation
  13. Architecture Overview
  14. Implementation Scope and Design Philosophy
  15. Chunking and Preprocessing
  16. ...and 46 more sections

Figures (3)

  • Figure 1: Document Length Scaling Effect. (a) F1 score degradation decreases from 9.7% on short documents to only 1.0% on long documents. (b) Memory savings increase from 15.5% to 72.4%. This demonstrates that BudgetMem's benefits scale dramatically with document length, making it ideal for research papers, legal documents, and other long-context applications.
  • Figure 2: Budget Sensitivity Analysis on long documents (5K-10K tokens). The plot shows F1 score (blue, left axis) and memory savings (purple, right axis) as functions of budget ratio. The sweet spot at 30% budget achieves 72.4% memory savings with 80.4% F1 score, providing an optimal quality-efficiency tradeoff.
  • Figure 3: Comparison of BudgetMem against naive selection strategies. All methods use 30% budget on long documents. BudgetMem's feature-based salience scoring (F1: 0.8042) outperforms Random (0.6892), First-N (0.7254), Last-N (0.6734), and TF-IDF-only (0.7689) approaches. The multi-feature approach provides a 4.6% improvement over the best naive baseline, validating the importance of combining entity density, position bias, and discourse markers.