当前位置: 首页 » 资讯 » 新科技 » 正文

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

IP属地 中国·北京 机器之心Pro 时间:2026-03-16 12:13:12



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

当 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 生成图片可视化

免责声明:本网信息来自于互联网,目的在于传递更多信息,并不代表本网赞同其观点。其内容真实性、完整性不作任何保证或承诺。如若本网有任何内容侵犯您的权益,请及时联系我们,本站将会在24小时内处理完毕。