Memory-Efficient Gradient Unrolling for Large-Scale Bi-level Optimization

2024-06-20

arXiv.org

AI Summary

Large-scale bilevel optimization faces critical bottlenecks: excessive memory consumption, high bias in meta-gradient estimation, and poor parallelizability. To address these, this paper proposes the Forward Gradient Unrolling ((FG)²U) method—a novel stochastic meta-gradient estimator that is provably unbiased, memory-efficient (empirically reducing memory usage by 70%), and amenable to distributed parallelization. (FG)²U supports both two-stage training and zeroth-order extensions. We establish its theoretical convergence guarantees under standard assumptions. Empirically, (FG)²U achieves state-of-the-art performance on hyperparameter optimization and neural architecture search tasks, significantly outperforming existing large-scale bilevel optimization approaches. Its design eliminates the need for costly backward-mode differentiation through the inner loop, thereby circumventing the memory and computational limitations inherent in conventional implicit or iterative differentiation strategies.

Technology Category

• Search and Optimization: Learning to Search
• Reasoning under Uncertainty: Stochastic Optimization
• Machine Learning: Optimization

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking

Abstract

Bi-level optimization (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 solutions has become increasingly critical. Traditional gradient-based bi-level optimization algorithms, due to their inherent characteristics, are ill-suited to meet the demands of large-scale applications. In this paper, we introduce $ extbf{F}$orward $ extbf{G}$radient $ extbf{U}$nrolling with $ extbf{F}$orward $ extbf{F}$radient, abbreviated as $( extbf{FG})^2 extbf{U}$, which achieves an unbiased stochastic approximation of the meta gradient for bi-level optimization. $( ext{FG})^2 ext{U}$ circumvents the memory and approximation issues associated with classical bi-level optimization approaches, and delivers significantly more accurate gradient estimates than existing large-scale bi-level optimization approaches. Additionally, $( ext{FG})^2 ext{U}$ is inherently designed to support parallel computing, enabling it to effectively leverage large-scale distributed computing systems to achieve significant computational efficiency. In practice, $( ext{FG})^2 ext{U}$ and other methods can be strategically placed at different stages of the training process to achieve a more cost-effective two-phase paradigm. Further, $( ext{FG})^2 ext{U}$ is easy to implement within popular deep learning frameworks, and can be conveniently adapted to address more challenging zeroth-order bi-level optimization scenarios. We provide a thorough convergence analysis and a comprehensive practical discussion for $( ext{FG})^2 ext{U}$, complemented by extensive empirical evaluations, showcasing its superior performance in diverse large-scale bi-level optimization tasks. Code is available at https://github.com/ShenQianli/FG2U.

Problem

• Large-scale Bilevel Optimization
• Deep Learning Models
• Efficiency and Accuracy

Innovation

• (FG)^2U Method
• Parallel Computing
• Large-Scale Optimization