Bilevel ZOFO: Bridging Parameter-Efficient and Zeroth-Order Techniques for Efficient LLM Fine-Tuning and Meta-Training

We thank the reviewer for their detailed review of our method. We will address the concerns raise below:

1- Notations and typo

Thank you for the comment. We have bolded the vectors in the revised version, marked in blue for easy reference. Please refer to the updated version. We have fixed the typo you mentioned in the revised version.

2- Why Tune both the PEFT and base model?

While PEFT methods effectively reduce training costs and memory usage, they do not always achieve the same level of task-specific performance as full model fine-tuning (as noted in studies such as Hu et al., 2022; Li & Liang, 2021; Zaken et al., 2022). Therefore, in our approach, we aim to improve performance by tuning the full model in addition to the PEFT module.

Directly tuning the full model with first-order methods, however, is computationally expensive. To address this, we considered using a zeroth-order (ZO) method, which allows tuning without direct gradient calculations. But ZO tuning alone typically relies on hard prompts, which can be suboptimal for fine-tuning. For single-task scenarios, tuning the PEFT module alongside ZO tuning of the full model helps overcome these limitations, enhancing performance compared to using either ZO or PEFT alone. Our experimental results (Table 1 and Table 2) confirm this hypothesis.

3- Computational Cost of First-Order and Efficiency Comparison

Please refer to the global responses on memory profiling and efficiency analysis, as well as Figure 5 and Table 6 in the revised version of our paper. Figure 5 demonstrates that computing first-order (FO) gradients for PEFT parameters is not significantly demanding in terms of memory usage. In fact, in some cases, zero-order (ZO) methods require more memory than FO PEFT. Also, our method is competitive to MeZO in terms of memory requirements, while addressing its limitations and converging much faster. Our method outperform MeZO by a large gap in half the number of iterations.

We calculate the gradient of PEFT parameters because using ZO for both the full model parameters and the PEFT parameters results in higher approximation errors. Our goal is to achieve optimal performance within the memory budget. Calculating the exact gradient of PEFT parameters is feasible since it is significantly less expensive than FO full fine-tuning. Therefore, we choose First-Order PEFT fine-tuning to enhance performance.

Regardless, if the memory constraints are extreme, our bilevel framework can incorporate doing ZO finetuning for updating the PEFT parameters too.

We greatly appreciate the time and effort you put into providing constructive feedback. Following your suggestions, we have carefully revised the manuscript to address your concerns. The updated version incorporates all recommended changes and improvements. We kindly request that you review the revised submission and our rebuttal and consider providing additional feedback. We would be especially grateful if you might consider adjusting your score in light of the improvements made. Thank you once again for your valuable input and for contributing to the quality of our work.

We appreciate the reviewer's continued feedback. We address your concerns as follows:

Regarding the memory consumption of FO vs. ZO

The reason FO sometimes uses less memory than ZO is related to the nature of PEFT (Parameter-Efficient Fine-Tuning). In PEFT, the base model remains frozen, so we do not need to store the gradients of the base model during FO computations. On the other hand, ZO, while avoiding backpropagation, still requires storing the gradient estimates for the base model, leading to comparable or even higher memory consumption in certain scenarios. Importantly, the higher memory in certain scenarios depends on the number of tokens in a task sample on average. For instance, in datasets with a large number of tokens per task (e.g., MultiRC, used for profiling MEZO in their paper), MEZO performs better in terms of memory. However, this is not always the case, as the memory consumption is task-dependent.

We would also like to address the reviewer’s concern, ‘Even if this is correct.’ Our claims are fully supported by evidence, and the code for MEZO and ZO-BENCH is publicly available. To further clarify, the memory profiling results can be independently verified, and we encourage you to explore the experiments to confirm these findings.

On the memory consumption being a merit

The observation that the memory consumption of our proposed algorithm is equivalent to the maximum of ZO and FO is, in fact, a strength of our method rather than a drawback. Our systematic design effectively leverages this characteristic to improve performance, enhancing either FO or ZO without exceeding their individual memory consumption limits. The algorithm’s performance is clearly very effective in largely improving the current methods for both single task and multi-task learning.

Hard prompt issue in ZO is problematic

Contrary to the reviewer’s comment that "the hard prompt issue itself is not so problematic," we would like to remind the reviewer that previous work has demonstrated that improper hard prompts can significantly impact performance of ZO. This is why we proposed tuning the prompts instead of relying on hard prompts. Below is a detailed comparison between results with and without prompts:

