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

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.

💡 Research Summary

The rapid advancement of Large Language Models (LLMs) has brought significant computational challenges, particularly regarding the high cost of fine-tuning for downstream tasks. Traditionally, two main approaches have emerged to mitigate these costs: Parameter-Efficient Fine-Tuning (PEFT) and Zeroth-Order (ZO) optimization. While PEFT reduces memory usage by training only a small subset of parameters, it often fails to match the performance of full fine-tuning in high-accuracy scenarios. Conversely, ZO methods offer extreme memory efficiency by eliminating back-propagation and relying solely on forward passes, but they are plagued by slow convergence and extreme sensitivity to prompt selection due to high gradient estimation variance.

This paper introduces “Bilevel-ZOFO,” a novel bilevel optimization framework designed to bridge the gap between these two methodologies. The core innovation lies in its dual-loop architecture that synergistically combines the strengths of both First-Order (FO) and Zeroth-Order (ZO) approaches.

In the inner loop, the framework employs FO-PEFT for rapid, local adaptation. This step is not merely about parameter updates; it serves as a critical variance reduction mechanism. By performing fast, low-memory adaptation, the inner loop stabilizes the search space and provides a more reliable gradient estimate for the outer loop, effectively mitigating the inherent instability of ZO methods.

In the outer loop, the framework utilizes ZO optimization to update the entire pre-trained backbone. This allows the model to leverage its full capacity, overcoming the performance ceiling typically associated with PEFT. By updating the full backbone through the stabilized guidance of the inner loop, Bilevel-ZOFO achieves a balance between the high-capacity learning of full fine-tuning and the memory efficiency of ZO methods.

Empirical results demonstrate that Bilevel-ZOFO significantly outperforms existing ZO and FO-PEFT methods, achieving a 2x to 4x acceleration in training speed while maintaining comparable memory efficiency. Furthermore, the authors provide theoretical convergence guarantees and demonstrate that the method is highly effective for meta-learning, as it can quickly adapt to new tasks by updating only a tiny FO-PEFT block while simultaneously refining the backbone. This breakthrough offers a highly efficient paradigm for large-scale model adaptation, making high-performance LLM fine-tuning more accessible and scalable.