摘要

视觉Transformer（Vision Transformers, ViTs）与MLP-Mixers标志着在用通用神经网络架构替代人工设计的特征或归纳偏置方面的进一步探索。现有方法通过大规模数据训练（如大规模预训练和/或反复的强数据增强）来提升模型性能，但仍存在优化相关问题（例如对初始化和学习率的敏感性）。为此，本文从损失函数几何结构的角度出发，研究ViTs与MLP-Mixers，旨在提升模型在训练阶段的数据效率以及在推理阶段的泛化能力。通过可视化分析与Hessian矩阵分析发现，收敛后的模型存在极其尖锐的局部极小值。通过引入一种近期提出的锐度感知优化器（sharpness-aware optimizer）以增强模型的平滑性，我们显著提升了ViTs与MLP-Mixers在多种任务上的准确率与鲁棒性，涵盖监督学习、对抗学习、对比学习及迁移学习等场景。例如，在仅采用简单Inception风格预处理的情况下，ViT-B/16与Mixer-B/16在ImageNet上的Top-1准确率分别提升了5.3%与11.0%。进一步分析表明，模型平滑性的提升主要归因于前几层中活跃神经元的稀疏化。由此得到的ViTs在从零开始训练（无大规模预训练或强数据增强）的情况下，其性能已超越同等规模与吞吐量的ResNet。模型检查点已公开于：\url{https://github.com/google-research/vision_transformer}。