HELENE: Hessian Layer-wise Clipping and Gradient Annealing for Accelerating Fine-tuning LLM with Zeroth-order Optimization

Thank you for the reviewers’ insightful comments. We acknowledge the broad comparison between various optimizers. However, this paper focuses specifically on designing a zeroth-order (ZO) optimization algorithm, which differs fundamentally in scope from Adalomo and Galore. Adalomo primarily addresses memory usage reduction for backpropagation in first-order (FO) algorithms, while Galore emphasizes low-rank approximations of gradients in FO algorithms. It is well established that FO and ZO optimization serve distinct purposes. For example, ZO optimization is particularly advantageous in scenarios where gradients are inaccessible or when working within black-box models. Therefore, this work remains centered on advancing ZO optimization methodologies.

Comparison Between Helene and Adam

As a ZO optimizer, Helene demonstrates a memory complexity comparable to MeZO-Adam, requiring approximately 2P+ memory for optimization and an additional P+ for storing the optimizer's internal states. Despite this similarity in memory usage, Helene significantly outperforms MeZO-Adam in both efficiency and effectiveness, especially for convergence speed as you can observe in the paper, making it a superior choice in practical applications.

Comparison Between Helene and Adalomo

Adalomo employs several strategies to reduce memory usage, such as lowering numerical precision and frequent checkpointing, layer-by-layer. While these techniques reduce memory requirements, they increase time complexity from O(n) to O(n^2) where n is the total number of layers requiring updates. This transforms the problem from a dynamic programming challenge to a brute-force computation, significantly slowing convergence. Empirical results show that fine-tuning a LLaMA-7B model requires approximately 70+ GB of GPU memory for Helene and MeZO-Adam, compared to around 65 GB for Adalomo. All three optimizers necessitate the use of three NVIDIA 3090 GPUs. While Adalomo achieves modest memory reductions, it does so at the cost of efficiency. In contrast, Helene focuses on algorithmic design and convergence improvements for ZO optimization, independent of memory-saving techniques like those used in Adalomo and LOMO, which target FO backpropagation. These approaches could potentially complement ZO methods but fall outside the scope of this paper.

Comparison Between Helene and Galore

Both Helene and Adalomo aim to enable memory-efficient, full-parameter fine-tuning. However, Galore is a first-order-based approach by leveraging low-rank projection, compressing gradient information to a lower-dimensional space and re-projecting it back for updates. While this strategy reduces memory requirements, it has two notable limitations. First, low-rank projections may lead to the loss of critical gradient information. Helene overcomes this limitation by maintaining the full fidelity of gradients across the entire model parameter space. Second, Galore may lack curvature informativeness when the loss landscape curvature becomes extreme during optimization. In contrast, Helene accelerates training by accurately estimating and leveraging second-order information, avoiding the convergence delays caused by checkpointing or low-rank projections. These advantages make Helene a highly competitive alternative for memory-efficient fine-tuning tasks.

We agree that the reviewer suggests that well-engineered first-order methods can achieve memory usage comparable. We acknowledge that thoughtful and efficient engineering, like the techniques used in LOMO and AdaLOMO, can effectively reduce the memory footprint of first-order methods, potentially benefiting zero-order methods as well. We appreciate the reviewer’s discussion on the theoretical aspects of memory costs.

Technically, as mentioned in the LOMO and AdaLOMO papers, the authors claim that their methods achieve O(1) memory cost, in contrast to the O(\sqrt{N}) requirement typically seen in backpropagation (BP) with gradient checkpointing. However, this appears to be inconsistent with the reviewer's statement that checkpointing requires storing every \sqrt{N} layers in LOMO or AdaLOMO. We respectfully argue that O(1) memory under such conditions seems implausible. We are familiar with the typical gradient checkpointing approach but are unclear if achieving O(1) memory cost aligns with typical gradient checkpointing practices.

Questions

Q1: Proof of A-GNB

We appreciate the reviewer's comments. We acknowledge that the A-GNB is a diagonal Hessian estimator and an unbiased estimator of the Gauss-Newton matrix. Accordingly, we have revised the derivation of the A-GNB estimator and clarified the statements to be more precise in the updated manuscript.

Q2: About ICL performances.

We conducted experiments to evaluate In-Context Learning (ICL) using the provided MeZO script, thoroughly validating the results across different GPU environments to ensure reliability. We also observed the case where the performance of ICL is lower than that of zero-shot learning in MeZO. We hypothesize that this phenomenon is attributable to the relatively small size of the OPT/1.3B model, which may limit its capacity to effectively leverage in-context information.

