电子科大提出注意力机制创新,改进 MobileViT变体在早期降采样阶段的注意力 QKV 操作研究 !

科技   2024-11-18 19:05   广东  


在本研究中,作者提出了一种改进的MobileViT变体,该变体在降采样阶段的早期阶段执行基于注意力的QKV操作。直接在高分辨率特征图上执行QKV操作由于其大尺寸和大量 Token 而计算上非常耗时。

为了解决这个问题,作者引入了一种使用卷积神经网络(CNN)生成重要性 Mask (Filter Mask)的滤波注意力机制,该 Mask 用于选择注意力计算中最有信息量的像素。

具体来说,卷积神经网络对特征图的像素进行评分,并按这些评分对像素进行排序,以选择前个像素(在不同网络层中,的值不同)。这种方法有效地减少了参与注意力计算的 Token 数量,从而降低了计算复杂性并提高了处理速度。

此外,作者发现重要性 Mask 具有可解释性,因为模型更关注对结果至关重要的图像区域。

实验结果显示,作者的模型在提高参数效率和计算速度的同时,也实现了高精度的提升。

与其他模型相比,作者的方法在降低计算资源消耗的同时,保持了高性能。

I Introduction

视觉 Transformer (ViT)[1]通过引入基于 Transformer 的架构,对计算机视觉领域进行了革新。ViT利用注意力机制在16x16图像块上进行操作,实现对图像的全局上下文学习。然而,尽管ViT在各种任务上都能实现最先进的结果,但其对大粒子的依赖使其在应用于高分辨率图像时计算成本高昂。QKV操作[2]的二次复杂性意味着,增加粒子的数量会极大地增加计算负担。

为解决此问题,MobileViT [3] 作为一种轻量级的 ViT 版本,在降采样图像后进行注意力计算。这减少了参与注意力的 Token 数量,从而提高了计算效率。然而,MobileViT 通过在降采样特征图(例如,56x56 或 28x28)上操作,牺牲了细粒度,这可能限制了其捕捉更细微细节的能力。此外,并非所有像素对最终预测做出相同贡献。许多像素是噪声或无关的,而其他像素对于决策至关重要。

在本文中,作者提出了一种名为FilterVit的移动ViT的新变体,该变体允许在不进行早期下采样的情况下实现更细粒度的注意力。作者的方法利用卷积神经网络(CNN)生成一个重要性 Mask ,该 Mask 确定特征图中的哪些像素对于注意力最为相关。通过按像素重要性分数对齐并选择前K个像素,作者仅对图像的最重要部分执行注意力,从而降低了QKV操作的计算复杂性。

不同于在GPU上实现稀疏矩阵操作的挑战性任务,作者的方法通过直接选择重要像素提供了一个实际解决方案。这确保了计算负载的降低,同时保持了准确性。

此外,作者的方法具有可解释性,重要性 Mask 突出了模型在注意力过程中关注的图像关键区域。例如,在羊的图片中,该 Mask 有效地突出了羊的轮廓,同时忽略了无关的背景特征。

总之,作者的贡献有三个方面:

  1. 作者提出了一种轻量级、更高效的注意力机制,可以在减少计算复杂度的同时实现更细粒度的注意力。
  2. 作者的方法天生具有可解释性,因为它专注于图像中最相关的区域,这一点可以通过重要性 Mask 的视觉化进行证明。
  3. 通过一次消融研究,作者引入了DropoutVIT,这是作者模型的一个变体,通过随机选择像素来模拟dropout,进一步验证了作者的方法的可扩展性和鲁棒性。

II Related Work

Vision Transformers

视觉 Transformer (ViT)[1]通过将注意力机制应用于图像块(通常为16x16块),在计算机视觉任务上取得了显著成功。

尽管ViT在捕捉图像中的全局依赖关系方面表现出色,但其QKV操作的二次复杂度[2]构成了计算挑战,尤其是在高分辨率图像方面。随着图像分辨率的增加, Token 的数量迅速增加,导致计算成本显著提高。

MobileViT

MobileViT [3] 针对这个问题,通过在应用基于 Transformer 的注意力之前对图像进行下采样来解决。

