[论文解读] When Vision Transformers Outperform ResNets without Pre-training or Strong Data Augmentations

Xiangning Chen, Cho‐Jui Hsieh+1 more|arXiv (Cornell University)|Jun 3, 2021

Advanced Neural Network Applications参考文献 60被引用 103

一句话总结

论文表明，使用尖锐度感知优化器（SAM）从头训练的 Vision Transformer（ViT）和 MLP-Mixer 在与同等规模的 ResNet 相配对时可以超过它们；且无需大规模预训练或强数据增强，并通过分析损失景观来解释原因。

ABSTRACT

Vision Transformers (ViTs) and MLPs signal further efforts on replacing hand-wired features or inductive biases with general-purpose neural architectures. Existing works empower the models by massive data, such as large-scale pre-training and/or repeated strong data augmentations, and still report optimization-related problems (e.g., sensitivity to initialization and learning rates). Hence, this paper investigates ViTs and MLP-Mixers from the lens of loss geometry, intending to improve the models' data efficiency at training and generalization at inference. Visualization and Hessian reveal extremely sharp local minima of converged models. By promoting smoothness with a recently proposed sharpness-aware optimizer, we substantially improve the accuracy and robustness of ViTs and MLP-Mixers on various tasks spanning supervised, adversarial, contrastive, and transfer learning (e.g., +5.3\% and +11.0\% top-1 accuracy on ImageNet for ViT-B/16 and Mixer-B/16, respectively, with the simple Inception-style preprocessing). We show that the improved smoothness attributes to sparser active neurons in the first few layers. The resultant ViTs outperform ResNets of similar size and throughput when trained from scratch on ImageNet without large-scale pre-training or strong data augmentations. Model checkpoints are available at \url{https://github.com/google-research/vision_transformer}.

研究动机与目标

• 研究从头训练的 ViT 和 MLP-Mixer，以理解与 ResNet 相比的优化和泛化差距。
• 分析损失景观和 Hessian 以表征卷积缺失结构的尖锐最小值。
• 评估尖锐度感知优化器（SAM）对准确性、鲁棒性和训练动态在不同任务中的影响。
• 证明 SAM 使在没有大规模预训练或强增强的情况下实现有竞争力甚至更优的表现。

提出的方法

• 在 ImageNet 收敛时研究 ViT 与 MLP-Mixer 的损失景观和 Hessian 谱。
• 将尖锐度感知最小化器（SAM）应用于 ViT 和 MLP-Mixer，以促进损失几何中的平坦区域。
• 使用相同的基础预处理，将 SAM 训练的 ViT/Mixer 与从头训练的 ResNet 进行比较，评估准确性和鲁棒性。
• 分析使用 SAM 时模型的内在变化，包括 Hessian 块、权重范数和神经元激活稀疏性。
• 可视化注意力图以评估由于 SAM 带来的可解释性提升。

实验结果

研究问题

• RQ1 在没有预训练或大量增强的情况下，ViT 和 MLP-Mixer 是否可通过面向损失几何的优化器超越从头训练的 ResNet？
• RQ2 SAM 如何影响损失景观、Hessian 曲率以及卷积缺失架构的可训练性？
• RQ3 伴随 SAM 出现的内在架构变化有哪些（如稀疏性、权重范数），它们与泛化和鲁棒性之间的关系如何？
• RQ4 在 ImageNet-R 和 ImageNet-C 的情况下，使用 SAM 训练的 ViT 和 Mixer 在准确性和鲁棒性方面相较 ResNet 的表现如何？

主要发现

• SAM 显著提升从头训练的 ViT 和 MLP-Mixer 的准确性和鲁棒性（例如 ViT-B/16 和 Mixer-B/16 在 ImageNet 上获得显著的 top-1 增益）。
• ViT 与 MLP-Mixer 在没有 SAM 的情况下收敛到更尖锐的局部极小值，具有较大的 Hessian 特征值，但 SAM 能平滑损失景观并降低平均和最差曲率。
• SAM 导致早期层中更稀疏的活跃神经元，尤其在 MLP-Mixer 中，这与更平滑的损失几何和更好的泛化相关。
• 在 ImageNet 从头训练且吞吐量相近的条件下，带 SAM 的 ViT 在与同等或更大规模的 ResNet 相比下可以取得更好的性能，且无需大规模预训练或强增强。
• SAM 相对于基线模型，在对损坏（ImageNet-C）和分布变换（ImageNet-R）的鲁棒性提升方面的效果比在干净准确性方面更明显。
• SAM 使使用 SGD（带动量）的 ViT 训练接近 Adam/SAM 组合的性能，缩小优化器之间的差距。