Second-Order Fine-Tuning without Pain for LLMs: A Hessian Informed Zeroth-Order Optimizer

Thank you for your thorough review and valuable comments on our manuscript. We appreciate the opportunity to address your concerns and clarify any misunderstandings.
Question 1: Treatment of $uu^{\top}$ as a Diagonal Matrix
[Revised answer]Firstly we would like to address the broader computational considerations. Due to the significant memory and computational overhead required to estimate the full Hessian matrix $d \times d$, we chose to compute the diagonal of the Hessian instead. This approach strikes a balance between computational feasibility and the accuracy needed for our analysis. As a result, in our code, we do indeed compute $diagonal(u_i u_i^{\top}) = [u_1^2, u_2^2, ... u_n^2]$, whose rank is not 1 and you are correct. The statement in the algorithm is included to illustrate the theoretical process of our approach. However, in the actual computation, particularly in Line 9, this step is a middle state and does not appear. It is meant purely for clarity in explaining the algorithm's underlying theory.

Question 2: Small Value of the Hessian Smooth Parameter
Due to the large variance of the Hessian, with some values being extremely large, so Hessian smooth term is very small. But it is still sufficient to have a meaningful impact on the convergence and stability of the optimization process. Our experimental results demonstrate that, under the same parameter configurations, HiZOO significantly improves the convergence rate compared to other methods.

Question 3: Experiments on Generation Tasks
Thank you for your insightful suggestion. We agree that including generation tasks would provide a more comprehensive evaluation of HiZOO's performance. We will conduct additional experiments on generation tasks and show the results to you as soon as possible to showcase HiZOO's capabilities in generative settings.

Question 4: some specific figures
We acknowledge that sometimes HiZOO exhibits poor training performance in certain cases, it's also the same for MeZO or other zeroth-order methods. This issue arises due to the specific characteristics of some tasks, instead of the algorithm. However, HiZOO generally demonstrates better convergence behavior than MeZO under the same conditions.

Question 5: Computational cost of calculating R and C matrix in Algorithm 2

It 's true that this calculation involves some additional computations. However, the time required for generating these matrices is minimal compared to the overall time of one forward pass. As shown in Table below (also Table 11 in Paper), we indicate that matrix decomposition increases only about 6% of the total training time. Therefore, this calculation does not significantly impact the efficiency of the algorithm.

Table: Wallclock time per step between MeZO, HiZOO, and HiZOO-L.
The increase in wallclock time per step for HiZOO compared to MeZO is less than 1.5 times across different model sizes. Also, HiZOO-L increases only about 6% of the total training time than HiZOO. To avoid introducing additional overheads such as inter-GPU communication, results are measured on the same dataset (SST-2) and GPUs (80GB A100), with each result averaged over 100 steps. "BS" refers to batch size. For the relatively smaller RoBERTa-large model, we used a BS=64, while for models larger than 1B parameters, we used a BS=16.

We appreciate your constructive feedback and believe that addressing these points will strengthen our manuscript. We are committed to improving our work and will incorporate the necessary revisions to enhance its clarity and comprehensiveness.

Under the assumptions of gradient descent and convex optimization, alongside Assumption $\frac{tr(\Sigma^{1/2}H\Sigma^{1/2}) }{\|\| \Sigma^{1/2}H\Sigma^{1/2}\|\|_2} < r$, we believe there is strong potential to demonstrate that the diagonal-Hessian based preconditioning in HiZOO offers significant advantages over MeZO. Such findings would underscore the promising practical applications of our method and its effectiveness. We are willing to share further insights and are committed to enhancing this aspect of the theory, trying to ensure a more comprehensive theoretical explanation is included in the final version of the paper.

Also, it is important to acknowledge that extending these advantages to scenarios involving stochastic gradient descent (SGD) or non-convex optimization presents significant challenges. To the best of our knowledge, there is few or even no existing work in the field that has successfully established improvements under these conditions. This highlights the value of our contributions within the specific scope of convex optimization and gradient descent.

We sincerely hope that you will reconsider our work and, if possible, provide further support for our research.

We deeply appreciate the time and effort you have invested in reviewing our manuscript. Your thoughtful comments and constructive suggestions have been invaluable in improving the quality of our work. Thank you for your insightful and detailed feedback.

