新加坡科研机构 DIMAP | 独立模块感知剪枝方法压缩分层 Transformer ，在图像分类基准上性能SOTA ！

集智书童公众号

发布于 2024-05-10 14:05:18

1240

发布于 2024-05-10 14:05:18

分层视觉 Transformer （ViTs）相对于传统ViTs有两个优势。首先，分层ViTs通过局部自注意力实现与图像大小成线性的计算复杂性。其次，分层ViTs通过在更深层次上合并图像块创建分层特征图，以进行密集预测。然而，现有的剪枝方法忽略了分层ViTs的独特性质，并使用幅度值作为权重重要性。这种方法导致了两个主要缺点。首先，“局部”注意力权重在“全局” Level 上进行比较，可能导致一些“局部”重要的权重因其在“全局”范围内相对较小的幅度而被剪枝。幅度剪枝的第二个问题是它没有考虑到网络中不同的权重分布，这对于在各个层次结构中提取从粗到细的特征至关重要。为了解决上述问题，作者开发了一种数据独立模块感知剪枝方法（DIMAP）来压缩分层ViTs。为了确保不同层次结构中的“局部”注意力权重在贡献方面公平比较，作者将它们视为一个模块，并通过分析其信息失真来考察其贡献。此外，作者引入了一种仅基于权重且不需要输入图像的新型权重度量，从而消除了对块合并过程的依赖。作者的方法在不同大小的Swin Transformer上验证了其实用性和优势，在ImageNet-1k分类任务上。值得注意的是，当作者移除了Swin-B的52.5% FLOPs和52.7%参数时，top-5准确率下降仅为0.07%。当作者减少了Swin-S的33.2% FLOPs和33.2%参数时，甚至可以实现比原始模型高0.8%的相对top-5准确率。代码可在以下链接找到：https://github.com/he-y/Data-independent-Module-Aware-Pruning。

1 Introduction

视觉 Transformer （Dosovitskiy等人(2020); Touvron等人(2020); Yuan等人(2021)）在计算机视觉领域，包括图像分类、检测和分割方面取得了最新的（SOTA）性能。然而，自注意力机制的利用和卷积的移除导致视觉 Transformer （Dosovitskiy等人(2020); Liu等人(2021)）的计算负担沉重且参数数量巨大。因此，有必要修剪模型以减少计算成本和所需的存储空间。

由于它们具有不同的结构和权重属性（Raghu等人(2021); Park和Kim(2022)），直接将传统的CNN权重剪枝方法应用于视觉 Transformer 是具有挑战性的。现代CNN（如VGGNet（Simonyan和Zisserman(2015)）和ResNet（He等人(2016)）最重要的结构是卷积层。如图1所示，ResNet-50的卷积层占整个网络92.0%的参数和99.9%的浮点运算（FLOPs）。相比之下，卷积层在Swin Transformer中仅占0.1%的参数和0.1%的FLOPs。因此在剪枝过程中应该考虑ViTs的特殊结构。

尽管已经有尝试剪枝传统的ViTs，但分层ViTs的独特特性尚未得到充分探索。局部注意力是分层ViTs的一个基本特性，它要求在窗口内而不是整个图像上进行注意力计算。虽然这在很大程度上减少了计算量，但也提出了一个问题：即使两个窗口属于同一图像，一个窗口内的注意力权重也不能与另一个窗口的权重进行比较。

作者用一个例子来解释之前幅度剪枝方法的问题。假设窗口

的注意力值范围从0到0.3，窗口

的范围从0到0.7。之前的幅度剪枝方法（Han等人(2015)）将做出这样的剪枝决策：在移除窗口

中大于0.3的权重之前，应该移除窗口

中的所有权重。然而，这种决策可能并不总是正确的，因为移除特定窗口（如窗口

）内的所有权重可能会导致图像中丢失重要的物体。因此，在确定图像中最不重要的注意力权重时，作者不应考虑其幅度，而应根据其贡献对它们进行排名。

