Adam-mini: Use Fewer Learning Rates To Gain More
Authors: Yushun Zhang
Congliang Chen
Ziniu Li
Tian Ding
Chenwei Wu
Diederik P Kingma
Yinyu Ye
Zhi-Quan Luo
Ruoyu Sun
Venue: ICLR 2025 Poster (2025)
Keywords: large language model
memory
optimizer
Primary area: optimization

Abstract: We propose Adam-mini, an optimizer that achieves on-par or better performance than AdamW with $50$% less memory footprint. Adam-mini reduces memory by cutting down the learning rate resources in Adam (i.e., $1/\sqrt{v}$). By delving into the Hessian structure of neural nets, we find Adam’s $v$ might not function at its full potential as effectively as we expected. We find that $\geq 99.9$% of these learning rates in $v$ could be harmlessly removed if we (1) carefully partition the parameters into blocks following our proposed 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$% higher throughput than AdamW when pre-training Llama 2-7B on $2\times$ A800-80GB GPUs, which saves 33% wall-clock time for pre-training.

Published as a conference paper at ICLR 2025
ADAM-MINI: USE FEWER LEARNING RATES TO GAIN MORE
Yushun Zhang Congliang Chen 1h3 ZZiiiu Li Tian Ding Chenwei Wu',
Diederik P. Kingma', Yinyu Ye Zhi-Quan Luo Ruoyu Sun 1123
1 The Chinese University of Hong Kong, Shenzhen, China;
2 Shenzhen International Center for Industrial and Applied Mathematies;
3 Shenzhen Research Institute of Big Data; Duke University;
5 Anthropie; 6 Stanford University
ABSTRACT
We propose Adam-mini, an optimizer that achieves on-par or better performance than AdamW with 50% less memory footprint. Adam-mini reduces memory by cutting down the learning rate resources in Adam (i.e., 1/) By delving into the Hessian structure of neural nets, we find Adam's v might not function at its full potential as effectively as we expected. We find that 2 99.9% of these learning rates in v could be harmlessly removed if we (I) carefully partition the parameters into blocks following our proposed 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-minn 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 Adam-mini also alleviates communication overheads among GPUs, thereby increasing throughput. For instance, Adam-minn achieves 49.6% higher throughput than Adam when pre- veeetrriiing Llama 2-7B on 2x A800-80GB GPUs, which saves 33% wall-clock time for pre-training
INTRODUCTION
Adam (Kingma & Ba, 2014) has become the de-facto optimizer for training large language models (LLMs) (e.g.. (Vaswani et al., 2017; Achiam et al., 2023; Touvron et al., 2023; Team et al., 2023)). Despite 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 2x the memory of the model size 2. This memory consumption has become major burden in LLM training. For instance, to train 7B model, Adam alone requires about 56 GB for m and v, and with the gradients included, a total of 86 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., 2020) must be used in practice which unfortunately increases the latency and slows down the training (Rajbhandari et al., 2021).
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.
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 's m and v (Zhang et al., 2020; Kunstner et al., 2023). remains uncertain which components Adam are indispensable for superior performance, and which components could be re-designed or improved. One notable attempt is Adafactor (Shazeer & Stern, 2018), 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., 2023) and Section 3.4). One possible reason is that the current v in Adam is crucial and cannot be simplified.
*: Equal contribution. f: Correspondence author.
'Our implementation of Adam-mini is available at https: /ttgttttttuu///////////tt hecchuu.coooaaeec coom/yyshun// Adammmiamm mini We restate the update rules of Adam and AdamW (Loshchilov & Hutter, 2017) in Appendix E.

---

Published as a conference paper at ICLR 2025
49.9%
49.6% 3725.99
Va
245
343
M
m
Adamw
AdamW
23% speed
Memory (1)
Throughpat (1)
(a) Memory (4) and Throughput (1)
(b) Loss V.S. iteration
(c) Loss V.S. time
(a) Adam-mini takes less memory and reach higher throughput (# tokens per second). The throughput is tested on 2x A800-80GB GPUs. (b, c) Adam-mini performs on-par with AdamW, but takes 33% less time to process the same # tokens.
Figure/ 1: Results for Llama 2-7B pre-training.
This is possible as most existing Adam variants that attempt to modify e to varying extents have been reported to perform worse than Adam (Orabona, 2020). Another possible reason is that there potential to cut down v, but Adafactor does not use the most suitable way: matrix factorization generie approach that could be applied broadly, but it does not leverage much problem-specific structure, thus does not work well on specific neural-net tasks.
In this work, we find it is possible to signifi- cantly reduce the usage of v. Currently, Adam assigns an individual learning rate for each pa- rameter, i.e., i-th parameter receives learning rate " where eú is the i-th component of yer For a billion- paatmoee 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 near-ble near-block-diage block-diagonal with several dense princi- ple sub-blocks (Collobert, 2004). We then find that, for each of these dense bub-blocks, there exists single high-quality learning rate that outperforms Adam, provided that we have enough resources to search it out. Since the number of dense sub- 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.
Mood
poanser
Hesslan noural
Figure 2: assigns learning rates (Irs) by Hessian structure. It uses more Irs than SGD but fewer than Adam.
An illustration of Adam-mini. Adam-mini
We then propose cheap and simple way to find good learning rates that 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 according to the dense Hessian sub-blocks, and call it ga for b E , For each we calculate the quantity below.
(2) mean: @ gb) B2 Ub. b = 1, 1..... B.
We then use /Tb as the learning rate for the parameters in the block associated with 96- Such 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 simple illustration in Figure 2 and relegate the complete form later Algorithm 2. We summarize our main contribution as follows.
New optimizer. We propose a new optimizer called Adam-mini First, Adam-min partitions the model parameters based on the principle we established upon the Hessian structure. Then, it chooses a single learning rate for each block using the average of Adam's 's e 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 99.9% proportion of Adam's v, which saves 50% of the memory cost of Adam.
Effectiveness: Despite the memory cut down, empirically verify that Adam-mini performs on par or even better than Adam W on various language models sized from 39M to 13B, including pre-training, supervised fine-tuning (SFT), and reinforcement learning from human feedback (RLHF). Adam-min mini also performs similarly to Adam on non-LLM tasks such as training diffusion models, vision models, and graph neural networks.
Efficiency: Adam-minn can reach higher throughput than AdamW. We observe that Adam- mini reaches 49.6% higher throughput of AdamW when pre- pre-training Llama 2-7B 2x 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.
2

---

Published as a conference paper at ICLR 2025
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 Hessian. This principle can apply generic problems with block diagonal Hessian:
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, single (but good) learning rate suffices 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 mip by output neurons; (3) embed and output by tokens. We emphasize that our Hessian-basee partition principle is crucial, as naive or default partitions (e.g. partitioning by layers) would cause training instability LLMs.
METHOD
2.1
MOTIVATIONS AND OBSERVATIONS
Now we discuss our observations that motivate the design of Adam-mini. We start by investigating the role of Adam's and explore possibilities for improvement. In Adam, provides an individual learning rate for each parameter, i.e., i-th parameter receives the learning rate -.. where e is the
i-th component of v. Very recently, Zhang et al. (2024) pointed out that such design is crucial for modern architectures such as Transformers. This is because these models often exhibit flesiinnblock heterogeneity, i.e., the Hessian of different parameter blocks have dramatically different eigenvalue distributions (We restate their findings in Appendix E.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., 2024) suggest that necessary to use different learning rate for each block. Nonetheless, Adam does much more than that: it assigns individual learning rate not just for each block, but for each parameter. Note that the number of parameters is much larger than the number of blocks. This begs the question:
(Q1) Is it necessary use customized learning rate for each parameter?
If not, how much can we save?
(a) Hessian of ML.P (Collobert, 2004)
(b) Hessian of MLP initialization
(c) Hessian of ML.P at 50% step
(d) Hessian of MLP
at 100% step
Figure 3:
The bear-bock-ddiggnne Hessian structure of neural nets. (a) is Hessian of an MLP after I training step reported in (Collobert, 2004). (b.c,d): the Hessians of a I-hidden-layer MLP on CIFAR- 100. The wearblock-ddiggnn structure maintains throughout training, where cach block corresponds to one neuron.
To answer (Q1), we delve into the Hessian structures of neural networks. First, we recall an important (but often overlooked) result: the Hessian of neural nets is near-block-diagone This is an old result that has been reported for two decades; see (Collobert, 2004, Section 7). The authors also provided theoretical explanations. We restate their analysis in Appendix C. We now provide some case studies.
Case study I: random quadratic problems. With the above observation mind, we now explore (Q1) generic optimization problems with block-diagonal Hessian. We consider the random quadratic minimization problem min 1 ww Hw where the Hessian H is random positive definite (PD) matrix and is visualized in Figure 4 (a). We compare the coordinate-ws wise learning-rate method, i..., Adam, with the minimization because the optimal learning rate has close form. We have the following findings.
method, i.c., gradient descent (GD). We choose quadratic
(1): shown in Figure 4 (a) and (b), Adam outperforms the optimal loggle-laaminggracc method.
This is expected since Adam deploys different learning rates to different parameters.
All experimental details in Section 2 shown in Appendix F.2.
3

---

