Adam-mini: Use Fewer Learning Rates To Gain More

## Abstract

Abstract We propose Adam-mini, an optimizer that achieves on par or better performance than AdamW with 50 % percent 50 50\% 50 % less memory footprint. Adam-mini reduces memory by cutting down the learning rate resources in Adam (i.e., 1 / v 1 𝑣 1/\sqrt{v} 1 / square-root start_ARG italic_v end_ARG ). By investigating the Hessian structure of neural nets, we find Adam’s v 𝑣 v italic_v might not function at its full potential as effectively as we expected. We find that ≥ 99.9 % absent percent 99.9 \geq 99.9\% ≥ 99.9 % of these learning rates in v 𝑣 v italic_v could be harmlessly removed if we (1) carefully partition the parameters into blocks following our new principle on Hessian structure; (2) assign a single but good learning rate to each parameter block. We then provide one simple way to find good learning rates and propose Adam-mini. Empirically, we verify that Adam-mini performs on par or better than AdamW on various language models sized from 39M to 13B for pre-training, supervised fine-tuning, and RLHF. The reduced memory footprint of Adam-mini also alleviates communication overheads among GPUs, thereby increasing throughput. For instance, Adam-mini achieves 49.6 % percent 49.6 49.6\% 49.6 % higher throughput than AdamW when pre-training Llama 2-7B on 2 × 2\times 2 × A800-80GB GPUs, which saves 33 % percent 33 33\% 33 % wall-clock time for pre-training. 1 1 1 Our implementation of Adam-mini is available at https://github.com/zyushun/Adam-mini .

## 1 Introduction

Adam (Kingma & Ba, ) has become the de-facto optimizer for training large language models (LLMs) (e.g., (Vaswani et al., ; Achiam et al., ; Touvron et al., ; Team et al., )). Despite its superior performance, Adam is expensive to use.
Specifically, Adam requires the memory for its optimizer states: the first-order momentum $m$, and the second-order momentum $v$. These in total take at least $2\times$ the memory of the model size (^2^22We restate the update rules of Adam and AdamW (Loshchilov & Hutter, ) in Appendix D.1.).
This memory consumption has become a major burden in LLM training. For instance, to train a 7B model, Adam alone requires about 56 GB for $m$ and $v$, and with the gradients included, a total of 84 GB is needed. This is expensive even for cutting-edge graphics cards (e.g., A100-80GB). To support training,
CPU-offload and optimizer state sharding (Rajbhandari et al., ) must be used in practice, which unfortunately increases the latency and slows down the training (Rajbhandari et al., ).

It is intriguing to design effective optimizers that require less memory.
First, it lowers the threshold of training LLMs and encourages participation from more diverse researchers, especially those with limited GPU resources.
Second, it requires fewer GPUs
to train a model with a desired size, leading to substantial savings in both cost and energy.
Third,
it can ease the burden of CPU offloading and model sharding, which in turn, can enhance the throughput and accelerate the training process.

However, it is challenging to modify Adam without sacrificing its performance. One primary reason is that we still lack understanding of the role of Adam’s $m$ and $v$ (Zhang et al., ; Kunstner et al., ). It remains uncertain which components in Adam are indispensable for superior performance, and which components could be re-designed or improved. One notable attempt is Adafactor (Shazeer & Stern, ), which cuts down memory by low-rank factorization on $v$. However, we find that Adafactor is not easy to tune and often performs worse than Adam (see evidence in (Luo et al., ) and Section 3.4). One possible reason is that the current $v$ in Adam is crucial and cannot be simplified.
This is possible as most existing Adam variants that attempt to modify $v$ to varying extents have been reported to perform worse than Adam (Orabona, ). Another possible reason is that there is potential to cut down $v$, but Adafactor does not use the most suitable way: matrix factorization is a generic approach that could be applied broadly, but it does not leverage much problem-specific structure, thus it does not work well on specific neural-net tasks.

Figure: Figure 3: An illustration of Adam-mini. Adam-mini assigns learning rates (lrs) by Hessian structure. It uses more lrs than SGD but fewer than Adam.
Refer to caption: extracted/5989783/images/illustration.png

In this work, we find it is possible to significantly reduce the usage
of $v$.
Currently, Adam assigns an individual learning rate for each parameter, i.e., $i$-th parameter receives learning rate $\frac{\eta}{\sqrt{v_{i}}}$, where $v_{i}$ is the $i$-th component of $v$. For a billion-parameter model, Adam requires billions of learning rates. We argue that it is possible to achieve on par or better performance with much fewer learning rates.
We first recall a classical result that the Hessian of neural nets is near-block-diagonal with several dense principal sub-blocks (Collobert, ).
We then find that, for each of these dense sub-blocks, there exists a single high-quality learning rate that outperforms Adam, provided that we have enough resources to search it out.
Since the number of dense sub-blocks is much fewer than the number of parameters,
our findings imply that it is possible to achieve good performance with much fewer learning rates.
The remaining question is how to find them efficiently.

We then propose a cheap and simple way to find good learning rates that are sufficient to perform on par or better than Adam. We introduce the proposed design principle here: we first partition the gradient vector into $B$ sub-vectors based on the smallest dense sub-block in the Hessian, and call it $g_{b}$ for $b\in\{1,\cdots,B\}$. Then for each $g_{b}$, we calculate the quantity below.

$$ $\displaystyle\boxed{v_{b}=(1-\beta_{2})*\texttt{mean}(g_{b}\odot g_{b})+\beta_ {2}*v_{b},\quad b=1,\cdots B.}$ $$

We then use $\eta/\sqrt{v_{b}}$ as the learning rate for the parameters in the block associated with $g_{b}$. This design changes almost all Adam’s $v$ to a negligible amount of scalars and thus reduces the memory.
We call the corresponding method Adam-mini. We provide a simple illustration in Figure 3 and relegate the complete form later in Algorithm 1.
We summarize our main contribution as follows.

- •
New optimizer. We propose a new optimizer called Adam-mini. First, Adam-mini partitions the model parameters based on our new principle: partition by the smallest dense sub-block in the Hessian. Then, it chooses a single learning rate for each block using the average of Adam’s $v$ in that block. Adam-mini has the following advantages.
–
Lightweightness:
By design,
Adam-mini
largely reduces the number of learning rates used in Adam.
For mainstream LLMs,
Adam-mini could cut down $\geq 99.9\%$ proportion of Adam’s $v$, which saves 50% of the memory cost of Adam.
–
Effectiveness:
Adam-mini performs on par or even better than AdamW on various language models sized from 39M to 13B, including pre-training, supervised fine-tuning (SFT), and reinforcement learning from human feedback (RLHF). Further, Adam-mini’s loss curves
closely resemble those of AdamW. Adam-mini also performs similarly to Adam on non-LLM tasks such as training diffusion models, vision models, and graph neural networks.
–
Efficiency: Adam-mini can reach a higher throughput than AdamW. We observe that Adam-mini reaches $49.6\%$ higher throughput of AdamW when pre-training Llama 2-7B on $2\times$ A800-80GB, which saves 33.1% wall-clock time for pre-training. The efficiency comes from two factors.
First, Adam-mini does not introduce extra computation in per-step updates. Second, the memory cut-down allows larger batch sizes per GPU, and at the same time, it eases the burden of communication among GPUs, which is usually a major overhead.
- •
Generic partition principle. A key component in Adam-mini is the strategy for parameter partition. We propose to partition parameters based on the smallest dense sub-block in the Hessian. This principle can apply to generic problems with a block diagonal Hessian: we find that more learning rates do not necessarily bring extra gain for these problems. In particular, for the problem associated with each dense sub-block, a single (but good) learning rate suffices to bring better performance.
- •
Hessian structure and partition principle of Transformers. We empirically apply the above principle to Transformers. We find that Transformer Hessian’s smallest dense blocks are: (1) query, key by heads; (2) value, attn.proj and mlp by output neurons; (3) embed_layer and output_layer by tokens.
We emphasize that our Hessian-based partition principle is crucial for good performance, as naive or default partitions (e.g., partitioning by layers) would cause training instability on LLMs.

## 2 Method

### 2.1 Motivations and Observations

Now we discuss our observations that motivate the design of Adam-mini. (^3^33 All experimental details in Section 2 are shown in Appendix E.2.)
We start by investigating the role of Adam’s $v$ and explore possibilities for improvement. In Adam, $v$ provides an individual learning rate for each parameter, i.e., $i$-th parameter receives the learning rate $\frac{\eta}{\sqrt{v_{i}}}$, where $v_{i}$ is the $i$-th component of $v$.
Very recently, Zhang et al. () pointed out that such design is crucial for modern architectures such as Transformers. This is because these models often exhibit Hessian-block heterogeneity, i.e., the Hessian of different parameter blocks have dramatically different eigenvalue distributions (We restate their findings in Appendix D.2).
This phenomenon suggests that different parameter blocks need different learning rates. This can be provided by Adam’s $v$.

The findings in (Zhang et al., ) suggest that it is necessary to use a different learning rate for each block. Nonetheless, Adam does much more than that: it assigns an individual learning rate not just for each block, but for each parameter. Note that the number of parameters (could be billions) is much larger than the number of blocks (usually hundreds). This raises the question:

(Q1) Is it necessary to use a customized learning rate for each parameter?

If not, how much can we save?

To answer (Q1), we investigate the Hessian structures of neural networks.
First, we recall an important (but often overlooked) result: the Hessian of neural nets is near-block-diagonal. This is an old result that has been reported for at least two decades; see, e.g., (Collobert, , Section 7). The authors also provide theoretical justification. Consider a standard supervised learning problem: minimizing $\ell(f(\theta,x),y)$ where $\ell(\cdot,\cdot)$ is the Cross-Entropy (CE) loss, $f(\theta,x)=\sum_{i=1}^{n}v_{i}\phi(w_{i}^{\top}x)$ is an 1-hidden-layer neural network with input $x\in\mathbb{R}^{d}$, weight $w_{i}\in\mathbb{R}^{d}$, $v_{i}\in\mathbb{R}$, and label $y\in\{0,1\}$, then the off-diagonal-block Hessian elements would contain

