PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Weiqi Sun, Qianyue Wu, Wenjia Wang, Tianyang Hu

arXiv:2610.13674v1stat.MLcs.AIcs.LG

TL;DR

Linear attention offers a way to reduce sequence-length computation, but preserving pretrained attention under limited adaptation remains challenging. OAL exposes dimension-wise query–key contributions and groups similar contributions to construct a compact quadratic approximation. It improves frozen-attention fidelity and converted-model performance in the reported evaluations.

  • Problem

    Linear attention can reduce sequence-length computation, but approximations still need to preserve pretrained attention under limited adaptation, and existing parameterizations do not explicitly organize unequal dimension contributions.

    Source: §1 · pp. 1–2

  • Method

    OAL makes individual query–key contributions explicit with a Hadamard-product quadratic formulation and groups dimensions with similar contributions to share coefficients.

    Source: §1 · p. 2

  • Results

    OAL improves frozen-attention fidelity on Qwen2.5-0.5B and model-level performance after LoRA adaptation on both backbones; it reduces attention-output relative MSE by approximately 8.9% compared with DiJiang.

    Source: §1 · p. 2, §7 · p. 11

  • Takeaways & Limitations

    Contribution-based grouping retains dimension-aware flexibility in a compact kernel that supports linear sequence-length evaluation for fixed head dimensions.

    Source: §1 · p. 2, §7 · p. 11

    Experiments cover 0.5B–1B backbones, so performance at larger scales remains unverified.

    Source: §7 · p. 11

Abstract

from arXiv · show

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

1. Introduction

Linear attention can reduce sequence-length computation, but approximations still need to preserve pretrained attention under limited adaptation. OAL makes dimension-wise contributions explicit and groups similar contributions to build a compact approximation; it improves frozen-attention fidelity and converted-model performance in the reported evaluations.

1. Introduction

  • A small subset of query–key dimensions contributes disproportionately to attention logits and approximation error, while existing parameterizations are not explicitly organized around those unequal contributions.This motivates treating outlier dimensions differently from the remaining dimensions.Source: §1 · pp. 1–2
  • OAL uses a Hadamard-product formulation to expose individual query–key contributions, enabling dimension-specific linear and cross-dimensional quadratic coefficients while retaining separable evaluation.This generalizes the globally shared coefficients of a quadratic approximation.Source: §1 · p. 2
  • Grouping dimensions with similar contributions lets OAL share coefficients within groups and across group pairs, accommodating outliers with fewer learned coefficients.The grouped positive quadratic kernel is calibrated against pretrained attention outputs using frozen queries, keys, and values.Source: §1 · p. 2
  • OAL improves frozen-attention fidelity and converted-model performance; on frozen Qwen2.5-0.5B states, it reduces attention-output relative MSE by approximately 8.9% compared with DiJiang (Chen et al., 2024b).Under matched LoRA adaptation, it also outperforms evaluated converted methods on WikiText-103 perplexity and all three downstream tasks on both Qwen2.5-0.5B and Llama-3.2-1B.Source: §1 · p. 2
  • The paper characterizes dimension-wise heterogeneity, develops a grouped dimension-aware quadratic formulation, and evaluates frozen-attention fidelity and adapted model quality under a matched conversion protocol.It also examines group-aware flexibility through a diagnostic budget study.Source: §1 · p. 2

2. Preliminaries

Softmax attention computes pairwise token interactions, with quadratic cost in sequence length. Linear attention uses separable query and key features for linear sequence-length evaluation, while this section presents a second-order polynomial kernel as a starting point for dimension-aware approximations.

2. Preliminaries

  • Direct softmax attention evaluation costs O(N^2(d_h + d_v)) operations and O(N^2) memory for explicit score-matrix storage.The cost grows quadratically with sequence length.Source: §2 · p. 3
  • Linear attention approximates the softmax kernel with separable query and key features, enabling incremental causal evaluation with linear sequence-length scaling for fixed feature dimension and memory size.Methods differ mainly in how they construct the feature map and approximate the softmax kernel.Source: §2 · p. 3
  • A second-order Taylor expansion of exp(s) leads to a global quadratic kernel p(s) = a + bs + cs^2, whose coefficients control its constant, linear, and quadratic terms.This global quadratic is the starting point for the paper’s dimension-aware kernel.Source: §2 · p. 3

3. Dimension-Wise Heterogeneity in Pretrained Attention

The paper examines whether query–key contributions are concentrated in a few dimensions and whether this concentration is associated with attention-approximation error. It finds concentrated contributions and reports that approximating a small high-contribution subset retains substantial output error.

