Neural Tangent Kernel Empowered Federated Learning

1 INTRODUCTION . Federated learning ( FL ) has emerged as a popular paradigm involving a large number of workers collaboratively solving a machine learning problem ( Kairouz et al. , 2021 ) . In a typical FL framework , a server broadcasts a global model to selected workers and collects model updates without needing to access the raw data . One popular algorithm is known as federated averaging ( FedAvg ) ( McMahan et al. , 2017 ) , in which workers perform stochastic gradient descent ( SGD ) to update the local models and upload the weight vectors to the server . A new global model is constructed on the server by averaging the received weight vectors . Li et al . ( 2020 ) characterized some unique challenges of FL . First , client data are generated locally and remain decentralized , which implies that they may not be independent and identically distributed ( IID ) . Prior works have shown that statistical heterogeneity can negatively influence the convergence of FedAvg ( Zhao et al. , 2018 ) . This phenomenon may be explained that local updating under data heterogeneity will cause cost-function inconsistency ( Wang et al. , 2020 ) . More challengingly , the learning procedure is susceptible to system heterogeneity , including the diversity of hardware , battery power , and network connectivity . Local updating schemes often exacerbate the straggler issue caused by heterogeneous system characteristics . Recent studies have proposed various strategies to alleviate the statistical heterogeneity . One possible solution is to share a globally available dataset with participants to reduce the distance between client-data distributions and the population distribution ( Zhao et al. , 2018 ) . In practice , though , such a dataset may be unavailable or too small to meaningfully compensate for the heterogeneity . Some researchers replaced the coordinate-wise weight averaging strategy in FedAvg with nonlinear aggregation schemes ( Wang et al. , 2020 ; Chen & Chao , 2021 ) . The nonlinear aggregation relies on a separate optimization routine , which can be elusive , especially when the algorithm does not converge well . Another direction is to modify the local objectives or local update schemes to cancel the effects of client drift ( Li et al. , 2020 ; Karimireddy et al. , 2020 ) . However , some studies reported that these methods are not consistently effective , and may perform worse than FedAvg when evaluated in various settings ( Reddi et al. , 2021 ; Haddadpour et al. , 2021 ; Chen & Chao , 2021 ) . In this work , we present a neural tangent kernel empowered federated learning ( NTK-FL ) paradigm . Given a fixed number of communication rounds , NTK-FL outperforms state-of-the-art methods in terms of test accuracy . We summarize our contributions as follows . • We propose a novel FL paradigm without requiring workers to perform gradient descent . To the best of our knowledge , this is the first work using the NTK method to replace gradient descent in FL algorithms . • Our scheme inherently solves the non-IID data problem of FL . Compared with FedAvg , it is robust to different degrees of data heterogeneity and has a consistently fast convergence speed . We verify the effectiveness of the paradigm theoretically and experimentally . • We add communication-efficient and privacy-preserving features to the paradigm and develop CP-NTK-FL by combining strategies such as random projection and data subsampling . We show that some strategies can also be applied to traditional FL methods . Although such methods cause performance degradation when applied to FedAvg , they only slightly worsen the model accuracy when applied to the proposed CP-NTK-FL . 2 RELATED WORK . Neural Tangent Kernel . Jacot et al . ( 2018 ) showed that training an infinitely wide neural network with gradient descent in the parameter space is equivalent to kernel regression in the function space . Lee et al . ( 2019 ) used a first-order Taylor expansion to approximate the neural network output and derived the training dynamics in a closed form . Chen et al . ( 2020 ) established the generalization bounds for a two-layer over-parameterized neural network with the NTK framework . The NTK computation has been extended to convolutional neural networks ( CNNs ) ( Arora et al. , 2019 ) , recurrent neural networks ( RNNs ) ( Alemohammad et al. , 2021 ) , and even to neural networks with arbitrary architectures ( Yang & Littwin , 2021 ) . Empirical studies have also provided a good understanding of the wide neural networks training ( Lee et al. , 2020 ) . Federated Learning . FL aims to train a model with distributed workers without transmitting local data ( McMahan et al. , 2017 ; Kairouz et al. , 2021 ) . FedAvg has been proposed as a generic solution with many theoretical analyses and implementation variants . Recent studies have shown a growing interest in improving its communication efficiency , privacy guarantees , and robustness to heterogeneity . To reduce communication cost , gradient quantization and sparsification were incorporated into FedAvg ( Reisizadeh et al. , 2020 ; Sattler et al. , 2019 ) . From the security perspective , Zhu et al . ( 2019 ) showed that sharing gradients may cause privacy leakage . To address this challenge , differentially private federated optimization and decentralized aggregation methods were developed ( Girgis et al. , 2021 ; Cheng et al. , 2021 ) . Other works put the focus on the statistical heterogeneity issue and designed various methods such as adding regularization terms to the objective function ( Li et al. , 2020 ; Smith et al. , 2017 ) . In this work , we focus on a novel FL paradigm where the global model is derived based on the NTK evolution . We show that the proposed NTK-FL is robust to statistical heterogeneity by design , and extend it to a variant with improved communication efficiency and enhanced privacy . Kernel Methods in Federated Learning . The NTK framework has been mostly used for convergence analyses in FL . Seo et al . ( 2020 ) studied two knowledge distillation methods in FL and compared their convergence properties based on the neural network function evolution in the NTK regime . Li et al . ( 2021 ) incorporated batch normalization layers to local models , and provided theoretical justification for its faster convergence by studying the minimum nonnegative eigenvalue of the tangent kernel matrix . Huang et al . ( 2021 ) directly used the NTK framework to analyze the convergence rate and generalization bound of two-layer ReLU neural networks trained with FedAvg . Su et al . ( 2021 ) studied the convergence behavior of a set of FL algorithms in the kernel regression setting . In comparison , our work does not focus on pure convergence analyses of existing algorithms . We propose a novel FL framework by replacing the gradient descent with the NTK evolution . 3 BACKGROUND AND PRELIMINARIES . We use lowercase nonitalic boldface , nonitalic boldface capital , and italic boldface capital letters to denote vectors , matrices , and tensors , respectively . For example , for column vectors aj ∈ RM , j ∈ { 1 , . . . , N } , A = [ a1 , . . . , aN ] is an M × N matrix . A third-order tensor A ∈ RK×M×N can be viewed as a concatenation of such matrices . We use a slice to denote a matrix in a third-order tensor by varying two indices ( Kolda & Bader , 2009 ) . Take tensor A , for instance : Ai : : is a matrix of the ith horizontal slice , and A : j : is its jth lateral slice ( Kolda & Bader , 2009 ) . Finally , the indicator function of an event is denoted by 1 ( · ) . 3.1 FEDERATED LEARNING MODEL . Consider an FL architecture where a server trains a global model by indirectly using datasets distributed among M workers . The local dataset of the mth worker is denoted by Dm = { ( xm , i , ym , i ) } Nmi=1 , where ( xm , i , ym , i ) is an input-output pair , drawn from a distribution Pm . The local objective can be formulated as an empirical risk minimization over Nm training examples : Fm ( w ) = 1 Nm ∑Nm i=1R ( w ; xm , j , ym , i ) , where R is a sample-wise risk function quantifying the er- ror of model with a weight vector w ∈ Rd estimating the label ym , i for an input xm , i . The global objective function is denoted by F ( w ) , and the optimization problem may be formulated as : min w∈Rd F ( w ) = 1 M M∑ m=1 Fm ( w ) . ( 1 ) 3.2 LINEARIZED NEURAL NETWORK MODEL . Let ( xi , yi ) denote a training pair , with xi ∈ Rd1 and yi ∈ Rd2 , where d1 is the input dimension and d2 is the output dimension . X , [ x1 , . . . , xN ] > represents the input matrix and Y , [ y1 , . . . , yN ] > represents the label matrix . Consider a neural network function f : Rd1 → Rd2 parameterized by a vector w ∈ Rd , which is the vectorization of all weights for the multilayer network . Given an input xi , the network outputs a prediction ŷi = f ( w ; xi ) . Let ` ( ŷi , yi ) be the loss function measuring the dissimilarity between the predicted result ŷi and the true label yi . We are interested in finding an optimal weight vector w ? that minimizes the empirical loss over N training examples : w ? = argmin w L ( w ; X , Y ) , 1 N N∑ i=1 ` ( ŷi , yi ) . ( 2 ) One common optimization method is the gradient descent training . Given the learning rate η , gradient descent updates the weights at each time step as : w ( t+1 ) = w ( t ) − η∇wL . To simplify the notation , let f ( t ) ( x ) be the output at time step twith an input x , i.e. , f ( t ) ( x ) , f ( w ( t ) ; x ) . Following Lee et al . ( 2019 ) , we use the first-order Taylor expansion around the initial weight vector w ( 0 ) to approximate the neural network output given an input x , i.e. , f ( t ) ( x ) ≈ f ( 0 ) ( x ) + J ( 0 ) ( x ) ( w ( t ) −w ( 0 ) ) , ( 3 ) where J ( 0 ) ( x ) = [ ∇f ( 0 ) 1 ( x ) , . . . , ∇f ( 0 ) d2 ( x ) ] > , with ∇f ( t ) j ( x ) , [ ∂ŷ ( t ) j /∂w ( t ) 1 , . . . , ∂ŷ ( t ) j /∂w ( t ) d ] > being the gradient of the jth component of the neural network output with respect to w ( t ) . Consider the halved mean-squared error ( MSE ) loss ` , namely , ` = 1d2 ∑d2 j=1 1 2 ( ŷj − yj ) 2 . Based on the continuous-time limit , one can show that the dynamics of the gradient flow are governed by the following differential equation : df dt = −ηH ( 0 ) ( f ( t ) ( X ) −Y ) , ( 4 ) where f ( t ) ( X ) ∈ RN×d2 is a matrix of concatenated output for all training examples , and H ( 0 ) is the neural tangent kernel at time step 0 , with each entry ( H ( 0 ) ) ij equal to the scaled Frobenius inner product of the Jacobian matrices : ( H ( 0 ) ) ij = 1d2 〈 J ( 0 ) ( xi ) , J ( 0 ) ( xj ) 〉 F . The differential equation ( 4 ) has the closed-form solution : f ( t ) ( X ) = ( I− e− ηt N H ( 0 ) ) Y + e− ηt N H ( 0 ) f ( 0 ) ( X ) . ( 5 ) The neural network state f ( t ) ( X ) can thus be directly obtained from ( 5 ) without running the gradient descent algorithm .

Summary: This paper presents a new federated learning (FL) framework that employs the neural tangent kernel (NTK) method to replace the widely used gradient descent algorithm for optimization. To improve communication efficiency and privacy-preserving features, data sampling and random projection techniques are used in the proposed FL-NTK framework. Experiments are conducted to demonstrate the advantages of the proposed FL-NTK in the robustness to data heterogeneity and communication efficiency, as compared to the baseline FedAvg.