Memory-efficient gradient unrolling for large-scale Bi-level optimization

Shen, Qianli, Wang, Yezhen, Yang, Zhouhao, Li, Xiang, Wang, Haonan, Zhang, Yang, Scarlett, Jonathan, Zhu, Zhanxing and Kawaguchi, Kenji (2024) Memory-efficient gradient unrolling for large-scale Bi-level optimization. Globerson, A., Mackey, L., Belgrave, D., Fan, A., Paquet, U., Tomczak, J. and Zhang, C. (eds.) In Advances in Neural Information Processing Systems 37 (NeurIPS 2024). Neural Information Processing Systems Foundation. 31 pp .

Abstract

Bi-level optimizaiton (BO) has become a fundamental mathematical framework for addressing hierarchical machine learning problems. As deep learning models continue to grow in size, the demand for scalable bi-level optimization has become increasingly critical. Traditional gradient-based bi-level optimizaiton algorithms, due to their inherent characteristics, are ill-suited to meet the demands of large-scale applications. In this paper, we introduce Forward Gradient Unrolling with Forward Gradient, abbreviated as FGU, which achieves an unbiased stochastic approximation of the meta gradient for bi-level optimizaiton. FGU circumvents the memory and approximation issues associated with classical bi-level optimizaiton approaches, and delivers significantly more accurate gradient estimates than existing large-scale bi-level optimizaiton approaches. Additionally, FGU is inherently designed to support parallel computing, enabling it to effectively leverage large-scale distributed computing systems to achieve significant computational efficiency. In practice, FGU and other methods can be strategically placed at different stages of the training process to achieve a more cost-effective two-phase paradigm. Further, FGU is easy to implement within popular deep learning frameworks, and can be conveniently adapted to address more challenging zeroth-order bi-level optimizaiton scenarios. We provide a thorough convergence analysis and a comprehensive practical discussion for FGU, complemented by extensive empirical evaluations, showcasing its superior performance in diverse large-scale bi-level optimizaiton tasks.

This record has no associated files available for download.

Contributors

Author: Zhanxing Zhu

Download statistics