Fast Training Methods for Stochastic Compositional Optimization Problems
Year: 2022

Abstract:
The stochastic compositional optimization problem covers a wide range of machine learning models, such as sparse additive models and model-agnostic meta-learning. Thus, it is necessary to develop efficient methods for its optimization. Existing methods for the stochastic compositional optimization problem only focus on the single machine scenario, which is far from satisfactory when data are distributed on different devices. To address this problem, we propose novel decentralized stochastic compositional gradient descent methods to efficiently train the largescale stochastic compositional optimization problem. To the best of our knowledge, our work is the first one facilitating decentralized training for this kind of problem. Furthermore, we provide the convergence analysis for our methods, which shows that the convergence rate of our methods can achieve linear speedup with respect to the number of devices. At last, we apply our decentralized training methods to the model-agnostic meta-learning problem, and the experimental results confirm the superior performance of our methods.

heading: 1 Introduction
text: The stochastic compositional optimization problem [20] plays an important role in the machine learning field, since it covers numerous applications, such as policy evaluation [19], sparse additive models [20], and model-agnostic meta-learning [6, 1]. Specifically, the stochastic compositional optimization problem is defined as follows:
min x∈Rd
Eζ [ f ( Eξ[g(x; ξ)]; ζ )] , (1)
where the outer-level function f(y) = Eζ [f(y; ζ)] : Rd ′ → R is a smooth and nonconvex function, the inner-level function g(x) = Eξ[g(x; ξ)] : Rd → Rd ′
is a smooth function. It is obvious that the loss function consists of two stochastic functions f and g, making it different from the noncompositional problem.
The compositional structure in Eq. (1) leads to more challenges in optimization compared with the non-compositional problem, since the stochastic gradient of the loss function in Eq. (1) is not an unbiased estimation of the full gradient. Thus, efficient training of the stochastic compositional optimization problem has attracted increasing attention in recent years. For example, [20] proposed the stochastic compositional gradient descent (SCGD) method to deal with the two-level stochasticity in Eq. (1). But it has a worse convergence rate than the standard stochastic gradient descent method. To improve it, a series of variance-reduced methods have been proposed. For example, [25] developed the composite randomized incremental gradient method based on the SAGA [4] technique. [27, 24] further improved the convergence rate based on the SPIDER [5] method.
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
Although aforementioned methods have achieved much progress in optimizing Eq. (1), all of them only focus on the single-machine scenario. In fact, data are usually distributed on different devices in many real-world applications. It is necessary to study the distributed training method for the stochastic compositional optimization problem. In the past few years, numerous distributed training methods have been proposed, such as asynchronous stochastic gradient descent (SGD) [16], decentralized SGD [13], local SGD [14]. But they are not applicable to the stochastic compositional problem. Especially, their convergence results do not hold for Eq. (1) because those distributed training methods require the stochastic gradient is an unbiased estimation of the full gradient. As far as we know, there are no existing works studying distributed training methods for Eq. (1) with theoretical guarantees.
To fill the aforementioned gap, we developed two novel decentralized training methods for the stochastic compositional optimization problem. In particular, we proposed a gossip-based decentralized stochastic compositional gradient descent (GP-SCGD) method and a gradient-tracking-based decentralized stochastic compositional gradient descent (GT-SCGD) method. With our methods, different devices can collaboratively optimize Eq. (1). However, establishing the convergence rate of our proposed methods is challenging. Existing convergence analysis techniques for decentralized SGD are not applicable to our methods because the stochastic compositional gradient is different from the stochastic gradient. To address this challenge, we proposed new convergence analysis techniques for the decentralized SCGD method, where we show how to bound the gradient variance and the consensus error under the decentralized setting. In particular, both GP-SCGD and GT-SCGD can achieve the convergence rate O( 1Kϵ2 ) to achieve ϵ-accuracy solution, where K is the number of devices. It indicates that our methods can achieve the linear speedup with respect to the number of devices, which is consistent with decentralized SGD. To the best of our knowledge, our work is the first one to establish the convergence rate of decentralized SCGD and achieving the linear speedup. Moreover, the sample complexity of our methods is O( 1Kϵ3 ). It indicates that our methods can achieve a better sample complexity O( 1ϵ3 ) than O( 1 ϵ4 ) of traditional SCGD [20] when K = 1, which further confirms the superiority of our methods. Finally, we applied our proposed methods to optimize the model-agnostic meta-learning problem. The empirical results confirm the effectiveness of our methods. In the following, we summarize the contributions of our work.
• We proposed two novel decentralized stochastic compositional gradient descent methods: GP-SCGD and GT-SCGD. This is the first work studying the decentralized training method for stochastic compositional optimization problems.
• We established the convergence rate of our decentralized stochastic compositional gradient descent methods with novel theoretical analysis. This is the first work establishing the convergence rate for decentralized SCGD and showing the linear speedup with respect to the number of devices. In addition, our methods can achieve better sample complexities than traditional SCGD.
• The extensive empirical evaluation on the model-agnostic meta-learning task confirms the effectiveness of our proposed methods.

