On the Convergence of Adam-Type Algorithms for Bilevel Optimization under Unbounded Smoothness

Authors: Xiaochuan Gong, Jie Hao, Mingrui Liu
Primary area: optimization
Keywords: Bilevel Optimization, Adam, Unbounded Smoothness

Abstract:
Adam has become one of the most popular optimizers for training modern deep neural networks, such as transformers. However, its applicability is largely restricted to single-level optimization problems. In this paper, we aim to extend vanilla Adam to tackle bilevel optimization problems, which have important applications in machine learning, such as meta-learning. In particular, we study stochastic bilevel optimization problems where the lower-level function is strongly convex and the upper-level objective is nonconvex with potentially unbounded smoothness. This unbounded smooth objective function covers a broad class of neural networks, including transformers, which may exhibit non-Lipschitz gradients. In this work, we first introduce AdamBO, a single-loop Adam-type method that achieves $\widetilde{O}(\epsilon^{-4})$ oracle complexity to find $\epsilon$-stationary points, where the oracle calls involve stochastic gradient or Hessian/Jacobian-vector product evaluations. The key to our analysis is a novel randomness decoupling lemma that provides refined control over the lower-level variable. Additionally, we propose VR-AdamBO, a variance-reduced version with an improved oracle complexity of $\widetilde{O}(\epsilon^{-3})$. The improved analysis is based on a novel stopping time approach and a careful treatment of the lower-level error. We conduct extensive experiments on various machine learning tasks involving bilevel formulations with recurrent neural networks (RNNs) and transformers, demonstrating the effectiveness of our proposed Adam-type algorithms.

Decision: reject
Reviewer scores: 6, 3, 6, 6

Reviews and discussion:
<thread>
creation date:
1730697668104

thread_type:
official_review

<thread_object>
creator: reviewer_x4qj

<content>
summary:
The paper presents AdamBO, a novel extension of the Adam optimizer designed for bilevel optimization problems, particularly when the upper-level function exhibits unbounded smoothness. The authors address the challenges inherent in bilevel optimization, such as the dependency between the upper-level hypergradient and lower-level variable. The work also introduces VR-AdamBO, a variance-reduced variant, to enhance convergence efficiency. Both algorithms demonstrate provable theoretical convergence guarantees and are validated through experiments on tasks involving RNNs and transformers, showing significant performance improvements.

The authors also proves that AdamBO converges to $\varepsilon$-stationary points with $O(\varepsilon^{-4})$ oracle complexity. The variance-reduced variant VR-AdamBO has an improved oracle complexity of $O(\varepsilon^{-3})$.

strengths:
1. The authors provide comprehensive and solid convergence proofs for both AdamBO and VR-AdamBO, incorporating novel elements such as a randomness decoupling lemma and stopping-time analysis for variance reduction. These theoretical insights are significant and may be of independent interest for future research.
2. The algorithms are tested on relevant machine learning tasks, including meta-learning and deep AUC maximization, showing strong empirical performance.

weaknesses:
1. Instead of stacking all the assumptions and lemmas, I think a better way to present your theoretical results is that placing the detailed proof sketch in the main text and moving some assumptions and lemmas to the appendix. Also, it would be better if you can clearly state the difficult part of the proofs as well as the comparison and difference with other relevant works. The current version is quite difficult for the audience to understand the entire line of the proof.

questions:
Please refer to the "weaknesses" part for the questions.
</content>
</thread_object>

</thread>
<thread>
creation date:
1731032906238

thread_type:
official_review

<thread_object>
creator: reviewer_nmmt

<content>
summary:
In this paper, the authors extend to Adam to solve bilevel optimization problems. The authors proposed two Adam-type algorithms, AdamBO and VR-AdamBO that is a variance-reduced version with an improved oracle complexity of the first one. The authors analysis a novel decoupling lemma that provides refined control over the lower-level variable. Finally, the authors test the effectiveness of their proposed algorithms on some machine learning tasks.

strengths:
1. the authors provide solid theoretical guarantees for the convergence of their proposed Adam-type algorithms under the unbounded smoothness assumption and other standard assumptions in current bilevel optimization analysis.