First, similar to Table 5 of the MeZO paper, we experimented with different choices of hard prompts. Table 1 below shows the MeZO results for tuning OPT 1.3b on SST-2 and COPA.

Tab. 1: MeZO's Prompt Sensitivity

We compare this with our method on OPT1.3b. Tab. 2 below summarizes our results:

Tab. 2: Bilevel-ZOFO's Prompt Sensitivity

We can see that our method effectively mitigates the sensitivity of MeZO to hard prompts. The difference between the results with and without a simple hard prompt in our experiment is much smaller than MeZO's.

We must clarify that FO is not intended to mitigate the hard prompt issue; the bilevel structure is. The bilevel framework we propose can flexibly use either ZO or FO for PEFT, depending on memory restrictions. This flexibility is one of the key merits of our contribution. As we explained previously, we choose to use FO because it provides better accuracy within the memory budget.

To improve clarity, we will add the above explanations to the main text. Based on these points, we respectfully encourage you to consider this contribution and reconsider your score. If you have any additional questions or require further clarification, we are happy to provide further explanations.

We sincerely thank the reviewer for their feedback. Below we address the concerns raised:

1- Mitigating the Sensitivity of ZO

We thank the reviewer for mentioning this. We now have added more experiments to show that our method mitigates the sensitivity of MeZO to hard prompts.

1.1- Sensitivity to different hard prompts

First, similar to Table 5 of the MeZO paper, we experiment with different choices of hard prompts. Tab. 1 below shows the MeZO results for tuning Opt 1.3b on SST2 and COPA.

Tab. 1: MeZO's Prompt Sensitivity

We compare this with our method on OPT1.3b. Tab. 2 below summarizes our results:

Tab. 2: Bilevel-ZOFO's Prompt Sensitivity

We can see that our method effectively mitigates the sensitivity of MeZO to hard prompts. The difference between the results with and without a simple hard prompt in our experiment is much less than MeZO's

1.2- The effectiveness is not the result of more parameters (Experiments of PEFT followed by ZO finetuning)

To also validate that the improved results are not because of tuning more parameters, we conducted an experiment on COPA using OPT1.3B.

Two-stage tuning: First, we performed first-order prompt tuning for a fixed number of steps (same as the number of lower-level updates in bilevel-ZOFO), followed by additional tuning using ZO for the same number of iterations as the upper level updates in bilevel-ZOFO:

As shown in our experiments (see Tab. 3 below), the two-stage method does not does not work at all because the accuracy after a first stage is 74.33, while at the end of the two-stage setting the accuracy becomes 51.66. Our bilevel structure makes the trained prompts dynamically optimal for the full ZO fine-tuning.

Tab. 3:

3- Unclear MeZO baseline

We sincerely thank the reviewer for this valuable feedback and apologize for the confusion. MeZO, introduced by Malladi et al., is the first method to apply ZO fine-tuning to LLMs. More specifically, MeZO is replacing the gradient in the model with the approximation (Equation 7 in our paper) and then doing SGD or Adam. We have revised the paper to clearly define and explain what MeZO refers to, please see the first paragraph in Section 2.1 as well as the second paraphraph in Section 4 in the revised version. Moreover, we have now included all details of the experimental setting in Appendix B.

Once again, we appreciate your constructive and thorough comments. We hope our explanations have addressed your concerns.

We greatly appreciate the time and effort you put into providing constructive feedback. Following your suggestions, we have carefully revised the manuscript to address your concerns. The updated version incorporates all recommended changes and improvements. We kindly request that you review the revised submission and our rebuttal and consider providing additional feedback. We would be especially grateful if you might consider adjusting your score in light of the improvements made. Thank you once again for your valuable input and for contributing to the quality of our work.

Thanks the authors for the additional experiments. However I still have some questions on the additional experiment results. For Tab. 3, it's not reasonable that the performance of stage 2 is much lower than the result from stage 1. The author should conduct basic hyperparameter tuning to make sure the result is reasonable or explain the counter-intuitive result.

Now I understand the MeZO baseline, but it seems the existence of MeZO weakens the novelty of this work. (This should be a minor point though).

Therefore, I'll keep my rating.

Thank you for your thoughtful feedback and for taking the time to share your insights. We're glad our explanation addressed your concerns and appreciate your consideration of our additional experiments.

Regarding your question about the choice of evaluating under a low-resource setting with only 1000 training examples: This work builds upon MeZO, which served as the foundation and baseline for our work. For single-task experiments, we aligned the number of data points with MeZO (See Appendix E of their paper) to ensure a direct and fair comparison, facilitating straightforward interpretation of improvements for readers. Similarly, in the multi-task setting, we based our sample distribution on MetaICL, which serves as another critical baseline for our work.