Question 1: $\Sigma$ is not required to be close to the inverse Hessian
We apologize for any confusion caused by our previous presentation. To clarify, the Hessian matrix in our method does not require the condition that the matrix be close to the inverse Hessian, which means that:

Drawing from the lemma presented in MiNES(Mirror Natural Evolution Strategies):

$\frac{1}{2} \cdot \mathbb{E}_u (u^{\top} \Sigma^{1/2} H \Sigma^{1/2} u \cdot (\Sigma^{-1/2}u u^{\top} \Sigma^{-1/2} - \Sigma^{-1}))=H$

where $H$ is the Hessian $\nabla^2 \mathcal{L}(\theta; \mathcal{B})$ and $\Sigma$ is a positive definite matrix.

Question 2: theoretical result of convergence
Thank you for your valuable feedback. We sincerely apologize for the typo in our manuscript. We will revise Theorem 3.2 to reflect this correction for clarity, shown as below:

Let the descent direction $g_\mu(\theta_t)$ defined as:

$g_\mu(\theta_t) = \sum_{i=1}^b \frac{\mathcal{L}(\theta_t + \mu \Sigma_t^{1/2} u_i; \mathcal{B}_t) - \mathcal{L}(\theta_t - \mu \Sigma_t^{1/2} u_i;\mathcal{B}_t)}{2b\mu} \Sigma_t^{1/2}u_i.$

Based on Assumption, if the update rule for $\theta$ is $\theta_{t+1} = \theta_t - \eta g_\mu(\theta_t)$ for a single step, then it's established that: $\mathbb{E} \left[\mathcal{L}(\theta_{t+1}) \mid \right] \leq \mathcal{L}(\theta_t) - \frac{\eta_t}{4} \|\nabla \mathcal{L}(\theta_t)\|_{\Sigma_t}^2 + 2\eta_t^2L\left( \mathrm{tr}(\Sigma_t) +\beta_u \right)\sigma^2 + O(\mu^2).$

Furthermore, given iteration number $T$, we choose the step size $\eta = \frac{1}{8\sqrt{T}L(\max_t\mathrm{tr}(\Sigma_t) +\beta_u)}$ and take $\theta_{\mbox{out}} = \theta_j$ with $j$ uniformly sampled from $\{1, \dots, T\}$. Then, we have

$\mathbb{E} \left[\|\nabla \mathcal{L}(\theta_{\mbox{out}})\|^2\right] \leq \frac{32L\left( \max_t \{\mathrm{tr}(\Sigma_t)\} +\beta_u \right) \mathcal{L}(\theta_1) - \mathcal{L}(\theta_*))}{\sqrt{T}\beta_\ell } + \frac{\sigma^2}{T^{3/2} \beta_\ell} + O\left(\mu^2\right),$ where $\mathcal{L}(\theta_* )$ minimizes the function $\mathcal{L}(\theta;)$. The above equation shows that as $T\to \infty$, HiZOO can converge to the stationary point.\

Question3 Theoretical concern
In fact, in the realm of non-convex optimization, it is difficult to show that a Newton-like algorithm with approximate Hessian matrices can help to achieve faster convergence rate. And we admit that our convergence analysis does not provide theoretical support that our method can achieve a fast convergence rate. Our theory mainly is used to show that our algorithm can converge to some stationary point.
We have provided an updated convergence proof for your reference. Due to its length, we have directly updated the PDF file, specifically in Line 763-884.

Question4: apply HiZOO in pre-training
Thank you for your insightful comment. Indeed, there have been studies that have applied zeroth-order optimizers for training neural networks from scratch, eg.[DeepZero: Scaling up Zeroth-Order Optimization for Deep Model Training]. However, due to the high computational cost associated with training from scratch, in this work we focus on fine-tuning. In the future, we are willing to explore using HiZOO for pre-training.
Additionally, in our test functions, we set the exactly same hyper-parameters for Adam, MeZO and our method. For example, in function: $f(x,y) = 10(x-1)^2*(2x^2+x+1)+2(y-4)^2$, we set the same start point$(-1,2)$, end point(1,4), learning rate $10^{-3}$, scaling factor $10^{-6}$, and we can see that in the same setting, our method performs better than MeZO.

Dear Reviewer MAUK, we greatly appreciate the time and effort you have dedicated to reviewing our submission.

As part of the ongoing discussion period, we have provided detailed responses to the points raised in your review. We would sincerely appreciate any additional feedback or clarifications you may have. With only a few days remaining in the discussion period, we are eager to ensure a productive and collaborative exchange.

