Adam-mini：使用更少的学习率获得更多

摘要

我们提出了 Adam-mini，这是一种优化器，它的性能与 AdamW 相当或更好，并且内存占用 $45\%$ 到 $50\%$ 更少。 Adam-mini 通过减少 Adam 中的学习率资源（即 $1/\sqrt{v}$)来减少内存。 我们发现，如果我们（1）按照我们提出的 Hessian 结构原则仔细地将参数划分为块，则可以无害地删除 $v$ 中的这些学习率中的 $\ge 90\%$； （2）为每个参数块分配一个单一但良好的学习率。 我们进一步发现，对于每个参数块，都存在一个可以超越 Adam 的高质量学习率，只要有足够的资源来搜索它即可。 然后，我们提供了一种经济高效的方法来找到良好的学习率，并提出了 Adam-mini。 根据经验，我们验证了 Adam-mini 在大小从 125M 到 7B 的各种语言模型上，在预训练、监督微调和 RLHF 方面的表现与 AdamW 相当或更好。 Adam-mini 内存占用的减少还减轻了 GPU 之间的通信开销，从而提高了吞吐量。 例如，在 $2\times$ A800-80GB GPU 上预训练 Llama2-7B 时，Adam-mini 比 AdamW 实现了$49.6\%$更高的吞吐量，从而节省了 33% 的预训练挂钟时间。

1简介

Adam(W) [25, 33] 已成为训练大型语言模型（例如，[61, 2, 60, 58])。 尽管 Adam 性能优越，但使用成本昂贵。 具体来说，Adam 需要其优化器状态的内存：一阶动量 $m$ 和二阶动量 $v$。这些总共至少占用 $2\times$ 模型大小的内存. 这种内存消耗已经成为大语言模型训练的一大负担。 例如，要训练 7B 模型，仅 Adam 就需要每张卡约 56 GB 的空间用于 $m$ 和 $v$，加上梯度，总共需要 86 GB。 即使对于最先进的显卡（例如 A100-80GB）来说，这也是昂贵的。 为了支持高内存算法，实践中必须使用CPU卸载训练和分片[49]，不幸的是这会增加延迟并减慢[50]。 当训练更大的模型（例如具有 540B 参数的 PaLM）[7] 时，情况会进一步恶化。 在这种情况下，仅 Adam 就占用了 50 多个 GPU，成为预训练的主要开销。

设计需要更少内存的有效优化器是很有趣的。 首先，内存的减少可以减轻CPU卸载的负担，并减轻对模型参数进行分片的需要。 这两者都减少了 GPU 和 CPU 之间的通信，从而可以提高吞吐量并加速训练过程。 其次，它允许从业者使用更少的 GPU 来训练具有所需大小的模型，从而大大节省成本和能源。 第三，降低了训练大语言模型的门槛，鼓励更多GPU资源有限的研究者参与。

在不牺牲 Adam 性能的情况下修改 Adam 是一项挑战。 一个主要原因是我们对 Adam 的 $m$ 和 $v$ [72, 27] 的作用仍然缺乏了解。 目前还不确定 Adam 中的哪些组件对于卓越性能是必不可少的，反之，哪些组件可以重新设计或改进。 一个值得注意的尝试是 Adafactor [56]，它通过对 $v$ 进行低阶分解来减少内存。然而，据广泛报道，Adafactor 在大语言模型训练中的性能下降（例如[35])。 一个可能的原因是 Adam 中当前的 $v$ 对于有效性至关重要，并且不能在不影响性能的情况下减少。 这是可能的，因为据报道，大多数试图在不同程度上修改 $v$ 的现有 Adam 变体的性能比 Adam [40] 更差。 另一个可能的原因是，有可能减少 $v$，但 Adafactor 没有使用最合适的方式：矩阵分解是一种可以广泛应用的通用方法，但它并没有利用太多针对特定问题的方法结构，因此它不能很好地处理特定的神经网络任务。

图3： 亚当迷你的插图。 Adam-mini 通过 Hessian 结构分配学习率（lrs）。 它使用的 lrs 多于 SGD，但少于 Adam。

在这项工作中，我们发现可以显着减少 $v$ 的使用。目前，Adam 的 $v$ 为每个参数分配单独的学习率，即第 $i$ 个参数接收学习率 $\frac{\eta }{\sqrt{v_i}}$，其中 $v_i$是 $v$ 的第 $i$ 组件。对于十亿规模的模型，Adam 设计了数十亿的学习率。 我们认为，用更少的学习率资源可以获得同等或更好的性能。 我们注意到 Transformer 的 Hessian 矩阵具有近块对角线结构，由不同大小的密集子块组成。 我们发现，对于每个密集子块，只要我们有足够的资源来搜索它，就存在一个优于 Adam 的高质量学习率。 由于密集子块的数量比参数的数量少得多，我们的研究结果意味着可以用更少的学习率实现良好的性能。 剩下的问题是如何有效地找到它们。

然后，我们提出了一种廉价且简单的方法来找到足以达到或优于 Adam 的良好学习率。 我们在这里介绍一下所提出的设计原理：我们首先根据密集的Hessian子块将梯度向量划分为$B$子向量，并将其称为$g_b$为$b=[B]$。 对于每个 $g_b$，我们计算以下数量。

$\boxed{v_b=\left(1-{\beta }_2\right)\ast \text{mean}\left(g_b\circ g_b\right)+{\beta }_2\ast v_b,b=1,\mathrm{\cdots }B}$

然后，我们使用 $\eta /\sqrt{v_b}$ 作为与 $g_b$ 相关的参数的学习率。 我们将相应的方法称为 Adam-mini。 我们在图3中提供了一个简单的说明，并在稍后的算法1中对完整的形式进行了降级。 因此，Adam-mini 将几乎所有 Adam 的 $v$ 更改为少数标量，从而显着降低了内存成本。 我们的主要贡献总结如下。

• 新的优化器。 我们提出了一个名为 Adam-mini 的新优化器。 首先，Adam-mini 根据我们在 Hessian 结构上建立的原则对模型参数进行划分。 然后，它使用该块中 Adam 的 $v$ 的平均值为每个块选择一个学习率。 Adam-mini具有以下优点。

  • 轻量级：按照设计，Adam-mini 很大程度上减少了 Adam 中使用的学习率数量。 对于主流大语言模型，Adam-mini可以减少Adam的$v$中$\ge 90\%$的比例，从而节省Adam 45%~50%的内存消耗。

  • 有效性：尽管内存减少了，但我们凭经验验证 Adam-mini 在 125M 到 7B 大小的各种语言模型上的表现与 AdamW 相当甚至更好，包括预训练、监督微调（ SFT）和来自人类反馈的强化学习（RLHF）。 Adam-mini 在训练扩散模型、视觉模型和图神经网络等非 LLM 任务上也取得了更好的性能。

  • 效率： Adam-mini 可以达到比 AdamW 更高的吞吐量。 我们观察到，在 $2\times$ A800-80GB 上预训练 Llama2-7B 时，Adam-mini 达到了 $49.6\%$ AdamW 更高的吞吐量，这节省了 33.1% 的预训练挂钟时间。 效率来自两个因素。 首先，Adam-mini 不会在每步更新中引入额外的计算。 其次，内存的减少允许每个 GPU 具有更大的批处理大小，同时，它减轻了 GPU 之间的通信负担，而这通常是一项主要开销。

