Paper: FedImpro: Measuring and Improving Client Update in Federated Learning (ICLR-2024)

Reviewer comment (Theory):
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.

Author response:
**Q1** More explanation on client drift:
>  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. 

***Ans for Q1):*** Thanks for your insightful comments. We provide detailed feedbacks as below. 
- **Case 1: The global label distribution is balanced, but local label distribution is different on different clients.** Although the overall label distribution is balanced, one client may have completely different conditional distribution on the label from other clients. Our analysis focuses on this case, and closes the conditional Wasserstein distance between these two clients.
- **Case 2: The global label distribution is balanced, and local label distribution is also same with different clients.** In this case, the client drift may still happen due to the feature distribution is different [1,2]. In this case, we may utilize the Wasserstein distance not conditioned on label in unsupervised domain adaptation [3,4] to analyse this problem. We will consider this case as the future work.

Thanks again for your constructive comments. We have added these discussion in Appendix B to enhance the quality of our paper.


**Q2** Experiment issues:
>  since the analysis is on the conditional distribution, the experiments should reflect that, by using a dataset with the natural partition, such as FEMNIST

***Ans for Q2):*** 
Thanks for your constructive comments. We conduct experiments with FEMNIST using its natural partition. New results are summarized in Table 4-1, showing that FedImpro can improve the model performance than other baselines.

**Table 4-1**
|                              | FedAvg | FedProx | FedNova | **FedImpro** |
|------------------------------|--------|---------|---------|--------------|
| Test Accuracy (%)            | 80.83  | 79.70   | 68.96   | **82.77**    |
| Comm. Round to attain Target Acc. | 82     | NaN     | NaN     | **45**       |

Reviewer comment (Experiments):
Since the feature distribution depends on the client selected in each round, unlike the standard FedAvg training where the number of clients per round only impacts the convergence rate, 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.

Author response:
**Q3** Ablation study on the number of clients:
>  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.

***Ans for Q3):*** 
Thanks for this constructive comments. We conduct more ablation studies with training CIFAR-10 and varying the number of selected clients per round. The result are summarized in Table 4-2, in which $a$ is the dirichlet parameter, E the local epochs, M the number of total clients, $M_c$ the number of selected clients per round. Results show that increasing the selected clients helps increase the performance of both FedAvg and FedImpro. It means that the feature estimation works well with more selected clients. We have added these results in Appendix of our revision.

**Table 4-2**
| $a$     | E   | $M$    | $M_c$ |FedAvg| FedImpro       |
|------|-------|-----|--|----|------|
|  0.1   | 1   | 10   | 5  | 83.65±2.03     |**88.45±0.43** |
| 0.1   | 1   | 10   | 10  | 86.65±1.26     | **90.25±0.93** |
| 0.05  | 1   | 10   | 5  | 75.36±1.92     |**81.75±1.03** |
| 0.05  | 1   | 10   | 10  | 78.36±1.58     |**85.75±1.21** |
| 0.1   | 5   | 10   | 5  | 85.69±0.57     |**88.10±0.20** |
| 0.1   | 5   | 10   | 10  | 87.92±0.31     |**90.92±0.25** |
| 0.1   | 1   | 100   | 10 | 73.42±1.19    |**77.56±1.02** |
| 0.1   | 1   | 100   | 20 | 77.82±0.94    |**85.39±1.15** |

Reviewer comment (Reproducibility):
Do you use h_hat to update the low-level model?

Author response:
**Q5** Interpreting $\hat{h}$
>  do you use h_hat to update the low-level model?

***Ans for Q5):*** Thanks for your insightful question. The low-level model will not receive gradients calculated from the $\hat{h}$. Thus, the low-level model is not explicitly updated by the $\hat{h}$. However, from another aspect, training the high-level model benefits from the closer-distance features and the reduced gradient heterogeneity. In FedImpro, the backward-propagated gradients on original features $h$ have lower bias than FedAvg. Therefore, **low-level model can implicitly benefit from $\hat{h}$**. We have added these discussion in Appendix B.

Reviewer comment (Novelty):
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.

