Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Huaqin Zhao, Jiaxi Li, Yi Pan, Shizhe Liang, Xiaofeng Yang, Wei Liu, Xiang Li, Fei Dou, Tianming Liu, Jin Lu

TL;DR

HELENE is a novel scalable and memory-efficient optimizer that integrates annealed A-GNB gradients with a diagonal Hessian estimation and layer-wise clipping, serving as a second-order pre-conditioner that improves convergence rates, particularly for models with heterogeneous layer dimensions, by reducing the dependency on the total parameter space dimension.

Abstract

Fine-tuning large language models (LLMs) poses significant memory challenges, as the back-propagation process demands extensive resources, especially with growing model sizes. Recent work, MeZO, addresses this issue using a zeroth-order (ZO) optimization method, which reduces memory consumption by matching the usage to the inference phase. However, MeZO experiences slow convergence due to varying curvatures across model parameters. To overcome this limitation, we introduce HELENE, a novel scalable and memory-efficient optimizer that integrates annealed A-GNB gradients with a diagonal Hessian estimation and layer-wise clipping, serving as a second-order pre-conditioner. This combination allows for faster and more stable convergence. Our theoretical analysis demonstrates that HELENE improves convergence rates, particularly for models with heterogeneous layer dimensions, by reducing the dependency on the total parameter space dimension. Instead, the method scales with the largest layer dimension, making it highly suitable for modern LLM architectures. Experimental results on RoBERTa-large and OPT-1.3B across multiple tasks show that HELENE achieves up to a 20x speedup compared to MeZO, with average accuracy improvements of 1.5%. Furthermore, HELENE remains compatible with both full parameter tuning and parameter-efficient fine-tuning (PEFT), outperforming several state-of-the-art optimizers. The codes will be released after reviewing.

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

TL;DR

HELENE is a novel scalable and memory-efficient optimizer that integrates annealed A-GNB gradients with a diagonal Hessian estimation and layer-wise clipping, serving as a second-order pre-conditioner that improves convergence rates, particularly for models with heterogeneous layer dimensions, by reducing the dependency on the total parameter space dimension.

Abstract

Fine-tuning large language models (LLMs) poses significant memory challenges, as the back-propagation process demands extensive resources, especially with growing model sizes. Recent work, MeZO, addresses this issue using a zeroth-order (ZO) optimization method, which reduces memory consumption by matching the usage to the inference phase. However, MeZO experiences slow convergence due to varying curvatures across model parameters. To overcome this limitation, we introduce HELENE, a novel scalable and memory-efficient optimizer that integrates annealed A-GNB gradients with a diagonal Hessian estimation and layer-wise clipping, serving as a second-order pre-conditioner. This combination allows for faster and more stable convergence. Our theoretical analysis demonstrates that HELENE improves convergence rates, particularly for models with heterogeneous layer dimensions, by reducing the dependency on the total parameter space dimension. Instead, the method scales with the largest layer dimension, making it highly suitable for modern LLM architectures. Experimental results on RoBERTa-large and OPT-1.3B across multiple tasks show that HELENE achieves up to a 20x speedup compared to MeZO, with average accuracy improvements of 1.5%. Furthermore, HELENE remains compatible with both full parameter tuning and parameter-efficient fine-tuning (PEFT), outperforming several state-of-the-art optimizers. The codes will be released after reviewing.

Paper Structure

This paper contains 28 sections, 13 theorems, 81 equations, 7 figures, 3 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Preliminaries
  3. Zeroth-Order Gradient Estimators and MeZO
  4. Diagonal Hessian Approximation
  5. Method
  6. Motivation
  7. HELENE: Hessian Layer-wise Clipping and Gradient Annealing
  8. EMA of diagonal Hessian estimates
  9. Annealing Mechanism
  10. Asymptotic Gauss-Newton-Bartlett (A-GNB) Estimator
  11. Laywerwise Clipped Diagonal Hessian to help Newton's method
  12. Convergence Analysis
  13. Experiments
  14. Masked Language Models
  15. Auto-Regressive LLMs
  16. ...and 13 more sections

Key Result

Theorem 1

Under Assumptions 1 and 2, let $\eta = \frac{1}{2}$ and $\lambda_i = \frac{R_i}{2\sqrt{d_i}}$. The update reaches a loss at most $\epsilon$ in steps, where $L$ is the loss function, $\boldsymbol{\theta}_{0,i}$ is the initial parameter vector for layer $i$, $\mu_i$ is the strong convexity constant for layer $i$, and $R_i$ is the bound on the distance between $\boldsymbol{\theta}_{0,i}$ and $\bolds

Figures (7)

  • Figure 1: The motivating toy example. HELENE can maintain stable updates when facing curvature issues, while other second-order optimizers are severely affected by them.
  • Figure 2: Comparison of HELENE with Newton's method and Sophia. The performance of this training loss cross-checks with the toy sample in Figure \ref{['fig:landscape']}.
  • Figure 3: Performance and convergence of MeZO and HELENE for fine-tuning, LoRA, and prefix-tuning of OPT-1.3B on different datasets. HELENE achieves approximate $10\times$ speedup and up to 15$\%$ accuracy improvement compared to MeZO.
  • Figure 3: Performance of LLM fine-tuning on SST2 over pre-trained Roberta-Large and OPT-1.3B. Best performance among ZO methods (including Forward-Grad) is highlighted in bold.
  • Figure 4: Validation losses for ZO-optimizers. MeZO:0.426, Adam:0.286, AdamW:0.351, Lion:0.343, HELENE:0.283.
  • ...and 2 more figures

Theorems & Definitions (25)

  • Theorem 1
  • Lemma 1: Divergence to Infinity
  • proof
  • Lemma 2: Parameter Bound
  • proof
  • Lemma 3: Gradient Norm Bound
  • proof
  • Lemma 4: Stability of Gradient Flow
  • proof
  • Lemma 5: Quadratic Form Integration
  • ...and 15 more