• 分区原则。 Adam-mini 的一个关键组成部分是参数划分策略。 我们建议基于 Hessian 中的最小密集子块来划分参数。 这个原则可以适用于一般问题：对于块对角结构的问题，我们发现更多的学习率不一定会带来额外的增益。 特别是，对于与每个密集子块相关的问题，单个（但良好的）学习率足以带来更好的性能。

• 变形金刚的 Hessian 结构。 我们根据经验将上述原理应用于变形金刚。 我们发现 Transformer Hessian 的最小密集块是：Query 和 Key（按头）；价值和预测（整体）；和 MLP（按层）。 我们强调，基于 Hessian 的划分原则对于良好的性能至关重要。 相反，朴素或默认分区会过度简化问题并导致大语言模型不稳定。

2方法

2.1 动机和观察

现在我们讨论激发 Adam-mini 设计的观察结果。 我们首先研究 Adam 的 $v$ 的作用并探索改进的可能性。 在 Adam 中，$v$ 为每个参数提供单独的学习率，即第 $i$ 个参数接收学习率 $\frac{\eta }{\sqrt{v_i}}$，其中 $v_i$ 是 $v$ 的第 $i$ 组件。最近，Zhang 等人[74]指出这样的设计至关重要，因为 Transformers 对于不同的块需要不同的学习率。 他们提供了两个证据：首先，变形金刚和各种神经网络的 Hessian 矩阵接近块对角线（在图 LABEL:fig:block_diagonal 中重述）。 其次，对于 Transformers，每个块都有显着不同的特征值分布（在附录 A.2 中重述）。 结合在一起，Transformers 需要不同块的不同学习率来处理特征值异质性。 这可以由 Adam 的 $v$ 提供。

[74]中的研究结果表明，有必要对每个块使用不同的学习率。 尽管如此，Adam 的作用远不止于此：它不仅为每个块分配单独的学习率，还为每个参数分配单独的学习率。 请注意，参数的数量（可能是数十亿）远大于块的数量（通常是数百）。 这就引出了一个问题：

是否有必要为每个参数使用单独的学习率？ 如果没有，我们能节省多少钱？

我们首先探讨一般优化问题的问题，然后讨论 Transformer。 根据图LABEL:fig:block_diagonal中的发现，我们对块对角 Hessian 的随机二次最小化问题进行了案例研究。 我们在图 LABEL:fig:random_quadratic (a) 中可视化 Hessian 矩阵，并将详细设置放在附录 B.2 中。 我们采用坐标学习率方法，即 Adam，和单学习率方法，即梯度下降（GD）。 我们选择二次最小化，因为最佳学习率具有接近的形式。 我们有以下发现。

• (1)：如图LABEL:fig:random_quadratic(a)和(b)所示，Adam优于最优单学习率方法。 这是预期的，因为 Adam 对不同的参数部署了不同的学习率。

• (2):如图LABEL:fig:random_quadratic(c)和(d)所示，我们考虑一个新问题，其Hessian是(的密集子块) A）。 我们考虑这个新问题的最佳单一学习率方法，发现它优于 Adam，尽管 Adam 分配了更多的学习率。 类似的现象适用于(a)的所有三个子块。

• (3)：如果我们在(2)中收集这些最佳学习率并将它们应用到GD的“分块”版本中，它将比原始版本的Adam更快问题（图LABEL:fig:random_quadratic (b) 中的绿线）。

综上所述，对于块对角 Hessian 的一般问题，更多的学习率并不一定会带来额外的增益。 特别是，对于每个密集子块，单个（但良好的）学习率足以比使用数十或数百个学习率带来更好的性能。 为什么会出现这种情况呢？ 一种可能的解释如下。 Adam可以看作是一种对角预处理方法，但它可能不是一个好的预处理器，因此不能有效地减少稠密子矩阵[15,69,57]的条件数。

我们在《变形金刚》中也观察到类似的现象。 我们考虑代码库中的 4 层 Transformer。 在 PyTorch 的默认分区下，我们随机选择一个参数块作为“留出”块，并将坐标学习率更改为单学习率对应部分。 我们使用 Adam 来处理其余的块。 我们对遗漏块的学习率进行网格搜索并应用余弦衰减时间表。 我们报告了最好的结果，并将这种方法称为“Adam（留一法）”。 图LABEL:fig:leave_one_out显示，对于所有随机挑选的留出块，Adam（留一出）可以实现与 Adam 相似或更好的性能。 当我们随机遗漏最多三个块并搜索三个学习率时，也会观察到类似的现象。 由于网格搜索的成本呈指数级增长，我们无法承担留下更多块的后果。

总结以上所有现象，我们发现：对于 Transformers，可以用比 Adam 少得多的学习率达到相似或更好的性能。 剩下的问题是如何在不进行网格搜索的情况下找到它们。 根据这一发现，在下一部分中，我们提出了一种简单且经济高效的方法，称为 Adam-mini，它可以带来与 Adam 相当甚至更好的性能，但学习率要低 90% 到 99%。

2.2 建议方法：Adam-mini

基于上面的讨论，我们在算法1中提出了Adam-mini。 Adam-mini 的目标是减少 Adam 中的学习率资源，而无需像 2.1 节那样费力地对学习率进行网格搜索。 Adam-mini 由两个步骤组成。 步骤1仅在初始化时需要。

步骤 1-1。 我们将模型参数划分为块。 对于 Transformers，我们使用算法 2：“Transformers 分区”，该算法按头对所有查询和键进行分区，并使用默认的 PyTorch 分区来处理其余部分。 对于其他网络，我们使用默认的 PyTorch 分区，并将其称为算法 3：“非 Transformers 分区”。 我们稍后在2.3节中讨论其设计原理。

步骤 1-2。 我们使用算法4选择embd_blocks。 对于 Transformer 来说，它包括嵌入层和输出层。 对于其他网络，将不选择任何参数。

第2步。 对于 embd_blocks 之外的每个参数块，我们使用单个学习率。 为了在每个块中有效地选择合适的学习率，Adam-mini 只需将 vanilla Adam 中的 $\text{g}\circ \text{g}$ 替换为其平均值。 我们采用这些平均值的移动平均值，就像 Adam 一样。

