Norm-Aware Linear Attention

• Norm-Aware Linear Attention is a family of efficient linearized transformer attention mechanisms that restore query norm sensitivity to recover softmax-like entropy dynamics.
• It employs dynamic kernel exponentiation and norm-preserving transformations to ensure linear computational complexity while enhancing selectivity across various domains.
• Empirical evaluations demonstrate significant improvements in accuracy, perplexity, and detection metrics in vision, language, and speech tasks.

Norm-Aware Linear Attention encompasses a family of linearized attention mechanisms for transformers that explicitly re-incorporate the norm (magnitude) of the query—often neglected in vanilla linear attention—in order to recover critical properties of softmax attention. The principal aim is to reconcile the computational scalability of linear attention ($O(N)$ in sequence length) with the desirable “spikiness” (entropy reduction) and expressivity that softmax enables, which is fundamentally tied to the norm of the query vector. Multiple instantiations exist, including NaLa (“Norm-Aware Linear Attention”), MALA (“Magnitude/Amplitude-Aware Linear Attention”), and norm-stabilized variants in gated recurrent frameworks. These architectures have demonstrated state-of-the-art results on a wide range of vision, language, and speech modeling tasks, with compelling theoretical and empirical justification for explicit query-norm modeling (Meng et al., 26 Jun 2025, Fan et al., 1 Jul 2025, Tang et al., 18 Nov 2025, Lu et al., 3 Feb 2025).

1. Mathematical Foundations and Query-Norm Coupling

Classic softmax attention yields quadratic complexity with respect to sequence length but achieves dynamic, norm-sensitive entropy reduction. Given queries $Q \in \mathbb{R}^{N\times d}$, keys $K \in \mathbb{R}^{N \times d}$, and values $V$, the softmax attention output at position $t$ is

$$o_t = \sum_{i=1}^N \mathrm{Softmax}_i\left(\frac{q_t^\top k_i}{\sqrt{d}}\right) v_i$$

where $q_t$ and $k_i$ can be factored as $q_t = \|q_t\| d(q_t)$ and $k_i = \|k_i\| d(k_i)$, with $d(q_t)$ and $d(k_i)$ unit direction vectors. The “spikiness” (sharpness) of the attention distribution, governed by entropy, increases with $\|q_t\|$, concentrating probability mass and enhancing selectivity—an essential property for effective modeling.

In contrast, classical linear attention (e.g., with kernel $\phi(x)$) approximates the softmax kernel by

$$\exp(q_t^\top k_i / \sqrt{d}) \approx \phi(q_t)^\top \phi(k_i)$$

leading to

$$\tilde{Y}_i = \frac{\phi(Q_i) \left(\sum_j \phi(K_j)^\top V_j\right)}{\phi(Q_i)\left(\sum_m \phi(K_m)^\top\right)}$$

However, the scaling property of $\phi$ means the query norm cancels in the normalization; all queries, regardless of magnitude, yield attention distributions of the same entropy. This undermines the controllable selectivity intrinsic to softmax attention (Fan et al., 1 Jul 2025, Meng et al., 26 Jun 2025).

Norm-aware mechanisms such as NaLa and MALA restore the influence of the query norm through feature map design, norm-and-direction decomposition, or explicit scaling and offset strategies. For example, in NaLa, a dynamic power-law exponent $p(\|q_t\|) = \lambda(1/2 + \tanh \|q_t\|)$ governs the spikiness by amplifying per-feature contributions as the query norm increases (Meng et al., 26 Jun 2025). In MALA, dedicated scaling ($\beta_i$) and offset ($\gamma_i$) factors, directly parameterized by the inner product $c_i = \phi(Q_i) \cdot \sum_j \phi(K_j)$, ensure that larger query norms translate to more concentrated (lower-entropy) attention patterns, mimicking the behavior of softmax (Fan et al., 1 Jul 2025, Tang et al., 18 Nov 2025).

2. Kernel Design, Nonnegativity, and Norm Preservation

Constructing a suitable kernel for linear attention necessitates both nonnegativity and retention of norm/direction information. “Norm-blind” kernels, such as $\operatorname{ReLU}(x) + 1$, fail to produce sharper score distributions as the query norm increases; all attention outputs exhibit constant entropy, limiting model expressivity (Meng et al., 26 Jun 2025).

