[论文解读] So-ViT: Mind Visual Tokens for Vision Transformer.

Jiangtao Xie, Ruiren Zeng+3 more|arXiv (Cornell University)|Apr 22, 2021

Advanced Neural Network Applications参考文献 14被引用 11

一句话总结

So-ViT 提出了一种新型视觉 Transformer 架构，通过结合视觉 token 的二阶交叉协方差池化与类别 token，同时引入一种轻量级、分层的基于卷积的 token 嵌入模块，在从零开始训练时显著提升了性能。该方法在 ImageNet-1K 上实现了最先进（SOTA）的准确率，无需预训练，优于所有 ViT 变体和 CNN 模型。

ABSTRACT

Recently the vision transformer (ViT) architecture, where the backbone purely consists of self-attention mechanism, has achieved very promising performance in visual classification. However, the high performance of the original ViT heavily depends on pretraining using ultra large-scale datasets, and it significantly underperforms on ImageNet-1K if trained from scratch. This paper makes the efforts toward addressing this problem, by carefully considering the role of visual tokens. First, for classification head, existing ViT only exploits class token while entirely neglecting rich semantic information inherent in high-level visual tokens. Therefore, we propose a new classification paradigm, where the second-order, cross-covariance pooling of visual tokens is combined with class token for final classification. Meanwhile, a fast singular value power normalization is proposed for improving the second-order pooling. Second, the original ViT employs the naive embedding of fixed-size image patches, lacking the ability to model translation equivariance and locality. To alleviate this problem, we develop a light-weight, hierarchical module based on off-the-shelf convolutions for visual token embedding. The proposed architecture, which we call So-ViT, is thoroughly evaluated on ImageNet-1K. The results show our models, when trained from scratch, outperform the competing ViT variants, while being on par with or better than state-of-the-art CNN models. Code is available at this https URL

研究动机与目标

• 解决视觉 Transformer 在 ImageNet-1K 上从零开始训练时性能较差的问题，尽管其在预训练下表现强劲。
• 克服 ViT 中类别 token 的局限性，通过利用高层视觉 token 中丰富的语义信息。
• 通过引入一种分层的、基于卷积的嵌入模块，改进 token 表示学习，以建模局部性与平移等变性。
• 开发一种快速的归一化技术用于二阶池化，以提升计算效率和稳定性。
• 在无需 ImageNet-1K 预训练的情况下，实现与最先进 CNN 模型或 ViT 变体相当或更优的准确率。

提出的方法

• 提出一种新型分类头，将类别 token 与所有视觉 token 的二阶交叉协方差池化相结合，以捕捉 token 之间的关系。
• 引入一种快速的奇异值幂次归一化方法，以稳定并加速二阶池化操作。
• 设计一种轻量级、分层的卷积模块，以替代标准的固定大小 patch 嵌入，从而更好地建模局部模式和平移等变性。
• 将分层卷积嵌入模块集成到 ViT 主干网络中，从输入图像生成更具信息量的视觉 token。
• 在 ImageNet-1K 上端到端从零开始训练完整的 So-ViT 模型，不使用任何 ImageNet 预训练。
• 使用可学习的类别 token，其注意力机制作用于所有视觉 token，最终预测结果由池化后的视觉 token 特征与类别 token 共同决定。

实验结果

研究问题

• RQ1 利用类别 token 之外的高层视觉 token 是否能提升零样本和从零开始训练的 ViT 性能？
• RQ2 对视觉 token 进行二阶交叉协方差池化在捕捉图像分类判别性特征方面的有效性如何？
• RQ3 轻量级卷积模块是否能在保持 ViT 自注意力效率的同时，改善视觉 token 表示学习？
• RQ4 所提出的快速奇异值幂次归一化是否能提升二阶池化的稳定性和速度？
• RQ5 在无需预训练的情况下，从零开始训练的 ViT 变体能否在 ImageNet-1K 上达到或超越最先进 CNN 模型的准确率？

主要发现

• So-ViT 在从零开始训练 ImageNet-1K 时实现了最先进准确率，优于现有 ViT 变体。
• 所提出的二阶交叉协方差池化显著提升了分类性能，通过利用高层视觉 token 中丰富的语义信息。
• 分层卷积嵌入模块增强了局部归纳偏置和平移等变性，从而带来更优的特征表示。
• 快速奇异值幂次归一化实现了高效且稳定的二阶池化，提升了训练收敛性。
• So-ViT 在无需任何预训练的情况下，实现了与最先进 CNN 模型相当或更优的 ImageNet-1K 准确率。
• 消融实验证实，池化后的视觉 token 分类头与分层嵌入模块均对最终准确率的提升有显著贡献。