一个简单的例子。 我们用一个简单的例子来说明Adam-mini的按键设计。 对于具有 5 个参数 $w\in {ℝ}^5$ 的问题，Adam 和 Adam-mini 都执行 $w=w-u\circ m$，其中 $m$ 是一阶动量，$u$ 有以下不同的形式：

• 对于亚当：$u_{\text{Adam}}=(\frac{\eta }{\sqrt{v_1}},\frac{\eta }{\sqrt{v_2}},\frac{\eta }{\sqrt{v_3}},\frac{\eta }{\sqrt{v_4}},\frac{\eta }{\sqrt{v_5}}).$

• 对于 Adam-mini：假设分区是 $(1,2,3)$ 和 $(4,5)$ 然后

  $$u_{\text{mini}}=(\frac{\eta }{\sqrt{\left(v_1+v_2+v_3\right)/3}},\frac{\eta }{\sqrt{\left(v_1+v_2+v_3\right)/3}},\frac{\eta }{\sqrt{\left(v_1+v_2+v_3\right)/3}},\frac{\eta }{\sqrt{\left(v_4+v_5\right)/2}},\frac{\eta }{\sqrt{\left(v_4+v_5\right)/2}}).$$

请注意，有效元素 $u_{\text{mini}}$ 的数量等于块的数量，该数量可能明显小于 $u_{\text{Adam}}$ 的数量，后者等于参数的数量。 对于大语言模型，我们将证明这将释放 $v$ 中的 $\ge 90\%$ 元素。

关于“embd_blocks”的备注。 对于 Transformer，“embd_blocks”指的是嵌入层和输出层。 在这里，我们解释为什么步骤 2 不适用于这些块。 这是因为这些块与步骤2中的平均操作不一致：通过嵌入层的设计，如果与这些行对应的标记没有出现在当前数据小批量中，则许多行将为零。 换句话说，唯一的非零行是表示当前输入单词[47]的行。 因此，对整个嵌入层取平均值将吸收大量的零，这会导致学习率出现高度偏差。 类似的论点也适用于输出层。 在图LABEL:fig:gpt2 (a)中，我们表明删除embd_blocks会导致训练不稳定（参见标记为“Adam-mini (embd_blocks_removed)”的曲线）。

2.3 分区策略原则

现在我们讨论如何选择 Adam-mini 的参数划分。 根据2.1节的分析，大致原理如下原理1所示。 基于这一原则，我们现在在神经网络的背景下提出相应的分区策略。

原则一： 我们应该将参数划分为块，使得每个参数块与 Hessian 中的最小密集子块相关联。

根据文献中报告的块对角线结构（图LABEL:fig:block_diagonal)，PyTorch 中的默认分区将是一个合理的候选。 在第 3 节中，我们将证明这种划分确实广泛适用于 ResNet、扩散模型和图模型等非 Transformer 任务。 我们在算法3中展示了这一策略：“非 Transformer 的分区”。

算法1 Pytorch 风格的 Adam-mini

1: Input weight-decay coefficient $\lambda$ and current step $t$

2: Choose param_blocks from Algorithm 2 or 3

3: Choose embd_blocks from Algorithm 4

4: for name, param in param_blocks do

5: g = param.grad

6: param = param - ${\eta }_t\ast \lambda \ast$ param

7: $\text{m}=(1-{\beta }_1)\ast \text{g}+{\beta }_1\ast \text{m}$

8: $\widehat{\text{m}}=\frac{\text{m}}{1-{\beta }_1^t}$

9: if param in embd_blocks then

10: $\text{v}=(1-{\beta }_2)\ast \text{g}\circ \text{g}+{\beta }_2\ast \text{v}$

11: else

12: $\text{v}=(1-{\beta }_2)\ast \text{mean}(\text{g}\circ \text{g})+{\beta }_2\ast \text{v}$

13: end if

14: $\widehat{\text{v}}=\frac{\text{v}}{1-{\beta }_2^t}$

15: param = param - ${\eta }_t$ * $\frac{\widehat{\text{m}}}{\sqrt{\widehat{\text{v}}}+ϵ}$

16: end for

算法2 Transformer 隔断

1: param_blocks = {}

2: for name, param in parameters do

3: if name is Query or Key then

4: Partition param by heads as param[0..heads-1]

5: for i = 0...heads-1 do

6: param_blocks[name+i]=param[i]

7: end for

8: else

9: param_blocks[name]=param

10: end if

11: end for

12: return param_blocks

算法3 非 Transformer 的分区

1: param_blocks = {}

2: for name, param in parameters do

3: param_blocks[name]=param

4: end for

5: return param_blocks

算法4 获取嵌入块

1: if Transformer then

2: return [Embedding layer, Output layer]

3: else

4: return []

5: end if

不幸的是，我们发现默认的 PyTorch 分区在 Transformers 上并不总是能很好地工作。 特别是，我们发现 Adam-mini 在 1B 模型上遇到训练不稳定的情况（参见图LABEL:fig:babygpt_hessian_plot (d)）。 我们怀疑这是因为默认的 PyTorch 分区没有完全捕获 Hessian 结构。 为了找到答案，我们探索了 4 层 Transformer 的 Hessian 矩阵，如图 LABEL:fig:leave_one_out 所示。 在默认的 PyTorch 分区下，我们在初始化时为每个参数块（或等效地，整个 Hessian 中的主要子块）计算精确的 Hessian 矩阵。 我们在图 LABEL:fig:babygpt_hessian_plot 中报告了一些代表，并在附录 A 中对其他代表进行了降级。 我们发现有两类 Hessian 子块。

• 第 1 类： 与整个Hessian类似，Hessian子块本身还具有块对角线结构，并且由更小的稠密矩阵组成。 该类包括Query和Key。 我们凭经验发现，小密集子块的数量等于多头注意力中头的数量。

• 第2类： Hessian 子块具有密集的结构，无法进一步划分为更小的子块。 此类包括价值层、注意力投影层和 MLP 层。 请注意，Value 的 Hessian 结构与 Query 和 Key 的 Hessian 结构不同，尽管它们都由四个头组成。 这可能是因为在 self-attention 设计中 Value 位于 softmax 算子之外，而 Query 和 Key 则不然。

基于以上发现，我们看到默认的 Pytorch 分区确实不是最适合 Transformers：查询和键可以按头进一步分区为不同的块。 根据原则1，不同的头应该属于不同的块。 直观上，不同的头部在理解标记时扮演着不同的角色，因此它们需要不同的学习率似乎是合理的。

然后，我们介绍由此产生的算法2：“Transformers 分区”，该算法按头对查询和键进行分区。 该过程遵循基于最小Hessian子块的原则。 如图LABEL:fig:babygpt_hessian_plot(d)所示。 这一策略确实稳定了训练并提高了表现。

