TCT: Convexifying Federated Learning using Bootstrapped Neural Tangent Kernels
Authors: Yaodong Yu
Alexander Wei
Sai Praneeth Karimireddy
Yi Ma
Michael Jordan
Venue: NeurIPS 2022 Accept (2022)
Keywords: Federated learning
Optimization
Decentralized learning
Distributed optimization

Abstract: State-of-the-art federated learning methods can perform far worse than their centralized counterparts when clients have dissimilar data distributions. For neural networks, even when centralized SGD easily finds a solution that is simultaneously performant for all clients, current federated optimization methods fail to converge to a comparable solution. We show that this performance disparity can largely be attributed to optimization challenges presented by nonconvexity. Specifically, we find that the early layers of the network do learn useful features, but the final layers fail to make use of them. That is, federated optimization applied to this non-convex problem distorts the learning of the final layers. Leveraging this observation, we propose a Train-Convexify-Train (TCT) procedure to sidestep this issue: first, learn features using off-the-shelf methods (e.g., FedAvg); then, optimize a convexified problem obtained from the network's empirical neural tangent kernel approximation. Our technique yields accuracy improvements of up to $+36\%$ on FMNIST and $+37\%$ on CIFAR10 when clients have dissimilar data. 

TCT: Convexifying Federated Learning using Bootstrapped Neural Tangent Kernels
Yaodong Yu UC Berkeley
Alexander Wei UC Berkeley
Sai Praneeth Karimireddy UC Berkeley
yyuêeecs. berkeley.ede edu
awei awei@berkeley.o edu edu
sp. karimmrreddybbarkeleyy edu
Yi Ma UC Berkeley
Michael L. Jordan UC Berkeley
yimaβeecs. berkeley.ede
jordanEcs. berkeley. edu
Abstract
State-of-the-art federated learning methods can perform far worse than their cen- tralized counterparts when clients have dissimilar data distributions. For neural networks, even when centralized SGD easily finds a solution that is simultaneously performant for all clients, current federated optimization methods fail to converge a comparable solution. We show that this performance disparity can largely be attributed to optimization challenges presented by nonconvexity. Specifically, we find that the early layers of the network do learn useful features, but the final layers fail to make use of them. That is, federated optimization applied to this non-convex problem distorts the learning of the final layers. Leveraging this observation, we propose Train-Conee Convexify-Truie (TCT) procedure to sidestep this issue: first, learn features using off-the-shelf methods (e.g.. FedAvg); then, optimize a convexified problem obtained from the network's empirical neural tangent kernel approxima- tion. Our technique yields accuracy improvements of up to +36% on FMNIST and +37% on CIFARIO 10 when clients have dissimilar data.
Introduction
Federated learning is a newly emerging paradigm for machine learning where multiple data holders (clients) collaborate to train a model on their combined dataset. Clients only share partially trained models and other statistics computed from their dataset, keeping their raw data local and private [53, 37]. By obviating the need for a third party to collect and store clients' data, federated learning has several advantages over the classical, centralized paradigm [14, 31, 23]: it ensures clients' consent is tied to the specific task at hand by requiring active participation of the clients training, confers some basic level of privacy, and has the potential to make machine learning more participatory in general [43, 36). Further, widespread legislation data portability and privacy requirements (such as GDPR and CCPA) might make federated learning a necessity [59].
Collaboration among clients most attractive when clients have very different subsets combined dataset (data heterogeneity). For example, different autonomous driving companies may only be able to collect data in weather conditions specific to their location, whereas their vehicles would need to function under all conditions. In such scenario, it would be mutually beneficial for companies in geographically diverse locations to collaborate and share data with each other. Further, in such settings, clients are physically separated and connected by ad-hoc networks with large latencies and limited bandwidth. This is especially true when clients are edge devices such as mobile phones, IoT sensors, etc. Thus, communication efficiency is crucial for practical federated learning. However, it is precisely under such circumstances (large data heterogeneity and low communication) that current
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

algorithms fail dramatically [27, 48, 39, 61, 71, 1, 46, 3, 72, etc.]. This motivates our central question: Why do current federated methods fail in the face of data heterogeneity- -nd how can we fix them?
Our solution. We make two main observations: (i) We show that, even with data heterogeneity, linear models can be trained in a federated manner through gradient correction techniques such as SCAFFOLD [39]. While this observation is promising, it alone remains limited, as linear models not rich enough to solve practical problems of interest (e.g.. those that require feature learning). (ii) We shed light on why current federated algorithms struggle to train deep, nonconvex models. We observe that the failure of existing methods for neural networks is not uniform across the layers. The early layers of the network do in fact learn useful features, but the final layers fail to make use of them. Specifically, federated optimization applied to this nonconvex problem results in distorted final layers.
These observations suggest truin-convexify-tree federated algorithm, which we call 7CT: first, use any off-the-shelf federated algorithm [such as FedAvg. 53] to train deep model to extract useful features; then, compute convex approximation of the deep model using its empirical Neural Tangent Kernel (cNTK) [34, 44, 20, 51, 75], and use gradient correction methods such as SCAFFOLD to train the final model. Effectively, the second-stage features freeze the features learned in the first stage and fit a linear model over them. We show that this simple strategy is highly performant on a variety of tasks and models- we obtain accuracy gains up to 36% points FMNIST with CNN, 37% points on CIFARIO with ResNet18 8-NGN, and 16% points on CIFAR with ResNet18 G-NGN. Further, its convergence remains unaffected even by extreme data heterogeneity. Finally, we show that given pre-trained model, our method completely closes the gap hetween centralized and federated methods.
2 Related Work
Federated learning. There are two main motivating scenarios for federated learning (FL). The first is where internet service companies (e.g., Google, Facebook, Apple, etc.) want to train machine learning models over their users' data, but do not want transmit raw personalized data away from user devices [60, 8]. This is the setting of cross-device federated learning and is characterized by an extremely large number of unreliable clients, each of whom has very little data and the collections of data are assumed to be homogeneous [37, 10, 38, 8J. The second motivating scenario is when valuable data split across different organizations, each of whom either protected by privacy regulation or is simply unwilling to share their raw data. Such "data islands" are common among hospital networks, financial institutions, autonomous vehicle companies, etc. This is known as cross-silo federated learning and is characterized by few highly reliable clients, who potentially have extremely diverse data. In this work, we focus on the latter scenario.
Metries in FI.. FL. research considers numerous metrics, such as fairness across users [55, 47, 62], formal security and privacy guarantees [9, 60, 21, 56), robustness to corrupted agents and corrupted training data [7, 64, 19, 40, 26), preventing backdoors at test time [6, 66, 69, 52], etc. While these concerns important, the main goal of FL (and our work) is to achieve high accuracy with minimal communication [53]. Clients are typically geographically separated yet need to communicate large deep learning models over unoptimized ad-hoc networks [37]. Finally, we focus the setting where all users are interested in training the same model over the combined dataset. This is in contrast to model-agnostic protocols [49, 58, 3] or personalized federated learning [16, 18. 78, 13, 42, 12]. Finally, we focus on minimizing the number of rounds required. Our approach can be combined with communication compression, which reduces bits sent per round [67, 4, 24, 65].
Federated optimization. Algorithms for FL proceed in rounds. In each round, the server sends model to the clients, who partially train this model using their local compute and data. The clients send these partially trained models back to the server who then aggregates them, finishing a round. FedAvg [53], which is the de facto standard FL algorithm, uses SGD to perform local updates on the clients and aggregates the client models by simply averaging their parameters. Unfortunately, however, FedAvg has been observed to perform poorly when faced with data heterogeneity across the clients [27. 48, 39, 61, 71, 1, 46, 3, 72, 17, etc.). Theoretical investigations of this phenomenon [39, 76] showed that this was a result of gradient heterogeneity across the clients. Consider FedAvg initialized with the globally optimal model. this model is not also optimal for each of the clients as well, the local updates will push it away from the global optimum. Thus, convergence would require a careful tuning of param-paraeeters To overcome this issue, SCAFFOLD [39] and FedDyn [1] propose to use control variates to correct for the biases of the individual clients akin to variance

---

