MKOR：基于 Kronecker 因子、使用秩-1 更新的动量赋能优化器
机器学习
2024-01-31 v2 人工智能
计算机视觉与模式识别
最优化与控制
摘要
本工作提出一种基于 Kronecker 因子、使用秩-1 更新的动量赋能优化器，称为 MKOR，其改善了深度神经网络（DNNs）的训练时间和收敛特性。二阶技术虽然比一阶方法享有更高的收敛速率，但相对于模型大小和/或训练批大小具有三次复杂度。因此它们在 transformer 模型（例如大语言模型（LLMs））中表现出较差的可扩展性和性能，因为这些模型中的批大小随注意力机制序列长度而缩放，导致较大的模型大小和批大小。MKOR 的复杂度相对于模型大小是二次的，缓解了二阶方法中的计算瓶颈。由于其高计算复杂度，二阶方法的最先进实现只能不频繁地更新二阶信息，因此未能充分利用这些更新更好收敛的潜力。通过降低二阶更新的通信复杂度并实现线性通信复杂度，MKOR 增加了二阶更新的频率。我们还提出了 MKOR 的混合版本（称为 MKOR-H），其在训练中期若二阶更新不再加速收敛则回退到一阶优化器。我们的实验表明，在 64 个 GPU 上的 BERT-Large-Uncased 上，MKOR 分别比最先进的一阶方法（例如 LAMB 优化器）和最佳二阶方法实现（即 KAISA/KFAC）高出至多 2.57 倍和 1.85 倍。
引用
@article{arxiv.2306.01685,
title = {MKOR: Momentum-Enabled Kronecker-Factor-Based Optimizer Using Rank-1 Updates},
author = {Mohammad Mozaffari and Sikan Li and Zhao Zhang and Maryam Mehri Dehnavi},
journal= {arXiv preprint arXiv:2306.01685},
year = {2024}
}
备注
Published at 37th Conference on Neural Information Processing Systems (NeurIPS 2023)