Federated Submodel Optimization for Hot and Cold Data Features
Abstract
We focus on federated learning in practical recommender systems and natural language processing scenarios. The global model for federated optimization typically contains a large and sparse embedding layer, while each client’s local data tend to interact with part of features, updating only a small submodel with the feature-related embedding vectors. We identify a new and important issue that distinct data features normally involve different numbers of clients, generating the differentiation of hot and cold features.
We further reveal that the classical federated averaging algorithm (FedAvg) or its variants, which randomly selects clients to participate and uniformly averages their submodel updates, will be severely slowed down, because different parameters of the global model are optimized at different speeds.
More speciﬁcally, the model parameters related to hot (resp., cold) features will be updated quickly (resp., slowly).
We thus propose federated submodel averaging (FedSubAvg), which introduces the number of feature-related clients as the metric of feature heat to correct the aggregation of submodel updates.
We prove that due to the dispersion of feature heat, the global objective is ill-conditioned, and FedSubAvg works as a suitable diagonal preconditioner.
We also rigorously analyze FedSubAvg’s convergence rate to stationary points.
We ﬁnally evaluate
FedSubAvg over several public and industrial datasets.
The evaluation results demonstrate that FedSubAvg signiﬁcantly outperforms FedAvg and its variants.
Introduction
Federated learning (FL) [1] allows a large number of clients (e g., millions of smartphone users) to collaborate in the training of a global machine learning (ML) model under the coordination of a cloud server without sharing raw data.
For the production use on the top, Google has deployed FL among the users of its Android keyboard, called Gboard, to polish language models [2].
As the default optimization algorithm of FL at the bottom, federated averaging (FedAvg) averages the participating clients’ local new models to update the global model.
Much effort of existing work was devoted to proving the convergence of FedAvg, and the key challenge is that the clients’ local data are normally non-independent and identically distributed (non-i.i.d.)
One line of work [7, 8, 9, 10, 11] established an O (1 / p NT ) convergence, where N denotes the total number of clients and T denotes the total number of iterations.
However, these work required all the clients to participant in each round of FL, which is not practical in cross-device FL.
Another line of work [12, 13, 14, 15, 16, 17] allowed partial client participation and proved an O (1 / p KT ) convergence of FedAvg, where K denotes the number of chosen clients in each round.
Some other work proposed variants of FedAvg to better deal with data heterogeneity.
For example, Li et al.
[18] proposed FedProx by adding a proximal term to local objective; and Karimireddy et al.
[19] proposed Scaffold by adding a control variate for each client to control local training.
Besides non-i.i.d.
from data distribution, the interactions between the data features and the clients are mutual, partial, and differentiated, especially in practical recommender systems (RS) and natural language processing (NLP) scenarios.
First, distinct data features are normally involved by different numbers of clients.
For example, the differentiation of popular and unpopular items in RS or hot and cold words in NLP is common.
We call such an observation feature heat dispersion .
Second, a certain client’s local data tend to involve a small subspace of the full feature space, which further implies that the client needs to download and update only the feature-related part of the full global model, called a submodel in [20]. For example, deep recommendation and language models are normally stacked with a large and sparse embedding layer and some other dense layers, while a client’s submodel comprises the full dense layers and the embedding vectors for the client’s few local items or words rather than the full and huge embedding layer.
The existing work on FL has not studied the issue of feature heat dispersion yet.
However, different parameters of the global model involve distinct features and will be optimized by different numbers of clients and at different speeds, severely deteriorating the performance of FedAvg and its variants.
We take an extreme example in RS for illustration.
For an unpopular item 1 that appears only in 1% of the clients’ (denoted as client group G 1 ) local datasets, the corresponding embedding vector for item 1 is involved in the submodels of those clients in G 1 .
Using conventional FedAvg and its variants, only the clients in G 1 will upload non-zero updates, and the update of the embedding vector for item 1 will be slowed down 100 times.
In contrast, for a popular item 2 that appears in all the clients’ local datasets, the update of the corresponding embedding vector will not be slowed down.
To deal with feature heat dispersion, we propose federated submodel averaging (FedSubAvg), which ﬁrst averages the local updates of the chosen clients, just like FedAvg, but then multiplies the aggregated update of each model parameter with the ratio between the total number of clients and the number of clients who involve this model parameter.
Such a small correction ensures that the expectation of each model parameter’s global update is equal to the average of the local updates of the clients who involve this parameter.
We theoretically demonstrate the advantage of FedSubAvg over
FedAvg and analyze the convergence of FedSubAvg.
We ﬁrst prove that the global objective is illconditioned, leading to the slow convergence of FedAvg, and FedSubAvg works as a suitable diagonal preconditioner.
We also obtain an O ( p N/ ( n min KT )) convergence with respect to stationary points, where n min denotes the minimum of the number of clients who involve each individual parameter.
We summarize the key contributions of this work as follows: • To the best of our knowledge, we are the ﬁrst to make an in-depth study of FL from the differentiation of hot and cold data features, which is common in practice. • We identify the defect of FedAvg and its variants in handling feature heat dispersion and propose a novel, effective, and efﬁcient FedSubAvg algorithm. • We theoretically show that the the global objective is ill-conditioned, and FedSubAvg essentially works as a preconditioner for acceleration.
We also give the convergence rate of FedSubAvg in the general non-convex case.
• Using the public MovieLens, Sentiment140, and Amazon datasets, as well as an industrial dataset from Alibaba, we extensively evaluate FedSubAvg 2 and compare it with FedAvg, FedProx, Scaffold, and FedAdam.
The evaluation results reveal the superiority of FedSubAvg from faster convergence and smaller train loss.
Problem Formulation
In this section, we formulate the federated submodel optimization problem with feature heat dispersion in RS and NLP scenarios.
Optimization Objective.
We consider a distributed optimization setting, in which N clients collaboratively solve the following consensus optimization problem:
[Equation]
where D i denotes client i ’s local empirical distribution and ⇠ i ⇠ D i denotes the local training data; X is the full global model; F ( X , ⇠ i ) is the train loss of X over the local data ⇠ i ; and f i ( X ) is the local empirical error, taking expectation over the randomness of the local data.
Model Structure and Submodel.
The full recommendation or language model X normally adopts the network structure of an embedding layer plus some other dense layers, where sparse input features are mapped into embedding vectors, concatenated, and fed into the upper layers.
In practice, a client’s local data involve only part of the full global model, namely, a submodel, which is related to the client’s local data features.
For example, in RS (resp., NLP) scenario, client i can use the data collected in the previous week as its local training set D i , and retrieve the submodel in a key-value lookup way, typically, by retrieving a few embedding vectors based on the local item ids (resp.
, word ids) and directly taking the other dense network layers.
Considering the full embedding layer is far beyond any mobile device’s capacity, such a submodel design makes FL in these industrial scenarios possible.
Formally, we use S = { 1 , 2 , · · · , M } to index the parameters of the full model X 2 R M and call it the full index set.
We let S ( i ) ✓ S denote the index set of client i ’s submodel X S ( i ) .
In other words, the full model excluding the submodel (i.e., X S \ S ( i ) ) does not affect the model output, and the local gradient of X S \ S ( i ) will always be zero.
This further implies that we can rewrite the global objective function in a distributed submodel way:
[Equation]
The gradient of f is [Equation]
and the Hessian is[-H (cid:52) = ∇ 2 f ( X ) = 1 N (cid:80) Ni =1 ∇ 
 f i (cid:0) X S ( i ) (cid:1)-]
