Addax: Utilizing Zeroth-Order Gradients to Improve Memory Efficiency and Performance of SGD for Fine-Tuning Language Models
Authors: Zeman Li
Xinwei Zhang
Peilin Zhong
Yuan Deng
Meisam Razaviyayn
Vahab Mirrokni
Venue: ICLR 2025 Poster (2025)
Keywords: Large Language Models
memory efficient Fine-tuning
Zeroth order optimization
Primary area: optimization

Abstract: Fine-tuning language models (LMs) with the standard Adam optimizer often demands excessive memory, limiting accessibility. The ``in-place'' version of Stochastic Gradient Descent (IP-SGD) and Memory-Efficient Zeroth-order Optimizer (MeZO) have been proposed as solutions to improve memory efficiency. However, IP-SGD still requires a decent amount of memory, and MeZO suffers from slow convergence and degraded final performance due to its zeroth-order nature. This paper introduces Addax, a novel method that improves both memory efficiency and algorithm performance of IP-SGD by integrating it with MeZO. Specifically, Addax computes the zeroth-order or first-order gradient of the data points in the minibatch based on their memory consumption and combines zeroth- and first-order gradient estimates to obtain the updated direction in each step.
By computing the zeroth-order order gradient of data points that require more memory and the first-order gradient of the ones that require less memory, Addax overcomes the slow convergence of MeZO and excessive memory requirement of IP-SGD. Additionally, the zeroth-order gradient acts as a regularizer for the first-order gradient, further enhancing the model's final performance.
Theoretically, we establish the convergence of Addax under mild assumptions, demonstrating faster convergence and less restrictive hyper-parameter choices than MeZO. Our extensive experiments with diverse LMs and tasks show that Addax consistently outperforms MeZO in terms of accuracy and convergence speed, while having a comparable memory footprint. 
In particular, our experiments using one A100 GPU on OPT-13B model reveal that, on average, Addax outperforms MeZO in terms of accuracy/F1 score by 14%, and runs $15\times$ faster, while having a comparable memory footprint to MeZO. In our experiments on the larger OPT-30B model, on average, Addax outperforms MeZO in terms of accuracy/F1 score by >16% and runs $30\times$ faster on a single H100 GPU. Moreover, Addax surpasses the performance of standard fine-tuning approaches, such as IP-SGD and Adam, in most tasks in terms of Accuracy/F1 score with significantly less memory requirement.

Published as a conference paper at ICLR 2025
ADDAX:
UTILIZING ZEROTH-ORDER GRADIENTS TO IMPROVE MEMORY EFFICIENCY AND PERFORMANCE OF SGD FOR VINE-TUNING LANGUAGE MODELS
Zeman
Xinwei Zhang
Peilin Zhong
USC
USC
Google Research
zemanli (isaaani@usceeuu
xinveiz@usc.edi
peilinzegoogle com
Yuan Deng
Meisam Razaviyayn
Vahab Mirrokni
Google Research
USC
Google Research mirrokni@google..om
dengyyuanggoggle.com com
razaviyalusc.
ABSTRACT
Fine-tuning language models (LMs) with the standard Adam optimizer often demands excessive memory, limiting accessibility. The "in-place" version of Stochastic Gradient Descent (IP-SGD) and Memory-Efficient Zeroth-order Op- timizer (MeZO) have been proposed as solutions improve memory efficiency. However, IP-SGD still requires decent amount of memory, and MeZO suffers from slow convergence and degraded final performance due to its zeroth- nature. This paper introduces Addax. novel method that improves both memory efficiency and algorithm performance Specifically, Addax computes the zeroth-order or first-order gradient of the data points in the minibatch based on their memory consumption and combines zeroth- and first-order gradient estimates to obtain the updated direction in each step. By computing the zeroth-ordes order gradient of data points that require more memory and the first-order gradient of the ones that require less memory, Addax overcomes the slow convergence of MeZO and excessive memory requirement of IP-SGD. Ad- ditionally, the zeroth-order gradient acts regularizer for the first-order gradient, further enhancing the model's final performance. Theoretically, we establish the convergence of Addax under mild assumptions, demonstrating faster convergence parameter choices than MeZO. Our extensive experiments with diverse LMS and tasks show that Addax consistently outperforms MeZO terms of accuracy and convergence speed, while having a comparable memory footprint. particular, our experiments using one A100 GPU on OPT-13B model reveal that, on average, Addax outperforms MeZO in terms of accuracy/F1 score by 14%, and runs 15x faster, while having a comparable memory footprint to MeZO. In our experiments MeZO in terms of accuracy/F1 score by > 16% and runs 30x faster on a single H100 GPU. Moreover, Addax surpasses the performance of standard fine-tuning approaches, such as IP-SGD and Adam, in most tasks in terms of Accuracy/F1 score with significantly less memory requirement.
IP-SGD by integrating it with Me7Ó.
and less restrictive hyper-parameter
the larger OPT-30B model, on average, Addax outperforms
INTRODUCTION
Fine-tuning pre-trained language models (LMs) a crucial step a wide range of natural language processing tasks, including text classification and sentiment analysis (Devlin et al.) 2019), well as their use different domains (Gururangan et al. 2020 Schwwrtz et al. 2019 Lu et al. 2019). However, standard fine-tuning approaches (with Adam optimizer) demand excessive memory re- quirement due to gradient and/or the optimizer state storage, presenting a challenge as LMs grow in scale (Brown et al. 2020; OpenAl 2023). For instance, fine-tuning 13-billion parameter --biltion-paramtter model like OPT (Zhang et al. 2022) in mixed precision requires over 316 GB of memory, hindering accessibility for researchers and practitioners with limited resources and specialized hardware.

---

Published as a conference paper at ICLR 2025
Recently, various memory -efnnorr-effiientttt methods for fine-tuning LMs have been proposed. In-context learning (ICL) utilizes a single inference pass; incorporating label examples in its context for prediction (Brown et al. 2020). Despite its limited success, ICL's performance is shown to be less effective than (Adam) fine-tuning for medium-sized LMs (Brown et al. 2020). As an alternative approach, Carameter-Ffficiee Fine-Tuning (PEFT) selectively tunes fraction of the network while freezing the rest of the parameters, and significantly reduces the parameters needed for fine-tuning (Hu) al. 2022 Li and Liang, 202T; Lester et al. 202T Despite its efficiency, fine-tuning LMs with PEFT may still require more memory than model inference. For example, fine-tuning OPT-13B with Adam with a batch size of requires 4 KH00 GPUs (316GB total), whereas utilizing PEFT decreases this to 2xH100 GPUs (158GB total) with a batch size of 16 (Brown et al. 2020). Nonetheless, this requirement is still 6x greater than 25GB needed for model inference.
Another approach for memory-efficien fine-tuning is to lessen the memory footprint of the optimizer. Recently, the Memory-Ffficien Zeroth-order Optimizer (MeZO) is proposed by Malladi et al. (2023). MeZO generates gradient estimators solely through forward passes with minimal memory overhead. Unlike classical zeroth-order optimization method ZO-SGD (Spall 1992). MeZO allows in in-place perturbation of model parameters avoid storing the perturbation vector. One desirable property of McZO (and in general approaches that reduce the memory overhead of the optimizers) is that can be combined with other methods, such as PEFT. Moreover, Malladi al. (2023) showed that the memory footprint of MeZO can be 12x lower than Adam. While being memory efficient, MeZO suffers from 1) slow convergence rate compared to standard fine-tuming methods such as Adam; and 2) possible degradation of the performance (e.g. accuracy) of the fine-tunel model compared to Adam (see the experiments in Malladi al. (2033 and Table 12p Observing these drawbacks, is natural to ask:
Question: Can we develop an optimizer for fine-tuning language models (LMs) that requires significantly less memory than standard Adam but still enjoys fast convergence and produces high-quality fine-tuned models?
In an effort to answer this question, this work proposes Addax (ADDition of grAdient estimates through memory-effiient eXecctiion) a method that: i) is memory efficient, ii) has fast convergence speed and iil) achieves the best performance across wide range of fine-tuning methods and tasks. Our specific contributions are as follows:
1. Algorithm design. We develop Addax, a novel approach that cleverly assigns different batches of data to McZO or in-place SGD (IP-SGD), and combines the computed gradients. This strategic assignment of data points, based on input length, allows Addax to maintain memory footprint comparable to MeZO while achieving significant improvements. Specifically, Addax accelerates MeZO's convergence rate and boosts the final model's performance, effectively overcoming the limitations of zeroth-order optimization. Our design is driven by two novel observations: 1) computing gradients for different data points requires varying memory, and 2) integrating zeroth-order updates with first-order methods enhances the quality of the fine-tuned model.
2. Theoretical analysis. Theoretically, analyze Addax's convergence under mild assumptions in two scenarios: with and without data assignment procedure. Unlike McZO, the convergence rate of Addax independent of model size without requiring the assumption of low-rankness of the Hessian. Additionally, we show that the hyperparameters for Addax are less restrictive compared to those of MeZO.
3. Numerical comparisons. We perform comprehensive experiments over a broad range of model architectures (c.g., masked LM and autoregressive LM), model scales ranging from 350M to 70B parameters, and tasks including classification, multiple-choice questions, and content generation. Compared to SGD and IP-SGD, Addax has a lower memory footprint and is on par with MeZO. Using single A100 (40GB) GPU, Addax successfully fine-tunes the OPT-13B model on all nine tasks, while SGD fails on all and IP-SGD fails on three due to memory limitations (see Figure I]. Addax surpasses MeZO by 14% in accuracy/F1 and converges 15x faster average, with similar memory footprint. Furthermore, Addax outperforms Adam in seven out of nine tasks while reduces up to 89% the memory. When fine-tume fine-tuning larger model like OPT-30B with single H100 (80GB) GPU, Addax achieves superior performance over SPD and MeZO on various tasks, e.g. test accuracy 16% higher than MeZO, 4.8% higher than IP-SGD when IP-SGD applicable
2

---

