Year: 2021

Abstract:
Federated learning (FL) is a privacy-preserving paradigm where multiple participants jointly solve a machine learning problem without sharing raw data. Unlike traditional distributed learning, a unique characteristic of FL is statistical heterogeneity, namely, data distributions across participants are different from each other. Meanwhile, recent advances in the interpretation of neural networks have seen a wide use of neural tangent kernel (NTK) for convergence and generalization analyses. In this paper, we propose a novel FL paradigm empowered by the NTK framework. The proposed paradigm addresses the challenge of statistical heterogeneity by transmitting update data that are more expressive than those of the traditional FL paradigms. Specifically, sample-wise Jacobian matrices, rather than model weights/gradients, are uploaded by participants. The server then constructs an empirical kernel matrix to update a global model without explicitly performing gradient descent. We further develop a variant with improved communication efficiency and enhanced privacy. Numerical results show that the proposed paradigm can achieve the same accuracy while reducing the number of communication rounds by an order of magnitude compared to federated averaging.

heading: 1 INTRODUCTION
text: 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.

heading: 2 RELATED WORK
text: 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.

heading: 3 BACKGROUND AND PRELIMINARIES
text: 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 (·).

heading: 3.1 FEDERATED LEARNING MODEL
text: 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)

heading: 3.2 LINEARIZED NEURAL NETWORK MODEL
text: 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.