By adhering to these established benchmarks, we aim to provide a clear and consistent basis for comparison between our work and the respective baselines.

We will also take your suggestion of including more comprehensive experiments into account for our revision to further strengthen the work.

Thank you again for your valuable comments and suggestions. We hope you would consider increasing your score, as it will greatly support the impact of this work. If you have further concerns, we will be more than happy to address them.

Thanks the authors for additional experiments. I'm more convinced about the reported results. I suggest the authors to include more comprehensive experiments in the revision. I've also lower my confidence as I'm not familiar with the MeZO baseline.

And another question about the experiments. Why do the authors choose to evaluate under a low-resource setting where the training set only contains 1000 examples? I understand this is an additional question not included in the initial review, but I feel this is important for evaluating this work.

Thank you for your continued feedback. We appreciate the opportunity to provide further clarification.

The training loss curves for both stages of a two-stage approach and our bilevel framework are provided at the following links for reference:

• First-stage PEFT training loss: link
• Second-stage MeZO training loss: link
• Our bilevel method training loss: link

When running MeZO in the second stage, the training loss exhibits oscillations and does not show improvement within 500–1000 iterations. This behavior is consistent with findings in the original MeZO paper, which notes that MeZO typically requires much longer to converge—on the order of 100k iterations. The oscillatory behavior observed within the shorter training duration is not surprising due to gradient approximation errors.

In contrast, our bilevel method effectively addresses the issues of MeZO and demonstrates improved performance over both MeZO and the PEFT baseline, even with the same number of ZO iterations.

We hope this additional clarification and the provided evidence address your concerns.

Thank you for your continued feedback and for highlighting the importance of including additional experiments to further support our claims.

1. Regarding Tab. 3 and the Counter-Intuitive Result in Stage 2

To address your concern, we conducted additional hyperparameter tuning to ensure our reported results accurately reflect that two-stage training is not effective. The updated results are derived from the same range of hyperparameters used in Bilevel-ZOFO as reported in Section B.1.3 of the appendix. We initially tested using only the best hyperparameters identified from the results in Table 1 of the paper. However, following the reviewer's suggestion, we re-ran the experiments with all hyperparameters. These results represent the mean of three seeds:

The results indicate the following:

1. Even with extensive hyperparameter tuning, the second stage does not improve the results achieved after the first stage and is highly likely to decrease performance.
2. Our method, however, improves performance when using the same number of steps in the upper and lower levels, respectively.

The observed performance drop after the second stage is indeed counter-intuitive at first glance. However, it is a limitation of MeZO as it approximates gradients. While further fine-tuning intuitively should improve performance, the inherent noise in gradient approximation can lead to suboptimal updates. This observation is consistent with the fact that MeZO typically requires a significant number of iterations to converge. This is a key contribution of our work: Our approach addresses MeZO's challenges, such as sensitivity to hard prompts and long convergence times, while outperforming both MeZO and PEFT and maintaining similar memory efficiency.

The intuition behind why our method is effective in enhancing both MeZO's full-model tuning and PEFT is in the nested bilevel structure. This structure encodes more information (as reflected in the training method) from the prompt tuning stage than only treating it as a first stage, thereby providing better guidance for MeZO.

We will also include an expanded explanation and additional experimental results in the manuscript to clarify the difference between our method and the two-stage strategy.

2. Regarding the Novelty in the Presence of MeZO:

We would like to clarify that the novelty and contributions of our work build upon and extend MeZO.

• More Robust: We address MeZO’s limitations by mitigating its sensitivity to hard prompts, as demonstrated in Tables 1 and 2 in the previous rebuttal.

• More Efficient and Effective Multitask Training Method: Compared to previous methods for multitask training, our approach is lightweight while achieving better performance. These improvements highlight our method's unique contribution to enhancing efficiency and scalability.

• Faster Convergence: Our approach significantly reduces the number of iterations required by MeZO, enabling faster and more reliable convergence.

• Optimizing the Merits of MeZO and PEFT: The proposed nested bilevel structure enhances the performance of both MeZO’s full-model tuning and PEFT. Our method outperforms MeZO and PEFT while maintaining similar memory efficiency.

We hope these additional experiments and explanations regarding our contributions address your concerns. If they resolve your concerns, we kindly ask for a reconsideration of the rating. If any additional issues remain, we would be happy to provide further clarifications.

We thank the reviewer for thier valuable comments and feedback. We address their concerns below:

1- Efficiency of the proposed bilevel method

