Sparse MeZO: Less Parameters for Better Performance in Zeroth-Order LLM Fine-Tuning

## Abstract

Abstract While fine-tuning large language models (LLMs) for specific tasks often yields impressive results, it comes at the cost of memory inefficiency due to back-propagation in gradient-based training. Memory-efficient Zeroth-order (MeZO) optimizers, recently proposed to address this issue, only require forward passes during training, making them more memory-friendly. However, the quality of gradient estimates in zeroth order optimization often depends on the data dimensionality, potentially explaining why MeZO still exhibits significant performance drops compared to standard fine-tuning across various tasks.
Inspired by the success of Parameter-Efficient Fine-Tuning (PEFT), this paper introduces Sparse MeZO, a novel memory-efficient zeroth-order optimization approach that applies ZO only to a carefully chosen subset of parameters. We propose a simple yet effective parameter selection scheme that yields significant performance gains with Sparse-MeZO. Additionally, we develop a memory-optimized implementation for sparse masking, ensuring the algorithm requires only inference-level memory consumption, allowing Sparse-MeZO to fine-tune LLaMA-30b on a single A100 GPU.
Experimental results illustrate that Sparse-MeZO consistently improves both performance and convergence speed over MeZO without any overhead. For example, it achieves a 9% absolute accuracy improvement and 3.5x speedup over MeZO on the RTE task.

## 1 Introduction

Fine-tuning large language models for specific tasks or datasets has become a prevalent practice in machine learning. However, a major obstacle in fine-tuning is the substantial memory requirements, which escalate as models increase in size and complexity, thereby limiting the scalability and accessibility for those with limited computational resources.

Figure: Figure 1: Performance of MeZO and Sparse-MeZO (S-MeZO) on RTE task. S-MeZO can achieve $3.5$x speedup compared with MeZO.
Refer to caption: x1.png

To mitigate the memory constraints, Parameter Efficient Fine-Tuning (PEFT) has been developed, allowing for the modification of only a subset of parameters and achieving comparable results to full model tuning (Zaken et al., ; Li & Liang, ; Lester et al., ; Hu et al., ; Zhang et al., ).
However, PEFT methods still necessitate the calculation of gradients for backpropagation and caching of numerous activations during training, which introduces additional memory overhead.
For instance, demonstrates that, even with PEFT, training still requires approximately 6 times more memory than the memory cost for inference.
This discrepancy raises a critical question: Can large language models be fine-tuned solely with the cost of inference?

In response to these challenges, zeroth-order (ZO) optimization presents a promising solution (Spall, ). ZO optimization is a gradient-free method that estimates gradients using only the forward pass of the model, eliminating the need for backpropagation and, consequently, reducing memory usage. MeZO (Malladi et al., ) is a recently proposed zeroth-order method for fine-tuning LLMs that has demonstrated impressive performance. However, the performance and convergence rate of ZO methods are usually highly dependent on the dimensionality of the parameters, with higher dimensionality leading to slower convergence.

Building on the insights from PEFT, we propose a novel sparse memory efficient zeroth-order method (Sparse-MeZO) to solve the issues from the high dimensionality of LLMs. In traditional parameter-efficient fine-tuning, there’s a tradeoff between memory efficiency and performance. But surprisingly, when we only tune a subset of parameters in MeZO, it significantly improves the performance while accelerating the convergence rate. This renders our method an improvement over MeZO without any overhead. Our contributions can be summarized as follows:

- •
In this paper, we propose a sparse Memory-Efficient Zeroth-Order optimization method Sparse-MeZO (S-MeZO) for large language model fine-tuning. We also provide theoretical analysis to show the convergence of Sparse-MeZO.
- •
To determine the sparse parameters that need to be updated, we investigate how the importance of each parameter in fine-tuning correlates with its magnitude.
Although existing sparsification methods usually prune out parameters with smaller magnitudes, our evaluations show that smaller parameters are more important than larger weights for Zeroth-order fine-tuning methods.
- •
Different from the efficient implementation with random seed in MeZO, we propose a novel memory-efficient implementation of Sparse-MeZO, which can compute the sparse mask and perturb parameters in the forward pass. The technique enables fine-tuning LLaMA-30b with Sparse-MeZO on a single A100 GPU.
- •
We conduct empirical studies on LLaMA and OPT. The experimental results demonstrate that Sparse-MeZO can improve the fine-tuning performance and yield a faster convergence rate compared with vanilla MeZO across a wide range of natural language processing tasks. For example, it achieves a 9% absolute accuracy improvement and 3.5x speedup over MeZO on the RTE task, as shown in Figure 1.

## 2 Preliminaries

### 2.1 Parameter-Efficient Fine-Tuning

Parameter-Efficient Fine-Tuning (PEFT) is designed to facilitate efficient adaptation by updating only a subset of the model’s parameters, rather than fine-tuning the entire model (Hu et al., ; Zaken et al., ). These PEFT approaches can be categorized in various ways. We mainly focus on the selective methods and additive methods.

#### 2.1.1 Selective Methods

Selective Methods try to selectively fine-tune a portion of a model and these methods have been explored in various studies. For example, ; focused on the model’s bias terms, finding that fine-tuning these terms alone could rival the results of fine-tuning the entire model. However, the effectiveness of this approach diminishes with larger datasets, as shown in further analysis by . Beyond static parameter adjustments, there has been an exploration into dynamically modifying parts of the model. For instance, FreezeOut, introduced by , suggests a gradual freezing of layers, effectively removing early layers from the backward pass to quicken training. This concept was later applied to language models, with AutoFreeze (Liu et al., ) confirming its viability. Nevertheless, these techniques still demand considerable computational resources and tend to yield less optimal final outcomes.

#### 2.1.2 Additive Methods

Additive methods, as an alternative to updating existing parameters, involve incorporating new layers into models, with the fine-tuning process focusing solely on these added layers (Houlsby et al., ; Rebuffi et al., ; Lin et al., ; Hu et al., ). Traditional techniques in this category, such as adapters (Houlsby et al., ), implemented layer additions in a sequential manner, which unfortunately led to increased inference latency. LoRA (Hu et al., ) has been proposed to mitigate this issue, which freezes the weights of the pre-trained model and introduces trainable matrices based on rank decomposition into each layer. Then, it can directly integrate the newly learned weights into the main model. Following this, IA3 (Liu et al., ) introduced innovative methods for adding parameters, balancing parameter count with accuracy, while LST (Sung et al., ) introduced a highway structure that learns only small, auxiliary channels, aiming to decrease memory demands. Despite these advancements, additive methods generally require meticulous design, and many fail to reduce the computational load during the backward pass, as seen in IA3, LoRA, and Adapter approaches. Additionally, some of these methods, including Adapter and LST, can incur additional computational overhead in the forward pass, making them less practical for certain applications.

### 2.2 Zeroth-Order Optimization

Unlike traditional gradient-based optimization methods that rely on derivatives to guide the search for optimal solutions, zeroth-order (ZO) optimization techniques do not require derivatives for optimization (Spall, ; Liu et al., , ). These methods utilize only the value of the objective function, denoted as $f(\boldsymbol{x})$, at any chosen point $\boldsymbol{x}$. To estimate the gradient in the direction of vector $\boldsymbol{z}$, the objective function is assessed at two points in close proximity, $f(\boldsymbol{x}+\epsilon\boldsymbol{z})$ and $f(\boldsymbol{x}-\epsilon\boldsymbol{z})$, with $\epsilon$ being a minimal value. Following this, conventional optimization algorithms, such as gradient descent or coordinate descent, are implemented using these approximated gradient values. Currently, ZO methods have been widely used in various applications, such as adversarial attack and defense (Chen et al., ; Tu et al., ; Ye et al., ; Ilyas et al., ), Auto-ML (Ruan et al., ; Wang et al., ), natural language processing (Sun et al., , ), reinforcement learning (Vemula et al., ), Signal Processing (Liu et al., ), and on-chip training (Gu et al., ).

#### 2.2.1 SPSA

Simultaneous Perturbation Stochastic Approximation (SPSA) (Spall, ) is a descent method that optimizes systems with multiple unknown parameters to identify global minima without relying on backpropagation. SPSA estimates gradients based on two measurements of the objective function, thus eliminating the need for automatic differentiation and backpropagation to obtain precise gradients. This approach notably avoids memory overhead associated with storing intermediate results for backpropagation. Specially, it employs a random vector $\boldsymbol{z}\in\mathbb{R}^{d}$, where $\boldsymbol{z}\sim\mathcal{N}(\boldsymbol{0},\boldsymbol{\mathcal{I}_{d}})$, to define to perturb the parameters and estimate the gradient based on the difference between $\mathcal{L}(\boldsymbol{\theta}+\epsilon\boldsymbol{z})$ and $\mathcal{L}(\boldsymbol{\theta}-\epsilon\boldsymbol{z})$:

$$ $\displaystyle\boldsymbol{g_{z}(\theta)}$ $\displaystyle=\frac{\mathcal{L}(\boldsymbol{\theta}+\epsilon\boldsymbol{z})- \mathcal{L}(\boldsymbol{\theta}-\epsilon\boldsymbol{z})}{2\epsilon}\boldsymbol {z}$ (1) $\displaystyle\approx\boldsymbol{z}\boldsymbol{z}^{\top}g(\boldsymbol{\theta}).$ $$

#### 2.2.2 MeZO

ZO-SGD employs SPSA to estimate the gradient. In general, conventional ZO-SGD algorithms utilizing SPSA consume twice the inference memory. MeZO (Malladi et al., ) is a memory-efficient variant of ZO-SGD. It circumvents the storage of gradients by saving the random seed and resampling the same random noise $\boldsymbol{z}$ with the seed during forward process. Therefore, it can reduce the training costs close to inference memory cost.

We try to further explain how MeZO can be memory-efficient in Figure 2. At step 1, it will sample a noise $\boldsymbol{z}$ to perturb $\boldsymbol{\theta_{t}}$ and then can obtain $\boldsymbol{\theta_{t}^{\prime}}$ to calculate $\mathcal{L}^{+}$. At step 2, MeZO can resample the same noise $\boldsymbol{z}$ with same seed $s_{t}$ to directly perturb $\boldsymbol{\theta_{t}^{\prime}}$ but not $\boldsymbol{\theta_{t}}$. Then, it can use $\boldsymbol{\theta_{t}^{\prime}}-2\epsilon\boldsymbol{z}$ to reconstruct $\boldsymbol{\theta_{t}}-\epsilon\boldsymbol{z}=\boldsymbol{\theta_{t}^{\prime
\prime}}$ without additional memory cost and finally calculate the loss value $\mathclap{L}^{-}$.

