Contextual Stochastic Bilevel Optimization
This work addresses a foundational problem in machine learning with applications in meta-learning, personalized federated learning, and distributionally robust optimization, but it is incremental as it builds upon existing bilevel optimization frameworks by adding contextual information.
The authors tackled the problem of contextual stochastic bilevel optimization (CSBO), which extends classical stochastic bilevel optimization by incorporating contextual information in the lower-level problem, and they introduced an efficient double-loop gradient method based on the Multilevel Monte-Carlo technique, achieving sample and computational complexities that match existing lower bounds for stochastic nonconvex optimization and do not depend on the number of tasks in meta-learning.
We introduce contextual stochastic bilevel optimization (CSBO) -- a stochastic bilevel optimization framework with the lower-level problem minimizing an expectation conditioned on some contextual information and the upper-level decision variable. This framework extends classical stochastic bilevel optimization when the lower-level decision maker responds optimally not only to the decision of the upper-level decision maker but also to some side information and when there are multiple or even infinite many followers. It captures important applications such as meta-learning, personalized federated learning, end-to-end learning, and Wasserstein distributionally robust optimization with side information (WDRO-SI). Due to the presence of contextual information, existing single-loop methods for classical stochastic bilevel optimization are unable to converge. To overcome this challenge, we introduce an efficient double-loop gradient method based on the Multilevel Monte-Carlo (MLMC) technique and establish its sample and computational complexities. When specialized to stochastic nonconvex optimization, our method matches existing lower bounds. For meta-learning, the complexity of our method does not depend on the number of tasks. Numerical experiments further validate our theoretical results.