用于 MoE 训练的稀疏反向传播
机器学习
2023-10-03 v1 人工智能
计算与语言
计算机视觉与模式识别
摘要
混合专家（MoE）模型的一个决定性特征是其通过专家路由进行稀疏计算的能力，从而带来卓越的可扩展性。然而，作为深度学习基石的反向传播需要稠密计算，因此在 MoE 梯度计算中带来挑战。在此，我们引入 SparseMixer，一种可扩展的梯度估计器，弥合了反向传播与稀疏专家路由之间的差距。与典型的 MoE 训练为稀疏计算与可扩展性而策略性地忽略某些梯度项不同，SparseMixer 为这些项提供可扩展的梯度近似，从而实现 MoE 训练中的可靠梯度估计。基于数值 ODE 框架，SparseMixer 利用中点法（一种二阶 ODE 求解器）以可忽略的计算开销提供精确的梯度近似。将 SparseMixer 应用于 Switch Transformer 的预训练与机器翻译任务，SparseMixer 展现出可观的性能提升，将训练收敛加速至多 2 倍。
引用
@article{arxiv.2310.00811,
title = {Sparse Backpropagation for MoE Training},
author = {Liyuan Liu and Jianfeng Gao and Weizhu Chen},
journal= {arXiv preprint arXiv:2310.00811},
year = {2023}
}
备注
Work in progress