Neural Tangent Kernel Empowered Federated Learning
Authors: Kai Yue
Richeng Jin
Ryan Pilgrim
Chau-Wai Wong
Dror Baron
Huaiyu Dai
Venue: ICLR 2022 Submitted (2022)
Keywords: Federated Learning
Neural Tangent Kernel

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. 

Under review as a conference paper at ICLR 2022
NEURAL TANGENT KERNEL EMPOWERED FEDERATED LEARNING
Anonymous authors Paper under double-blind review
ABSTRACT
Federated learning (FL) is a privacy-preseviing paradigm where multiple partic- ipants jointly solve a machine learning problem without sharing raw data. Un- like traditional distributed learning, a unique characteristic of FL is statistical het- erogeneity, 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 generaliza- tion analyses. In this paper, we propose 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 con- structs an empirical kernel matrix to update a global model without explicitly per- forming gradient descent. We further develop variant with improved communi- cation efficiency and enhanced privacy. Numerical results show that the proposed paradigm can achieve the same accuracy while reducing the number of communi- cation rounds by an order of magnitude compared to federated averaging.
INTRODUCTION
Federated learning (FL) has emerged as popular paradigm involving large number of workers col- laboratively solving a machine learning problem (Kairouz al. [202]). In a typical FL framework, server broadcasts global model to selected workers and collects model updates without needing to access the raw data. One popular algorithm known as federated averaging (FedAvg) (MeMa- han 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 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 lo- cally 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 chal- lengingly, the learning procedure 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 possi- ble solution to share a globally available dataset with participants to reduce the distance between client-data distributions and the population distribution (Zhao et al., (0188). In practice, though, such a dataset may unavailable or too small to meaningfully compensate for the heterogeneity. Some researchers replaced the coordinate- srrdinate-wiie wise weight averaging strategy in FedAvg with nonlinear ag- gregation schemes (Wang et al., 2020 Chen & Chao 2021). The nonlinear aggregation relies on separate optimization routine, which can be elusive, especially when the algorithm does not con- verge 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 al. 2021 HHddddppour et al.
2021 Chen & Chao 202T).

---

Under review as a conference paper at ICLR 2022
In this work, we present neural tangent kernel empowered federated learning (NTK-FL) paradigm. Given a fixed number of commumication rounds, NTK-FL outperforms state-of-the methods in terms of test accuracy. We summarize our contributions 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-llI data problem of FL. Compared with FedAvg, it is robust 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-efficiic and peraacy-peeservnng 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.
RELATED WORK
Neural Tangent Kernel. Jacot al. (2018) showed that training infinitely wide neural net- work with gradient descent in the parameter space is equivalent to kernel regression in the function space. [ee et al. (2099) used a first-order Taylor expansion to approximate the neural network output and derived the training dynamics in closed form. Chen et al. (2200) established the generalization bounds for two-layer peeeeaaraatteeeed neural network with the NTK framework. The NTK com- putation has been extended to convolutional neural networks (CNNs) (Arora et al. (2019), recurrent neural networks (RNNs) (Alemohammad al. 202T). 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 model with distributed workers without transmitting local data (McMahan et al., 2017 Kairouz et al., 202T). 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 hetero- geneity. To reduce communication cost, gradient quantization and sparsification were incorporated into FedAvg (Reisizadch 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, differen- tially private federated optimization and decentralized aggregation methods were developed (Girgis) al. 2021 [Ceng al. (021). 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 al., 2017). p. In this work, we focus on 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 a variant with improved communication efficiency and enhanced privacy.
Kernel Methods in Federated Learning. The NTK framework has been mostly used for con- vergence analyses in FL. Sco 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 the- oretical 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 con- vergence rate and generalization bound two-layer ReLU neural networks trained with FedAvg. Su et al. 2021) studied the convergence behavior of set of FL algorithms in the kernel regression set- ting. In comparison, our work does not focus on pure convergence analyses of existing algorithms. We propose novel FL framework by replacing the gradient descent with the NTK evolution.
2

---

Under review as a conference paper at ICLR 2022
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 a, R* j e {I. {l....,A N]. A [a....., aN] is M x N matrix. A third-order tensor A E RKXMXX can be viewed as concatenation of such matrices. We use slice to denote a matrix a third-order tensor by varying two indices (Kolda & Bader 2009). Take tensor A, for instance: A is a matrix of the äth horizontal slice, and A.j is its jth lateral slice (Kolda & Bader 2009). Finally, the indicator function of an event is denoted by 11 (.).
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 mtth worker is denoted by Dm where (Xm,i, xm,)) is input-output pair, drawn from distribution P... The local objective can be formulated as an empirical risk minimization over Nm training examples: (w)() N. Nm R(w; Xm.jj Y Mm,;), where R is sample-wise risk function quantifying the er- ror of model with a weight vector w E Rd estimating the label Ym.s for an input Xm,i- The global objective function is denoted by F(w), and the optimization problem may be formulated as:
((xm,in Y mm,.)) / i=1
M
min F(w)= weRd
Σ (w)
M
3.2 LINEARIZED NEURAL NETWORK MODEL
Let (x, Y:) denote a training pair, with x; E Rd: and y; E Rd',, where d is the input dimenssin and d2 is the output dimension. X 2 , xN] represents the input matrix and Y [y,,..., yN] represents the label matrix. Consider a neural network function f. Rd Rd parameterized by vector w E Rdd which is the vectorization of all weights for the multilayer network. Given an input x, the network outputs a prediction y; f(w: x,). Let E(}..y: y:) be the loss function measuring the dissimilarity between the predicted result y; and the true label y;. We are interested in finding an optimal weight vector w' that minimizes the empirical loss over N training examples:
N
w argmin L(w: X,
((......
(2)
N
One common optimization method is the gradient descent training. Given the learning rate 7), gra- dient descent updates the weights each time step w(++)) w(() V To simplify the
notation, let f (*) (x) be the output time step t with input x, i.e., (x) & f(w(0):x) Following Lee al. (2019), we use the first-order Taylor expansion around the initial weight vector w()) to approximate the neural network output given an input x, i.e.,
(x)(x) ~ f(0) (x) + ( (w(( - (0)),
33)
where (x) = [Vrl""
V d2 Tx)]_. with vI.h) (x) A [09) /w 00 /0
being the gradient of the jth component of the neural network output with respect to w Consider the halved mean-squared error (MSE) loss L, namely, é = d Σラe { )² Based the continuous-time limit, one show that the dynamics of the gradient flow are governed by the following differential equation:
df = --
(PP) -Y),
dt Nxda is matrix of concatenated output for all training examples, and ()) is the neural tangent kernel at time step 0, with each entry (H()), equal to the scaled Frobenius inner
where p()) (X)
(X)
product of the Jacobian matrices: (H()));; da (J() (()))) The differential equation
4 has the closed-form solution:
(()) (X)
(X) = (I-e (H))
-
Y + H (()) (X).
(5)
3

---

