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: This paper proposes a sketching-based federated learning algorithm to address the communication and privacy concerns in the federated learning setting. In the proposed algorithm, each client projects its local update (gradient) to a lower dimension via a sketching matrix before communicating the update to the server. The server aggregates the sketched local updates from different clients and communicates the aggregated update to the clients. The clients then de-sketch the message received from the server to update their local model.   The paper shows that if the sketching matrix satisfies the "coordinate-wise embedding" property, then the proposed algorithm leads to provable convergence for smooth and convex objective functions.   The paper also discusses the issue of privacy and proposes to add Gaussian noise to the sketched gradients to claim differential privacy of the resulting learning algorithm.