Figure: Figure 2: Memory Efficient Implementation in MeZO. At step 1, we will sample a random noise $\boldsymbol{z}$ to perturb $\boldsymbol{\theta_{t}^{\prime}}=\boldsymbol{\theta_{t}}+\epsilon\boldsymbol{z}$ and save the seed $s_{t}$. At step 2, MeZO can resample the same noise $\boldsymbol{z}$ with same $s_{t}$ to calculate $\boldsymbol{\theta^{\prime\prime}}=\boldsymbol{\theta^{\prime}}-2\epsilon \boldsymbol{z}=\boldsymbol{\theta_{t}}-\epsilon\boldsymbol{z}$ and release the memory of $\boldsymbol{\theta^{\prime}}$.
Refer to caption: x2.png

#### 2.2.3 Sparsity for Zeroth-order Optimization

The hypothesis proposed by , known as the lottery ticket hypothesis, showed that within a densely connected neural network that is randomly initialized, there exists a subnetwork of sparse yet high-quality connections. Several related works have tried to apply the sparsity to zeroth-order optimization (Wang et al., ; Cai et al., , ; Balasubramanian & Ghadimi, ; Ohta et al., ; Gu et al., ; Chen et al., ). For example, DeepZero proposes a novel ZO training protocol with model pruning guided sparsity. However, these methods mainly focus on the neural network training from scratch with random initialization, while the application of sparse zeroth-order optimization in fine-tuning tasks remains an area of ongoing exploration. In this paper, we will analyze the effects of sparse ZO methods on large language model fine-tuning tasks and propose a novel sparse zeroth-order method to obtain a better performance and faster convergence.

## 3 Proposed Method

Parameter-Efficient Fine-Tuning (PEFT) has been widely used in various areas, which only updates a fraction of full parameters and can yield comparable performance. In addition, it’s widely recognized that both the performance and the convergence rate are significantly influenced by the dimensionality of the weights to ZO methods. For example, illustrates that an increase in dimensionality tends to proportionally slow down the convergence rate. These motivate us to propose a sparse zeroth-order optimization method for large language model fine-tuning, which only updates a subset of full parameters and yields a comparable performance and faster convergence.

### 3.1 Sparse-MeZO

Consider a labelled dataset $\mathcal{D}=\{(\boldsymbol{x}_{i},\boldsymbol{y}_{i})\}_{i\in\mathcal{[|D|]}}$ and let $\mathcal{L}(\boldsymbol{\theta};\mathcal{B})$ denotes the loss on a mini-batch $\mathcal{B}$. We can define a sparse mask $\boldsymbol{m}\in\{0,1\}^{d}$ to selectively sample the random noise $\boldsymbol{z}\in\mathbb{R}^{d}\text{ with }\boldsymbol{z}\sim\mathcal{N}(
\boldsymbol{0},\boldsymbol{I}_{d})$ on the sub-net of pre-trained model. A sparsified version of random perturbation can be defined as $\boldsymbol{\hat{z}}\in\mathbb{R}^{d}$:

$$ $\boldsymbol{\hat{z}}=\boldsymbol{m}\odot\boldsymbol{z}.$ (2) $$

Based on this sparse perturbation $\boldsymbol{\hat{z}}$, we can redefine MeZO algorithm on Section 2.2.2 as Sparse-MeZO. The main difference is from the estimated gradient $\boldsymbol{g_{\hat{z}}(\boldsymbol{\theta})}$, which can be defined as :

$$ $\displaystyle\boldsymbol{g_{\hat{z}}(\boldsymbol{\theta})}$ $\displaystyle=\frac{\mathcal{L}(\boldsymbol{\theta}+\epsilon\boldsymbol{\hat{z }};\mathcal{B})-\mathcal{L}(\boldsymbol{\theta}-\epsilon\boldsymbol{\hat{z}}; \mathcal{B})}{2\epsilon}$ (3) $\displaystyle=\frac{\mathcal{L}(\boldsymbol{\theta}+\epsilon\boldsymbol{m} \odot\boldsymbol{z};\mathcal{B})-\mathcal{L}(\boldsymbol{\theta}-\epsilon \boldsymbol{m}\odot\boldsymbol{z};\mathcal{B})}{2\epsilon},$ $$

where $\epsilon$ represents the perturbation scale. Note that although working with a sparse perturbation, we still need to use full parameters $\boldsymbol{\theta}$ in the forward process to calculate the loss value $\mathcal{L}(\boldsymbol{\theta+\epsilon\boldsymbol{\hat{z}}})$.

### 3.2 Convergence Analysis of Sparse-MeZO

In this section, we will explain why Sparse-MeZO can accelerate the convergence. We can define a sub-network in pre-trained large language models, which is determined by the sparse mask $\boldsymbol{m}$. The main idea of our proof is that if we follow the updated role in Sparse-MeZO, the gradient norm on the sub-network can be smaller than $\sigma^{2}$ after $\mathcal{O}(\frac{\hat{d}L}{\sigma^{2}})$ steps, where $\hat{d}$ is the number of parameters in the sub-network. Therefore, ZO can use fewer steps to converge when we only focus on a sub-network. Some related work has illustrated that only tuning the sub-network can achieve comparable performance, which will be empirically verified in our experiments.

Firstly, we assume the loss function $\mathcal{L}(\boldsymbol{\theta};\boldsymbol{x})$ is Lipschitz Continuous:

###### Assumption 3.1 (Lipschitz Continuous) .

$$ $\|\nabla\mathcal{L}(\boldsymbol{\theta};\boldsymbol{x})-\nabla\mathcal{L}( \boldsymbol{\theta^{\prime}},\boldsymbol{x})\|\leq\frac{L(l)}{2}\|\boldsymbol{ \theta}-\boldsymbol{\theta^{\prime}}\|^{2},$ (4) $$

where $\nabla\mathcal{L}(\boldsymbol{\theta};\boldsymbol{x})$ denotes the true first-order gradient of $\boldsymbol{\theta}$ on $\boldsymbol{x}$ and $L(l)$ represents the Lipschitz constant of $\mathcal{L}(\cdot)$.
Given $\mathcal{L}_{\hat{z}}(\boldsymbol{\theta})=\mathbb{E}_{\boldsymbol{\hat{z}}}[
\mathcal{L}(\boldsymbol{\theta}+\epsilon\boldsymbol{\hat{z}})]$ and the above Assumption 4,
we can obtain the relationship between sparse gradient $\widehat{\nabla}_{\boldsymbol{\theta}}\mathcal{L}_{\hat{z}}(\boldsymbol{\theta})$ and the expectation of estimated sparse ZO gradient $\boldsymbol{g_{\hat{z}}(\theta)}$:

###### Lemma 3.2 .

ZO gradient $\boldsymbol{g_{\hat{z}}(\boldsymbol{\theta})}$ is unbiased estimation of $\widehat{\nabla}_{\boldsymbol{\theta}}\mathcal{L}_{\hat{z}}(\boldsymbol{\theta})$:

$$ $\displaystyle\widehat{\nabla}_{\boldsymbol{\theta}}\mathcal{L}_{\hat{z}}( \boldsymbol{\boldsymbol{\boldsymbol{\theta}}})$ $\displaystyle=\boldsymbol{m}\odot\nabla_{\boldsymbol{\theta}}\mathcal{L}_{\hat {z}}(\boldsymbol{\theta})$ (5) $\displaystyle=\boldsymbol{m}\odot\nabla_{\boldsymbol{\boldsymbol{\theta}}} \mathbb{E}_{\boldsymbol{\hat{z}}}[\mathcal{L}(\boldsymbol{\theta}+\epsilon \boldsymbol{\hat{z}})]$ $\displaystyle=\mathbb{E}_{\boldsymbol{\hat{z}}}[\frac{\mathcal{L}(\boldsymbol{ \theta}+\epsilon\boldsymbol{\hat{z}})-\mathcal{L}(\boldsymbol{\theta}-\epsilon \boldsymbol{\hat{z}})}{2\epsilon}\boldsymbol{\hat{z}}]$ $\displaystyle=\mathbb{E}_{\boldsymbol{\hat{z}}}[\boldsymbol{g_{\hat{z}}(\theta )}],$ $$

where $\boldsymbol{g_{\hat{z}}(\theta)}=\frac{\mathcal{L}(\boldsymbol{\theta}+
\epsilon\boldsymbol{\hat{z}})-\mathcal{L}(\boldsymbol{\theta}-\epsilon
\boldsymbol{\hat{z}})}{2\epsilon}\boldsymbol{\hat{z}}$. We can find that $\boldsymbol{g_{\hat{z}}(\theta)}$ is unbiased estimation of $\widehat{\nabla}_{\boldsymbol{\theta}}\mathcal{L}_{\hat{z}}(\boldsymbol{\theta})$.
Then, based on the equation $\widehat{\nabla}_{\boldsymbol{\boldsymbol{\boldsymbol{\theta}}}}\mathcal{L}_{
\hat{z}}(\boldsymbol{\boldsymbol{\theta}})=\mathbb{E}_{\hat{z}}[\boldsymbol{g_
{z}(\theta)}]$ in Lemma 5, we can use the distance $\|\widehat{\nabla}_{\boldsymbol{\theta}}\mathcal{L}_{\hat{z}}(\boldsymbol{
\theta})-\nabla_{\boldsymbol{\theta}}\mathcal{L}_{m}(\boldsymbol{\boldsymbol{
\theta}})\|$ to analyze the the relationship between the true sparse gradient $\nabla_{\boldsymbol{\theta}}\mathcal{L}_{m}(\boldsymbol{\theta})=\boldsymbol{m
}\odot\nabla_{\boldsymbol{\theta}}\mathcal{L}(\boldsymbol{\theta})$ and sparse gradient $\widehat{\nabla}_{\boldsymbol{\theta}}\mathcal{L}_{\hat{z}}(\boldsymbol{\theta})$:

###### Lemma 3.3 .

Let $\mathcal{L}$ be Lipschitz Continuous, we have:

$$ $\|\nabla_{\boldsymbol{\theta}}\mathcal{L}_{m}(\boldsymbol{\boldsymbol{\theta}} )\|^{2}\leq 2\|\widehat{\nabla}_{\theta}\mathcal{L}_{\hat{z}}(\boldsymbol{ \theta})\|^{2}+\frac{\epsilon^{2}L^{2}(l)}{2}(\hat{d}+4)^{3}.$ (6) $$

where $\nabla_{\boldsymbol{\theta}}\mathcal{L}_{m}(\boldsymbol{\theta})=\boldsymbol{m
}\odot\nabla_{\boldsymbol{\theta}}\mathcal{L}(\boldsymbol{\theta})$, $\hat{d}=\sum_{i=1}^{i=d}m_{i}$ is the number of selected parameters in mask $\boldsymbol{m}$, $L(l)$ represents the Lipschitz constant. Finally, we can obtain the convergence rate of Sparse-MeZO.

###### Theorem 3.4 .

Assuming a sequence of generated parameters $\{\boldsymbol{\theta_{t}}\}_{t\geq 0}$ in Sparse-MeZO. We can have:

$$ $\mathbb{E}_{\hat{z},x}[\|\nabla_{\boldsymbol{\theta}}\mathcal{L}_{m}( \boldsymbol{\theta_{T}})\|^{2}]\leq\sigma^{2}$ (7) $$

