# Learning What to Remember: Adaptive Probabilistic Memory Retention for Memory-Efficient Language Models

arXiv: https://arxiv.org/abs/2510.08798 · HTML: https://mlsanity.com/paper/2510.08798/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 37 of its 46 references that are in the map.
- First posted on arXiv 9 Oct 2025.

## Builds on (most cited first, 30 of 37)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Llama 2: Open Foundation and Fine-Tuned Chat Models](https://mlsanity.com/paper/2307.09288/) (2023)
- [GLUE: A Multi-Task Benchmark and Analysis Platform for Natural Language Understanding](https://mlsanity.com/paper/1804.07461/) (2018)
- [Longformer: The Long-Document Transformer](https://mlsanity.com/paper/2004.05150/) (2020)
- [Mistral 7B](https://mlsanity.com/paper/2310.06825/) (2023)
- [FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning](https://mlsanity.com/paper/2307.08691/) (2023)
- [Convex optimization](https://mlsanity.com/paper/2106.01946/) (2021)
- [Efficient Streaming Language Models with Attention Sinks](https://mlsanity.com/paper/2309.17453/) (2023)
- [Big Bird: Transformers for Longer Sequences](https://mlsanity.com/paper/2007.14062/) (2020)
- [Reformer: The Efficient Transformer](https://mlsanity.com/paper/2001.04451/) (2020)
- [Rethinking Attention with Performers](https://mlsanity.com/paper/2009.14794/) (2020)
- [Generating Long Sequences with Sparse Transformers](https://mlsanity.com/paper/1904.10509/) (2019)
- [The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables](https://mlsanity.com/paper/1611.00712/) (2016)
- [Token Merging: Your ViT But Faster](https://mlsanity.com/paper/2210.09461/) (2022)
- [SnapKV: LLM Knows What You are Looking for Before Generation](https://mlsanity.com/paper/2404.14469/) (2024)
- [Efficient Transformers: A Survey](https://mlsanity.com/paper/2009.06732/) (2020)
- [Compressive Transformers for Long-Range Sequence Modelling](https://mlsanity.com/paper/1911.05507/) (2019)
- [Learning Sparse Neural Networks through L0 Regularization](https://mlsanity.com/paper/1712.01312/) (2017)
- [The Falcon Series of Open Language Models](https://mlsanity.com/paper/2311.16867/) (2023)
- [Adaptive Computation Time for Recurrent Neural Networks](https://mlsanity.com/paper/1603.08983/) (2016)
- [A Dataset of Information-Seeking Questions and Answers Anchored in Research Papers](https://mlsanity.com/paper/2105.03011/) (2021)
- [A Discourse-Aware Attention Model for Abstractive Summarization of Long Documents](https://mlsanity.com/paper/1804.05685/) (2018)
- [DeeBERT: Dynamic Early Exiting for Accelerating BERT Inference](https://mlsanity.com/paper/2004.12993/) (2020)
- [Benchmarking Retrieval-Augmented Generation for Medicine](https://mlsanity.com/paper/2402.13178/) (2024)
- [CUAD: An Expert-Annotated NLP Dataset for Legal Contract Review](https://mlsanity.com/paper/2103.06268/) (2021)
- [FastBERT: a Self-distilling BERT with Adaptive Inference Time](https://mlsanity.com/paper/2004.02178/) (2020)
- [Memorizing Transformers](https://mlsanity.com/paper/2203.08913/) (2022)
- [On the Use of ArXiv as a Dataset](https://mlsanity.com/paper/1905.00075/) (2019)

## Built on by (most cited first, 1 of 1)

- [Linguistics and Human Brain: A Perspective of Computational Neuroscience](https://mlsanity.com/paper/2602.08275/) (2026)

## Abstract (arXiv)

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.