heading: 4 PROPOSED FL PARADIGM VIA THE NTK FRAMEWORK
text: In this section, we present the NTK-FL paradigm (Figure 1) and then extend it to the variant CPNTK-FL (Figure 2) with improved communication efficiency and enhanced privacy. The detailed algorithm descriptions are presented as follows.
4.1 NTK-FL PARADIGM
NTK-FL follows four steps to update the global model in each round. First, the server will select a subset Ck of workers and broadcast to them a model weight vector w(k) from the kth round. Here, the superscript k is the communication round index, and it should be distinguished from the gradient descent time step t in Section 3.2. Second, each worker will use its local training data Dm to compute a Jacobian tensor J (k)m ∈ RNm×d2×d, ∀ m ∈ Ck, which is a concatenation of Nm sample-wise Jacobian matrices (J (k)m )i:: = [∇f (k) 1 (xm,i), . . . ,∇f (k) d2 (xm,i)] >, i ∈ {1, . . . , Nm}. The worker will then upload the Jacobian tensor J (k)m , labels Ym, and initial condition f (k)(Xm) to the server. The transmitted information corresponds to the variables in the state evolution of f (t) in (5). Third, the server will construct a global Jacobian tensor J (k) ∈ RN×d2×d based on received J (k)m ’s, with each worker contributing Nm horizontal slices to J (k).
We use a toy example to explain the process as follows. Suppose the server selects worker 1 and worker 2 in a certain round. Workers 1 and 2 will compute the Jacobian tensors J (k)1 and J (k) 2 , respectively. The global Jacobian tensor is constructed as:
J (k) i:: =
{ J
(k) 1,i:: , if i ∈ {1, . . . , N1},
J (k) 2,j:: , j = i−N1, if i ∈ {N1 + 1, . . . , N1 +N2}.
(6)
After obtaining the global Jacobian tensor J (k), the (i, j)th entry of the global kernel H(k) is calculated as the scaled Frobenius inner product of two horizontal slices of J (k), i.e., (H(k))ij = 1 d2 〈J (k)i:: ,J (k) j:: 〉F. For simplicity, Fourth, the server will perform the NTK evolution to obtain the updated neural network function f (k+1) and weight vector w(k+1). With a slight abuse of notation, let f (k,t) denote the neural network output at gradient descent step t in communication round k. The neural network function evolution dynamics and weight evolution dynamics are given by:
f (k,t) = ( I− e− ηt N H (k) ) Y(k) + e− ηt N H (k) f (k), (7a)
w(k,t) = d2∑ j=1 (J (k) :j: ) >R (k,t) :j + w (k), (7b)
where J (k):j: is the jth lateral slice of J (k), and R(k,t):j is the jth column of the residual matrix R (k,t) defined as follows:
R(k,t) , η
Nd2 t−1∑ u=0 [ Y(k) − f (k,u)(X(k)) ] . (8)
Note that X(k) = [X>1 , . . . ,X > Ck ] > denotes a concatenation of worker training examples, and Y(k) = [Y>1 , . . . ,Y > Ck ] > denotes a concatenation of worker labels. The weight evolution in (7b) is derived by unfolding the gradient descent steps. To update the global weight, the server performs the evolution with various integer steps {t1, . . . , tΨ} and selects the best one with the smallest loss value. Our goal is to minimize the training loss with a small generalization gap (Nakkiran et al., 2020). The updated weight is decided by the following procedure:
w(k+1) , w(k), t(k) = argmin tj L(f(w(k,tj); X(k),Y(k))). (9)
Alternatively, if the server has an available validation dataset, the optimal number of update steps can be selected based on the model validation performance. In practice, such a validation dataset can be obtained from held-out workers (Wang et al., 2021). Based on the closed-form solution in (7b), the search of t(k) over the grid {t1, . . . , tΨ} can be completed in O(Ψ) time.
Comparison of NTK-FL and Huang et al. (2021). Huang et al. (2021) presented the details of FedAvg by letting clients use local updates and upload gradients to train a two-layer neural network. In contrast, NTK-FL let each client transmit Jacobian matrices without performing local SGD steps. The model weight is updated via NTK evolution in (7b). The main differences include: (1) clients transmit more expressive Jacobian matrices to improve model performance in the non-IID FL setting; (2) the computation is shifted to the server.
Robustness Against Statistical Heterogeneity. In essence, statistical heterogeneity comes from the decentralized data of heterogeneous distributions owned by individual workers. If privacy is not an issue, the non-IID challenge can be readily resolved by mixing all workers’ datasets and training a centralized model. In NTK-FL, instead of building a centralized dataset, we use Jacobian matrices to construct a global kernel H(k), which is a concise representation of gathered data points from all selected workers. This representation is yet more expressive/less compact than that of a traditional FL algorithm. More precisely, the update being sent for NTK-FL regarding the ith training example of the mth worker for NTK-FL is Jm = [∇f1(xm,i), . . . ,∇fd2(xm,i)]>, whereas the gradient update being sent for FedAvg is ∇L(w; xm,i,ym,i) = 1d2 ∑d2 j=1 (ŷm,i,j − ym,i,j)∇fj(xm,i), a weighted sum of coordinates of Jm. By sending Jacobian matrices Jm and jointly processing them on the server, NTK-FL delays the more aggressive data aggregation step after the communication stage and therefore better approximates the centralized learning setting than FedAvg does.

heading: 4.2 CP-NTK-FL VARIANT
text: Compared to FedAvg, NTK-FL does not incur additional client computational overhead since calculating the Jacobian tensor enjoys the same communication efficiency with computing aggregated gradients. Without locally updating weight vectors, NTK-FL is faster than FedAvg on the client side. In this section, we focus on the perspectives of the communication efficiency and security in terms of data confidentiality and membership privacy.
For communication, we follow the widely adopted analysis framework in wireless communication to examine only the client uplink overhead, assuming that the downlink bandwidth is much larger and the server will have enough transmission power (Tran et al., 2019). In NTK-FL, the client uplink communication cost and space complexity are dominated by a third-order tensor J (k)m , i.e., an O (Nmd2d) complexity compared to O (d) in FedAvg. For security, we investigate a threat model where a curious server may perform membership inference attacks (Nasr et al., 2018) or data reconstruction attacks (Zhu et al., 2019). Compared to the averaged gradient, sample-wise Jacobian matrices are more expressive, which may facilitate such attacks from the aggregation server. We extend NTK-FL by combining various tools to solve the aforementioned problems without jeopardizing the performance severely. Although it is possible to incorporate these tools into FedAvg, we will show that overall it will lead to more severe accuracy drop.
Jacobian Dimension Reduction. First, we let themth worker sample a subset Bm from its dataset Dm uniformly for the training. Let β ∈ (0, 1) denote the sampling rate, Bm contains N ′m = βNm data points, with the training pairs denoted by (X′m,Y ′ m). Next, we consider using a random projection to preprocess the input data via a seed shared by a trusted key server. Formally, the sampled training examples are projected into Z′m, i.e., Z ′ m = X ′ mP, where P ∈ Rd1×d ′ 1 is a projection matrix generated based on a seed ρ with IID standard Gaussian entries. In general, we have d′1 < d1 and an non-invertible projection operation. The concept of trusted key server follows the trusted third party in cryptography (Van Oorschot, 2020), and we assume it will not be compromised.
These two steps can already reduce the uplink communication overhead and enhance privacy. We first examine the current Jacobian tensor J ′(k)m ∈ RN ′ m×d2×d ′ . Compared with its original version J (k)m , it has reduced dimensionality at the cost of certain information loss. Meanwhile, the random projection will defend against the data reconstruction attack, as the Jacobian tensor is now evaluated at the projected data Z′m. We empirically verify their impact on the test accuracy in Figure 3. We set d′1 = 100 and sampling rate β = 0.4, and train a multilayer perceptron with 100 hidden nodes on the Fashion-MNIST dataset (Xiao et al., 2017). The joint effect of these strategies is a slight accuracy drop in NTK-FL and a nonnegligible performance degradation in FedAvg.
Jacobian Compression and Shuffling. We use a compression scheme to reduce the size of the Jacobian tensor by zeroing out the coordinates with small magnitude (Alistarh et al., 2018). In addition to the communication efficiency, this compression scheme is empirically effective against the data reconstruction attack (Zhu et al., 2019). To further ensure the confidentiality and membership privacy, we introduce a shuffling server, inspired by some recent frameworks (Girgis et al., 2021; Cheng et al., 2021), to permute Jacobian tensors J (k)m ’s, neural network states f (k) m ’s, and labels Ym’s.
Based on (7b), we denote the model update by ∆w(k) , w(k+1) − w(k) = ∑d2 j=1(J (k) :j: ) >R (k,t) :j , which is a sum of matrix products. If rows and columns are permuted in synchronization, the weight update ∆w(k) will remain unchanged. Considering the high dimensionality of the neural network weight, the reconstruction attack becomes computationally infeasible. As provable differential privacy guarantee does not explicitly protect against the reconstruction attack (Zhang et al., 2020), we leave a thorough privacy study for future work.

heading: 5 ANALYSIS OF ALGORITHM
text: In this section, we analyze the loss decay rate between successive communication rounds in NTKFL and make comparisons with FedAvg. Similar to Du et al. (2019) and Dukler et al. (2020), we
consider a two-layer neural network f : Rd → R of the following form to facilitate our analysis:
f(x; V, c) = 1√ n n∑ r=1 crσ(v > r x), (10)
where x ∈ Rd1 is an input, vr ∈ Rd1 is the weight vector in the first layer, V = [v1, · · · ,vn], cr ∈ R is the weight in the second layer, and σ(·) is the rectified linear unit (ReLU) function, namely σ(z) = max(z, 0), applied coordinatewise. We state two assumptions as prerequisites. Assumption 1 The first layer vr’s are sampled fromN (0, α2I). The second layer cr’s are sampled from {−1, 1} with equal probability and are kept fixed during training. Assumption 1 gives the initial distribution of the neural network parameters. Similar assumptions can be found in Dukler et al. (2020). We add restrictions to the input data in the next assumption. Assumption 2 (Normalized input). The input data are normalized, i.e., ‖xi‖2 6 1,∀ i.
For this neural network model, the (i, j)th entry of the empirical kernel matrix H(k) given in (3.2) can be calculated as: (H(k))ij = 1nx > i xj ∑n r=1 1 (k) ir 1 (k) jr , where 1 (k) ir , 1{〈v (k) r ,xi〉 > 0}, and the term c2r is omitted according to Assumption 1. Define H ∞, whose (i, j)th entry is given by:
(H∞)ij , Ev∼N (0,α2I) [ x>i xj 1(v >xi > 0) 1(v >xj > 0) ] . (11)
Let λ0 denote the minimum eigenvalue of H∞, which is restricted in the next assumption. Assumption 3 The kernel matrix H∞ is positive definite, namely, λ0 > 0.
In fact, the positive-definite property of H∞ can be shown under certain conditions (Dukler et al., 2020). For simplicity, we omit the proof details and directly assume the positive definiteness of H∞ in Assumption 3. Next, we study the residual term ‖f (k)(X)−y‖22 in communication round k, where X = [X>1 , . . . ,X > M ] > ∈ RN×d1 denote a concatenation of client inputs and y = [y>1 , . . . ,y>M ]> ∈ RN denote a concatenation of client labels. We give the convergence result by analyzing how the residual term decays between successive rounds. Theorem 1 For the NTK-FL scheme under Assumptions 1 to 3, let the learning rate η = O ( λ0 N
) and the neural network width n = Ω ( N2
λ20 ln N
2
δ
) , then with probability at least 1−δ, the one-round
loss decay of NTK-FL is∥∥f (k+1)(X)− y∥∥2 2 6 ( 1− ηλ0
2N
)t(k) ∥∥f (k)(X)− y∥∥2 2 , (12)
where t(k) is the number of NTK update steps defined in (9).
The proof of Theorem 1 can be found in Appendix B. By studying the asymmetric kernel matrix caused by local update (Huang et al., 2021), we have the following theorem for FedAvg, where the proof can be found in Appendix C.
Theorem 2 For FedAvg under Assumptions 1 to 3, let the learning rate η = O (
λ0 τN |Ck|
) and the
neural network width n = Ω ( N2
λ20 ln N
2
δ
) , then with probability at least 1 − δ, the one-round loss
decay of FedAvg is ∥∥f (k+1)(X)− y∥∥2 2 6 ( 1− ητλ0
2N |Ck|
)∥∥f (k)(X)− y∥∥2 2 , (13)
where τ is the number of local iterations, and |Ck| is the cardinality of the worker set in round k. Remark 1 (Fast Convergence of NTK-FL). The convergence rate of NTK-FL is faster than FedAvg. To see this, we compare the Binomial approximation of the decay coefficient in Theorem 1 with the decay coefficient in Theorem 2, i.e., 1− η1t
(k)λ0 2N + O ( η21 ) < 1− η2τλ02N |Ck| , where η 1 for a large
N 1. The number of NTK update steps t(k) is chosen dynamically in (9), which is on the order of 102 to 103, whereas τ is often on the order of magnitude of 10 in literature (Reisizadeh et al., 2020; Haddadpour et al., 2021). One can verify that η1t(k)λ0 is larger than η2τλ0/|Ck| and draw the conclusion in (1).
1For example, if we have 100 clients, each of which has more than 100 data points, then N is on the order of 104. Considering the choice of the learning rate η1, the Binomial approximation holds in (1).

heading: 6 EXPERIMENTAL RESULTS
text: Federated Settings. We use three datasets, namely, MNIST (LeCun et al., 1998), FashionMNIST (Xiao et al., 2017), and EMNIST (Cohen et al., 2017) digits. All of them contain C = 10 categories. For MNIST and Fashion-MNIST, we follow Hsu et al. (2019) to simulate non-IID data with the symmetric Dirichlet distribution (Good, 1976). Specifically, for the mth worker, we draw a random vector qm ∼ Dir(α), where qm = [qm,1, . . . , qm,C ]> belongs to the (C − 1)-standard simplex. Images with category k are assigned to the mth worker in proportional to (100 · qm,k)%. The heterogeneity in this setting mainly comes from label skewness. For the EMNIST dataset, it has a federated version that splits the dataset into shards indexed by the original writer of the digits (Kairouz et al., 2021). The heterogeneity mainly comes from feature skewness. A multilayer perceptron model with 100 hidden nodes is chosen as the target neural network model. We consider a total of 300 workers and select 20 of them with equal probability in each round.
Convergence. We empirically verify the convergence rate of the proposed method. For FedAvg, we use the number of local iterations from {10, 20, . . . , 50} and report the best results. For NTK-FL, we choose t(k) over the set {100, 200, . . . , 2000}. We use the following methods that are robust to the non-IID setting as the baselines: (i) Data sharing scheme suggested by Zhao et al. (2018), where a global dataset is broadcasted to workers for local training; the size of the global dataset is set to be 10% of the total number of local data points. (ii) Federated normalized averaging (FedNova) (Wang et al., 2020), where the workers transmit normalized gradient vectors to the server. (iii) Centralized training simulation, where the server collects the data points from subset Ck of workers and performs gradient descent to directly train the global model. Clearly, scheme (iii) achieves the performance that can be considered as an upper bound of all other algorithms. The training curves over three repetitions are shown in Figure 4. More implementation details and the results on CIFAR-10 (Krizhevsky, 2009) can be found in Appendix A. Our proposed NTK-FL method shows consistent advantages over other methods in different non-IID scenarios.
Degree of Heterogeneity. In this experiment, we select the Dirichlet distribution parameter α from {0.1, 0.2, 0.3, 0.4, 0.5} and simulate different degrees of heterogeneity on Fashion-MNIST dataset. A smaller α will increase the degree of heterogeneity in the data distribution. We evaluate NTK-FL, DataShare, FedNova, and FedAvg model test accuracy after training for 50 rounds. The mean values over three repetitions are shown in Figure 5, where each point is obtained over five repetitions with standard deviation less than 1%. It can be observed that NTK-FL achieves stable test accuracy in different heterogeneous settings. In comparison, FedAvg and FedNova show a performance drop in the small α region. NTK-FL has more advantages over baselines methods when the degree of heterogeneity is larger.
Effect of Hyperparameters. We study the effect of the tunable parameters in CP-NTK-FL. We change the local data sampling rate β and dimension d′1, and evaluate the model test accuracy on the non-IID Fashion-MNIST dataset (α = 0.1) after 10 communication rounds. The results are shown in Figure 6. A larger data sampling rate β or a larger dimension d′1 will cause less information loss,
and are expected to achieve a higher test accuracy. The results also show that the scheme is robust to different combinations of hyperparameters.
Uplink Communication. We evaluate the uplink communication efficiency of CP-NTKFL (d′1 = 200, β = 0.3) by measuring the number of rounds and cumulative uplink communication cost to reach a test accuracy of 85% on non-IID Fashion-MNIST dataset (α = 0.1). The results over three repetitions are shown in Table 1. Compared with federated learning with compression (FedCOM) (Haddadpour et al., 2021), quantized SGD (QSGD) (Alistarh et al., 2017), and FedAvg, CP-NTK-FL achieves the goal within an order of magnitude fewer iterations, which is particularly advantageous for applications with nonnegligible encoding/decoding delays or network latency.

heading: 7 CONCLUSION AND FUTURE WORK
text: In this paper, we have proposed an NTK empowered FL paradigm. It inherently solves the statistical heterogeneity challenge. By constructing a global kernel based on the local sample-wise Jacobian matrices, the global model weights can be updated via NTK evolution in the parameter space. Compared with traditional algorithms such as FedAvg, NTK-FL has a more centralized training flavor by transmitting more expressive updates. The effectiveness of the proposed paradigm has been verified theoretically and experimentally.
In future work, it will be interesting to extend the paradigm for other neural network architectures, such as CNNs, residual networks (ResNets) (He et al., 2016), and RNNs. It is also worthwhile to further improve the efficiency of NTK-FL and explore its savings in wall-clock time. We believe the proposed paradigm will provide a new perspective to solve federated learning challenges.

heading: A ADDITIONAL EXPERIMENTAL RESULTS
text: We give the detailed setting of the learning rate and batch size. For the learning rate η, we search over the set {10−3, 3× 10−3, 10−2, 3× 10−2, 10−1}. The learning rate is fixed during the training. For the client batch size, we set it to 200 for all datasets. We evaluate different methods, including the centralized training simulation, data sharing method (Zhao et al., 2018), FedNova (Wang et al., 2020), FedAvg (McMahan et al., 2017), and the proposed NTK-FL on the non-IID CIFAR-10 dataset (Krizhevsky, 2009) and present the results in Figure 7. NTK-FL outperforms other FL algorithms and shows test accuracy close to the centralized simulation. The observation is consistent with the results in Figure 4.

heading: B PROOF OF THEOREM 1
text: Let Im denote a set of indices such that for i ∈ Im, (xi, yi) ∈ Dm. We first present some lemmas to facilitate the convergence analysis. In communication round k, define S(k)i as the set of indices corresponding to neurons whose activation pattern is similar to its initial state for an input xi:
S(k)i , { r ∈ {1, . . . , n} ∣∣ ∃ v, ‖v − v(0)r ‖2 6 R, 1(0)ir 6= 1 (v>xi > 0) }. (14) We upper bound the cardinality of Si in Lemma 1. Lemma 1 Under Assumption 1 to 2, with probability at least 1− δ, we have
|Si| 6 √ 2
π
nR δα , ∀ i ∈ {1, . . . , N}. (15)
Proof. To bound |Si| = n∑ r=1 1 (r ∈ Si), consider an event Air defined as follows:
Air , {∃ v, ‖v − v(0)r ‖2 6 R, 1 (0) ir 6= 1 ( v>xi > 0 ) }. (16)
Clearly, 1 (r ∈ Si) = 1 (Air). According to Assumption 2, ‖xi‖ 6 1, it can be shown that the event Air happens if and only if |(v(0)r )>xi| 6 R based on a geometric argument. Based on Assumption 1, we have (v(0)r )>xi ∼ N (0, α2). The probability of event Air is
P[Air] = P [ |(v(0)r )>xi| 6 R ] (17a)
= erf (
R√ 2α
) 6 √ 2
π
R α . (17b)
By Markov’s inequality, we have with probability at least 1− δ, n∑ r=1 1 (r ∈ Si) 6 √ 2 π nR δα . (18)
The proof is complete. We bound the perturbation of the kernel matrix H(k,t) in Lemma 2.
Lemma 2 Under Assumption 1 to 2, if ∀ r ∈ {1, . . . , n}, ‖v(k,t)r − v(0)r ‖2 6 R, then
‖H(k,t) −H(0)‖2 6 2 √
2NR√ πδα . (19)
Proof. We have
‖H(k,t) −H(0)‖22 6 ‖H(k,t) −H(0)‖2F (20a)
= N∑ i=1 N∑ j=1 [ (H(k,t))ij − (H(0))ij ]2 (20b)
= 1
n2 N∑ i=1 N∑ j=1 (x>i xj) 2 ( n∑ r=1 1 (k,t) ir 1 (k,t) jr − 1 (0) ir 1 (0) jr )2 . (20c)
Consider the event Air defined in (16). Let φ (k,t) ijr , 1 (k,t) ir 1 (k,t) jr − 1 (0) ir 1 (0) jr . If ¬Air and ¬Ajr happen, clearly we have |φ(k,t)ijr | = 0. Therefore, the expectation of |φ (k,t) ijr | can be bounded as
E [∣∣∣φ(k,t)ijr ∣∣∣] 6 P(Air ∪Ajr) · 1 (21a)
6 P(Air) + P(Ajr) (21b) ¬ 6 2 √ 2
π
R α , (21c)
where ¬ comes from (17b). By Markov’s inequality, we have with probability at least 1− δ, |φ(k,t)ijr | 6 2 √ 2
π
R δα . (22)
Plugging (22) into (20c) yields
‖H(k,t) −H(0)‖22 6 N2 n2 8n2R2 πδ2α2 = 8N2R2 πδ2α2 . (23)
Taking the square root on both sides completes the proof.
Lemma 3 With probability at least 1− δ, ‖H(0) −H∞‖2 6 N √ ln (2N2/δ)
2n . (24)
Proof. We have
‖H(0) −H∞‖22 6 ‖H(0) −H∞‖2F = N∑ i=1 N∑ j=1 [ (H(0))ij − (H∞)ij ]2 . (25)
Note that (H(0))ij = 1nx > i xj n∑ r=1 1 (0) ir 1 (0) jr , (H (0))ij ∈ [−1, 1]. By Hoeffding’s inequality, we have with probability at least 1− δ/n2,∣∣∣(H(0))ij − (H∞)ij∣∣∣ 6√ ln (2N2/δ) 2n . (26)
Applying the union bound over i, j ∈ [N ] yields ‖H(0) −H∞‖2 6 N √ ln (2N2/δ)
2n . (27)
The proof is complete.
Now we are going to prove Theorem 1. Theorem 1 For the NTK-FL scheme under Assumptions 1 to 3, let the learning rate η = O ( λ0 N ) and
the neural network width n = Ω ( N2
λ20 ln 2N
2
δ
) , then with probability at least 1 − δ, the one-round
loss decay of NTK-FL is ‖f (k+1)(X)− y‖22 ≤ (
1− ηλ0 2N
)t(k) ‖f (k)(X)− y‖22. (28)
Proof. Taking the difference between successive terms yields
f (k,t+1)(xi)− f (k,t)(xi) = 1√ n n∑ r=1 [ crσ ( (v(k,t+1)r ) >xi ) − crσ ( (v(k,t)r ) >xi )] . (29)
We decompose the difference term to the sum of dIi and d II i , based on the set Si:
dIi , 1√ n ∑ r/∈Si [ crσ ( (v(k,t+1)r ) >xi ) − crσ ( (v(k,t)r ) >xi )] , (30a)
dIIi , 1√ n ∑ r∈Si [ crσ ( (v(k,t+1)r ) >xi ) − crσ ( (v(k,t)r ) >xi )] . (30b)
Consider the residual term∥∥f (k,t+1)(X)− y∥∥2 2
(31a) = ∥∥f (k,t+1)(X)− f (k,t)(X) + f (k,t)(X)− y∥∥2
2 (31b) = ∥∥f (k,t)(X)− y∥∥2
2 + 2
〈 dI + dII, f (k,t)(X)− y 〉 + ∥∥f (k,t+1)(X)− f (k,t)(X)∥∥2
2 . (31c)
We will give upper bounds for the inner product terms 〈 dI, f (k,t)(X)− y 〉 , 〈 dII, f (k,t)(X)− y 〉 ,
and the difference term ∥∥f (k,t+1)(X) − f (k,t)(X)∥∥2
2 , separately. Based on the property of the set
Si, we have
dIi = − η√ n ∑ r/∈Si cr 〈∇vrL,xi〉1 (k,t) ir (32a)
= − η nN N∑ j=1 ( f (k,t)(xj)− yj ) x>j xi ∑ r/∈Si c2r 1 (k,t) ir 1 (k,t) jr (32b)
= − η N N∑ j=1 ( f (k,t)(xj)− yj )( (H(k,t))ij − (H⊥(k,t))ij ) , (32c)
where (H⊥(k,t))ij is defined as
(H⊥(k,t))ij , 1
n x>i xj n∑ r∈Si 1 (k,t) ir 1 (k,t) jr . (33)
For the inner product term 〈 dI, f (k,t)(X)− y 〉 , we have〈
dI, f (k,t)(X)− y 〉
= − η N (f (k,t)(X)− y)>(H(k,t) −H⊥(k,t))(f (k,t)(X)− y). (34)
Let T1 and T2 denote the following terms
T1 , −(f (k,t)(X)− y)>H(k,t)(f (k,t)(X)− y), (35a) T2 , (f (k,t)(X)− y)>H⊥(k,t)(f (k,t)(X)− y). (35b)
With probability at least 1− δ, T1 can be bounded as:
T1 = −(f (k,t)(X)− y)>(H(k,t) −H(0) + H(0) −H∞ + H∞)(f (k,t)(X)− y) (36a) 6 −(f (k,t)(X)− y)>(H(k,t) −H(0))(f (k,t)(X)− y)
− (f (k,t)(X)− y)>(H(0) −H∞)(f (k,t)(X)− y)− λ0 ∥∥f (k,t)(X)− y∥∥2
2 (36b)
¬ 6
( 2 √
2NR√ πδα +N
√ ln (2N2/δ)
2n − λ0
)∥∥f (k,t)(X)− y∥∥2 2 , (36c)
where ¬ comes from Lemma 2 and Lemma 3. To bound the term T2, consider the `2 norm of the matrix H⊥(k,t). With probability at least 1− δ, we have:
‖H⊥(k,t)‖2 6 ‖H⊥(k,t)‖F (37a)
=  N∑ i=1 N∑ j=1 ( 1 n ∑ r∈Si x>i xj1 (k,t) ir 1 (k,t) jr )2 12 (37b) 6 N
n |Si|
¬ 6
√ 2
π
NR
δα , (37c)
where ¬ comes from Lemma 1. Therefore, with probability at least 1− δ, we have
T2 6
√ 2
π
NR
δα
∥∥f (k,t)(X)− y∥∥2 2 . (38)
Combine the results of (36c) and (38):〈 dI, f (k,t)(X)− y 〉 6 η ( 3 √
2R√ πδα +
√ ln (2N2/δ)
2n − λ0 N )∥∥f (k,t)(X)− y∥∥2 2 . (39)
For the inner product term 〈 dII, f (k,t)(X)− y 〉 , we first bound ‖dII‖22 as follows:
‖dII‖22 = N∑ i=1 ( 1√ n ∑ r∈Si [ crσ ( (v(k,t+1)r ) >xi ) − crσ ( (v(k,t)r ) >xi )])2 (40a)
¬ 6 η2
n N∑ i=1 |Si| ∑ r∈Si (cr〈∇vrL,xi〉)2 (40b)
 6 η2
n N∑ i=1 |Si| ∑ r∈Si ‖∇vrL‖22 ‖xi‖22 (40c)
6 η2N
n |Si|2 max r∈[n] ∥∥∇vrL∥∥22 (40d) 6 η2|Si|2
n2 ∥∥f (k,t)(X)− y∥∥2 2 , (40e)
where ¬ comes from the Lipschitz continuity of the ReLU function σ(·),  holds due to Cauchy– Schwartz inequality. Plug (18) into (40e), we have with probability at least 1− δ:
‖dII‖22 6 2η2R2 πδ2α2 ∥∥f (k,t)(X)− y∥∥2 2 . (41)
The inner product term 〈 dII, f (k,t)(X)− y 〉 can be bounded as〈
dII, f (k,t)(X)− y 〉 6
√ 2ηR√ πδα ∥∥f (k,t)(X)− y∥∥2 2 . (42)
Finally, the bound for the difference term is derived as∥∥f (k,t+1)(X)− f (k,t)(X)∥∥2 2 6 N∑ i=1 ( η√ n n∑ r=1 cr〈∇vrL,xi〉 )2 6 η2 ∥∥f (k,t)(X)− y∥∥2 2 . (43)
Combine the results of (39), (42) and (43):∥∥f (k,t+1)(X)−y∥∥2 2 6 [ 1 + 8 √ 2ηR√ πδα + 2η √ ln (2N2/δ) 2n − 2ηλ0 N + η2 ]∥∥f (k,t)(X)−y∥∥2 2 . (44)
Let R = O ( δαλ0 N ) , n = Ω ( N2 λ20 ln N 2 δ ) , and η = O(λ0N ), we have∥∥f (k,t+1)(X)− y∥∥2
2 6
( 1− ηλ0
2N
)∥∥f (k,t)(X)− y∥∥2 2 . (45)
Summing up over the selected number t(k) iterations completes the proof.

heading: C PROOF OF THEOREM 2
text: Theorem 2 For FedAvg under Assumptions 1 to 3, let the learning rate η = O (
λ0 τN |Ck|
) and the
neural network width n = Ω ( N2
λ20 ln 2N
2
δ
) , then with probability at least 1 − δ, the one-round loss
decay of FedAvg is ∥∥f (k+1)(X)− y∥∥2 2 6 ( 1− ητλ0
2N |Ck|
)∥∥f (k)(X)− y∥∥2 2 . (46)
Proof. We first construct a different set of kernel matrices {Λ(k),Λ(k,τ)m } similar to Huang et al. (2021). Let 1(k,u)imr , 1{〈v (k,u) m,r ,xi〉 > 0}, the (i, j)th entry of Λ(k,u)m and Λ(k,u) is defined as
(Λ(k,u)m )ij , 1
n x>i xj n∑ r=1 1 (k,0) imr 1 (k,u) jmr , (47a)
(Λ(k,u))ij , (Λ (k,u) m )ij , if (xj , yj) ∈ Dm. (47b)
Taking the difference between successive terms yields
f (k+1)(xi)− f (k)(xi) = 1√ n n∑ r=1 [ crσ ( (v(k+1)r ) >xi ) − crσ ( (v(k)r ) >xi )] . (48)
We decompose the difference term to the sum of dIi and d II i , based on the set Si and its complement:
dIi , 1√ n ∑ r/∈Si [ crσ ( (v(k+1)r ) >xi ) − crσ ( (v(k)r ) >xi )] , (49a)
dIIi , 1√ n ∑ r∈Si [ crσ ( (v(k+1)r ) >xi ) − crσ ( (v(k)r ) >xi )] . (49b)
Consider the residual term∥∥f (k+1)(X)− y∥∥2 2
(50a) = ∥∥f (k+1)(X)− f (k)(X) + f (k)(X)− y∥∥2
2 (50b) = ∥∥f (k)(X)− y∥∥2
2 + 2
〈 dI + dII, f (k)(X)− y 〉 + ∥∥f (k+1)(X)− f (k)(X)∥∥2
2 . (50c) We will give upper bounds for the inner product terms 〈 dI, f (k)(X)− y 〉 , 〈 dII, f (k)(X)− y 〉 , and
the difference term ∥∥f (k+1)(X) − f (k)(X)∥∥2
2 , separately. For an input x ∈ Rd1 , let f (k,u)m (x) ,
1√ n n∑ r=1
crσ(〈v(k,u)m,r ),x〉). By the update rule of FedAvg, the relation between the weight vector v(k)r in successive communication rounds is:
v(k+1)r = v (k) r −
η |Ck| ∑ m∈Ck τ−1∑ u=0 ∇L v (k,u) r
(51a)
= v(k)r − ηcr
N √ n|Ck| ∑ m∈Ck τ−1∑ u=0 ∑ j∈Im (f (k,u)m (xj)− yj)xj1 (k,u) jmr . (51b)
Based on the property of the set Si, we have
dIi = − 1√ n ∑ m∈Ck τ−1∑ u=0 ∑ r/∈Si cr 〈 v(k+1)r − v(k)r ,xi 〉 1 (k) ir (52a)
= − η Nn|Ck| ∑ m∈Ck τ−1∑ u=0 ∑ r/∈Si ∑ j∈Im (f (k,u)m (xj)− yj)x>i xj1 (k) ir 1 (k,u) jmr (52b)
= − η N |Ck| ∑ m∈Ck τ−1∑ u=0 ∑ j∈Im (f (k,u)m (xj)− yj) [ (Λ(k,u)m )ij − (Λ⊥(k,u)m )ij ] . (52c)
For the inner product term 〈 dI, f (k)(X)− y 〉 , we have
〈 dI, f (k)(X)− y 〉 = − η
N |Ck| τ−1∑ u=0 (f (k)(X)− y)>(Λ(k,u) −Λ⊥(k,u))(f (k,u)m (X)− y). (53)
Let T1 and T2 denote the following terms
T1 , −(f (k)(X)− y)>Λ(k,u)(f (k,u)g (X)− y), (54a)
T2 , (f (k)(X)− y)>Λ⊥(k,u)(f (k,u)g (X)− y), (54b)
where f (k,u)g (X) , [f (k,u) 1 (X1) >, · · · , f (k,u)|Ck| (X|Ck|) >]>. We are going to bound T1 and T2 separately. T1 can be written as:
T1 = −(f (k)(X)− y)>(Λ(k,u) −H(0) + H(0) −H∞ + H∞)(f (k,u)g (X)− y) (55a)
= −(f (k)(X)− y)>(Λ(k,u) −H(0))(f (k,u)g (X)− y)
− (f (k)(X)− y)>(H(0) −H∞)(f (k,u)g (X)− y)
− (f (k)(X)− y)>H∞(f (k)(X)− y) − (f (k)(X)− y)>H∞(f (k,u)g (X)− f (k)(X)). (55b)
First, we bound the norm of f (k,u)g (X)− y. It can be shown that
‖f (k,u)m (Xm)− ym‖2 = ‖f (k,u)m (Xm)− f (k,u−1)m (Xm) + f (k,u−1)m (Xm)− ym‖2 (56a) 6 ‖f (k,u)m (Xm)− f (k,u−1)m (Xm)‖2 + ‖f (k,u−1)m (Xm)− ym‖2 (56b) ¬ 6 (1 + η)‖f (k,u−1)m (Xm)− ym‖2, (56c)
where ¬ holds based on the derivation of (43). Applying (56c) recursively yields
‖f (k,u)m (Xm)− ym‖2 6 (1 + η)u‖f (k)(Xm)− ym‖2. (57)
The bound for ‖f (k,u)g (X)− y‖22 can thus be derived as
‖f (k,u)g (X)− y‖22 = N∑ i=1 [ f (k,u)g (xi)− yi ]2 (58a)
= ∑ m∈Ck ∥∥f (k,u)m (Xm)− ym∥∥22 (58b) 6 (1 + η)2u ∥∥f (k)(X)− y∥∥2 2 . (58c)
Second, following the steps in Lemma 2, it can be shown that with probability at least 1− δ,
‖Λ(k,t) −H(0)‖2 6 2 √
2NR√ πδα . (59)
We also bound the difference between f (k,u)g (X) and f (k)(X) as follows:
‖f (k,u)g (X)− f (k)(X)‖2 ¬ 6 u∑ v=1 ‖f (k,v)g (X)− f (k,v−1)g (X)‖2 (60a)
 6 u∑ v=1 η‖f (k,v−1)g (X)− y‖2 (60b)
® 6 u∑ v=1 η(1 + η)v−1‖f (k)(X)− y‖2 (60c)
= [(1 + η)u − 1] ‖f (k)(X)− y‖2, (60d)
where ¬ holds due to triangle inequality,  comes from (43), ® comes from (58c). Plugging the results from (58c), (59), and (60d) into (55b), we have with probability at least 1− δ,
T1 6
[ (1 + η)u ( 2 √
2NR√ πδα +N
√ ln (2N2/δ)
2n + κλ0
) − (1 + κ)λ0 ] ‖f (k)(X)− y‖22, (61)
where κ is the condition number of the matrix H∞. Next, consider the bound for T2. The `2 norm of Λ⊥(k,u) can be bounded as
‖Λ⊥(k,u)‖2 6 ‖Λ⊥(k,u)‖F (62a)
=  N∑ i=1 ∑ m∈Ck ∑ j∈Im ( 1 n ∑ r∈Si x>i xj1 (k) ir 1 (k,u) jmr )2 12 (62b) 6 N
n |Si|
¬ 6
√ 2
π
NR
δα , (62c)
where ¬ comes from Lemma 1. Therefore, we have with probability at least 1− δ,
T2 6 (1 + η) u
√ 2
π
NR
δα
∥∥f (k)(X)− y∥∥2 2 . (63)
Combine the results of (61) and (63):〈 dI, f (k)(X)− y 〉 6 τ
|Ck|
[( η +
(τ − 1) 2 η2 + o(η2) )3√2R√ πδα + √ ln ( 2N2 δ ) 2n + κλ0 N  − (1 + κ)ηλ0
N
]∥∥f (k)(X)− y∥∥2 2 .
(64)
For the inner product term 〈 dII, f (k)(X)− y 〉 , we first bound ‖dII‖22 with probability at least 1− δ:
‖dII‖22 = N∑ i=1 ( 1√ n ∑ r∈Si [ crσ ( (v(k+1)r ) >xi ) − crσ ( (v(k)r ) >xi )])2 (65a)
6 1
n N∑ i=1 |Si| ∑ r∈Si ( cr〈v(k+1)r − v(k)r ,xi〉 )2 (65b)
6 1
n N∑ i=1 |Si| ∑ r∈Si  ηcr N √ n|Ck| ∑ m∈Ck τ−1∑ u=0 ∑ j∈Im (f (k,u)m (xj)− yj)1 (k,u) jmr 2 (65c) 6 η2
N2n2|Ck|2 N∑ i=1 |Si| ∑ r∈Si ∑ m∈Ck τ−1∑ u=0 ∑ j∈Im ∣∣∣f (k,u)m (xj)− yj∣∣∣ 2 (65d)
6 η2 N2n2|Ck|2 N∑ i=1 |Si| ∑ r∈Si ( ∑ m∈Ck τ−1∑ u=0 |Im| ∥∥∥f (k,u)m (Xm)− ym∥∥∥ 2 )2 (65e)
¬ 6
η2 N2n2|Ck|2 N∑ i=1 |Si| ∑ r∈Si ( ∑ m∈Ck τ−1∑ u=0 (1 + η)u|Im| ∥∥∥f (k)(Xm)− ym∥∥∥ 2 )2 (65f)
 6
1 N2n2|Ck|2 N∑ i=1 |Si| ∑ r∈Si ( ∑ m∈Ck ((1 + η)τ − 1) |Im| ∥∥∥f (k)(Xm)− ym∥∥∥ 1 )2 (65g)
® 6
1 Nn2|Ck|2 N∑ i=1 |Si| ∑ r∈Si ( ((1 + η)τ − 1) ∥∥∥f (k)(X)− y∥∥∥ 2 )2 (65h)
¯ 6
2R2
πδ2α2|Ck|2
( τη +
τ(τ − 1) 2 η2 + o(η2) )2 ∥∥f (k)(X)− y∥∥2 2 . (65i)
where ¬ comes from (57),  holds due to ‖a‖1 6 ‖a‖2, ® holds due to ‖a‖1 6 √
dim(a)‖a‖2, ¯ is from Lemma 1. With probability at least 1− δ, the inner product term can thus be bounded as〈
dII, f (k)(X)− y 〉 6 √ 2τR√
πδα|Ck|
( η +
(τ − 1) 2 η2 + o(η2) )∥∥f (k)(X)− y∥∥2 2 . (66)
The bound for the difference term is derived as
∥∥f (k+1)(X)− f (k)(X)∥∥2 2 6 N∑ i=1 ( η√ n n∑ r=1 cr〈v(k+1)r − v(k)r ,xi〉 )2 (67a)
6 1
|Ck|2
( τη +
τ(τ − 1) 2 η2 + o(η2) )2 ∥∥f (k)(X)− y∥∥2 2 . (67b)
Combine the results of (64), (66) and (67b):
∥∥f (k+1)(X)− y∥∥2 2 6 { 1 + 2ητ
|Ck| [4√2R√ πδα + √ ln ( 2N2 δ ) 2n + κλ0 N  − (1 + κ)λ0
N
] + η2τ2
|Ck|2 + o(η2)
}∥∥f (k)(X)− y∥∥2 2 .
(68)
Let R = O ( δαλ0 N ) , n = Ω ( N2 λ20 ln N 2 δ ) , and η = O ( λ0 τN |Ck| ) , we have
∥∥f (k+1)(X)− y∥∥2 2 6 ( 1− ητλ0
2N |Ck|
)∥∥f (k+1)(X)− y∥∥2 2 . (69a)

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

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. 

This paper proposes a federated learning (FL) paradigm empowered by the neural tangent kernel (NTK) framework. The NTK under the infinite-width regime allows us to analyze a dynamic of the corresponding neural network without a gradient descent algorithm. The authors utilize it for predicting the best parameters aggregated from multiple workers. However, this framework requires transmitting Jacobian matrices between workers and the aggregation server, which results in increasing computational overhead and exposing more private information. To this end, they adopt dimensionality reduction via randomness-sharing projection, zeroing out compression as well as shuffling. They also address that their NTK-based FL scheme has a faster convergence rate compared to that of FedAvg, for a two-layer network under specific assumptions. Finally, empirical results support that the proposed FL method performs better than other FL algorithms and achieves similar test accuracy to the ideal centralized case.

This paper introduces new algorithms for the federated learning (FL) farmwork using neural tangent kernel (NTK) paradigm. Two algorithms are referred to as NTK-FL and CP-NTK-FL, where the latter is a variant of the former for improved communication efficiency and privacy preserving. The proposed algorithms are aimed to address statistical heterogeneity across the workers. An important point is that unlike typical training algorithms for the FL setting, here the workers upload the labels and sample Jacobian matrices (representing NTK) to the server. The server then uses tools from NTK to obtain the trained neural network (instead of using gradient descent). Theoretical and empirical validations (on MNIST, Fashion-MNIST and EMNIST), and comparisons with FedAvg are provided.