FedImpro: Measuring and Improving Client Update in Federated Learning
Authors: Zhenheng Tang
Yonggang Zhang
Shaohuai Shi
Xinmei Tian
Tongliang Liu
Bo Han
Xiaowen Chu
Venue: ICLR 2024 poster (2024)
Keywords: Federated Learning
Deep Learning
Primary area: general machine learning (i.e., none of the above)
TL;DR: This work measures the generalization contribution of local training in federated learning, and utilizes estimated features to improves them to accelerate the convergence.

Abstract: Federated Learning (FL) models often experience client drift caused by heterogeneous data, where the distribution of data differs across clients. To address this issue, advanced research primarily focuses on manipulating the existing gradients to achieve more consistent client models. In this paper, we present an alternative perspective on client drift and aim to mitigate it by generating improved local models. First, we analyze the generalization contribution of local training and conclude that this generalization contribution is bounded by the conditional Wasserstein distance between the data distribution of different clients. Then, we propose FedImpro, to construct similar conditional distributions for local training. Specifically, FedImpro decouples the model into high-level and low-level components, and trains the high-level portion on reconstructed feature distributions. This approach enhances the generalization contribution and reduces the dissimilarity of gradients in FL. Experimental results show that FedImpro can help FL defend against data heterogeneity and enhance the generalization performance of the model.

Published as a conference paper at ICLR 2024
FEDIMPRO: MEASURING AND IMPROVING CLIENT UP- DATE IN FEDERATED LEARNING
Zhenheng Tang !
Yonggang Zhang' Shaohuai Shi?
Xinmei Tian
Tongliang Liu Bo Han Xiaowen S
1 Department of Computer Science, Hong Kong Baptist University
2 Harbin Institute of Technology, Shenzhen
3 University of Science and Technology of China 4 Sydney Al Centre, The University of Sydney
5 DSA Thrust, The Hong Kong University of Science and Technology (Guangzhou)
ABSTRACT
Federated Learning (FL) models often experience client drift caused by hetero- geneous data, where the distribution of data differs across clients. To address this issue, advanced research primarily focuses manipulating the existing gra- dients to achieve more consistent client models. In this paper, we present an alternative perspective local models. First, we analyze the generalization contribution of local training and conclude that this generalization contribution is bounded by the conditional Wasserstein distance between the data distribution of different clients. Then, we propose Ied construct similar conditional distributions for local train- ing. Specifically, F Fed Impro decouples the model into high-level and low-level components, and trains the high-level portion on reconstructed feature distributions. This approach enbances the generalization contribution and reduces the dissimi- larity of gradients in FL.. Experimental results show that Fed I m can help FI. defend against data heterogeneity and enhance the generalization performance of the model.
client drift and aim to mitigate it by generating improved
INTRODUCTION
The convergence rate and the generalization performance of FL suffers from heterogeneous data distributions across clients (Non-IID data) (Kairouz et al. 2019). The FL community theoretically and empirically found that the "client drift" caused by the heterogeneous data the main reason of such a performance drop (Guo et al.] Wang et al., 2020a). The client drift means the far distance between local models on clients after being trained on private datasets.
Recent convergence analysis (Reddi et al. 2021 Woodworth et al. 2020) of FedAvg shows that the degree of client drift is linearly upper bounded by gradient dissimilarity. Therefore, most existing works (Karimireddy al. 2020 Wang et al.) (020a) focus on gradient correction techniques to accelerate the convergence rate of local training. However, these techniques rely on manipulating gradients and updates to obtain more similar gradients (Woodworth et al., 2020 Wang al. 2020a]]
Sun et al., (023a)). However, the empirical results of these methods show that there still exists performance gap between FL and centralized training.
In this paper, we provide a novel view to correct gradients and updates. Specifically, we formulate the objective of local training in FL systems generalization contribution problem. The generalization contribution means how much local training on one client can improve the generalization performance on other clients" distributions for server models. We evaluate the generalization performance of local model on other clients' data distributions. Our theoretical analysis shows that the generalization contribution of local training bounded by the conditional Wasserstein distance between clients' distributions. This implies that even if the marginal distributions on different clients are the same, it is insufficient to achieve a guaranteed generalization performance of local training. Therefore, the key
Corresponding author (xwchu @ hk). This work is partially done during the internship in The Hong Kong University of Science and Technology (Guangzhou).

---

