Bilevel ZOFO: Efficient LLM Fine-Tuning and Meta-Training

TL;DR

Bilevel-ZOFO introduces a bilevel optimization approach combining first-order PEFT and zeroth-order updates to efficiently fine-tune large language models with improved convergence, stability, and generalization.

Contribution

This paper proposes Bilevel-ZOFO, a novel bilevel optimization method that integrates FO-PEFT and ZO techniques for efficient, stable, and generalizable LLM fine-tuning and meta-training.

Findings

Achieves 2-4 times faster training than existing methods.

Maintains similar memory efficiency while improving convergence.

Effectively combines full-model capacity with few-shot adaptation.

Abstract

Fine-tuning pre-trained Large Language Models (LLMs) for downstream tasks using First-Order (FO) optimizers presents significant computational challenges. Parameter-Efficient Fine-Tuning (PEFT) methods address these by freezing most model parameters and training only a small subset. However, PEFT often underperforms compared to full fine-tuning when high task-specific accuracy is required. Zeroth-Order (ZO) methods fine-tune the entire pre-trained model without back-propagation, estimating gradients through forward passes only. While memory-efficient, ZO methods suffer from slow convergence and high sensitivity to prompt selection. We bridge these two worlds with Bilevel-ZOFO, a bilevel optimization method that couples fast, local FO-PEFT adaptation at the inner level with stable, memory-efficient ZO updates of the full backbone at the outer level. The FO-PEFT inner loop performs fast, low-memory local adaptation that reduces the variance of ZO estimates and stabilizes the search, guiding the outer ZO updates of the full backbone and reducing prompt sensitivity. In the mean time, the outer ZO provides better generalization ability for PEFT. We provide theoretical convergence guarantees and empirically demonstrate that Bilevel-ZOFO significantly outperforms existing ZO and FO-PEFT methods, achieving 2-4 times faster training while maintaining similar memory efficiency. Additionally, we show by updating the backbone with ZO and adapting only a tiny FO-PEFT block per task, Bilevel-ZOFO combines full-model capacity with few-shot efficiency, making it a very efficient meta-learning algorithm that quickly adapts to new tasks.