Published as a conference paper at ICLR 2025
600 1000
(a) Hessian matrix
(b) Total loss
(c) Dense sub-block (d) Sub-problem loss
Figure 4: (a): The Hessian of three-block random quadratic problem. (b): Training curves for the problem
associated with the full Hessian in (a). The optimal single (blockwise) learning rate chosen based the full (blockwise) Hessian in (a), (c): The 1st dense Hessian sub-blocks in (a). (d): Training curves for the new problem associated with the Hessian (c).
(2): as shown in Figure 4 (c) and (d), consider new problem whose Hessian is 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 uub-blocks (a).
(3): If we collect these optimal learning rates in (2) and apply them to "blockwise" version of GD, it would be faster than Adam on the original problem (the green line in Figure (b)).
In summary, for generic problems with -ickoaiiaoonal Hessian, we find that more learning rates do not necessarily bring extra gain. In particular, for each dense sub-block, single (but good) learning rate suflices 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 diagonal preconditioned method, i.e., at the 1-th step:
1P++1 = we - TIe Dame:
(I)
where D, = Diag /ve) is a diagonal matrix, my is the sst-order momentum, w,, and V 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 x(DH) reduced over x( (H)", where H usually refers to the Hessian matrix and s(-) is the condition number (smaller is better). Unfortunately, there is no guarantee of &s s(DH) I <(H) and this inequality often requires strict assumptions on both D and H. For instance, ss s(DH)) would be small if H is close diagonal and D is cleverly designed compressor of H (Forsythe & Straus, 1955; Young, 1954; Sun & Ye, 2021: Qu al., 2022).
Here, we numerically explore the effectiveness of Adam's preconditioner within each dense Hessian sub-block. We generate a random dense PD matrix Hy E gaxd and use proxy for the dense Hessian sub-block of neural nets in Figure 3. We define D = Diag 1/ /V) where O
Hyæ E Rd,, and each entry æ; ~ NO, 1/Vd) follows Xavier initialization. We explore the interplay between the following two metrics:
xC(Damm Hb) <(Ho)
(2)
5, Ho.asl
where E [0, 1] is the -are:onnl-over--f--iaoooo ratio", and we use it to measure how dense Ho is (H& is pure diagonal when 1). measures the effectiveness of Adam's preconditioner D when operating on the Hessian-block. Hy (the smaller the better). We investigate the change of r when changing the structure of H₁, including changing т, dimension d, and also x x(H)). We emphasize that for a fixed d or (HH)), we change by only rotating the eigenvectors, but not changing the eigenvalues of Hg. This ensures is the only changing factor in the experiments.
We summarize the key findings in Figure 5: for H₁, with most dimension d and < n(Hs). r decreases as 7 That is, Ddamm is effective when H, is close to diagonal, and Adamm is not so effective when H₂ is dense. This aligns with the convergence rates in Figure 4. It is intriguing to provide lower bound on x² Adam H7)) to ground the observation in Figure 5, and we are not aware of any existing lower bound of this kind. Note that it is rather difficult to characterize x(Dadam partially because the extreme eigenvalues are neither sub-additive nor uub-mulipplicative (Kittanch, 2006). We leave 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 Pyrorc default partition, and randomly choose one parameter block as
4

---

Published as a conference paper at ICLR 2025
(a) dimension d
(b) dimension k(Hb)
Figure 5: The effectiveness of Adam's preconditioner Dadmm on different matrix structures of Ho. (a): for most dimension d, is large when 7 is small r and are defined in Eq. (2)), This indicates that Adam might not be effective when H, is dense. We fix x <(H)) 500 here. (b): We use the same setups (a), except that we fix the dimension 50 and change the x-axis to K( ((H)).
the "left-out" block and change the coordinate wise learning rate to a single-learning rate counter- pprt.. 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 shows that Adam (leave-one-out) can achieve similar better performance than Adam. A similar phenomenon is also observed when we randomly leave out up 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 possible 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 99.9% fewer learning rates.
2.2 PROPOSED METHOD: ADAM-MINT
We now introduce Adam-mini. We will first state the "general principled form" of Adam-ming and then introduce the "the realization" of Adam-mini on specific architectures. In this section, present the general form of Adam-mini in Algorithm 1. Following this general principled form, Adam-mini will have different realizations on different architectures, and the concrete example on Transformers shown in Appendix B. As shown in Algorithm 1, Adam-mini contains two steps.
Algorithm Adam-mini (General form)
Step 1 Partition the model parameters into blocks by Hessian structure. We discuss Prin- ciple later in Section 2.3. For different architectures, the principle will be realized in different forms; see Algorithm 3: "Parti- tion for non- Tranrfoonorrr" and Algorithm 3: "Partition for Transformers".
Input weight-decay coefficient À and current step
2: Partition params into param param_blocks blocks
by Principle 1 in Section 2.3
3: for param in param param_blocks do
blocks
4: g param. grad
5: param param- 1/1 * As param
6: m [1 (1-81 * 3 #m
{1)
7:
f T-
Step 2. For each parameter block, single learning rate. To efficiently choose suitable learning rate in each block, Adam- mini simply replaces g in vanilla Adam by its mean value. We adopt the moving average on these mean values as in Adam.
8:
v (1 32) *mean(g@q) +
9:
à 1-3
v
10: param param 7/: Vo+c
end for
simple example of Adam-mini. We simple example illustrate the key design of Adam- mini. For a problem with parameters w E IR$, Adam and Adam-mini both perform w = w G m, where m the 1st-order momentum and has different forms as follows:
For Adam: UAdam
( ( vua
veg vua (
For Adam-mini: suppose the partition is (1,2,3) and (4,5) then
Umini = (( V (ms tva+3 (vi (va+to)//2 (va+ta)/22
(
Note that the number of effective elements Umini equals the number blocks, which could be significantly smaller than that of Addam- which equals the number of parameters. For LLMs, this will free 2 99.9% elements in
5

---

Published as a conference paper at ICLR 2025
25
(a) Randomly leave one (b) Randomly leave two (c) Randomly leave three block out
(d) Performance gap
blocks out
blocks out
Figure 6: (a) (b) (c) Adam (leave-r-out) can reach a similar or better performance than Adam all randomly picked left-out blocks. x = 1, 2, 3. (d) The performance gap between Adam and Adam (leave-one-out) for all possible blocks. We find Adam (leave-onc-out, always performs on par with Adam, and for most blocks, Adam (lcave-onc-out) performs better. It seems possible to perform well with much fewer learning rates than Adam.
2.3 PRINCIPLE FOR THE PARTITION STRATEGY
We now discuss how choose the parameter partition for Adam-mini. A straightforward way is to use PyTorch default partition. Unfortunately, we find that the Py PyTorch default partition does not work well on larger-scaled tasks. In particular, we find that Adam-mini encounters training instability on IB models (see Figure 7 (i)). We suspect this is because the default Py PyTorh partition did not fully capture the Hessian structure. We propose a general principle Principle 1 below.
Principle 1: We should partition parameters into blocks, such that each parameter block is associated with the smallest dense sub-block in 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 since it is necessary to use (at least) one distinct learning rate for each Hessian block (as evident in Appendix E.2).
Does the Py Torch default partition follow Principle 1? To find out, we explore the Hessian of a small Transformer as in Figure 7. Under the default Py PoTohh partition, we compute the Hessian for each parameter block after training step. We find four classes of Hessian sub-blocks.
Class 1: query and key. The Hessian of query and key have hou-bbock-diaggna structures.
The number of blocks equals the number of heads.
Class 2: attn. proj and MLPs. The Hessian of attn. .innnppre proj and MLPs have block-diagonal structures. The number of blocks equals the number output neurons.
Class 3: value. For value, the structure of Hessian 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 significantly different from that of query ry and key, although they all consist of four heads. The Hessian entries of va lue are also about 10 larger than those of query and key One possible reason is that value is positioned outside the softmax operator in the self-attention design, while query and key are not.
Class 4: embed and output. For these two layers, the Hessian sub-block has a near-block- diagonal structure and the number of 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; va 1 ve, attn. proj. and MLPs should be partitioned by output neurons; embed and output 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 va value as whole" (see discussions in Appendix D.6). By default, we will partition value by output neurons.
We then introduce the resulting Algorithm 3: "Partition for Transformers" in Appendix B. As shown in Figure 7 (i). This strategy indeed stabilizes the training and boosts the performance.
2.4 SOME CHARACTERISTICS OF ADAM-MINI AND DISCUSSIONS
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 result, Adam-mini cuts down more than 99.9% of Adam's 's v, which saves 50% of Adam's 's memory.
"This might be one source of the heterogeneity of Hessian eigenvalues as reported by (Zhang et al., 2024).

---

Published as a conference paper at ICLR 2025
300
(a) query (4 heads)
(b) key (4 heads)
(c) value (4 heads)
40
(d) attn proj (16 neurons) (c) m1pp fc__ (32 neurons)
(f) mlp proj (16 neurons)
Anim Pytrem deloch sertion
400
314
306
200
(g) embod (8 tokens)
(h) ootput (8 tokens)
(i) Llama 2-1B
Figure 7: (a-h): The Hessian of different parameter blocks in a small Transformer at 1 % training step. Here, neuron refers to the "output neuron". , We find that these Hessian sub-blocks have wear-block-diagona structure and consists of smaller dense matrices. Different parameter blocks have different numbers of small dense matrices. (i) Training curves on Llama 2- B.B When using Py Torch default partition, Adam-mini could suffer loss spikes. The spike disappears when we use the partition strategy in Algorithm 3.
Table I: Memory cost of AdamW v.s. Adam-mini. Table 2: Throughput (f) test on 2x A800-80GB GPUs Calculation is based float 32, which is a standard for Llama 2-7B pre-training. Xmeans out of memory. choice for optimizer states.
GPU hours (J) to pre- ve-train Llama 2-7B with the optimal token amount by Chinchila's 's law.
Model
Optimizer
Memory (GB)
GPT-2- 1.5B
Adam
12.48
Optimizer
bs GPU total bs
Throughput (1)
Adam-mini
256
5372.19 49.6%)
GPT-2-1.5B
Adam-min mini
6.24 (50% 4)
AdamW
256
Llama 2-1B
Adam W
8.80
AdamW
1
256
3725.59
Llama 2-1B
Adam-mint
4.40 (50% 4)
Optimizer
Tokens (B)
GPU hours (h) ({)
Llama 2-7B
AdamW W
53.92
AdamW
74.56
Llama 2-7B
Adam-mini
26.96 (50% 4)
Adam-mini
49.85 (1 33.1%)
Llama 3-8B
Adam
64.24
Adam W
52T9.16
Llama 3-8B
Adam-mini
32.12 (50% 4)
Adam-mini
70
3489.55 (1 33.1%)
Llama 2-13B
Adam
104.16
Adam W
140
10438.32
Llama 2-13B
Adam-mini
52.08
(50% 4)
Adam-mini
140
6979.10 (1 33.1%)
Higher throughput. Adam-mini can reach a higher throughput than Adam especially under limited GPU resources. There are two reasons. First, Adam-mint does not introduce extra computation in its update rules. The averaging operation in Algorithm 1 incurs negligible cost and significantly reduces the number of vector-square-rood 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 major overhead (Rajbhandari al., 2021). We report evidence in Table 2. When pre-training Llama 2-7B on 2x A800-80GB GPUs, we find Adam-mini could reach 49.6% higher throughput than AdamW. This translates 33.1 1 reduction of wall-clock time on processing the same amount of tokens for pre-training.
Why using mean(v) learning rates. Due limited space, move the discussions Appendix C.
Has room to improve. Adam-min designs the learning rate for each dense Hessian sub-block using the average of Adam's in that block. Such design achieves cheap computation, but it might not be
7

---

Published as a conference paper at ICLR 2025
12"
Atar mbs (7yhmbh dedrzit partiio))
Valtarton foee
1
Adam mire 3.20
Adam
Walatan
SM4 [AMB
5
12
25.0 6.
10.5
125
15.
1755
25
6.0
16.
155
15.0
17.5
Thloos Dellor)
Joloes Diiloo
(a) GPT-2-122M
(b) GPT-2-330M and 1,5
Figure 8: For GPT-2 series pre-training, Adam-mini performs similarly to AdamW with 50% less memory, while other methods perform worse.
optimal. We believe there is great room to improve the learning rate design. As shown in Figure we can reach much faster convergence if we utilize more information in the dense block to design the learning rate (e.g., using eigenvalues each block), However, such a design requires expensive computation. We leave it as an important future direction.
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 (RL.F).. (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 D.5. All LLM experiments are conducted four NVIDIA A800-80GB GPUs and the rest are conducted on four V100 GPUs. All the experimental details are explained in Appendix F.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., 2019) series (125M to I.5B) on Openwebtext (Gokaslan al., 2019). We train Llama series (20M 13B) (Touvron et al., 2023) on C4 (Raffel et al., 2020), We compare Adam-min with AdamW (Loshchilov Hutter, 2017) as well as popular memory- efffcietff methods including Adafactor (Shazeer & Stern, 2018), CAME (Luo et al., 2023), and SM3 (Anil et al., 2019). For Adafactor and SM3, we incorporate momentum with 31 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 8 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 8 (a), we run Adam-mini (PyTorch default partition), which partition parameters by PyTorch default partition. We find that Adam-min mini (PyTorch default partition) performs poorly. We stop the trial since shows clear unstable behavior. In Figure 9, we further present the training loss curves. We find that the loss curves of Adam-min closely resemble those of Adam
4.0
AdamW-125M
3.8-
AdumW-3300
3.8
AdammW 1.58
8 9.4-
mini Adam
3.2-
Arhem-mini-11
3.0
2.8
2.6
0.0
2's
7.5
12.5
15.0
17.5
20.0
(a) GPT-2 series
(b) Trajectory comparison
Figure 9: (a): The training loss curves of Adam-mini closely resemble those of AdamW. (b): The trajectory of Adam-mini stays close to the trajectory of AdamW (in terms of the 6-2 distance of model checkpoints).
Llama series. Figure 10 shows the results for pre-training Llama series. We also train Llama 2-7B as shown in Figure (c) in Section 1. We find that Adam-min Miammmiiijj performs on par with Adam W, 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 (b) and the detailed description is in Appendix F. This might be because Adam-min makes fewer modifications over Adam W.
Sensitivity analysis. On GPT-2-125M, pre- pre-training task, test the sensitivity of Adam-min mini to hyperparameters. We report the validation loss after training with 2.5B tokens (by Chinchilla's 's law). As shown in Figure 12 (c), Adam-mini seems not overly sensitive to hyperparameters.
8

---

Published as a conference paper at ICLR 2025
7.0
AdamW-I111
GAME
Man
Naam 13B
Atadbcterr
Adam--mini-B
ooe sonsing
85.3
Adam-mini-1330 135
5.00
4.5
3.5
3.00
2000
6000
8000
10000
Itemitoon
(a) Llama 2-1B
(b) Llama 3-8B and Llama 2-13B
Figure 10: (a, b): Loss curves of pre-training training Llama series from IB to 13B. Adam-mini performs on par or better than A with 50% less memory, while other methods perform worse. Further, loss curves of Adam-mini closely resemble the curves by AdamW.
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 IB. For the model with size paraam. we train the model with about 20 * 71 param tokens, which is suggested to be the optimal amount by Chinchilla's law (Hoffmann et al., 2022). The largest- largest-scaled experiment we conducted is Llama 2- 1B pre-training with 26.2B tokens, which takes about 170 GPU hours on 4x A800-80G8 GPUs. The total running time for the scaling law experiments is about 300 GPU hours.
Liatta Serses 2he traaning Da Chundtila's Lawe
Limma Series Pre tranning by Chinchilla's Law)
Adaes
Finull Vall dation loss Gog)
W
Adem-minii
Villatti Jas Dog)
Adam-mini
Airedhiiee 2014
2x
10
FLOPs (lgg)
Parsmeters Org)
(a) Scaling laws in terms of compute
(b) Scaling laws in terms of parameters
Figure 11: (a, b): Scaling laws of Adam-mini. We pre-train Llama 2 architectures by Chinchilla's law. For all models sized from 39M to IB, Adam-mini's loss curves are consistently similar to Ad but Adam-mini uses 50% less memory. Further, as shown in (b), Adam-mini reaches a lower final loss than AdamW for all models. The fitted lines in (b) suggest that Adam-mini can be scaled up to larger models (if the scaling law holds).
As shown in Figure 11, ddam-mini's loss curves are consistently similar to A. We also present the final validation perplexity and find that Adam-min mini reaches slightly lower perplexity than Adam W for all models (see Figure 11 (b), also see Table 4 in Appendix D). The fitted lines in Figure 11 (b) suggest that Adam-mini can be scaled up to larger models (if the scaling law holds).
Another advantage of Adam mini is its ability to reduce computational costs for scaling law experiments. Scaling law experiments are typically used to predict the optimal configurations for large-scale models by fitting the performance of smaller-scale proxy models. To accelerate the development of large-scale models, it is crucial to minimize costs during the fitting process (Hägele et al., 2024). Adam-mini achieves this by delivering the same scaling results while using significantly less memory and time cost (c.g., 33% less GPU hours, as shown in Figure 1).
3.3 SUPERVISED FINE-TUNING AND RLHF
We now test Adam-mini on SFT and RLHF We use the Llama 2-7B pretrained model (Touvron et al., 2023) for our study. We use the afteeeebbcck dataset and implement the RLHF workflow from (Ouyang et al., 2022). We use ReMax (Li et al., 2023), a memory memory-effiiint alternative to PPO
(a) SFT
(b) RLHF
(c) Sensitivity analysis
Figure 12: (a, b): SFT, and RLHF when aligning Llama 2-7B, Adam-mini reaches better performance (smaller perplexity, higher reward) than AdamW with 50% less memory. (c) Sensitivity analysis of Adam-mini on GPT-2-125M pre-training (by Chinchilla's law). Adam-mini seems not overly sensitive to hyperparameters.

---

Published as a conference paper at ICLR 2025
124866
7.03
40% SCO9L
28.5%
2.3499
Adar-miiti
Thrugghutt (7)
(a) Llama 2-20M
Llama 2-1B
(c) Throughput comparison
Figure 13: (a, b): Training curves of Adafactor Adafactor Adafactor-Zhai-verr Zhaioversion Llama 2-20M & Llama 2-1B pre-training. We find these two methods consistently underperform Adam-mini. (c): On Llama 2-1B, Adam-mint achieves 40% higher throughput than Adafactor (tested on 2x A800-80GB GPUs).
(Schulman et al., 2017), to optimize the preference reward. As shown in Figure 12, Adam-mini performs par better than AdamW. Adam-mini also achieves better alignment performance MT-Bench using GPT-4 as judge. The results are shown later Table 5 in Appendix D.4.
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., 2022), which we call "Adafactor-Zhal" -eiioversion"". For both versions, we use momentum with B = 0.9.
We first conduct learning rate grid-search on Llama 2-20M and train it following Chinchilla's 's law. As shown in Figure 13 (a), we find that Adafactor -haiitor---hivaae 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) 83 = 0.95; (2) {10-30, 10-10 10-8 10-6); (3) warm-up steps = (1%, 2%, 3%, 4%, 5%, 10%) total steps. The results are shown in Appendix D.7. 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 Llama 2-1B, In contrast to the case of Llama 2-20M, we find that the Adafactor- chai-version now suffers from training instability and the original version performs better. Nevertheless, they still underperform Adam-mini. In Appendix D.8, we conduct similar hyperparameter search for Lion (Chen et al., 2024b) and we find it also underperforms Adam-ming
About hyperparameter tuning. We acknowledge that might be possible improve these methods we spend more resources on grid search (as claimed by a recent work (Zhao et al., 2024b)). However, based
experience far, it is 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 find the correct combination. In contrast, Adam-mini much easier to use. In all our experiments, Adam-mini performs well using the same hyperparameters as Adam W (including learning rate, B1, $2, e, etc.).
Throughput comparison. Besides the performance comparison, we further find that Adafactor has higher latency than Adam-mini (Figure 13 (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 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 latency issues also apply other variants of Adafactors such as CAME. In contrast, Adam- saves computation when taking the square root of c. As such, Adam-min reaches a higher throughput.
the input and output dimension. Note that similar
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.
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 simple and cost- effective way design a learning rate for each dense Hessian sub-block, but might not be an optimal way. We leave the development of stronger designs as a future direction.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGEMENT
Yushun Zhang would like to sincerely thank Less Wright, Andrew Gu, and other open-source contributors for their invaluable support in integrating Adam-mini into the mainstream codebases, The work of Ruoyu Sun was supported by NSFC (No. 12326608); Hetao Shenzhen- Hong Kong Science and Technology Innovation Cooperation Zone Project (No. HZQSWS-KCCYB-.. (-CCYB-2022016), together with Tian Ding: University Development Fund UDF01001491. the Chinese University of Hong Kong, Shenzhen; Guangdong Provincial Key Laboratory of Mathematical Foundations for Artificial Intelligence (2023B 12 12010001). The work of Z.-Q. Luo supported by the Guangdong Major Project of Basic and Applied Basic Research NN0.202300303300000), the Guangdong Provincial Key Laboratory of Big Data Computing, and the National Key Research and Development Project under grant 2022YFA 1003900. The work of Tian Ding is also supported by Internal Project of Shenzhen Research Institute of Big Data under Grant J00220240005.
BROADER IMPACT
This work designs a new algorithm to reduce the training cost for LLMs. Our algorithm can help the community better train large Al models. However, it could be a potential threat if the Al models are used for illegal usage.
REFERENCES
Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, lige Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. Gpt-4 technical report. arXiv preprint arXiv:2303.0877 087744 2023.
Kwangjun Ahn, Zhiyu Zhang, Yunbum Kook, and Yan Dai. Understanding adam optimizer via online learning of updates: Adam is ftrl in disguise. In Forty-first International Conference on Machine Learning.
Rohan Anil, Vineet Gupta, Tomer Koren, and Yoram Singer. Memory efficient adaptive optimization. Advances in Neural Information Processing Systems, 32, 2019.
Anonymous authors.
Deconstructing what makes a good optimizer for language models. http://opparrveevrrre////ffsssfs 2024.
Pratik Chaudhari, Anna Choromanska, Stefano Soatto, Yann LeCun, Carlo Baldassi, Christian Borgs, Jennifer Chayes, Levent Sagun, and Riccardo Zecchina. Entropy-sgd: Biasing gradient descent into wide valleys. Journal of Statistical Mechanics: Theory and Experiment, 2019(12):124018, 2019.
Junyu Chen, Han Cai, Junsong Chen, Enze Xie, Shang Yang, Haotian Tang. Muyang Li, Yao Lu, and Song Han. Deep compression autoencoder for efficient high-resolution diffusion models. arXiv preprint arXiv:2410.100 10733, 2024a.
Tianqi Chen, Bing Xu, Chiyuan Zhang, and Carlos Guestrin. Training deep nets with sublinear memory cost. arXiv preprint arXiv: 1604.06/076 2016.
Xiangning Chen, Chen Liang, Da Huang, Esteban Real, Kaiyuan Wang, Hieu Pham, Xuanyi Dong, Thang Luong, Cho-Jui Hsieh, Yifeng Lu, et al. Symbolic discovery of optimization algorithms. Advances in neural information processing systems, 36, 2024b.
Ronan Collobert. Large scale machine learning. Technical report, Université de Paris VI, 2004.
Ganqu Cui, Lifan Yuan, Ning Ding, Guanming Yao, Wei Zhu, Yuan Ni, Guotong Xie, Zhiyuan Liu, and Maosong Sun. Ultrfeedbacck: Boosting language models with high-quality feedback, 2023,
André Belotto Da Silva and Maxime Gazeau. A general system of differential equations to model first-order adaptive algorithms. The Journal of Machine Learning Research, 21(1):502-55113 2020.
11

---

Published as a conference paper at ICLR 2025
Rudrajit Das, Naman Agarwal, Sujay Sanghavi, and Inderjit S Dhillon. Towards quantifying the preconditioning effect of adam. arXiv preprint arXiv:2402.07114, 2024.
Yann N Dauphin, Atish Agarwala, and Hossein Mobahi. Neglected hessian component explains mysteries in sharpness regularization. arXiv preprint arXiv:2401. 10809, 2024.
Aaron Defazio, Harsh Mchta, Konstantin Mishchenko, Ahmed Khaled, Ashok Cutkosky, et al. The road less scheduled. arXiv preprint arXiv:2405. 2024.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In 2009 IEEE conference on computer vision and pattern recognition. pp. 248-255. Ieee, 2009.
Tim Dettmers, Mike Lewis, Sam Shleifer, and Luke Zettlemoyer. 8-bit optimizers via block-wise quantization. In International Conference on Learning Representations, 2021.
John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of machine learning research, 12(7), 2011.
George E Forsythe and Ernst G Straus. On best conditioned matrices. Proceedings of the American Mathematical Sociery, 6(3):340-33) 1955.
Behrooz Ghorbani, Shankar Krishnan, and Ying Xiao. An investigation into neural net optimization via hessian eigenvalue density. In International Conference Machine Learning, pp. 2232-2241.
PMLR, 2019.
Boris Ginsburg, Patrice Castonguay, Oleksii Hrinchuk, Oleksii Kuchaiev, Vitaly Lavrukhin, Ryan Leary, Jason Li, Huyen Nguyen, Yang Zhang, and Jonathan M Cohen. Training deep networks with stochastic gradient normalized by layerwise adaptive second moments. 2019.
Aaron Gokaslan, Vanya Cohen, Ellie Pavlick, and Stefanic Tellex. Openwebtext corpus, 2019.
Guy Gur-Ari, Daniel A Roberts, and Ethan Dyer. Gradient descent happens in tiny subspace. arXiv preprint arXiv: /0t::/8/22. 18/2.04754 2018
Alexander Hägele, Elie Bakouch, Atli Kosson, Loubna Ben Allal, Leandro Von Werra, and Martin Jaggi. Scaling laws and compute-optimal training beyond fixed training durations. arXiv preprint arXiv:2405. 18392, 2024.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition. pp. 770-778, 2016.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840-688 2020.
Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, et al. Training compute-optimal large language models. arXiv preprint arXiv:2203. 2022.
Edward J Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. Lora: Low-rank adaptation of large language models. arXiv preprint 210.:2200600685,. 2021.
Adam Ibrahim, Benjamin Thérien, Kshitij Gupta, Mats Richter, Quentin Anthony, Timothée Lesort, Eugene Belilovsky, and Irina Rish. Simple and scalable strategies to continually pre-train large language models. arXiv preprint arXiv:2403.088 080333633.8 2024.
Kaiqi Jiang, Dhruv Malik, and Yuanzhi Li. How does adaptive optimization impact local neural network geometry? Advances in Neural Information Processing Systems, 36. 2023.
Diederik P Kingma and Jimmy Ba. Adam: method for stochastic optimization. arXiv preprint arXiv: 1412.6980, 2014.
12

---

Published as a conference paper at ICLR 2025
Thomas N Kipf and Max Welling. Semi- Semi-supervised classification with graph convolutional networks. In International Conference on Learning Representations, 2016.
Fuad Kittanch. Spectral radius inequalities for hilbert space operators. Proceedings of the American Mathematical Society, pp. 385-390, 2006.
Frederik Kunstner, Jacques Chen, Jonathan Wilder Lavington, and Mark Schmidt. Noise not the main factor behind the gap between sgd and adam on transformers, but sign descent might be. arXiv preprint arXiv:2304. 13960, 2023.
Frederik Kunstner, Robin Yadav, Alan Milligan, Mark Schmidt, and Alberto Bietti. Heavy-tailed class imbalance and why adam outperforms gradient descent on language models. arXiv preprint arXiv:2402. 19449, 2024.
Bingrui Li, Jianfei Chen, and Jun Zhu. Memory efficient optimizers with 4-bit states. Advances in Neural Information Processing Systems, 36, 2024.
Ziniu Li, Tian Xu, Yushun Zhang, Yang Yu, Ruoyu Sun, and Zhi- Thi-Quan Luo. Remax: simple, effec-
tive, and efficient method for aligning large language models. arXiv preprint arXiv:2310.10505, 2023.
Zhenyu Liao and Michael W Mahoney. Hessian eigenspectra of more realistic nonlinear models. Advances in Neural Information Processing Systems, 34:20104-2011 17, 2021.
Hong Liu, Zhiyuan Li, David Hall, Percy Liang, and Tengyu Ma. Sophia: scalable stochastic second-order optimizer for language model pre-training. arXiv preprint arXiv:2305. /4342, 2023.
Yang Liu, Jeremy Bernstein, Markus Meister, and Yisong Yue. Learning by turning: Neural architecture aware optimisation. In International Conference on Machine Learning, pp. 6748- 6758. PMLR, 202
Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows. In Proceedings of the IEEE/CVF international conference on computer vision, pp. 10012-10022, 2021b.
Iya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint arXiv:1711.01101, 2017.
Qijun Luo, Hengxu Yu, and Xiao Li. Badam: memory efficient full parameter training method for large language models. arXiv preprint vrXiv:2404.02827, 2024.
Yang Luo, Xiaozhe Ren, Zangwei Zheng, Zhuo Jiang, Xin Jiang, and Yang You. Came: Confidence- guided adaptive memory efficient optimization. arXiv preprint arXlv:2307.02047, 2023.
Kai Lv, Hang Yan. Qipeng Guo, Haijun Lv, and Xipeng Qiu. Adalomo: Low-memory optimization with adaptive learning rate. arXiv preprint 2.I::2310010195 2023a.
Kai Lv. Yuqing Yang. Tengxiao Liu, Qinghui Gao, Qipeng Guo, and Xipeng Qiu. Full parameter fine-tuning for large language models with limited resources. arXiv preprint arXiv:2306.00 09782, 2023b.
Sadhika Malladi, Tianyu Gao, Eshaan Nichani, Alex Damian, Jason D Lee, Danqi Chen, and Sanjeev Arora. Fine-tuning language models with just forward passes. Advances in Neural Information Processing Systems, 05553038-33773 2023.
James Martens and Roger Grosse.
Optimizing neural networks with factored approximate curvature. In International conference on machine learning, pp. 2408-2417 2417. PMLR, 2015.
Francesco Orabona. Neural networks (maybe) evolved to make adam the best wtt.:efeeee....e./eemmemee neural -ne work wwwewwor-mmabbeeeee evol htw.::/---e.eeee--ootooeeaeeeeee
optimizer. 2020.
URL
zer/
13

---

Published as a conference paper at ICLR 2025
Long Ouyang, Jeffrey Wu, Xu Jiang. Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang. Sandhini Agarwal, Katarina Slama, Alex Ray, et al. Training language models to follow instructions with human feedback. Advances in neural information processing systems, 35:27730- 27744, 2022.
Yan Pan and Yuanzhi Li. Toward understanding why adam converges faster than sgd for transformers. arXiv preprint arXiv:2306.0022 2023.
Vardan Papyan. The full spectrum of deepnet hessians at scale: Dynamics with sgd training and sample size. arXiv preprint arXiv: 1.07062, 2018.
Vardan Papyan. Measurements of three-level hierarchical structure in the outliers in the spectrum of deepnet hessians. arXiv preprint arXiv: 2019.
Vardan Papyan. Traces of class/cross-class structure pervade deep learning spectra. The Journal of Machine Learning Research, 21(1):1017-11026, 2020.
Barak A Pearlmutter. Fast exact multiplication by the hessian. Neural computation, 6(1):147-160 1994.
William Peebles and Saining Xie. Scalable diffusion models with transformers. Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 4195-4205, 2023.
Zhaonan Qu, Wenzhi Gao, Oliver Hinder, Yinyu Ye, and Zhengyuan Zhou. Optimal diagonal preconditioning: Theory and practice. arXiv preprint arXiv:2209.00800 2022.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, llya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
Colin Raffel, Noam Shazeer, Adam Roberts. Katherine Lee, Sharan Narang. Michael Matena, Yanqi Zhou, Wei Li, and Peter Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of Machine Learning Research, 21(140):1-67, 2020. URL http::// //mrr//mmrr...// wtggppaprrrseee.ooeeeeee /22////1122222 html.
Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. Zero: Memory optimizations toward training trillion parameter models. In SC20: International Conference for High Performance Computing, Networking, Storage and Analysis, pp. 1-16. IEEE, 2020.
Samyam Rajbhandari, Olatunji Ruwase, Jeff Rasley, Shaden Smith, and Yuxiong He. Zero-infinity: Breaking the gpu memory wall for extreme scale deep learning. In Proceedings of the international conference for high performance computing, networking, storage and analysis, pp. 1-14, 2021.
Nicolas Roux, Pierre- Antoine Manzagol, and Yoshua Bengio. Topmoumoute online natural gradient algorithm. Advances neural information processing systems, 20, 2007.
Levent Sagun, Leon Bottou, and Yann LeCun. Eigenvalues of the hessian in deep learning: Singularity and beyond. arXiv preprint arXiv:1611.07470 2016.
Levent Sagun, Utku Evci, V Ugur Guney, Yann Dauphin, and Leon Bottou. Empirical analysis of the hessian of over-parametrized neural networks. arXiv preprint arXiv: 04454, 2017.
Adepu Ravi Sankar, Yash Khasbage, Rahul Vigneswaran, and Vineeth N Balasubramanian. A deeper look at the hessian eigenspectrum of deep neural networks and its applications to regularization. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pp. 9481-9488, 2021.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv: wwXXiv:107.06.30 1707.06347, 2017.
Noam Shazeer and Mitchell Stern. Adafactor: Adaptive learning rates with sublinear memory cost. In International Conference on Machine Learning, pp. 4596-4604. PMLR, 2018.
Naichen Shi, Dawei Li, Mingyi Hong, and Ruoyu Sun. Rmsprop converges with proper hyper- parameter. In International Conference on Learning Representations, 2020.
14

---

Published as a conference paper at ICLR 2025
Ruoyu Sun and Yinyu Ye. Worst-case complexity of cyclic coordinate descent: O (n" 2) (n 2) gap with randomized version. Mathematical Programming. 185:487-520, 2021.
Gemini Team, Rohan Anil, Sebastian Borgeaud, Yonghui Wu, Jean- -aan-Bppistte Alayrac, Jiahui Yu, Radu Soricut, Johan Schalkwyk, Andrew M Dai, Anja Hauth, et Gemini: family highly capable multimodal models. arXiv preprint arXiv:2312. 1805, 2023.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307. 092888 2023.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
Petar Veličković, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pietro Lio, and Yoshua Bengio. Graph attention networks. arXiv preprint arxiv:/710.10900 2017.
Petar Velickovic, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pietro Lio, Yoshua Bengio, et al. Graph attention networks. stat, 1050(20):10 48550, 2017.
Bohan Wang, Yushun Zhang, Huishuai Zhang, Qi Meng, Zhi-Ming Ma, Tie- Liu, and Wei Chen. Provable adaptivity in adam. arXiv preprint arXiv:2208. 09900, 2022.
Yikai Wu, Xingyu Zhu, Chenwei Wu, Annie Wang, and Rong Ge. Dissecting hessian: Understanding common structure of hessian in neural networks. arXiv preprint arXiv:2010.04261 2020.
Xingyu Xie, Pan Zhou, Huan Li, Zhouchen Lin, and Shuicheng Yan. Adan: Adaptive nesterov momentum algorithm for faster optimizing deep models. arXiv preprint arXiv:2208.06677 2022.
Zhewei Yao, Amir Gholami, Qi Lei, Kurt Keutzer, and Michael W Mahoney. Hessian-based analysis of large batch training and robustness to adversaries. Advances in Neural Information Processing Systems, 31, 2018.
Zhewei Yao, Amir Gholami, Kurt Keutzer, and Michael W Mahoney. Pyhessian: Neural networks through the lens of the hessian. In 2020 IEEE international conference big data (Big data), pp. 581-590. IEEE, 2020.
Yang You, Jing Li, Sashank Reddi, Jonathan Hseu, Sanjiv Kumar, Srinadh Bhojanapalli, Xiaodan Song, James Demmel, Kurt Keutzer, and Cho-Jui Hsieh. Large batch optimization for deep learning: Training bert in 76 minutes. arXiv preprint
2019.
David Young. Iterative methods for solving partial difference equations of elliptic type. Transactions of the American Mathematical Society, 76(1):92-111, 1954.
Xiaohua Zhai, Alexander Kolesnikov, Neil Houlsby, and Lucas Beyer. Scaling vision transformers. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 12104-12113, 2022.
Guodong Zhang, Lala Li, Zachary Nado, James Martens, Sushant Sachdeva, George Dahl, Chris Shallue, and Roger B Grosse. Which algorithmic choices matter which batch sizes? insights from a noisy quadratic model. Advances in neural information processing systems, 32, 2019a.
Jingzhao Zhang, Tianxing He, Suvrit Sra, and Ali Jadbabaie. Why gradient clipping accelerates training: theoretical justification for adaptivity. arXiv preprint arXiv: 1905. 1881, 2019b.
Jingzhao Zhang, Sai Praneeth Karimireddy, Andreas Veit, Seungyeon Kim, Sashank Reddi, Sanjiv Kumar, and Suvrit Sra. Why are adaptive methods good for attention models? Advances in Neural Information Processing Systems, 33:15383-15393, 2020.
Yushun Zhang, Congliang Chen, Naichen Shi, Ruoyu Sun, and Zhi- Zhi-Quan Luo. Adam converge without any modification on update rules. Advances in Neural Information Processing Systems, 35:
2022.
15

---

Published as a conference paper at ICLR 2025
Yushun Zhang, Congliang Chen, Tian Ding, Ziniu Li, Ruoyu Sun, and Zhi Quan Luo. Why trans- formers need adam: A hessian perspective. arXiv preprint arXiv22402.16788 2024.
Jiawei Zhao, Zhenyu Zhang, Beidi Chen, Zhangyang Wang, Anima Anandkumar, and Yuandong Tian. Galore: Memory-efficient IIm training by gradient low-rank projection. arXiv preprint arXiv:2403.03507, 2024a.
Rosie Zhao, Depen Morwani, David Brandfonbrener, Nikhil Vyas, and Sham Kakade. Deconstructing what makes good optimizer for language models. arXiv preprint arXiv:2407.0799 077.7777.77 2024b.
Lianmin Zheng, Wei-Lin Chiang, Ying Sheng, Siyuan Zhuang. Zhanghao Wu, Yonghao Zhuang, Zi Lin, Zhuohan Li, Dacheng Li, Eric Xing, et al. Judging Ilm-as-a-judge with mt-bench and chatbot arena. Advances in Neural Information Processing Systems, 36, 2024.
Shuai Zheng and James T Kwok. Blockwise adaptivity: Faster training and better generalization in deep learning. arXiv preprint arXiv:1905.0 09899, 2019.
16

---

Published as a conference paper at ICLR 2025
A RELATED WORKS
Understanding of Adam. There is an active line of works trying to understand why Adam works well (Zhang et al., 2019b; Wu et al., 2020; Shi et al., 2020; Zhang et al., 2022; Wang et al., 2022; Pan & Li, 2023; Jiang et al., 2023; Kunstner et al., 2023; Zhang et al., 2024; Ahn et al.; Kunstner et al., 2024). In contrast these works, we point out that Adam's 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 al., 2024) also explores the effectiveness of Adam's 's preconditioner Ddam from linear algebra perspective. They focus on (a variant of) Adam and prove the following result: First, for diagonal dominant (DD) matrix H, when the dimension d is less than K k//3, their modified version of Adam exhibits faster convergence rate compared to gradient descent (GD); Second, for non-DD matrix, the constant terms in Adam's 's upper bound can be much larger than that of GD. Their results take valuable and important step towards understanding D Adam- However, these results cannot fully support our numerical findings in Figure 5. This is because they only provide an upper bound for the non-DD case, while we need a lower bound. We note that is rather difficult to derive the desired lower bound, and we leave it as 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, 2004; Roux et al., 2007; Martens & Grosse, 2015; Sagun et al., 2016; 2017; Chaudhari et al., 2019; Papyan, 2020; Wu et al., 2020; Liao & Mahoney, 2021; Papyan, 2018; 2019; Sankar et al., 2021; Gur-Ari et al., 2018; Yao et al., 2018; Zhang et al., 2019a; Ghorbani et al., 2019; Yao et al., 2020; Dauphin et al., 2024). Inspired by these works, 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, 2018) and its variant CAME (Luo et al., 2023) conduct nonnegative low-rank factorization over Adam's v. SM3 (Anil et al., 2019) is a lightweight version of AdaGrad (Duchi al., 2011). 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 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 Adam W (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, 2019) and NovoGrad (Ginsburg et al., 2019). Both of them use biock-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 states that layer-by-layer parameter partition can lead minimum 6-2 norm solutions. Aligning with the theory, they find that the PyTooch default partition (BAGM-B.) 1) indeed brings overall the best performance on both CIFAR-I0 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 7 and 8. We then propose a new partition strategy Algorithm 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 different design to lst-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-miin mini. It remain unclear whether this design can work on large-scale tasks like I.LMs.
17

---

Published as a conference paper at ICLR 2025
Besides algorithmic design, work also provides understandings of Adam, and particularly, how Adam behaves on generie optimization problems with near-block-diee 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.
Gal (Zhao et al., 2024a) new memory-efficien optimizer for LLMs. Given a gradient matrix g. GaLore calculates low-rank gradient estimator 9 and then calculates m and based on this 9. Adam-mini can potentially be combined with GaLore to reach further memory reduction on The combined method, e.g.. "GaL_oee can further reduce about 40% memory on GaLore and about 81% on Adam in total. Additionally, GaLore-mini can ease the offload burden and enhance the throughput of GaLore, especially when training customer-level GPUs with limited memory.
2. Sophia (Liu et al., 2023) 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., 2021) is a memory-efficient method for SFT tasks. This method fine-tuner 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 result, Adam-mini brings better performance (Figure 22). In parallel to our work, BAdam (Luo et al., 2024) conducts SFT in a block-coordinate-deseee Adam steps to solve the sub bobbrroblem 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., 2021a) also cuts down the memory of Adam. removes the Ist-order momentum and uses neuron-specific projected gradient-style update. According to (Liu et al., 2021a), 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., 2023) uses zeroth-order methods to approximate the gradient information. It is possible combine this idea with Adam-mini to further save memory for SFT. Adam-min can also potentially be combined with other diagonal preconditioned methods (such AdaGrad (Duchi al., 201 1) and Adan (Xie et al., 2022)) as well recent schedule-free optimizers such as SchedulefreeAdam/ W (Defazio et al., 2024).
(BCD) fashion. This method requires repeated
There several other tricks that ease GPU memory burden but are orthogonal optimizer design. These tricks include gradient checkpointing (Chen et al., 2016), model offloading and sharding (Rajbhandari al., 2020; 2021), quantization (Dettmers et al., 2021; Li et al., 2024), and fused update (Ly et al., 2023a;b). Adam-mini can be implemented upon these tricks.
Finally, we discuss another popular adaptive optimizer called LAMB (You et al., 2019) (see Algo- rithm 7 in Appendix E) LAMB might be misunderstood as a similar optimizer to Adam-mine mami, but actually, it not. We emphasize that Adam-mini significantly different from LAMB. First, LAMB still keeps the same coordinate-wise learning-rate caatin design 1/ V as in Adam. Second, in addition to this 1/ /v. v. LAMB re-scales the parameters in layer- fashion. This re-scaling design is often known the "layer-wise learning rates", but to be precise, it is actually an additional "layer-wise scaling" besides the "coordinate-wise learning rates 1/V0". 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 GPT2 pre-training (Figure 10).
'These results are calculated based on (Zhao al., 2024a, Table 1). We consider Llama 2-7B and r 1024 Gal.ore.
18

---

Published as a conference paper at ICLR 2025
B THE COMPLETE FORM OF ADAM-MINI
now present the specific realization of Adam-mini on Transformers and other architectures. To be precise, Algorithm should be renamed as "Partition for CNNs, Diffusion models, and Graph Neural Networks", since we have only tested Algorithm on these models. In the future, is possible that we will have more complicated non- Transsfoorrrrr architectures 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 Principle 1 in Section 2.2.
which Algorithm 3 fails. In
Algorithm 2 Adam-mini in Pytorch style 1: Input weight-decay coefficient À and current step
Algorithm Partition for Transformers
1: @atam_bboccss = {
for name, param in parameters do
2: Choose param hasram__looks blocks from Algorithm or
3: if embed' or output' / in name then Partition param by tokens
3: for param in paramaaaaaaaaa blocks do
5: for i = 1...tokens-Da do
4: g param grad
6: param_block [name+i] -aperiampppaamii
5: param pww.aaaaa---.. 1/t A param
7: end for
6: m (1 31) + m
8:
else if query' or key' in name then
7:
f -
9: Partition param by heads
I-37
8:
v (1-32) m))
10:
for i = 0. heads-1 do param param_blocks blocks [name+1 ]=param [i]
11:
9:
G I-
12:
end for
10: param = param - T/: * Vo+c
13:
else if 'value', 'attn.proj" or 'mlp'
11: end for
in name then
14: Partition param by output neurons
15:
for i = 0
pttappuutnuuunsst_pee do
16:
param_bl bocks [name+il =param
Algorithm 3 Partition for non- con-Transformers
17:
end for
1: param blocks
param_blocks
(
18: else
2: for name, param in parameters do
19:
param blocks [name] "param
3: param_blocks [name] =param
20: end if
4: end for
21: end for
5: return param_blocks
22: return param_blocks
C MORE DISCUSSIONS
Analysis in (Collobert, 2004). We now briefly restate the analysis in (Collobert, 2004) on "why does the Hessian neural networks exihibts wear-block-diagonn structure?", Consider a standard supervised learning problem: minimizing {(,(0, 20,,),, y) where ((.-)) is the Cross-Entropy (CE) loss, f0,,2) En ( œ)) an I-hidden-layer neural net with input æ E Rd, weight w; E Rdd 0/ E and label y E {0, 1}. then the off-diagonal-locc Hessian elements would contain
fe(S(e, x). v) = p(x) (1 p(x)) (w, ø'
if
(3)
Ow,Dwy
where =1/11exp((=///// x))) and d' is the derivative of o(-). Since the training objective to maximize p(x), the term p(z) (1 p(z)) will quickly shrink to zero. This term will push the structure where each block corresponds to one output neuron. The
Hessian to saar-block-diagome authors report that this can happen just after 1 training step, restated in Figure 3 (a). We also numerically reproduce this result on a small MLP on CIFAR-100 in Figure (b,c,d), and show that the wear-block-diagone Hessian structure maintains along training.
Finally, we note that this analysis is rather informal. We leave a more rigorous theoretical study is left for future work.
Why using average v as learning rates. In Line 9 of Algorithm 1, we use the average of v in block the learning rate for that block. We choose such design due the following reasons.
19

---

Published as a conference paper at ICLR 2025
First: grid-search is too expensive. Optimal blockwise learning rates be powerful (as evident in Figure 6), but they are too expensive to search. Such a searching procedure is not scalable.
average of can be borrowed from Adam. Compared to searching all the learning rates from scratch, it is much easier to "borrow" them from the current design of Adam. The average of v is the most natural quantity to "borrow" and performs the best among other candidates such as the maximum of v (see the ablation studies in Appendix D.2), We find that the average of v helps Adam-mini to as effective as Adam (though not significantly surpassing it).
Second:
Third: average of v keeps us close Adam. For neural nets, we find that the average of is a good representative for the whole v in the block, and can help Adam-mini keep close to Adam. The reason comes from backpropagation (BP) rule: for one data sample, the gradient of the weight matrix W E Rdxd can be expressed as G == aW ez E Rd where is certain BP error vector and all share the same BP error term éi, which is usually non-negligible when G + O. Therefore, G usually has similar entries within a row (which associates with the same output neuron), and its mean value can be good representative of the whole row. As result, we find that Adam-mini's trajectory closely resembles that of Adam (see the curves in Figure and Figure 10). One resulting advantage is that Adam-mini can maintain the scaling laws of LLMs trained by Adam, while substantially saving the training cost (see evidence in Figure 11).
Dé
the input feature the current weight. For all entries in the i-th row of G. they
D MORE EXPERIMENTAL RESULTS
D.1 MORE RESULTS FOR MOTIVATION
In Section 2, we showed that "Adam's 's v is redundant on dense Hessian subblock". We provide experiments on random quadratic functions to support the claim. Here, we conduct the following new experiments on a 1-layer Transformer. We will show that "using single learning rate per block" is also sufficient for Transformers.
The following Exp 1 and 2 extend the random quadratic experiments in Figure 4 and to Trans- formers.
Exp 1: Adam's learning rate is redundant on the dense Hessian subblock. We take some small dense blocks in the Hessian I-layer Transformer and denoted as H. We compare x(H) and <<D Adam H) as in the paper. We find Adam is not effective in reducing the kappa of these blocks, and many Irs in Adam can be redundant.
Table 3: Comparison of k(H) and k(D H) for the dense blocks in the Hessian of 1-layer Transformer.
Hessian Block
k(H)
K Adam H)
1st head in Query
103.80
176.88
1st head in Key
103.46
213.82
1st head in Value
165.66
332.76
1st neuron in attn.proj
39.92
94.56
1st neuron MLP fel
22.04
70.92
1st neuron in MLP ________________________________ proj
63.85
236.71
Exp 2: Single learning rate per block is sufficient. We conduct the "block-wise GD" and grid-search the learning rate for each block. The result is shown in the following figure. We find that block-wise GD outperforms AdamW. This extends the setting from the random quadratic problem in Figure
that Adam is redundant on the dense Hessian subblocks (Exp 1), and a single Ir for each block can work well (Exp 2). These experiments show that our conclusions on random quadratic problems can be extended to Transformers.
Combining Exp and 2,
20

---

Published as a conference paper at ICLR 2025
Results on a 1-aaer Transformer
1.8
Adam
Adam mini
1.6
Blockwise GD
Validation Loos
1.4
1.2
20
40
60
80
100
lteralion
On 1-layer Transformer, we conduct the "blockwise GD" and we grid-search the learning rate for each block. We find that blockwise GD outperforms AdamW.
Figure 14:
D.2 ABLATION STUDIES ON THE DESIGN OF ADAM-MINI
We here provide more reasons why choose mean(v) as the blockwise learning rates. We conduct ablation studies on different choices quantities that we can borrow from Adam, including 2-norm(v), 1 1-norm(v), max(v), and min(r). we found that all these candidates perform worse than mean (V) in Adam-mini.
Llama 2-20M Pre- Pre-training
5.0
4.8
Validation loss
4.6
1.4
(1-norm(v))-Sed
Adam-mini (2-norm(v)--5e-
(ma
Adam-mini (minn())--5e33
Adam mini---in---
4.0
500
1000
1500
2000
2500
3000
3500
4000
Iteration
Figure 15: Ablation studies on the design of Adam-mini. We find that mean(v) performs better than other candidates. The blue curve does not show because the algorithm diverges and the curve is of range.
D.3 MORE RESULTS ON THE SCALING LAW EXPERIMENTS
The complete loss curves Llama 2-1B, We here present the complete validation loss curve of Llama 2-1B, training 20B tokens, which corresponds to the rightmost curve in the scaling law experiments Figure 11 (a). We note that this is a complete pre-training run under the definition of Chinchila's law (Hoffmann al., 2022). We find that Adam-mini's loss curves closely resemble those of AdamW.
The final validation perplexity. In Table present the final validation perplexity for all models after training on the token amount suggested by Chinchilla's 's law (Hoffmann et al., 2022), For all models from 39M IB, find that Adam-min reaches a lower validation perplexity than AdamW,
D.4 GPT-4 EVALUATION SCORE OF SET AND RLHF
On the pre-trained Llama 2-7B (released by Meta (Touvron et al., 2023)), we run Adam-mini downstream tasks including SFT and RLHF. We evaluate the alignment performance in terms of chat ability using the MT-Bench (Zheng al., 2024), where GPT-4 assesses multi-tum chatting capabilities and assigns a score from 0 to 10 (higher is better). Our results, presented in Table demonstrate that Adam-mini outperform AdamW.
21

---

Published as a conference paper at ICLR 2025
Adami Adam -in-22---1-118
Val loss
25000 50000
75000
100000
150000 Iteration
125000
175000 200000
Loss curves of pre-training Llama 2-1B on 26B tokens. This is complete pre-training run under the definition Chinchila's ' law. We find that Adam-minn performs similarly to AdamW throughout the training, but with 50% less memory.
Figure 16:
Table 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
IB
Total tokens
1.02B
1.76B
2.67B
4.25B
7.10B
26.21B
Adam W
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
Table 5: Averaged GPT-4 evaluation score (1) of SFT and RLHF on the MT-Bench.
(LoRA)
SFT
RLHF
AdamW
Adam-mini
AdamW
Adam-mini
AdamW
Adam-mini
MTT-Bench
4.23
4.41
5.37
5.40
5.54
5.68
D.5 NON-LLM TASKS
We now evaluate Adam-mini on non-LLM tasks. Table 6 shows the results for the training of ResNet18(He et al., 2016), Swin-Transformer (Liu et al., 2021b), DiT-XL-2 (Peebles & Xie, 2023), DC-AE-Diffusion (Chen et al., 2024a) on ImageNet, DDPM diffusion model (Ho et al., 2020) on CelebA, Graph Convolution Net (GCN) (Kipf & Welling, 2016), and Graph Attention Net (GAT) (Veličković et al., 2017) on OGB-arxi The training curves are shown in Figure 17 and 18. We find the performance of Adam-mini be comparable or better than AdamW, but with less memory.
Table 6: On popular non-LLM tasks, Adam-mini performs on par or better than AdamW.
Domain
Model
Optimizer AdamW W
Metric
25% steps
50% steps
75% steps
100% steps
Vision
DDPM
Train loss ())
0.0529
0.0497
0.0420
0.0394
Vision
DDPM
Adam-mine mini
Train loss ())
0.0525
0,0495
0,0416
0.0388
Vision
ResNer18
AdamW
Val ace (t)
0.6149
0.6478
0.6613
0.6669
Vision
ResNet18
Adam-mini
Val acc (1)
0.6140
0.6501
0.6629
0.6667
Vision
Swin-Trane Transformer
AdamW W
Val acc (f)
0.6290
0.6940
0.7180
0.7310
Vision
Swin- Transformer
Adam-mini
ace (f)
0.6230
0.6960
0.7
0.7300
Vision
DiT-XL-2
AdamW
Train loss (J)
0.1605
0.1696
0.1607
0.1431
Vision
DiT-XL-22
Adam-mini
Train (
0.1601
0.1693
0.1605
0.1430
Vision
DC-AE-Diflusion
AdamW
Train loss (.)
0.2860
0.2820
0,2800
0.2780
Vision
DC-AE-Diftssion
Adam-mini
Train loss (U)
0.2860
0.2830
0.2800
0.2780
Graph
GAT
AdamW
Val ace(f)
0.7277
0.7367
0.7399
0.7421
Graph
GAT
Adam-mini
Val acc (f)
0.7378
0.7394
0.7403
0.7429
Graph
GCN
AdamW
Val acc (f)
0.7347
0.7428
0.7379
0.7374
Graph
GCN
Adam-mini
Val ace (f)
0.7406
0.7427
0.7380
0.7423
22

---

Published as a conference paper at ICLR 2025
In Table 7, further evaluate the image quality from the model trained by Adam-ming We find that the Adam-mini performs on par with AdamW.
Table 7: Evaluation scores: Adam Adam-mini performs on par with AdamW.
Domain
Model
Optimizer
FID (J)
Inception Score (1)
Vision
DiT-XL-2
AdamW
91.83
12.38
Vision
DiT-XI--2
Adam-mini
88.20
13,90
Vision
DC-AE-Diffusion
AdamW
34.72
41.79
Vision
DC-AE-Diffusion
Adam-mini
33.15
44.38
0.7
03
0.00
6,05
0.50
$6323
1.10
Denns
333g
Epacs
Swin-Trrasformes
(b) DiT-XL-2
(c) DC-AE-Diffusion
Figure 17: The training curves of Swin- Transformer, DiT-XL-2, and DC-AE-Diffusion, We find that Adam-mini performs on par with Adam W
2.7
RewNeee 16 en Emaps
0.047
2,060
0.0-
5.045
0.3
2.040
5.0
2.2
Uporh
Steration
105
(a) ResNet-18
(b) DDPM
Figure 18: The training curves of ResNet-18 and DDPM diffusion model. We find that Adam-mini performs par with AdamW.
D.6 MORE DISCUSSIONS ON THE PARTITION STRATEGIES OF VA LUB
As shown in Figure 7, the Hessian structure of val ue is less clear compared to other blocks: 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 8 (with more than 50k total steps) and the scaling law experiments of Llama models in Figure 11 (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 10. Based on these findings, we recommend using strategy (I) when the total number of training steps large, and using strategy (II) if otherwise. Note that strategy (II) can be used simply by adding one line code after creating the optimizer: opt imi zer. WV names { }.
D.7 DETAILED COMPARISON WITH ADAFACTOR
In this section, conduct more hyperparameter search for Adafactor on Llama 2-20M pre-training. We will focus on tuning Adafactor Adafactor-Zhai-verieee Zhai-version since it performs better than the original Adafactor (see Figure 13). We consider the following three setups.
23

---

Published as a conference paper at ICLR 2025
Setup 1: We change the default B2 0.999 to B2 0.95 and sweep over learning rates.
Setup 2: We use learning rate 5e-3, A2 0.95 and sweep over warm-up step
(1%, 2%, 3%, 4%, 5%, }0% total steps.
Setup 3: We use learning rate 5e-3, P2 0.95 and warm-up step 1% total steps and sweep over {10- 30 10-16 10-5 10-6).
375
3.5
540
8522
A22
3XC
15:0
384 25K
3R1
483
58 3X0
1560
3006
3R1 3590
409
320
1500
300
XX0
35K
(a) Llama 2-20M
(b) Llama 2-20M
(c) Llama 2-20M
Figure 19: The training curves of Adafactor-Zhai- -erssonn on Llama 2-20M pre-training. (a.b.c) corresponds the aforementioned Setup 1, 2, 3, respectively. We find that Adafactor consistently underperforms Adam-mini.
The results are shown in Figure 19, In these cases, Adafactor-Zhann Zhai-version consistently underper- forms Adam-mini and the change of hyperparameters does not help much.
D.8 DETAILED COMPARISON WITH LION
We now conduct the hyperparameter grid search over Lion. We find that Lion is not easy to tune and have not managed to make Lion work. We consider the following settings.
Tuning strategies in (authors, 2024), authors (2024) carefully tune Lion on Llama models (150M, 300M, 600M, 1.2B). We will adopt their optimal tuning strategies in (authors, 2024, Table 1) 6. We here summarize their key messages.
Message 1: The optimal learning rate (Ir) of Lion is usually 10 times smaller than AdamW.
Message 2: The magical number Ir = 16e-4 works the best for most models (Llama 150M, 300M, 600M).
Message 3: ;3 = {0.95, 0.9) perform similarly and perform significantly better than other 81 candidates including 31 = (0.99, 0.98, 0.8, 0.5, 0).
Message 4: 3g (0.99, 0.98, 0.95) perform similarly and perform significantly better than other 82 candidates including B2 = (0.9999,0.. 0.999, 0.995, 0.9, 0.8}.
In the following, we will use the above messages to tune the hyperparameters of Lion.
Architecture. We consider Llama 2-20M, which is the same architecture as the ones investigated in (authors, 2024), but with different model size. We also consider GPT-2-125M. which is a task that
Lion is not tested before (neither (authors, 2024) nor other literatures to our knowledge).
Our tuning strategies. Following the above Message 1 and 2 from (authors, 2024), we will use the following tuning strategies for Lion on Ilama 2-20M and 1PT-22125MM
Learning rate for Llama 2-20M: The standard Ir is Se-3, so we try Ir = [5e-4, 6e-4, 7e-4, 8e-4, 9e-4, 1e-3, 2c-3, 3e-3, 4e-3, 5e-3]. For completeness, we also investigate Ir = [4e-4, 3. 3.16--4, 2e-4, le-4].
Learning rate for GPT-2-125M: The standard 6e-4, so we try Ir [6e-5, 7e-5, 8c-5, 9c-5, le-4, 2e-4, 0.000316, 4e-4, 5e-4, 6e-4]
As for (81, 32), we will use (8,, 02) (0.95, 0.98). We use these hyperparameters for two reasons. First, they are the optimal choice among other candidates by Message and 4. Second, (31, 32) =
'We would like to mention that (authors, 2024) is concurrent work us, and their tuning strategies in is not public available hy the time we submitted this script.
24

---

Published as a conference paper at ICLR 2025
(0.95, 0.98) is recommended by the authors of Lion to be "helpful in mitigating instability during training"
Liama 223 teshipg
10'
GPF2-125M Pre-training
7.5
7.0
Train 1.00 6x
Lise
Lis 4a0-440
Unm a Hol
x 10
Lia-ded
Lio
4.
3.0
500
1303
2000
Beesban
2030
2300
3000
3:00
4002
16a
10
Iteration
10
(a) Llama 2-20M
(h) GPT-2-125M
Figure 20: The training curves of Lion on Llama 2-20M and GPT-2-125M pre-training. The hyperparameters are chosen under the optimal strategies by (authors, 2024). We find that Lion consistently underperforms Adam-mini on Llama 2-20M, and encounters loss spikes on GPT-2- 125M.
The results are shown in Figure 20, After using the above tuning strategies, we find that Lion still underperforms Adam-mini and AdamW on Llama 2-20M and GPT-2-125M, In particular, Lion encounters loss spikes on GPT-2-1 125M for all the learning rate candidates above.
We summarize findings on Lion below.
First: worse performance. With all the effort above, haven't managed make Lion work, and we haven't t been able to reproduce (authors, 2024) on Llama 2-20M and GPT-2-125M (different model size and architectures from (authors, 2024)). One possible reason is that Lion might work under their specific setup (dataset, architecture, batch size, etc.), but the effectiveness is not easily transferable.
Second: no general tuning guidance. We find that there are no general tuning strategies for Lion. We emphasize that authors (2024) only focuses on Llama architectures, and their resulting tuning strategy seems not robust and transferable to other architectures. In particular, their optimal strategy on Llama causes loss spikes on GPT-2. To our knowledge, Table in (authors, 2024) is the only public tuning strategy for Lion, so it seems unclear how to tune Lion in general.
In contrast, Adam-mini is much easier to use. In all our experiments (a wide range of tasks and models), Adam-mini performs on par with Adam using the same hyperparameters as AdamW (including learning rate, 81, B2, etc.). We believe that "easy adaptation of the hyperparameters" can serve as one advantage of Adam-mini over Lion, apart from the performance superiority.
Third: A Adam-mini is more principled and explainable. We emphasize that Lion is designed by symbolic search, and design principle is largely unclear. In contrast, the Design principle of Adam-mini is much more understandable: we remove the redundant Irs in Adam according to the Hessian structure. We believe Adam-mint is more "white-box" than Lion and more explainable to users.
D.9 ADDITIONAL FINDINGS ON GPT-2-330M
In Figure 21, we report some unexpected results on GPT-2-330M. When using the recommended learning rate 3 by (Liu al., 2023), find that AdamW encounters loss spikes, while Adam-mini does not. The loss spike of Adam can be mitigated by changing c from le-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.
D.10 COMBINING WITH LORA
We note that Adam-mini can be combined with LoRA (Hu et al., 2021): : we can change the Adam steps in LoRA by Adam-mini. The results are shown in Figure 22. We find that such changes with
Phtpps://4 thub. com www.//uciiaaaaaaaa wwwpii///nnnnnnpoooooo
25

---

Published as a conference paper at ICLR 2025
AdamW-3e 4--------
AdamW-3e-4-ee 30-4-4p6-1e-8
Adom mins 3e e-s-- 6--6
Validation Loss
6-
Adam-mini-S3 3e---n---------1------------
1.0
15
20
(billion)
Figure 21: On GPT-2-330M, AdamW encounters loss spikes, while Adam-min mini does not. The loss spike of AdamW can be mitigated by changing from 1c-8 to 1e-6.
boost performance. We grid-search learning rates for both methods and report the best performance. The detailed learning rate contiguration is shown in Section F.
Alm-im
2.45
2.43
2.59
257
245
Figure 22: Training curves of SFT (LoRA). We find that LoRA performs better when replacing the Adam steps by Adam-mini.
D.11 SAMPLE RESPONSES FROM LLMS TRAINED ADAM-MINI
The sample response in shown in Figure 23.
Compose engaging travel blog post about a recent trip Hawaii, highlighting cultural experiences and must-see aftractions.
Prompt:
Sure! Here's 's travel blog post about a recent trip to Hawaii: Whether you'te 're interested in history, nature, simply soaking up the sun, Hawaii has something for everyone. Until next time, Aloha! Adam-minis Sure! Here's travel blog post about a recent trip to Hawaii: highly recommend visiting this bcautiful archipolago for a truly unforgettable adventure.
AdamW:
Prompt: Rewrite your previous response. Start every sentence with the letter A.
AdamW: A trip Hawaii is must-visit destination for any traveler, offering variety of cultural experiences and must-see attractions. From Polynesian Cultural Center to 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
must-do for any travel enthusiast. trip the beautiful Hawaïan Islands a paradise on earth. A visit Oahu 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 great way to learn about the history and traditions of the islands. A traditional luau is must-attend event. Overall, trip Hawaii is an unforgettable experience.
Adam-mini: ni: trip Hawan
Figure 23: Response Sample of Llaaa 2-7B finedtuned by AdamW and Adam-mini to the #81 test question from the MT-Bench (Zheng al., 2024) dataset. Since there is little information the first round response, omit the answers from the models.
26

---

Published as a conference paper at ICLR 2025
E SOME PRELIMINARY RESULTS
E. 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 along the training. Both and vectors of the same size as # model parameter.
Algorithm 5 Adam in Pytorch style
1: Let À be the weight decay coefficient
2: for paran in parameter_blocke do
3: g= p=paarm...... grad
4: i 0 then
5:
g = g + À param
6:
end if
7:
param = param -
8:
m (1 *
9:
m 1-
10:
(1 (1-3) 82) 1 82 v
11:
G 1-3
12: param param=param- 1/1 Vo+c
13: end for
Algorithm 6 AdamW in Pytorch style
1: Let À be the weight decay coefficient
2: for param parameter_blockk do
3: g param.. grad
4: param = param 7/1 * X=
5:
(1- - (31) q +
6: f I-3
7:
v (1-32) (22 + 32 v
8:
V T-8
9: param = param 1/: Vo+c
10: end for
Algorithm 7 LAMB in Pytorch style
1: Let À be the weight decay coefficient, let á be scaling function.
2: for param
_all___ayerr___ do
3: g= par.m. grad
4: paara-= --r-aaaammmmm 1/t * X
5:
[1 +
6:
m 1-
7: v (1-32) -+2)* +
8:
& 1-37
9:
Vote
param ttra=paaree param T/ * a Allpa
Ir Xeparam
end for
E.2 PRELIMINARY RESULTS IN (ZHANG ET AL., 2024)
We here restate (Zhang et al., 2024, Figure 3). This figure shows that: for Transformers, different parameter blocks have different Hessian eigenvalue distributions, while for CNNs, the eigenvalue
27

---

Published as a conference paper at ICLR 2025
distributions are similar among blocks. This suggests that Transformers need different learning rates for different blocks to handle the heterogeneity in eigenvalue distributions.
(a) VGGI6
(b) VGGI6
BERT
BERT
(a) VGG16
(b) BERT
Figure 24: Figure 3 in (Zhang et al., 2024). The eigenvalues distribution are similar among blocks for CNNs, while they differ significantly across blocks for Transformers. This indicates Transformers need different learning rates for different blocks to handle the heterogeneity in eigenvalues.
F EXPERIMENTAL DETAILS
F.1 TRAINING CONFIGURATIONS FOR SECTION
Unless mentioned otherwise, 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, use the recommended learning rates by (Liu al., 2023), which reported to be optimal by grid search. Unless mentioned otherwise, Adam-mini, Adafactor, CAME, SM3, and LAMB use the same learning rate 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 31 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" to train GPT2 sized 125M (small), 330M (medium), and 1.5B (XL. on Openwebtext. For all models,
sea 1 en 1024, batch size 480, weight decay coefficient = 0.1, e = le-8, 31 = 0.9, 3g = 0.95. We use cosine-decay learning rate schedule with 2000 iterations of warm-up. For GPT2-small and medium, recommended peak learning rate by (Liu et al., 2023), 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?, The chosen peak learning rates are 6e-4, 3e-4, le-4 for emall, medium, XL, respectively. The minimal learning rate is chosen as 3e-5, 6e-5, le-5 for these models.
use the
Llama pre-training. For all experiments the Llama series (from 20M to 13B), Torchtitan codebase 10 and C4 dataset (Raffel et al., 2020). For all experiments, we use weight decay 0.95. For Llama 2-1B Llama 3-8B, we learning
the
coefficient À= = 0.1, 1e-8, 31 0.9, rate = 3e-4. For Llama 2-13B, we use learning rate = Ic-4. As for the learning rate schedule, we use warm-up step = % total step and use linear decay schedule after the warm-up (this the default setting in the Torchtitan codebase). For Figure 10 (a) and all the experiments of Adafactor and Lion, we use seq seq_ien len = 512 and batch size 128. For Figure 10 (b), we use seq len 2048 and batch size 8. For Figure 10 (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.
"https: / thub. http://www/aaatttyyo//////eee/ee 9htpps http://///stttaooooorr////eeaaaa config/gpt2_ yaml "hhtps:///o thub. himbb..om/ggtorch/ttrcttttt
ob/e 83c88 Cec5 971 7111 2d4 ab08a 602 6806632232
28

---

Published as a conference paper at ICLR 2025
SFT and RLHE. We use the Llama 2-7B pretrained model (Touvron et al., 2023) for our study. We use the ultrafeedback dataset 11 The implementation of SFT and RLHF code is based on the ReMax codebase 2 Specifically, we train 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., 2023), memory-efficient alternative to PPO (Schulman et al., 2017); to optimize the preference reward.
We use DeepSpeed ZeRO-2 in our training. GPT-4 evaluation template in Table is from the codebase In the reward optimization stage, We use ReMax, a memory-efficient alternative to PPO. We use UltraFedbbckk dataset Cui al. (2023) 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 (1c-6, 2c-6, 3e-6, 4e-6, 5e-6, 1c-5, 2c-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 as recommended by LlamaFactory" The rest of the hyperparameters of Adam W and Adam-mini are 1e-8, 31 = -09, B2 = 0.95.
ReMax. We use 48 samples in batch and train the model for epoch. By searching the peak learning rate from (5e-7, 1c-6, 2e-6) based on validation reward, AdamW uses 1e-6 while Adam-mini selects 5e-7 the peak learning rate. The weight decay coefficient is set to O. The rest of the hyperparameters of Adam W and Adam-mini are c 1e-8, B = 0.9, 32 = 0.95.
ResNet. We use the PyTorch official implementation codebase to train ResNet18 (He et al., 2016) on ImageNet (Deng et al., 2009). We use cosine- -ecay learning rate, epoch =90, B3 = 0.9, B2 = 0.999, e = 1c-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 Adam W.
DDPM. We use the codebase 16 to train DDPM diffusion models (Ho et al., 2020). The image size is 64 and the training objective is to predict the noise as in (Ho et al., 2020). We use the default U-Net and the dimension multiply in U-Net is (1, 2, 4, 8). We use the CelebA " 17 and train the diffusion model with learning rate 5 x 10-5 with cosine decay. The batch size is 128 and the training epoch is 50.
archiecture hyper-paramend
Tronforranformmrs. We use the official implementation 18 of Sran-Trrnssormrrs (Liu et al., 2021b) on ImageNet. All configurations as default. For both Adam-mii mini and AdamW, we use the default leamning rate 5e-4.
D We use the official implementation 19 of DiT-XL-2 (Peebles & Xie, 2023) on ImageNet. All configurations as default. For both Adam-mini and AdamW, we use the default learning rate le-4.
DC-AF-Diffusion We use the official implementation" 20 of DC-AE-Diffusion (Chen et al., 2024a) on ImageNet. All configurations as default. For both Adam-mini and AdamW, we use the default learning rate 2e-4.
cleaned
do do_ae/RHAIMM..
29

---

Published as a conference paper at ICLR 2025
Graph Neural Networks. We use the DGL implementation" of Graph Convolution Networks (GCN) (Kipf& & Welling, 2016) and Graph Attention Networks (GAT) (Velickovic et al., 2017) for OGBN-arxiv dataset. All configurations default. For both Adam-mint and AdamW, default learning rate = 0.005 for GCN and the default learning rate = 0.002 for GAT.
the
Scaling law experiments. We use the codebase Torchtitan to train Llama models of different sizes. All the model configurations shown in Table and all the training configurations shown in Table 9. The experimental setups are inspired by (Hägele et al., 2024). In all experiments, we fix the warm-up steps to be 1% the total steps, as suggested by (Ibrahim et al., 2024).
Model Size
dmodel
Mayers
Nheads
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 8: The model configurations in 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. 0.13M
20.4K
2.67B
26.17
162M
6e-4
0.13M
32.4K
4.25B
26.23
271M
6c-4
0.13M
54.2K
7.10B
26.21
IB
2e-4
0.13M
200K
26.21B
26.21
Table Training configurations for the scaling law experiments.
Trajectory comparison in Figure 10 (c). We train a 8-layer Transformer sized I 11M on Openwebtext and launch AdamW, Adam-mini, and other memory-efficient optimizers under the same random seed and same learning rate le-5. We save the model weights for every 250 iterations and compare their Euclidean distance to the weights along AdamW's trajectory.
F.2 DETAILED SETUP FOR OTHER EXPERIMENTS
Configurations for Figure 3. We train 1-hidden-layer MLP with 8 neurons CIFAR-100. We Adam with le-4 learning rate with Cosine decay schedule. We train the network with 128 batch size and 20 epochs (i.e., about 7900 total training steps). With the help of auto-differentiation framework, we calculate the Hessian with two passes of backpropagation (Pearlmutter, 1994) and the calculation exact.
Configurations for Figure For each dense sub-block H 1, 2,3, we use random positive definite matrices. We fix the choose the cigenvalues of each Hi follows: for = 1, we independently sample from (1,2,3) for times; for 1 2, we repeat this procedure for 199,100,101 ): for 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 + p). where L, ir are the largest and smallest eigenvalue of the (blockwise) Hessian. We use Adam with B1 0. This helps focus on the effect of coordinatewise learning rate in Adam. We also set B2 = to the time-varying learning rate. This is necessary because, for any 82 1, Adam with constant learning rate will oscillate quadratic functions. This is theoretically proved in (Da Silva Gazeau, 2020, Proposition 12, Figure 1) and empirically observed in (Zhang et al., 2024, Section 3.3).
thub. http://ttwwotooeeeeeeee///e////a ogbn-arxiv bhttps://ogb. www.sstttaaaaoooooooooo staafforddeeu//oo/ss//eeeeeooo thub. httrpss///ttbbb.o.//rrrrrrrrrrr
30

---

Published as a conference paper at ICLR 2025
Configurations for Figure 5. To generate positive definite matrix Ha, first uniformly sample did-1) independent angles 0, from the interval -- } where j. Starting with the identity matrix, we perform a rotation of the i-th and j-th rows by the angle 02, for each sampled pair. Through did-1) rotation operations, we obtain the orthogonal matrix Q. We define A diag(x, L....,
... 1). and the matrix H₁ is generated using the expression Hp = QAQ". The python code for Hs generation listed as follows:
def generate_Hb (theta, kappa, d)
Q np. np.eye(d)
for i in range
for in range (i+1,d)
p np.
P[i,i] math. cos (1,
e[i,j] math. math. (tt
P[j, -math. [
P[j,j] math.cos (theta (i, 31)
10
Q
11
Lamboa np. up.eye (d)
12
Lamboa 0, 01 = kappa
return @ Lanbda Q. Q.transpose
We note that approaches 0, the diagonal-overrof--diaaeee the sampled values of 0, we utilize RM to produce H,, with varying ratios, where R { k 1,,,...55p0 50). For each matrix, we sample 100 initial points from the Xavier initialization distribution to compute the resulting к of Adam algorithm. For each pair of d and K, we sample 40 different values. By averaging the results obtained, we plot the Figure 5.
ratio of the matrix Q decreases. For
Configurations for Figure 7. We use the codebase We enlarge the vocabulary size from 2 to and we change the objective function to the standard pre- pre-training objective function (i.e., predict the next token). We consider 1-layer Transfomer with emb 16, head 4, and the width (i.e., the number of output neurons) of mlp. Ec __ equals 32. The Hessian calculation follows the same procedure in Figure 3.
is worth mentioning that the weight matrices might have different shapes depending on the different codebases. For instance, the weight matrix W in codebase might be implemented as W 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 7, we re-order the Hessian rows & columns such that the block diagonal structure visible.
Throughput Comparison in Table 2. turn on CPU offload. We report the throughput from the summary file the Wandb log.
The results are tested on 2x A800-80GB GPUs. We did not
/https://colal research ottpg///caoooo////////oooooaooen usp-sharing
31


Reviews:
[{"review_id": "TLTMPELH2Z", "reviewer": "Reviewer_AcGY", "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.", "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?\n\n2. 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?\n\n3. 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?\n\n4. 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?\n\n5. 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?\n\n6. 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?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 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.  \n\n2. 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.  \n\n3. 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.", "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.\n\n- 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.\n\n- 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.\n\n- 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.\n\n- 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.\n\n- 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.\n\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "YiVceWWwJj", "reviewer": "Reviewer_waYV", "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.", "questions": "See weaknesses.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. This paper studies an important problem and provides an interesting and novel solution. \n2. The paper is well-written with clear motivation and good presentation flow.\n3. Extensive experiments are conducted to justify Adam-mini's effectiveness.", "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? \n2. It will be good to add discussions/insights about how to adapt hyper-parameters when switching from AdamW to Adam-mini.\n3. 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:\n- https://github.com/facebookresearch/DiT\n- https://github.com/mit-han-lab/efficientvit/blob/master/applications/dc_ae/README.md", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "3PckT3Yu0c", "reviewer": "Reviewer_8UzS", "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.\n\nThe 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. \n\nThe 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.", "questions": "Minor referencing errors and typos:\n\n- Line 294: instability on 1B models (see Figure 7 (d)) → instability on 1B models (see Figure 8)\n- Line 349-350: As shown in Figure 7 (d). This strategy → As shown in Figure 8, this strategy \n- Lines 1053-1054: The results are shown in Figure 15 → The results are shown in Figure 16\n- Lines 418-419: Figure 9 (b) shows the loss curves → Figure 9 shows the loss curves\n- Line 483: As shown in Figure 12 → As shown in Figure 12 (a)\n- Table1 does not seem to be referenced except in the caption of Figure 14 in the appendix \n- There appears to be no reference to Figure 19 in Section D.2\n- What is “Adam-mini-2e-4-1B-v” in Figure 12 (a)?\n- Line 374: BP in “where e is certain BP error” is defined 2 lines later. \n- Lines 449-450: This might because → This might be because", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "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.", "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.\n\nMoreover, 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "QAQw51Ho7U", "reviewer": "Reviewer_rJgs", "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.", "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.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. This paper focus on an interesting and important problem. memory-efficient optimization is very important for llm training. \n\n2. The paper provides the results from diverse domains to verify the efficiency and performance of Adam-mini. \n\n3. This paper provides the results on LLM pre-training, SFT and RLHF to verify the performance of proposed method Adam-mini.", "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\ninitialization 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. \n\n2. 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. \n\n3. 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. \n\n4. 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. \n\n5. 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. \n\n6. 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.\n\nAlthough 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. \n\n\n[1] Deconstructing What Makes a Good Optimizer for Language Models", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)