Learning What to Remember: Adaptive Probabilistic Memory Retention for Memory-Efficient Language Models (2025)

Abstract: Transformer attention scales quadratically with sequence length O(n^2) , limiting long-context use. We propose Adaptive Retention , a probabilistic, layer-wise token selection mechanism that learns which representations to keep under a strict global budget M . Retention is modeled with Bernoulli gates trained via a Hard-Concrete/variational relaxation and enforced with a simple top- M rule at inference, making the method differentiable and drop-in for standard encoders. Across classification, extractive QA, and long-document summarization, keeping only 30–50% of tokens preserves \geq 95\% of full-model performance while cutting peak memory by \sim 35 –45% and improving throughput by up to \sim 1.8\times . This architecture-agnostic approach delivers practical long-context efficiency without modifying base attention or task heads.

## introduction transformer-based language models have achieved remarkable success across a wide range of nlp tasks @xcite @xcite , but their memory requirements grow quadratically with sequence length, posing significant challenges for long-context processing @xcite . to address these limitations, various approaches have explored sparse attention patterns to reduce computational and memory overhead @xcite @xcite ) and compressed memory representations to extend effective context lengths @xcite @xcite . recent work has further proposed dynamic memory and computation strategies, such as internet-scale memory compression @xcite , token merging for efficient inference @xcite , and streaming memory management @xcite . however, existing methods often either rely on architectural modifications or apply fixed heuristics without explicitly modeling adaptive token retention under strict memory budgets. inspired by advances in adaptive computation time @xcite @xcite and improvements in selective attention @xcite , our approach formulates memory retention as a probabilistic learning problem, enabling language models to dynamically retain the most informative token representations while respecting a global memory constraint. through this formulation, our approach achieves significant memory savings with minimal performance degradation, offering a flexible and efficient solution applicable to standard transformer architectures. ## conclusion adaptive retention learns which tokens to keep via bernoulli gating with hard-concrete relaxation and a lagrangian budget, closely matching fullsequence accuracy while reducing memory and latency. ablations validate each component (