分层ViTs的第二个特性是图像块在更高层次上合并。相比之下，传统的ViTs（Dosovitskiy等人(2020); Touvron等人(2020)）在不同层次上具有相同数量的图像块。分层ViTs的块合并特性旨在为密集预测提取从粗到细粒度的特征。为了实现这一特性，不同层的权重分布应适应那些层的图像块。为了最小化不同图像块大小的影响，权重重要性度量应该是数据独立的。作为权重度量最简单的方法是使用幅度值，这可以确保数据的独立性。然而，在跨窗口比较注意力值或跨层比较权重值时，这种方法并不有效。

在本文中，作者提出了一种数据独立模块感知剪枝方法（DIMAP）来解决上述两个问题。首先，作者在模块内包含不同的层，并在模块 Level 而不是窗口或层 Level 评估权重的重要性。通过分析单个权重引起的Frobenius扭曲（Park等人(2020)），作者可以获得它们的权重重要性，并安全地比较所有权重的贡献。这使得作者可以在模块内评估权重的重要性，而不管它们位于窗口或层内的位置。此外，作者在模块 Level 分析信息扭曲，并开发了一种新颖的数据独立权重重要性度量。DIMAP中提出的数据独立模块感知权重重要性度量具有几个理想的特性。首先，它为模块 Level 的权重重要性提供了一个“全局”排名，同时不改变层内的“局部”排名，使其成为层状幅度剪枝方法的泛化。其次，它不使用任何额外的参数，完全基于权重值，使其简单且高效。最后，它可以不需要复杂的剪枝敏感性分析就能应用，允许有效的单次剪枝。

贡献总结如下：

作者分析了单个权重和模块 Level 剪枝引起的Frobenius扭曲。
作者提出了一种数据独立的权重重要性度量来解决之前剪枝方法的问题。
实验验证了作者的方法在图像分类基准上的结果与最先进的结果相当。

2 Related Works

卷积神经网络剪枝。1) 权重剪枝。 权重剪枝旨在修剪网络的细粒度权重。例如，Han等人（2015年）提出移除绝对值低于阈值的较小权重。Guo等人（2016年）提出在训练过程中动态丢弃权重。Lebedev和Lempitsky（2016年）在损失函数上利用组稀疏正则化，使一些整个权重组向零收缩。然而，这些权重剪枝方法是为CNN设计的，因此很难将这些方法直接应用于具有不同性质的视觉 Transformer （Raghu等人，2021；Park和Kim，2022）。2) 滤波器剪枝。 滤波器剪枝；Suau等人，完全移除滤波器以获得具有结构化稀疏性的模型，因此被剪枝的卷积模型可以实现更好的加速。Li等人（2017年）使用L1范数来评估网络中滤波器的重要性。He等人（2018年）利用强化学习自动寻找每一层的冗余。然而，在视觉 Transformer 中并没有那么多的“滤波器”，所以作者无法利用这些算法。

视觉 Transformer 的变体。 视觉 Transformer 最初由ViT Dosovitskiy等人（2020年）提出，它不是在像素上使用注意力，而是在图像的小块上使用注意力。此后，出现了不同的视觉 Transformer 变体Srinivas）引入了一个蒸馏标记，使学生网络通过注意力从教师网络学习。BoTNet Srinivas等人（2021年）提出用全局自注意力替换ResNet最后三个瓶颈块中的空间卷积。T2T-ViT Yuan等人（2021年）引入了逐层Tokens-To-Token转换，从零开始训练视觉 Transformer 。PVT Wang等人（2021年）利用金字塔结构进行密集预测任务。CrossViT Chen等人（2021年）用ViT学习多尺度特征表示。TNT Han等人（2021年）引入了Transformer in Transformer，以挖掘不同尺度和位置目标的特征。CoaT Xu等人（2021年）为图像 Transformer 赋予了丰富的多尺度上下文建模能力。LV-ViT Jiang等人（2021年）提出了一种新的ViT训练损失，通过利用所有块标记。Swin Transformer Liu等人（2021年）因其分层结构和高效的移位窗口注意力而受到关注。

