N EURAL T ANGENT K ERNEL E MPOWERED F EDERATED L EARNING
Anonymous authors
Paper under double-blind review
A BSTRACT
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.
Speciﬁcally, 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 efﬁciency 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.
I NTRODUCTION
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.
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 inﬂuence 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 ﬁxed 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 ﬁrst 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 c ommunication-efﬁcient and p rivacy-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.
R ELATED W ORK
Neural Tangent Kernel.
Jacot et al.
(2018) showed that training an inﬁnitely 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 ﬁrst-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 efﬁciency, privacy guarantees, and robustness to heterogeneity.
To reduce communication cost, gradient quantization and sparsiﬁcation 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 efﬁciency 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 justiﬁcation 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.
B ACKGROUND AND P RELIMINARIES
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 a j ∈ R M , j ∈ { 1 , . . .
, N } , A =
[ a 1 , . .
. , a N ] is an M × N matrix.
A third-order tensor A ∈ R K × 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: A i :: is a matrix of the i th horizontal slice, and A : j : is its j th lateral slice (Kolda & Bader, 2009).
Finally, the indicator function of an event is denoted by 1 ( · ) .
3. F EDERATED L EARNING M ODEL
Consider an FL architecture where a server trains a global model by indirectly using datasets distributed among M workers.
The local dataset of the m th worker is denoted by D m = { ( x m,i , y m,i ) }
N m i =1 , where ( x m,i , y m,i ) is an input-output pair, drawn from a distribution P m .
The local objective can be formulated as an empirical risk minimization over N m training examples:
F m ( w ) =
1 N m (cid:80) N m i =1 R ( w ; x m,j , y m,i ) , where R is a sample-wise risk function quantifying the error of model with a weight vector w ∈ R d estimating the label y m,i for an input x m,i .
The global objective function is denoted by F ( w ) , and the optimization problem may be formulated as:
[Equation]
3. L INEARIZED N EURAL N ETWORK M ODEL
Let ( x i , y i ) denote a training pair, with x i ∈ R
d 1 and y i ∈ R d 2 , where d 1 is the input dimension and d 2 is the output dimension.
X (cid:44) [ x 1 , . . .
, x N ] (cid:62) represents the input matrix and Y (cid:44)
[ y 1 , . . .
, y N ] (cid:62) represents the label matrix.
Consider a neural network function f : R d 1 → R d 2 parameterized by a vector w ∈ R d , which is the vectorization of all weights for the multilayer network.
Given an input x i , the network outputs a prediction ˆ y i =
f ( w ; x i ) .
Let (cid:96) (ˆ y i , y i ) be the loss function measuring the dissimilarity between the predicted result ˆ y i and the true label y i .
We are interested in ﬁnding an optimal weight vector w (cid:63) that minimizes the empirical loss over N training examples:
[Equation]
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 ) − η ∇ w L. To simplify the notation, let f ( t ) ( x ) be the output at time step t with an input x , i.e., f ( t ) ( x ) (cid:44) f
( w ( t ) ; x ) .
Following
Lee et al.
(2019), we use the ﬁrst-order Taylor expansion around the initial weight vector w (0) to approximate the neural network output given an input x , i.e.,
[Equation]
where J (0) ( x ) =
[ ∇ f (0)1 ( x ) , . .
, ∇ f (0) d 2 ( x )] (cid:62) , with ∇ f ( t ) j
( x ) (cid:44) [
∂ ˆ y ( t ) j /
[Equation]
( t ) 1 , . . . ,
∂ ˆ y ( t ) j /∂w ( t ) d ] (cid:62) being the gradient of the j th component of the neural network output with respect to w ( t ) .
Consider the halved mean-squared error (MSE) loss (cid:96) , namely, (cid:96) = 1 d 2 (cid:80) d 2 j =1 12 (ˆ y j − y j ) 2 .
Based on the continuous-time limit, one can show that the dynamics of the gradient ﬂow are governed by the following differential equation:
[Equation]
where f ( t ) ( X ) ∈ R N ×
d 2 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 = 1 d 2 (cid:10) J (0) ( x i ) , J (0) ( x j ) (cid:11) F .
The differential equation (4) has the closed-form solution:
[Equation]
The neural network state f ( t ) ( X ) can thus be directly obtained from (5) without running the gradient descent algorithm.
P ROPOSED FL P ARADIGM VIA THE NTK F RAMEWORK
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 efﬁciency and enhanced privacy.
The detailed algorithm descriptions are presented as follows.
[Equation]
[Figure]
NTK-FL follows four steps to update the global model in each round.
First , the server will select a subset C k of workers and broadcast to them a model weight vector w ( k ) from the k th 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
D m to compute a Jacobian tensor J ( k ) m
[Equation]
[Figure]
× d , ∀ m ∈ C k , which is a concatenation of N m sample-wise Jacobian matrices ( J ( k ) m ) i :: = [ ∇ f ( k ) 1 ( x m,i ) , . .
. , ∇ f ( k ) d 2 ( x m,i )] (cid:62) , i ∈ { 1 , . .
N m } .
The worker will then upload the Jacobian tensor J ( k ) m , labels Y m , and initial condition f ( k ) ( X m ) 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 ) ∈
R N × d 2 × d based on received J ( k ) m ’s, with each worker contributing N m 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 )and J ( k ) 2 , respectively.
The global Jacobian tensor is constructed as:
[Equation]
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 d 2 (cid:104) J ( k ) i :: , J ( k ) j :: (cid:105) 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:
[Equation]
where J ( k ) :
j : is the j th lateral slice of J ( k ) , and R ( k,t ) :
j is the j th column of the residual matrix R ( k,t ) deﬁned as follows:
[Equation]
Note that X ( k ) =
[ X (cid:62) 1 , . . .
, X (cid:62)C k ] (cid:62) denotes a concatenation of worker training examples, and
Y ( k ) =
[ Y (cid:62) 1 , . . .
, Y (cid:62)C k ] (cid:62) 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 { t 1 , . .
. , 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:
[Equation]
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 { t 1 , . . .
, t Ψ } can be completed in O (Ψ) time.
Comparison of NTK-FL and Huang et al.
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 i th training example of the m th worker for NTK-FL is J m =
[ ∇ f 1 ( x m,i ) , . .
. , ∇ f d 2 ( x m,i )] (cid:62) , whereas the gradient update being sent for FedAvg is ∇ L ( w ; x m,i , y m,i ) = 1 d 2 (cid:80) d 2 j =1 (ˆ y m,i,j − y m,i,j ) ∇ f j ( x m,i ) , a weighted sum of coordinates of J m .
By sending Jacobian matrices J m 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.
4. CP-NTK-FL VARIANT
Compared to FedAvg, NTK-FL does not incur additional client computational overhead since calculating the Jacobian tensor enjoys the same communication efﬁciency 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 efﬁciency and security in terms of data conﬁdentiality 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 ( N m d 2 d ) 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.
[Figure]
E ( k
[Figure]
) (cid:173)
[Figure]
w ( k ) & decrypts
[Figure]
get the seed ρ w ( k ) ,ρ B m ⊂D m (cid:174)
[Figure]
Z (cid:48) m
[Figure]
C ( J (cid:48) ( k ) m ) , f ( k ) m
( Z (cid:48) m ) , Y m
[Figure]
[Figure]
Jacobian Dimension Reduction.
First, we let the m th worker sample a subset B m from its dataset
D m uniformly for the training.
Let β ∈
(0 , 1) denote the sampling rate, B
m contains N (cid:48) m =
βN m data points, with the training pairs denoted by ( X (cid:48) m , Y (cid:48) 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 (cid:48) m , i.e., Z (cid:48) m
=
X (cid:48)
m P , where P ∈ R d 1 × d (cid:48) 1 is a projection matrix generated based on a seed ρ with IID standard Gaussian entries.
In general, we have d (cid:48) 1 < dand 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 ﬁrst examine the current Jacobian tensor J (cid:48)
( k ) m ∈ R N (cid:48) m
× d 2 × d (cid:48) .
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 (cid:48) m .
We empirically verify their impact on the test accuracy in Figure 3.
We set d (cid:48) 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 Shufﬂing.
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 efﬁciency, this compression scheme is empirically effective against the data reconstruction attack (Zhu et al., 2019).
To further ensure the conﬁdentiality and membership privacy, we introduce a shufﬂing 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 Y m ’s.
Based on (7b), we denote the model update by ∆ w ( k ) (cid:44) w ( k +1) − w ( k )
= (cid:80) d 2 j =1 ( J ( k ) :
j : ) (cid:62) 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.
A NALYSIS OF A LGORITHM
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.
and Dukler et al.
(2020), we consider a two-layer neural network f : R d → R of the following form to facilitate our analysis:
[Equation]
where x ∈ R d 1 is an input, v r ∈ R d 1 is the weight vector in the ﬁrst layer, V =
[ v 1 , · · · , v n ] , c r ∈ R is the weight in the second layer, and σ ( · ) is the rectiﬁed linear unit (ReLU) function, namely σ ( z ) =
max( z, 0) , applied coordinatewise.
We state two assumptions as prerequisites.
Assumption 1 The ﬁrst layer v r ’s are sampled from N (0 , α 2 I ) .
The second layer c r ’s are sampled from {− 1 , 1 } with equal probability and are kept ﬁxed during training.
Assumption 1 gives the initial distribution of the neural network parameters.
Similar assumptions can be found in Dukler et al.
We add restrictions to the input data in the next assumption.
Assumption 2 (Normalized input).
The input data are normalized, i.e., (cid:107) x i (cid:107) 2 (cid:54) 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 = 1 n x (cid:62) i x j (cid:80) nr =1 1 ( k ) ir 1 ( k ) jr , where 1 ( k ) ir (cid:44) 1 {(cid:104) v ( k ) r , x i (cid:105) (cid:62) 0 } , and the term c 2 r is omitted according to Assumption 1.
Deﬁne H ∞ , whose ( i, j ) th entry is given by:
[Equation]
Let λ 0 denote the minimum eigenvalue of H ∞ , which is restricted in the next assumption.
Assumption
The kernel matrix H ∞ is positive deﬁnite, namely, λ 0 > 0 .
In fact, the positive-deﬁnite 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 deﬁniteness of H ∞ in Assumption 3.
Next, we study the residual term (cid:107) f ( k ) ( X ) − y (cid:107) 22 in communication round k , where
X =
[ X (cid:62) 1 , . . .
, X (cid:62) M ] (cid:62)
∈ R N × d 1 denote a concatenation of client inputs and y =
[ y (cid:62) 1 , . . .
, y (cid:62) M ] (cid:62) ∈
R N 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 (cid:0) λ 0 N (cid:1) and the neural network width n = Ω (cid:16) N 2 λ 20 ln N 2 δ (cid:17) , then with probability at least 1 − δ , the one-round loss decay of NTK-FL is
[Equation]
(12) where t ( k ) is the number of NTK update steps deﬁned 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 (cid:16) λ 0 τN
|C k | (cid:17) and the neural network width n = Ω (cid:16) N 2 λ 20 ln N 2 δ (cid:17) , then with probability at least 1 − δ , the one-round loss decay of FedAvg is
[Equation]
where τ is the number of local iterations, and |C k | 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 coefﬁcient in Theorem 1 with the decay coefﬁcient in Theorem 2, i.e., 1 − η 1 t ( k ) λ 0 2 N + O (cid:0) η 21 (cid:1) < 1 − η 2 τλ 0 2 N |C k | , where η (cid:28) 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 10 2 to 10 3 , whereas τ is often on the order of magnitude of 10 in literature (Reisizadeh et al., 2020;
Haddadpour et al., 2021).
One can verify that η 1 t ( k ) λ 0 is larger than η 2 τλ 0 / |C k | and draw the conclusion in (1) .
[Figure]
E XPERIMENTAL R ESULTS
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 =categories.
For MNIST and Fashion-MNIST, we follow Hsu et al.
(2019) to simulate non-IID data with the symmetric Dirichlet distribution (Good, 1976).
Speciﬁcally, for the m th worker, we draw a random vector q m ∼ Dir ( α ) , where q m = [ q m, 1 , . . .
, q m,C ] (cid:62) belongs to the ( C − 1) -standard simplex.
Images with category k are assigned to the m th worker in proportional to (100 · q m,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 C k 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-
(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 ﬁve 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 (cid:48) 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 (cid:48) 1 will cause less information loss,
[Figure] [Figure]
and are expected to achieve a higher test accuracy scheme is robust to different combinations
CP-NTK-FL can achieve the target goal within the fewest communication rounds without incurring communication cost signiﬁcantly.
optimization algorithms comm.
rounds comm.
cost (MB)
CP-NTK-FL
FedCOM
QSGD (4 bit)
FedAvg
Uplink Communication.
We evaluate the uplink communication efﬁciency of CP-NTKFL ( d (cid:48) 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.
C ONCLUSION AND F UTURE W ORK
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 ﬂavor by transmitting more expressive updates.
The effectiveness of the proposed paradigm has been veriﬁed 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 efﬁciency 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
A A DDITIONAL E XPERIMENTAL R ESULTS
[Figure]
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 ﬁxed 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 CIFARdataset (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.
B P ROOF OF T HEOREM
Let I m denote a set of indices such that for i ∈
I m , ( x i , y i ) ∈
D m .
We ﬁrst present some lemmas to facilitate the convergence analysis.
In communication round k , deﬁne S ( k ) i as the set of indices corresponding to neurons whose activation pattern is similar to its initial state for an input x i :
S ( k ) i (cid:44) (cid:8) r ∈ { 1 , . . . , n } (cid:12)(cid:12) ∃ v , (cid:107) v − v (0) r
(cid:107) 2 (cid:54)
[Equation]
We upper bound the cardinality of S i in Lemma 1.
Lemma 1 Under Assumption 1 to 2, with probability at least 1 − δ , we have
[Equation]
Proof.
To bound |S i | = n (cid:80) r = 1 ( r ∈ S i ) , consider an event A ir deﬁned as follows:
[Equation]
Clearly, 1 ( r ∈ S i ) =
1 ( A ir ) .
According to Assumption 2, (cid:107) x i (cid:107) (cid:54) 1 , it can be shown that the event
A ir happens if and only if | ( v (0) r ) (cid:62) x i |
(cid:54) R based on a geometric argument.
Based on Assumption 1, we have ( v (0) r ) (cid:62) x i ∼ N (0 , α 2 ) .
The probability of event
A ir is
[Equation]
17b
[Equation]
By Markov’s inequality, we have with probability at least 1 − δ ,
[Equation]
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 } , (cid:107) v ( k,t
[Equation]
r
[Equation]
, then
[Equation]
Proof.
We have
[Equation]
Consider the event A ir deﬁned in (16).
Let φ ( k,t ) ijr
[Equation]
−
[Equation] [Equation]
jr .
If ¬ A ir and ¬ A jr happen, clearly we have | φ ( k,t ) ijr | = 0 .
Therefore, the expectation of | φ ( k,t ) ijr | can be bounded as
[Equation]
φ
[Equation]
where (cid:172) comes from (17b).
By Markov’s inequality, we have with probability at least 1 − δ , | φ ( k,t ) ijr | (cid:54) 2 (cid:114)
[Equation]
Plugging (22) into (20c) yields
[Equation]
Taking the square root on both sides completes the proof.
Lemma 3 With probability at least 1 − δ ,
[Equation]
Proof.
We have
[Equation]
Note that ( H (0) ) ij = 1 n x (cid:62) i x j n (cid:80) r = 1 (0)
ir 1 (0) jr , ( H (0) ) ij ∈ [ − 1 , 1] .
By Hoeffding’s inequality, we have with probability at least 1 − δ/n 2 ,
[Equation]
Applying the union bound over i, j ∈ [ N ] yields
[Equation]
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 (cid:0) λ 0 N (cid:1) and the neural network width n = Ω (cid:16) N 2 λ 20 ln 2 N 2 δ (cid:17) , then with probability at least 1 − δ , the one-round loss decay of NTK-FL is
[Equation]
Proof.
Taking the difference between successive terms yields
[Equation]
We decompose the difference term to the sum of d
I i and d II i , based on the set S i :
[Equation]
Consider the residual term
[Equation]
31a
[Equation]
31b
[Equation]
31c
[Equation]
We will give upper bounds for the inner product terms (cid:10) d
I , f ( k,t ) ( X ) − y (cid:11) , (cid:10) d II , f ( k,t ) ( X ) − y (cid:11) , and the difference term (cid:13)(cid:13)
f
( k,t +1)
( X ) − f ( k,t ) ( X ) (cid:13)(cid:13) 2 2 , separately.
Based on the property of the set
S i , we have
[Equation]
where ( H ⊥ ( k,t ) ) ij is deﬁned as
[Equation]
For the inner product term (cid:10) d I , f ( k,t ) ( X ) − y (cid:11) , we have
[Equation]
Let T 1 and T 2 denote the following terms
[Equation]
(35b)
With probability at least 1 − δ
[Equation]
T 1 can be bounded as:
[Equation]
where (cid:172) comes from Lemma 2 and Lemma 3.
To bound the term T 2 , consider the (cid:96) 2 norm of the matrix
H ⊥
( k,t ) .
With probability at least 1 − δ , we have:
[Equation]
where (cid:172) comes from Lemma 1.
Therefore, with probability at least 1 − δ , we have
[Equation]
Combine the results of (36c) and (38):
[Equation]
For the inner product term (cid:10) d II , f ( k,t ) ( X ) − y (cid:11) , we ﬁrst bound (cid:107) d II (cid:107)
[Equation]
follows:
[Equation]
where (cid:172) comes from the Lipschitz continuity of the ReLU function σ ( · ) , (cid:173) holds due to Cauchy–
Schwartz inequality.
Plug (18) into (40e), we have with probability at least 1 − δ :
[Equation]
The inner product term (cid:10) d II , f ( k,t ) ( X ) − y (cid:11) can be bounded as
[Equation] [Equation]
Finally, the bound for the difference term is derived as
[Equation]
Combine the results of (39), (42) and (43):
[Equation]
Let R
=
O (cid:0) δαλ 0 N (cid:1) , n = Ω (cid:16) N 2 λ 20 ln N 2 δ (cid:17)
[Equation]
and η = O ( λ 0 N ) , we have
[Equation] [Equation]
Summing up over the selected number t ( k ) iterations completes the proof.
C P ROOF OF T HEOREM
Theorem 2 For FedAvg under Assumptions 1 to 3, let the learning rate
η
= O (cid:16) λ 0 τN
|C k | (cid:17) and the neural network width n = Ω (cid:16) N 2 λ 20 ln 2 N 2 δ (cid:17) , then with probability at least 1 − δ , the one-round loss decay of FedAvg is
[Equation]
Proof.
We ﬁrst construct a different set of kernel matrices { Λ ( k ) , Λ ( k,τ ) m } similar to Huang et al.
Let 1 ( k,u ) imr (cid:44) 1 {(cid:104) v ( k,u ) m,r , x i (cid:105) (cid:62) 0 } , the ( i, j ) th entry of Λ ( k,u ) m and Λ ( k,u ) is deﬁned as
[Equation]
k
[Equation]
Taking the difference between successive terms yields
[Equation]
We decompose the difference term to the sum of d
I i and d II i , based on the set S i and its complement:
[Equation]
Consider the residual term
[Equation]
We will give upper bounds for the inner product terms (cid:10) d
I , f ( k ) ( X ) − y (cid:11) , (cid:10) d II , f ( k ) ( X ) − y (cid:11) , and the difference term (cid:13)(cid:13) f
( k +1)
( X ) − f ( k ) ( X ) (cid:13)(cid:13) 2 2 , separately.
For an input x ∈ R d 1 , let f ( k,u ) m ( x ) (cid:44) 1 √ n n (cid:80) r =1 c r σ ( (cid:104) v ( k,u ) m,r ) , x (cid:105) ) .
By the update rule of FedAvg, the relation between the weight vector v ( k ) r in successive communication rounds is:
[Equation]
51b
[Equation]
Based on the property of the set S i , we have
[Equation]
52b
[Equation]
For the inner product term (cid:10) d I , f ( k ) ( X ) − y (cid:11) , we have
[Equation]
Let T 1 and T 2 denote the following terms
[Equation]
where f ( k,u ) g ( X ) (cid:44)
[ f ( k,u )( X 1 ) (cid:62) , · · · , f ( k,u ) |C k | ( X |C k | ) (cid:62) ]
(cid:62) .
We are going to bound T 1 and T 2 separately.
T 1 can be written as:
[Equation]
First, we bound the norm of f ( k,u ) g
[Equation]
X
[Equation]
− y .
It can be shown that
[Equation]
where (cid:172) holds based on the derivation of (43).
Applying (56c) recursively yields
[Equation]
The bound for (cid:107) f
( k,u ) g
( X ) − y (cid:107) 22 can thus be derived as
[Equation]
Second, following the steps in Lemma 2, it can be shown that with probability at least 1 − δ ,
[Equation]
We also bound the difference between f ( k,u ) g ( X ) and f ( k ) ( X ) as follows:
[Equation]
where (cid:172) holds due to triangle inequality, (cid:173) comes from (43), (cid:174) comes from (58c).
Plugging the results from (58c), (59), and (60d) into (55b), we have with probability at least 1 − δ ,
T 1 (cid:54)
[Equation]
where κ is the condition number of the matrix H ∞ .
Next, consider the bound for T 2 .
The (cid:96) 2 norm of Λ ⊥ ( k,u ) can be bounded as
[Equation]
where (cid:172) comes from Lemma 1.
Therefore, we have with probability at least 1 − δ ,
[Equation]
[Equation]
Combine the results of (61) and (63):
[Equation]
For the inner product term (cid:10) d II , f ( k ) ( X ) − y (cid:11) , we ﬁrst bound (cid:107) d II (cid:107) 22 with probability at least 1 − δ :
[Equation]
=
[Equation]
65b
[Equation]
n
[Equation]
m ∈C k τ − 1 (cid:88) u
= (cid:88)
j ∈I m ( f ( k,u ) m ( x j ) − y j ) 1 ( k,u ) jmr 
 (65c) (cid:54) η
N 2 n 2 |C k |
N (cid:88) i =
|S i | (cid:88) r ∈S
i   (cid:88) m ∈C k τ − 1 (cid:88)
u
= (cid:88)
j ∈I m (cid:12)(cid:12) (cid:12) f
( k,u ) m ( x j ) − y j
(cid:12)(cid:12) (cid:12) 
 (65d) (cid:54) η
N 2 n 2 |C k |
N (cid:88) i =
|S i | (cid:88) r ∈S
i (cid:32) (cid:88) m ∈C k τ − 1 (cid:88) u = |I
m | (cid:13)(cid:13)(cid:13)
f ( k,u ) m ( X m ) − y m (cid:13)(cid:13)(cid:13)(cid:33)(65e) (cid:172) (cid:54) η
N 2 n 2 |C k |
N (cid:88) i =
|S i | (cid:88) r ∈S
i (cid:32) (cid:88) m ∈C k τ − 1 (cid:88) u
= (1 + η ) u |I
m | (cid:13)(cid:13)(cid:13) f
( k ) ( X m ) − y m (cid:13)(cid:13)(cid:13)(cid:33)(65f) (cid:173) (cid:54)
N 2 n 2 |C k |
N (cid:88) i =
|S i | (cid:88) r ∈S
i (cid:32) (cid:88) m
[Equation]
k ((
[Equation]
η )
[Equation]
1) |I
m | (cid:13)(cid:13)(cid:13) f ( k ) ( X m ) − y m
(cid:13)(cid:13)(cid:13)(cid:33)(65g) (cid:174) (cid:54)
Nn
[Equation]
k |
N (cid:88) i = |
[Equation]
i | (cid:88) r ∈S i (cid:16) ((1 + η ) τ − 1) (cid:13)(cid:13)(cid:13) f
( k ) ( X ) − y (cid:13)(cid:13)(cid:13) 2 (cid:17)(65h) (cid:175)
[Equation]
πδ
[Equation]
(cid:13)(cid:13)
[Equation]
(65i) where (cid:172) comes from (57), (cid:173) holds due to (cid:107) a (cid:107) 1 (cid:54) (cid:107) a (cid:107) 2 , (cid:174) holds due to (cid:107) a (cid:107) 1 (cid:54) (cid:112) dim( a ) (cid:107) a (cid:107) 2 , (cid:175) is from Lemma 1.
With probability at least 1 − δ , the inner product term can thus be bounded as
[Equation]
The bound for the difference term is derived as
[Equation]
67b
[Equation]
Combine the results of (64), (66) and (67b):
[Equation]
Let R
=
O (cid:0)
[Equation]
0 N (cid:1) , n = Ω (cid:16) N 2 λ 20 ln N 2 δ (cid:17) , and η = O (cid:16) λ 0 τN
|C k | (cid:17) , we have
[Equation]