[Equation]
Note that when doing any operation (e.g., summation) over multiple submodels, gradients, and Hessians, they will be automatically aligned according to the indices.
Feature/Parameter Heat Dispersion.
We ﬁnally introduce the metric of feature heat dispersion (resp., the resulting parameter heat dispersion), which is deﬁned as the ratio between the maximum and the minimum of the number of clients who involve each individual feature (resp., parameter).
Considering the fact that a parameter may involve one or multiple data features in RS and NLP 3 , the feature heat dispersion is a lower bound of the parameter heat dispersion.
In other words, high feature heat dispersion inevitably leads to high parameter heat dispersion.
For convenience and clarity in the submodel-level analysis, we mainly take the metric of parameter heat dispersion.
We let n m denote the number of clients involving the parameter with index m 2 S and set n max = max m 2 S n m , n min = min m 2 S n m .
Then, the parameter heat dispersion is n max /n min .
Algorithm Design
In this section, we ﬁrst show that FedAvg suffers from high parameter heat dispersion in distributed submodel optimization and then propose FedSubAvg to remedy the defect. {+[Figure] Bob 7 0 0 
 [Figure]+}
[Figure] 1.5 3 3 1= # [Figure] 3.1 Slow Convergence of FedAvg
Example 1.
We consider a special distributed convex optimization problem with two model parameters, denoted as w 1 and w 2 .
Each client i ’s local data ⇠ i ⇠ D i are with mean e i = E
[ ⇠ i ] = 0 .
In addition, w 1 involves only client 1, while w 2 involves all the N clients.
Then, the parameter heat dispersion is n 2 /n 1 = N .
We formulate this learning problem as minimizing the mean square error:
[Equation]
where f ( w 1 , w 2 ) S (1) [Equation] ⇤ = w
[Equation]
w
[Equation] E ⇠ 1 ⇠ D 1 [ k ⇠ 1  e 1 k 2 ] ; and for i = 2 , 3 , · · · , N , f i  [Equation] i
[Equation]
For this example, the optimal model is ( w ⇤ 1 , w ⇤ 2 ) =
We leverage FedAvg with only one local iteration and let each client compute the exact (not stochastic) gradient.
The learning rate is denoted as ⌘ , and the model is initialized as ( w 01 , w 02 ) .
After r rounds, the model will become
[Equation]
[Equation]
[Equation]
r  w
[Equation]
, w > . By choosing ⌘ = 0 . 5 , ( w r 1 , w r 2 ) =
((1  1 /N )
r w 01 , 0) , we can ﬁnd that in the FL scenario with high parameter heat dispersion N , w 1 will converge at a quite low speed.
We also depict the optimization process of FedAvg in Figure 2.
3.2 Federated Submodel Averaging
To mitigate parameter heat dispersion, we propose FedSubAvg.
As shown in Figure 1, compared with FedAvg, the key principle of FedSubAvg is to further multiply the aggregated update of each model parameter with the ratio between the total number of clients and the number of clients who involve this model parameter 4 (i.e., for the model parameter with index m , the correction coefﬁcient is N/n m ).
We still examine Example 1 for illustration.
FedSubAvg will multiply the aggregated update of w 1 and w 2 with N and 1 , respectively.
By correction, the model at the r -th round is:
[Equation]
where  is the learning rate of FedSubAvg.
Therefore, FedSubAvg converges quickly to the optimal model (0 , 0) .
We also depict the optimization processes of FedSubAvg and FedAvg for Examplein Figure 2, when the parameter heat dispersion is 100 .
We can observe that FedSubAvg greatly outperforms FedAvg from convergence speed and loss.
Please refer to Appendix F for details.
Require:
The total number of clients N , the number of local iterations I , the number of clients involving the m -th model parameter n m , the number of selected clients K per round. [Equation]
We now present the design details of FedSubAvg in Algorithm 1.
In each round r of FL, the cloud server ﬁrst selects K clients to participate (Line 3), denoted as C r .
Each selected client i 2 C r determines its index set S ( i ) of submodel based on the local training set (Line 12).
Then, client i uses S ( i ) to download the submodel X rS ( i ) from the cloud server and initializes the local submodel x r, 1 i (Lines13–14).
Client i locally trains its submodel by doing I iterations of stochastic gradient descent (SGD) (Lines 15–17) and uploads the submodel update  x rS ( i ) (Line 18).
After receiving the submodel updates from the selected clients, the cloud server performs aggregation for each index in the union of the participating clients’ index sets and updates the global full model (Lines 7–10).
In particular, for the global model parameter with index m , namely, X r { m } , its expected update, after being corrected with the coefﬁcient N/n m , is
[Equation]
which is equal to the average of the local updates of the clients involving it, as required.
Theoretical Analysis
In this section, we ﬁrst prove that in FL with high parameter heat dispersion, the global objective f is ill-conditioned.
We then prove that optimizing f with FedSubAvg approximates to optimizing a preconditioning objective ˆ f with gradient descent (GD), thereby remedying ill conditioning.
In particular, the diagonal preconditioning matrix comprises of the correction coefﬁcients { N/n m | m 2 S } in FedSubAvg and can be obtained without needing to access any client’s local raw data and without any expensive computing, which keeps the tenet of FL.
Further by analyzing the Hessian of ˆ f , we demonstrate that the superiority of optimizing ˆ f over f .
We also obtain the convergence guarantee for FedSubAvg.
4.1 FedAvg with Ill-Conditioned Global Objective
We ﬁrst make an assumption about the eigenvalues of each client’s local Hessian.
Assumption 1 (Bounded Hessian) .
For any model X , the global Hessian is non-singular, and for each client i , the Hessian of f i satisﬁes:
[Equation]
[-(cid:22) − ρ 
-]
[-(cid:22) ρ 
-]
Further, if X is in a locally convex area of f , for each parameter m , more than 1  ↵ of the related clients satisﬁes {+⇢ 1 I i  r 
+}
f i ( X S ( i ) )
⇢
I i , where ↵ is a constant with (1  ↵ ) ⇢ 1  ↵⇢ 2 > 0 .
Assumption 1 bounds the eigenvalues of [-∇ 
-]r 2 f i ( X S ( i ) ) and ensures that when X is in a locally convex area of f , X S ( i ) is also in a locally convex area of f i for most of the clients, which helps to obtain the bounds of the global Hessian.
Under Assumption 1, we can analyze the condition number of H :  ( H ) 4 =
max ( H ) /  min ( H ) , where  max ( H ) and  min ( H ) denote the biggest and smallest singular values of H , respectively.
Theorem 1.
Under Assumption 1, for any model X in a locally convex area, the condition number of the global Hessian of f satisﬁes:
[Equation]
Proof.
Please refer to Appendix A.
Theorem 1 reveals that when the parameter heat dispersion n max /n min is high, then the global objective is ill-conditioned.
Speciﬁcally, in the large-scale RS and NLP scenarios, since the fully connected layer involves all the N clients, while the embedding vector of a cold item or word normally involves a few clients, we have  ( H )  ⇥ ( N ) .
This further implies with a large N in practice, the global objective is extremely ill-conditioned.
Therefore, the conventional FedAvg and its variants, which are the approximations to GD, will converge at a quite slow speed, even when the global model is in a locally convex area.
4.2 FedSubAvg as a Preconditioner Better than FedAvg
For brevity, we introduce t to denote the global iteration index, thus replacing the round index r and the local iteration index j , where t =
( r  1) ⇥ I
+ j .
We consider a single iteration of FedSubAvg:
[Equation]
where X t {+
+} = D [Equation] Ni =1 x ti ) is deﬁned as the global model at iteration t , and D = diag { N/n 1 , N/n 2 , · · · , N/n M } .
By treating D as a preconditioning matrix [21], we can construct a new objective ˆ f ( ˆ X ) 4 = f
( D 12 ˆ X ) =
f ( X ) with variable ˆ X
=
D  12 X .
The gradient and the
Hessian 5 of ˆ f are:
[Equation]
Dauphin et al.
[21] showed that one SGD update for ˆ f corresponds to equation 1.
Therefore, we can compare FedSubAvg and FedAvg by comparing the characteristics of ˆ f and f .
Theorem 2.
Under Assumption 1, for any X , the upper bound of the condition number of the corresponding Hessian, ˆ H 4 = r 2 ˆ f ( ˆ X ) , is always smaller than that of H :
[Equation]
Further, if X in a locally convex area, the condition number of r 2 ˆ f ( ˆ X ) , satisﬁes:
[Equation]
Proof.
Please refer to Appendix B.
Theorem 2 indicates that ˆ H always has a smaller condition number upper bound compared with
H .
Further, in a locally convex area, the proposed FedSubAvg signiﬁcantly reduces the condition number: ˆ H is well-conditioned with  ( H )  ⇥ (1) .
Therefore, the proposed FedSubAvg, which is an approximation to GD for objective ˆ f , will maintain the efﬁciency of FedAvg in any case and converge much faster than FedAvg in a locally convex area of f .
4.3 Convergence Guarantee for FedSubAvg
In this section, we analyze the convergence rate of FedSubAvg in the general non-convex case based on some standard assumptions.
We use k r ˆ f ( ˆ X ) k 2 = r f ( X ) > D r f ( X ) rather than k r f ( X ) k 2 to characterize the convergence rate.
This is because the curvature of the original global objective f has been severely “diluted”.
In particular, most model parameters involve only a small number of clients, while the zero gradients contributed by many non-involved clients are inaccurately counted in when computing the global gradient 6 .
We next make the following assumptions about the objective functions, as well as the variance and the feasible space of the stochastic gradients.
Assumption 2 (Smoothness) .
f i ( · ) is L -smooth if
[Equation]
Assumption 3 (Bounded Variance) .
During local training, the variance of stochastic gradients on each client is bounded by  2 : 8 i, t : E ⇠ i ⇠ D i
[Equation]
Assumption 4 (Bounded Gradient Norm) .
During local training, the expected l 2 -norm of the stochastic gradients is bounded by a constant G 2 : 8 i, t : E ⇠ i ⇠ D i
[Equation]
Assumption 2 is standard.
Assumptions 3 and 4 were widely made in the literature [8, 9, 11, 22, 13].
Under these assumptions, we can bound the gradient norm.
Theorem 3.
Under Assumptions 2, 3, and 4, we can bound the expected average of the squared gradient norm
[Equation]
When T > n min K 3 /N and  = ⇥
[Equation]
n min K NT ◆ , we have
[Equation]
Proof.
Please refer to Appendix C.
Evaluation
In this section, we extensively evaluate the performance of FedSubAvg over several datasets with different feature heat dispersion.
5.1 Experimental Setups
We choose the following tasks and models for evaluation.
The statistics about clients and samples, as well as the feature heat dispersion are shown in Table 1.
LR for Rating Classiﬁcation. We perform a rating classiﬁcation task over the MovieLens-1M dataset [23], which contains 6,040 clients, 3,883 movies, and 1,000,209 samples.
We preprocess the [Table] dataset to be suitable for binary classiﬁcation.
In particular, the original user ratings of movies range from 0 to 5.
We label the samples with the ratings of 4 and 5 to be positive and label the rest to be negative.
We randomly select 20% of the samples as the test dataset and leave the remaining 80% as the training dataset for FL.
The task is to predict whether users will rate a given movie to be positive based on the user’s gender and age and on the movie ID.
We ﬁrst encode gender, age, movie, gender cross movie, and age cross movie, based on the one-hot encoder.
Next, the features are input into a logistic regression (LR) model to predict the label.
LSTM for Sentiment Analysis. We perform a text sentiment classiﬁcation task on the Sentiment dataset [24], which comprises 1,600,000 tweets collected from 659,775 twitter users.
In this task, we use a two-layer long short-term memory (LSTM) network with 100 hidden units and an embedding layer as the binary classiﬁer, where the embedding dimension is set to 25.
We naturally partition this dataset by letting each Twitter account correspond to a client.
We keep only the clients who hold more than 40 samples and get 1,473 clients in total.
We randomly select 20% of the samples as the test dataset and leave the remaining 80% as the training set for FL.
DIN for CTR Prediction. We perform a click-through rate (CTR) prediction task on the Amazon electronics dataset and an Alibaba industrial dataset.
The Amazon dataset contains 1,689,188 reviews contributed by 192,403 users for 63,001 items.
The ratings range from 0 to 5.
We label the samples with the rating of 5 to be positive and label the rest to be negative.
We naturally partition this dataset by letting each Amazon user correspond to a client.
For each client, we take the user ID, the historical sequence of positively rated product as the input to predict the label.
We keep only the clients who hold more than 40 samples and get 1,870 clients in total.
We select samples with the timestamps more than 1,385,000,000 as the test dataset and leave the remaining samples as the training set.
The
Alibaba dataset is built from 30-day impression and click logs of 49,023 Taobao clients from June 15, 2019 to July 15, 2019.
For a certain Taobao user, we leverage its click behaviors in previous 14 days as historical data to predict its click and non-click behaviors in the following 1 day.
We leave out the behaviors within the last 1 day as the target items of the test set while putting the other samples into the training set.
For both datasets, we take the deployed deep interest network (DIN) [25] in Alibaba as the model, where the embedding dimension is set to 18.
We use the following ﬁve baselines for comparison.
• FedAvg averages the local model updates from the participating clients to update the global model.
• FedProx is the ﬁrst variant of FedAvg.
The main difference from FedAvg is that FedProx adds a quadratic proximal term to explicitly limit the local model updates.
We set the coefﬁcient of the proximal term to 0.01. • Scaffold is another important variant of FedAvg.
The key difference from FedAvg is that each client keeps a variate to control the local model updates in Scaffold.
However, the size of the control variate is equal to size of the full model, which is prohibitively inefﬁcient for the learning tasks with large-scale full models.
Therefore, for the CTR prediction tasks, we make an approximation to Scaffold.
In particular, the cloud server performs the controlled update step every round by weighted averaging the historical updates.
Please refer to Appendix D.2 for details.
• FedAdam is also an important FL algorithm, which adapts an adaptive optimizer [26].
We make more comparisons and discussions in Appendix E. • CentralSGD runs the standard SGD algorithm to train the global model using the whole dataset, sets the number of iterations in each round as I , and sets the batch size to the sum of the selected clients’ local batch sizes in each round.
This ensures the same amount of data per round with the distributed algorithms.
[Figure]
[Table]
Regarding the experimental settings, we choose mini-batch SGD as the optimization algorithm.
For the tasks of rating classiﬁcation and sentiment analysis, K = 50 clients are randomly chosen per round as default; and for the CTR prediction tasks, K is set to 100 as default.
The settings of the other hyperparameters are deferred to Appendix D.3.
5.2 Evaluation Results
We ﬁrst present the results of FedSubAvg and the baselines under the default K .
We then vary K to show its impact.
FedSubAvg vs. Baselines. For the rating classiﬁcation on the MovieLens dataset and the sentiment analysis on the Sent140 dataset, we plot the train loss in Figure 3(a) and Figure 3(b); and for the
CTR prediction on Amazon and the Alibaba dataset, we plot the test area under the curve (AUC)in Figure 3(c) and 3(d).
In addition, we measure the convergence rates of different algorithms by counting the communication rounds to reach a target train loss or test AUC.
We set the target loss in the rating classiﬁcation (resp., the sentiment analysis) to be the minimum loss of CentralSGD, which is 0.325 (resp., 0.380); and we set the target test AUC to be 0.6 in two CTR prediction tasks.
The results are listed in Table 2.
From Figure 3 and Table 2, we observe that FedSubAvg consistently outperforms FedAvg and its variants.
Speciﬁcally, (1) in the rating classiﬁcation, FedSubAvg always has the smallest train loss during FL and reaches the target at the 100-th communication round, 1 . 7 ⇥ faster than FedAvg,
FedProx, and FedAdam, and 1 . 8 ⇥ faster than Scaffold and CentralSGD; (2) in the sentiment analysis,
FedSubAvg still has the smallest train loss during FL and reaches the target at the 260-th round, 3 .
77 ⇥ faster than CentralSGD, while FedAvg, FedProx, Scaffold, and FedAdam cannot reach the target even in 1,000 rounds; (3) in the CTR prediction on the Amazon dataset, FedSubAvg achieves the highest AUC while FedAdam achieves the target AUC ﬁrst among all the FL algorithms.
FedSubAvg achieves the highest AUC of 0.641 in 200 rounds, decreasing by 0.013 in terms of AUC compared with the ideal CentralSGD.
In contrast, FedAvg achieves the highest AUC of 0.523, FedProx achieves the highest AUC of 0.519, Scaffold achieves the highest AUC of 0.514, and FedAdam achieves the highest AUC of 0.633.
In addition, FedSubAvg reaches the target test AUC at the 53-th round, while FedAvg, FedProx, and Scaffold cannot reach the target even after 200 rounds; and (4) in the
CTR prediction on the Alibaba dataset, FedSubAvg still outperforms all the other FL algorithms.
FedSubAvg achieves the highest AUC of 0.626 in 5,000 rounds, decreasing by 0.016 compared with [Figure] [Table]
CentralSGD.
In contrast, FedProx achieves the highest AUC of 0.514, FedAvg achieves the highest
AUC of 0.509, Scaffold achieves the highest AUC of 0.507, and FedAdam achieves the highest AUC of 0.554.
Moreover, FedSubAvg reaches the target test AUC at the 610-th round, whereas the other
FL algorithms cannot reach the target even after 5,000 rounds.
Impact of participating clients. We next evaluate the impact of the number of selected clients K per round on FedSubAvg.
We set K = 10, 30, and 50 for the rating classiﬁcation and the sentiment analysis, and set K = 20, 60, and 100 for the CTR prediction.
We plot the results in Figure 4 and record the minimum number of rounds to reach the target in Table 3.
We observe that FedSubAvg with a larger K generally converges much faster, which validates the speedup with respect to K .
(1) In the rating classiﬁcation, FedSubAvg with different K behaves somewhat uniformly.
This is because a larger K improves the convergence by reducing the variance of the global model update, while for such a simple convex optimization scenario, the variance is already small enough with
K = 10 ; (2) in the sentiment analysis, FedSubAvg with K = 50 reaches the target train loss 1 .
62 ⇥ faster than FedSubAvg with K = 10 ; (3) in the CTR prediction on the Amazon dataset , FedSubAvg with K = 100 reaches the target test AUC 1 . 79 ⇥ faster than FedSubAvg with K = 20 ; and (4) in the CTR prediction on the Alibaba dataset, FedSubAvg with K = 100 reaches the target test AUC 5 .
69 ⇥ faster than FedSubAvg with K = 20 .
Conclusion
In this work, we studied federated submodel optimization over non-i.i.d.
data with feature heat dispersion.
We proposed FedSubAvg, which ensures the expectation of the global update of each model parameter to be equal to the average of the local updates of the clients who involve it.
We also proved that FedSubAvg works as a preconditioner to improve collaborative optimization and thoroughly analyzed the convergence.
Empirical studies demonstrated the remarkable superiority of
FedSubAvg over FedAvg and its variants.
Acknowledgment
This work was supported in part by National Key R&D Program of China No. 2019YFB2102200, in part by China NSF grant No. 62202296, 62025204, 62072303, 61972252, 61972254, and 61832005, in part by Alibaba Group through Alibaba Innovation Research (AIR) Program, and in part by
Tencent Rhino Bird Key Research Project.
The opinions, ﬁndings, conclusions, and recommendations expressed in this paper are those of the authors and do not necessarily reﬂect the views of the funding agencies or the government.