0.8
Top-1 Accuracy
Top-1 Accuracy
0.7-
0.30-
0.5
0:25-
0.4
Feddvg Train to 0.1]
SCaTFoLD- 1.
Pedtvgg to 0.1)
SCAFFOLD-Tranda #
Fodkug-Test 0.3)
SCAFFOLD
Festuu--Tsst Tist 0.11
SCAFFOLD-Tcstte 0.3)
0.15
FedAug Train (&C 2)
Fectug-Train
SCATFOLD TAAroLL-TrannCCC-
Fodhng (eC
SCATOLL-Testee-22
0.1
Feckug-Test (AC-
SCAFFOLD TestreC
0.10
Communication Round
Communication Round
(a) Using linear model.
(b) Using ResNet-18.
Figure 1: Performance of FedAvg and SCAFFOLD on CIFARIO when data are split among ten
clients in two ways (#C=2 and a=0.1). The #C=2 split is more non-i.i.d. than the a=0.1 split. For convex problems (left), gradient correction methods such as SCAFFOLD are relatively unaffected by data heterogeneity. and consistently outperform FedAvg. However, for nonconvex problems (right), FedAvg and SCAFFOLD perform very similarly and both are strongly negatively affected by data heterogeneity.
reduction [35, 15]. This gradient correction applied every local update by the client and provably nullifies the effect of gradient heterogeneity [39, 54, 12]. However, we show here, such methods are insufficient to overcome high data heterogeneity especially for deep learning. Other, more heuristic approaches to combat gradient heterogeneity include using regularizer [48] and sophisticated server aggregation strategies such as momentum [28, 70, 50] or adaptivity [61, 38, 11].
A second line of work pins the blame on performance loss due to averaging nonconvex models. To overcome this, Singh and Jaggi [63], Yu et al. [81] propose learn mapping between the weights of the client models before averaging. Afonin and Karimireddy [3] advocates a functional perspective and replaces the averaging step with knowledge distillation, and Wang et [74], Li et al. [46], Tan et al. [68] attempt to align the internal representations of the client models. However, averaging is unlikely to be the only culprit since FedAvg does succeed under low heterogeneity, and averaging nonconvex models can lead to improved performance [33, 77].
Neural Tangent Kernels (NTK) and neural network linearization. NTK was first proposed to analyze the limiting behavior of infinitely wide networks [34, 44]. While NTK with MSE may be bad approximation of real-world finite networks in general [22]. it approximates the fine-tuning of a pre-trained network well [57], especially with some minor modifications [2]. That is, NTK cannot capture feature learning but does capture how a model utilizes learnt features better than last/mid layer activations.
3 The Effect of Nonconvexity
In this section, we investigate the poor performance of FedAvg [53] and SCAFFOLD [39] empirically in the setting of deep neural networks, focusing on image classification with a ResNet-18. To construct our federated learning setup, we split the CIFAR-10 dataset in a highly heterogeneous manner among clients. We either assign each client two classes (denoted by #C=2) or distribute samples according to Dirichlet distribution with a 0.1 (denoted by a=0.1). 1). For more details, see Section 5.1.
Insufficiency of gradient correction methods. Current theoretical work [e.g., 39, 61, 1, 73] attributes the slowdown from data heterogeneity the individual clients having varying local optima. If no single model is simultaneously optimal for all clients, then the updates of different clients can compete with and distort each other, leading to slow convergence. This tension is captured by the variance of the updates across the clients [client gradient heterogeneity, see 72]. Gradient correction methods such as SCAFFOLD [39] and FedDyn [1] explicitly correct for this and are provably unaffected by gradient heterogeneity for both convex and nonconvex losses.
3

---

Table 1: Feature learning by FedAvg. We report the test accuracy of ResNet-18 after (centralized) retraining of the last & layers on CIFARIO. The earlier (7 - f) layers are frozen to either random initialization or to the weights of FedAvg-trainee model. The difference measures utility of the E) layers learnt by FedAvg. The baseline FedAvg model without additional training gets 56.9% accuracy. We see that all layers of the FedAvg model contain useful information.
Layers retrained
Accuracy (%) Random init
Accuracy (%) FedAvg init
Improvement (%)
(FedAvg Random)
1/7 last layer
35.37
77.93
42.56
2/7 last layers
67.33
87.04
19.71
3/7 last layers
80.18
89.28
9.10
4/7 last layers
88.03
90.57
2.54
5/7 last layers
91.34
91.61
0.27
6/7 last layers
91.78
91.91
0.13
These theoretical predictions are aligned with the results of Figure I(a), where the loss landscape is convex: SCAFFOLD is relatively unaffected by the level of heterogeneity and consistently outperforms FedAvg. In particular, performance is largely dictated by the algorithm and not the data distributions. This shows that client gradient heterogeneity captures the difficulty of the problem well. On the other hand, when training a ResNet-18 model with nonconvex loss landscape, Figure ((b) shows that both FedAvg and SCAFFOLD suffer from data heterogeneity. This is despite the theory of gradient correction applying to both convex and nonconvex losses. Further, the train and test accuracies in Figure (b) match quite closely, suggesting that the failure lies in optimization (not fitting the training data) rather than generalization. Thus, while the current theory makes no qualitative distinctions between convex and nonconvex convergence, the practical behavior of algorithms in these settings is very different. Such differences between theoretical predictions and practical reality suggests that black-box notions such as gradient heterogeneity are insufficient for capturing the difficulty of training deep models.
Ease of feature learning. now now dive into how a ResNet- 18 trained with FedAvg (56.9% accuracy) differs from the centralized baseline (91.9% accuracy). We first apply linear probing to the FedAvg model (i.e., retraining with all but the output layer frozen). Note that this equivalent to (convex) logistic regression over the last-layer activations. This simple procedure produces striking jump from 56.9% 77.9% accuracy. Thus, of the 35% gap accuracy between the FedAvg and centralized models, 21% % may be attributed to a failure to optimize the linear output layer. We next extend this experiment towards probing the information content of other layers.
Given a FedAvg-trained model, we can use centralized training to retrain only the last layers while keeping the rest of the (7 f) layers (or ResNet blocks) frozen. We can also perform this procedure starting from a randomly initialized model. The performance difference between these two models can be attributed the information content of the frozen (7 E) layers of the FedAvg model. Table I summarizes the results of this experiment. The large difference in accuracy (up to 42.65 42.6%) indicates the initial layers of the FedAvg model have learned useful features. There continues to be a gap between the FedAvg features and random features in the earlier layers as well,' meaning that all layers of the FedAvg model learn useful features. We conjecture this is because from the perspective of earlier layers which perform simple edge detection, the tasks are independent of labels and the clients are i.i.d. However, the higher layers are more specialized and the effect of the heterogeneity is stronger.
Method
Based on the observations in Section 3, we propose train-convexify-traii (TCT) as a method for overcoming data heterogeneity when training deep models in a federated setting. Our high-level
'The significant decrease in the gap as we go down the layers may be because of the skip connections in the lower ResNet blocks which allow the random frozen layers to be sidestepped. This underestimates the true utility and information content in the earlier FedAvg layers.

---

intuition is that we want to leverage both the features learned from applying FedAvg to neural networks and the effectiveness of convex federated optimization. More specifically, we perform several rounds of "hootstrap" FedAvg to learn features before solving a convexified version of the original optimization problem.
4.1 Computing the Empirical Neural Tangent Kernel
To sidestep the challenges presented by nonconvexity, we describe how we approximate neural network by its "linearization." Given neural network (())0)) with weights 0, E R² mapping inputs x E R to R we replace it by its empirical neural tangent kernel (eNTK) approximation at Oo given by
0 00 fx: Bu),
f(x;0) f(x; 0o) + (0 Bo)
at each x DR Under this approximation, f(x; 0) is a linear function of the "feature vector" (f(1; 50 f(x; (o)) and the original nonconvex optimization problem becomes (convex) linear regression with respect to these features. Leveraging NTK for solving federated optimization problems has also been studied previous work [29, 82].
To reduce the computational burden working with the eNTK approximation, we make two further approximations: First, we randomly reinitialize the last layer of 9 and only consider 0 fa: 0b) with respect to a single output logit. Over the randomness of this reinitialization, E[f 80) = 0. Moreover, given the random reinitialization, all the output logits of f(x; are symmetric. These observations mean each data point x can be represented by a P-dimensional feature vector o fi (x:; fo), where f1( 0)) refers to the first output logit. Then, we apply dimensionality reduction by subsampling p random coordinates from this / P-ddimensional fcaturization. In our setting, this sub-sampling has the added benefit of reducing the number of bits communicated per round.
In summary, we transform our original (nonconvex) optimization problem over a neural network initialized at 0g into convex optimization problem in three steps: (i) reinitialize the last layer of do: (ii) for each data point æ, compute the gradient ΦeNTK (x: 0o) == as fi (æ: 8o); (iii) subsample the coordinates of δeNTK (2; fb) for each obtain reduced-diie fiiieeddddimmiionaliyy eNTK representation. Let
8: T + TRP denote this subsampling operation. Finally, we solve the resulting linear regression
problem over these eNTK representations."
4.2 Convexifying Federated Learning via eNTK Representations
The eNTK approximation lets us convexify the neural net optimization problem: following Section 4.1, we may extract (from a model trained with FedAvg) eNTK representations of inputs from each client. It remains to fit an overparameterized linear model using these eNTK features in federated manner. For ease of presentation, we denote the subsampled eNTK representation of input x by RP, where is the eNTK feature dimension after subsampling. We use 2 to represent the eNTK feature of the i-th sample from the k-th client. Then, for K the number of clients, Y the one-hot encoded labels, n& the number of data points of the k-th client, E Ng the number of data points across all clients, and p n& n, we can approximate the nonconvex neural net optimization problem by the convex linear regression problem
1
nc
min L(W)=
Σ Pk (W),,
where L&(W):
W
Y:*
(1)
Thk
To obtain the eNTK representation z of an input x, we take Op in Section 4.1 to be the weights of a model trained with FedAvg. As we will show Section 5, the convex reformulation in Eq. (1) significantly reduces the number of communication rounds needed to find an optimal solution.
2 For classification problems, we one-hor hot encoded labels and fit a linear model using squared loss. 'That such representations empirically have low effective dimension due fast eigenvalue decay [see, e.g.. 75] means that such random projection approximately preserves geometry the data points [5. 83). For all of our experiments, we set = 100, 000.
"Given fitted linear model with weights W p the prediction at is arg max, [WT 8(PeNTK(=)))
5

---

4.3 Train-Con Corin--onvexity-Trainn (TCT)
We now present our algorithm rainnconvvxiff-traa (TCT), with convexification done via the neural tangent kernel, for federated optimization.
TCT - train-convexify-train with eNTK representations
Extract eNTK features from a TrdAvg-traines model. FedAvg is first used train the model for T1 communication rounds. Let 8y, denote the model weights after these Ti rounds. Then, each client locally computes subsampled eNTK features, i.e., - = 8(gENTK (x7; (rr)) )) for l E [K] and E [mx]-
Stage 1:
Stage 2: Decentralized linear regression with gradient correction. Given samples {[2,, ) on each client k, first normalize the eNTK inputs all clients with a single communication round." " Then, solve the linear regression problem defined in Eq. (1) by SCAFFOLD with local learning rate and local steps M.
"For every feature in the eNTK representation, subtract the mean and scale to unit variance. "The detailed description of SCAFFOLD for solving linear regression problems can be found in Algorithm Appendix has the same communication and computation cost as FedAvg.
To motivate TCT, recall that in Section we found that FedAvg learns "useful" features despite its poor performance, especially in the earlier layers. By taking eNTK approximation, TCT optimizes a convex approximation while using information from all layers of the model. Empirically, we find that these extracted eNTK features significantly reduce the number of communication rounds needed to learn performant model, even with data heterogeneity.
Experiments
We now study the performance of TCT for the decentralized training of deep neural networks in the presence of data heterogeneity. We compare TCT to state-of-the-art federated learning algorithms three benchmark tasks in federated learning. For each task, apply these algorithms client data distributions with varying degrees data heterogeneity. We find that our proposed approach significantly outperforms existing algorithms when clients have highly heterogeneous data across all tasks. For additional experimental results and implementation details, see Appendix B. Our code is available at https: //gtttb....... com/ Somm//ooddngy//TT..
5.1 Experimental Setup
Datasets and degrees of data heterogeneity. We assess the performance of federated learning algorithms the image classification tasks FMNIST [80], CIFARIO, and CIFAR 100 [41]. FMNIST and CIFARIO each consist of 10 classes, while CIFAR 100 includes images from 100 classes. There 60,000 training images in FMNIST, and 50,000 training images in CIFAR10/100.
To vary the degree of data heterogeneity, we follow the setup of Li et al. [45]. We consider two types of non-i.i.d. data distribution: (i) Data heterogeneity sampled from a symmetric Dirichlet distribution with parameter [49, 71]. That is, we sample P. Dir x (a) from K K-dimensional symmetric Dirichlet distribution and assign p. -fracciio the class e samples to client k. (Smaller corresponds to more heterogeneity.) (ii) Clients get samples from fixed subset of classes [53]. That is, each client is allocated subset of classes; then, the samples of each class are split into non non-overlapping subsets and assigned to clients that were allocated this class. We use #C to denote the number of classes allocated to each client. For example, #C=2 means each client has samples from 2 classes. To allow for consistent comparisons, all of our experiments are run with 10 clients.
Models. For FMNIST, we use a convolutional neural network with ReLU activations consisting of two convolutional layers with max pooling followed by two fully connected layers (SimpleCNN), For CIFAR10 and CIFAR 100, we mainly consider an 18-layer residual network [25] with 4 basic residual blocks (ResNet-15) 18). In Appendix B.2, we present experimental results for other architectures.
Algorithms and training schemes. We compare TCT to state-of-the-art federated learning algo- rithms, focusing on the widely-used algorithms FedAvg [53], FedProx [48], and SCAFFOLD [39].

---

Table The top-1 accuracy (%) of our algorithm (TCT) vs. state-of-the tfrt---rrtt--tttt federated learning algo- rithms evaluated on FMNIST, CIFARIO, and CIFAR100. We vary the degree of data heterogeneity by controlling the a parameter of the symmetric Dirichlet distribution Dir (ax) and the #C parameter for assigning how many labels each client owns. Higher accuracy is better. The highest top-1 accuracy in each setting highlighted in bold.
Datasets
Architectures
Methods
Non-i..d. degree
#C 1
#C
0.1
0.5
FedAvg
35.10%
85.18%
86.18%
90.09%
FMNIST
SimpleCNN
FedProx
50.04%
84.91 %
86.31%
89.77%
SCAFFOLD
12.80%
42.80%
83.87%
89.40%
TCT
86.32%
90.33%
90.78%
91.13%
Centralized
91. 11.40%
#C =
#C =
0.1
a 0.5
FedAvg
11.27%
56.86%
82.60%
90.43%
CIFAR-10
ResNet-18
FedProx
12.30%
56.87%
83.31%
90.68%
SCAFFOLD
10.00%
46.75%
80.46%
90.72%
TCT
49.92%
83.02%
89.21 %
91.10%
Centralized
91.90%
0.001
0.01
0.1
0 0.5
FedAvg 53.89%
54.22%
63.49%
67.65%
CIFAR-100
ResNet-18
FedProx
52.87%
54.32%
63.47%
67.54%
SCAFFOLD
49.86%
54.07%
65.67%
71.07% %
TCT
68.42%
69.07%
69.66%
69.68%
Centralized
73.61%
Act
Train
Test
Accurscy
Top-1 Accuracy
2-001
Trin Acc
Test
Communication Round
Communication Round
Communication Round
FedAvg.
(b) SCAFFOLD.
(c) TCT.
Figure 2: Training/test accuracy communication round for FedAvg (left), SCAFFOLD (middle), and our algorithm TCT (right) on the CIFAR100 dataset with various degrees of non-iid-ness (Dirx(a) (a) with E (0.1, 0.01, 0.001)). Dotted lines represent the training accuracy, and dashdot lines with markers represent the test accuracy.
(For comparisons additional algorithms, see Appendix B.1.) Each client uses SGD with weight decay 10-5 and batch size 64 by default. For each baseline method, we run it for 200 total communi- cation rounds using 5 local training epochs with local learning rate selected from (0.1, 0.01, 0.001) by grid search. For TCT, we run 100 rounds of FedAvg in Stage following the above and use 100 communication rounds in Stage 2 with M 500 local steps and local learning rate n 5 10
5.2
Main Results
Table 2 displays the top-1 accuracy of all algorithm the three tasks with varying degrees of data heterogeneity. We evaluated each algorithms on each task under four degrees data heterogeneity. Smaller #C and in Table 2 correspond to higher heterogeneity.
7

---

Trariin Accuracy 1
Top Traning Accuracy
Top Traning
Teddep
Tackig
Tedlup
(Mo1)000
1030)
2
Communication Round
Communication Round
Communicaiin Round
(a) FMNIST.
(b) CIFARIO.
(c) CIFARI00.
Figure 3: Training accuracy vs. communication round for full batch gradient descent (GD) and TCT
on FMNIST-[#C=2] (a), CIFARIO-[HC=2] (b), and CIFAR100-(a= 0.01] (c). Each dotted line with square markers represents the training accuracy of GD with some learning rate. Dashed lines with circle markers represent the training accuracy of TCT with different numbers of local steps. We also include the training accuracy results of FedAvg with learning rate 0.1. We use TCT-GD to denote the variant of TCT which replaces SCAFFÓLD with GD in Stage 2.
We find that the existing federated algorithms all suffer when data heterogeneity high across all three tasks. For example, the top-1 accuracy of FedAvg on CIFAR-10 is 56.86% when #C=2, which much worse than the 90.43% achieved in a more homogeneous setting (e.g. a = 0.5). In contrast, TCT achieves consistently strong performance, even in the face of high data heterogeneity. More specifically, TCT achieves the best top-1 accuracy performance across all settings except CIFAR-100 with 0.5, where TCT does only slightly worse than SCAFFOLD.
In absolute terms, we find that TCT is not affected much by data heterogeneity, with performance dropping by less than 1.5% on CIFAR 100 as goes from 0.5 to 0.001. Moreover, our algorithm improves over existing methods by at least 15% in the challenging cases, including FMNIST with 0.01 and a 0.001. And, perhaps surprisingly, our algorithm still performs relatively well in the extreme non-i.i.d. setting where each client sees only single class.
#C-1, CIFAR-10 10 with #C=1 and #C=2, and CIFAR- 100 with a
Figure 2 compares the performances of FedAvg, SCAFFOLD, and TCT in more detail on CIFARI 100 dataset with different degrees of data heterogeneity. We consider the Dirichlet distribution with parameter a E {0.1, 0.01, 0.001} } and compare the training and test accuracy of these three algorithms. As shown in Figures 2(a) and 2(b), both FedAvg and SCAFFOLD struggle when data heterogeneity high: for both algorithms, test accuracy drops significantly when decreases. In contrast, we see from Figure 2(c) that TCT maintains almost the same test accuracy for different Furthermore, the same set of default parameters for our algorithm, including local learning rate and the number of local steps, is relatively robust to different levels of data heterogeneity.
5.3 Communication Efficiency
To understand the effectiveness of the local steps in our algorithm, we compare SCAFFOLD (used in TCT-Stage 2) to full batch gradient descent (GD) applied to the overparameterized linear regression problem in Stage 2 of TCT on these datasets. For our algorithm, we set local steps M E (10?, 10²) and use the default local learning rate. For full batch GD, we vary the learning rate from 10-5 5 to 10-1 and visualize the ones that do not diverge.
The results are summarized in Figure 3. Each dotted line with square markers in Figure 3 corresponds to full batch GD with some learning rate. Across all three datasets, our proposed algorithm consistently outperforms full batch GD. Mcanwhile, we find that more local steps for our algorithms lead faster convergence across all settings. In particular, our algorithm converges within 20 communication rounds on CIFAR100 (as shown in Figure 3(c)). These results suggest that our proposed algorithm can largely leverage the local computation and improve communication efficiency.
5.4 Ablations
Gradient correction. We investigate the role of gradient correction when solving overparameterized linear regression with eNTK features in TCT. We compare SCAFFOLD (used TCT) to FedAvg on solving the regression problems and summarize the results in Figure 4. We use the default
8

