详解超强 ResNet 变体 NFNet(二):匹敌 ViT 性能的 JFT-4B 大规模预训练
极市导读
决定一个视觉骨干模型性能的最重要因素是用于训练的计算预算和训练的数据量。 >>加入极市CV技术交流群,走在计算机视觉的最前沿
本文目录
1 NFNet 视觉大模型:匹敌 ViT 性能的大规模预训练
(来自 Google DeepMind)
1.1 背景:视觉社区从中等规模数据集到超大规模数据集的性能评估
1.2 JFT-4B 预训练的 NFNet 服从对数-对数缩放定律
1.3 ImageNet 微调之后的 NFNet 匹配了 Vision Transformer 的性能
太长不看版
许多研究人员认为 ConvNet 在小型或中等规模的数据集上表现良好,但是在面对超大规模的数据集时不如 Vision Transformer。在本文中作者希望挑战这个观点,通过评估在 JFT-4B (大型有标注数据集,通常用于训练基础模型) 上预训练的 ConvNet。
选取的 ConvNet 模型:NFNet (解读详见如下链接)。NFNet 由于不含 BN 带来的隐式正则化效果,除非显式正则化,否则对于像 ImageNet 这样的数据集往往会过拟合。但是当在极大规模的数据集 (比如本文使用的 JFT-4B) 上进行预训练时,这种正则化可能不仅是不必要的,反而会对性能有损害,因为妨碍了模型将其全部的参数应用于拟合训练集的能力。因此,不含 BN 的 NFNet 天然地适合于超大规模数据集的预训练。
详解超强 ResNet 变体 NFNet:抛弃归一化后,性能却达到了最强!
因此,作者在本文中选取 NFNet 模型家族训练了一系列增加深度和宽度逐渐增加的模型,还观察到 loss 和 compute budget 之间的对数-对数缩放定律 (Log-Log Scaling Law)。在 ImageNet 上进行微调后,NFNet 与计算预算相当的 Vision Transformer 的精度相匹配,NFNet 最强的微调之后模型达到了 90.4% 的 Top-1 精度。
1 NFNet 视觉大模型:匹敌 ViT 性能的大规模预训练
论文名称: ConvNets Match Vision Transformers at Scale (Arxiv 2023)
论文地址:
http://arxiv.org/pdf/2310.16764.pdf
1.1 背景:视觉社区从中等规模数据集到超大规模数据集的性能评估
卷积神经网络 (ConvNet) 对于深度学习的许多早期成功有巨大的推动作用。Deep ConvNet 于 20 年前首次商业部署,而 AlexNet 在 2012 年 ImageNet 挑战方面取得了成功。近十年的 ConvNet 主导了计算机视觉基准。但是近年来,它们越来越多地被 Vision Transformer (ViT) 所取代。
同时,计算机视觉社区在评估模型性能时,已经从中等规模数据集 (如 ImageNet) 转向了从网上收集的通用超大规模数据集 (如 JFT-4B) 预训练网络的性能。这就提出了一个重要问题:Vision Transformer 是否优于使用相似计算资源预训练的 ConvNet 架构?
尽管社区中大多数研究人员认为 Vision Transformer 表现出了比 ConvNet 更好的缩放属性,但是几乎没有证据来证明这一点。许多的研究将 ViT 与较弱的 ConvNet 基线 (通常是原始的 ResNet 架构) 进行比较。而最强的 ViT 模型已经使用了超过 500k TPU-v3 core hours 的计算资源,这已经大大超过了用于预训练 ConvNet 的计算预算。
许多研究人员认为 ConvNet 在小型或中等规模的数据集上表现良好,但是在面对超大规模的数据集时不如 Vision Transformer。在本文中作者希望挑战这个观点,通过评估在 JFT-4B (大型有标注数据集,通常用于训练基础模型) 上预训练的 ConvNet。
作者评估的 ConvNet 模型是 NFNet 模型家族,一种与第一个ViT论文同时发布的纯卷积架构。NFNet 由于不含 BN 带来的隐式正则化效果,除非显式正则化,否则对于像 ImageNet 这样的数据集往往会过拟合。但是当在极大规模的数据集 (比如本文使用的 JFT-4B) 上进行预训练时,这种正则化可能不仅是不必要的,反而会对性能有损害,因为妨碍了模型将其全部的参数应用于拟合训练集的能力。因此,不含 BN 的 NFNet 天然地适合于超大规模数据集的预训练。
作者没对模型架构或训练过程进行任何更改 (除了调整学习率或者 Epoch 数等简单的超参数)。作者的计算预算最多为 110k TPU-v4 core hours,并在 JFT-4B 数据集上预训练,该数据集包含来自 30k 类别的大约 4 billion 的有标注图片。
1.2 JFT-4B 预训练的 NFNet 服从对数-对数缩放定律
作者在 JFT-4B 上训练一系列不同深度的 NFNet 模型。每个模型都使用余弦衰减学习率计划训练,Epoch 预算在了 0.25 到 8 之间,基本学习率针对每个 Epoch 预算作单独调整。作者在图1中绘制了 JFT-4B 的验证集损失与计算预算之间的关系曲线。
模型: F7 的宽度与 F3 相同,但深度加倍。同样,F3 是 F1 深度的两倍,F1 是 F0 深度的两倍。F3+ 和 F7+ 的深度与 F3 和 F7 相同,但宽度较大。
训练策略: 使用批量大小为 4096 的 Momentum 和 Adaptive Gradient Clipping (AGC) 的 SGD 优化器进行训练,在训练期间使用 224×224 的图像分辨率,在推理时使用 256 × 256。在训练期间从 JFT-4B 中删除了 ImageNet 训练集和验证集中图像的近乎重复的部分。
作者图1中展示出的趋势是清晰的 TPU-v4 Core Hours 的对数和 JFT-4B 验证集精度之间的线性关系,也就是 TPU-v4 Core Hours 和 JFT-4B 验证集精度之间的对数-对数缩放定律 (Log-Log Scaling Law)。
最佳模型大小和最佳 Epoch 预算 (达到最低验证集损失) 都随着计算预算的增加而增加。作者注意到对于大于约 5k TPU-v4 Core Hours 的整体计算预算而言,最佳的 Epoch 预算大于 1。
如图2所示是不同尺寸的3个模型在不同 Epoch 预算的最优学习率 (最小化验证损失)。所有的模型都显示出相似的最佳学习率:当 Epoch 预算较小时最佳学习率约为1.6,随着计算预算的增加而逐渐减小。对于大模型,最优的学习率下降得更快。在实践中,可以通过假设最优学习率随着模型大小和 Epoch 预算的增加而缓慢单调地下降,从而在2次试验内有效地调整学习率。
作者注意到图1中的一些预训练模型的表现不如预期好,比如,不同预训练预算下的 NFNet-F7+ 模型的曲线不够平滑。作者认为这是因为如果训练运行被抢占或者重启,数据加载管道并不能保证每个训练样本在每个 Epoch 中被采样一次,如果多次重新启动训练运行,可能会导致一些训练样本被欠采样。
1.3 ImageNet 微调之后的 NFNet 匹配了 Vision Transformer 的性能
在下图3中,作者在 ImageNet 上微调 JFT-4B 上预训练的 NFNets,并根据预训练期间使用的计算预算绘制 Top-1 Error。正则化方法包括 Stochastic Depth,Dropout,和 Sharpness Aware Minimization (SAM)[1],对每个模型进行微调 50 个 Epoch,作者以 384×384 的分辨率进行训练,并在 480×480 上进行评估。
如图3所示,随着计算预算的增加,ImageNet Top-1 的精度不断提高。最昂贵的预训练模型,一个经过 8 个 Epoch 预训练的 NFNet-F7+,实现了 90.3% 的 ImageNet Top-1 精度,同时需要大约 110k TPU-v4 Core Hour 来预训练和 1.6k TPU-v4 Core Hour 进行微调。此外,如果在微调期间额外引入 Repeated Augmentation,就实现了 90.4% 的 Top-1 的精度。
为了观察和比较大规模 JFT-4B 预训练的作用,作者列出了这个结果:在没有额外数据的情况下,NFNet 在 ImageNet 上报告的最佳 Top-1 精度仅为 86.8%,这表明 NFNet 从 JFT-4B 的预训练中受益匪浅。
作者也在 TPU-v4 上评估了这些模型的预训练速度,并估计:
ViT-g/14 需要 120k TPU-v4 Core Hour
ViT-G/14 需要 280k TPU-v4 Core Hour
SoViT-400m/14 需要 130k TPU-v4 Core Hour
注意到 NFNets 是针对 TPU-v4 优化的,估计 NFNet-F7+ 将需要 250 个 TPU-v3 Core Hour 预训练 8 个 Epoch。
最后,作者注意到在 JFT-4B 上实现最低验证集损失的预训练 checkpoint 在微调后并不总是在 ImageNet 上实现最高的 Top-1 精度的。特别是作者发现,在固定的预训练计算预算下,微调总偏爱稍大的模型和略小的 Epoch 预算。
最后讨论
本文最终的结论是:决定一个视觉骨干模型性能的最重要因素是用于训练的计算预算和训练的数据量。尽管 ViT 在视觉任务上的成功让人印象深刻,但是在本文看来,没有强有力的证据表明预训练的视觉 Transformer 模型的性能在计算预算相当的条件下优于预训练的 ConvNet 模型。
但不可否认的是,Transformer 模型在特定的上下文中具有实际优势,比如能够在多种模态中使用相似的组件的能力。
参考
-
^Sharpness-Aware Minimization for Efficiently Improving Generalization
公众号后台回复“极市直播”获取100+期极市技术直播回放+PPT
极市干货
# 极市平台签约作者#
科技猛兽
知乎:科技猛兽
清华大学自动化系19级硕士
研究领域:AI边缘计算 (Efficient AI with Tiny Resource):专注模型压缩,搜索,量化,加速,加法网络,以及它们与其他任务的结合,更好地服务于端侧设备。
作品精选