Weakness

W1: About mismatch between theory and experiment

Thank you for highlighting the inaccurate statements. We have revised the algorithm description to specify that the gradient is estimated using SPSA, rather than computed exactly.

W2: Exact gradient and Hessian diagonal in Lemma 10

We appreciate the reviewer for highlighting the confusion in the Appendix. We have revised the proof to provide a more detailed derivation of the lemma, incorporating the clear use of ZO estimation for the gradient and Hessian. Please refer to the updated Lemma 10 for clarification.

W3: About HiZOO

Thanks for your suggestions. HiZOO did a great job. We provide comparisons between HiZOO and HELENE in the following two perspectives of both performances and algorithm efficiency (Wallclock time per step).

1. About performances

We recognize the superior accuracy of Hessian estimation in the HIZOO algorithm. However, we contend that the rank-1 update for Hessian computation in line 9 of the HIZOO algorithm incurs a memory cost of $O(n^2)$, where $n$ represents the number of model parameters, leading to higher memory requirements. In contrast, our method achieves significantly lower theoretical computational complexity and reduces memory consumption.

We provide comparisons of performances between HiZOO and HELENE below. We run model RoBERTa-large and report the average performances on six datasets (SST-2, SST-5, SNLI, MNLI, RTE, and TREC) with three different tuning methods, i.e., Full-parameter fine-tuning (FT), LoRA and Prefix-tuning (Prefix).

2. About algorithm efficiency (Wallclock time per step)

We report wallclock time per step between MeZO, HiZOO, and HELENE in the following table. We conduct experiments with RoBERTa-large on the dataset SST-2, with batch size of 64. We run the experiments on a single GPU (80GB A100) to avoid time calculation errors caused by commnication costs between multiple GPUs.

Additionally, since experimental settings between HiZOO and HELENE are different, the results reported in the paper HiZOO may not be directly compared with ours. We will do more comprehensive and comparative experiments in the future.

W4: About performance comparison with Adam

For the performances comparation with Adam, we summarize data into a tabular to illustrate.

The following table shows the comparison between MeZO-Adam and HELENE on the SST-2 dataset with RoBERTa-large and OPT-1.3B.

The following table shows the comparison between MeZO-Adam and HELENE for full-parameter fine-tuning on six different datasets with model RoBERTa-large. For figure 5a, we have re-plotted it.

Thanks for the response!

1. & 2. Thanks for the update in Lemma 10. I didn't check the proof in details but is the new Lemma 10 used in Lemma 11 or the actual convergence analysis?
2. Thanks for the new results. My concern with HiZZO is addressed.
3. Thanks for the experiment results, but my question is actually on the comparison of Adam with the same momentum annealing vs. Helene.

For a more or less degree, I think (1) clipping second-moment term, (2) update second-moment term in less frequency is not really helping Helene in ZO updates in general. It might be useful when combining with the exact gradients, but with ZO estimated gradient (effectively as a scalar times random Gaussian vector) their help is not clear. The discussion on the paper is really focused on the case of exact gradient but not in ZO noisy estimated gradient's case. We cannot simply borrow the idea from FO or second-order optimizer to the ZO without a clear intuition or justification.

5. It is better to provide a brief proof of A-GNB's unbiasedness to the Gauss-Newton matrix.
6. Thanks for the explanation.

I am still concerned with the first and the forth point, so I will still retain my previous score at this moment.

We highly appreciate the reviewer's careful review of the mathematical statements. For the 1st and 4th points, yes, Lemma 10 is used in Lemma 11 and the convergence proof. To address your query, we want to clarify that it offers theoretical understanding of the descent properties of the proposed ZO optimizer. Specifically, Lemma 10 quantifies the bound on the descent step and provides insights into how the ZO updates deviate from exact gradient updates, influenced by factors like $\delta_g$ and $\delta_H$, which are determined by the step size $h$ and the number of perturbations $m$. Like SGD, when h is sufficiently small, as our bound proved in Lemma 10, we can prove that our algorithm offers updates that converges to the desired update in expectation with 0 error.

