[论文解读] Large batch size training of neural networks with adversarial training and second-order information

Zhewei Yao, Amir Gholami+5 more|arXiv (Cornell University)|Oct 2, 2018

Advanced Neural Network Applications参考文献 54被引用 35

一句话总结

该论文提出了一种基于二阶 Hessian 信息和对抗训练的自适应批量大小训练框架（ABSA），以加速大规模批量神经网络的训练。通过基于 Ray 的弹性聚类动态调整批量大小，并利用曲率感知的学习率调整，ABSA 在 ImageNet 上将 SGD 迭代次数减少最多 5 倍，训练时间减少 8.78 倍，实现了最先进的加速效果，且超参数调优极少。

ABSTRACT

The most straightforward method to accelerate Stochastic Gradient Descent (SGD) computation is to distribute the randomly selected batch of inputs over multiple processors. To keep the distributed processors fully utilized requires commensurately growing the batch size. However, large batch training often leads to poorer generalization. A recently proposed solution for this problem is to use adaptive batch sizes in SGD. In this case, one starts with a small number of processes and scales the processes as training progresses. Two major challenges with this approach are (i) that dynamically resizing the cluster can add non-trivial overhead, in part since it is currently not supported, and (ii) that the overall speed up is limited by the initial phase with smaller batches. In this work, we address both challenges by developing a new adaptive batch size framework, with autoscaling based on the Ray framework. This allows very efficient elastic scaling with negligible resizing overhead (0.32\% of time for ResNet18 ImageNet training). Furthermore, we propose a new adaptive batch size training scheme using second order methods and adversarial training. These enable increasing batch sizes earlier during training, which leads to better training time. We extensively evaluate our method on Cifar-10/100, SVHN, TinyImageNet, and ImageNet datasets, using multiple neural networks, including ResNets and smaller networks such as SqueezeNext. Our method exceeds the performance of existing solutions in terms of both accuracy and the number of SGD iterations (up to 1\% and $5 imes$, respectively). Importantly, this is achieved without any additional hyper-parameter tuning to tailor our method in any of these experiments.

研究动机与目标

• 解决因损失曲面陡峭而导致的大批量 SGD 训练中泛化性能差和加速有限的问题。
• 克服自适应批量大小训练中初始小批量阶段的瓶颈，尽管具备弹性扩展能力，但整体加速仍受限制。
• 开发一种可扩展、低开销的框架，用于动态批量大小和聚簇大小的调整，以实现高效的大型批量训练。
• 整合二阶 Hessian 信息与对抗训练，以在训练早期更早、更安全地增加批量大小。
• 在多种数据集和网络架构上实现更快收敛与更高准确率，且超参数调优极少。

提出的方法

• 提出一种自适应批量大小（ABS）方法，利用 Hessian 特征值谱分析，根据损失曲面曲率动态调整批量大小和学习率。
• 通过最小-最大优化将 ABS 与对抗鲁棒性相结合，提出自适应批量大小对抗训练（ABSA），隐式正则化以避免尖锐极小值。
• 实现基于 Ray 的分布式框架，支持弹性扩展，聚簇大小调整的开销接近于零（仅 0.32%），显著降低通信与重启成本。
• 在训练初期使用梯度累积，仅需 1–2 个 GPU 高效计算 Hessian 信息，将 Hessian 计算成本降至总训练时间的 9.3%。
• 采用预热阶段，逐步增加批量大小和学习率，仅在最终实验中调优，进一步加速收敛。
• 利用对抗训练带来的隐式正则化，在早期增加批量大小时稳定训练，提升泛化性能与鲁棒性。

实验结果

研究问题

• RQ1 二阶 Hessian 信息是否能实现训练过程中更早、更安全的批量大小增加，从而减少初始小批量阶段的瓶颈？
• RQ2 将基于 Hessian 的自适应批量大小与对抗训练相结合，是否能提升大规模批量训练中的泛化性能与鲁棒性？
• RQ3 基于 Ray 的框架是否能实现在训练过程中弹性、低开销的聚簇大小调整，从而消除先前方法的通信与重启开销？
• RQ4 在不牺牲准确率的前提下，尤其是扩展到大规模批量时，SGD 迭代次数最多可减少多少？
• RQ5 所提方法是否能在多种数据集与网络架构上实现最先进的加速效果，且仅需极少的超参数调优？

主要发现

• ABSA 在 CIFAR-10 上将 SGD 迭代次数减少最多 5 倍，在 ImageNet（ResNet18）上减少 28.8 倍，仅用 14.8K 次更新即达到 70.04% 的 ImageNet 顶部-1 准确率。
• 仅通过调优预热阶段（超参数调优极少），ABSA 实现了 ResNet18 在 ImageNet 上 8.78 倍的加速，训练时间从 125K 秒减少至 14.2K 秒。
• 基于 Ray 的框架在聚簇大小调整时仅引入 0.32% 的开销，实现了高效弹性扩展，性能损失可忽略不计。
• Hessian 计算仅占总训练时间的 9.3%，主要瓶颈在于早期阶段使用少量 GPU 进行梯度累积。
• ABSA 在 CIFAR-10 上使用 ResNet18 与 90 个周期、最终批量大小为 16K 时，达到 84.24% 的准确率，优于标准大批次 SGD（76.82%）和小批次 SGD（83.05%）。
• 该方法在多种架构（ResNet、SqueezeNext、AlexNet）和数据集（CIFAR-10/100、SVHN、TinyImageNet、ImageNet）上均具有一致优势，准确率与迭代效率全面超越基线方法。