for any $T=\mathcal{O}(\frac{\hat{d}L}{\sigma^{2}})$

where $L(l)\leq L$ for all $\mathcal{L}(\boldsymbol{\theta_{t}})$. This theorem illustrates that the presence of pronounced sparsity patterns, along with the smoothness of the objective function, can significantly enhance the rate of convergence, potentially achieving a linear acceleration.

###### Assumption 3.1 (Lipschitz Continuous) .

###### Lemma 3.2 .

###### Lemma 3.3 .

###### Theorem 3.4 .

### 3.3 Less Parameters for Better Performance

From the analysis above, we find that Sparse-MeZO can speed up the convergence of zeroth-order optimization.
After that, the main problem is how to determine the sparse mask $\boldsymbol{m}$. Network pruning is a popular method to prune the unimportant parameters and only use these selected important parameters to make the inference process fast and efficient (LeCun et al., ; Hanson & Pratt, ; Hassibi & Stork, ; Ström, ; Han et al., , ).
Specifically, it typically removes connections with small parameters and retains those with large values. Drawing from these recent studies in network pruning, we first explore how the value of parameters affects zeroth-order optimization. This leads us to examine if perturbations to larger parameters are more critical.

Is the perturbation on larger parameters more important than smaller parameters ?

We conduct experiments to examine how parameter values of varying magnitudes impact performance. Specifically, we set the median value of each layer as a threshold to separate parameters into two groups. Then, we only perturb and update the parameters in one of these groups. For instance, we experiment with updating either just the smaller or larger parameters, using the same settings as the standard MeZO, on SuperGLUE. Then, we compared their performance to that of vanilla MeZO.
The results in Figure 3 show that updating only the parameters smaller than the median value leads to better performance on BoolQ, RTE, WIC, and comparable results on SST-2. However, updating only the larger parameters results in significantly worse performance compared to the standard MeZO. This suggests that the key to MeZO’s performance gain lies in the smaller weights.

The reason updating only smaller parameters leads to better performance might be due to the larger weights being well-trained already. Moving them in a random direction is unlikely to enhance performance. On the other hand, since smaller weights are less trained, there’s a higher chance of improvement regardless of the direction of the update.

Figure: Figure 3: Comparing performance across various parameter adjustments. MeZO: the vanilla MeZO approach, MeZO-S: Fine-Tuning smaller parameters with MeZO, MeZO-L: Fine-Tuning larger parameters with MeZO.
Refer to caption: x3.png

Based on the observations from Figure 3, we can create a sparse mask, $\boldsymbol{m}$, determined by parameter magnitudes. Specifically, we employ $\boldsymbol{m}$ to identify and only modify parameters of smaller magnitude. These targeted parameters are defined as $\boldsymbol{\hat{\theta}}=\boldsymbol{m}\odot\boldsymbol{\theta}$. It’s important to note that we still preserve the complete set of parameters, but we apply sparse perturbations and gradient estimations only to the selected ones. This approach allows us to integrate the sparse mask into the standard MeZO method as a straightforward, adaptable tool. Then, we will introduce when and where to calculate the mask.

- $\bullet$
Constant Mask: Setting the Mask Before Training. We compare the parameter values to a threshold for each layer to set the mask before training begins. However, a significant downside of this approach is the extra memory required to store a sparse mask, which is as large as the pre-trained model itself. Our goal is for our method to enhance performance without using more GPU memory or causing extra overhead.
- $\bullet$
Dynamic Mask: Determining Mask at Each Iteration. We can establish a threshold for each layer before training and then generate the mask by comparing parameter values to this threshold during each iteration. This method avoids the necessity of storing a large mask, $\boldsymbol{m}$.

Figure: Algorithm 1 Sparse-MeZO (S-MeZO)

In this paper, we’ll employ a dynamic mask to choose which parameters to perturb and update, addressing the issue of memory constraints.

The pseudo-code is provided in Algorithm 1. This algorithm outlines that we first establish the threshold $h_{i}$ for each layer before beginning training. We then use GetMask (Algorithm 3) to compare each parameter against its threshold $h_{i}$ and create the mask $\boldsymbol{m}$. Following this, we introduce the function PerturbParameters (Algorithm 2) to generate a Gaussian noise sample $\boldsymbol{z}\sim\mathcal{N}(\boldsymbol{0},\boldsymbol{I_{d}})$ and apply the mask $\boldsymbol{m}$ to produce a sparse perturbation $\boldsymbol{\hat{z}}=\boldsymbol{m}\odot\boldsymbol{z}$. With $\boldsymbol{\hat{z}}$, we perturb the current parameters $\boldsymbol{\theta_{t}}$ to get new parameters $\boldsymbol{\theta_{t}+\epsilon\hat{z}}$ and $\boldsymbol{\theta_{t}-\epsilon\hat{z}}$. This allows us to compute two distinct loss values: $l^{+}=\mathcal{L}(\boldsymbol{\theta_{t}}+\epsilon\boldsymbol{\hat{z}})$ and $l^{-}=\mathcal{L}(\boldsymbol{\theta_{t}}-\epsilon\boldsymbol{\hat{z}})$. From these losses, we calculate the estimated sparse gradient $\boldsymbol{g_{m}}(\boldsymbol{\theta_{t}})=\text{proj_grad}*\boldsymbol{\hat
{z}}$, where $\text{proj_grad}=\frac{l^{+}-l^{-}}{2\epsilon}$.
Finally, this gradient can be used with a learning rate $\eta_{t}$ to update $\boldsymbol{\theta_{t}}$.

Figure: Algorithm 2 PerturbParameters

Figure: Algorithm 3 GetMask

### 3.4 Memory-Efficient Implementation of Sparse-MeZO

In this paper, our primary aim is to introduce an efficient method for fine-tuning language models using zeroth-order optimization, enhancing performance on downstream tasks. As outlined in Algorithm 1, our approach involves perturbing the parameters $\boldsymbol{\theta_{t}}$ twice to generate two distinct sets of parameters, $\boldsymbol{\theta_{t}^{\prime}}=\boldsymbol{\theta_{t}}+\epsilon\boldsymbol{z}$ and $\boldsymbol{\theta_{t}^{\prime\prime}}=\boldsymbol{\theta_{t}}-\epsilon
\boldsymbol{z}$. We then use the estimated gradient to update the original parameters $\boldsymbol{\theta_{t}}$. This step typically requires storing two separate sets of parameters, leading to increased memory usage during fine-tuning.

As depicted in Figure 2, recently proposed MeZO, conserves memory by saving random seeds $s$ and using it to resample $z$ for calculating $\boldsymbol{\theta_{t}^{\prime}}$, $\boldsymbol{\theta_{t}^{\prime\prime}}$, and reconstructing $\boldsymbol{\theta_{t}}$ without needing extra memory. However, applying a sparse mask $\boldsymbol{m}$ for calculating sparse perturbation $\boldsymbol{\hat{z}}$ in MeZO poses a memory issue. We cannot simply reconstruct $\boldsymbol{\hat{z}}$ by saving the random seed because the sparse mask, determined by parameter magnitudes, changes when parameters are altered by the perturbation. To address this, we propose potential solutions for the memory issue.

Figure: Figure 4: Efficient Implementation of Sparse-MeZO. (a) represents vanilla implementation of S-MeZO, (b) represents our proposed efficient implementation of S-MeZO.
Refer to caption: x4.png

1-bit Quantization: We can apply 1-bit quantization to store the mask $\boldsymbol{m}$, as it consists solely of 0s and 1s. However, this method still increases memory usage, which isn’t our goal. As a solution, we introduce a novel, memory-saving approach for zeroth-order optimization that calculates the mask $\boldsymbol{m}$ on the fly during the forward pass.

**Table 1: Accuracy of Fine-Tuning LLaMA-7b on SuperGLUE (1,000 examples). ICL: In-Context Learning, FT: full-parameter fine-tuning with Adam, R-MeZO: MeZO with Random Mask.**
| Model | Method | BoolQ | RTE | WIC | MultiRC | SST-2 | COPA | Average |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| LLaMA-7b | Zero-Shot | $65.1$ | $49.5$ | $50.6$ | $55.8$ | $79.7$ | $59.7$ | $60.1$ |
| LLaMA-7b | ICL | $67.4$ | $54.5$ | $52.7$ | $58.7$ | $81.2$ | $84.4$ | $66.5$ |
| LLaMA-7b | FT | $84.5\pm 0.0$ | $83.6\pm 0.9$ | $68.4\pm 1.3$ | $80.2\pm 1.4$ | $95.7\pm 0.3$ | $85.0\pm 0.8$ | $82.9\pm 0.8$ |
| LLaMA-7b | MeZO | $75.9\pm 1.1$ | $71.7\pm 1.5$ | $61.4\pm 1.8$ | $69.8\pm 0.7$ | $94.6\pm 0.3$ | $86.3\pm 0.9$ | $76.6\pm 1.1$ |
| LLaMA-7b | R-MeZO | $76.9\pm 0.7$ | $75.2\pm 1.7$ | $62.1\pm 0.4$ | $68.1\pm 2.0$ | $94.6\pm 0.2$ | $84.3\pm 1.7$ | $76.9\pm 1.1$ |
| LLaMA-7b | S-MeZO | $\bf{80.9\pm 1.6}$ | $\bf{80.7\pm 1.4}$ | $\bf{64.9\pm 1.5}$ | $\bf{73.3\pm 1.2}$ | $\bf{95.0\pm 0.3}$ | $\bf{86.7\pm 0.7}$ | $\bf{80.3\pm 1.2}$ |

Calculating the Mask During the Forward Pass: By computing the mask and perturb parameters in the forward pass, we eliminate the need to store perturbed parameters $\boldsymbol{\theta_{t}^{\prime}}$ and $\boldsymbol{\theta_{t}^{\prime\prime}}$. This means we only have to keep the original parameters $\boldsymbol{\theta_{t}}$ throughout training. We try to compare the difference between vanilla implementation and our proposed efficient implementation in Figure 4. For vanilla implementation, we first need to calculate the perturbed parameters with mask $\boldsymbol{m}$: $\boldsymbol{\theta_{t}^{\prime}}=\boldsymbol{\theta_{t}}+\epsilon\boldsymbol{m
}\odot\boldsymbol{z}$. After that, we can use perturbed parameters $\boldsymbol{\theta_{t}^{\prime}}$ to calculate the loss value $l^{+}$ with the forward process. For example, the output of layer $i$ can be defined as $\boldsymbol{y^{(i)}}=\boldsymbol{\theta_{t}^{\prime(i)}}\boldsymbol{x^{(i)}}+
\boldsymbol{b^{(i)}}$. Noted that we need to save the vanilla parameters $\boldsymbol{\theta_{t}}$ and mask $\boldsymbol{m}$ for vanilla implementation. However, for our proposed efficient implementation, we only need to save vanilla parameters $\boldsymbol{\theta_{t}}$. More specially, we can calculate the mask $\boldsymbol{m^{(i)}}$ of layer $i$ during the forward process and then obtain the output of this layer: $\boldsymbol{y^{(i)}}=(\boldsymbol{\theta_{t}^{(i)}}+\epsilon m(\boldsymbol{
\theta_{t}})\boldsymbol{z^{(i)}})\boldsymbol{x^{(i)}}+\boldsymbol{b^{(i)}}$, where $m(\cdot)$ represents the function GetMask to calculate mask $\boldsymbol{m^{(i)}}$. Then, we can release the memory about mask $\boldsymbol{m}^{(i)}$ and calculate the output and mask of next layer.