heading: 2 Related Works

heading: 2.1 Stochastic Compositional Optimization Problem
text: The stochastic compositional optimization problem is common in the machine learning field. However, it is more difficult to optimize than the standard non-compositional problem. The reason is that its stochastic gradient is not an unbiased estimation of the full gradient, which is shown as follows:
Eξ,ζ [∇g(x; ξ)T∇gf(g(x; ξ); ζ)] ̸= ∇g(x)T∇gf(g(x)) . (2)
To train the stochastic compositional problem efficiently, numerous methods have been proposed in the past few years. Specifically, [20] developed the stochastic compositional gradient descent (SCGD) method, which is defined as follows:
ut = (1− βt−1)ut−1 + βt−1g(xt;Bξ,t) , xt+1 = xt − ηt∇g(xt;Bξ,t)T∇gf(ut;Aζ,t) ,
(3)
where Bξ,t and Aζ,t denote the random samples, 0 < βt < 1 is a hyperparameter, and ηt is the learning rate. Here, ut serves as the estimator of the inner function g(xt), and then the gradient
∇gf(g(xt)) is estimated by ∇gf(ut;Aζ,t). As a result, this stochastic compositional gradient ∇g(xt;Bξ,t)T∇gf(ut;Aζ,t) has a smaller variance to improve the accelerate the convergence speed. However, it still has a worse theoretical convergence rate than SGD. Furthermore, to improve the convergence speed of SCGD for nonconvex problems, advanced variance reduction techniques have been incorporated into SCGD. For example, [25] incorporated the variance reduction technique SAGA [15] to SCGD and obtained a better convergence rate. Later, [27] employed another variance reduction technique SPIDER [5] for SCGD, while [23] accelerated SCGD with the STORM [3] variance reduction technique. More recently, some works [2, 10] disclosed the connection between stochastic bilevel optimization problems and stochastic compositional problems and developed the alternating SGD method to solve Eq. (1). However, all of these methods only studied the convergence rate for the single machine case, ignoring the distributed setting.

heading: 2.2 Decentralized Training Methods
text: Recently, with the emergence of large-scale data, numerous distributed training methods [16, 13, 14, 18, 11, 7, 9, 8] have been proposed. Among them, the decentralized training method, where different devices conduct the peer-to-peer communication, has attracted a lot of attention. In particular, there are two kinds of communication strategies in a decentralized training system. They are the gossip strategy and the gradient tracking strategy. Based on the gossip strategy, [13] studied the convergence rate of decentralized SGD and disclosed how the topology of the communication graph affects the convergence rate. Additionally, numerous works aim to reduce the communication cost of the gossip-based decentralized SGD by compressing gradients or skipping the communication step. Meanwhile, based on the gradient-tracking strategy, [17] developed the distributed stochastic gradient tracking method and established its convergence rate. Recently, some works [18, 22] combined the variance reduction strategy and the gradient-tracking strategy to further accelerate decentralized SGD. However, all of these methods are not applicable to SCGD. Thus, in this paper, we will develop efficient decentralized SCGD methods for optimizing the stochastic compositional optimization problem.

heading: 3 Decentralized Stochastic Compositional Optimization Methods

