Fast Training  Method for  Stochastic Compositional Optimization Problems

1 Introduction . 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 . 2 Related Works . 2.1 Stochastic Compositional Optimization Problem . 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 . 2.2 Decentralized Training Methods . 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 . 3 Decentralized Stochastic Compositional Optimization Methods . 3.1 Problem Definition . 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 . 3.2 Gossip-based Decentralized Stochastic Compositional Gradient Descent Method . 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

Summary: 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.