将Adam算法扩展至双层优化，解决Transformer训练难题
On the Convergence of Adam-Type Algorithm for Bilevel Optimization under Unbounded Smoothness
- 提出单循环Adam型算法AdamBO处理双层优化问题
- 在非凸上层目标下达到ε-平稳点的˜O(ε⁻⁴)查询复杂度
- 适用于RNN与Transformer等具有非Lipschitz梯度的模型
Adam已成为现代深度神经网络（如Transformer）训练中最流行的优化器之一，但其应用主要局限于单层优化问题。本文旨在将原始Adam扩展至双层优化场景，该问题在机器学习中具有重要应用，如元学习。我们研究的是下层函数强凸、上层目标非凸且可能具有无界光滑性的随机双层优化问题，此类目标涵盖广泛神经网络（包括Transformer），其梯度可能不满足Lipschitz条件。本文提出AdamBO，一种单循环的Adam型方法，可在˜O(ε⁻⁴)的查询复杂度内找到ε-平稳点，查询操作包括随机梯度或海森/雅可比-向量积的评估。分析核心在于一个新颖的随机性解耦引理，实现了对下层变量更精细的控制。我们在涉及RNN和Transformer的多种双层学习任务上进行了大量实验，验证了所提算法的有效性。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。