Norm$\times$Direction: Restoring the Missing Query Norm in Vision Linear Attention

Linear attention mitigates the quadratic complexity of softmax attention but suffers from a critical loss of expressiveness. We identify two primary causes: (1) The normalization operation cancels the query norm, which breaks the correlation between a query's norm and the spikiness (entropy) of the attention distribution as in softmax attention. (2) Standard techniques for enforcing non-negativity cause destructive information loss by nullifying valid inner-product interactions. To address these challenges, we introduce **NaLaFormer**, a novel linear attention mechanism built upon a norm$\times$direction (ND) decomposition of the query and key vectors. We leverage each component to solve a distinct problem: The *query norm* is injected into our kernel to create a query-norm-aware map that restores the attention distribution's spikiness. The *direction vectors* are processed by a geometric, cosine-based similarity metric that guarantees non-negativity while preserving the rich, fine-grained information of the inner product. We validate NaLaFormer through a comprehensive multi-modal evaluation, where it sets new state-of-the-art benchmarks for linear attention. Our model achieves up to a 7.5\% accuracy gain on ImageNet-1K and a 4.7\% mIoU improvement on ADE20K over comparable baselines. It demonstrates profound efficiency, reducing peak memory by a transformative 92.3\% in token-intensive super-resolution tasks (70K+ tokens). NaLaFormer's versatility is further confirmed as it surpasses strong baselines like Mamba on common-sense reasoning and sets a new state-of-the-art on the Long Range Arena (LRA) benchmark. Source code can be found in the supplementary materials.

Recent artificial intelligence models widely rely on attention mechanisms to help determine which parts of the input are more important. While traditional attention mechanisms perform well, their quadratic computational complexity makes them significantly slow when processing high-resolution images or long texts. To accelerate computation, researchers proposed "linear attention", however, this approach often suffers from a loss of critical expressiveness, leading to an inability to focus on key information. This paper identifies that the root cause lies in linear attention neglecting the norm of each query vector, which acts as an essential importance signal, while also destructively losing valuable interaction details to enforce non-negativity. To resolve these issues, this paper propose NaLaFormer, a novel method that reintroduces query norm information to help the model "refocus" on critical content, while utilizing a clever trigonometric function-based technique to preserve fine-grained details under non-negativity constraints. Empirical results show that NaLaFormer improves accuracy by up to 7.5% on image classification and reduces peak memory usage by 92.3% in image super-resolution tasks. It also outperforms current linear models in language understanding and long-sequence benchmarks, offering a new paradigm for designing linear attentions that balance both efficiency and performance.