即插即用外部注意力机制EA,涨点起飞起飞了!

文摘   2024-12-24 17:20   中国香港  

论文介绍

题目:https://arxiv.org/abs/2105.02358

论文地址:Beyond Self-attention: External Attention using Two Linear Layers for Visual Tasks

QQ深度学习交流群:994264161

扫描下方二维码,加入深度学习论文指南星球!

加入即可获得,模块缝合、制作、写作技巧,学会“结构”创新、“创新点”创新,从三区到顶会,小论文大论文,毕业一站式服务

创新点

  • 提出外部注意力机制(External Attention)

    • 提出了一个新的注意力机制,称为外部注意力,与传统的自注意力相比,它具有线性复杂度(O(n)),并通过利用两个外部的、可学习的共享记忆单元来捕获全数据集的潜在关系。

    • 不同于自注意力只关注单样本内的关系,外部注意力能够隐式地捕获整个数据集中的特征关联,从而增强模型的泛化能力。

  • 轻量化实现

    • 外部注意力只需要两个线性层和两个归一化层实现,与自注意力的计算和内存成本相比,大幅降低。

    • 实验表明,该方法可以在低计算成本下达到或超越原始自注意力及其变体的效果。

  • 多头外部注意力机制

    • 将多头机制引入外部注意力,构建了一种完全基于多层感知机(MLP)的新型架构,称为EAMLP(External Attention MLP)。该架构在图像分类任务中达到了与卷积神经网络(CNN)和Transformer相媲美的性能。

  • 适用多种视觉任务

    • 外部注意力被成功应用于多种视觉任务,包括图像分类、目标检测、语义分割、实例分割、图像生成以及点云分析。

    • 实验结果表明,外部注意力在这些任务中具有普适性,并在需要低计算成本的场景下表现出色。

  • 实验验证

    在多个数据集上进行了广泛的实验(如ImageNet、Pascal VOC、ADE20K、Cityscapes等),验证了外部注意力在多种任务中的有效性和高效性。

方法

整体架构

     这篇论文提出了一种基于外部注意力机制(External Attention)的轻量化模型架构,核心通过两个可学习的共享记忆单元捕获全数据集的全局上下文关系,具有线性复杂度(O(n))。该机制可灵活替代传统自注意力模块,并被集成到语义分割(如EANet)、图像分类(如EAMLP)、目标检测、点云分析等任务中,通过减少计算成本显著提升模型性能,展示了在多种视觉任务中的通用性和高效性。

模型的整体架构设计可以分为以下两类:

(1)语义分割任务的架构(EANet)

  • 如论文中图3(Page 4) 所示,语义分割架构基于FCN(Fully Convolutional Network),并在骨干网络(如ResNet)后添加外部注意力模块。

  • 具体流程

  1. 图像特征通过骨干网络提取基本特征。

  2. 外部注意力模块对特征进行全局上下文建模。

  3. 使用全连接层或MLP生成最终的分割结果。

(2)完全基于MLP的架构(EAMLP)

  • 在图像分类任务中,提出了基于外部注意力和多层感知机的全MLP架构,称为EAMLP。

  • 模块设计

    • 结合多头外部注意力机制,将输入特征划分为多个通道,分别计算注意力。

    • 将所有头的输出特征拼接,通过线性变换生成最终输出。


即插即用模块作用

EA 作为一个即插即用模块

  • 低计算资源限制的场景

    • 外部注意力具有线性计算复杂度(O(n)),相比自注意力(O(n²)大幅降低了计算成本,适合应用于需要高效计算的场景,例如移动端和嵌入式设备。

  • 需要全局上下文建模的视觉任务

    • 图像分类:捕获全局特征关系,提升分类精度。

    • 语义分割:增强模型对全局语义关系的理解,精确分割目标区域。

    • 目标检测与实例分割:通过全局上下文信息改善边界识别和细节处理。

    • 图像生成:提升生成器和判别器对长程依赖关系的建模能力。

    • 点云处理:在点云分类和分割任务中捕获三维点之间的全局关系。

  • 多任务场景

    • 外部注意力可以灵活嵌入各种任务框架(如FCN、Transformer、Faster R-CNN等),显著提升模型的跨任务表现。

消融实验结果


内容

在PASCAL VOC数据集上的消融实验结果。

  • 对比了自注意力(SA)与外部注意力(EA)的性能,并探讨了以下因素的影响:

    • 使用的归一化方法(Softmax与DoubleNorm)。

    • 外部记忆单元的数量(#S)。

    • 骨干网络的不同配置(ResNet-50与ResNet-101)。

  • 结果

    • 外部注意力在所有设置下都优于自注意力,且DoubleNorm归一化方法显著提升了性能。

    • 增加记忆单元数量可以进一步提升性能,但收益递减。

    • 骨干网络性能提升后,外部注意力的效果更为显著。



内容

在ImageNet分类任务中不同模型的消融实验。

  • 对比了外部注意力、多头外部注意力(MEA)、以及全MLP架构(EAMLP)的效果。

  • 结果

    • 多头外部注意力比单头外部注意力性能更强。

    • 在全MLP架构中,外部注意力的引入使性能接近或超过传统Transformer模型。



内容

外部注意力与自注意力及其多种变体的参数和计算量对比。

  • 结果

    • 外部注意力参数量仅为自注意力的一半,计算复杂度比自注意力低32倍。

    • 与其他变体相比,外部注意力在计算效率上依然更具优势。

即插即用模块

import numpy as np
import torch
from torch import nn
from torch.nn import init

# 论文地址:https://arxiv.org/abs/2105.02358
# 论文:Beyond Self-attention: External Attention using Two Linear Layers for Visual Tasks


class ExternalAttention(nn.Module):

    def __init__(self, d_model,S=64):
        super().__init__()
        self.mk=nn.Linear(d_model,S,bias=False)
        self.mv=nn.Linear(S,d_model,bias=False)
        self.softmax=nn.Softmax(dim=1)
        self.init_weights()


    def init_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                init.kaiming_normal_(m.weight, mode='fan_out')
                if m.bias is not None:
                    init.constant_(m.bias, 0)
            elif isinstance(m, nn.BatchNorm2d):
                init.constant_(m.weight, 1)
                init.constant_(m.bias, 0)
            elif isinstance(m, nn.Linear):
                init.normal_(m.weight, std=0.001)
                if m.bias is not None:
                    init.constant_(m.bias, 0)

    def forward(self, queries):
        attn=self.mk(queries) #bs,n,S
        attn=self.softmax(attn) #bs,n,S
        attn=attn/torch.sum(attn,dim=2,keepdim=True) #bs,n,S
        out=self.mv(attn) #bs,n,d_model

        return out


if __name__ == '__main__':
    input=torch.randn(50,49,512)
    block = ExternalAttention(d_model=512,S=8)
    output=block(input)    print(output.shape)

便捷下载方式

浏览打开网址:https://github.com/ai-dawang/PlugNPlay-Modules

更多分析可见原文


ai缝合大王
聚焦AI前沿,分享相关技术、论文,研究生自救指南
 最新文章