用可学习的令牌淘汰机制，解决线性注意力遗忘问题。
Alleviating Forgetfulness of Linear Attention by Hybrid Sparse Attention and Contextualized Learnable Token Eviction
- 引入可学习的令牌淘汰策略，恢复对历史信息的直接访问。
- 在保持线性复杂度前提下，显著提升长序列检索任务性能。
- 适合需要高效处理长序列且关注历史记忆的场景。
线性注意力模型通过将整个输入序列压缩为固定大小的循环状态，提供了比Transformer更高效的替代方案，但其有限的内存会导致遗忘问题，影响检索密集型任务的表现。为缓解这一问题，我们探索了一系列混合模型，以恢复对过去令牌的直接访问。这些模型在线性注意力与全注意力之间引入了中间时间与空间复杂度的稀疏注意力，包括基于令牌淘汰的稀疏注意力和查询感知的原生稀疏注意力。特别地，我们提出一种新颖的可学习令牌淘汰方法。结合滑动窗口注意力，一个端到端可训练的轻量级CNN能从前后相邻的令牌中聚合信息，自适应地每头保留一组关键的键值对，同时保持线性注意力的恒定时间和空间复杂度。还提供了针对稀疏注意力机制的高效Triton内核。在多个检索密集型基准上的实证评估验证了我们方法的有效性。
原文摘要 · Abstract (English)
Linear-attention models that compress the entire input sequence into a fixed-size recurrent state offer an efficient alternative to Transformers, but their finite memory induces forgetfulness that harms retrieval-intensive tasks. To mitigate the issue, we explore a series of hybrid models that restore direct access to past tokens. We interleave token mixers with intermediate time and space complexity between linear and full attention, including sparse attention with token eviction, and the query-aware native sparse attention. Particularly, we propose a novel learnable token eviction approach. Combined with sliding-window attention, an end-to-end trainable lightweight CNN aggregates information from both past and future adjacent tokens to adaptively retain a limited set of critical KV-pairs per head, maintaining linear attention's constant time and space complexity. Efficient Triton kernels for the sparse attention mechanisms are provided. Empirical evaluations on retrieval-intensive benchmarks support the effectiveness of our approaches.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。