heading: 3.1 Problem Definition
text: Formally, we consider the distributed stochastic compositional optimization problem, which is defined as follows:
min x∈Rd
1
K K∑ k=1 Eζ [ f (k) ( Eξ[g(k)(x; ξ)]; ζ )] , (4)
where x ∈ Rd denotes the model parameter, K is the total number of devices, f (k) = Eζ [f (k)(y; ζ)] : Rd′ → R and g(k) = Eξ[g(k)(x; ξ)] : Rd → Rd ′ are smooth functions on the k-th device, and
accordingly Eζ [ f (k) ( Eξ[g(k)(x; ξ)]; ζ )] is the loss function on the k-th device. In other words, each
device has its own data and all devices collaboratively learn the model parameter x by utilizing their local data.
To efficiently train this distributed stochastic compositional optimization problem, we consider the decentralized training method. In particular, each device connects with its neighbors, composing a communication network. All devices conduct the peer-to-peer communication based on this communication network. Formally, we represent the communication network with G = (V,W ) where V = {v1, v2, · · · , vK} denotes all devices and W = [wij ] ∈ RK×K is the adjacency matrix which indicates whether two devices are connected or not. Following [11, 13], the adjacency matrix W is assumed to satisfy the following assumption.
Assumption 1. The adjacency matrix W has the following properties:
• Nonnegative: wij ≥ 0, ∀i, j.
• Symmetric: wij = wji, ∀i, j.
• Doubly Stochastic: ∑K i=1 wij = 1, ∑K j=1 wij = 1.
• The eigenvalues of W can be sorted as |λn| ≤ · · · ≤ |λ2| < |λ1| = 1.
In our convergence analysis, we denote λ = |λ2| for simplicity. In terms of the aforementioned decentralized training setting, we proposed two descentralized stochastic compositional gradient descent methods. The details are described in the following subsection.

heading: 3.2 Gossip-based Decentralized Stochastic Compositional Gradient Descent Method
text: In Algorithm 1, we developed the gossip-based decentralized stochastic compositional gradient descent (GP-DSCGD) method. Following the standard SCGD method [20], each device estimates the inner-level function g(k)(x(k)t ) with u (k) t as follows:
u (k) t = (1− γβt−1)u (k) t−1 + γβt−1g (k)(x (k) t ;B (k) t ) , (5)
where γ > 0 and βt > 0 are two hyperparameters, γβt < 1, k is the index of devices and t is the index of iterations. Different from the standard SCGD method in Eq. (3), there is an additional hyperparameter γ when computing u(k)t , which can help controlling the estimation variance. We will show it in our theoretical analysis in Appendix. Based on this estimation, each device computes the stochastic compositional gradient z(k)t as shown in Line 9. Then, each worker uses the gossip communication strategy to update the local model parameter as follows:
x̃ (k) t+1 = ∑ j∈Nvk wkjx (j) t − ηz (k) t , (6)
where η > 0 denotes the learning rate, Nvk = {j|wkj > 0} is the neighboring devices of the k-th device. Here, the first term on the right-hand side incurs the communication to average the model parameter of neighboring devices, and the second term indicates to update the model parameter with the local gradient z(k)t . Then, instead of using x̃ (k) t+1 as the new model parameter, we compute the model parameter of the (t+ 1)-th iteration as follows:
x (k) t+1 = x (k) t + βt(x̃ (k) t+1 − x (k) t ). (7)
In fact, this strategy couples the update of u(k)t and x (k) t by the sharing hyperparameter βt, which can benefit controlling the estimation variance of u(k)t . All devices repeat the aforementioned steps until the iterates converge.
Algorithm 1 Gossip-based Decentralized Stochastic Compositional Gradient Descent (GP-DSCGD)
Input: x(k)0 = x0, βt > 0, γ > 0, η > 0. 1: for t = 0, · · · , T − 1, each device k do 2: Sample a subset of samples B(k)t to compute: 3: if t = 0 then 4: u(k)t = g (k)(x (k) t ;B (k) t )
5: else 6: u(k)t = (1− γβt−1)u (k) t−1 + γβt−1g (k)(x (k) t ;B (k) t ) 7: end if 8: v(k)t = ∇g(k)(x (k) t ;B (k) t ) 9: Sample a subset of samples A(k)t to compute: z (k) t = (v (k) t ) T∇f (k)(u(k)t ;A (k) t )
10: x̃(k)t+1 = ∑
j∈Nvk wkjx
(j) t − ηz (k) t
11: x(k)t+1 = x (k) t + βt(x̃ (k) t+1 − x (k) t ) 12: end for