通过在较小分辨率(例如,56x56或28x28)上执行注意力,可以减少 Token 的数量,从而降低计算复杂性。然而,这种下采样牺牲了注意力的细粒度,当注意力应用于较低分辨率的特征图时,可能错过一些细微的细节。

Efficient Attention Mechanisms

几种基于注意力的模型已经被提出以优化计算复杂性。Longformer [4] 采用稀疏注意力机制,限制注意力在每一个 Token 周围的一个局部窗口,从而减少了计算负载,但在图像中可能难以捕捉全局依赖性。

Performer [5] 使用基于核的方法近似softmax注意力,实现了与序列长度成正比的线性复杂度,这是减少视觉任务中注意力成本的有前景的方法。Linformer [6] 进一步通过将 Key和Value 映射到较低维数来降低复杂度,在保持线性复杂度的同时节省了内存。Reformer [7] 引入局部敏感哈希(LSH)来近似注意力空间中的最近邻居,使计算更高效,尽管当需要高度详细的结构时,这种方法可能面临挑战。

Lightweight Convolutional Architectures

除了注意力机制外,轻量级卷积架构也被开发出来以平衡效率和性能。MobileNetV2 [8] 引入了反向残差模块以优化信息流,所需参数更少。MobileNetV3 [9] 对此进行了改进,通过神经架构搜索(NAS)设计出高效的结构,优化了延迟和准确性。 

EfficientNet [10] 通过平衡深度、宽度和分辨率来扩展网络,实现了跨配置的最优性能。GhostNet [11] 通过生成更少的特征图来减少特征图冗余,从而提高速度并减少参数数量。最后, LCNet [12] 专注于实现低复杂度的卷积操作,以实现更快的推理时间,同时保持合理的准确性,使其非常适合边缘计算任务。

Filter Attention Mechanism

作者的工作在这些进展的基础上,引入了Filter Attention机制。与那些降采样或线性化注意力计算的模型不同,作者的方法使用由卷积神经网络生成的滤波器 Mask ,选择性地将注意力应用于最重要的像素。根据像素的重要性对像素进行排序,并选择前K个像素,从而减少了QKV计算中涉及的 Token 数量。

这保持了全局上下文,同时显著降低了计算复杂度,而没有牺牲注意力粒度。作者的方法还具有可解释性,因为可视化滤波器 Mask 可以显示模型有效地关注图像的最相关部分。

III Method

Structure

在作者提出的网络架构中,作者将卷积层与 transformer 层相结合,以实现局部特征提取和全局注意力的平衡。作者方法的关键创新在于引入了一种 Filter Attention 机制,该机制可以选择性地关注特征图中的重要区域,从而显著降低基于注意力的计算复杂性。

网络通过卷积神经网络(CNN)处理输入图像,产生特征图。这一特征图被视为一组可 Reshape 的 Token ,通过Transformer Encoder [2]进行传递。与传统注意力机制不同,作者的Filter Attention根据卷积神经网络生成的过滤 Mask 选择性地识别重要 Token ,从而减少在注意力计算中使用的 Token 数量。

这种机制使模型能够忽略图像中不相关的区域,因为并非每个像素对最终预测的贡献相等。理论上,某些像素可能表示噪声或无用的细节,甚至在有用的像素中,它们对预测的贡献也不是均匀的[13]。

Filter Attention Block

作者的方法的核心是Filter Attention块,该块将卷积块与transformer编码器集成在一起。关键思想是将特征图中的像素视为transformer编码器的token,但不是处理所有token,而是应用一个滤波器 Mask 来选择最重要的几个。

首先,通过卷积神经网络(CNN)模块对图像进行处理,生成一个特征图。CNN同时也产生一个卷积 Mask ,为每个像素分配重要性分数。然后,特征图与 Mask 进行逐元素乘法操作,以过滤掉不太重要的 Token ,公式表示如下:

在这里,表示元素乘法,这确保了只有最相关的像素被保留下来进行进一步的基于注意力的计算。

接下来, 会被 Reshape 成一个适合 Transformer 结构的形状 ,其中 ,其中 是过滤后选定的 Token 数量。过滤后的 Token 通过一个 Transformer 编码器 [2] 传递,产生输出

