Linearizing softmax attention offers a promising route to efficient Transformer inference, yet preserving the fidelity of pretrained attention remains difficult. We identify a major source of error that standard linearization overlooks: a handful of outlier dimensions dominate the query‑key dot products and account for a large share of the approximation error. Motivated by this heterogeneity, we introduce Outlier‑Aware Attention Linearization (OAL), a dimension‑aware method for constructing quadratic attention approximations.
Starting from a second‑order Taylor expansion, OAL adopts a Hadamard‑product formulation that makes each query‑key contribution explicit:
$$f(q,k) \approx f(0,0) + \nabla f|_{0}\cdot(q\odot k) + \frac12 (q\odot k)^{\top} \mathbf{W} (q\odot k)$$
This relaxes globally shared coefficients to allow dimension‑specific linear terms and cross‑dimensional quadratic interactions while keeping the evaluation separable. To retain this flexibility with limited parameters, OAL groups dimensions with similar contribution scales to share coefficients, reserving distinct treatment for the few outliers. The resulting positive quadratic kernel can be evaluated in linear time with respect to sequence length for a fixed head dimension.
Experiments on Qwen2.5‑0.5B show that OAL reduces the relative MSE of frozen‑attention outputs by about 8.9% compared with DiJiang. Replacing middle‑layer attention modules and jointly fine‑tuning the kernel coefficients with LoRA further lowers WikiText‑103 perplexity from 31.71 to 27.50 and improves accuracy on PIQA, HellaSwag, and ARC‑E. These results demonstrate the value of outlier‑aware dimension grouping for both attention approximation and pretrained‑model conversion.
Review: OAL’s fine‑grained dimension modeling and grouping effectively mitigate outlier‑driven errors while preserving computational efficiency, offering a practical solution for high‑performance inference in large language models.