We acknowledge that it is indeed non-trivial to directly adopt FO optimization convergence analysis for ZO methods, given the inherent differences in their update mechanisms. Therefore, the theorems we prove including Lemma 10 aim to establish the unique convergence properties of the proposed ZO method. Lemma 10 specifically analyze the effect of approximation errors (e.g., finite difference step size and perturbation sampling) on convergence, similar to the role of gradient noise in SGD-based analyses. The insights provided by the lemma are used to guide our convergence proofs, similar to how SGD is studied when noise in gradient estimation is incorporated in many SGD analyses.

As mentioned, we have proven that the Hessian clipping module ensures convergence even with noisy gradient estimation. However, the reviewer raised an important question about how the annealing module might impact convergence. We believe our specific annealing design effectively mitigates the influence of noisy gradients, particularly as optimization approaches the stationary point, where the gradient magnitude becomes smaller than the noise magnitude in gradient estimation. In such cases, our proposed annealing, which assigns less weight to the gradient, can help overcome barriers caused by significant gradient noise. This opens a promising direction for future work, where we aim to focus exclusively on refining annealing techniques in zero-order optimization.

Thank you for your comments. We have addressed this issue by revising the proof of A-GNB unbiasedness for the Gauss-Newton (GN) matrix. The updated brief proof has been incorporated into Section 3.4 of the main text, with the revisions highlighted in blue for clarity.

We sincerely thank the reviewer for their insightful comments and the valuable references provided. We greatly appreciate the feedback and will carefully incorporate the suggestions and cited literature into our manuscript.

Questions

Q1: About Zeroth-optimization and LoRA

We have addressed our memory usage above. For example, in OPT-1.3B, MeZO/zero-shot requires 4GB, ICL needs 6GB, Prefix Fine-Tuning uses 19GB, and full-parameter fine-tuning consumes 27GB, while HELENE requires only 14GB. Our optimizer is fully compatible with LoRA, our method is not focused on parameter-efficient tuning and should not be compared with LoRA. Our optimizer is compatible with Lora, if we choose to fine-tune OPT-1.3B with HELENE(Lora,rank=8), the memory usage is approximately 330MB, 2.3% of 14GB.

Q2: About training time

We provide our actual training time of Helene as well as MeZO in the following table. Helene shows faster convergence rates compared with MeZO.

Q3: About the annealing mechanism

Thank you for pointing out the confusion. Unlike Gradient Descent methods, Zero-Order methods require a stronger emphasis on momentum due to the increasing noise in SPSA gradient estimation as optimization progresses, as shown in Appendix Figure 7. When the gradient is given more weight relative to past momentum, training becomes unstable and may even diverge. The observed increase in noise in SPSA likely stems from the greater difficulty in optimization as model parameters approach a local minimum, where the noise from sampling perturbations can surpass the signal from the true gradient. To address this, we propose a novel annealing strategy specifically designed for Zero-Order methods, where the factor $\alpha$ dynamically reduce the gradient's contribution in the momentum update.

Q4: About Algorithm2

Thank you for the reviewer’s comments highlighting the unclear statement. We have updated the manuscript to include a step-by-step derivation for clarity.

Q5: For HELENE in first-order optimization

MeZO is a framework with zeroth-order optimization designed to reduce memory cost of fine-tuning LLMs. HELENE is a MeZO method estimating second-order information to accelerate convergence.

As stated in our previous answer, HELENE introduces three key components for its framework: Diagonal Hessian Estimation, the Annealing Mechanism, and Layerwise Clipping. For first-order optimizers, Hessian information is not applicable, so we did not conduct experiments in this context. The effect of Annealing can be observed in Figure 5a, where the difference between the blue and orange lines highlights its impact on the first-order optimizer. To demonstrate the effect of Clipping, we performed an experiment on image classification tasks.

We apply layerwise clipping on MeZO-SGD for Image classification task on the dataset CIFAR-10, the following result demonstrates the effectiveness of the layerwise clipping.

Weakness

W1: About Presentation and structure of paper

Apologies for the confusion caused by our paper structure. In the Method section, we first present our motivation to address three key issues: slow convergence, higher noise in ZO gradient estimation, and extreme curvature. To tackle these challenges, we propose the following three key components from HELENE, the diagonal Hessian Estimation (Section 3.3), the Annealing Mechanism (Section 3.3.1), and Layerwise Clipping (Section 3.5).

W2: About memory usage

