通过混合稀疏注意力和上下文化可学习 Token 驱逐缓解线性注意力模型的遗忘问题

将整个输入序列压缩到固定大小循环状态中的线性注意力模型为 Transformer 提供了高效替代方案，但其有限的内存会导致遗忘现象，损害检索密集型任务。为缓解此问题，我们探索了一系列混合模型，以恢复对过去 token 的直接访问。我们将 token 混合器与介于线性注意力和完全注意力之间的中间时间和空间复杂度交错，包括带有 token 驱逐的稀疏注意力，以及查询感知的原生稀疏注意力。特别地，我们提出了一种新颖的可学习 token 驱逐方法。结合滑动窗口注意力，一个端到端可训练的轻量级 CNN 聚合来自过去和未来相邻 token 的信息，以适应地保留每个头中有限数量的关键 KV 配对，保持线性注意力的恒定时间和空间复杂度。提供高效的 Triton kernel 用于稀疏注意力机制。在检索密集型基准测试上的经验评估支持我们方法的有效性。