本文约3300字，建议阅读10分钟
本文一作孟维康是哈尔滨工业大学（深圳）与鹏城实验室联合培养的博士生，本科毕业于哈尔滨工业大学，主要研究方向是大规模基础模型的高效训练和推理算法研究。
通过引入极性感知计算和可学习幂函数，PolaFormer在保持线性复杂度的同时，显著提升了模型的性能和效率，并在图像分类、目标检测和语义分割等任务上取得了优异的实验结果。
本文一作孟维康是哈尔滨工业大学（深圳）与鹏城实验室联合培养的博士生，本科毕业于哈尔滨工业大学，主要研究方向是大规模基础模型的高效训练和推理算法研究。
通讯作者张正教授，哈尔滨工业大学（深圳）的长聘教授及博士生导师，教育部青年长江学者，广东特支计划青年珠江学者，深圳市优青。长期从事高效能多模态机器学习的研究，专注于高效与可信多模态大模型。
课题组：Big Media Intelligence (BMI) 欢迎校内外优秀学者的加入以及来访交流。
课题组主页：
https://cszhengzhang.cn/BMI/
-
论文标题：PolaFormer: Polarity-aware Linear Attention for Vision Transformers
-
论文链接：
https://arxiv.org/pdf/2501.15061
-
GitHub 链接：
https://github.com/ZacharyMeng/PolaFormer
-
Huggingface 权重链接：
https://huggingface.co/ZachMeng/PolaFormer/tree/main
尽管 Vision Transformer 及其变种在视觉任务上取得了亮眼的性能，但仍面临着自注意力机制时空间平方复杂度的挑战。为了解决这一问题，线性自注意力通过设计新的核函数替换标准自注意力机制中的 softmax 函数，使模型复杂度降低为线性。这篇论文中，研究者提出了一个新的「极性感知线性注意力」模块，使模型达到了更高的任务性能与计算效率。
具体来说，本工作从线性自注意力方法需要满足注意力权重矩阵的两个特性（即正值性和低信息熵）入手。首先，指出了现有的做法为了满足正值性，牺牲了 Q 矩阵和 K 矩阵元素中负值的缺陷，提出了极性感知的计算方式可以保证 Q 矩阵和 K 矩阵中所有元素可以平等地进行相似度的计算，使计算结果更准确，模型表示能力更强。其次，本文提出只要采用一族具有特殊性质的映射函数，就可以有效降低注意力权重分布的信息熵，并给出了数学上的证明。
大量的实验表明，本文提出的线性注意力模块可以直接替换现有 Vision Transformer 框架中的自注意力模块，并在视觉基础任务和 LRA 任务上一致地提升了性能。
Transformer 模型已经在广泛的视觉任务中展现出亮眼的性能。其核心模块 —— 通过 softmax 归一化的点积自注意力机制，让 Transformer 模型可以有效地捕捉长距离依赖关系。然而，这带来了模型 O (N^2) 复杂度，在处理长序列视频或高分辨率图像时，会导致相当大的计算开销和显存占用。这限制了它们在资源受限环境中的效率，使得在这些场景下的实际部署变得困难。
线性注意力，作为一种更具可行性的解决方案使用核化特征映射替换 q,k 点积中的 Softmax 操作，有效地将时间和空间复杂度从 O (N²d) 降低到 O (Nd²)。尽管线性注意力在计算效率上有所提升，但在表达能力方面仍不及基于 Softmax 的注意力，我们的分析确定了造成这种不足的两个主要原因，它们都源于 Softmax 近似过程中的信息丢失：
-
负值丢失。
依赖非负特征映射（如 ReLU）的线性注意力模型无法保持与原始 q,k 点积的一致性。这些特征映射仅保留了正 - 正交互作用，而关键的正 - 负和负 - 负交互作用则完全丢失。这种选择性表示限制了模型捕获全面关系范围的能力，导致注意力图的表达能力减弱和判别力降低。
-
注意力分布高信息熵。
没有 softmax 的指数缩放，线性注意力会导致权重分布更加均匀且熵更低。这种均匀性削弱了模型区分强弱 q,k 对的能力，损害了其对重要特征的关注，并在需要精细细节的任务中降低了性能。
在这项工作中，作者提出了一种极性感知线性注意力（PolaFormer）机制，旨在通过纳入被忽略的负交互作用来解决先前线性注意力模型的局限性。与此同时，为了解决线性注意力中常见的注意力权重分布信息熵过高的问题，他们提供了数学理论基础，表明如果一个逐元素计算的函数具有正的一阶和二阶导数，则可以重新缩放 q,k 响应以降低熵。这些增强功能共同提供了一个更稳健的解决方案，以缩小线性化和基于 Softmax 的注意力之间的差距。
可见，自注意力的复杂度是 O (N2d)。这种复杂度使得自注意力机制在处理长序列时效率低下，导致计算成本急剧上升。目前，降低自注意力的复杂度的主要方法包括但不限于稀疏注意力、线性化注意力以及基于核的注意力等。
通过利用矩阵乘法的结合律，每个头的复杂度可以降低到 O (Nd’2)，其中 d’是特征映射后的维度，与序列长度成线性关系。
极性感知注意力背后的核心思想是为了解决现有线性注意力机制的局限性，这些机制通常会丢弃来自负成分的有价值信息。
PolaFormer 在处理负成分时，极性感知注意力将 query 和 key 向量分解为它们的正部和负部。这种分解允许机制分别考虑正相似度和负相似度对注意力权重的影响。具体来说，对于查询向量 q 和键向量 k，可以将它们分解为：
前两项捕捉了同号成分之间的相似性，而后两项则代表了异号成分之间的相互作用。之前的线性注意力方法，如基于 ReLU 的特征映射，通过将负成分映射到零来消除它们，这在近似 q,k 点积时会导致显著的信息丢失。
为了解决这个问题，极性感知注意力机制根据 q,k 的极性将它们分开，并独立计算它们之间的相互作用。注意力权重的计算方式如下：
PolaFormer 根据极性明确地将 q,k 对分开，处理在内积计算过程中维度的同号和异号交互作用。这些交互作用在两个流中处理，从而能够更准确地重建原始的 softmax 注意力权重。为了避免不必要的复杂性，作者沿着通道维度拆分 v 向量，在不引入额外可学习参数的情况下处理这两种类型的交互作用。然后，将输出进行拼接，并通过一个可学习的符号感知矩阵进行缩放，以确保准确重建 q,k 关系。
作者统计分析了两个 G 矩阵的特性，存在一个明显的负相关和价值差异。这证明了本文提出的可学习混合策略补偿了松弛减法操作所带来的影响。
为了解决线性注意力中常见的注意力权重分布信息熵过高的问题，作者提供了数学理论基础，表明如果一个逐元素计算的函数具有正的一阶和二阶导数，则可以重新缩放 q,k 响应以降低熵。
这一理论有助于阐明为什么先前的特征映射会提高信息熵，从而导致注意力分布过于平滑。为了简化，作者采用通道级可学习的幂函数进行重新缩放，这保留了 Softmax 中固有的指数函数的尖锐性。这使得模型能够捕获尖锐的注意力峰值，提高了其区分强弱响应的能力。与此同时，为了区分开不同通道之间的主次关系，作者设计了可学习的幂次来捕捉每个维度的不同重要性