Thank you once again for your valuable input and consideration.

Thank you to the authors for their responses.

However, from a theoretical perspective, the convergence of HiZOO cannot outperform vanilla SGD due to the trace term, which fails to justify the necessity of using this inverse Hessian estimator. Additionally, after reading the MiNES paper, the proposed estimator appears to estimate not the inverse Hessian but rather something closer to the inverse Fisher matrix. While Fisher and Hessian may share similarities in estimating curvature information, they cannot be considered equivalent in large language models (LLMs) that do not use ReLU activations. This raises a fundamental question about whether HiZOO can truly be considered a Hessian-informed method. Moreover, even if it is correct to say that it estimates the inverse Fisher matrix, works like [1] have shown that natural gradient descent (using K-FAC, approximate NGD) achieves linear convergence. Thus, HiZOO should at least provide theoretical insights or intuition in this regard.

For these reasons, I will maintain my current score

[1] Fast Convergence of Natural Gradient Descent for Overparameterized Neural Networks, G. Zhang et al., NeurIPS 2019.

First and foremost, we sincerely appreciate your effort in raising new questions for further exploration. We will address each of your concerns in detail below, and we hope our responses will earn your understanding and approval.

(Also we want to kindly remind you that this work primarily focuses on experimental contributions. As a result, we conduct lots of experiment result to show the performance of our method and do not extensively delve into theoretical analyses.)

Question:convergence of HiZOO cannot outperform vanilla SGD due to the trace term, fails to justify the necessity of inverse Hessian estimator.

Answer: Thank you for your insightful comment on the convergence properties of HiZOO. As indicated in the theoretical analysis of MeZO[1] (HiZOO is an improvement based on MeZO), the loss decrease for MeZO can be expressed as:

$\mathbb{E}[L(\theta_{t+1}) | \theta_t] - L(\theta_t) \leq \frac{1}{\gamma} \left[-\eta_{SGD} \|\nabla L(\theta_t)\|^2 + \frac{1}{2} \eta_{SGD}^2 \ell \cdot \mathbb{E}[\|\nabla L(\theta; \mathcal{B})\|^2] \right],$

where $\gamma^{-1} = O(n/r)$, $n$ denotes the number of queries per iteration, and $r$ is the effective rank of the Hessian. For MeZO/HiZOO, $n = 1$, which implies that its convergence rate depends solely on $r$, the effective rank of the Hessian, rather than $d$, the parameter dimension, as in traditional SGD.

In traditional zeroth-order SGD, the efficiency is scaled by $d$, as the loss decrease is proportional to $1/d$ due to the dimensionality of the parameter space. However, since $r$ is much smaller than $d$, the factor $1/r$ in MeZO and HiZOO results in a faster loss decrease compared to $1/d$ in SGD. A detailed analysis of this can be found in the convergence analysis section of MeZO.

This distinction is significant because, as shown in Adahessian[3], Sophia[4], Shampoo[5], employing curvature information (Hessian approximations) enables algorithms to converge faster than first-order methods like SGD, especially in high-dimensional optimization problems. Specifically, the effective rank $r$ is typically much smaller than the parameter dimension $d$, owing to the low-rank property of the Hessian matrix observed in many practical machine learning scenarios.

Moreover, HiZOO improves upon MeZO by introducing second-order information through preconditioning methods that estimate the diagonal of the Hessian.

Question:HiZOO can truly be considered a Hessian-informed method or Fisher-informed method?

Answer: HiZOO is definitely a Hessian-informed method. Thank you for raising this insightful question. Regarding the convergence properties of HiZOO, we would like to emphasize that our method indeed estimates the inverse Hessian matrix, as demonstrated in the theoretical analysis of MiNES[2]. Specifically, Lemma 11 in MiNES establishes that the sequence $\Sigma_k^{-1}$, generated by their algorithm, satisfies:

$\|\Sigma_k^{-1} - \Pi_{\mathcal{S}} (\nabla^2 f(\mu^*))\|^2 \leq \frac{C}{k},$

