补全Query Norm缺失！哈工深团队重构线性注意力，显存直降92.3%

当 Transformer 席卷计算机视觉领域高分辨率图像、超长序列任务带来的算力与显存瓶颈愈发凸显标准 Softmax 注意力的二次复杂度让 70Ktoken 的超分辨率任务直接显存爆炸高分辨率图像分割、检测的推理延迟居高不下。线性注意力虽通过核函数重构实现了线性复杂度完美解决了算力开销问题却始终无法摆脱性能退化的问题与原生 Softmax 注意力的精度差距始终难以弥合。近日哈工深张正团队、联合鹏城实验室、昆士兰大学等团队发布重磅论文 Norm×Direction提出 NaLaFormerNorm-aware Linear Attention Transformer框架。首次通过模长-方向分解精准定位并解决了线性注意力的两大核心缺陷在保持线性复杂度的同时实现了视觉任务的精度全面超越更在 70Ktoken 超分任务中实现 92.3% 的峰值显存降低为线性注意力在视觉领域的落地开辟了全新范式。论文标题Norm×Direction: Restoring the Missing Query Norm in Vision Linear Attention论文链接https://arxiv.org/pdf/2506.21137作者团队哈尔滨工业大学深圳SMULL Group、鹏城实验室、昆士兰大学 UQMM Lab核心作者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、1ELU 等激活函数直接抹除向量中的负值却直接导致了 q-k 内积中有效交互信息的丢失——原本具有区分度的负向语义关联被直接清零最终让相似度表征变得稀疏且缺乏细粒度信息。〓 图2. 不同非负策略的内积贡献对比ReLU、1ELU 均丢失了原始内积的尖峰性与细粒度信息而本文提出的余弦方向方法完美保留了原始分布特征核心创新Norm×Direction 分解双管齐下补全线性注意力短板针对上述两大核心痛点SMULL 团队提出了基于 Norm×DirectionND向量分解的解决方案将向量的模长Norm表征信息重要性与方向Direction编码语义信息解耦。分别针对性解决两大缺陷最终实现了对 Softmax 注意力核心特性的完美复刻同时完整保留线性复杂度。1. Query-Norm-Aware Feature Map重建丢失的模长-信息熵关联团队首先通过 ND 分解对线性注意力的计算过程进行了数学重构。精准定位了 Query Norm 被抵消的核心环节清晰揭示了传统线性注意力的输出仅受 Key Norm 影响Query Norm 在归一化中被完全抵消导致模型对 Query Norm的无感。对此团队设计了Query-Norm-Aware 特征映射将 Query Norm 显式编码进核函数中通过 Norm 依赖的锐化函数动态调控注意力熵其中锐化函数可根据 Query Norm 动态调整幂次完美复刻了 Softmax 注意力中 “Query 模长越大注意力分布越尖锐” 的核心特性重建了 Query Norm 与注意力熵的负相关关系。2. 余弦方向相似度无损实现非负性零信息丢失针对传统非负约束的信息丢失问题团队基于三角同构理论提出了全新的余弦方向相似度机制对方向分量进行几何变换在保证非负性的同时完整保留内积的细粒度信息。团队为每个标量方向分量设计了二维向量映射经过变换后Query 与 Key 的内积可转化为通过 tanh 映射将方向分量缩放至即可保证余弦项始终非负同时完整保留了原始内积的方向交互信息彻底避免了传统激活函数带来的信息损失。3. NaLaFormer统一的 Norm 感知线性注意力架构基于上述两大核心创新团队构建了端到端的 NaLaFormer 架构将 Norm 感知线性注意力模块与门控架构深度融合实现了对视觉 Transformer 的无缝替换。〓 图3. NaLaFormer 整体架构与核心机制aNaLaFormer 基础模块bNorm 感知方法恢复了 Query Norm - 熵负相关c余弦方向机制实现无损非负性约束最终的注意力输出公式为其中与为融合了 Norm 感知与余弦方向映射的 Query/Key 特征映射G 为门控矩阵实现了对注意力输出的自适应调制。✅ 图像分类刷新线性注意力 ImageNet 纪录在 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% 检测 AP、43.0% 分割 AP全面超越同量级经典视觉 Transformer 骨干。ADE20K 语义分割任务中较同量级基线最高提升 4.7% mIoUCityScapes 城市场景分割也以 82.5% mIoU 领跑同规模模型。✅ 超分辨率70Ktoken 峰值显存直降 92.3%在 70Ktoken 的高分辨率超分任务中NaLaFormer 的线性复杂度优势被彻底释放。在保持重建精度与主流方案持平的同时Urban100 4× 超分任务中实现 36.4% 的推理延迟降低峰值显存从 69G 降至 5.3G降幅高达 92.3%彻底解决了高分辨率视觉任务的显存爆炸痛点。✅ 长序列与跨模态能力LRA 新 SOTA语言建模超越 Mamba。团队进一步验证了 NaLaFormer 的跨模态泛化性与长序列建模能力在长序列建模基准 Long Range ArenaLRA上NaLaFormer 实现 61.2% 的平均精度刷新了线性注意力模型的 SOTA 纪录同时保持 827.7 的高吞吐量与极低的显存占用。从零训练 340M 参数量的语言模型在常识推理任务中NaLaFormer 平均得分超越 Mamba、RetNet、GLA 等强基线证明了其在语言模态的超强适配性。✅ 扩散生成DiT/SiT 性能稳步提升团队进一步在扩散 TransformerDiT框架中验证了 NaLaFormer 的生成建模能力。在 ImageNet-1K 256×256 生成任务中将原始注意力模块替换为 NaLaFormer 后NaLaDiT 相比原版 DiTFID 从 68.40 降至 61.64IS 提升至 23.24生成质量稳步提升NaLaSiT 在 SiT 变体中的表现FID 低至 53.08sFID 8.94IS 27.63实现基于 SiT 架构的高效能生成。更多阅读#投 稿 通 道#让你的文字被更多人看到如何才能让更多的优质内容以更短路径到达读者群体缩短读者寻找优质内容的成本呢答案就是你不认识的人。总有一些你不认识的人知道你想知道的东西。PaperWeekly 或许可以成为一座桥梁促使不同背景、不同方向的学者和学术灵感相互碰撞迸发出更多的可能性。PaperWeekly 鼓励高校实验室或个人在我们的平台上分享各类优质内容可以是最新论文解读也可以是学术热点剖析、科研心得或竞赛经验讲解等。我们的目的只有一个让知识真正流动起来。稿件基本要求• 文章确系个人原创作品未曾在公开渠道发表如为其他平台已发表或待发表的文章请明确标注• 稿件建议以markdown格式撰写文中配图以附件形式发送要求图片清晰无版权问题• PaperWeekly 尊重原作者署名权并将为每篇被采纳的原创首发稿件提供业内具有竞争力稿酬具体依据文章阅读量和文章质量阶梯制结算投稿通道• 投稿邮箱hrpaperweekly.site• 来稿请备注即时联系方式微信以便我们在稿件选用的第一时间联系作者• 您也可以直接添加小编微信pwbot02快速投稿备注姓名-投稿△长按添加PaperWeekly小编现在在「知乎」也能找到我们了进入知乎首页搜索「PaperWeekly」点击「关注」订阅我们的专栏吧·