Stochastic Bilevel Optimization with Heavy-Tailed Noise
This work addresses optimization challenges in machine learning applications like large language models and reinforcement learning, but it is incremental as it builds on existing bounded variance results.
This paper tackles stochastic bilevel optimization with heavy-tailed noise, proposing a method (N^2SBA) that achieves an ε-stationary point with a stochastic first-order oracle complexity of Õ(κ^{(7p-3)/(p-1)} σ^{p/(p-1)} ε^{-(4p-2)/(p-1)}), and extends it to minimax optimization with similar complexity bounds.
This paper considers the smooth bilevel optimization in which the lower-level problem is strongly convex and the upper-level problem is possibly nonconvex. We focus on the stochastic setting that the algorithm can access the unbiased stochastic gradient evaluation with heavy-tailed noise, which is prevalent in many machine learning applications such as training large language models and reinforcement learning. We propose a nested-loop normalized stochastic bilevel approximation (N$^2$SBA) for finding an $ε$-stationary point with the stochastic first-order oracle (SFO) complexity of $\tilde{\mathcal{O}}\big(κ^{\frac{7p-3}{p-1}} σ^{\frac{p}{p-1}} ε^{-\frac{4 p - 2}{p-1}}\big)$, where $κ$ is the condition number, $p\in(1,2]$ is the order of central moment for the noise, and $σ$ is the noise level. Furthermore, we specialize our idea to solve the nonconvex-strongly-concave minimax optimization problem, achieving an $ε$-stationary point with the SFO complexity of $\tilde{\mathcal O}\big(κ^{\frac{2p-1}{p-1}} σ^{\frac{p}{p-1}} ε^{-\frac{3p-2}{p-1}}\big)$. All above upper bounds match the best-known results under the special case of the bounded variance setting, i.e., $p=2$.