3.2. Approximation Error Concentration on High-Contribution Dimensions

  • The analysis uses coordinate-wise query–key products across causal pairs to define dimension contributions, then selects high-contribution dimensions and applies approximations only to that subset while leaving the rest exact.The retained-error ratio compares subset-only output MSE with the full-head error.Source: §3.1 · p. 4, §3.2 · p. 4

3.1. Query–Key Outliers and Dimension-Wise Contributions

  • The contribution profile is concentrated on a few dimensions; dimensions 26 and 58 have high contribution scores and correspond to prominent query or key coordinates.The figure shows query magnitude, key magnitude, and contribution score across dimensions.Source: §3.1 · p. 4, Figure 1 · p. 4

3.2. Approximation Error Concentration on High-Contribution Dimensions

  • About 5.5% of dimensions retained a substantial fraction of full-head attention-output MSE across evaluated methods, with ratios above Gaussian controls; none of 1,000 random-coordinate null replicates exceeded the observed ratios (empirical p < 0.001).The table reports retained-error ratio R_S (%) and notes that higher values indicate stronger error concentration.Source: §3.2 · p. 4, Table 1 · p. 5

4. Outlier-Aware Attention Linearization

OAL makes query–key dimension contributions explicit, groups dimensions to share coefficients, and evaluates a positive quadratic kernel without pairwise token interactions. Its grouped parameterization preserves linear sequence-length scaling for fixed head dimensions.

4.1. Dimension-Aware Linearization

  • The Hadamard–diagonal formulation supports dimension-specific linear terms and cross-dimensional quadratic interactions while evaluating the specified kernel exactly in linear sequence-length time.The quadratic term avoids forming an N × N matrix; the softmax approximation, not the recurrence, introduces approximation error.Source: §4 · p. 5, §4.1 · pp. 5–6

4.2. Outlier-Aware Dimension Grouping

  • For a representative head, a small number of groups captures most error reduction, with diminishing gains as the group budget increases.Dimensions are sorted by calibration contribution scores, and large gaps in the high-contribution tail determine group boundaries; Qwen2.5-0.5B uses 2–5 groups per head and Llama-3.2-1B uses 2–4.Source: §4.2 · p. 6

4.3. The OAL Algorithm

  • OAL parameterizes a positive quadratic kernel by sharing coefficients within contribution-based groups and group pairs.With groups fixed, its factors are calibrated against pretrained softmax attention outputs using frozen queries, keys, and values, then initialize joint adaptation with LoRA.Source: §4.3 · pp. 6–7

4.4. Complexity Analysis

  • OAL's causal evaluation has linear sequence-length cost for fixed head dimensions, with state space independent of sequence length.Grouping reduces learned parameters to O(G^2), while the implementation retains dimension-pair states.Source: §4.4 · p. 7
  • With the same block-GEMM backend, OAL (G=5) has latency and peak memory comparable to the single-group learned global quadratic, while FlashAttention remains faster and more memory-efficient.The reported cross-method comparisons reflect evaluated implementations with differing numerical precision and optimization levels.Source: §4.4 · pp. 7–8

5. Experiments

Experiments evaluate frozen-attention fidelity and model quality after attention replacement with limited LoRA adaptation. OAL leads the evaluated converted methods on both backbones across language modeling and downstream tasks.

5.1. Frozen-Attention Approximation Fidelity

  • Frozen-attention comparisons use identical pretrained states, with approximation error measured against softmax outputs before and after output projection Wo.The comparison covers all query heads in the specified layers after RoPE.Source: §5.1 · p. 8
  • OAL achieves the lowest MSE and relative MSE before and after output projection Wo, while DiJiang is the strongest evaluated baseline across all four metrics.The evaluation uses frozen Qwen2.5-0.5B queries, keys, and values from Layers 4, 12, and 20 at sequence length 4096.Source: §5.1 · pp. 8–9

5.2. End-to-End Conversion and Downstream Performance

  • Conversion experiments retain softmax attention in the initial and final layers and replace intervening layers, with all methods receiving matched LoRA configurations, training data, and update budgets.OAL jointly optimizes calibrated kernel coefficients with LoRA while keeping dimension groups fixed.Source: §5.2 · p. 9
  • OAL leads converted methods in WikiText-103 perplexity and all three downstream accuracies on both backbones; versus the learned global quadratic, perplexity falls from 31.71 to 27.50 on Qwen2.5-0.5B and from 36.03 to 32.99 on Llama-3.2-1B.The LoRA-adapted softmax reference remains stronger on every metric.Source: §5.2 · pp. 9–10, Table 3 · p. 10

6. Related Work

Related work covers efficient kernel approximations and pretrained-model conversion, alongside research on nonuniformity in pretrained Transformers. OAL distinguishes itself by organizing kernel coefficients according to dimension-wise contributions.

