[论文解读] A Closer Look at Codistillation for Distributed Training

Shagun Sodhani, Olivier Delalleau+4 more|arXiv (Cornell University)|Oct 6, 2020

Topic Modeling参考文献 20被引用 4

一句话总结

本文提出了一种名为 co-distillation 的通信高效替代方案，用于标准的数据并行训练。在 co-distillation 中，多个模型并行训练，并通过知识蒸馏损失而非参数平均来实现同步。实验表明，2路 co-distillation 在使用 2 倍设备的情况下，可达到与标准数据并行训练相当的准确率，同时将通信量减少高达 1000 倍；然而，当扩展到 n>2 个模型时，其性能受到模型和数据集特性的限制。

ABSTRACT

Codistillation has been proposed as a mechanism to share knowledge among concurrently trained models by encouraging them to represent the same function through an auxiliary loss. This contrasts with the more commonly used fully-synchronous data-parallel stochastic gradient descent methods, where different model replicas average their gradients (or parameters) at every iteration and thus maintain identical parameters. We investigate codistillation in a distributed training setup, complementing previous work which focused on extremely large batch sizes. Surprisingly, we find that even at moderate batch sizes, models trained with codistillation can perform as well as models trained with synchronous data-parallel methods, despite using a much weaker synchronization mechanism. These findings hold across a range of batch sizes and learning rate schedules, as well as different kinds of models and datasets. Obtaining this level of accuracy, however, requires properly accounting for the regularization effect of codistillation, which we highlight through several empirical observations. Overall, this work contributes to a better understanding of codistillation and how to best take advantage of it in a distributed computing environment.

研究动机与目标

• 探究 co-distillation 是否可作为可扩展的、通信高效的替代方案，用于标准的数据并行训练。
• 评估 2 路 co-distillation 相较于基于 all-reduce 的传统数据并行训练在性能和通信效率方面的表现。
• 理解将 co-distillation 扩展至超过 2 个模型（n>2）在分布式训练中所面临的局限性和挑战。
• 分析 co-distillation 如何影响训练动态，包括其类似正则化的效应。
• 探索 n 路 co-distillation 在何种条件下可优于 2 路 co-distillation 或标准训练。

提出的方法

• Co-distillation 并行训练多个模型，每个模型在数据子集上训练，并通过惩罚模型间预测不一致性的蒸馏损失进行优化。
• 与通过 all-reduce 同步模型参数的方式不同，模型通过知识蒸馏实现功能上的对齐。
• 训练目标结合了标准交叉熵损失与对称蒸馏损失，以促使模型输出趋于一致。
• 通过仅传输蒸馏后的预测结果或软标签，而非完整梯度或参数，从而降低通信开销。
• 该方法在 ResNet50（ImageNet）和基于 Transformer 的机器翻译模型上进行了评估，并对 n 路 co-distillation 进行了消融研究。
• 超参数 α 控制蒸馏损失的权重，并经过调优以获得最佳性能。

实验结果

研究问题

• RQ1 2 路 co-distillation 是否能在将通信量降低数个数量级的同时，实现与标准数据并行训练相当的准确率？
• RQ2 随着每个模型的设备数量增加，co-distillation 的扩展性如何？是否能保持线性加速？
• RQ3 为何将 co-distillation 扩展至超过两个模型（n>2）时，无法在不同工作负载中一致地提升性能？
• RQ4 模型架构和数据集特性在 n 路 co-distillation 的有效性中起到何种作用？
• RQ5 co-distillation 如何影响训练动态？其引入了何种正则化效应？

主要发现

• 当使用 2 倍设备时，2 路 co-distillation 在每轮迭代中可将通信量减少高达 1000 倍，同时达到与标准数据并行训练相当的准确率。
• 该方法在每个模型的设备数量上表现出线性扩展性：当设备数量加倍时，收敛所需的训练更新次数可减半。
• 将 co-distillation 扩展至 n>2 个模型时，结果参差不齐：某些情况下性能有所提升，但并不一致，且通常不优于 2 路 co-distillation。
• 在超过 2 个模型时扩展失败的主要原因与模型架构和数据集特定特性（如泛化行为和损失曲面）密切相关。
• Co-distillation 引入了类似正则化的效应，影响了训练动态，必须在调优中加以考虑以实现最佳性能。
• 尽管通信量减少，但端到端加速并非必然，其效果取决于模型大小和通信带宽，其中大模型受益更明显。