过滤注意力机制(FilterAttention mechanism)详细描述如下。通过计算重要性图并应用Transformer Encoder,该机制能够选择性地关注特征图中最有信息量的像素。这种方法在降低计算复杂度的同时,保留了关键的空间信息。

Inverted Residual Block

在作者的网络中,作者采用倒置残差模块[8]来提高特征提取的同时保持模型效率。这个模块首先扩展输入通道,应用逐点卷积,然后将输出投影回原始尺寸,创建一个残差连接。这种结构确保了重要空间信息的保留,同时减少了模型中的参数数量,使其在小型和大型视觉任务上具有计算效率。

Global Self-Attention with Pooling

为了进一步降低注意力机制的计算复杂度,作者在将特征图输入自注意力层之前对其进行平均池化。池化降低了输入特征图的空间维度,从而降低了进入注意力模块的 Token 数量,从而减小了QKV操作的二次成本。尽管进行了这种减少,但全局上下文仍得到保持,确保模型仍能捕捉到图像中最相关区域之间的长期依赖关系。

Transformer Encoder

选自过滤后的特征图的 Token 被传递给Transformer Encoder [2]进行处理。transformer层应用自注意力机制来捕获 Token 之间的全局依赖关系。通过transformer之后, Token 被 Reshape 回其原始的空间维度,并经过后续的CNN层进行进一步的优化。

最后,输出特征图用于分类或其他下游任务,利用从图像中提取的局部和全局信息。

IV Experiment

作者使用五个不同的img-100子集进行了实验,每个子集都从ImageNet-1K数据集[14]中随机采样。本文中呈现的结果基于第一个img-100子集。为确保公平比较,所有模型的训练设置和超参数都采用了统一的设置。

本文研究中使用的五个img-100子集的具体类别选择已在GitHub  https://github.com/BobSun98/FilterVIT 中提供,以供参考和可重复性。

Dataset

img-100子集是从ImageNet-1K数据集[14]中的100个随机选择的类别构建的。在这些子集上进行训练和验证。每个模型在这些五个不同的子集上进行训练,以提高鲁棒性和减少过拟合,确保报告的结果不会因特定子集的选择而产生偏差。较小的img-100数据集允许更快地进行原型设计和算法验证,同时仍然在减少过拟合方面具有挑战性。如果轻量级模型在较小数据集上表现良好,那么它们具有更好的泛化能力的潜力。

Iv-A1 Data Augmentation

为提高所有模型的泛化能力,作者在训练过程中应用了统一的数据增强技术。这些技术包括:随机缩放(224)以随机缩放和调整图像大小至224x224像素,水平翻转(RandomHorizontalFlip)以随机水平翻转图像,引入广泛变化性的简单增强(TrivialAugmentWide [15])以及使用[0.5, 0.5, 0.5]均值和[0.5, 0.5, 0.5]标准差的归一化。在验证阶段,图像被缩放到256x256像素,然后进行中心裁剪至224x224像素,使用与之前相同的均值和标准差值进行归一化。

Training and Hyperparameters

所有模型都使用相同的超参数进行训练。将批量大小设置为64,学习率为0.0005,该学习率使用余弦退火计划[16]衰减至最小学习率1e-5。

使用AdamW优化器[17]训练模型120个周期,其中β1设置为0.9,β2设置为0.999,权重衰减为0.01。学习率调度器CosineAnnealingLR的最大周期为120个周期,最小学习率1e-5。

Hardware Configuration

训练在特斯拉P40 GPU上进行。在CUDA上的性能(以帧每秒(FPS)为单位)在RTX 4090 GPU上测量,而CPU性能则在苹果M1 Pro芯片上评估。这些硬件配置允许对每个模型在高性能和资源受限环境下的效率进行比较。

V Experiment

作者使用来自ImageNet-1K数据集[14]的五种不同的img-100子集进行了实验。本文中呈现的结果基于第一种img-100子集。为了确保公平比较,所有模型的训练设置和超参数都采用了统一的设置。

本文研究中使用的五种img-100子集的具体类别选择已在GitHub上提供,以便参考和可重复实现:https://github.com/BobSun98/FilterVIT。

Dataset

