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

Fine-tuning large language models (LLMs) is necessary for specific downstream tasks, but the classic adaptive first-order optimizer entails prohibitive GPU memory because of backpropagation. Recent works such as MeZO have turned to zerothorder optimizers for fine-tuning, which reduce substantial memory by using just two forward passes. However, heterogeneous curvatures across different parameter dimensions in LLMs often cause convergence instability or even failure. In this work, we propose HiZOO, a diagonal Hessian informed Zeroth-Order Optimizer , which is the first to leverage the diagonal Hessian to enhance ZOO for fine-tuning LLMs. We provide the theoretical proof for HiZOO and visualize the optimization trajectories on the test functions. Extensive experiments on various models (RoBERTa, OPT, Phi-2, and LLama3, with 350M∼66B parameters) indicate that HiZOO significantly reduces the number of training steps and improves model accuracy. For example, on the SST2 task, HiZOO achieves an 8× speed-up and better accuracy. Even when scaled to 66B-model, HiZOO outperforms MeZO with up to 5.1% absolute improvement. We also propose HiZOO-L, which reduces the Hessian memory cost to 10% of the MeZO, while maintaining almost same performance. Compared with ZO-Adam, HiZOO-L achieves a 4.3% absolute improvement, just using 50% of the GPU memory. Code is available at https://github.com/Yanjun-Zhao/HiZOO.