To clarify, while HELENE introduces additional memory overhead compared to MeZO, its requirements remain efficient relative to other methods. For example, on OPT-1.3B, MeZO/zero-shot requires 4GB, in-context learning (ICL) needs 6GB, Prefix Fine-Tuning uses 19GB, and full-parameter fine-tuning consumes 27GB, whereas HELENE requires only 14GB—comparable to the memory usage of Adam. It is also worth noting that HELENE's improvements and memory consumption are best compared to other optimizers. Additionally, similar to other optimizers, HELENE is fully compatible with parameter-efficient methods such as LoRA, which could further mitigate memory requirements if needed.

W3: About experiments on larger models

Due to resource limitations, we were unable to include evaluations for larger models such as LLaMA2-7B or LLaMA2-13B in this submission, but we plan to address this in future work.

Regarding the ablation study, we recognize the importance of jointly analyzing Hessian and clipping. The value of the Hessian fluctuates significantly, making it challenging to use directly. Most algorithms employ smoothing techniques to address this issue-for example, HiZZO[1] applies EMA and square root, AdaHessian[2] uses spatial averaging, and Sophia[3] clips the Newton update $\mathbf{H}^{-1}\mathbf{g}$. The existing optimizers leveraging Hessians have rarely conducted detailed ablation studies to isolate the standalone impact of the Hessian. Thus, we jointly analyzing Hessian and clipping. Figure 5a presents an analysis of HELENE's three key components: From the blue curve to the orange curve, adding the annealing mechanism eliminates EMA-induced bias, improving stability. HELENE (purple) incorporates clipping into the diagonal Hessian, which significantly reduces the number of steps required for convergence (3000 steps).These results demonstrate that each of the components contributes meaningfully to accelerating zeroth-order optimization.

W4: About typos and wording errors Thanks you for pointing out the typos, we have corrected it.

[1] Zhao, Yanjun, et, al. "Second-order fine-tuning without pain for llms: A hessian informed zeroth-order optimizer" arXiv preprint arXiv:2402.15173 (2024).

[2] Yao, Zhewei, et, al. "ADAHESSIAN: An Adaptive Second Order Optimizer for Machine Learning" AAAI 2021.

[3] Liu, Hong, et, al. "Sophia: A Scalable Stochastic Second-order Optimizer for Language Model Pre-training" ICLR 2024.

Weakness

W1：About Slower than direct gradient-based method:

HELENE, our proposed zero-order method, is designed to enhance convergence speed compared to other zero-order methods like MeZO and MeZO-Adam. This paper specifically focuses on zero-order methods for memory-efficient training of large models. Therefore, comparing HELENE with first-order (i.e., gradient-based) methods falls outside the scope of this paper. The purpose of the annealing mechanism is to mitigating the bias introduced by native gradient EMA. The result can be found in Figure5a.

W2：About annealing

We revised the figure to further clarify how ablation study was organized. The annealing mechanism mitigates the gradient noise as optimization progress. When approaching the local minimum, the magnitude of noise in gradient estimation using SPSA could be greater than the magnitude of gradient, which can not be addressed well by native gradient EMA. As shown in the ablation study in our revised Figure 5(a), the EMA with annealing mechanism (orange curve) maintains a lower loss as steps increase compared to the EMA (blue curve), effectively mitigating the noise influence in ZO optimization process.

W3: About additional hyperparameters

We acknowledge that layer-wise clipping introduces additional hyperparameters. However, Figure 6 demonstrates that these hyperparameters are not highly sensitive. To address concerns about added tuning complexity, we found that using a percentile-based criterion for each layer—selecting parameters based on their absolute values within a top percentage (e.g., top 10%)—yields good performance. This approach reduces the tuning process to a single key hyperparameter, the clipping ratio, thereby significantly simplifying the procedure.

W4: About performances and convergence rates

1. In terms of training accuracy, Table 3, shown below, may address your concerns. Helene generally outperforms most Zero-Order optimizers or at least matches their performance, while also achieving significant speed improvements. As shown in Figure 3, our algorithm consistently accelerates different downstream tasks and parameter-efficient fine-tuning, demonstrating its robustness and efficiency.
2. Regarding the convergence rate, as illustrated in Figure 4, our algorithm consistently outperforms Adam, AdamW, Lion, and other optimizers in fine-tuning the OPT/1.3B model on the SST-2 dataset.

We also report and compare the performances of MeZO-Adam and HELENE on RoBERTa-Large in the following table. Helene shows an improvement of 0.5% compared with MeZO-Adam.

Question

Q1: About actual running time

To evaluate empirical convergence performance, we repeated the experiments using the LLM OPT-1.3B model on the SST-2 dataset with the same batch size and summarized the time costs and empirical convergence rates in the table below.