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 proposes an algorithm for federated learning that leverages recent advances in the NTK framework. Participants share sample-wise jacobian matrices instead of model weights or gradients, based on the intuition that the NTK is able to capture useful statistical information when learning under statistical heterogeneity. The paper presents numerical results showing that the proposed approach maintains efficiency across different heterogeneity levels. Further, the authors introduce a practical implementation that reduces communication, and compare this approach with FedAvg and FedNova.