论文

核化线性注意力:用对称锥体打破容量墙

Kernelized Linear Attention: Breaking the Capacity Wall with Symmetric Cones

模型架构Transformer

摘要

线性注意力保证了恒定时间的循环推理,但在联想回忆方面却急剧退化。我们将注意力回忆表述为球形填充问题,并引入核化线性注意力激活(KATA),这是一个框架,其特征图是通过自对偶同质锥体证明非负注意力权重从第一原理导出的。基于这一观察,我们表明,一阶正半定(PSD)特征提供了有利的能力——干扰权衡。 KATA 恢复了无参数凸输出门,并通过 Welch 干涉层表征关联容量。对于高于此底线的公差,KATA 在不添加参数的情况下扩大了状态,并允许在投影维度中具有指数级数量的键的球形代码。我们在两个操作点将 KATA 实现为融合 Triton 内核:闪存注意力式转发高达 ${\sim}1.6\times$ FlashAttention-2 吞吐量,以及精确的 $O(T)$ 分块状态形式,在 $131$k 词元时达到 ${\sim}11\times$ FlashAttention-2 转发吞吐量。一阶特征的关联扫描将块大小 $C$ 的块间递归深度降低到 $O(\log(T/C))$,并将匹配的顺序线性注意基线的吞吐量平均为 ${\sim}2.4\times$。在远程 MQAR 和重复密钥覆盖方面,多个 KATA 变体的性能优于门控 DeltaNet,同时报告了参数计数和状态大小以及准确性。归纳法保留了近乎完美的召回率,而内核基准测试表明可以有效地实现映射。 KATA 在 $16\times$ 的分布外长度上保留了 $0.985$ MQAR,大约四分之一的 KV 缓存条目接近 softmax。 340M 参数 LLM 上的实验揭示了依赖于特征的流畅性权衡,并阐明了位置嵌入、增量规则和衰减门如何与特征几何体相互作用。