KATA突破线性注意力容量瓶颈,实现高效长序列关联记忆。
Kernelized Linear Attention: Breaking the Capacity Wall with Symmetric Cones

- 基于对偶锥构造非负注意力特征,从原理上保证权重合理性。
- 在13.1万词长序列上达到FlashAttention-2 11倍吞吐,且无参数扩展状态。
- 适合需要高容量长程记忆的模型,如超长文本生成与推理任务。
线性注意力虽可实现常数时间递归推理,但在关联记忆任务中性能急剧下降。本文将注意力回忆建模为球面打包问题,提出核化线性注意力激活(KATA)框架,其特征映射源自自对偶齐次锥的严格推导,确保非负注意力权重。研究发现,秩一半正定特征在容量与干扰间取得最优平衡。KATA无需参数即可实现凸输出门,并以Welch干扰下界刻画关联容量;当容忍度高于该下界时,可在不增加参数情况下扩大状态空间,支持投影维度呈指数级增长的球面码。通过融合Triton内核实现两种运行模式:一种类FlashAttention前向达约1.6倍速度提升,另一种精确的O(T)分块状态形式在13.1万词序列上达约11倍速度提升。一阶特征的关联扫描将跨块递归深度降至O(log(T/C)),平均提速2.4倍于匹配的顺序线性注意力基线。在长程MQAR与重复键覆盖任务中,多个KATA变体超越Gated DeltaNet,同时报告参数量与状态大小。归纳推理保持近乎完美召回,核基准测试表明特征可高效实现。在16倍分布外长度下仍保持0.985的MQAR,仅需软注意力约四分之一的键值缓存条目。3.4亿参数大模型实验揭示特征依赖的流畅性权衡,澄清位置嵌入、增量规则与衰减门如何与特征几何交互。
原文摘要 · Abstract (English)
Linear attention promises constant-time recurrent inference but degrades sharply on associative recall. We formulate attention recall as a spherical-packing problem and introduce Kernelized Linear Attention Activations (KATA), a framework whose feature maps are derived from first principles by certifying nonnegative attention weights through a self-dual homogeneous cone. Building on this observation, we show that rank-one positive semi-definite (PSD) features offer a favorable capacity--interference tradeoff. KATA recovers a parameter-free convex output gate and characterizes associative capacity through the Welch interference floor. For tolerances above this floor, KATA enlarges the state without adding parameters and admits spherical codes with exponentially many keys in the projection dimension. We implement KATA as fused Triton kernels at two operating points: a flash-attention-style forward up to ${\sim}1.6\times$ FlashAttention-2 throughput, and an exact $O(T)$ chunked-state form that reaches ${\sim}11\times$ FlashAttention-2 forward throughput at $131$k tokens. An associative scan of the first-order feature lowers the inter-chunk recurrence depth to $O(\log(T/C))$ for chunk size $C$ and averages ${\sim}2.4\times$ the throughput of a matched sequential linear-attention baseline. On long-range MQAR and repeated-key overwrite, several KATA variants outperform Gated DeltaNet, with parameter counts and state sizes reported alongside accuracy. Induction preserves near-perfect recall, while kernel benchmarks show that the maps can be implemented efficiently. KATA retains $0.985$ MQAR at a $16\times$ out-of-distribution length, approaching the softmax with roughly one quarter of the KV-cache entries. Experiments on 340M-parameter LLMs reveal a feature-dependent fluency trade-off and clarify how positional embeddings, delta rules, and decay gates interact with feature geometry.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。