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