• 问答
  • 技术
  • 实践
  • 资源
无需额外参数的注意力模块!中大开源 SimAM:10 行代码快速实现解析解
技术讨论

作者丨happy
编辑丨极市平台

标题&作者团队

code: https://github.com/ZjjConan/SimAM

本文是中山大学在注意力机制方面的尝试,从神经科学理论出发,构建了一种能量函数挖掘神经元重要性,并对此推导出了解析解以加速计算。通过ImageNet分类、COCO检测与分割等任务验证了所提SimAM的灵活性与有效性。值得一提的是,所提SimAM是一种无参数注意力模块

Abstract

本文提出一种概念简单且非常有效的注意力模块。不同于现有的通道/空域注意力模块,该模块无需额外参数为特征图推导出3D注意力权值。具体来说,本文基于著名的神经科学理论提出优化能量函数以挖掘神经元的重要性。本文进一步针对该能量函数推导出一种快速解析解并表明:该解析解仅需不超过10行代码即可实现。该模块的另一个优势在于:大部分操作均基于所定义的能量函数选择,避免了过多的结构调整。最后,本文在不同的任务上对所提注意力模块的有效性、灵活性进行验证。

本文主要贡献包含以下几点:

  • 受启发于人脑注意力机制,本文提出一种3D注意力模块并设计了一种能量函数用于计算注意力权值;
  • 本文推导出了能量函数的解析解加速了注意力权值的计算并得到了一种轻量型注意力模块;
  • 将所提注意力嵌入到现有ConvNet中在不同任务上进行了灵活性与有效性的验证。

Method

在正式介绍本文所提注意力模块之前,我们先对现有代表性注意力模块(比如SE、CBAM、GC)进行简要总结;然后,我们再引出本文所提完全不同架构的注意力模块。

Overview of Existing Attention Modules

上图a与b列出了现有两种类型的注意力模块:

  • 通道注意力:1D注意力,它对不同通道区别对待,对所有位置同等对待;
  • 空域注意力:2D注意力,它对不同位置区别对待,对所有通道同等对待。

以下图为例,SE缺失了关于"grey_whale"的某些重要成分。我们认为3D注意力比1D和2D更佳,进而提出了上图c的3D注意力模块。

现有注意力模块的另一个重要影响因素:权值生成方法。现有注意力往往采用额外的子网络生成注意力权值,比如SE的GAP+FC+ReLU+FC+Sigmoid。更多注意力模块的操作、参数量可参考下表。总而言之,现有注意力的结构设计需要大量的工程性实验。我们认为:注意力机制的实现应当通过神经科学中的某些统一原则引导设计

Our Attention Module

已有研究BAM、CBAM分别将空域注意力与通道注意力进行并行或串行组合。然而,人脑的两种注意力往往是协同工作,因此,我们提出了统一权值的注意力模块。

为更好的实现注意力,我们需要评估每个神经元的重要性。在神经科学中,信息丰富的神经元通常表现出与周围神经元不同的放电模式。而且,激活神经元通常会抑制周围神经元,即空域抑制。换句话说,具有空域抑制效应的神经元应当赋予更高的重要性。最简单的寻找重要神经元的方法:度量神经元之间的线性可分性。因此,我们定义了如下能量函数:
$$
e_t(w_t, b_t, y, x_i) = (yt - \hat{t})^2 + \frac{1}{M-1}\sum{i=1}^{M-1}(y_o - \hat{x_i})^2
$$
其中,$\hat{t} = w_t t + b_t, \hat{x_i} = w_t x_i + b_t$。最小化上述公式等价于训练同一通道内神经元t与其他神经元之间的线性可分性。为简单起见,我们采用二值标签,并添加正则项,最终的能量函数定义如下:
$$
e_t(w_t, b_t, y, xi) = \frac{1}{M-1} \sum{i=1}^{M-1}(-1-(w_t x_i+b_t))^2 + (1-(w_t t + b_t))^2 + \lambda w_t^2
$$
理论上,每个通道有$M=H \times W$个能量函数。幸运的是,上述公式具有如下解析解:
$$
w_t = -\frac{2(t-u_t)}{(t-u_t)^2 + 2\sigma_t^2 + 2\lambda} \
b_t = -\frac{1}{2}(t+u_t) w_t
$$
其中,$ut = \frac{1}{M-1}\sum{i=1}^{M-1} x_i, \sigmat^2 = \frac{1}{M-1} \sum{i=1}^{M-1}(x_i - u_t)^2$。因此,最小能量可以通过如下公式得到:
$$
e^{}_t = \frac{4(\hat{\sigma^2} + \lambda)}{(t-\hat{u})^2 + 2\hat{\sigma}^2 + 2\lambda}
$$
上述公式意味着:能量越低,神经元t与周围神经元的区别越大,重要性越高。因此,神经元的重要性可以通过$1/e^{
}_t$得到。

到目前为止,我们推导了能量函数并挖掘了神经元的重要性。按照注意力机制的定义,我们需要对特征进行增强处理:
$$
\tilde{X} = sigmoid(\frac{1}{E}) \odot X
$$
下图给出了SimAM的pytorch风格实现code。

Experiments

image-20210721110733957

上表给出了Cifar数据集上的性能对比,从中可以看到:

  • 相比其他注意力机制,所提SimAM取得了最佳性能;
  • 在小网络方面,PreResNet56在CIFAR10数据集上取得了最佳性能(92.47/69.13),显著优于其他注意力;
  • 在大网络方面,所提SimAM同样优于SE与CBAM等注意力机制;
  • 上述实验结果均表明:所提无参注意力SimAM是一种通用注意力机制,并不局限于特定网络。

上表给出了ImageNet数据集上不同注意力机制的性能对比,从中可以看到:

  • 所有注意力模块均可以提升基线模型的性能;
  • 所提SimAM在ResNet18与ResNet101基线上取得了最佳性能提升;
  • 对于ResNet34、ResNet50、ResNeXt50、MobileNetV2,所提SimAM仍可取得与其他注意力相当性能;
  • 值得一提的是,所提SimAM并不会引入额外的参数
  • 在推理速度方面,所提SimAM与SE、ECA相当,优于CBAM、SRM。

上图对比了不同$\lambda$的性能对比,从中可以看到:

  • 所提注意力机制对于超参非常鲁棒;
  • $\lambda=0.0001$时取得了最佳的性能均衡。

上图采用Grad-CAM对所提SimAM进行了可视化,可以看到:经由SimAM提炼的特征可以更好的聚焦主体目标

上表对比了COCO检测、分割任务上的性能,可以看到:

  • SE与SimAM均可大幅提升基线模型的性能;
  • 对于检测任务,两种注意力模块取得了非常相似的性能;
  • 对于实例分割任务,所提SimAM取得了稍优性能。
  • 值得一提,SimAM不会引入额外参数量,而SE则会引入额外的参数量。比如,SE-ResNet50引入了2.5M参数量,SE-ResNet101引入了4.7M参数量。
  • 0
  • 0
  • 287
收藏
暂无评论