Our new X account is live! Follow @wizwand_team for updates
WorkDL logo mark

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

About

Fine-tuning large language models (LLMs) with classic first-order optimizers entails prohibitive GPU memory due to the backpropagation process. Recent works have turned to zeroth-order optimizers for fine-tuning, which save substantial memory by using two forward passes. However, these optimizers are plagued by the heterogeneity of parameter curvatures across different dimensions. In this work, we propose HiZOO, a diagonal Hessian informed zeroth-order optimizer which is the first work to leverage the diagonal Hessian to enhance zeroth-order optimizer for fine-tuning LLMs. What's more, HiZOO avoids the expensive memory cost and only increases one forward pass per step. Extensive experiments on various models (350M~66B parameters) indicate that HiZOO improves model convergence, significantly reducing training steps and effectively enhancing model accuracy. Moreover, we visualize the optimization trajectories of HiZOO on test functions, illustrating its effectiveness in handling heterogeneous curvatures. Lastly, we provide theoretical proofs of convergence for HiZOO. Code is publicly available at https://anonymous.4open.science/r/HiZOO27F8.

Yanjun Zhao, Sizhe Dang, Haishan Ye, Guang Dai, Yi Qian, Ivor W.Tsang• 2024

Related benchmarks

TaskDatasetResultRank
Image ClassificationCIFAR-100
Accuracy64.9
302
Text ClassificationBoolQ
Accuracy73.9
84
Text ClassificationRTE
Accuracy71.8
78
ClassificationSST2
Accuracy92.1
58
Sentence CompletionCOPA
Accuracy88
48
ClassificationCB
Accuracy69.6
46
GenerationSQuAD
F1 Score83.8
44
ClassificationWSC
Accuracy63.5
41
Word-in-Context ClassificationWiC
Accuracy60.2
34
ClassificationMultiRC
Accuracy64.8
29
Showing 10 of 19 rows

Other info

Follow for update