2.4 Adam-mini的一些特性

内存减少了。 根据我们的设计，对于 Transformers，Adam-mini 将学习率的数量从总参数数量减少到嵌入层大小、输出层大小以及其他层中的块数量之和。对 Transformer 进行分区的过程。 因此，减少的内存比例取决于模型中非嵌入参数的比例。 在 Llama2-7B 上，该比例为 96.2%。 对于主流的大语言模型，这个分数通常是$\ge 90\%$（参见附录A表5中的总结）。 请注意，随着越来越多的层堆叠到架构中，该分数单调接近 100%，例如，对于 Llama3-70B，该比例为 99.25%。 综上所述，Adam-mini 可以减少超过 $90\%$ 的 $v$，最多节省 $45\%$ 到 $50\%$ 的 Adam 内存。

表1： AdamW 与 AdamW 的内存消耗 Llama2-7B 上的 Adam-mini。 计算基于float32。

Model	Optimizer	Memory (GB)
GPT2-1.5B	AdamW	12.48
GPT2-1.5B	Adam-mini	6.56 $(\mathbf{47.4}\%\downarrow )$
TinyLlama-1B	AdamW	8.80
TinyLlama-1B	Adam-mini	4.92 $(\mathbf{44.10}\%\downarrow )$
Llama2-7B	AdamW	53.92
Llama2-7B	Adam-mini	28.04 $(\mathbf{48.04}\%\downarrow )$

表2： 在 $2\times$ A800-80GB GPU 上进行 Llama2-7B 预训练的吞吐量 ($\uparrow$) 测试。 ✗表示内存不足。

Optimizer	Micro_bs	Total_bs	Throughput ($\uparrow$)
Adam-mini	4	256	5572.19 $(\uparrow \mathbf{49.6}\%)$
AdamW	2	256	✗
AdamW	1	256	3725.59

更高的吞吐量。 Adam-mini可以达到比AdamW更高的吞吐量，特别是在硬件资源有限的情况下。 有两个原因。 首先，Adam-mini 的更新规则中没有引入额外的计算。 算法1中的平均操作产生的成本可以忽略不计。 此外，Adam-mini 显着减少了 AdamW 中张量平方根和张量除法运算的数量。 其次，由于内存的减少，Adam-mini 可以支持每个 GPU 更大的批量大小，同时，它减少了 GPU 之间的通信，这是众所周知的主要开销[50].

由于这些特性，Adam-mini 可以减少预训练的总时间。 我们在表2中提供了证据。 这里，Micro_bs 和 Total_bs 表示每个设备和总批量大小。 当在 $2\times$ A800-80GB GPU 上预训练 Llama2-7B 时，我们发现 Adam-mini 的吞吐量比 AdamW 高 49.6%。 回想一下，吞吐量是指每秒处理的 Token 数，因此 Adam-mini 可以在处理相同数量的 Token 进行预训练时节省 $33.1\%$ 时间。

有改进的空间。 Adam-mini 使用该块中 Adam 的 $v$ 的平均值来设计每个密集 Hessian 子块的学习率。 这样的设计实现了廉价的计算，但它可能不是最优的。 幸运的是，当前的设计足以以更少的内存实现与 Adam 相同或稍好一些的性能。 我们相信学习率设计还有很大的改进空间，但它需要对每个密集的 Hessian 子块进行更细粒度的分析。 我们将其作为未来的重要方向。

一些正交组合。 Adam-mini的思想可以与各种现有方法正交结合。 我们在这里提供了一些示例，并将更多讨论放在第 4 节中。

• GaLore [75] 是大语言模型的一个新的内存高效优化器。 给定梯度矩阵 $g$，GaLore 计算出一个低阶梯度估计器 $\widehat{g}$，然后根据该估计器计算出 $m$ 和 $v$ $\widehat{g}$。 Adam-mini 可以与 GaLore 结合使用，以进一步减少 $v$ 上的内存。组合方法，例如“GaLore-mini”，总共可以进一步减少 GaLore 上的大约 $40\%$ 内存和 AdamW 上的大约 $81\%$ 内存。 此外，GaLore-mini 可以减轻卸载负担并提高 GaLore 的吞吐量，特别是在内存有限的客户级 GPU 上进行训练时。

• Sophia [31] 是另一个最新的对角预处理优化器。 正如 Adam 一样，Sophia 也需要 $m$ 和 $v$ 的内存。可以将 Adam-mini 和 Sophia 组合起来得到“Sophia-mini”，这可以节省 Sophia 高达 50% 的内存。 Sophia-mini 还可以提高吞吐量并进一步加快 Sophia 的挂钟时间，如表 2 所示。

3实验

我们现在验证 Adam-mini 的功效。 我们主要关注大语言模型任务，包括预训练、监督微调（SFT）和人类反馈强化学习（RLHF）。 我们还将在非 LLM 任务上评估 Adam-mini。 所有大语言模型实验均在四块NVIDIA A800-80GB GPU上进行，其余实验在四块V100 GPU上进行。 所有实验细节均在附录B.1中解释。

3.1预训练

设置。 我们预训练了开源大语言模型，包括GPT2系列和Llama系列。 我们从头开始在主流英语语料库上训练这些模型。 特别是，我们在 Openwebtext [18] 上训练 GPT2 [48] 系列。 我们在 CommonCrawl 上训练 TinyLlama-1B、Llama2-7B [60] .我们用 2B、3B 和 25B 词符（根据钦奇拉定律 [22]的建议）来训练模型。 我们根据标准协议选择模型配置（例如上下文长度）。 我们将 Adam-mini 与 AdamW 以及流行的内存高效方法进行比较，包括 Adafactor [56]、CAME [35] 和 SM3[4]。 对于 Adafactor 和 SM3，我们将动量与 ${\beta }_1=0.9$ 结合起来，以确保与其他方法进行公平比较。 我们在相同的预算内调整所有方法的学习率并报告最佳性能。

骆驼系列。 图LABEL:fig:tinyllama(a)显示了预训练TinyLlama-1B的验证损失曲线。 Llama2-7B 的训练曲线如图 1 中的图 LABEL:fig:intro (c) 所示。 对于 TinyLlama-1B 和 Llama2-7B，我们发现 Adam-mini 的性能与 AdamW 相当，但内存较少。

GPT2系列。 图LABEL:fig:tinyllama（b）显示了GPT2-125M的验证损失曲线，图LABEL:fig:gpt2显示了大小从330M到GPT2的验证损失曲线。 1.5B。 我们发现 Adam-mini 在内存较少的情况下与 AdamW 的良好性能相匹配，而其他方法的性能较差。 在图LABEL:fig:tinyllama（b）中，我们进一步考虑Adam-mini（embd_blocks_removed），它设置了embd_blocks$=\mathrm{\varnothing }$。 也就是说，我们对整个嵌入（输出）层使用一个学习率。 我们发现 Adam-mini (embd_blocks_removed) 表现不佳，正如 2.1 节中的分析所预期的那样。 我们停止试验，因为它显示出明显的不稳定行为。