2. A novel random decoupling lemma that is proposed by the authors provides a technical tool to control the lower-level error without concerns about the dependency issues from the upper-level randomness. This lemma can be extended to other analysis of Adam-type algorithms for bilevel optimizaition problems.

weaknesses:
1. The author did not provide a clear comparison of the assumptions and the complexity of the newly proposed algorithms with popular algorithms that appeared in previous literature in the paper.

2. The theoretical convergence results provided by the author, whether in AdamBO or VR-AdamBO algorithms, cannot guarantee that the first-order momentum coefficient is a constant, that is, the coefficient depends on the accuracy level $\epsilon$.

questions:
1. In your assumption 3.5: “….. almost surely ”, what does “almost surely” mean?

2. In the part of VR-AdamBO, you state that “VR-AdamBO can be regarded as a generalization of the AccBO….., with the key difference being that AccBO uses normalized SGD with momentum for the upper-level update,…”, so do you mean that you adopted the same variance reduction technique as AccBO, the only difference being the way you update the upper level variables?
</content>
</thread_object>

</thread>
<thread>
creation date:
1731119027512

thread_type:
official_review

<thread_object>
creator: reviewer_sze1

<content>
summary:
This paper extends the Adam optimizer to address bilevel optimization problems. The authors introduce AdamBO, a single-loop, Adam-based algorithm for bilevel problems with a strongly convex lower-level function and a nonconvex upper-level function, achieving an iteration complexity of O(epsilon^{-4}). They also propose VR-AdamBO, a variance-reduced version that further improves oracle complexity to O(epsilon^{-3}). The proof utilized a randomness decoupling lemma and stopping time approach. The paper demonstrates the effectiveness of these algorithms on one meta-learning dataset.

strengths:
Extend Adam to bilevel optimization problems, and provide some convergence anlysis.

weaknesses:
i) The paper requires the coefficient of the learning rate to be 1/(sqrt{v} + lambda) where lambda >0 (it is often called epsilon in Adam papers).  
  i.1) This makes the analysis rather weak. In fact, with a positive lambda, the analysis will avoid a major challenge of lower bounding sqrt{v}, making the analysis much easier. 
  i.2) There are quite a few constants that depend on lambda >0:
   a) beta < O(lambda);
   b) stepsize eta < O(lambda^1.5 beta) = O(lambda^2.5). 
  This makes the effective stepsize always super small (it can be proved that the effective stepsize eta/(sqrt{v} + beta) < eta/beta < lamda^1.5). Then it is somewhat easy to prove Adam converges with an upper bounded stepsize. 

ii) The paper requires beta (which is 1 - beta1) to be very small, i.e. the typical beta1 is very close to 1, and the new gradient information comes very slowly. This is a very special variant of Adam, making the result not applicable to general Adam. This significantly weakens the claim of the paper that "Adam-based  bilevel algorithm is proved to converge". 

iii) Experiments are relatively weak. AUC minimization of binary text classification is not a common task,and it only uses 3-layer RNN and transformers. Meta-learning on  SNLI with 3-layer RNN and transformers is not that convincing.

questions:
i) Why do you require beta to be very small? Can you get around the assumption? 

ii) Can you get around the asumption of lambda > 0, or remove the dependence of eta on lambda?
</content>
</thread_object>

</thread>
<thread>
creation date:
1731131539765

thread_type:
official_review

<thread_object>
creator: reviewer_8c3x

<content>
summary:
In this paper, the authors propose a variant of Adam called AdamBO, specifically designed for bilevel optimization problems. AdamBO updates the parameters of both the lower-level and upper-level functions simultaneously within a single loop. The authors provide a theoretical convergence guarantee for the hypergradient under an unbounded smoothness condition, which os common in recent theoretical analyses. Additionally, the authors introduce a variance reduction version of AdamBO that achieves improved oracle complexity under stronger conditions.

strengths:
- This paper gives the first study to investigate the convergence of Adam-type algorithms in the context of bilevel optimization problems. The assumptions made for the theoretical convergence analysis are mostly mild and standard. The results presented are convincing and well-supported by the explanations in the proof sketch.

- The writing in the introduction is relatively clear and comprehensive. Readers can easily follow the motivation behind this study, and the proposed algorithms are explained very clearly.

