通过自适应选择保留部分标记，显著降低长文本模型内存占用。
Learning What to Remember: Adaptive Probabilistic Memory Retention for Memory-Efficient Language Models
- 用可学习的伯努利门控机制按层筛选关键标记，保持全局预算
- 仅保留30%-50%标记仍保持95%以上性能，内存减少35%-45%
- 无需修改原模型结构，适合各类需长文本处理的任务
Transformer注意力随序列长度呈平方级增长（O(n²)），限制了长文本应用。本文提出自适应保留机制（Adaptive Retention），一种基于概率的、逐层的标记选择方法，在严格全局预算M下学习保留哪些表示。通过硬-混凝土/变分松弛训练伯努利门控，并在推理时采用简单top-M规则，使方法可微且可直接嵌入标准编码器。在分类、抽取式问答和长文档摘要任务中，仅保留30%-50%的标记即可保持≥95%的完整模型性能，同时将峰值内存降低约35%-45%，吞吐量提升最高达1.8倍。该方法与架构无关，无需修改基础注意力或任务头，实现实用的长上下文效率。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。