Provably Convergent Federated Trilevel Learning
Yang Jiao, Kai Yang, Tiancheng Wu, Chengtao Jian, Jianwei Huang
TL;DR
The paper tackles privacy-preserving trilevel optimization by introducing AFTO, an asynchronous federated framework that uses a novel hyper-polyhedral approximation built from $μ$-cuts to handle $μ$-weakly convex objectives. It establishes non-asymptotic convergence with a rate of $T(ε)=O\left(\frac{1}{ε^2}\right)$ to obtain an $ε$-stationary point and demonstrates substantial empirical acceleration (up to ~80%) on real-world tasks. By reformulating the TLO as a consensus problem and layering two sets of cutting planes, the approach enables distributed optimization with stragglers and bounded staleness, while refining the approximation via iterative $μ$-cuts and pruning inactive planes. The method is validated on distributed robust hyperparameter optimization and domain adaptation, showing faster convergence and superior performance against state-of-the-art baselines. Overall, the work advances scalable, privacy-preserving TLO with concrete convergence guarantees and practical impact for complex ML systems.
Abstract
Trilevel learning, also called trilevel optimization (TLO), has been recognized as a powerful modelling tool for hierarchical decision process and widely applied in many machine learning applications, such as robust neural architecture search, hyperparameter optimization, and domain adaptation. Tackling TLO problems has presented a great challenge due to their nested decision-making structure. In addition, existing works on TLO face the following key challenges: 1) they all focus on the non-distributed setting, which may lead to privacy breach; 2) they do not offer any non-asymptotic convergence analysis which characterizes how fast an algorithm converges. To address the aforementioned challenges, this paper proposes an asynchronous federated trilevel optimization method to solve TLO problems. The proposed method utilizes $μ$-cuts to construct a hyper-polyhedral approximation for the TLO problem and solve it in an asynchronous manner. We demonstrate that the proposed $μ$-cuts are applicable to not only convex functions but also a wide range of non-convex functions that meet the $μ$-weakly convex assumption. Furthermore, we theoretically analyze the non-asymptotic convergence rate for the proposed method by showing its iteration complexity to obtain $ε$-stationary point is upper bounded by $\mathcal{O}(\frac{1}{ε^2})$. Extensive experiments on real-world datasets have been conducted to elucidate the superiority of the proposed method, e.g., it has a faster convergence rate with a maximum acceleration of approximately 80$\%$.