[论文解读] Vision Xformers: Efficient Attention for Image Classification

Pranav Jeevan, Amit Sethi|arXiv (Cornell University)|Jul 5, 2021

Advanced Neural Network Applications参考文献 18被引用 12

一句话总结

本文提出视觉X-former（ViX），一种改进的视觉Transformer架构，用线性注意力机制（如Performer、Linformer、Nyströmformer）替代二次方自注意力机制，引入卷积嵌入层以提供归纳偏置，并采用旋转位置编码（RoPE）。这些改进使GPU显存使用量最高降低七倍，同时保持或提升图像分类准确率，尤其在小样本数据集上表现更优，使Transformer在资源受限环境下更具可及性。

ABSTRACT

Although transformers have become the neural architectures of choice for natural language processing, they require orders of magnitude more training data, GPU memory, and computations in order to compete with convolutional neural networks for computer vision. The attention mechanism of transformers scales quadratically with the length of the input sequence, and unrolled images have long sequence lengths. Plus, transformers lack an inductive bias that is appropriate for images. We tested three modifications to vision transformer (ViT) architectures that address these shortcomings. Firstly, we alleviate the quadratic bottleneck by using linear attention mechanisms, called X-formers (such that, X in {Performer, Linformer, Nyströmformer}), thereby creating Vision X-formers (ViXs). This resulted in up to a seven times reduction in the GPU memory requirement. We also compared their performance with FNet and multi-layer perceptron mixers, which further reduced the GPU memory requirement. Secondly, we introduced an inductive bias for images by replacing the initial linear embedding layer by convolutional layers in ViX, which significantly increased classification accuracy without increasing the model size. Thirdly, we replaced the learnable 1D position embeddings in ViT with Rotary Position Embedding (RoPE), which increases the classification accuracy for the same model size. We believe that incorporating such changes can democratize transformers by making them accessible to those with limited data and computing resources.

研究动机与目标

• 为解决视觉Transformer（ViT）因二次方注意力复杂度导致的高计算与内存开销问题。
• 通过引入卷积嵌入层提供归纳偏置，提升在数据有限情况下的图像分类准确率。
• 通过将可学习的一维位置嵌入替换为旋转位置编码（RoPE），提升模型泛化能力与性能。
• 证明这些架构改进可普遍应用于多种基于Transformer的视觉模型，包括LeViT、CCT、CvT、PiT、FNet和MLP混合器。
• 通过降低数据与计算需求，实现视觉Transformer的普惠化。

提出的方法

• 将ViT中的标准自注意力机制替换为Performer、Linformer、Nyströmformer等线性注意力机制，将计算复杂度从O(n²)降低至O(n)。
• 用卷积层堆叠替代初始的线性图像块嵌入层，以提供2D图像结构的归纳偏置。
• 将可学习的一维位置嵌入替换为旋转位置编码（RoPE），其能更有效地编码相对位置信息，并在不同序列长度间具有更好的泛化能力。
• 在低数据与低GPU显存条件下，在标准图像分类基准（如ImageNet）上训练并评估修改后的ViX模型。
• 将所提改进扩展至其他视觉Transformer架构（如LeViT、CCT、CvT、PiT）及高效模型（如FNet、MLP混合器），以评估其泛化能力。
• 对比所有变体的性能与资源消耗（GPU显存、FLOPs、推理时间），量化效率提升效果。

实验结果

研究问题

• RQ1 用线性注意力机制替代二次方自注意力机制，是否能显著降低视觉Transformer的GPU显存与计算成本，同时不损失准确率？
• RQ2 用卷积层替代线性嵌入层，是否能通过引入2D归纳偏置显著提升图像分类准确率？
• RQ3 在视觉Transformer中，旋转位置编码（RoPE）是否能优于可学习的一维位置嵌入，尤其在低数据场景下？
• RQ4 这些架构改进在多大程度上可泛化至不同视觉Transformer架构及FNet、MLP混合器等高效模型？
• RQ5 这些改进是否能实现显著降低数据与硬件需求下的SOTA图像分类性能？

主要发现

• 在ViX中使用线性注意力机制（如Performer、Nyströmformer）使GPU显存使用量最高降低七倍，同时保持或提升准确率。
• 将线性嵌入层替换为卷积层显著提升了图像分类准确率，即使在模型更小、训练数据更少的情况下也表现出色，证明了2D归纳偏置的优势。
• 采用旋转位置编码（RoPE）进一步提升了分类准确率，且未增加模型参数或显著增加显存使用。
• 所提改进在LeViT、CCT、CvT、PiT等多种视觉Transformer架构中均有效，展现出广泛的适用性与一致的性能增益。
• FNet与MLP混合器模型也从改进中受益：FNet在结合卷积嵌入后表现出显著的准确率提升，而MLP混合器则未从这种混合设计中获益。
• 线性注意力、卷积嵌入与RoPE的结合，使在降低数据与计算需求的前提下实现高性能图像分类成为可能，显著提升了Transformer在低资源环境下的可及性。