Not all parameters are equal: a Hessian informed differential learning rate for deep learning

We thank the reviewer for the comments! We address them below and welcome more feedbacks. We would appreciate it if the reviewer could raise the score if satisfied.

Evaluation of Hi-DLR

Thank you for mentioning this. We have added a detailed complexity analysis in Appendix B in our revised version. In a full-parameter training on a single GPU, the relative training time of Hi-DLR is $\frac{1}{1+\frac{4K}{3\Phi}}$. For instance when $K=3,\Phi=10$, Hi-DLR is 70% as fast as a base optimizer.

The authors state that they consistently observe existing PET methods selecting highly influential parameters, ... Why is this the case? The PPI metric is an estimation of parameter influence, so what is the corresponding ground truth for parameter influence? Without ground truth, how can we be certain that PET methods have indeed selected the highly influential parameters?

It remains an open question (that may be out of the scope of this paper) why certain PET methods are effective and select highly influential parameters. Some theories have reasoned that the fine-tuning information is low-rank so a small portion of trainable model parameters can capture it. We are happy to include some references if the reviewer is interested. In deep learning, unfortunately we don't have the ground truth of the true parameter influence. Our approach is to first identify the highly influential parameters through our PPI metric, then determine which PET methods include these parameters. E.g. in Figure 7 third sub-plot, LoRA includes such parameters but BitFit does not. We then validate our determination on the small models by directly trying it on larger models (see Table 3,4), and if the new PET is effective (i.e. comparable to the full fine-tuning), then we are more confident about the selection.

We thank the reviewer for all the constructive feedback and the well-grounded questions. We will address your comments one by one. We would sincerely appreciate it if the reviewer could provide more feedback or questions!

1."The idea is very similar to the derivation of Newton’s method"

we have built the connection between our method and Newton in line 179-181 but also stated the differences in line 182-190. The Newton’s method requires the inverse of Hessian, which is hard to compute in large-scale model trainings.

"the approximation of the Hessian with the diagonal matrix is not novel either"

We would like to stress that we’re not diagonalizing Hessian/Fisher. We have stated many classic methods like Adam that diagonalize Hessian/Fisher by introducing preconditionings in line 182-190. As a comparison, our diagonalization trick is applied to gHg, which is an important componant of our Hessian-informed adaptive learning rate. Hence it can be combined with any general optimizer that may or may not use the diagonalization of Hessian/Fisher. We provide the following simplified formula to highlight the difference between our framework and others. (Equation 2.2, 2.3 give a more sophisticated expression of our diagonalization trick.)

SGD: $\eta * I$

Newton: $1*H^{-1}$

AdaHessian: $\eta*diag(H)^{-1}$

Adagrad/Adam: $\eta*diag(P)^{-1}$

———————————————————————————

Hi-DLR SGD: $\frac{Gg}{ diag(gHg)}*I$

Hi-DLR Adam: $\frac{Gg}{diag(gHg)}*diag(P)^{-1}$

2. From the mathematic viewpoint, our implementation on AdamW is $Gg/diag(gHg)*diag(P)^{-1}$. The entry-wise adaptive learning rates of AdamW is presented in $diag(P)$ and our method adds another level of control (by grouping the d entry-wise learning rates into K groups, each with a meta-learning rate that we design). In contrast, the traditional learning rate for Adam is one single meta-learning rate that lacks the degree of freedom.

From the algorithmic viewpoint, our meta-framework takes an optimizer as a black box, as presented in our algorithm 1. We only use some extra forward functions to decide the optimal learning rates given the parameter grouping.

3. Sorry we missed this. We kindly refer to Appendix A.6 for the experimental details of ViT classification. Please let us know if this is sufficient or if you have questions about the details of other experiments. We will release the code upon acceptance.

4. We did experiment on pretraining one regression and one classification tasks in section 4.3. However, we are limited in computational resource for further pretraining experiments, which we leave for future work.

We thank the reviewer for the comments! We address them below and welcome more feedbacks. We would appreciate it if the reviewer could raise the score if satisfied.

1. This method includes other time-consuming steps and is incomplete without measurement.

Thank you for mentioning this. We have added a detailed complexity analysis in Appendix B in our revised version.

Under the same conditions, would a standard AdamW approach achieve a similar result?