We now compare the memory requirement and the wall-clock-time efficiency of our method to the baseline. (see the global response and Table 6 and Figure 5 now added to the appendix of the revised version of the paper), our method is as computationally efficient as the baselines.

We also mentioned in the method section (lines 216-221) that traditional bilevel optimization methods tend to be computationally intensive, primarily due to the need for multiple Hessian-vector products. However, recent advances have introduced a new approach to solving bilevel problems that requires only first-order information, significantly reducing computational costs. Our method aligns with this newer approach and we further reduce the computational cost using the zeroth order approximation of the base model. This allows us to achieve more efficient computation.

2- Can the two-level optimizations share the same tuning dataset?

In this work, we consider two settings: single-task fine-tuning and multitask learning. For single-task fine-tuning, PEFT models (such as prompts) can be treated as hyperparameters for the full model. Typically, we use a separate dataset for hyperparameter tuning to avoid overfitting the data used in the upper level.

In multitask learning, especially in few-shot learning, the lower level focuses on training task-specific parameters, while the upper level trains a meta-model that generalizes well across all tasks. After meta-training, the PEFT model can be further trained on a few new data points when applied to an unseen task. Our bilevel training process is designed to simulate this application process, which is why the datasets used in the lower and upper levels are distinct.

3- Any new technical contributions From Eq. 3 to Eq. 7?

The technical challenge in this work lies in incorporating zeroth-order gradient estimation into a penalty-based bilevel optimization framework. Most existing studies on bilevel optimization focus on providing theoretical convergence guarantees for first-order methods. Given the nested structure of the upper and lower-level problems, it is not immediately clear whether a zeroth-order approach would achieve both theoretical and empirical convergence.

In this work, we demonstrate that our proposed method not only performs effectively but also has theoretical convergence guarantees. For LLM applications, the contribution of this work is a more effective framework to advancing zeroth-order fine-tuning and enhancing multitask learning.

4- Lack of efficiency comparisons

Please refer to the global responses for the efficiency comparison. We have now added memory profiling and wall-clock-time analysis.

5- Comparisons with improvements over MeZO

We thank the reviewer for pointing this out. We are aware of other methods to improve MeZO, as noted in the related work section (please let us know if we missed any). Since our proposed structure is orthogonal the choice of the zeroth-order tuning methods, our method can integrate other zeroth-order variations in the upper-level steps. However, due to time and resource limitations—and, more importantly, to maintain focus on the bilevel structure and its unique multitask learning capabilities—we used the classic MeZO as a baseline to demonstrate the effectiveness of our approach.

6- Why not consider tuning soft prompts in this paper?

Thank you for your question. In our lower level, one of our PEFT models involves tuning the soft prompt. Unlike [1], we use a first-order method for tuning prompts. We didn't use a zeroth-order approach because soft prompts have significantly fewer trainable parameters than the full pre-trained model, making it computationally inexpensive to calculate their exact gradient as opposed to tuning the whole model. This allows for more accurate estimation compared to zeroth-order methods.

The setting in [1] differs from ours. [1] addresses a black-box scenario where access to the full model parameters is restricted, whereas we consider a white-box scenario where we have full access to the model parameters, enabling zeroth order fine-tuning within pre-trained LLMs. We have now included [1] in the second paragraph on page 1 when discussing soft prompt tuning.

7- Question about the theoretical analysis

Thank you for your comment. We included the theoretical convergence analysis because bilevel optimization is challenging to train due to its nested structure between the upper and lower level problems. Given this complexity, it is not immediately obvious that the new method incorporated with zeroth order approximation will work effectively in practice. Our theoretical analysis ensures that the proposed method will not diverge when applied in experiments, providing a foundation of reliability for real applications.

We thank Reviewer tCaA for their valuable feedback. Below we address the concerns raised by the reviewer.

1- Larger Scale Models

We thank the reviewer for highlighting this concern.

While we recognize that scaling to much larger models for single-task experiments, such as OPT13B, would further strengthen our work, our testing was limited by resource constraints. As with most academic labs, we can afford and focus on researching 1B-7B parameter models, which should be sufficient for testing and validating research ideas.

Our experiments have already demonstrated the effectiveness of our methods on models ranging from 1B to 7B parameters. Similarly, MeZO has shown strong performance on 1B-7B models, as reported in the original paper. Our findings are consistent with MeZO, and Bilevel-ZOFO consistently outperforms other baseline methods across this model range. Based on these results, we believe our method would yield similar improvements on even larger models.

Nonetheless, we are actively working to extend our method to larger-scale models and will share those results as soon as they are available.