where $\nabla^2 f(\mu^*)$ is the Hessian matrix of the loss function $f$ at the optimal point $\mu^*$. The projection $\Pi_{\mathcal{S}}$ removes the extreme eigenvalues (both maximum and minimum) of the original Hessian matrix $\nabla^2 f(\mu^*)$, effectively regularizing it within a spectral range defined by $\mathcal{S}$. When $\mathcal{S}$ is chosen to be sufficiently large, $\Pi_{\mathcal{S}} (\nabla^2 f(\mu^*))$ becomes equivalent to $\nabla^2 f(\mu^*)$ itself, as the projection no longer affects the spectral properties of the Hessian. This result clearly demonstrates that $\Sigma_k^{-1}$ approximates $\nabla^2 f(\mu^*)$ as $k \to \infty$, at a convergence rate of $O\left(\frac{\log k}{k}\right)$.

While the MiNES paper also discusses the Fisher matrix $F_\mu$ in the context of natural gradient descent, it is important to note that this section is unrelated to our analysis. HiZOO focuses solely on the inverse covariance matrix $\Sigma_k^{-1}$, which approximates the Hessian matrix and does not involve the Fisher matrix. The distinction is particularly evident in the update rules for $\Sigma_k$, which are derived based on Hessian-specific properties rather than Fisher information.

[1]Malladi, Sadhika, et al. "Fine-tuning language models with just forward passes." Advances in Neural Information Processing Systems 36 (2023): 53038-53075.

[2]Ye, Haishan and Tong Zhang. “Mirror Natural Evolution Strategies.” ArXiv abs/1910.11490 (2019): n. pag.

[3]Yao, Z., Gholami, A., Shen, S., Keutzer, K., & Mahoney, M.W. (2020). ADAHESSIAN: An Adaptive Second Order Optimizer for Machine Learning. ArXiv, abs/2006.00719.

[4]Liu, H., Li, Z., Hall, D.L., Liang, P., & Ma, T. (2023). Sophia: A Scalable Stochastic Second-order Optimizer for Language Model Pre-training. ArXiv, abs/2305.14342.

[5]Anil, R., Gupta, V., Koren, T., Regan, K., and Singer, Y. Scalable second order optimization for deep learning. arXiv preprint arXiv:2002.09018, 2020.

I appreciate the authors' response.

However, the local effective rank condition (presented in MeZO) could also be applied to the analysis for HiZOO. Also, the proof technique in current version of HiZOO does not give theoretical superiority over the MeZO even with the local effective rank condition. So, I think that the authors should improve this part for future direction.

I also understand this work does more focus on the empirical studies, so I raised the score up to 5. But I strongly recommend that the authors should revise the theory related to trace of estimators.

We sincerely appreciate the time and effort you have dedicated to reviewing our manuscript. Your valuable comments and suggestions are instrumental in enhancing the quality of our work. Thank you for your insightful feedback.

Question 1 Thank you for your suggestion and this reference paper provides a very solid theoretical introduction, which is worth learning for us. By Theorem 3.2, we can obtain a query complexity $\left(\frac{\mathrm{tr}(\Sigma)}{\beta_\ell}\right)^2\cdot\frac{1}{\epsilon^4}$ for our method. Our query complexity depends on $\epsilon^{-4}$ while the reference paper depends $\epsilon^{-2}$, our baseline MeZO is also $\epsilon^{-4}$. This is because our method samples function (namely Stochastic Gradient Descent) while the this paper uses all functions per iteration (namely Gradient Descent). So we could not conduct direct comparison. However, we will cite this reference paper and built upon their methodology in our work.

Question 2 We will simplify the detailed derivations in Section 3 as recommended. Additionally, we will move the definitions of the three test functions in Section 3.5 to the main text to enhance clarity and readability.

Question 3 Thank you for pointing out the importance of extending our experiments. On Roberta-large we conduct the few-shot learning experiment and on OPT, Llama, Phi models we have conducted the full-shot learning, which have included the supervised situations. Also, we will consider including more datasets or models to evaluate our method.

Dear Reviewer RuKr, we greatly appreciate the time and effort you have dedicated to reviewing our submission.

We have read this paper [Stochastic Two Points Method for Deep Model Zeroth-order Optimization], https://arxiv.org/abs/2402.01621. This paper is truly impressive, with solid theoretical foundations and remarkable contributions. I greatly admire this work and see it as a valuable resource for my learning.

As part of the ongoing discussion period, we have provided detailed responses to the points raised in your review. We would sincerely appreciate any additional feedback or clarifications you may have. With only a few days remaining in the discussion period, we are eager to ensure a productive and collaborative exchange.

Thank you once again for your valuable input and consideration.