🤖 机器学习

MKOR: Momentum-Enabled Kronecker-Factor-Based Optimizer Using Rank-1 Updates

本文介绍了 MKOR，这是一种具备动量机制的二阶优化器，它利用秩-1 更新和基于 Sherman-Morrison 的求逆方法，将计算复杂度和通信复杂度分别从三次降低至二次和线性，从而与最先进的一阶和二阶方法相比，显著加速了大语言模型的训练与收敛。

原作者： Mohammad Mozaffari, Sikan Li, Zhao Zhang, Maryam Mehri Dehnavi

发布于 2024-01-30✓ Author reviewed ⓘ

📖 1 分钟阅读☕ 轻松阅读

原作者： Mohammad Mozaffari, Sikan Li, Zhao Zhang, Maryam Mehri Dehnavi

✨ 这是对下方论文的AI生成解释。它不是由作者撰写的。如需技术准确性，请参阅原始论文。阅读完整免责声明

在现代人工智能的广阔版图中，最强大的系统是建立在深度神经网络之上的，这些复杂的数学结构能够学习识别模式、翻译语言并生成文本。为了教导这些系统，研究人员必须解决一个巨大的优化问题：找到一组完美的内部数值，即权重，使网络能够进行准确的预测。几十年来，标准方法一直是向着能最快减少误差的方向微调这些权重，这种方法类似于登山者通过感受脚下的坡度来摸索下山路径。这种一阶方法可靠且应用广泛，但它可能很慢，通常需要数千步才能到达谷底。

一种更具野心的策略被称为二阶优化，它试图为登山者提供一张完整的地形图，包括地面的曲率，从而实现更快速的下降。通过理解坡度的变化方式，这些方法在理论上可以用更少的步数达到最优解。然而，从历史上看，这种方法对于当今使用的巨型模型来说，计算成本过高，难以投入实际应用。为大型网络计算必要的曲率信息需要巨大的处理能力和内存，这往往使得该方法比它旨在取代的更简单的方法还要慢。挑战在于，如何在不承担携带地图之重担的情况下，捕捉到那位阅读地图的登山者的速度。

来自多伦多大学、罗格斯大学和德克萨斯高级计算中心的科研团队开发了一种名为 MKOR 的新工具，解决了这一困境。他们创建了一个具备动量能力的优化器，该优化器使用一种特定的数学捷径来近似地形形状，而无需计算每一个细节。MKOR 并没有尝试计算损失函数景观中完整的、复杂的曲率，而是使用了一种称为“秩-1 更新”（rank-1 updates）的技术。这种方法通过专注于最重要的变化方向来简化计算，有效地将计算成本从立方关系降低到了二次方关系。从实际意义上讲，这意味着随着模型规模的增大，该方法具有更好的扩展性，避免了此前阻碍二阶方法应用于大语言模型的瓶颈。

研究人员在名为 BERT-Large-Uncased 的大语言模型上测试了这种新的优化器，并在 64 个高性能图形处理器上进行了实验。他们将 MKOR 与现有的最佳一阶方法以及最先进的二阶实现进行了对比。结果令人瞩目。在涉及问答的任务中，标准的一阶方法需要 1,536 步才能达到特定的准确度水平。而新的 MKOR 方法仅用 600 步就达到了相同或更高的准确度。在测量完成训练所需的总时间时，MKOR 用时仅略超过三小时便完成了任务，而标准方法则耗时近八小时。这代表了 2.57 倍的加速，有效地将训练时间缩短了一半以上。

除了纯粹的速度之外，该研究还解决了一个经常困扰二阶方法的关键稳定性问题。由于这些方法根据复杂的数学因素调整学习过程，它们容易出现“梯度爆炸”现象，即调整幅度变得过大，导致训练过程崩溃或发散。以往的解决方案通常需要显著降低学习率以保持安全，这抵消了速度优势。MKOR 内置了一个稳定器，用于监测这些调整的大小。如果数值变得过大，系统会自动对其进行重新缩放，在保持二阶方法激进节奏的同时，防止训练因数值过大而崩溃。这使得优化器能够在没有不稳定性风险的情况下快速移动。

研究人员还观察到，二阶方法的优势在训练早期最为显著，即模型正在学习数据的宏观轮廓时。随着模型接近最优解，复杂曲率计算带来的收益就会减少。为了解决这个问题，他们引入了一个混合版本的优化器，称为 MKOR-H。这个版本以强大的二阶方法开始，以加速初始学习阶段，但一旦训练损失停止快速下降，就会自动切换到更简单的一阶方法。这种混合方法使系统能够以更少的步数收敛，在不到一半的时间内达到与标准方法相同的最终准确度。