Published as a conference paper at ICLR 2024
Clients
Server
Add noise
Moving average update
Jer4s
Local training
A
Synthelic feature
Update
Output
30
Broadcast
Feature
Concat
(m))
Dm
Synthetic feature &
Model aggregation
Fenture
Concat
(4.3))
Figure 1: Training process of our framework. On any m-th Client, the low-level model uses the raw data I' as input, and outputs feature hnn The high-level model uses hm and samples h from a shared distribution H" as input for forward and backward propagation. Noises will be added to the locally estimated H before aggregation on the Server to update the global H Model parameters follow the FedAvg aggregation or other Fl. aggregation algorithms.
to promoting generalization contribution to leverage the same or similar conditional distributions for local training.
However, collecting data to construct identical distributions shared across clients is forbidden due to privacy concerns. To avoid privacy leakage, we propose decoupling a deep neural network into a low-level model and a high-level one, i.e., feature extractor network and a classifier network. Consequently, we can construct shared identical distribution in the feature space. Namely, on each client, we coarsely coarsely"]sstinn estimate the feature distribution obtained by the low-level network and send the noised estimated distribution to the server model. After aggregating the received distributions, the server broadcasts the aggregated distribution and the server model to clients simultancously. Theoretically, show that introducing such simple decoupling strategy promotes the generalization contribution and alleviates gradient dissimilarity. Our extensive experimental results demonstrate the effectiveness of F FedImpro, where we consider the global test accuracy of four datasets under various FL settings following previous works (He et al. 20206 Li et al. 20206 Wang et al. 2020a).
Our main contributions include: (1) We theoretically show that the critical, yet far overlooked generalization contribution of local training is bounded by the conditional Wasserstein distance between clients' distributions (Section [. (2) We are the first theoretically propose that sharing similar features between clients can improve the generalization contribution from local training, and significantly reduce the gradient dissimilarity, revealing new perspective of rectifying client drift (Section 4.2). (3) We propose FedImpro to efficiently estimate feature distributions with privacy protection, and train the high-level model on more similar features. Although the distribution estima- tion is approximate, the generalization contribution of the high-level model is improved and gradient dissimilarity is reduced (Section 4.3). (4) We conduct extensive experiments to validate gradient dissimilarity reduction and benefits generalization performance of F FedImpro (Section
2 RELATED WORKS
We review FL algorithms aiming to address the Non-III problem and introduce other works related to measuring client contribution and decoupled training. Due limited space, we leave more detailed discussion of the literature review in Appendix D
2.1 ADDRESSING NON-IID PROBLEM IN FL
Model Regularization focuses on calibrating the local models to restrict them not to be excessively far away from the server model. A number of works like FedProx (Li et al. 2020b). FedDyn (Acar et al. 2021), SCAFFOLD (Karimireddy et al. 2020). VHL (Tang et al. (022b) utilizes shared noise data to calibrate feature distributions. FedETF (Li et 2023b) proposed a synthetic and fixed ETF
classified to resolve the classifier delemma. SphereFed (Dong et al. 2022 constrains the learned representations to be a unit hypersphere.
'Considering the computation and communication cnsts, avoiding privacy Icakage, the distrihution estimation approximate. And the parameters will be with noise.
2

---

Published as a conference paper at ICLR 2024
Reducing Gradient Variance tries to correct the directions of local updates at clients via other gradient information. This kind of method aims to accelerate and stabilize the convergence, like
FedNova (Wang al. 2020a), FedOPT (Reddi et al. (22TT) and FedSpeed (Sun et al. 2023b). Our theorem provides a new angle to reduce gradient variance.
Sharing data methods proposes sharing the logits data between clients (Yang et al. 2023 Cronus (Chang et al.) 2019) shares the logits to defend the poisoning attack. CCVR (Luo et al. 2021 transmit the logits statistics of data samples to calibrate the last layer of Federated models.
CCVR (uoo et al. 2021 also share the parameters of local feature distribution. FedDF (Lin et al. 2020 finetunes on the aggregated model via the knowledge distillation on shared publica data. The supriority of Fedlmpro may result from the difference between these works. Specifically, FedImpro focuses on the training process and reducing gradient dissimilarity of the high-level layers. In contrast, CCVR is a post-hoe method and calibrates merely the classifiers. Furthermore, the FedDF fails to distill knowledge when using the random noise datasets as shown in the results.
2.2
MEASURING CONTRIBUTION FROM CLIENTS
Clients' willingness to participate in FL training depends on the rewards offered. Hence, it is crucial to evaluate their contributions to the model's performance (Yu et al.) 2020 Ng et al. 2020 Liu et al.,
2022; Sim et al. 2020). Some studies (Yuan et al., 2022) propose experimental methods to measure
performance gaps from unseen client distributions. Data shapley (Ghorbani & Zou 2019) Sim et al.
2020: Liu et al. 2022) is introduced to assess the generalization performance improvement resulting from client participation. These approaches evaluate the generalization performance with or without certain clients engaging in the entire FL process. However, hope understand the contribution of clients at each communication round. Consequently, our theoretical conclusions guide a modification on feature distributions, improving the generalization performance of the trained model.
2.3 SPLIT TRAINING
Some works propose Split FL (SFL) to utilize split training to accelerate federated learning (Oh et al. 2022 Thapa et al., 2020). In SFL, the model is split into client-side and server-side parts. At each communication round, the client only downloads the client-side model from the server, and conducts forward propagation, and sends the hidden features to the server to compute the loss and conduct backward propagation. These methods aim to accelerate the training speed of FL on the client side and cannot support local updates. In addition, sending all raw features could introduce a high risk of data leakage. Thus, we omit the comparisons to these methods.
2.4 PRIVACY CONCERNS
There are many other works (Luo et al. 2021, Li & Wang, 2019 He al., 2020a Liang al., 2020: Thapa et al., 20200 Oh et al. 2022) that propose to share the hidden features to the server or other clients. Different from them, our decoupling strategy shares the parameters of the estimated feature distributions instead of the raw features, avoiding privacy leakage. We show that FedImpro successfully protect the original data privacy in Appendix F.7
PRELIMINARIES
3.1 PROBLEM DEFINITION
Suppose we have set of clients M (1,2, (1, with M being the total number of participating
clients. FL aims to make these clients with their own data distribution D., cooperatively learn machine learning model parameterized as 0 E Rd. Suppose there are C' classes in all datasets UmeM indexed by [C]. sample in D,, is denoted by (x, y) E x [C], where x is model input in the space X' and y is its corresponding label. The model is denoted by p(0: x)) X - R
Formally, the global optimization problem of FL can be formulated as (McMahan et al. 2017):
M
M
min F(O) = Pm F (0) = Σ BERA
Σ m=1
Pm Pm DD (0;x,,
(1)
m
3

---

Published as a conference paper at ICLR 2024
where (0) (xx.)))DD.. (8; is the local objective function of client with f(0; 2; y) CE((((;z).)) v). CE denotes the cross-entropy loss, Pm > 0 and > mm Pm = 1. Usually, Pm is set
N. where nmm. denotes the number of samples on client m and N Σ M
2m=1 Pmm--
The clients usually have a low communication bandwidth, causing extremely long training time. To address this issue, the classical FL algorithm FedAvg (McMahan et al. (217)) proposes to utilize local updates. Specifically, at each round r, the server sends the global model ø to a subset of clients S M which randomly chosen. Then, all selected clients conduct obtain new client models 10mh which are sent back to the server. Finally, the server averages local models according to the dataset size of clients obtain new global model 0".
iterations of updates to
3.2 GENERALIZATION QUANTIFICATION
Besides defining the metric for the training procedure, we also introduce metric for the testing phase. Specifically, we define criteria for measuring the generalization performance for a given deep model. Built upon the margin theory (Koltchinskii & Panchenko 2002 Elsayed et al. 2018), for given model p(f:-) parameterized with 6, we use the worst-case margin to measure the generalizability on the data distribution D:
Definition 3.1. (Worst-case margin.) Given a distribution D, the worst-case margin of model p(0;-) is defined as Wa(p(0), D) (xyy)>> (tg:aa,,,(000)"))_...) d(x', x) with d being specific distance, where the argmax, (00; x/) + means the p(0; x²) mis-classifies the x'.
This definition measures the expected largest distance between the data æ with label y and the data x' that is mis-classified by the model Thus, smaller margin means higher possibility to mis-classify the data æ. Thus, we can leverage the defined worst-case margin to quantify the generalization performance for a given model p and a data distribution D under specific distance. Moreover, the defined margin is always not less than zero. It is clear that if the margin is equal to zero, the model mis-classifies almost all samples of the given distribution.
4
DECOUPLED TRAINING AGAINST DATA HETEROGENEITY
This section formulates the generalization contribution in FL systems and decoupling gradient dissimilarity.
4.1
GENERALIZATION CONTRIBUTION
Although Eq. quantifies the performance of model with parameter 0, it focuses more on the training distribution. In FL, we cooperatively train machine learning models because of a belief that introducing more clients seems to contribute to the performance of the server models. Given client m, we quantify the "belief", i.e., the generalization contribution, FL systems follows: (P.. Wa(p(O + A), D (mm),
where A a pseudo gradient obtained by applying learning algorithm L(-) to a distribution Dm Wa is the quantification of generalization, and D DD means the data distribution of all clients except for client m. Eq. depiets the contribution of client m to generalization ability. Intuitively, we prefer the client where the generalization contribution can be lower bounded.
Definition 4.1. The Conditional Wasserstein distance Ca(D, D¹) between the distribution D and D):
1 Ca(D, D') BL.iiD inf JE.T(P)p.P B(nr))))(( x') + E (e.(())
1
inf
Jdl x').
Built upon Definition 3.1 and Eq. we are ready to state the following theorem (proof in Appendix C.II.
Theorem 4.1. With the pseudo gradient A obtained by L(D..), the generalization contribution is
lower bounded:
EaI(P Wa(plO D).)) 2BAT Walp + A). A). D... EAILAD... Walp(0 + Dm)) Walp(0 A), D. 2Ca(Pm, Dm),,,
where D' represents the dataset sampled Dm
"The similar definition is used in the literature (Franceschi et al. 2018
"The pseudo gradient round calculated as: A 07 with the maximum local iterations T.
4

---

Published as a conference paper at ICLR 2024
Remark 4.1. Theorem 4.1 implies that three terms are related to the generalization contribution. The first and second terms are intuitive, showing that the generalization contribution of a distribution Dm is expected to be large and similar to that of a training dataset Dm- The last term is also intuitive, which implies that promoting the generalization performance requires constructing similar conditional distributions. Both the Definition 4. and Theorem 4.1 use distributions conditioned on the label y. so we write the feature distribution Hyy as H for brevity in rest of the paper.
Built upon the theoretical analysis, it is straightforward to make all client models trained on similar distributions to obtain higher generalization performance. However, collecting data to construct such distribution is forbidden in FL due to privacy concerns. To address this challenge, we propose decoupling a deep neural network into a feature extractor network (Pt parameterized by Otow E Rd and a classifier network parameterized by thigh E Rd,, and making the classifier network trained on the similar conditional distributions Hly with less discrepancy, as shown in Figure | Here, di and de represent the dimensions of parameters flo and Oxigh, respectively.
Specifically, client can estimate its own hidden feature distribution as Hmm using the local hidden features h = PP (x) (x..)))) and send H.m to the server for the global distribution approximation. Then, the server aggregates the received distributions obtain the global feature distribution H. and broadcasts it, being similar to the model average in the FedAvg. Finally, classifier networks of all clients thus performs local training on both the local hidden features ((a.y)-D... and the shared H during the local training. To protect privacy and reduce computation and communication costs, we propose an approximate but efficient feature estimation methods in Section 4.3 To verify the privacy protection effect, following (Luo et al. (02T), we reconstruct the raw images from features by model inversion attack (Zhao et al., 2021 1: Zhou et al. 2023) in Figure I1
12 and
13 in Appendix F.7
showing FedImpro successfully protect the original data privacy.
In what follows, we show that such a decoupling strategy can reduce the gradient dissimilarity, besides the promoted generalization performance. To help understand the theory, We also provide an overview of interpreting and connecting our theory to the FedImpro in Appendix C.4
4.2 DECOUPLED GRADIENT DISSIMILARITY
The gradient dissimilarity in FL resulted from heterogeneous data, i.e., the data distribution on client
m, Dm. is different from that on client k, Dz (Karimireddy et al. 2020). The commonly used quantitative measure of gradient dissimilarity is defined as inter-client gradient variance (CGV).
Definition 4.2, Inter-client Gradient Variance (CGV): (Kairouz et al.
2019
Karimireddy et al. 2020
Woodworth et al. 2020 Koloskova al. 2020) CGV(F,
16 (r.u)
D.,.
fm, 2, y) V
CGV IS usually assumed to be upper bounded (Kairouz et al.
2019 Woodworth et al. 2020 Lian
ct al.
2017 i.e., CGV(F,6) ((..))D....
|| fm (B;
F(()) of with a constant Ø.
Upper bounded gradient dissimilarity benefits the theoretical convergence rate WWoodoothh et al. 2020). Specifically, lower gradient dissimilarity directly causes higher convergence rate (Karmıreddy al. 2020 Li al., 2020b Woodworth et al. (020). This means that the decoupling strategy can also beneht the convergence rate the gradient dissimilarity can be reduced. Now, we are ready to demonstrate how to reduce the gradient dissimilarity CGV with our decoupling strategy. With
representing fm (0; x, y) as {Vor.. fm (f) x. y), Vonagn fm;(; x. }) we propose that the CGV can be divided into two terms of the different parts of (see Appendix for details):
CGV(F, 8) =B(r.)) fm (0; x, v) 2(0))
(3)
[B V fm (0; x. v) Vol... oro. F(O) + || Onaga f (0; v) V Onaga [(0)]
According to the chain rule of the gradients of deep model, we can derive that the high-level part of gradients that are calculated with the raw data and labels (z, y) D,, is equal to gradients with the hidden features and labels (h- P0.. ()),y) (proof in Appendix C.2
Brigh fm (0; x, xx,y y) = Vonigh fm h, y),
M
V Vough F(0)= = Σ (a.)) Dm V Venegr (0:h,
h
m=1
in which fm (0; h y) is computed by forwarding the h = Peroo (x) through the high-level model without the low-level part. In FedImpro, with shared H. client m will sample h H and
5

---

Published as a conference paper at ICLR 2024
hm P0rom (z.)(()))) train their classifier network, then the objective function becomes as
M
M
min F(0) = Σ DER"
Pm D. (0; P, E hwH
Pm (x.() h~7
[000 (x). [(0; h w) (5)
Here, Pm PaattBa with nm and Åm being the sampling size of (I, Dm and h H
N+N
respectively. and N = =E m==1 Rm. Now, we are ready to state the following theorem of reducing gradient dissimilarity by sampling features from the same distribution (proof in Appendix C.3
Theorem 4.2. Under the gradient variance meusure CGV (Definition
(42). with om satisfying
Byg
N+N the objective function F(0) causes a tighter bounded gradient dissimilarity,
L.e CGV(F, 0) ((E,,))DD...l Virw Im (0; xr, y) Blow F(O)112
N²
II Puigh Im (n,00;00 (0; V)
(N+N)²
V Wonigh P(O) CGV(F,O).
Remark 4.2. Theorem 4.2 shows that the high-level gradient dissimilarity can be reduced as
N²
+N+
times by sampling the same features between clients. Hence, estimating and sharing feature distribu- tions is the key to promoting the generalization contribution and the reduction of gradient dissimilarity. Note that choosing Ñ impractical to sample infinite features h. First, the distribution is estimated using limited samples, leading to biased estimations. Second, infinite sampling will dramatically increase the calculating cost. We set N = N in our experiments.
can eliminate high-level dissimilarity. However, two reasons make it
4.3 TRAINING PROCEDURE
As shown in Algorithm | Fed Impro merely requires two extra steps compared with the vanilla FedAvg method and can be easily plugged into other FL algorithms: a) estimating and broadcasting global distribution H: per- forming local training with both the local data (xr, y) and the hidden features (h Hly. v).
Algorithm 1 Framework of FedImpro.
server input: initial 8", maximum communication
round R client m''s input: local iterations T
Initialization: server distributes the initial model 0
to all clients, and the initial global 7
Moreover, sampling h Hly has two additional advantages as follows. First, directly sharing the raw hidden features may ineur privacy concerns. The raw data may be reconstructed by feature inversion methods (Zhao al. 202I). One can use different distribution approximation meth- ods to estimate {hm E M) avoid expos- ing the raw data. Second, the hidden features usually have much higher dimensions than the raw data (Lin et al.] [021). Hence, communicat- ing and saving them between clients and servers may not be practical. We use different dis- tribution approximation methods to obtain H. Transmitting the parameters of HL can consume less communication resource than hidden fea- tures (/mm./m Im E M).
Server Executes:
for each round r 0, 1,... server samples a set of clients S, C (1....,MM) M).
server communicates 0. and 7 to all clients E S.
for each client S in parallel do do
gr+ - Hr+1 +- CliemtUpdatemmm H").
end for gr+1
MM Pm 0r+1 Tm.E-1
Update H using { (H Im E S
end for
ClientUpdate(mn 0, H): for each local iteration & with = 0.... .T do
Sample raw data (ar, y) D.,. and h H/y.
Following previous work (Knddall & Gal 2017), exploit the Gaussian distribution approx- imate the feature distributions. Although it is an inexact estimation of real features, Gaussian distribution is computation-fiise frienddl. And the mean and variance of features used estimate it widely exist in BatchNorm layers (loffe & Szegedy 2015), which are also communicated between server and clients.
0m.tt1 + 0m.. Mm.. V ; x, h, u) (Eq.
Update Hm. using hu.. = &i (x).
end for Return 0 and Hmm to server.
Wee reuse here for brevity, the input of f can he the input the hidden feature h 400 (z).

---

Published as a conference paper at ICLR 2024
Table 1: Best test accuracy (%) of all experimental results. "Cent." means centralized training. "Acc." means the test accuracy. Each experiment is repeated 3 times with different random seed. The standard deviation of each experiment shown its mean value.
Cent. Acc.
FI. Setting
FL. Test Accuracy FedNova
Dataset
E
M
FodAvg
FedProx
SCAFFOLD
FodDyn
FedImpro
0.1
1
10
83.6542.03
83.22:+1.52
82.33+1.13
84:97:0887
84.73:1..81
88.45:10.43
0.05
1
10
75.3611
77.49:11.24
33.61:887
73.49:11.42
77.21:1..25
81.75.11.6 1.03
CIFAR-
92.53
0.1
10
85.69+0.57
85.33+0.39
84.4-041
86.92-40.28
86.59+00399
0.1
1
100
1..42+1.119
68, 18.59911.03
59.2243.
74,94+0.98
75.29++0.95
77.56-+1.02
Average
79.53
78.66
64.89
80.08
80.96
83,97
0.1
1
10
88.67:10.34
88.92.10.25
87.8110.336
87.97:10.41
89.01.10.25
90.83:1 0.19
FMNIST
93.7
0.05
0.1
5
10
10
82.73-0.98
87 11..-t0..5 52
83.66-10..22
88 88.41+0.38
88 88.4440.29
76.16: 1.29
81989
87 0.66626662
13.20-t1.199
88.50++0.52
86.42:140 0.71
89.87--021
0.1
100
90.12:0.19
90.39: 12
88.2420.31
90.40:t0111
90.57:1:.21
90.98:0015
Average
87,28
87.85
85.16
86.98
87.82
89.53
0.1
1
10
88.20 1,21
87.04-10..8
83.87-42.15
88.48-11.31
90.82:1:1 1.09
92.37-10.82 0.82
0.05
10
80.674110 1,92
$2.39-+1..3
8..29-+1881
84.01
84.1241 1.28
90.25-40.69
SVHN
95.27
0.1
5
10
86. 1.19
86.05:00.72
83,14
88.10:0,91
89.92=00.50
91.58=+0.2
0.1
100
92.42:40.21
92.29:10.18
92.06:10.20
92,44:0.31
92.82:0.133
93.42:10.15
Average
86.90
86.94
85.34
88.26
89,42
91.91
0.1
1
10
69.38:1
69.78:10.91
65.74:11.52
69.52:10.75
69.5911.552
70.28-10.38
0.05
10
63.80 1.29
64.75-t 1.25
16
64.57 1.27
64.90
66.60- 0.91
CIFAR-100
74,25
0.1
10
68.
6 1.02
68.71+40888
1.267
67.99: 1.04
68.52+0441
68.79-+052
0.1
100
53.22 1.20
54. 14.10:11.32
23.7743.21
55 40:0.81
54.82::0.81
0..07-1075
Average
64.34
$4.92
65.44
On each client m, a Gaussian distribution N (M om) parameterized with pmm and om is used to approximate the feature distribution. On the server-side, .N ((ug: 7g) estimate the global feature distributions. As shown in Figure and Algorithm [ during the local training, clients update Jmm and om using the real feature h.. following moving average strategy which is widely used in the literature (Toffe & Szegedy, 2015 Wang et al. 2021):
((++1) pmp + (1 3m) x mean(hm).
a (t+1) Bmo m + (1 Bm) variance(hmm)
where t is the iteration of the local training, 3m is the momentum coefficient. To enhance privacy protection, on the server side, Bg and og are updated with local parameters plus noise e::
11 (r+1) = -ggu?" + +(1 (g) x
(uf + +++),
St
d(+1)) (gg -3) x
1
(of + +
St
where e:- ~ NO, a.). Appendix F.2 and Table F.2 show the performance of FedImpro with different noise degrees de. And Appendix F.7 show that the model inversion attacks fail to invert data based on shared Pm, m. Pgg and og. illustrating that FedImpro can successfully protect the original data privacy.
EXPERIMENTS
5.1 EXPERIMENT SETUP
Federated Datasets and Models. We verify FedImpro with four datasets commonly used in the FL community, i.e., CIFAR- Krizhevsky Hinton 2009). FMNIST (Xiao et al. 2017). SVHN (Netzer et al. (011). and CIFAR- 100 (Krizhevsky & Hinton] 2009). : We use the Latent Dirichlet Sampling (LDA) partition method to simulate the Non-IID data distribution, which is the most used partition method in FL (He et al.) 2020b Li al.) 2021c| Luo al.) 2021). We train Resnet-18 18 on CIFAR-10, FMNIST and SVHN, and 55 on CIFAR-100 We conduct experiments with two different Non-IID degrees,
0.1 and 0.05. We simulate cross-silo FL with M 10 and drvsse FL with M = 100. To simulate the partial participation each round, the number of sample clients 5 for M 10 and 10 for M 100. Some additional experiment results shown in Appendix
7

---

Published as a conference paper at ICLR 2024
Baselines and Metrics. We choose the classical FL algorithm, FedAvg (McMahan et al.
2017). and
recent effective FL algorithms proposed to address the client drift problem. including FedProx (Li
al. 2020b), SCAFFOLD (Karimireddy et al. (020), and FedNova (Wang et al. 2020a). as baselines. The detailed hyper-parameters of all experiments are reported in Appendix E We use two
metrics, the best accuracy and the number of communication rounds to achieve target accuracy. which is set to the best accuracy of FedAvg. We also measure the weight divergence (Karimireddy 0: |I. as reflects the effect on gradient dissimilarity reduction
al. 2020 |5 CiEs 110
5.2
EXPERIMENTAL RESULTS
Basic FL setting. As shown in Table I using the classical FL training setting, i.e. a 0.1, = and . = 10, for CIFAR- 10, FMNIST and S SVHN, FedImpro achieves much higher generalization performance than other methods. We also find that, for CIFAR- 100, the performance of Fedlmpro is similar to FedProx. We conjecture that CIFAR- -10AR-110 dataset has more classes than other datasets, leading to the results. Thus, a powerful feature estimation approach instead of a simple Gaussian assumption can be a promising direction to enhance the performance.
Impacts of Non-IID Degree. As shown in Table I for all datasets with high Non-IID degree (a = 0.05). FedImpro obtains more performance gains than the case of lower Non-IID degree (a = 0.1). For example, we obtain 92.37% test accuracy on SVHN with a 0.1, higher than the FedNova by 3.89%. Furthermore, when Non-IID degree increases
0.05, obtain 90.25% test accuracy, higher than FedNova by 6.14%. And for CIFAR- 100, Fedlmpro shows benefits when 0.05, demonstrating that Fedlmpro can defend against more severe data heterogeneity.
Different Number of Clients. We also show the results of 100-client FL setting in Table FedImpro works well with all datasets, demonstrating excellent scalability with more clients.
Different Local Epochs. More local training epochs E could reduce the communication rounds, saving communication cost in practical scenarios. In Table the results of E 5 CIFAR-10, FMNIST, and SVHN verify that FedImpro works well when increasing local training time.
Weight Divergence. We show the weight di- vergence of different methods in Figure (b). where FedNova is excluded due its significant weight divergence. The divergence is calculated by |Sa
1 Eves 0-l. which shows the dissim-
ilarity of local client models after local training. At the initial training stage, the weight diver- gence is similar for different methods. During this stage, the lewelevel model is still unstable and the feature estimation is not accurate. Af- ter about 500 communication rounds, FedImpro shows lower weight divergence than others, indicating that it converges faster than others.
(a) Test Accuracy (b) Weight Divergence
Figure 2: CIFARIO with a = 0.1, E. = 1, M 10.
Convergence Speed. Figure 2 (a) shows that FedImpro can accelerate the convergence of F And we compare the communication rounds that different algorithms need to attain the target accuracy in Table[]] Results show Fed- Impro significantly improves the convergence speed.
Table 3: Splitting at different convolution layers in R8
Layer
5-th
9-th
13-th
17-th
Test Acc. (%) 87.64
88.45
87.86
84.08
5.3 ABLATION STUDY
To verify the impacts of the depth of gradient decoupling, we conduct experiments by splitting at different layers, including the 5-th, 9-th, 13-th and 17-th layers. Table 3 demonstrates that FedImpro can obtain benefits at low or middle layers. Decoupling at the 17-th layer will decrease the performance, which consistent with our conclusion in Sec. 4.2 Specifically, decoupling at a
STo ensure the repraducibility of experimental results, we open source code at attttps:///tthub com/ wi zard 1203/red Impro
Due to high instability of training with severe heterogeneity, the actual test accuracy are semitransparent lines and the smoothed accuracy are opaque lines for better visualization. More results in Appendix F.
8

---

Published as a conference paper at ICLR 2024
Table 2: Communication Round attain the target accuracy. "NaN." means achieving the target.
FL Setting
Target Acc.
Communication Round attain the target accuracy
Datasel
E
M
FedAvg
FedProx
SCAFFOLD
FedNova
FedDyn
FedImpro
0.1
1
10
82.0
142
128
863
142
291
128
0.05
10
73.0
247
121
NaN
407
195
CIFAR-I0
0.1
84.0
128
128
360
80
109
100
73.0
957
NaN
NaN
992
706
0.1
1
10
87.0
83
76
275
83
62
32
0.05
10
81.0
94
94
NaN
395
82
52
FMNIST
0.1
10
87.0
147
31
163
88
29
0.1
100
90.0
375
470
NaN
317
592
441
0.1
10
87.0
292
247
NaN
251
162
50
0.05
10
80.0
578
68
358
242
120
50
SVHN
0.1
5
10
86.0
251
350
NaN
NaN
92
11
0.1
100
92.0
471
356
669
356
429
346
0.1
10
69.0
712
857
NaN
733
682
614
0.05
1
10
61.0
386
386
755
366
416
313
CIFAR-100
0.1
5
10
68.0
335
307
182
282
329
300
0.1 1
100
53.0
992
939
NaN
910
951
854
very high layer may not be enough resist gradient dissimilarity, leading weak data heterogeneity mitigation. Interestingly, according to Theorem [4.2] decoupling at the 5-th layer should diminish more gradient dissimilarity than the 9-th and 13-th layers; but it does not show performance gains. We conjecture that it is due to the difficulty of distribution estimation, since biased estimation leads to poor generalization contribution. As other works (Lin al., 2021] indicate, features the lower level usually are richer larger than at the higher level. Thus, estimating the lower-level features is much more difficult than the higher-level.
5.4 DISCUSSION
In this section, we provide some more experi- mental supports for FedImpro. All experiment results of this section are conducted on CIFAR- 10 with ResNet-18, a = 0.1, E = 1 and M 10. And further experiment result are shown in Appendix due to the limited space.
FedImpro only guarantees the reduction of high- level gradient dissimilarity without considering low-level part. We experimentally find that low-level weight divergence shrinks faster than high-level. Here, we show the layer-wise weight divergence in Figure 3 We choose and show the divergence of 10 layers in Figure 3 (a), and the different stages of ResNet- 18 in Figure (b). As we hope to demonstrate the divergence trend, we normalize each line with its maximum value. The results show that the low-level divergence shrinks faster than the high-level divergence. This means that reducing the high-level gradient dissimilarity is more important than the low-level.
1000
(a) Some Choosen layers.
(b) Different Parts.
Figure 3: Layer divergence of FedAvg.
(b).
CONCLUSION
In this paper, correct client drift from novel perspective of generalization contribution, which bounded by the conditional Wasserstein distance between clients' distributions. The theoretical conclusion inspires to propose decoupling neural networks and constructing similar feature distributions, which greatly reduces the gradient dissimilarity by training with shared feature distribution without privacy breach. We theoretically verify the gradient dissimilarity reduction and empirically validate the benefits of Fed IedImpr on generalization performance. Our work opens new path of enhancing FL performance from a generalization perspective. Future works may exploit better feature estimators like generative models (Goodfellow et al.] 2014; Karras et al. 2019 to sample higher-quality features, while reducing the communication and computation costs.

---

Published as a conference paper at ICLR 2024
7 ACKNOWLEDGMENT
This work was partially supported by National Natural Science Foundation of China under Grant No. 62272122, a Hong Kong RIF grant under Grant No. R6021-20, and Hong Kong CRF grants under Grant No. C2004-21G and C7004-22G. TL is partially supported by the following Aus- tralian Research Council projects: FT220100318, DP220102121, LP220100527, LP220200949, IC190100031. BH was supported by the NSFC General Program No. 62376235, Guangdong Basic and Applied Basic Research Foundation No. 2022A151 1652, HKBU Faculty Niche Research Areas No. RCCFFNRAA-IG/22-22//CCII NSFC No. 62222117 17, the Fundamental Research Funds for the Central Universities under contract WK3490000005, and KY2100000117. SHS was supported in part by the National Natural Science Foundation of China (NSFC) under Grant No. 62302123 and Guangdong Provincial Key Laboratory of Novel Security Intelligence Technologies under Grant 2022B 2212010005.
and CCF-Baidu Open Fund. XMT was supported in part by
We thank the area chair and reviewers for their valuable comments.
ETHICS STATEMENT
This paper does not raise any ethics concerns. This study does not involve any human subjects, practices to data set releases, potentially harmful insights, methodologies and applications, potential conflicts of interest and sponsorship, diseriminnttoo/biaaaaainss concerns, privacy and security issues, legal compliance, and research integrity issues.
10

---

Published as a conference paper at ICLR 2024
REFERENCES
Durmus Alp Emre Acar, Yue Zhao, Ramon Matas, Matthew Mattina, Paul Whatmough, and Venkatesh Saligrama. Federated learning based on dynamic regularization. International Conference Learning Representations, 2021.
net/ net/forum?
Id=B7v4QMR629W
Yoshua Bengio, Salem Lahlou, Tristan Deleu, Edward J. Hu, Mo Tiwari, and Emmanuel Bengio. Gflownet foundations. Journal of Machine Learning Research, 24(210):1-55, 2023. URL http?
0364.
Sameer Bibikar, Haris Vikalo, Zhangyang Wang, and Xiaohan Chen. Federated dynamic sparse training: Computing less, communicating less, yet learning better. In Proceedings of the AAAI Conference Artificial Intelligence, volume 36, pp. 6080-6088 2022.
Ilii Bistritz, Ariana Mann, and Nicholas Bambos. Distributed distillation for on-device learning. Advances in Neural Information Processing Systems, 33:22593-22604, 2020.
Kuntai Cai, Xiaoyu Lei, Jianxin Wei, and Xiaokui Xiao. Data synthesis via differentially private markov random fields. Proc. VLDB Endow., 14(11))2100-2202, jul 2021. ISSN 2150-8097, doi: 6249 3476272
0111///3447624..4777777 URL Btthstts:: http:::://ooo.ggggg///
Sebastian Caldas, Peter Wu, Tian Li, Jakub Konečný, H Brendan McMahan, Virginia Smith, and Ameet Talwalkar. Leaf: A benchmark for federated settings. arXiv preprint arXiv: 18/2000977 2018.
Hongyan Chang, Virat Shejwalkar, Reza Shokri, and Amir Houmansadr. Cronus: Robust and heteroge- neous collaborative learning with black-box knowledge transfer. arXiv preprint arXiv:1912.1127, 2019.
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.
Chatalic, V Schellekens, F Houssiau, de Montjoye, I. Jacques, and R Gribonval. Compressive learning with privacy guarantees. Information and Inference: A Journal of the IMA, 05 2021. ISSN 2049-8772. doi: 10. 00.1009//aaia//aaa0000. URL https::/doi. org 1093/ 1093/imaiai imaiai/ 1aab005 iaab005.
Hong-You You Chen and Wei-Lun Chao. On bridging generic and personalized federated learning for image classification. In International Conference Learning Representations, 2021.
Yae Jee Cho, Jianyu Wang, and Gauri Joshi. Client selection in federated learning: Convergence analysis and power-of-choice selection strategies. arXiv preprint arXiv:2010.0/243, 2020.
Liam Collins, Hamed Hassani, Aryan Mokhtari, and Sanjay Shakkottai. Exploiting shared representa- tions for personalized federated learning. Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 2089-2099. PMLR, 18-24 Jul 2021.
Xin Dong, Sai Qian Zhang. Ang Li, and HT Kung. Spherefed: Hyperspherical federated learning. In European Conference on Computer Vision, pp. 165-184. Springer, 2022.
Ron Dorfman, Shay Vargaftik, Yaniv Ben-Itzhak, and Kfrr Y. Levy. Docoft: downlink compression for cross-device federated learning. In Proceedings of the 40th International Conference on Machine Learning, 2023.
Gamaleldin Elsayed, Dilip Krishnan, Hossein Mobahi, Kevin Regan, and Samy Bengio. Large margin deep networks for classification. In Advances in Neural Information Processing Systems, 2018.
Jean- Yves Franceschi, Alhussein Fawzi, and Omar Fawzi. Robustness of classifiers uniform E, and gaussian noise. In International Conference on Artificial Intelligence and Statistics, pp. 1280-1288. PMLR, 2018.
11

---

Published as a conference paper at ICLR 2024
Amirata Ghorbani and James Zou. Data shapley: Equitable valuation of data for machine learning.
In Interaational Conference on Machine Learning, pp. 2242-2251. PMLR, 2019.
Jack Goetz and Ambuj Tewari. Federated arXiv:2008. 04489, 2020.
learning via synthetic data. arXiv preprint
Jack Goetz, Kshitiz Malik. Duc Bui, Seungwhan Moon, Honglei Liu, and Anuj Kumar. Active federated learning. arXiv preprint urXiv:I909 1909.,9999,, 2019.
lan J. Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. In Proceedings of the 27th International Conference on Neural Information Processing Systems - Volume 2, NIPS' 14, pp. 2672-2680, Cambridge, MA, USA, 2014. MIT Press.
Yongxin Guo, Xiaoying Tang, and Tao Lin. FedBR: Improving federated learning heterogeneous data via local learning bias reduction. In Proceedings of the 40th International Conference on Machine Learning.
Kartik Gupta, Marios Fournarakis, Matthias Reisser, Christos Louizos, and Markus Nagel. Quantiza- tion robust federated learning for efficient inference on heterogeneous devices. Transactions on Machine Learning Research, 2023.
Weituo Hao, Mostafa El-Khamy, Jungwon Lee, Jianyi Zhang, Kevin J Liang, Changyou Chen, and Lawrence Carin Duke. Towards fair federated learning with zero-shot data augmentation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp.
3310-3319, 2021.
Moritz Hardt and Guy N. Rothblum. A multiplicative weights mechanism for privacy-preservine data analysis. In 2010 IEEE 51 Annual Symposium on Foundations of Computer Science, pp. 61-70, 2010. doi: 10. 10.1109/0CS22000.33
Moritz Hardt, Katrina Ligett, and Frank Mesherry. A simple and practical algorithm for dif- ferentially private data release. F. Pereira, C. C. Burges, L. Bottou, and K. Q. Wein- berger (eds.), Advances in Neural Information Processing Systems, volume 25. Curran Asso- ciates, Inc., 2012. URL htttttt:///oooeeeetteeeeeeee
neurips wwwpurips..cpppppe//e...a http:ccpapp///22/EII... pdf
Chaoyang He, Murali Annavaram, and Salman Avestimehr. Group knowledge transfer: Federated learning of large cnns at the edge. In Advances in Neural Information Processing Systems 34, 2020a.
Chaoyang He, Songze Li, Jinhyun So, Mi Zhang, Hongyi Wang. Xiaoyang Wang. Praneeth Vepakomma, Abhishek Singh, Hang Qiu, Li Shen, Peilin Zhao, Yan Kang, Yang Liu, Ramesh Raskar, Qiang Yang. Murali Annavaram, and Salman Avestimehr. Fedml: A research library and benchmark for federated machine learning. arXiv preprint arXiv:2007.13516 2020b.
Chaoyang He, Alay Diliphhai Shah, Zhenheng Tang, Fan Adarshan Naiynar Sivashunmugam, Keerti Bhogaraju, Mita Shimpi, Li Shen, Xiaowen Chu, Mahdi Soltanolkotabi, and Salman Avestimehr, Fedev: A federated learning framework for diverse computer vision tasks. arXiv preprint /1.::::1:111///00 2021.
T. Hsu, Hang Qi, and Matthew Brown. Measuring the effects of non-identical data distribution for federated visual classification. ArXiv, htw/1909.000 999906939,,...1 2019.
Tzu-Ming Harry Hsu, Hang Qi, and Matthew Brown. Federated visual classification with real- world data distribution. In Computer Vision-ECCV 2020: 16th European Conference, Glasgow, UK, August 23-28, 2020, Proceedings, Part X 16, pp. 96-92. Springer, 2020.
Zhouyuan Huo, Bin Gu, and Heng Huang. Training neural networks using features replay. Advances in Neural Information Processing Systems, 31, 2018.
Sergey Ioffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In International conference on machine learning, pp. 448-456 PMLR, 2015.
12

---

Published as a conference paper at ICLR 2024
Max Jaderberg, Wojciech Marian Czamecki, Simon Osindero, Oriol Vinyals, Alex Graves, David Silver, and Koray Kavukcuoglu. Decoupled neural interfaces using synthetic gradients. In International conference machine learning, pp. 1627-1635, PMLR, 2017.
Eunjeong Jeong, Seungeun Oh, Hyesung Kim, Jihong Park, Mehdi Bennis, and Seong- Kim. Communication-effiee efficieen on-device machine learning: Federated distillation and augmentation under non-iid private data. NeurIPS, 2018.
Noah Johnson, Joseph P Near, and Dawn Song. Towards practical differential privacy for sql queries. Proceedings of the VLDB Endowment, 11(5):526-555 5399,,5336900 2018.
Peter Kairouz, H Brendan McMahan, Brendan Avent, Aurélien Bellet, Mehdi Bennis, Arjun Nitin Bhagoji, Keith Bonawitz, Zachary Charles, Graham Cormode, Rachel Cummings, et al. Advances and open problems in federated learning. arXiv preprint arXiv: 2019.
Sai Praneeth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank Reddi, Sebastian Stich, and Ananda Theertha Suresh. SCAFFOLD: Stochastic controlled averaging for federated learning. In ICMI,, 2020.
Tero Karras, Samuli Laine, and Timo Aila. A style-based generator architecture for generative adversarial networks. In Proceedings of the IEEE/CVF Conference on Computer Vision and Puttern Recognition, pp. 4401-4410, 2019.
Alex Kendall and Yarin Gal. What uncertainties do we need in bayesian deep learning for computer vision? Advances in neural information processing systems, 30, 2017.
Anastasia Koloskova, Nicolas Loizou, Sadra Boreiri, Martin Jaggi, and Sebastian Stich. A unified theory of decentralized SGD with changing topology and local updates. In Hal Daumé III and Aarti Singh (eds.), Proceedings of the 37th International Conference on Machine Learning. volume 119 of Proceedings of Machine Learning Research, pp. 5381-5393, PMLR, 13-18 Jul 2020. URL https::/proceeddnng mlr prwss/ http:///l///k/kk2aaaa.aa html
Vladimir Koltchinskii and Dmitry Panchenko. Empirical margin distributions and bounding the generalization error of combined classifiers. The Annals of Statistics, 30(1):1-55)) 2002.
Jakub Konečný, H. Brendan McMahan, Felix X. Yu, Peter Richtárik, Ananda Theertha Suresh, and
Dave Bacon. Federated Learning: Strategies for Improving Communication Efficiency. arXiv e-prints, art. 1.X.::6610005400 October 2016.
A. Krizhevsky and G. Hinton. Learning multiple layers of features from tiny images. Master's 's thesis, Department of Computer Science, University of Toronto, 2009.
Fan Lai, Xiangfeng Zhu, Harsha V. Madhyastha, and Mosharaf Chowdhury. Oort: Efficient fed-
erated learning via guided participant selection. In 15th USENIX Symposium on Operating Systems Design and Implementation (OSDI 21), pp. 19 35. USENIX Association, July 2021.
ISBN 978-1-939133-221
URL https: /wwww usenix org ottg//oofereece///s/222
presentation/lai
Fan Lai, Yinwei Dai, Sanjay Singapuram, Jiachen Liu, Xiangfeng Zhu, Harsha Madhyastha, and Mosharaf Chowdhury. Fedscale: Benchmarking model and system performance of federated learning at scale. In International Conference on Machine Learning, pp. 11814-11827, PMLR, 2022.
Chen- Yuen-Yuloe Lee, Tanmay Batra, Mohammad Haris Baig, and Daniel Ulbricht. Sliced wasserstein discrepancy for unsupervised domain adaptation. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 10285-10295 2019.
Ang Li, Jingwei Sun, Binghui Wang, Lin Duan, Sicheng Li, Yiran Chen, and Hai Li. Lotteryfl: Empower edge intelligence with personalized and communication- 2021 IEEF/ACM Symposium Édge Computing (SEC), pp. 68-79, 2021a. doi: 10.1 145/345314.. 3492909.
federated learning. In
Daliang Li and Junpu Wang. Fedmd: Heterogenous federated learning via model distillation. arXiv preprint arXiv: 1910.0358/ 2019.
13

---

Published as a conference paper at ICLR 2024
Mengxue Li, Yi-Ming Zhai, You-Wei Luo, Peng-Fei Ge, and Jhuan-Xian Ren. Enhanced transport distance for unsupervised domain adaptation. In 2020 JEEE/CVE Conference on Computer Vision and Pattern Recognition (CVPR), pp. 1993--19941, 2020a. doi: 10.1099CCVP222602220000199999999
Qinbin Li, Yiqun Diao, Quan Chen, and Bingsheng He. Federated learning non-iid data silos: An experimental study, 2021b.
Qinbin Li, Bingsheng He, and Dawn Song. Model-contrastive federated learning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 10713-10722, 2021e.
Qinbin Li, Zeyi Wen, Zhaomin Wu, Sixu Hu, Naibo Wang. Yuan Li, Xu Liu, and Bingsheng He. survey on federated learning systems: Vision, hype and reality for data privacy and protection. IEEE Transactions on Knowledge and Data Engineering, 2021d.
Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith. Federated optimization in heterogeneous networks. In Proceedings of Machine Learning and Sys- tems, volume 2, pp. 429 450, 2020b. tttp: ngs , m sys. org/ /appprr////// / Paw.aaeeee-aa.......ae Paper. pdd
2200
Xiang Li, Kaixuan Huang, Wenhao Yang. Shusen Wang. and Zhihua Zhang. On the convergence of fedavg on non-iid data. International Conference on Learning Representations, 2020c. URL
https: http::://cpeereeeee wttp//peeeeeeeeeeeo//e net/ http://wwwweeeoooemeeeeeeeeeeeee
Xuan Li, Zhanke Zhou, Jianing Zhu, Jiangchao Yao, Tongliang Liu, and Bo Han. Deepinception: Hypnotize large language model to be jailbreaker. arXiv preprint 1xXi::3311003199 2023a.
Zexi Li, Xinyi Shang, Rui He, Tao Lin, and Chao Wu. No fear of classifier biases: Neural collapse inspired federated learning with synthetic and fixed classifier. In Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV), pp. 5319-5329, October 2023b.
Xiangru Lian, Ce Zhang, Huan Zhang, Cho-Jui Hsich, Wei Zhang, and Ji Liu. Can decentralized algorithms outperform centralized algorithms? case study for decentralized parallel stochastic gradient descent. In Advances in Neural Information Processing Systems, pp. 5330-5340, 2017.
Paul Pu Liang, Terrance Liu, Liu Ziyin, Nicholas B Allen, Randy P Auerbach, David Brent, Ruslan Salakhutdinov, and Louis Phillipp Morency. Think locally, act globally: Federated learning with local and global representations. arXiv preprint arXiv:2001.01523, 2020.
Ji Lin, Wei-Ming Chen, Han Cai, Chuang Gan, and song han. Memory-efficient patch-based inference for tiny deep learning. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, 2021. URL hattps://operreiiit net./ wttorrmmmdccccpppp 7uKNp
Tao Lin, Lingjing Kong, Sebastian Stich, and Martin Jaggi. Ensemble distillation for robust model fusion in federated learning. In NeurlPS, 2020.
Zelei Liu, Yuanyuan Chen, Han Yu, Yang Liu, and Lizhen Cui. Gtg-shapley: Efficient and accurate participant contribution evaluation in federated learning. ACM Trans. Intell. Syst. Technol., 13(4), may 2022. ISSN 2157-6904.
Zemin Liu, Yuan Li, Nan Chen, Qian Wang, Bryan Hooi, and Bingsheng He. survey of imbalanced learning graphs: Problems, techniques, and future directions. arXiv preprint arXiv:2308. 13821, 2023.
Yunhui Long, Boxin Wang, Zhuolin Yang, Bhavya Kailkhura, Aston Zhang, Carl Gunter, and Bo Li. G-pate: Scalable differentially private data generator via private aggregation of teacher discriminators. Advances in Neural Information Processing Systems, 34, 2021.
Sindy Löwe, Peter O' O'Conoo, and Bastiaan Veeling. Putting an end to end-to-end: Gradient-isolated learning of representations. Advances in neural information processing systems, 32, 2019.
Bingqiao Luo, Zhen Zhang, Qian Wang, Anli Ke, Shengliang Lu, and Bingsheng He. Ai Ai-powered fraud detection in decentralized finance: A project life cycle perspective. arXiv preprint ar0iv:2308. 15992, 2023.
14

---

Published as a conference paper at ICLR 2024
Mi Luo, Fei Chen, Dapeng Hu, Yifan Zhang. Jian Liang, and Jiashi Feng. No fear of heterogeneity; Classifier calibration for federated learning with non-lID data. In A. Beygelzimer, Y. Dauphin, P. Liang. and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, 2021. URL https:ts: httpp/s/////eeevvviieee ew net/ forum? d=AF Ba____N_____h__
Enrique S Marquez, Jonathon S Hare, and Mahesan Niranjan. Deep cascade learning. /EEE fransactions on neural networks and learning systems, (9(1):5475-5485 2018.
Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. Communication-effical learning of deep networks from decentralized data. In Artificial Intelli- gence and Statistics, pp. 1273-1282, 2017.
Omar Montasser, Steve Hanneke, and Nathan Srebro. Vc classes adversarially robustly learnable, but only improperly. In Conference on Learning T'eeory pp. 2512-2530. PMLR, 2019.
Yuval Netzer, Tao Wang, Adam Coates, Alessandro Bissacco, Bo Wu, and Andrew Y. Ng. Reading digits in natural images with unsupervised feature learning. In NIPS Workshop on Deep Learning and Unsupervised Feature Learning 2011, 2011. URL http:://fldd standfrdddaaafodd..coo
boossenubeers/ripss2//
hwusenumbers.... pdf
Kang Loon Ng, Zichen Chen, Zelei Liu, Han Yu, Yang Liu, and Qiang Yang. multi-player game for studying federated learning incentive schemes. In IJCAI International Joint Conference on Artificial Intelligence. pp. 5279, 2020.
Ngoc-Hieu Nguyen, Tuan-Anh Nguyen, Tuan Nguyen, Vu Tien Hoang, Dung D Le, and Kok-Seng Wong. Towards efficient communication federated recommendation system via low-rank training. arXiv preprint arXiv:2401.04. 003748, 2024.
Arild Nokland and Lars Hiller Eidnes. Training neural networks with local error signals. In International conference on machine learning, pp. 4839-4850, PMLR, 2019.
Seungeun Oh, Jihong Park, Praneeth Vepakomma, Sihun Back, Ramesh Raskar, Mehdi Bennis, and Scong-Lyun Kim. Locfedmix-sl: sl: Localize, federate, and mix for improved scalability, convergence. and latency in split learning. In Proceedings of the ACM Web Conference 2022, pp. 3347-3357, 2022.
Xinchi Qiu, Javier Fernandez- -arauuss,,,, Pedro PB Gusmao, Yan Gao, Titouan Parcollet, and Nicholas Donald Lane. ZeroFL: Efficient on-device training for federated learning with local sparsity. In International Conference on Learning Representations, 2022.
Sashank J. Reddi, Zachary Charles, Manzil Zaheer, Zachary Garrett, Keith Rush, Jakub Konečný, Sanjiv Kumar, and Hugh Brendan McMahan. Adaptive federated optimization. In International Conference on Learning Representations, 2021. URL hat.httpss.. httttss://operevveew opwnreview. net/ forum?
id- LdLLLGF33B13US
Amirhossein Reisizadeh, Aryan Mokhtari, Hamed Hassani, Ali Jadbabaie, and Ramtin Pedarsani. Fedpaq: A communication--fficieet federated learning method with periodic averaging and quanti- zation. In International Conference Artificial Intelligence and Statistics, pp. 2021-203 PMLR, 2020.
Monica Ribero and Haris Vikalo. Communication--fficcieen federated learning via optimal client sampling. arXiv preprint 1rX11:2007.15197, 2020.
Shaohuai Shi, Xiaowen Chu, and Bo Li. Mg-wibp: Efficient data communication for distributed synchronous sgd algorithms. In IEEE INFOCOM 2019-IEEE Conference on Computer Communi- cations, pp. 172-180. IEEE, 2019.
MyungJac Shin, Chihoon Hwang. Joongheon Kim, Jihong Park, Mehdi Bennis, and Scong- Lyun Kim. Xor mixup: preprint arXiv:2006.09 000000051/4.... 2020.
data augmentation for one-shot federated learning. arXiv
Rachael Hwee Ling Sim, Yehong Zhang, Mun Choon Chan, and Bryan Kian Hsiang Low. Col- laborative machine learning with incentive-aware model rewards. In Proceedings of the 37th International Conference Machine Learning, ICML'20. JMI..orrg 2020.
15

---

Published as a conference paper at ICLR 2024
Yan Sun, Li Shen, Shixiang Chen, Liang Ding, and Dacheng Tao. Dynamic regularized sharpness aware minimization in federated learning: Approaching global consistency and smooth landscape. In Proceedings of the 40th International Conference Machine Learning. 2023a.
Yan Sun, Li Shen, Tiansheng Huang, Liang Ding, and Dacheng Tao. Fedspeed: Larger local interval, less communication round, and higher generalization accuracy. In The Eleventh International Conference on Learning Representations, 2023b.
1d=b2 xxx22KXXXKKK
Alysa Ziying Tan, Han Yu, Lizhen Cui, and Qiang Yang. Towards personalized federated learning. IEEE Transactions on Neural Networks and Learning Systems, pp. 1-17, 2022. doi: 10.1
TNNLS.2022 2022.3160699
Zhenheng Tang, Shaohuai Shi, Xiaowen Chu, Wei Wang, and Bo Li. Communication-efficieenn distributed deep learning: A comprehensive survey. arXiv preprint arXiv:2003.06307 2020.
Zhenheng Tang, Zhikai Hu, Shaohuai Shi, Yiu-ming Cheung, Yilun Jin, Zhenghang Ren, and Xiaowen Chu. Data resampling for federated learning with non-iid labels. In International Workshop on Federated and Transfer Leurning for Data Sparsity and Confidentiality in Conjunction with IJCAI 202I(FTL-IJCAI21) 2021.
Zhenheng Tang. Shaohuai Shi, Bo Li, and Xiaowen Chu. Gossipfl: A decentralized federated learning framework with sparsified and adaptive communication. IEEE Transactions on Parallel and Distributed Systems, pp. 1-13, 2022a. doi: (0.1100//DDS
Zhenheng Tang, Yonggang Zhang, Shaohuai Shi, Xin He, Bo Han, and Xiaowen Chu. Virtual homogeneity learning: Defending against data heterogeneity in federated learning. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato (eds.), Pro- ceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pp. 21111-21132. PMLR, 17-23 Jul 2022b.
Zhenheng Tang, Xiaowen Chu, Ryan Yide Ran, Sunwoo Lee, Shaohuai Shi, Yonggang Zhang, Yuxin Wang, Alex Qiaozhong Liang, Salman Avestimehr, and Chaoyang He. Fedml parrot: A scalable federated learning system via heterogeneity-ave aware scheduling on sequential and hierarchical training, arXiv preprint arXiv:2303.01778, 2023a.
Zhenheng Tang, Yuxin Wang, Xin He, Longteng Zhang, Xinglin Pan, Qiang Wang, Rongfei Zeng, Kaiyong Zhao, Shaohuai Shi, Bingsheng He, al. Fusionai: Decentralized training and deploying Iims with massive consumer-level gpus. arXiv preprint arXiv:2309.01 172, 2023b.
Rajeev Thakur, Rolf Rabenseifner, and William Gropp. Optimization of collective communi- cation operations in mpich. Int. High Perform. Comput. Appl., 19(1):49 66, feb 2005. ISSN 1094-3420. doi: 10.1177/1943342005551 1094342005051522
Chandra Thapa, Mahawaga Arachchige Pathum Chamikara, Seyit Camtepe, and Lichao Sun. Splitfed: When federated learning meets split learning. arXiv preprint arXiv:2004. 12088, 2020.
Jianyu Wang, Qinghua Liu, Hao Liang, Gauri Joshi, and H. Vincent Poor. Tackling the objective inconsistency problem in heterogeneous federated optimization. In Advances in Neural Information Processing Systems, volume 33, pp. 7611-7623, 2020a.
URL hattp::/ neurips wttppper//02////////eeppp file/ .dff
564127cu 3ww.3aaaaaaaaacec http://cccc..a.eeeeeeeeeeeeeeee-
Yulin Wang, Zanlin Ni, Shiji Song, Le Yang, and Gao Huang. Revisiting locally supervised learning: an alternative to end-to-end training. In International Conference on Learning Representations, 2020b.
Yulin Wang, Gao Huang, Shiji Song, Xuran Pan, Yitong Xia, and Cheng Wu. Regularizing deep networks with semantic data augmentation. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2021.
Blake E Woodworth, Kumar Kshitij Patel, and Nati Srebro. Minibatch vs local sgd for heterogeneous distributed learning. Advances in Neural Information Processing Systems, 13:6281-6299 2020.
16

---

Published as a conference paper at ICLR 2024
Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-mmist: novel image dataset for benchmarking machine learning algorithms. arXiv preprint arXiv:1008.07747, 2017.
Haibo Yang, Minghong Fang, and Jia Liu. Achieving linear speedup with partial worker participation in non-iid federated learning. In International Conference on Learning Representations, 2020.
Zhiqin Yang. Yonggang Zhang, Yu Zheng, Xinmei Tian, Hao Peng, Tongliang Liu, and Bo Han. Fedfed: Feature distillation against data heterogeneity in federated learning. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.
Tehrim Yoon, Sumin Shin, Sung Ju Hwang, and Eunho Yang. Fedmix: Approximation mixup under
mean augmented federated learning. In International Conference Learning Representations, 2021.
Han Yu, Zelei Liu, Yang Liu, Tianjian Chen, Mingshu Cong, Xi Weng, Dusit Niyato, and Qiang Yang. A sustainable incentive scheme for federated learning. IEEE Intelligent Systems, 35(4): 58-69, 2020. doi: 10. 109/MIS. D00.IIS..22..997777..
Honglin Yuan, Warren Richard Morningstar, Lin Ning, and Karan Singhal. What do we mean by generalization in federated learning? In International Conference on Learning Representations, 2022. URL httplttttppppeteeeeeetet
net wet.
Yufeng Zhan, Peng Li, Zhihao Qu, Deze Zeng, and Song Guo. A learning- based incentive mechanism for federated learning. IEEE Internet Things Journal, 7(7):6330-6368, 2020. doi: 10.1 109/JIOT. 2020.2967772.
Yonggang Zhang, Mingming Gong, Tongliang Liu, Gang Niu, Xinmei Tian, Bo Han, Bernhard Schölkopf, and Kun Zhang. Adversarial robustness through the lens of causality. In International Conference on Learning Representations, 2021.
Nanxuan Zhao, Zhirong Wu, Rynson W. H. Lau, and Stephen Lin. What makes instance discrimina- tion good for transfer learning? In International Conference Learning Representations, 2021. URL ptttpps://pperevvine
net 2tte:/orrormi
Yue Zhao, Meng Li, Liangzhen Lai, Naveen Suda, Damon Civin, and Vikas Chandra. Federated learning with non-iid data. arXiv preprint arXiv:1806.0055: 2018
Zhanke Zhou, Chenyu Zhou, Xuan Li, Jiangchao Yao, Quanming Yao, and Bo Han. strengthening and defending graph reconstruction attack with markov chain approximation. In ICML, 2023.
Huiping Zhuang, Zhenyu Weng, Fulin Luo, Toh Kar-Ann, Haizhou Li, and Zhiping Lin. Accumulated decoupled learning with gradient staleness mitigation for convolutional neural networks. In International Conference Machine Learning, pp. 2235-12944 PMLR, 2021.
17

---

Published as a conference paper at ICLR 2024
SUPPLEMENTARY MATERIAL
A BROADER IMPACT
Measuring Client Contribution During Local Training. As discussed in the section 2 current works mainly focus on measuring generalization contribution from clients from participating during the whole training process. We consider measuring this contribution during each communication round, which opens a new angle toward the convergence analysis of FL. Future works may fill the generalization gap between FL and centralized training with all datasets.
Relationship between Privacy and Performance. We analyze the relationship between the sharing features and the raw data in section 4.2 and Appendix. However, we do not deeply investigate how sharing features or parameters of estimated feature distribution threatens the privacy of private raw data. Sharing features at lower level may reduce gradient dissimilarity and high generalization performance of FL, yet leading to higher risks of data privacy. Future works may consider figuring out the trade-off between data privacy and the generalization performance with sharing features.
Connections of our work to knowledge distillation and domain generalization. The approximation of features generated based on the client data and low-level models can be seen as kind of knowledge distillation of other clients. More in-depth analyses of this problem would be an exciting direction, which will be added to our future works. The domain generalization is also an exciting connection to federated learning. It is interesting connect the measurements of client contribution to the domain generalization.
B MORE DISCUSSIONS
Extra computational overhead. The experiment hardware and software are described in Section F. FedImpro requires little more computing time of FedAvg in simulation. Almost all current experi- ments are usually simulated in FL community (He et al. 2020b Lai et al. [022;) Li et al. (021d). We provide a quantitative comparison between FedAvg and FedImpro based on this test setup: Training ResNet- 18 on Dirichlet partitioned CIFAR- 10 dataset with
0.1, .M = 10, sampling 5 clients
each round, with total communication rounds as 1000 and local epoch as 1.
The simulation time of the FedAvg around 8h 52m 12s, while FedImpro consumes around 10h 32m 32s. According to the Table in the main text, FedImpro achieves the 82% acc in similar cpu time with FedAvg. When E 5, the simulation time of FedAvg increases to 20h 10m 25s. The convergence speed of FedImpro with F = 1 is better than FedAvg with F = 5. Note that in this simulation, the communication time ignored, because the worldworld communication does not happen.
Extra communication overhead. Communiation round is a more important real-world metric compared to cpu time in simulation. In real-world federated learning. the bandwidth with the Internet (around 1 ~ 10 MB/s) is very low comparing to the cluster environemnt (around 10 ~ 1000 GB/s). Taking ResNet18 with 48 MB as example, each communication with 10 clients would consume 48 ~ 480s, 1000 rounds requires 48000 ~ 480000s (15 ~ 150 hours). And some other factors like the communication latency and unstability will further increase the communication time. Thus, the large costs mainly exist in communication overhead instead of the computing time (McMahan et al., 2017 Kairouz al. 2019). As shown in Table in original text, FedImpro achieve mush faster convergence than other algorithms.
The communication time in real-world can be characterized by the alpha-beta model (Thakur et al. 2005; (hi et al.) 2019): Teomm communication), 33 the inverse of the communication speed, and M the message size. In FedImpro, the estimated parameters are communicated along with the model parameters. Thus (a) the number of communications not increased, and the communication time from a will not increase: (b) The M increased as Mmodel + Mestimators where Mmodet is the model size and Mestimator is the size of estimated parameters. Taking ResNet-20 as an example, the size of estimated parameters is equal to double (mean and variance) feature size (not increased with batch size) as 4KB (ResNet-20), which is greatly less than the model size around 48 MB. Therefore, the additional communication cost is almost zero
+ 3MM where a is the latency per message (or say each
18

---

Published as a conference paper at ICLR 2024
Does the low-level model benefits from the h? The low-level model will not receive gradients calculated from the h. Thus, the low-level model is not explicitly updated by the h. However, from another aspect, training the high-level model benefits from the closer-distance features and the reduced gradient heterogeneity. FedImpro, the parooaardpropaggated gradients on original features have lower bias than FedAvg. Therefore, low-level model can implicitly benefit from h.
It non-trivial to conduct the ablation study on training low-level model with h. The implicit benefits, iee better backward gradients from the high-level models, require the high-level models are also better ("better" means that they are good for model generalization and defending data heterogeneity). Thus, the benefits high-level and low-level model are coupled. Thus, how to build direct bridge to benefit training low-level model from is still under exploring.
When label distribution the same on different clients, can we still benefit from Fed In this case, the client drift may still happen due to the feature distribution is different (Li et al. [b22bb] Kairouz et al., 2019). In this case, we may utilize the Wasserstein distance not conditioned on label unsupervised domain adaptation (Lee al. 2019 Li et al. 2020a) to analyse this problem. We will consider this case as the future work.
FedAvging the low-level model keeping the low-level model locally for each clients? FedImpro focuses on how to better learning global model instead of personalized or heterogeneous models. Thus, we conduct FedAvg on the low-level model. However, The core idea of FedImpro, sharing estimated features would be very helpful in personalized split-FL (keep the low-level model locally), with these reasons: (a) our Theorem 4. and 4.2 can also be applied into personalized FL with sharing estimated features; (b) The large communication overhead in personalized split-FI. is large due to the communication each forward and backward propagation (low-leve) and high-level models are typically deployed in clients and server, respectively (I (Liang et al.) 2020) Collins et al.) 2021; Thapa et al., 2020: Chen & Chao 202T)). Now, sharing estimated features would help reduce this communication lot by decoupling the forward and backward propagations.
C PROOF
C.I BOUNDED GENERALIZATION CONTRIBUTION
Given client m, quantify the generalization contribution, in FL systems follows: EA.L(D,.) Wa(p(O + D(,m))
(8)
where A is a pseudo gradient obtained by applying learning algorithm L(D..) distribution Pm Wa is the quantification of generalization, and D (D))) means the distribution of all clients except for client m. Theorem C.1. With the pseudo gradient obtained by L(Dm)), the generalization contribution is lower bounded:
EAL(D Wal(() A).)(D.)) ELL(P..) A), D..) [BaL [EAL W) D..)..)) Wa(p(0 A), Dm)) D(,,),
where Dm represents the dataset sampled from Dm-
Proof. To derive the lower bound, we decompose the conditional quantification of generalization, ie.. Wa(p(0 + DDmm)
Walp(0 + D(,( = P)) WWl(p[W A)), Dm,) + Wa(p(O+ +0 Dm.)
A),
Wa(p(O + D..) Wa(p(e + D..)
where we denote p as p(0 + A) for brevity and D,m stands for the dataset sampled from Dm Built upon the decomposition, we have:
BAL(P.)) D(,)) 2 EAL(P.) Wa(p A)
A), +), Wa(p(0 A,
-
- EA.L(D.) Wa(p(e W ()) D - + AA),
(10)
19

---

Published as a conference paper at ICLR 2024
The first term in Eq. 10 represents the empirical generalization performance. The second term in Eq. means that the performance gap between the model trained on sampled dataset and that trained on the distribution, rigorous analysis can be found in (Montasser et al. 2019). Note that, the first two terms are independent on the distribution D D'D) (Dm). so the focus of generalization contribution is mainly the last term, i.e., (...(....) Wa (e( A) D ()))). Wa(p A), D,,)).
The proof is relatively straightforward, long derive the upper bound of Wa(p(0 + Dm)) and Wa(p(e + Dm))) For Wa(p(O + A), Dm). we have:
Wa(p(0 + A), mm)
-1 BO-Do
inf argmax, ((O:z)) d(z, x')
-R(,lu) (
inf
d(r, x²)
ungmax,, p(0;x)) ):7
inf (10( )oo aremax,
d(x', x²") + d(x, x")
inf (1)(- )o,, argmax, ((8:"))
d(x', x (l)))~D Ble, d(x, x")
= -E (10))> By" Pm my
inf argmax, ((O:r))
d(r', x²) + (1))) B (( d( x"),
where Jy stands for the optimal transport between the conditional distribution Dm and D D Similarly, we have:
Walp(O A), SB (-1)))D E inf
d(x'x²")
BB. ,)D d(x,
Combining these inequality, we have: Wa(p(0 A), Dm)) Wa(p(0 D() <2Ca(D D)Dm))
+ max {Ó(Dm, D (Dm), (DD D)Dm)
(11)
where
8(Dm,, D)D,,. =B BBp-lu)DD.. E. DD Pimlo
inf d(a', x²))
inf B(le) Dm Di angman, p(0:z
d(x', ((x'xx"))
and
= -Er (((,b) En Didl inf ppm=))((0==))))) d(x', x²))
E D) BD argmas,p(0:=))e...... ((
inf d(x' d(x',a").
The upper bound is straightforward. For example, if the label distributions the same, i.e. y ~ D(Dm is equal to y ~ Dm, we have:
Walp(O + A), D.)) Wa(p(O+ D.D,.) (C,
According to Eq. 11 the last term in Eq. 10 is bounded:
EAL(RR) W7 (p( + A), Pm)) A))
(12)
<EAL..D) Wa(p((( + A), Dm.)) - W2(p(O A), Dm)))
which is further upper bounded by conditional Wasserstein distance when the label distributions are the same:
EALPP.) Wa + A), W2 (),, < (C,,D D D(.)..) + max (6(D., D(Dm), (D,n, D(mm)
(13)
Thus, the label distribution will have additional impact the bound. If the label distributions the same, then we have
EALTP Wa(p(0 D..)) Wa(p(0 + D.,)) (C. D(D...).
(14)
which completes the proof.
n
20

---

Published as a conference paper at ICLR 2024
C.2 DERIVATION OF DECOUPLING GRADIENT VARIANCE
The derivation of Equation 3
Because vm = {V Veow mm Pm E Rd,, Viroe Sm E lRd and
V Onsignmm E TR we have
(( (0
(15)
d
fm (0:2,0) -
V m (0; v) VF(0)()) +
(VIm(0:2,1)))) vF()))))
i=d++1
-B( V fm (0: y) V 8tow F(O) + V Vnigh (m 2, y) - V Vonigh F(0]
The derivation of Equation 4 Assuming multi-layers neural network consists L linear layers, each of which is followed by an activation function. And the loss function is C'E()). The forward function can be formulated as:
f(0, 100
(16)
Then the gradient on I-th weight should be:
0ʃ
0s
(z)) Izn θTn-1 0 Ozn-1 0741(2441)
Ozi
9i
00,
Ör (2n)
den (Tn-1(2n-1) dzn-1 0Tn-2(2n-2) 02411
08,
(17)
0f
(2(2n) (2n- 1.. ((22)))))))))))
(18)
(2n)
0s
II r((:)))) ((41(242)7()) 11-4+2
(19)
(zn)
in which I ni zt, is the weight, activation function, output of the 1-th layer, respectively. Thus, can see that the gradient of I-th layer is independent of the data, hidden features, and weights before I-th layer if we directly input a z (-th layer.
C.3 PROOF OF THEOREM
We restate the optimization goals of using the private raw data (z,z) 3) of clients and the shared hidden features h H y as following:
M
min F(0) = DeRd
Prm Fule 10 h,
E Pmr F.e.DDD. [M(;a,u) S(;: (.,hhh w)]
ha7H
hw7t
(20)
Theorem C.2. Under the gradient variance measure CGV (Definition 4.2). with i. satisfying
N Ñ+ the objective function F(0) causes tighter bounded gradient dissimilarity. i.e.,
Rom +ho
the CGV(F. B(2,))DD..l I| Ver.o Im (0; x, Von. F(0) + (N+NY=
Naa
1I Vnegr I' (0; x 1)
Vorigh 2(0) CG)(F,
Proof.
CGV(F,0) = -E (0; v ((0)112
hn7
-B Dm [| Blo (0; Vew F3
+ B( hw7
|I| Onege, Jm (0; x. 3) + V Vonga smn (0; 1) - Vonign F(O)12. (21)
(22)
21

---

Published as a conference paper at ICLR 2024
On m-th client, the number of samples of (z., y) is nm and the hm is fim- Then the high-level gradient variance becomes:
Pum
am V Brsgp fm (0: h v) Venngn F(
E (r.y))D,, hw7
V Vnnghh fm (, +
7m
Am
Tmm + fim
= (a.y)>DDm hwHt
Rom
V Bnigh Soassmm(Bee am(00
im
V Vusge fm (0: h v)
Tmm
am
Nmm Pem
M
7m + fm
im
V Vnnggn Sm (0;
V Onigh Im (0: h v))
N Am imm
nm
pyy
m=1
=IE (m.
Bhigh Jm (0; N0
Nmm V Vanigh Sm mn (0; y)
Nmn + Pom
N + N
N2 (N Ny2
E D V Vision Sm (0;
Tym
V Bhigh S (0; x. y)1l
(23)
Combining Equation 23 and
21 we obtain
CGV(F 0) E (z.z))
V Blo fm (0; - V Velow Nlow F(0)112
N
Onioo fm (0: x, v) Vonigh F(0)112
(N N)2
which completes the proof.
For the convergence analysis, there have been many convergence analyses of FedAvg from gra- dient dissimilarity viewpoint (Woodworth et al. 2020: Lian et al.] 2017 Karimireddy et al. (2020). Specifically, the convergence dissimilarity plays a crucial role in the bound. In this work, we propose a novel approach inspired by the generalization view reduce the gradient dissimilarity, thus provide tighter bound regarding the convergence rate. This is consistent with our experiments, see Table 2
is upper bounded by many factors, among which the gradient
C.4 INTERPRETING AND CONNECTING THEORY WITH ALGORITHMS
We summarize our theory and how it motivates our algorithm as following.
Our work involves two theorems, i.e., Theorem and Theorem 4.2 where Theorem motivates the proposed method and Theorem 4.2 indicates another advantage of the proposed method.
Here, Definition measures the model's generalizability data distribution from the margin theory, and Definition 4.1 is the conditional Wasserstein distance for two distributions.
Theorem
is built upon two definitions, i.e., Definition 3.1 and Definition 4.
Theorem 4.1 shows that promoting the generalization performance requires constructing similar conditional distributions. This motivates our method, i.e., aiming at making all client models trained on similar conditional distributions to obtain higher generalization performance.
In our work, inspired by Theorem 4.1 we regard latent features as the data discussed in Theorem 4.1 Accordingly, we can construct similar conditional distributions for the latent features and train models using these similar conditional distributions.
Theorem 4.2 is built upon Definition
4.2 where gradient dissimilarity in FL quantitatively
measured by Definition 4.2
Theorem 4.2 shows that our method can reduce gradient dissimilarity to benefit model convergence.
For advanced architectures, e.g., a well-trained GFlowNet (Bengio et al. 2023). reducing the distance in the feature space induce the distance in the data space. However, the property could be hard to maintain for other deep networks, e.g., ResNet.
22

---

Published as a conference paper at ICLR 2024
D MORE RELATED WORK
D.1 ADDRESSING NON-IID PROBLEM IN FL
The convergence and generalization performance of Federated Learning (FI.) McMMhan et al. 2017 suffers from the heterogeneous data distribution across all clients (Zhao et al. 2018 Li et al. 2020c Kairouz et al., 2019). There exists severe divergence between local objective functions of clients, making local models of FL diverge (Li et al. 2020b Karimireddy et al. 2020). which is called client drift.
Although researchers have designed many new optimization methods to address this problem. is still open problem. The performance of federated learning under severe Non-IID data distribution The previous methods that address Non-III data problems can be classified into the following directions.
is far behind the centralized training.
Model Regularization focuses on calibrating the local models to restrict them not to be excessively far away from the server model. A number of works (Li et al. 2020b Acar et al.. 2021;; Karimireddy et al. 2020) add a regularizer of local-global model difference. FedProx (Li et al. (222bb) adds a penalty of the L2 distance between local models to the server model. SCAFFOLD (Kariirredyy et al. 2020) utilizes the history information to correct the local updates of clients. FedDyn (Acar et al. 2021) proposes to dynamically update the risk objective to ensure the device optima is asymptotically consistent. FedIR (Hsu et al. 2020) applies important weight to the client's local objectives to obtain an unbiased estimator of loss. MOON (Li et al. 2021c) adds the local-global contrastive loss to learn similar representation between clients. CCVR (Luo et al. 202I) transmits the statistics of logits and label information of data samples to calibrate the classifier. FedETF et al. 2023b) proposed synthetic and fixed ETF classified to resolve the classifier delemma, which orthogonal to method. SphereFed (Dong et al., 2022) proposed constraining learned representations of data points to be unit hypersphere shared by clients. Specifically, in SphereFed, the classifier is fiexed with weights spanning the unit hypersphere, and calibrated by a mean squared loss. Besides, SphereFed discovers that the non-overlapped feature distributions for the same class lead weaker consistency of the local learning targets from another perspective. FedImpro alleviate this problem by estimating and sharing similar features.
Reducing Gradient Variance tries to correct the local updates directions of clients via other gradient information. This kind (Wang et al.) 2020al| Hsu et al.) 2019:; Reddi et al. (202]) of methods aims to accelerate and stabilize the convergence. FedNova (Wang et al.) (020a) normalizes the local updates to eliminate the inconsistency between the local and global optimization objective functions. Adjusting data sampling order in local clients is verified to accelerate convergence (Tang et al.) 2021). FedAvgM (Hsu et al., 2019)) exploits the history updates of the server model to rectify clients' updates. FEDOPT (Reddi et al. [202 proposes a unified framework of FL. It considers the clients' updates as the gradients in centralizer training to generalize the optimization methods in centralized training into FL. FedAdaGrad and FedAdam are FL versions of AdaGrad and Adam. FedSpeed (Sun et al.) (023b) utilized a prox-correction term on local updates to reduce the biases introduced by the prox-term, as a necessary regularizer to maintain the strong local consistency. FedSpeed further merges the vanilla stochastic gradient with a perturbation computed from an extra gradient ascent step in the neighborhood, which can be seen as reducing gradient heterogeneity from another perspective.
Sharing Features. Personalized Federated Learning hopes to make clients optimize different personal models learn knowledge from other clients and adapt their own datasets (Tan al. 2022). The knowledge transfer of personalization is mainly implemented by introducing personalized
parameters (Liang al. 2020 Thapa et al. 2020 Li et al., 2021a). or knowledge distillation (He
et al. 2020a]] |in et al.] 2020 Li & Wang 2019 Bistritz. et al. 2020) on shared local features or
extra datasets. Due to the preference for optimizing local objective functions, however, personalized
federated models do not have comparable generic performance (evaluated on global test dataset) to normal FL (Chen & Chao 2021). Our main goal is to learn a better generic model. Thus, we omit comparisons to personalized FL algorithms.
Except Personalized Federated Learning, some other works propose to share features to improve federated learning. Cronus (Chang al. (2019) proposes sharing the logits to defend the poisoning attack. CCVR (Luo et al.) 2021) transmit the logits statistics of data samples to calibrate the last layer of Federated models. CCVR (Luo et al. 2021) also share the parameters of local feature distribution.
23

---

Published as a conference paper at ICLR 2024
However, do not need to share the number of different labels with the server, which protects the privacy of label distribution of clients. Moreover, our method acts as framework for exploiting the sharing features to reduce gradient dissimilarity. The feature estimator does not need to be the Gaussian distribution of local features. One may utilize other estimators or even features of some extra datasets rather than the private ones.
The original cause of client drift is data heterogeneity. Some researchers find that sharing part of private data can significantly improve the convergence speed and generalization performance (Zhao et al. 2018), yet it sacrifices the privacy of clients' data.
Sharing Data.
Thus, both reduce data heterogeneity and protect data privacy, series of works (Hardt & Rothblun um
2010; Hardt et al. 2012 Chatalic et al., [021] Johnson et al., 2018 11 Cai et al. [221]] Li et al. 2203a,
Yang et al. 2023) add noise on data to implement sharing data with privacy guarantee some degree.
Some other works focus on sharing a part of synthetic data(Jeong et al.
2018 Long et al. 2021
Goetz & Tewaril 2020 Hao et al., (02T) or data statisties (Shin et al., 2020 Yoon et al. 2021 to help
reduce data heterogeneity rather than raw data.
FedDF (Lin et al. 2020) utilizes other data and conducts knowledge distillation based on these data to transter knowledge of models between server and clients. 'hee core idea of FedDF is to conduct finetuning on the aggregated model via the knowledge distillation with the new shared data.
Communication compressed FL Communication compression methods aim to reduce the commu- nication size in each round. Typical methods include sparsifying most of unimportant weights (Dorf man et al. 2023 Bibikar et al. 2022; Qiu et al. 2022: Tang et al. 2022a 2020, 2023b). quantizing model updates with fewer bits than the conventional 32 bits (Reisizadeh et al. 2020 Gupta et al. 2023), and low-rank decomposition to communicate smaller matrices (guuyen et al.]| [2024] Konečny et al. 2016). Despite reducing communication costs, these methods are not as cost-effective as one- shot HL. and our methods, due to the necessary of enormous communication rounds for convergence.
D.2 MEASURING CONTRIBUTION FROM CLIENTS
Generalization Contribution. Clients are only willing to participate a FL training when given enough rewards. Thus, it is important to measure their contributions to the model performance Yu
et al.
2020
Ng et al. 2020 Liu et al. 2022, Sim al. 2020).
There have been some works (Yuan et al. 2022 Yu et al. 2020 Ng et al. 2020 Liu et al. 2022
Sim et al.] 2020) proposed to measure the genera generalization lizaiin contribution from clients in FL. Some
works (Yuan et al. 2022) propose to experimentally measure the performance gaps from the unseen
client distributions. Data shapley (Ghorbani & Zou 2019: Yu al., 2020 Luo et al. 2023 Liu et al.
2023) is proposed to measure the generalization performance gain of client participation. Liu et al. 2022) improves the calculation efficiency of Data Shapley. And there some other work that proposes to measure the contribution by learning-based methods (Zhan et al. 2020). Our proposed questions are different from these works. Precisely, these works measure the generalization performance gap with or without some clients that never join the collaborative training of clients. However, we hope to understand the contribution of clients at each communication round, Based on this understanding, we can further improve the FL training and obtain a better generalization performance.
has been empirically verified that a large number of selected clients introduces new challenges to optimization and generalization of FL (Charles et al. (221), although some theoretical works show the benefits from (Yang et al.) 2020). This encourages us to understand what happens during the local training and aggregation. Causality is also a potential way to explore the client contribution (Zhng et al.] 2021).
Client Selection. Several works (Cho et al., 2020 Goetz et al., 2019 Ribero & Vikalo, 2020 Lai
et al. 2021) propose new algorithms to strategically select clients rather than randomly. However, these methods only consider the hardware resources or local generalization ability. How local training affects the global generalization ability has not been explored.
D.3 SPLIT TRAINING
To efficiently train neural networks. split training instead of end-to-end training is proposed to break the forward, backward, model updating dependency between layers of neural networks.
24

---

Published as a conference paper at ICLR 2024
Table 4: Demystifying different FL. algorithms related the sharing data and features.
Shared Thing
Low-level Model
Objective
Chatalic et al.
2021
Cai al. 202
Raw Data With Noise
Shared
Others
Long cral
2021
Hio crile 2044
Params, of Data Generator
Shared
Global Model Performance
Yoon 55
2021
(UXU)
STAL of raw Data
Shared
Global Model Performance
etb 2021
STAT. of Logis, Label Distribution Hidden Features legits
Shared
Global Model Performance
Chane
2019
Shared
Defend Poisoning Attack
Li& Wang 2019; ISstiiiiz al.
2020
Private
Personalized FL
20120 Laing
Hidklen Features
Private
Personalized Fl.
napa of dr. 2020 Oi Oftera craie 2022
Hidden Features
Shared
Accelerate Training
FedImpro
Params, Estimated Feat. Distribution
Shared
Global Model Performance
Note: "STAT." means statistic information, like mean or standard deviation, "Feat." means hidden features, "Params." means parameters.
To break the backward dependency on subsequent layers, hidden features could be forwarded to another loss function to obtain the Local Error Signals (Marquez et al., 2018 Nøkland & Eidnes 2019 Löwe al. 2019 Wang et al., 2020b Zhuang et a 2021). 1). How to design suitable local error still remains as an open problem. Some works propose to utilize extra modules to synthesize gradients (Jaderberg et al. 2017). so that the backward and updates of different layers can be decoupled. Features Replay (Huo et al.) 2018) is to reload the history features of the preceding layers into the next layers. By reusing the history features, the calculation on different layers could be asynchronously conducted.
Some works propose Split FL (SFL) to utilize split training to accelerate federated learning (Oh et al., 2022, Thapa et al., 2020). In SFL, the model is split into client-side and server-side parts. At each communication round, the client only downloads the client-side model from the server, conducts forward propagation, and sends the hidden features to the server for computing loss and backward propagation. This method aims to accelerate FL's training speed on the client side and cannot support local updates. In addition, sending all raw features could introduce high data privacy risk. Thus, omit the comparisons to these methods.
We demystify different FL algorithms related to the shared features in Table
E DETAILS OF EXPERIMENT CONFIGURATION
E. HARDWARE AND SOFTWARE CONFIGURATION
Hardware and Library. We conduct experiments using GPU GTX-2080 Ti, CPU Intel(R) Xeon(R) Gold 5115 CPU @ 2.40GHz. The operating system is Ubuntu 16.04.6 LTS. The Pytorch version is 1.8.1. The Cuda version is 10.2.
Framework. The experiment framework is built upon a popular FL. framework FedML (He et al. 20205 Tang et al. (023aa). And we conduct the standalone simulation to simulate FL. Specifically, the computation of client training is conducted sequentially, using only one GPU, which a mainstream
experimental design of FL (He et al. 2020b Sun et al. 2023b Reddi et al. 2021: Luo et al. 202T
Liang et al. 2020; Li et al. 22021b] Caldas et al.J 2018). due it is friendly to simulate large number of clients. Local models are offloaded to CPU when it is not being simulated. Thus, the real-world computation time would be much less than the reported computation time. Besides, the
communication does not happen, which would be the main bottleneck in real-world FL..
E.2 SYPER-PARAMERE
The learning rate configuration has been listed in Table We report the best results and their learning rates (grid search in {0.0001, 0.001, 0.01, 0.1,
And for all experiments, we use SGD as optimizer for all experiments, with batch size of 128 and weight decay of 0.0001. Note that we set momentum as 0 for baselines, as we find the momentum of 0.9 may harm the convergence and performance of FedAvg in severe Non-IID situations. We also report the best test accuracy of baselines that are trained with momentum of 0.9 in Table 6 The client-side momentum in FL training does not always commit better convergence because the momentum introduces larger local updates, increasing the client drift, which is also observed in
25

---

Published as a conference paper at ICLR 2024
recent benchmark (He et al. 2021). And the server-side momentum (Hsu al. 2019) may improve the performance. The compared algorithms including FedAvg, FedProx, SCAFFOLD, FedNova do not the server-side momentum. For the fair comparisons we did not use the server-side momentum for all algorithms.
For K 10 and K 100, the maximum communication round is 1000, For K 10 and E the maximum communication round 400 (due to the 5 increase the calculation cost). The number of clients selected for calculation 5 per round for K 10, and 10 for K 100.
Table 5: Learning rate all experiments.
FL Setting
Dataset
FedAvg
FedProx
SCAFFOLD
FedNova
FedImpro
a
E
K
0.1
1
10
0.1
0.1
0.1
0.1
0.1
0.05
10
0.1
0.1
0.01
0.1
0.1
CIFAR-10
0.1
0.1
1
100
0.1
0.1
0.01
0.1
0.1
0.1
1
10
0.1
0.1
0.1
0.1
0.1
0.05
1
10
0.1
0.1
0.001
0.1
0.1
FMNIST
0.1
5
10
0.1
0.1
0.1
0.1
0.1
0.1
100
0.1
0.01
0.1
0.1
10
0.1
0.01
100
0.1
0.05
1
10
0.1
0.1
0.01
0.1
0.1
SVHN
0.1
5
10
0.1
0.01
0.01
0.01
0.1
0.1
1
100
0.1
0.1
0.001
0.1
0.1
0.1
10
0.1
0.1
0.1
0.1
0.1
0.05
0.1
0.1
0.1
0.1
CIFAR-100
0.1
5
10
0.1
0.1
0.1
0.1
0.1
0.1
1
100
0.1
0.1
0.1
0.1
0.1
E.3
MORE CONVERGENCE FIGURES
Except the Figure 2 in the main paper, we provide more convergence results as Figures 4 5 and
These results show that method accelerate FL training and obtain higher generalization performance.
Moddl Ann M
Aoll
223
125
800
130 130
300
300
300
100
199
40
6.6
852
2000
204
652
1200
(a)
0.1, K 10. (b)
0.1, K
(0
0.1, K 100.
0.05, K 10.
F7
E
E=1
Figure 4: Convergence comparison of CIFAR-10.
(8)
0.1. K
10, 0.1, (0) 0.1. 100.
0.05, 10.
E 5
E =
Figure 5: Convergence comparison of FMNIST.
26

---

Published as a conference paper at ICLR 2024
Table 6: Baselines with Momentum-SGD.
FL Setting
Dataset
FedAvg
FedProx
SCAFFOLD
FodNova
E
K
0.1
1
10
79.98
83.56
83.58
81.35
0.05
1
10
69.02
78.66
38.55
64.78
CIFAR-10
0.1
5
10
84.79
82,18
86.20
86,09
0.1
1
100
49.61
49.97
52.24
46.53
0.1
1
10
86.81
87.12
86.21
86.99
0.05
1
10
78.57
81.96
76.08
79.06
FMNIST
0.1
5
10
87.45
86.07
87.10
87.53
0.1
100
90.11
90.71
85.99
87.09
10
88.56
86.51
80.61
89.12
0.05
1
10
82.67
78.57
74.23
82.22
SVHN
0.1
5
10
87.92
78.43
81.07
88.17
0.1
1
100
89.44
89.51
89.55
82.08
0.1
10
67.95
65.29
67.14
68.26
0.05
1
10
62.07
61.52
59.04
60.35
CIFAR-100
0.1
5
10
69.81
62,62
70.68
70.05
0.1
1
100
48,33
48.14
51.63
48.12
(8)
0.1, (b,
0.1,
(c)
0.1,
100.
0.05, 10.
E 5
E
Figure 6: Convergence comparison of SVHN.
(a)
0.1, K 1b, 0.1, 10.(c)
0.1, 100, (d)
0.05,
10.
E=1
=
E ==1
E=1
Figure 7: Convergence comparison of CIFARI00.
F ADDITIONAL EXPERIMENTS
80
0.53
0..05
Dessseed Learning
0.00
10
20.0
4000
60000
80000
10500
400
600
1050
Round
Round
(a) Long Training Time
(b) Our method with different de-
grees noise
Figure 8: Additional experiments on CIFAR-I0 with a 0.1, K 10, =
27

---

Published as a conference paper at ICLR 2024
0.10
2.0
1.0
joev
21d
Diergence
0.28
Weeph Cheegeete
2.B
0.8
0.6-
0.6-
24t
240 coe
0.24
12.4
2011 crey
2.4
ALB giey
44th coev
0.-
0.00
450
850
0.0
0.0
1000
Round
Round
Round
ResNet-11 with FMNIST
(b) ResNet--18 with SVIIN
(c) ResNet-50 with CIFAR-100
Figure 9: Layer divergence of FedAvg.
Table
Test Accuracy of our method with different degrees of noise.
Hr
0.0
0.001
0.005
0.01
0.05
0.1
Test Accuracy (%)
88.45
88.43
88.23
88.26
88.07
88.11
88.3
Table 8: Test Accuracy of different algorithms on FEMNIST.
FedAvg
FedProx
FedNova
FedImpro
Test Accuracy (%) Comm. Round to attain the Target Acc.
80.83
79.70
68.96
82.77
82
NaN
NaN
45
F.1 TRAINING WITH LONGER TIME
To demonstrate the difficulty of optimization of FedAvg in heterogeneous-data environment, we show the results of training 10000 rounds, as shown in Figure (a). During this 10000 rounds, the highest test accuracy of FedAvg with fixed learning rate 88.5% and of the FedAvg with decayed learnign rate is 82.65%. Note that we set the learning rate decay exponentially decay at each communication round, wich rate 0.997. Even after 2000 rounds, the learning rate becomes as the around 0.0026 times the original learning rate. The results show that the longer training time cannot fill the generalization performance gap between FedAvg and centralized training, encouraging us to develop new optimization schemes improve it.
10w. 203
10.rd
Figure 10: Convergence comparison of FEMNIST with 3400 clients.
F.2 SHARING ESTIMATING PARAMETERS WITH NOISES OF DIFFERENT DEGREES
To enhance the security the sharing feature distribution, we add the noise e ~ (O, σr) on the d. and Pm. The privacy degree could be enhanced by the larger pe. We show the results of our method with different a. in Figure & (b) and Table 7 The results show that under the high perturbation the estimated parameters, method attains both high privacy and generalization gains.
F.3 MORE EXPERIMENTS OF THE REAL-WORLD DATASETS
To verify the effect of our methods on the world FL datasets, we conduct experiments with Federated (FENNS() (Caldas et al.) 2018 He et al. 2020b). which has 3400 users, 67 1585 training samples and 77483 testing samples. We sample 20 clients per round, and conduct local training with 10 epochs. We search the learning rate for algorithms in (0.01, 0.05, 0.1} and find the 0.05 is the best for all algorithm. Figure 10 and Table 8 show that method converges faster
28

---

Published as a conference paper at ICLR 2024
and attains better generalization performance than other methods. Note that the SCAFFOLD is not included the experiments, as it has a very high requirement (storing the control variates) of simulating 3400 clients with few machines.
More Results the Layer-wise Divergence. We conduct more experiments of the layer divergence of FedAvg with different datasets including FMNIST, SVHN and CIFAR-100, training with ResNet- 18 and ResNet-50. As BanureBand@Jsoows the divergence of the low-level model divergence shrinks faster than the high-level. Thus, reducing the high-level gradient dissimilarity more crucial than the low-level.
F.4 MORE RESULTS OF DIFFERENT MODEL ARCHITECTURE
To verify the effect of our method on different model architectures, we conduct additional experiments of training VGG-9 on CIFAR-10, instead of the ResNet architectures. The experimental results shown in Table 9 The target accuracy is 82%. We can see that our method can outperform baseline methods with different architectures.
Table Test Accuracy different algorithms with VGG-9 on CIFAR-10.
FedAvg
FedProx
SCAFFOLD
FedNova
FedImpro
Test Accuracy (%)
82.58
82.92
82.43
82.91
84.52
Comm. Round attain Target Acc.
836
512
664
F.5 MORE BASELINES
We further compare our method with FedDyn Acar al. 2021), CCVR (Luo al. 2021 Fed- Speed (Sun et al. 2023b) and FedDF (Lin et al. 2020). The results are given in Table TO as below. As original experiments in these works do not utilize the same FL setting, we report their test accuracy of CIFAR- 10/100 dataset and the according settings in original papers. To make the comparison fair, we choose the same or the harder settings of their methods and report their original results. The larger dirichlet parameter means more severe data heterogeneity. The higher communication round means training with longer time. We can see that our method can outperform these methods with the same setting even the more difficult settings for us.
Table 10: Comparisons with more baselines.
Method FedDyn FedSpeed
Dataset
Dirichlet
Client Number
Communication Round
Test ACC.
0.3
1400
77.33
0.3
1400
82.68
CIFAR-10
FedImpro
0.3
100
1400
85.14
FedImpro
0.1
100
1400
82.76
CCVR
0.1
100
62.68
CIFAR-10
FedDF (no extra data)
0.1
100
38.6
FedImpro
0.1
100
68.29
SphereFed
0.1
10
100
69.19
CIFAR-100
FedImpro
0.1
10
100
70.28
F.6 MORE ABLATION STUDY ON NUMBER OF SELECTED CLIENTS
The feature distribution depends on the client selected in each round. Thus, to analyze the number of clients on the feature distribution estimation, we conduct ablation study with training CIFAR-10 datasets with varying the number of selected clients. Table 11 shows the effect of varying the number of selected clients per round. More clients can improve both FedAvg and FedImpro. The FedImpro can work well with more selected clients.
F.7 RECONSTRUCTED Raw DATA BY MODEL INVERSION ATTACKS
We utilize the model inversion method (Zhao et al. 2021 Zhou et al. 2023) used in (Luo et al. 2021 verify the privacy protection of our methods. The original private images shown in Figure
29

---

Published as a conference paper at ICLR 2024
a
E
M
M
FedAvg
Fed
0.1
1
10
83.65+2.00
88.45:10.43
0.1
1
10
10
8..65=1.226
90.25+10.3
0.05
1
10
75.36+1.92
81.75++1.03
0.05
1
10
10
78.36=1.58
85.75+11.21
0.1
5
85.69+0.57
88.10::0.20
0.1
5
10
10
87.92+0.31
0.22
0.1
1
100
10
73.42+1.19
77.56:11.02
0.1
1
100
20
77.82+0.94
85.39-11.15
Table 11: Ablation study with training on CIFAR-10 on different number of selected clients.
reconstructed images using the features of each image are shown in Figure 12 reconstructed images using the mean of features of all images are shown in Figure 13 Based on the results, we observe that feature inversion attacks struggle to reconstruct private data using the shared feature distribution parameters om and Pm- Thus, our method is robust against model inversion attacks.
Figure 11: Original private images that fed into the model.
Figure 12: Reconstructed images using the raw features. Note that each image is reconstructed by the feature of each private image, i.e. sharing the features. shows that the feature inversion method can successfully reconstruct the original image based on the raw features.
(a) No Noise
(b) 0.001
(c) 0.01
(d) 14 = 0.1
(c) p4 = 0.5
Figure 13: Reconstructed images using the mean of features (our methods) of all images with noises of different degrees. Me is the variance of the Gaussian noise. Now, feature inversion method cannot reconstruct the original images..
30


Reviews:
[{"review_id": "0kHHzMiTe3", "reviewer": "Reviewer_Kzyf", "summary": "The paper addresses the issue of client drift arising in heterogeneous federated settings. As shown by recent literature, local models drifting away from the convergence points of the global model lead to slower and unstable convergence. This work first shows how the generalization contribution is bounded by the conditional Wasserstein distance between clients’ local data distributions. To reduce the client drift, FedImpro decouples local training into two stages: the lower part of the model is trained on the local data, while the higher part is also trained on shared reconstructed features. Empirical results on FL benchmark datasets show the efficacy of FedImpro in terms of final performance, convergence speed, and reduced weights divergence.", "questions": "1. How does FedImpro compare with more recent FL methods addressing data heterogeneity? Please refer to Weaknesses for examples on related works.\n2. How much bigger is FedImpro's impact on the communication both in terms of number of communications between clients and server, and message size w.r.t. FedAvg and the other introduced FL baselines?\n3. How does FedImpro compare with CCVR and SphereFed w.r.t. centralized performances?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The paper addresses a very relevant issue for the FL community, i.e. limiting the negative effects of the client drift in heterogeneous settings.\n- The paper is well written and easy to follow\n- Very detailed discussion of related works\n- Theoretical claims supported by proofs\n- Extensive empirical analysis. FedImpro is compared with some state-of-the-art approaches in terms of final performance, convergence speed, weight divergence. Interesting ablation study on the depth of gradient decoupling. Results compared under different levels of client participation, local training epochs, data heterogeneity, and settings (cross-device vs cross-silo FL).\n- Limitations regarding communication and privacy concerns are explicitly addressed. It is shown how FedImpro does not lead to privacy leakage in the gradient inversion attack.\n- All details for reproducing the experiments are provided.", "weaknesses": "- My main concern regards the feasibility of deploying FedImpro in real-world contexts. FedImpro notably increases both the number of communications between clients and server, and the message size. The paper points out how the global distribution can be estimated using methods which impact the communication network less, but that does not eliminate the need for additional communication. \n- Some relevant related works are not discussed: ETF [1], SphereFed [2], FedSpeed [3]. \n- FedImpro is compared with relatively old baselines (published 3 years ago). More recent and relevant baselines are for instance FedDyn [4], FedSpeed, ETF. Also, I believe CCVR and SphereFed are highly correlated with FedImpro and it is relevant to understand how they behave w.r.t. FedImpro in terms of client drift reduction, final performance, communication costs, even if they are only applied at the end of training.\n- From Fig. 2, the actual impact of FedImpro on reducing the weight divergence appears very limited.\n\n**Minor weaknesses and typos:**\n- Lack of NLP datasets, e.g. Shakespeare and StackOverflow, or large scale datasets, e.g. Landmarks-User-160k.\n- Typo (?) page 7: I believe the explanation of cross-device or cross-silo setting should be switched, i.e. M=100 for cross-device FL and M=10 for cross-silo FL.\n\n**References**\n[1] Li, Zexi, et al. \"No Fear of Classifier Biases: Neural Collapse Inspired Federated Learning with Synthetic and Fixed Classifier.\" ICCV (2023).\n[2] Dong, Xin, et al. \"Spherefed: Hyperspherical federated learning.\" European Conference on Computer Vision. Cham: Springer Nature Switzerland, 2022.\n[3] Sun, Yan, et al. \"Fedspeed: Larger local interval, less communication round, and higher generalization accuracy.\" ICLR (2023).\n[4] Acar, Durmus Alp Emre, et al. \"Federated learning based on dynamic regularization.\" ICLR (2021).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "mZaVJChRqf", "reviewer": "Reviewer_AtvR", "summary": "This paper proposed FedImpro, a framework which constructs similar conditional distributions for local training. And FedImpro decouples the model into high-level and low-level parts and trains the high-level part on reconstructed feature distributions, causing promoted generalization contribution and alleviated gradient dissimilarity of FL.", "questions": "See Weaknesses.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "1. This paper has a good level of writing and it is easy to follow. The idea is easy to follow and understand.\n2. This paper combine split training with feature sharing to improve the generalization of the model.", "weaknesses": "1. I notice that the author ignore a very related and state-of-art baesline FedDyn [1], could the author conduct comparion experiments with FedDyn?\n\n2. The timecomsuming for training the model increases for FedImpro. Could the author list the cpu-time cost comparion experiments to reach the target accuracy?\n\n\n[1] Acar, Durmus Alp Emre, et al. \"Federated learning based on dynamic regularization.\" *arXiv preprint arXiv:2111.04263* (2021).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "KK9LD4PAuK", "reviewer": "Reviewer_9583", "summary": "The paper proposes FedImpro to construct similar features across clients for better federated training. The paper theoretically shows that the \"generalization contribution\" of local training is bounded by the conditional Wasserstein distances.", "questions": "I don't really follow how \"generalization contribution is bounded by conditional Wasserstein distance\" leads to the proposed FedImpro algorithm. I think the connection and logic flow can be made more clear.", "limitations": "", "rating": 8, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The idea of generalization contribution in FL sounds novel. \n2. Experimental performances of FedImpro look superior.", "weaknesses": "The idea of having a lower-level and a higher-level neural network in FL is not new, i.e. the feature extraction network idea. I don't see many comparisons to these previous work in the experimental section.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "nLMiwLd4ie", "reviewer": "Reviewer_LhTp", "summary": "The paper proposes FedImpro to mitigate the client drift problem in FL. It provides a new way to correct gradients and updates. The method mainly has 2 parts. The first is the study of the generalization contribution of each client, then proposes to leverage the same or similar conditional distributions for local training. The second is to propose decoupling a deep neural network into a low-level model (features extraction) and a high-level model (classifier network).", "questions": "- do you use h_hat to update the low-level model?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Even though decoupling a neural network into a feature extractor network and a classifier network is not novel, the paper proposed to combine the decoupling method with the features distribution estimation method with privacy protection is quite novel for mitigating client drift. Also, the sampling part with the synthetic features to increase the distribution similarity is quite interesting too.", "weaknesses": "- All the distributions in the theoretical analysis considered conditional distribution conditioned on the label y, and then the paper said that ' it is straightforward to make all client models trained on similar distributions to obtain higher generalization performance'. But, for a dataset such as CIFAR10, when we partition it among clients, the non-IID is introduced by the label imbalanced across clients, which means that the conditional distribution on the label is the same. However, we would still experience client drift in this case. I think more explanation/analysis is required on this aspect. Also, since the analysis is on the conditional distribution, the experiments should reflect that, by using a dataset with the natural partition, such as FEMNIST. \n\n- since the feature distribution depends on the client selected in each round. Unlike the standard FedAvg training, the number of clients per round only impacts the convergence rate. Here for the FedImprov method, it will also impact the distribution estimation. Therefore, an ablation study analyzing the number of clients on the feature distribution estimation is lacking.\n\nMinor:\n- a small typo in Sec 5.1, M=10 should be cross-silo and M=100 should be cross-device.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)