BudgetMem: Learning Selective Memory Policies for Cost-Efficient Long-Context LLM Processing

Forum topic · QianXun · 2025-11-10

Summary

BudgetMem is a memory-efficient architecture for long-context language model processing, proposed by engineers from AT&T, Bank of America, and Ford. Instead of storing everything like conventional RAG systems, BudgetMem uses a trainable salience scorer (entity density, TF-IDF, positional bias, discourse markers, numeric content) to selectively retain only the most important text chunks under a memory budget. It combines a dual-tier memory design—episodic memory for recent chunks and semantic memory storing 100-token distilled summaries with 768-dim embeddings—with a three-stage retrieval pipeline (hybrid dense+BM25 search, cross-encoder reranking, episodic integration). On long documents (5K-10K tokens), BudgetMem cuts memory usage by 72.4% (from 100% to 27.6%) with only a 1.0% F1 drop (0.8123 to 0.8042), outperforming naive selection baselines by 3.5 F1 points; a 30% budget is the identified sweet spot. Trade-offs include a ~20% latency increase and poor cost-benefit on short documents. The approach democratizes long-context AI, running on consumer hardware such as Google Colab Pro.

BudgetMem (arXiv:2511.04919) is an architecture that lets large language models handle long documents by selectively remembering only salient information, rather than storing everything as traditional RAG systems do.

Key points

Motivation

Processing 100K+ token contexts can require ~40GB of GPU memory, making long-context processing expensive.
Conventional RAG stores all chunks indiscriminately, drowning in low-value content. BudgetMem instead asks: *what should be remembered?*

Architecture

What to write

s_i = \sigma(w^T f_i + b)

Chunks are selected via TopK under a budget $$B$$ : $S = \text{TopK}(\{s_i\}, K = B)$ . 2. How to store: A dual-tier memory — episodic memory holds the 10–20 most recent chunks; semantic memory stores 80–120-token distilled summaries (Llama-3.2-3B LoRA summarizer) each tagged with a 768-dim embedding. 3. What to retrieve: A three-stage pipeline — hybrid dense + BM25 search (7:3 fusion) yields 40 candidates; a 300M-parameter cross-encoder reranks to 5–8 chunks; recent episodic memory is always appended. A ranking loss

\mathcal{L}_{\text{rank}} = \sum_{i \in P, j \notin P} \max(0, \gamma + s_j - s_i)

trains the scorer to rank answer-bearing chunks above distractors. Summaries are trained with a multi-task loss (ROUGE-L content coverage + answerability).

Experimental results (700 QA cases)

Short docs (SQuAD v2.0, 237 tokens avg): F1 dropped 0.8011 → 0.7232 (−9.7%), only 15.5% memory savings — little benefit on short text.
Long docs (200 synthetic academic papers, ~7200 tokens): F1 0.8123 → 0.8042 (−1.0%) while memory fell from 100% to 27.6% (72.4% savings); latency rose 20.8% (2.45s → 2.96s).
Budget sweep: 10% budget → F1 0.6245; 30% budget → F1 0.8042 (sweet spot); 40% → 0.8156 with 60.1% savings; diminishing returns beyond 50%.
Baselines (30% budget, long docs): random 0.6892, First-N 0.7254, Last-N 0.6734, TF-IDF-only 0.7689 vs. BudgetMem's full features at 0.8042 (+3.5 points over the best naive baseline).

When to use / limitations

Strong for: long documents (5K+ tokens), structured content (papers, contracts), localized Q&A, hardware-constrained deployments (mobile, IoT, cost-sensitive cloud).
Weak for: short documents, answers spanning multiple low-salience chunks, queries about appendices/footnotes, latency-critical real-time systems.

Future work

Validation on real benchmarks (Qasper, GovReport, LongBench); learned write policies instead of hand-tuned weights; adaptive per-document budgets; multimodal salience scoring (tables, charts, code); human-in-the-loop evaluation.

Practical takeaways

Use 30–40% budget on resource-constrained hardware for long texts; use full-context models below ~1K tokens.
Zero-shot deployment works with hand-tuned feature weights; the whole pipeline runs on Google Colab Pro (~$10/month).

References

[1] Alla, C. V. K., Gaddam, H. N., & Kommi, M. (2025). *BudgetMem: Learning Selective Memory Policies for Cost-Efficient Long-Context Processing in Language Models*. arXiv:2511.04919v1.
[2] Borgeaud, S., et al. (2022). *Improving language models by retrieving from trillions of tokens*. ICML.
[3] Lewis, P., et al. (2020). *Retrieval-augmented generation for knowledge-intensive NLP tasks*. NeurIPS.
[4] Karpukhin, V., et al. (2020). *Dense passage retrieval for open-domain question answering*. EMNLP.
[5] Raffel, C., et al. (2020). *Exploring the limits of transfer learning with a unified text-to-text transformer*. JMLR.