从ImageNet-1K数据集中的100个随机选择的类别构建了img-100子集[14]。在这些子集上进行了训练和验证。每个模型都在五个不同的子集中进行训练,以提高鲁棒性和降低过拟合,确保报告的结果不因特定子集的选择而产生偏差。较小的img-100数据集允许更快的原型设计和算法验证,同时仍然具有减少过拟合的挑战。如果轻量级模型在较小的数据集上表现良好,那么它们具有更好的泛化潜力的可能性更大。

V-B1 Data Augmentation

为了提高所有模型的泛化能力,作者在训练过程中应用了统一的数据增强技术。这些技术包括:随机缩放(224)以随机缩放和裁剪图像至224×224像素,随机翻转(水平方向)以随机翻转图像,引入广泛的增强变化[15](TrivialAugmentWide),以及使用[0.5,0.5,0.5]的平均值和[0.5,0.5,0.5]的标准差进行归一化。在验证阶段,图像被重新缩放到256×256像素,然后进行中心裁剪至224×224像素,使用与之前相同的平均值和标准差值进行归一化。

Models and Baselines

为了评估作者提出的FilterMobileViT的有效性,作者将它与几种 Baseline 模型进行比较,这些模型代表了高效神经网络领域的各种架构。这些包括MobileNetV2[8],以其倒置残差块和线性 Bottleneck 而闻名,以及MobileNetV3[9],它通过神经架构搜索和硬Swish激活函数改进了其前身。作者还包括EfficientNet-Lite0[10],该模型在网络深度、宽度和分辨率之间实现最佳效率平衡,以及GhostNet[11],它通过减少特征图中的冗余来提高计算性能。

除了卷积架构,作者还与轻量级的Transformer模型进行了比较,例如TinyViT[18],这是一种专为高性能和较少的参数设计的小型视觉Transformer,以及LeViT[19],这是一种结合卷积和Transformer层的混合模型,用于高效的推理。作者还包括了TinyNet[20],这是一个优化小型参数大小和计算成本的模型家族,以及LCNet[12],它在轻量级卷积网络设计中优先考虑速度。最后,MobileViTS 被纳入其中,因为它将Transformer层集成到一种适合移动设备架构中,类似于作者的方法。

这些模型被选中,以在不同的架构策略之间进行全面比较,突出作者FilterMobileViT在准确性和效率方面的优势。

Training Details

所有模型均使用一致的超参数进行训练:批量大小为64,使用AdamW优化器[17](,权重衰减为0.01)。初始学习率设置为0.0005,并在遵循余弦退火计划[16]后衰减到最小值,在120个epoch后完成。CosineAnnealingLR被用作学习率调度器,个epoch。训练在Tesla P40 GPU上进行,推理性能则分别测量在CUDA的RTX 4090 GPU和CPU上的Apple M1 Pro芯片上。这些配置允许对不同硬件环境下的模型效率进行全面比较,从高性能GPU到资源受限的CPU。

Experimental Results

表1显示,FilterMobileViT在准确性和计算效率之间取得了很好的平衡。仅需189万参数,其准确率达到了0.861,超越了除Tiny-ViT[18]之外的大多数模型。Tiny-ViT的准确率略高0.876,但参数数量更多(10.59M)。这表明尽管Tiny-ViT的性能略好,但FilterMobileViT在模型大小方面提供了更高效的解决方案,使其更适合在资源受限的环境中部署。

滤波移动ViT在每秒帧数(FPS)方面也表现出竞争力的性能。尽管它在CPU和CUDA平台上略逊于像LCNet_100这样的模型,后者的FPS更高,但LCNet在速度和准确性之间更注重速度,导致准确性为0.808,而滤波移动ViT在速度和准确性之间实现了更好的折中,保持了高准确性和合理的计算效率。

此外,与轻量级模型如MobileNetV2[8]和GhostNet[11]相比,FilterMobileViT不仅减少了参数数量,还提高了准确性。例如,MobileNetV2具有235万参数,准确率为0.833,而GhostNet具有402万参数,实现了0.842的准确率。这种改进凸显了作者在选择性关注图像最具有信息性的区域方面,Filter Attention机制的有效性,从而在不增加计算成本的情况下实现了更好的性能。

在推理速度方面,尽管EfficientNet_Lite0MobileNetV3_Small在CUDA上的FPS更高,但与FilterMobileViT相比,它们的准确率较低。这进一步突显了作者方法在实现平衡性能方面的优势,该性能适用于在准确性和效率均为核心需求的应用场景中。

