核化线性注意力:用对称锥突破容量限制
Kernelized Linear Attention: Breaking the Capacity Wall with Symmetric Cones
AI总结:
研究线性注意力在关联记忆上的退化问题,提出核化线性注意力激活框架KATA,通过自对偶齐次锥认证非负权重导出特征图,实现高效计算,在长距离记忆等任务中表现出色,揭示了特征相关权衡及相互作用。
AI中文摘要:
线性注意力保证了恒定时间的循环推理,但在关联记忆方面急剧退化。我们将注意力记忆表述为一个球形填充问题,并引入核化线性注意力激活(KATA)框架,其特征图通过自对偶齐次锥认证非负注意力权重从第一原理导出。基于此观察,我们表明秩一正半定(PSD)特征提供了有利的容量-干扰权衡。KATA恢复了无参数凸输出门,并通过韦尔奇干扰下限表征关联容量。对于高于此下限的容差,KATA在不增加参数的情况下扩大状态,并允许在投影维度中具有指数多个键的球形码。我们在两个操作点将KATA实现为融合的Triton内核:一个闪存注意力风格的前向操作,吞吐量高达约1.6倍的FlashAttention-2;以及一个精确的O(T)分块状态形式,在131k令牌时达到约11倍的FlashAttention-2前向吞吐量。一阶特征的关联扫描将分块间循环深度降低到O(log(T/C)),对于分块大小C,平均吞吐量约为匹配顺序线性注意力基线的2.4倍。在长距离MQAR和重复键覆盖方面,几个KATA变体优于门控DeltaNet,并报告了参数数量、状态大小和准确率。归纳保留了近乎完美的记忆,而内核基准测试表明这些映射可以有效实现。KATA在16倍的分布外长度下保留0.985的MQAR,使用大约四分之一的KV缓存条目接近softmax。在340M参数的语言模型上的实验揭示了特征相关的流畅性权衡,并阐明了位置嵌入、增量规则和衰减门如何与特征几何相互作用。
英文摘要:
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.