heading: 3.3 Gradient-Tracking-based Decentralized Stochastic Compositional Gradient Descent Method
text: Besides the gossip-based communication strategy, the gradient tracking strategy is also widely used in the decentralized training method. Therefore, for completeness, we further proposed the gradienttracking-based descentralized stochastic compositional gradient descent (GT-DSCGD) method in
Algorithm 2. Same with Algorithm 1, we compute u(k)t and z (k) t as shown in Lines 2-9. Different from Algorithm 1, GT-DSCGD introduces an auxiliary variable s(k)t for each device, which is used to track the global stochastic compositional gradient 1K ∑K k=1 z (k) t [18, 17]. Specifically, s (k) t is updated as follows: s (k) t =
∑ j∈Nvk wkjs (j) t−1 + z (k) t − z (k) t−1 . (8)
It is easy to verify that 1K ∑K k=1 s (k) t = 1 K ∑K k=1 z (k) t . Then, each device uses s (k) t instead of z (k) t to update the model parameter as follows:
x̃ (k) t+1 = ∑ j∈Nvk wkjx (j) t − ηs (k) t . (9)
To sum up, the difference between our two methods is that GT-DSCGD has an additional step to track the gradient with s(k)t and accordingly updates the model parameter with s (k) t .
Algorithm 2 Gradient-Tracking-based Decentralized Stochastic Compositional Gradient Descent (GT-DSCGD)
Input: x(k)0 = x0, z (k) −1 = 0, s (k) −1 = 0, βt > 0, γ > 0, η > 0.
1: for t = 0, · · · , T − 1 , each device k do 2: Sample a subset of samples B(k)t to compute: 3: if t = 0 then 4: u(k)t = g (k)(x (k) t ;B (k) t ) 5: else 6: u(k)t = (1− γβt−1)u (k) t−1 + γβt−1g (k)(x (k) t ;B (k) t ) 7: end if 8: v(k)t = ∇g(k)(x (k) t ;B (k) t ) 9: Sample a subset of samples A(k)t to compute: z (k) t = (v (k) t ) T∇f (k)(u(k)t ;A (k) t )
10: s(k)t = ∑
j∈Nvk wkjs
(j) t−1 + z (k) t − z (k) t−1 // gradient tracking 11: x̃(k)t+1 = ∑ j∈Nk wkjx (j) t − ηs (k) t 12: x(k)t+1 = x (k) t + βt(x̃ (k) t+1 − x (k) t ) 13: end for

heading: 4 Convergence Analysis

heading: 4.1 Assumptions
text: To establish the convergence rate of our proposed two methods, we first introduce some standard assumptions, which are also used in existing stochastic compositional gradient descent methods [20, 21, 26, 24]. Assumption 2. (Smoothness) For ∀y1, y2 ∈ Rd ′
and ∀x1, x2 ∈ Rd, there exist two constant values

heading: Lf > 0 and Lg > 0 such that
text: ∥∇f (k)(y1)−∇f (k)(y2)∥ ≤ Lf∥y1 − y2∥, ∥∇g(k)(x1)−∇g(k)(x2)∥ ≤ Lg∥x1 − x2∥ . (10)
Assumption 3. (Bounded gradient) For ∀x ∈ Rd and ∀y ∈ Rd′ , there exist two constant values Cg > 0 and Cf > 0 such that
E[∥∇g(k)(x; ξ)∥2] ≤ C2g , E[∥∇f (k)(y; ζ)∥2] ≤ C2f . (11)
Assumption 4. (Bounded variance) For x ∈ Rd and y ∈ Rd′ , there exist three constant values σf > 0, σg > 0, σg′ > 0, such that
E[∥∇f (k)(y; ζ)−∇f (k)(y)∥2] ≤ σ2f , E[∥∇g(k)(x; ξ)−∇g(k)(x)∥2] ≤ σ2g′ ,
E[∥g(k)(x; ξ)− g(k)(x)∥2] ≤ σ2g . (12)
To study the convergence rate of our two methods, we denote F (x) = 1K ∑K k=1 F (k)(x) =
1 K ∑K k=1 f (k)(g(k)(x)). Then, F (x) is LF -smooth with LF = C2gLf + CfLg. In addition, we
denote x̄t = 1K ∑K k=1 x (k) t and x∗ as the optimal solution.