## 4 Experiments

Following a similar setting to MeZO, we evaluate the performance of our proposed method on SuperGLUE (Wang et al., ). The experimental results show that our proposed method can achieve better performance while also attaining faster convergence.

### 4.1 Experimental Setting

Datasets. To verify the performance gain of our proposed method, we conduct experiments on various fine-tuning tasks include SST-2 (Socher et al., ), RTE (Dagan et al., ; Haim et al., ; Giampiccolo et al., ; Bentivogli et al., ), BoolQ (Clark et al., ), WIC (Pilehvar & Camacho-Collados, ), MultiRC (Khashabi et al., )) and multi-class task COPA (Roemmele et al., ).

Models. We primarily use pre-trained LLaMA-7b (Touvron et al., ) to evaluate the performance of our proposed method on downstream tasks. To further demonstrate our method’s versatility, we also test it with OPT-13b (Zhang et al., ). Additionally, to examine our method’s scalability, we evaluate it on larger models, such as LLaMA-30b.

Baselines. First, we compare our method to the vanilla MeZO to demonstrate how sparsification enhances MeZO’s convergence speed and overall performance. Additionally, to show that our proposed S-MeZO effectively identifies and modifies crucial parameters, we contrast it with R-MeZO (a version of MeZO applying a random mask to select parameters for optimization). In addition, we also explore the impact of zero-shot optimization on improving a pre-trained language model’s capabilities through experiments with zeroth-shot learning and in-context learning techniques (Brown et al., ). Lastly, to understand the performance gap between zeroth-order and first-order optimization in fine-tuning large language models, we present results from conventional full-parameter fine-tuning (FT) using the Adam optimizer, the most widely used method for such tasks.

Training Procedure. We adopt most of the training hyperparameters from the standard MeZO, including dataset configuration, batch size, training epochs, epsilon value, and task prompts, with the key difference being a higher learning rate for S-MeZO due to updating only a subset of the parameters. The primary goal of our training is the next token prediction. For the dataset, we use MeZO’s approach, randomly selecting 1,000 examples for training and testing the model on another set of 1,000 examples (Zhou et al., ). We perform the experiments using three different seeds and report the average of the outcomes. In addition, the total training steps for LLaMA and OPT is 20,000 and we evaluate its performance on the test dataset every 100 steps.

### 4.2 Performance

To evaluate the performance of our proposed method S-MeZO, we initially tested it on the SuperGLUE benchmark using the LLaMA-7b model. The fine-tuning results, presented in Table 1, reveal that our S-MeZO method outperforms other zero-order (ZO) techniques like MeZO and R-MeZO. For instance, S-MeZO boosts MeZO’s accuracy from $71.7\%$ to $80.7\%$ on RTE ($\uparrow 9\%$) and from $75.9\%$ to $80.9\%$ on BoolQ ($\uparrow 5\%$). Furthermore, all zeroth-order-based methods surpassed the performance of Zero-shot learning and in-context learning, demonstrating that zeroth-order optimization significantly enhances the pre-trained model’s effectiveness on downstream tasks.
Finally, we can find that S-MeZO significantly bridges the performance gap between zero-order and first-order optimization methods.

Figure: (a) RTE
Refer to caption: x5.png

### 4.3 Convergence Rate

To verify that S-MeZO converges faster than MeZO, we carried out multiple experiments for comparison. The accuracy over steps is plotted in Figure 5, which shows that S-MeZO can use fewer steps to achieve a better performance than vanilla MeZO. For example, S-MeZO only needs about 5,000 steps to achieve $70\%$ accuracy but vanilla MeZO needs 17,500 steps. Finally, S-MeZO can achieve about $3.5$x speedup on RTE and $3$x speedup on BoolQ.

### 4.4 Memory Usage

Table 2 shows the memory consumption for MeZO, S-MeZO, and traditional full-parameter fine-tuning of LLaMA-7b. The data reveal that S-MeZO does not require more memory than MeZO and offers a substantial saving of roughly $12$ times less GPU memory compared to full-parameter fine-tuning. For instance, S-MeZO with Efficient Implementation (S-MeZO-EI) cuts down the memory needed from $158.6$ GB for full tuning to just $14.6$ GB on MultiRC task. In addition, S-MeZO with efficient implementation can reduce the memory cost from $28.3$ GB of vanilla S-MeZO to $14.6$ GB across all five tasks, which also illustrates the efficiency of our proposed implementation method: Calculating the Mask During the Forward Pass. Finally, we can only use inference memory cost to fine-tune large language models.

**Table 2: Memory Usage (batch size = 1) of Fine-Tuning LLaMA-7b on SuperGLUE (1,000 examples). EI represents the Efficient Implementation in section 3.4.**
| Method | SST-2 | RTE | BoolQ | WIC | MultiRC |
| --- | --- | --- | --- | --- | --- |
| FT | $114.7$ | $123.7$ | $128.7$ | $115.3$ | $158.6$ |
| MeZO | $14.6$ | $14.6$ | $14.6$ | $14.6$ | $14.6$ |
| S-MeZO | $28.3$ | $28.3$ | $28.3$ | $28.3$ | $28.3$ |
| S-MeZO-EI | $14.6$ | $14.6$ | $14.6$ | $14.6$ | $14.6$ |

### 4.5 Sparse Rate

For S-MeZO, we need to define the sparsity of the pre-trained model before starting to fine-tune it. To analyze the effects of sparsity value on the performance, we conduct experiments with various sparsity values (from $0.0$ to $0.85$). Figure 6 summarizes these experimental results with different sparsity values. We can find that a significant performance gain can be obtained when we use the sparsity value from $0.5$ to $0.8$. In addition, for most tasks, a sparsity value of $0.8$ or $0.75$ usually means a better performance. For example, S-MeZO can improve the accuracy from $71.7\%$ (when $r=0.0$) to $82.3\%$ (when $r=0.8$). It can also obtain a performance gain of $6.6\%$ for WIC (from $75.9\%$ to $82.5\%$).

Figure: Figure 6: The effects of Sparsity for Fine-tuning LLaMA-7b with S-MeZO.
Refer to caption: x8.png

**Table 3: Accuracy of Fine-Tuning LLaMA-7b and LLaMA-30b on SuperGLUE (1,000 examples).**
| Model | Method | BoolQ | RTE | WIC |
| --- | --- | --- | --- | --- |
| LLaMA-7b | MeZO | $75.9$ | $71.7$ | $61.4$ |
| LLaMA-7b | S-MeZO | $80.9$ | $80.7$ | $64.9$ |
| LLaMA-30b | MeZO | $83.8$ | $76.9$ | $63.3$ |
| LLaMA-30b | S-MeZO | 85.7 | 82.1 | 67.3 |

### 4.6 Scalability

In Table 1, we mainly introduce the performance of our methods on LLaMA-7b. A direct question is whether our proposed method can scale to larger language models. Therefore, in this section, we further explore our proposed method S-MeZO on LLaMA-30b. As shown in Table 3, we can see that the a larger model usually can obtain a better fine-tuned performance. For example, the accuracy on RTE with MeZO can be improved from $71.1\%$ on LLaMA-7b to $76.9\%$ on LLaMA-30b. Our method S-MeZO can further improve the performance on RTE to $82.1\%$ on LLaMA-30b. In addition, S-MeZO can further improve the accuracy on BoolQ to $85.7\%$ on LLaMA-30b.

**Table 4: Accuracy of Fine-Tuning OPT on SuperGLUE (1,000 examples). ICL: In-Context Learning, R-MeZO: MeZO with Random Mask.**
| Model | Method | BoolQ | RTE | WIC |
| --- | --- | --- | --- | --- |
| OPT-13b | Zero Shot | $59.0$ | $59.6$ | $55.0$ |
| OPT-13b | ICL | $66.9$ | $62.1$ | $50.5$ |
| OPT-13b | MeZO | $72.1$ | $75.5$ | $62.2$ |
| OPT-13b | R-MeZO | $72.3$ | $75.2$ | $61.7$ |
| OPT-13b | S-MeZO | $\bf{73.8}$ | $\bf{77.6}$ | $\bf{63.7}$ |

### 4.7 The Analysis about Efficient Implementation

In section 3.4,we present the efficient implementation of S-MeZO, which enables our proposed method to require only the inference memory cost for fine-tuning large language models. To analyze the actual GPU memory usage during the training process, we provide these results in Table 2. We can find that S-MeZO needs the same GPU memory for all five tasks, which can also save about $50\%$ memory compared to sparse-mezo. That also illustrates the efficiency of our proposed efficient implementation.

## 5 Conclusion

In this paper, we propose a novel memory-efficient zeroth-order fine-tuning method Sparse-MeZO, which can use a similar memory cost to the inference process. We evaluate the performance of fine-tuning LLaMA and OPT with Sparse-MeZO on SuperGULE benchmark and the experimental results illustrate that Sparse-MeZO can achieve a higher accuracy and faster convergence. Finally, we can fine-tune LLaMA-30b in a single A100 GPU. However, there is still a performance gap between our proposed method Sparse-MeZO and first-order fine-tuning methods. We plan to address these limitations and enhance Sparse-MeZO’s capabilities in our future research.

## Appendix A The Prompts in LLaMA and OPT

**Table 5: The prompts of the datasets we used in our LLaMA experiments.**
| Dataset | Type | Prompt |
| --- | --- | --- |
| SST-2 | cls. | {premise} |
|  |  | Does this mean that “{hypothesis}” is true? Yes or No? |
|  |  | Yes/No |
| RTE | cls. | Suppose “{premise}” Can we infer that “{hypothesis}”? Yes, No, or Maybe? |
|  |  | Yes/No/Maybe |
| BoolQ | cls. | {passage} {question} ? |
|  |  | Yes/No |
| WIC | cls. | Does the word “{word}” have the same meaning in these two sentences? Yes, No? |
|  |  | {sent1} |
|  |  | {sent2} |
|  |  | Yes/No |
| MultiRC | cls. | {paragraph} |
|  |  | Question: {question} |
|  |  | I found this answer “{answer}”. Is that correct? Yes or No? |
|  |  | Yes/No |
| COPA | mch. | {premise} so/because {candidate} |

## Appendix B Hyperparameters

We will introduce the hyperparameters searching grids in Table 6, which can help people to reproduce our results.