Kernel Linearization and Pretrained-Model Conversion

  • Kernel methods such as Performer, SpectraFormer, and DiJiang approximate softmax attention with feature representations for linear-time evaluation; OAL instead structures coefficients by dimension-wise contribution.The cited work also includes T2R, Hedgehog, and LoLCATs as approaches to pretrained-model conversion or attention preservation.Source: §6 · p. 10

Polynomial Attention

  • Polynomial-attention methods include second-order Taylor kernels, normalized and affine-transformed quadratic kernels, and sketches for higher-degree attention; OAL targets pretrained softmax recovery under limited adaptation.OAL retains coordinate products and learns linear-term and quadratic-interaction coefficients, with sharing based on contribution statistics.Source: §6 · p. 10

Activation and Attention Outliers

  • Prior work studies large hidden activations, attention sinks, and query/key alignment; OAL examines how concentrated query–key contributions relate to linearization error.That association guides OAL's coefficient sharing.Source: §6 · p. 10

7. Discussion

The discussion presents OAL as a grouped quadratic approach to linearizing pretrained softmax attention. It reports improved fidelity and adapted model performance, while limiting the evidence to small backbones.

7. Discussion

  • OAL uses a compact grouped quadratic kernel and supports exact recurrent evaluation with linear sequence-length complexity.Its design is motivated by approximation error concentrated on a small subset of high-contribution query–key dimensions.Source: §7 · p. 11
  • The discussion reports improved frozen-attention fidelity on Qwen2.5-0.5B and improved model-level performance after LoRA adaptation on Qwen2.5-0.5B and Llama-3.2-1B.Source: §7 · p. 11
  • Evaluation covers only 0.5B–1B backbones, so performance at larger scales remains unverified.The authors identify larger and longer-context models as future evaluation targets.Source: §7 · p. 11

AI Use Statement

The authors report using generative AI for manuscript editing and literature discovery, with all assisted content reviewed and verified by the authors.

AI Use Statement

  • Generative AI tools aided writing, including language editing and improving manuscript clarity and organization.Source: AI Use Statement · p. 11
  • The authors also used generative AI tools for retrieval and discovery, including related-work searches.Source: AI Use Statement · p. 11
  • The authors state that they reviewed and verified all AI-assisted content and take responsibility for the final manuscript.Source: AI Use Statement · p. 11

A. Additional Analysis of Query–Key Heterogeneity

Additional analyses test whether query–key contribution heterogeneity generalizes across models and controls, affects attention-output approximation error, and remains stable across inputs and sequence lengths.

A.1. Cross-Model and Control Evidence

  • GPT-2 Medium also shows concentrated dimension-wise contributions, though its dominant contribution is spread across more dimensions than in Qwen2.5-0.5B.The isotropic Gaussian controls indicate that the observed concentration is not explained solely by dimensionality or coordinate selection.Source: §A.1 · p. 14, Figure 5 · p. 15

A.2. High-Contribution-Dimension Intervention

  • The selected high-contribution dimensions total 147 across 42 query heads, approximately 5.5% of the dimensions evaluated.The intervention approximates the selected dimensions and leaves the complement exact at the unnormalized-kernel level; error is assessed on attention outputs.Source: §A.2 · pp. 15–16

A.3. Stability of Dimension-Wise Contributions

  • Calibration and held-out contribution profiles are highly consistent, with Spearman correlation ρ=0.991 and identical top-five dimensions.The comparison uses a representative Qwen2.5-0.5B attention head.Source: §A.3 · p. 17
  • Dimensions 26 and 58 remain the largest contributors across evaluated sequence lengths and jointly account for approximately 46–49% of total contribution.Together with the held-out-input result, the authors report stability under changes in input content and sequence length.Source: §A.3 · p. 17, Table 4 · p. 17

B. Grouping and Coefficient-Calibration Details

This section specifies how contribution-based groups are selected and kept fixed, how OAL coefficients are calibrated and evaluated, and how the grouped kernel relates to a global quadratic and separable realization.

B. Grouping and Coefficient-Calibration Details

  • Grouping sorts calibration contribution scores, searches for boundaries in the high-contribution tail, and typically separates a shared low-contribution bulk from refined high-contribution groups.Groups are determined independently for each layer and query head before coefficient fitting and remain fixed during calibration, adaptation, and evaluation.Source: B. Grouping and… · pp. 18–19
  • OAL coefficients are fitted by matching normalized causal outputs to pretrained softmax outputs with groups and pretrained query, key, and value states fixed.The calibrated factors initialize model-level adaptation, where factors are jointly optimized with LoRA while groups remain fixed; coefficient-fitting data are separate from the short-prompt grouping data.Source: B. Grouping and… · p. 19

