Not All Dimensions Are Equal: Outlier-Aware Attention Linearization with Grouped Quadratic Kernels
Abstract
Linearizing softmax attention offers a route to efficient Transformer inference, but preserving pretrained attention remains challenging. We identify an important source of approximation error overlooked by standard linearization methods: a small number of outlier dimensions contribute disproportionately to query-key dot products and account for a large fraction of the approximation error. Motivated by this heterogeneity, we propose Outlier-Aware Attention Linearization (OAL), a dimension-aware method for constructing quadratic attention approximations. Starting from a second-order Taylor approximation, we introduce a Hadamard-product formulation that makes individual query-key contributions explicit. This formulation relaxes globally shared coefficients to support dimension-specific linear terms and cross-dimensional quadratic interactions while preserving separable evaluation. To retain this flexibility with fewer parameters, OAL groups dimensions with similar contribution scales to share coefficients, while allowing distinct treatment of the few outliers. The resulting positive quadratic kernel is calibrated against pretrained attention outputs and admits linear-time evaluation with respect to sequence length for fixed head dimensions. On Qwen2.5-0.5B, OAL reduces frozen-attention output relative MSE by approximately 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 compared with a learned global quadratic, while improving 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. The code is available at https://github.com/wuqianyue09/OAL
Related
- A Method for Linearization of Systems of Quadratic Equations
- Quadratic Scalar Curvature Decay and Uniform Positivity in Dimensions Four through Seven
- One-Cut Risk Profiles under Quadratic Loss: Discrete Convexity, Continuous Limits, and Higher Dimensions
- Almost global existence for quadratic Schrödinger equations in two dimensions
- Quadratic distances in even dimensions over prime fields
BibTeX
@misc{sun2026dimensions,
title = {{Not All Dimensions Are Equal: Outlier-Aware Attention Linearization with Grouped Quadratic Kernels}},
author = {Weiqi Sun and Qianyue Wu and Wenjia Wang and Tianyang Hu},
howpublished = {arXiv},
year = {2026},
month = oct,
doi = {10.48550/arxiv.2610.13674},
eprint = {2610.13674},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2610.13674},
}
RIS
TY - UNPB
TI - Not All Dimensions Are Equal: Outlier-Aware Attention Linearization with Grouped Quadratic Kernels
AU - Weiqi Sun
AU - Qianyue Wu
AU - Wenjia Wang
AU - Tianyang Hu
PY - 2026
DA - 2026/10/09
PB - arXiv
DO - 10.48550/arxiv.2610.13674
UR - https://arxiv.org/abs/2610.13674
AB - Linearizing softmax attention offers a route to efficient Transformer inference, but preserving pretrained attention remains challenging. We identify an important source of approximation error overlooked by standard linearization methods: a small number of outlier dimensions contribute disproportionately to query-key dot products and account for a large fraction of the approximation error. Motivated by this heterogeneity, we propose Outlier-Aware Attention Linearization (OAL), a dimension-aware method for constructing quadratic attention approximations. Starting from a second-order Taylor approximation, we introduce a Hadamard-product formulation that makes individual query-key contributions explicit. This formulation relaxes globally shared coefficients to support dimension-specific linear terms and cross-dimensional quadratic interactions while preserving separable evaluation. To retain this flexibility with fewer parameters, OAL groups dimensions with similar contribution scales to share coefficients, while allowing distinct treatment of the few outliers. The resulting positive quadratic kernel is calibrated against pretrained attention outputs and admits linear-time evaluation with respect to sequence length for fixed head dimensions. On Qwen2.5-0.5B, OAL reduces frozen-attention output relative MSE by approximately 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 compared with a learned global quadratic, while improving 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. The code is available at https://github.com/wuqianyue09/OAL
ER -