Abstract
Multi-head self-attention (MSA) is a key component of Vision Transformers (ViTs), which have achieved great success in various vision tasks. However, their high computational cost and memory footprint hinder their deployment on resource-constrained devices. Conventional pruning approaches can only compress and accelerate the MSA module using head pruning, although the head is not an atomic unit. To address this issue, we propose a novel graph-aware neuron-level pruning method, Structured Neuron-level Pruning (SNP). SNP prunes neurons with less informative attention scores and eliminates redundancy among heads. Specifically, it prunes graphically connected query and key layers having the least informative attention scores while preserving the overall attention scores. Value layers, which can be pruned independently, are pruned to eliminate inter-head redundancy. Our proposed method effectively compresses and accelerates Transformer-based models for both edge devices and server processors. For instance, the DeiT-Small with SNP runs 3.1$\times$ faster than the original model and achieves performance that is 21.94\% faster and 1.12\% higher than the DeiT-Tiny. Additionally, SNP combine successfully with conventional head or block pruning approaches. SNP with head pruning could compress the DeiT-Base by 80\% of the parameters and computational costs and achieve 3.85$\times$ faster inference speed on RTX3090 and 4.93$\times$ on Jetson Nano.
Abstract (translated)
多头自注意（MSA）是 Vision Transformers（ViTs）的关键组件，已经在各种视觉任务中取得了巨大的成功。然而，它们的高计算成本和内存开销阻碍了它们在资源受限设备上的部署。传统的剪枝方法仅通过头剪枝来压缩和加速 MSA 模块，尽管头不是原子的单位。为解决这个问题，我们提出了一种新的基于图的神经元级剪枝方法，结构化神经元级剪枝（SNP）。SNP 通过剪除具有较低信息性注意分数的神经元来截断图形连接的查询和键层，同时保留整体注意分数。价值层（可以独立剪除）被剪除以消除头之间的冗余。我们提出的方法既有效地压缩了基于 Transformer 的模型，又加速了边缘设备和服务器处理器的运行速度。例如，使用 SNP 的 DeiT-Small 比原始模型快 3.1 倍，实现了性能比原模型快 21.94% 和更高的速度。此外，SNP 成功地与传统的头或块剪枝方法结合使用。使用头剪枝的 SNP 可以将 DeiT-Base 的参数压缩 80%，并在 RTX3090 上实现 3.85 倍的推理速度，在 Jetson Nano 上实现 4.93 倍的推理速度。
URL
https://arxiv.org/abs/2404.11630
