本文一作孟維康是哈爾濱工業(yè)大學（深圳）與鵬城實驗室聯(lián)合培養(yǎng)的博士生，本科畢業(yè)于哈爾濱工業(yè)大學，主要研究方向是高效能基礎模型。通訊作者張正教授，哈爾濱工業(yè)大學（深圳）的長聘教授及博士生導師，教育部青年長江學者，廣東特支計劃青年珠江學者，深圳市優(yōu)青。長期從事高效能多模態(tài)機器學習的研究，專注于高效與可信多模態(tài)大模型。

當 Transformer 席卷計算機視覺領域，高分辨率圖像、超長序列任務帶來的算力與顯存瓶頸愈發(fā)凸顯：標準 Softmax 注意力的二次復雜度，讓 70K+token 的超分辨率任務直接顯存爆炸，高分辨率圖像分割、檢測的推理延遲居高不下。

線性注意力雖通過核函數重構實現了線性復雜度，完美解決了算力開銷問題，卻始終無法擺脫性能退化的問題，與原生 Softmax 注意力的精度差距始終難以彌合。

近日，哈工深張正團隊、聯(lián)合鵬城實驗室、昆士蘭大學等團隊，發(fā)布重磅論文《Norm×Direction: Restoring the Missing Query Norm in Vision Linear Attention》，提出 NaLaFormer（Norm-aware Linear Attention Transformer） 框架，首次通過模長 - 方向分解，精準定位并解決了線性注意力的兩大核心缺陷，在保持線性復雜度的同時，實現了視覺任務的精度全面超越，更在 70K+token 超分任務中實現 92.3% 的峰值顯存降低，為線性注意力在視覺領域的落地開辟了全新范式！

• 論文標題：Norm×Direction: Restoring the Missing Query Norm in Vision Linear Attention
• 論文鏈接：https://arxiv.org/pdf/2506.21137
• 核心作者：Weikang Meng、Yadan Luo、Liangyu Huo、Yingjian Li、Yaowei Wang、Xin Li、Zheng Zhang

痛點直擊：

線性注意力性能崩塌的兩大致命傷

線性注意力的核心目標，是通過線性可分核函數替代 Softmax 中的指數算子，利用矩陣結合律將計算復雜度從 O (N2) 降至 O (N)。但在 NaLaFormer 的研究中，首次系統(tǒng)性揭示了現有方案始終無法逼近 Softmax 性能的兩大根源：

1、Query Norm 被抵消，注意力尖峰性徹底丟失

團隊通過數學推導與實驗驗證發(fā)現：在 Softmax 注意力中，Query 向量的模長（Norm）與注意力分布的熵呈強負相關 ——Query 模長越大，注意力分布越尖銳（熵越低），模型越能精準聚焦于語義關鍵的 token，這是 Softmax 注意力具備強表征能力的核心原因。但在傳統(tǒng)線性注意力中，歸一化操作會直接抵消 Query Norm 的影響，讓線性注意力徹底失去了對注意力分布尖峰性的動態(tài)調控能力，最終輸出過于平滑的注意力分布，無法聚焦關鍵信息，表征能力大幅退化。

圖 1 Softmax 注意力中熵與模長的相關性：Query 模長（x 軸）與注意力熵（y 軸）呈強負相關，Key 模長則無穩(wěn)定影響

2、非負性約束導致不可逆的信息損失

線性注意力的核函數必須滿足非負性，才能將注意力得分解釋為歸一化的概率分布。現有方案普遍通過 ReLU、1+ELU 等激活函數直接抹除向量中的負值，卻直接導致了 q-k 內積中有效交互信息的丟失 —— 原本具有區(qū)分度的負向語義關聯(lián)被直接清零，最終讓相似度表征變得稀疏且缺乏細粒度信息。

圖 2 不同非負策略的內積貢獻對比：ReLU、1+ELU 均丟失了原始內積的尖峰性與細粒度信息，而本文提出的余弦方向方法完美保留了原始分布特征

核心創(chuàng)新：

Norm×Direction 分解，補全線性注意力機制短板

針對上述兩大核心痛點，NaLaFormer 團隊提出了基于 Norm×Direction（ND）向量分解的解決方案，將向量的模長（Norm，表征信息重要性）與方向（Direction，編碼語義信息）解耦，分別針對性解決兩大缺陷，最終實現了對 Softmax 注意力核心特性的完美復刻，同時完整保留線性復雜度。

1、Query-Norm-Aware Feature Map：重建丟失的模長 - 信息熵關聯(lián)

團隊首先通過 ND 分解，對線性注意力的計算過程進行了數學重構，精準定位了 Query Norm 被抵消的核心環(huán)節(jié)，清晰揭示了傳統(tǒng)線性注意力的輸出僅受 Key Norm 影響，Query Norm 在歸一化中被完全抵消，導致模型對 Query Norm 的無感。