敏感性分析。 在GPT2-125M预训练任务上，我们测试了Adam-mini对超参数的敏感性。 我们报告使用 2.5B Token 训练后的验证损失。 如图LABEL:fig:tinylllama所示，Adam-mini似乎对超参数不敏感。

3.2有监督微调和 RLHF

在本节中，我们评估 Adam-mini 对于下游微调任务的有效性。 具体来说，我们考虑两个代表性任务：监督微调（SFT）和来自人类反馈的强化学习（RLHF）。 我们使用 Llama2-7b 预训练模型（由 Meta [60] 发布）进行研究。 我们使用ultrafeedback数据集 并实现 [41] 中的 RLHF 工作流程。 具体来说，我们使用 40% 的所选数据训练 SFT 模型，并使用剩余的 60% 训练奖励模型。 然后，我们应用强化学习算法 ReMax [29]（PPO [55] 的内存高效替代方案）来优化偏好奖励。

我们使用 MT-Bench [76] 评估聊天能力方面的对齐性能，其中 GPT-4 评估多轮聊天能力并分配 0 到 10 之间的分数（越高越好）。 Table 3 中显示的结果表明 Adam-mini 可以优于 AdamW，即使 Adam-mini 使用单一学习率并且更节省内存。

表3： SFT 和 RLHF 在 MT-Bench [76] 上的 GPT-4 评估分数 ($\uparrow$)。

	SFT (LoRA)			SFT			RLHF
	AdamW	Adam-mini		AdamW	Adam-mini		AdamW	Adam-mini
MT-Bench	4.23	4.41		5.37	5.40		5.54	5.68

3.3 非 LLM 任务

我们现在在非 LLM 任务上评估 Adam-mini。 表4显示了在ImagNet上训练ResNet18的结果； CelebA 上的扩散模型； OGB-arxiv 上的图卷积网络（GCN）、图注意力网络（GAT）。 我们发现 Adam-mini 达到了与 AdamW 相当或更好的性能，但内存更少。

表4： 在流行的非 LLM 任务上，Adam-mini 的表现与 AdamW 相当或更好。

Domain	Model	Optimizer	Metric	25% steps	50% steps	75% steps	100% steps
Vision	ResNet18	AdamW	Test acc ($\uparrow$)	0.6149	0.6478	0.6613	0.6669
Vision	ResNet18	Adam-mini	Test acc ($\uparrow$)	0.6140	0.6501	0.6629	0.6667
Vision	Diffusion model	AdamW	Train loss ($\downarrow$)	0.0529	0.0497	0.0420	0.0394
Vision	Diffusion model	Adam-mini	Train loss ($\downarrow$)	0.0525	0.0495	0.0416	0.0388
Graph	GAT	AdamW	Val acc($\uparrow$)	0.7277	0.7367	0.7399	0.7421
Graph	GAT	Adam-mini	Val acc ($\uparrow$)	0.7378	0.7394	0.7403	0.7429
Graph	GCN	AdamW	Val acc ($\uparrow$)	0.7347	0.7428	0.7379	0.7374
Graph	GCN	Adam-mini	Val acc ($\uparrow$)	0.7406	0.7427	0.7380	0.7423

4相关作品

对亚当的理解。 有一系列积极的工作试图理解为什么 Adam 表现良好[71, 64, 73, 63, 42, 24, 27, 74, 3]。 与这些工作相比，我们指出 Adam 的 $v$ 可能无法像我们预期的那样有效地发挥其全部潜力：有时更少的学习率可以达到相同或更好的结果（由于密集的 Hessian 子结构）块）。 我们的发现可能会激发更强大的优化器，更好地适应神经网络 Hessian 结构。

关于神经网络的 Hessian 矩阵。 Hessian 矩阵对于梯度方法的行为至关重要。 研究 MLP 和 CNN 的 Hessian 矩阵有几个重要的尝试[8, 51, 39, 52, 53, 5, 45, 64, 30, 43, 44, 54, 19, 66, 70, 16, 67 ，11]。 受这些作品的启发，我们探索了变形金刚的 Hessian 结构，并将其与亚当的行为联系起来。 然后我们找到了改进的空间，并建议将 Adam 瘦身为 Adam-mini。

适用于一般任务的轻量级优化器。 有多种尝试来降低 Adam 的记忆成本。 Adafactor [56] 及其变体 CAME [35] 对 Adam 的 $v$ 进行非负低秩分解。 SM3 [4] 是 AdaGrad [14] 的轻量级版本。 SM3通过取某个候选集中的最小值来选择第$i$个参数的学习率，候选集中的每个元素都与预定覆盖下的最大平方梯度相关。 所有这些方法都可以释放 $v$ 的几乎所有内存，并节省 Adam 约 48% 的内存。 然而，我们发现它们的性能在各种实验中退化，而 Adam-mini 保持与 AdamW 一样有效（图LABEL:fig:tinyllama)。

完成这项工作后，我们注意到有两种方法分享了 Adam-mini 的一些想法：BAGM [77] 和 NovoGrad [17]。 它们都使用逐块或逐层自适应学习率来实现稳健的性能和更好的泛化。 我们总结了它们与 Adam-mini 的主要区别。 BAGM 对参数进行分区以达到最小范数解决方案并实现可证明的稳健性。 根据他们的理论，最终的策略是：（I）他们使用默认的 PyTorch 分区； (II)他们没有引入embd_blocks。 尽管这些分区设计可能对鲁棒性有好处，但它们忽略了 Hessian 结构并过度简化了训练问题（正如我们在 2.2 和 2.3 节中讨论的）。 因此，这些设计将导致大规模实验中训练的不稳定。 这在我们失败的 Adam-mini 基础知识版本中很明显，如图 LABEL:fig:babygpt_hessian_plot (d) 和 LABEL:fig:tinyllama (b) 所示。 至于 NovoGrad，它使用分层学习率设计，比 BAGM 中的块式设计更粗糙，因此与 BAGM 具有类似的缺点。 此外，NovoGrad 对一阶动量引入了不同的设计：它不是对过去的梯度进行加权和，而是对“当前梯度除以二阶动量”进行加权和。 这样的设计与AdamW和Adam-mini有很大不同。 目前还不清楚这样的设计是否适用于像大语言模型这样的大规模任务。

