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

Thanks for your insightful comments, we carefully address your concerns below.

W1:

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).

Thank you for this insightful suggestion. We agree that a more rigorous comparison is needed. Our new analysis compares SGD and MeZO using the same batch, as shown in Figure 5. The results show MeZO has a higher probability of loss increment (0.4-0.6) compared to SGD (0.2-0.3). While both methods can increase loss on different batches, MeZO exhibits this behavior more frequently.

W2:

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.

Thank you for your thoughtful review. Let me clarify the gradient estimation in MeZO.

You are correct that MeZO randomly samples a direction to estimate the gradient. To be more precise, following MeZO's definition and Equation (2), the estimated gradient is: $g = p_g \cdot z$, where $p_g = (L(\theta + \epsilon z) - L(\theta - \epsilon z))/(2\epsilon)$ is a scalar value that's shared across all parameters, while $z$ is the randomly sampled direction vector.

We apologize if our wording was unclear. When we mentioned "shared value", we were referring to the scalar multiplier $p_g$, not the final gradient estimates. Since this scalar multiplier is applied uniformly across dimensions, the resulting update may not align well with the true gradient direction, which could be considered a limitation of the method.

Thanks for your constructive comments, we carefully address your concerns below.

W1:

The overall contribution is incremental as the work adds additional sparsification steps on top of the MeZO optimization algorithm.

We appreciate this feedback and would like to highlight several substantial contributions of our work beyond MeZO:

• Technical Innovation: We identified and addressed a fundamental limitation in ZO optimization - that gradient noise impacts large weights more severely than small weights. This insight led to our selective parameter updating approach.
• Performance Gains: Our method achieves substantial improvements over MeZO. For example, (1) 9% absolute accuracy improvement on RTE, (2) 3.5x faster convergence, (3) Comparable performance to full fine-tuning while maintaining MeZO's memory efficiency.
• Novel Implementation: We developed a memory-optimized implementation for sparse masking that maintains inference-level memory consumption, enabling fine-tuning of LLaMA-30b on a single A100 GPU.
• Theoretical Analysis: We provided rigorous theoretical analysis proving faster convergence rates with sparsification, establishing that our method isn't just an empirical improvement but has sound theoretical foundations.

W2:

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.

Thank you for this constructive comment. We have extended our evaluation to more recent models, including LLaMA2-7b and LLaMA3-8b. As shown in the table, Sparse-MeZO demonstrates consistent performance improvements on these newer models across RTE, BoolQ, and WIC tasks. For the final version, we will (1) Include comprehensive results on these newer models, (2) Add instruction tuning experiments on Alpaca/OASST datasets, (3) Expand evaluation to additional contemporary models as suggested (Qwen-2.5, Gemma-2).

W3:

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.

Thank you for these valuable suggestions about additional baselines and memory analysis. We've expanded our evaluation:

• Performance Comparison: Our experiments with LoRA on LLaMA-7b show: While LoRA achieves better performance, Sparse-MeZO significantly narrows the gap between zeroth-order and first-order methods.

• Memory Analysis: With batch size = 1, Sparse-MeZO reduces memory usage by 25% compared to LoRA. We will include comprehensive memory comparisons with LoRA, DoRA, and gradient accumulation approaches in Table 3 of the final version.

Q1:

Given a fixed training budget does SparseMeZO outperform fine-tuning with adapters (+ gradient accumulation)?

Thank you for this insightful comment. While Sparse-MeZO achieves our primary goal of memory efficiency, we acknowledge that LoRA with gradient accumulation would likely achieve faster convergence given the same training budget. Our method's key advantage is enabling fine-tuning of very large models (e.g., LLaMA-30b) on a single GPU where memory constraints are the primary bottleneck, even if this comes at the cost of longer training time. For example, Sparse MeZO can save 25% memory compared with LoRA (with gradient accumulation).

Thanks for your constructive feedback, we carefully address your concerns below.

Q1:

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]

Thanks for your insightful comment. In our implementation, while the sparsity threshold remains fixed, the mask updates dynamically each iteration using: mask = (torch.abs(weight) < threshold).float().