- In addition to the theoretical guarantees for gradient convergence, the authors provide convincing experimental results that support the performance of AdamBO in bilevel optimization problems.

weaknesses:
- The condition in Theorems 4.1 and 5.1 that states $\beta=\tilde \Theta(\epsilon^2)$ is unusual. The default setting for $\beta$ of Adam is 0.1 in Pytorch, and this condition appears to imply poor convergence results for such a default choice of $\beta$. 

- Some variables are not provided a definition when introduced, which causes difficulties in understanding. For instance, notations related to stochastic gradients are not explained clearly: $\zeta^{(q,j)}$'s at line 151 seems undefined and a formal definition for this notation needs to be added, particularly regarding the meanings of the superscripts $q$ and $j$.

questions:
See weakness.
</content>
</thread_object>

</thread>
<thread>
creation date:
1731836720785

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response to All Reviewers

comment:
Thank you to all the reviewers for taking the time to review our paper and provide valuable feedback. We have addressed each of your concerns individually and summarize the key changes made during the rebuttal phase below. Major modifications are highlighted in blue in Appendices G, H, and I of the revised paper.

1. We have explained why we need a small $\beta$ (i.e., $\beta$ depends on the target gradient norm $\epsilon$) in theory. First, small $\beta$ is indeed a common assumption for convergence analysis of Adam in the literature ([Li et al., 2023](https://arxiv.org/pdf/2304.13972); [Guo et al., 2021b](https://arxiv.org/pdf/2104.14840)). In addition, the purpose of using small $\beta$ is to control the hypergradient bias caused by the lower-level estimation error in the bilevel optimization setting as in ([Chen et al., 2023](https://arxiv.org/pdf/2306.12067); [Hao et al., 2024](https://arxiv.org/pdf/2401.09587.pdf); [Gong et al., 2024a](https://openreview.net/pdf?id=36rWa8zVkh); [b](https://arxiv.org/pdf/2409.19212)). Second, existing convergence analyses of Adam that do not require this choice of $\beta$ rely on other strong assumptions for the objective function, such as the bounded gradient assumption ([De et al., 2018](https://arxiv.org/pdf/1807.06766); [Défossez et al., 2020](https://arxiv.org/pdf/2003.02395)) or the strong growth condition ([Zhang et al., 2022](https://arxiv.org/pdf/2208.09632); [Wang et al., 2022](https://arxiv.org/pdf/2208.09900)). These assumptions are incompatible with our setting.

2. In Appendix G, we have added two tables: **Table 1** (on page 68) compares Adam-related papers under different settings and assumptions, and **Table 2** (on page 68) compares various bilevel optimization algorithms under the unbounded smoothness setting.

3. In Appendix H, we have conducted additional meta-learning experiments using an 8-layer BERT ([Devlin et al., 2018](https://arxiv.org/pdf/1810.04805)) on the widely-used question classification dataset TREC ([Li & Roth, 2002](https://aclanthology.org/C02-1150.pdf)). The results, shown in **Figure 4** (on page 69), demonstrate that our proposed algorithm AdamBO achieves fast convergence to the best training and testing performance among all baselines.

4. In Appendix H, we have also conducted experiments in **Figure 5** (on page 69) to show the empirical performance of our algorithm is not very sensitive to the choice of $\lambda$. Although the default choice of $\lambda$ is $10^{-8}$ ([Kingma & Ba, 2014](https://arxiv.org/pdf/1412.6980)), increasing it up to $10^{-4}$ only causes minor differences in AUC maximization, and increasing it up to $10^{-3}$ leads to minor changes in hyper-representation performance with BERT ([Devlin et al., 2018](https://arxiv.org/pdf/1810.04805)).

5. In Appendix I, we have provided a more detailed proof sketch for VR-AdamBO (Theorem 5.1).

Please let us know if you have additional concerns, and we are more than willing to address any further questions.

Best,

Authors
</content>
</thread_object>

</thread>


Metareview:
metareview:
After careful consideration of the reviews and subsequent author responses, the consensus among the reviewers remains that this work, while presenting an attempt to extend Adam-type methods to bilevel optimization, does not fully address the core challenges of analyzing Adam in bilevel settings.