《Zero-TPrune: Zero-Shot Token Pruning through Leveraging of the Attention Graph》CVPR2024

夏日的盒盒

于 2024-10-10 21:50:38 发布

阅读量766

点赞数 21

文章标签：剪枝算法机器学习深度学习 transformer 计算机视觉

本文链接：https://blog.csdn.net/qq_46981910/article/details/142832549

版权

摘要

这篇论文介绍了一种名为Zero-TPrune的方法，旨在解决预训练Transformer模型在边缘设备部署时面临的挑战。这些挑战主要是由于模型大小和推理成本的指数级增长，特别是输入序列中的令牌数量导致的计算复杂度呈二次方增加。Zero-TPrune是一种零样本（zero-shot）令牌修剪方法，它利用预训练Transformer模型的注意力图来对令牌进行重要性排名，并移除信息量较少的令牌。这种方法不需要在剪枝后进行昂贵的微调（fine-tuning）过程。

概述

拟解决的问题：Transformer模型因其强大的迁移学习能力和显著的学习容量，在多个应用领域表现出色。然而，随着模型参数数量的激增，这些模型在边缘设备上的部署变得越来越困难，尤其是在能源、内存和计算资源受限的情况下。

创新之处：

零样本修剪：Zero-TPrune是第一个在不需要微调的情况下，同时考虑令牌的重要性和相似性进行修剪的方法。
注意力图的利用：通过将注意力矩阵视为有向图的邻接矩阵，并应用图移位算子迭代获得重要性得分分布。
相似性测量：在修剪过程中，Zero-TPrune不仅考虑令牌的重要性，还考虑它们之间的相似性，以进一步减少模型的计算成本。
高效性能：与需要微调的最先进的修剪方法相比，Zero-TPrune在修剪后减少了约0.3%的准确率损失，同时消除了微调的需要。

方法

Zero-TPrune 的总体框架。可以在任何 Transformer 块之间插入修剪层以减少令牌的数量。每个修剪层由一个 I' 阶段、I 阶段和 S 阶段组成。

Zero-TPrune方法分为三个阶段：I-stage和S-stage。

I′-stage（预重要性评估阶段）：在这一阶段，对令牌进行初步的重要性评估，但不进行修剪。
I-stage（重要性评估阶段）：在这一阶段，使用加权PageRank（WPR）算法来计算令牌的重要性得分，并保留最重要的令牌。
S-stage（相似性评估阶段）：在这一阶段，基于令牌的嵌入向量测量它们之间的相似性，并在相似的令牌对中只保留一个令牌。

Zero-TPrune的核心思想是利用预训练Transformer模型的注意力图来评估令牌的重要性和相似性。注意力图可以被视为一个有向图，其中节点代表令牌，边代表令牌之间的注意力权重。通过在这个图上应用图信号处理技术，可以迭代地计算每个令牌的重要性得分。

Zero-TPrune的修剪过程包括以下几个关键步骤：

预处理：将注意力矩阵视为有向图的邻接矩阵。
重要性评估：使用WPR算法迭代地计算每个令牌的重要性得分。
令牌划分：根据重要性得分将令牌划分为两组，一组包含较重要的令牌，另一组包含较不重要的令牌。
相似性评估：在较不重要的令牌组内，测量令牌之间的相似性，并在相似的令牌对中选择性地修剪令牌。
修剪：根据上述步骤的结果，修剪掉选定的令牌，从而减少模型的计算复杂度。

3.1 I-stage:基于重要性的剪枝

为了保留前 k 个重要标记，我们首先需要定义一个将不重要的标记与重要标记区分开来的度量。该指标称为重要性分数。为了获得重要性分数，我们将注意矩阵 A(h,l) 视为有向图的邻接矩阵，称为注意力图，如图所示。我们希望使用这个图中的信息来为每个标记（节点）分配一个重要性分数并删除不重要的分数。这一过程面临的挑战包括：(i) 注意力图很复杂。该图是有向和全连接的，当输入是图像时，通常包括数百个节点。(ii) 我们对算法产生的计算开销有严格的约束。对于不同的输入，在推理过程中需要修剪一组不同的标记。

为了解决上述挑战，论文提出了一种受PageRank算法启发的加权PageRank (WPR) 算法来派生令牌的重要性得分。PageRank算法最初用于Google搜索中对网页进行排名，而在Zero-TPrune中，每个节点的图信号被视为每个令牌的重要性。

我们统一初始化图信号。然后，我们使用邻接矩阵作为图移位算子 (GSO)。当 GSO 应用于图信号时，每个节点通过分配给输出边的权重投票哪个节点更重要，即令牌关注其他令牌。请注意，每个节点在投票时可能产生不同的影响。如果节点本身更重要，则该节点的投票更为重要。如图所示，可以看到重要性的排名。

我们得到第 l 层每个节点的重要性分数（即令牌）的表达式，第 h 个头如下：

3.2 S-stage：基于相似性的剪枝

以前的大多数令牌修剪方法仅依赖于令牌重要性来进行修剪，而忽略了令牌之间的相似性的影响。然而，高重要性分数并不总是表明令牌的必要性。图(a) 显示了鱼图像，其红色框部分非常重要。如果这部分被屏蔽，如图 3(b) 所示，我们无法识别该图像表示的内容。这也是我们的基于图的 WPR 算法获得的重要性分数分布所证实的。然而，如果我们屏蔽这条鱼的中心，如图3(c)所示，将这个图像识别为鱼仍然很容易，尽管中心中的标记在我们的算法中具有很高的重要性分数。这表明一旦选择了一些重要的标记，其他一些重要的标记就不再需要。

S-stage的核心思想是，即使某些令牌具有高的重要性得分，它们在模型进行有效推理时也可能是不必要的。这是因为在图像识别等视觉任务中，输入图像的一些部分即使被遮挡，也不会显著影响模型的识别能力。因此，测量重要令牌之间的相似性，并基于这种相似性进行修剪，可以进一步减少模型的计算成本。

修剪流程：

划分（Partition）：根据令牌的重要性得分，将令牌划分为两个大致相等大小的组，通常称为Group A和Group B。论文中提到了几种不同的划分策略，包括交替划分、顺序划分和随机划分。
匹配（Match）：对于Group A中的每个令牌，识别Group B中与之最相似的令牌，并记录这些对的相似度。
修剪（Prune）：选择相似度最高的几对令牌，并在每对中修剪掉一个令牌。这里选择修剪而不是合并令牌，是因为相似的令牌只包含很少的不同信息，修剪一个令牌只会带来最小的信息损失。

结论

Zero-TPrune通过在不同的Transformer骨干网络上实现和评估，展示了在不需要修剪后微调的情况下，与最先进的Transformer修剪方法相比，只有微小的准确率降低。此外，与不需要微调的最先进方法相比，Zero-TPrune在保持性能的同时，最多可以减少45%的准确率损失。这强调了Zero-TPrune在平衡模型压缩和性能保持方面的有效性，使其成为高效、准确修剪Transformer模型的有前景的方法。未来的工作可以进一步增强Zero-TPrune的能力，包括将其应用于图像重建、分割和生成等任务。