[论文解读] Improve Vision Transformers Training by Suppressing Over-smoothing

Chengyue Gong, Dilin Wang+3 more|arXiv (Cornell University)|Apr 26, 2021

Advanced Neural Network Applications参考文献 4被引用 22

一句话总结

该论文提出了一种无需架构修改的视觉变换器训练稳定化技术，解决了自注意力层中过度平滑导致的特征崩溃和性能下降问题。通过引入促进多样性的损失函数以及用于Cutmix的图像块分类头，该方法使更深更宽的视觉变换器在ImageNet上实现了85.0%的top-1准确率，且无需额外参数或教师模型。

ABSTRACT

Introducing the transformer structure into computer vision tasks holds the promise of yielding a better speed-accuracy trade-off than traditional convolution networks. However, directly training vanilla transformers on vision tasks has been shown to yield unstable and sub-optimal results. As a result, recent works propose to modify transformer structures by incorporating convolutional layers to improve the performance on vision tasks. This work investigates how to stabilize the training of vision transformers \emph{without} special structure modification. We observe that the instability of transformer training on vision tasks can be attributed to the over-smoothing problem, that the self-attention layers tend to map the different patches from the input image into a similar latent representation, hence yielding the loss of information and degeneration of performance, especially when the number of layers is large. We then propose a number of techniques to alleviate this problem, including introducing additional loss functions to encourage diversity, prevent loss of information, and discriminate different patches by additional patch classification loss for Cutmix. We show that our proposed techniques stabilize the training and allow us to train wider and deeper vision transformers, achieving 85.0\% top-1 accuracy on ImageNet validation set without introducing extra teachers or additional convolution layers. Our code will be made publicly available at this https URL .

研究动机与目标

• 解决仅使用自注意力机制的原始视觉变换器在训练过程中出现的不稳定性和次优性能问题。
• 识别自注意力层中的过度平滑是深层变换器中特征崩溃和性能下降的根本原因。
• 在不修改变换器架构或添加卷积层的前提下，提升训练稳定性和模型容量。
• 通过保持特征多样性和信息完整性，实现更深更宽的视觉变换器的训练。
• 仅通过损失正则化技术，在ImageNet上实现最先进水平的准确率。

提出的方法

• 引入辅助图像块分类头，以在Cutmix增强过程中提升特征判别能力。
• 应用促进多样性的损失函数，以鼓励图像块之间表示的差异性。
• 使用类似对比学习的损失函数，通过保持块间差异来防止特征崩溃。
• 将图像块分类损失与Cutmix数据增强相结合，以提升泛化能力。
• 在使用标准自注意力机制的同时，通过这些损失组件稳定视觉变换器的训练。
• 避免采用残差连接或卷积归纳偏置等架构修改。

实验结果

研究问题

• RQ1 仅使用自注意力机制时，视觉变换器训练不稳定的根源是什么？
• RQ2 自注意力层中的过度平滑如何影响特征表示和模型准确率？
• RQ3 是否可以在不修改视觉变换器架构的前提下改善训练稳定性？
• RQ4 仅通过损失正则化技术，能够将更深更宽的视觉变换器训练到何种程度？
• RQ5 是否可以在不增加参数或知识蒸馏的情况下，在ImageNet上实现高准确率？

主要发现

• 所提方法在不修改架构的前提下稳定了视觉变换器的训练，使更深更宽的模型能够可靠训练。
• 该模型在ImageNet验证集上实现了85.0%的top-1准确率，与最先进结果相当，且无需知识蒸馏或额外参数。
• 增加图像块分类头和多样性损失显著减少了过度平滑现象，并提升了特征判别能力。
• 该方法在模型深度和宽度增加时仍能保持性能增益，展现出对架构扩展的鲁棒性。
• 该方法优于原始视觉变换器，并与依赖残差连接或卷积归纳偏置的模型表现相当。
• 代码已公开发布，有助于实现可复现性，并推动训练稳定视觉变换器的进一步研究。