综上所述，这两种方法具有不同的设计，并且它们的划分策略过度简化了训练问题。 因此，它们会导致大规模实验中的训练不稳定，如图 LABEL:fig:babygpt_hessian_plot (d) 和 LABEL:fig:tinyllama (b) 所示。 相比之下，Adam-mini 按照我们提出的 Hessian 结构原则仔细分配学习率。 这样的设计原则对于训练的稳定性至关重要，并且在包括7B模型在内的各种大语言模型上都表现良好。

其他正交方法。 Adam-mini的思想可以与各种现有方法正交结合以获得更多优点。 在2.4节中，我们已经讨论了将 Adam-mini 与 GaLore [75] 和 Sophia [31] 结合起来可能带来的好处。 我们在这里列出了更多潜在的组合。 LoRA [23] 是一种用于 SFT 任务的内存高效方法。 该方法通过附加低阶适配器微调模型，并使用 Adam 更新这些适配器。 请注意，LoRA 中的 Adam 步骤可以用 Adam-mini 代替。 因此，Adam-mini 带来了更好的性能（图LABEL:fig:sft)。 与我们的工作并行，BAdam [34] 以块坐标下降 (BCD) 方式进行 SFT。 该方法需要重复 Adam 步骤来解决 BCD 中的子问题。 与 LoRA 类似，BAdam 中的 Adam 步骤可以用 Adam-mini 代替，以进一步减少内存。 Nero 优化器 [32] 也会减少 Adam 的内存。 它消除了一阶动量并使用神经元特定的投影梯度式更新。 根据[32]，他们的设计对权重矩阵施加了约束，具有“平衡激励和抑制”的优点。 这种设计有可能与 Adam-mini 相结合，以进一步提高性能。 为了节省微调大语言模型的内存成本，MeZO [38]使用零阶方法来近似梯度信息。 可以将这个想法与 Adam-mini 结合起来，进一步节省 SFT 的内存。 Adam-mini 还可以与其他对角线预处理方法结合使用，例如 AdaGrad[14] 和 Adan[65]。

还有其他一些技巧可以减轻 GPU 内存负担，但与优化器设计正交。 这些技巧包括梯度检查点 [6]、模型卸载和分片 [49, 50]、量化 [13, 28] 和融合更新[36, 37]。 Adam-mini 可以根据这些技巧来实现。

最后，我们讨论另一种流行的自适应优化器，称为 LAMB [68]（参见附录A.1中的算法7) ）。 LAMB 可能会被误解为与 Adam-mini 类似的优化器，但实际上并非如此。 我们强调 Adam-mini 与 LAMB 显着不同。 首先，LAMB 仍然保持与 Adam 相同的坐标学习率设计$1/\sqrt{v}$。 其次，除了这个$1/\sqrt{v}$之外，LAMB还进一步以逐层的方式重新缩放参数。 这种重新缩放设计通常被称为“逐层学习率”，但准确地说，它实际上是除了“坐标学习率$1/\sqrt{v}$”之外的额外“逐层缩放” 。 因此，LAMB 并不比 Adam 节省内存，而且其整体设计与 Adam-mini 有很大不同。 这是可以理解的，因为 LAMB 是为大批量训练而设计的，而不是为了节省内存。 从数值上看，我们发现 LAMB 在 GPT2 预训练上的表现比 Adam-mini 差（图LABEL:fig:tinyllama (b)）。

5结束语

我们提出了 Adam-mini，这是一个优化器，可以节省 Adam 45% 到 50% 的内存。 我们指出，Adam-mini 的设计还有很大的改进空间：目前 Adam-mini 使用一种简单且经济有效的方法来设计每个密集 Hessian 子块的学习率，但这可能不是最佳方法。 我们将开发更强大的设计作为未来的方向。

更广泛的影响

我们提出了一种训练人工智能模型的新方法。 我们的方法可以帮助节省大型人工智能模型训练的能量。 然而，如果人工智能模型被用于非法用途，这将是一个潜在的威胁。

附录 A更多结果

A.1 Adam、AdamW 和 LAMB 的初步介绍

我们在算法5、6和7中引入了Adam、AdamW和LAMB。 这些方法需要沿着训练跟踪 $m$ 和 $v$。 $m$ 和 $v$ 都是与 # 模型参数大小相同的向量。

算法5 Pytorch 风格的 Adam

1: Let $\lambda$ be the weight decay coefficient

2: for param in parameter_blocks do

3: g = param.grad

4: if $\lambda >0$ then

5: $\text{g}=\text{g}+\lambda \ast \text{param}$

6: end if

7: param = param - ${\eta }_t\ast \lambda \ast$ g

8: $\text{m}=(1-{\beta }_1)\ast \text{g}+{\beta }_1\ast \text{m}$

9: $\widehat{\text{m}}=\frac{\text{m}}{1-{\beta }_1^t}$

10: $\text{v}=(1-{\beta }_2)\ast \text{g}\odot \text{g}+{\beta }_2\ast \text{v}$

11: $\widehat{\text{v}}=\frac{\text{v}}{1-{\beta }_2^t}$

12: param = param - ${\eta }_t$ * $\frac{\widehat{\text{m}}}{\sqrt{\widehat{\text{v}}}+ϵ}$

13: end for

算法6 Pytorch 风格的 AdamW

1: Let $\lambda$ be the weight decay coefficient

2: for param in parameter_blocks do

3: g = param.grad

4: param = param - ${\eta }_t\ast \lambda \ast$ g

5: $\text{m}=(1-{\beta }_1)\ast \text{g}+{\beta }_1\ast \text{m}$

6: $\widehat{\text{m}}=\frac{\text{m}}{1-{\beta }_1^t}$

7: $\text{v}=(1-{\beta }_2)\ast \text{g}\odot \text{g}+{\beta }_2\ast \text{v}$

8: $\widehat{\text{v}}=\frac{\text{v}}{1-{\beta }_2^t}$

9: param = param - ${\eta }_t$ * $\frac{\widehat{\text{m}}}{\sqrt{\widehat{\text{v}}}+ϵ}$

10: end for

算法7 Pytorch 风格的 LAMB

1: Let $\lambda$ be the weight decay coefficient, let $\varphi$ be a scaling function.

2: for param in all_layers do

3: g = param.grad

4: param = param - ${\eta }_t\ast \lambda \ast$ g

5: $\text{m}=(1-{\beta }_1)\ast \text{g}+{\beta }_1\ast \text{m}$

6: $\widehat{\text{m}}=\frac{\text{m}}{1-{\beta }_1^t}$

7: $\text{v}=(1-{\beta }_2)\ast \text{g}\odot \text{g}+{\beta }_2\ast \text{v}$

8: $\widehat{\text{v}}=\frac{\text{v}}{1-{\beta }_2^t}$

9: $\text{r}=\frac{\widehat{\text{m}}}{\sqrt{\widehat{\text{v}}}+ϵ}$

10: param = param - ${\eta }_t$ * $\frac{\varphi \left(\Vert \text{param}\Vert \right)}{\Vert r+\lambda \ast \text{param}\Vert }$ * r

