在训练大型语言模型（LLM）时，Adam(W) 基本上已经成为了人们默认使用的优化器。
Adam 尽管性能优异，但使用成本很高。具体来说，Adam 需要内存来保存其优化器状态：一阶动量 m 和二阶动量 v^2。这总共需要模型大小至少 2 倍的内存。这样的内存消耗已经成为了 LLM 训练的一大主要负担。
举个例子，要训练一个 7B 模型，只是 Adam 就需要每张卡有大约 56 GB 来保存 m 和 v；而如果再加上梯度，则总共需要 86 GB。即使使用最先进的 A100-80GB，成本也过高了。
为了支持这样的高内存算法，实践中必须要使用 CPU 卸载与分片，但这又会增加延迟，减慢训练速度。在训练 PaLM （有 5400 亿参数）这样的更大型模型时，情况还会更糟。在这种情况下，Adam 自身就要占用超过 50 GB，并且这也是预训练阶段的一大主要开销。
因此，人们希望设计出内存需求更少又有效的优化器。首先，减少内存可以减轻 CPU 卸载的负担并能减轻对模型参数执行分片的需求。这些都能减少 GPU 和 CPU 之间的通信量，并进一步提升训练过程的吞吐量和速度。其次，这允许实践者使用更少的 GPU 来训练所需大小的模型，从而极大地节省成本和能源。第三，这能降低训练 LLM 的门槛并鼓励更多 GPU 资源有限的研究者参与进来。
但是，修改 Adam 但不影响其性能的难度非常大。其中一个主要原因是我们仍然不太理解 Adam 的 m 和 v 的作用。我们还不清楚 Adam 的哪些组件对其卓越性能而言必不可少，也就更不知道可以重新设计和改进哪些组件了。
Adafactor 是一个颇受欢迎的尝试，其能降低在 v 上低秩分解所用的内存。但是，很多研究发现 Adafactor 训练 LLM 的性能不佳。可能的原因有两个。第一，当前 Adam 中的 v 可能对有效性来说非常关键，已经没有缩减空间。第二，v 是有可能缩减的，只是 Adafactor 使用的方法并不是最合适的：矩阵分解是一种可以广泛应用的通用方法，但它没有利用太多特定于问题的结构，因此它在特定的神经网络任务上效果不佳。
近日，香港中文大学（深圳）、深圳市大数据研究院、杜克大学和斯坦福大学的一个联合研究团队发现：可以通过一个简单技巧来降低 v 的使用量。
目前而言，Adam 的 v 会为每个参数单独分配一个学习率，即第 i 个参数获得的学习率为 
，，其中 v_i 是 v 的第 i 个组件。对于十亿参数量的模型，Adam 就要设计十亿个学习率。
该团队认为这些学习率资源可以大幅降低，同时还不会影响模型性能，甚至能让模型获得更优性能。该团队注意到，Transformer 的 Hessian 有一种接近块对角线的结构，其由不同大小的密集子块构成。他们发现，对于每一个这样的密集子块，都存在性能优于 Adam 的单个高质量学习率——只要有足够的资源将其搜索出来。由于密集子块的数量远少于参数数量，这个发现表明使用远远更少的学习率是有可能取得优良性能的。剩下的问题如何高效地找到它们。
为了找到性能足以比肩甚至优于 Adam 的优良学习率，该团队提出了一种低成本的简单方法 Adam-mini。另外，该团队也发布了 Adam-mini 的一种实现。
论文标题：Adam-mini: Use Fewer Learning Rates To Gain More
论文地址：https://arxiv.org/pdf/2406.16793
实现代码：https://github.com/zyushun/Adam-mini
在实验中，Adam-mini 的表现非常优秀。如图 1 所示， 在预训练 Llama2-7B 时，Adam-mini 的内存占用可以大幅降低，同时吞吐量也有明显提升。在保证性能与 Adam 相当的同时，其能带来 33% 的速度提升。
图 2 展示了 Adam-mini 的大致图示和学习率使用情况。可以看到，Adam-mini 是根据 Hessian 结构来分配学习率，因此其使用的学习率比 SGD 多，但远少于 Adam。
本研究的通讯作者，香港中文大学（深圳）数据科学学院副教授孙若愚（Ruoyu Sun）表示：「如果你喜欢 Adam(W)，那 Adam-mini 就是你的不二之选！」
方法
动机和观察
为了应对特征值的异构性，Transformer 的每一参数块都需要不同的学习率。Adam 的 v 可以提供这一点。
论文《Why Transformers Need Adam: A Hessian Perspective》发现：必须为每个块使用不同的学习率。但是，Adam 所做的却远不止此：其并不只是为每个块分配不同的学习率，而是为每个参数都分配不同的学习率。而参数量（可能超过数十亿）要远远多于块的数量（通常数百计）。那问题就来了：
有必要为每个参数都使用单独的学习率吗？如果没必要，又能节省多少？
该团队进行了一番探索，得到了如下发现：
如图 4 (a) 和 (b) 所示，Adam 的表现优于最优的单学习率方法。这符合预期。
如图 4 (c) 和 (d) 所示，这里研究了其 Hessian 是 (a) 的一个密集子块的新问题。该团队探究了针对这个问题的最优单学习率方法，结果发现其表现优于 Adam，即使 Adam 分配的学习率要多得多。(a) 的所有三个子块都有类似的现象。
如果收集 (2) 中的那些最优学习率，并将其用于梯度下降的一个「逐块」版本，那么其在原始问题上的速度会比 Adam 快，见图 4(b) 中的绿线。
也就是说，对于带有块对角 Hessian 的一般问题而言，许多学习率并不一定会带来额外收益。
该团队也在 Transformer 中观察到了类似的现象。
总结起来，该团队发现：对于 Transformer，有可能使用比 Adam 远远更少的学习率实现与之相当或更好的性能。剩下的问题就是在不使用网格搜索的情况下如何找到这些学习率。
新方法：Adam-mini
基于上述讨论，该团队提出了 Adam-mini，见算法 1。Adam-mini 的目标是降低 Adam 中的学习率资源，而无需费力地对学习率执行网格搜索。Adam-mini 包含两个步骤。步骤 1 只在初始化时进行。
步骤 1-1：将模型参数切分为块。对于 Transformer，该团队选择的策略是「Partition for Transformers」，即根据头来切分所有的 Query 和 Key，并使用 PyTorch 来切分其余部分。对于其它网络，则使用默认的 PyTorch 切分即可，他们将其称之为「Partition for non-Transformers」。
步骤 1-2：使用 embd_blocks。对于 Transformer，其包含嵌入层和输出层。对于其它网络，不选择任何参数。
步骤 2：对于每个位于 embd_blocks 之外的参数块，都使用单个学习率。为了高效地为每个块选择合适的学习率，Adam-mini 的做法是直接将原始 Adam 中的 g◦g 替换成其均值。Adam-mini 像 Adam 一样在这些均值上使用了移动平均值。
一个简单示例。为了说明 Adam-mini 的关键设计，该团队在论文中给出了一个简单的示例。假设有一个问题，其有 5 个参数 w ∈ ℝ^5，Adam 和 Adam-mini 都会执行 w = w − u ◦ m，其中 m 是一阶动量，而 u 有不同的形式，如下所示：
对于 Adam：
对于 Adam-mini：假设分片方式为 (1, 2, 3) 和 (4, 5)，则
请注意，有效元素 u_mini 的数量就等于块的数量，这远远小于 u_Adam，其等于参数数量。结果表明，这能让 LLM 的 v 中的元素数量减少 90% 以上。
分片策略的原则
现在讨论如何为 Adam-mini 选择参数分片。该团队基于前述分析得出了一个广义原则：
原则 1：应当将参数切分成不同的参数块，使得每个参数块都关联了 Hessian 中最小的密集子块。
由于参数的块对角结构，PyTorch 的默认分片方法是一个合理的候选项。实际上，这种分片确实能很好地用于非 Transformer 任务，比如 ResNet、扩散模型和图模型。算法 3 Partition for non-Transformers 展示了该策略。
但不幸的是，该团队发现，默认的 PyTorch 切片并不总是能很好地应对 Transformer。比如他们发现 Adam-mini 在 1B 模型上会出现训练不稳定问题（见图 6(d)）。
他们猜测这是因为 PyTorch 切片无法完整地理解 Hessian结构。经过一番探索后，他们发现 Hessian 子模块分为两类：
第一类：类似于 Hessian 整体，这种 Hessian 子块本身也有进一步的块对角结构，由更小的密集矩阵构成。这一类包含 Query 和 Key。他们通过实验发现，小型密集子块的数量就等于多头注意力中头的数量。
第二类：这种 Hessian 子块有密集的结构，无法进一步分成更小的块。这一类包含 Value、注意力投射和 MLP 层。请注意 Value 的 Hessian 结构不同于 Query 和 Key 的 Hessian 结构，但它们全都由 4 个头组成。这是因为 Value 位于自注意力设计的 softmax 算子之外，而 Query 和 Key 却不是。
基于上述发现，该团队发现默认的 PyTorch 分片确实不是最适合 Transformer 的。可以根据头将 Query 和 Key 进一步切分成不同的块。根据原则 1，不同的头应当属于不同的块。根据直觉就能知道，不同的头在理解 token 时的作用也不同，因此它们需要不同的学习率就很合理了。
由此，该团队得到了算法 2：Partition for Transformers，其可根据头来切分 Query 和 Key。该过程遵循基于最小 Hessian 子块的原则如图 6 (d) 所示。该策略确实能稳定训练过程并提升性能。
Adam-mini 的一些特点
减少内存用量
根据设计，Adam-mini 可为 Transformer 减少学习率的数量——从参数总数减少到嵌入层的大小、输出层的大小和其它层中块的数量的总和。因此，减少的内存比例取决于模型中非嵌入参数的比例。
在 Llama2-7B 上，这个比例是 96.2%。对于主流的语言模型，这个比例通常 ≥ 90%。请注意，如果层数更多，这个比例也会更加接近 100%；比如对于 Llama3-70B，这个比例为 99.25%。因此，Adam-mini 可将 v 减少至少 90%，由此可为 Adam 带来 45% 到 50% 的内存节省。
吞吐量更高
Adam-mini 可取得比 AdamW 更高的吞吐量，尤其是当硬件资源有限时。
基于此，Adam-mini 可以减少预训练的总时间。如表 2 所示。当在 2 台 A800-80GB 上预训练 Llama2-7B 时，Adam-mini 的吞吐量比 AdamW 高 49.6%。召回吞吐量 (↑) 指的是每秒处理的 token 数量，因此在预训练时，Adam-mini 处理相同数量的 token 可节省 33.1% 的时间。
实验
预训练
设置
该团队预训练了一些开源 LLM，包括 GPT2 系列和 LLM 系列。他们在主流英语语料库上从头开始训练了这些模型。
具体来说，他们在 Openwebtext 上训练了 GPT2 系列，在 CommonCrawl 上训练了 TinyLlama-1B 和 Llama2-7B。他们使用 2B、3B 和 25B token 训练了模型。模型配置（比如上下文长度）是根据标准协议选择的。
该团队将 Adam-mini 与 AdamW 以及 Adafactor、CAME 和 SM3 等常用内存高效型方法进行了比较。对于 Adafactor 和 SM3，为了确保与其它方法进行公平的比较，这里整合了 β_1 = 0.9 的动量。他们在同样的预算内对所有方法的学习率进行了调节，并报告了最佳性能。
Llama 系列
图 7 (a) 展示了预训练 TinyLlama-1B 的验证损失曲线。Llama2-7B 的训练曲线已经在图 1(c) 中给出。对于 TinyLlama-1B 和 Llama2-7B，可以发现 Adam-mini 使用更少的内存便能达到与 AdamW 相近的效果。
GPT2 系列
图 7 (b) 展现了 GPT2-125M 的验证损失曲线，图 8 展示了大小从 330M 到 1.5B 的 GPT2 的验证损失曲线。可以看到，Adam-mini 使用更少的内存就能取得与 AdamW 相当的优良性能，而其它方法的表现要更差一些。
敏感度分析
该团队在 GPT2-125M 预训练任务上测试了 Adam-mini 对超参数的敏感度。这里报告了使用 2.5B token 训练之后的验证损失。如图 7 所示，Adam-mini 看起来对超参数并不敏感。
监督式微调和 RLHF
该团队评估了 Adam-mini 对下游微调任务的有效性。
具体来说，他们考虑了两种有代表性的任务：监督式微调（SFT）和根据人类反馈的强化学习（RLHF）。这里使用的预训练模型是 Llama-2-7B。
结果见表 3，可以看到 Adam-mini 的表现优于 AdamW，即便 Adam-mini 使用了单个学习率并且内存效率更高。
非 LLM 任务
该团队也在非 LLM 任务评估了 Adam-mini。表 4 给出了在 ImageNet 上训练 ResNet18、在 CelebA 上训练扩散模型、在 OGB-arxiv 上训练图卷积网络（GCN）和图注意力网络（GAT）的结果。结果发现 Adam-mini 使用更少的内存便能取得与 AdamW 相当或更好的性能。