研究还探讨了为什么这种特定的数学捷径如此有效。研究人员分析了神经网络内部的数据结构，发现描述网络不同部分如何相互作用的协方差矩阵通常是低秩的。这意味着尽管数据具有高维性，但信息集中在少数几个关键方向上。通过使用秩-1 更新来近似这些矩阵，MKOR 捕捉到了优化所需的本质信息，而无需承担全矩阵计算的开销。团队通过广泛的测试验证了这种近似引入的误差极小，证实了这种简化并不会以牺牲准确度为代价。

除了速度和稳定性，新方法还显著减少了训练期间不同处理器之间需要传输的数据量。在成千上万个处理器协同工作的分布式计算中，通信可能成为主要的瓶颈。MKOR 将这种通信复杂度从二次方关系降低到了线性关系，这意味着随着更多处理器的加入，系统可以高效扩展，而不会被数据传输延迟所拖累。这种效率使该方法能够处理现代大语言模型所要求的大规模批次大小，而这曾是此前二阶方法的局限。

这项工作的意义不仅限于训练某一个特定的模型。通过使二阶优化在大规模系统中变得切实可行，MKOR 为更高效地训练人工智能提供了一条路径。研究人员指出，他们的方法可以帮助机器学习从业者节省大量的能源和时间。虽然技术本身是一种中性的工具，但其在从医疗诊断到气候建模等领域的应用，可能会加速科学发现。然而，作者也承认，高效的优化算法可以应用于任何模型，包括那些具有负面社会影响的模型，并强调此类技术的负责任使用取决于使用它的人的选择。

最终，这篇论文为深度学习中的一个长期问题提出了一个具体的解决方案。通过结合动量、秩-1 近似和混合切换机制，研究人员创建了一个比前代产品更快、更稳定且更具扩展性的优化器。实验证明，在不承担令其望而却步的高昂计算成本的前提下，利用二阶方法的力量是完全可能的。随着人工智能领域不断向更大、更复杂的模型推进，像 MKOR 这样的工具提供了一种更高效地训练这些系统的方法，将曾经的理论优势转化为了现实。

技术摘要：MKOR —— 基于秩-1 更新的动量增强型 Kronecker 分解优化器

1. 问题陈述

二阶优化方法，如自然梯度下降 (NGD) 和 Kronecker 分解近似曲率 (KFAC)，通过利用曲率信息（Hessian 或 Fisher 信息矩阵近似），提供了优于一阶方法（如 SGD、Adam、LAMB）的收敛速率。然而，由于以下三个主要瓶颈，这些方法在训练大规模深度神经网络 (DNN)（特别是 Transformer 和大语言模型 LLM）时受到了阻碍：

计算复杂度： 传统的 KFAC 方法需要对大小为 $d \times d$ （其中 $d$ 为层维度）的协方差矩阵进行求逆，导致 $O(d^3)$ 的复杂度。类似地，基于 Sherman-Morrison-Woodbury (SMW) 的方法 (SNGD) 通常随 batch 大小 $b$ 扩展（ $O(b^3)$ ）。在 Transformer 中，由于 batch 大小随序列长度扩展，这些复杂度变得难以承受。
通信开销： 分布式训练需要在各个工作节点（workers）之间同步二阶信息（协方差矩阵及其逆矩阵）。KFAC 需要同步 $O(d^2)$ 的数据，而 SNGD 需要 $O(bd + b^2)$ ，这造成了显著的通信瓶颈，限制了可扩展性。
更新频率与稳定性： 由于高昂的求逆成本，最先进的实现（如 KAISA）使用“陈旧”的因子进行低频更新（例如每 100–1000 次迭代更新一次）。这降低了曲率信息的有效性。此外，二阶方法容易出现数值不稳定（梯度爆炸、奇异矩阵），通常需要阻尼因子或高精度算术，这进一步增加了计算成本。

2. 方法论：MKOR 算法

作者提出了 MKOR（基于秩-1 更新的动量增强型 Kronecker 分解优化器），通过四个核心机制解决了上述问题：

