哈工深张正团队提出模长感知线性注意力！显存直降92.3%！

本文一作孟维康是哈尔滨工业大学（深圳）与鹏城实验室联合培养的博士生，本科毕业于哈尔滨工业大学，主要研究方向是高效能基础模型。通讯作者张正教授，哈尔滨工业大学（深圳）的长聘教授及博士生导师，教育部青年长江学者，广东特支计划青年珠江学者，深圳市优青。长期从事高效能多模态机器学习的研究，专注于高效与可信多模态大模型。

当 Transformer 席卷计算机视觉领域，高分辨率图像、超长序列任务带来的算力与显存瓶颈愈发凸显：标准 Softmax 注意力的二次复杂度，让 70K+token 的超分辨率任务直接显存爆炸，高分辨率图像分割、检测的推理延迟居高不下。

线性注意力虽通过核函数重构实现了线性复杂度，完美解决了算力开销问题，却始终无法摆脱性能退化的问题，与原生 Softmax 注意力的精度差距始终难以弥合。

近日，哈工深张正团队、联合鹏城实验室、昆士兰大学等团队，发布重磅论文《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 (N²) 降至 O (N)。但在 NaLaFormer 的研究中，首次系统性揭示了现有方案始终无法逼近 Softmax 性能的两大根源：

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

团队通过数学推导与实验验证发现：在 Softmax 注意力中，Query 向量的模长（Norm）与注意力分布的熵呈强负相关 ——Query 模长越大，注意力分布越尖锐（熵越低），模型越能精准聚焦于语义关键的 token，这是 Softmax 注意力具备强表征能力的核心原因。但在传统线性注意力中，归一化操作会直接抵消 Query Norm 的影响，让线性注意力彻底失去了对注意力分布尖峰性的动态调控能力，最终输出过于平滑的注意力分布，无法聚焦关键信息，表征能力大幅退化。

图 1 Softmax 注意力中熵与模长的相关性：Query 模长（x 轴）与注意力熵（y 轴）呈强负相关，Key 模长则无稳定影响

2、非负性约束导致不可逆的信息损失

线性注意力的核函数必须满足非负性，才能将注意力得分解释为归一化的概率分布。现有方案普遍通过 ReLU、1+ELU 等激活函数直接抹除向量中的负值，却直接导致了 q-k 内积中有效交互信息的丢失 —— 原本具有区分度的负向语义关联被直接清零，最终让相似度表征变得稀疏且缺乏细粒度信息。

图 2 不同非负策略的内积贡献对比：ReLU、1+ELU 均丢失了原始内积的尖峰性与细粒度信息，而本文提出的余弦方向方法完美保留了原始分布特征

核心创新：

Norm×Direction 分解，补全线性注意力机制短板

针对上述两大核心痛点，NaLaFormer 团队提出了基于 Norm×Direction（ND）向量分解的解决方案，将向量的模长（Norm，表征信息重要性）与方向（Direction，编码语义信息）解耦，分别针对性解决两大缺陷，最终实现了对 Softmax 注意力核心特性的完美复刻，同时完整保留线性复杂度。

1、Query-Norm-Aware Feature Map：重建丢失的模长 - 信息熵关联

团队首先通过 ND 分解，对线性注意力的计算过程进行了数学重构，精准定位了 Query Norm 被抵消的核心环节，清晰揭示了传统线性注意力的输出仅受 Key Norm 影响，Query Norm 在归一化中被完全抵消，导致模型对 Query Norm 的无感。

对此，团队设计了 Query-Norm-Aware 特征映射，将 Query Norm 显式编码进核函数中，通过 Norm 依赖的锐化函数动态调控注意力熵：

其中锐化函数 f (x)=λ∗(τ+tanh (x))，可根据 Query Norm 动态调整幂次，完美复刻了 Softmax 注意力中 “Query 模长越大，注意力分布越尖锐” 的核心特性，重建了 Query Norm 与注意力熵的负相关关系。

2、余弦方向相似度：无损实现非负性，零信息丢失

针对传统非负约束的信息丢失问题，团队基于托勒密理论，提出了全新的余弦方向相似度机制，对方向分量进行几何变换，在保证非负性的同时，完整保留内积的细粒度信息。

团队为每个标量方向分量设计了二维向量映射：

经过变换后，Query 与 Key 的内积可转化为：

通过 tanh 映射将方向分量缩放至 [−π/4, π/4]，即可保证余弦项始终非负，同时完整保留了原始内积的方向交互信息，彻底避免了传统激活函数带来的信息损失。

3. NaLaFormer：统一的 Norm 感知线性注意力架构

基于上述两大核心创新，团队构建了端到端的 NaLaFormer 架构，将模长感知线性注意力模块与门控架构深度融合，实现了对视觉 Transformer 的直接替换。

图 3 NaLaFormer 整体架构与核心机制：(a) NaLaFormer 基础模块；(b) Norm 感知方法恢复了 Query Norm - 熵负相关；(c) 余弦方向机制实现无损非负性约束

最终的注意力输出公式为：

实验结果

✅ 图像分类任务 ImageNet1K：

在 ImageNet-1K 图像分类基准上，NaLaFormer 实现了对现有线性注意力模型的全面超越。轻量级 NaLaFormer-XT 仅 8M 参数量就拿下 79.1% Top-1 精度，同量级下较基线最高提升 7.5%；大规模 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 领跑同规模模型。

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

在 70K+token 的高分辨率超分任务中，NaLaFormer 的线性复杂度优势被彻底释放。在保持重建精度与主流方案持平的同时，Urban100 4× 超分任务中实现 36.4% 的推理延迟降低，峰值显存从 69G 降至 5.3G，降幅高达 92.3%，彻底解决了高分辨率视觉任务的显存爆炸痛点。

✅ 长序列与跨模态能力：LRA 新 SOTA，语言建模超越 Mamba。

团队进一步验证了 NaLaFormer 的跨模态泛化性与长序列建模能力：在长序列建模基准 Long Range Arena（LRA）上，NaLaFormer 实现 61.2% 的平均精度，刷新了线性注意力模型的 SOTA 纪录，同时保持 827.7 的高吞吐量与极低的显存占用；从零训练 340M 参数量的语言模型，在常识推理任务中，NaLaFormer 平均得分超越 Mamba、RetNet、GLA 等强基线，证明了其在语言模态的超强适配性。

✅ 扩散生成：DiT/SiT 性能稳步提升

团队进一步在扩散 Transformer（DiT）框架中验证了 NaLaFormer 的生成建模能力。在 ImageNet-1K 256×256 生成任务中，将原始注意力模块替换为 NaLaFormer 后：

• NaLaDiT 相比原版 DiT，FID 从 68.40 降至 61.64，IS 提升至 23.24，生成质量稳步提升；
• NaLaSiT 在 SiT 变体中的表现，FID 低至 53.08，sFID 8.94，IS 27.63，实现基于 SiT 架构的高效能生成。

NaLaDiT 生成图片可视化