Author response:
Thanks for the insightful comments and constructed suggestions. We firstly provide more discussions on differences between FedImpro and these related works. Then we provide an experimental comparison between the most related works that also aim to enhance the global model performance with utilizing features.
- There are some related works that focus on feature extraction or spliting lower and higher level model. We have discussed them in related works. We summarize the differences between FedImpro and these works and in Table 3-1 as below. Works [1,2] share raw data with noise. FedDF [3] finetunes on the aggregated model via the knowledge distillation relying on shared publica data. Methods [4,5,6,7] enhance model performance by data generation or sharing data statistics. CCVR [8] is a post-hoc method, transmiting the logit statistics of data samples to calibrate the last layer of Federated models. Different from them. [9] aims to defend poisoning attacks. [10-15] directly share hidden features to enhance personal or global federated learning. Different from these methods, FedImpro focuses on **the training process**, constructing **similar feature distribution**, **reducing gradient dissimilarity** of the high-level layers, **without any extra data**, and **sharing estimator parameters instead of raw hidden features**.

Table 3-1
|                                              | Shared Thing                     | Low-level Model | Objective                  |
|----------------------------------------------|----------------------------------|-----------------|----------------------------|
| [1,2] | Raw Data With Noise               | Shared          | Others                     |
| [3] | Extra Data               | Shared          | Others                     |
| [4,5] | Params. of Data Generator        | Shared          | Global Model Performance   |
| [6,7] | STAT. of raw Data                | Shared          | Global Model Performance   |
|  [8] | STAT. of Logis, Label Distribution | Shared        | Global Model Performance   |
| [9] | Hidden Features                  | Shared          | Defend Poisoning Attack    |
| [10,11] | logits                           | Private         | Personalized FL            |
| [12,13]| Hidden Features                  | Private         | Personalized FL            |
| [14,15] | Hidden Features                  | Shared          | Accelerate Training        |
| **FedImpro**                                 | Params. of Estimated Feat. Distribution | Shared   | Global Model Performance   |

- Specifically, we conducted experiments comparing our method with feature sharing methods, such as CCVR [8] and FedDF [3]. As CCVR and our methods do not need extra public training data. FedDF utilizes the noise dataset to conduct knowledge distillation following their original paper [3]. The communication round of our method is reduced to 100 for fair comparison. The results are reported in Table 3-2 as shown below, where our method outperforms existing feature sharing methods. FedDF severely relies on the extra meaningful data to enhance the model performance.

**Table 3-2**
|               |  CCVR | FedDF with noise data  | Ours   | 
|:-:|:-:|:-:|:-:|
| Test Accuracy (%)| 62.68 | 38.6 | **68.29** | 

We have incorporated corresponding discussion and experimental comparisons in our revision to enhance the quality of our work.

Reviewer comment (Theory):
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.

Author response:
Thank you for the suggestion aimed at making our work clear. Besides the previous response, we will add the following explanations to our revision. We will provide more detailed explanations if the reviewer finds some points **hard to follow or improper**.

- Our work involves two theorems, i.e., Theorem 4.1 and Theorem 4.2, where Theorem 4.1 motivates the proposed method and Theorem 4.2 indicates another advantage of the proposed method. 
- Theorem 4.1 is built upon two definitions, i.e., Definition 3.1 and Definition 4.1. Here, Definition 3.1 measures the model's generalizability on a data distribution from the margin theory, and Definition 4.1 is the conditional Wasserstein distance for two distributions. 
- 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 is 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 [3], reducing the distance in the feature space can induce the distance in the data space. However, the property could be hard to maintain for other deep networks, e.g., ResNet.

We have updated the above explainations in appendix for easier understanding. Thanks again for the valuable suggestions, which make our work easier to follow.

Reviewer comment (Experiments):
I notice that the author ignore a very related and state-of-art baesline FedDyn [1], could the author conduct comparion experiments with FedDyn?

Author response:
We would like to thank the reviewer for taking the time to review. We appreciate that you find that our paper is well written and easy to follow, the idea is easy to follow and understand, the generalization of the model is improved. In light of your insightful comments, we offer comprehensive responses below. We also add them into our main text or appendix in the revision. We hope that these responses and the revision have effectively resolved your concerns and enhanced the overall excellence of our work.


**Q1** Baseline issues:
> I notice that the author ignore a very related and state-of-art baesline FedDyn [1], could the author conduct comparion experiments with FedDyn?

