A Stochastic Linearized Augmented Lagrangian Method for Decentralized Bilevel Optimization
Abstract
Bilevel optimization has been shown to be a powerful framework for formulating multi-task machine learning problems, e.g., reinforcement learning (RL) and metalearning, where the decision variables are coupled in both levels of the minimization problems.
In practice, the learning tasks would be located at different computing resource environments, and thus there is a need for deploying a decentralized training framework to implement multi-agent and multi-task learning.
We develop a stochastic linearized augmented Lagrangian method (SLAM) for solving general nonconvex bilevel optimization problems over a graph, where both upper and lower optimization variables are able to achieve a consensus.
We also establish that the theoretical convergence rate of the proposed SLAM to the Karush-Kuhn-Tucker (KKT) points of this class of problems is on the same order as the one achieved by the classical distributed stochastic gradient descent for only single-level nonconvex minimization problems.
Numerical results tested on multi-agent RL problems showcase the superiority of SLAM compared with the benchmarks.
Introduction
In this paper we consider the following general decentralized bilevel optimization (DBO) framework with applications to machine learning problems.
Suppose that there are n nodes over a connected graph G = {E , V} , where E and V represent the edges and vertices.
Let N i denote the set of neighboring nodes for node i .
Then the goal of DBO is to have these nodes jointly minimize two levels of optimization problems.
More formally, DBO is expressed as
[Equation]
where vector x i is the upper level (UL) optimization variable at each node i , vector y i,k denotes the lower level (LL) decision variable for the k th learning task at node i , f i (; ) is a (smooth) UL loss function and possibly nonconvex with respect to (w.r.t.)
both the UL and LL variables, g i,k ( , ) denotes the LL objective function of the k th task at node i , m represents the total number of LL optimization problems, the consensus constraints x i = x
j , y i,k = y j,k , j ∈ N i , ∀ i ∈ [ n ] , ∀ k ∈ [ m ] , enforce the model agreements at each level of the problems and for each LL learning task, and y ∗ k
= [ y ∗ 1 ,k , . . .
, y ∗ n,k ] T is the optimal solutions of the k th LL problem under the consensus constraints.
Applications of Bilevel Optimization.
Many machine learning problems can be formulated mathematically as a form of bilevel optimization or, more precisely, a special case of problem (1), e.g., meta-learning or meta reinforcement learning (RL), actor-critic (AC) schemes in RL, hyperparameter optimization (HPO), and so on.
Classical bilevel optimization is referred to as the case where there is no consensus constraint but with only two levels of the minimization subproblems, i.e., min x f ( x , y ∗ ( x )) , s.t.
y ∗ ( x ) =
arg min y g ( x , y ) , which is also known as Stackelberg games [1] with the UL decision variable as the leader and the LL decision variable as the follower.
It turns out that this class of optimization problems is useful in formulating a wide range of hierarchical or nested structured machine learning problems.
For example, one of the most popular domain adaption learning models, model-agnostic meta-learning (MAML) [2, 3], can be written as a special case of bilevel programming [4], where the UL model provides a good initialization for accelerating learning procedures by implementing the LL algorithms.
The idea behind the model design is that the UL model is considered as the meta learner that searches for a permutation-invariant subspace over multiple task-specific learners at the
LL so that the performance of the MAML model can be generalized well for unseen or testing data samples.
The theoretical analysis of the generalization performance of this class of bilevel problems has shown that MAML can indeed decrease the generalization error as the number of tasks increases, at least for strongly convex loss functions [5].
Subsequently, a thorough ablation study from the latent representation perspective shows that feature reuse is the actual dominant factor in improving the generalization performance of MAML [6], and the authors propose a neural network-oriented algorithm with almost no inner loop (ANIL) that splits the neural network parameters into two parts corresponding to the UL and LL optimization problems, respectively.
Extensive numerical experiments illustrate that ANIL achieves almost the same accuracy as the classical MAML but with significant computational savings.
This example further strengthens the necessity of variable splitting in the learning structure by optimizing two levels of objective functions to enhance the generalization performance.
Beyond the traditional supervised meta-learning scenarios, MAML has also been applied to increasing the generalization ability of agents in RL problems by replacing the (stochastic) gradient with the (natural) policy gradient (PG) [3] under the same two-level structure.
Besides meta-learning problems, AC structure in RL is another class of common learning frameworks that can be formulated by a bilevel optimization problem in nature [7, 8, 9], where the actor step at the UL aims at optimizing the policy while the critic step at the LL is responsible for value function evaluation.
In addition, as the expressiveness of neural networks increased sharply over the past decades, the reuse of large models with adaptation to multi-task learning problems presents promising solutions by leveraging the pre-train and fine-tune strategy, such as in applications of HPO [10, 11] where the hyperparameters are trained at the UL problem so that the downstream learning tasks are learned with low costs including the expense of both computation and memory.
Applications of Multi-agent Settings.
When multiple computational resources are available and connected, it is well motivated that exploring them solves distributed large-scale problems with a reduced amount of training time or performs multi-task learning.
The bilevel structure of the meta-learning (ML) is a good fit in this scenario as either UL/LL or both levels may need to access the networked data samples rather than local ones.
For example, a federated learning setting of
MAML [12] and bilevel optimization [13] have been built up over multiple nodes recently, where the meta/UL learner finds an initial shared model while the local/LL learners leverage it for adapting data distributions of individual users.
In such a way, the federated MAML model can realize personalized learning without sharing heterogeneous data over numerous clients.
Once there is no central controller for coordinating the model aggregation, a diffusion-based MAML (Dif-MAML) [14] is proposed by spreading the model parameters over a network, where the UL parameter is updated by one step of stochastic gradient descent (SGD) based on a combination of the parameters of neighbors as the initialization for local model updates.
Decentralized hierarchical structured learning is even more stringent in the multi-agent RL (MARL) setting [15] as the learning tasks are essentially located at scattered sensors and/or controllers.
Under this setting, MARL problem becomes a multi-objective optimization problem under provided (approximate) value functions, where the policy of each agent needs to be learned locally by certain efficient iterative methods, such as multi-agent deep deterministic policy gradient (MADDPG) [16], trust region methods [17], optimal baseline based variance reduced policy gradient [18], and/or improved by more advanced techniques, e.g., constrained policy optimization [19] and large sequence models [20].
In such a way, the total reward can be maximized over the distributed agents through optimizing the networked policy.
In a fully collaborative setting, the team-based value function is even required to be shared over all the agents such that each agent is able to improve its policy based on the estimated total reward.
For example, the decentralized AC (DAC) scheme has been investigated widely [15, 21, 22], where each agent uses the actor step to optimize its policy while the critic step performs one step [23] or multiple steps of temporal difference learning with mini-batch sampling (MDAC)
[22, 24] and communications so that the team-based reward over the network is obtained by each agent.
It turns out that DAC can be formulated as a special case of problem (1) as there is no consensus at the UL.
Recently, it has been revealed that if there exists homogeneity of the state and action spaces, decentralized policy consensus (or a partial policy parameter sharing strategy) provides significant merits to the centralized training and decentralized execution paradigm in terms of learning scalability and efficiency[23, 25], which motivates the consensus process at both
UL and LL DBO problems.
Related Theoretical Works.
Given the fruitful results across these many applications, the corresponding theoretical analysis has been developing very fast as well for variants of bilevel optimization problems.
For example, the convergence behaviors of classical inexact MAML (iMAML) methods have been quantified for both convex [26, 27] and nonconvex [28] cases of the UL loss function, where the LL algorithm only performs one step of stochastic gradient descent (SGD) based on the LL objective functions as the adaptation step.
Moreover, the iteration complexity of ANIL with multiple iterations for minimizing the LL problems have been studied in [29], which justifies the significant computational advantages of ANIL compared with MAML in theory.
Furthermore, the finite-time analysis of AC algorithms has shown [30] that, once the learning rates at both the actor and critic sides are chosen properly, a two timescale AC algorithm can achieve an O ( ϵ − 2 . 5 ) iteration complexity for finding the first-order stationary points (FOSPs) of general nonconvex reward functions.
Besides these theoretical analyses in a specific learning setting, the algorithm design and corresponding convergence analysis for general bilevel optimization solvers have been recently advancing at a rapid pace under certain assumptions that the UL objective function is general nonconvex while the LL objective functions are strongly convex, which covers the existing convergence results shown for AC algorithms.
The typical algorithms include those with double-loop structure, those with two timescale or single timescale but single-loop, and those with error-correction or accelerated/variance-reduction.
To be more specific, double-loop algorithms, such as bilevel stochastic approximation (BSA) methods [31] and stochastic bilevel optimizers (stoBiO) [32], mainly request an inner loop to solve the LL problem up to a certain error tolerance or with a certain number of iterations and then switch back to optimize the UL problem, which can achieve an O ( ϵ − 2 ) convergence rate to the ϵ -FOSPs.
In practice, single-loop algorithms are implemented more efficiently in terms of computational complexity and hyperparameter tuning compared to double-loop algorithms.
A two-timescale stochastic approximation (TTSA) was analyzed in [33], but it is shown that TTSA needs O ( ϵ − 2 . 5 ) number of iterations to achieve the ϵ -FOSPs.
Later, an error correction method, named the Single-Timescale stochAstic BiLevEl optimization (STABLE) method [34], improves the convergence rate of the single-loop algorithm to O ( ϵ − 2 ) and a tighter analysis for ALternating Stochastic gradient dEscenT (ALSET) shows that the single-loop algorithm can also achieve a convergence of O ( ϵ − 2 ) without the error correction technique.
When more advanced momentum-assisted or variance reduction methods are adopted in the algorithm design, the subsequent works, such as the momentum-based recursive bilevel optimizer (MRBO) [35] and the single-timescale double-momentum stochastic approximation (SUSTAIN) [36] and the variance reduced BiAdam (VR-BiAdam) [37], can sharpen the convergence rate of bilevel algorithms to O
( ϵ − 1 . 5 ) .
For the theoretical works on MAML/MARL, it is shown in [22, 24] that when the critic side is allowed the consensus step at each agent to approximate the networked rewards, MDAC algorithms can achieve an O ( ϵ − 2 ) convergence rate to FOSPs, but both of them require an inner loop procedure for the LL problem which makes the algorithms double loop.
Dif-MAML [14] is able to perform the UL consensus-based meta learning, but iMAML considered in Dif-MAML is only a very special case of bilevel.
Thus, the applicability of Dif-MAML is restrictive.
One of the closest works to ours is coordinated AC (CAC) [23], which can realize the consensus on both UL and LL problems with
O ( ϵ − 2 . 5 ) number of iterations and is only for DAC problems.
A theoretical comparison between our
[Table] [Table]
work and closely related previous works on bilevel programming is shown in Table 1.
There is a line of independent work on decentralized optimization.
But the existing works are only suitable for single-level minimization of only nonconvex problems, such as distributed SGD [38, 39], stochastic gradient tracking [40, 41] and stochastic primal dual algorithm [42], which can achieve an O (1 / ( nϵ 2 )) convergence rate to FOSPs for general nonconvex objective function optimization problems.
Main Contributions of This Work.
In this work, we consider a very general DBO setting, where both UL and LL problems can include a consensus constraint for model parameter sharing and there would be multiple LL problems coupled with the UL problem.
To solve this problem efficiently in a fully decentralized way, we propose a S tochastic L inearized A ugmented Lagrangian M ethod (SLAM) for dealing with both of the two levels of the optimization processes and the consensus constraints at each level.
Leveraging the linearized augmented Lagrangian function as a surrogate, the design of
SLAM is simple and easily implemented as it is a single-loop algorithm with only step sizes to be tuned for convergence.
We make the standard assumptions on Lipschitz continuity and convexity for both the UL and LL optimization problems as shown in the existing literature.
We establish the conditions of SLAM w.r.t.
convergence to ϵ -Karush-Kuhn-Tucker (KKT) points of problem (1) at a rate of O (1 / ( nϵ 2 )) , matching the standard convergence rate achieved by decentralized SGD type of algorithm to FOSPs for only single-level nonconvex minimization problems.
Remarkably, through numerical experiments on MARL problems, it is observed that SLAM can converge faster than the existing MARL methods and even achieve higher rewards in most cases.
To summarize, the main contributions of this work are highlighted as follows:
▶ Our proposed SLAM algorithm is generic, and thus generalizes the single agent-based bilevel algorithms to the multi-agent setting and is amnable to be specialized to solve multiple consensusbased DBO problems.
▶ SLAM is a single-timescale and single-loop algorithm that can find the ϵ -KKT points at a rate of O (1 / ( nϵ 2 )) , which shows a linear speedup w.r
t. the number of nodes.
To the best of our knowledge, this is the first work that shows a decentralized stochastic algorithm can achieve this rate under the constraints where any level or both levels of the DBO problem requires the consensus process.
▶ Numerical results that illustrate the proposed SLAM outperforms the state-of-the-art MARL algorithms over heterogeneous networks in terms of both convergence speed and achievable rewards.
Due to space limitations, all technical proofs are deferred to the supplement.
Decentralized Bilevel Optimization Framework
Problem formulation of DBO.
One of the main motivations for performing decentralized joint learning is dealing with large-scale dataset or scattered data samples.
At each node, the loss function can be written as f i ( x i , y ∗
i, 1 ( x i ) , . . .
, y ∗ i,m ( x i ))
≜
E ξ ∈D U i [ F i ( x i , y ∗
i, 1 ( x i ) , . . .
, y ∗ i,m ( x i ); ξ )] , where D U i denotes the local data distributions at the UL optimization problem, and
F i ( x i , y ∗ i, 1 ( x i ) , . . .
, y ∗ i,m ( x i ); ξ ) represents the estimation error of the UL learning model on data ξ ∈ D U i .
Similarly, the LL learning tasks also include randomly sampled data from a local distribution D L i,k for task k , so the LL cost function at each node can be expressed as g i,k
( x i , y i,k )
≜
E ζ ∈D L i,k
[ G i,k ( x i , y i ; ζ )] , ∀ k , where G i,k denotes the estimation error of the LL learning model over y k,i on data ζ ∈ D L i,k .
It is well known that SGD is one of the most efficient algorithms for tackling large amounts of data samples.
Before showing the algorithm design, we first reformulate problem (1) in a concise and compact way from a global view of the variables.
Let x ≜ [ x 1 , . . .
, x n ] T and y k ≜
[ y 1 ,k , . . .
, y n,k ] T .
Then, problem (1) can be rewritten by concatenated variables as
[Equation]
where g k ( x , y k ) denotes the k th LL loss function, A ∈ R |E|× n represents the incidence matrix 1 and f i ( x i , y ∗ i,k ( x i ))
abbreviates f i ( x i , y ∗ i, 1 ( x i ) , . .
, y ∗ i,m ( x i )) for notational brevity.
Algorithm Design.
Towards this end, it is straightforward to construct a variant of the classical augmented Lagrangian function for the UL optimization problem as
[Equation]
where λ denotes the dual variable (Lagrangian multiplier) for the consensus constraint, ρ > 0 , and γ is a scaling factor (which will be determined later).
Motivated by the primal-dual optimization framework [44], one step of gradient descent based on the linearized objective function with a following gradient ascent step is sufficient for the minimization of the general nonconvex loss function under the linear constraints, which means that there is no need to solve an inner optimization problem before updating the Lagrangian multiplier as is done in the classical augmented Lagrangian method.
When both the UL and LL objective functions are differentiable and the inverse of the Hessian matrix at the LL problem exists, i.e., ∇ 2 y k y k g k ( x , y ∗ k ( x )) is invertible, then there exists a closed form for ∇ f i ( x i , y ∗ i,k ( x i )) .
Following the existing works on bilevel algorithm designs, replacing y ∗ i,k ( x i ) by y i,k in the gradient of f i ( x i , y ∗ i,k
( x i )) w.r.t.
x i can serve as an efficient surrogate for the stochastic gradient estimate.
However, in the decentralized setting, only individual loss functions are observable at each agent, therefore, the local UL implicit gradient is computed through replacing g k ( x , y k ) by g i,k
( x i , y i,k ) , denoted as ∇ f i ( x i , y i,k ) .
Let h rg,k
and h rf respectively denote the distributed stochastic gradient estimate of the LL and UL objective functions at points ( x r , y rk ) and ( x r , y r +1 k ) , ∀ k , w.r.t.
y k and x , where r represents the index of iterations.
Thus, our proposed
SLAM can be expressed as
[Equation]
where ω k is the dual variable for ensuring the LL consensus process for each learning task, α and β are the parameters of the quadratic penalization terms, and ρ/γ here is the step-size for the updates of the dual variables.
Implementation of SLAM.
Noting that the objective functions in each subproblem, i.e., (4a) and (4c), are quadratic, we can easily have the updates of both UL and LL optimization variables as
[Equation]
where 1 /α and 1 /β serve as the step-sizes of updating both UL and LL learning models.
Subtracting the equality with the same one from the previous iteration for both (5a) and (5b) ends up with efficient model updates of both the UL and LL learning problems as follows:
[Equation]
where the mixing matrices, with τ g = β/γ and τ f = α/γ , are defined as
[Equation]
According to (6a) and (6b), it can be readily observed that SLAM is amenable to a fully decentralized implementation.
The detailed algorithm description is provided in Algorithm 1 from a local view of the model update, where [ W ] ij denotes the ij th entry of matrix W , [ h rg ] i,k is the gradient estimate of ∇ g i,k ( x ri , y ri,k ) (i.e., h rg,k
=
[[ h rg ] 1 ,k , . .
, [ h rg ] n,k ] T ), and similarly [ h rf ] i is the local gradient estimate of ∇ f i ( x ri , y r
+1 i,k ) (i.e., h rf =
[[ h rf ] 1 , . . .
, [ h rf ] n ] T ).
Algorithm
Initialization: α, β, γ , x 1 i , y
i,k , ∀ i, k , and set λ
1 = ω 1 k = 0 , ∀ k ; 1: for r = 1 , 2 , · · · , T do 2: for i = 1 , 2 , · · · , n in parallel over the network do 3:
Estimate gradient ∇ g
i,k ( x ri , y ri,k ) for each task and ∇ f i ( x ri , y r
+
i,k ) locally 4: y r
+
i,k = 2 (cid:80) j ∈N i
[ W g ] ij y rj,k − [ W ′ g ] ij y r − 1 j,k − β − 1 (cid:0) [ h rg ] i,k − [ h r − 1 g ] i,k (cid:1)
▷ LL models 5: x r
+1 i = 2 (cid:80) j ∈N
i
[ W f ] ij x rj − [ W ′ f ] ij x r − 1 j − α − 1 (cid:16) [ h rf ] i − [ h r − 1 f ] i (cid:17)
▷ UL model 6: end for 7: end for
Besides, if there is a consensus requirement at only one level of the optimization problem, then the problem at the other level becomes one with multiple objective functions.
Our proposed SLAM can also be applied for solving any of these problems by a minor revision of the generic SLAM formulation.
To be more specific, we provide the following discussion.
A Special Case of DBO (1) (with only consensus in the LL problems).
If there is only a need for consensus of LL model parameters, then problem (2) reduces to the following DBO problem.
For example, in solving multi-agent actor-critic RL problems, the UL optimization problem consists of improving the policy for each agent while the LL problem requires all the agents to jointly evaluate the value function over the whole network.
The DBO problem is then expressed as
[Equation]
The major difference between problem (2) and (8) is that the UL optimization problem includes multiple objectives over the model parameters x i , ∀ i ∈ [ n ] .
In this case, the updating rule of variable x in (6b) reduces to x r +
= x
r − h rf /α by forgoing the dual update w.r.t.
λ .
The detailed implementation is summarized in Algorithm 2, where we name this special case of SLAM by
SLAM-L as the LL consensus process is the main feature in this setting.
Initialization: α, β, γ , x 1 i , y
i,k , ∀ i, k , and set ω 1 k = 0 , ∀ k ; 1: for r = 1 , 2 , · · · , T do 2: for i = 1 , 2 , · · · , n in parallel over the network do 3:
Estimate gradient ∇ g
i,k ( x ri , y ri,k ) for each task and ∇ f i ( x ri , y r
+
i,k ) locally 4: y r
+
i,k = 2 (cid:80) j ∈N i
[ W g ] ij y rj,k − [ W ′ g ] ij y r − 1 j,k − β − 1 (cid:0) [ h rg ] i,k − [ h r − 1 g ] i,k (cid:1)
▷ LL models 5: x r
+1 i
= x
[Equation]
−
[Equation] [Equation]
rf
[Equation]
, ∀ i ▷ UL models 6: end for 7: end for
A Special Case of DBO (1) (with only consensus in the UL problem).
The other special is analogous to the first one with the difference being the absence of the LL consensus process in comparison to (2), which is written as follows:
[Equation]
where there are multiple objectives in the LL optimization problems.
Problem (9) also covers a wide range of applications in machine learning, e.g., multi-task and/or personalized learning, and so on.
In this case, the update of variable y k shown in (5a) is changed to y r +
k
=
y rk − h
rg /β as there is no consensus constraint involved.
Analogous to the previous case, the implementation of this algorithm is presented in Algorithm 3 and termed as SLAM-U. y
[Equation]
▷ LL models 5: x
[Equation]
UL model 6: for
Theoretical Convergence Results
Before showing the theoretical results about the convergence guarantees of SLAM, we first need five main classes of assumptions used in showing the descent of some quantifiable function so that SLAM can reach the ϵ -KKT points of the DBO problems.
More detailed definitions and properties regarding these assumptions are deferred to the supplement.
3.1 Assumptions
Our theoretical results are based on the following assumptions on the properties of the loss functions in both the UL and LL optimization problems, which are mainly related to the continuity of the objective function and stochasticity of the gradient estimates.
A1.
(Lipschitz continuity of both UL and LL objective functions) Assume that functions
f i ( · ) , ∇ f i ( · , ) , ∇ g i,k ( · ) , ∇ 2 g i,k ( · ) , ∀ i , are (block-wise)
Lipschitz continuous with constants
L f, 0 , L f, 1 , L g, 1 , L g, 2 for both x and y k , ∀ k , and ∇ 2 x i
y i,k g
i,k ( · ) ,
A3.
(Quality of the stochastic gradient estimate)
The stochastic estimates of ∇ f i ( x i , y i,k ) , ∇ y
i
g
i,k ( x i , y i,k ) , ∀ i, k , are unbiased and their variances are bounded by σ 2 f , σ 2 g . A4.
Assume that the UL objective functions
f i ( x i , y ∗ i,k
( x i )) , ∀ i, k are lower bounded.
A5.
(Strong convexity of g i,k ( · ) w.r.t.
y i,k ) Function g i,k ( · ) is µ g -strongly convex w.r.t.
y i,k , ∀ i, k .
Note that these assumptions are commonly used in the convergence analysis for bilevel and decentralized optimization algorithms.
Given these assumptions, we are now in a position to provide the following theoretical convergence guarantees.
3.2 Convergence Rates of SLAM
Theorem 1.
(Convergence rate of SLAM to ϵ -KKT points) Suppose that A1-A5 hold and assume ∥∇ 2 y i
y
i g i,k
[Equation]
, y
[Equation]
) − n
[Equation]
ni =1 ∇ 2 y
i
y
i g i,k ( · , y ′ i )
∥
≤ L g ∥
y i − y ′ i ∥ , ∀ i, k if ∇ 2 g i,k ( · ) , ∀ i, k are required in computing the UL implicit gradient.
When step-sizes are chosen as 1 /α ∼ 1 /β ∼
O ( (cid:112) n/T ) , τ f , τ g ≥ O ( ρσ max ( A T A )) , the mini-batch size of h rf is O (log( nT )) , then the iterates { x r , λ r , y rk , ω rk , ∀ k, r } generated by SLAM satisfy
[Equation]
where x = n − 1 1 T x , and T denotes the total number of iterations.
Remark 1 .
It is noted in Theorem 1 that the convergence rate achieved by SLAM to find the ϵapproximate KKT points of (1) (including both the first-order stationarity of the solutions and the violation of constraints) is on the order of 1 / ( nϵ 2 ) .
Therefore, it follows that a linear speedup w.r.t.
the number of learners can be achieved by SLAM for DBO, matching the classical results of distributed
SGD for only single-level general nonconvex problems.
Remark 2 .
In comparison with existing bilevel algorithms, SLAM is a single timescale algorithm since the learning rates can be chosen as 1 /α ∼ 1 /β , which is consistent with ALSET [43].
Remark 3 .
The major novelty of obtaining these theoretical results relies on the developed variant of the augmented Lagrangian function and subsequently derived recursion of the successive dual variables, which quantify well the consensus errors resulting from both UL and LL optimization processes in terms of primal variables.
Note that this is distinct from the existing theoretical analysis of stochastic algorithms, such as distributed SGD [38, 39], stochastic gradient tracking [40, 41], stochastic primal-dual algorithms [42, 45], etc.
Corollary 1.
(Convergence rate of SLAM-L to ϵ -KKT points) Suppose that A1-A5 hold and assume ∥∇ 2 y i
y
i g i,k ( x i , y i ) − ∇
yy g k ( x , y ′ i )
∥
≤ L g ∥
y i − y ′ i ∥ , ∀ i, k if ∇ 2 g i,k ( · ) , ∀ i, k are required in computing the UL implicit gradient.
When step-sizes are chosen as 1 /α ∼ O (1 / √ T ) , 1 /β ∼ O ( [Equation] n/T ) , τ f , τ g ≥ O ( ρσ max ( A T A )) , the mini-batch size of h rf is O (log( nT )) , the iterates { x r , y rk , ω rk , ∀ k, r } generated by SLAM-L satisfy
UL:
[Equation]
and
LL: (
[Equation] [Equation]
Remark 4 .
Different from Theorem 1, the stationarity of the UL model parameters requires the shrinkage of the gradient size over each individual UL problem as shown in Corollary 1, so there is no speedup on the convergence rate guarantee at UL.
Corollary 2.
(Convergence rate of SLAM-U to ϵ -KKT points) Suppose that A1-A5 hold.
Given the conditions on 1 /α, 1 /β, τ f , τ g and the mini-batch size of h r f shown in Theorem 1, the iterates { x r , λ r , y rk , ∀ k, r } generated by SLAM-U satisfy
UL: (10a) and
LL
[Equation]
[Figure] [Figure]
Numerical Results
In this section
[Figure]
we evaluate our proposed algorithm using two MARL environments: 1) the cooperative navigation task [16], which is built on the OpenAI Gym platform [46]; and 2) the pursuit-evasion game [47], which is built on the PettingZoo platform [48].
Detailed experimental settings and additional numerical results are provided in the supplement
[Figure]
Cooperative Navigation Task.
In this game, we consider that the n agents are aiming to jointly reach n different landmarks as soon as possible, where the Erdos Renyi Graph is used.
We assume that each agent can observe the global state and has 5 possible actions: stay, left, right, up, and down.
This task consists of a shared common goal of avoiding collision among the agents while they navigate to the targeting landmarks.
In the simulations, each agent locally maintains two fully connected neural networks as the actor network (at UL w.r.t.
x i ) and the critic network (at LL w.r.t.
y i ), respectively.
Moreover, each agent shares its critic network with its neighbors to cooperatively estimate the global value function and independently train its actor network to complete its local task.
We compare the performance of our proposed SLAM with application to the DAC setting, named
SLAM-AC, with two benchmark algorithms: DAC [15] and mini-batch DAC (MDAC) [22].
Theoretically, MDAC needs an O ( ϵ − 1 ln ϵ − 1 ) batch size in its inner loop to update critic parameters before each update in policy parameters, which is not practical.
Here, we set a small batch B = 10 in the inner loop for MDAC to achieve fast convergence.
The simulation results on this coordination game are presented in Figure1, where the performance is averaged over 5 independent Monte Carlo (MC) trials for each algorithm.
Pursuit-Evasion Game.
In the pursuit-evasion game, there are two groups of nodes: pursuers (agents) and evaders.
The agents are connected through a ring graph.
Pursuers could observe the global state of the video game.
An evader is considered caught if two pursuers simultaneously arrive at the evader’s location.
As each pursuer should learn to cooperate with other pursuers to catch the evaders, the pursuers share certain similarities with each other since they need to follow similar strategies to achieve their local tasks: simultaneously catching an evader with other pursuers.
We follow the experimental set up in [23], where all agents partially share their actor networks with neighbors for collaborations in their policy spaces and fully share their critic network to cooperatively learn the global value function.
In Figure 2, we compare SLAM-AC with two benchmarks, CAC [23] and MDAC [22], with 5 MC trials again.
To ensure a fair comparison, all algorithms use the same parameter sharing scheme mentioned above.
Note that CAC [23] is a variant of DAC [15] and the only difference is that CAC can partially share its policy parameters while the policy parameters are not shared in DAC.
In the experiment, each agent maintains two convolutional neural networks, one for the actor and one for the critic (Please refer to the supplement for detailed structures).
Concluding Remark
In this paper, we studied a generic form of the DBO problem, which is shown to have three major variants that formulate multiple hierarchical machine learning problems.
Targeting these DBO problems, we proposed SLAM – a simple and elegant algorithm to solve DBO in a fully decentralized way.
Under mild conditions, we establish theoretical results showing that our proposed SLAM is able to find the ϵ -KKT points with a convergence rate of O (1 / ( nϵ 2 )) , which matches the standard convergence rate achieved by the classical distributed SGD algorithms for solving only single-level general nonconvex optimization problems.
We tested the performance of SLAM numerically on a
MARL scenario and found that SLAM outperformed the traditional AC algorithms w.r.t.
convergence speed and (in most cases) achievable rewards.
Societal impact.
To the best of our knowledge, we do not see any ethical or negative immediate societal consequence of this work.
Acknowledgments
M. Hong and S. Zeng are partially supported by NSF grants CIF-1910385 and CMMI-1727757.
Checklist
The checklist follows the references.
Please read the checklist guidelines carefully for information on how to answer these questions.
For each question, change the default [TODO] to [Yes] , [No] , or [N/A] .
You are strongly encouraged to include a justification to your answer , either by referencing the appropriate section of your paper or providing a brief inline description.
For example:
▶ Did you include the license to the code and datasets?
[Yes] See ▶
Did you include the license to the code and datasets?
[No] The code and the data are proprietary
▶ Did you
include the
license
to the
code
and datasets?
[N/A]
Please do not modify the questions and only use the provided macros for your answers.
Note that the Checklist section does not count towards the page limit.
In your paper, please delete this instructions block and only keep the Checklist section heading above along with the questions/answers below.
For all authors... (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contribu- tions and scope?
[Yes] (b) Did you describe the limitations of your work?
[Yes] (c) Did you discuss any potential negative societal impacts of your work?
[Yes] (d) Have you read the ethics review guidelines and ensured that your paper conforms to them?
[Yes] 2.
If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?
[Yes] (b) Did you include complete proofs of all theoretical results?
[Yes] 3.
If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
[No] (b) Did you specify all the training details (e
g., data splits, hyperparameters, how they were chosen)?
[Yes] (c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[Yes] (d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?
[Yes] 4.
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?
[Yes] (b) Did you mention the license of the assets?
[Yes] (c) Did you include any new assets either in the supplemental material or as a URL?
[No] (d) Did you discuss whether
and how consent was obtained from people whose data you’re us- ing/curating?
[N/A] (e) Did you discuss whether the data you are using/curating contains personally identifiable informa- tion or offensive content?
[N/A] 5.
If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if applicable?
[N/A] (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?
[N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
[N/A]