**Table 6: The hyperparameter searching grids for LLaMA-7b experiments.**
| Experiment | Hyperparameters | Values |
| --- | --- | --- |
| MeZO | Batch size | $16$ |
|  | Learning rate | $\{5\mathrm{e}{-7},1\mathrm{e}{-6},2\mathrm{e}{-6}\}$ |
|  | $\epsilon$ | $1\mathrm{e}{-3}$ |
| MeZO-Random | Batch size | $16$ |
|  | Learning rate | $\{1\mathrm{e}{-6},2\mathrm{e}{-6},3\mathrm{e}{-6},4\mathrm{e}{-6},5\mathrm{e}{<br>-6}\}$ |
|  | $\epsilon$ | $1\mathrm{e}{-3}$ |
| S-MeZO | Batch size | $16$ |
|  | Learning rate | $\{1\mathrm{e}{-6},2\mathrm{e}{-6},3\mathrm{e}{-6},4\mathrm{e}{-6},5\mathrm{e}{<br>-6}\}$ |
|  | $\epsilon$ | $1\mathrm{e}{-3}$ |
| FT with Adam | Batch size | $8$ |
|  | Learning Rates | $\{1\mathrm{e}{-5},5\mathrm{e}{-5},8\mathrm{e}{-5}\}$ |

We then introduce the sparsity of each task in SuperGULU when we fine-tune LLaMA-7b.

**Table 7: Sparsity in SuperGULU when we fine-tune LLaMA-7b.**
| Method | SST-2 | RTE | BoolQ | WIC | MultiRC |
| --- | --- | --- | --- | --- | --- |
| Sparse MeZO | $0.70$ | $0.75$ | $0.80$ | $0.80$ | $0.80$ |

## Appendix C The Proof of Lemma 3.2

Let $\mathcal{L}_{z}(\theta)$ be the expectation of $\mathcal{L}(\theta+\epsilon m\odot z)$:

$$ $\displaystyle\mathcal{L}_{\hat{z}}(\theta):$ $\displaystyle=\mathbb{E}_{z}[\mathcal{L}(\theta+\epsilon m\odot z)]$ (8) $\displaystyle=\mathbb{E}_{\hat{z}}[\mathcal{L}(\theta+\epsilon\hat{z})]$ $$

We can obtain the Lemma:

$$ $\displaystyle\widehat{\nabla}_{\theta}\mathcal{L}_{\hat{z}}(\theta)$ $\displaystyle=m\odot\nabla_{\theta}\mathcal{L}_{\hat{z}}(\theta)$ (9) $\displaystyle=m\odot\mathbb{E}_{z}[\nabla_{\theta}\mathcal{L}(\theta+\epsilon m \odot z)]$ $\displaystyle=\mathbb{E}_{z}[\frac{\mathcal{L}(\theta+\epsilon m\odot z)- \mathcal{L}(\theta-\epsilon m\odot z)}{2\epsilon}m\odot z]$ $\displaystyle=\mathbb{E}_{\hat{z}}[\frac{\mathcal{L}(\theta+\epsilon\hat{z})- \mathcal{L}(\theta-\epsilon\hat{z})}{2\epsilon}\hat{z}]$ $$

Proof:

$$ $\displaystyle\widehat{\nabla}_{\theta}\mathcal{L}_{\hat{z}}(\theta)$ $\displaystyle=\widehat{\nabla}_{\theta}\mathbb{E}_{\hat{z}}[\mathcal{L}(\theta +\epsilon\hat{z})]$ (10) $\displaystyle=\widehat{\nabla}_{\theta}\int_{\hat{z}}\text{pdf}_{\hat{z}}(z) \mathcal{L}(\theta+\epsilon z)dz$ $\displaystyle=m\odot\nabla_{\theta}\int_{\hat{z}}\text{pdf}_{\hat{z}}(z) \mathcal{L}(\theta+\epsilon z)dz$ $\displaystyle=m\odot\int_{\hat{z}}\nabla_{\theta}\text{pdf}_{\hat{z}}(z) \mathcal{L}(\theta+\epsilon z)dz$ $\displaystyle=\frac{1}{k}m\odot\int_{\hat{z}}\nabla_{\theta}e^{-\frac{1}{2}\|z \|^{2}}\mathcal{L}(\theta+\epsilon z)dz$ $\displaystyle=\frac{1}{k}m\odot\int_{\hat{y}}\nabla_{\theta}e^{-\frac{1}{2}\| \frac{y-\theta}{\epsilon}\|^{2}}\mathcal{L}(y)\frac{1}{\epsilon^{n}}dy$ $\displaystyle=\frac{1}{k}m\odot\int_{\hat{y}}\frac{y-\theta}{\epsilon^{2}}e^{- \frac{1}{2\epsilon^{2}}\|y-\theta\|^{2}}\mathcal{L}(y)\frac{1}{\epsilon^{n}}dy$ $\displaystyle=\frac{1}{k}m\odot\int_{\hat{z}}\frac{z}{\epsilon}e^{-\frac{1}{2} \|z\|^{2}}\mathcal{L}(\theta+\epsilon z)dz$ $\displaystyle=m\odot\int_{\hat{z}}\text{pdf}_{\hat{z}}(z)\mathcal{L}(\theta+ \epsilon z)\frac{z}{\epsilon}dz$ $\displaystyle=\mathbb{E}_{\hat{z}}[m\odot\frac{\mathcal{L}(\theta+\epsilon\hat {z})}{\epsilon}\hat{z}]$ $\displaystyle=\mathbb{E}_{\hat{z}}[\frac{\mathcal{L}(\theta+\epsilon\hat{z})}{ \epsilon}\hat{z}]$ $$

where we can define $y=\theta+\epsilon z$, $\hat{y}=\theta+\epsilon m\odot z$, $k=\sqrt{(2\pi)^{\hat{d}}}$ and $\hat{d}$ is the number of 1 in $\boldsymbol{m}$.

Therefore, we can obtain the gradient $\nabla_{\theta}\mathcal{L}_{m}(\theta)$ is equal to $\mathbb{E}_{\hat{z}}[\frac{\mathcal{L}(\theta+\epsilon\hat{z})}{\epsilon}\hat{
z}]$.

In addition, we will prove $\mathbb{E}_{\hat{z}}[\frac{\mathcal{L}(\theta+\epsilon\hat{z})}{\epsilon}\hat{
z}]$ is also equal to $\mathbb{E}_{\hat{z}}[\frac{\mathcal{L}(\theta+\epsilon\hat{z})-L(\theta)}{
\epsilon}\hat{z}]$:

$$ $\displaystyle\mathbb{E}_{\hat{z}}[\frac{\mathcal{L}(\theta+\epsilon\hat{z})- \mathcal{L}(\theta)}{\epsilon}\hat{z}]$ (11) $\displaystyle=\frac{1}{k}\int_{\hat{z}}\frac{\mathcal{L}(\theta+\epsilon z)- \mathcal{L}(\theta)}{\epsilon}ze^{-\frac{1}{2}\|z\|^{2}}dz$ $\displaystyle=\frac{1}{k}\int_{\hat{\epsilon}}\frac{\mathcal{L}(\theta+ \epsilon z)}{\epsilon}ze^{-\frac{1}{2}\|z\|^{2}}dz-\frac{\mathcal{L}(\theta)}{ \epsilon k}\int_{\hat{z}}ze^{-\frac{1}{2}\|z\|^{2}}dz$ $\displaystyle=\mathbb{E}_{\hat{z}}[\frac{\mathcal{L}(\theta+\epsilon\hat{z})}{ \epsilon}\hat{z}]$ $$

After that, we can get the relationship between $\mathbb{E}_{\hat{z}}[\frac{\mathcal{L}(\theta)-\mathcal{L}(\theta-\epsilon\hat
{z})}{\epsilon}\hat{z}]$ and $\mathbb{E}_{\hat{z}}[\frac{\mathcal{L}(\theta+\epsilon\hat{z})}{\epsilon}\hat{
z}]$:

$$ $\displaystyle\mathbb{E}_{\hat{z}}[\frac{\mathcal{L}(\theta)-\mathcal{L}(\theta -\epsilon\hat{z})}{\epsilon}\hat{z}]$ $\displaystyle=\mathbb{E}_{\hat{z}}[\frac{\mathcal{L}(\theta+\epsilon(-\hat{z}) )-\mathcal{L}(\theta)}{\epsilon}(-\hat{z})]$ (12) $\displaystyle=\mathbb{E}_{\hat{z}}[\frac{\mathcal{L}(\theta+\epsilon\hat{z}- \mathcal{L}(\theta))}{\epsilon}\hat{z}]$ $\displaystyle=\mathbb{E}_{\hat{z}}[\frac{\mathcal{L}(\theta+\epsilon\hat{z})}{ \epsilon}\hat{z}].$ $$

Based on the Equation 11 and Equation 12, we can obtain:

$$ $\displaystyle\mathbb{E}_{\hat{z}}[\frac{\mathcal{L}(\theta+\epsilon\hat{z})- \mathcal{L}(\theta-\epsilon\hat{z})}{2\epsilon}\hat{z}]$ (13) $\displaystyle=\frac{1}{2}(\mathbb{E}_{\hat{z}}[\frac{\mathcal{L}(\theta+ \epsilon\hat{z})}{\epsilon}\hat{z}-\frac{\mathcal{L}(\theta)}{\epsilon}\hat{z} +\frac{\mathcal{L}(\theta)}{\epsilon}\hat{z}-\frac{\mathcal{L}(\theta-\epsilon \hat{z})}{\epsilon}\hat{z}])$ $\displaystyle=\frac{1}{2}(\mathbb{E}_{\hat{z}}[\frac{\mathcal{L}(\theta+ \epsilon\hat{z})-\mathcal{L}(\theta)}{\epsilon}\hat{z}]+\mathbb{E}_{\hat{z}}[ \frac{\mathcal{L}(\theta)-\mathcal{L}(\theta-\epsilon\hat{z})}{\epsilon}\hat{z }])$ $\displaystyle=\frac{1}{2}(\mathbb{E}_{\hat{z}}[\frac{\mathcal{L}(\theta+ \epsilon\hat{z})}{\epsilon}\hat{z}]+\mathbb{E}_{\hat{z}}[\frac{\mathcal{L}( \theta+\epsilon\hat{z})}{\epsilon}\hat{z}])$ $\displaystyle=\mathbb{E}_{\hat{z}}[\frac{\mathcal{L}(\theta+\epsilon\hat{z})}{ \epsilon}\hat{z}]$ $\displaystyle=\widehat{\nabla}_{\theta}\mathcal{L}_{\hat{z}}(\theta)$ $$

Finally, we can obtain the relationship between $\mathbb{E}_{\hat{z}}[\frac{\mathcal{L}(\theta+\epsilon\hat{z})-\mathcal{L}(
\theta-\epsilon\hat{z})}{2\epsilon}\hat{z}]$ and $\widehat{\nabla}_{\theta}\mathcal{L}_{\hat{z}}(\theta)$ and finish the proof.

