[论文解读] AutoAssist: A Framework to Accelerate Training of Deep Neural Networks

Jiong Zhang, Hsiang‐Fu Yu+1 more|arXiv (Cornell University)|May 8, 2019

Advanced Neural Network Applications参考文献 31被引用 6

一句话总结

AutoAssist 通过引入一个轻量级的 Assistant 网络，根据梯度大小动态过滤低影响的训练样本，使主模型（Boss）专注于有信息量的样本，从而加速深度学习训练。该方法在图像分类任务中使 ResNet 减少 40% 的训练轮次，在机器翻译任务中使 Transformer 模型训练速度提升 30%，且通过异步 CPU/GPU 训练实现极低开销。

ABSTRACT

Deep neural networks have yielded superior performance in many applications; however, the gradient computation in a deep model with millions of instances lead to a lengthy training process even with modern GPU/TPU hardware acceleration. In this paper, we propose AutoAssist, a simple framework to accelerate training of a deep neural network. Typically, as the training procedure evolves, the amount of improvement in the current model by a stochastic gradient update on each instance varies dynamically. In AutoAssist, we utilize this fact and design a simple instance shrinking operation, which is used to filter out instances with relatively low marginal improvement to the current model; thus the computationally intensive gradient computations are performed on informative instances as much as possible. We prove that the proposed technique outperforms vanilla SGD with existing importance sampling approaches for linear SVM problems, and establish an O(1/k) convergence for strongly convex problems. In order to apply the proposed techniques to accelerate training of deep models, we propose to jointly train a very lightweight Assistant network in addition to the original deep network referred to as Boss. The Assistant network is designed to gauge the importance of a given instance with respect to the current Boss such that a shrinking operation can be applied in the batch generator. With careful design, we train the Boss and Assistant in a nonblocking and asynchronous fashion such that overhead is minimal. We demonstrate that AutoAssist reduces the number of epochs by 40% for training a ResNet to reach the same test accuracy on an image classification data set and saves 30% training time needed for a transformer model to yield the same BLEU scores on a translation dataset.

研究动机与目标

• 通过减少在低影响训练样本上的计算浪费，加速深度神经网络训练。
• 解决随机梯度下降（SGD）中均匀 mini-batch 采样带来的低效问题，即所有样本被同等对待，尽管其对模型性能提升的贡献各不相同。
• 设计一个轻量级、动态自适应的 Assistant 网络，实时评估样本重要性，无需事先了解数据难度。
• 通过在 CPU/GPU 上异步、非阻塞地训练 Boss 和 Assistant，最小化训练开销。
• 在强凸设置下，理论上保证收敛，即使存在因样本筛选导致的偏差梯度。

提出的方法

• 提出双模型框架：Boss（主深度网络）与轻量级 Assistant 网络，后者基于当前模型性能预测样本重要性。
• 使用梯度大小作为样本效用的代理指标，通过样本筛选操作过滤掉边际提升低的样本。
• 在 CPU 和 GPU 上异步训练 Assistant 和 Boss，以最小化计算开销。
• 在批量生成器中应用筛选操作，优先处理高梯度样本，从而提升每次 GPU/TPU 周期的收益。
• 将训练过程建模为带偏差的随机梯度方法，并证明在强凸问题下收敛速度为 $O(1/k)$。
• 采用受对偶坐标下降启发的机制指导样本选择，灵感来自支持向量机（SVM）优化技术。

实验结果

研究问题

• RQ1 基于梯度大小的样本级重要性过滤能否提升深度学习训练效率？
• RQ2 使用轻量级 Assistant 网络进行样本筛选后，其收敛速度是否与原始 SGD 相当？
• RQ3 Boss 与 Assistant 的异步训练能否在极低开销下实现显著加速？
• RQ4 与课程学习（curriculum learning）和自 paced learning 等现有方法相比，AutoAssist 在收敛速度和准确率方面表现如何？
• RQ5 在不牺牲模型性能的前提下，训练时间与训练轮次最多可减少多少？

主要发现

• 在图像分类任务中，AutoAssist 使 ResNet 的训练轮次减少 40%，在仅 13.7 轮内达到与基线模型 24 轮相同的测试准确率。
• 在 WMT14 英德机器翻译数据集上，AutoAssist 在训练 11.8 亿词后达到 BLEU 得分为 26，相比基线模型（需 18.2 亿词）节省了 35% 的训练词量。
• 在 Multi30K 数据集上，采用异步 CPU/GPU 训练的 AutoAssist 将时间开销从顺序训练的 61% 降低至 5%，同时保持更快的收敛速度和更高的 BLEU 得分。
• 在机器翻译任务中，AutoAssist 使 Transformer 模型的训练时间减少 30%，同时达到与基线相同的 BLEU 得分。
• 在 MNIST 数据集上，AutoAssist 仅用 13.7 轮即达到 99.17% 的测试准确率，而基线模型需 24 轮达到 98.67%，表明收敛性能更优。
• 异步训练流水线将开销降低至 10% 以下，使该框架可高效集成至现有深度学习工作流中。