Contextual Stochastic Bilevel Optimization

TL;DR

This paper introduces a new stochastic bilevel optimization framework that incorporates contextual information, enabling applications like meta-learning and federated learning, and proposes an efficient double-loop gradient method with proven convergence.

Contribution

It extends classical stochastic bilevel optimization to include context, and develops a novel double-loop gradient method with complexity guarantees.

Findings

The proposed method converges efficiently with proven complexity bounds.

It applies to various applications including meta-learning and WDRO-SI.

Numerical experiments validate the theoretical convergence results.

Abstract

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.