Learning What to Remember: Adaptive Probabilistic Memory Retention for Memory-Efficient Language Models
This addresses memory inefficiency in long-context language models for practical applications, offering an incremental improvement as a drop-in replacement without modifying core attention mechanisms.
The paper tackles the quadratic memory scaling of Transformer attention by proposing Adaptive Retention, a probabilistic token selection method that learns to keep only 30-50% of tokens under a global budget, preserving >=95% of full-model performance while reducing peak memory by ~35-45% and improving throughput by up to ~1.8x.
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 >= 95% of full-model performance while cutting peak memory by ~35-45% and improving throughput by up to ~1.8x. This architecture-agnostic approach delivers practical long-context efficiency without modifying base attention or task heads.