arXiv cs.AI精选76

并非所有维度都同等重要:基于分组二次核的异常值感知注意力线性化

Not All Dimensions Are Equal: Outlier-Aware Attention Linearization with Grouped Quadratic Kernels

· 4天前

精选理由

识别并单独处理少数关键维度,在保持线性推理效率的同时提升注意力近似精度与模型表现。

AI 摘要

将 softmax 注意力线性化为高效的 Transformer 推理提供了一条路径,但保留预训练注意力的行为仍然具有挑战性。我们发现了一个被标准线性化方法忽视的重要近似误差来源:少数异常维度对查询与键的点积贡献不成比例地大,并造成了很大一部分近似误差。受这种异质性的启发,我们提出 Outlier-Aware Attention Linearization(OAL),这是一种用于构建二次注意力近似的维度感知方法。从二阶 Taylor 近似出发,我们引入了一种 Hadamard 积形式,显式呈现查询与键各维度的贡献。该形式放宽了系数全局共享的限制,支持维度特定的线性项和跨维度的二次交互,同时保留可分离计算的性质。为了以更少的参数保留这种灵活性,OAL 将贡献尺度相近的维度分组以共享系数,同时允许对少数异常维度进行区别处理。所得正值二次核以预训练注意力输出为参照进行校准,并且在注意力头维度固定时,可按序列长度的线性时间复杂度计算。在 Qwen2.5-0.5B 上,与 DiJiang 相比,OAL 将冻结注意力的输出相对均方误差降低了约 8.9%。替换中间层注意力模块,并结合 LoRA 联合微调核系数后,与学习得到的全局二次核相比,WikiText-103 上的困惑度进一步从 31.71 降至 27.50,同时在 PIQA、HellaSwag 和 ARC-E 上的准确率也有所提升。这些结果表明,异常值感知的维度分组对于注意力近似和预训练模型转换均具有价值。代码已公开于 https://github.com/wuqianyue09/OAL

查看原文(arXiv cs.AI)