[论文解读] Early Convolutions Help Transformers See Better

Tete Xiao, Mannat Singh+4 more|arXiv (Cornell University)|Jun 28, 2021

Advanced Neural Network Applications参考文献 56被引用 352

一句话总结

将 ViT 的 patchify stem 替换为一个轻量级卷积干（ViT C）显著提高优化稳定性，在不损失准确度的前提下使 SGD 或 AdamW 可用，并在 ImageNet 的不同模型规模和数据规模下实现 1–2% 的 top-1 提升。

ABSTRACT

Vision transformer (ViT) models exhibit substandard optimizability. In particular, they are sensitive to the choice of optimizer (AdamW vs. SGD), optimizer hyperparameters, and training schedule length. In comparison, modern convolutional neural networks are easier to optimize. Why is this the case? In this work, we conjecture that the issue lies with the patchify stem of ViT models, which is implemented by a stride-p p*p convolution (p=16 by default) applied to the input image. This large-kernel plus large-stride convolution runs counter to typical design choices of convolutional layers in neural networks. To test whether this atypical design choice causes an issue, we analyze the optimization behavior of ViT models with their original patchify stem versus a simple counterpart where we replace the ViT stem by a small number of stacked stride-two 3*3 convolutions. While the vast majority of computation in the two ViT designs is identical, we find that this small change in early visual processing results in markedly different training behavior in terms of the sensitivity to optimization settings as well as the final model accuracy. Using a convolutional stem in ViT dramatically increases optimization stability and also improves peak performance (by ~1-2% top-1 accuracy on ImageNet-1k), while maintaining flops and runtime. The improvement can be observed across the wide spectrum of model complexities (from 1G to 36G flops) and dataset scales (from ImageNet-1k to ImageNet-21k). These findings lead us to recommend using a standard, lightweight convolutional stem for ViT models in this regime as a more robust architectural choice compared to the original ViT model design.

研究动机与目标

• 动机与理解为何 ViT 的可优化性落后于 CNN。
• 研究早期视觉处理（patchify stem）是否会在 ViT 中引发优化问题。
• 评估一个最小的卷积干作为可替换的 drop-in 方案，并将其训练动态与原始 ViT 干进行比较。
• 评估收敛速度、优化器兼容性、超参数稳定性以及跨不同模型规模和数据规模的峰值性能。

提出的方法

• 将 ViT 的 patchify stem（步幅为 stride-p，p×p）替换为大约 5 层的轻量级卷积干。
• 移除一个 Transformer 块，以保持 FLOPs 与运行时间的对齐。
• 在多种规模（1G–36G FLOPs）和 ImageNet 规模（1k–21k）上训练 ViT P（原始 patchify stem）和 ViT C（卷积干）。
• 通过误差分布函数（EDF）评估优化行为，包含训练长度稳定性、优化器稳定性（AdamW 与 SGD）以及超参数稳定性（学习率 lr，权重衰减 wd）。
• 在受控训练设置下，结合 EMA 权重和标准数据增强方案，比较峰值性能。

实验结果

研究问题

• RQ1 将 patchify stem 替换为卷积干是否会影响 ViT 的优化稳定性和收敛速度？
• RQ2 具备卷积干的 ViT 模型是否也能以 SGD 以及 AdamW 进行有效训练而不显著损失准确度？
• RQ3 与 patchify stem 相比，卷积干对学习率和权重衰减稳定性有何影响？
• RQ4 在不同模型复杂度和数据集规模下，对 peak ImageNet top-1 的影响是什么？
• RQ5 在标准训练方案下，卷积干是否足以使 ViT 与 CNN 竞争？

主要发现

• ViT C 在不同模型复杂度下收敛速度快于 ViT P。
• ViT C 使得既可用 AdamW 也可用 SGD 进行训练且不会产生显著的准确损失，与之相比 ViT P 在 SGD 下表现吃力。
• ViT C 对学习率和权重衰减更稳定，超参数敏感度接近 CNN 水平。
• 在 1G–36G FLOPs 与 ImageNet 1k–21k 上，ViT C 相对于 ViT P 在 ImageNet-1k 的 top-1 错误大约低 1–2 个百分点。
• 在公平的峰值性能比较中，ViT C 在中等复杂度范围内优于 ViT P 和具竞争力的 CNN；在 ImageNet-21k 预训练下，ViT C 超越所有。