This work proposes a Momentum-Enabled Kronecker-Factor-Based Optimizer Using Rank-1 updates, called MKOR, that improves the training time and convergence properties of deep neural networks (DNNs). Second-order techniques, while enjoying higher convergence rates vs first-order counterparts, have cubic complexity with respect to either the model size and/or the training batch size. Hence they exhibit poor scalability and performance in transformer models, e.g. large language models (LLMs), because the batch sizes in these models scale by the attention mechanism sequence length, leading to large model size and batch sizes. MKOR's complexity is quadratic with respect to the model size, alleviating the computation bottlenecks in second-order methods. Because of their high computation complexity, state-of-the-art implementations of second-order methods can only afford to update the second order information infrequently, and thus do not fully exploit the promise of better convergence from these updates. By reducing the communication complexity of the second-order updates as well as achieving a linear communication complexity, MKOR increases the frequency of second order updates. We also propose a hybrid version of MKOR (called MKOR-H) that mid-training falls backs to a first order optimizer if the second order updates no longer accelerate convergence. Our experiments show that MKOR outperforms state -of-the-art first order methods, e.g. the LAMB optimizer, and best implementations of second-order methods, i.e. KAISA/KFAC, up to 2.57x and 1.85x respectively on BERT-Large-Uncased on 64 GPUs.

翻译：本文提出一种基于动量与秩1更新的Kronecker因子优化器（MKOR），旨在提升深度神经网络（DNNs）的训练速度与收敛性能。二阶优化方法虽具有比一阶方法更高的收敛率，但其计算复杂度随模型规模和/或训练批次大小呈三次方增长。因此，在Transformer模型（如大语言模型LLMs）中，由于批次大小随注意力机制序列长度扩展，导致模型规模与批次大小均显著增大，二阶方法存在可扩展性差与性能低下的问题。MKOR的复杂度与模型规模呈二次关系，有效缓解了二阶方法的计算瓶颈。受限于高计算复杂度，当前最先进的二阶方法实现仅能低频更新二阶信息，无法充分发挥其收敛优势。通过降低二阶更新的通信复杂度并实现线性通信复杂度，MKOR提高了二阶更新的频率。我们进一步提出混合版本MKOR-H，该版本在训练中后期二阶更新不再加速收敛时，自动退化为一阶优化器。实验表明，在64块GPU上针对BERT-Large-Uncased模型训练时，MKOR相较最先进一阶方法（如LAMB优化器）和最优二阶方法实现（如KAISA/KFAC）分别实现高达2.57倍和1.85倍的加速。