BooNTK: Convexifying Federated Learning using Bootstrapped Neural Tangent Kernels
Anonymous Author(s)
AfﬁliationAddress email
Abstract
State-of-the-art federated learning methods can perform far worse than their centralized counterparts when clients have dissimilar data distributions.
For neuralnetworks, even when centralized SGD easily ﬁnds a solution that is simultaneouslyperformant for all clients, current federated optimization methods fail to convergeto a comparable solution.
We show that this performance disparity can largely beattributed to optimization challenges presented by non-convexity .
Speciﬁcally, weﬁnd that the early layers of the network do learn useful features, but the ﬁnal layersfail to make use of them.
That is, federated optimization applied to this non-convexproblem distorts the learning of the ﬁnal layers.
Leveraging this observation, wepropose a simple, two-stage training procedure
Bootstrapped eNTK (BooNTK) 1 tosidestep 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 neuraltangent kernel approximation.
Our technique yields accuracy improvements of upto +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 trainedmodels and other statistics about their dataset, keeping their raw data local and private [46, 32].
Byobviating the need for a third party to collect and store clients’ data, federated learning has severaladvantages over the classical, centralized paradigm [13, 27, 22]: it ensures clients’ consent is tied tothe speciﬁc task at hand by requiring active participation of the clients in training, confers some basiclevel of privacy and has the potential to make machine learning more participatory in general [38, 31].
Further, widespread legislation of data portability and privacy requirements (such as GDPR and
CCPA) might even make federated learning a necessity [51].
Collaboration among clients is most attractive when clients have very different subsets of the combineddataset ( data heterogeneity ).
For example, different autonomous driving companies may only be ableto collect data in weather conditions speciﬁc to their location, whereas their vehicles would needto function under all conditions.
In such a scenario, it would be mutually beneﬁcial for companiesin geographically diverse locations to collaborate and share data with each other.
Further, in suchsettings, clients are physically separated and connected by ad-hoc networks with large latencies andlimited bandwidth.
This is especially true when clients are edge devices such as mobile phones, IoTsensors, etc.
Thus, communication efﬁciency is crucial for practical federated learning.
However, it isprecisely under such circumstances (large data heterogeneity and low communication) that currentalgorithms fail dramatically [25, 43, 34, 53, 63, 1, 41, 2, 64, etc.].
This motivates our central question:
Why do current federated methods fail in the face of data heterogeneity—and how can we ﬁx them?
Pronounced boon-T-K
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 [34].
While this observation is promising, it alone remains limited, as linear modelsare 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, non-convex models.
Weobserve that the failure of existing methods for neural networks is not uniform across the layers.
Theearly layers of the network do in fact learn useful features, but the ﬁnal layers fail to make use ofthem.
Speciﬁcally, federated optimization applied to this non-convex problem results in distortedﬁnal layers.
These observations suggest a two-stage federated training algorithm, which we call BooNTK : ﬁrst,use any off-the-shelf federated algorithm (such as FedAvg [46]) to train a deep model to extract usefulfeatures; then, compute a convex approximation of the deep model using its empirical Neural Tangent
Kernel (eNTK) [29, 39, 20, 66], and use gradient correction methods such as SCAFFOLD [34] totrain the ﬁnal model.
Effectively, the second stage features freezes the features learnt in the ﬁrst stageand learns a linear model over them.
We show that this simple strategy is highly performant on avariety of tasks and models—we see 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.
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 machinelearning models over their users’ data, but do not want to transmit raw personalizable data away fromuser devices [52, 7].
This is the setting of cross-device federated learning and is characterized by anextremely large number of unreliable clients, each of whom has very little but homogeneous data [32,9, 33, 7].
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 vehiclecompanies, etc.
This is known as cross-silo federated learning and is characterized by a few highlyreliable clients, who potentially have extremely diverse data.
In this work, we focus on the latter.
Metrics in FL.
FL research considers numerous metrics, such as fairness across users [48, 42, 54],formal security and privacy guarantees [8, 52, 21, 49], robustness to corrupted agents and dataduring training [6, 56, 19, 35], preventing backdoors at test time [5, 58, 61, 45], etc.
While theseconcerns are important, the main goal of FL (and our work) is to achieve high accuracy with minimalcommunication [46].
Clients are typically geographically separated yet need to communicate largedeep learning models over unoptimized ad-hoc networks [32].
Furthermore, we focus on the settingwhere all users are interested in training the same model over the combined dataset.
This is in contrastto model-agnostic protocols [44, 50, 2] or personalized federated learning [15, 18, 69, 11, 37].
Finally, we focus on minimizing the number of rounds required.
Our approach can be combined withcommunication compression, which reduces bits sent per round [59, 3, 23, 57].
Federated optimization.
Algorithms for FL proceed in rounds.
In each round, the server sends amodel to the clients, who partially train this model using their local compute and data.
The clientssend these partially trained models back to the server who then aggregates them, ﬁnishing a round.
FedAvg [46], which is the de-facto standard FL algorithm, uses SGD to perform local updates on theclients and aggregates the client models by simply averaging their parameters.
However, FedAvg wasobserved to perform terribly when faced with data heterogeneity across the clients [25, 43, 34, 53, 63,1, 41, 2, 64, etc.].
Theoretical investigations of this phenomenon [34, 67] hypothesized that this wasa result of gradient heterogeneity across the clients.
Consider FedAvg initialized with the globallyoptimal model.
If this model is not also optimal for each of the clients as well, the local updateswill push it away from the global optimum.
Thus, convergence would require a careful tuning ofhyper-parameters.
To overcome this issue, SCAFFOLD [34] and FedDyn [1] propose to use controlvariates to correct for the biases of the individual clients akin to variance reduction [30, 14].
Thisgradient correction is applied in every local update by the client and provably nulliﬁes the effectof gradient heterogeneity [34, 47].
However, as we show here, such methods seem insufﬁcient toovercome high data heterogeneity especially for deep learning.
Other, more ad-hoc approaches tocombat gradient heterogeneity include using a regularizer FedProx [43] and sophisticated server
[Figure]
aggregation strategies such as momentum [26, 62] or adaptivity [53, 10].
A second line of work pinsthe blame on performance loss due to averaging non-convex models.
To overcome this, [55, 71]propose to learn a mapping between the weights of the client models before averaging, [2] advocatesa functional perspective and replaces the averaging step with knowledge distillation, and [65, 41, 60]attempt to align the internal representations of the client models.
However, averaging is unlikely tobe the only culprit since FedAvg does succeed under low heterogeneity, and averaging non-convexmodels can lead to improved performance [28, 68].
The Effect of Non-convexity
In this section, we empirically investigate the poor performance of FedAvg [46] and SCAFFOLD [34]when applied to deep neural networks, focusing on image classiﬁcation with a ResNet-18.
Toconstruct our federated learning setup, we split the CIFAR-10 dataset in a highly heterogeneousmanner among 10 clients.
We either assign each client two classes (denoted by #C=2 ) or distributesamples 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.
Prior theoretical work (e.g. [34, 53, 63, 1, 64])attributes the slowdown from data heterogeneity to the individual clients having wildly varying local
[Figure]
optima.
If no single model is simultaneously optimal for all clients, then the updates of differentclients can compete with and distort each other, leading to slow convergence.
This tension is capturedby the variance of the updates across the clients (a.k.a.
client gradient heterogeneity [64]).
Gradientcorrection methods such as SCAFFOLD [34] and FedDyn [1] explicitly correct for this and areprovably unaffected by gradient heterogeneity for both convex and non-convex losses.
These theoretical predictions are perfectly aligned with the results of Figure 1(a), where the loss landscape is convex: SCAFFOLD is relatively unaffected by the level of heterogeneity and consistentlyoutperforms FedAvg.
In particular, performance is largely dictated by the algorithm and not the datadistributions.
This shows that client gradient heterogeneity captures the difﬁculty of the problem well.
On the other hand, when training a ResNet-18 model with non-convex loss landscape, Figure 1(b)shows that both FedAvg and SCAFFOLD suffer from data heterogeneity.
This is despite the theoryof gradient correction applying to both convex and non-convex losses [34].
Further, the train andtest accuracies in Figure 1(b) match quite closely, suggesting that the failure lies in optimization (notﬁtting the training data) rather than generalization.
The difference between theoretical predictions andpractical reality in the non-convex setting suggests black-box notions such as gradient heterogeneityare insufﬁcient for capturing the difﬁculty of training deep models. [Table]
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 FedAvgmodel (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 from56 .
9 % to 77 .
9 % accuracy.
Thus, of the 35 % gap in accuracy between the FedAvg and centralizedmodels, 21 % may be attributed to a failure to optimize the linear output layer.
We next extend thisexperiment towards probing the information content of other layers.
Given a FedAvg-trained model, we can use centralized training to retrain only the last ` layers whilekeeping the rest of the (7  ` ) layers (or ResNet blocks) frozen.
We can also perform this procedurestarting from a randomly initialized model.
The performance difference between these two modelscan be attributed to the information content of the frozen (7  ` ) layers of the FedAvg model.
Table 1summarizes the results of this experiment.
The large difference in accuracy (up to 42 .
6 %) indicatesthe initial layers of the FedAvg model have learnt useful features.
There continues to be a gap betweenthe FedAvg features and random features in the earlier layers as well 2 , meaning that all layers of the
FedAvg model learn useful features.
We conjecture this is because from the perspective of earlierlayers which perform simple edge detection, the tasks are independent of labels and the clients arei.i.d.
However, the higher layers are more specialized and the effect of the heterogeneity is stronger.
Method
Based on our observations in Section 3, we propose Bootstrapped eNTK as a method for overcomingdata heterogeneity when training deep models in a federated setting.
Our high-level intuition isthat we want to leverage both the features learned from applying FedAvg to neural networks andthe effectiveness of convex federated optimization.
We thus 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 non-convexity, 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 non-convex optimization problem becomes (convex) linearregression with respect to these features.
[Table] [Table] [Table]
To reduce the computational burden of working with the eNTK approximation, we make two furtherapproximations: First, we randomly reinitialize the last layer of ✓ 0 and only consider @@✓ f ( x ; ✓ 0 ) withrespect to a single output logit.
Over the randomness of this reinitialization, E [ f ( x ; ✓ 0 )] = 0 .
Andgiven the random reinitialization, all the output logits of f ( x ; ✓ 0 ) are symmetric.
These observationsmean each data point x can be represented by a P -dimensional feature vector @@✓ f 1 ( x ; ✓ 0 ) , wheref 1 ( · ; ✓ 0 ) refers to the ﬁrst output logit.
Then, we apply a dimensionality reduction by subsampling prandom coordinates from this P -dimensional featurization.
In our setting, this sub-sampling has theadded beneﬁt of reducing the number of bits communicated per round.
In summary, we transform our original (non-convex) optimization problem over a neural networkinitialized 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 thecoordinates of  eNTK ( x ; ✓ 0 ) for each x to obtain a reduced-dimensionality eNTK representation.
Let 161 S : R P !
R p denote this subsampling operation.
Finally, we solve the resulting linear regressionproblem 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 , wherep is the eNTK feature dimension after subsampling.
We use z ki to represent the eNTK feature of thei -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 171 clients, and p k := n k /n , we can approximate the non-convex neural net optimization problem by theconvex 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 ofa model trained with FedAvg.
As we will show in Section 5, the convex reformulation in Eq.
(1)signiﬁcantly reduce the number of communication rounds needed to ﬁnd an optimal solution.
4.3 Bootstrapped eNTK (BooNTK)
We now present our algorithm Bootstrapped eNTK (BooNTK) for federated optimization.
Bootstrapped eNTK (BooNTK)
• Stage 1: Extracting eNTK features from a FedAvg-trained model.
FedAvg is ﬁrst used to train the model for T 1 communication rounds.
Let ✓ T 1 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
For
in the eNTK representation, subtract
the mean and scale to unit variance description of SCAFFOLD for solving linear regression problems be found in Algorithm 1,
It has the same communication and computation cost as FedAvg
To motivate BooNTK, recall that in Section 3 we found that FedAvg learns “useful” features despiteits poor performance, especially in the earlier layers.
By taking an eNTK approximation, BooNTKoptimizes 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 roundsneeded to learn a performant model, even with data heterogeneity.
Experiments
We now study the performance of BooNTK for the decentralized training of deep neural networksin the presence of data heterogeneity.
We compare BooNTK to state-of-the-art federated learningalgorithms on three benchmark tasks in federated learning.
For each task, we apply these algorithmson client data distributions with varying degrees of data heterogeneity.
We ﬁnd that our proposedapproach signiﬁcantly outperforms existing algorithms when clients have highly heterogeneous dataacross all tasks.
For additional experimental results and implementation details, see Appendix B.
5.1 Experimental Setup
Datasets and degrees of data heterogeneity.
We assess the performance of federated learningalgorithms on the image classiﬁcation tasks FMNIST [70], CIFAR10, and CIFAR100 [36].
FMNISTand CIFAR10 each consist of 10 classes, while CIFAR100 includes images from 100 classes.
Thereare 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 twotypes of non-i.i.d.
data distribution: (i) Data heterogeneity sampled from a symmetric Dirichletdistribution with parameter ↵ [44, 63].
That is, we sample p c
⇠ Dir K ( ↵ ) from a K -dimensionalsymmetric 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 [46].
That is, each client is allocated a subset of classes; then, the samples of each class are split intonon-overlapping subsets and assigned to clients that were allocated this class.
We use #C to denotethe number of classes allocated to each client.
For example, #C=2 means each client has samplesfrom 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 oftwo convolutional layers with max pooling followed by two fully connected layers (SimpleCNN).
For
CIFAR10 and CIFAR100, we mainly consider an 18-layer residual network [24] with 4 basic residualblocks (ResNet-18).
In Appendix B.2, we present experimental results for other architectures.
Algorithms and training schemes.
We compare BooNTK to state-of-the-art federated learning algorithms, focusing on the widely-used algorithms FedAvg [46], FedProx [43], and SCAFFOLD [34].(For comparisons to additional algorithms, see Appendix B.1.)
Each client uses SGD with weightdecay 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 BooNTK, we run 100 rounds of FedAvg in Stage 1 following the above and use100 communication 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 dataheterogeneity.
We evaluated each algorithms on each task under four degrees of data heterogeneity.
Smaller #C and ↵ in Table 2 correspond to higher heterogeneity.
We ﬁnd that the existing federated algorithms all suffer when data heterogeneity is high across allthree tasks.
For example, the top-1 accuracy of FedAvg on CIFAR-10 is 56 .
86% when #C=2 , whichis much worse than the 90 .
43% achieved in a more homogeneous setting (e.g. ↵ = 0 . 5 )
In contrast,
BooNTK achieves consistently strong performance, even in the face of high data heterogeneity.
More speciﬁcally, BooNTK achieves the best top-1 accuracy performance across all settings except
CIFAR-100 with ↵ = 0 .
5 , where BooNTK does only slightly worse than SCAFFOLD.
In dropping by less than 1 . 5% on CIFAR100 as ↵ goes from 0 .
5 to 0 .
Moreover, our algorithmimproves over existing methods by at least 15% in the challenging cases, including FMNIST with
[Table]
[Figure]
#C=1 , CIFAR-10 with #C=1 and #C=2 , and CIFAR-100 with ↵ = 0 .
01 and ↵ = 0 .
And, perhapssurprisingly, our algorithm still performs relatively well in the extreme non-i.i.d.
setting where eachclient sees only a single class.
Figure 2 compares the performances of FedAvg, SCAFFOLD, and BooNTK in more detail on
CIFAR100 dataset with different degrees of data heterogeneity.
We consider the Dirichlet distributionwith parameter ↵ 2 { 0 . 1 , 0 . 01 , 0 . 001 } and compare the training and test accuracy of these threealgorithms.
As shown in Figures 2(a) and 2(b), both FedAvg and SCAFFOLD struggle when dataheterogeneity is high: for both algorithms, test accuracy drops signiﬁcantly when ↵ decreases.
Incontrast, we see from Figure 2(c) that BooNTK maintains almost the same test accuracy for different↵ .
Furthermore, the same set of default parameters for our algorithm, including local learning rateand 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 (usedin BooNTK-Stage 2) to full batch gradient descent (GD) applied to the overparameterized linear [Figure] [Figure] regression problem in Stage 2 of BooNTK 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 learningrate 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 somelearning 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 allsettings.
In particular, our algorithm converges within 20 communication rounds on CIFAR100 (asshown in Figure 3(c)).
These results suggest that our proposed algorithm can largely leverage thelocal computation and improve communication efﬁciency.
5.4 Ablations
Gradient correction.
We investigate the role of gradient correction when solving overparameterizedlinear regression with eNTK features in BooNTK.
We compare SCAFFOLD (used in BooNTK)to FedAvg on solving the regression problems and summarize the results in Figure 4.
We use thedefault 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 thenumber of local steps is large ( M  100 ) across three datasets.
We also ﬁnd that the performanceof 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 convergesfaster when the number of local steps increases.
These observations highlight the importance ofgradient correction in our algorithm.
[Figure]
Model weights for computing eNTK features.
To understand the impact of the model weightstrained in Stage 1 of BooNTK, we evaluate BooNTK 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 atrandom initialization.
Speciﬁcally, without FedAvg training, the eNTK (at random initialization)performs worse than standard FedAvg.
In contrast, BooNTK signiﬁcantly outperforms FedAvg by alarge margin (by ~ 20% in test accuracy) when eNTK features are extracted from a FedAvg-trainedmodel.
Also, we ﬁnd that BooNTK is stable with respect to the choice of communication rounds T 1in Stage 1.
For example, models trained by BooNTK with T 1  60 achieve similar performance.
Effect of normalization.
In Figure 5(b), we investigate the role of normalization on BooNTK bycomparing BooNTK run with normalized and unnormalized eNTK features.
The same number oflocal steps ( M = 500 ) is applied for both settings.
We tune the learning rate ⌘ for each setting and plotthe run that performs best (as measured in training accuracy).
The results in Figure 5(b) suggest thatthe normalization step in BooNTK signiﬁcantly improves the communication efﬁciency by increasingconvergence speed.
In particular, BooNTK with normalization converges to ~ 100% training accuracyin ~40 communication rounds, which is much faster than BooNTK without normalization.
Conclusion
Through experiments, we showed that non-convexity poses a signiﬁcant challenge for federatedlearning algorithms.
We found that a neural network trained in such a manner does learn someuseful features, but fails to use them and thus has poor overall accuracy.
To side-step this issue, weproposed a two-step training procedure: ﬁrst, train the neural network using FedAvg; then, optimize(using SCAFFOLD)
a convex approximation of the model obtained using its empirical neural tangentkernel.
We showed that the ﬁrst stage extracts meaningful features, whereas the second stage learnsto utilize these features to obtain a highly performant model.
Our algorithm obtains much highertest accuracies, is signiﬁcantly faster, and more stable to hyper-parameters than existing federatedlearning approaches.
In this work, we leveraged the idea of convexifying a deep model to mainly improve federated training.
The same idea could potentially also improve differential privacy guarantees, since convex losses havetighter bounds on privacy leakage [16].
More generally, we believe our approach can be widely usedto develop principled solutions for FL problems where directly analyzing deep models is intractable,e.g. generalizing to unseen clients [72], continual learning [17], personalization [12], etc.