[论文解读] Skip-Attention: Improving Vision Transformers by Paying Less Attention

Shashanka Venkataramanan, Amir Hossein Ghodrati+3 more|arXiv (Cornell University)|Jan 5, 2023

CCD and CMOS Imaging Sensors被引用 6

一句话总结

Skip-Attention (SkipAt) 通过重用早期层中的自注意力计算，避免在深层网络中重复进行昂贵的注意力计算，从而提升视觉Transformer（ViT）的性能。通过引入一种轻量级参数化函数（受深度可分离卷积启发），SkipAt 在图像分类、语义分割、去噪和视频任务中均实现了更高的吞吐量和更高的准确率，相较于基线ViT，在ADE20K数据集上实现最高达40%的推理加速，在自监督预训练中提速26%。

ABSTRACT

This work aims to improve the efficiency of vision transformers (ViT). While ViTs use computationally expensive self-attention operations in every layer, we identify that these operations are highly correlated across layers -- a key redundancy that causes unnecessary computations. Based on this observation, we propose SkipAt, a method to reuse self-attention computation from preceding layers to approximate attention at one or more subsequent layers. To ensure that reusing self-attention blocks across layers does not degrade the performance, we introduce a simple parametric function, which outperforms the baseline transformer's performance while running computationally faster. We show the effectiveness of our method in image classification and self-supervised learning on ImageNet-1K, semantic segmentation on ADE20K, image denoising on SIDD, and video denoising on DAVIS. We achieve improved throughput at the same-or-higher accuracy levels in all these tasks.

研究动机与目标

• 解决视觉Transformer（ViT）中自注意力计算成本过高的问题，其计算复杂度随序列长度呈二次方增长。
• 发现不同层之间的多头自注意力（MSA）输出高度相关，尤其是对于[CLS] token，表明存在冗余。
• 设计一种即插即用模块，通过重用早期层的表示来替代深层网络中的昂贵MSA计算。
• 设计一种轻量级参数化函数，在保持性能的同时提升效率，实现更快的推理与训练速度。
• 在多种视觉任务中验证方法的泛化能力，包括图像分类、语义分割、图像与视频去噪，以及自监督学习。

提出的方法

• 提出SkipAt，一种即插即用模块，通过重用早期层（如第3–8层）的注意力特征，替代深层网络中的MSA模块。
• 引入可学习的参数化函数Φ，将重用的特征映射到目标层的表示空间。
• 默认使用5×5卷积核的深度可分离卷积（DwC），通道扩展比为2，以在表达能力与效率之间取得平衡。
• 将跳过机制应用于分类任务与密集预测任务，包括图像级与像素级输出。
• 使用与基线ViT相同的训练数据与训练协议，进行端到端联合训练。
• 使该方法可适用于任意ViT架构，包括小模型、基础模型与大模型，适配多种数据集与任务。

Figure 2 : Attention correlation . Mean of the attention heads from the CLS token of a pretrained ViT-T/16 at different layers from the validation set of ImageNet-1K. Numbers below each attention map indicates the cosine similarity of $A_{l}^{\texttt{[CLS]}}$ with $A_{l-1}^{\texttt{[CLS]}}$ .

实验结果

研究问题

• RQ1 在不造成性能下降的前提下，深层ViT网络中的自注意力计算在多大程度上可被重用早期层的特征来近似？
• RQ2 一个简单的参数化函数是否能有效适配重用的注意力特征，从而在性能上超越直接重用或基线ViT？
• RQ3 SkipAt在不同视觉任务（包括密集预测与视频理解）中的推理速度与FLOPs如何变化？
• RQ4 SkipAt在自监督预训练中是否保持或提升性能，能否在不损失准确率的前提下缩短预训练时间？
• RQ5 在准确率与吞吐量之间取得最佳平衡时，参数化函数（卷积核大小、通道扩展比）的最优配置是什么？

主要发现

• 在ADE20K语义分割基准上，SkipAt实现40%的推理加速，同时将Top-1准确率提升0.6%（相比基线ViT-T/16）。
• 在ImageNet-1K上，SkipAt实现67.7%的Top-1准确率，吞吐量达6.9K张/秒（6.9×10³ img/sec），在准确率上比基线ViT高出1.9%，吞吐量提升19%。
• 在DAVIS视频去噪任务中，SkipAt达到SOTA方法的PSNR（35.16），但FLOPs减少17%（77.1 GFLOPs vs. 93.2 for UniFormer）。
• 在ImageNet-1K上，SkipAt将自监督预训练时间减少26%，且下游任务准确率无损失，证明其在表示学习中的效率优势。
• 最优SkipAt配置为使用5×5深度可分离卷积，通道扩展比为2，在ImageNet-1K上实现67.7%的Top-1准确率与6.9K张/秒的吞吐量。
• 消融实验表明，若使用恒等函数进行特征重用，准确率下降4.7%；而SkipAt的参数化函数在准确率与效率上始终优于所有基线模型。

Figure 3 : CKA analysis of $A^{[\texttt{CLS}]}$ and $Z^{\text{MSA}}$ across different layers of pretrained ViT-T/16 on the validation set of Imagenet-1K. Vanilla ViT-T/16 has high correlation across both attention maps (layer 3 to 10) and $Z^{\text{MSA}}$ (layer 2 to 8)