哈工深团队突破线性注意力瓶颈！Norm×Direction分解实现视觉任务精度全面超越，70K+token超分任务显存降低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

问题分析：线性注意力的两大核心缺陷
该研究系统性地揭示了现有线性注意力方案性能不足的两大根源：

1. 查询向量模长被抵消，注意力尖峰性丢失
   在 Softmax 注意力中，查询向量的模长与注意力分布的熵呈强负相关——模长越大，注意力分布越尖锐，模型越能聚焦于关键信息。然而，传统线性注意力中的归一化操作会抵消查询向量模长的影响，导致注意力分布过于平滑，表征能力退化。
   
   图1：Softmax 注意力中，查询向量模长与注意力熵呈强负相关。

2. 非负性约束导致信息损失
   为保证注意力得分可解释为概率分布，线性注意力核函数需满足非负性。现有方法通常使用 ReLU、1+ELU 等激活函数直接去除负值，但这导致了查询-键值对中负向语义关联信息的丢失，使得相似度表征变得稀疏且缺乏细粒度。
   
   图2：不同非负化策略的内积贡献对比。传统方法丢失了原始分布的尖峰与细粒度信息。

解决方案：Norm×Direction 分解
针对上述缺陷，研究团队提出了基于“模长-方向”分解的解决方案，将向量的模长（表征信息重要性）与方向（编码语义信息）解耦，分别针对性解决问题。

1. 查询模长感知的特征映射
   通过数学重构，团队定位了查询模长被抵消的环节，并设计了查询模长感知的特征映射，将查询模长显式编码进核函数。通过一个依赖于模长的锐化函数动态调控注意力分布的熵，成功复现了 Softmax 注意力中“查询模长越大，注意力越尖锐”的核心特性。
   
   

2. 余弦方向相似度机制
   为解决非负性约束下的信息损失问题，团队基于几何变换提出了余弦方向相似度机制。通过对方向分量进行特定的二维向量映射和缩放，在保证内积结果非负的同时，完整保留了原始向量间的细粒度方向交互信息，实现了无损的非负化约束。
   
   

3. 统一的 NaLaFormer 架构
   整合上述两大创新，团队构建了端到端的 NaLaFormer 架构。该架构将模长感知线性注意力模块与门控机制深度融合，可直接替换现有视觉 Transformer 中的注意力模块。
   
   图3：NaLaFormer 整体架构与核心机制。
   最终的注意力输出公式整合了模长感知与余弦方向映射。
   
   
   

实验结果

• 图像分类任务 (ImageNet-1K)
  NaLaFormer 在 ImageNet-1K 分类基准上全面超越现有线性注意力模型。其中，轻量级 NaLaFormer-XT 仅用 8M 参数量即达到 79.1% 的 Top-1 准确率；大规模 NaLaFormer-L 以 95M 参数量实现了 85.7% 的 Top-1 准确率。
  

• 密集预测任务
  在目标检测、实例分割等密集预测任务上，NaLaFormer 同样在同参数量级模型中取得领先性能。

在 COCO 目标检测、ADE20K/CityScapes 语义分割等密集预测任务中，NaLaFormer 展现出极强的细粒度表征能力。在 COCO 数据集上，基于 Mask R-CNN 框架的 NaLaFormer-T 实现了 47.6% 的边界框 AP 和 43.0% 的掩码 AP，全面超越了同量级的经典视觉 Transformer 骨干网络。在 ADE20K 语义分割任务中，其较同量级基线模型的最高提升达 4.7% mIoU；在 CityScapes 城市场景分割任务中，则以 82.5% mIoU 的成绩领跑同规模模型。

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

在处理高达 70K+token 的高分辨率超分任务时，NaLaFormer 的线性复杂度优势得到充分发挥。在保持重建精度与主流方案持平的同时，其在 Urban100 数据集上的 4× 超分任务中实现了 36.4% 的推理延迟降低，峰值显存占用从 69GB 大幅降至 5.3GB，降幅高达 92.3%，有效解决了高分辨率视觉任务中的显存瓶颈问题。

长序列与跨模态能力：刷新 LRA 基准，语言建模表现优异

研究团队进一步验证了 NaLaFormer 的跨模态泛化性与长序列建模能力。在长序列建模基准 Long Range Arena (LRA) 上，NaLaFormer 实现了 61.2% 的平均精度，刷新了线性注意力模型的最佳纪录，同时保持了 827.7 的高吞吐量与极低的显存占用。此外，一个从零开始训练的 3.4 亿参数语言模型在常识推理任务中，其平均得分超越了 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 生成图片可视化

关注“鲸栖”小程序，掌握最新AI资讯