$$ $\frac{\partial^{2}\ell(f(\theta,x),y)}{\partial w_{i}\partial w_{j}}={\color[ rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}p_{\theta}(y|x)\left (1-p_{\theta}(y|x)\right)}v_{i}v_{j}\phi^{\prime}\left(w_{i}^{\top}x\right) \phi^{\prime}\left(w_{j}^{\top}x\right)xx^{\top}\quad\text{for }i\neq j,$ (1) $$

where $p_{\theta}(y|x)=1/(1+\exp(-yf(\theta,x)))$ and $\phi^{\prime}(\cdot)$ is the derivative of $\phi(\cdot)$. Since the training objective is to maximize $p_{\theta}(y|x)$, the term $p_{\theta}(y|x)\left(1-p_{\theta}(y|x)\right)$
will quickly shrink to zero. This term will push the Hessian to near-block-diagonal structure where each block corresponds to one output neuron. The authors report that this can happen just after 1 training step, as restated in Figure LABEL:fig:block_diagonal (a). We also numerically reproduce this result on a small MLP in Figure LABEL:fig:block_diagonal (b,c,d), and show that the near-block-diagonal Hessian structure maintains along the training.

Case study I: random quadratic problems. With the above observation in mind, we now explore (Q1) on generic optimization problems with block-diagonal Hessian.
We consider the random quadratic minimization problem $\min_{w}\frac{1}{2}w^{\top}Hw$ where the Hessian $H$ is a random positive definite (PD) matrix and is visualized in Figure LABEL:fig:random_quadratic (a).
We compare the coordinate-wise learning-rate method, i.e., Adam, with the single-learning-rate method, i.e., gradient descent (GD). We choose quadratic minimization because the optimal learning rate has a close form. We have the following findings.

- •
(1): as shown in Figure LABEL:fig:random_quadratic (a) and (b), Adam outperforms the optimal single-learning-rate method. This is expected since Adam deploys different learning rates to different parameters.
- •
(2): as shown in Figure LABEL:fig:random_quadratic (c) and (d), we consider a new problem whose Hessian
is a dense sub-block of (a). We consider the optimal single learning-rate method for this new problem and find it outperforms Adam, even though Adam assigns much more learning rates. Similar phenomena apply to all the three sub-blocks of (a).
- •
(3): If we collect these optimal learning rates in (2) and apply them to a “blockwise" version of GD, it would be faster than Adam on the original problem (the green line in Figure LABEL:fig:random_quadratic (b)).

In summary, for generic problems with block-diagonal Hessian, we find that more learning rates do not necessarily bring extra gain. In particular, for each dense sub-block, a single (but good) learning rate suffices to bring better performance than using tens or hundreds more.

More discussions on case study I.
Why would this happen? We provide
one possible explanation
from a linear algebra perspective.
Adam can be viewed as a diagonal preconditioned method
, i.e., at the $t$-th step:

$$ $w_{t+1}=w_{t}-\eta_{t}D_{t}m_{t},$ (2) $$

where $D_{t}=\operatorname{Diag}(1/{\sqrt{v_{t}}})$ is a diagonal matrix, $m_{t}$ is the 1st-order momentum, $w_{t}$ and $\eta_{t}$ are model parameters and learning rate.
However, Adam may not be an optimal preconditioner and thus
cannot effectively reduce the condition number of the dense
sub-matrix.
In the field of optimization, the effectiveness of a diagonal preconditioner $D$
is often measured by “how much is $\kappa(DH)$ reduced over $\kappa(H)$", where $H$ usually refers to the Hessian matrix and $\kappa(\cdot)$ is the condition number (smaller is better).
Unfortunately, there is no guarantee of
$\kappa(DH)\leq\kappa(H)$ and this inequality often requires strict assumptions on both $D$ and $H$. For instance, $\kappa(DH)$ would be small if $H$ is close to diagonal and $D$ is a cleverly designed compressor of $H$
(Forsythe & Straus, ; Young, ; Sun & Ye, ; Qu et al., ).

Here, we numerically explore the effectiveness of Adam’s preconditioner within each dense Hessian sub-block. We generate a random dense PD matrix $H_{b}\in\mathbb{R}^{d\times d}$ and use it as a proxy for the dense Hessian sub-block of neural nets in Figure LABEL:fig:block_diagonal.
We define $D_{\text{Adam}}=\operatorname{Diag}({1}/{\sqrt{v}})$, where $v=g\odot g$, $g=H_{b}x\in\mathbb{R}^{d}$, and each entry $x_{i}\sim\mathcal{N}(0,1/{\sqrt{d}})$ follows Xavier initialization. We explore the interplay between the following two metrics:

$$ $\tau=\frac{\sum_{i}|H_{b,i,i}|}{\sum_{i,j}\left|H_{b,i,j}\right|},\quad r= \frac{\kappa(D_{\text{Adam}}H_{b})}{\kappa(H_{b})},$ (3) $$

where $\tau\in[0,1]$ is the “diagonal-over-off-diagonal ratio", and we use it to measure how dense $H_{b}$ is ($H_{b}$ is pure diagonal when $\tau=1$).
$r\geq 0$ measures the effectiveness of Adam’s preconditioner $D_{\text{Adam}}$ when operating on the Hessian-block $H_{b}$ (the smaller the better). We investigate the change of $r$ when changing the structure of $H_{b}$, including changing $\tau$, dimension $d$, and also $\kappa(H_{b})$.
We emphasize that for a fixed $d$ or $\kappa(H_{b})$, we change $\tau$ by only rotating the eigenvectors, but not changing the eigenvalues of $H_{b}$. This ensures that $\tau$ is the only changing factor in the experiments.

We summarize the key findings in Figure LABEL:fig:kappa_phase_transition: for $H_{b}$ with most dimension $d$ and $\kappa(H_{b})$, $r$ decreases as $\tau\rightarrow 1$. That is, $D_{\text{Adam}}$ is effective when $H_{b}$ is close to diagonal, and $D_{\text{Adam}}$ is not so effective when $H_{b}$ is dense. This aligns with the convergence rates in Figure LABEL:fig:random_quadratic.
It is intriguing to provide a lower bound on $\kappa(D_{\text{Adam}}H_{b})$ to ground the observation in Figure LABEL:fig:kappa_phase_transition, and we are not aware of any existing lower bound of this kind. Note that it is rather difficult to characterize $\kappa(D_{\text{Adam}}H)$, partially because the extreme eigenvalues are neither sub-additive nor sub-multiplicative (Kittaneh, ).
We leave it as an important but challenging future direction. To summarize, for the dense Hessian-blocks,
it is possible to outperform Adam with only one good learning rate.

Case study II: Transformers.
The above analysis suggests there is room to cut down the number of learning rates.
We also observe similar phenomena in Transformers.
We consider a 4-layer Transformer
and under the PyTorch default partition, and we
randomly choose one parameter block as the “left-out" block and change the coordinate-wise learning rate to a single-learning rate counter-part. We use Adam for the rest of the blocks. We grid-search the learning rate for the left-out block and apply the cosine decay schedule like the rest of the blocks. We report the best result and call this method “Adam (leave-one-out)". Figure LABEL:fig:leave_one_out shows that Adam (leave-one-out) can achieve similar or better performance than Adam. A similar phenomenon is also observed when we randomly leave out up to three blocks and search three learning rates.
We do not explore the possibility of leaving more blocks out since the cost of grid search grows exponentially.

To summarize this section, we find that it is possible to reach similar or better performance with much fewer learning rates than Adam. The remaining issue is how to find them without grid-search.
In the next part, we propose a simple and effective method called Adam-mini, which could bring comparable or even better performance than Adam, but with 90% to 99% fewer learning rates.

### 2.2 Proposed Method: Adam-mini

Adam-mini aims to cut down the learning rate resources in Adam without laboriously grid-searching the learning rates as in Section 2.1. Since the procedure of Adam-mini is relatively long, we relegate the complete form of Adam-mini
in Algorithm 1, 2, 3 in Appendix B. We summarize the key designs here.
Adam-mini consists of two steps.

Step 1 Partition the model parameters into blocks
by Hessian structure.
We discuss the partition principle later in Section 2.3.
For different architectures, the principle will be realized in
different forms; see Algorithm 2: “Partition for non-Transformers".
and Algorithm 3: “Partition for Transformers".

Step 2. For each parameter block, we use a single learning rate. To efficiently choose a suitable learning rate in each block, Adam-mini simply replaces $\texttt{g}\odot\texttt{g}$ in vanilla Adam by its mean value.
We adopt the moving average on these mean values as in Adam.

A simple example of Adam-mini.
We use a simple example to illustrate the key design of Adam-mini.
For a problem with 5 parameters $w\in\mathbb{R}^{5}$,
Adam and Adam-mini both perform $w=w-u\odot m$,
where $m$ is the 1st-order momentum and $u$ has different forms as follows:

- •
For Adam: $u_{\text{Adam}}=\big{(}\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}}}
\big{)}.$
- •
For Adam-mini: suppose the partition is $(1,2,3)$ and $(4,5)$ then
$u_{\text{mini}}=\big{(}{\color[rgb]{1,0.4,0.0}\definecolor[named]{
pgfstrokecolor}{rgb}{1,0.4,0.0}\frac{\eta}{\sqrt{(v_{1}+v_{2}+v_{3})/3}},\frac
{\eta}{\sqrt{(v_{1}+v_{2}+v_{3})/3}},\frac{\eta}{\sqrt{(v_{1}+v_{2}+v_{3})/3}}
},{\color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}\frac{\eta
}{\sqrt{(v_{4}+v_{5})/2}},\frac{\eta}{\sqrt{(v_{4}+v_{5})/2}}}\big{)}.$