heading: 4.2 Convergence Rate
text: In terms of aforementioned definitions and assumptions, we establish convergence rates of our proposed Algorithm 1 and Algorithm 2, respectively.
Theorem 1. Assume Assumptions 1-4 hold, for Algorithm 1, by setting |Akt | = |Bkt | = B, γ > 0, βt = β ≤ min{ 18γ , 1 2ηLF , 1}, η ≤ min{η1, η2, η3} where
η1 = γ
3 √ 6C2gLf
, η2 = 1− λ2
2 /(27C4gL2f + γ2L2F γ2 + γ(24C2fL 2 g + 99C 4 gL 2 f )
2C2gL 2 f
) ,
η3 =
√ b2 + 4ac− b
2a , a =
(24C2fL 2 g + 3C 4 gL 2 f )
8(1− λ2)C2gL2f +
12C2g 1− λ2 + 27C4gL 2 f 2γ2 , b =
2
1− λ2 , c =
γ
4C2gL 2 f
,
(13) we can get
1
T T−1∑ t=0 E[∥∇F (x̄t)∥2] ≤ 2E[F (x0)− F (x∗)] ηβT + 2(3C2gσ 2 f + 3C 2 fσ 2 g′ + C 2 gL 2 fσ 2 g + 24γσ 2 g) B
+ 48C2fσ 2 g′ + 16C 2 gσ 2 f + 99C 2 gL 2 fσ 2 g
16βC2gL 2 fB
+ 6C2gL 2 fσ 2 g
βγB +
48σ2
βB .
(14)
Corollary 1. Assume Assumptions 1-4 hold, for Algorithm 1, by setting η = O( √ K√ T ), B = O( √ KT ), β and γ to be two positive constant values, we can get
1
T T−1∑ t=0 E[∥∇F (x̄t)∥2] ≤ 2E[F (x0)− F (x∗)] β √ KT + 2(3C2gσ 2 f + 3C 2 fσ 2 g′ + C 2 gL 2 fσ 2 g + 24γσ 2 g)√ KT
+ 48C2fσ 2 g′ + 16C 2 gσ 2 f + 99C 2 gL 2 fσ 2 g
16βC2gL 2 f
√ KT
+ 6C2gL 2 fσ 2 g
βγ √ KT
+ 48σ2
β √ KT .
(15)
Remark 1. For sufficiently large T , Corollary 1 indicates that GP-DSCGD can achieve linear speedup with respect to the number of devices, which is consistent with decentralized SGD.
Remark 2. To achieve the ϵ-accuracy solution such that 1T ∑T−1
t=0 E[∥∇F (x̄t)∥2] ≤ ϵ, the communication complexity is O ( 1
Kϵ2
) and the sample complexity is T ×B = O ( 1
Kϵ3
) . When K = 1, the
sample complexity of our method is better than O( 1ϵ3 ) of traditional SCGD [20].
Note that when T is sufficiently large, 12ηLF could be greater than 1 so that β and γ do not affect the order of the convergence rate in Corollary 1. Additionally, our method does not use the acceleration or variance reduction techniques as [21, 25, 24] so that the sample complexity of our method is inferior to those methods.
Theorem 2. Assume Assumptions 1-4 hold, for Algorithm 2, by setting |Akt | = |Bkt | = B, γ > 0, βt = β ≤ min{ 18γ , 1 2ηLF , 1}, and η ≤ min{η1, η2, η3}, where
η1 = 4γ(1− λ2)
8γ2L2F + 289C 4 gL 2 f + 16C 2 fL 2 g
, η2 = γ 2 √
19C4gL 2 f + C 2 fL 2 g
,
η3 =
√ b2 + 4ac− b
2a , a =
27C4gL 2 f
4γ2 +
3C2fL 2 g
8γ2 +
3C4gL 2 f
128γ2 , b =
6
1− λ2 , c = (1− λ)2 ,
(16)
we can get
1
T T−1∑ t=0 E[∥∇F (x̄t)∥2] ≤ 2E[F (x0)− F (x∗)] ηβT + 6C2gσ 2 f + 6C 2 fσ 2 g′ + 3C 2 gL 2 fσ 2 g B
+ 24η2(C2fσ 2 g′ + C 2 gσ 2 f )(19C 4 gL 2 f + C 2 fL 2 g)
γ2B +
48η2(C2fσ 2 g′ + C 2 gσ 2 f )(19C 4 gL 2 f + C 2 fL 2 g)
γ2(1− λ2)B
+ 48η(C2fσ 2 g′ + C 2 gσ 2 f )
γ(1− λ2)B +
96η(C2fσ 2 g′ + C 2 gσ 2 f )
γ(1− λ2)2B +
8C2gL 2 fσ 2 g + 96(1− λ)C2gL2fσ2
γβB .
(17)
Corollary 2. Assume Assumptions 1-4 hold, for Algorithm 2, by setting η = O( √ K√ T ), B = O( √ KT ), β and γ to be two positive constant values, we can get
1
T T−1∑ t=0 E[∥∇F (x̄t)∥2] ≤ 2E[F (x0)− F (x∗)] β √ KT + 48(C2fσ 2 g′ + C 2 gσ 2 f ) γ(1− λ2)T + 96(C2fσ 2 g′ + C 2 gσ 2 f ) γ(1− λ2)2T
+ 24(C2fσ 2 g′ + C 2 gσ 2 f )(19C 4 gL 2 f + C 2 fL 2 g)
γ2K1/2T 3/2 +
48(C2fσ 2 g′ + C 2 gσ 2 f )(19C 4 gL 2 f + C 2 fL 2 g)
γ2(1− λ2)K1/2T 3/2
+ 6C2gσ 2 f + 6C 2 fσ 2 g′ + 3C 2 gL 2 fσ 2 g√
KT +
8C2gL 2 fσ 2 g + 96(1− λ)C2gL2fσ2
γβ √ KT
.
(18) Remark 3. For sufficiently large T , Corollary 2 also indicates that GT-DSCGD can achieve linear speedup with respect to the number of devices, which is consistent with the decentralized SGD method. Similar to GP-DSCGD, to achieve the ϵ-accuracy solution such that 1T ∑T−1 t=0 E[∥∇F (x̄t)∥2] ≤ ϵ,
the communication complexity of GT-DSCGD is O (
1 Kϵ2
) and the sample complexity is O ( 1
Kϵ3
) .