---

Recurccy
Rocumny
Accuracy
ICTesshi
SChfecligg (P-332) SCtidecleg (M-222201
ICTiostigg pleiin
Icreasiig 14-1000
TCTraskog (4-1000
24-20
SCT (2-32391
TCT (4-100))
(a) FMNIST.
(b) CIFAR1O.
(c) CIFARI0O.
Figure Comparing TCT TO TCT-FedAvg for solving the overparameterizee linear regression problem on (a) FMNIST-(#C-2), (b) CIFARIO-[#C-2], and (c) CIFAR100-[a 0.01]. We use TCT-FedAvg denote variant of TCT that uses FedAvg instead of SCAFFOLD to perform linear regression in TCT-Stage 2. Dotted red lines with square markers represent the training accuracy of TCT-Fecl ICT-FedAvg with different numbers of local steps. Dashed blue red lines with circle markers represent the training accuracy of TCT with different numbers of local steps. A darker color means more local steps.
16-
Fedhug Train Acc
[oo-1 Training Accuracy
Fedhwg Acc
Top-1 Accuracy
TCT-laett Ace
Round T3 =
(random initl
Round
Round T2 40
Round
FedAug-Train Acc
MM
Round 100
Round
TCT (woN) Train Acc
03
Communication
Communication Round
(a) Effect of FedAvg communication rounds in Stage
(b) Effect of normalization.
Figure 5: (a). We evaluate TCT on using checkpoints save at different communication rounds Ti in Stage 1. Ti = 0 corresponds the randmon initialized model weights scenario (without FedAvg training). Dash lines with square markers represent the training accuracy, and dotted lines with circle makers represent the test accuracy. (b). We study the effect of pre-conditioning on TCT. TCT (wN) corresponds to the setting where eNTK features are normalized, and TCT (woN) corresponds to the without normalization step setting.
local learning and consider three different numbers of local steps for both algorithms, i.e., M E (10, 100, 1000). As shown Figure 4, our approach largely outperforms FedAvg when the number of local steps large (M 100) across three datasets. We also find that the performance of FedAvg can even degrade when the number of local steps increases. For example. FedAvg with M 1000 performs the worst across all three datasets. In contrast FedAvg, SCAFFOLD converges faster when the number of local steps increases. These observations highlight the importance of gradient correction in our algorithm.
Model weights for computing eNTK features. To understand the impact of the model weights trained in Stage of TCT, we evaluate TCT run with different Ti parameters. We consider Ti E {0, 20, 40, 60, 80, 1000. where Ti = corresponds to randomly initialized weights. From Figure 5()), we find that weights after FedAvg training are much more effective than weights at random initialization. Specifically, without FedAvg training, the eNTK (at random initialization) performs worse than standard FedAvg. In contrast, TCT significantly outperforms FedAvg by large margin (roughly 20% in test accuracy) when eNTK features are extracted from BedAvg-traing model. Also, we find that TCT stable with respect to the choice of communication rounds T, in Stage 1. For example, models trained by TCT with Ti Z 60 achieve similar performance.
Effect normalization. In Figure 5(b), we investigate the role of normalization on TCT by comparing TCT run with normalized and unnormalized eNTK features. The same number of local

---

steps (M = 500) is applied for both settings. We tune the learning rate y for each setting and plot the run that performs best (as measured in training accuracy). The results in Figure 5(b) suggest that the normalization step in TCT significantly improves the communication efficiency by increasing convergence speed. In particular, TCT with normalization converges to nearly 100% training accuracy in approximately 40 communication rounds, which is much faster than TCT without normalization.
Pre-training vs. Bootstrapping. In Appendix B.4, we explore the effect of starting from a pre- trained model instead of relying on bootstrapping to learn the features. We find that pre-training further improves the performance of TCT and completely erases the gap between centralized and federated learning.
Additionally, we conduct experiments on investigating the role of training loss function and subsam- pling approximation in TCT-Stage 2. For TCT-Stage we find that neither using the cross-entropy loss as the training objective nor applying full eNTK representations significantly improves the performance of TCT. On the other hand, applying subsampling approximation in TCT-Stage 2 can largely improve the communication efficiency compared to the full eNTK representations approach. See Appendix B.7 for detailed experimental results.
6 Conclusion
We have argued that nonconvexity poses significant challenge for federated learning algorithms. We found that a neural network trained in such a manner does learn useful features, but fails to use them and thus has poor overall accuracy. To sidestep this issue, we proposed train-convexify-tree procedure: first, train the neural network using FedAvg; then, optimize (using SCAFFOLD) a convex approximation of the model obtained using its empirical neural tangent kernel. We showed that the first stage extracts meaningful features, whereas the second stage learns to utilize these features to obtain a highly performant model. The resulting algorithm is significantly faster and more stable to hyper-pare parameters than previous federated learning methods. Finally, we also showed that given a good pre- pre-pretrained feature extractor, our convexify-tree train procedure fully closes the gap between centralized and federated learning.
Our algorithm adds to the growing body of work using eNTK to linearize neural networks and obtain tractable convex approximations. However, unlike most of these past works which only work with pre-trained models, our bootstrapping allows training models from scratch. Finally, we stress that the success of our approach underscores the need to revisit theoretical understanding heterogencous federated learning. Nonconvexity seems to play an outsized role but its effect in FL has hitherto been unexplored. In particular, black-box notions of difficulty such as gradient dissimilarity or distances between client optima seem insufficient to capture practical performance. It is likely that further progress in the field (e.g. federated pre- pre-training of foundational models), will require tackling the issue of nonconvexity head on.
Acknowledgments and Disclosure of Funding
We would like to thank the anonymous reviewers for their constructive suggestions and comments. Yaodong Yu acknowledges support from the joint Simons Foundation-No DMS grant #2031899, Alexander Wei acknowledges support from an NSF Graduate Research Fellowship under grant DGE2146652.. Sai Praneeth Karimireddy acknowledges support of an SNSF postdoc mobility fellowship. Yi Ma acknowledges support from ONR grants N00014-20-1- 2102 and the joint Simons Foundation-NSF DMS grant #2031899. Michael Jordan acknowledges support of the ONR Mathematical Data Science program.
and N00014-22-i-
References
[1] Durmus Alp Emre Acar, Yue Zhao, Ramon Matas, Matthew Mattina, Paul Whatmough, and Venkatesh Saligrama. Federated learning based dynamic regularization. In International Conference on Learning Representations, 2021. URL https:///penreeven wwwpeerrriiieeeeeererueua net/forum? id-B7v40MR629w.
10

---

[2] Alessandro Achille, Aditya Golatkar, Avinash Ravichandran, Marzia Polito, and Stefano Soatto. Lqf: Linear quadratic fine-tuning. In Proceedings of the IEEE/CVE Conference on Computer Vision and Puttern Recognition, pages 15729-15739, 2021.
[3] Andrei Afonin and Sai Praneeth Karimireddy. Towards model agnostic federated learning using knowledge distillation. arXiv preprint arXiv:21 10. 152/0, 2021.
[4] Dan Alistarh, Demjan Grubic, Jerry Li, Ryota Tomioka, and Milan Vojnovic. QSGD: ceatmiicaatio-fffiiienn SGD via gradient quantization and encoding. Advances in Neural Information Processing Systems, 30, 2017.
[5] Haim Avron, Kenneth L. Clarkson, and David P. Woodruff. Sharper bounds for regularized data fitting. Approximation, Randomization, and Combinatorial Optimization. Algorithms and Techniques.
[6] Eugene Bagdasaryan, Andreas Veit, Yiqing Hua, Deborah Estrin, and Vitaly Shmatikov. How to backdoor federated learning. In International Conference on Artificial Intelligence and Statistics, pages 2938-294442 PMLR, 2020.
[7] Peva Blanchard, El Mahdi El Mhamdi, Rachid Guerraoui, and Julien Stainer. Machine learn- ing with adversaries: Byzantine tolerant gradient descent. Advances Neural Information Processing Systems, 30, 2017.
[8] Kallista Bonawitz, Peter Kairouz, Brendan McMahan, and Daniel Ramage. Federated learning and privacy: Building preseviig systems for machine learning and data science on decentralized data. Quene, 19(5):87-114, 2021.
[9] Keith Bonawitz, Vladimir Ivanov, Ben Kreuter, Antonio Marcedone, H Brendan McMahan, Sarvar Patel, Daniel Ramage, Aaron Segal, and Karn Seth. Practical secure aggregation for privacy-preserving machine learning. In Proceedings of the 2017 ACM SIGSAC Conference on Computer and Communications Security, pages 1175-1191, 2017.
[10] Keith Bonawitz, Hubert Eichner, Wolfgang Grieskamp, Dzmitry Huba, Alex Ingerman, Vladimir Ivanov, Chloe Kiddon, Jakub Konečný, Stefano Mazzoechi, Brendan McMahan, et al. Towards federated learning at scale: System design. Proceedings of Machine Learning and Systems, 1: 374-388, 2019.
[11] Zachary Charles, Zachary Garrett, Zhouyuan Huo, Sergei Shmulyian, and Virginia Smith. On large-cohort training for federated learning. Advances in Neural Information Processing Systems, 34, 2021.
[12] El Mahdi Chayti and Sai Praneeth Karimireddy. Optimization with access to auxiliary informa- tion. arXiv preprini arXiv:2206.00395. 2022.
[13] Liam Collins, Hamed Hassani, Aryan Mokhtari, and Sanjay Shakkottai. Exploiting shared representations for personalized federated learning. In International Conference Machine Learning, pages 2089-2099. PMLR, 2021.
[14] Jeffrey Dean, Greg Corrado, Rajat Monga, Kai Chen, Matthieu Devin, Mark Mao, Marc aurelao Ranzato, Andrew Senior, Paul Tucker, Ke Yang, et al. Large scale distributed deep networks. Advances in Neural Information Processing Systems, 25, 2012.
[15] Aaron Defazio, Francis Bach, and Simon Lacoste-Julien. Saga: fast incremental gradient method with support for non-strongly convex composite objectives. Advances in Neural Information Processing Systems, 27, 2014.
[16] Yuyang Deng, Mohammad Mahdi Kamani, and Mehrdad Mahdavi. Adaptive personalized federated learning. arXiv preprint arXiv:2003. 13461, 2020.
[17] Jean Ogier du Terrail, Samy- Samy-Safwan Ayed, Edwige Cyffers, Felix Grimberg, Chaoyang He, Regis Loeb, Paul Mangold, Tanguy Marchand, Othmane Marfoq, Erum Mushtaq, et al. Flamby: Datasets and benchmarks for cross-silo federated learning in realistic settings. 2022.

---

[18] Alireza Fallah, Aryan Mokhtari, and Asuman Ozdaglar. Personalized federated learning: A meta-learning approach. arXiv preprint arXiv:2002.07948, 2020.
[19] Minghong Fang, Xiaoyu Cao, Jinyuan Jia, and Neil Gong. Local model poisoning attacks to Blyyzantine-oobstt federated learning- In 29th USENIX Security Symposium (USENIX Security 20), pages 1605-1622, 2020.
[20] Stanislav Fort, Gintare Karolina Dziugaite, Mansheej Paul, Sepideh Kharaghani, Daniel M Roy, and Surya Ganguli. Deep learning versus kernel learning: an empirical study of loss landscape geometry and the time evolution of the neural tangent kernel. In Advances in Neural Information Processing Systems, volume 33, pages 5850-5861, 2020.
[21] Clement Fung, Chris JM Yoon, and Ivan Beschastnikh. Mitigating sybils in federated learning poisoning. arXiv preprint arXiv: 008.8888..0886 2018.
[22] Micah Goldblum, Jonas Geiping, Avi Schwarzschild, Michael Moeller, and Tom Goldstein. Truth backpropaganda? empirical investigation of deep learning theory. arXiv preprint arXiv:1910.00359, 2019.
[23] Priya Goyal, Piotr Dollár, Ross Girshick, Pieter Noordhuis, Lukasz Wesolowski. Aapo Kyrola, Andrew Tulloch, Yangqing Jia, and Kaiming He. Accurate, large minibatch SGD: Training Imagenet in 1 hour. arXiv preprint arXiv: (Xxi::706.00677 2017.
[24] Farzin Haddadpour, Mohammad Mahdi Kamani, Aryan Mokhtari, and Mehrdad Mahdavi. Federated learning with compression: Unified analysis and sharp guarantees. In International Conference on Artificial Intelligence and Statistics, pages 2350-2356 PMLR, 2021.
[25] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770-778, 2016.
[26] Lie He, Sai Praneeth Karimireddy, and Martin Jaggi. eerantine-robus. decentralized learning via self-centered clipping. arXiv preprint arXiv:2202.01544 2022.
[27] Kevin Hsieh, Amar Phanishayee, Onur Mutlu, and Phillip Gibbons. The non-iid data quagmire of decentralized machine learning. In International Conference on Machine Learning, pages 4387-4398. PMLR, 2020.
[28] 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.060 06335, 2019.
[29] Baihe Huang, Xiaoxiao Li, Zhao Song, and Xin Yang. Fl-ntk: A neural tangent kernel- based framework for federated learning analysis. International Conference on Machine Learning. pages 4423 4434. PMLR, 2021.
[30] Like Hui and Mikhail Belkin. Evaluation of neural architectures trained with square loss vs cross-entropy in classification tasks. arXiv preprint 2xXi:20006003322, 2020.
[31] Forrest N landola, Matthew W Moskewicz, Khalid Ashraf, and Kurt Keutzer. Firecaffe: near- linear acceleration of deep neural network training on compute clusters. In Proceedings of the IEEE Conference Computer Vision and Pattern Recognition, pages 2592-2600, 2016.
[32] Sergey Ioffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In International conference on machine learning, pages 448-456. PMLR, 2015.
[33] Pavel Izmailov, Dmitrii Podoprikhin, Timur Garipov, Dmitry Vetrov, and Andrew Gordon Wilson. Averaging weights leads to wider optima and better generalization. arXiv preprint 180. 2018.
[34] Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. Advances Neural Information Processing Systems, 31, 2018.
12

---

[35] Rie Johnson and Tong Zhang. Accelerating stochastic gradient descent using predictive variance reduction. Advances in Neural Information Processing Systems, 26, 2013.
[36] Charles I Jones and Christopher Tonetti. Nonrivalry and the economics of data. American Economic Review, 10(())2819-55)) 2020.
[37] Peter Kairouz, H. Brendan McMahan, et al. Advances and open problems in federated learning. Foundations and Trends Machine Learning, 14(1-2):1-210) 2021.
[38] Sai Praneeth Karimireddy, Martin Jaggi, Satyen Kale, Mehryar Mohri, Sashank J Reddi, Sebastian U Stich, and Ananda Theertha Suresh. Mime: Mimicking centralized stochastic algorithms in federated learning. arXiv preprint arXiv:2008. 2020.
[39] 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 on Machine Learning, pages 5132-5143. PMLR, 2020.
[40] Sai Praneeth Karimireddy, Lie He, and Martin Jaggi. Byzantine-robust learning hetero- geneous datasets via bucketing. In International Conference on Learning Representations. 2021.
[41] Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
[42] Viraj Kulkarni, Milind Kulkarni, and Aniruddha Pant. Survey of personalization techniques for federated learning. 2020 Fourth World Conference on Smart Trends in Systems, Security and Sustainability (WorldS4), pages 794-77 797. IEEE, 2020.
[43] Bogdan Kulynych, David Madras, Smitha Milli, Inioluwa Deborah Raji, Angela Zhou, and Richard Zemel. Participatory approaches to machine learning. International Conference on Machine Learning Workshop, 2020.
[44] Jachoon 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. Advances in Neural Information Processing Systems, 32, 2019.
[45] Qinbin Li, Yiqun Diao, Quan Chen, and Bingsheng He. Federated learning on non-iid data silos: An experimental study. arXiv preprint arXiv:2102.02079, 2021.
[46] Qinbin Li, Bingsheng He, and Dawn Song. Model-contrastive federated learning. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pages 10713-10722, 2021.
[47] Tian Li, Maziar Sanjabi, Ahmad Beirami, and Virginia Smith. Fair resource allocation in federated learning, arXiv preprint arXiv: 19099999997 2019.
[48] Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith. Federated optimization in heterogeneous networks. Proceedings of Machine Learning and Systems, 2:429 450, 2020.
[49] Tao Lin, Lingjing Kong, Sebastian U Stich, and Martin Jaggi. Ensemble distillation for robust model fusion federated learning. Advances in Neural Information Processing Systems, 33: 2351-2363, 2020.
[50] Tao Lin, Sai Praneeth Karimireddy, Sebastian U Stich, and Martin Jaggi. Quasi-global mo- mentum: Accelerating decentralized deep learning on heterogeneous data. arXiv preprint arXiv:2102.04761 2021.
[51] Philip M Long. Properties of the after kernel. arXiv preprint arXiv:2105. 10585, 2021.
[52] Lingjuan Lyu, Han Yu, and Qiang Yang. Threats to federated learning: survey. arXiv preprint arXiv:2003.02133. 2020.
13

---

[53] Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. Communication-effiee efficient learning of deep networks from decentralized data, In Artificial Intelligence and Statistics, pages 1273-1282. PMLR, 2017.
[54] Konstantin Mishchenko, Grigory Malinovsky, Sebastian Stich, and Peter Richtárik, Proxskip: Yes! local gradient steps provably lead to communication acceleration! finally! arXiv preprint arXiv:2202.093:. 09355. 2022.
[55] Mehryar Mohri, Gary Sivek, and Ananda Theertha Suresh. Agnostic federated learning. In International Conference on Machine Learning, pages 4615-4625. PMIR, 2019.
[56] Viraaji Mothukuri, Reza M Parizi, Seyedamin Pouriyeh, Yan Huang, Ali Dehghantanha, and Gautam Srivastava. survey on security and privacy of federated learning. Future Generation Computer Systems, 1 115:619-640, 2021.
[57] Fangzhou Mu, Yingyu Liang, and Yin Li. Gradients as features for deep representation learning. arXiv preprint arXiv:2004.05529, 2020.
[58] Kaan Ozkara, Navjot Singh, Deepesh Data, and Suhas Diggavi. Quped: Quantized personaliza- tion via distillation with applications to federated learning. Advances in Neural Information Processing Systems, 34, 2021.
[59] Alex Pentland, Alexander Lipton, and Thomas Hardjono. Building the New Economy: Data as Capital. MIT Press, 2021.
[60] Swaroop Ramaswamy, Om Thakkar, Rajiv Mathews, Galen Andrew, H Brendan McMahan, and Françoise Beaufays. Training production language models without memorizing user data. arXiv preprint arXiv:2009.10033 2020.
[61] Sashank 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. URL https: //tterev///eeeevieew wwtpeerrrviee..ne/foormrm net/forum? id=LkFG31B13U5.
[62] Yuxin Shi, Han Yu, and Cyril Leung. A survey of fairness-aware federated learning, arXiv preprint arXiv:211 0182 2021.
[63] Sidak Pal Singh and Martin Jaggi. Model fusion via optimal transport. Advances in Neural Information Processing Systems, 33:22045-22055, 2020.
[64] Jinhyun So, Başak Güler, and A Salman Avestimehr.
Byzantine-resilient secure federated learning. IEEE Journal on Selected Areas in Communications, 3977)22188-2111, 2020.
[65] Sebastian U Stich and Sai Praneeth Karimireddy. The error-feedback framework: Better rates for sgd with delayed gradients and compressed updates. Journal of Machine Learning Research, 21:1-36, 2020.
[66] Ziteng Sun, Peter Kairouz, Ananda Theertha Suresh, and Brendan McMahan. Can you really backdoor federated learning? arXiv preprint 1XXiv:9911.0766: 2019.
[67] Ananda Theertha Suresh, X Yu Felix, Sanjiv Kumar, and H Brendan McMahan. Distributed mean estimation with limited communication. In International Conference on Machine Learn- ing. pages 3329-3337 PMLR, 2017.
[68] Yue Tan, Guodong Long, Lu Liu, Tianyi Zhou, Qinghua Lu, Jing Jiang, and Chengqi Zhang. Fedproto: Federated prototype learning over heterogeneous devices. arXiv preprint 105.0223, 2021.
[69] Hongyi Wang, Kartik Sreenivasan, Shashank Rajput, Harit Vishwakarma, Saurabh Agarwal, Jy-yong Sohn, Kangwook Lee, and Dimitris Papailiopoulos. Attack of the tails: Yes, you really can backdoor federated learning. Advances in Neurul Information Processing Systems, 33: 16070-16084, 2020.
14

---

[70] Jianyu Wang, Vinayak Tantia, Nicolas Ballas, and Michael Rabbat. proving communication-effin efficciet distributed sgd with slow momentum. arXiv preprint arXiv: 1XXiv:9910000443 2019.
Slowmo:
Im-
[71] Jianyu Wang, Qinghua Liu, Hao Liang, Gauri Joshi, and H Vincent Poor. Tackling the objective inconsistency problem heterogeneous federated optimization. Advances in Neural Information Processing Systems, 33:7611-7622 2020.
[72] 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 federated optimization. arXiv preprini arXiv:2107.06917 2021.
[73] Jianyu Wang, Rudrajit Das, Gauri Joshi, Satyen Kale, Zheng Xu, and Tong Zhang. On the unreasonable effectiveness of federated averaging with heterogeneous data. arXiv preprint arXiv:2206.00 04723, 2022.
[74] Shiqiang Wang, Tiffany Tuor, Theodoros Salonidis, Kin K Leung, Christian Makaya, Ting He, and Kevin Chan. Adaptive federated learning in resource constrained edge computing systems. IEEE Journal on Selected Areas in Communications, 37(6))1205-122 2019.
[75] Alexander Wei, Wei Hu, and Jacob Steinhardt. More than a toy: Random matrix models predict how real-world neural representations generalize. arXiv preprint arXiv:2203.06/77 2022.
[76] Blake Woodworth, Kumar Kshitij Patel, and Nati Srebro. Minibatch local sgd for heteroge- neous distributed learning. Advances in Neural Information Processing Systems, 33:6281-6292
2020.
[77] Mitchell Wortsman, Gabriel Ilharco, Samir Yitzhak Gadre, Rebecca Roelofs, Raphael Gontijo- Lopes, Ari S Morcos, Hongseok Namkoong, Ali Farhadi, Yair Carmon, Simon Kornblith, et al. Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time. arXiv preprint arXiv:2203.00448 2022.
[78] Qiong Wu, Kaiwen He, and Xu Chen. Personalized federated learning for intelligent loT applications: A cloud-edge based framework. IEEE Open Journal of the Computer Society, 1: 35-44. 2020.
[79] Yuxin Wu and Kaiming He. Group normalization. In Proceedings of the European Conference on Computer Vision (ECCV), pages 3-19, 2018.
[80] Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-MNIST: novel image dataset for benchmarking machine learning algorithms. arXiv preprint arXiv: 1708.077477 2017.
[81] Fuxun Yu, Weishan Zhang, Zhuwei Qin, Zirui Xu, Di Wang, Chenchen Liu, Zhi Tian, and Xiang Chen. Fed2: Feature-aligned federated learning. In Proceedings of the 27th ACM SIGKDD Conference on Knowledge Discovery & Data Mining, pages 2066-2074, 2021.
[82] Kai Yue, Richeng Jin, Ryan Pilgrim, Chau-Was Wai Wong, Dror Baron, and Huaiyu Dai. Neural tangent kernel empowered federated learning. In International Conference Machine Learning, pages 25783-25803. PMLR, 2022.
[83] Luca Zancato, Alessandro Achille, Avinash Ravichandran, Rahul Bhotika, and Stefano Soatto. Predicting training time without training. Advances in Neural Information Processing Systems, 186136-6146, 2020.
15

---

Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope? [Yes]
(b) Did you describe the limitations of your work? [Yes]
(c) Did you discuss any potential negative societal impacts of your work? [Yes] Our work furthers the goal of federated learning which aims make ML more participatory. We see this overall as positive outcome. However, it can potentially also be used to circumvent existing data portability and privacy regulations. E.g. hospitals with previously confidential patient health records could circumvent privacy regulations by using FL. They can sell their information to an insurance company, without directly revealing any raw patient data. We believe preventing such bad outcomes requires direct regulatory measures (e.g. prohibiting insurance companies from discriminating among their consumers) are necessary rather than relying on vague notions of privacy. More generally, laws and mechanisms need to be designed to ensure that companies don't t capture all the benefits arising from the access data through FL, and instead these benefits are passed on to all sections of society.
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [N/A]
(b) Did you include complete proofs of all theoretical results? [N/A]
3. If you ran experiments..
(a) Did you include the code, data, and instructions needed to reproduce the main ex- perimental results (either in the supplemental material URL)? [Yes] See supplemental material.
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Yes] See Section 5. and Appendix A.
(c) Did you report error bars (e.g., with respect to the random seed after running experi- ments multiple times)? [Yes]
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)? [Yes]
4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators? [Yes]
(b) Did you mention the license of the assets? [N/A]
(c) Did you include any new assets either in the supplemental material or as URL? [No]
(d) Did you discuss whether and how consent was obtained from people whose data you're rou're using/curating? [N/A]
(e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [N/A]
5. If you used crowdsourcing conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if applicable? [N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable? [N/A]
(c) Did you include the estimated hourly wage paid participants and the total amount spent on participant compensation? [N/A]
16


Reviews:
[{"review_id": "_vC1nFDhB8d", "reviewer": "Reviewer_rWVN", "summary": "This paper focuses on difficulty introduced by non-convexity and data heterogeneity in federated learning.\nAuthors first show that, with data heterogeneity, linear models and convex optimization problem can be trained efficiently with gradient correction techniques such as SCAFFOLD, while non convex problem cannot. \nThen, in order to sidestep the non-convexity for neural network, authors substitute original model with a linear approximation and original loss with a quadratic loss, such that non-convex optimization problem turns into a convex regression problem based on NTK. \nSince feature learning is necessary for NTK align with data to provide good result, authors introduce two-stage training, where feature learning only happens in first stage, and convex optimization with gradient correction happens in second stage. \nAuthors conduct experiments to show that such two-stage optimization can fit the data faster and produce better test accuracy.", "questions": "If I understand correctly, the final training result for BookNTK is a neural network plus a linear model. Any inference requires a backward propagation to derive the feature for empirical NTK to apply a linear model. This is different from common workflow. Is there a way to compile the final network and the linear model into a single model that support inference with a single forward pass?\n\nSection 3 mentioned \"the train and test accuracies in Figure 1(b) match quite closely, suggesting that the failure lies in optimization\". However Figure 2 shows a different picture, where train acc and test acc have a large gap. What's the reason behind this discrepancy?\n\nWhy quadratic loss instead of cross-entropy is used in (1)?\n\nHow much worse is BookNTK than centralized training? Authors only mentioned one centralized baseline in section 3, where they show that \"21% (out of the 35% gap in accuracy) may be attributed to a failure to optimize the linear output layer\". However, for BookNTK, none of results in section 5 and appendix contains centralized training. I think this question is as important as \"How much better is BookNTK than FedAvg\". I strongly suggest authors to add centralized training result for a comparison.", "limitations": "Yes.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strength:\n\nThis paper introduced a novel two-stage scheme to combine the feature learning capacity of neural network and for efficient optimization linear model.\n\nWeakness:\n\nThe convexified problem is introduced mainly due to an optimization consideration: SCAFFOLD perform well on convex method. However, the convex formulation clearly sacrifice the model capacity. The feature learning can only happen in first stage (or non-convex part), thus BookNTK learns less feature than centralized training, which doesn't involve two-stage training.\n\n", "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "YRDX79Rnge", "reviewer": "Reviewer_UZ2q", "summary": "This paper makes two main contributions: 1) empirically observing that the early layers of neural networks trained by FedAvg learn useful features even in heterogeneous data settings, while the performance degradation due to data heterogeneity is due to ineffective later layers, and 2) based on this observation, proposing a novel two-stage cross-silo federated learning method, BooNTK, that first uses FedAvg to learn early-layer features, applies a particular transformation based on these features to all data points, then learns a linear classifier on top of the transformed data using SCAFFOLD. \n\nTo verify that FedAvg learns useful early-layer features, FedAvg is run on a heterogeneous image dataset, then the last $\\ell$-many layers are retrained in a centralized manner while the earlier layers are held fixed. It is observed that using the FedAvg-pretrained early layers leads to huge performance improvement over using random weights for the early layers when the last layer is retrained centrally. Also, there is some improvement due to using the FedAvg-pretrained vs random weights when the last $\\ell$-many layers are retrained centrally for all $\\ell$, suggesting that the early layers learned by FedAvg are extracting useful information. This observation inspires the first stage of the proposed method: FedAvg pretraining to learn the early layer weights. The paper also observes that SCAFFOLD is robust to data heterogeneity when the client losses are convex (but is not robust when the losses are nonconvex). This inspires the second stage of the proposed method: applying SCAFFOLD to convex losses to learn the last linear layer of the network. In particular, the second stage approach consists of computing the neural tangent kernel (NTK) representation of each data point using the pretrained weights in the NTK computation. Then, SCAFFOLD is employed to solve a multi-output linear regression with the transformed data points as inputs and the one-hot encoded labels as the target vectors. Empirical results are provided showing large improvement in training and testing accuracy of the proposed method over FedAvg, FedProx and SCAFFOLD on CIFAR100, FEMNIST, and MNIST with very heterogeneous partitions.\n", "questions": "Why is the proposed approach limited to the cross-silo setting, with a small number of clients? It seems to me that it should also work for settings with many clients. Or, if there are many clients and fewer samples per client, does locally optimizing the high-dimensional linear regression diverge?\n", "limitations": "The authors should discuss the computational complexity of the proposed method.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths\n\n- The paper makes interesting empirical observations that are relevant to NeurIPS and may inspire future work. In particular the observation that FedAvg learns useful features even in data heterogeneous settings is novel and helps to explain the empirical success of FedAvg plus fine-tuning.\n\n- These observations motivate a new federated learning method with promising empirical performance.\n\n- The experimental evaluation is mostly thorough, with multiple datasets tested and helpful ablations.\n\n- The writing is clear.\n\nWeaknesses\n\n1. Not enough intuition or empirical evidence is provided for why the second stage of the algorithm should consist of linear regression on the NTK-transformed data points rather than simply fixing the first L-1 layers and running lear regression to learn the paramters of the last layer with MSE loss, or learn it with multi-class logistic regression and cross entropy loss as is conventional. Both of the latter approaches maintain convexity of the loss functions, are much simpler to implement, and involve an optimization over far fewer parameters (presumably the output of the last layer has dimension far less than p=100,000 as in the NTK approach). These approaches should be compared against as baselines and intuition should be provided as to why they are not used. \n\n2. On a similar note, results on the computational cost of the proposed method should be included. It seems to be much larger than the baselines due to mapping every data point to a high dimension via parameter derivative computation.\n\n3. Personalized FL has been shown to be an effective alternative to learning a single global model in data heterogeneous settings. As such, some personalized FL methods should also be compared against. Also, the experimental results would be strengthened by comparison against more recent FL methods that learn a single global model, e.g. [23,63]. Comparing with only 3 baselines is low for an empirical paper.\n\nMinor notes\n- Missing related works: Huang et al., 2021 and Yue  et al., 2021 employ the NTK for FL.\n\n- Intuition on why it makes sense to run linear regression for classification problems rather than logistic regression would be helpful.\n\n- Since BootNTK gets to pretrain using FedAvg for T_1 rounds, fair comparison with the other methods should allow them to train for an extra T_1 rounds.\n\n-Cross-device FL may still have data heterogeneity\n\n- SCAFFOLD has higher communication and computational cost than fedavg by constant factor due to computing and communicating the gradient correction terms (footnote b).\n\n- “By taking an eNTK approximation, BooNTK optimizes a convex approximation while using information from all layers of the model.” - besides the last layer\n\nHuang et al., FL-NTK: A Neural Tangent Kernel-based Framework for Federated Learning Convergence Analysis, https://arxiv.org/pdf/2105.05001.pdf, 2021.\nYue et al., Neural Tangent Kernel Empowered Federated Learning, https://arxiv.org/pdf/2110.03681.pdf, 2021.\n\n", "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "CoYcAC3Ml3y", "reviewer": "Reviewer_7oaK", "summary": "This work proposes BooNTK; a two stage approach for FL that allows for better performance in cross-silo settings. The main idea is to 1) perform standard FedAvg training on a non-convex model, such as a neural network, then 2) approximate it with a first order Taylor approximation around the parameters found and finally 3) optimize the parameters of this linear model with federated training using optimisers with gradient correction, such as SCAFFOLD. The authors motivate this through the optimization difficulties of non-convex models in the non-iid setting; empirically, the loss of performance when moving to non-idd data is larger for non-convex model compared to convex.\n\nThe authors describe several further approximations that they do to make the method more efficient (i.e., removing the bias term of the Taylor approximation and only considering a subset of the coordinates of the gradient vector) and then demonstrate BooNTK’s performance on three label skew non-iid settings on FMNIST, CIFAR10 and CIFAR100.\n", "questions": "Most of my questions revolve around the weaknesses described above.\n\n1. The authors do not take into account that the linear model is less flexible by design, therefore it is harder to fit the non-iid peculiarities (and could explain its improved performance). I believe a control experiment would highlight whether this can be an issue in practice. For example, one could do a heavily regularised non-convex model to see if the gap is shrinking in Figure 1. How would, e.g., the performance of a pipeline similar to BooNTK fare if instead of the linear model fitting step (i.e., stage 2), one just switched from FedAvg to FedProx with a strong regularisation strength for stage 2?\n\n2. I believe that the effect of various sources of non-iidness on BooNTK should be investigated, as the experiments consider only label skew (where generally adjusting the classifier only is sufficient). What happens when there is, e.g., covariate shift? Intuitively, this should affect the earlier layers more.\n\n3. I believe some more investigation on the eNTK part of BooNTK is required. For example, how does finetuning just the classification layer (while keeping the rest of the network frozen) with SCAFFOLD work, relative to the eNTK approach (which requires further approximations)? Furthermore, how does BooNTK perform without the approximations to the eNTK (i.e., do the approximations have beneficial regularisation effects or are they detrimental for performance)?\n\n4. The authors consider classification tasks, however for the second stage of BooNTK, they consider an MSE loss on an one-hot representation of the targets. What is the motivation of the MSE loss? Intuitively, you should be able to apply a softmax on the linear model of the first eq. at sec. 4.1 to get a logistic regression model, which is more appropriate for classification. \n", "limitations": "The authors could have spent a bit more time discussing potentially negative aspects of their method (e.g., the linearity of the model in the second stage).", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths\n- Good results on the cross-silo setting\n    - There is significant improvement by the BooNTK pipeline on all tasks involving a varying amount of label skew.\n- Intuitive ablations and method exposition\n    - The method is explained / motivated well and I liked the layer importance investigation. The ablation studies are also useful and highlight the sensitivity of the method on the choice of (some) hyper parameters.\n- Simple method\n    - The method is quite simple and straightforward. As a bonus, the second step is also computationally more efficient than training the original neural network.\n\nWeaknesses:\n- No discussion about how different non-iid ness settings affect the method\n    - Given the claims of the work about the negative effects of data heterogeneity in the non-convex setting, I would have expected that the authors experimented with more diverse non-iid settings (i.e., not just label skew). As the current label-skew experiments concern mostly different marginals over the labels, $p(y)$, at each client, a better adjusted classification layer is important. I believe this is one of the reasons that BooNTK is effective in these scenarios. However, not all non-iidness is just label skew; for example, you could consider an image recognition system running in different mobile phones. As each phone comes with its own camera sensor, covariate shift (i.e., different $p(x)$ across clients) can also manifest as a source of non-iidness. In this case, I would intuitively expect that the distribution of the features also differs among clients, therefore, adjusting just the classifier might not be enough.\n- No discussion about how the linear model is less flexible \n    - The linear model is less flexible than the non-convex neural network. As a result, its improved performance (on a given feature set) could be just because it has less degrees of freedom to adapt to the non-i.i.d. peculiarities. I believe a discussion around this is missing.\n- There are several approximations involved in the computation of the eNTK features, the impact of which is unclear\n    - In section 4.1 the authors describe a series of approximations to eNTK in order to reduce the computational burden, however there is no (empirical) evaluation on how each one affects the final performance. \n", "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept