ICML 2023 | UPop: 使用统一渐进剪枝压缩视觉-语言Transformers

科技

2023-08-04 05:08

TL;DR: UPop 是首个针对 Vision-Language Transformers 的结构化剪枝框架，此外也适用于单模态 Transformers。UPop 在各种多模态和单模态任务上进行了广泛的实验，相关代码已开源。

论文题目：

UPop: Unified and Progressive Pruning for Compressing Vision-Language Transformers

文章链接：

https://proceedings.mlr.press/v202/shi23e/shi23e.pdf

代码链接：

https://github.com/sdc17/UPop

网站链接：

https://dachuanshi.com/UPop-Project

随着 AI 模型（例如各种 Transformer [1] 模型）的飞速发展，它们的参数量和计算代价也在急剧增长。要让这些大模型能在更多的消费级设备上使用和部署，模型压缩至关重要。虽然此前的研究者们提出了多种方法用于模型压缩（例如模型剪枝、量化、知识蒸馏等等），但是这些方法主要针对于单模态模型（例如单模态的 CV 和 NLP 模型）。至于如何压缩多模态模型，相对来说是有待探究的。

本文将介绍使用结构化剪枝方法 UPop 实现对于 Vision-Language Transformer 模型的压缩。

简介

是什么：UPop 是首个用于 Vision-Language Transformer 模型的结构化剪枝框架，它能够对各种多模态和单模态任务（视觉推理、图像描述、视觉问答、图像-文本检索，文本-图像检索、图像分类和图像分割）、数据集（NLVR2、COCO Caption、VQAv2、COCO、Flickr30K、ImageNet 和 ADE20K）和模型结构（BLIP [2]、CLIP [3]、DeiT [4] 和 Segmenter [5]）进行有效的结构化剪枝。

解决了什么问题：视频中演示了：（1）不同于以往使用网格搜索进行重复实验，UPop 使用的 Unified Search 能更高效地为不同模态和结构搜索给定总压缩率下各自的最优压缩率。（2）此外，UPop 使用的 Progressive Pruning 消除了以往的剪枝方法中存在着的，搜索后的模型与待重新训练的子网络之间的参数权重 gap，从而获得了更好的收敛性和在高压缩率下更好的表现。

表现如何：例如在多模态任务上，UPop 能在视觉问答的 VQAv2 数据集和视觉推理的 NLVR2 数据集上分别以 1.2% 和 2.0% 准确度的损失实现 2 倍的压缩。而例如在单模态任务上，UPop 能在图像分类的 ImageNet 数据集和图像分割的 ADE20K 数据集上分别实现 1.5 倍和 1.2 倍不降低准确度的无损压缩，下表展示了 UPop 在 vector-level structured 剪枝粒度下的部分实验结果。

需要注意的是 UPop 的实验结果是使用切实更小的参数矩阵得到的。例如剪枝前的模型参数矩阵为，那么剪枝后模型的参数矩阵则为，，其中至少有一个小于号是严格的。而对于非结构化剪枝（unstructured/weight pruning）的方法来说，剪枝后模型的参数矩阵为，也即剪枝后的模型参数矩阵由原有的模型参数矩阵乘上相同形状 element-wise 的 01 掩码矩阵得到，剪枝后总体参数形状仍为（因为矩阵中每行每列非 0 元素数量不同，所以不能实际剪出来一个小矩阵）。

方法

2.1 剪枝粒度

UPop 实现的是 vector-level 的结构化剪枝，也即最小的剪枝单元为模型参数矩阵的整行或者整列。相比于非结构化的剪枝方法，vector-level 的结构化剪枝使得UPop 可以将剪枝后的小模型真正地从原模型中抽取出来（矩阵消除整行或整列后仍然是一个合法的矩阵，然而非结构化剪枝会导致矩阵中每行或者每列需要剪掉的参数数量不同，从而剪枝后的参数无法构成一个合法的矩阵）。

总体而言，非结构化剪枝由于具有更精细的粒度，因此剪枝后的模型准确度会相对更高，但是部署会相对更困难。而结构化剪枝由于粒度更粗糙，因此剪枝后的模型的准确度会相对更低，但部署会更简单。UPop 方法本身并不受限于某一种粒度，而具体使用结构化还是非结构化来实现剪枝算法实际上是在模型准确性和可部署性之间的一个 tradeoff，因而在比较剪枝后的模型准确度时应限定在相同的粒度下进行比较才是公平的。

2.2 Unified Search

Unified Search 的使用解决了以往的方法中存在着的第一个 challenge：当给定一个总的压缩率时，我们需要手动地为不同模态中的不同的可压缩成分分配各自合适的压缩率。然而手动分配是十分低效的，尤其是当模型存在着多种类型的可压缩组件时（例如模型的视觉和语言分支中各自的 Self-Attentions, Cross-Attentions 以及 Feed-Forward Nets）。

下面以 50% 的总压缩率为例对比以往的方法和 UPop 所使用的 Unified Search：

2.1.1 分配相同的压缩率（Before）

首先最 naive 的方法是为所有不同模态中的不同的可压缩成分分配相同的压缩率。这种方法虽然简单，但是在多模态模型中，不同模态以及不同模态内部的不同成分的实际最优压缩率可能方差很大。在这种情况下，简单地给所有模态和成分分配相同的压缩率只能取得次优的表现。

2.1.2 使用Grid Search多次搜索合适的压缩率（Before）

表现更好的方法是使用多次重复实验（例如网络搜索）去测试为不同的模态和成分分配不同的压缩率，从而找到更优的分配方案。虽然这种策略能取得更好的表现，但与之相对地，时间成本也大大增加了。

2.1.3 使用Unified Search自适应地分配合适的压缩率（UPop）

Unified Search 将所有的模态中的所有的可压缩成分视为统一的空间，进行统一搜索和排序，给出全局的压缩率分配方案，自适应地为不同模态中的不同成分挑选出需要被剪枝的参数。Unified Search 将我们从为寻找合适的压缩率分配方案所进行的重复实验中解放了出来。

▲ UPop: Unified Search for adaptive assignments.

这里有个小 trick 是我们需要在统一排序前，先对不同结构所使用的重要性指标（在 Unified Search 中是可学习的 mask ，而在 Progressive Pruning 中是的累计梯度，其中：

）做 z-score standardization（下标 a 和 m 分别代表 Attention 和 FFN 结构）：

这是因为不同结构中重要性指标的分布可能差异很大，直接统一排序会对某些结构产生 bias，例如会导致绝大部分被剪枝的参数都集中在 Self-Attentions 中，而 z-score standardization 的使用减轻了这种 bias。

2.3 Progressive Pruning

Progressive Pruning 的使用解决了以往的方法中存在着的第二个 challenge：在搜索阶段结束后，模型中不重要的参数会被移除。然而，这些被移除的参数中很多具有非 0 的权重，将它们突然置为 0 会损害剪枝后的模型的收敛性。

换句话说，存在于搜索后的模型与待重新训练的剪枝后的模型之间的，参数权重上的显著的 gap，使得剪枝后的模型更难以收敛且损害了模型的表现。尤其是在高压缩率下，当模型参数中大量的非 0 权重突然被置为 0 后，剪枝后的模型很可能将无法收敛。

下面以剪除 50% 的模型参数为例对比以往的方法和 UPop 所使用的 Progressive Pruning：

2.3.1 一次性剪除50%的模型参数（Before）

首先最 naive 的方法是一次性剪除所要求的总参数量。这种方法虽然简单，但是剪枝前后的模型存在着前文所述的巨大的参数 gap，因此只能取得次优的表现，且剪枝后的模型在高压率下的收敛性较差。

▲ Before: Prune the required total number of parameters at once.

2.3.2 使用Iterative Pruning分多次剪枝（Before）

一个更优的解决方案是使用 Iterative Pruning [6] 分多次进行剪枝，例如下图所示的分为 4 次，每次剪枝 12.5% 的参数量。Iterative Pruning 可以有效地减小剪枝前后的模型之间的参数 gap，从而相比单次的剪枝有更优的表现。

▲ Before: Iterative Pruning for dividing pruning into multiple times.

2.3.3 使用Progressive Pruning消除剪枝前后的模型之间的参数gap（UPop）

Iterative Pruning 减小了剪枝前后的模型之间的参数 gap，而 Progressive Pruning 则是彻底消除了该参数 gap。Progressive Pruning 不仅将众多的参数分为多次剪枝，同时还将每个待剪枝的参数的剪枝过程分为了多次，以确保在搜索阶段结束后所有待剪枝的参数的权重正好为 0，从而消除了剪枝前后的模型之间的参数 gap，获得了更好的表现和在高压缩率下更好的收敛性。

▲ UPop: Progressive Pruning eliminates the parameter gap.

这里有个 trick 是 progressively 收缩参数权重直到为 0。通常来说模型中参数的更新是由优化器控制的，例如使用常规的梯度下降更新：

但它无法满足我们的要求——自由地控制每个迭代 t 时参数权重的值，例如我们要将一个初值为 1 的参数在 100 次迭代后刚好收缩到 0：

因此这里我们需要使用自定义的规则来进行部分参数的更新。具体来说我们使用常规的优化器来更新模型原有的参数，使用自定义的规则更新插入到模型中的 mask 的参数。例如在每个时间步 t 中，先更新当前期望达到的压缩率（其中为最终需要达到的压缩率，为搜索阶段的总迭代次数）：

这里我们使用上述的余弦更新会比前文中的均匀更新效果更好，再通过统一排序确定当前要剪枝的权重对象：

最后更新当前时间步 t 的参数权重：

2.4 算法实现

实验

UPop 的实验包括各种多模态和单模态任务（视觉推理、图像描述、视觉问答、图像-文本检索，文本-图像检索、图像分类和图像分割）、数据集（NLVR2、COCO Caption、VQAv2、COCO、Flickr30K、ImageNet 和 ADE20K）和模型结构（BLIP、CLIP、DeiT 和 Segmenter）。

3.1 在视觉推理任务的NLVR2数据集上压缩BLIP模型的实验

下表中 √ 和 × 分别代表在该压缩率下模型是否收敛。可见 UPop 在相同的压缩率下能取得更好的表现，同时在更高的压缩率下也具有更好的收敛性。参数量 Params 和 FLOPs 的单位分别为 M 和 G。

▲ Multimodal example1: Compress BLIP on the NLVR2 dataset of the Visual Reasoning task.

3.2 验证Unified Search有效性的实验

以 NLVR2 数据集为例，下表中 Manual 代表给定总压缩率为 50% 时，手动实验多组不同的分配方案，而 Adaptive 代表使用 Unified Search 获得自适应的分配方案。可见使用 Unified Search 更加省时，剪枝后的模型表现也更好。

3.3 验证Progressive Pruning有效性的实验

以 NLVR2 数据集为例，下表展示了给定总压缩率为 50% 时，在搜索阶段结束后，不做任何 retrain 直接用被剪枝的模型进行测试，可见 Progressive Pruning 可以使得模型在无需 retrain 的情况下也能收敛。

▲ Searching only and without any retraining

类似地，下表展示了给定总压缩率为 50% 时，在搜索阶段结束后，对被剪枝的模型只进行一个 epoch 的 retrain，可见 Progressive Pruning 能取得显著更好的效果。

▲ Retraining only one epoch after searching

3.4 在图像描述任务的COCO Caption数据集和视觉问答任务的和VQAv2数据集上压缩BLIP的实验

符号和单位同前文，指标 CIDEr 和 SPICE 为越高越好。可见 UPop 在不同的任务和数据集上也有良好的通用性。

▲ Multimodal example2：Compress BLIP on the COCO Caption dataset of Image Caption task and VQAv2 dataset of Visual Question Answering task.

3.5 在图像-文本检索和文本-图像检索任务的COCO和Flickr30K数据集上压缩BLIP模型的实验

符号和单位同前文，指标 R@1, R@5 和 R@10 为越高越好。实验结果进步一验证了 UPop 的通用性。

▲ Multimodal example3: Compress BLIP on the COCO and Flickr30K datasets of Image-Text Retrieval and Text-Image Retrieval tasks.

3.6 在图像-文本检索和文本-图像检索任务的COCO和Flickr30K数据集上压缩CLIP模型的实验

符号和单位同前文，实验结果进一步验证了 UPop 在不同模型结构上的通用性。

3.7 在图像分类任务的ImageNet数据集上压缩DeiT模型的实验

值得一提的是虽然 UPop 的实验主要在多模态 Transformers 上完成，但论文中也报告了 UPop 在单模态任务上的表现。如下列 tradeoff 图中的蓝色三角所示，UPop 在单模态 Transformer 的压缩上也能取得非常有竞争力的表现。

▲ Unimodal example1: Compress DeiT on the ImageNet dataset of Image Classification task.

3.8 在图像分割任务的ADE20K数据集上压缩Segmenter模型的实验

在图像分割任务上的实验结果表明：与原始模型相比，使用 UPop 剪枝后的模型可以在图像分割任务上得到 1.2 倍的无损压缩和约 1.4 倍的损失低于 1%mIoU 的压缩；与其他参数量或计算量相近的模型相比，UPop 剪枝后的模型可以得到很有竞争力的 trade-offs。例如 1.1 倍压缩后的模型在 multi-scale testing（MS）上可以超过所有其他模型，在 single-scale testing（SS）上也能在只使用第一名 54% 计算量的情况下取得第二；

与在图像分类任务上的实验结果对比，在图像分类任务上，UPop 可以取得约 1.5 倍的无损压缩和 2 倍的损失低于 1% Acc 的压缩，显然 UPop 在图像分类任务上有更好的表现。我们认为这种表现上的差异应归结于任务不同的内在属性。对于分类任务来说由于很多 pixel 是冗余的，对于最终结果影响很小，因此能被更容易地剪枝。而对于分割这种密集预测型的任务来说，由于需要输出每个 pixel 对应的类别，因此存在的冗余更少，对模型的剪枝也更困难。

▲ Unimodal example2: Compress Segmenter on the ADE20K dataset of Image Segmentation task.

3.9 各个模态中不同的可压缩成分的占比随压缩率增长而变化的可视化

下图展示了随着压缩率的增长，在剪枝后的模型中，vision 和 language 两个模态的 Self-Attention 模块的参数量占总参数量的比例逐渐上升，而 FFNs 模块的占比则逐渐减小。这表明当模型的参数量受限时，两个模态中 Self-Attentions 模块将比 FFNs 模块更加重要。

此外我们也能观察到在进行实验的视觉推理任务上，vision 模态相比 language 模态更加重要。Cross-Attentions 占比的变化趋势则是总体上先下降后上升。这表明在低压缩率下，vision 和 language 两个模态的参数量是较为充足的，此时 Cross-Attentions 的重要性更低，而在高压缩率下 vision 和 language 两个模态缺乏足够的参数量，此时 Cross-Attentions 会变得更为重要。

3.10 各层的占比随着压缩率增长而变化的可视化

下图展示了随着压缩率的增长，在剪枝后的模型中，模型的中间层将占据越来越大的比重，这揭示了模型中大部分有价值的信息是在模型的中间层趋于完善。在更浅的层，信息还不够完善，而在更深的层，对于中间层的信息做进一步的 refinements 在参数量受限时并不那么重要。

▲ Variation of layers as the compression ratio increases.

总结

UPop 是一个支持多种多模态和单模态任务、数据集和模型结构的通用的结构化剪枝框架，同时也是第一个针对 Vision-Language Transformers 模型结构化剪枝的框架。UPop 具有更加简洁高效的搜索，自适应的压缩率分配，消除了模型剪枝前后的参数 gap，在高压缩率下有更好的收敛性等优点。充分的实验展示了 UPop 方法具有良好的有效性和通用性。

参考文献

[1] Attention is All you Need https://proceedings.neurips.cc/paper_files/paper/2017/file/3f5ee243547dee91fbd053c1c4a845aa-Paper.pdf

[2] Blip: Bootstrapping language-image pre-training for unified vision-language understanding and generation https://proceedings.mlr.press/v162/li22n/li22n.pdf

[3] Learning Transferable Visual Models From Natural Language Supervision https://arxiv.org/pdf/2103.00020.pdf

[4] Training data-efficient image transformers & distillation through attention https://arxiv.org/pdf/2012.12877.pdf

[5] Segmenter: Transformer for Semantic Segmentation https://openaccess.thecvf.com/content/ICCV2021/papers/Strudel_Segmenter_Transformer_for_Semantic_Segmentation_ICCV_2021_paper.pdf

[6] Learning both Weights and Connections for Efficient Neural Networks https://proceedings.neurips.cc/paper/2015/file/ae0eb3eed39d2bcef4622b2499a05fe6-Paper.pdf

更多阅读