显存暴降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：重塑线性注意力的双轮驱动

基于上述发现，研究团队提出了NaLaFormer（Norm-aware Linear Attention Former），通过两大创新来解决线性注意力的核心缺陷：查询范数感知特征映射和余弦方向相似性。

查询范数感知特征映射的设计目标是将查询范数重新注入到注意力计算中。

研究团队设计了一个范数依赖的锐化函数。

通过将查询范数作为幂函数的指数因子，模型能够动态调节注意力分布的锐度：当查询向量重要时，注意力会更加集中；当查询向量不太重要时，注意力会保持一定的分散性。

这一设计在数学上恢复了softmax注意力中查询范数与注意力锐度之间的关联，使得线性注意力也能够像softmax注意力一样动态聚焦。

余弦方向相似性的设计则是为了解决非负性保证的信息损失问题。

传统线性注意力采用ReLU或1+ELU等激活函数来强制非负性，但这种做法会直接丢弃所有负值分量，导致大量有效信息被清除。

当一个查询分量和一个键分量符号相反时，它们的交互信息会完全消失，这对于捕捉复杂的语义关系是一个巨大的损失。

研究团队借鉴三角恒等式的思想，提出了一种基于余弦相似度的方向映射方法。

通过这种方式，相似度计算转化为余弦函数的计算，其结果天然非负。

更重要的是，这种方法保留了原始方向分量的所有信息：当两个方向一致时，余弦值接近1；当两个方向相反时，余弦值接近0（而非传统方法中的直接丢弃）。

实验可视化结果清晰地表明，这种方法能够保留原始点积中的锐度信息，而ReLU和1+ELU都会导致信息的显著损失。

NaLaFormer将这两项创新整合在一个统一的架构中。

特征映射同时考虑范数和方向：对于查询向量，使用查询范数感知的幂函数来调节锐度；对于键向量，使用固定指数的幂函数进行缩放；方向分量则通过余弦映射保证非负性。

整个计算过程保持了线性复杂度，同时恢复了softmax注意力的关键特性。

多模态验证成果斐然

研究团队在多个视觉和语言任务上对NaLaFormer进行了全面评估，涵盖图像分类、目标检测、语义分割、超分辨率、扩散模型和语言建模等多个领域。

NaLaFormer在各项任务上都取得了显著的性能提升，充分验证了其设计的有效性。

在ImageNet-1K图像分类任务上，NaLaFormer展现了全面的领先优势。

下表展示了不同规模模型与现有方法的对比：

从表中可以看出，NaLaFormer在不同规模的模型上都取得了显著的性能提升。

这些结果表明，查询范数感知机制对线性注意力的性能提升具有普遍性和可扩展性。

在语义分割任务上，NaLaFormer同样表现出色。

在ADE20K数据集上，NaLaFormer-T达到了46.9%的mIoU，相比同类模型提升了4.7%；在Cityscapes数据集上，NaLaFormer-T达到了82.5%的mIoU，在相同规模的模型中名列前茅。

可视化结果进一步证实了NaLaFormer的优势：相比SegNeXt，NaLaFormer能够捕捉到更清晰的边界和更丰富的结构细节，这对于自动驾驶等实际应用具有重要意义。

然而，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.21137

END