[论文解读] How to train your ViT? Data, Augmentation, and Regularization in Vision Transformers

Andreas Steiner, А. И. Колесников+4 more|arXiv (Cornell University)|Jun 18, 2021

Advanced Neural Network Applications被引用 228

一句话总结

该论文系统地研究数据规模、增强、正则化、模型大小和计算预算如何相互作用以影响 Vision Transformers，显示 AugReg 可以替代大数据量，并且从更大上游数据集进行迁移学习通常在成本上更具性价比。

ABSTRACT

Vision Transformers (ViT) have been shown to attain highly competitive performance for a wide range of vision applications, such as image classification, object detection and semantic image segmentation. In comparison to convolutional neural networks, the Vision Transformer's weaker inductive bias is generally found to cause an increased reliance on model regularization or data augmentation ("AugReg" for short) when training on smaller training datasets. We conduct a systematic empirical study in order to better understand the interplay between the amount of training data, AugReg, model size and compute budget. As one result of this study we find that the combination of increased compute and AugReg can yield models with the same performance as models trained on an order of magnitude more training data: we train ViT models of various sizes on the public ImageNet-21k dataset which either match or outperform their counterparts trained on the larger, but not publicly available JFT-300M dataset.

研究动机与目标

• 理解训练数据量、增强和正则化如何在不同模型大小的 ViTs 中相互作用
• 量化对 ViTs 的预训练和迁移到下游任务的计算与数据权衡
• 在多样化数据集上评估迁移学习与从头训练结合 AugReg 的效果
• 为在计算约束下选择预训练检查点与训练策略提供实际建议

提出的方法

• 在 Unified 设置下，在 TPU/JAX/Flax 上对 ImageNet-1k 和 ImageNet-21k 预训练多种尺寸的 ViT 模型（Ti、S、B、L）以及 ResNet+ViT 混合模型
• 应用正则化（Dropout、随机深度）和增强（Mixup、RandAugment），并进行 28 配置的超参数扫描
• 在下游数据集上使用 SGD 与余弦调度进行微调，比较从头训练与从上游模型迁移
• 系统性地改变数据规模、计算预算和增强强度以绘制性能景观

实验结果

研究问题

• RQ1 数据规模、增强和正则化如何相互作用以影响不同模型大小的 ViTs 的性能？
• RQ2 从更大上游数据集进行迁移学习是否在成本上更具性价比于在 AugReg 条件下从头训练以应对常见下游任务？
• RQ3 增加预训练数据与增强对跨 VTAB 任务的迁移性能有何影响？
• RQ4 在计算约束下，选取预训练检查点和训练策略可获得哪些实际指导？

主要发现

• AugReg 可以在可比任务上达到相当于训练数据量提升 10 倍的收益
• 对于实际数据集，使用预训练模型进行迁移学习通常在成本上更高效且效果更好，而非从头训练
• 更多的预训练数据能产生更通用的模型，从而在多样化的 VTAB 任务上实现更好的迁移
• 在中等规模的数据集上，增强效果更明显；但在非常大的上游数据（如 ImageNet-21k，受限计算）下，若不增加计算，AugReg 往往会带来负面影响
• 通过上游验证准确度来选择上游模型通常对迁移有效；推荐使用 ImageNet-21k 的检查点
• 将补丁大小增大到 32 可以超越同速率的更薄模型，表明补丁大小是一个超越参数量级之外的关键容量因素。