Note that the number of effective elements $u_{\text{mini}}$ equals the number of blocks, which could be significantly smaller than that of $u_{\text{Adam}}$, which equals the number of parameters. For LLMs, we will show that this would free $\geq 99.9\%$ elements in $v$ (will be discussed later in Section 2.4).

### 2.3 Principle for the Partition Strategy

We now discuss how to choose the parameter partition for Adam-mini. When developing the method, we tried the PyTorch default partition. In Section 3, we will show that this simple partition works well on small and middle-scaled non-Transformer tasks, including ResNet, diffusion model, and graph models. We present this strategy in Algorithm 2: “Partition for non-Transformers".

Unfortunately, we find that the PyTorch default partition does not always work well on larger-scaled tasks including Transformer training. In particular, we find that Adam-mini encounters training instability on 1B models (see Figure LABEL:fig:babygpt_hessian_plot (i)).
We suspect this is because the PyTorch default partition did not fully capture the Hessian structure.
We propose a general principle in Principle 1 below (^4^44We assume the Hessian rows & columns are properly re-ordered s.t. the block-diagonal structure is visible.).

Principle 1: We should partition parameters into blocks, such that each parameter block is associated with the smallest dense sub-blocks in the Hessian.

Principle 1 comes from the analysis in Section 2.1: it is possible to harmlessly reduce the number of Adam’s learning rates within each dense Hessian block. However, if the partition is too coarse and violates Principle 1, we might accidentally remove some crucial learning rates and oversimplify the problem, causing training failure.
It is important to follow Principle 1 since it is necessary to use (at least) one distinct learning rate for each Hessian block (as evident in Appendix D.2).

Does the PyTorch default partition follow Principle 1?
To find out, we explore the Hessian of a small Transformer as in Figure LABEL:fig:babygpt_hessian_plot. Under the PyTorch default partition, we compute the exact Hessian for each parameter block at 1% training step.
We find there are four classes of Hessian sub-blocks.

- •
Class 1: query and key. For query and key, the Hessian sub-block itself has a block-diagonal structure and consists of smaller dense matrices. We empirically find that the number of small dense sub-blocks equals the number of heads in multi-head attention.
- •
Class 2: value. For value, the structure of Hessian sub-blocks seems less clear. It seems to have the hint of 16 diagonal blocks (16 is the number of output neurons), but the pattern is less obvious. This Hessian structure is significantly different from that of query and key, although they all consist of four heads. The Hessian entries of value are also about $10^{5}$ larger than those of query and key (^5^55This might be one source of the heterogeneity of Hessian eigenvalues as reported by (Zhang et al., ).). One possible reason is that value is positioned outside the softmax operator in the self-attention design, while query and key are not.
- •
Class 3: attn.proj and MLPs. For attn.proj, mlp.fc_1, and mlp.proj, the Hessian sub-block has a block-diagonal structure and the number of small dense sub-blocks equals the number of output neuron. This observation suggests that the calculation in Eq. (1) not only holds in pure MLPs, but also can generalize to various building blocks in Transformers.
- •
Class 4: embed_layer and output_layer. For these two layers, the Hessian sub-block has a block-diagonal structure and the number of small dense sub-blocks equals the number of tokens.

Based on the above findings, we find that the PyTorch default partition is indeed not the best fit for Transformers. By Principle 1, query and key should be further partitioned by heads; value, attn.proj, and MLPs should be partitioned by output neurons; embed_layer and output_layer should be partitioned by tokens. As for value, the Hessian shows the hint of 16 diagonal blocks (where 16 is the number of output neurons), but the pattern is less clear. Our experiments show that “partition value by output neurons" works well in general, yet there are also some special cases where it is better to “treat value as a whole" (see discussions in Appendix C.3). By default, we will partition value by output neurons.

We then introduce the resulting Algorithm 3: “Partition for Transformers".
As shown in Figure LABEL:fig:babygpt_hessian_plot (d). This strategy indeed stabilizes the training and boosts the performance.

### 2.4 Some Characteristics of Adam-mini

Memory cut down. Adam-mini reduces the number of learning rates from the number of model parameters to
the number of total number of blocks by our partition strategies.
As a result,
Adam-mini cuts down more than $99.9\%$ of Adam’s $v$, which saves $50\%$ of Adam’s memory.

Higher throughput.
Adam-mini can reach a higher throughput than AdamW, especially under limited GPU resources. There are two reasons.
First, Adam-mini does not introduce extra computation in its update rules. The averaging operation in Algorithm 1 incurs negligible cost and it significantly reduces the number of vector-square-root and vector-division operations in AdamW. Second, thanks to the memory cut-down, Adam-mini can support larger batch sizes per GPU. It also reduces the communication among GPUs, which is known to be a major overhead (Rajbhandari et al., ).

Owing to these properties, Adam-mini could reduce the overall time for pre-training.
We provide evidence in Table 2.
When pre-training Llama 2-7B on $2\times$ A800-80GB GPUs, we find Adam-mini could reach 49.6% higher throughput than AdamW. This translates to ${\bf 33.1\%}$ reduction of wall-clock time on processing the same amount of tokens for pre-training.

Why averaging over $v$ can work. We now provide an intuitive explanation of why “averaging $v$ by output neuron" can be a reasonable choice.
For one data sample, the gradient of the weight matrix $W\in\mathbb{R}^{d\times d}$ can be expressed as $G:=\frac{\partial\ell}{\partial W}=ez^{\top}\in\mathbb{R}^{d\times d}$, where $e$ is certain backpropagation (BP) error vector and $z$ is the input feature to the current weight. For all entries in the $i$-th row of $G$, they all share the same BP error term $e_{i}$, which is usually non-negligible when $G\neq 0$. Therefore, $G$ usually has similar entries within a row (which associates with the same output neuron), and its mean value can be a good representative of the whole row.

**Table 1: Memory cost of AdamW v.s. Adam-mini. Calculation is based on float32, which is a standard choice for optimizer states.**
| Model | Optimizer | Memory (GB) |
| --- | --- | --- |
| GPT-2-1.5B | AdamW | 12.48 |
| GPT-2-1.5B | Adam-mini | 6.24 $({\bf 50\%}\downarrow)$ |
| Llama 2-1B | AdamW | 8.80 |
| Llama 2-1B | Adam-mini | 4.40 $({\bf 50\%}\downarrow)$ |
| Llama 2-7B | AdamW | 53.92 |
| Llama 2-7B | Adam-mini | 26.96 $({\bf 50\%}\downarrow)$ |
| Llama 3-8B | AdamW | 64.24 |
| Llama 3-8B | Adam-mini | 32.12 $({\bf 50\%}\downarrow)$ |
| Llama 2-13B | AdamW | 104.16 |
| Llama 2-13B | Adam-mini | 52.08 $({\bf 50\%}\downarrow)$ |

**Table 2: Throughput ($\uparrow$) test on $2\times$ A800-80GB GPUs for Llama 2-7B pre-training. ✗means out of memory. GPU hours ($\downarrow$) to pre-train Llama 2-7B with the optimal token amount by Chinchila’s law.**
| Optimizer | bs_per_GPU | total_bs | Throughput ($\uparrow$) |
| --- | --- | --- | --- |
| Adam-mini | 4 | 256 | 5572.19 $(\uparrow{\bf 49.6\%})$ |
| AdamW | 2 | 256 | ✗ |
| AdamW | 1 | 256 | 3725.59 |

Has room to improve. Adam-mini designs the learning rate for each dense Hessian sub-block using the average of Adam’s $v$ in that block. Such a design achieves cheap computation, but it might not be optimal.
We believe there is great room to improve the learning rate design. As shown in Figure LABEL:fig:random_quadratic, we can reach much faster convergence if we utilize more information in the dense block to design the learning rate (e.g., using eigenvalues of each block), However, such a design requires expensive computation. We leave it as an important future direction.

## 3 Experiments

We now verify the efficacy of Adam-mini on two types of neural-net tasks: (1) LLM tasks including pre-training, supervised fine-tuning (SFT), and reinforcement learning from human feedback (RLHF). (2) Non-LLM tasks including vision, graph, and diffusion model training. Due to the limited space, we primarily focus on LLM tasks in this section, and we relegate the non-LLM tasks to Appendix C.1. All LLM experiments are conducted on four NVIDIA A800-80GB GPUs and the rest are conducted on four V100 GPUs.
All the experimental details are explained in Appendix E.1.

### 3.1 Pre-training

Setups. We pre-train LLMs including GPT-2 series and Llama series. We train these models on mainstream English Corpus from scratch. In particular, We train GPT-2 (Radford et al., ) series (125M to 1.5B) on Openwebtext (Gokaslan et al., ). We train Llama series (20M to 13B) (Touvron et al., ) on C4 (Raffel et al., ).
We compare Adam-mini with AdamW (Loshchilov & Hutter, ) as well as popular memory-efficient methods including Adafactor (Shazeer & Stern, ), CAME (Luo et al., ), and SM3 (Anil et al., ). For Adafactor and SM3, we incorporate momentum with $\beta_{1}=0.9$ to ensure a fair comparison with other methods.
We tune the learning rate for all methods, using the same tuning budget for each, and report the best performance.

GPT-2 series. Figure LABEL:fig:gpt shows the results for GPT-2 series pre-training. We find that Adam-mini performs similarly to AdamW with 50% less memory, while other methods perform worse.
In Figure LABEL:fig:gpt (a),
we run *Adam-mini (PyTorch default partition)*, which partition parameters by PyTorch default partition.
We find that *Adam-mini (PyTorch default partition)* performs poorly. We stop the trial since it shows clear unstable behavior. In Figure LABEL:fig:gpt_train, we further present the training loss curves. We find that the loss curves of Adam-mini closely resemble those of AdamW.

Llama series. Figure LABEL:fig:tinyllama shows the results for pre-training Llama series.
We also train Llama 2-7B as shown in Figure LABEL:fig:intro (c) in Section 1.
We find that Adam-mini performs on par with AdamW, while other methods do not.
Further, Adam-mini’s loss curves closely resemble the curves by AdamW.