heading: 4.3 Proof Sketch
text: In this subsection, we present the main idea of establishing the convergence rate of Algorithm 1 and Algorithm 2. More details are deferred to Appendix.
Proof sketch of Theorem 1 To establish the convergence rate of Algorithm 1, we proposed a novel potential function as follows:
Pt = E[F (x̄t)] + 3ηC2gL 2 f
γ
1
K K∑ k=1 E[∥u(k)t − g(k)(x (k) t )∥2] + 4η K K∑ k=1 E[∥u(k)t − ūt∥2]
+ γηβ
4C2gL 2 f
1
K K∑ k=1 E[∥z(k)t − z̄t∥2] + 1 K K∑ k=1 E[∥x(k)t − x̄t∥2] .
(19)
Then, the task boils down to bound each term in this potential function. In particular, we obtain these bounds in Lemmas 3, 4, 5, 6, respectively in Appendix. With these lemmas, by setting the hyperparameter as shown in Theorem 1, we can get
Pt+1 − Pt ≤ − ηβ
2 E[∥∇F (x̄t)∥2] +
ηβ(3C2gσ 2 f + 3C 2 fσ 2 g′ + C 2 gL 2 fσ 2 g + 24γσ 2 g)
B
+ 24C2fσ 2 g′ + 8C 2 gσ 2 f + 3C 2 gL 2 fσ 2 g
B
γηβ
4C2gL 2 f
.
(20)
By summing t, we can complete the proof.
Proof sketch of Theorem 2 Similar to the convergence analysis of Algorithm 1, we also proposed a novel potential function to facilitate the theoretical analysis. Here, different from the proof of Theorem 1, we didn’t directly use E[∥s(k)t − s̄t∥2] in the potential function. Instead, we introduce an additional variable h(k)t , which is defined as follows:
h (k) −1 = 0 , h (k) t = ∑ j∈Nvk wkjh (j) t−1 + q (k) t − q (k) t−1 , (21)
where q(k)t = ∇g(k)(x (k) t ) T∇gf (k)(u(k)t ) and q (k) −1 = 0. Here, the difference between h (k) t and s (k) t is that h(k)t uses the full gradient of each function while s (k) t uses the stochastic gradient. Note that h (k) t is just used for theoretical analysis. It does NOT need to be computed in practice. Based on this new variable, we define the following potential function:
Pt = E[F (x̄t)] + 4ηC2gL 2 f
γ
1
K K∑ k=1 E[∥u(k)t − g(k)(x (k) t )∥2] + η(1− λ) γ 1 K K∑ k=1 E[∥h(k)t − h̄t∥2]
+ 1
γ
1
K K∑ k=1 E[∥x(k)t − x̄t∥2] .
(22) Similar to the proof of Theorem 1, the task boils down to bound each term in the potential function, respectively. The detailed bounds can be found in Lemmas 9, 13, 11. Then, we can also get the bound for Pt+1 − Pt and complete the proof by summing t over all iterations like the proof of Theorem 1.

