线性化 softmax 注意力是提升 Transformer 推理效率的关键路径,但在保持预训练注意力特性的同时仍面临显著误差。我们发现,标准线性化方法忽略了一类关键因素:少数异常维度在 query‑key 点积中贡献异常大,导致近似误差占比过高。基于此异质性,我们提出 异常感知注意力线性化 (OAL),一种面向维度的二次注意力近似构造方法。
从二阶泰勒展开出发,OAL 引入 Hadamard 积形式,使每个 query‑key 贡献显式化:
$$f(q,k) \approx f(0,0) + \nabla f|_{0}\cdot(q\odot k) + \frac12 (q\odot k)^{\top} \mathbf{W} (q\odot k)$$
该表述将全局共享系数放宽为维度特定的线性项和跨维度的二次交互项,同时保持可分离计算。为在参数预算内保留这种灵活性,OAL 将贡献幅度相近的维度分组共享系数,仅对少数异常维度使用独立系数。得到的正二次核在固定头维度下可实现相对于序列长度的线性时间评估。
在 Qwen2.5‑0.5B 上实验表明,OAL 将冻结注意力输出的相对 MSE 比 DiJiang 下降约 8.9%。进一步在中层注意力模块替换并结合 LoRA 联合微调,WikiText‑103 的困惑度从 31.71 降至 27.50,且在 PIQA、HellaSwag、ARC‑E 等基准上取得精度提升。这些结果验证了异常感知维度分组在注意力近似与预训练模型迁移中的价值。
点评:OAL 通过细粒度的维度建模和分组策略,有效抑制了异常维度带来的误差,兼顾了近似精度和计算效率,为大模型的高效推理提供了实用路径。