C. Frozen-Attention Error Metrics

  • RelMSE normalizes each held-out sequence’s squared output error by the energy of its exact softmax output before averaging across sequences.For post-output-projection metrics, the exact and approximate head outputs are concatenated and passed through the frozen output projection before applying the same definitions.Source: C. Frozen-Attention Error Metrics · pp. 20–21

D. Global Reference and HD Realization

  • The learned global quadratic is a special case of OAL’s grouped parameterization, recovered by sharing the same linear and quadratic coefficients across all groups.The global reference uses a positive-semidefinite coefficient matrix with ε>0, ensuring strictly positive weights.Source: D. Global Reference and HD Realization · p. 21
  • The grouped-kernel identity follows from separable query and key factors, but pointwise positivity of OAL weights does not imply Mercer positive semidefiniteness.The implementation can pack symmetric pair states rather than storing both ordered pairs; off-diagonal pairs have multiplicity two during readout.Source: D. Global Reference and HD Realization · pp. 21–22

E. Runtime and Memory Benchmark

The runtime and memory measurements benchmark causal attention operators under specified hardware and implementation settings, while the conversion setup targets intermediate layers with pronounced activations and retains softmax attention elsewhere.

E. Runtime and Memory Benchmark

  • Figure 3 measures forward-only causal attention on an A100 with batch size 1, 14 query heads, 2 key/value heads, and head dimensions 64.Sequence lengths range from 2^10 to 2^15, and the measurements are operator-level rather than end-to-end or autoregressive decoding latency.Source: E. Runtime and Memory Benchmark · p. 22
  • OAL uses a fixed five-group configuration and shares an optimized HD block-GEMM backend with the learned global quadratic, unlike the other methods’ causal execution paths.Implementations differ in optimization, and arithmetic precision also varies; measured cross-method gaps therefore cannot be attributed solely to kernel formulations.Source: E. Runtime and Memory Benchmark · p. 22
  • Peak GPU memory includes resident inputs, method state, and workspace; SpectraFormer and math-backend SDPA sweeps stop at sequence length 8192 under a pre-run memory guard.The 16384-token cases were skipped because estimated peaks exceeded 80% of free GPU memory, not because observed out-of-memory failures occurred.Source: E. Runtime and Memory Benchmark · p. 22

F. Conversion and LoRA Setup

  • Pronounced intermediate-layer activation outliers motivate the replacement ranges, which retain softmax attention in the first and last layers of each backbone.The converted models are hybrid, so retained softmax layers still contribute quadratic prefill computation.Source: F. Conversion and LoRA Setup · pp. 22–23

F.1. Baseline Implementations and Adaptations

The paper adapts several attention approximations within a common pretrained-backbone conversion protocol, then applies matched LoRA training and evaluates language modeling and downstream accuracy.

F.1. Baseline Implementations and Adaptations

  • The baseline comparison integrates adapted kernels into pretrained backbones while retaining native projections, RoPE, and GQA mapping, rather than reproducing each method’s original full architecture and training recipe.The learned global quadratic is a controlled reference sharing OAL’s output-calibration protocol, isolating global versus contribution-based coefficient sharing within the quadratic family.Source: §F.1 · pp. 23–24
  • Performer uses 250 positive orthogonal random features per head, with fixed random projections and FP64 feature computation and prefix accumulation.Only LoRA parameters are trained for Performer.Source: §F.1 · p. 23
  • SpectraFormer uses 65 positive features and three trainable FastFood vectors per layer, totaling 192 kernel parameters per layer.The vectors are shared across query heads within a layer and use FP64 feature and accumulation arithmetic.Source: §F.1 · p. 24
  • DiJiang adds trainable positive feature weights per layer and query head, while omitting the released model’s positional decay, GroupNorm, and output gate.Its reported row is the normalized causal adaptation of the released feature map.Source: §F.1 · p. 24
  • TaylorShift retains a fixed second-order polynomial and learns a per-head temperature, using causal prefix evaluation without optional output-amplitude normalization.The temperature is trained jointly with LoRA.Source: §F.1 · p. 24

F.2. LoRA Adaptation

  • All methods use rank-8 LoRA on query, key, value, and output projections, with backbones frozen and 512 updates on WikiText-103.OAL’s kernel coefficients are jointly optimized with LoRA while its groups remain fixed; the runs use one training seed.Source: §F.2 · p. 25

F.3. Evaluation Protocol

  • Evaluation uses WikiText-103 test perplexity and zero-shot accuracy on PIQA, HellaSwag, and ARC-Easy under a shared tokenizer and task protocol within each backbone.Perplexity is computed over non-overlapping 4096-token blocks; downstream scoring uses candidate-token conditional log-likelihood with task-specific normalization.Source: §F.3 · p. 25
Loading 2610.13674v1…