We appreciate the connection to Dynamic Sparse Training (DST) [1,2] and provide the discussion in the revision. While DST's findings about update frequency are valuable, we chose an update frequency of 1 to maintain memory efficiency. This allows us to compute masks on-the-fly without storing them, using only the model weights and threshold values. Our experiments comparing update frequencies on LLaMA-7b show that frequency=1 outperforms frequency=5, as shown in the table. We will incorporate a discussion of DST and mask update frequency in our revision.

Q2:

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

Thanks for your constructive comment. Sparse-MeZO can indeed be integrated with LoRA by applying the masking mechanism to the weight matrices: $\theta = \theta_0 + m \odot (AB)$. We also conduct experiments to compare these methods:

Our experimental results show that Sparse-MeZO with LoRA achieves comparable performance to vanilla Sparse-MeZO while outperforming MeZO+LoRA.

Q3:

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?

We use a larger learning rate for sparse-mezo. For example, sparse mezo can use a 5x learning rate when only update 20% or 25% parameters on RTE task. As demonstrated in Figure 2(a), vanilla MeZO diverges with larger learning rates, while Sparse-MeZO remains stable. This ability to use larger learning rates is a key factor enabling the 3.5x speedup.

Q4:

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?

Thank you for this insightful comment about advanced pruning methods. We also provide the discussion in the revision.

• Our experimental results demonstrate that updating large-magnitude weights, which is a key component of many advanced pruning methods, actually leads to worse performance in zeroth-order optimization compared to focusing on small-magnitude weights.
• While Wanda and SparseGPT offer sophisticated pruning strategies, implementing their scoring mechanisms during training would add computational overhead compared to our simple magnitude-based threshold approach.
• Our method achieves strong performance while maintaining memory efficiency, suggesting that more complex pruning strategies may not be necessary for effective zeroth-order optimization.

Thanks for your constructive feedback, we carefully address your concerns below.

W1 & Q2:

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. How sensitive is Sparse-MeZO’s performance to the choice of layer-wise sparsity threshold, and what considerations guided the threshold selection?

Thank you very much for your valuable comment. We determine thresholds using a principled sparsity-based approach. Specifically, we use a percentile-based method where the threshold is set based on a target sparsity level. For example, with 80% sparsity, we sort the weight values of each layer and set the threshold at the 80th percentile. Importantly, this threshold is determined once before training begins and remains fixed throughout the optimization process.

Our empirical results in Figure 4 demonstrate that a sparsity level of 75% consistently yields strong performance across diverse tasks (RTE, BoolQ, and WIC), highlighting the generalizability of our threshold selection method. Furthermore, we observe robust performance gains compared to vanilla MeZO across a range of sparsity values (70-80%), indicating that our method is not sensitive to the exact threshold choice. This stability across different tasks eliminates the need for task-specific threshold tuning.

To further illustrate with a concrete example: if we have a layer with 1000 weights and set 75% sparsity, we would sort all weight magnitudes and set the threshold at the 750th value (75th percentile). Any weight with magnitude below this threshold is considered small. This percentile-based approach automatically adapts to each layer's weight distribution, eliminating the need for manual tuning across different tasks or model sizes. As demonstrated in Figure 4, this single approach with 75% sparsity works robustly across RTE, BoolQ, and WIC tasks, yielding consistent improvements over vanilla MeZO.

W2:

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.

Thank you for this comment about distinguishing the specific benefits of small-magnitude weight selection across optimization methods. Our additional experiments compare this selection strategy's impact on both zeroth-order (MeZO) and first-order (Adam) optimization, with results shown in the table. While weight selection provides some generalization benefits for first-order methods, our findings demonstrate that the improvements are markedly larger for zeroth-order optimization, suggesting this approach particularly addresses zeroth-order optimization challenges.

The results in the table demonstrate three key findings:

1. For zeroth-order optimization, Sparse MeZO improves MeZO's average performance from 74.7 to 79.0, showing clear benefits of weight masking.
2. For first-order optimization with Adam using a fixed learning rate (1e-4), masking actually reduces performance. With 0.7 sparsity, average performance drops from 82.5 to 80.9, and with 0.8 sparsity, it further decreases to 80.1.
3. While increasing the learning rate (from 1e-4 to 2e-4) with masking in first-order optimization does improve results (from 80.9 to 82.6 at 0.7 sparsity, and 80.1 to 82.7 at 0.8 sparsity), the gains are minimal (0.1-0.2 improvement). This confirms that masking provides significantly greater benefits for zeroth-order optimization.