The NaLa approach employs a decoupling of norm and direction, mapping $q_t$ and $k_i$ to

$$\phi_q(q_t) = |d(q_t)|^{p(\|q_t\|)} \begin{bmatrix} \cos d(q_t) \ \sin d(q_t) \end{bmatrix} \in \mathbb{R}^{2d}$$

$$\phi_k(k_i) = |k_i|^{\lambda} \begin{bmatrix} \cos d(k_i) \ \sin d(k_i) \end{bmatrix}$$

followed by an inner product structure that guarantees nonnegativity via $$\langle\phi_\mathrm{dir}(u), \phi_\mathrm{dir}(v)\rangle = \sum_m \cos(u_m - v_m) > 0$$, preserving vector norm fidelity and enabling nonnegative interactions without the drawbacks of elementwise ReLU (Meng et al., 26 Jun 2025).

Similarly, MALA leverages a combination of additive and multiplicative norm modulation, directly controlling the sharpness of the kernel. ReGLA further integrates bounded and variance-normalized exponential feature maps with explicit LayerNorm to manage gradient stability and ensure controlled norm propagation (Lu et al., 3 Feb 2025).

3. Algorithmic Structure and Computational Complexity

All major norm-aware linear attention mechanisms retain $O(N)$ time and memory complexity. NaLa decomposes each query and key, computes dynamic kernel exponents, and applies the norm-preserving cosine-sine mapping, followed by kernelized dot-product summaries and normalization—all in linear time:

• Query/key decoupling and feature transformation: $O(N d)$
• Attention summaries and normalization: $O(N d_v d)$ No $N \times N$ matrix or loop over all pairs is required (Meng et al., 26 Jun 2025).

MALA and variants employ batchwise precomputation of key-value summaries and the necessary scaling/offset factors for each query, preserving amortized linear throughput:

• $O(N d_k d_v + N d_k + N d_v)$ total cost, matching the efficiency of other linear kernels (Tang et al., 18 Nov 2025, Fan et al., 1 Jul 2025).

ReGLA introduces minimal $O(d)$ overhead for LayerNorm, exponentiation, and gating refinements, while retaining a recurrent state-of-size $O(d^2)$ per token (Lu et al., 3 Feb 2025). In all cases, the resulting computation is scalable to long contexts and large batch sizes.

4. Theoretical Properties: Entropy Control and Robustness

A central theoretical guarantee is the reinstatement of query-norm-dependent entropy reduction. In NaLa, the kernel design ensures that the entropy of the resulting positive sequence decreases monotonically with increasing query norm, closely paralleling softmax’s behavior (Meng et al., 26 Jun 2025). This monotonicity arises from the concavity of positive sequence entropy and the exponential tilting provided by the dynamic power-law kernel.

MALA’s scaling and offset ($\beta_i$, $\gamma_i$) establish algebraically that, as $\|\phi(Q_i)\|$ increases, both the top–average score difference and overall attention sharpness increase, restoring softmax-like adaptive concentration. The resulting attention distribution interpolates between the flatness of standard linear attention and the spikiness of softmax, offering a more balanced structure with controlled smoothness (Fan et al., 1 Jul 2025, Tang et al., 18 Nov 2025).

Norm-stabilized mechanisms such as ReGLA additionally introduce variance normalization and LayerNorm on Q/K inputs, empirically demonstrating drastic improvements in perplexity and stability, especially on long-sequence or deep-stack tasks (Lu et al., 3 Feb 2025). The use of nonnegative, bounded feature maps (e.g., normalized exponentials) prevents outlier activations and norm drift.

5. Empirical Performance and Applications

Norm-aware linear attention models consistently outperform “norm-blind” linear baselines across domains:

• Vision: On ImageNet-1K, NaLa-S achieves 84.3% and NaLa-B 85.2% top-1 accuracy, exceeding linear baselines by 1.6–2.2%. MALA-based MAViT models reach 85.7–86.0%, surpassing previous linear and sometimes even softmax-based benchmarks (Meng et al., 26 Jun 2025, Fan et al., 1 Jul 2025).
• Object Detection and Segmentation: NaLa-T and MALA improve AP metrics by 2–3 points over linear baselines and match or surpass heavier transformer/convolutional backbones on COCO, ADE20K, and UPerNet (Meng et al., 26 Jun 2025, Fan et al., 1 Jul 2025).
• Language Modeling: NaLa-DeltaNet and MALA-based models reduce perplexity by 1–3 points compared to standard linear attention, improving zero-shot and commonsense-QA accuracy (Meng et al., 26 Jun 2025, Fan et al., 1 Jul 2025).
• Speech and Generative Vision: In ultra-lightweight speech enhancement (IMSE), MALA reduces parameter count by up to 16.8% while maintaining or improving PESQ metrics over prior Taylor-based and deformable embedding methods (Tang et al., 18 Nov 2025). MALA attains lower FID scores and higher IS on diffusion-based image generation (Fan et al., 1 Jul 2025).
• Efficiency: Linear or better throughput is preserved; wall-clock times on modern hardware are up to 4× faster than softmax transformers (Meng et al., 26 Jun 2025).

A summary of empirical claims is presented in the table below:

Application Domain	Norm-Aware Method	Key Metric(s)	Performance Gain
ImageNet-1K	NaLa, MALA	Top-1 Accuracy	+1.6–2.2% over linear baselines
COCO Object Detection	NaLa-T, MALA	AP^b, AP^m	+2–3 points over PolaFormer, MILA
Semantic Segmentation	NaLa-S, MALA	mIoU	+0.6–2.7% over ViG-S, PolaFormer
Speech Enhancement	MALA (IMSE)	PESQ Score	14.6–16.8% param. reduction, stable
Language Modeling	NaLa-DeltaNet, MALA	Perplexity, QA Acc	–1–3 PPL, +0.7% QA accuracy
Diffusion Gen. Vision	MALA	FID, IS	FID ↓, IS ↑ over prior methods

(Meng et al., 26 Jun 2025, Fan et al., 1 Jul 2025, Tang et al., 18 Nov 2025)

6. Norm-Aware Design in Broader Linear Attention Architectures

ReGLA demonstrates how norm-aware design principles apply in Gated Linear Attention frameworks, introducing normalized exponential kernels, explicit variance scaling, and an auxiliary LayerNorm to stabilize model state. Experimental ablations confirm that omitting normalization increases language modeling perplexity by 18–60%, directly confirming the necessity of norm stabilization (Lu et al., 3 Feb 2025). The “refined gate” further enables robust gradient backpropagation, addressing vanishing signal in saturated gating regimes.

Practical guidelines extracted from these works include:

• Always use bounded, nonnegative feature maps for stability.
• Explicitly rescale feature norms to maintain desired variance.
• Retain normalization in the denominator, unless alternative controls are implemented.
• Insert explicit LayerNorms prior to kernel mapping.
• Consider dynamic (query-norm-dependent) kernel exponents or explicit scaling/offsets to recuperate softmax-like entropy dynamics.

7. Strengths, Limitations, and Open Challenges

Strengths of norm-aware linear attention include:

• Restored norm-sensitive entropy modulation and nonnegativity, closing the expressivity gap of linear models compared to softmax.
• Consistent empirical improvements on vision, language, and speech tasks—up to +4.2% over non-norm-aware baselines.
• Linear complexity and implementation simplicity for integration into transformer backbones and U-Net architectures.

Limitations and remaining open issues:

• Residual performance gap relative to full softmax attention for fine-grained tasks; empirical evidence mainly covers discriminative rather than generative regimes.
• Theoretical flexibility may be restricted by patterning norm-dependence solely through tanh-regularized power-law exponents or scalar scaling, which could be further generalized.
• Hyperparameter choices (e.g., $\lambda$ in NaLa, feature map selection) may be ad hoc, and richer, possibly learned entropy control mechanisms are plausible directions (Meng et al., 26 Jun 2025, Fan et al., 1 Jul 2025).
• Extension to large-scale generative modeling tasks (autoregressive LMs, diffusion) is incompletely explored.

Norm-Aware Linear Attention models constitute a principled, scalable, and empirically validated class of attention mechanisms, enabling high-fidelity global modeling in transformers with linear resource demands, and providing a template for future kernel and normalization innovations (Meng et al., 26 Jun 2025, Fan et al., 1 Jul 2025, Tang et al., 18 Nov 2025, Lu et al., 3 Feb 2025).