Contextual Stochastic Bilevel Optimization

Summary

This work studies the so-called contextual stochastic bilevel optimization (CSBO) problem, in which the lower-level problem is a conditional expectation problem under some contextual information. Some applications in distributional robust optimization fall into this category. The paper proposes a double-loop Monte-Carlo based method, which leverage EpochSGD in Hazan and Kale, 2014 to approximate the lower-level solution. A approximate hypergradient is provided based on an explicit form via implicit function theorem. Experiments on meta-learning and instrumental variable regression are provided.

Strengths

1. This work considers a problem that is different from existing studies by introducing the contextual information. It covers some important cases in, e.g., DRO. 2. The algorithms are reasonable to me. Using MLMC to further improve the performance of DL-SGD is good. Theoretical complexity and convergence are analyzed. Experiments seem to support the design principles.

Weaknesses

1. The studied problem seems a little bit artificial. The example in (3) seems to artificially change a single-level problem into a bilevel one. In addition, the hyper-gradient form is almost the same as in the non-contextual case and hence the tools therein may be used with some adaptations. 2. Some important components such as MLMC, EpochSGD, Neumann series expansion are existing techniques. Some challenges like variance control (e.g., in $\hat v$) and hyper-gradient computation can be well coped with by these techniques. Thus, the novelty is not that significant. 3. In experiments, in Fig. 1, why does MAML stop at this large loss value? Is it because of a non-satisfactory hyperparameter tuning or something else. It seems that the stepsizes for MAML may be chosen to be too large. The comparison seems not fair. 4. The experiments do not show the important of bilevel optimization. For example, no baselines other than the proposed CSBO solvers are provided. Some baselines (e.g., single-level ones, or some standard baselines) in Wasserstein DRO may need to be included. Overall, I am not fully convinced that this is an important problem, and given the above concerns, I lean toward the negative side.

Questions

See Weakness part.

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

Limitations

See Weakness part.

Summary

Contextual stochastic bilevel optimization (CSBO) is introduced in this paper. An efficient double-loop gradient method based on the Multilevel Monte-Carlo (MLMC) is proposed. The proposed framework captures important applications such as meta-learning, Wasserstein distributionally robust optimization with side information (WDRO-SI), and instrumental variable regression (IV).

Strengths

An interesting problem, that is, Contextual Stochastic Bilevel Optimization Problem is proposed in this work. And an efficient algorithm is proposed to solve this problem. The proposed framework captures important applications such as meta-learning, Wasserstein distributionally robust optimization with side information (WDRO-SI), and instrumental variable regression (IV). However, I have some concerns as follows.

Weaknesses

1. I think it's necessary to emphasize the difficulty of solving the Contextual Stochastic Bilevel Optimization Problem compared with traditional bilevel optimization problems. The presentation of this work is poor. I believe this is an excellent work, I suggest that you should modify the presentation of this work to better clarify the contributions of this work. 2. In the experiment, the results are limited. For example, in the meta-learning application, I suggest the authors compare the proposed method with the state-of-the-art bilevel optimization methods [1][2][3], which are shown to be able to address the meta-learning task. Alternatively, you need to specify why these methods [1][2][3] are not applicable to this application. Furthermore, the authors should conduct experiments on more datasets to better evaluate the proposed method, for example, Omniglot dataset. 3. I suggest the author briefly introduce some existing bilevel optimization works in machine learning, for example, hyper-gradient-based methods [1, 4] and approximation-based methods [2], and then discuss why these methods fail to be applied to the Contextual Stochastic Bilevel Optimization Problem. [1] Bilevel optimization: Convergence analysis and enhanced design, ICML, 2021 [2] Asynchronous Distributed Bilevel Optimization, ICLR 2023 [3] Bilevel Programming for Hyperparameter Optimization and Meta-Learning, ICML 2018 [4] Provably faster algorithms for bilevel optimization, NeurIPS 2021

Questions

1. I think it's necessary to emphasize the difficulty of solving the Contextual Stochastic Bilevel Optimization Problem compared with traditional bilevel optimization problems. The presentation of this work is poor. I believe this is an excellent work, I suggest that you should modify the presentation of this work to better clarify the contributions of this work. 2. In the experiment, the results are limited. For example, in the meta-learning application, I suggest the authors compare the proposed method with the state-of-the-art bilevel optimization methods [1][2][3], which are shown to be able to address the meta-learning task. Alternatively, you need to specify why these methods [1][2][3] are not applicable to this application. Furthermore, the authors should conduct experiments on more datasets to better evaluate the proposed method, for example, Omniglot dataset. 3. I suggest the author briefly introduce some existing bilevel optimization works in machine learning, for example, hyper-gradient-based methods [1, 4] and approximation-based methods [2], and then discuss why these methods fail to be applied to the Contextual Stochastic Bilevel Optimization Problem. [1] Bilevel optimization: Convergence analysis and enhanced design, ICML, 2021 [2] Asynchronous Distributed Bilevel Optimization, ICLR 2023 [3] Bilevel Programming for Hyperparameter Optimization and Meta-Learning, ICML 2018 [4] Provably faster algorithms for bilevel optimization, NeurIPS 2021

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

2 fair

Presentation

1 poor

Contribution