***Ans for Q1):*** Following your valuable suggestions, we compare our method with FedDyn [1] as below. 
- FedDyn proposes to dynamically update the risk objective to ensure the device optima is asymptotically consistent. Different from FedDyn, FedImpro enhances convergence and model performance by constructing more similar feature distribution and reducing gradient heterogeneity, which is orthogonal to FedDyn or other FL algorithms discussed in related works.
- We conduct experiment comparisons with FedDyn for all settings. The results of FedDyn together with FedAvg and FedImpro are given in Table 2-1 as below. The results show that **FedImpro achieves better final accuracy than FedAvg and FedDyn**. We have added these new results in Table 1 of the main text. FedDyn is observed with much poorer performance when data heterogeneity is severe, which is also observed in [4].


**Table 2-1**
| Dataset     | a     | E   | M    | FedAvg | FedDyn | FedImpro       |
|-------------|-------|-----|------|----------------|----------------|----------------|
| **CIFAR-10** |  0.1   | 1   | 10   | 83.65±2.03     |84.73±0.81 | **88.45±0.43** |
|             |   0.05  | 1   | 10   | 75.36±1.92     |77.21±1.25| **81.75±1.03** |
|             |  0.1   | 5   | 10   | 85.69±0.57     |86.59±0.39| **88.10±0.20** |
|             |    0.1   | 1   | 100  | 73.42±1.19     | 75.29±0.95| **77.56±1.02** |
|  **FMNIST** | 0.1   | 1   | 10   | 88.67±0.34     |89.01±0.25| **90.83±0.19** |
|              | 0.05  | 1   | 10   | 82.73±0.98     |83.20±1.19| **86.42±0.71** |
|              | 0.1   | 5   | 10   | 87.6±0.52      |88.50±0.52| **89.87±0.21** |
|              | 0.1   | 1   | 100  | 90.12±0.19     | 90.57±0.21 | **90.98±0.15** |
| **SVHN**     | 0.1   | 1   | 10   | 88.20±1.21 | 90.82±1.09 | **92.37±0.82**|
|           | 0.05   | 1   | 10   | 80.67±1.92  | 84.12±1.28 | **90.25±0.69**|
|        | 0.1   | 5   | 10   | 86.32±1.19 | 89.92±0.50 | **91.58±0.72**|
|        | 0.1   | 1   | 100   | 92.42±0.21 |92.82±0.13 | **93.42±0.15**|
| **CIFAR-100**     | 0.1   | 1   | 10   | 69.38±1.02 | 69.59±0.52 | **70.28±0.33**|
|           | 0.05   | 1   | 10   | 63.80±1.29  |64.90±0.69 | **66.60±0.91**|
|        | 0.1   | 5   | 10   | 68.39±1.02 |68.52±0.41 | **68.79±0.52**|
|        | 0.1   | 1   | 100   | 53.22±1.20 | 54.82±0.81 | **56.07±0.75**|

Reviewer comment (Experiments):
The timecomsuming for training the model increases for FedImpro. Could the author list the cpu-time cost comparion experiments to reach the target accuracy?

Author response:
**Q2** Efficiency issue:
> The timecomsuming for training the model increases for FedImpro. Could the author list the cpu-time cost comparion experiments to reach the target accuracy?


***Ans for Q2):*** Thanks for your insightful comments that are strongly related to the practical deployments. We provided the detailed feedbacks as below.

- **Computational Overhead.** FedImpro requires little more computing time of FedAvg in simulation. Almost all current experiments are usually simulated in FL community [5,6,7]. When setting *local epoch = 1*, *communication rounds = 1000*, the simulation time of the FedAvg is around 8h 52m 12s, while FedImpro consumes around 10h 32m 32s. According to the Table 2 in the main text, FedImpro achieves the 82% acc in similar cpu time with FedAvg. Interestingly, when *local epoch = 5*, the simulation time of FedAvg increases to 20h 10m 25s. The convergence speed of FedImpro with *local epoch = 1* is better than FedAvg with *local epoch = 5*. **Note that in this simulation, the communication does not happen. However, the communication overhead dominates the system costs in the real-world deployment**. 