2- Memory Profiling and Wall Clock Analysis

Please see the global comment. We have now included the memory profiling and wall-clock-time analysis of our method compared to the baselines.

3- Quantitative Evaluations for Convergence

We now added Figure 4 in the revised version to show that Bilevel-ZOFO converges smoothly. Figure 4 presents the training loss for the lower-level objective of the bilevel framework with Lora as the PEFT model. As shown, consistent with the guarantees provided by our theoretical analysis, Bilevel-ZOFO converges.

4-Experimental Setting and Hyperparameters

We now have provided all details on hyperparameters and our computational environment in Appendix B of the revised paper and include them here for your reference:

Training Data: For each task, we use 1,000 training examples, 500 validation samples, and 1,000 test samples. In Bilevel-ZOFO experiments, the training data is split into upper-level and lower-level subsets using a 1:2 ratio.

Hyperparameter Search:

• For FO PEFT and MeZO, we explored learning rates from ${1e-2, 1e-3, 1e-4, 1e-5, 1e-6}$ across all variants. The baseline hyperparameters are consistent with those used in the MeZO experiments. We also set $\epsilon=0.001$, following MeZO.

• For Bilevel-ZOFO, we tuned both upper- and lower-level learning rates with {1e-4, 1e-5, 1e-6, 1e-7} and {1e-2, 1e-3, 1e-4, 1e-5}, respectively.

Training Details:

We use Adam optimizer with betas=(0.9, 0.999) for the baselines and for both upper and lower updates of our method.

All experiments used a batch size of 8 and were conducted in bfloat16 precision on a single A6000 Ada 48GB GPU. MeZO was run for 10,000 steps, while FO and Bilevel-ZOFO methods were run for 5,000 steps. Our implementation builds upon MeZO’s codebase, and memory profiling as well as latency calculations are based on their framework.

If there is any details missing please point them out, and we would be happy to provide additional information. We thank you again for your feedback.

We greatly appreciate the time and effort you put into providing constructive feedback. Following your suggestions, we have carefully revised the manuscript to address your concerns. The updated version incorporates all recommended changes and improvements. We kindly request that you review the revised submission and our rebuttal and consider providing additional feedback. We would be especially grateful if you might consider adjusting your score in light of the improvements made. Thank you once again for your valuable input and for contributing to the quality of our work.

Questions, Concerns of Reviewer DSjq and How We Addressed Them:

1. Computational Cost Comparisons with Other Methods: We included detailed efficiency comparisons, demonstrating competitive performance and reduced computational demands (Table 6, Figure 5). Additionally, we explained how the proposed bilevel method is not as expensive as classical bilevel methods.

2. Using the Same Dataset for Both Levels of Optimization: We explained that distinct datasets are used to avoid overfitting in single-task fine-tuning and to simulate few-shot multitask learning.

3. Novelty in the Method and Contribution of Theoretical Analysis: We elaborated on the theoretical and empirical difficulties of integrating ZO gradient estimation into bilevel optimization. We emphasized the importance of theoretically ensuring convergence and reliability in real-world applications.

4. Why not Use Improved MeZO Variants: We clarified that our method focuses on bilevel structures and multitask capabilities, making it orthogonal to the choice of zeroth-order tuning methods. Additionally, we highlighted that the proposed method is compatible with other ZO variants, facilitating future integrations.

Outcome: The reviewer acknowledged the responses and the revisions, noting that their concerns were addressed and raising their score.

Questions, Concerns of Reviewer mLL4 and How We Addressed Them:

1. The need for simultaneous tuning full model and PEFT: We explained that the bilevel structure with ZO for the full model and FO for PEFT enhances task-specific performance, supported by empirical results. We also showed that the memory usage of our approach does not exceed the memory usage of PEFT or MeZO.

2. ZO motivation and memory profiling: We added efficiency comparisons (Fig. 5, Table 6), showing our method balances memory and convergence faster than MeZO, with FO sometimes requiring less memory due to frozen base models.

3. How this method mitigates Hard Prompt Sensitivity: We provided experimental evidence showing significant sensitivity in MeZO versus significant robustness in our approach due to the bilevel structure.

Outcome: We addressed all concerns with supportive evidence, including clarifications on notation, efficiency profiling, and detailed justifications for our method’s design choices. While the reviewer acknowledged our responses, they believed that incorporating the FO method in our approach was intended to address the hard prompt issue. However, we explained that this is incorrect --the bilevel structure deals with that aspect. Instead, using FO in PEFT tuning in our approach ensures performance improvements within the memory budget. Despite us resolving all concerns, the reviewer did not further engage.