Published as a conference paper at ICLR 2025
(see Figure 2 and T. Similar results observed when fine-tuning Addax OPT-66B and Llama-2- 2-70B using three H100 (240GB total) GPUs (See Table 2 and
Next, we will discuss preliminaries, and will leave further discussions on prior work to Appendix
NOTATIONS AND PRELIMINARIES
2.1 NOTATIONS
We are interested in optimizing the (smooth and possibly non-convex) loss function
min BER
[(0) = Ezwp [0(0:z))]
(1)
parameterized by 0 E Rd, where D denotes the (fine-tuning) data distribution and ær denotes the (fine-tuning) data point. Throughout, mark the values related zeroth- and first-order gradient with (.),, (.) respectively, and denote the iteration and coordinate indices as (-),, (-)s, where e (0. ................................ T (l.....d}.
addition, we assume parameterized an M-layer network, i.e. where 6m E Rd and Σm= dm = d. The maximum sequence length in the dataset D is denoted by Lmaa- We slightly abuse the notation and denote the loss evaluated minibatch B L(0; B) & ErEB ((0;z))
= (0
2.2 MEMORY-
SARRTT--ORDER OPTIMIZER
MeZO Malladi et al. 2023) is a memory- efficien fine-tuning approach based on the zeroth-order
stochastic gradient descent method (ZO-SGD). The update step of MeZO follows SGD, i.c. 0
0, V (0.; ) with the key difference of using zeroth-order gradient V (0;; BT) estimated on
minibatch B%, instead of first-order gradient V B,) The zeroth-order gradient is estimated using Simultancous Perturbation Stochastic Approximation (SPSA, (Spall 1992))
(( C(e + ez; B) L(O c2;
(2)
2x
where z is a random search direction, c.g., z ~ (N. I). In ZO-SGD, SPSA requires generating and storing z IRd for multiple times of model perturbation, 0 ez, ez, which is memory inefficient. MeZO reduces memory consumption in the implementation of SPAS by only storing the random seed that generates z, reducing memory consumption from (d) to O(1).
However, due to the inherent bias and higher noise level in zeroth- order gradient estimate (Nesterov and Spokoiny 2017), MeZO suffers from slower convergence and worse final model performance compared to first-order methods, c.f. Malladi et al. (2023 Table 1), and the experiments in section g
2.3 PRELIMINARIES AND MAJOR OBSERVATIONS
In this subsection, provide the key observations behind the development of Addax:
SGD can match the performance of Adam in fine-tuning tasks. Although Adam and AdamW (Kingma and Ba 2015 Loshchilov and Hutter 2019) have demonstrated better performance than SGD in training deep learning models from scratch, they require additional memory to store optimizer states, which unfavorable in fine-tuning LMs. On the other hand, it has been observed that SGD achieves comparable performance to Adam in fine-tuning tasks (see, e.g., ZZang et al. (2024a) and that fine-tuning LMs using (16-bit) SGD can have comparable performance to fin-ttnnig with (32-bit) Adam. For example,
al.) (2023). ). This observation also re-confirmed our experiments showing
the right panel of Figure 3 in three fine-tuning tasks of RTE, CB, and COPA (Wang et al. 2019). SGD leads to higher performance while requiring significantly less memory. The success of using SGD in fine-tuning LMs attributed to the relative "nice" landscape of fine-tuning tasks of LMs (Hao et al. 2019 Zhang et al. 2024a).
Memory-efficient implementation of SGD via "in-place" updates. To reduce the memory footprint of SGD, several studies have explored the utilization of in-place (IP) updates during backward propagation (Zhao et al.] 2024] [v et al.) (202)). Instead of separating the backward propagation and weight update steps, which require storing the gradients for all layers, in- in-place SGD (IP-SGD)
3

---

Published as a conference paper at ICLR 2025
100
80
Ausmany:
40
SST-2
RIE
CB
BoolQ
WSC
WIC
MuluRC
COPA
SQuAD
330
0PU Nenr 1095
3000
250
100
50
SST-2
RTE
CB
BaolQ
WSC
WIC
MultiCC
COPA
SQcAD
10
Thoe salo)
10'
10"
SST-2
RTR
BoolQ
WSC
WIC
MulsiRC
COPA
SQaAD
Zert-ehor
Mr20
9GDb
IP-SGD
Addae
Adam
Figure 1: Accuracy/F-1 score, memory, and convergence time resulted from fine-tuning OPT-13B model with various algorithms on one A100 (40GB) GPU, except for Adam, which runs five GPUs. The label during fine-tuning
"OOM" the run encounters aut-of-memory Addax consistently outperforms other methods terms of Accuracy, with GPU memory consumption comparable to MeZO. Except for Adam, all other methods are running in 16-bit mode. We do not report the time for Adam as it requires five GPUs. The exact numbers can be found in Table 12 Appendix F.I
with smallest batch size.
100
(80
80
40
SST-22
RTE
BlalQ
WSC
WIC
MultRCC
SQuAD
Mrnnot (63)
70
50
158-122
RTE
Bool()
WSC
WIC
MultiRC
SQUAD
10'
109
SST-2
RTE
BooiQ
WSC
WTC
MuluRC
SQUAD
Addax
Figure 2: Accuracy/F-1 score, memory, and convergence time resulted from fine-tuning OPT-30B model with various algorithms on one H100 (80GB) GPU. The label "OOM" means the run encounters out- of of-memory error during fine-tuning. Addax Icads to the hest final accuracy in all experiments, have comparable memory fooiprint to MeZO, while converges orders of magnitude faster. The exact numbers related to this figure be found in Table 15 Appendix 122

---

Published as a conference paper at ICLR 2025
40
IP-SGD Mr20
90
35
8G
Mennry (GB)
TASK COPA
Arcuracy
SGD
30
80
Jula
ET SIGT
75
25
113
2
6
10
14
18
22
26
30
50
100
150
200
250
300
Batch Size
GPU Metnory (GB)
Figure 3: Left: Memory profie of fine-tuning OPT-13B with IP-SGD and Mc70 on synthetic dataset with fixed sequence length of 300. Right: Fine-tuning OPT-13B using IP-SGD and small batch sizes (BS) can outperform Adam, while consuming significantly lower memory.
combines the two steps by updating the weights in each layer as soon as the gradients are computed, and immediately discards the gradient after it has been used. In-place update avoids storing the gradients of the full model and, thus, significantly reduces the memory footprint of SGD.
Though improving over SGD, IP-SGD may still require more memory than MeZO. To compare the memory consumption of IP-SGD and Me70, we record the memory consumption the batch size in the left panel of Figure As illustrated in the figure, with a memory constraint of batch size of 18 for running McZO, while only use the batch size of 2 for
function of
30GB, running IP-SGD.
Memory required for computing gradient de- pends on the input sequence length. We ob- serve that in fine-tuning LM models, the memory required for first-order gradient estimation de- pends the maximum input sequence length (in tokens). In Figure 4 we record the memory con- sumption of SGD, IP-SGD, and MeZO with fixed batch size and varying input sequence length. As illustrated in the figure, the memory consumption of all algorithms increases as the input sequence length increases, and the memory increase of IP- SGD is much faster than MeZO. Given the fact that the input sequence lengths vary a lot a given fine-tuning task across different data points (see F Appendix D). we conclude that the memory inefficiency of IP-SGD compared with MeZO mainly attributes to the data points with longer sequence length in the dataset.
70
SGD
IP-SGD
60
McZO
Memory (CB)
50
40
30
100
200
300
400
500
Sequence Length
Figure 4: Memory profiling of SGD, IP-SGD, and MeZO OPT-13B fine-tuning with syuthetic datasets with varying sequence lengths (fixing batch size 8).
With the above observations in hand, discuss algorithm, Addax, in the next section.
ADDAX
Building on the insights from Section we propose Addax, method that assigns data batches to either MeZO or IP-SGD based memory needs, and then combines their gradients. In memory- limited scenarios, Addax reduces memory usage by assigning shorter sequences to IP-SGD and longer ones to MeZO, leading similar convergence rate IP-SGD, much faster than MeZO, while keeping similar memory footprint, significantly less than IP-SGD. Furthermore, incorporating zeroth-order gradients with first-order gradients enhance the final model performance. We first propose the Addax algorithm, provide its convergence properties, and finally discuss how incorporating the zeroth-order gradient benefits SGD in the final model performance.
5

---

Published as a conference paper at ICLR 2025
3.1 ALGORITHM OVERVIEW
Addax first starts by partitioning the dataset into data points with long/short sequence length, i.e., D = D U 'D' where D = {" E D length(#) > | and D (z E D length(2)) S Lr} for some given threshold hyper-parameter LT, where length(-) measures the sequence length of data point. As discussed in Section 2.3 computing gradient on data points in D requires much less memory than computing the gradients in D because the maximum input sequence length of 'D' is capped by L1 which is smaller than Immax- With this data partition, Addax computes the zeroth-order gradient on the data points with long sequence length in D while computing the first-order gradients for the data points with short sequence length in D', resulting in a considerable memory reduction compared with IP-SGD. In particular, each iteration of the algorithm, Addax first draws a random batch 5 (with | = K") of data points from D and random search direction z E I with Then, using the drawn samples, it utilizes SPSA to estimate the directional
E|z] 0 and E Jzzz" derivative of the objective function in the direction z at the point 0 with a small perturbation size e:
( + cz; xx) E(O cz:x)
KJ
2
DEB"
Then, Addax draws random batch B (with |B'| = K' ) of data from D and computes g' 1 KT ExEs: vf(0;z). Finally, updates the model parameters by:
0 x- 0 -7 (azg) + -01-00 +(1-a)g))
where 1 is the learning rate and a E [0, 1] mixing constant for combining the two gradient estimates. If one naīvely implements the update rule in equation 3 it needs to store the gradient g' and the direction z. However, as discussed in Section 2 one can perform the update rule in equation in-place, without storing the values of g² or Such memory-efficiee
implementation of Addax We leave further detailed discussions on this algorithm to Appendix A
described in algorithm [
Addax without data assignment (Addax- AA). When sufficient memory is available to perform IP- SGD, one might argue that there is no need to use zeroth-order gradients. However, as demonstrated in our extensive experiments (see Figure 5 and Fine-tuning results for datasets SST-2, RTE, WSC of Addax (Lr = 320) in Table incorporating zeroth-order gradients still results in better final accuracy for the fine-tuned model. Therefore, even when memory constraints are not a concern, we continue to utilize zeroth-order gradients by randomly selecting a data batch and combining its order gradient with IP-SGD. This is achieved by setting D t- D and DJ +- D in step 3 of Algorithm This version of the algorithm is referred to as Addax- WA.
Having introduced Addax, we would like to comment on a recent related work. Concurrently with our work, /Zhang et al. (2024b) propose "hybrid ZO-FO" fine-tuning scheme for LLMss, which also integrates zeroth- and first-order gradient estimates. However, their method differs significantly from ours. They limit backpropagation to the deeper layers and use zeroth-order optimization in the shallower layers. This approach prevents them from taking advantage of the memory savings offered by in-place update rules and does not harness the benefits of zeroth-order methods for improving the final model accuracy. In contrast, our approach utilizes in- in-place update rules, ensuring that memory usage does not scale significantly with model size and removing the need to limit backpropagation to specific layers. Furthermore, we allocate data to optimizers differently, leading to additional major memory savings.
3.2 THEORETICAL ANALYSIS
Depending on whether step 3 or step is executed in Algorithm [ the algorithm follows two distinct trajectories. We analyze each case separately. We begin by discussing the convergence result of Addax when step is applied, i.e., the Addax- WA algorithm. We provide only the informal version of the results, while the formal theorems can be found in Appendix G We start by presenting the convergence of Addax in the general nonconvex setting:
Theorem 3.1 (Informal), Assume that the loss E is Lipschitz smooth, and the first-order stochastic 1/2) and e
gradient is unbiased and has bounded variance. Choosing Tt 1
Old 11/44-1144 in Addax leads to the convergence rate:
E [[VCC(ON1))=
1
ad
=
K1
K
6

---

Published as a conference paper at ICLR 2025
Algorithm 1 Addax: ADDition of grAdient estimates through memory-efficient eXecution
1: Input: with M Layers, T, C, LT, Lmazy K K' perturbation scale mixing parameter a E [0, 11, and dataset D.
2: if LT 2 Liaaz then
D D D
4: else
5:
D +- (r D length(x)) [T). D t {a E (length(()) S IT) } | do
6: (
7:
Randomly draw mini-batchin 8o, B uniformly from Do D.D' with Ko, K samples.
8:
(g),s) ZerothGrad B" (algorithm 2
# Estimate zeroth-order gradient
9:
for m M....,dd
10:
Im t KJ E D1EB' Vo..
# Estimate layer first-order gradient # Update model parameters
11:
On 0n. (1 a)
12:
9m < None
# Clear gradients
Reset random number generator with seed
14:
for m 1,... M do
2m N0, Id...)
16:
6m t 6m mg 2m
Update model parameters
17:
2m f- None
18:
Output: 0
K Ko4dKı, we obtain the convergence rate O (V TIKOLaKI)
Further, by choosing the optimal
The formal version of this Theorem and its proof can be found in Appendix G.3
Remark 1. Compared with the convergence rate of zeroth-order methods Ghadimi and Lan (0013); Nesterov and Spokoiny (2017), the above convergence rate is nearly invension-iddependen. In
particular, the factor +ak") upper bounded by TKT.
The restrictions the choice of parameters for Addax more relaxed than the ones in zeroth-order methods. In particular, ZO-SGD (or MeZO) requires choosing smaller parameters (d--1T-1/2) and (11 (0(1/vdl /dT) for guaranteeing convergence. These choices are clearly more restrictive than the choice of parameters Theorem 3.1 Therefore, one can choose larger learning rate in Addax compared to MeZO. This is also observed in our experiments (see Appendix D.5 for details).
Remark 2.
Remark 3. Prior works observed that pre tre-nrained LMs have a low effective rank Hessian (Aghajanyan al., 2020 Li al.) 2018 Papyan| 2018 2020). Under the assumption that the Hessian has low effective rank, the dependency on the parameter dimension d can be further improved (see
Theorem G.12 in Appendix (
Besides assuming the loss is nonconvex and smooth, we also provide the convergence of Addax for strongly convex and smooth loss functions:
Theorem 3.2 (Informal). Assume that C is strongly convex, per-sample loss is Lipschitz smooth, and the first-order stochastic gradients unbiased and have bounded variance. Choosing T = O (T)) and = O(a11220 1/2) in Addax leads to the convergence rate:
O[ - 0. 0-11==
In(T)) T
(
(1- a²d add)
Ko xO)
))
KI
Further, by choosing the optimal a
K Kot Addax converges with rate
In(T)d T(K*+aK))
The above results cover the case where step is executed in Algorithm I When step 5 executed, obtain similar results (see Theorem G.10 Appendix G.4 for details).
3.3 FURTHER DISCUSSIONS ON THE BENEFITS OF UTILIZING ZEROTH-ORDER GRADIENTS
Beyond enhancing memory efficiency IP-SGD by eliminating the need to compute gradients for input data with longer sequence lengths that more memory-intensive, using zeroth-order

---

Published as a conference paper at ICLR 2025
RIE BlooQ
(2
86
2(0)
Test Accuracy (56)
82
IP-SCD
80
-2
0
4
6
8 K
10
12
14
16
-2
0
2
4
Figure 5: Left: An illustration of loss function C()) alongside its Gaussian smoothed version (()). Minimizing
(1 ()) + a() can help escape sharp local minima and find higher quality solutions. Right: The regularization effect of zeroth-order gradient estimates on first-order gradient estimates. We fix K in Addax across experiments while varying K from to the special case where K Addax reduces to IP-SGD.
updates alongside first-order updates improve the final performance and accuracy of the fine-tuned model compared with vanilla IP-SGD. We hypothesize the following reasons for this improvement in performance:
Zeroth-order updates may help escape spurious and sharp local minima. It is known that zeroth-order gradient estimates are noisy estimates the gradient (see Nesterov and Spokoiny (2017) and Lemma G.5 in Appendix G). It has also been observed that injecting noise into the gradient can be beneficial in nonconvex optimization. For example, Ge et al. (2015) and JJin et al.) (2017) showed that adding noise the gradient direction can help escape saddle points in nonconvex optimization. Moreover, [Zhou et al. (2019) showed (both experimentally and theoretically) that injecting noise in help the algorithm in escaping bad/spurious local minima. One also
the gradient direction argue that noise would help the algorithm in finding flat minima and avoid sharp local minima Liu al. ()2IT. l,. All these insights suggest that Addax can converge to better local minima. Our extensive experiments in Section 4 also show that Addax outperforms SGD in terms of the final performance of the fine-tuned model.
Zeroth-order updates act as regularizer. To simplify the presentation, consider the case where step 3 is executed Algorithm T Recall that the zeroth-order gradient an unbiased estimator of smoothed version of the actual loss. That is, E[gg = V Z(O), where 2(0) A E, [C + ez) the Gaussian smoothed version of the original loss function L(0) (see Nesterov and Spokoiny (2017 Section 1) and Nemirovsky et al. (1983) Section 9.3)). Thus, the Addax update rule in equation aims solve the optimization problem
min a1-a) aZ(().
Such regularization, illustrated in Figure help escape sharp local minima and find higher- quality solutions. When Step executed, then this regularization is only done for some of the data points, but still be effective. This could be the reason that Addax finds higher-quality models in all our experiments.
EXPERIMENTS
Experiment Settings. We conduct fine-tuning experiments on five different models: the masked LM RoBERTaa-lare of (L.uu et al. 2019) (350M), the OPT-13B, OPT-30B, OPT-66B (Zhang al.) 2022) and Llama-2-70B models (Touvron et al.) (023). We also explore the impact of parer--rramtterss and the batch size on Addax's performance, detailed in Appendix E.2 Further details on the datasets, prompts used, and implementation be found in Appendix
The code for experiments
available at httth:::///// http://gttttttt..eeg t www.:/t.i.mmmmmmicmm...... ott.::aii.iiiiiiooooooooaaaaaaaa Addax
Observations on OPT-13B experiments Following Malladi et al. 2023), we fine-tune the OPT-13B model using single A100 GPU with Addax, Me7O, IP-SGD, and SGD. The results presented
8

---