heading: 5 Experiment

heading: 5.1 Model-Agnostic Meta-Learning
text: Model-Agnostic Meta-Learning (MAML) [6] is to learn a meta-initialization model parameter for a set of new tasks. Specifically, it is to optimize the following model:
min x∈Rd
Ei∼p,Ditest li ( EDitrain ( x− α∇li ( x,Ditrain )) , Ditest ) , (23)
where p is the task distribution, α is the learning rate, li is the loss function for the i-th task, Ditrain is the training set of the i-th task and Ditest is the testing set of the i-th task. Note that we omit the superscript k for simplication. MAML can be reformulated as a stochastic compositional optimization problem as follows:
f(y; ζ) = li ( y,Ditest ) , where ζ = ( i,Ditest ) g(x; ξ) = x− α∇li ( x,Ditrain ) , where ξ = Ditrain .
(24)
Then, we can use the stochastic compostional gradient to optimize MAML. In particular, we can use ut to track g(xt) as follows:
ut = (1−γηt−1)ut−1+γηt−1g(xt;Bξ,t) = (1−γηt−1)ut−1+γηt−1(xt−1−α∇li ( xt−1, D i train ) ) . (25) In our experiment, we will use GP-DSCGD and GT-DSCGD to solve the distributed MAML problem, where tasks are distributed on different devices.

heading: 5.2 Regression Task
text: Following [6], we verify the performance of our methods with the regression problem. In particular, each task in MAML is a regression task, which maps the input to a sine wave. Different sine waves have different amplitudes and phases. Specifically, to generate a set of sine waves, the amplitude is uniformly drawn from [0.1, 5.0] and the phase is uniformly drawn from [0, π]. Additionally, the input a is randomly drawn from [−5.0, 5.0]. Then, we use a parameterized neural network f(x; a) to approximate the sine wave. In our experiment, f(x; a) has two hidden layers and the dimensionality of different layers is [1, 40, 40, 1]. In each hidden layer, we use the ReLU function as the activation function. There are numerous regression tasks, the goal of MAML is to learn a good initialization for the model parameter x such that it can be quickly adapted to new regression tasks and have good performance.
In our experiment, we use four GPUs where each GPU is viewed as a device. These devices compose a ring graph. On each device, the number of tasks in each meta-batch is set to 200. The number of samples in each task for training is set to 10. As for the testing set, the number of tasks is set to 500 and the number of samples in each task is also set to 10. In addition, the number of iterations for adaptation in the training phase is set to 1 while it is set to 10 in the testing phase. To verify the performance of our methods, we compare them with the gossip-based decentralized SGD (DSGD)
method. Specifically, the decentralized SGD method directly uses ut = g(xt). In other words, DSGD directly uses xt−1−α∇li ( xt−1, D i train ) for adaptation rather than that in Eq. (25). In our experiment, α is set to 0.01. For adaptation, the learning rate η of DSGD is set to η = 0.001. Sine Adam has better performance in MAML [6], we actually use the decentralized Adam in our experiment as the baseline method. Accordingly, for fair comparison, we also use the adaptive learning rate for our two methods based on the stochastic compositional gradient z(k)t and s (k) t . In addition, we set γ = 3.0, η = 0.03, β = 0.33. The reason for this setting is that x(k)t+1 ≈ x (k) t − βηz (k) t . Therefore, the effective learning rate βη of our methods is approximately equal to that of the baseline method, which is a fair comparison.
In Figure 1(a), we show the average and standard deviation of the training loss across different devices. It can be seen that our two methods converge much faster than DSGD. The reason is that our methods use ut to track the g(xt) so that the estimation variance for g(xt) is smaller than DSGD. Meanwhile, it can be seen that our two methods achieve similar performance. Moreover, in Figure 1(b), we plot the loss function value when using the learned model parameter as the initialization for new tasks and then conducting 10 gradient descent steps. It can be seen that our two methods can achieve better performance in adaptation to the new task.