总的来说,FilterMobileViT证明,将Filter注意力机制整合进来可以显著提高准确率和效率,在领域内超越了几个现有的模型。它的参数量较少且速度竞争力强,使其成为具有吸引力的选项,可以在计算资源有限的设备上部署。

例如,[21]提出了注意力展开方法,用于可视化注意力图,提供了模型在预测过程中关注的地方的洞察。同样,[22]开发了通过聚类其学习特征来分析ViTs内部表示的技术。此外,Vit-CX [13]利用ViTs的最后一层表示,并采用聚类来表示模型的关注区域。在作者的实验中,作者发现FilterVIT天生具有可解释性,无需额外的可视化技术。

为了更好地理解FilterMobileVit的行为,作者可视化了每个Filter Attention层生成的滤波器 Mask 。这是通过在预训练模型的测试图像集上进行推理,并提取不同层上的滤波器 Mask 来实现的。然后,这些 Mask 被重新缩放以匹配原始图像的大小,并叠加在输入图像上,以提供模型关注区域清晰的可视化表示[23]。

在第一个滤波器注意力层中,模型倾向于同时关注前景和背景,有时会突出物体的边缘。这表明第一层负责识别图像的广泛和一般特征。随着作者深入到第二个层,关注点更多地转向图像的主要目标或前景,说明模型已经开始优化其注意力并集中关注最显著的特征。在第三个和最后一个层中,注意力转向背景,可能是为了优化整体上下文,并确保无关区域得到适当的加权。

这段合作行为表明,过滤注意力机制不仅选择计算重要区域,还帮助模型在不同粒度上理解图像。通过逐步优化其注意力从边缘到主要物体再到背景,过滤移动Vit在其决策过程中表现出可解释性,如图5所示。

这些观察结果表明,模型的滤波 Mask 并非任意设置,而是反映了一种理解图像的结构性方法。每个层都贡献了图像的独特方面,使模型能够有效地在前景物体和更广泛的环境之间平衡注意力。这种分层注意力结构不仅提高了性能,还提供了模型处理视觉信息的方式的更可解释的理解,同时验证了设计正确性。

Ablation Study: DropOutVIT

在本研究中,作者介绍了一种名为DropOutVIT的FilterVIT模型的变体。DropOutVIT与FilterVIT之间的关键区别在于模型如何选择像素进行注意力计算。FilterVIT利用CNN生成的过滤器 Mask 来确定性选择最重要的像素,而DropOutVIT用一种随机取样的方法替换了这种机制,即在每次训练步骤中,随机选择一部分像素通过Transformer Encoder。这可以被认为类似于在注意力计算中应用dropout,引入了选择过程的一定程度的随机性[24]。

DropOutVIT的背后的理由是探索随机取样像素是否能导致相似或改进的性能,从而防止对特定特征的过度拟合。通过引入随机性,作者假设模型可能在大规模训练阶段之前(FilterVIT的滤波器 Mask 尚未完全优化)表现更好,特别是泛化能力。

V-1 DropOutVIT vs. FilterVIT

作者在几个训练运行中比较了DropOutVIT和FilterVIT的性能(见图6)。最初,两种模型表现出相似的性能,甚至在前几个epoch中,DropOutVIT的表现略好于FilterVIT。这可能归因于随机采样过程中选择的特征的多样性增加,这可能潜在地防止了过拟合。

然而,随着训练的进行,DropOutVIT似乎在第50个周期左右达到了性能的平稳期,而FilterVIT继续改进。

这种现象的一个可能解释是,FilterVIT中的CNN逐渐学会了生成更有意义的滤波器 Mask ,这使得模型能够专注于计算最重要的像素以进行注意力计算。这种精炼的选择过程可能使得FilterVIT能够实现更好的整体性能,因为模型越来越受益于有针对性的注意力。

参考文献

[0]. FilterVIT and DropoutVIT.

机器学习算法那些事
号主是大厂人工智能专家,专注于机器学习,深度学习以及计算机视觉等研究方向,每天会更新人工智能最前沿知识和分享自己的论文总结和学习笔记,让你系统化的学习每个知识点,每天进步一点点。
 最新文章