显存暴降92%！哈工大为线性注意力开辟了新道路

2017年Transformer模型横空出世在自然语言处理领域掀起了AI革命。此后视觉领域也迎来了Transformer的浪潮从图像分类到目标检测从语义分割到图像生成注意力机制几乎渗透到了计算机视觉的每一个角落。然而这个强大的机制始终背负着二次复杂度的致命弱点。当处理高分辨率图像或长序列文本时计算成本会呈指数级增长让许多实际应用望而却步。线性注意力机制的出现原本被寄予厚望期望在保持效率的同时继承softmax注意力的强大表达能力。然而现实却让人失望。线性注意力虽然解决了复杂度问题却始终无法匹敌softmax注意力的性能原因何在哈尔滨工业大学深圳、鹏城实验室和昆士兰大学的研究团队深入剖析了这一问题发现了一个被长期忽视的关键因素查询范数的丢失。基于这一发现他们提出了NaLaFormer在ImageNet-1K上实现了最高7.5%的准确率提升在超分辨率任务中峰值内存降低达92.3%为线性注意力机制的发展开辟了新道路。注意力机制的进化之路要理解NaLaFormer的创新我们需要先回顾注意力机制的演进历程。在传统的Transformer模型中核心的自注意力机制通过计算查询向量与键向量的点积相似度再经过softmax归一化最终加权求和得到输出。这个看似简单的过程实际上蕴含着深刻的数学原理点积相似度衡量了两个向量在语义空间中的接近程度而softmax归一化则将相似度转换为概率分布使得模型能够聚焦于最相关的信息。然而这个优雅的机制存在一个致命的计算瓶颈。对于长度为N的序列自注意力需要计算N×N的相似度矩阵复杂度为O(N²)。当处理高分辨率图像时一张1024×1024的图片会产生超过100万个图像块计算量将变成天文数字。更糟糕的是这种二次复杂度不仅在训练阶段造成困扰在推理阶段同样令人头疼严重限制了Transformer在资源受限场景中的应用。线性注意力的提出正是为了解决这一困境。其核心思想是利用核函数将softmax中的指数运算替换为线性可分离的形式从而重新排列计算顺序将复杂度从O(N²)降低到O(N)。这一思路在理论上堪称完美但在实践中却始终存在一个难以逾越的鸿沟线性注意力的表达能力明显弱于softmax注意力在各类任务上的性能都有不同程度的下降。研究者们尝试了各种方法来弥补这一差距包括使用不同的激活函数、引入幂函数来锐化注意力分布等但效果始终有限。问题的根源究竟在哪里研究团队没有像前人那样在工程层面修修补补而是选择从数学原理层面深入剖析。他们发现线性注意力性能不足的根源在于两个核心问题一是归一化操作取消了查询范数导致注意力分布失去了应有的锐度二是传统的非负性保证方法会破坏信息完整性。这两个问题交织在一起共同导致了线性注意力的表达能力下降。查询范数一个被遗忘的关键变量研究团队的核心发现可以用一句话概括在softmax注意力中查询向量的范数与注意力分布的锐度存在密切关联而线性注意力的归一化操作切断了这种关联。为了理解这一发现的重要性我们需要深入分析注意力机制的数学本质。研究团队提出了范数×方向(ND)分解的概念将任意向量分解为其范数和方向两个分量。在softmax注意力中查询范数被保留在指数函数内部可以自然地充当温度参数当查询范数较大时注意力分布会变得更加尖锐模型能够更专注地聚焦于关键信息当查询范数较小时注意力分布趋于平滑模型会对更多信息保持开放态度。这种动态调节机制是softmax注意力强大表达能力的重要来源。然而在线性注意力中情况截然不同。由于采用了归一化操作查询范数在分子和分母中被相互抵消导致注意力输出完全不受查询范数的影响。研究团队将这种现象称为查询范数无意识(query-norm-unaware)。为了验证这一发现团队在Swin Transformer上进行了实验分析结果令人振奋在softmax注意力中查询范数与注意力分布的熵呈现强烈的负相关关系。范数越大熵越小注意力越集中而键范数则没有这种相关性。这一发现揭示了softmax注意力的一个重要特性查询范数天然承担着调节注意力锐度的功能。更深入的理论分析表明对于查询向量的每一个方向分量注意力分布的熵会随着查询范数的增大而单调递减。这意味着查询范数越大模型就越能够从众多候选中筛选出最相关的信息。这一特性对于视觉任务尤其重要因为图像中往往包含大量背景噪声只有聚焦于关键区域才能获得准确的识别结果。线性注意力丢失了这一机制自然难以匹敌softmax注意力的性能。NaLaFormer重塑线性注意力的双轮驱动基于上述发现研究团队提出了NaLaFormerNorm-aware Linear Attention Former通过两大创新来解决线性注意力的核心缺陷查询范数感知特征映射和余弦方向相似性。查询范数感知特征映射的设计目标是将查询范数重新注入到注意力计算中。研究团队设计了一个范数依赖的锐化函数。通过将查询范数作为幂函数的指数因子模型能够动态调节注意力分布的锐度当查询向量重要时注意力会更加集中当查询向量不太重要时注意力会保持一定的分散性。这一设计在数学上恢复了softmax注意力中查询范数与注意力锐度之间的关联使得线性注意力也能够像softmax注意力一样动态聚焦。余弦方向相似性的设计则是为了解决非负性保证的信息损失问题。传统线性注意力采用ReLU或1ELU等激活函数来强制非负性但这种做法会直接丢弃所有负值分量导致大量有效信息被清除。当一个查询分量和一个键分量符号相反时它们的交互信息会完全消失这对于捕捉复杂的语义关系是一个巨大的损失。研究团队借鉴三角恒等式的思想提出了一种基于余弦相似度的方向映射方法。通过这种方式相似度计算转化为余弦函数的计算其结果天然非负。更重要的是这种方法保留了原始方向分量的所有信息当两个方向一致时余弦值接近1当两个方向相反时余弦值接近0而非传统方法中的直接丢弃。实验可视化结果清晰地表明这种方法能够保留原始点积中的锐度信息而ReLU和1ELU都会导致信息的显著损失。NaLaFormer将这两项创新整合在一个统一的架构中。特征映射同时考虑范数和方向对于查询向量使用查询范数感知的幂函数来调节锐度对于键向量使用固定指数的幂函数进行缩放方向分量则通过余弦映射保证非负性。整个计算过程保持了线性复杂度同时恢复了softmax注意力的关键特性。多模态验证成果斐然研究团队在多个视觉和语言任务上对NaLaFormer进行了全面评估涵盖图像分类、目标检测、语义分割、超分辨率、扩散模型和语言建模等多个领域。NaLaFormer在各项任务上都取得了显著的性能提升充分验证了其设计的有效性。在ImageNet-1K图像分类任务上NaLaFormer展现了全面的领先优势。下表展示了不同规模模型与现有方法的对比从表中可以看出NaLaFormer在不同规模的模型上都取得了显著的性能提升。这些结果表明查询范数感知机制对线性注意力的性能提升具有普遍性和可扩展性。在语义分割任务上NaLaFormer同样表现出色。在ADE20K数据集上NaLaFormer-T达到了46.9%的mIoU相比同类模型提升了4.7%在Cityscapes数据集上NaLaFormer-T达到了82.5%的mIoU在相同规模的模型中名列前茅。可视化结果进一步证实了NaLaFormer的优势相比SegNeXtNaLaFormer能够捕捉到更清晰的边界和更丰富的结构细节这对于自动驾驶等实际应用具有重要意义。然而NaLaFormer最令人印象深刻的成就体现在超分辨率任务的效率提升上。超分辨率是一个典型的token密集型任务一张高分辨率图像可能产生超过7万个图像块传统注意力机制的内存消耗令人望而却步。NaLaFormer在这一任务中展现了惊人的效率优势峰值内存从69GB降至5.3GB降低了92.3%使得原本需要高端服务器才能运行的任务现在可以在普通工作站上完成延迟从195ms降至124ms降低了36.4%使得实时应用成为可能。而在PSNR和SSIM指标上NaLaSR与ESRT相当没有因为效率提升而牺牲重建质量。在长序列建模能力上NaLaFormer在Long Range Arena基准测试中取得了61.2%的平均准确率超越了所有参与对比的线性注意力模型。这一结果表明查询范数感知机制对于捕捉长距离依赖关系具有重要价值。在语言建模任务上研究团队从零开始训练了一个3.4亿参数的语言模型在常识推理任务上超越了包括Mamba在内的强基线模型充分证明了NaLaFormer的跨模态适用性。研究团队通过严谨的理论分析和创新的工程设计成功地将这一发现转化为实际的技术进步为线性注意力机制的发展注入了新的活力。在人工智能快速发展的今天效率与性能的平衡始终是一个核心挑战NaLaFormer的研究成果为这一挑战提供了一个令人振奋的答案。参考资料https://arxiv.org/pdf/2506.21137END