Gefen: Optimized Stochastic Optimizer
本文介绍了 Gefen，一种内存高效的优化器，它通过自动二阶矩共享和学习型一阶矩量化，将 AdamW 的内存占用减少了约 8 倍，从而在保持同等性能的同时，实现了更大的批次大小和更高的吞吐量。
原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。 阅读完整免责声明
想象一下你正在训练一个巨型机器人（深度学习模型）学习一项新技能。为了做到这一点，机器人需要一个“教练”（优化器）来告诉它在每次练习结束后如何调整其内部设置。目前最受欢迎的教练叫做 AdamW。
AdamW 是一个很棒的教练，但它背着一个沉重的背包。对于机器人的每一个设置，AdamW 都会携带两个额外的笔记本（称为“一阶和二阶动量状态”）来记录过去的错误以及变化的速度。如果你的机器人有 10 亿个设置，AdamW 的背包会增加另外 20 亿个设置的重量。这使得背包变得如此沉重，以至于你无法在标准的训练计算机上装下这么大的机器人，或者你不得不让机器人在极小的、缓慢的步骤中进行练习。
这篇论文介绍了一个新的教练叫做 Gefen。Gefin 是一个“内存高效型”优化器，它将那个沉重的背包缩小了 8 倍，同时保持机器人的学习速度和质量与 AdamW 完全相同。
以下是 Gefen 如何实现这一点的，我们使用了一些简单的类比：
1. “集体拥抱”策略（共享笔记）
问题： AdamW 为机器人大脑中的每一个数字都写了一份单独的笔记。
Gefen 的解决方案： Gefen 意识到许多这些数字实际上是“好朋友”。它们的反应方式非常相似。与其为每一个单独的朋友写一份独特的笔记，Gefen 将他们分组为团队（块），并给整个团队 一份共享的笔记本。
- 它是如何知道如何分组的？ 通常，你需要一张超级复杂的地图（称为 Hessian 矩阵）来观察谁和谁是朋友，但对于巨型机器人来说，这张地图实在太大了，无法绘制出来。
- Gefen 的窍门： Gefen 查看机器人的第一次练习课程。它观察哪些数字以相似的模式一起移动。它假设：“如果你在开始时一起移动，那么以后你也可能会一起移动。”然后，它会自动将这些数字分组。不需要人类告诉它应该创建哪些组；它自己就能搞定。
2. “素描艺术家”策略（量化）
问题： 即便有了共享笔记本，里面的数字仍然是用高精度编写的（例如 10 位小数），这非常占用空间。
Gefen 的解决方案： Gefen 使用了一种“素描艺术家”的方法。它不再写下精确的数字（例如 0.123456789），而是将其舍入到来自预制列表（码本）的最近的“标准值”。
- 智能列表： 大多数优化器使用通用的、固定的列表（就像一把带有固定刻度的尺子）。然而，Gefen 会观察正在指导的特定机器人，并学习专门针对该机器人的完美刻度列表。它使用一种聪明的数学方法（动态规划）来创建一个完美的定制尺子，以完美契合数据并最大限度地减少误差。
- 稳定性： 一旦它在开始阶段学会了这个列表，它就会在整个过程中一直使用同一个列表。它不需要每天重新绘制尺子，这节省了时间并保持了稳定性。
结果：更轻的背包，同样的性能
作者在各种模型上测试了 Gefen，从小型图像分类器到大型语言模型（如 Llama 3）。
- 内存： 与 AdamW 相比，Gefen 将优化器所需的内存减少了约 8 倍。对于一个 130 亿参数的模型，它将内存从约 97 GB 降低到了仅 1.5 GB。
- 速度： 因为背包变轻了，机器人可以一次携带更大的“微批次”（练习组）。在多台计算机（FSDP 和 DDP）的测试中，这使得训练速度提升了 56%，因为计算机不再需要等待内存释放。
- 质量： 尽管进行了大量的压缩和共享，机器人的学习效果与使用沉重的 AdamW 背包时完全一样。最终的性能（准确率或损失）是完全一致的。
为什么这很重要
把它想象成一次旅行时的打包过程。AdamW 就像是一个行李箱，其中每一只袜子都有一个独立的盒子。Gefen 则像是意识到你可以把所有的袜子卷成一个紧凑的捆，然后放进一个小袋子里。你节省了大量的空间，但你仍然拥有所有的袜子，并且你仍然可以一样快地到达目的地。
论文声称 Gefen 是一个“即插即用型替代品”，这意味着你可以直接将其替换到现有的训练代码中，无需更改任何设置，它就会立即为你节省内存并可能提高你的训练速度。
您所在领域的论文太多了？
获取与您研究关键词匹配的最新论文每日摘要——附技术摘要，使用您的语言。