對此，團隊設計了 Query-Norm-Aware 特征映射，將 Query Norm 顯式編碼進核函數中，通過 Norm 依賴的銳化函數動態(tài)調控注意力熵：

其中銳化函數 f (x)=λ?(τ+tanh (x))，可根據 Query Norm 動態(tài)調整冪次，完美復刻了 Softmax 注意力中 “Query 模長越大，注意力分布越尖銳” 的核心特性，重建了 Query Norm 與注意力熵的負相關關系。

2、余弦方向相似度：無損實現非負性，零信息丟失

針對傳統(tǒng)非負約束的信息丟失問題，團隊基于托勒密理論，提出了全新的余弦方向相似度機制，對方向分量進行幾何變換，在保證非負性的同時，完整保留內積的細粒度信息。

團隊為每個標量方向分量設計了二維向量映射：

經過變換后，Query 與 Key 的內積可轉化為：

通過 tanh 映射將方向分量縮放至 [?π/4, π/4]，即可保證余弦項始終非負，同時完整保留了原始內積的方向交互信息，徹底避免了傳統(tǒng)激活函數帶來的信息損失。

3. NaLaFormer：統(tǒng)一的 Norm 感知線性注意力架構

基于上述兩大核心創(chuàng)新，團隊構建了端到端的 NaLaFormer 架構，將模長感知線性注意力模塊與門控架構深度融合，實現了對視覺 Transformer 的直接替換。

圖 3 NaLaFormer 整體架構與核心機制：(a) NaLaFormer 基礎模塊；(b) Norm 感知方法恢復了 Query Norm - 熵負相關；(c) 余弦方向機制實現無損非負性約束

最終的注意力輸出公式為：

實驗結果

? 圖像分類任務 ImageNet1K：

在 ImageNet-1K 圖像分類基準上，NaLaFormer 實現了對現有線性注意力模型的全面超越。輕量級 NaLaFormer-XT 僅 8M 參數量就拿下 79.1% Top-1 精度，同量級下較基線最高提升 7.5%；大規(guī)模 NaLaFormer-L 以 95M 參數量實現 85.7% Top-1 精度，刷新了高效視覺模型在該基準上的新紀錄。

? 密集預測：檢測、分割全面領跑同量級模型

在 COCO 目標檢測、ADE20K/CityScapes 語義分割等密集預測任務中，NaLaFormer 展現出極強的細粒度表征能力。COCO 數據集上，NaLaFormer-T 基于 Mask R-CNN 框架實現 47.6% APb、43.0% 分割 APm，全面超越同量級經典視覺 Transformer 骨干；ADE20K 語義分割任務中，較同量級基線最高提升 4.7% mIoU，CityScapes 城市場景分割也以 82.5% mIoU 領跑同規(guī)模模型。

? 超分辨率：70K+token 峰值顯存直降 92.3%

在 70K+token 的高分辨率超分任務中，NaLaFormer 的線性復雜度優(yōu)勢被徹底釋放。在保持重建精度與主流方案持平的同時，Urban100 4× 超分任務中實現 36.4% 的推理延遲降低，峰值顯存從 69G 降至 5.3G，降幅高達 92.3%，徹底解決了高分辨率視覺任務的顯存爆炸痛點。

? 長序列與跨模態(tài)能力：LRA 新 SOTA，語言建模超越 Mamba。

團隊進一步驗證了 NaLaFormer 的跨模態(tài)泛化性與長序列建模能力：在長序列建?；鶞?Long Range Arena（LRA）上，NaLaFormer 實現 61.2% 的平均精度，刷新了線性注意力模型的 SOTA 紀錄，同時保持 827.7 的高吞吐量與極低的顯存占用；從零訓練 340M 參數量的語言模型，在常識推理任務中，NaLaFormer 平均得分超越 Mamba、RetNet、GLA 等強基線，證明了其在語言模態(tài)的超強適配性。

? 擴散生成：DiT/SiT 性能穩(wěn)步提升

團隊進一步在擴散 Transformer（DiT）框架中驗證了 NaLaFormer 的生成建模能力。在 ImageNet-1K 256×256 生成任務中，將原始注意力模塊替換為 NaLaFormer 后：

• NaLaDiT 相比原版 DiT，FID 從 68.40 降至 61.64，IS 提升至 23.24，生成質量穩(wěn)步提升；
• NaLaSiT 在 SiT 變體中的表現，FID 低至 53.08，sFID 8.94，IS 27.63，實現基于 SiT 架構的高效能生成。

NaLaDiT 生成圖片可視化