Trajectory comparison. On a small Transformer, Adam-mini generates similar trajectories
to that of AdamW, while other methods cannot. This can be seen in Figure LABEL:fig:gpt_train (b) and the detailed description is in Appendix E. This might be because Adam-mini
makes fewer modifications over AdamW.

Sensitivity analysis.
On GPT-2-125M pre-training task, we test the sensitivity of Adam-mini to hyperparameters. We report the validation loss after training with 2.5B tokens (by Chinchilla’s law). As shown in Figure LABEL:fig:sft (c), Adam-mini seems not overly sensitive to hyperparameters.

### 3.2 Scaling Laws of Adam-mini

We now show the efficacy of Adam-mini through scaling law experiments. We use C4 dataset to pre-train the Llama 2 architecture from 39M to 1B. For the model with size $n_{\text{param}}$, we train the model with about $20*n_{\text{param}}$ tokens, which is suggested to be the optimal amount by Chinchilla’s law (Hoffmann et al., ). The largest-scaled experiment we conducted is Llama 2-1B pre-training with 26.2B tokens, which takes about 170 GPU hours on $4\times$ A800-80GB GPUs. The total running time for the scaling law experiments is about 300 GPU hours.

As shown in Figure LABEL:fig:scaling_law, Adam-mini’s loss curves are consistently similar to AdamW. We also present the final validation perplexity and find that Adam-mini reaches a slightly lower perplexity than AdamW for all models (see Figure LABEL:fig:scaling_law (b), also see Table 5 in Appendix C).
The fitted lines in Figure LABEL:fig:scaling_law (b) suggest that Adam-mini can be scaled up to larger models (if the scaling law holds).

### 3.3 Supervised Fine-tuning and RLHF

We now run Adam-mini on downstream tasks including SFT and RLHF. We use the Llama 2-7B pretrained model (Touvron et al., ) for our study. We use the ultrafeedback dataset
and implement the RLHF workflow from (Ouyang et al., ). We use ReMax (Li et al., ), a memory-efficient alternative to PPO (Schulman et al., ), to optimize the preference reward.

We evaluate the alignment performance in terms of chat ability using the MT-Bench (Zheng et al., ), where GPT-4 assesses multi-turn chatting capabilities and assigns a score from 0 to 10 (higher is better). Our results, presented in Table 3, demonstrate that Adam-mini can outperform AdamW.

**Table 3: Averaged GPT-4 evaluation score ($\uparrow$) of SFT and RLHF on the MT-Bench.**
|  | 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.4 Detailed Comparison with Adafactor

We now carefully compare Adam-mini and the popular memory-efficient optimizer Adafactor.
Besides the original Adafactor, we also consider a modified version in (Zhai et al., ), which we call “Adafactor-Zhai-version". For both versions, we use momentum with $\beta_{1}=0.9$.

We first conduct learning rate grid-search on Llama 2-20M and train it following Chinchilla’s law. As shown in Figure LABEL:fig:adafactor
(a), we find that Adafactor-Zhai-version improves over the original version, but both versions of Adafactor are still consistently worse than Adam-mini. We further sweep over other hyperparameters including (1) $\beta_{2}=0.95$; (2) $\epsilon=\{10^{-30},10^{-16},10^{-8},10^{-6}\}$; (3) warm-up steps = $\{1\%,2\%,3\%,4\%,5\%,10\%\}$ total steps. The results are shown in Appendix C.4.
We find that the change of hyperparameters does not significantly boost the performance of Adafactor, and both versions still underperform Adam-mini.

We further sweep hyperparameters on Llama 2-1B. In contrast to the case of Llama 2-20M, we find that the Adafactor-Zhai-version now suffers from training instability and the original version performs better. Nevertheless, they still underperform Adam-mini. In Appendix C.5, we conduct a similar hyperparameter search for Lion (Chen et al., ) and we find it
also underperforms Adam-mini.

About hyperparameter tuning. We acknowledge that it might be possible to improve these methods if we spend more resources on grid search (as claimed by a recent work (Zhao et al., )). However, based on our experience so far, it is not easy to tune these methods, and to our knowledge, there is no much open-source guidance. Recall that there are 9 tunable hyperparameters in Adafactor, so it is rather non-trivial to find the correct combination. In contrast, Adam-mini is much easier to use. In all our experiments, Adam-mini performs well using the same hyperparameters as AdamW (including learning rate, $\beta_{1},\beta_{2},\epsilon$, etc.).

Throughput comparison. Besides the performance comparison, we further find that Adafactor has a higher latency than Adam-mini (Figure LABEL:fig:adafactor (c)). This is primarily due to two reasons. First, Adam-mini only requires computing the mean by rows of the weight matrix, whereas Adafactor needs to sum across both the rows and the columns. Second, the dimension of $v$ in Adam-mini equals the output dimension or the number of heads, which is significantly smaller than the dimension of $v$ in Adafactor, which equals the product of the input and output dimension. Consequently, Adam-mini saves computation when taking the square root of $v$. As such, Adam-mini reaches a higher throughput.

Summary of Section 3. Finally, we summarize three key observations from all the experiments above. For all the models we tried, we observed that:

1. Adam-mini performs on par with AdamW with 50% less memory.

2. Adam-mini performs well using the same hyperparameters as AdamW.

3. Adam-mini’s loss curves closely resemble those of AdamW.

## 4 Concluding Remarks

We proposed Adam-mini, an optimizer that saves 50% memory of Adam. We remark that there is great room to improve the design of Adam-mini: currently Adam-mini uses a simple and cost-effective way to design a learning rate for each dense Hessian sub-block, but it might not be an optimal way. We leave the development of stronger designs as a future direction.

## Appendix A Related works

Understanding of Adam. There is an active line of works trying to understand why Adam works well (Zhang et al., ; Wu et al., ; Shi et al., ; Zhang et al., ; Wang et al., ; Pan & Li, ; Jiang et al., ; Kunstner et al., ; Zhang et al., ; ; Kunstner et al., ). In contrast to these works, we point out that Adam’s $v$ might not function at its full potential as effectively as we expected: sometimes fewer learning rates can reach the same or better results (due to the dense Hessian sub-blocks). Our findings might motivate stronger optimizers that better fit the neural-net Hessian structure.

Similarly to in our Section 2.1, a recent work (Das et al., ) also explores the effectiveness of Adam’s preconditioner $D_{\text{Adam}}$ from a linear algebra perspective. They focus on (a variant of) Adam and prove the following result: First, for diagonal dominant (DD) matrix $H_{b}$ when the dimension $d$ is less than $\kappa^{1/3}$, their modified version of Adam exhibits a faster convergence rate compared to gradient descent (GD); Second, for non-DD matrix, the constant terms in Adam’s upper bound can be much larger than that of GD. Their results take a valuable and important step towards understanding $D_{\text{Adam}}$. However, these results cannot fully support our numerical findings in Figure LABEL:fig:kappa_phase_transition. This is because they only provide an upper bound for the non-DD case, while we need a lower bound. We note that it is rather difficult to derive the desired lower bound, and we leave it as a future direction.

On the Hessian of Neural Nets. Hessian matrix is crucial for the behaviors of gradient methods.
There are several important attempts to study the Hessian of MLPs and CNNs (Collobert, ; Roux et al., ; Martens & Grosse, ; Sagun et al., ; ; Chaudhari et al., ; Papyan, ; Wu et al., ; Liao & Mahoney, ; Papyan, ; ; Sankar et al., ; Gur-Ari et al., ; Yao et al., ; Zhang et al., ; Ghorbani et al., ; Yao et al., ; Dauphin et al., ). Inspired by these works, we explore the Hessian structure of Transformers and
connect it to the behaviors of Adam. We then find room to improve and propose to slim down Adam into Adam-mini.

Lightweight optimizers for general tasks.
There are several attempts to reduce the memory cost of Adam. Adafactor (Shazeer & Stern, ) and its variant CAME (Luo et al., ) conduct nonnegative low-rank factorization over Adam’s $v$.
SM3 (Anil et al., ) is a lightweight version of AdaGrad (Duchi et al., ). SM3 chooses the learning rate of the $i$-th parameter by taking
the minimal value in a certain candidate set, and each element in the candidate set is related to the maximal squared gradient under a predetermined cover.
All these aforementioned methods could release almost all memory for $v$ and save about 48% of Adam’s memory. However, we find that their performance degenerate in various experiments, while Adam-mini maintains as effective as AdamW (as shown in Section 3).

After completing this work, we noticed two methods that share some of the ideas of Adam-mini:
BAGM (Zheng & Kwok, ) and NovoGrad (Ginsburg et al., ). Both of them use block-wise or layer-wise adaptive learning rates to achieve robust performance and better generalization. We summarize their key differences with Adam-mini.
BAGM partitions parameters to reach minimal-norm solutions and achieve provable robustness. In particular, their theory in Proposition 1 states that layer-by-layer parameter partition can lead to minimum $\ell$-2 norm solutions. Aligning with the theory, they find that the PyTorch default partition (BAGM-B.1) indeed brings overall the best performance on both CIFAR-10 and ImageNet. Although the PyTorch default partition may have benefits on robustness, we find that it overlooks the Hessian structure and oversimplifies the training problem for Transformers (as we discussed in Section 2.3). As a result, the PyTorch default partition will lead to training instability in large-scale LLMs, and this is evident in our failed preliminary versions of Adam-mini in Figure LABEL:fig:babygpt_hessian_plot and LABEL:fig:gpt. We then propose a new partition strategy Algorithm 3 which partition parameters by the smallest dense Hessian sub-blocks. For Transformers, Algorithm 3 uses different strategies for different building blocks (e.g., partition the embedding layer by tokens, and partition Query by heads) and we find that Algorithm 3 is necessary to stabilize the training.

As for NovoGrad, it also uses a layer-wise learning rate design (by PyTorch default partition)
and thus face the same training instability issues as BAGM. Further, NovoGrad introduces a different design to 1st-order momentum: instead of performing weighted-sum on the past gradients, it performs weighted-sum on “the current gradient divided by the 2nd-order momentum". Such design is largely different from AdamW and Adam-mini. It remain unclear whether this design can work on large-scale tasks like LLMs.