## Appendix D The Proof of Lemma 3.3

$$ $\|\nabla_{\boldsymbol{\theta}}\mathcal{L}_{m}(\boldsymbol{\boldsymbol{\theta}} )\|^{2}\leq 2\|\widehat{\nabla}_{\boldsymbol{\theta}}\mathcal{L}_{\hat{z}}( \boldsymbol{\theta})\|^{2}+\frac{\epsilon^{2}L^{2}(l)}{2}(\hat{d}+4)^{3}.$ (14) $$

Proof:

We can first define the distance between $\widehat{\nabla}_{\boldsymbol{\boldsymbol{\boldsymbol{\theta}}}}\mathcal{L}_{
\hat{z}}(\boldsymbol{\boldsymbol{\theta}})=\mathbb{E}_{\hat{z}}[\boldsymbol{g_
{\hat{z}}(\theta)}]$ and sparse FO gradient $\nabla\mathcal{L}_{m}(\theta)$ as:

$$ $\displaystyle\|\widehat{\nabla}_{\theta}\mathcal{L}_{\hat{z}}(\theta)-\nabla_{ \theta}\mathcal{L}_{m}(\theta)\|$ (15) $\displaystyle=\|\frac{1}{k}\int_{z}(\frac{\mathcal{L}(\theta+\epsilon z)- \mathcal{L}(\theta-\epsilon z)}{2\epsilon}-\langle\nabla_{\theta}\mathcal{L}_{ m}(\theta),z\rangle)ze^{-\frac{1}{2}\|z\|^{2}}d\hat{z}\|$ $\displaystyle=\|\frac{1}{k}\int_{z}(\frac{\mathcal{L}(\theta+\epsilon z)- \mathcal{L}(\theta)}{\epsilon}-\langle m\odot\nabla_{\theta}\mathcal{L}(\theta ),z\rangle)ze^{-\frac{1}{2}\|z\|^{2}}d\hat{z}\|$ $\displaystyle\leq\frac{1}{k\epsilon}\int_{z}|\mathcal{L}(\theta+\epsilon z)- \mathcal{L}(\theta)-\epsilon\langle\nabla_{\theta}\mathcal{L}(\theta),\epsilon \rangle|\|m\odot z\|e^{-\frac{1}{2}\|z\|^{2}}d\hat{z}$ $\displaystyle\leq\frac{\epsilon L(l)}{2k}\int_{\epsilon}\|z\|^{2}\|m\odot z\|e ^{-\frac{1}{2}\|z\|^{2}}d\hat{z}$ $\displaystyle=\frac{\epsilon L(l)}{2}\mathbb{E}_{\hat{z}}[\|\hat{z}\|^{3}]$ $\displaystyle\leq\frac{\epsilon L(l)}{2}(\hat{d}+3)^{\frac{3}{2}}$ $$

where $\hat{d}$ is the number of selected parameters with mask $\boldsymbol{m}$. In addition, $\|\boldsymbol{a}+\boldsymbol{b}\|^{2}\leq 2\|\boldsymbol{a}\|^{2}+2\|
\boldsymbol{b}\|^{2}$, we can define $\boldsymbol{a}=\boldsymbol{a}-\boldsymbol{b}$ and obtain that $\|\boldsymbol{a}\|^{2}\leq 2\|\boldsymbol{a}-\boldsymbol{b}\|^{2}+2\|
\boldsymbol{b}\|^{2}$. Let $\boldsymbol{a}=\nabla_{\theta}\mathcal{L}_{m}(\theta)$ and $\boldsymbol{b}=\widehat{\nabla}_{\theta}\mathcal{L}_{\hat{z}}(\theta)$, we can obtain:

$$ $\displaystyle\|\nabla_{\theta}\mathcal{L}_{m}(\theta)\|^{2}$ $\displaystyle\leq 2\|\widehat{\nabla}_{\theta}\mathcal{L}_{\hat{z}}(\theta)- \nabla_{\theta}\mathcal{L}_{m}(\theta)\|^{2}+2\|\widehat{\nabla}_{\theta} \mathcal{L}_{\hat{z}}(\theta)\|^{2}$ (16) $\displaystyle\leq\frac{\epsilon^{2}L^{2}(l)}{2}(\hat{d}+3)^{3}+2\|\widehat{ \nabla}_{\theta}\mathcal{L}_{\hat{z}}(\theta)\|^{2}$ $\displaystyle\leq\frac{\epsilon^{2}L^{2}(l)}{2}(\hat{d}+4)^{3}+2\|\widehat{ \nabla}_{\theta}\mathcal{L}_{\hat{z}}(\theta)\|^{2}$ $$

## Appendix E The Proof of Theorem 3.4

Proof:

$$ $\displaystyle\mathcal{L}_{\hat{z}}(\theta)-\mathcal{L}(\theta)$ $\displaystyle=\mathbb{E}_{\hat{z}}[\mathcal{L}(\theta+\epsilon\hat{z})- \mathcal{L}(\theta)]$ (17) $\displaystyle=\mathbb{E}_{\hat{z}}[\mathcal{L}(\theta+\epsilon\hat{z})- \mathcal{L}(\theta)-\epsilon\langle\nabla\mathcal{L}(\theta),\hat{z}\rangle]$ $\displaystyle=\frac{1}{k}\int_{\hat{z}}[\mathcal{L}(\theta+\epsilon z)- \mathcal{L}(\theta)-\epsilon\langle\nabla\mathcal{L}(\theta),z\rangle]e^{- \frac{1}{2}\|z\|^{2}}dz$ $\displaystyle\leq\frac{1}{k}\int_{\hat{z}}\frac{\epsilon^{2}L(l)}{2}\|z\|^{2}e ^{-\frac{1}{2}\|z\|^{2}}dz$ $\displaystyle=\frac{\epsilon^{2}L(l)}{2}\mathbb{E}_{\hat{z}}[\|\hat{z}\|^{2}]$ $\displaystyle\leq\frac{\epsilon^{2}L(l)}{2}\hat{d}$ $$

The first inequality holds because Lipschitz Continuous: $|\mathcal{L}(\theta^{\prime})-\mathcal{L}(\theta)-\langle\nabla\mathcal{L}(
\theta),\theta^{\prime}-\theta\rangle|\leq\frac{L(l)}{2}\|\theta^{\prime}-
\theta\|^{2}$, where $\theta^{\prime}=\theta+\epsilon z$. The second inequality holds because $\mathbb{E}_{\hat{z}}[\|\hat{z}\|^{2}]=\hat{d}$, where $\hat{d}$ is the number of $1$ in mask $\boldsymbol{m}$.

$$ $\displaystyle[(\mathcal{L}_{\hat{z}}(\theta)-\mathcal{L}(\theta))-(\mathcal{L} _{\hat{z}}(\theta+\epsilon\hat{z})-\mathcal{L}(\theta+\epsilon\hat{z}))]^{2}$ (18) $\displaystyle\leq 2[\mathcal{L}_{\hat{z}}(\theta)-\mathcal{L}(\theta)]^{2}+2[ \mathcal{L}_{\hat{z}}(\theta+\epsilon\hat{z})-\mathcal{L}(\theta+\epsilon\hat{ z})]^{2}$ $\displaystyle\leq\frac{\epsilon^{4}L^{2}(l)}{2}\hat{d}^{2}+\frac{\epsilon^{4}L ^{2}(l)}{2}\hat{d}^{2}$ $\displaystyle=\epsilon^{4}L^{2}(l)\hat{d}^{2}$ $$

The first inequality is due to $\|\boldsymbol{a}+\boldsymbol{b}\|^{2}\leq 2\|\boldsymbol{a}\|^{2}+2\|
\boldsymbol{b}\|^{2}$, where $\boldsymbol{a}=\mathcal{L}_{\hat{z}}(\theta)-\mathcal{L}(\theta),\boldsymbol{b
}=\mathcal{L}_{\hat{z}}(\theta+\epsilon\hat{z})-\mathcal{L}(\theta+\epsilon
\hat{z})$. The second inequality is due to the Equation 17.

$$ $\displaystyle{}^{2}$ $\displaystyle\leq 2[\mathcal{L}_{\hat{z}}(\theta+\epsilon\hat{z})-\mathcal{L}_ {\hat{z}}(\theta)-\epsilon\langle\widehat{\nabla}_{\theta}\mathcal{L}_{\hat{z} }(\theta),\hat{z}\rangle]^{2}+2[\epsilon\langle\widehat{\nabla}_{\theta} \mathcal{L}_{\hat{z}}(\theta),\hat{z}\rangle]^{2}$ (19) $\displaystyle\leq\frac{\epsilon^{4}L^{2}(l)}{2}\|\hat{z}\|^{4}+2\epsilon^{2} \langle\widehat{\nabla}_{\theta}\mathcal{L}_{\hat{z}}(\theta),\hat{z}\rangle^{2}$ $\displaystyle\leq\frac{\epsilon^{4}L^{2}(l)}{2}\|\hat{z}\|^{4}+2\epsilon^{2}\| \widehat{\nabla}_{\theta}\mathcal{L}_{\hat{z}}(\theta)\|^{2}\|\hat{z}\|^{2}$ $$

The first inequality is due to $\|\boldsymbol{a}+\boldsymbol{b}\|^{2}\leq 2\|\boldsymbol{a}\|^{2}+2\|
\boldsymbol{b}\|^{2}$. The second inequality holds because Lipschitz Continuous: $|\mathcal{L}(\theta^{\prime})-\mathcal{L}(\theta)-\langle\nabla\mathcal{L}(
\theta),\theta^{\prime}-\theta\rangle|\leq\frac{L(l)}{2}\|\theta^{\prime}-
\theta\|^{2}$, where $\theta^{\prime}=\theta+\epsilon\hat{z}$.

$$ $\displaystyle[\mathcal{L}(\theta+\epsilon\hat{z})-\mathcal{L}(\theta)]^{2}$ (20) $\displaystyle\leq 2[(\mathcal{L}_{\hat{z}}(\theta)-\mathcal{L}(\theta))-( \mathcal{L}_{\hat{z}}(\theta+\epsilon\hat{z})-\mathcal{L}(\theta+\epsilon\hat{ z}))]^{2}+2[\mathcal{L}_{\hat{z}}(\theta+\epsilon\hat{z})-\mathcal{L}_{\hat{z} }(\theta)]^{2}$ $\displaystyle\leq 2\epsilon^{4}L^{2}(l)\hat{d}^{2}+\epsilon^{4}L^{2}(l)\|\hat{ z}\|^{4}+4\epsilon^{2}\|\widehat{\nabla}_{\theta}\mathcal{L}_{\hat{z}}(\theta) \|^{2}\|\hat{z}\|^{2}$ $$

The first inequality is due to $\|\boldsymbol{a}+\boldsymbol{b}\|^{2}\leq 2\|\boldsymbol{a}\|^{2}+2\|
\boldsymbol{b}\|^{2}$. The last inequality holds because Equation 18 and Equation 19.

