OAL: A New Outlier-Aware Approach to Attention Linearization

The Old Problem and a New Finding in Attention Linearization

The softmax attention mechanism in Transformers carries quadratic computational overhead with respect to sequence length during inference, making efficiency a critical bottleneck for long-sequence scenarios. Linearization has long been a key strategy for efficient inference, but a persistent challenge remains: how to approximate attention with linear complexity while preserving the behavior of the original pretrained model as closely as possible.

A new study titled Not All Dimensions Are Equal targets an error source that standard linearization methods have broadly overlooked — outlier dimensions. The researchers found that in query-key dot products, a small number of dimensions contribute a disproportionate share of the energy and account for a large portion of the approximation error. In other words, not all dimensions should be treated equally, and using globally shared coefficients to uniformly fit all dimensions will naturally produce larger distortions in these outlier dimensions.

The basic idea of linearized attention is to replace $\exp(q^\top k)$ in softmax with a kernel function $\phi(q)^\top \phi(k)$, reordering the computation so that the cumulative sum of $\phi(k)^\top v$ is computed first, then multiplied by $\phi(q)$. This makes the computation per token independent of sequence length, reducing overall complexity from $O(n^2 d)$ to $O(nd^2)$. Early representative works include Performer (random feature approximation) and Linear Transformer. However, these approaches have consistently struggled in practice with rising perplexity and degraded long-range dependency modeling, primarily because the sharp distribution of softmax is difficult to accurately reconstruct with low-dimensional feature maps. Subsequent work like DiJiang attempted to improve approximation quality by introducing quadratic terms, but still assumed uniform treatment across dimensions. The outlier dimension problem identified by OAL is precisely the systematic underestimation of error hidden within this uniform assumption.

OAL's Core Idea: Dimension-Aware Quadratic Kernels

Building on the observation of dimensional heterogeneity, the authors propose Outlier-Aware Attention Linearization (OAL) — a dimension-aware method for constructing quadratic attention approximations. Its technical approach can be understood in several steps.

From Second-Order Taylor Expansion to Hadamard Products

OAL doesn't directly swap in a standard kernel function replacement. Instead, it starts from a second-order Taylor approximation of softmax attention. The researchers introduce a Hadamard product (element-wise product) formulation that makes each dimension's contribution to the query-key interaction explicitly visible. This reformulation is crucial — it decomposes the dot product that was previously packed together, providing the mathematical foundation for treating different dimensions differently downstream.

The second-order Taylor expansion approximates $\exp(x)$ as $1 + x + x^2/2$. Applied to dot-product attention, $\exp(q^\top k)$ can be expanded into a combination of a constant term, a linear term $q^\top k$, and a quadratic term $(q^\top k)^2/2$. Critically, $(q^\top k)^2$ can be written as $(q \odot q)^\top (k \odot k)$ (where $\odot$ denotes element-wise multiplication, i.e., the Hadamard product), which means the quadratic term can still be computed in a separable fashion — accumulating on the key side first, then taking a dot product with the query side — thereby maintaining linear time complexity. The advantage of the Hadamard formulation is that it makes each dimension's contribution explicitly separable in the expression, providing a direct mathematical entry point for dimension-differentiated treatment (i.e., outlier dimension grouping), rather than mixing all dimensions into a single global coefficient.

From Globally Shared Coefficients to Dimension-Specific Terms

Traditional methods typically use a set of globally shared coefficients to construct the approximation. OAL relaxes this constraint, supporting dimension-specific linear terms and cross-dimension quadratic interaction terms, while maintaining a separable evaluation structure — a prerequisite for preserving linear time complexity. In other words, flexibility increases without sacrificing efficiency.

Controlling Parameter Count Through Grouping

Fully dimension-specific treatment would cause parameter explosion. OAL's cleverness lies in grouping dimensions by contribution scale: dimensions with similar contributions share coefficients, while the few outlier dimensions are handled individually. This preserves fine-grained modeling of outlier terms while keeping overall parameter count manageable. The resulting positive quadratic kernel is calibrated against pretrained attention outputs and, for a fixed head dimension, achieves linear-time evaluation relative to sequence length.

Experimental Results: Reduced Error and Downstream Gains

The paper validates OAL on Qwen2.5-0.5B, providing evidence on two fronts.

In terms of frozen-attention approximation error, OAL reduces the relative mean squared error (MSE) of the output by approximately 8.9% compared to DiJiang. This confirms that dimension-aware grouping does capture error sources that standard methods miss.

At the model conversion and fine-tuning level, the researchers replaced attention modules in intermediate layers with OAL and jointly fine-tuned the kernel coefficients using LoRA. Compared to a learned global quadratic kernel, WikiText-103 perplexity dropped from 31.71 to 27.50, with accuracy improvements also observed on commonsense reasoning benchmarks including PIQA, HellaSwag, and ARC-E.

The signal from these numbers is consistent: the outlier-aware grouping strategy both improves pure attention approximation and facilitates the smooth conversion of pretrained models into their linearized counterparts.

Implications for Efficient Inference

OAL's value lies not just in a specific error percentage, but in its re-examination of the error structure in linearization approximation. Past research has largely focused on kernel function forms and feature mappings while implicitly assuming all dimensions are equal. OAL demonstrates that this assumption does not hold in real models — a small number of outlier dimensions dominate the error, and targeting them yields significant gains.

This line of thinking resonates with how the quantization community approaches outliers: in quantization, outlier channels are similarly a primary source of precision loss, and selectively protecting them has become standard practice. OAL transfers analogous insight to attention linearization, providing a new design dimension for efficiently converting pretrained models.

For developers looking to convert existing Transformers to linear-complexity inference without full retraining, the "grouping + special outlier handling + LoRA calibration" combination that OAL offers is worth paying attention to. The paper's code is open-sourced on GitHub (github.com/wuqianyue09/OAL) for reproduction and further exploration.

In the LLM quantization space, the outlier problem has been thoroughly studied. LLM.int8() found that certain channels in Transformers have activation values far exceeding others, directly impacting integer quantization precision; methods like SmoothQuant and AWQ address this by migrating outlier channels through scaling or mixed-precision protection. These quantization works and OAL collectively reveal a more general phenomenon: the intermediate representations of pretrained large models are naturally highly non-uniform in their dimensional distribution, with a small number of dimensions carrying a disproportionate amount of information. Ignoring this non-uniformity — whether in quantization or attention linearization — will produce concentrated precision losses at these critical dimensions. OAL transfers insights already validated in quantization to linearization approximation, suggesting that "identifying and specially handling outlier dimensions" may be a broadly applicable principle for efficiently adapting pretrained models across methods.

Summary

OAL's core contribution is operationalizing the observation that "not all dimensions are equal" into a practical quadratic kernel construction method. Starting from a second-order Taylor expansion, it uses the Hadamard formulation to make dimensional contributions explicit, applies grouping to control parameter count, and models outlier dimensions individually — achieving improvements in both approximation error and downstream task performance while maintaining linear-time evaluation. For researchers focused on efficient Transformer inference and pretrained model conversion, this is a technically well-motivated and empirically grounded approach.