Besides algorithmic design, our work also provides new understandings of Adam, and particularly, how Adam behaves on generic optimization problems with near-block-diagonal Hessian.
We also provide new findings on the Hessian structure of Transformers and provide new principles for designing better algorithms.

Other orthogonal methods.
The idea of Adam-mini can be orthogonally combined with various existing methods. We list two most relevant examples here.

- 1.
GaLore (Zhao et al., ) is a new memory-efficient optimizer for LLMs. Given a gradient matrix $g$, GaLore calculates a low-rank gradient estimator $\hat{g}$ and then calculates $m$ and $v$ based on this $\hat{g}$. Adam-mini can potentially be combined with GaLore to reach further memory reduction on $v$. The combined method, e.g., “GaLore-mini", can further reduce about $40\%$ memory on GaLore and about $81\%$ on AdamW in total.(^6^66These results are calculated based on (Zhao et al., , Table 1). We consider Llama 2-7B and $r=1024$ in GaLore.) Additionally, GaLore-mini can ease the offload burden and enhance the throughput of GaLore, especially when training on customer-level GPUs with limited memory.
- 2.
Sophia (Liu et al., ) is another recent diagonal preconditioned optimizer. Just as Adam, Sophia requires memory for $m$ and $v$. It is possible to combine Adam-mini and Sophia to get “Sophia-mini", which saves up to 50% of memory in Sophia. Sophia-mini can also enhance throughput and further speed up Sophia on wall-clock time as in Table 2.

We list more potential combinations here.
LoRA (Hu et al., ) is a memory-efficient method for SFT tasks. This method fine-tunes the model via additive low-rank adaptors and uses Adam to update these adaptors. Note that the Adam steps in LoRA can be replaced by Adam-mini. As a result, Adam-mini brings better performance (Figure 16).
In parallel to our work, BAdam (Luo et al., ) conducts SFT in a block-coordinate-descent (BCD) fashion. This method requires repeated Adam steps to solve the sub-problem in BCD. Similarly as in LoRA, the Adam steps in BAdam can be replaced by Adam-mini to further reduce memory. Nero optimizer (Liu et al., ) also cuts down the memory of Adam. It removes the 1st-order momentum and uses a neuron-specific projected gradient-style update. According to (Liu et al., ), their design imposes constraints on weight matrices and has the advantage of “balanced excitation and inhibition". Such design can potentially be combined with Adam-mini to further boost performance. To save the memory cost for fine-tuning LLMs, MeZO (Malladi et al., ) uses zeroth-order methods to approximate the gradient information. It is possible to combine this idea with Adam-mini to further save memory for SFT. Adam-mini can also potentially be combined with other diagonal preconditioned methods (such as AdaGrad (Duchi et al., ) and Adan (Xie et al., )) as well as recent schedule-free optimizers such as SchedulefreeAdamW (Defazio et al., ).

There are several other tricks that ease GPU memory burden but are orthogonal to optimizer design. These tricks include gradient checkpointing (Chen et al., ), model offloading and sharding (Rajbhandari et al., ; ), quantization (Dettmers et al., ; Li et al., ), and fused update (Lv et al., ; ).
Adam-mini can be implemented upon these tricks.

Finally, we discuss another popular adaptive optimizer called LAMB (You et al., ) (see Algorithm 7 in Appendix D.1). LAMB might be misunderstood as a similar optimizer to Adam-mini, but actually, it is not.
We emphasize that Adam-mini is significantly different from LAMB.
First, LAMB still keeps the same coordinate-wise learning-rate design $1/\sqrt{v}$ as in Adam. Second, in addition to this $1/\sqrt{v}$, LAMB re-scales the parameters in a layer-by-layer fashion. This re-scaling design is often known as the “layer-wise learning rates", but to be precise, it is actually an additional “layer-wise scaling" besides the “coordinate-wise learning rates $1/\sqrt{v}$". As a result, LAMB does not save memory over Adam and its overall design is quite different from Adam-mini. This is understandable because LAMB is designed for large-batch training, not for memory saving. Numerically, we find that LAMB performs worse than Adam-mini on GPT2 pre-training (Figure LABEL:fig:tinyllama).

## Appendix B The Complete Form of Adam-mini

We now present the complete form of Adam-mini. To be precise, Algorithm 2 should be renamed as “Partition for CNNs, Diffusion models, and Graph Neural Networks", since we have only tested Algorithm 2 on these models. In the future, it is possible that we will have more complicated non-Transformer architectures on which Algorithm 2 fails. In those cases, we need to investigate the Hessian structure of these new architectures (like what we did for Transformers) and then develop the concrete partition algorithms following our Principle 1 in Section 2.2.

Figure: Algorithm 1 Adam-mini in Pytorch style

Figure: Algorithm 2 Partition for non-Transformers

Figure: Algorithm 3 Partition for Transformers

## Appendix C More Results

### C.1 Non-LLM Tasks

We now evaluate Adam-mini on non-LLM tasks.
Table 4 shows the results for training ResNet18 on ImageNet, a diffusion model on CelebA, a Graph Convolution Net (GCN), and a Graph Attention Net (GAT) on OGB-arxiv. We find that Adam-mini performs comparable or better than AdamW, but with 50% less memory. Here, we simply apply Algorithm 2 to Adam-mini and find it is sufficient to bring good performance. It is possible that Adam-mini would perform better under a more fine-grained partition strategy of CNNs, GNNs, etc., but it requires investigation into the Hessian structure. We leave it for future investigation.

**Table 4: On popular non-LLM tasks, Adam-mini performs on par or better than AdamW with 50% less memory.**
| Domain | Model | Optimizer | Metric | 25% steps | 50% steps | 75% steps | 100% steps |
| --- | --- | --- | --- | --- | --- | --- | --- |
| 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 |
| 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 |
| 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 |

### C.2 More Results on the Scaling Law Experiments

The complete loss curves of Llama 2-1B. We here present the complete validation loss curve of Llama 2-1B, training on  20B tokens, which corresponds to the rightmost curve in the scaling law experiments in Figure LABEL:fig:scaling_law (a). We note that this is a complete pre-training run under the definition of Chinchila’s law (Hoffmann et al., ). We find that Adam-mini’s loss curves closely resemble those of AdamW.

Figure: Figure 12: Loss curves of pre-training Llama 2-1B on 26B tokens. This is a complete pre-training run under the definition Chinchila’s law. We find that Adam-mini performs similarly to AdamW throughout the training, but with 50% less memory.
Refer to caption: x1.png

The final validation perplexity. In Table 5, we present the final validation perplexity for all models after training on the token amount suggested by Chinchilla’s law (Hoffmann et al., ). For all models from 39M to 1B, we find that Adam-mini reaches a lower validation perplexity than AdamW.

**Table 5: Final validation perplexity for all models pre-trained by Adam and Adam-mini. The token amount follows Chinchila’s law. After pre-training, Adam-mini reaches a lower validation perplexity than AdamW.**
| Model size | 39M | 67M | 102M | 162M | 271M | 1B |
| --- | --- | --- | --- | --- | --- | --- |
| Total tokens | 1.02B | 1.76B | 2.67B | 4.25B | 7.10B | 26.21B |
| AdamW | 40.795 | 29.319 | 24.670 | 20.360 | 17.178 | 12.452 |
| Adam-mini | 40.407 | 29.014 | 24.192 | 20.172 | 17.035 | 12.372 |

### C.3 More Discussions on the Partition Strategies of Value

As shown in Figure LABEL:fig:babygpt_hessian_plot, the Hessian structure of value is less clear compared to other blocks: it shows the hint of 16 diagonal blocks (where 16 is the number of output neurons), but the pattern is not that clear. This gives rise to two potential partition strategies: (I) partition by output neuron; (II) treat as a whole. Numerically, we find that strategy (I) works well when the number of total training steps is large. This includes most of our experiments such as GPT-2 in Figure LABEL:fig:gpt (with more than 50k total steps) and the scaling law experiments of Llama models in Figure LABEL:fig:scaling_law (e.g., Llama 2-1B is trained with more than 200k total steps). On the other hand, we find that strategy (II) works better when the number of total training steps is small. This includes our Llama experiments with 10k total steps in Figure LABEL:fig:tinyllama.

Based on these findings, we recommend using strategy (I) when the total number of training steps is large, and using strategy (II) if otherwise. Note that strategy (II) can be used simply by adding one line of code after creating the optimizer: optimizer.wv_names = {}.

### C.4 Detailed Comparison with Adafactor

In this section, we conduct a more hyperparameter search for Adafactor on Llama 2-20M pre-training. We will focus on tuning Adafactor-Zhai-version since it performs better than the original Adafactor (see Figure LABEL:fig:adafactor). We consider the following three setups.

- •
Setup 1: We change the default $\beta_{2}=0.999$ to $\beta_{2}=0.95$ and sweep over learning rates.
- •
Setup 2: We use learning rate $=$ 5e-3, $\beta_{2}=0.95$ and sweep over warm-up step = $\{1\%,2\%,3\%,4\%,5\%,10\%\}$ total steps.
- •
Setup 3: We use learning rate $=$ 5e-3, $\beta_{2}=0.95$ and warm-up step $=1\%$ total steps and sweep over $\epsilon=\{10^{-30},10^{-16},10^{-8},10^{-6}\}$.

The results are shown in Figure LABEL:fig:adafactor_appendix. In all these cases, Adafactor-Zhai-version consistently underperforms Adam-mini and the change of hyperparameters does not help much.

### C.5 Detailed Comparison with Lion

Similarly to Adafactor, we conduct the hyperparameter grid search over Lion. We consider the following three setups.