heading: 5.3 Classification Task
text: Besides the regression task, we also verify the performance of our distributed training methods with the classification task. Here, we use the Omniglot image dataset [12]. It has 1623 characters and each character has 20 images. Under the MAML setting, recognizing each character corresponds to a task so that there are 1623 tasks totally. In our experiment, we use 1200 tasks as the training set and the rest tasks as the testing set. To recognize the character, we use the same convolutional neural network as [6]. To get the meta-initialization for the model parameter, we use the 5-way 1-shot setting. As
for the distributed training, we still use four devices. The meta batch-size on each device is 8. The number of iterations for adaptation in the training phase is set to 1 while it is set to 3 in the testing phase. Additionally, we set the inner learning rate α = 0.4, the learning rate of DSGD η = 0.001. As for our methods, we keep the same learning rates as the regression task.
From Figure 1(c), we also plot the loss function value for the classification task. We can still find that our two methods converge much faster than DSGD, which further confirms the effectiveness of our proposed methods. Additionally, we show the classification accuracy after adaptation for new tasks in Figure 1(d). It is obvious that our two methods can achieve higher accuracy and smaller variance compared with DSGD, which further confirms the effectiveness of our methods.

heading: 5.4 Additional Experiment
text: In this experiment, we verify the performance of our methods with different number of devices. In particular, we compare the performance using four GPUs and that using eight GPUs for the regression task. For these two settings, we use the same overall meta-batch size as that in Section 5.2, and then uniformly distribute the tasks on all devices. As for other hyperparameters, they are the same for these two settings. In Figure 2, we plot the training loss function value with respect to the consumed time. It can be observed that our two methods converge much faster when using more devices, which confirms the speedup of the decentralized training.

heading: 6 Conclusion
text: In this paper, we studied the distributed training methods for the stochastic compositional optimization problem. In particular, we proposed two novel decentralized stochastic compositional gradient descent method based on the gossip communication mechanism and the gradient tracking communication mechanism. As far as we know, our work is the first one facilitating distributed training for large-scale stochastic compositional optimization problem. Meanwhile, we established the convergence rate of our proposed methods with novel theoretical analysis. The theoretical results indicate that our two methods can achieve linear speedup with respect to the number of devices. This is the first work disclosing this favorable result. At last, extensive experimental results on the MAML task demonstrate the superior performance of our methods.
Acknowledgments and Disclosure of Funding
This work was partially supported by NSF IIS 1845666, 1852606, 1838627, 1837956, 1956002, OIA 2040588. In addition, this research includes calculations carried out on HPC resources supported in part by the National Science Foundation through major research instrumentation grant number 1625061 and by the US Army Research Laboratory under contract number W911NF-16-2-0189.

Review summaries:
The paper presents two algorithms for distributed optimization of stochastic composite objectives where one is interested in optimizing one outer objective that is computed on points that depend on another function (called the inner objective in the paper). The authors adapt two strategies for distributed optimization to the stochastic composite optimization problem and derive convergence rates for the general non-convex outer objective problems, that show linear rates w.r.t. the number of devices used (this is probably the main contribution of this paper). Finally, they present a set of experiments in meta-learning using MAML as meta-learning algorithm.

This paper studies the stochastic compositional optimization problem under the distributed learning setting. Two decentralized optimization methods are proposed for solving this problem, both of which are based on SCGD method. Theoretical analysis is conducted to provide convergence bound for both methods. Experiments are also carried out accordingly.

This paper studies decentralized stochastic compositional optimization problem, and proposes two algorithms: GP-DSCGD and GT-DSCGD. They also have convergence guarantee for both algorithms. This problem has significant application in model agnostic meta learning.