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

全新线性注意力范式!哈工深张正团队提出模长感知线性注意力!显存直降92.3%!
2026年03月15日 11:30 机器之心Pro

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

当 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 模长则无稳定影响图 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) 余弦方向机制实现无损非负性约束

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

其中 

为融合了 Norm 感知与余弦方向映射的 Query/Key 特征映射,G 为门控矩阵,实现了对注意力输出的自适应调制。

 与

实验结果

✅ 图像分类任务 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 生成图片可视化
NaLaDiT 生成图片可视化NaLaDiT 生成图片可视化

注意力张正
新浪科技公众号
新浪科技公众号

“掌”握科技鲜闻 (微信搜索techsina或扫描左侧二维码关注)

创事记

科学探索

科学大家

苹果汇

众测

专题

官方微博

新浪科技 新浪数码 新浪手机 科学探索 苹果汇 新浪众测

公众号

新浪科技

新浪科技为你带来最新鲜的科技资讯

苹果汇

苹果汇为你带来最新鲜的苹果产品新闻

新浪众测

新酷产品第一时间免费试玩

新浪探索

提供最新的科学家新闻,精彩的震撼图片