Under review as a conference paper at ICLR 2022
The neural network state f() can thus be directly obtained from (5) without running the gradient descent algorithm.
PROPOSED FL PARADIGM VIA THE NTK FRAMEWORK
In this section, we present the NTK-FL paradigm (Figure I and then extend it to the variant CP- NTK-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 update the global model in each round. First, the server will select subset Cg of workers and broad- cast to them a model weight vector ((k) from the kth round. Here, the superscript k is the communication round index, and it should be distinguished from the gradient de- scent time step t in Section 3.2 Second, each worker will use its local training data Dm compute a Jacobian tensor ( E RN,, xdy V m Ck. which is concate- nation of Nm sample-wise Jacobian matrices (J(H) ):: [vr(<) (xm,)) (xm,.).....
aggregation server
builds
kernel H & peforms weight evolution
server evaluates w (gdd)) & selects the best one
worker sends (( (X))(X.)) and Y.,
Ulice
1Dmm
server updates
w(lk)
& hroadcats
w (
vrla (xm.:))²
worker receives weight w(A)
EE Nm... The worker will then up- load the Jacobian tensor ((k), labels Ym. and initial condition p(k) (X... to the server. The transmitted information corresponds to the vari- ables in the state evolution of ((() in Third, the server will construct global Jacobian ten- sor (*) Nxdxxd based on received J (kk)s. m with each worker contributing Nn, horizontal
the mnt worker
Figure 1: first receives the weight ((k), and then uploads the Jacobian tensor J m label Ym. and initial condition p(k) (Xm). The server builds global kernel (A) and performs the weight evolution with (t.....,ty We use (9) to find the best t and update the weight accordingly.
Schematic of NTK-FL. Each worker
slices to (A),
We use toy example to explain the process as follows. Suppose the server selects worker and worker in a certain round. Workers and 2 will compute the Jacobian tensors [(A) and (k) respectively. The global Jacobian tensor is constructed as:
1. if E (1....,NI}, J
2,3 .J ------i N1, if (N1 1,.... N1 No).
(k)
After obtaining the global Jacobian tensor (k)) the jihh entry of the global kernel H is cal- culated as the scaled Frobenius inner product of two horizontal slices of ( i.e. (H()),, ) 1
ju: )F. For simplicity, Fourth, the server will perform the NTK evolution to obtain the updated neural network function p(++1) and weight vector (k+1), With slight abuse of notation, let f (lk.t) denote the neural network output at gradient descent step in communication round k. The neural network function evolution dynamics and weight evolution dynamics are given by:
dg (f ((k))
f (l.et)
(I (I-e
H( (()) N
+e (( (())
(7a)
do
w
(1(
(J
(k)
T (k.t)
+ (K)
(7b)
:j:
5=1
4

---

Under review as a conference paper at ICLR 2022
where J is the jth lateral slice of j(h), and (R.t) R. is the jihh column of the residual matrix R (k.0)
defined as follows:
R
Nd ( - p(k,a)
( (x(1))
(8)
Note that (() = I XL denotes a concatenation of worker training examples, and ( ........ T denotes a concatenation of worker labels. The weight evolution in (7b) derived by unfolding the gradient descent steps. To update the global weight, the server performs 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:
the evolution with various integer steps (41,..., (v)
w (k+1) A (k),
{(k) = argmin L(f(w (1.ss) X Y(h))).
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 (b). the search of (()) over the grid (ti.....fy can be completed in ( 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 two-layer neural net- work. In contrast, NTK-FL. let each client transmit Jacohian 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 shifted to the server.
Robustness Against Statistical Heterogeneity. In essence, statistical heterogeneity comes from the decentralized data of heterogencous 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 global kernel Hik), 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 tradi- tional FL algorithm. More precisely, the update being sent for NTK-FL regarding the ith training example of the math worker for NTK-FL is Jmm = [Vh (xmi(x. ..... (x,, whereas the gra- dient update being sent for FedAvg is VL(W; WW: Fm,)) = 1 Σ}= (8m,k. Vm....) Vf, (Xm,,), weighted sum of coordinates of J m. By sending Jacobian matrices Jmm 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.
de
4.2 CP-NTK-FL VARIANT
Compared to FedAvg, NTK-FL does not incur additional client computational overhead since cal- culating 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, 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 up- link communication cost and space complexity are dominated by a third-order tensor (R) i.e., an O dNdd complexity compared to O (d) in FedAvg. For security, we investigate threat model where a curious server may perform membership inference attacks (Nasr et al. 2018) or data re- construction 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 -TK-FL by combining various tools to solve the aforementioned problems without jeopar- dizing the performance severely. Although is possible to incorporate these tools into FedAvg, we will show that overall it will lead to more severe accuracy drop.
follow the widely adopted analysis framework in wireless communication
5

---

Under review as a conference paper at ICLR 2022
aggregation server
shuffling
aggregngtion
performs permutatiio
server builds
kernel H
obtains Aw
Text Accumcy
shuftling server
trusted key server
0.7
key server
enerypts p
NIK-FL
transmits
0.6
NIK-FL
Foddng
E(ka, p)
FedAug
0.5
12 Round
worker gets Zi sends C(J ). Im (Z.)): Y.
Bmn C Den
(k) Ap
worker receives weight w(f & decrypts to get the seed p
Figure 3: Training results of 300 workers via NTK-FL and FedAvg, along with variants with the local dataset subsampling and random pro- jection, denoted as NTK-FL' and FedAvg', re- spectively. We train a two-layer multilayer per- ceptron on the Fashion-MNIST dataset. The joint effect causes more accuracy degradation in FedAvg (red) than in NTK-FL (black).
the m worker
Figure 2: Schematic of CP-NTK-FL. A trusted key server (orange) sends an encrypted seed E(kk, with the public key kt for random pro- jection. The client transmits the required message to the shuffling server (blue) for permutation.
Jacobian Dimension Reduction. First, we let the mth worker sample a subset B,m from its dataset D.,. uniformly for the training. Let 3 E (0, 1) denote the sampling rate, Bm, contains N BN.,. data points, with the training pairs denoted by (X', Y²) Next, we consider using a random pro- jection preprocess the input data via seed shared by trusted key server. Formally, the sampled training examples are projected into Z'n i.c., Z'm X,,, where P E Rd xd; projection matrix generated based on seed with IID standard Gaussian entries. In general, we have d' < 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) E TR,,", xd,xd', Compared with original version J (k) it has reduced dimensionality at the cost of certain information loss. Meanwhile, the random projection will defend against the data reconstruction attack, the Jacobian tensor is now evaluated at the projected data Zzz- We empirically verify their impact on the test accuracy in Figure 3 We d' 100 and sampling rate 0.4, and train multilayer perceptron with 100 hidden nodes on the Bishion-MNIS" dataset (Xiao et al. 2017). The joint effect of these strategies is slight accuracy drop NTK-FI. and a nonnegligible performance degradation FedAvg.
m.
Jacobian Compression and Shuffling, We use compression scheme to reduce the size of the Ja- cobian 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 al. [009)) To further ensure the confidentiality and membership pri- vacy, we introduce shuffling server, inspired by some recent frameworks (Girgis et al. 2021 Cheng 2021). to permute Jacobian tensors J (R) neural network states p(()) 's, and labels Y 's. Based on (7b). we denote the model update by (k ((++1) (k) E (k) TR. (k.1)
which is a sum matrix products. rows and columns are permuted in synchronization, the weight update (4) will remain unchanged. Considering the high dimensionality of the neural network weight, the reconstruction attack becomes computationally infeasible. As provable differential pri- vacy guarantee does not explicitly protect against the reconstruction attack [Zhang et al. 2020 we leave thorough privacy study for future work.
ANALYSIS OF ALGORITHM
In this section, we analyze the loss decay rate between successive communication rounds NTK- FL and make comparisons with FedAvg. Similar to Du al. (2019) and Dukler et al. (2020).

---

Under review as a conference paper at ICLR 2022
consider a two-layer neural network f R + R of the following form to facilitate our analysis:
J Σ (v) vn
"
f(x: c))
(10)
where x E IRd is an input, v, E Rd is the weight vector in the first layer, V = [V1, vnl.
er E I is the weight in the second layer, and a-)) is the rectified linear unit (ReLU) function, namely (2) = max(z, 0), applied coordinatewise. We state two assumptions as prerequisites.
Assumption The first layer Vr 's are sampled from (00 a²T). The second layer oy '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 imput). The input data are normalized, i.e,, ||x=|12 < 1, i.
For this neural network model, the (i. jtth entry of the empirical kernel matrix ((k given in 3.2
can be calculated as: (H()); )3 = x x; Σn
(8)
(1t)) where 1 (1) & 1((v))) (h) .x)) > 0}, and the
ir
term e, is omitted according to Assumption
Define H whose (i. entry is given by:
(H*) by - Ev-.N(0.aa1) xx x > 0) 1(v x, > 0)]
(11)
Let Ao denote the minimum eigenvalue of H which is restricted in the next assumption. Assumption 3 The kernel matrix a is positive definite, namely, do > O.
In fact, the positive-definite property 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 (A) (X) -yl1ž in communication round k. where
X x=
X E Nxd denote concatenation of client inputs and y = [y],.... y R. denote concatenation of client labels. We give the convergence result by analyzing how the
residual term decays between successive rounds.
Theorem For the NTK-FL scheme under Assumptions 1 t0 let the learning rate r = O Ag
and the neural network width = 5 loss decay of NTK-FL is
N Ng
In
N²) 8
then with probability at least 1 1, the one-round
(((+))() - v1 < (1
17
(x) - -1112
(12)
2N
where {8)) is the number of NTK update steps defined in
The proof of Theorem I 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 3 let the learning rate 17 O
An and the
(FNICET
neural network width S2
2²
In N²2 then with probability at least 1 5, the one-round loss X7
decay of FedAvg is
(()(t) - (1
nTAo 2N|CH|
(x)
(13)
where T is the number of local iterations, and |C&l is the cardinality of the worker set in round Remark (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 ! with the
decay coefficient in Theorem i.e., 1 - Nit de + O (ni) < <1
Darda where 1) < 1 for a large 2N
2N
N D The number of NTK update steps {(k) is chosen dynamically in (9). which is on the order of
10 to 10a, whereas T often on the order of magnitude of 10 literature (Reisizadeh et al. 2020 Haddadpour et al. 2021). One can verify that (h) Xo is larger than TgT [Cx] and draw the conclusion
For example, if we have 100 clients, each of which has more than 100 data points, then N is on the order of 10 Considering the choice of the learning rate 77, the Binomial approximation holds in I1.
7

---

Under review as a conference paper at ICLR 2022
0.00
0.33
0.34
0.40
1.000
0.26
0.03
TwNino 0.22
Tolly
26
29
20
(c)
Figure 4: Test accuracy versus communication round of different methods evaluated on: (a) EM- NIST dataset, where the heterogeneity comes from feature skewness. (b) non-IID MNIST dataset with label skewness, where the Dirichlet distribution parameter
0.5. (c) non-IID Fashion-
MNIST dataset with label skewness, where the Dirichlet distribution parameter a = 0.5. NTK-FL
outperforms all baseline FL algorithms in different scenarios, and achieves similar test performance compared with the ideal centralized training case.
6 EXPERIMENTAL RESULTS
Federated Settings. We use three datasets, namely, MNIST (Xiao et al. 2017), and EMNIST (Cohen et categories. For MNIST and Fashion-MNIST, we follow
MNIST (eCCun et al. 1998). Fashion-
2017) digits. All of them contain C' 10 Hsu et al. (2019) to simulate non-lID data with the symmetric Dirichlet distribution (Good 1976). Specifically, for the mtth worker, we draw
random vector mm Dir(ar), where 9m [9m.1..... (tm.l--.-99..)) belongs to the (C 1)-sanddard simplex. Images with category are assigned to the mith worker in proportional to (100- qm. )%. 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 per- ceptron 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.
We empirically verify the convergence rate of the proposed method. For FedAvg, use the number of local iterations from (10, 20, ....... 50) and report the best results. For NTK-FL choose {(4) over the set {100, 200, ..... 2000}. We 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 global dataset is broadcasted workers for local training; the size of the global dataset is be 10% of the total number of local data points. (ii) Federated normalized averaging (FedNova) (Wang al.) (020), where the workers transmit normalized gradient vectors to the server. (iii) Central- ized training simulation, where the server collects the data points from subset C₁ of workers and performs gradient descent to directly train the global model. Clearly, scheme (iii) achieves the per- formance that can be considered as an upper bound of all other algorithms. The training curves over three repetitions are shown Figure More implementation details and the results on CIFAR- 10 (Kizhevsky 2009) can be found in Appendix A Our proposed NTK-FL. method shows consistent advantages over other methods in different non-lII scenarios.
Convergence.
In this experiment, we select the Dirichlet distribution parameter from 1, 0.3, 0.4, 0.5) and simulate different degrees of heterogeneity on Fashion-MNISTT dataset. A smaller a 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%. can be observed that NTK-FL achieves stable test accuracy in different heterogeneous settings. In comparison, FedAvg and FedNova show performance drop in the small a region. NTK-FL has more advantages over baselines methods when the degree of heterogeneity is larger.
Degree of Heterogeneity.
Effect of Hyperparameters. We study the effect of the tunable parameters in CP-NTK-FL. We change the local data sampling rate 33 and dimension d', and evaluate the model test accuracy on the non-llD Fashion-MNIST dataset (a = 0.1) after 10 communication rounds. The results are shown in Figure 6 A larger data sampling rate 33 larger dimension d' will cause less information loss,
8

---

Under review as a conference paper at ICLR 2022
0.86
0.85
83
83.8
84.8
54.4
84.5
85.0
Test Text Accumecy
84
0.84
$2.8
$3.9
84.4
$4.5
84.4
84.9
Acoomcy Mo
0.83
NTK-FL
0.82
DutaShare
FedNova
83
FedAvg
R2.4
63.5
84.2
54.4
84.1
64.2
0.1
0.2
0.3 0.4
0.5
$1.9
$3.2
83.7
94.2
84.5
84.6
-82
Figure 5: Test accuracy versus the Dirichlet distri- bution parameter a for different methods evaluated on the non-IlD Pashion-MNIS, dataset. Reducing the value of will increase the degree of hetero- geneity in the data distribution. NTK-FL is robust different heterogeneous data distributions, and shows more advantages over FedAvg and FedNova when the degree of heterogeneity is larger.
Figure 6: CP-NTK-FL test accuracy for dif- ferent hyperparams. A larger data sampling rate / and a larger dimension d' are expected to give higher test accuracy. In general, the scheme is robust to different combinations of hyperparameters.
and are expected achieve higher test accuracy. The results also show that the scheme is robust to different combinations of hy hyperparameters.
We evaluate the uplink communication efficiency of CP-NTK- FL (d'- 200, B = 0.3) by measuring the num- ber of rounds and cumulative uplink commu- accuracy of 85% on non-IID Fashion-MNIST dataset (a = 0.1). The results over three repetitions shown in Table [ Compared with federated learn- ing with compression (FedCOM) (Haddadpour et al.] (0211, quantized SGD (QSGD) (Alis- tarh et al., 2017), and FedAvg, CP-NTK-FL. achieves the goal within an order of magnitude fewer iterations, which particularly advanta- geous for applications with nonnegligible en- coding/decoding delays or network latency.
Uplink Communication.
Table I: Uplink communication cost to reach 85% on non-IID Fashion-MNIST dataset (a 0.1). CP-NTK-FL
achieve the target goal within the fewest communication rounds without incur- ring communication cost significantly.
nication cost reach
optimization
comm.
comm. cost (MB)
algorithms
rounds
CP-NTK-FL.
26
386
FedCOM
250
379
QSGD (4 bit)
614
465
FedAvg
284
1720
CONCLUSION AND FUTURE WORK
In this paper, we have proposed an NTK empowered FL paradigm. inherently solves the statistical heterogeneity challenge. By constructing global kernel based the local sample-wise Jacobian matrices, the global model weights can be updated via NTK evolution in the parameter space. Com- pared with traditional algorithms such as FedAvg, NTK-FL has 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 CNNs, residual networks (ResNets) (He al.) (016)), 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.
REFERENCES
Sina Alemohammad, Zichao Wang, Randall Balestriero, and Richard Baraniuk. The recurrent neural tangent kernel. In International Conference Learning Representations, 2021.

---

Under review as a conference paper at ICLR 2022
Dan Alistarh, Demjan Grubic, Jenry Li, Ryota Tomioka, and Milan Vojnovie. QSGD: Communication-effitie efficcet SGD via gradient quantization and encoding. Advances Neural In- formation Processing Systems, 2017.
Dan Alistarh, Torsten Hoefler, Mikael Johansson, Nikola Konstantinov, Sarit Khirirat, and Cedric Renggli. The convergence of sparsified gradient methods. In Advances in Neural Information Processing Systems, 2018.
Sanjeev Arora, Simon S. Du, Wei Hu, Zhiyuan Li, Ruslan Salakhutdinov, and Ruosong Wang. On exact computation with an infinitely wide neural net. In Advances Neural Information Processing Systems, 2019.
Hong-You You Chen and Wei-Lun Chao. Fedbe: Making Bayesian model ensemble applicable to feder- ated learning. In International Conference on Learning Representations, 2021.
Zixiang Chen, Yuan Cao, Quanquan Gu, and Tong Zhang. A generalized neural tangent kernel analysis for two-layer neural networks. In Advances in Neural Information Processing Systems, 2020.
Pau-Chen Cheng, Kevin Eykholt, Zhongshu Gu, Hani Jamjoom, KR Jayaram, Enriquillo Valdez, and Ashish Verma. Separation of powers in federated learning, arXiv preprint 1050:2005500400, 2021.
Gregory Cohen, Saeed Afshar, Jonathan Tapson, and Andre Van Schaik. EMNIST: Extending MNIST to handwritten letters. In International Joint Conference on Neural Networks, 2017.
Simon S. Du, Xiyu Zhai, Barnabas Poczos, and Aarti Singh. Gradient descent provably optimizes wwerpparameterized neural networks. In International Conference Learning Representations, 2019.
Yonatan Dukler, Guido Montufar, and Quanquan Gu. Optimization theory for ReLu neural networks trained with normalization layers. International Conference on Machine Learning, 2020.
Antonious M. Girgis. Deepesh Data, Suhas N. Diggavi, Peter Kairouz, and Ananda Theertha Suresh. Shuffled model of differential privacy in federated learning. In International Conference on Arti- ficial Intelligence and Statistics, 2021.
Irving J Good. On the application of symmetric dirichlet distributions and their mixtures to contin- gency tables. The Annals of Statistics, 4(6): 159-1189 189, 1976.
Farzin Haddadpour, Mohammad Mahdi Kamani, Aryan Mokhtari, and Mchrdad Mahdavi. Feder- ated learning with compression: Unified analysis and sharp guarantees. In International Confer- on Artificial Intelligence and Statistics, 2021.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In International Conference on Computer Vision and Pattern Recognition, 2016.
Tzu-Ming Harry Hsu, Hang Qi, and Matthew Brown. Measuring the effects of non-identical data distribution for federated visual classification. arXiv preprint arXiv:1909.063:3 2019.
Baihe Huang, Xiaoxiao Li. Zhao Song, and Xin Yang. FL-NTK: A neural tangent kerne--base framework for federated learning convergence analysis. arXiv preprint arXiv:2105.0500. 2021.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and gen- eralization in neural networks. In Advances in Neural Information Processing Systems, 2018.
Peter Kairouz, H Brendan McMahan, Brendan Avent, Aurélien Bellet, Mehdi Bennis, et al. Ad- vances and open problems in federated learning. Foundations and Trends in Machine Learning, 2021.
Sai Praneeth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank Reddi, Sebastian Stich, and Ananda Theertha Suresh. Scaffold: Stochastic controlled averaging for federated learning. In International Conference Machine Learning, 2020.
10

---

Under review as a conference paper at ICLR 2022
Tamara G Kolda and Brett W Bader. Tensor decompositions and applications. SIAM review, 51(3): 455-500, 2009.
Alex Krizhevsky. Learning multiple layers of features from tiny images. Master thesis, Dept. of Comput, Sci., Univ. of Toronto, Toronto, Canada, 2009.
Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient- based learning applied to document recognition. Proceedings of the IEEE, 96(11):2278-22) 23244, 1998.
Jaehoon Lee, Lechao Xiao, Samuel Schoenholz, Yasaman Bahri, Roman Novak, Jascha Sohl- Dickstein, and Jeffrey Pennington. Wide neural networks of any depth evolve as linear models under gradient descent. In Advances in Neural Information Processing Systems, 2019.
Jachoon Lee, Samuel S. Schoenholz, Jeffrey Pennington, Ben Adlam, Lechao Xiao. Roman Novak, and Jascha Sohl-Dickstein. Finite versus infinite neural networks: an empirical study. In Advances in Neural Information Processing Systems, 2020.
Tian Li, Anit Kumar Sahu, Ameet Talwalkar, and Virginia Smith. Federated learning: Challenges, methods, and future directions. IEEE Signal Processing Magazine, 37(3):50 60, 2020.
Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith. Federated optimization in heterogeneous networks. Proceedings of Machine Learning and Sys- lems, 2020.
Xiaoxiao Li, Meirui Jiang, Xiaofei Zhang, Michael Kamp, and Qi Dou. FedBN: Federated learn- ing on non-iid features via local batch normalization. In International Conference on Learning Representations, 2021.
Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. Communication-effccceett learning of deep networks from decentralized data. In Artificial Intelli- gence and Statistics, 2017.
Preetum Nakkiran, Gal Kaplun, Yamini Bansal, Tristan Yang, Boaz Barak, and Ilya Sutskever. Deep double descent: Where bigger models and more data hurt. In ICLR 2020 Eighth International Conference on Learning Representations, 2020.
Milad Nasr, Reza Shokri, and Amir Houmansadr. Machine learning with membership privacy us- ing adversarial regularization. In ACM SIGSAC Conference Computer and Communications Security, pp. 634-646, 2018.
Sashank J. Reddi, Zachary Charles, Manzil Zaheer, Zachary Garrett, Keith Rush, Jakub Konečný, Sanjiv Kumar, and Hugh Brendan McMahan. Adaptive federated optimization. In International Conference on Learning Representations, 2021.
Amirhossein Reisizadeh, Aryan Mokhtari, Hamed Hassani, Ali Jadbabaie, and Ramtin Pedarsani. Fedpaq: A communication-effiiii federated learning method with periodic averaging and quan- tization. In International Conference on Artificial Intelligence and Statistics, 2020.
Felix Sattler, Simon Wiedemann, Klaus-Robert Müller, and Wojeiech Samek. Robust and caaamunniaaion-efffiieeeee federated learning from non-iid data. IEEE Transactions on Neural Ner- works and Learning Systems, 3 (19))3400-334 14100-313,,201 2019.
Hyowoon Seo, Jihong Park, Seungeun Oh, Mehdi Bennis, and Seong-Lyun Kim. Federated knowl- edge distillation. arXiv preprini arXiv:2011.00 1.02667, 2020.
Virginia Smith, Chao-Kai Chiang, Maziar Sanjabi, and Ameet Talwalkar. Federated multi-task learning. Advances in Neural Information Processing Systems, 2017.
Lili Su, Jiaming Xu, and Pengkun Yang. Achieving statistical optimality of federated learning: Beyond stationary points. arXiv preprint arXiv:2106.12211 2021.
Nguyen H Tran, Wei Bao, Albert Zomaya, Minh NH Nguyen, and Choong Seon Hong. Federated learning over wireless networks: Optimization model design and analysis. In IEEE INFOCOM 2019-IEEE Conference on Computer Communications, pp. 1387-1395. IEEF, 2019.

---

Under review as a conference paper at ICLR 2022
Paul C Van Oorschot. Computer Security and the Internet: Tools and Jewels. Springer Nature, 2020.
Hongyi Wang, Mikhail Yurochkin, Yuekai Sun, Dimitris Papailiopoulos, and Yasaman Khazaeni. Federated learning with matched averaging. In International Conference on Learning Represen- fations, 2020.
Jianyu Wang. Qinghua Liu, Hao Liang, Gauri Joshi, and H. Vincent Poor. Tackling the objective in- consistency problem in heterogeneous federated optimization. In Advances in Neural Information Processing Systems, 2020.
Jianyu Wang. Zachary Charles, Zheng Xu, Gauri Joshi, H Brendan McMahan, Maruan Al-Shedivat, Galen Andrew, Salman Avestimehr, Katharine Daly, Deepesh Data, et al. A field guide feder- ated optimization. arXiv preprint arXiv:2107.06917, 2021.
Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-MNIST. novel image dataset for bench- marking machine learning algorithms. arXiv preprint arXiv: 077477777777 2017.
Greg Yang and Etai Littwin. Tensor programs iib: Architectural universality neural tangent kernel training dynamics. In International Conference on Machine Learning, 2021.
Yuheng Zhang, Ruoxi Jia, Hengzhi Pei, Wenxiao Wang, Bo Li, and Dawn Song. The secret vealer: Generative model-inversion attacks against deep neural networks. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp- 253 233-261, 2020.
Yue Zhao, Meng Li, Liangzhen Lai, Naveen Suda, Damon Civin, and Vikas Chandra. Federated learning with non-iid data. arXiv preprint arXiv: 00582, 2018.
Ligeng Zhu, Zhijian Liu, and Song Han. Deep leakage from gradients. In Advances in Neural Information Processing Systems, 2019.
12

---

Under review as a conference paper at ICLR 2022
A ADDITIONAL EXPERIMENTAL RESULTS
We give the detailed setting of the learning rate and batch size. For the learning rate ?/. search over the set {10-3, 3 x 10-3, ,0-,,10-2 2,10-23 333 x 10-2,10 10 The learning rate is fixed dur- ing the training. For the client batch size, we it 200 for all datasets. We eval- uate different methods, including the central- ized 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
0.45
0.40
Text Asemaey
Centralzzed
NTK-FL.
DataShare
0.30
FodNova
FodAvg
0
10
20
30
40
50
Commmiiattio Round
dataset (Krizheskyy 2009 and present the re- sults 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
Figure 7: Test accuracy versus communication round of different methods evaluated on the non- IID CIFAR-10 10 dataset, where the Dirichlet distri- bution parameter a = 0.1.
B PROOF OF THEOREM
Let Im denote a set of indices such that for i E Im. (x:, Vi) E Dm- We first present some lemmas facilitate the convergence analysis. In communication round k, define s!A) as the set of indices corresponding to neurons whose activation pattern is similar to its initial state for an input x
S (r (1.....,m)| m
||v
R.
(0) 1 + 1 (v
x, 0)}.
(14)
We upper bound the cardinality of S in Lemma T
Lemma Under Assumption
with probubility least 1 8, we have
nnR ViE
|S|| < V
NJ.
(15)
π da
Proof. To bound |S:| = Σ 1 (r 8.), consider event Air defined follows:
Arr A 13
v R. 1 (y x, 0)}.
(16)
Clearly, 1 (r S,) = 1 (Arr). According to Assumption 2 | < 1, it can be shown that the event Air happens if and only (v. Tx, < based on geometric argument, Based Assumption | we have (0)) x; ~ NO, a²). The probability of event A,, is
P[AN1 = P [[ Txi
x, < R
(17a)
R
R
erf
V a
(17b)
22xx
By Markov's inequality, have with probability at least 1 5,
2 nk
Σ 1( (ireS) 8.) <
(18)
5a
The proof is complete.
n
We bound the perturbation of the kernel matrix (k.0) in Lemma 2
Lemma Under Assumption to vfr {l,...,nk n). (k.tt
( l|2 & R, then
2v 2NN
(k.2) H (0) 1/2 <
(19)
vōoa
13

---

Under review as a conference paper at ICLR 2022
Proof. We have
12.0
H
H (k.0)
H
(20a)
N
N
H (H.))))ss
(H )()
(20b)
[x] x2)²
Σ (k.t)
(4x-0) 1 1//
(20c)
ir
Consider the event Air defined in 16). Let (k:t) happen, clearly we have (k..t)
& (1.A) (R.)) - (0) 1o If AAir and -Afr
Pgr
can be bounded as 0. Therefore, the expectation of 10 sjjr
(k.t))
(k.0) Pijr
< IP(Air U A -1
(21a)
< (P(A)))) + P(Arr)
(21b)
2 R
2
(21c)
π a
where ! comes from (7bb). By Markov's inequality, we have with probability at least - 8,
2 R
2
(22)
δa
Plugging
22
into 20c yields
N2 8n R2
8N2
(k.))
(23)
n x
x
Taking the square root both sides completes the proof.
Lemma 3 With probability at least 1 8,
In (2N2/8)
H (0) H
N
(24)
2n
Proof. We have
N
H H
H
Σ
(0) (H)))
(25)
Note that (H) x x,
Σ
ir jr
(0) (0)
(H yg E [_1, By Hoeffding's inequality, we have
with probability at least 8/n2,
(2N2/8)
(H()); )ij
(26)
2n
Applying the union bound over i,j [N] yields
In (2N2/8) 2
H (0) H 2 <
NV
(27)
2n
The proof is complete.
Now we are going to prove Theorem
Theorem For the NTK-FL scheme under Assumptions 1 to the neural network width n = S x7 2
let the learning rate 17 O ( and
then with probability at least 1 - 8, the one-round
loss decay of NTK-FLI is
804)
[((++)((X) - y113 S
n7o
(x) - -113-
(28)
2N
14

---

Under review as a conference paper at ICLR 2022
Proof.
Taking the difference between successive terms yields
((x)+t)(xx)) (xx)) 1 Σ Go ((v(h:tt1)) xi) c,a ((v(hi0)) x3) )
"
(29)
We decompose the difference term to the sum of dł and d.t based on the set S₁:
df
A
G,
(v((,.++1))xx4 xx) c. ((v!4:0)) xi)
(30a)
rgS,
d
G
(v(l: xx) Go ((vSti0)) T xi)
(30b)
rES;
Consider the residual term
| (X) y11
(31a)
-[(((h+t)) - -((h.)) (X) + (x))0()) v12 =[1 - -1122 + 2 d + d", f(h.)) (X) -y) + [|((((,tt)) - f(h.0) (x)
(31b)
(31c)
We will give upper bounds for the inner product terms (d', (X) (x)-)) y): (d . ((Ht)) y)
and the difference term ( (X) (k.1) (x) separately. Based on the property of the set Sü we have
df =
Σ G.
LV (1.1))
(32a)
Vn
N
7
(fth0) (xx)) - -30) x3 x, E
(k:t) (h.1) cx ir
(32b)
nN
fr
j=1
r4s,
N
(x,) -w;) ((H(K))))s - (H-
(32c)
where (H (H- (k.0))s is defined as
H (1:0))
A x, Xj
1 (k.t) 16.1)
(33)
n
rES,
For the inner product term (d', (x)
- y)) we have
(dl, ((h0) (X) y) = " (p(1100) ->)
(T - (((0)))(((()))))) -))
(34)
Let T, and Ta denote the following terms
T _ ((((t.))))) - HH ( (X) y):
(35a)
T2 A (f(h (X) y) H ((())) (X) y).
(35b)
With probability at least 1 - 8. Ti can be bounded as:
T = (f((h)) (X) - y) T) ((H(k. - (0) + ( 0 + H (F(t,)()) yy
(36a)
(p(k.)) (X) T) ((H(k.)) - -H (((k)t) (X) yy)
-((KK0 (X) T) ((H(") - -HO) (K.)e) (X) -y) - - 112
(36b)
2V2NR In (2N2/8) 2o [|((K)((x) -yl12-
VTda NNV
(36c)
2n
15

---

Under review as a conference paper at ICLR 2022
where ① comes from Lemma 2 and Lemma To bound the term T2, consider the l norm of the
matrix H (k:9) With probability at least
6, we have:
H 1|2 &
| ||
(37a)
2 :
N
(1:))) x, x ( 5r
x
(37b)
rES,
4
NR da
S
(37c)
n
π
where ! comes from Lemma] Therefore, with probability at least 1 - 6, we have
NR da
T2
(k.0) X) -11
(38)
Combine the results of 36c and (38
d'' (1,0 (X) 8) E
3 2R
In (2N2/8)
Ao [x)
- 1112
(39)
VTδa
2n
For the inner product term (d" ((k,t) (X) y). we first bound ||d as follows:
Nv'
|
Go (v.l.ctt))T Tx. -c, ((v!h.0)) Txi)
(40a)
t=1
rES:
N
17
S
Vv,
(40b)
rES;
1
S
Σ
|
T.12 1x
(40c)
rES,
11 N
S
max rE[n]
V L
(40d)
77 S 2
1
(40e)
n2
where ① comes from the Lipschitz continuity of the ReLU function o-),, @ holds due Cauchy- Schwartz inequality. Plug (18) into (40e), we have with probability at least 1 - 6:
2r²2R²
ldd < Tr ( (x) -
(41)
The inner product term (d", ((k (X) y) can be bounded as
V2nR (kk.0 (X) -112 VTδa
(a". p(ht) (X) y)
(42)
Finally, the bound for the difference term is derived as
2 G L,,
(x)
((..0)
< 5
-11 (43)
-116-
Combine the results of
42 and
(43
( (X) - y112-
1+
8v 22R
+ +2
In (2N2/8)
2n1do
+ n
[x) -y1?. (44)
-
vida
2n
N
Let R == 000 5S
N2 In )²²) and 17 1 00 A2)) we have
2g
[|( - y112 < (1 n7
2N (x) -yl
(45)
Summing up the selected number (k) iterations completes the proof.
16

---

Under review as a conference paper at ICLR 2022
C PROOF OF THEOREM! 2
Theorem 2 For FedAvg under Assumptions
and the then with probability least 6, the one-round loss
1 to 3
let the learning rate 1l 0
An
[FN[CE]
neural network width n 52 In
N²
2N
26
0
decay of FedAvg is
[[((+t)(()
-13)--llo y12 < (1
Ao 2N|CE
[[f()(()) -x112
(46)
Proof. We first construct a different set of kernel matrices (A),, A (l( similar to Huang et al.
2021 Let 1 (k:,)) A 1 (((v).))
(k:a))
x;)
> 0}. the (i, entry of Am
(k..)) and A (k.u) is defined as
imr
(A (k.a)) وe
1
xj k55 if (x,, Uj) E Dm-
iimr I jmm''s
(R.0)
(47a)
11
(A (krw)) )s
-
(47b)
m
Taking the difference between successive terms yields
((++)(xx)) - l (x,) 1
((v.)t)))(xx)- -cr C,O ((v(1)) xxxx)
G,O
(48)
We decompose the difference term to the sum of dt and a!! based on the set S; and its complement:
Go (vv x:) - cra ((v!4))T& xx ]
d! A
(49a)
r&S,
d A
CG ((v(+t2)) - Go ((v.tt) x, )
(49b)
rES;
Consider the residual term
|| (X) y1
(50a)
= - - + ((0) y112 = (x) - -1122 + (at + )()) (x)(
(50b)
(x) - -y) + [|](((t1(())))))))))(()({()) - ((8) 11 (50c)
We will give upper bounds for the inner product terms (d', ((0 - y). (dt f(k) (X) y). and the difference term [| f(k) [x) separately. For an input x E Rd, let (k.u) (x) 2
- (co((vvi (k:w)) x)). By the update rule of FedAvg, the relation between the weight vector V
in successive communication rounds is:
v (k+1)
(b)
LL (B,M)
(51a)
Cal meCi u=0
(1))
TJC
Σ
(x() Uj W;)xx jmr
(51b)
N /
mEC+ u= jeIm
Based on the property of the set S₁, we have
-
d! d- 1 Σ ΣΣ
(v(l+t)) - -v()) ,,,,,x))))))))),) 11th)
(52a)
mEC₁ a=0 S.
" Σ
ΣΣΣΣ
(F(0 - UU xx, x, (/) (k.w)
(52b)
Nn|Cal mECa u= rgS
jmr
r-1 " Σ
(x())))(x)))) - wi) [AAtiao) (A
(52c)
-NcC meC, u=u jEI
17

---

Under review as a conference paper at ICLR 2022
For the inner product term (d', (X) -y), we have
ad', ((b) (X) y) -x) - "
- -x)) (kaa)) A ( (1th:0))))))))) (X) - y). (53)
NI ==
Let T1 and T2 denote the following terms
Ti (( (X) y) ( (
(X) (X)- -y),
(54a)
Te - (f(h) - -y)
T A ((k f( (X) - -y),
(54b)
where f/ s(4,,0(( (X) A |s( (x, T S/20 (X (X(ce))TTT We are going to bound T₁ and T2 sepa-
rately. Ti can be written as:
Ti = -(((()) (X) - y) T (k.u) H + + (0) H H (( - y)
(55a)
=-(119)XX (X) (()((A)))))) H (1
(X) y)
- (X) - -y) T(H()) )
(f (x)- y)
(X) y) H ( y)
- -(1 - y) T (1 (s(
(k) (X)).
(55b)
First, we bound the norm of j p())a)) (X) y. It can be shown that
|| (l0.)) (Xm)
ymml2 (1 (Xm) plk 1) (Xm) + ((k.u-1) (Xm) Vm
(56a)
|| (k. (Xm) - -( 1) (X 112 + (( (X,m) ymll2
(56b)
(1 + 1) ((k. (Xm) - ymmlz 1l2
(56c)
where holds based on the derivation of Applying (56c) recursively yields
ll (Xm) - mmlll (1 + m) (Xm) - vmlle.
(57)
The bound for 11/ (X) y13 can thus be derived as
[1 (X) -
plle.) (xx) -Ub 1
(58a)
|| (Xm) - -mml 11
(58b)
mECt
< (1 + m) | (X)
-y12-
(58c)
Second, following the steps in Lemma 2 it can be shown that with probability at least 1 - 5,
2/2NR
| (k.1) H <
(59)
VTōa
We also bound the difference between f(k.)) (X) and j ((k) (X) as follows:
8
u
|1 (k) (X)
((k) (X) 1|2
(kk,)
(X) f( 1) (X)l12
(60a)
g
fll. 1) (X) - 12
(60b)
g
u=1
@
711 + 11) "- p|)) (x) -1l/2
(60c)
V=1
[(1 + m)* (x)
- y y/2,,
(60d)
18

---

Under review as a conference paper at ICLR 2022
where ①D holds due to triangle inequality. 2 comes from (43). @ comes from (8c). results from (5c). (59). and (60d) into (55)). we have with probabiiiy at least 8,
Plugging the
2V/2N
In (2N2/6) + kλO (1 + k)
T1 <
(1 + m)
+ N
(X)
-yll. (61)
VTōa
2n
where a. is the condition number of the matrix H Next, consider the bound for T2. The Eg norm of A (kk,,u) can be bounded as
|| (k,
12 < ||
L(k.a)) ||F
(62a)
2) 7
)
Σ
Σ rES,
x x, (1b) (A:.u)
(62b)
1=
mEC+ J =E
N
NR da
(62c)
n
where comes from Lemma
1 Therefore, we have with probability at least 6,
N2 π da
T2 (1+mm)
( (X)
(63)
Combine the results of 61 and 663
(() - x) x
(5 - 1)
11² + o(19))
3/2R VTδa
In (2N2
+
+
+
Cc
2
2n
N
(64)
(1 + k
||((((()) -y1
N
For the inner product term (d ((k) (X)
y)), we first bound |dd with probability at least --:: 5:
((v(44411) xx, Go ((v{0)]xx))] xx)]
2
GoO
(65a)
rES.
N
S Σ
(cr(v(tt))
(k) xi) 2
(65b)
1=1
TES:
N
1
S Σ
ne,
(Fl)))((()) )- - v; (k.
(65c)
N
jm
1=1
TES:
meCi 4=0 jEIm
1
| Σ
Σ mECa Jmm
plk.u)) (xx) - Mi
(x,))(x,))
(65d)
N2 CP 2
reS;
1/2
| Σ
Σ
Iml p (X - Vm Xm112
(65e)
N2 C2
1=1
rES,
\mEC u=0
| Σ
Σ
(1 + m)" (()) -8mm
(65l)
N2n² CA12
reS;
\mEC, u=0
Σ - 1) 5(0((()))mm ym 111 2
1
Σ
| Σ
(65g)
N2
rES;
AmEC
N
((1 + n)² - -1)++)*-11|pp -1) - =x11)²
|S||
(65h)
Nn Cal2
i=1
rES.
2R
T(T - 1) 17 + o(12)) (x))) -y112
<82a?1c112
(T)
2
(65i)
19

---

Under review as a conference paper at ICLR 2022
where @ comes from (57]- 2 holds due to ||a]]] < ||a||2 @ holds due to ||a|| & v dim a
| $
from Lemma
I
With probability at least - 6, the inner product term can thus bounded as
V VT
1 (T 1 + ulr))) [|)00x) - y1122
(d"
(66)
2
The bound for the difference term is derived as
(x)
(v (k+1)
vpr xxi
(k)
(67a)
1=1
pm
T(T 1)
2 112 0(1))) [|{(2)(x) y112
T1j +
(67b)
C 2
2
Combine the results of 64
66
and (67b
||
v112 <
+ 2//T
4/2R
In ( (2N2
kdU
[c
VTδa
N
(68)
(1 + Ao
11²7
111100020-- -11
cx12 + o(()))
N
Let R 0 0 (
n S2
N2 xg
In
N² 8
and 1j = 0 >N FN/Caa) ) we have
|||( - 2112 < (1
A
[1)( - -11
(69a)
2N|CE
20


Reviews:
[{"review_id": "fQgIhSJaei", "reviewer": "Reviewer_uPVB", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "- This paper uses the NTK method as a tool for predicting the output of neural networks. As in eq (6), the dynamics of the gradient flow can be expressed by the neural tangent kernel (NTK). However, the NTK does not change over time under the **infinite-width regime** (please see [1]). The paper does not mention anything about the width limit of networks. Moreover, the neural network used in the proposed FL framework cannot have infinite width, hence, the evolution dynamics (Eq. 8, 10(b)) may be not valid to predict the network outputs.\n\n    [1] Chizat et al., On Lazy Training in Differentiable Programming. NeurIPS. 2019\n\n\n- The paradigm requires each worker to send a Jacobian tensor with respect to its assigned dataset, which increases communication overhead. Indeed, the transmitted parameters are used for updating the global parameter governed by Eq. (10b). Can the communication bottleneck improve by transmitting only essential factors for the updates? For example, instead of sending J_m^{(k)}, f^{(k)}(X_m), Y_m, one can transmit the corresponding sub-matrix of w^{(k,t)}, which leads to reduce the communication overhead. Is there any reason that workers send the Jacobian matrix?\n\n\n- CP-FL-NTK variant uses dimension reduction via a random projection with a shared seed. A larger compression ratio (beta) reduces the communication overhead better but the performance also does so. How does degradation relate to the compression ratio? Is there any provable guarantee or practical guidance to choose a proper value of beta? And what is the exact random projection? Is it Gaussian random projection? or random sketching? Moreover, can the corresponding privacy guarantee be analyzed?\n\n\n- In Remark 1, the authors argue that Eq. (19) implies that convergence in Eq. (17) is faster than that in Eq. (18). But, precisely, it is not true because of $1- t \\cdot (\\eta \\lambda_0 / (2N)) \\leq 1- (\\eta \\lambda_0 / (2N))^t$ for integers $t \\geq 1$.\n\n- Minor Issues:\n\n  - The big-O notations usually do not include constants, hence writing $\\Omega(N^2 / \\lambda^2 * \\log(N/d))$ in Theorem 1 seems more natural.\n\n  - It is good to briefly describe FedAvg algorithms for better understanding.\n\n  - A d-dimensional Dirichlet distribution requires a d-dimensional parameter vector (alpha), but it is chosen by scalar values in section 6. Does it mean that the vector with same values? Do the authors have a chance to try other parameter (alpha) settings?\n\n  - It would be great if graph colors and markers in Figures 4 and 5 match each other.\n", "summary_of_the_paper": "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.", "summary_of_the_review": "The paper is well-written and the methodology is very clear. However, it seems that the paper needs to address the usage of NTK more concretely, e.g., clarifying the width limit. In addition, the second algorithm with communication efficiency and privacy protection also needs more rigorous analyses. Overall, this paper is under the bar of acceptance.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "1: The contributions are neither significant nor novel."}, "extra_text": {}}, {"review_id": "e9WekRkqAzF", "reviewer": "Reviewer_KsCY", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n+ The idea of using NTK for model optimization without gradient descent in the FL paradigm is interesting and somewhat novel. \n+ Inherent issues such as communication cost and data privacy of the proposed method are properly discussed. \n\nWeaknesses:\n\nA. Method\n- In a way, the proposed method shifts much of the computation to the server-side (e.g., weight update and selecting the best one). The various integer update steps for the weight evolution are not clearly specified in the experiments (reproducibility). And the time cost for the grid search of $t^{(k)}$ should be discussed. \n- To address the risk of data leakage and the high communication cost of FL-NTK, random data projection, and Jacobian shuffling are used. Although it seems that these tricks will do the job, additional resources (i.e., the trusted key server and the shuffling server) are required, which may not be available for some applications. \n\nB. Experiment evaluation\n- The network used in the experiments, i.e., a multilayer perception with 100 hidden nodes, is too simple. More standard network architectures like ResNets (e.g., ResNet-18) should be employed for evaluation and comparison. \n- The datasets used in the experiments are also small, and not diverse (basically MNIST and its variants). It would be more convincing to show results and comparison on larger datasets (more samples and classes), such as CIFAR-100 and Tiny-ImageNet.\n- In the data heterogeneity experiments, a few popular baselines such as FedProx and SCAFFOLD are not compared. \n- In Fig. 5, the performance of FL-NTK does not increase as the level of data non-IIDness goes down. An accuracy drop is observed when $\\alpha=0.3$. Can you provide an explanation for this? \n- For CP-NTK-FL, communication with the trusted key server and the shuffling server is also required, in addition to the communication between clients and server. Therefore, a breakdown of the comm. cost is helpful information. ", "summary_of_the_paper": "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. ", "summary_of_the_review": "Overall, the proposed NTK approach for FL is interesting and has novelty. However, there are a few key limitations/weaknesses: (1) FL-NTK has a higher risk of data leakage and communication cost compared to FedAvg. The proposed CP-FL-NTK may address these issues but at the cost of requiring additional resources (i.e., the trusted key server and the shuffling server), which could limit the practical usage of the method. (2) As a general approach, a more comprehensive and rigorous experimental evaluation is expected to fully demonstrate the effectiveness of the proposed method.\n\n\n", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "Wba3eLq9r-W", "reviewer": "Reviewer_C84w", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "This paper combines two popular relevant machine learning frameworks, NTK and federated learning, with the aim of addressing statistical heterogeneity in federated settings. The paper proposes a practical implementation that addresses potential communication burdens. The experimental section provides interesting experimentation showing this framework outperforms FedNova and FedAvg on three datasets. \n\nThe paper is well organized and clear. Previous work is discussed and necessary background is introduced. The experimental section also presents interesting results. Plots and experiment details are discussed and explained. \n\nWeaknesses / clarifications.  \n- My main criticism is that throughout the paper it is mentioned that the NTK “inherently solves the non-IID data problem”, when learning global model under statistical heterogeneity of workers. Yet, from the formulation of the NTK this is not evident and it is not clearly developed in the paper, for example by  explaining why the Jacobian enhances generalization while gradients don’t. Why is this framework more robust? I acknowledge that at the end of section 4.1 there is a paragraph on this, however it does not respond to my question as it affirms that the global kernel $H^{(k)}$ is more expressive than gradients without an explanation of why. \n\n- It is not very clear for me how this algorithm is different from Algorithm 1 in [HLSY21]. Further, it would also help to differentiate the paper from this previous to include a discussion on how the presented convergence results are different from  the ones in [HLSY21]. \n\n- Given that several $t^{k}$ are tested at each round, how does this affect the convergence result?\n\n\n- Throughout the paper a”privacy” is used in different contexts. I would suggest a clarification of what privacy means in each context and what attacks are prevented by taking a specific action. For example, the random projections do not protect from an \"honest but curious\" server, since all workers share the same projection matrix. This step does not provide Differential privacy. Some differential privacy guarantees can be provided from the shuffling model, but again this is not very clearly explained, and what it is protecting from. \n\n\nMinor:\n- I understand NTK intrinsic difficulties to try on larger models, but given that the contribution of this paper is a practical algorithm it would be interesting to see results in more realistic FL settings, perhaps a small language task. \n- Given that FL is an applied field, It would be good to have a small discussion on practical aspects of FL like what to expect under user intermittency, privacy (see above), computation and memory costs, and formal communication costs, e.g. $O(poly(M, N_m, \\beta, ...))$ and how it compares with SotA algorithms. \n", "summary_of_the_paper": "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. ", "summary_of_the_review": "This paper proposes a practical algorithm for FL based on recent advances on the NTK, and FL-NTK. It provides interesting preliminary and promising results. However, there are several aspects related to clarifying intuitions and differentiating this paper from previous work that still need to be addressed. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "-T2HEGVHkny", "reviewer": "Reviewer_4JCT", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n\n- The paper provides a novel approach to FL using NTK paradaigm\n- The text is well written and clear to follow\n\nWeaknesses:\n\n- The notations can improve to make both the algorithm and the analysis more clear. It seems that $X_m$ and $Y_m$ are used for worker m's samples, and $X^{(k)}$ and $Y^{(k)}$ are used for the sample in round $k$. In the analysis and equation (11) $X$ is used without subscript/superscript. Could you clarify if $X^{(k)}$ and $Y^{(k)}$ are concatenations of $X_m$ and $Y_m$ at round $k$? What do X and y in Theorem 1 represent?  \n\n- I may be missing something here, but it seems that the closed form expressions for the trained network using NTK, for example, the ones given in equations (8), (10) and (11) of (Lee et al., 2019) also depend on the inverse kernel matrix. That seems to be standard in kernel methods in general. In fact, kernel methods are known to be computationally prohibitive due to the matrix inverse step. It is not clear to me why and how the closed form solution given in equation (8) does not depend on the kernel inverse. I augest to use clear references for all the results which are not proven in this paper, to make the paper self-contained. \nIt seems X in equation (8) is only the sample. Has this become possible because (8) does not actually represent a closed form solution, but rather only the output values at the sample?\n\n- While the novel framework is nicely contextualized through a literature review, the information that is uploaded by the workers in this framework violates the privacy preserving feature of FL. Noting that privacy preserving is the main motivation for the FL framework, this can be a problematic aspect of the novel framework. Although CP-NTK-FL is introduced to improve the privacy by compression and shuffling, without a clear discussion and characterization of privacy guarantee, the framework may not be effective. I think authors should include a clear discussion on this issue.  ", "summary_of_the_paper": "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.  ", "summary_of_the_review": "In summary, I appreciate the novel framework introduced in this paper for FL problem. Some of the results are however not completely clear to me. I also suggest a clear discussion on privacy issue that is central to the FL problem. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper proposes a novel Federated Learning (FL) framework that leverages the Neural Tangent Kernel (NTK), to replace the gradient-descent algorithm for optimization. Specifically, the workers upload the labels and the Jacobian matrices to the server, and the server uses the tools from the NTK to obtain a trained neural network. However since this could lead to increased communication cost and compromise of data privacy, the authors propose data sampling and random projection techniques to alleviate the problem. The authors provide a theoretical analysis that the proposed scheme has a faster convergence than FedAvg under specific assumptions, and experimentally validate that it significantly outperforms previous FL algorithms, achieving similar test accuracy to ideal centralized cases.

Pros
- The idea of using NTK for model optimization without gradient descent and use of it in the FL setting is both interesting and novel.
- The paper properly discusses and tackles the new challenges posed by the introduction of the new method.
- The paper is well-organized and clearly written, with sufficient discussion of related works and backgrounds. 

Cons

- The proposed method puts heavy computational burdens on the server-side.
- The method violates the privacy preserving feature of FL by its nature, and while the proposed compression shuffling alleviates the concern, more discussion is necessary.
- Missing comparison against popular baselines such as FedProx and SCAFFOLD. 
- The faster convergence of the proposed method in comparison to FedAvg depends on the learning rate and is not always true.
- There is a gap between the theory and practice, which makes the practicality of the algorithm still questionable. 

Although the reviewers found the idea as novel, the proposed techniques for alleviating communication cost and privacy concerns convincing, and considered both the theoretical analysis and experimental validation thorough, all reviewers leaned toward rejection due to critical concerns unanswered. During the discussion period, the authors alleviate many of the minor concerns from the reviewers, but there were still remaining concerns on the gap between the theory and practice on its convergence behavior, and insufficient discussion of the privacy-preserving feature of the proposed method, as well as shifting of computation burdens to the server. Thus, the reviewers reached a consensus that the paper is not yet ready for publication. 

Despite the low average score, the novelty of the idea and the quality of the paper is much higher than those of the accepted papers in my batch, and I strongly believe that this will become a high impact paper, if remaining concerns from the reviewers are properly resolved.
Decision: Reject