in Figure I and Table 12 In this configuration, we aim to select the largest possible batch sizes for McZO, IP-SGD, and SGD and optimize (K"; K') and LT for Addax using one GPU. Averaging over nine tasks, Addax outperforms MeZ0 in this configuration by 14% in accuracy/FI score and Published as a conference paper at ICLR 2025
converges 15x faster. To highlight the differences in convergence speed, we plot the convergence curves of Addax-WA and Me70 using the same batch size in Figure [1]
Remarkably, while Addax can successfully run on all datasets, SGD encounters out-of-memory errors on all nine tasks and IP-SGD fails in three out of nine tasks. For the six datasets where IP-SGD is able to fine-tune, Addax achieves an average accuracy/F1 score of 81.7 compared to 80.3 for IP-SGD. Finally, Addax even outperforms Adam in seven out of nine tasks while reduces up to 89% the memory. We also conducted additional experiments with Addax and IP-SGD, incorporating gradient checkpointing technique as described by Chen et al. (016). Please refer to Appendix for details.
Observations on OPT-30B, OPT-66B and Llama-2-70B experiments. We also perform experi- ments on fine-tuning the larger-size models OPT-30B, OPT-66B and Llama-2-700 In particular, we fine-tune OPT-30B model using one H100 GPU and report the results in Figure Table To summarize the results, we report the averaged performance metrics (time, accuracy, and memory) in Table Our experiment on OPT-30B shows that while Addax has comparable memory consumption to McZO, on average, Addax outperforms McZO by more than 16% in terms of final accuracy and converges x0x faster. Moreover, SGD and IP-SGD failed in three out of seven tasks due to out-of-memory, while Addax can run on all tasks. In terms of the final accuracy of the fine-tuned model, Addax outperforms Me7O, SGD, and IP-SGD in all experiments in OPT-30B.
II
Similar observations are made in fine-tuning ОPT-66B and 2-70B models on three H100 GPUs: SGD and IP-SGD fail in some all tasks due out-of-memory, while Addax can efficiently fine-tune on all tasks. Addax outperforms other methods in six of seven tasks for the OPT-66B model and all six tasks for the Llama-2-70B model. Moreover, Addax outperforms McZO in terms of final accuracy while being orders of magnitude faster. The results are summarized in Table 2 and Table 3 repectively.
Observations on RoBERTa-large experiments. Beside large autoregressive language model, we also perform experiments on smaller language model: RoBERTa-large with 350M parameters. The results can be found in Figure 7 As can be seen in this figure, 16/32-bit Addax outperforms zero-shot and MeZO across six different tasks and surpasses Adam in four out of six tasks. We also investigate
K' the two important hyper-parametten parameters, a and KOTKT on Addax's 's performance. The results can be found in Figure g and Figure | As shown in the top row of the heatmaps of the two figures, it is observed that an increase in the ratio KT-KT correlates with improved accuracy across tasks for both 16-bit and 32-bit Addax configurations. We did not identify a consistent trend for a across different tasks for both 16-bit and 32-bit Addax, suggesting that the optimal a requires tuning and could be task-specific.
K''
Zeroth-order gradient estimates improve model performance even when K' is small. We report the detailed choice of the batch size for different algorithms in Table 12 Notably, fine-tuning OPT- 13B using Addax with a smaller first-order batch size K surpasses the performance of SGD with larger batch sizes. For example, Addax achieves an accuracy of 68.3 on the WIC task with KJ = 4, while IP-SGD achieves 66.0 with a batch size of 12. This suggests that the zeroth-order gradient estimate in Addax provides stability (and regularization of the gradient) even when K' is small and can effectively reduce memory usage. Additional experimental results are given in Appendix E
It is important to note that reported memory usage should be interpreted with caution, as it depends on the selected batch size. Reducing the batch size can lower memory usage, but it may come at the cost of accuracy and convergence speed. In tables, when entry is marked with x, it means that even with the smallest batch size in our grid (i.e. BS = 2), the algorithm still results in out-of-memory.
Hyperparameter tuning. For MeZO, SGD, and SGD, we select the largest possible batch size from the hyperparameter search grid that maximizes GPU memory usage without causing out-of- memory. For Addax, we choose K K1, and LT values that optimize GPU usage during fine-tuning on the MultiRC dataset, as it is the task with the longest sequence length. Successful fine-tuning on MultiRC implies Addax can handle smaller tasks as well. For detailed procedures, please refer to Appendix D.6.

---

Published as a conference paper at ICLR 2025
Table 1: Summary of OPT-30B fine-tuning results on one H100 GPU (80GB): The METRIC, representing the average performance across short datasets (SST-2, RTE, WSC, WIC) and long datasets (BoolQ, MultiRC, SQuAD), Lmaz is the maximum sequence length within dataset. Detailed results are in Table IS
Short Datasets (Lmar S 260))
Long Datasets (Lmax > 260)
Physical Memory
Wall-clock time to Accuarey/F1(9) (uninn)
Physical Memory
Wall-clock time to Accuarcy/F1 the best validation
Method
the best validation
(Fine-Tuning)
MeZO
66GB
655.7min
69.3
78GB
776.0min
68.7
SGD
IP-SGD
70GB
3.0min
75.1
Addax
68GB
14,5min
78.7
77GB
28. 28.5min
82.0
Table 2: Summary of OPT-66B fine-tuning results on three H100 GPUs (240GB total); The METRIC, representing the average performance across short datasets (SST-2, RTE, BoolQ. WSC, WIC, SQuAD) and long dataset (MultiRC). Lmaz the maximum sequence length within a dataset. Detailed results in Table 16
Short Datasets (Lmar I 420))
Long Dataset (Lmar 420)
Physical Memory
Wall-clock time to the best validation
Accuarcy/F1(94 (Fine Tuning)
Physical Memory
Wall-clock time to the best validation
Accuarcy/F1 (%)
Method
(Fine (Fine-Tuning Tuning)
McZO
170GB
511.5min
72.4
197GB
379.6min
61.1
SGD
IP-SGD
170GB
3.8min
77.1
Addax
173GB
21. 11.9min
80.6
215GB
76. 9mii
80.6
Table 3: Summary of Llama-2-70B fine-tuning results on three H100 GPUs (240GB total): The nota- tion METRIC represents the average performance across short datasets (RTE, WSC, WIC) and long datascts (BoolQ, MultiRC, SQuAD). Linas is the maximum sequence length within dataset. Detailed results in
Table
Short Datasets (Lmaz S 260))
Long Dataset (Loar > 260)
Physical Memury
Wall-clock time to the best validation
Accuarcy/FII (Fine- -uninnu
Physical Memory
Wall-clock time to the best validation
Accuarcy/Fi I (%)
Method
(Fine-Tuning)
McZO
149GB
4609min
61.1
186GB
792.3min
73.3 *
SGD
x
*
*
IP-SGD
189.5GB
6.2min
78.2
*
*
Addax
190.
21.2min
80.1
218GB
37.2min
88.9
5
CONCLUSION, BROADER IMPACT, AND LIMITATIONS
This paper introduces Addax, enemory-efficient fine-tuning method for Language Models (LMs). By leveraging both first- and zeroth-order stochastic gradient estimates, Addax demonstrates improved memory efficiency without sacrificing convergence speed model performance, as validated by our extensive experiments across various models, tasks, and datasets. Addax has the potential to impact language model fine-tuning tasks for researchers and machine learning practitioners with limited resources. With a convergence time comparable to first-order methods and memory usage similar to zeroth-order methods, Addax provides resource- -effrcceet approach to optimizing higher- quality fine-tuned models. Furthermore, Addax has proven effective with large-scale models, and its memory-efficiee nature can make large-scale fine-tuning more feasible by requiring fewer resources.
Limitation: Addax introduces an additional hyper-parameter parameter, which requires tuning for best performance. While the search grid for a is small and includes only five different values across all OPT experimees, this is still an additional burden for institutions with limited resources.
Future works: As future work, Addax may also have potential application in tasks other than fine-tuning. For example, Addax may be used for pre-training tasks or even combined with the Adam algorithm by passing the combined first- and zeroth-order gradients to Adam. From a theoretical perspective, conducting regularize the performance of first-order ones is of particular interest. We did not examine Addax combined with memory-efficien methods like PEFT (Hu et al. 2022 Li and Liang 2021 Lester et al. 202T), quantization Dettmers et al. (2022). or memory-efticient attention Dao et al. Zandieh et al. (2023): Han al. (2023). but hope to explore these in future work.
in-depth theoretical analysis of how zeroth-order gradient estimates
2022):
10

---

Published as a conference paper at ICLR 2025
REFERENCES
A. Agarwal, M. J. Wainwright, P. Bartlett, and P. Ravikumar. Information-theoree lower bounds on the oracle complexity of convex optimization. Advances in Neural Information Processing Systems, 22, 2009.
A. Aghajanyan, L. Zettlemoyer, and S. Gupta. Intrinsic dimensionality explains the effectiveness of language model fine-tuning. arXiv: 2012.13255, 2020.
S. H. Bach, V. Sanh, Z.-X. Yong, A. Webson, C. Raffel, N. V. Nayak, A. Sharma, T. Kim, M. S. Bari, T. Fevry, et al. Promptsource: An integrated development environment and repository for natural language prompts. arXiv preprint arXiv:2202.01229, 2022.
K. Balasubramanian and S. Ghadimi. Zeroth-ordes nonconvex stochastic optimization: Handling constraints, high dimensionality, and saddle points. Foundations of Computational Mathematics, 22(1):35-76, 2022.
R. Bar-Haim, I. Dagan, B. Dolan, L. Ferro, and D. Giampiccolo. The second pascal recognising textual entailment challenge. Proceedings of the Second PASCAL Challenges Workshop Recognising Textual Entailment, 01 2006.
L. Bentivogli, Clark, I. Dagan, and D. Giampiccolo. The fifth pascal recognizing textual entailment challenge. TAC, 7:8, 2009.
S. R. Bowman, G. Angeli, C. Potts, and C. D. Manning. A large annotated corpus for learning natural language inference. In Proceedings the 2015 Conference on Empirical Methods in Natural Language Processing, pages 632-642, Lisbon, Portugal, Sept. 2015. Association for Computational Linguistics. doi: 10.18653/11/D15-11 1075.
T. Brown, B. Mann, N. Ryder, M. Subbiah, J. D. Kaplan, P. Dhariwal, A. Neelakantan, P. Shyam, G. Sastry, A. Askell, S. Agarwal, A. Borbert-Vosss. G. Krueger, T. Henighan, R. Child, A. Ramesh, D. Ziegler, J. Wu, C. Winter, C. Hesse, M. Chen, E. Sigler, M. Litwin, S. Gray, B. Chess, J. Clark, C. Berner, S. McCandlish, A. Radford, I. Sutskever, and D. Amodei. Language models are few-shot learners. In H. Larochelle, M. Ranzato, R. Hadsell, M. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 1877-1901, 2020.
Chen, B. Xu, C. Zhang, and Guestrin. Training deep nets with sublinear memory cost, arXiv preprint arXiv:1604 0.44..61444 2016.
C. Clark, K. Lee, M.-W. Chang, T. Kwiatkowski, M. Collins, and K. Toutanova. Boolq: Exploring the surprising difficulty of natural yes/no questions. arXiv preprint arXiv: 1905. 10044, 2019.
I. Dagan, O. Glickman, and B. Magnini. The pascal recognising textual entailment challenge. In Proceedings of the First International Conference on Machine Learning Challenges: Evaluat- ing Predictive Uncertainty Visual Object Classification, and Recognizing Textual Entailment, MLCW 05, page 177-190, Berlin, Heidelberg, 2005. Springer- Springer-Verlag. ISBN 3540334270, doi: 10.1007/1 1736790 9. URL https: / 101010000000010111366666611
T. Dao, D. Fu, S. Ermon, A. Rudra, and C. Ré. Flashattention: Fast and memory-efficient exact attention with io-awareness. Advances in Neural Information Processing Systems, 35:16344-16359, 2022.
M.-C. De Marneffe, M. Simons, and Tonhauser. The commitmentbank: Investigating projection in naturally occurring discourse. In proceedings of Sinn und Bedeutung, volume 23, pages 107-124, 2019.
Dettmers, M. Lewis, Y. Belkada, and Zettlemoyer. Llm. int8 (): 8-bit matrix multiplication for transformens at scale. arXiv preprint arXiv:2208.07339, 2022.
T. Dettmers, A. Pagnoni, A. Holtzman, and L. Zettlemoyer. Qlora: Efficient finetuning of quantized lims. arXiv preprint arXiv:2305. /4311 2023.
11

---

Published as a conference paper at ICLR 2025
J. Devlin, M.-W. Chang, K. Lee, and K. Toutanova. BERT: Pre-training of deep bidirectional trans- formers for language understanding. Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguisties: Human Language Technologies (Long and Short Papers),
2019.
J. Duchi, E. Hazan, and Y. Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of machine learning research, 12(7), 2011.
i. Fairscale. Fairscale: A general purpose modular pytorch library for high performance and large scale training, 2021.
T. Gao, A. Fisch, and D. Chen. Making pre-trained language models better few-shot learners. arXiv preprint ariv:22012.15728, 2020.
X. Gao, B. Jiang, and S. Zhang, On the anformation-adaatiii variants of the admm: an iteration complexity perspective. Journal of Scientific Computing, 76:327-363, 2018.
R. Ge, F. Huang, C. Jin, and Y. Yuan. Escaping from saddle points online stochastic gradient for tensor decomposition. In Conference on learning theory, pages 797-842. PMLR, 2015.
S. Ghadimi and G. Lan. Stochastic first-and zeroth-order methods for nonconvex stochastic program- ming. SIAM journal on oplimization, (3(4))2341-2368 2013.
D. Giampiccolo, B. Magnini, I. Dagan, and B. Dolan. The third pascal recognizing textual entailment challenge. In Proceedings of the ACL-PASCAL Workshop on Textual Entailment and Paraphrasing, RTE '07, page 1-9, USA, 2007. Association for Computational Linguisties.
W. Guo, J. Long, Y. Zeng, Z. Liu, X. Yang, Y. Ran, J. R. Gardner, O. Bastani, C. De Sa, X. Yu, et al. Zeroth-order fine-tuning of lims with extreme sparsity. arXiv preprint arXiv:2406.02 02913, 2024.
S. Gururangan, A. Marasović, S. Swayamdipta, K. Lo, I. Beltagy, D. Downey, and N. A. Smith. Don't stop pretraining: Adapt language models to domains and tasks. Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pages
2020.
I. Han. R. Jayaram, A. Karbasi, V. Mirrokni, D. P. Woodruff. and A. Zandich. Hyperattention: Long-context attention in near-linear time. arXiv preprint arXiv:2310.05880 2023.
Y. Hao, L. Dong, F. Wei, and K. Xu. Visualizing and understanding the effectiveness of bert. arXiv preprint arXiv: 1900: 056200 2019.
G. Hinton, N. Srivastava, and K. Swersky. Neural networks for machine learning lecture 6a overview of mini-batch gradient descent. Cited on, 14(8):2, 2012.
E. J. Hu, Y. Shen, P. Wallis, Z. Allen-Zhu, Y. Li, S. Wang, L. Wang, and W. Chen. Lora: Low-rank adaptation of large language models. In In International Conference on Learning Representations, 2022.
C. Jin, R. Ge, Netrapalli, S. M. Kakade, and M. I. Jordan. How escape saddle points efficiently. In International conference on machine learning, pages 1724 1732. PMLR, 2017.
D. Khashabi, S. Chaturvedi, M. Roth, S. Upadhyay, and D. Roth. Looking beyond the surface: A challenge set for reading comprehension over multiple sentences. In Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume I (Long Papers), pages
2018.
D. P. Kingma and J. Ba. Adam: A method for stochastic optimization. In International Conference on Learning Representations, 2015.
N. Landro, Gallo, and R. La Grassa. Mixing adam and sgd: A combined optimization method. arXiv preprint arXiv:201 10/22011.000 1.08042, 2020.
B. Lester, R. Al-Rfou, and N. Constant. The power of scale for parameter-efficien prompt tuning. In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, pages 3045-3059, 2021.
12

---

Published as a conference paper at ICLR 2025
H. Levesque, E. Davis, and L. Morgenstem. The winograd schema challenge. In Thirteenth international conference on the principles of knowledge representation and reasoning, 2012.
C. Li, H. Farkhoor, R. Liu, and J. Yosinski. Measuring the intrinsic dimension of objective landscapes, 2018.
X. L. Li and P. Liang. Prefix- Trefix-tuning: Optimizing continuous prompts for generation. In Proceed- ings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference Natural Language Processing (Volume I: Long Pupers), pages 3582-4597 2021.
T. Liu, Y. Li, S. Wei, E. Zhou, and T. Zhao. Noisy gradient descent converges to flat minima for nonconvex matrix factorization. International Conference on Artificial Intelligence and Statistics, pages 1891-1899. PMLR, 2021.
Y. Liu, M. Ott, Goyal, Du, M. Joshi, D. Chen, O. Levy, M. Lewis, L. Zettlemoyer, and V. Stoy- anov. Roberta: A robustly optimized bert pretraining approach. arXiv preprint arXiv: 1692. 2019.
I. Loshchilov and F. Hutter. Decoupled weight decay regularization. In International Conference on Learning Representations, 2019.
J. Lu, D. Batra, D. Parikh, and S. Lee. Vilbert: Pretraining task-agnostic visiolinguistic representations for rision-and-aanguage tasks. Advances in neural information processing systems, 32, 2019.
K. Lv, Y. Yang. T. Liu, Q. Gao, Q. Guo, and X. Qiu. Full parameter fine-tuning for large language models with limited resources. arXiv preprint urXiv:2306.09782, 2023.
Malladi, T. Gao, E. Nichani, A. Damian, J. D. Lee, D. Chen, and S. Arora. Fine-tuning language models with just forward passes. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.
A. Nemirovsky, D. Yudin, and E. DAWSON. Wiley-interscience series in discrete mathematics, 1983.
Y. Nesterov and V. Spokoiny. Random gradient- free minimization of convex functions. Foundations of Computational Mathematics, 17(2):527 -66, 2017.
OpenAl. Gpt-4 technical report. arXiv: 2303. 303.08774, 2023.
V. Papyan. The full spectrum of deepnet hessians at scale: Dynamics with sgd training and sample size. arXiv preprint arXiv:1871.07060 2018.
V. Papyan. Traces of class/cross-class structure pervade deep learning spectra. Journal of Machine Learning Research, 21(252):1-64, 2020.
A. Paszke, S. Gross, F. Massa, A. Lerer, J. Bradbury, G. Chanan, T. Killeen, Z. Lin, N. Gimelshein, L. Antiga, et al. Pytorch: An imperative style, nigh-performance deep learning library. Advances in neural information processing systems, 32, 2019.
M. T. Pilehvar and J. Camacho-Collados. Wie: the word-in-context dataset for evaluating context- sensitive meaning representations. arXiv preprint 1808.0888.09) 2018.
P. Rajpurkar, J. Zhang, K. Lopyrev, and P. Liang. Squad: 100,000+ questions for machine compre- hension of text. arXiv preprint arXiv: 05250, 2016.
M. Riedmiller and H. Braun. Rprop: fast adaptive learning algorithm. In Proc. of the Int. Symposium on Computer and Information Science VII. 1992.
H. Robbins and S. Monro. A stochastic approximation method. The annals of mathematical statistics, pages 400-407, 1951.
M. Roemmele, C. A. Bejan, and A. S. Gordon. Choice of plausible alternatives: An evaluation of commonsense causal reasoning. In 201 AAAI Spring Symposium Series, 2011.
13

