[论文解读] Interlocking Backpropagation: Improving depthwise model-parallelism

Aidan N. Gomez, Oscar Key+5 more|arXiv (Cornell University)|Oct 8, 2020

Advanced Neural Network Applications参考文献 8被引用 5

一句话总结

本文提出了一种称为互锁反向传播的混合训练策略，通过在深度神经网络中定期插入辅助分类头，改进深度模型并行性，限制梯度流以提升计算效率。该方法在固定时间约束下，训练速度显著快于端到端反向传播，同时恢复了纯局部学习中损失的大部分性能，在ResNets和Transformer模型上均实现了更高的训练效率和测试准确率。

ABSTRACT

The number of parameters in state of the art neural networks has drastically increased in recent years. This surge of interest in large scale neural networks has motivated the development of new distributed training strategies enabling such models. One such strategy is model-parallel distributed training. Unfortunately, model-parallelism can suffer from poor resource utilisation, which leads to wasted resources. In this work, we improve upon recent developments in an idealised model-parallel optimisation setting: local learning. Motivated by poor resource utilisation in the global setting and poor task performance in the local setting, we introduce a class of intermediary strategies between local and global learning referred to as interlocking backpropagation. These strategies preserve many of the compute-efficiency advantages of local optimisation, while recovering much of the task performance achieved by global optimisation. We assess our strategies on both image classification ResNets and Transformer language models, finding that our strategy consistently out-performs local learning in terms of task performance, and out-performs global learning in training efficiency.

研究动机与目标

• 解决全局模型并行训练中资源利用率低下以及局部学习导致性能下降的问题。
• 开发一种介于局部与全局反向传播之间的中间策略，以保持计算效率并提升任务性能。
• 评估大规模分布式训练中训练速度与模型准确率之间的权衡。
• 提供一个通用的开源框架，用于研究互锁反向传播策略。

提出的方法

• 提出n路互锁反向传播，即在深度神经网络中定期插入辅助分类头，以生成局部训练信号。
• 限制模块之间的梯度流，防止完全的反向传播依赖，从而实现更早的权重更新。
• 利用辅助头的局部损失独立训练每个模块，减少空闲时间和通信瓶颈。
• 采用固定时间训练视角评估性能，优先选择步长时间更短的策略。
• 将该方法应用于ResNets和基于Transformer的语言模型，使用Adam优化器训练并测量困惑度。
• 采用微批次处理和基于TPU的分布式训练，在多个加速器上扩展实验规模。

实验结果

研究问题

• RQ1 在深度模型并行性中，介于局部学习与全局学习之间的混合训练策略是否能同时提升训练效率和模型性能？
• RQ2 在训练速度和测试准确率方面，互锁反向传播与端到端训练及1路局部学习相比表现如何？
• RQ3 辅助分类头在多大程度上能恢复纯局部训练中损失的性能？
• RQ4 在固定时间约束下，互锁反向传播能否有效扩展到如Transformer等大模型？

主要发现

• 2路互锁反向传播将训练步长时间减少至标准端到端训练的一半以下，同时在4模块Transformer模型上实现了相似的测试困惑度。
• 在固定时间约束下，互锁反向传播在模型性能上优于端到端训练，因为它能在相同时间内完成更多次权重更新。
• 在CIFAR-10和CIFAR-100数据集上，2路互锁反向传播在测试准确率上均优于端到端和1路训练，尽管训练速度更快。
• 对于Transformer语言模型，互锁反向传播显著缩小了1路与端到端训练之间的性能差距，尽管其困惑度仍未超越端到端训练。
• 该方法表明，调节n路互锁中的n参数，可在计算效率与建模性能之间实现平滑的权衡。
• 位于 https://github.com/oscarkey/interlocking-backprop 的开源框架支持对这一类优化算法的可复现研究。