- **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\sim10$ MB/s) is very low comparing to the cluster environemnt (around $10\sim 1000$ GB/s). Taking ResNet18 with 48 MB as example, each communication with 10 clients would consume $48\sim 480$s, 1000 rounds requires $48000\sim 480000$s ($15\sim 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 [2,3]. The extra communication size of the estimation parameters in FedImpro is only 4KB, which is almost zero compared with the model size 48MB**. As shown in Table 2 in original text, FedImpro achieve mush faster convergence than other algorithms.




> ***Reference***
> 
> [1] Acar, Durmus Alp Emre, et al. "Federated learning based on dynamic regularization." ICLR (2021).
> 
> [2] Communication-efficient learning of deep networks from decentralized data. In AISTATS 2017.
> 
> [3] Advances and Open Problems in Federated Learning. In Foundations and Trends® in Machine Learning 2021.
> 
> [4] Sun, Yan, et al. "Fedspeed: Larger local interval, less communication round, and higher generalization accuracy." ICLR (2023). 
> 
> [5] FedML: A Research Library and Benchmark for Federated Machine Learning. In Arxiv 2020.
> 
> [6] FedScale: Benchmarking Model and System Performance of Federated Learning at Scale. In ICML 2022.
> 
> [7] A Survey on Federated Learning Systems: Vision, Hype and Reality for Data Privacy and Protection. In TKDE 2023.

Reviewer comment (Evaluation):
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.

Author response:
***Ans for Q1):*** Thanks for your insightful comments for the real-world deployments. We provide following analysis of the extra communication overhead. Specifically, the communication time in real-world can be characterized by the alpha-beta model [1,2]: $T_{comm}=\alpha + \beta M$, where $\alpha$ is the latency per message (or say each communication), $\beta$ 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 is not increased**, and the communication time from $\alpha$ will not increase; 
- **(b)** The $M$ increased as $M_{model} + M_{estimator}$. Taking ResNet-20 as an example, the size of estimated parameters is equal to doubled (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**.

Reviewer comment (Novelty):
Some relevant related works are not discussed: ETF [1], SphereFed [2], FedSpeed [3].

Author response:
***Ans for Q2):*** Thanks for the provided references from the reviewer, which are definitely related to our work. We have added discussions into the related works in revision. The following detailed discussions are listed in revised appendix due to the limited space of the main text. 
- FedETF [3] proposed a synthetic and fixed ETF classified to resolve the classifier delemma, which is orthogonal to FedImpro. 
- SphereFed [4] proposed constraining learned representations of data points to be a 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 to weaker consistency of the local learning targets from another perspective. FedImpro alleviate this problem by estimating and sharing similar features. 
- FedSpeed [5] 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.

Reviewer comment (Experiments):
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.

Author response:
***Ans for Q3):*** Following your valuable suggestions, we compare our method with SphereFed[4], FedDyn[6], CCVR [7], and FedSpeed [5]. The results are given in Table 1-1 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 $\alpha$ 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 or even the more difficult settings for us**. 

**Table 1-1**
| Method | Dataset | Dirichlet $\alpha$ |  Client Number | Communication Round | Test ACC. |
|:-:|:-:|:-:|:-:|:-:|:-:|
| FedDyn[6] | CIFAR-10 | 0.3 | 100 | 1400  | 77.33 |
| FedSpeed[5] | CIFAR-10 |0.3 | 100 | 1400 | 82.68 |
| FedImpro | CIFAR-10 |0.3 | 100 | 1400 | **85.14** |
| FedImpro | CIFAR-10 |0.1 | 100 | 1400 | **82.76** |
|--------|--------|---------|---------|--------|------|
| CCVR[7] | CIFAR-10 | 0.1 | 10 | 100 |  62.68 |
| FedImpro | CIFAR-10 | 0.1 | 10 | 100 | **68.29** |
|--------|--------|---------|---------|--------|------|
| SphereFed [4] | CIFAR-100 | 0.1 | 10 | 100 |  69.19 |
| FedImpro | CIFAR-100 | 0.1 | 10 | 100 | **70.28** |

Those previously baselines in the main paper are commonly compared in the recent FL works. Thanks again, we have added these results and analysis to our revision.

Reviewer comment (Evaluation):
From Fig. 2, the actual impact of FedImpro on reducing the weight divergence appears very limited.

Author response:
***Ans for Q4):*** Thanks for your careful reading of our paper and insightful comments.
- During the previous trainig stage, the weight divergence redution is limited, because the **low-level model and the feature estimation are still unstable**. After about 500 communication rounds, FedImpro shows lower weight divergence than others.
-  As Fig. 2 (a), Fig. 4, 5, 6 and 7 (in Appendix) show,  although in the early training stage the weight divergence reduction is limited, **the training convergence can  benefit from FedImpro**. Future work can think about how to fasten convergence of low-level models.

Reviewer comment (Writing):
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.

Author response:
***Ans for Q6):*** Thanks for your careful reviewing and the kind comments. Yes, the cross-silo FL and cross-device FL should be switched. We have modified the typo and update the text in the revision.