论文

遗忘注意力:重要性感知注意力就是全部所需

Forget Attention: Importance-Aware Attention Is All You Need

模型架构混合架构

摘要

将注意力的全局检索与状态空间模型(SSM)的顺序重要性信号相结合是混合语言建模的公开挑战。变形金刚无所不在,但无法确定优先顺序; SSM 知道什么是重要的,但无法重新审视。现有的混合体——Jamba(块级)和 Hymba(头级)——将两者放在不同的隔间中,因此在注意力计算本身期间两者都不会通知对方。我们提出了 SISA(SSM-Informed Softmax Attention),它将 SSM 派生的重要性项直接添加到注意力分数中,并将完整操作实现为对增强查询/键向量的单个 SDPA 调用——没有循环状态,没有自定义内核。在 152M / 5B 词元下,SISA 达到 LAMBADA 贪婪 17.3%(相对于 Transformer 13.9 和 Mamba-3 15.5),并从步骤 1K 达到 NIAH 100%,比 Transformer 的检索收敛速度快 7 倍;在 369M 上,Mamba-3 领先 LAMBADA,而 SISA 保留完美的 NIAH 和股票 SDPA 执行。因此,SISA 定义了 SSM-注意力混合体的第三个设计轴——分数级融合——超越了主导该领域的块级和头级范例。

遗忘注意力:重要性感知注意力就是全部所需:论文配图
图 1:SSM-注意力融合的三个级别。块级和头级混合结合了独立计算的输出。 SISA 将 SSM 信号注入注意力分数本身,从而减少了增强 Q/K 和一次 SDPA 调用的操作。