$$ $\displaystyle\mathbb{E}_{z,x}[\|g_{\hat{z}}(\theta)\|^{2}]$ $\displaystyle=\mathbb{E}_{\hat{z}}[\|\frac{\mathcal{L}(\theta+\epsilon\hat{z}) -\mathcal{L}(\theta-\epsilon\hat{z})}{2\epsilon}\hat{z}\|^{2}]$ (21) $\displaystyle=\mathbb{E}_{\hat{z}}[\|\frac{\mathcal{L}(\theta+\epsilon\hat{z}) -\mathcal{L}(\theta)}{2\epsilon}\hat{z}+\frac{\mathcal{L}(\theta)-\mathcal{L}( \theta-\epsilon\hat{z})}{2\epsilon}\hat{z}\|^{2}]$ $\displaystyle\leq\mathbb{E}_{\hat{z}}[2\|\frac{\mathcal{L}(\theta+\epsilon\hat {z})-\mathcal{L}(\theta)}{2\epsilon}\hat{z}\|^{2}+2\|\frac{\mathcal{L}(\theta) -\mathcal{L}(\theta-\epsilon\hat{z})}{2\epsilon}\hat{z}\|^{2}]$ $\displaystyle=\mathbb{E}_{\hat{z}}[\frac{1}{2\epsilon^{2}}[\mathcal{L}(\theta+ \epsilon\hat{z})-\mathcal{L}(\theta)]^{2}\cdot\|\hat{z}\|^{2}+\frac{1}{2 \epsilon^{2}}[\mathcal{L}(\theta)-\mathcal{L}(\theta-\epsilon\hat{z})]^{2} \cdot\|\hat{z}\|^{2}]$ $\displaystyle\leq\mathbb{E}_{\hat{z}}[2\epsilon^{2}L^{2}(l)\hat{d}^{2}\|\hat{z }\|^{2}+\epsilon^{2}L^{2}(l)\|\hat{z}\|^{6}+4\|\widehat{\nabla}\mathcal{L}_{ \hat{z}}(\theta)\|^{2}\|\hat{z}\|^{4}]$ $\displaystyle\leq 2\epsilon^{2}L^{2}(l)\hat{d}^{3}+\epsilon^{2}L^{2}(l)(\hat{d }+6)^{3}+4(\hat{d}+4)^{2}\|\widehat{\nabla}\mathcal{L}_{\hat{z}}(\theta)\|^{2}$ $\displaystyle\leq 3\epsilon^{2}L^{2}(l)(\hat{d}+4)^{3}+4(\hat{d}+4)^{2}\| \widehat{\nabla}\mathcal{L}_{\hat{z}}(\theta)\|^{2}$ $$

The first inequality holds because $\|\boldsymbol{a}+\boldsymbol{b}\|^{2}\leq 2\|\boldsymbol{a}\|^{2}+2\|
\boldsymbol{b}\|^{2}$, where $\boldsymbol{a}=\frac{\mathcal{L}(\theta+\epsilon\hat{z})-\mathcal{L}(\theta)}{
2\epsilon}\hat{z}$, $\boldsymbol{b}=\frac{\mathcal{L}(\theta)-\mathcal{L}(\theta-\epsilon\hat{z})}{
2\epsilon}\hat{z}$. The second inequality is due to the Equation 20. The third inequality holds because $\mathbb{E}_{\hat{z}}[\|\hat{z}\|^{2}]=\hat{d}$, $\mathbb{E}_{\hat{z}}[\|\hat{z}\|^{p}]\leq(\hat{d}+p)^{\frac{p}{2}}$ for $p\geq 2$. The last inequality holds because $2\hat{d}^{3}+(\hat{d}+6)^{3}\leq 3(\hat{d}+4)^{3}$.

Based on the assumption about Lipschitz Continuous, we can obtain: $|\mathcal{L}(\theta_{t+1})-\mathcal{L}(\theta_{t})-\langle\nabla\mathcal{L}(
\theta_{t}),\theta_{t+1}-\theta_{t}\rangle|\leq\frac{L(l)}{2}\|\theta_{t+1}-
\theta_{t}\|^{2}$.

Then, we can obtain:

$$ $\displaystyle\mathcal{L}_{\hat{z}}(\theta_{t+1})-\mathcal{L}_{\hat{z}}(\theta_ {t})-\langle\widehat{\nabla}\mathcal{L}_{\hat{z}}(\theta_{t}),\theta_{t+1}- \theta_{t}\rangle\leq|\mathcal{L}_{\hat{z}}(\theta_{t+1})-\mathcal{L}_{\hat{z} }(\theta_{t})-\langle\widehat{\nabla}\mathcal{L}_{\hat{z}}(\theta_{t}),\theta_ {t+1}-\theta_{t}\rangle|\leq\frac{L(l)}{2}\|\theta_{t+1}-\theta_{t}\|^{2}$ (22) $$

Based on the equation, we can follow the update rule: $\theta_{t+1}=\theta_{t}-\eta_{t}g_{\hat{z}}(\theta_{t})$ and we can find:

$$ $\displaystyle\mathcal{L}_{\hat{z}}(\theta_{t+1})$ $\displaystyle\leq\mathcal{L}_{\hat{z}}(\theta_{t})+\langle\widehat{\nabla} \mathcal{L}_{\hat{z}}(\theta_{t}),\theta_{t+1}-\theta_{t}\rangle+\frac{L(l)}{2 }\|\theta_{t}-\theta_{t+1}\|^{2}$ (23) $\displaystyle=\mathcal{L}_{\hat{z}}(\theta_{t})-\eta_{t}\langle\widehat{\nabla }\mathcal{L}_{\hat{z}}(\theta_{t}),g_{\hat{z}}(\theta_{t})\rangle+\frac{(\eta_ {t})^{2}L(l)}{2}\|g_{\hat{z}}(\theta_{t})\|^{2}$ $$

where $\eta_{t}$ represents the learning rate at step $t$. Then, we can take the expectation of Equation 23 for $\hat{z}$ and input $x$:

$$ $\displaystyle\mathbb{E}_{\hat{z},x}[\mathcal{L}_{\hat{z}}(\theta_{t+1})]$ (24) $\displaystyle\leq\mathbb{E}_{\hat{z},x}[\mathcal{L}_{\hat{z}}(\theta_{t})]- \eta_{t}\mathbb{E}_{\hat{z},x}[\|\widehat{\nabla}\mathcal{L}_{\hat{z}}(\theta_ {t})\|^{2}]+\frac{(\eta_{t})^{2}L(l_{z})}{2}\mathbb{E}_{\hat{z},x}[\|g_{z}( \theta_{t})\|^{2}]$ $\displaystyle\leq\mathbb{E}_{\hat{z},x}[\mathcal{L}_{\hat{z}}(\theta_{t})]- \eta_{t}\mathbb{E}_{\hat{z},x}[\|\widehat{\nabla}\mathcal{L}_{\hat{z}}(\theta_ {t})\|^{2}]+\frac{(\eta_{t})^{2}L(l)}{2}(4(\hat{d}_{t}+4)\mathbb{E}_{\hat{z},x }[\|\widehat{\nabla}\mathcal{L}_{\hat{z}}(\theta_{t})\|^{2}]+3\epsilon^{2}L^{2 }(l)(\hat{d}_{t}+4)^{3})$ $$

