提出HELENE优化器,用二阶信息加速大模型微调,内存占用低且收敛快。
HELENE: Hessian Layer-wise Clipping and Gradient Annealing for Accelerating Fine-tuning LLM with Zeroth-order Optimization
- 结合对角海森矩阵估计与分层梯度裁剪,提升零阶优化稳定性。
- 在多个任务上相比MeZO提速20倍,准确率平均提升1.5%。
- 适合大模型全参数或高效微调,特别适配层间维度差异大的架构。
微调大语言模型面临巨大内存压力,因反向传播需大量资源,尤其随模型规模增大而加剧。近期的MeZO采用零阶(ZO)优化方法,通过匹配推理阶段的内存消耗来缓解此问题,但存在因参数曲率差异导致收敛缓慢的缺陷。为此,本文提出HELENE,一种可扩展、内存高效的优化器,融合了渐进式A-GNB梯度与对角海森矩阵估计,并引入分层梯度裁剪,作为二阶预条件器。该设计实现更快更稳定的收敛。理论分析表明,HELENE改善了收敛速率,尤其适用于层维度异质的模型,其复杂度由最大层维度决定而非总参数量,契合现代大模型架构。在RoBERTa-large和OPT-1.3B上的实验显示,相较于MeZO,HELENE在多个任务上实现最高20倍加速,平均准确率提升1.5%。此外,HELENE兼容全参数微调与参数高效微调(PEFT),性能超越多个先进优化器。代码将在评审后发布。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。