SRIBD 2025-04-11 18:05 广东
Adam-mini 在降低 50% 优化器内存开销的同时，保持了与 Adam 优化器相当甚至更优的性能，提升吞吐量可达50%。
在这项工作中，研究团队通过分析 Transformer 模型的 Hessian 结构，提出了一种轻量化优化器 Adam-mini 。Adam-mini 大幅减少了 Adam 优化器中学习率的数量，在降低 50%优化器内存开销的同时保持了与 Adam 优化器相当甚至更优的性能，提升吞吐量可达50%。
相关研究成果已发表于 ICLR 2025，并在Medium、YouTube及 X 等平台获多次关注解读。
当前，大语言模型的广泛应用使得训练大型模型的内存开销成为研究者和企业面临的一大挑战。这主要归因于大语言模型广泛采用的 Adam 优化器（Adaptive Moment Estimation），尽管其性能卓越，但内存需求较高。具体而言，Adam 优化器在训练过程中需要额外存储两个状态变量：一阶动量 m 和二阶动量 v（即每个参数需要存储两组变量），从而显著增加了内存负担。
举例而言，训练一个 Llama 2-7B（70亿参数量）模型的内存占用大致如下：
权重参数：28 GB（ 7B X 4 byte）
Adam 状态变量：58 GB（ 2 x 7B X 4 byte）
其他：激活值和梯度缓存等
如此高的内存需求对硬件造成了巨大的压力，即便是先进的 A100-80GB 显卡也难以直接承载，必须依赖内存分片或CPU卸载技术来分担负载。然而，这些额外的处理方式不仅显著增加了训练过程的复杂性，还加重了通信开销。
为了改进Adam，我们希望理解为什么现在的Transformer架构需要用Adam训练。我们注意到前人发现了神经网络Hessian 有近似块对角结构（如图 2 所示）。受这一现象的启发，我们考虑分析不同块对应的 Hessian阵的特征谱（即: 该矩阵所有特征值的分布）。我们于 2024 年发现了Transformer 架构中存在的 Hessian 块异质性 （Block Heterogeneity）现象，并发现这一性质和Adam密切相关。相关研究成果已发表于 NeurIPS 2024 。
图 2：Hessian 近似块对角结构
Definition [Hessian Block Heterogeneity]
Hessian 块异质性指的是，神经网络的 Hessian 矩阵在不同参数块之间的特征谱具有显著差异的现象。
如图 3(a) 所示，以 CNN 架构的 VGG16 为例，不同块（卷积层）的 Hessian 谱较为相似。然而，图 3(b) 表明，在 Transformer 架构的 BERT 中，不同块的 Hessian 谱差异很大 (比如特征值范围差异可超200倍)。
(a) VGG16 (b) BERT
图 3：VGG16（ CNN）和 BERT（Transformer）在初始化时的分块 Hessian 谱
众所周知，在 CNN 模型中， SGD 与 Adam 性能相当；但是在 Transformer 模型中，SGD 的性能则显著逊色于 Adam。我们推测，这种性能差异在一定程度上源于 Hessian 块异质性。在大量真实任务和合成的异质性任务上进行的实验表明， SGD 的性能会随着块异质性程度的增加而下降，从而进一步支持了这一推测。
具体来讲，我们将每个参数块的 Hessian 谱归一化为概率分布，并采用 JS distance 来度量分布之间的差异，从而定量描述块异质性的程度。如表 1 和图 3 所示，我们对若干模型在初始化时的 JS distance 进行了计算，得出了以下发现：
图 3 的趋势表明：
JS distance 越大，块异质性程度越高， SGD 优化器的性能相较于 Adam 越差；
Transformer 架构主导的模型，其块异质性程度是 CNN 的数百倍。
我们进一步将块异质性和优化器的表现构建了联系。块异质性意味着不同的参数块有不同的优化需求，使用单一学习率的 SGD 难以适应多样优化需求，因此我们认为块异质性正是 Adam 在 Transformer 架构中优于 SGD 的关键原因（总结如图）。
这一发现在本文的主题 Adam-mini 优化器的设计过程中起到了关键作用。
上文我们的发现表明 Adam能满足Transformer 架构中不同参数块的 Hessian 矩阵的不同的优化需求，这说明了Adam是“充分的”，但还不足以说明它是“必要”的。我们在此基础上开展了进一步的消融实验，我们发现：在同一个参数块内，Adam目前的“独立分配学习率”可能并没有发挥出应有的功效。具体原因涉及到计算数学中一个悠久的话题，即关于diagonal-preconditioner的局限性分析，在此省略。总结来说，我们认为Adam目前的设计可能存在冗余。
基于上述动机和观察，我们对Adam的冗余部分进行了“瘦身”，并提出了 Adam-mini 优化器，相关研究成果已发表于 ICLR 2025，其核心设计思想如下：
1. 基于 Hessian 结构（如图4）分块：对于 Transformer 架构， query 和 key 按 head 分块；value 、attn.proj 和 MLP 按输出神经元分块；embed_layer 和 output_layer 按 token 分块。
2. 简化学习率：取子块 b 梯度的均方值为唯一学习率，且在块内共享。
简单的理论分析表明， Adam-mini 可以无损地移除超过 99.9% 的二阶动量 v ，从而为 Adam 优化器节省接近 50% 的内存。
我们从多方面对 Adam-mini 进行了实验评估，实验语言模型的预训练和微调任务，以下是主要结论：
1. 大幅降低内存开销：在 Llama 系列（图 5，6）模型的预训练中， Adam-mini 相比 AdamW 减少了 50% 的优化器状态内存开销。与此同时， Adam-mini 在多项指标上都与 AdamW 的表现持平，在某些任务中甚至略优。
2. 提升吞吐量：如图 8 所示，由于减少了通信开销和内存占用， Adam-mini 在预训练 Llama 2-7B 模型时吞吐 量提升了 49.6% ，训练时间缩短了 33%。
3. 无需额外调参：Adam-mini 用AdamW原版超参数即可达到一样性能，可实现算法无脑迁移。
4. 可扩展到更大的模型：我们进行了大量Scaling law的实验（图7）并发现：对于不同的尺度的模型，Adam-mini的训练曲线都与AdamW高度重合。这为Adam-mini在更大规模实验中的可扩展性和有效性提供了证据。
近年来，低秩方法已成为内存优化的一个有效方案，其中的代表作是 GaLore。我们注意到 GaLore 和Adam-mini 的技术思路在本质上是正交的，因此融合两者的核心思想，提出了全新的优化器 GaLore-mini（被 NeurIPS 2024 Workshop接收）。
GaLore-mini 通过利用梯度的低秩特性，将梯度投影到低维空间来更新优化器状态，同时结合 Adam-mini 的策略，显著减少了优化器的存储需求，从而极大程度降低了内存开销。相较于 AdamW，GaLore-mini 在不损失性能的情况下实现了 81% 的内存节省。
各种方法的内存开销对比如图 8 所示。
Adam-mini 与现有框架高度兼容，用户仅需一行代码即可轻松安装：
使用方法与 AdamW 类似：
论文相关参考文献
[1] Zhang, Y., Chen, C., Ding, T., Li, Z., Sun, R.*, & Luo, Z. Why transformers need Adam: A Hessian perspective. Advances in Neural Information Processing Systems (NeurIPS) 2024, 37, 131786-131823, https://arxiv.org/pdf/2402.16788
[2] Zhang, Y., Chen, C., Li, Z., Ding, T., Wu, C., Kingma, D. P., ... & Sun, R.. Adam-mini: Use fewer learning rates to gain more. International Conference on Learning Representations (ICLR) 2025, https://arxiv.org/pdf/2406.16793
[3] Huang, W., Zhang, Z., Zhang, Y., Luo, Z. Q., Sun, R., & Wang, Z. Galore-mini: Low rank gradient learning with fewer learning rates. In NeurIPS 2024 Workshop on Fine-Tuning in Modern Machine Learning: Principles and Scalability, https://openreview.net/pdf?id=45nZPmfNN6.
如果您对本文内容感兴趣，欢迎访问 GitHub 项目主页或扫描下方二维码了解更多关于Adam-mini的代码使用介绍。
如果你对本文介绍的研究内容感兴趣，期待进一步深入学习并参与相关科研工作，欢迎加入我们的课题组！
课题组隶属于香港中文大学(深圳)数据科学学院，研究氛围活跃，并与多所国内外高校及企业保持紧密合作。研究方向涵盖大模型的算法与理论、学习优化以及大模型在垂直领域的应用。目前我们正在招收：
博士后
MPhil / 博士研究生
实习生 / 科研助理
我们希望你：
对 LLM / 优化算法有浓厚兴趣
具备扎实的数学与编程基础
自驱力强，有强烈的研究热情
欢迎将简历 & 简要介绍发送至：sunruoyu@cuhk.edu.cn，邮件标题请注明“职位申请+姓名+感兴趣的研究方向”。
未经「AI科技评论」授权，严禁以任何方式在网页、论坛、社区进行转载！
公众号转载请先在「AI科技评论」后台留言取得授权，转载时需标注来源并插入本公众号名片。
UCL强化学习派：汪军与他的学生们
为什么中国只有一个 DeepSeek？
为什么是梁文锋做出了DeepSeek？