The first inequality is due to the Equation 9 and Equation 23. The second inequality holds because Equation 21 provides the result about $\mathbb{E}_{\hat{z},x}[\|g_{z}(\theta_{t})\|^{2}$.

Then, we can select learning rate $\eta_{t}=\frac{1}{4(\hat{d}_{t}+4)L(l)}$ and obtain:

$$ $\mathbb{E}_{\hat{z},x}[\mathcal{L}_{\hat{z}}(\theta_{t+1})]\leq\mathbb{E}_{ \hat{z},x}[\mathcal{L}_{\hat{z}}(\theta_{t})]-\frac{1}{8(\hat{d}_{t}+4)L(l)} \mathbb{E}_{\hat{z},x}[\|\widehat{\nabla}\mathcal{L}_{\hat{z}}(\theta_{t})\|^{ 2}]+\frac{3\epsilon^{2}}{32}L(l)(\hat{d}_{t}+4)$ (25) $$

Then, taking the sum of Equation 25 over the index from $T+1$ to 0 0, we can have that :

$$ $\mathbb{E}_{\hat{z},x}[\|\widehat{\nabla}\mathcal{L}_{\hat{z}}(\theta_{T})\|^{ 2}]\leq 8(\hat{d}+4)L[\frac{\mathcal{L}_{\hat{z}}(\theta_{0})-\mathcal{L}_{ \hat{z}}^{*}}{T+1}+\frac{3\epsilon^{2}}{32}L(\hat{d}+4)]$ (26) $$

where $L(l)\leq L$ for all $\mathcal{L}(\boldsymbol{\theta_{t}})$. Thus, based on Lemma 6, we can have:

$$ $\displaystyle\mathbb{E}_{\hat{z},x}[\|\nabla\mathcal{L}_{m}(\theta_{T})\|^{2}]$ $\displaystyle\leq\frac{\epsilon^{2}L^{2}}{2}(\hat{d}+4)^{3}+2\mathbb{E}_{\hat{ z},x}[\|\widehat{\nabla}\mathcal{L}_{\hat{z}}(\theta_{T})\|^{2}]$ (27) $\displaystyle\leq 16(\hat{d}+4)L\frac{\mathcal{L}_{\hat{z}}(\theta_{0})- \mathcal{L}_{\hat{z}}^{*}}{T+1}+\frac{\epsilon^{2}L^{2}}{2}(\hat{d}+4)^{2}( \hat{d}+\frac{11}{2})$ $$

The second inequality is due to the Equation 26. To obtain $\sigma$-accurate solution: $\mathbb{E}_{\hat{z},x}[\|\nabla\mathcal{L}_{m}(\theta_{T})\|^{2}]\leq\sigma^{2}$, we can define $\epsilon=\Omega(\frac{\sigma}{\hat{d}^{\frac{3}{2}}L})$.

$$ $\displaystyle 16(\hat{d}+4)L\frac{\mathcal{L}_{\hat{z}}(\theta_{0})-\mathcal{L }_{\hat{z}}^{*}}{T+1}+\mathcal{O}(\epsilon^{2}L^{2}\hat{d}^{3})$ $\displaystyle=16(\hat{d}+4)L\frac{\mathcal{L}_{\hat{z}}(\theta_{0}-\mathcal{L} _{\hat{z}}^{*})}{T+1}+\mathcal{O}(\sigma^{2})$ (28) $\displaystyle T$ $\displaystyle=\mathcal{O}(\frac{\hat{d}L}{\sigma^{2}})$ $$

Finally, we can finish the proof of the theorem. This theorem illustrates that the presence of pronounced sparsity patterns, along with the smoothness of the objective function, can significantly enhance the rate of convergence, potentially achieving a linear acceleration.

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 5
presentation: 4
questions: 1. Could the authors clarify the computational cost associated with generating the dynamic mask at each iteration and how it compares to MeZO in practice? such as a wall-clock time wise comparison similar to Figure 1

2. How sensitive is Sparse-MeZO’s performance to the choice of layer-wise sparsity threshold, and what considerations guided the threshold selection?

3. Can you provide detailed threshold hyperparameters for reproducibility?
soundness: 3
strengths: 1. Incorporating sparsity and MeZO is an interesting direction for performance improvement. 

2. The paper offers good empirical validation, showing clear improvements over baseline methods across a range of fine-tuning tasks.

3. The paper is structured logically, making it easy to follow the motivation and methodology
summary: The paper proposes Sparse-MeZO, a memory-efficient zeroth-order optimization (ZO) technique for fine-tuning LLM by selectively optimizing a subset of parameters, rather than all. Based on MeZO, Sparse-MeZO introduces a sparse mask that targets smaller, noise-resilient weights, thereby mitigating gradient estimation noise, improving convergence speed, and reducing memory usage. Key results demonstrate Sparse-MeZO’s efficiency, achieving faster convergence and improved performance over MeZO, with the ability to fine-tune models like LLaMA-30b on limited hardware resources (e.g., a single A100 GPU).
weaknesses: 1. The paper lacks details on how thresholds for selecting small-magnitude weights are determined, how they vary across tasks, and whether they require re-tuning for different settings.

2. It’s not clear if masking and selecting small-magnitude weights specifically benefit zeroth-order optimization more than generic fine-tuning (first-order methods). Since subset selection often improves generalization, it would be needed to evaluate this effect.

3. There is no specific numbers regarding the computational overhead introduced by dynamic masking and threshold calculation.

4. The motivation for why zeroth-order optimization particularly benefits from small-magnitude weights lacks enough theoretical or empirical support. An ablation study showing this effect more clearly would strengthen the argument.

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 2
questions: - Mask Updates per Iteration: Are the masks updated in each iteration? The concept of dynamic masking appears similar to Dynamic Sparse Training (DST), where the frequency of mask updates is a critical factor affecting performance [1,2]

- Can Sparse MeZO be integrated with LoRA, similar to how MeZO has been combined with LoRA in previous work?

-  When claiming the 3.5× speed-up in Figure 1, are both methods using the same learning rate, or does the speed-up depend on differing learning rates?

-  Determining the masks is a key component of this work, and the authors use a simple magnitude-based method. Do you believe that employing more advanced methods like Wanda [3] or SparseGPT [4] could further enhance performance?

[1] Do we actually need dense over-parameterization? in-time over-parameterization in sparse training, ICML 2021

[2] Rigging the Lottery: Making All Tickets Winners， ICML 2020

[3] A Simple and Effective Pruning Approach for Large Language Models, ICLR2024

[4] SparseGPT: Massive Language Models Can Be Accurately Pruned in One-Shot
soundness: 3
strengths: - Interesting Idea: Proposes Sparse MeZO, selectively applying zeroth-order optimization to improve memory efficiency.- 
- Good Performance: Achieves significant gains in accuracy and convergence speed over existing methods.
- Simple but Efficient Algorithm: Offers a straightforward yet effective approach for fine-tuning large models on limited hardware.
summary: The paper introduces Sparse MeZO, a memory-efficient optimization method for fine-tuning large language models. By selectively applying zeroth-order updates to a carefully chosen subset of parameters, it reduces estimation errors inherent in zeroth-order methods. This approach improves both performance and convergence speed without increasing memory usage, enabling efficient fine-tuning of large models like LLaMA-30b on limited hardware resources.
weaknesses: Some details need to be clarify, as explained in Questions

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: * Given a fixed training budget does SparseMeZO outperform fine-tuning with adapters (+ gradient accumulation)?

* I think it would be insightful to add a comparison between the gradient updates with standard SGD optimization and SparseMeZO on a smaller scale model (say of size ~1B parameters).
soundness: 2
strengths: * The fact that instability is caused by the noisy optimization of weights with large magnitude appears to be a useful practical insight. 

* SparseMeZO consistently improves upon dense MeZO optimization in different setups. 

* The proposed memory-efficient implementation of S-MeZO incurs almost zero memory overheads relative to original MeZO and allows tuning large models on a single GPU.
summary: This work proposes zeroth order memory-efficient optimization with sparse updates of model parameters for fine-tuning. Sparse updates are shown to facilitate better convergence than vanilla MeZO. The introduced method is evaluated on several tasks from SuperGLUE and several model architectures - Llama-1, Mistral, OPT.
weaknesses: * The overall contribution is incremental as the work adds additional sparsification steps on top of the MeZO optimization algorithm. 

* Evaluation setup is outdated (as for Fall 2024). Llama-1 and Mistral-v0.1 were released more than a year ago, and current models (Llama-3, Llama-3.1, Qwen-2.5, gemma-2) have advanced dramatically in terms of quality since then. In addition, I would suggest testing the approach on a more challenging and representative fine-tuning setup, such as instruction tuning on Alpaca/OASST or any other instruction fine-tuning dataset to fully appreciate the efficacy of SparseMeZO. 

* Some important baselines are not taken into consideration. PEFT adapters (LoRA, DoRA) allow for significant memory reduction on optimizer states. Memory footprint on activations and gradients is not very large for relatively short sequences (which is the case for SuperGLUE tasks) and small batches. In addition, it can be further decreased by gradient accumulation at the cost of additional compute. I would suggest adding comparison with these options and reporting memory usage in Table 3.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 4
questions: See above.
soundness: 2
strengths: The paper is well motivated, the approach is novel and the results are convincing. I particularly enjoyed the writing style and clear structure of the paper, it was mostly easy to follow, enjoyable to read and the reader was not left with many questions. The proposed method is interesting, however there are open questions, which I will discuss below. Overall, I think this is a good paper that however needs to resolve several questions before being suited for publication.
summary: The paper builds on the Memory-efficient Zeroth-order optimizer (MeZO) and introduces a new method, called Sparse MeZO. While MeZO is successful in training a neural network only with forward passes and by estimating the batch gradient using a finite difference scheme, there is still a significant gap between MeZO and common first-order methods. The authors aim to narrow this gap by updating only a sparse subset of the model parameters in each iteration. In particular, they propose to only update the smallest weights in each layer. They provide efficient implementations of doing so and show that Sparse MeZO performs better than multiple baselines, including MeZO.
weaknesses: While I appreciate research in this important field, I have several concerns regarding soundness, clarity and contribution of the work, which I explain in detail below. I hope that these remarks are helpful in improving the work and I am happy to discuss my evaluation.

### Soundness
- In Figure 2b, you compare the potential loss increase on a different batch than the one used by MeZO, you then conclude that "zeroth-order gradient estimates suffer from overfitting or noise" (Line 189). I think to really conclude this, you would need to compare the MeZO gradient on a batch with the actual gradient on the same batch. The very same statement could be true for e.g. SGD, especially if the batch size is small and the variance of the gradient is high. I would greatly appreciate to see this figure instead or in addition to Figure 2b (potentially in appendix).
- In Line 197, you write that "all parameters share the same value". What is meant by this? The formula you say that "is used to estimate the gradient" only gives the magnitude of the multiplier, the actual gradient vector is given by $z$, where definitely not all elements are the same. Apart from that being a false premise, you then conclude that not all parameters are optimized in the true gradient direction, which I found a surprising statement. As far as I understand it, MeZO or variants randomly sample a direction along which the gradient is estimated, constrained to that particular direction. It will with very high probability not be anywhere close to the true gradient direction, by virtue of the fact that you are random sampling the directional derivative.  Please clarify this, maybe I am mistaken here. 
- In the same paragraph, you somehow derive that "small weights are less impacted by noise corruption and can generalize better" - it is not clear to my why you derive this conclusion, is it just by the fact that you empirically find this to work better? What is the mathematical intuition or even justification for this statement? I think to resolve this soundness issue, you have to answer to very important questions: First of all, why does it make sense to use less weights in the first place? For First-order Optimization this would be absurd, what is the difference here? Why should this work better than using all weights? And secondly, why would exactly the smallest weights have the highest impact? I could potentially follow this argument if you were actually sparsifying the computed gradient, then one could argue with Taylor approximation or the like, but you are not doing this, you are instead sparsifying the weights. Why?
- Figure 4 highlights the effect of sparsity on the final performance. It is not clear why this curve looks like this and why that would make sense. The authors should clarify why they think that sparsity in the weights during gradient estimation helps (apart from providing empirical evidence that is performs better). What sparsity 70% performs bad, then sparsity 75% performs pretty good, and sparsity 80% performs worse again seems very arbitrary and not intuitive. In the worst case, this is just a noise artifact.


### Clarity and Experimental Validation
- Line 123: What is "SPSA"? Is this defined anywhere in the paper?
- In Line 260, you explain how you determine the mask. Why don't you set the threshold dynamically in each iteration? It's efficient and seems like the obvious thing to do.
- In Table 1, you report values for MeZO-LoRA. While I might have a vague idea of what that could be, this seems to be nowhere explained nor defined. Is this known from the literature or was this just missed?
- Are the numbers in the tables reported with respect to the same overall runtime or number of iterations? If MeZO variants are faster than e.g. FT, it would be nice to see how they compare on the same runtime.
- In the appendix (e.g. F), you refer to Lemma 3.2. What exactly is this, where is this defined? The same holds for the following appendix sections, I think these are referring to theorems/lemmata that do not exist or are not defined with the same numbering style. Please clarify.

### Contribution
- As outlined above, I think the contribution of the work is a major issue here. I fully acknowledge that Sparse MeZO achieves better results and is in that sense a meaningful contribution. However, the derivation of the method seems to be at least somewhat vague, it lacks justification apart from achieving better results. There is not much insight to gain since the paper lacks mathematical justifications, or at least intuitions. The derivation from gradient noise seems to be not very rigorous, at least to me. I hope that the authors can convince me otherwise.


### Minor Remarks
- Section 4.1 uses \citet everywhere where I think \citep is intended, this hinders readability.
- Line 483: There is a typo, I guess it should be "on a single A100 GPU".

Meta-review: This paper proposes Sparse-MeZO, a variant of a recent zeroth-order optimizer (MeZO) used in fine tuning LLMs. Different from MeZO, the authors' method involves only performing zeroth order optimization over a subset of the learnable parameters. Since reducing parameter dimensionality is dramatically beneficial to all zeroth order optimization routines, this results in faster convergence and superior performance. Strengths obviously include the good empirical evaluation with clear improvements over MeZO. While I think this paper is well on its way to acceptance given the strong empirical results, there were significant concerns about (1) the clarity of certain key details that I think *still* remain unsatisfactorily addressed even in the updated paper (see below), and (2) key details above the algorithm.

Decision: Reject