- •
Setup 1: We consider Llama 2-20M pre-training (the same setting as in Appendix C.4). We use the default hyperparameter $\beta_{1}=0.9,\beta_{2}=0.99$ and sweep over learning rates.
- •
Setup 2: We consider Llama 2-20M pre-training (the same setting as in Appendix C.4). We change the hyperparameters to $\beta_{1}=0.95$, $\beta_{2}=0.98$ and sweep over learning rates. This set of hyperparameters is claimed to be “helpful in mitigating instability during training" and is recommended by the authors(^7^77[https://github.com/lucidrains/lion-pytorch](https://github.com/lucidrains/lion-pytorch)).
- •
Setup 3: We consider Llama-1B pre-training (the same setting as in Section 3.4). We use the default hyperparameter $\beta_{1}=0.9,\beta_{2}=0.99$ and sweep over learning rates.

The results are shown in Figure LABEL:fig:adafactor_appendix. In all these cases, Lion consistently underperforms Adam-mini.

### C.6 Additional Findings on GPT-2-330M

In Figure 15, we report some unexpected results on GPT-2-330M. When using the recommended learning rate 3e-4 by (Liu et al., ), we find that AdamW encounters loss spikes, while Adam-mini does not. The loss spike of AdamW can be mitigated by changing $\epsilon$ from 1e-8 to 1e-6. We have not fully understood the cause of the loss spike and how Adam-mini prevents it in this experiment. We are not sure whether such benefit of Adam-mini maintains on larger models. We leave more investigation as an interesting future direction.

Figure: Figure 15: On GPT-2-330M, AdamW encounters loss spikes, while Adam-mini does not. The loss spike of AdamW can be mitigated by changing $\epsilon$ from 1e-8 to 1e-6.
Refer to caption: x2.png

### C.7 Combining with LoRA

We note that Adam-mini can be combined with LoRA (Hu et al., ): we can change the Adam steps in LoRA by Adam-mini. The results are shown in Figure 16. We find that such changes with boost performance. We grid-search learning rates for both methods and report the best performance. The detailed learning rate configuration is shown in Section E.

Figure: Figure 16: Training curves of SFT (LoRA). We find that LoRA performs better when replacing the Adam steps by Adam-mini.
Refer to caption: x3.png

### C.8 Sample Responses from LLMs trained by Adam-mini

The sample response in shown in Figure 17.

Figure: Figure 17: Response Sample of Llama 2-7B finedtuned by AdamW and Adam-mini to the #81 test question from the MT-Bench (Zheng et al., ) dataset. Since there is little information in the first round response, we omit the answers from the models.

## Appendix D Some Preliminary Results

### D.1 Preliminaries on Adam, AdamW and LAMB

We introduce Adam, AdamW, and LAMB in Algorithm 5, 6, and 7. These methods need to track $m$ and $v$ along the training. Both $m$ and $v$ are vectors of the same size as # model parameter.

Figure: Algorithm 5 Adam in Pytorch style

Figure: Algorithm 6 AdamW in Pytorch style

Figure: Algorithm 7 LAMB in Pytorch style

### D.2 Preliminary results in (Zhang et al., 2024 )

We here restate (Zhang et al., , Figure 3). This figure shows that: for Transformers, different parameter blocks have different Hessian eigenvalue distributions, while for CNNs, the eigenvalue distributions are similar among blocks. This suggests that Transformers need different learning rates for different blocks to handle the heterogeneity in eigenvalue distributions.

## Appendix E Experimental Details

### E.1 Training configurations for Section 3

Unless mentioned otherwise, we choose
the model configurations by their standard protocols. We choose the learning rates by the recommendation from open-source platforms
if applicable.
For instance, for GPT2 series, we use the recommended learning rates by (Liu et al., ), which are reported to be optimal by grid search.
Unless mentioned otherwise, Adam-mini, Adafactor, CAME, SM3, and LAMB use the same learning rate as the recommended ones of AdamW.
If there is no public recommended learning rate for AdamW, we tune the learning rate for all optimizers within the same computational budget and report the best performance.
For other hyperparameters, we follow the recommendation from open-source platforms or by their default setting.
For SM3 and Adafactor, we incorporate momentum with $\beta_{1}=0.9$ to offer a fair comparison with other optimizers and the rest of the hyperparameters are set as default. The detailed configurations are explained as follows.

#### GPT2 pre-training.