---

Published as a conference paper at ICLR 2025
D. Schwartz, M. Toneva, and L. Wehbe. Inducing brain-relevant bias in natural language processing models. Advances in neural information processing systems, 32, 2019.
R. Socher, A. Perelygin, J. Wu, J. Chuang, C. D. Manning, A. Ng. and C. Potts. Recursive deep models for semantic compositionality over a sentiment treebank. In Proceedings of the 2013 Conference on Empirical Methods in Natural Language Processing, pages 1631-1644 Seattle, Washington, USA, Oct. 2013. Association for Computational Linguistics. URL ht.p:s::sss..........p/ //acl aatthhnnnthhoooooo logy brg/D13-1170
J. Spall. Multivariate stochastic approximation using a simultancous perturbation gradient approxi- mation. IEEE Transactions Automatic Control, 37(3):332-341 341, 1992. doi: 10.1109/9.1 11962..
H. Touvron, L. Martin, K. Stone, P. Albert, A. Almahairi, Y. Babaci, N. Bashlykov, S. Batra, P. Bhargava, S. Bhosale. al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint (.Xi:2337.00928800 2023.
A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and L. Polosukhin. Attention all you need. Advances in neural information processing systems, 30, 2017.
E. M. Voorhees and D. M. Tice. Building question answering test collection. In Proceedings of the 23rd anmual international ACM SIGIR conference on Research and development in information retrieval, pages 200-207, 2000.
A. Wang. Y. Pruksachatkun. N. Nangia, A. Singh, J. Michael, F. Hill, O. Levy, and S. Bowman. Supergluc: A stickier benchmark for general-purpose language understanding systems, Advances in neural information processing systems, 32, 2019.
Z. Wang. K. Balasubramanian, S. Ma, and M. Razaviyayn. Zeroth-order algorithms for nonconvex minimax problems with improved complexities. arXiv preprint arXiv:2001.07819 2020.
Z. Wang. K. Balasubramanian, S. Ma, and M. Razaviyayn. Zeroth-order algorithms for nonconvex- strongly-concave minimax problems with improved complexities. Journal of Global Optimization, 87(2):709-744 2023.
A. Williams, N. Nangia, and S. Bowman. A broad croer-coverage challenge corpus for sentence under- standing through inference. In Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume (Long Papers), pages 1112 1122, New Orleans, Louisiana, June 2018. Association for Computational Linguistics. doi: 00.18653/v1/118 101. URL wttpthtt:///mm http:///acaaattooooooooo org/oig 1101
T. Wolf, L. Debut, V. Sanh, J. Chaumond, C. Delangue, A. Moi, P. Cistac, T. Rault, R. Louf, M. Funtowicz, et al. Transformers: State-of-the the-th-e--nn art natural language processing. In Proceedings of the 2020 conference on empirical methods in natural language processing: system demonstrations, pages 38 38-45, 2020.
A. Zandieh, I. Han, M. Daliri, and A. Karbasi. Kdeformer: Accelerating transformers via kernel density estimation. arXiv preprini arXiv:2302.0245// 2023.
L. Zhang, K. K. Thekumparampil, S. Oh, and N. He. Dpzero: Dimension-indeen independent and differen- tially private zeroth-order optimization. In International Workshop on Federated Learning in the Age of Foundation Models Conjunction with NeurlPS 2023, 2023.
S. Zhang, S. Roller, N. Goyal, M. Artetxe, M. Chen, S. Chen, C. Dewan, M. Diab, X. Li, X. V. Lin,
T. Mihaylov, M. Ott, S. Shleifer, K. Shuster, D. Simig, P. S. Koura, A. Sridhar, T. Wang, and I. Zettlemoyer. Opt: Open pre-trained transformer language models. arXiv: 2205.01068, 2022.
Y. Zhang, C. Chen, T. Ding, Z. Li. R. Sun, and Z.-Q. Luo. Why transformers need adam: A hessian perspective. arXiv preprint arXiv:2402. 16788, 2024a.
Y. Zhang. P. Li, J. Hong, J. Li, Y. Zhang, W. Zheng, P.-Y. Chen, J. D. Lee, W. Yin. M. Hong, et al. Revisiting zeroth-order optimization for memory-efficien Ilm the-ttnnng: A benchmark. arXiv preprint arXiv:2402. 1,52, 2024b.
14

---

Published as a conference paper at ICLR 2025
J. Zhao, Z. Zhang, B. Chen, Z. Wang, A. Anandkumar, and Y. Tian. Galore: Memory-efficieer Ilm training by gradient low low-rank projection. arXiv preprint arXiv:2403.03507, 2024.
M. Zhou, T. Liu, Y. Li, D. Lin, E. Zhou, and T. Zhao. Toward understanding the importance of noise in training neural networks. In International Conference on Machine Learning, pages 7594-7602. PMLR, 2019.
A MORE DISCUSSION ON ADDAX
Algorithm [ Noutlines the steps of Addax. The process starts by determining whether the dataset requires partitioning based the sequence length threshold LT. There are two possible scenarios: If memory constraints are not concern, LT can exceed Lmax, the maximum sequence length in the dataset D. this case, Addax retains the dataset as whole, with D and D both being equivalent to 'D (Step() 3). However, when memory is limited and running IP-SGD on the entire dataset 'D is not feasible, Addax saves memory by partitioning the dataset according to I.T- assigns samples with sequence lengths shorter than L.T to D and the remaining samples to DD' (Step 5). This enables Addax to run in memory-constrained settings where IP-SGD would otherwise be infeasible.
In Step 8 the zeroth-order gradient estimator g" and random seed are obtained using the samples 8 with batch size Ky, which are drawn uniformly from the dataset D Similarly, Step 10 gets the layer / first-order gradients g. E && from samples B through backward propagation.
A major step in algorithm T is the computation of zeroth-order directional derivative go. done in 8 which is the subroutine call in Algorithm [ Algorithm is also used in MeZO where the directional derivative obtained through the classical zeroth-order gradient estimate SPSA; see equation equation To get the zeroth-order gradient estimates, Algorithm[]
the evaluation of the loss function L through two forward passes at points 0+ez and eez. The naïve implementation of the SPSA algorithm costs twice the memory of inference because of the need store the value of E I2 Algorithm removes this overhead by generating random seed and resetting the random number generator each time model parameters are perturbed (see Step 3 7 in Algorithm [ This approach guarantees that Algorithm 3 maintains a consistent direction for the random vector z across the two perturbations. Employing this in-place operation results in Algorithm consumption comparable to that of inference.
2 having memory
Steps 9 to 12 in Algorithm I are the main update steps based on the first-order gradient estimates for Addax. In Step 10 we back-propagate the first-order gradient estimates g for layer and update the model parameters in TT In Step 12 Addax frees up the calculated first-order gradient estimates 9.m This step is crucial because freeing up per-layer gradients ensures that memory requirements do not scale with the number of model parameters. More discussion on the related literature on this in-place update rule can be found in Appendix
Steps to 17 Algorithm describes the updates Addax based on zeroth-order gradient estimates. We use the same idea as in Malladi et al. (2023). where the seed is stored instead of the random vector 2. The random generator is reset before updating the components (see Step [3 in Algorithm [ described before. For each component O, in O where m ranges from to M, the process begins by generating a random direction Zm ~ (0, I in Step 15 Subsequently, each On is updated using zeroth-order gradient estimates. When iteration reaches T, Addax outputs the final model parameters 0.
In general, the in-place update operations have the same output as the update rule 0 < 0 ()(2gg (1 - agg²) ). This fine mme-grained control of dynamically allocated gradients ensures that Addax remains memory-effictee during fine-tuning.
B MORE DISCUSSION ON THE IN-PLACE UPDATES
In this section, we provide more detailed discussion on in-place updates. The technique of in- in-place gradient updates during backward propagation, as referenced in our approach, has been previously used in Zhao al. (2024); Lv al. (2023). In modern deep learning training frameworks, such as
15

---

Published as a conference paper at ICLR 2025
Algorithm 2 ZerothGrad (Malladi et al. 2023)
1: Input: parameters RR loss C RR + R, samples B, perturbation scale c.
2: Generate random seed
3: 0 t- PertubParammters((..
4: 4 + [(0; B)
5: 0 + PertubParameters((. -2e, s)
6: L- t L(0; B)
7: 0 + PertubParameters((.. c)
8: g < (4+ (2e)
9: Output: 9.8
Algorithm 3 PertubParameters
1: Input: parameters with M Layers, perturbation scale e, random seed
2: Reset random number generator with seed
3: for m 1, I................. MM do
4: 2m (0, Ia..)
5: Om f- On + 6:m
Update model parameters
6: fm <- None
7: Output: 0
PyTorch (Paszke al. 2019 they store the gradient tensor for computing optimizer states and update the model weights after all layers of gradients are computed. This approach is perfect for models with a small number of parameters; however, fine-tuning large model, like OPT-13B with 13 billion parameters, requires significant memory because the gradient tensor has the same size as the number model parameters. For example, as for OPT-13B model, each parameter needs 2 bytes bytes for gradient storage, totaling 26 GB or 52 GB of memory, respectively. Since Addax does not require any optimizer states, such memory overhead can be avoided by combining the computation of first-order gradient estimates and parameter updates into single step. As described in algorithm lines 9 12, we sequentially iterate over the Mh laver to the 1st layer, compute the gradient L.. (line 10). and perform in-place update On. (line 11.. Right after that, free the memory for gradient 9m (line 12). The loss computation and the update of zeroth-order gradient do remain the same as algorithm It is important to note that while the main update rule of mixing first-order and zeroth-order gradients remains unchanged, the implementation details significantly impact the memory consumption of fine-tuning tasks.
It is also worth noting that the in-place update rule has its own limitations. Firstly, it prevents the opti- mizer from using gradient accumulation, technique that scales batch size by accumulating gradients over several batches and only updating the optimizer after specified number of batches. Secondly, prevents gradient normalization, as the norm of the gradient must be known for normalization. In our paper, distinguish between SGD and SPDSGD: SGD uses gradient normalization during fine-tuning, while IP-SGD does not. This distinction leads to differences final performance and convergence time. For all experiments, except those using Adam, we do not employ the gradient accumulation technique.
C DISCUSSION ON RELATED WORKS
Stochastic First-order Optimizers in Deep Learning. SGD (Robbins and Monro 1951) has long been used in training deep neural networks due its convergence rate that is independent of the number model parameters. However, adaptive first-order optimizers have shown advantages over tuning, final model performance, and convergence speed. The concept of adaptive first-order optimizers dates back to RPROP (Riedmiller and Braun 1992). AdaGrad (Duchi al. 201T) adjusts the learning rate based on the estimated geometry, assigning higher rates less frequent features. RMSProp (Hinton et al. (012)) builds on RPROP, making it effective for small batch sizes. Adam (Kingma and Ba] 2015). inspired by AdaGrad and RMSProp. incorporates running average of gradients and has become the preferred method for training neural networks due to its fast
SGD in
https://pytorcch.org/
16

---

Published as a conference paper at ICLR 2025
convergence and reduced need for cyyper-paramedd tuning. Numerous studies have demonstrated Adam's success (Gururangan et al. 2020 Liu et al., 22199) Of (023), and researchers continue to investigate its efficacy with Transformer architectures (Vaswani al. 2017).
Zeroth-order optimization. Zeroth-order optimization has been extensively studied in convex, strongly convex, and non-convex settings in the optimization literature (Nesterov and Spokoiny
2017 Nemirovsky et al.] 1983, Wang et al.] 2023 Agarwal et al. 2009; Gao et al.) 2018] Wang
al.
2020). It is known that the convergence rate zeroth-order methods generally scales with the
number of parameters d. This property makes zeroth-order methods less effective for training deep
neural networks for which the number of parameters d can be very large. Recently, MeZO (Malladi et al. 2023) demonstrated that in language fine-tuning tasks, ZO-SGD can perform comparably to first-order methods. By using in-place perturbation, Me70 applies ZO-SGD memory- manner, keeping memory usage comparable to inference. The success of fine-tuning with zeroth-order methods may be due the fact that LM fine-tuning can occur in a very oow-iimeniiona subspace (Aghajanyan et al. 2020 Li et al. 2018). However, MeZO suffers from significantly slow convergence speed and slightly worse performance compared to first-order optimizers. Balasubramanian and Ghadimi (2022) estimate the Hessian to perform ZO optimization along important directions. Guo et al. (2024) focus on fine-tuning minimal subset of LLM parameters using zeroth-order methods, incorporating sparsity and quantization to overcome memory limitations.
Mixing update directions from different optimizers. There are recent studies mixing the update directions of different optimizers to enhance the performance of trainingffin-tuningg MAS (Landro et al. 2020) integrates SGD and Adam by assigning constant weights to balance the contributions of gradient estimates from both optimizers. Concurrent our work. Zhang al. (2024b) explores integrating zeroth- and first-order gradient estimates through a "hybrid ZO-FO fine-tuning scheme for LLMs". However, their method is completely different than ours. In particular, they restrict backpropagation to the deeper layers of the model, using zeroth-order optimization for the shallower layers to update parameters. Although this technique enhances memory efficiency, it neglects the possibility of using in-place first-order updates, leading to significant memory usage. In contrast, our approach employs in-place first-order gradient estimates, ensuring that memory requirements do not scale with the number of model parameters, eliminating the need to restrict backpropagation to specific layers. As we discussed in the main body, the memory demands of fine-tuning LMs with small first-order batch sizes are comparable those using zeroth-order batch sizes for many fine-tuning tasks. By utilizing both gradient estimates, our proposed method, Addax, not only proves be more memory-efficient but also surpasses competing methods in terms of the final accuracy of the fine-tuned model.
D EXPERIMENT SETUP
The code is available at httpp://tttthoo.oeeecccietttoror ence/ Addax
D.1 DATASETS
Our setup mainly follows the experiments in (Malladi et al. 2023). We employ the same datasets utilized in Malladi et al. (2023). Unless otherwise noted, we apply the same data processing procedures and settings for both validation and training.
For the RoBERTa-larte large model, we utilize the following datasets: SST-2 Socher et al. 2013). SST- (Socher al. 2013), SNLI (Bowman et al., 2013). MNLI (Williams et al. 2018). RTE (Dgan et al.) 2005) Hai- et al. 2006 Giampiccolo et al. 2007 Bentivogli et al.] 2009). and TREC (Voorhees) and Tice| (000). The test set is limited to 1,000 examples for both training and validation purposes. In our few-shot learning experiments, set A 16, where k represents the number of examples per class for training and validation.
For the OPT experiments,
employ the SuperGLUE dataset /Wang al. 2019 comprising BoolQ (Clark et al. 2019), CB (De Marneffe et al. 2019). COPA (Roemmele et al. 2011), Mul-
tiRC (Khashabi et al. 2018). RTE (Dagan et al. 2005 Bar HarrHaim et al. 2006 Giampiccolo et al. 2018), and WSC (Levesque) al. 2012). Following the approach of Malladi et al. 2023), we also include SST-2 (Socher et al.
2007 Bentivogli et al. 2009). WIC (Pilehva and Gumacho-Colladee
17

---

Published as a conference paper at ICLR 2025
2013) for development purposes, along with one question answering (QA) datasets: SQuAD Ra jpurkar et al. (016)). For each dataset, we randomly select 1, 000 examples for training, 500 examples for validation, and 1, 000 examples for testing.
D.2 DATASETS OVERVIEW
In this section, provide brief overview of the datasets used in our experiments. For both training and fine-tuning with transformers, peak memory usage is usually determined by the dataset's longest sequence length Lmaz. Samples shorter than Lmaz are padded this length. Figure 6 shows the histograms of sequence lengths for six different datasets (SST-2, RTE, WSC, WIC, MultiRC) tokenized by the OPT-13B tokenizer. The distributions of these data points are mostly right-skewed normal distributions, meaning that there are relatively a small number of samples with long sequence lengths in each dataset.
D.3 PROMPTS
To ensure a fair comparison, we employ the same prompts as those used by Malladi et al. (2023). which were initially adapted from Gao et al. (2020), GPT-3 (Brown al. 2020), and PromptSource (Bach) et al., 2022). Table 4 presents the prompts employed in our RoBERTa-large experiments, while Table details the prompts utilized for the OPT experiments.
Table 4: The prompts for each dataset used in our RoBERTa-large experiments. These prompts are identical to those used by Malladi et al. (2023). There are three different task types: sentiment classification (sentiment cls.), topic classification (topic cls.) and natural language inference (NLI). C' is the number of classes for each dataset. The label words can be filled in the [MASK] token of the prompt template. <S1> and <S2> are the first and second (if any) input sentences.
Dataset
C
Type
Prompt
Label words
SST-2
2
sentiment cls.
<S1> It was [MASK].
(great, terrible] [great, good, okay, bad, terrible) [Description, Expression, Entity,
SST-5
5
sentiment cls.
<S1> It was [MASK].
TREC
6
topic cls.
[MASK] <S1>
Human, Location, Number) (Yes, Maybe, No) Yes, Maybe, No)
MNLI
3
NLI
<S1> [MASK], <S2>
SNLI
3
NLI
<SI>>
? [MASK]. <S2>
RTE
2
NLI
<SI>>
[MASK], <S2>
(Yes, No)
D.4
IMPLEMENTATION
For the RoBERTa-large experiments, we run Addax in two separate computational precision settings: one using 16-bit floating point caleulations (FP16), referred to as 16-bit Addax, and the other using 32-bit foating-poing calculations (FP32), denoted as 32-bit Addax for clarity. For fine-tuning RoBERTa-Iarge For all RoBERTa-large experiments, MeZO and Adam loaded in the FP32 setting. Since fine-tuning RoBERTa-land models does not require significant memory. we set Lt sufficiently large for running Addax, ensuring that D and D are equivalent to the total dataset D. For all RoBERTa-large experiments, we use V100 GPUs for fine-tuning.
For the OPT and Llama experiments, Addax is used solely in the FP16 setting. We run SGD and IP-SGD in the FP16 setting and Adam in the FP32 setting. Unless otherwise noted, Addax, SGD, IP- SGD, and MeZO are trained the FP16 setting, while Adam uses the FP32 setting. If the maximum sequence length within the dataset Jamaz is less than Lr. Addax does not partition the dataset and DD and DJ are equivalent the total dataset D. In the scenarios where IP-SGD fails, Addax further reduces the memory usage by assigning different batches of data to ZO-SGD and FO-SGD based on LT. Specifically, Addax partitions the training data based on the length of data points, given threshold Le. The data is divided into D and D where D {a D length(x:) Lr) and D = { E D length(x) S }.
We do not employ advanced quantization techniques such as LLM. int 8 () (Dettmers al. 2022) and QLoRA (Dettmers et al.] 2023), nor do we integrate Addax with Serameter-Efficiine Fine- Tuning methods (PEFT) (Hu a 2022 Li and Liang 2021 Lester et al. 2021). For model inference,
18

---

Published as a conference paper at ICLR 2025
200
200
150
Frequency
150
Frequency
100
100
50
50
0
20
40
60
100
200
Sequence Length
Sequence Length
(a) Task: SST2
Task: RTE
15
150
Frequency
10
Frequency
100
50
0
40
60
80
40
60
80
Sequence Length
Sequence Length
(c) Task: WSC
(d) Task: WIC
150
100
Frequency
100
Frequency
75
50
50
25
0
200
400
200
400
600
Sequence Length
Sequence Length
(c) Task: BoolQ
(f) Task: MultiRC.
Figure 6: Histogram of sequence lengths for different datasets tokenized by the OPT-13B tokenizer. Among these datasets, MultiRC: has the longest sequence length, Iymaz = 739.
utilize the standard PyTorch (Paszke et al. memory-efficient approaches such as Attenntto (Dao al. (022). KDEformer (Zandich et al.) 2023), and HyperAttention Han et al. (203).. Although the combination between Addax and these methods remains unexplored, we posit that their combination could significantly enhance Addax by further reducing the memory demands and augmenting performance. We leave the exploration of Addax's 's interplay with various memory-efficient methods to future work.
2019) implementation of transformer. We do not use the
19

---

Published as a conference paper at ICLR 2025
The prompts used in the OPT experiments identical to those used in Malladi et al. (2023). There are three types of tasks: classification (cls.), multipl--chocee (mch.), and -anestionnnnswriing (QA). <text> the input from the dataset and blue text are the label words. We follow the same practice as in Malladi et al. (2023): for the inference task, we incorporate different candidates into the prompt, compute the average log- oog-likelhhood for each, and select the candidate with the highest score. For question-answering (QA) tasks, answers are produced through greedy decoding.
Table 5:
Dataset
Type
Prompt
SST-2 RTE
cls.
<text> It was terrible/great
cls.
<premise>
Does this mean that "<hypothesis>" is true? Yes or No?
Yes/No
CB
cls.
Suppose <premise> Can we infer that "chypothesis>"? Yes, No, or Maybe?
Yes/No/Maabe
BoolQ
cls.
<passage> <question>?
Yes/No
WSC
cls.
<lexi> In the previous sentence, does the pronoun "<span2>" refer to <spanl>? Yes No?
Yes/No
WIC
cls.
Does the word "<word>" have the same meaning in these two sentences? Yes, No?
<sent>>
<sent2>
Yes/No
MultiRC
cls.
<paragraph>
Question: <question>
found this answer "<answer>". Is that correct? Yes No?
Yes/No
COPA
mch.
<premise> so/because <candidate>
SQuAD
QA
Tille: <tille>
Context: <context>
Question: <question>
Answer:
DROP
QA
Passage: <context>
Question: <question>
Answer:
D.5 EYPER-PARAMETER
We present the hyper- ssper-parameters for all experiments conducted with RoBERTa-large in Table and those for OPT-13B, OPT-30B, OPT-66B, Llama-2-70B in Table Table Table 9 and Table T0 respectively. It is important to note that for both models, the hyper-parameters grid utilized for MeZO and Adam adheres to the specifications set forth in Malladi et al. (2023).
For RoBERTa-lary experiments, both Addax and MeZO employ constant learning schedule, while Adam uses linear scheduling. For the training process, Addax and Adam are set for 1, 000 steps, while MeZO extends 100, 000 steps. We check validation performance every 50 training step and save the best validation checkpoint for testing.
For OPT experiments, Addax, SGD, IP-SGD and MezO similarly adopt constant learning rate schedule, with Adam maintaining its linear scheduling. Here, Adam is configured for 100 steps, whereas Addax. SGD, IP-SGD are set for 1, 000 steps, and MeZO for 20, 000 steps. For the SGD and IP-SGD experiments in OPT- 13B, we choose the batch size from our grid that maximizes the memory usage of a single A 110 GPU, as larger batch sizes tend yield better results. For OPT-30B, OPT-66B experiments and Llama-2-70B experiments, same or less than K' of Addax. We check validation performance every 1/20 training step and save the best validation checkpoint for testing. We leave detailed discussion how selected batch size, K! KO, and Ly for reporting accuracy in Appendix D.6
set the batch size of SGD and IP-SGD
As explained in the main body, Addax a faster convergence. For the RoBERTa-large experiments, Addax uses the learning rate of [ee- 5, 5e 5, le 4), while Me70 uses the learning rate 7) of {le 7, le 6, le 20-55 5). For the
larger learning rates than MeZO, resulting in
20

---

(K² K Learning Rate r) Published as a conference paper at ICLR 2025
OPT experiments, fix the learning rate 17 of Addax le 4. while MeZO magnitude smaller learning rate of {le 6, 11-7).
Table The hyper-parameter grids used for RoBERTa-large experiments. Addax and MeZO use constant learning rate schedule, and Adam uses linear scheduling. Addax and Adam use 1K steps and MeZO uses 100K steps. We check validation performance every 50 training steps and save the best for testing.
Experiment
cypper-arameters
Values
16-bit/32-bit Addax
R + K
64
KCl
(0.1,0.2,03300.,0033
K
Learning Rate
4] 1e-3
a
(3e - 4,1e 3, 3e
3, 4e 3.
5e 3, 7e 3, le
2, le 1}
McZO
Batch size Learning Rate
64 {le le 6, le 5}
1e 1e- 3
32-bit Adam
Batch size
(2,4, 5e 10-55 5}
Learning Rate 77
{le 5, 3e
Table 7: The syper--prramete grids used for OPT-13B experiments in one A100 GPU (40GB). Addax, SGD, IP-SGD and MeZO use a constant learning rate schedule, and Adam uses linear scheduling. Adam uses 200 steps. Addax, IP-SGD, SGD use IK steps and MeZO 20K steps. We check validation performance every 1/20 training steps and save the best for testing. Note that for IP-SGD, SGD, and MeZO, some runs may have encountered out-of -emoory errors during training when fine-tuning with one A100 GPU.
Experiment
Hyper- pagper-prammterss
Values
Addax
(4, 6)
le 4
le 3
a Lr
(1c 4,3c - 4,5G 5e - 4, 7e
4, 9c - 4}
{150, 155, 160, 165, 170
MeZO
Batch size Learning Rate 1)
{2,4, 6, 8, 10, 12, 14, 16, 20, 24, 28, 32} {le- 6, 6
77
le
SGD
Batch size
[2, 4, 6, 8, 10, 12, 14, 16, 20, 24, 28, 28,328)) 32)
Learning Rate y)
(5e 3, le
2, 5e
2}
IP-SGD
Batch size
(2,4,6,8, 10, 12, 14, 16, 20, 24, 28, 32}
Learning Rate 7
{le- 1.25e 7.5e 4)
Adam
Batch size
Learning Rate
{le - 5,5e 5e - 5, 8e - 5}
D.6 DETAILED METHODS IN SELECTING BS, K", K AND LT FOR REPORTING
D.6.1 OPT-13B EXPERIMENTS
In our OPT-13B experiments, we run Addax, MeZO, IP-SGD, and SGD using a single A100 (40GB) GPU. For MeZO, IP-SGD, and SGD, we select a batch size from the grid {2, 4, 6, 8, 10, 12, 14, 16, 20, 24, 28, 32). choosing the largest possible batch size that maximizes GPU memory usage without encountering out-of-memory errors. A method is considered to have failed to fine-tune the dataset if it cannot run with the smallest batch size from the grid. For example, SGD fails to fine-tune all datasets with the smallest batch size, and SPD fails to fine-tur me-tune the BoolQ, MultiRC, and SQuAD datasets single A A100 (40GB) GPU.
21

