本文一作孟维康是哈尔滨工业大学(深圳)与鹏城实验室联合培养的博士生,本科毕业于哈尔滨工业大学,主要研究方向是高效能基础模型。通讯作者张正教授,哈尔滨工业大学(深圳)的长聘教授及博士生导师,教育部青年长江学者,广东特支计划青年珠江学者,深圳市优青。长期从事高效能多模态机器学习的研究,专注于高效与可信多模态大模型。
当 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 生成图片可视化
热门跟贴