11: end for

A.2 [74]中的初步结果

我们在此重申[74，图3]。 该图表明：对于 Transformers，不同的参数块具有不同的 Hessian 特征值分布，而对于 CNN，块之间的特征值分布相似。 这表明 Transformer 需要不同块的不同学习率来处理特征值分布的异质性。

A.3 图LABEL:fig:babygpt_hessian_plot中更多Hessian子块

在图LABEL:fig:babygpt_hessian_plot中，我们在第二层展示了Query、Key、Value和MLP_proj的Hessian矩阵。 我们现在在图 LABEL:fig:babygpt_hessian_plot 中使用的 4 层 Transformer 中的所有层中全面提供 Hessian 注意力和 MLP 块。 我们在初始化时绘制 Hessian 矩阵。 实验详情参见附录B.2。 如下所示，几乎所有这些 Hessian 子块都属于 Class 1 或 Class 2。 唯一的例外是第 4 层中的 MLP_fc 块，它显示对角线图案。

这就引出了一个问题：Adam 在这个对角块上是否明显优于单学习率方法？ 我们发现答案仍然是否定的。 这是因为该块中的 Hessian 特征值比较集中。 我们在图 LABEL:fig:babygpt_mlp_fc 中提供了证据。 这表明该块中的优化问题是条件良好的，并且单学习率方法可以很好地工作。 如图LABEL:fig:babygpt_mlp_fc所示，最优单学习率仍然明显优于 Adam，与其他密集块中发生的情况类似。 另一个证据在于图LABEL:fig:leave_one_out（d）：当第4层的MLP_fc块被选为留出块时，Adam（留一出）实现了0.1416的较低损失比亚当。

A.4 大语言模型中的嵌入比例

表5展示了GPT2[48]、Llama系列[60]、Phi等主流大语言模型嵌入输出参数的占比-3 [1] 和杰玛 [59]。 我们发现嵌入和输出层占总参数的比例相当小，通常是$\le 10\%$。

表5： 主流大语言模型中嵌入和输出参数的比例。 请注意，GPT2 和 Gemma 使用权重绑定，因此嵌入层和输出层的参数是共享的。 我们发现嵌入和输出层通常需要 $\le 10\%$ 的总参数。

Model	Vocabulary size	Embedding dimension	# Parameters for the embedding & output layer	# Total parameters	Proportion
GPT2-1.5B	50304	1600	50304 * 1600	1.5B	5.17%
TinyLlama-1B	32000	2048	32000*2048*2	1.1B	11.80%
Llama2-7B	32000	4096	32000*4096*2	6.7B	3.80%
Llama2-13B	32000	4096	32000*4096*2	13B	2.02%
Llama2-33B	32000	4096	32000*4096*2	33B	0.79%
Llama2-65B	32000	4096	32000*4096*2	65B	0.40%
Llama3-8B	128256	4096	128256*4096*2	8B	13%
Llama3-70B	128256	4096	128256*4096*2	70B	0.15%
Phi-3	32064	3072	32064*3072*2	3.8B	5.18%
Gemma-7B	256128	3072	256128*3072	8.5B	9.83%

A.5 与流行的内存高效优化器的进一步比较

与 AdamW 类似的轨迹。

作为一个额外的发现，Adam-mini 可以生成与 AdamW 类似的轨迹，而其他流行的内存高效方法却不能。 我们在 Openwebtext 上训练大小为 11M 的 6 层 Transformer，并在相同的随机种子和相同的学习率 1e-5 下启动 AdamW、Adam-mini 和其他内存高效优化器。 我们每 250 次迭代保存一次模型权重，并将它们的欧几里德距离与 AdamW 轨迹上的权重进行比较。 如图16所示，在整个训练过程中，Adam-mini获得的模型权重与AdamW接近。 相比之下，其他优化器偏离了 AdamW。 这是因为 Adam-mini 对 AdamW 的修改比其他方法少。

图16： 在 6 层 Transformer 上，我们比较了不同方法轨迹上的模型检查点。 我们发现 Adam-mini 训练获得的模型权重在整个过程中与 AdamW 获得的模型权重接近。

A.6 来自 Adam-mini 训练的大语言模型的示例响应

示例响应如图 17 所示。

Prompt: Compose an engaging travel blog post about a recent trip to Hawaii, highlighting cultural experiences and must-see attractions.

AdamW: Sure! Here’s a travel blog post about a recent trip to Hawaii: …… Whether you’re interested in history, nature, or simply soaking up the sun, Hawaii has something for everyone. Until next time, Aloha!

Adam-mini: Sure! Here’s a travel blog post about a recent trip to Hawaii: …… I highly recommend visiting this beautiful archipelago for a truly unforgettable adventure.

Prompt: Rewrite your previous response. Start every sentence with the letter A.

AdamW: A trip to Hawaii is a must-visit destination for any traveler, offering a variety of cultural experiences and must-see attractions. From the Polynesian Cultural Center to the USS Arizona Memorial, and from Lanikai Beach to the lively night markets, every corner of the islands reveals a new surprise and a new way to connect with the spirit of Hawaii.

Adam-mini: A trip to Hawaii is a must-do for any travel enthusiast. A trip to the beautiful Hawaiian Islands is a paradise on earth. A visit to Oahu is a great way to start your adventure. A helicopter tour of the island is both educational and breathtaking. A visit to Maui is a must-see. A visit to Haleakala National Park is a highlight of your trip. A visit to the Polynesian Cultural Center is a great way to learn about the history and traditions of the islands. A traditional luau is a must-attend event. Overall, a trip to Hawaii is an unforgettable experience.

图 17： Llama2-7B 的响应样本由 AdamW 和 Adam-mini 针对 MT-Bench [76] 数据集的 #81 测试问题进行了微调。 由于第一轮响应中的信息很少，因此我们省略了模型中的答案。

附录 B更多实验细节

B.1 第 3 部分的训练配置

对于所有实验，我们根据标准协议选择模型配置（例如上下文长度）。 如果适用，我们会根据开源平台的推荐来选择学习率。 例如，对于 GPT2 系列，我们使用 [31] 推荐的学习率，据报道该学习率通过网格搜索是最佳的。 除非另有说明，Adam-mini、Adafactor、SM3 和 LAMB 使用与 AdamW 推荐的相同的学习率。 我们根据作者。 如果 AdamW 没有公开推荐的学习率，我们会在相同的计算预算内调整所有优化器的学习率并报告最佳性能。 对于其他超参数，我们遵循开源平台的建议或默认设置。 对于 SM3 和 Adafactor，我们将动量与 ${\beta }_1=0.9$ 结合起来，以提供与其他优化器的公平比较，其余超参数设置为默认值。 详细配置说明如下。

