Iterative Sketching and its Application to Federated Learning

1 INTRODUCTION . Federated learning enables multiple parties to collaboratively train a machine learning model without directly exchanging training data . This has become particularly important in areas of artificial intelligence where users care data privacy , security , and access rights , including healthcare ( Li et al. , 2020b ; 2019 ) , internet of things ( Chen et al. , 2020 ) , and fraud detection ( Zheng et al. , 2020 ) . Given the importance and popularity of federated learning , it has become an important research topic for academia and industry , mostly focusing on two central themes . One is on data privacy . Federated learning seemingly protects clients ’ privacy , since it only communicates gradient information . Unfortunately , recent studies ( Geiping et al. , 2020 ; Zhu & Han , 2020 ; Wang et al. , 2019 ) have demonstrated that attackers can recover the input data from the shared gradients . The reason why the attacks work is the gradients carry important information about the training data ( Ateniese et al. , 2015 ; Fredrikson et al. , 2015 ) . The second one is on communication efficiency . Machine learning models are becoming increasingly larger but client devices that carry private data only have limited network bandwidth . The size of the gradient is the same as the size of the model , and the amount of data to communication between the clients and the servers thus is large . This becomes even more problematic when conducting federated learning on mobile and edge devices , where the bandwidth of the network is low . The communication cost is one of the most important the key performance bottlenecks in federated learning systems Goga & Teixeira ( 2012 ) ; Konečnỳ et al . ( 2016 ) . Many works try to address this challenge through local optimization method , such as local gradient descent ( GD ) , local stochastic gradient descent ( SGD ) and their variants ( Konečnỳ et al. , 2016 ; McMahan et al. , 2017 ; Stich , 2018 ) . Despite existing efforts , no work addresses both challenges simultaneously as far as we concern . Therefore , we ask the following question : Is there a FL framework that protects the local privacy and tackles the communication challenge ? In this paper , we achieve these goals by using an old but powerful idea — the Johnson-Lindenstrauss transform and its fast variations ( Fast-JL , see Ailon & Chazelle ( 2006 ) ) . If we view the transform as a sketch matrix , and its transpose as de-sketch , then our main idea is to iteratively apply the sketch and de-sketch matrices to gradients . Instead of running the vanilla gradient descent w ( t+1 ) ← w ( t ) − η · g ( t ) using true gradient g ( t ) ∈ Rd , we apply sketch and de-sketch to the gradient : w ( t+1 ) ← w ( t ) − η ·R > ·R · g ( t ) . Here R ∈ Rbsketch×d denotes a sketching matrix that sketches the true gradient to a lower dimension andR > ∈ Rd×bsketch denotes the de-sketching process that maps the sketched gradient back to the true gradient dimension . The coordinate-wise embedding property ( Song & Yu , 2021 ) ensures R > Rg ( t ) being an unbiased estimator of g ( t ) with bounded second moments , implying the new sketched gradient descent scheme preserving the original convergence properties . Hence , all clients will only communicate sketched gradients to the server , the server averages the sketched gradients and broacasts it back to all clients . Finally , each client de-sketches the received gradients and perform local updates . Since the sketching dimension is always small compared to original dimension , we save communication cost per iteration via Johnson-Lindenstrauss . Though the communication problem has been addressed , such framework still faces privacy challenge : applying Johnson-Lindenstrauss “ masks ” the communicated gradients , but in order to give a provable privacy guarantee on this framework , we introduce additive low-dimensional Gaussian noises to make sure that the communicated vectors themselves are differential private . We summarize the contributions in this work as follows : Framework contribution : We propose a new federated learning framework that iteratively applies sketch and de-sketch matrices to gradients , which enjoys the following advantages : • Privacy : it preserves the privacy via additive low-dimensional Gaussian noise . • Communication : it reduces communication per round , since at each synchronization step , only a sketched gradient of lower-dimensional is communicated . • Convergence : it preserves the convergence rate of vanilla local gradient descent method . Technical contribution : Our analysis technique is also of independent interest in the relating areas : • Unlike classical sketch-and-solve paradigm , our iterative sketch and de-sketch method can be combined with gradient-based methods and extended to broader optimization problems . • We provide rigorous analysis on the impact of introducing sketching through coordinate-wise embedding , which can be generalized to other areas ( Song & Yu , 2021 ) . • As a by-product , we give a novel linear convergence result of local GD under the strongly-convex and smooth scheme . 2 RELATED WORK . Federated Learning Federated learning ( FL ) is an emerging framework in distributed deep learning . FL allows multiple parties or clients collaboratively train a model without data sharing . Fl let the local client perform most of the computation and a central sever update the model parameters through aggregation then transfers the parameters to local models ( Dean et al. , 2012 ; Shokri & Shmatikov , 2015 ; McMahan et al. , 2016 ; 2017 ) . In this way , the details of the data are not disclosed in between each party . Unlike the standard parallel setting , FL has three unique challenge ( Li et al. , 2020a ) , including communication cost , data heterogeneity and client robustness . In our work , we focus on the first two challenges . The training data are massively distributed over an incredibly large number of devices , and the connection between the central server and a device is slow . A direct consequence is the slow communication , which motivated communication-efficient FL algorithm . Federated average ( FedAvg ) ( McMahan et al. , 2017 ) firstly addressed the communication efficiency problem by introducing a global model to aggregate local stochastic gradient descent updates . Later , different variations and adaptations have arisen . This encompasses a myriad of possible approaches , including developing better optimization algorithms ( Wang et al. , 2020a ) and generalizing model to heterogeneous clients under special assumptions ( Zhao et al. , 2018 ; Kairouz et al. , 2019 ; Li et al. , 2021 ) . Local GD and Local SGD To seek the communication efficiency of Federated learning , local SGD has been proposed ( Konečnỳ et al. , 2016 ; McMahan et al. , 2017 ) , where each client does a few SGD iterations locally before the server averages the local estimators . Different variants of local SGD algorithms has been explored , including with momentum ( Yu et al. , 2019b ; Wang et al. , 2018 ) , with quantization ( Basu et al. , 2019 ; Reisizadeh et al. , 2020 ) , and with various variance-reduction methods ( Liang et al. , 2019 ; Karimireddy et al. , 2020 ) . Convergence analysis for local SGD mainly focuses on two regimes : identical data regime ( Stich , 2018 ; Basu et al. , 2019 ; Stich & Karimireddy , 2019 ; Haddadpour & Mahdavi , 2019 ; Khaled et al. , 2020 ) and heterogeneous data regime ( Jiang & Agrawal , 2018 ; Yu et al. , 2019a ; Basu et al. , 2019 ; Haddadpour & Mahdavi , 2019 ; Khaled et al. , 2019 ; 2020 ) . In this work , we propose our framework based upon vanilla local GD and our analysis focus on the heterogeneous data regime . Sketching Sketching has many applications in numerical linear , such as linear regression , lowrank approximation ( Clarkson & Woodruff , 2013 ; Nelson & Nguyên , 2013 ; Meng & Mahoney , 2013 ; Boutsidis & Woodruff , 2014 ; Song et al. , 2017 ; Andoni et al. , 2018 ) , distributed problems ( Woodruff & Zhong , 2016 ; Boutsidis et al. , 2016 ) , reinforcement learning Wang et al . ( 2020b ) , tensor decomposition ( Song et al. , 2019 ) , clustering ( Esfandiari et al. , 2021 ) , cutting plane method ( Jiang et al. , 2020 ) , generative adversarial networks ( Xiao et al. , 2018 ) and linear programming ( Lee et al. , 2019 ; Jiang et al. , 2021 ; Song & Yu , 2021 ) . Notations For a positive integer n , we use [ n ] to denote the set { 1 , 2 , · · · , n } . We use E [ · ] to denote expectation ( if it exists ) , and use Pr [ · ] to denote probability . For a vector x , we use ‖x‖2 : = ( ∑n i=1 x 2 i ) 1/2 to denote its ` 2 norm . We denote 1 { x=l } for l ∈ R to be the indicator function which equals to 1 if x = l and 0 otherwise . Let f : A → B and g : C → A be two functions , we use f ◦ g to denote the composition of functions f and g , i.e. , for any x ∈ C , ( f ◦ g ) ( x ) = f ( g ( x ) ) . We denote Id to be the identity mapping . 3 PROBLEM SETUP . Consider a federated learning scenario with N clients and corresponding local losses fc : Rd → R , our goal is to find min x∈Rd f ( x ) : = 1 N N∑ c=1 fc ( x ) ( 1 ) In this work , we consider the following classical convex and smooth setting for our objectives . Assumption 3.1 . Assume that the set of minimizers of ( 1 ) is nonempty . Each fc is µ-strongly convex for µ ≥ 0 and L-smooth . That is , for all x , y ∈ Rd , µ 2 ‖y − x‖22 ≤ fc ( y ) − fc ( x ) + 〈y − x , ∇fc ( x ) 〉 ≤ L 2 ‖y − x‖22 . Note in the case µ = 0 , this assumption reduces back to convexity and smoothness . Apart from the above assumption , we allow local losses to have arbitrary heterogeneity . In other words , we allow fc ’ s to be arbitrary functions . 4 OUR ALGORITHM . In this section , we propose a federated learning framework that addresses the communication efficiency issue . When the learning gradients are of high dimension , classical federated learning framework which sends the exact gradient could incur a heavy communication cost per round . Sketching technique , which emerges as an effective way to reduce the dimension of vector while preserving significant amount of information ( Sarlós , 2006 ; Woodruff , 2014 ) , is highly preferred in this setting . It enables us to compress the gradient vector into a lower dimension while preserving convergence rates , and greatly saves the communication cost per round . Motivated by above discussion , we propose the iterative sketching-based federated learning algorithm , which builds upon vanilla local gradient descent : we start with a predetermined sequence of independent sketching matrices shared across all clients . In each round , local clients accumulate and sketch its change over K local steps , then transmit the low-dimensional sketch to the server . Server then averages the sketches and transmits them back to all clients . Upon receiving , each client de-sketches to update the local model . Algorithm 1 Iterative Sketching-based Federated Learning Algorithm with K local steps 1 : procedure ITERATIVESKETCHINGFL 2 : Each client initializes w0 using the same set of random seed 3 : for t = 1→ T do . T denotes the total number of global steps 4 : / * Client * / 5 : parfor c = 1→ N do . N denotes the total number of clients 6 : if t = 1 then 7 : ut,0c ← w0 . ut,0c ∈ Rd 8 : else 9 : ut,0c ← wt−1 + deskt ( ∆w̃t−1 ) . deskt : Rbsketch → Rd de-sketch the change 10 : end if 11 : wt ← ut,0c 12 : for k = 1→ K do 13 : ut , kc ← ut , k−1c − ηlocal · ∇fc ( w ) |w=ut , k−1c 14 : end for 15 : ∆wc ( t ) ← ut , Kc − wt 16 : Client c sends skt ( ∆wc ( t ) ) to server . skt : Rd → Rbsketch sketch the change 17 : end parfor 18 : / * Server * / 19 : ∆w̃t ← ηglobal · 1N ∑N c=1 skt ( ∆wc ( t ) ) . ∆w̃ t ∈ Rd 20 : Server sends ∆w̃t to each client 21 : end for 22 : end procedure We highlight several distinct features of our algorithm : •Communication : In each sync step , we only communicates a low-dimensional sketched gradients , indicating a smaller communication cost per round . This property is particularly valuable in a smallbandwidth setting . • De-sketch : i : We emphasize that unlike the classical sketch-and-solve paradigm that decreases the problem dimension , our algorithm applies sketching in each round , combined with a de-sketching process which recovers back to the true gradient dimension . • Simple server task : : Server only needs to do simple averaging , indicating no need of a trustworthy party as the server . • Decentralization : : Our algorithm can be generalized to decentralized learning settings , where local clients can only communicate with neighboring nodes . In this case , it requiresO ( diam ) rounds to propagate the sketched local changes , where diam is the diameter of the network graph . 4.1 sk/desk VIA COORDINATE WISE EMBEDDING In this section , we discuss the concrete realization of the skt/deskt operators in Algorithm 1 through random sketching matrices . Note we should require any processed gradient deskt ◦ skt ( g ) to `` be close '' to the true gradient g to avoid breaking the convergence property of the algorithm . To achieve this , we first introduce the following property for a broad family of sketching matrices , namely coordinate-wise embedding ( Jiang et al. , 2021 ; Song & Yu , 2021 ) , that naturally connects with skt/deskt operators . Definition 4.1 ( a-coordinate-wise embedding ) . We say a randomized matrix R ∈ Rbsketch×d satisfying a-coordinate wise embedding if for any vector g , h ∈ Rd , we have ER∼Π [ h > R > Rg ] = h > g and ER∼Π [ ( h > R > Rg ) 2 ] ≤ ( h > g ) 2 + absketch ‖h‖ 2 2 · ‖g‖22 . iWe elaborate the difference of our iterative sketch and de-sketch approach compared to classical sketchand-solve approaches in Section 4.2 . In general , well-known sketching matrices have their coordinate-wise embedding parameter a being a small constant ( See Section D ) . Note that if we choose h to be one-hot vector ei , then the above conditions translate to ER∼Π [ R > Rg ] = g and ER∼Π [ ‖R > Rg‖22 ] ≤ ( 1 + a · dbsketch ) · ‖g‖ 2 2 . This implies that by choosing skt = Rt ∈ Rbsketch×d ( sketching ) , deskt = R > t ∈ Rd×bsketch ( de-sketching ) ( 2 ) for any iteration t ≥ 1 , where Rt ’ s are independent random matrices with sketching dimension bsketch , we obtain an unbiased sketching/de-sketching scheme with bounded variance as state in the following Theorem 4.2 . Theorem 4.2 . Let skt and deskt be defined by Eq . ( 2 ) using a sequence of independent sketching matrices Rt ∈ Rbsketch×d satisfying a-coordinate wise embedding property ( Def . 4.1 ) . Then the following properties hold : 1 . Independence : For different iterations , ( skt , deskt ) ’ s are independent of each other . 2 . Linearity : Both skt and deskt are linear operators . 3 . Unbiased estimator : For any fixed vector h ∈ Rd , it holds E [ deskt ( skt ( h ) ) ] = h. 4 . Bounded second moment : For any fixed vector h ∈ Rd , it holds E [ ‖deskt ( skt ( h ) ) ‖22 ] ≤ ( 1 + a · d/bsketch ) · ‖h‖22 . We will use the above property to instantiate the convergent proof and communication complexity in section 5 . We remark that unlike traditional sketching matrix R , one can intuitively think of matrix R > as a “ de-sketch ” matrix , it undoes sketching and recovers the sketched vector to original dimension .

Summary: The paper studies the application of linear sketching algorithms to reduce the communication complexity in federated Learning problems. In particular, the paper proposes that the clients send a low-dimensional linear sketch of their updates to the server. Convergence rate of local SGD when deployed with such a method for aggregating the updates are discussed for various function classes. Finally, the paper also discusses a method for integrating differentially privacy guarantees in their method.