NeurIPS 2022 | 四分钟内就能训练指标检测器，商汤基模型团队是怎么做到的？

本文提出了一种大批量训练算法 AGVM (Adaptive Gradient Variance Modulator)，不仅可以适配于指标检测恣意，同时也可以适配各类朋分恣意。AGVM 可以把指标检测的训练批量巨细扩大到 1536，帮助研讨职员四分钟训练 Faster R-CNN，3.5 小时把 COCO 刷到 62.2 mAP，均打破了指标检测训练速度的世界纪录。

论文地址：https://arxiv.org/pdf/2210.11078.pdf

代码地址：https://github.com/Sense-X/AGVM

在当前的机器学习社区中，有三个普遍的趋势。首先，神经收集模型会愈来愈大。在 NLP 领域中最大规模的模型已经达到了上万亿级别。在视觉领域，最大规模的模型也达到了三百亿的量级。其次，训练的数据集也变得愈来愈大。比如，ImageNet 21k 和谷歌的 JFT 数据集都具有相当规模的数据集。另外，由于数据集变得愈来愈大，训练 SOTA 模型的开销愈来愈大。

因此，提升训练效率就变得愈发重要。而分布式训练因为其适应于数据并行、模型并行和流水线并行的加速训练步骤的同时，也具备较高的 Deep Learning 通信效率而被广泛认为是一个有效的解决方案。

随着大模型时代的到来，指标检测器的训练速度愈来愈成为学术界和工业界的瓶颈，例如，在 COCO 的标准 setting 上把 mAP 训到 62 以上大概需要三天的时间，算上调试成本，这在业界几乎是不可接受的。那么，我们能不能把这个训练时间压到小时级别呢？事实上，在图片分类和自然语言处理恣意上，先前的研讨职员借助 32K 的批量巨细(batch size)，只需 14 分钟就可以完成 ImageNet 的训练，76 分钟完成 Bert 的训练。但是，在指标检测领域，还很欠缺这类研讨，导致研讨职员无法充分利用当前的算力，数据集和大模型。

大批量训练算法 AGVM 便是这个问题的最佳解决方案之一。为了支持如此大批量的训练，同时保持模型的训练精度，本研讨提出了一套全新的训练算法，根据麋集猜测不同模块的梯度方差(gradient variance)，动态调整每一个模块的学习率。作者在大量的麋集猜测收集和数据集上进行了试验，并且证实了该步骤的合理性。

步骤介绍

大批量训练是加速大型分布式系统中深度神经收集训练的关键。尤其是在如今的大模型时代，如果不采用大批量训练，一个收集的训练时间几乎是难以接受的。但是，大批量训练很难，因为它会产生泛化差距(generalization gap)，直接训练会导致其准确率降低。此前的大批量工作往往针对于图像分类以及一些自然语言处理的恣意，但麋集猜测恣意（包括检测朋分等），同样在视觉中处于举足轻重的位置，此前的步骤并不能在麋集猜测恣意上有很好的表现，甚至结果比基准线更差，这导致我们难以快速训练一个指标检测器。

为了解决这个问题，研讨职员进行了大量的试验。最后发现，相较于传统的分类收集，利用麋集猜测收集一个很重要的特性：麋集猜测收集往往是由多个组件组成的，以 Faster R-CNN 为例：它由四个部分组成，主干收集 (Backbone)，特性金字塔收集(FPN)，区域生成收集(RPN) 和检测头收集(head)，我们可以发现一个很有效的指标：麋集猜测收集不同组件的梯度方差，在训练批量很小时（例如 32），几乎是相同的，但当训练批量很大时（例如 512），它们呈现出很大的区别，如下图所示：

那么，能不能直接把这些拉平呢？这直接引出了 AGVM 算法。以随机梯度下降算法为例，上角标 i 代表第 i 个收集模块（例如 FPN 等），上角标 1 代表主干收集，代表学习率，锚定主干收集，可以直接将不同收集组件的梯度 g 的方差：

梯度的方差可以由以下式子估计：

方差的具体求解细节可以参考原文，本研讨同样引入了滑动平均机制，防止收集训练发散。同时，研讨证明了 AGVM 在非凸情况下的收敛性，讨论了动量以及衰减的处理方式，具体实现细节可以参考原文。

试验过程

本研讨首先在指标检测、实例朋分、全景朋分和语义朋分的各种麋集猜测收集上进行了测试，通过下表可以看到，当用标准批量巨细训练时，AGVM 相较传统步骤没有明显优势，但当在超大批量下训练时，AGVM 相较传统步骤拥有压倒性的优势，下图第二列从左至右分别表示指标检测，实例朋分，全景朋分和语义朋分的表现，AGVM 超越了有史以来的所有步骤：

下表详细对比了 AGVM 和传统步骤，体现出了本研讨步骤的优势：

同时，为了说明 AGVM 的优越性，本研讨进行了以下三个超大规模的试验。研讨职员把 Faster R-CNN 的 batch size 放到了 1536，这样利用 768 张 A100 可以在 4.2 分钟内完成训练。其次，借助 UniNet-G，本研讨可以在利用 480 张 A100 的情况下，3.5 个小时让模型在 COCO 上达到 62.2mAP（不包括主干收集预训练的时间），极大的减小了训练时间：

甚至，在 RetinaNet 上，本研讨把批量巨细扩展到 10K。这在指标检测领域是从未见的批量巨细，在如此大的批量下，每一个 epoch 只有十几个迭代次数，AGVM 在如此大的批量下，仍然能展现出很强的稳定性，性能如下图所示：

结果分析

本研讨探究了一个很重要的问题：以 RetinaNet 为例，如下图第一列所示，探究为什么会出现梯度方差不匹配这一现象。

本研讨认为，这一现象来自于：收集不同模块间的有效批量巨细 (effective batch size) 是不同的。例如，RetinaNet 的头收集的输入是由特性金字塔的五层收集输出的，特性金字塔的 top-down 和 bottom-up pathways，以及像素维度的损失函数计算会导致头收集和主干收集的等效批量巨细不同，这一原理导致了梯度方差不匹配的现象。

为了验证这一假设，本研讨依次给每一层特性使用单独的头收集，移去特性金字塔收集，随机忽略掉 75% 的用于计算损失函数的像素，最终，本研讨发现主干收集和头收集的梯度方差曲线重合了，本研讨也对 Faster R-CNN 做了类似的试验，如下图第二列所示，更多的讨论请参见原文。