2 fair

Limitations

NA

Summary

**Summary:** The paper introduces a novel algorithm for solving contextual stochastic bilevel optimization (CSBO) problems. The authors develop DL-SGD and RT-MLMC gradient estimators for the problem along with a SGD-based algorithm for solving the CSBO problem. The authors analyze the properties of the gradient estimators and provide finite-time convergence guarantees for the proposed method. The theoretical results are corroborated by experiments on MAML, DRO with side information, and IV regression problems.

Strengths

**Strengths:** - The problem formulation considered in the paper is challenging and is of significant interest to the ML community. - The paper is well written, and the ideas are presented clearly with discussions. - The authors provide theoretical guarantees for the proposed approaches with an analysis of the proposed gradient estimators. - The authors have conducted experiments on multiple tasks including MAML, DRO with side information, and IV regression problem to evaluate the performance of the proposed framework.

Weaknesses

**Weaknesses:** 1. A major confusion I have is about the unbiasedness of the gradient estimator stated after line 127. Specifically, the expressions $\nabla_1 y^\ast$ and $\nabla_2 f$ both depend on the random variable $\xi$. This implies that the two expressions are dependent on each other which further means that the expectation given in the equation (after line 127) will not be the same as $\nabla F$. Can the authors clarify why the given expression will be true? The same discussion holds for the approximate gradient expressions derived in eq (5) and the rest of the paper. This is a major issue since the proofs and the results are based on the independence of $\nabla_1 y^\ast$ and $\nabla_2 f$. \ If this issue is resolved, I am willing to raise my score. 2. The authors should clarify the intuition behind using EpochSGD to estimate $y^\ast$. In the discussion, the authors simply explain the algorithm without explaining the intuition. 3. The gradient inverse estimator utilized from Ghadhimi’s work is clear, however, why the iterative estimator RT-MLMC works is difficult to understand. Can the authors please explain the working of the proposed estimator, i.e., why it works? 4. The experiments on MAML and DRO do not compare the proposed approach with baseline algorithms in the area. - Numerous bilevel algorithms solve MAML, the authors should compare the performance of the proposed approach against at least a few of them. - Similarly, the authors need to show the performance of the proposed approach for solving the DRO problem against popular baselines. 5. It would be easier if the authors keep the notations consistent in the experiment and the theory section of the paper. 6. To motivate the problem better the authors should include some examples in the introduction section. This way it will be easier for the reader to appreciate the considered formulation. ---- Updated the score after the rebuttal.

Questions

Please see the limitations above.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

n/a

Summary

The paper investigates a generalization of the stochastic bilevel optimization model, in which the lower and upper optimization levels share a random variable. The authors design two gradient-based approaches named RT-MLMC and DL-SGD for solving this problem and analyze their performance. The two methods differ in the way they estimate the gradient of the objective: DL-SGD is rather straightforwardly based on existing results, while RT-MLMC improves upon the performance of DL-SGD by a carefully tailored sampling method. Finally, numerical examples are presented to demonstrate the generality and performance of the proposed model and methods.

Strengths

* The paper introduces a general model that is applicable in a wide range of situations. Although the model is more general than previous approaches, the proposed approximations appear to be as efficient as the methods developed for the special cases. * The paper is clear, focused, and well-written. * I have verified most of the math; it is easy to follow and, aside from a few typos, appears to be correct.

Weaknesses

There are several typos in critical parts of the proofs (see "Questions" below). I believe that these typos do not affect the correctness of the proofs, however, this somewhat lowers my confidence in the results.

Questions

1. Is the dependence by $f$ on $\eta$ in the upper level of (1) necessary? I believe that removing it does not hurt the generality of the model and may avoid some confusion (see typos below). 2. In Lemma 1, $\alpha$ should be defined. 3. Typos: * In (5), $\eta'$ is used as the argument for the $\nabla_1 f$ and $\nabla^2_{12}g$ terms while $\eta''$ is used in the $\nabla_2 f$ term. This is in contrast with the second formula for $\nabla F$ on Page 4, where an independent RV is used in $\nabla^2_{12}g$. * A similar typo appears on the bottom of Page 15, in the proof of Lemma 2, where $\eta$ in the term $\nabla^2_{12}g$ should be changed to $\eta'$ (as in the definition of $V(x)$ above). * A similar typo appears on Page 16, where $H_K(1)$ and the last term in $H_K(2)$ should share the same variable $\eta$ (compare the formula for $V(x)$ on page 16 versus the definition of $V(x)$ on Page 15).

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

4 excellent

Contribution

4 excellent

Limitations

N/A

Thank you for the response

I thank the authors for the detailed response. Most of my concerns have been addressed by the authors, especially, the ones concerning the unbiasedness of the gradient estimator. Consequently, I have updated my original rating of the paper.

Thanks for the response

Dear Authors, Thanks so much for the response! My questions have been answered satisfactorily, so I increase my score. Best, Reviewer

Acknowledgment

I thank the authors for their responses. Regarding question #1, I see no loss in generality in removing $\xi$ from the upper level since we can always include a copy of $\xi$ as part of $\eta$ (i.e., set $\eta' = (\eta, \xi)$). In any case, this is a very minor point and the authors should choose the form they prefer best.