TCT: Convexifying Federated Learning using Bootstrapped Neural Tangent Kernels
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 ﬁnds 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 .
Speciﬁcally, we ﬁnd that the early layers of the network do learn useful features, but the ﬁnal layers fail to make use of them.
That is, federated optimization applied to this non-convex problem distorts the learning of the ﬁnal layers.
Leveraging this observation, we propose a T rain-
C onvexify- T rain (TCT) procedure to sidestep this issue: ﬁrst, learn features using off-the-shelf methods (e.g., FedAvg); then, optimize a convexiﬁed 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.
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 speciﬁc task at hand by requiring active participation of the clients in training, confers some basic level of privacy, and has the potential to make machine learning more participatory in general [43, 36].
Further, widespread legislation of data portability and privacy requirements (such as
GDPR and CCPA) might even make federated learning a necessity [59].
Collaboration among clients is most attractive when clients have very different subsets of the combined dataset ( data heterogeneity ).
For example, different autonomous driving companies may only be able to collect data in weather conditions speciﬁc to their location, whereas their vehicles would need to function under all conditions.
In such a scenario, it would be mutually beneﬁcial 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 efﬁciency 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 (NeurIPS 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—and how can we ﬁx 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 are 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 ﬁnal layers fail to make use of them.
Speciﬁcally, federated optimization applied to this nonconvex problem results in distorted ﬁnal layers.
These observations suggest a train-convexify-train federated algorithm, which we call TCT : ﬁrst, use any off-the-shelf federated algorithm [such as FedAvg, 53] to train a deep model to extract useful features; then, compute a convex approximation of the deep model using its empirical Neural Tangent
Kernel (eNTK) [34, 44, 20, 51, 75], and use gradient correction methods such as SCAFFOLD to train the ﬁnal model.
Effectively, the second-stage features freeze the features learned in the ﬁrst stage and ﬁt 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 on FMNIST with a CNN, 37 % points on CIFAR10 with ResNet18-GN, and 16 % points on CIFAR100 with ResNet18-GN.
Further, its convergence remains unaffected even by extreme data heterogeneity.
Finally, we show that given a pre-trained model, our method completely closes the gap between centralized and federated methods.
Related Work
Federated learning.
There are two main motivating scenarios for federated learning (FL).
The ﬁrst 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 to 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, 8].
The second motivating scenario is when valuable data is split across different organizations, each of whom is either protected by privacy regulation or is simply unwilling to share their raw data.
Such “data islands” are common among hospital networks, ﬁnancial institutions, autonomous-vehicle companies, etc.
This is known as cross-silo federated learning and is characterized by a few highly reliable clients, who potentially have extremely diverse data.
In this work, we focus on the latter scenario.
Metrics in FL.
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 are 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 on 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 a 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, ﬁnishing 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.
If 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 hyper-parameters.
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 [Figure] reduction [35, 15].
This gradient correction is applied in every local update by the client and provably nulliﬁes the effect of gradient heterogeneity [39, 54, 12].
However, as we show here, such methods are insufﬁcient to overcome high data heterogeneity especially for deep learning.
Other, more heuristic approaches to combat gradient heterogeneity include using a 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 to learn a 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 al. [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 ﬁrst proposed to analyze the limiting behavior of inﬁnitely wide networks [34, 44].
While NTK with MSE may be a bad approximation of real-world ﬁnite networks in general [22], it approximates the ﬁne-tuning of a pre-trained network well [57], especially with some minor modiﬁcations [2].
That is, NTK cannot capture feature learning but does capture how a model utilizes learnt features better than last/mid layer activations.
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 classiﬁcation with a ResNet-18.
To construct our federated learning setup, we split the CIFAR-10 dataset in a highly heterogeneous manner among ten clients.
We either assign each client two classes (denoted by #C=2 ) or distribute samples according to a Dirichlet distribution with ↵ = 0 .
1 (denoted by ↵ =0.1).
For more details, see
Section 5.1.
Insufﬁciency of gradient correction methods.
Current theoretical work [e.g., 39, 61, 1, 73] attributes the slowdown from data heterogeneity to 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. [Table]
These theoretical predictions are aligned with the results of Figure 1(a) [Table] 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 difﬁculty of the problem well.
On the other hand [Table] when training a ResNet-18 model with nonconvex loss landscape, Figure 1(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 [Table] the train and test accuracies in Figure 1(b) match quite closely, suggesting that the failure lies in optimization (not ﬁtting 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 insufﬁcient for capturing the difﬁculty of training deep models.
Ease of feature learning.
We now dive into how a ResNet-18 trained with FedAvg ( 56 . 9 % accuracy) differs from the centralized baseline ( 91 . 9 % accuracy).
We ﬁrst apply linear probing to the FedAvg model (i.e., retraining with all but the output layer frozen).
Note that this is equivalent to (convex) logistic regression over the last-layer activations.
This simple procedure produces a striking jump from 56 .
9 % to 77 .
9 % accuracy.
Thus, of the 35 % gap in 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  ` ) 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 to the information content of the frozen (7  ` ) layers of the FedAvg model.
Tablesummarizes the results of this experiment.
The large difference in accuracy (up to 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, 1 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-train (TCT) as a method for overcoming data heterogeneity when training deep models in a federated setting.
Our high-level 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 speciﬁcally, we perform several rounds of “ bootstrap ” FedAvg to learn features before solving a convexiﬁed 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 a neural network by its “linearization.”
Given a neural network f ( · ; ✓ 0 ) with weights ✓ 0 2 R P mapping inputs x 2 R D to R C , we replace it by its empirical neural tangent kernel (eNTK) approximation at ✓ 0 given by
[Equation]
at each x 2 R D .
Under this approximation, f ( x ; ✓ ) is a linear function of the “feature vector” ( f ( x ; ✓ 0 ) , @@✓ f ( x ; ✓ 0 )) and the original nonconvex optimization problem becomes (convex) linear regression with respect to these features.
2 Leveraging NTK for solving federated optimization problems has also been studied in previous work [29, 82].
To reduce the computational burden of working with the eNTK approximation, we make two further approximations: First, we randomly reinitialize the last layer of ✓ 0 and only consider @@✓ f ( x ; ✓ 0 ) with respect to a single output logit.
Over the randomness of this reinitialization, E [ f ( x ; ✓ 0 )] = 0 .
Moreover, given the random reinitialization, all the output logits of f ( x ; ✓ 0 ) are symmetric.
These observations mean each data point x can be represented by a P -dimensional feature vector @ @ ✓ f 1 ( x ; ✓ 0 ) , where f 1 ( · ; ✓ 0 ) refers to the ﬁrst output logit.
Then, we apply a dimensionality reduction by subsampling p random coordinates from this P -dimensional featurization.
In our setting, this sub-sampling has the added beneﬁt of reducing the number of bits communicated per round.
In summary, we transform our original (nonconvex) optimization problem over a neural network initialized at ✓ 0 into a convex optimization problem in three steps: (i) reinitialize the last layer of ✓ 0 ; (ii) for each data point x , compute the gradient  eNTK ( x ; ✓ 0 ) :
= @@✓ f 1 ( x ; ✓ 0 ) ; (iii) subsample the coordinates of  eNTK ( x ; ✓ 0 ) for each x to obtain a reduced-dimensionality eNTK representation.
Let S : R P !
R p 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 ﬁt an overparameterized linear model using these eNTK features in a federated manner.
For ease of presentation, we denote the subsampled eNTK representation of input x by z 2 R p , where p is the eNTK feature dimension after subsampling.
We use z ki to represent the eNTK feature of the i -th sample from the k -th client.
Then, for K the number of clients, Y ki the one-hot encoded labels, n k the number of data points of the k -th client, n :=
P k
[ K ] n k
the number of data points across all clients, and p k := n k /n , we can approximate the nonconvex neural net optimization problem by the convex linear regression problem
[Equation]
To obtain the eNTK representation z of an input x , we take ✓ 0 in Section 4.1 to be the weights of a model trained with FedAvg.
As we will show in Section 5, the convex reformulation in Eq.
(1) signiﬁcantly reduces the number of communication rounds needed to ﬁnd an optimal solution
Given a ﬁtted linear model with weights W 2 R p ⇥
C , the prediction at x is arg max
j [ W > S (  eNTK ( x ))]
j .
4.3 Train-Convexify-Train (TCT)
We now present our algorithm train-convexify-train (TCT), with convexiﬁcation done via the neural tangent kernel, for federated optimization.
TCT — train-convexify-train with eNTK representations
• Stage 1: Extract eNTK features from a FedAvg-trained model.
FedAvg is ﬁrst used to train the model for T 1 communication rounds.
Let ✓ T 1 denote the model weights after these T 1 rounds.
Then, each client locally computes subsampled eNTK features,
i.e., z ki = S (  eNTK ( x ki ; ✓ T 1 )) for k 2 [ K ] and i 2 [ n k ] .
• Stage
: Decentralized linear regression with gradient correction.
Given samples { ( z ki , Y ki ) } n k i =1 on each client k , ﬁrst normalize the eNTK inputs of all clients with a single communication round.
a Then, solve the linear regression problem deﬁned in Eq.
(1) by SCAFFOLD with local learning rate ⌘ and local steps M .
b
To motivate TCT, recall that in Section 3 we found that FedAvg learns “useful” features despite its poor performance, especially in the earlier layers.
By taking an eNTK approximation, TCT optimizes a convex approximation while using information from all layers of the model.
Empirically, we ﬁnd that these extracted eNTK features signiﬁcantly reduce the number of communication rounds needed to learn a 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 on three benchmark tasks in federated learning.
For each task, we apply these algorithms on client data distributions with varying degrees of data heterogeneity.
We ﬁnd that our proposed approach signiﬁcantly 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://github.com/yaodongyu/TCT .
5.1 Experimental Setup
Datasets and degrees of data heterogeneity.
We assess the performance of federated learning algorithms on the image classiﬁcation tasks FMNIST [80], CIFAR10, and CIFAR100 [41].
FMNIST and CIFAR10 each consist of 10 classes, while CIFAR100 includes images from 100 classes.
There are 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.
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 c
⇠ Dir K ( ↵ ) from a K -dimensional symmetric Dirichlet distribution and assign a p kc -fraction of the class c samples to client k .
(Smaller ↵ corresponds to more heterogeneity.)
(ii) Clients get samples from a ﬁxed subset of classes [53].
That is, each client is allocated a subset of classes; then, the samples of each class are split into 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 CIFAR100, we mainly consider an 18-layer residual network [25] with 4 basic residual blocks (ResNet-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 algorithms, focusing on the widely-used algorithms FedAvg [53], FedProx [48], and SCAFFOLD [39]. [Table] [Figure] (For comparisons to 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 communication 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 1 following the above and usecommunication rounds in Stage 2 with M = 500 local steps and local learning rate ⌘ = 5 · 10  5 .
5.2 Main Results
Table 2 displays the top-1 accuracy of all algorithm on the three tasks with varying degrees of data heterogeneity.
We evaluated each algorithms on each task under four degrees of data heterogeneity.
Smaller #C and ↵ in Table 2 correspond to higher heterogeneity. [Figure]
We ﬁnd that the existing federated algorithms all suffer when data heterogeneity is high across all three tasks.
For example, the top-1 accuracy of FedAvg on CIFAR-10 is 56 .
when #C=2 , which is much worse than the 90 .
43% achieved in a more homogeneous setting (e.g. ↵ = 0 . 5 )
In contrast,
TCT achieves consistently strong performance, even in the face of high data heterogeneity.
More speciﬁcally, TCT achieves the best top-1 accuracy performance across all settings except CIFARwith ↵ = 0 .
5 , where TCT does only slightly worse than SCAFFOLD.
In absolute terms, we ﬁnd that TCT is not affected much by data heterogeneity, with performance dropping by less than 1 . 5% on CIFAR100 as ↵ goes from 0 . 5 to 0 .
Moreover, our algorithm improves over existing methods by at least 15% in the challenging cases, including FMNIST with #C=1 , CIFAR-10 with #C=1 and #C=2 , and CIFAR-100 with ↵ = 0 .
01 and ↵ = 0 .
And, perhaps surprisingly, our algorithm still performs relatively well in the extreme non-i.i.d.
setting where each client sees only a single class.
Figure 2 compares the performances of FedAvg, SCAFFOLD, and TCT in more detail on CIFAR dataset with different degrees of data heterogeneity.
We consider the Dirichlet distribution with parameter ↵ 2 { 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 is high: for both algorithms, test accuracy drops signiﬁcantly 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 Efﬁciency
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 2 { 10 2 , 10 3 } and use the default local learning rate.
For full batch GD, we vary the learning rate from 10  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.
Meanwhile, we ﬁnd that more local steps for our algorithms lead to 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 efﬁciency.
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 in TCT) to FedAvg on solving the regression problems and summarize the results in Figure 4.
We use the default [Figure] [Figure] local learning rate and consider three different numbers of local steps for both algorithms, i e.,
M 2 { 10 , 100 , 1000 } .
As shown in Figure 4, our approach largely outperforms FedAvg when the number of local steps is large ( M  100 ) across three datasets.
We also ﬁnd 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 to 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 1 of TCT, we evaluate TCT run with different T 1 parameters.
We consider
T 1 2 { 0 , 20 , 40 , 60 , 80 , 100 } , where T 1 = 0 corresponds to randomly initialized weights.
From
Figure 5(a), we ﬁnd that weights after FedAvg training are much more effective than weights at random initialization.
Speciﬁcally, without FedAvg training, the eNTK (at random initialization) performs worse than standard FedAvg.
In contrast, TCT signiﬁcantly outperforms FedAvg by a large margin (roughly 20% in test accuracy) when eNTK features are extracted from a FedAvg-trained model.
Also, we ﬁnd that TCT is stable with respect to the choice of communication rounds T 1 in Stage 1.
For example, models trained by TCT with T 1  60 achieve similar performance.
Effect of 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 ⌘ 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 signiﬁcantly improves the communication efﬁciency 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 pretrained model instead of relying on bootstrapping to learn the features.
We ﬁnd 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 subsampling approximation in TCT-Stage 2. For TCT-Stage 2, we ﬁnd that neither using the cross-entropy loss as the training objective nor applying full eNTK representations signiﬁcantly improves the performance of TCT. On the other hand, applying subsampling approximation in TCT-Stage 2 can largely improve the communication efﬁciency compared to the full eNTK representations approach.
See Appendix B.7 for detailed experimental results.
We have argued that nonconvexity poses a signiﬁcant 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 a train-convexify-train procedure: ﬁrst, 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 ﬁrst stage extracts meaningful features, whereas the second stage learns to utilize these features to obtain a highly performant model.
The resulting algorithm is signiﬁcantly faster and more stable to hyper-parameters than previous federated learning methods.
Finally, we also showed that given a good pre-pretrained feature extractor, our convexify-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 of heterogeneous federated learning.
Nonconvexity seems to play an outsized role but its effect in FL has hitherto been unexplored. In particular, black-box notions of difﬁculty such as gradient dissimilarity or distances between client optima seem insufﬁcient to capture practical performance. It is likely that further progress in the ﬁeld (e.g. federated 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-NSF DMS grant #2031899. Alexander Wei acknowledges support from an NSF Graduate Research Fellowship under grant DGE2146752.
Sai Praneeth Karimireddy acknowledges support of an SNSF postdoc mobility fellowship.
Yi Ma acknowledges support from ONR grants N00014-20-1-2002 and N00014-22-12102 and the joint Simons Foundation-NSF DMS grant #2031899.
Michael Jordan acknowledges support of the ONR Mathematical Data Science program.