秩-1 近似与基于 SM 的求逆：
MKOR 不计算并求逆完整的协方差矩阵，而是使用从样本平均值导出的秩-1 矩阵来近似激活值和梯度的协方差（ $E[aa^T] \approx \bar{a}\bar{a}^T$ ）。通过对这些秩-1 更新应用 Sherman-Morrison 恒等式，MKOR 直接更新 Kronecker 因子的逆。这使因子求逆的计算复杂度从 $O(d^3)$ 降低到 $O(d^2)$ 。
- 更新规则： 逆因子使用动量 ( $\gamma$ ) 和 Sherman-Morrison 公式进行更新，从而在无需完整矩阵求逆成本的情况下实现频繁更新。
线性通信复杂度：
在分布式设置中，MKOR 仅在工作节点之间同步两个秩-1 近似向量（大小为 $2d$ ），而非完整的协方差矩阵（ $4d^2$ ）。这使通信复杂度从二次方 $O(d^2)$ 降低到 线性 $O(d)$ 。此外，MKOR 对这些向量使用半精度 (FP16)，进一步使通信成本减半，由于数值敏感性，这一特性对于其他二阶方法来说很难实现。
稳定性机制：
为了防止二阶方法中常见的数值不稳定和梯度爆炸：
- 基于范数的稳定器： 监控逆因子的无穷范数。如果超过阈值，则将因子与单位矩阵进行混合（ $\hat{L} = \zeta L + (1-\zeta)I$ ）以防止发散。
- 梯度重缩放： 对预处理后的梯度进行重缩放，使其与原始梯度的范数匹配，确保学习率调度器保持有效。
- 理论稳定性： 作者证明（引理 3.1）了因子在不需要人工阻尼因子 ( $\mu I$ ) 的情况下保持正定，并且更新规则中的单标量除法保证不为零。
混合 MKOR (MKOR-H)：
由于意识到随着训练收敛（即损失函数地形趋向于单位矩阵），二阶优势往往会减弱，MKOR-H 引入了切换机制。它监控损失下降率，如果二阶更新不再提供加速，则自动退回到一阶优化器，从而优化成本与收敛速度之间的权衡。

3. 核心贡献

复杂度降低： MKOR 在二阶更新方面实现了 $O(d^2)$ 的计算复杂度和 $O(d)$ 的通信复杂度，缓解了 KFAC 中的 $O(d^3)$ 和 $O(d^2)$ 以及 SNGD 中的 $O(b^3)$ 瓶颈。
高频更新： 通过降低求逆成本，MKOR 能够实现每轮迭代（或非常频繁地）进行二阶更新，消除了对“陈旧”因子的需求，并充分利用曲率信息。
鲁棒性： 引入基于范数的稳定化和梯度重缩放，使 MKOR 能够在无需阻尼因子或双精度算术的情况下稳定运行，从而实现对半精度硬件的高效利用。
混合策略： MKOR-H 变体在二阶和一阶方法之间动态切换，以在整个训练生命周期内保持效率。

4. 实验结果

作者在 64 块 A100 GPU 上对 BERT-Large-Uncased 进行了评估，并将其与最先进的一阶优化器 LAMB、二阶实现 KAISA (KFAC) 和 HyLo (SNGD) 以及 Eva 进行对比。

加速比： 在预训练任务上，MKOR 比 LAMB 实现了 2.57 倍加速，比 KAISA 实现了 1.85 倍加速。具体而言，它将 BERT-Large-Uncased 的训练时间从 8 小时缩短至 3 小时。
收敛性： 在 SQuAD v1.1 问答任务上，MKOR-H 在 600 步内收敛（而 LAMB 需要 1,536 步），同时实现了相当或更好的 F1 分数（90.64 vs 90.44）。
GLUE 基准测试： MKOR 在 1,500 步内就在 GLUE 数据集上达到了新的 SOTA 指标。值得注意的是，尽管 KAISA 在 1,563 步后仍未能收敛到基准平均准确度，但 MKOR 成功收敛。
可扩展性： 与因高内存开销 ( $O(bd + b^2)$ ) 导致发散或无法装入 GPU 显存的 HyLo 不同，MKOR 能够成功使用大 batch size（通过累积实现高达 64k 的有效 batch size）进行训练，且没有内存问题。
近似有效性： 在 ResNet-50 和 BERT 上的实验表明，协方差矩阵的秩-1 近似产生的误差微乎其微，验证了核心近似假设。

5. 重要性与主张

本文声称 MKOR 成功弥合了二阶优化的理论优势与训练现代大规模 DNN 实际约束之间的差距。通过将二阶更新的复杂度降低到与一阶方法相当的水平（在随模型规模扩展方面），同时保持收敛优势，MKOR 实现了：

更快的训练： 显著减少 LLM 的实际训练时间。
更好的收敛： 能够在其他二阶方法（如 KAISA）失败或停滞的地方收敛到更高的准确度指标。
可扩展性： 能够扩展到大 batch size 和模型维度，而不会受到以往限制二阶方法的通信和内存瓶颈的影响。

作者强调，他们的工作使二阶优化在“大模型”时代变得可行，特别是解决了 Transformer 架构中由依赖序列长度的 batch size 引入的可扩展性问题。他们总结道，MKOR 代表了高效且有效的深度神经网络训练的新 SOTA。