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

當(dāng) Transformer 席卷計(jì)算機(jī)視覺(jué)領(lǐng)域，高分辨率圖像、超長(zhǎng)序列任務(wù)帶來(lái)的算力與顯存瓶頸愈發(fā)凸顯：標(biāo)準(zhǔn) Softmax 注意力的二次復(fù)雜度，讓 70K+token 的超分辨率任務(wù)直接顯存爆炸，高分辨率圖像分割、檢測(cè)的推理延遲居高不下。

線性注意力雖通過(guò)核函數(shù)重構(gòu)實(shí)現(xiàn)了線性復(fù)雜度，完美解決了算力開(kāi)銷問(wèn)題，卻始終無(wú)法擺脫性能退化的問(wèn)題，與原生 Softmax 注意力的精度差距始終難以彌合。

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

• 論文標(biāo)題：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

痛點(diǎn)直擊：

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

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

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

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

圖 1 Softmax 注意力中熵與模長(zhǎng)的相關(guān)性：Query 模長(zhǎng)（x 軸）與注意力熵（y 軸）呈強(qiáng)負(fù)相關(guān)，Key 模長(zhǎng)則無(wú)穩(wěn)定影響

2、非負(fù)性約束導(dǎo)致不可逆的信息損失

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

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

核心創(chuàng)新：

Norm×Direction 分解，補(bǔ)全線性注意力機(jī)制短板

針對(duì)上述兩大核心痛點(diǎn)，NaLaFormer 團(tuán)隊(duì)提出了基于 Norm×Direction（ND）向量分解的解決方案，將向量的模長(zhǎng)（Norm，表征信息重要性）與方向（Direction，編碼語(yǔ)義信息）解耦，分別針對(duì)性解決兩大缺陷，最終實(shí)現(xiàn)了對(duì) Softmax 注意力核心特性的完美復(fù)刻，同時(shí)完整保留線性復(fù)雜度。

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

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

對(duì)此，團(tuán)隊(duì)設(shè)計(jì)了 Query-Norm-Aware 特征映射，將 Query Norm 顯式編碼進(jìn)核函數(shù)中，通過(guò) Norm 依賴的銳化函數(shù)動(dòng)態(tài)調(diào)控注意力熵：

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

2、余弦方向相似度：無(wú)損實(shí)現(xiàn)非負(fù)性，零信息丟失

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

團(tuán)隊(duì)為每個(gè)標(biāo)量方向分量設(shè)計(jì)了二維向量映射：

經(jīng)過(guò)變換后，Query 與 Key 的內(nèi)積可轉(zhuǎn)化為：

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

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

基于上述兩大核心創(chuàng)新，團(tuán)隊(duì)構(gòu)建了端到端的 NaLaFormer 架構(gòu)，將模長(zhǎng)感知線性注意力模塊與門控架構(gòu)深度融合，實(shí)現(xiàn)了對(duì)視覺(jué) Transformer 的直接替換。

圖 3 NaLaFormer 整體架構(gòu)與核心機(jī)制：(a) NaLaFormer 基礎(chǔ)模塊；(b) Norm 感知方法恢復(fù)了 Query Norm - 熵負(fù)相關(guān)；(c) 余弦方向機(jī)制實(shí)現(xiàn)無(wú)損非負(fù)性約束

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

實(shí)驗(yàn)結(jié)果

? 圖像分類任務(wù) ImageNet1K：

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

? 密集預(yù)測(cè)：檢測(cè)、分割全面領(lǐng)跑同量級(jí)模型

在 COCO 目標(biāo)檢測(cè)、ADE20K/CityScapes 語(yǔ)義分割等密集預(yù)測(cè)任務(wù)中，NaLaFormer 展現(xiàn)出極強(qiáng)的細(xì)粒度表征能力。COCO 數(shù)據(jù)集上，NaLaFormer-T 基于 Mask R-CNN 框架實(shí)現(xiàn) 47.6% APb、43.0% 分割 APm，全面超越同量級(jí)經(jīng)典視覺(jué) Transformer 骨干；ADE20K 語(yǔ)義分割任務(wù)中，較同量級(jí)基線最高提升 4.7% mIoU，CityScapes 城市場(chǎng)景分割也以 82.5% mIoU 領(lǐng)跑同規(guī)模模型。

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

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

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

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

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

團(tuán)隊(duì)進(jìn)一步在擴(kuò)散 Transformer（DiT）框架中驗(yàn)證了 NaLaFormer 的生成建模能力。在 ImageNet-1K 256×256 生成任務(wù)中，將原始注意力模塊替換為 NaLaFormer 后：

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

NaLaDiT 生成圖片可視化