Attention linearization stops treating all dimensions alike: grouping the few outliers cuts Qwen2.5-0.5B perplexity from 31.71 to 27.50

Linearizing attention replaces softmax attention with separable query and key features so computation scales linearly with sequence length, but preserving the behavior of the pretrained attention it replaces has stayed difficult. Earlier linearization methods build vector-level feature maps or low-order polynomials in the aggregated query–key score, and their parameterization is not organized around how much individual query–key dimensions contribute. On frozen Qwen2.5-0.5B attention states, OAL reaches 0.432 relative MSE on the attention output, below DiJiang at 0.474, TaylorShift at 0.490, Performer at 0.569, and SpectraFormer at 0.585.

The method splits the quadratic kernel's coefficients by dimension, so the few dimensions that contribute most get different coefficients from the rest while separable evaluation is retained. A global quadratic shares one coefficient across all linear terms and one across all quadratic interactions, whereas OAL relaxes that sharing and gives each dimension and dimension pair its own coefficients. On Qwen2.5-0.5B and Llama-3.2-1B, after replacing middle-layer attention and fine-tuning with a matched LoRA budget, OAL reaches WikiText-103 perplexity of 27.50 and 32.99, below the learned global quadratic at 31.71 and 36.03 and below Performer, SpectraFormer, DiJiang, and TaylorShift.

Dimensions are sorted by calibration contribution and grouped so that similar dimensions share coefficients while the high-contribution tail is subdivided, with the groups fixed during fitting and adaptation. A fully dimension-specific quadratic needs a large number of coefficients, while the contribution distribution shows much of that flexibility can be shared. Qwen2.5-0.5B yields 2 to 5 groups per head and Llama-3.2-1B 2 to 4; on the representative head, a small number of groups captures most of the error reduction, with diminishing gains as the budget grows.

The reported scope is 0.5B to 1B backbones, Qwen2.5-0.5B and Llama-3.2-1B, under a hybrid replacement protocol that keeps softmax attention in the first and last layers and adapts with a matched LoRA budget on WikiText-103. Frozen-attention fidelity is evaluated on all query heads of Layers 4, 12, and 20 of Qwen2.5-0.5B Base, with queries and keys taken after RoPE. Runtime and memory measurements are forward-only attention-operator measurements, not end-to-end model latency or per-token decoding latency. The authors propose evaluating OAL on larger and longer-context models and exploring more compact implementations of the quadratic recurrent state.

Experiments cover only 0.5B to 1B backbones, so performance at larger scales remains unverified. A gap to the LoRA-adapted softmax reference remains on both backbones. Runtime and memory comparisons reflect the evaluated implementations, which differ in numerical precision and optimization level, so cross-method gaps cannot be attributed solely to the kernel formulations. The group-budget diagnostic covers one representative head, and the runtime benchmark uses a fixed five-group configuration rather than executing the heterogeneous group configurations of all calibrated heads in a converted model.