GPT2 预训练。

我们使用nanoGPT代码库 在 Openwebtext 上训练 GPT2 大小为 125M（小）、330M（中）、770M（大）和 1.5B（XL）。 对于所有模型，我们使用上下文长度= 1024，批量大小= 480，权重衰减系数$\lambda =0.1$， $$ε= 1e-8，${\beta }_1=0.9,{\beta }_2=0.95$。 我们使用余弦衰减学习率计划并进行 2000 次预热迭代。 对于 GPT2-小、中、大，我们使用 [31] 推荐的峰值学习率，据报道这是通过网格搜索找到的最佳学习率。 对于 GPT2-XL，我们使用 Levanter 。 对于 GPT2-small、medium、large、XL，选择的峰值学习率分别为 6e-4、3e-4、2e-4、1e-4。 这些模型的最小学习率选择为 3e-5、6e-5、1e-5、1e-5。

骆驼预训练。

我们使用 TinyLlama 代码库 在 CommonCrawl 上训练 TinyLlama-1B 和 Llama2-7B 。 我们使用批量大小 = 40。 我们将 CommonCrawl 的方向“2019-30”预处理为 10850 个 bin，总共提供 85GB 的 Token 。 我们留下最后 20 个 bin 作为验证集。 对于这两个模型，我们使用权重衰减系数$\lambda =0.1$， $$ε= 1e-8，${\beta }_1=0.9,{\beta }_2=0.95$。 我们使用余弦衰减学习率计划并进行 2000 次预热迭代。 对于 TinyLlam-1B，我们使用上下文长度 = 2048，批量大小 = 512，峰值学习率 = 2e-4，最小学习率 = 2e-5。 对于 Llama2-7B，我们使用上下文长度 = 4096，批量大小 = 256，峰值学习率 = 8e-5，最小学习率 = 8e-6。

SFT 和 RLHF。

SFT和RLHF代码的实现基于ReMax代码库。 我们在训练中使用 DeepSpeed ZeRO-2。 Table 3 中的 GPT-4 评估模板来自代码库。 在奖励优化阶段，我们使用 ReMax，这是 PPO 的内存高效替代品。 我们使用 UltraFeedback 数据集 [9] 并使用 40% 的数据用于 SFT，60% 的数据用于 ReMax。

SFT。 我们批量使用 80 个样本，并对模型进行 3 轮训练。 对于完整的参数调整，我们根据验证损失从 {1e-6, 2e-6, 3e-6, 4e-6, 5e-6, 1e-5, 2e-5} 中搜索学习率，并使用 2e AdamW 和 Adam-mini 均采用余弦退火 -6。 对于LoRA，我们将LoRA应用于除嵌入层之外的所有层。 LoRA的等级设置为128。 从同一组中选择学习率作为全参数调整后，当应用 LoRA 时，我们对 AdamW 和 Adam-mini 都使用 5e-6。 根据 LlamaFactory 的建议，权重衰减系数设置为 0 。 AdamW 和 Adam-mini 的其余超参数是 $$ε= 1e-8，${\beta }_1=0.9,{\beta }_2=0.95$。

ReMax。 我们在一批中使用 48 个样本，并对模型进行 1 轮训练。 根据验证奖励从{5e-7,1e-6,2e-6}中搜索峰值学习率，AdamW使用1e-6，而Adam-mini选择5e-7作为峰值学习率。 权重衰减系数设置为0。 AdamW 和 Adam-mini 的其余超参数是 $$ε= 1e-8，${\beta }_1=0.9,{\beta }_2=0.95$。

残网。

我们使用PyTorch官方实现代码库 在 ImageNet [12] 上训练 ResNet18 [20]。 我们使用余弦衰减学习率，epoch=90，${\beta }_1=0.9,{\beta }_2=0.999,ϵ=$1e-8。 对于 ResNet18，我们使用批量大小 = 256，峰值学习率 = 0.005。 对于 ViT-base，我们使用批量大小 = 128，峰值学习率 = 0.0001。 这些配置适用于 Adam-mini 和 AdamW。

扩散模型。

我们使用代码库 训练扩散模型。 图像大小为 64，训练目标是预测噪声，如 [21] 中所示。 我们使用默认的 U-Net 架构超参数，U-Net 中的维度乘法为 (1, 2, 4, 8)。 我们使用 CelebA 数据集 并使用余弦衰减的学习率 $5\times {10}^{-5}$ 训练扩散模型。 批量大小为 128，训练周期为 50。

图神经网络。

我们使用 DGL 实现 用于 OGBN-arxiv 的图卷积网络 (GCN) [26] 和图注意网络 (GAT) [62] 数据集。 所有配置均为默认。 对于 Adam-mini 和 AdamW，我们对 GCN 使用默认学习率 = 0.005，对 GAT 使用默认学习率 = 0.002。

B.2 其他实验的详细设置

图 LABEL:fig:random_quadratic 的配置。

对于每个密集子块$H_l,l=1,2,3$，我们使用随机正定矩阵。 我们将每个$H_l$的特征值选择固定如下：对于$l=1$，我们从{1,2,3}独立采样30次；对于$l=2$，我们对{99,100,101}重复此过程；对于 $l=3$，我们对 {4998 ,4999, 5000} 重复此过程。 对于单一（块式）学习率方法，我们使用具有最佳恒定学习率 $2/(L+\mu )$ 的 GD，其中 $L,\mu$ 是（块式）Hessian 的最大和最小特征值。 我们将 Adam 与 ${\beta }_1=0$ 一起使用。 这有助于我们关注 Adam 中坐标学习率的影响。 我们还将 ${\beta }_2=1$ 设置为随时间变化的学习率。 这是必要的，因为对于任何 ，具有恒定学习率的 Adam 将在二次函数上振荡。 这在[10，命题12，图1]中得到了理论上的证明，并在[74，第3.3节]中得到了经验观察。

图 LABEL:fig:babygpt_hessian_plot 的配置。

我们使用 。 我们考虑一个 4 层 Transfomer，其中 n_emb = 32，n_head = 4。 数据集是随机生成的长度为 42 的二进制序列。 借助自动微分框架，我们通过两遍反向传播[46]计算整个数据集上的Hessian矩阵，并且计算是准确的。 在图LABEL:fig:babygpt_hessian_plot中，我们首先将所有块的颜色条vmax的最大值设置为1e-6，这大致是大多数块中的最大值。 但随后我们发现整个查询和关键块都是黑色的，因此我们将这些块的 vmax 减小到 1e-8，然后块对角线图案变得清晰。

吞吐量比较见表2。

结果是使用 Tinyllama 代码库在 2$\times$ A800-80GB GPU 上进行测试的。 我们没有打开 CPU 卸载。 我们从 Wandb 日志的摘要文件中报告吞吐量。