跳到正文
arXiv cs.AI· Weiqi Sun, Qianyue Wu, Wenjia Wang, Tianyang Hu·· 6 天前AI 评分20

OAL:面向离群维度的分组二次核注意力线性化方法

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

AI 导读

研究者提出 Outlier-Aware Attention Linearization(OAL),一种维度感知的二次注意力近似方法,通过分组共享系数并单独处理少数离群维度,在固定头维度下实现相对序列长度的线性时间评估。

正文 · 原文

View PDF HTML (experimental)

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 this https URL
Subjects: Machine Learning (stat.ML); Artificial Intelligence (cs.AI); Machine Learning (cs.LG)
Cite as: arXiv:2610.13674 [stat.ML]
  (or arXiv:2610.13674v1 [stat.ML] for this version)
  https://doi.org/10.48550/arXiv.2610.13674

arXiv-issued DOI via DataCite (pending registration)

Submission history

From: Tianyang Hu [view email]
[v1] Fri, 9 Oct 2026 14:40:19 UTC (680 KB)

来源:arXiv cs.AI · arxiv.org