Empirically we have shown that a standard AdamW can achieve similar (but maybe worse) results in many tasks (e.g, table 1 and 2). However, we note that this only holds when AdamW is using a proper learning rate, which takes time to search or tune. Under the same condition (i.e. using K-group DLR), if we use grid search, the learning rate searching time for a standard AdamW would be $O(K^2)$. Our algorithm only takes $O(K/ \Phi)$ and can be further compressed to $O(1)$ if we choose $\Phi=O(k)$. Besides, Hi-DLR can analyze the parameter importance by PPI while using standard AdamW cannot achieve this.

2. In the paper, parameter groups are adjusted manually, which requires explanation. Additionally, for large models, the hyperparameter K warrants further discussion.

Our main focus is given a grouping, how to design proper DLR for it, instead of how to design the grouping. In this work, we also explore the second question in two ways: (1) We chose application areas that the grouping strategies come naturally, such as multi-task learning, NAM and PET, e.g. PET naturally has two groups, the frozen one and the trainable one. (2) We construct a large enough set in Section 5 (say K=6, which gives $2^6$ sub-groupings) and use PPI to select the optimal sub-grouping. These two ways can extend to larger models without reconsidering K. For instance, in Table 7, we select the grouping on GPT-small and directly apply to GPT-large.

We have added a new section to discuss our limitations and future directions in Appendix C.

3. When group number K is a large number, the update period will increase as its learning rate update every \phi iterations, which might influence the convergence.

We agree the training time will increase as K becomes large, if we don't use other tricks. In our paper, we have experimented with K up to 40 in CelebA. For even larger K, one remediation is to use linearly larger $\Phi$ as we explained in Section 3 "When to derive".

Could the authors provide the source code for reproduction? We will release the code upon acceptance.

We thank the reviewer for the comments! We address them below and welcome more feedbacks. We would appreciate it if the reviewer could raise the score if satisfied.

Although everything diagonalized, the computational cost of the search is relatively high when fitting the second-order model, especially when $K$ is large.

We have added a detailed complexity analysis in Appendix B. In a full-parameter training on a single GPU, the relative training time of Hi-DLR is $\frac{1}{1+\frac{4K}{3\Phi}}$. We agree the training time will increase as K becomes large, if we don't use other tricks. In our paper, we have experimented with K up to 40 in CelebA. For even larger K, one remediation is to use linearly larger as we explained in Section 3 "When to derive".

I wonder if the algorithm can actually be designed more directly, for example instead of randomly sampled $\eta$s, fix several points on K lines, and than estimate curvature, make the learning rate on a line the negative curvature, etc.

Yes, there are some flexibility of designing the algorithm. We have tested fitting points $[-2,-1,1,2]*\eta_k$ for $k=1,...,K$ and the results are indistinguishable. We'll add this to algorithm 1.

The authors discussed the empirical comparison with other adaptive learning rate algorithms like prodigy and adaptation, etc. however, they did not discuss the relationship of their adaptation with previous adaptive algorithms in principle.

We have mentioned Prodigy and Dadaptation in paragraph "Automatic ULR" in line 105. These methods in their current form only work in ULR. Hence, they cannot solve the DLR problem as we proposed.

$d_k$ is used both as an update direction, and as a notation of dimensionality in your paper.

Sorry for the confusion. $d_k$ represents the number of parameters in group k but $\mathbf{d_k}$ (in bold) is an update direction. We'll change the notation in camera ready.

In Equation (3.1), and when fitting the quadratic function...

We use the corrected/pre-conditioned gradient when using Adam. Yes, the second order Taylor expansion is on any direction. We haven't studied the relationship but we know our adaptive gradient is positively correlated to the corrected gradient, regardless of Adam's momentum, and adaptive schedule. To see this, assume $[g_1, g_2]$ is the gradient of Adam grouped into two groups. Then Hi-DLR leads to $[\eta_1 g_1,\eta_2 g_2]$. It is obvious that the inner product $\eta_1 ||g_1||^2+\eta_2 ||g_2||^2>0$ which always holds because Hi-DLR gives positive learning rates.

For LoRA, the optimal learning rates ratios for matrix A and B can be fixed (Hayou et al., 2024) . How does your schedule interact with this ratio? If $\lambda$ is fixed, then your algorithm cannot work further on Lora+?

Our algorithm is compatible with Lora+. If the $\lambda$ is fixed, the optimal $\eta$ can be found by Hi-ULR because the task changes from finding optimal $[\eta_1,\eta_2]$ to $[1, \lambda]*\eta$.