Bilevel optimization solved better with second-order stationarity methods

To Solve Bilevel Optimization with Nonconvex Lower Levels, We Need Second-Order Stationarity

Machine Learning

Summary

Bilevel optimization is a way to solve complex problems where one problem is nested inside another. Usually, solving these problems is easier when the inside problem is shaped nicely (convex), but many real-world problems are not that simple. The authors showed that using a more advanced approach that looks at how the inside problem bends (second-order information) helps find better solutions and avoids tricky points that can confuse simpler methods. They created a new method called PROBE that efficiently solves these tough problems and tested it on tasks involving language models and learning from multiple tasks.

What this means in practice

  • For machine learning engineers: Improve training reliability for models involving nested learning tasks by avoiding poor solutions in nonconvex lower-level problems.
  • For data curation teams: Enhance automatic data selection processes for large language models by better navigating complex optimization landscapes.

Authors

Zhiyao Zhang, Menglu Yu, Alvaro Velasquez, Nathaniel D. Bastian, Jia Liu

Abstract

Although bilevel optimization (BLO) has emerged as a powerful framework for addressing many complex and nested machine learning problems in recent years, most existing studies are confined to the lower-level strongly convex (LLSC) or lower-level generally convex (LLGC) settings (i.e., the lower-level objective function is assumed to be, at least, convex). While the LLSC/LLGC assumptions render more tractable algorithmic design and theoretical analysis, they are too rigid to encompass many machine learning problems in practice. The limitations of LLSC/LLGC assumptions in BLO motivate us to investigate solving the BLO problem in the general lower-level nonconvex (LLNC) settings, which remains in its infancy. In the literature on LLNC-BLO, most of the existing works either require additional structures in the lower-level objective function for tractable theoretical analysis, or adopt the first-order stationarity reformulation as a lower-level surrogate problem, which is inherited from the LLSC/LLGC settings but could lose their effectiveness in the LLNC setting. To bridge this gap, we propose to reformulate the nonconvex lower-level problem using a second-order stationarity-based surrogate, the solution of which guarantees a local optimal solution at the lower level. Based on this reformulation, we propose the PROBE (Perturbed gradient algorithm for bilevel problem) and show that it overcomes the limitations of prior works by probing and escaping lower-level saddle points. We prove that PROBE achieves a finite-time convergence rate of $O(T^{-2/5})$, where T denotes iterations. To our knowledge, this work is the first to establish the finite-time convergence for achieving lower-level second-order stationary solutions in general LLNC-BLO. Our experiments on both a large language model-based data curation task and a meta-learning task also show that PROBE outperforms SOTA methods.