We use the nanoGPT codebase(^8^88[https://github.com/karpathy/nanoGPT/tree/master](https://github.com/karpathy/nanoGPT/tree/master)) to train GPT2 sized 125M (small), 330M (medium),
and 1.5B (XL) on Openwebtext. For all models, we use seq_len $=$ 1024, batch size = 480, weight decay coefficient $\lambda=0.1$, $\epsilon=$ 1e-8, $\beta_{1}=0.9,\beta_{2}=0.95$. We use cosine-decay learning rate schedule with 2000 iterations of warm-up.
For GPT2-small and medium, we use the recommended peak learning rate by (Liu et al., ), which are reported to be the optimal ones found by grid search. For GPT2-XL, we use the recommended peak learning rate by the Levanter(^9^99[https://github.com/stanford-crfm/levanter/blob/e183ec80ec5971b12d4a3fb08a160268de342670/config/gpt2_xl.yaml](https://github.com/stanford-crfm/levanter/blob/e183ec80ec5971b12d4a3fb08a160268de342670/config/gpt2_xl.yaml)).
The chosen peak learning rates are 6e-4, 3e-4, 1e-4 for GPT2-small, medium, XL, respectively. The minimal learning rate is chosen as 3e-5, 6e-5, 1e-5 for these models.

#### Llama pre-training.

For all experiments on the Llama series (from 20M to 13B), we use the Torchtitan codebase(^10^1010[https://github.com/pytorch/torchtitan](https://github.com/pytorch/torchtitan)) and C4 dataset (Raffel et al., ). For all experiments, we use weight decay coefficient $\lambda=0.1$, $\epsilon=$ 1e-8, $\beta_{1}=0.9,\beta_{2}=0.95$. For Llama 2-1B, we use learning rate = 3e-4. For Llama 3-8B and Llama 2-13B, we use learning rate = 1e-4. As for the learning rate schedule, we use warm-up step = 1% total step and use linear decay schedule after the warm-up (this is the default setting in the Torchtitan codebase). For Figure LABEL:fig:tinyllama (a) and all the experiments of Adafactor and Lion, we use seq_len $=$ 512 and batch size $=$ 128. For Figure LABEL:fig:tinyllama (b), we use seq_len $=$ 2048 and batch size $=$ 8. For Figure LABEL:fig:tinyllama (b), we shrink the batch size due to the limited hardware.
For all the scaling law experiments, we use seq_len $=$ 512 and batch size $=$ 256. We summarize the detailed setups for the scaling law experiments in later paragraphs.

For the experiments in Figure LABEL:fig:tinyllama, we use one single learning rate per value block. This is done by adding one line of code after creating the optimizer: optimizer.wv_names = {}. We find that this version performs better when the total number is small (e.g., 10k steps as in Figure LABEL:fig:tinyllama). More discussions can be seen in Appendix C.3.

#### SFT and RLHF.

We use the Llama 2-7B pretrained model (Touvron et al., ) for our study. We use the ultrafeedback dataset
(^11^1111[https://huggingface.co/datasets/argilla/ultrafeedback-binarized-preferences-cleaned](https://huggingface.co/datasets/argilla/ultrafeedback-binarized-preferences-cleaned)). The implementation of SFT and RLHF code is based on the ReMax codebase(^12^1212[https://github.com/liziniu/ReMax](https://github.com/liziniu/ReMax)). Specifically, we train a SFT model with 40% of the chosen data and train a reward model using the remaining 60%. Then,
we apply the reinforcement learning algorithm ReMax (Li et al., ), a memory-efficient alternative to PPO (Schulman et al., ), to optimize the preference reward.

We use DeepSpeed ZeRO-2 in our training. GPT-4 evaluation template in Table 3 is from the codebase(^13^1313[https://github.com/lm-sys/FastChat/tree/main/fastchat/llm_judge](https://github.com/lm-sys/FastChat/tree/main/fastchat/llm_judge)). In the reward optimization stage, We use ReMax, a memory-efficient alternative to PPO. We use UltraFeedback dataset Cui et al. () and use 40% data for SFT and 60% data for ReMax.

SFT. We use 80 samples in a batch and train the model for 3 epochs. For the full parameter tuning, we search the learning rate from {1e-6, 2e-6, 3e-6, 4e-6, 5e-6, 1e-5, 2e-5} based on validation loss, and we use 2e-6 with cosine annealing for both AdamW and Adam-mini. For LoRA, We apply LoRA for all layers except the embedding layer. The rank of LoRA is set to 128. After selecting the learning rate from the same set as the full parameter tuning, we use 2e-5 for both AdamW and Adam-mini when LoRA is applied. The weight decay coefficient is set to 0 as recommended by LlamaFactory(^14^1414[https://github.com/hiyouga/LLaMA-Factory](https://github.com/hiyouga/LLaMA-Factory)). The rest of the hyperparameters of AdamW and Adam-mini are $\epsilon=$ 1e-8, $\beta_{1}=0.9,\ \beta_{2}=0.95$.

ReMax. We use 48 samples in a batch and train the model for 1 epoch. By searching the peak learning rate from {5e-7, 1e-6, 2e-6} based on validation reward, AdamW uses 1e-6 while Adam-mini selects 5e-7 as the peak learning rate. The weight decay coefficient is set to 0. The rest of the hyperparameters of AdamW and Adam-mini are $\epsilon=$ 1e-8, $\beta_{1}=0.9,\ \beta_{2}=0.95$.

#### ResNet.

We use the PyTorch official implementation codebase(^15^1515[https://github.com/pytorch/examples/blob/main/imagenet/main.py](https://github.com/pytorch/examples/blob/main/imagenet/main.py)) to train ResNet18 (He et al., ) on ImageNet (Deng et al., ). We use cosine-decay learning rate, epoch =90, $\beta_{1}=0.9,\beta_{2}=0.999,\epsilon=$1e-8. For ResNet18, we use batch size = 256, peak learning rate = 0.005. For ViT-base, we use batch size = 128, peak learning rate = 0.0001.
These configurations are used for both Adam-mini and AdamW.

#### Diffusion models.

We use the codebase(^16^1616[https://github.com/lucidrains/denoising-diffusion-pytorch](https://github.com/lucidrains/denoising-diffusion-pytorch)) to train diffusion models. The image size is 64 and the training objective is to predict the noise as in (Ho et al., ). We use the default U-Net archiecture hyper-parameters and the dimension multiply in U-Net is (1, 2, 4, 8). We use the CelebA dataset(^17^1717[https://cseweb.ucsd.edu/~weijian/static/datasets/celeba/](https://cseweb.ucsd.edu/~weijian/static/datasets/celeba/)) and train the diffusion model with a learning rate $5\times 10^{-5}$ with cosine decay. The batch size is 128 and the training epoch is 50.

#### Graph Neural Networks.

We use the DGL implementation(^18^1818[https://github.com/dmlc/dgl/tree/master/examples/pytorch/ogb/ogbn-arxiv](https://github.com/dmlc/dgl/tree/master/examples/pytorch/ogb/ogbn-arxiv)) of Graph Convolution Networks (GCN) (Kipf & Welling, ) and Graph Attention Networks (GAT) (Velickovic et al., ) for OGBN-arxiv(^19^1919[https://ogb.stanford.edu/docs/nodeprop/](https://ogb.stanford.edu/docs/nodeprop/)) dataset. All configurations as default. For both Adam-mini and AdamW, we use the default learning rate = 0.005 for GCN and the default learning rate = 0.002 for GAT.

#### Scaling law experiments.

We use the codebase Torchtitan(^20^2020[https://github.com/pytorch/torchtitan](https://github.com/pytorch/torchtitan)) to train Llama models of different sizes. All the model configurations are shown in Table 6 and all the training configurations are shown in Table 7. The experimental setups are inspired by (Hägele et al., ). In all experiments, we fix the warm-up steps to be 1% of the total steps, as suggested by (Ibrahim et al., ).

**Table 6: The model configurations in the scaling law experiments.**
| Model Size | $\text{d}_{\text{model}}$ | $\text{n}_{\text{layers}}$ | $\text{n}_{\text{heads}}$ | seq_len |
| --- | --- | --- | --- | --- |
| 39M | 384 | 8 | 6 | 512 |
| 67M | 512 | 10 | 8 | 512 |
| 102M | 640 | 12 | 10 | 512 |
| 162M | 768 | 16 | 12 | 512 |
| 271M | 1024 | 16 | 16 | 512 |
| 1B | 2048 | 18 | 16 | 512 |

**Table 7: Training configurations for the scaling law experiments.**
| Model | LR | Batch size (# tokens) | Steps | Tokens | Token/Params Ratio |
| --- | --- | --- | --- | --- | --- |
| 39M | 6e-4 | 0.13M | 7.8K | 1.02B | 26.15 |
| 67M | 6e-4 | 0.13M | 13.4K | 1.76B | 26.27 |
| 102M | 6e-4 | 0.13M | 20.4K | 2.67B | 26.17 |
| 162M | 6e-4 | 0.13M | 32.4K | 4.25B | 26.23 |
| 271M | 6e-4 | 0.13M | 54.2K | 7.10B | 26.21 |
| 1B | 2e-4 | 0.13M | 200K | 26.21B | 26.21 |

Trajectory comparison in Figure LABEL:fig:tinyllama (c).
We train a 8-layer Transformer sized 11M on Openwebtext and launch AdamW, Adam-mini, and other memory-efficient optimizers under the same random seed and same learning rate 1e-5. We save the model weights for every 250 iterations and compare their Euclidean distance to the weights along AdamW’s trajectory.

### E.2 Detailed Setup for Other Experiments

Configurations for Figure LABEL:fig:block_diagonal.
We use a synthetic binary classification dataset with 100 samples
using the data generation process as shown below. We use a 1-hidden-layer neural network with an input dimension of 64, a width of 10, and with Tanh activation function. We train the model for 500 steps using Adam with a learning rate of 1e-4, and the model reaches 100% classification accuracy.
With the help of auto-differentiation framework, we calculate the Hessian with two passes of backpropagation (Pearlmutter, ) and the calculation is exact. The experimental setup is motivated by (Zhang et al., ). The Hessian entries are ordered by the row of the weight matrix. For instance, given a linear model $f_{W}(x)=Wx\in\mathbb{R}^{2}$ with input $x\in\mathbb{R}^{3}$ and weight matrix $W=\begin{bmatrix}w_{1}&w_{2}&w_{3}\\
w_{4}&w_{5}&w_{6}\end{bmatrix}\in\mathbb{R}^{2\times 3}$, we first re-order weights as $(w_{1},w_{2},w_{3},w_{4},w_{5},w_{6})$, and then calculate the Hessian by this order.

#### Configurations for Figure LABEL:fig:random_quadratic .

For each dense sub-block $H_{l},l=1,2,3$, we use random positive definite matrices. We fix the choose the eigenvalues of each $H_{l}$ as follows: for $l=1$, we independently sample from {1,2,3} for 30 times; for $l=2$, we repeat this procedure for {99,100,101}; for $l=3$, we repeat this procedure for {4998 ,4999, 5000}.
For the single (blockwise) learning rate method, we use GD with optimal constant learning rate $2/(L+\mu)$, where $L,\mu$ are the largest and smallest eigenvalue of the (blockwise) Hessian.
We use Adam with $\beta_{1}=0$. This helps us focus on the effect of coordinatewise learning rate in Adam. We also set $\beta_{2}=1$ to the time-varying learning rate. This is necessary because, for any $\beta_{2}<1$, Adam with constant learning rate will oscillate on quadratic functions. This is theoretically proved in (Da Silva & Gazeau, , Proposition 12, Figure 1) and empirically observed in (Zhang et al., , Section 3.3).

#### Configurations for Figure LABEL:fig:kappa_phase_transition .

To generate a positive definite matrix $H_{b}$, we first uniformly sample $\frac{d(d-1)}{2}$ independent angles $\theta_{i,j}$ from the interval $[-\frac{\pi}{2},\frac{\pi}{2}]$, where $i<j$. Starting with the identity matrix, we perform a rotation of the $i$-th and $j$-th rows by the angle $\theta_{i,j}$ for each sampled pair. Through $\frac{d(d-1)}{2}$ rotation operations, we obtain the orthogonal matrix $Q$. We define $\Lambda=\text{diag}(\kappa,1,\ldots,1)$, and the matrix $H_{b}$ is generated using the expression $H_{b}=Q\Lambda Q^{\top}$. The python code for $H_{b}$ generation is listed as follows:

We note that as $\theta$ approaches 0, the diagonal-over-off-diagonal ratio $\tau$ of the matrix $Q$ decreases. For the sampled values of $\theta$, we utilize $R\theta$ to produce $H_{b}$ with varying ratios, where $R\in\{\frac{k}{50}|k=0,1,\ldots,50\}$. For each matrix, we sample 100 initial points from the Xavier initialization distribution to compute the resulting $\kappa$ of Adam algorithm. For each pair of $d$ and $\kappa$, we sample 40 different $\theta$ values. By averaging the results obtained, we plot the Figure LABEL:fig:kappa_phase_transition.

#### Configurations for Figure LABEL:fig:babygpt_hessian_plot .

We use the codebase(^21^2121[https://colab.research.google.com/drive/1SiF0KZJp75rUeetKOWqpsA8clmHP6jMg?usp=sharing](https://colab.research.google.com/drive/1SiF0KZJp75rUeetKOWqpsA8clmHP6jMg?usp=sharing)). We enlarge the vocabulary size from 2 to 8, and we change the objective function to the standard pre-training objective function (i.e., predict the next token). We consider a 1-layer Transfomer with n_emb = 16, n_head = 4, and the width (i.e., the number of output neurons) of mlp.fc_1 equals 32. The Hessian calculation follows the same procedure as in Figure LABEL:fig:block_diagonal.

It is worth mentioning that the weight matrices might have different shapes depending on the different codebases. For instance, the weight matrix $W$ in one codebase might be implemented as $W^{\top}$ in another codebase (this can happen for the embedding layer), and the same $W$ might also be stretched into an one-dimensional vector (this can happen for all weight matrices). In these cases, some extra re-ordering is needed. For all the Hessian calculations in Figure LABEL:fig:babygpt_hessian_plot, we re-order the Hessian rows & columns such that the block-diagonal structure is visible.

#### Throughput Comparison.

The results on throughput are tested on 2$\times$ A800-80GB GPUs. We did not turn on CPU offload. We report the throughput from the summary file of the Wandb log.

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: 1. How does Adam-mini make sure that its reduced learning rate partitions, relying on Hessian structures, can deliver a comparable optimization efficiency as full adaptive optimizers like Adam or AdamW, specially in dealing with sparse gradients and complex non-convex landscapes?

2. In which ways does Adam-mini maintain stability and avoid training instabilities, given its loss of per-parameter adaptation? Are there any specific mechanisms to prevent oscillations or divergence, particulary in deeper and more complex neural networks?

3. To what degree is the claimed memory savings consistent across varied model sizes and architectures, and does Adam-mini reliably reach faster convergence and throughput compared to the traditional adaptive optimizers?

4. How does Adam-mini approach hyperparameter sensitivity without the presence of per-parameter learning rates, and does this require additional hyperparameter tuning or raise computational costs?

5. What generalization ability does Adam-mini show when used on non-transformer models, such as CNNs or GNNs, and could its block-wise learning rate method limit performance in architectures where Hessian matrices are less structured?

6. How effectively does Adam-mini perform when compared to other recent memory-saving optimizers like Sophia or Lion, and is there potential for combining Adam-mini’s technique with finer adaptive methods to reach higher efficiency?
soundness: 3
strengths: 1. Adam-mini cuts memory use by almost half (up to 50%) vs AdamW, by smartly dividing params based on Hessian structures and using block-level learning rates; this gives major memory savings without losing much performance on big language models.  

2. By reducing comunication overheads between GPUs, Adam-mini boosts throughput by about 49.6% during Llama 2-7B training and lowers training time by 33%, which is really helpful in setups with limited resources.  

3. Adam-mini provides a more efficient computational approach by skipping per-param learning rates while still matching AdamW's performance, closing the gap between memory efficiency and optimizer effectiveness across different neural net types, like LLMs and other deep learning architectures.
summary: Adam-mini presents a memory-saving approach to optimization, targetting large transformer models by cutting down the number of learning rates needed for each parameter. It does this by using Hessian-based partitioning to group parameters, so it assigns just one learning rate per block instead of for every single parameter, which cuts down memory use a lot. This setup holds up well in performance next to more traditional optimizers like AdamW, especially in terms of throughput and training speed, as it lowers GPU communication needs. Adam-mini shows it can work well on various deep learning tasks, optimizing models efficiently while saving both computing resources and training time.
weaknesses: - This study mostly benchmarks on transformer models, leaving its effectivness on other architectures like CNNs and RNNs a bit underlooked; more benchmarking could either show it’s limitations or confirm its ability to adapt across other models types.

- The influence of the Hessian-based learning rate grouping on different gradient structures wasn’t deeply investigated, and comparing it with fully adaptive methods would make it’s efficiency clearer.

- Optimizer stability over long training durations hasn’t really been tested, so longer training runs would give more confidence in its reliability, especially for deep networks or tasks that need extended training.

- Adam-mini’s sensitivity to hyperparameters hasn’t been very explored, which might make tuning it more difficult, and more understanding here could show if it raises computational cost.

- Even though they briefly mention it, there’s no strong comparison with newer optimizers like Sophia and Lion; a head-to-head test would help show Adam-mini's specific advantages or where it could improve.

- Performance of the optimizer on tasks like fine-tuning and transfer learning is not talked about, and testing it on practical applications could confirm its strengths beyond just primary training.

- The grouping of parameters lacks deep mathematical support in the study, and adding a stronger theoretical basis would back up the partitioning method chosen and possibly lead to more improvements.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: See weaknesses.
soundness: 3
strengths: 1. This paper studies an important problem and provides an interesting and novel solution. 
2. The paper is well-written with clear motivation and good presentation flow.
3. Extensive experiments are conducted to justify Adam-mini's effectiveness.
summary: This paper introduces Adam-Mini, a new Adam variant for reducing the training memory cost. It changes Adam's parameter-wise learning rate scaling to block-wise learning rate scaling to reduce the cost of the second-order optimizer state. Extensive experiments show that Adam-mini can achieve on-par or better performance than AdamW with less memory footprint.
weaknesses: 1. Integrating Adam-mini into existing training frameworks seems non-trivial. For example, according to Algo 2, we need to specify the partitions manually for large models. Is there any solution to automatically get a good partition given general Pytorch models? 
2. It will be good to add discussions/insights about how to adapt hyper-parameters when switching from AdamW to Adam-mini.
3. In addition to training curves, it will be better to have comparisons after model convergence. Additionally, it will be interesting to see results on diffusion transformers for diffusion model experiments. Some public diffusion transformer training codebases:
- https://github.com/facebookresearch/DiT
- https://github.com/mit-han-lab/efficientvit/blob/master/applications/dc_ae/README.md

---

confidence: 3
contribution: 3
initial_rating: 8
presentation: 3
questions: Minor referencing errors and typos:

- Line 294: instability on 1B models (see Figure 7 (d)) → instability on 1B models (see Figure 8)
- Line 349-350: As shown in Figure 7 (d). This strategy → As shown in Figure 8, this strategy 
- Lines 1053-1054: The results are shown in Figure 15 → The results are shown in Figure 16
- Lines 418-419: Figure 9 (b) shows the loss curves → Figure 9 shows the loss curves
- Line 483: As shown in Figure 12 → As shown in Figure 12 (a)
- Table1 does not seem to be referenced except in the caption of Figure 14 in the appendix 
- There appears to be no reference to Figure 19 in Section D.2
- What is “Adam-mini-2e-4-1B-v” in Figure 12 (a)?
- Line 374: BP in “where e is certain BP error” is defined 2 lines later. 
- Lines 449-450: This might because → This might be because
soundness: 4
strengths: The paper is well-written, and the proposed idea is novel. The approach, though simple, achieves a substantial reduction in memory usage, higher throughput, and shorter wall-clock time without compromising performance. Given the increasing adoption of LLMs and their high demands on memory resources, this method has the potential to be impactful in memory-constrained LLM applications.
summary: This paper presents Adam-mini, a variant of Adam with significantly reduced memory requirements. The approach leverages the near-block-diagonal structure observed in the Hessian of neural networks, and assigns a single learning rate to groups of parameters within a block, thereby reducing memory usage substantially.

The Authors observe limitations in PyTorch’s default partition, particularly for Transformers. To address this the Authors examine the Hessian structure of a small Transformer and identify three distinct classes of Hessian sub-blocks within Transformers and propose a tailored partitioning strategy based on this. Additionally, the authors introduce a simple method to determine a suitable learning rate for each block by assigning the average of Adam’s $\upsilon$ to that block. 

The claims are supported by various experiments including pre-training experiments on GPT-2 and Llama series where Adam-mini is compared against AdamW and other SOTA memory-efficient Adam variants, which demonstrate the effectiveness of Adam-mini in reducing memory overhead while maintaining performance.
weaknesses: The method proposed in the paper does not seem to address large-scale non-Transformer models. Although Section C.2 includes experiments on various non-LLM tasks, it would be helpful to discuss whether Adam-mini offers advantages for large-scale non-LLM applications or if AdamW is preferred in these cases.

Moreover, given the success of the method on GPT-2 and Llama series and its potential impact, demonstrating the effectiveness of Adam-mini on more diverse models, such as BERT-like and vision transformers would be highly beneficial.

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 3
questions: 1. Figure 6, training loss is not a great metric to compare the performance of different methods, a lower training loss may mean overfitting. It could be better id the authors can provide results on downstream tasks.
soundness: 2
strengths: 1. This paper focus on an interesting and important problem. memory-efficient optimization is very important for llm training. 

2. The paper provides the results from diverse domains to verify the efficiency and performance of Adam-mini. 

3. This paper provides the results on LLM pre-training, SFT and RLHF to verify the performance of proposed method Adam-mini.
summary: This paper proposes a novel memory-efficient optimizer Adam-mini to reduce the memory cost by cutting down the learning rate resources in Adam. I tries to partition the parameters into different blocks and assign a single learning rate to each parameter block to reduce the memory cost from `v` in adam. The experimental results illustrate that the proposed method can achieve 49.6% higher throughput than AdamW for LLaMA2-7b pre-training.
weaknesses: 1. The paper is not very easy to follow for me. For example, the paper first aims to partition the model parameters into blocks by Hessian structure. However, in the algorithm part, the paper first tried the PyTorch default partition to construct the blocks for middle-scaled non-transformer tasks and then said the default partition can not work well for transformers. Therefore, they analyzed the exact Hessian at
initialization of transformer training and define the method of partitioning for different layers (such as attention layers, MLP layers). That make me feel a little confused about your proposed method. 

2. The motivation is also not very clear to me. I can understand that different blocks need different learning rates. But why a single learning rate can work well for the parameters of the whole block is not very clear to me. I noticed that the authors provided a toy example (random quadratic minimization problem) to show that, but I am not clear whether the conclusion can be general (for neural network training). Because the conclusion is very important and the proposed algorithm is also based on this conclusion, I think the authors could provide more analysis. 

3. The most visualization in this paper (Figure 2, 3, 4, 7) can just say different blocks show different patterns, but we cannot obtain the conclusion that the whole parameters of each block should use the same learning rate. 

4. The introduction about using the mean of (v * v) to calculate the learning rate is also not very clear. That is also the main contribution of your proposed method. Since you use the mean of (v * v) to change second-order momentum. 

5. I can understand that the paper proposes many valuable experimental findings, but the proposed method make me feel a little tricky. In the proposed method, the authors design too many components and make me consider the generality and scalability of the method. 

6. The paper should also provide some comparisons with strong baselines, such as Lion. These memory-efficient optimizers is very easy to use and general for different models (resnet and transformers). Therefore, you do not need to consider different algorithm design for different model architecture. They can also save the memory cost and obtain higher flops. Based on that, actually, I don't think adam-mini has too many advantages compared with them.

Although the paper provides the results of Lion, I am not sure whether the authors fully tune the hyper-parameters, because Lion usually need a very different LR. In addition, [1] provides thorough comparisons between different optimizers (memory-efficient optimizers) and the results illustrate that the memory-efficient optimizers can obtain comparable results with Adam.  Therefore, I guess these memory-efficient optimizers can finally achieve a very similar performance. 


[1] Deconstructing What Makes a Good Optimizer for Language Models

Meta-review: Adam-mini introduces a memory-efficient optimization approach for large-scale transformer models by employing Hessian-based partitioning to group parameters and assigning a single learning rate to each block, significantly reducing memory usage. The method is conceptually reasonable, easy to implement, and supported by extensive experimental validation. However, the paper underwent substantial revisions and numerous experiments were added during the rebuttal phase, creating a noticeable gap between the initial submission and the revised version. Additionally, a key issue remains unclear: under FSDP training, model weights are often partitioned across GPUs in a manner that does not align with layer-wise boundaries. For example, weights from a single layer can be split across multiple GPUs, causing the Hessian blocks used for parameter grouping to become fragmented and influenced by the framework’s weight partitioning scheme. This introduces a potential inconsistency between the theoretical motivation of Hessian-based grouping and its practical implementation. While the rebuttal partially addressed this concern, further clarification is needed in the final version. Nonetheless, the paper tackles an important challenge in reducing the memory overhead of adaptive optimization methods for large language models, enabling improved token throughput, which justifies a positive recommendation for acceptance.

Decision: Accept (Poster)