Distributed Stochastic Bilevel Optimization: Improved Complexity and Heterogeneity Analysis
Youcheng Niu, Jinming Xu, Ying Sun, Yan Huang, Li Chai
TL;DR
This work tackles distributed stochastic bilevel optimization with personalized inner problems, formulating $\Phi(x)=\frac{1}{m}\sum_i f_i(x,\theta_i^*(x))$ where $\theta_i^*(x)=\arg\min_\theta g_i(x,\theta)$. It introduces LoPA, a loopless, communication-efficient algorithm with two variants: LoPA-LG (local gradient) and LoPA-GT (gradient tracking), including a gradient momentum mechanism to mitigate hypergradient bias. The authors provide a comprehensive heterogeneity-aware convergence analysis, deriving rate bounds that explicitly depend on condition number $\kappa$, network gap $\rho$, heterogeneity $b$, and gradient variances $\sigma_p, \sigma_c$; gradient tracking further reduces heterogeneity effects and improves rates. They prove that LoPA attains $\mathcal{O}(\epsilon^{-2})$ complexity in terms of Hessian evaluations, surpassing prior DSBO methods, and validate the theory with numerical experiments on distributed classification and hyperparameter optimization.
Abstract
This paper consider solving a class of nonconvex-strongly-convex distributed stochastic bilevel optimization (DSBO) problems with personalized inner-level objectives. Most existing algorithms require computational loops for hypergradient estimation, leading to computational inefficiency. Moreover, the impact of data heterogeneity on convergence in bilevel problems is not explicitly characterized yet. To address these issues, we propose LoPA, a loopless personalized distributed algorithm that leverages a tracking mechanism for iterative approximation of inner-level solutions and Hessian-inverse matrices without relying on extra computation loops. Our theoretical analysis explicitly characterizes the heterogeneity across nodes (denoted by $b$), and establishes a sublinear rate of $\mathcal{O}( {\frac{1}{{{{\left( {1 - ρ} \right)}}K}} \!+ \!\frac{{(\frac{b}{\sqrt{m}})^{\frac{2}{3}} }}{{\left( {1 - ρ} \right)^{\frac{2}{3}} K^{\frac{2}{3}} }} \!+ \!\frac{1}{\sqrt{ K }}( {σ_{\operatorname{p} }} + \frac{1}{\sqrt{m}}{σ_{\operatorname{c} }} ) } )$ without the boundedness of local hypergradients, where ${σ_{\operatorname{p} }}$ and ${σ_{\operatorname{c} }}$ represent the gradient sampling variances associated with the inner- and outer-level variables, respectively. We also integrate LoPA with a gradient tracking scheme to eliminate the impact of data heterogeneity, yielding an improved rate of ${\mathcal{O}}(\frac{1}{ (1-ρ)^2K } \!+\! \frac{1}{\sqrt{K}}( σ_{\rm{p}} \!+\! \frac{1}{\sqrt{m}}σ_{\rm{c}} ) )$. The computational complexity of LoPA is of ${\mathcal{O}}({ε^{-2}})$ to an $ε$-stationary point, matching the communication complexity due to the loopless structure, which outperforms existing counterparts for DSBO. Numerical experiments validate the effectiveness of the proposed algorithm.