---

Published as a conference paper at ICLR 2025
Table 8: The haper--pppaaeee parem--ar grids used for OPT-30B experiments in one H100 GPU (80GB). Addax, SGD, IP-SGD and MeZO use constant learning rate schedule Addax, IP-SGD, SGD use IK steps and MeZO 20K steps. We check validation performance every 1/20 training steps and save the best for testing.
Experiment
sapper-prramters
Values
Addax
Learning Rate r (K* _
6), (4,
le 4 4
c
le
{le - 3.
4, 5e - 4, 7e 4, 9e - 4}
LT
{320, 180)
Me70
Batch size
{2,4.6. 8, 10, 12, 14, 16)
Learning Rate y
{1e 6, le 7}
-
SGD
Batch size
{2,
Learning Rate rj
(5e 3, le 2, 5e - 2}
IP-SGD
Batch size
(2,4) 43
Learning Rate 1)
[le 4, 1.25e 4, 7.5e 4)
Table 9: The hyper-parameter grids used for OPT-66B experiments in three H100 GPU (240GB total). Addax, SGD, IP-SGD and MeZO use constant learning rate schedule Addax, IP-SGD, SGD 1K steps and MeZO 20K steps. We check validation performance every 1/20 training step and save the best for testing.
Experiment
Hyper- peremetrrseterss
Values
Addax
(K2, K Learning Rate r) e
(4,6) 6)
le 4
le - 3
(le 4,
4, 5e 4,
4, 9e 4}
Lr
260
MeZO
Batch size Learning Rate r)
2244,6.8 8, 10, 12, 14, 16) {le - 7}
1e-3
SGD
Batch size
{2,4)
Learning Rate y
(5e - 3, le 2, 5e - 2}
IP-SGD
Batch size
(2, 4)
Learning Rate
(1e - 4, 1.20 4, 7.5e 4)
To further reduce the memory requirements of IP-SGD, we first fine-tune Addax on the longest dataset, MultiRC (Figure(o). as successfully running Addax the longest dataset ensures it can also run on shorter ones. We found that with (K1, k = (4, 6) and LT = 170, Addax achieves optimal performance on the MultiRC dataset while staying within the 40GB memory constraints. Thus, use the same (K', K = (4, 6) configuration for Addax when finettuning the remaining datasets, searching for optimal combinations of and Lr S 170.
D.6.2 OPT-30B, OPT-66B, LLAMA-2-70B EXPERIMENTS
For the OPT-30B experiments, we follow the same methodology as described in the OPT-13B experiments. For MeZO, IP-SGD, and SGD, we choose the largest possible batch size from the grid that maximizes GPU memory usage without encountering out-of-memory errors, using single H100 GPU (80GB). For Addax in the OPT-30B experiments, we fine-tune with two different settings: (K', KT) (4, 6) when L1 180 and (K' KD) (2, 6) when LT 320. In both settings, Addax does not encounter out-of -fmoorrrry issues when fine-tuning the MultiRC dataset in one GPU, indicating that it can also fine-tune the remaining datasets.
22

---

Published as a conference paper at ICLR 2025
Table 10: The hyper-paramettr grids used for Llama-2-70B experiments in three H100 GPU (240GB total). Addax, SGD, IP-SGD and MeZO use constant learning rate schedule Addax, IP-SGD, SGD use IK steps and MeZO 20K steps. We check validation performance every 1/20 training step and save the best for testing.
Experiment
saper-parameters
Values
Addax
_ (K+
(4, 6)
Learning Rate
le - 4
c
le
{le 3,
4, 5e
4, 7e
4, - 4}
Lr
240
Me70
Batch size
{2,4. 6, 8, 10, 12, 14, 16) {1e 6, le 7}
Learning Rate y
-
SGD
Batch size
(2,4)
Learning Rate fj
(5e 3, le 2, 5e - 2}
IP-SGD
Batch size Learning Rate 1)
(2,4)
(le 4, 1.25e 4, 7.5e 4)
80
Accuracy (%)
60
40
20
SST-2
SST-5
SNLI
MNLI
RTE
TREC
Zero-shot
32-bit Me20
32-bit Addax
16-bit Addax
32-bis Adam
Figure 7: Experiments on RoBERTa- large: 16/32-br bit Addax outperform zero-shot and McZO across all tasks and outperform Adam in four out of six tasks. Detailed numbers can be found in Table
In the OPT-66B and Llama-2-70B experiments, we fine-tune MeZO, IP-SGD, and SGD by selecting the largest possible batch size from the grid that maximizes GPU memory usage without encountering out-of-memory errors, using three H100 GPUs (a total of 240GB). For simplicity. Addax uses the
same (K1 K (4, 6) configuration from the OPT-13B experiments, with ly 260 for OPT-66B and LT = 240 for Llama-2-70B, allowing Addax to successfully run without out-of- offoooorryyyyooooooooooooooooyye errors.
D.7 MEMORY PROFILING
In our memory profiling, we conform to the methodologies previously established in Mal ladi et al. (2023). Our implementation utilizes the default configuration provided by the trans cransforme formers (Wolf et al. 2020) package. We do not turn any advanced memory opti- mization technique such as gradient checkpointing. For multi-GPU backpropagation, we utilize the Fully Sharded Data Parallel (FSDP) (Fairscale) 2021) by Py Porch (Paszke al.) 2019). We use Nvidia's nvidia-smi command to monitor the GPU memory usage. We report the maximum GPU memory consumption observed throughout all experiments.
23

---

Published as a conference paper at ICLR 2025
Table 11: Experiments on RoBERTa-large (350M parameters). 16-bit Addax and 32-bit Addax outperform zero-shot and MeZ0 across the board on 6 tasks, while surpassing Adam in four out of six tasks. All experiments use prompts (Appendix) D.3). For the accuracy of 32-bit MeZO and 32-bit Adam, we report the results from Malladi et al. 2023).
Task
SST-2
SST-5
SNLI
MNLI
RTE
TREC topic-
Type
-enntmmnnt
natural language inference
Zero-shot
79.0
33.5
50.2
48.8
51.4
32.0
Samples per classes: k
16
bit MeZO
90.5
45.5
68.5
58.7
64.0
76.9
32-bit Addax
90.6
49.1
79.3
69.9
64.6
89.6
16-bit Addax
91.4
50.4
79.3
68.2
67.2
90.8
32-bit Adam
91.9
47.5
77.5
70.0
66.4
85.0
E ROBERTA-LARGE EXPERIMENTS
E.I ROBERTA-LARGE EXPERIMENTS MAIN RESULTS
Table 11 reports the detailed numbers of the accuracy on the RoBERTa-large model across different fine-tuning methods that are shown in Figure 7 For the accuracy of MeZO and Adam, directly
report the results from Malladi et al. (2023).
E.2 INVESTIGATION ON THE HYPER- PARAMETERS OF ADDAX
In this section, we explore the effect of different accuracy of both 32-bit and 16-bit Addax across various tasks utilizing the RoBERTa-large model using different combinations of hyper-parameterrs We include the combinations of a and K² KT
specifically reporting the
for 32-bit and 16-bit Addax in Figure and Figure Generally, it is observed that an increase in the ratio K correlates with improved accuracy across tasks for both 16-bit and 32-bit Addax configurations, as evidenced by the top row of the heatmaps for each task in Figure gldnd Figure 9 We did identify a consistent trend for across different tasks for both 16-bit and 32-bit Addax, suggesting that the optimal a could be task-specific.
K²
F
OPT AND LLAMA EXPERIMENTS
F.1 OPT-13B EXPERIMENTS MAIN RESULTS
Table 12 reports the detailed numbers of the accuracy the OPT-13B model across different fine- tuning methods shown in Figure II Details on batch size, time to the best validation checkpoint and memory for different algorithms
also available in Table 12 We also report GPU memory consumption across tasks and different fine-tuning methods for the OPT-13B model Figure with the exact number reported in Table 12 See Appendix fo7 memory profiling details.
F.2 OPT-30B, OРT-66B AND LLAMA-2-77- EXPERIMENTS MAIN RESULTS
Tables 15
16 present the complete results of fine-tuning OPT-30B and OPT-66B with Addax-
P. SGD, PP.GG. and MeZO, including metrics such accuracy, memory usage, and batch size/(K' k)) Figures 2 and 10 show the corresponding data.
F.3 COMPARING IP-SGD AND ADDAX WITH GRADIENT CHECKPOINTING
In these experiments, we enabled the Gradient Checkpointing technique (Chen et al. 2016). which reduces memory requirements for storing activations during backward computation. Gradient check- pointing benefits IP- and Addax, as Addax operates on shorter sequences. In particular, we observed that when incorporate gradient checkpointing, Addax demonstrates lower memory
24

---

Published as a conference paper at ICLR 2025
90.6
0.0003
8900
91.11
11.4
55.4
90
0.0003
35.4
41.3
19.5
41.6
61.1
0.001
16.4
1933
16.3
0.001
36.1
38.4
56.6
01.3
46.1
0.003
88
0.003
0.004
1h-4
1944
55.33
10.7
16.3
35.0
44.9
40
0.003
16.9
1865
10.33
10.9
86
0.005
15.3
12.
1809
62.3
63.2
0.007
0.007
25.4
36.0
41.5
35
84
0.01
0.01
05.1
1947
95.4
82
10.3
61.7
1011
61.5
0.1
0.2
0.3
0.4
0.5
0.1
0.2
0.3
0.4
0.5
(a) Task: SST-2
(b) Task: SST-5
24-1
71.5
73.3
77.3
76.5
Best
66.3
00.9
05.1
65.7
0.0003
61.0
181
15.7
0.0003
85.1
0.001
83.7
136
73.3
21.4
71.
0.001
66.1
68.4
91.9
10.3
08.
0,003
76.1
0.003
96-1
0,004
0.001
60
0,005
31.7
5-.5
08.5
76.7
71.9
0.005
57.5
04.11
07.4
07
0,007
71.4
694
20.2
75.9
71.7
50
0.007
66.0
05.1
65-9
16.1
08
0.01
75.5
0.01
0.2
0.5
0.1
0.3
0.4
(c) Task: SNLI
(d) Task: MNLI
Best
61.2
61.6
MLLS
64.
K2L5
Best
86.2
72.6
26.2
88.
21.6
0.0003
62.5
0.001
72.8
51.2
M6.2
80
0.003
0.003
27.6
05.65
03:11
1361
81.0
Accunne (N
0.004
61.7
53.5
61
6:.6
60.0
0.004
28.8
65.3
0.005
00.6
12.8
06.1
20.6
0.005
0.007
824
57.5
0.007
36.1
0.01
0.01
27.6
15
79.1
16.6
10.2
10
55.00
0.1
15.5
0.3
(c) Task: RTE
(f) Task: TREC
Figure 8: The accuracy (%) of the 32-bit Addax across different tasks on the RoBERTa-large model, with variable combinations of and K¹
requirements than IP-SGD. To illustrate this further, conducted two sets of experiments. To illustrate this, we conducted two sets of experiments.
In the first experiment, repeated the fine-tuning process with OPT-13B the CB dataset, enabling gradient checkpointing for IP-SGD and Addax. We used the same batch size and (K1 K settings in previous experiments. The results of these experiments enabling gradient checkpointing reduces memory usage but increases convergence time without affecting the final model quality. With gradient checkpointing enabled, the memory usage of IP-SGD decreased from 37.3 GB to 33.3 GB, while the memory usage of Addax decreased from 39.2 GB to 29.2 GB. This demonstrates that Addax indeed benefits more from gradient checkpointing regarding memory reduction. This advantage arises because Addax stores intermediate activations only for shorter sequences, which require less memory. On the other hand, IP-SGD requires storage of intermediate activations for longer sequences, which require more memory.
presented in Table 13 Notice that
25

---

Published as a conference paper at ICLR 2025
Besl
40.31
0.0003
6316
N0.7
9:11
0.0003
35.3
37.4
14.4
91.2
65.2
50
0.001
82.6
87.0
89.2
80.0
9:5
90.0
0.001
41.5
20.
44.5
0.003
68.1
80.6
96.7
87.5
0.003
0.001
61.1
65.1
85.0
0.005
11.1
29.5
96.4
67.95
66.2
0.007
27.8
67.3
-35
0.01
80.0
60
99.0
9:99
82.5
0.01
-30
0.1
1013
93.6
146.
1911
$:.4
205
41.0
0.1
0.2
03.3
0.1
0.5
0.1
0.2
0.3
0.4
0.5
(a) Task: SST-2
(b) Task: SST-5
23.3
71.6
26.3
76.3
70.7
Best
6117
96.1
07.7
61.1
0.0003
183
12.6
0.0003
03.1
0,001
73.9
7-1
75.8
74.1
11-5
0.001
29.6
65.01
90.5
16.3
671
0,003
71.6
70
0.003
65-9
13.
AOccIIN (9)
0.001
60
0,003
76-3
11.4
71.7
0.005
50.6
61.9
55.1
07:1
0,007
11.3
73-7
71.6
76.1
71.1
0.007
484
69.1
65.7
67.7
05
-55
0.01
71.7
50
0.01
65.4
0.2
0.4
0.5
0.1
0.3
0.4
(c) Task: SNLI
(d) Task: MNLI
a
67.1
66.7
69.5
56.5
06.4
Best
71.0
80.6
M0S
2100
0.0003
66
0.001
61.7
6166
51.1
13.5
0.001
2042
36.1
760
0.003
82.
61.0
0331
06.4
18.1
64
0.003
47.2
8011
06.5
16.
18.6
Accurne N9
0.004
66.6
0.004
27.4
0.005
0.007
83.9
15.3
08:1
0.007
00.2
63.1
0.01
57.6
56.1
67.1
00.3
67.0
60
0.01
37.6
19.8
96.1
62.
60.6
40
01.1
0.1
75.1
0.5
0.3
(c) Task: RTE
(f) Task: TREC
Figure 9: The accuracy (%) of the 16-bit Addax across different tasks on the RoBERTa-large model, with variable combinations of and K²
KK
In the second experiment, with gradient checkpointing enabled for Addax and IP-SGD, selected the largest possible batch sizes to maximize GPU memory utilization on a single A100 (40GB) GPU. presented in Table [4] In this experiment, Addax achieves an
The results of these experiments accuracy of 91.1, while IP-SGD achieves only 87.5. Due to its data partitioning strategy based on the varying memory requirements of data points, Addax can accommodate larger first-order batch size than IP-SGD, leading to significantly better final model quality. This result further illustrates that enabling gradient checkpointing benefits Addax as well.
In conclusion, the use of gradient checkpointing would benefit Addax in terms of memory require- ments. Gradient checkpointing trades convergence time for reduced memory usage without altering the optimization process or the final model quality. This is true for both Addax and IP-SGD.
26

---

Published as a conference paper at ICLR 2025
Table 12: Experiments on OPT-13B (with 1,000 examples) using single A100 (40GB) GPU except for Adam, which is run on five H100 GPUs (350GB total). Addax consistently outperforms zero-shot, MeZO, and IP-SGD across nine tasks, and surpasses Adam in seven of the nine tasks. An asterisk (*) indicates runs that encountered out-of-memory enrors during training, even with the smallest possible batch size.
Task
SST-2
RTE
CB
BoolQ classification
WSC
WIC
MultiRC
ReCoRD multiple choice
SQUAD
Metries
Task type
generation
AccuracyF11 (%)
Zero-shot
58.8
59.6
46.4
59.0
38.5
55.0
46.9
80.0
46.2
McZO
91.9
65.3
69.6
66.5
61,5
59.7
59.4
86.0
82.6
IPD
94.5
82.3
85.7
63,5
66.0
90.0
Adam
92.1
79.1
71.4
77.0
63.5
69.6
76.2
81.0
84.5
Addax
94.5
84.8
89.3
81.0
63.5
68.3
71.2
90.0
88.4
Memory (GB)
McZ0
29.7
39.0
38.7
39.6
31,6
31.4
36,9
27.6
36.8
SGD
IP-SGD
38.3
37.7
Adam
248.4
252.3
275.2
315.0
251.7
250.1
349.4
247.7
259.8
Addax
28.7
35.6
39.2
38.0
29.4
29.3
39.2
27.7
33.3
Batch Size
McZO
32
14
32
SGD
IP-SGD
16
2
2
12
12
32
32-bit
(K), K
Addax
(4,
Time (Min)
MeZO
222.5
289.2
182.8
255.4
40.3
103.9
363.8
31.7
245.5
SGD
IP-SGD
2.8
4.2
2.2
7.6
Addax
10.2
23.2
13.5
35,5
2.1
17.4
5,3
0.9
10.8
Table 13: Experiments on OPT-133 fine-tuning on dataset CB with gradient checkpointing enabling or not. Using gradient checkpointing would benefit Addax more than IP-SGD in terms memory footprint.
Metrics
IP-SGD
IP-SGD (Gradient Checkpointing)
Addax
Addax (Gradient Checkpointing)
Accuracy (%)
85.7
85.7
89.3
89.3
Batch Size/( K )
2
2
(4.6)
(4.6)
Memory (GB)
37.7
33.3
39.2
29.2
Time (min)
2.2
3.63
13.5
15.9
Table 14: Experiments on OPT-13B fine-tuning on dataset CB with gradient checkpointing enabling using a single A 100 (40GB) GPU. The batch size/(K'1 K is chosen as the largest within the 40 GB memory constraint.
Metrics
IP-SGD (Gradient Checkpointing)
Addax (Gradient Checkpointing)
Accuracy (%)
87.5
91.1
Batch Size/(K K)
10
(24.12)
Memory (GB)
35.5
38.8
Time (min)
6.1
24.6
F.4
CONVERGENCE SPEED OF DIFFERENT TUNING METHODS ON THE OPT- 13B MODEL
In this section, we demonstrate that 16-bit Addax reaches convergence speed comparable 16-bit SGD, despite SGD using 4x more first-order samples for backward propagation. Mcanwhile, Addax memory consumption is comparable MeZ0. The comparison of convergence speeds across the three methods is illustrated in Figure 11] For MeZO and SGD, the batch size is set to 16, while for Addax, we configure (K1 K
) as (4, 12). The learning rate for Addax set at 7) = le 4. For SGD, the learning rates are 77 {5e 3, le 2.5e 5e 2}. For MeZO, we utilize learning rates of 7= (1e - 6, 1e - 7}. We select the hyper-parameee parameters that yield the best validation accuracy across three methods. We utilize single H100 GPU (80GB total) for rumning both Addax and McZO, whereas SGD requires two H100 GPUs (160GB total). MeZO requires significantly more steps (20K steps) to converge compared to Addax and SGD (1K steps). Addax with smaller first-order batch
27

---

Published as a conference paper at ICLR 2025
Table 15: Experiments on OPT-30B (with 1000 examples) on a one H100 (80GB) GPU. * means the
run encounters out-of-memory enrors during training.
Task
SST-2
RTE
BoolQ
WSC
WIC
MultiRC
SQuAD generation
Metrics
Task type
classification
Accuracy/FI (%)
Zero-shot
56.7
$2.0
39.1
38.5
50.2
44.2
46.5
SGD
MeZO
90.6
66.9
56.3
79.9
IP-SGD BS 2
89.6
77.6
63.5
68.0
IP-SGD BS = 4
91.2
63.55
66.5
Addax (Lr 320)
93.9
83.4
80.8
63.5
66.8
75,8
85.9
Addax (Ly 180)
95.1
85.9
82.3
63.5
70.2
67.8
88.0
Batch Size
MeZO
16
10
16
12
SGD BS = 2
IP SGD BS
(K)
Addax (Lr 320)
(2,
Addax (l, 180)
(4,
Memory (GB)
MeZO
62.0
75.0
79.8
64.6
63.8
76.0
78.3
IP-SGD BS 2
62.5
80.0
64.4
62.9
IP-SGD B$ 4
65.2
66.3
66.5
Addax (Ly 320)
64.9
78.4
62.6
81.0
69.1
Addax (Ir 180)
64.4
79.5
79.5
71.3
Time (min)
MeZO
719.3
980.0
499.0
116.9
762.6
962.8
866.2
IP-SGD BS
1.9
1.0
7.9
IP-SGD
Addax (Lr 320)
4.0
9.8
32,3
1.4
19.7
11.4
3.7
Addax (Lr 180)
9.7
23.1
25.5
1.5
23.5
48.6
11.3
100)
Auuratyy (53
60
40
SST-2
RTE
BoolQ
WSC
WIC
MulaRC
SQuAD
225
Menary (G8)
2010
175
150
125
SST-2
RTE
BoolQ
WSC
WIC
MuliRC
SQuAD
10a
Then (Mia)
102
101
SST-2
RTF
BoolQ
WSC
WIC
MubiRC
SQQuAD
M:Z0
Figure 10: Accuracy/F-1 score, memory, and convergence time resulted from fine-tuning OPT-66B model with zero-shot. MeZO, SGD, IP-SGD, and Addax H100 GPU (240GB total). The label "OOM" means the encounters oft-orymemory error during fine-tuning. The exact numbers related this figure can be found in Table 16 in Appendix F2
28

---

Published as a conference paper at ICLR 2025
Table 16: Experiments on OPT-66B (with 1000 examples) on three GPUs (240GB total). * means the run encounters out-of-memory enrors during training.
Task
SST-2
RTE
BoolQ
WSC
WIC
MultiRC
SQuAD
Metries
Task type
classification
generation
Accuracy/F1 (%)
Zero-shot
57.5
67.2
66.8
43.3
50.6
49.4
48.1
SGD
MeZO
91.2
65.7
72.7
63.5
58.9
61.1
82.5
IP-SGD (BS=2)
89.1
82.3
67.0
63.5
65.8
87.0
IP-SGD (BS=4)
92.4
78.7
63.5
70.1
Addax (Lr 260)
95.5
85.2
84.0
63,5
66.9
88.3
Batch Size
MeZO
16
16
IP-SGD (BS=2)
IP-SGD (BS=4)
(K)
Addax (Ir 260)
(4,
Memory (GB)
MeZO
139.8
177.0
204.2
144.0
143.2
197.3
210.2
IP-SGD
136.5
166.2
203.6
145.4
139.4
215.4
IP-SGD
141.2
213.3
8
147.8
145.0
Addax (Ly 260)
141.9
204.6
228.7
145.9
144.3
215.4
173.6
Time (min)
MeZO
439.1
980.5
286.6
152.4
173.7
379.6
1036.2
IP-SGD
0.4
2.8
0.7
IP-SGD
3.9
6
9.1
Addax (Lr 260)
7.6
36.3
31.7
15.1
14.2
76.9
26.7
Table 17: Experiments Llama-2-708 (with 1000 examples) on three H100 (80GB) GPUs.
means
the run encounters out-of-memory errors during training.
Task
RTE
BoolQ
WSC
WIC
MultiRC
SQuAD generation
Metries
Task type
classification
Accuracy/F1 (%)
Zero-shot
60.6
75.9
55.8
49.8
45.8
70.5
SGD
McZO
52.7
63.1
75.0
55.6
64.4
92.3
IP-SGD BS
85.2
75.0
73.4
*
IP- BS
75.0
74.3
Addax (Ly 240)
89.9
87.9
76.0
74.5
85.3
93.4
Batch Size
McZO
16
6
SPD BS =
IP BS
Addax (Lr 240)
(4, 6)
Memory (GB)
MeZO
159.4
195.9
143.6
143.6
169.3
192.9
SPD BS
235.2
150.8
151.6
*
IPDGGD BS =
164.0
182.5
Addax (LT 240)
239.5
231.7
162.9
167.9
236.1
187.3
Time (min)
McZO
1288.7
565.0
6133.7
6405.5
879.9
932.0
SPD BS 2
2.6
5.0
9.5
SPD BS
1.3
11.0
Addax (LT 240)
31.7
28.0
5.0
27.0
30.0
53.7
size K- = 4 achieves convergence speed similar to SGD with batch size of 16, despite requiring significantly less memory.
G THEORETICAL RESULTS AND PROOFS
G.1 LIST OF ASSUMPTIONS
Assumption G.I. ((0:z) is L-Lipschitz smooth, i.e.,
[V ve(0'; :0]0;
< L 0 0'll 10, a
ERR E
29

---

Published as a conference paper at ICLR 2025
0.8
1.00
0.6
0.75
Lass
MeZO (42.5 GB)
0.4
0.50
16-bit Addor (53.4 GB)
16-bit SGD (151.9 GB)
0.2
MeZO (37.0 GB)
0.25
16-bit Addox (45.1 GB)
0.0
16-bit SGD (105.9 GB)
0.00
1000
2000
3000
4000
5000
0
500
1000
1500
2000
Steps
Steps
(a) Task: RTE
(b) Task: CB
Figure 11: Convergence speed of three fine-tuning methods [Adda--Wa MeZO, and SGD) on two fine-tuning datasets with the OPT- model. set batch size to for MeZO and SGD and fix (K', ) (4, 12) for Addax-WA We utilize single H100 GPU (80GB total) for running both Addax- WA and MeZO, whereas SGD requires two H100 GPUs (160 GB total) to run with BS = 16. MeZO requires significantly more steps converge compared to Addax-WA and SGD. Addax-WA with 4x less first-order samples achieves a convergence speed similar to SGD, despite requiring significantly less memory.
Assumption G.2. The stochastic gradient is unbiased and has bounded variance, i.e.,
Eg [V((;==) VL((), E= u) VC(O) S O², vo Rd.
Assumption G.3 (Low effective rank Hessian). There exists matrix 0 H = L. La such that V²L(x) ~ H, and the effective rank is at most r, i.e.
tr (H) < Lr.
Assumption G.4. [()) is μu-convx, i.e.,
L()) 2 L(0") (VL(0)) 0¹) + 11
0'11, 40, e cRd
2
G.2 USEFUL LEMMAS
Lemma G.5 (Gao et al. (2018). Lemma 4.1 (b)). Suppose Assumption G.1 holds, then the expected gradient estimated with SPSA a biased estimation of & VL( and satisfies
LL?
Lemma G.6 (Malladi et al. (2023), Lemma 2). Suppose Assumption
G.1 and Assumption G.2 holds,
then the variance of the gradient estimated with SPSA satisfies
d S 0²
Var Var( (ar(VL(0; B)) = Eg | - C B3)
K
Lemma G.7 (Zhang al. 2023). Lemma C.1 (iv)). Let N(O, v E Rd be some fixed vector, H. E Rad be some fixed matrix independent of z, we have:
Ez u)²z Hz) d+2
d
2v
Hu ||l|| tr (H)
G.3 CONVERGENCE ANALYSIS OF ADDAX SMOOTH NONCONVEX SETTING WHEN DO AND D' ARE SAME
The following Theorem the precise statement of the result behind Theorem 3.1 presented in the main body.
30

---

Published as a conference paper at ICLR 2025
Theorem G.8. Under Assumptions C.1
G.2
by running algorithm For T iterations with 0 m
? - AL
2- vt, the output satisfies
4(C(00) L.)
E|| [(
(212--00
a(1 2(12)
4nL
(11-a)²
a²d
+
+ 2K
0².
2(2-a)
(2 a)
2K
Proof: By Assumption G.T
L E, [((((0++)] & ((0,) + [. ( 01)
- BB 0 - -0,11²
2
(a)
[(0,) "/1 (1=(0))(--)VV((1) o.ggo [L((0; (3)]
LafE (1 - (1-aa) v)) B [V[( |0]] |||
2
LufE (1 - a)
Lnja²
2 B3 (L(0,) VL(0; B¹) )1)
Var(VL(00; (5)
2
(6)
S C(O.) - o)) VL(00)] am (VC(O,), [VC(O,;
LaK
+ 2
(1 - (1-a)V VL) + gg A L(0,; B3]
LijE (1 - a)
anj
0²
2K o²,
2K
where (a) substitutes the update of 0 and takes expectations to g" "oogg"g.=0 g²:(b) follows from the Lemma G.6 The third term on the Right- Hand-Side (RHS) can be further bounded as
V ((,,),
Eg 4 [(00;:)))]
odl
VL(00))
OT
Ego [VL B
|
+ 2
OT
VL(O.) - Eg Eg (
(6)
2
2
(6)
2
OFT 6 L2d
[VC(0,)
Eg (V B
2
2
8
where (a) uses the fact that ||u + v/ = |||| | + 2 (u, v); (b) applies Lemma G.5 to the last term. The fourth term on the RHS of equation can be bounded as
Lng
2 ([
al-v)) + aEgz
(VC(; B"]
Li
[V((,)) + / C(,,; 13) -0) 1
2
(a)
|
2
Lnd VC(0,) 1 + a22 Lof
Tay2 [V.) 2 L2
where (a) applies Cauchy-Schwarz inequality: (b) applies Lemma G.5 to the last term. Substitute equation to equation we have
equation 6
E, [C( < ((O.) (1
2 Lmd)rr [VV
007
Eg
(8)
am LL2 20m,L) L (1
a)² 0-2
Lnia²
0²
8
2K1
2K'
Choose T). <
2- 4T.
we have 1
-Lm 2- 0, + 20mp S a o² and
a)
|C(00))
a77 2
Ego (( 801l|² < L(0,) E, (C(0(+1))
4
(9)
ampe? 2 2 (1 + a Q²/2)
+ Ln a)²
Lnfa²
0²
0²
8
2K1
2K
31

---

Published as a conference paper at ICLR 2025
Sum from t to T, we have
2 -
am
VL(00)1
aTh E
B V ((1;; ((100:191)111 - ((0)) - [E(((r++1]
4
2
t=0
(10)
T
a a a2/2)
L(1 - [1)=
Lard
+
Σ
Σ
0²
0²)
8
2K1
2k
t=0
t=0
Choose 7/. = 1| S
2- 41
Vt, and divide both side by
(2-a)//
we have
4
2a
4(C(0)) L.)
E (0,)11 + 2-a E
Ego /
S
12 a)
(11)
a(1 + -22 22
4nL
((1-a)
+ ad
2(2 -d) a)
(2 -a)
2
2k
which completes the proof.
Corollary G.9. By choosing 1 = min
2-
2(C(0))
and
4L
TLo²
K
+
RJ
1/4
2(c(0,) L2) ((1 a1- //2/K' + a2dd/K9)
es
T
L3/dd/ aa- 02212
algorithm
converges with rate
(1-a)2
a²d
S ST.
K
+
K
L(00) - C.
2
T
(1 a)
add
=0
VT
K3
K
G.4 CONVERGENCE ANALYSIS OF ADDAX SMOOTH NONCONVEX SETTING WHEN D AND DJ ARE DIFFERENT
Assume that datasets D and D3 are different. Consider loss functions L, L and L where L and L are evaluated on datasets D and DD! of sizes No and N respectively. The combined loss function L is defined as:
(N + ( No + N1
Theorem G.10. Assuming that the loss functions L, CD, and C satisfy Assumptions and Lo & IImaz running Algorihtm for T iterations with learning rate 0 7tt "
G.1
and G.2
20
Vi.
1/2 < [CL and choose a NI the output satisfies
C(0,) L. 1/T(1 a) 1
(12)
1/2 Ldc + 2/2
nL (1 a)
((1 a)
a²d
à
4(1 a)
2K
2K
32

---

Published as a conference paper at ICLR 2025
Proof: By Assumption G.1
E,[C(0++1)] _ ((0.) (V 01-1 0.)1 + E 11111 -111
((0,) 7/ (V.(0)) (1-)) + alggo (0;:
Li (1 ".VV' (0.) a. (V2
Lnd o02² Egi [| (0)) VL' (0,: B3) )19
Lnfa² Var(V
2
2 [(()) (VC(O)). (1-))()) + a.
(0,;
LiR (1[1 ( + (.) EB --- 2
2
LnF(1 a)²² 0²
LIÑa²d
2K1
2K
3 ((0,) (VC(O))) (1 - V.V + a. (0,;
Lnk ||VC(0))) + cdd
| 21, dd
Lud ax² 0²
anfadd 0²
2K
2K
(13)
where (a) substitutes the update of 8 and takes expectation of gb: g (b) add and subtract a ( and Sahway-Schwartz inequality the third term and the last two terms follow from the Lemma G.6 (c) follows from the Lemma G.5 The second term on the Right-Hand-Sinn (RHS) can be further bounded by
(1 (1 a) + gl (B,;
TI
a:) (V) -
aBga (0:)
(14)
= IVC(0))) am VL(O)),
- (0.)
The second term in equation 14 can be further bounded by
om (V(())) Ey (0,; (0); (", -
am Vl.
Egu -
(15)
Ln
VC()))]
2
where (a) applies Cauchy-Schwards inequality: (b) applies Lemma G.5 to equation we have
Substitute equation
16 back
T/ (VL(O.), v)
( (0.) + alg / - (0,; )²))
(16)
<-m =<m0)=(())) aydLd VC(0.]]
33

---

Published as a conference paper at ICLR 2025
Substitute equation 16 back to equation 13 we have
Lip
E, [,( (().) (1
mm |VV
eLd(a + Lite)
2LL
Lif a)²
Lnf
111 VV((,)) +
11 +
o² +
o²
2
8
2K1
2K
(a)
[(0,) (1 Lite mm [C(())]
2/2 Ld[a + Lmpe)
m
LAT a)²
Lii
Lnf
0²
o²
2
8
2K
2h
(17)
where (a) choose that e//2 - [V2 then 2/2 || 1.
Choose Tt I 20 we have 1 - Lan 2 1 a > 0, MeL < 2a and
L
(1 a)m VC()))l - [(e,) E, [((++1))
+ 1/2 L e3/2 Lda²))
(n - a)²
LTf
(18)
1]| +
0² +
0²
2K1
2K
Sum from have
Σ (1 (1- 1l-=a)m)11lV -(00) -[((((+1]
1=0
T 1/2 Lda(6 + 2/2
T Σ
(L(1-
a)) Ladd
+
m
σ² +
2k 0²
4
2K1
t=
(19)
Choosing 1): S x. vt and dividing both sides by - (1-o) we have
[]] < ((0)) - L.
nT a))
(20)
1/2 Ldcx(6 + 3/2 4(1 - a)
n
a)²
add
(1 (1-a)
2K1
2K
where the expectation is taken over (uniformly) and the randomness of the algorithm. This completes the proof.
Corollary G.11. By choosing 17 = min
2a
((C)))) (C(0))-C.) C.)
and
TLo2
[1-0]
[
+a²d/K") 1/4
(C(O))-2.)) (1-0)) // /K1
e S min
27 2(C(0)) )o (1-a)²
3/2
2/// 1KK a2d/K²)
Addax converges with rate
(1-a)²
add
E| V
V2L
KT
Kd
C(Oo) L.
2
a
1
a)²
add
K1
K
34

---

Published as a conference paper at ICLR 2025
G.5 CONVERGENCE OF ADDAX WITH LOW EFFICIENT RANK HESSIAN
Theorem G.12. Under Assumption
G.1 Assumption G.3 by running algorithm 7 for T iterations 2
with 1/2 = S min( [T-a]T, 1-at2n -a+22"I(227) }, Vt, the output satisfies
E]]VC())11) ((0,) C.
ae? L2 (1 2naLr)
+
n11T
8C1
nLo²
( (1-aa)
2(2 + r)a²
(21)
201
K1
KO
where C1 = 1
nL (1 a 222
2
Proof:
Using the Taylor expansion with Lagrange remainder, we have:
1
C(0,+1) [(0.) + (VL(0.)) 0++1 0t) + + 20 0.) v((())(0++ 0z), (22)
where 0' = 10; + (1 1)0 for some À E [0, 1]. Taking expectation conditioned on everything
until 6, we have:
E, [E(0,+1)) ((),) + (VC(0)), - 0.)) 2 E. [0.+] V2) - 0.)]
(a) < ((0,) NL (VZ()))11-//V/(000))==) (V
Th E ( (1--)/C00,,3))4444001111
H (1--))
Ego C(0.) ( (V T/0 Ego [V
+ (1 - a)2
EB, [VVC [
2
02
Ego
2
+n a a) (v) HE [L((0;
(6)
((0,) (1 a) VV Tha (V
(1 - a)2T. (IV2(0))) d2
+
K Eg [
2
n2a²
HV
2
1 a1 a) (V.())), HE LL, ("))
(23)
where (a) substitute the update rule of algorithm 1 (b) uses the fact that 8 and B are independent to the last term, and applies Assumption G.2 to the fourth term. Next, we bound the two terms separately. For the fifth term, have:
Eg
CCO,; B,,, HVL(, (3))));))))))))))3))))))))))() = Eg
HVL(0, ((4C00:))).H&(LL:55)))
Eg
HVL(O; BO) L(O,; Ir
(b)
tr HEgo
C(0,: L(,;; B8))
(24)
where the (a) uses the property to trace that ir (ABC) = Lr (BCA), and (b) uses the fact that trace is a linear operator so E(tr (-)] tr (E).)) We have:
C(0 + ez; BO) - C(O, ez; B0)
2
Eg
TT = Eg
zz
26
35

---

Published as a conference paper at ICLR 2025
(a)
EBG (2ez T V VL((;; 'zzz'
2ca
(C(0, + ez; ez; B²) - ((;,-ez;B)) - 2Ez VLio;; B²))² zz
+ 2
(6)
< 2 80 2 (x (C)0::B"))
+ Egg,, x
(L(0, + cz; tz) - L(0,;
cz T VL(0,; Ba)) 2² 2z
2
(C(()): B²) - [(0, e2; ) ez V VL( )²)
+ Rga
2
Assumption G.I
2
dT.c
2Ega, (z VL(0;; 13222
+
Egg,a
2z
2²
2
dL Ia.
B [z V2(0;;
(25)
2
where (a) extracts the constant, add and subtract 2e V c )²4) and uses the Cauchy-Schwarz
inequality; (b) add and subtract L(0;; Bb) to the second term, then applies the Cauchy-
inequality; by Assumption G.1 we have |C( + ez: B) - L( B) ez V C(0; BB] < Ledd, and (c)
uses the fact that E[zzz 1 Id as 2 ~ (0, Ia). Substitute equation 25 to equation 24
we have:
Egg
(VL(0;;)) HV ((O,;
L2
s 2tr (HEgo, a [z T VL(;;; (z)) )zz '1) +
2 H
L
=2Ega, [z v L(,;; )² Hz
Lemma
chi
2d
E
2V B²) HV (VL;(0, (a) + [V [[V to (r()
d+2
2
Assumption
G.3
2dI(2++)
rd2
S
Es.
VL(0;: 13)
d+2
2
Assumption
G2
22L(2+r)
o²
rd L
|V
(26)
2+2
K
2
For the last term in equation
23
applies the Cauchy-Schward inequality:
(VL(O)), H (VL(0;
[IVC(0))] + (VL(0;
(27)
where (a) applies the fact that (a, b) s [[1a|] + []6|]3). Substitute equation
26 equation
27 and
equation 6 back to equation 23 we have:
Eg E,(4(0+11] s C(O,) - 171
1 - a
MeL
(1
a 2a 2a² VL())))
2
2
ma(1 1/(1 - aa)L
|ggo [VL((.:
aaa2 Ldd (1 + 2m, Th Lo²
(1-a)²
2(2 + r)a?
8
2
K1
KU
By setting 7/- = 17 < minf (1 (I-a)L: 2-at2a-((2+)) }. summing from t = ...... T 1, and divide
1
2-a
both side by n/TT, with C = 2
(1 a + 2a²L(2 + r)). we have:
2
||] ((0o) - L. L2d (1 + 21aLr)
n11T
8C1
noo²
((1
(1-a)
2(2 + r)a²
+
(28)
20, K
K"O
The theorem proved.
36

---

Published as a conference paper at ICLR 2025
Corollary G.13. By choosing
-
2(C(Oo) L.)
min
(1- (1-o)L - + 20²
TLo² (1-a)²
2a2+r)aa
+
Ki
K
and
1/4
(2((((0)) - L.) ((1 a))/// +2(2+7)2/KK2)
1
T
L3/Ad/a
algorithm
converges with rate
(1-a)² Kt + 212-r)a
((0o) L.
[[]VC(0))] S V2L.
2 111. 20 222 (2(2+r))
T
[-aa²
2(2 + r
K1 +
K0
G.6 CONVERGENCE ANALYSIS OF ADDAX IN SMOOTH STRONGLY CONVEX SETTING
Theorem G.14. Under Assumptions G.
and G.4
by running algorithm for T iterations with
0 l) 17 S zT: Vt, the output satisfies
0, 0.11 < (1
7/44 T
100 0.1
2
(29)
a
(1/p no L2?
2m (1 a)² am 0²
0² +
1
By Assumption G.4 with μ 0, have:
E, [||1
0.1
E,[l 0 +1 0: 0 - -
|10, 0.11 + E. 0 0.11 + 20.41 00,0, 0. -0.)
110, 0.1 + 1 (1 (1-a)VC(O.) + a.g. (CC,, B¹]]
+ n (1 - a) 2 (L(0,) - VL B 11 1/ VVr(VV
2n: 0, 0., (1 a-)VV + agg (((e; 154)
Th L2d n/ a) σ²
T/ naadd o²
S 11, 0.1² 2nd |VC))
K1
KU
- 27k (0. 0., vL(0.)) 200/ (0. - 00,Egu Ego [L (0); - vL(O.))
(c)
I 110 0.11 + 4n L (L(0,) - C.) m (2(C(0)) - L.) + 01 0.11?
20n' (0. 0., Ega (L B3)] -V
ni 262 (1-
- a)² 0² K1
naaad 0²
(30)
KD
where (a) substitutes the update of 0 and takes expectations to g" gg and conditions (,:(b) plugs in equation 7 to the second term, and the third term follows from Lemma G.6 (c) uses the fact that L( < L(O.) + [V()) for convex and smooth C(.). This implies that [|)( < (L) L(O.)) and applies to the second term. We apply Assumption G.4 to the sixth term in (c) by setting 0, 0' to O, and 0,, respectively. The middle term can be further bounded as:
-2 0, 0., Eg (L00; - vL(..))
37

---

Published as a conference paper at ICLR 2025
(a) 1
2a
10:
0.112
+
Ego
2a
1
ac² d2
10, 0.11
(31)
2x
2μ
where (a) applies Young's s incquality for inner products, i.c., -2(a, b) - #||a|| + T||b|| with
2a in our case. Then we have:
11110++1 0.112 s (1 111 10. 0.11 22/ 27 (C(O.) L.)
+ "1100 2 (1/111111mmm 1z + m: I.2
1f(11 - a)
111 a
0²
0²
(32)
K1
K
By setting S
1 2I. we have 1 - 1-212 2 0. Then we have:
naa ma2(11/ (1)
B.110 - (1
100 0.12+
2
2
+ 1/2 - ))²
0² + naadd
o²²
(33)
K1
Ku
Recursively apply the above equation 33 and sum from t = to T - 1 by setting = = S
22. we
1
have
E|| 0r 1 0.11] - (1 (1-
T
110 0. 0,12 12
T-1
Σ
1
1--2)
M
n02 + +n)221.2
2
2
2
j=0
T-1
Σ
1
)
2
172(1 a)2
0²
n 0²
K1
Ku
j=0
110g 0.112 + a1 + L2d Le
[1
2
2m(1 - a)
0² 2ma²t K" o²
(34)
where (a) comes from 0 < 1 1-nu
1. This completes the proof
#= 000 0. 1|2
}}
Corollary G.15. By choosing y) min
2 In
and
2L nf
(--o²² add
10 02
kl
(1-a)2
a²d
Ki
+
100
Ld
Ta(1/11 + +)(x))
algorithm ! converges with rate
4² 10 0.11
(1-a) a)²
a²d
0|
0.1121
s
In
T
o²
nT
4
(1-0)²
+ a²d
o²
K1
Kō
R
In(T)
(11 [1=
+ ko0 ))
a²d
0
T
K1
38


Reviews:
[{"review_id": "8QKkow22jC", "reviewer": "Reviewer_covw", "summary": "This paper presents Addax, a novel approach to fine-tuning large language models (LMs) with reduced memory usage and improved convergence speed. Addax aims to bridge the gap between two existing methods—IP-SGD, which is memory-intensive, and MeZO, which suffers from slow convergence due to its zeroth-order gradient estimation. By combining zeroth-order and first-order gradient estimates based on the memory consumption of different data points, Addax attempts to balance the benefits of each approach, reducing memory requirements without sacrificing performance.", "questions": "- the experiments done here are all done with high-end h100 gpus, I am curious who are actually the targeted users for Addax? does this work on lower-end GPUs? say 3090s?\n- with the empirical improvement presented, and the complexity of the methods. I am curious if the authors could further list your No.1 selling point vs. well-established first-order finetuning methods, what would be your pitch to convince the community to adopt Addax in practise?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The concept behind Addax is promising. By selectively applying zeroth-order and first-order updates depending on memory requirements, Addax cleverly attempts to sidestep some of the limitations of both IP-SGD and MeZO. The added benefit of using zeroth-order gradients as a regularizer for first-order gradients is an interesting approach that theoretically could lead to better model generalization.\n\n- it assigns data points with higher memory demands to a zeroth-order gradient estimation, while data points with lower memory demands use a more standard first-order gradient approach.\n\n- the wallclock time details provided are very helpful.", "weaknesses": "- Addax introduces an additional layer of complexity by dynamically selecting gradient types based on memory consumption, which may require extra engineering effort to implement effectively in various environments. \n-  The algorithm involves splitting data into two subsets, resetting random seeds, and clearing gradients for each update step. This is quite complicated.\n- While the hybrid gradient approach is novel, it’s unclear from the presented results to what extent this technique actually enhances performance, or to be specific, in results shown in figure 1 e.g. some of the tasks's improvement seem very marginal compared to MEZO.", "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": "YtM6xW1p9i", "reviewer": "Reviewer_nhLb", "summary": "This paper proposes a novel LLM fine-tuning methods, which try to combine the first-order (sgd) and zeroth-order (mezo) to efficiently fine-tune LLM. The experiments on OPT and LLaMA Illustrate that the proposed method can achieve a better performance than SGD and faster convergence than MeZO.", "questions": "1. From my experience about MeZO, it usually begin to work after a large number of epochs. So I would like to ask whether Addax and SGD use a same epoch value?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The intuition and motivation are very easy to understand and the paper is also easy to follow. \n\n2. The paper provides the results on different models with different sizes to verify the performance of the proposed method, which also makes the paper more convincing.", "weaknesses": "1. The opinion that \"SGD can match the performance of Adam in fine-tuning tasks.\" is not very reasonable to me. I think the main reason why the authors reached this conclusion is because the tasks in this paper are too easy (mainly focusing on classification tasks). If the authors conduct some experiments on math and code generation tasks, Adam could be better than vanilla SGD. \n\n2. Ablation Study. We can notice that the paper provides the results of ip-sgd and addax in Table 1, 2, 3. And the main difference of these two methods is some long sequences use different optimizers. Specially, for addax, the long sequences use zero-order optimization. What is the result if we set the learning rate of these zeroth order optimization to 0 (That means we can skip these long sequences and no update with these data). I just want to see whether the zeroth-order optimization really works. I guess update too much times with sgd may cause overfitting and just skip some training steps (long sequences in addax) can improve the performance (even if don't use zeroth-order on this data). So I would like to see the results if we just skip these long sequence data in addax. \n\n3. The paper mainly focuses on the tasks on SuperGulu. These tasks are too easy to LLM fine-tuning with first-order methods. It would be better if the authors could provide the results of complex tasks. This is just a suggestion and should not be a weakness.", "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": "M3cTD3A3bE", "reviewer": "Reviewer_5pxQ", "summary": "This paper proposes Addax, which assigns different mini-batch of data to compute zeroth-order and first-order gradients respectively based on the sequence length of the data, and combines the two kinds of gradients to update the model parameter. Addax outperforms existing zeroth-order method MeZO in accuracy/F1 score and achieves comparable memory footprint with MeZO. Furthermore, it surpasses IP-SGD and Adam with significantly less memory requirement.", "questions": "Please refer to the Weaknesses.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "-\tThe motivation of the proposed method is well supported by the preliminary experiments.\n-\tThis paper provides theoretical guarantee that the propose method enjoys faster convergence and less restrictive hyper-parameter choices than MeZO.\n-\tThis paper conducts extensive experiments to demonstrate its effectiveness on models with various sizes ranging from 350M to 70B, achieving superior performance both in test accuracy, memory and time efficiency.", "weaknesses": "- In Algorithm 1 step 7, the mini-batch $\\beta^1$ is sampled from $D^1$, while in Line 298, it is claimed that $\\beta^1$ is sampled from $\\beta^0$, which is inconsistent. \n- The performance of the proposed method is sensitive to the mixing constant $\\alpha$ for combining the zeroth-order and first-order gradients. Also, it may be sensitive to the length threshold $L_T$ and the random seed. \n- It would be better to further demonstrate the effectiveness of the proposed method on more challenged tasks such as commonsense reasoning tasks and mathematics tasks.", "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": "XaaAgu16fq", "reviewer": "Reviewer_dTUU", "summary": "Authors propose Addax, algorithm which adaptively combined zero-order (MeZO) and in place first-order (IP-SGD) optimization methods, to achieve LLM fine-tuning with low memory requirement.", "questions": "How would results look like when we run IP-SGD with gradient checkpointing?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "An excellent description of problems with MeZO and IP-SGD. \nBetter results than MeZO while keeping memory consumtion low.", "weaknesses": "Authors do memory profiling without gradient checkpointing, a very common technique to lower memory requirements for storing activations for backward computation. I believe this overestimates the memory requirements of IP-SGD. \nThe main selling point is that Addax has lower memory requirements than IP-SGD. However, this selling point vanishes the moment we turn gradient checkpointing on. \n\nA minor selling point of Addax is that it sometimes produces better results even when zero-order gradients are computed concurrently with first-order ones (Addax-WA). However, it is hard to find which results actually belong to Addax-WA in the paper. E.g. section 4 says \"we plot the convergence curves of Addax-WA and MeZO using the same batch size in Figure 11\", but figure 11 label says \"Addax with 4×\nless first-order samples achieves a convergence speed similar to SGD\", which suggest we are not running Addax-WA here.\n\nNitpick: Figure 4 just feels wrong, because SGD should always have higher memory requirements than IP-SGD).", "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)