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

TL;DR

This work addresses bilevel optimization under unbounded smoothness by introducing AdamBO, a single-loop Adam-type method that updates the upper-level variable with Adam while warm-starting and updating the lower-level variable via SGD. The authors establish that AdamBO achieves ~O(b5^{-4}) oracle complexity to reach an b5-stationary point, despite non-negligible hypergradient bias arising from lower-level estimation, by proving a novel randomness decoupling lemma that separates upper- and lower-level randomness. The method is complemented by a refined analysis of hypergradient bias and a stopping-time framework, enabling high-probability convergence guarantees. Empirical results on meta-learning with hyper-representations and deep AUC maximization with RNNs/Transformers demonstrate faster convergence and substantial performance gains over existing bilevel baselines, underscoring the practical impact of adaptive, Adam-style updates in bilevel contexts. The work also discusses the dependence on the smoothing parameter bbeta and blambda, pointing to future directions to reduce this dependence while preserving strong empirical performance.

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 introduce AdamBO, a single-loop Adam-type method that achieves $\widetilde{O}(ε^{-4})$ oracle complexity to find $ε$-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. 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 algorithm.

Figures (5)

• Figure 1: Comparison with bilevel optimization baselines on BERT for hyper-representation.
• Figure 2: Comparison with bilevel optimization baselines on RNN for hyper-representation.
• Figure 3: Transformer for AUC maximization on Sentiment140 dataset with imbalance ratio of 0.9.
• Figure 4: RNN for AUC maximization on Sentiment140 dataset with imbalance ratio of 0.8. Figures (a), (b) are the results over epochs, Figures (c), (d) are the results over running time.
• Figure 5: Test accuracy of different models on AUC maximization and hyper-representaion using AdamBO with $\beta=0.1, \beta_{\mathop{\mathrm{sq}}\nolimits}=0.001$ and different $\lambda$s. (a) a 2-layer RNN model on AUC maximization (data imbalanced ratio = 0.8); (b) a 2-layer Transformer model on AUC maximization (data imbalanced ratio = 0.9); (c) an 8-layer BERT model on hyper-representation.

Theorems & Definitions (70)

• Definition 3.1: $(L_{x,0}, L_{x,1}, L_{y,0}, L_{y,1})$-Smoothness hao2024bilevel
• Theorem 4.1
• Lemma 4.2: Randomness Decoupling Lemma
• Lemma 4.3: Warm-Start
• Lemma 4.4
• Lemma 4.5
• Lemma 4.6
• Proposition 1.1
• proof : Proof of \ref{['prop:bi-adam-equiv']}
• Lemma 2.1