[论文解读] Pruning Self-attentions into Convolutional Layers in Single Path

Haoyu He, Jianfei Cai+5 more|arXiv (Cornell University)|Nov 23, 2021

Advanced Neural Network Applications参考文献 64被引用 16

一句话总结

SPViT 提出了一种单路径剪枝框架，通过在多头自注意力（MSA）与卷积操作之间共享权重，联合优化视觉变换器（Vision Transformers）中的自注意力与前馈网络剪枝，实现了高效、自动化的架构搜索，计算成本极低。该方法在 DeiT-B 上实现了 52.0% 的 FLOPs 减少，同时在 ImageNet-1k 上将 top-1 准确率提升了 0.6%，达到当前最优剪枝性能。

ABSTRACT

Vision Transformers (ViTs) have achieved impressive performance over various computer vision tasks. However, modeling global correlations with multi-head self-attention (MSA) layers leads to two widely recognized issues: the massive computational resource consumption and the lack of intrinsic inductive bias for modeling local visual patterns. To solve both issues, we devise a simple yet effective method named Single-Path Vision Transformer pruning (SPViT), to efficiently and automatically compress the pre-trained ViTs into compact models with proper locality added. Specifically, we first propose a novel weight-sharing scheme between MSA and convolutional operations, delivering a single-path space to encode all candidate operations. In this way, we cast the operation search problem as finding which subset of parameters to use in each MSA layer, which significantly reduces the computational cost and optimization difficulty, and the convolution kernels can be well initialized using pre-trained MSA parameters. Relying on the single-path space, we introduce learnable binary gates to encode the operation choices in MSA layers. Similarly, we further employ learnable gates to encode the fine-grained MLP expansion ratios of FFN layers. In this way, our SPViT optimizes the learnable gates to automatically explore from a vast and unified search space and flexibly adjust the MSA-FFN pruning proportions for each individual dense model. We conduct extensive experiments on two representative ViTs showing that our SPViT achieves a new SOTA for pruning on ImageNet-1k. For example, our SPViT can trim 52.0% FLOPs for DeiT-B and get an impressive 0.6% top-1 accuracy gain simultaneously. The source code is available at https://github.com/ziplab/SPViT.

研究动机与目标

• 为解决视觉变换器中因多头自注意力层导致的高计算成本及缺乏归纳偏置的问题。
• 实现预训练 ViT 的高效、自动剪枝，生成紧凑模型，并引入局部归纳偏置。
• 通过在单路径空间中统一 MSA 与卷积操作，降低基于 NAS 的剪枝方法的搜索成本与优化难度。
• 在目标 FLOP 约束下，联合优化 MSA 层与 FFN 扩展比例的剪枝。
• 通过对手动搜索架构的实证分析，探究 ViT 组件的结构重要性。

提出的方法

• 提出一种参数共享方案，将卷积操作表示为 MSA 参数的子集，实现注意力与卷积之间的共享参数化。
• 构建一个单路径搜索空间，将所有候选操作（MSA 与卷积）编码在每层的单一 MSA 中，降低搜索复杂度。
• 在 MSA 层中引入可学习的二值门，以选择使用 MSA 参数的哪一子集，从而有效剪枝为卷积操作。
• 将可学习门扩展至 FFN 层，以联合剪枝隐藏维度并优化 MLP 扩展比例。
• 端到端优化可学习门，以在目标 FLOP 约束下自动发现高效架构。
• 使用预训练的 MSA 参数初始化卷积核，实现高效且稳定的训练。

Figure 1: Multi-path vs. single-path search space. (a) Multi-path search space. MSA and different types of convolutional operations are added as separate trainable paths, leading to an expensive search cost. $\mathcal{P}$ and $g_{p}$ denote the ordered operation set and the learnable gate for the $p

实验结果

研究问题

• RQ1 统一的单路径搜索空间是否能降低基于 NAS 的 ViT 剪枝的计算成本与优化难度？
• RQ2 在 MSA 与卷积操作之间共享权重，是否能实现卷积核的有效且稳定的初始化？
• RQ3 在统一搜索空间中联合剪枝 MSA 与 FFN 层，是否能实现优于独立剪枝的效率-准确率权衡？
• RQ4 所搜索架构中的结构选择（例如，在浅层或深层块中将 MSA 剪枝为卷积）如何与模型的功能行为相关联？
• RQ5 FFN 容量在不同 ViT 块中的作用是什么？其在标准 ViT 与分层 ViT 架构中的表现如何变化？

主要发现

• SPViT 在 DeiT-B 上实现了 52.0% 的 FLOPs 减少，同时在 ImageNet-1k 上将 top-1 准确率提升了 0.6%，创下后训练剪枝的 SOTA 新纪录。
• 该方法成功将浅层 MSA 层剪枝为卷积层或跳跃连接，与显示浅层中存在局部注意力模式的注意力模式一致。
• 所有模型中的深层 MSA 层均被保留，证实其在全局表征学习中的重要性。
• 在标准 ViT 的深层 FFN 层中表现出更高的冗余性，优先被剪枝，表明后期阶段对容量需求较低。
• 在分层 ViT（如 Swin-Ti）中，每个阶段的浅层保留了更多的隐藏维度，表明在 token 合并后对更高容量有更强需求。
• 与多路径 NAS 相比，单路径搜索空间显著降低了搜索成本与优化难度，实现了对 MSA 与 FFN 剪枝的高效联合搜索。

Figure 2: Overview of our weight-sharing scheme, which expresses the bottleneck convolution with MSA parameters. First, we restrict the value and output projections in the MSA layer to only process the features within local regions. Then, we learn to ensemble the value projections with a parameteriz