视觉Transformer压缩。一些并行方法，如Heo等人；Chuanyang等人（2022年）被提出用于视觉Transformer的压缩和加速。1) 这一方向聚焦于网络的冗余性，并大多保持了原始网络的结构。一个重要的方向是减少输入图像标记（Lee等人，2023年）。例如，DynamicViT（Rao等人，2021年）逐步剪除多余的标记。EViT（Liang等人，2022年）重新组织标记以减少多头自注意力计算的成本。SVTE（Chen等人，2021年）提出了一种带有训练标记选择器的稀疏ViT。2) 另一个方向是处理网络本身。VTP（Zhu等人，2021年）使用稀疏正则化剪除ViT中不重要的特征。AutoFormer（Chen等人，2021年）使用一种架构搜索框架进行视觉Transformer搜索。As-ViT（Chen等人，2022年）自动扩展ViT。UVC（Yu等人，2021年）将三种有效技术结合在一起，包括剪枝、层跳过和知识蒸馏，用于高效的ViT。NViT（Yang等人，2021年）使用结构剪枝，对所有视觉Transformer的参数施加延迟感知正则化。大多数之前的压缩文献（Chen等人，2021年；Zhu等人，2021年；Chen等人，2022年；2021b）专注于剪除像DeiT（Touvron等人，2020年）这样的常规ViT，这些常规ViT使用全局自注意力。相比之下，分层ViT如Swin Transformer利用局部自注意力来节省计算成本。因此，压缩分层ViT比压缩常规ViT要困难得多。

3 Module-aware Pruning

Preliminaries

假设Swin Transformer有

层，且

\mathbf{W}^{(l)}\in\mathbb{R}^{K\times K\times N_{l}^{(l)}\times N_{O}^{(l)}}

是第

l_{th}

层的权重，其中

是核大小，

N_{l}^{(l)}

和

N_{O}^{(l)}

分别是输入和输出通道的数量。Swin Transformer的第一层是块嵌入层。它是一个卷积层，所以

K>1

且此块嵌入层的维度为4。对于后文的关注相关层和全连接层，这些层的维度为2，因此作者可以将

的值视为

K=1

。LayerNorm（LN）层是一维向量，这些层仅占FLOPs和存储的一小部分。

权重分布。不同功能层的权重分布如图2所示。第一个观察点是，第一块和第二块之间QKV层的权重分布有显著差异：第一块的QKV层包含的零明显多于第二块。分布的不相似性支持作者的观点，即直接比较关注度的幅度是不恰当的。此外，即使是相同的块，不同的功能层也有各种分布。应用0.03的幅度剪枝阈值，将导致从QKV和PRJ层移除大约30%的权重，但从FC1和FC2层移除多达80%的权重。很明显，幅度剪枝导致这些层的剪枝结果不平衡。为了确保跨不同层权重之间的公平比较，作者采用信息失真分析作为剪枝的权重度量，如第3.2节所示。

模块定义模块的定义如图3所示。更多细节可以在附录B中找到。

Distortion Analysis and Weight Metric

单层分析。 首先，作者使用一个单层的全连接层1从最小化修剪后

\ell_{2}

失真的角度考虑修剪问题，如Neyshabur等人（2015年）；李等人（2020年）所述。假设一个输入

x\in\mathbb{R}^{n}

和一个权重张量

W\in\mathbb{R}^{m\times n}

，修剪操作可以被看作是

\widetilde{W}:=M\odot W

，其中

是一个二元 Mask ，其条目为0或1。

是稀疏性约束

\|M\|_{0}\leq s

。作者的目标是找到一个修剪 Mask

，以最小化修剪层和未修剪层的输出之间的差异：

\min_{\|M\|_{0}\leq s}\sup_{\|x\|_{2}\leq 1}\|Wx-(M\odot W)x\|_{2} \tag{1}

辅助模块（AUX-M）。 辅助模块包括Swin Transformer块内外的辅助层。在Swin Transformer块内部，如图3(a)所示，有两个LN层。Swin Transformer块外的其他一些辅助层也包括在AUX-M中，这些层包括一个 Patch 嵌入层和几个 Patch 合并层。由于这些层对网络的贡献重要且参数数量相对较少，作者不修剪这些层。