Paper: Federated Learning under Label Shifts with Guarantees (ICLR-2024)

Reviewer comment (Experiments):
Empirical comparison: a closely related work is [1] as cited by the authors, which also considers adding a regularization term to perform for the label density ratio estimation. I think the experiments will be more convincing by also taking [1] as a compared method.  

[1]  Azizzadenesheli et al.  Regularized learning for domain adaptation under label shifts. In ICML 2019.

Author response:
**Empirical comparison**

 For your reference, we upload two plots anonymously [here](https://imgur.com/a/bCFhOTi). Regularized Learning under Label Shift (RLLS) [1] is a variation of BBSE, based on the confusion matrix approach. This method has been already compared with MLLS in [2]. I also implemented this method and obtained consistent results as Figure 1 (b)(e) in [2]. To sum up, RLLS, which is based on the confusion matrix, improved based on BBSE, underperforms divergence-based and EM-based methods. The potential reason has been studied by [2]: “This concludes MLLS’s superior performance is not because of differences in loss function used for distribution matching but due to differences in the granularity of the predictions, caused by crude confusion matrix aggregation”, and “MLLS’s superior performance (when applied with more granular calibration techniques) is not due to its objective but rather to the information lost by BBSE via confusion matrix calibration”.

Reviewer comment (Theory):
The connection between the optimization problem in Eq.(Reg-Est) and that shown by Eq.(3.1) and Eq.(3.2) is not immediately clear to me. In Eq.(Reg-Est), the regularization term is incorporated into the training of the density ratio, whereas in Eq.(3.1), it appears to be a part of the loss function to train the classifier $\theta^*$. It would be beneficial if the author could clarify this by providing a more formal statement that establishes the equivalence between these two optimization problems.

Author response:
**Formal statement that establishes the equivalence between these two optimization problems**
 As you correctly pointed out, the IW-ERM framework under federated learning can indeed be seen as an extension of the single-client IW-ERM to a multi-client scenario. This generalization is succinctly encapsulated in the optimization problem outlined in section C.3 of the Appendix.

From an optimization perspective, the IW-ERM's adaptation to multiple clients translates to the aggregation of the numerator of the weight ratio for each individual client. This modification not only preserves the essence of the IW-ERM approach but also effectively scales it for federated learning environments.

Regarding the convergence rate, the established bounds, and communication costs, our study presents several theorems that address these aspects in depth. These theorems provide a solid theoretical foundation, demonstrating how the IW-ERM under federated learning efficiently navigates these challenges. Specifically, they illustrate that despite the inherent complexities of federated learning, our approach remains both robust and scalable, effectively managing the balance between communication efficiency and convergence rate.  Additionally, the updated PDF includes various theorems concerning the lower and upper bounds of convergence rate and communication guarantees. Here is a summary: To sum up, we have:

1. Negligible impact on Convergence Rate and Communication Guarantees (Theorem 2):  $O(r_{max} h(T))$.
2. Upper Bound on Convergence Rate for Convex and Smooth Optimization (Theorem 3):
\begin{align}
E[l(h_{w_T})-l(h_{w^\star})]
\lesssim \frac{r_{max}b D^2}{\tau R}+\frac{(r_{max}b D^4)^{1/3}}{(\sqrt{\tau}R)^{2/3}}+\frac{D}{\sqrt{K\tau R}}.
\end{align}
3. Lower Bound on Convergence Rate for for Convex and Second-order Differentiable Optimization (Theorem 4):
\begin{align}
E[l(h_{w_T})-l(h_{w^\star})]
\gtrsim \frac{r_{\max}b D^2}{\tau R}+\frac{(r_{max}b D^4)^{1/3}}{(\sqrt{\tau}R)^{2/3}}+\frac{D}{\sqrt{K\tau R}}.
\end{align}
4. Convergence and Communication Bounds for Nonconvex Optimization for  PL with Compression (Theorem 5):
\begin{align}
R\lesssim \Big(\frac{q}{K}+1\Big)\kappa\log\Big(\frac{1}{\epsilon}\Big)\quad and \quad \tau\lesssim\Big(\frac{q+1}{K(q/K+1)\epsilon}\Big).
\end{align}
5. Convergence and Communication Guarantees for Nonconvex Optimization with Adaptive Step-sizes for Non-convex Optimization with Adaptive Step-sizes Assumptions (Theorem 6):
\begin{align}
T\lesssim  \frac{r_{max}}{K\epsilon^3}\quad and \quad R\lesssim\frac{r_{max}}{\epsilon^2}.
\end{align}
6. Oracle Complexity of Proximal Operator for  Composite Optimization with Proximal Operator (Theorem 7):  $O\big(r_{max}\sqrt{\kappa}\log(1/\epsilon)\big)$

7. High-probability Bound for Nonconvex Optimization for Sub-Gaussian Assumption (Theorem 8):

\begin{align}
\frac{1}{T}\sum_{t=1}^T\|\nabla_{w}l(h_{w_t})\|_2^2= O(\sigma \sqrt{\frac{r_m \beta}{T}} + \frac{\sigma^2\log(1/\delta)}{T}).
\end{align}

Reviewer comment (Theory):
I think this paper would be more appealing if the authors could show the advantage of the proposed method over MLLS from a theoretical view. (Please refer to the third point of weaknesses for more details.)

Author response:
**The advantage of the proposed method over MLLS from a theoretical view** 
The generalization bounds in Theorem 1 are there to  highlight the dependence on the complexity of the function class in which the optimization is conducted. The presence of an additional regularizer in our method induces a reduction in the size of the model. We have not been able to characterize exactly this reduction yet, but we believe this reduced complexity is behind our improved results. It is an interesting direction for future work.

Reviewer comment (Experiments):
While the authors have conducted experiments across a variety of settings, the scope of their datasets remains limited. To more convincingly demonstrate the robustness and practicality of the proposed method, it would be beneficial to extend these experiments to larger-scale datasets and real-world application scenarios.

Author response:
**Larger-scale datasets and real-world application scenarios**

We appreciate your recommendation to incorporate larger-scale datasets and explore real-world application scenarios. We extended our experiments from 100 clients to 200 clients on CIFAR10, obtained uniformly increased accuracy among methods, and our VRLS still dominated. 

------------

| CIFAR-10 | Our IW-ERM | FedAvg | FedBN |
|----------|------------|--------|-------|
| Avg. accuracy (100 clients) | **0.5354** | 0.3915 | 0.1537 |
| Avg. accuracy (200 clients) | **0.6216** | 0.5942 | 0.1753 |
----------

While we acknowledge the value of conducting experiments in real-world application scenarios, there are specific constraints and challenges that influenced our decision to not pursue it in this study.
For instance, as indicated in [1], while larger, more real-world datasets provide a valuable benchmark, $D(p_{te}(x|y) \Vert p_{tr}(x|y)) < \epsilon$ is hard to quantitatively assess in these contexts. An attempt to use test samples from CIFAR10.1-v6 [2, 3] was made. However, due to significant feature changes in this dataset, all methods including ours experienced a substantial increase in estimation errors. This led to comparisons that were not meaningful, as the estimations were poor and did not provide reliable insights.
Given these challenges and the nature of our study, we chose to focus on datasets where the constraints and feature stability allowed for more meaningful and accurate analysis. We believe that this approach, while more constrained, provides valuable insights within its scope and serves the objectives of our research effectively. In contrast, we can turn to directly focus on the conditional probability \( p(x|y) \) to address the relaxed label shift problem, which would be more efficient and straightforward to handle real-world problems.

Reviewer comment (Theory):
I am not sure of the relevance of the statistical results (Theorem 1) and optimization results (Theorem 2) in the context of federated learning (FL). In FL, there is only limited communication and this critical aspect is not categorized by these theorems. I am wondering how many communication rounds the algorithm requires to obtain a statistical and computational guarantee.

Author response:
**Theorems 1 and 2. Elaborate on the communication rounds theoretically and empirically?**

Theorem 1 primarily provides assurances for the accuracy of ratio estimation on an individual client, establishing that the upper bound of the estimation error is contingent upon the size of the test sample and the calibration quality of the pre-trained predictor. In the context of Federated Learning (FL), this theorem applies to each client by leveraging all available unlabeled test samples. By doing so, we locally estimate the ratio on each client, which allows us to closely approximate the true ratios, thereby ensuring that our estimation is as accurate as feasible. As a result, the implementation of Importance Weighted Empirical Risk Minimization (IWERM) will be as unbiased as possible with respect to the test distribution.

Theorem 2 gives the convergence rates in the context of our ratios with data-dependent constant terms which increase linearly with negligible communication overhead over the ERM-solver baseline without importance weighting. In Appendix F, we establish tight convergence rates and communication guarantee for Eq.(IW-ERM) with a broad range of importance optimization settings.

What’s more, the estimation of ratios and the implementation of IW-ERM introduce negligible additional communication costs during the global training phase. The estimation process is conducted prior to global training, which takes place locally on each client with an off-the-shelf predictor, which can potentially be MLP or other simple networks already existing. Subsequently, the estimated ratios, rather than the raw data, are shared across clients, entailing N^2 exchanges of K-dimensional vectors within one round of communication. Once global training commences, these ratios remain static with iterations T—they are neither updated nor exchanged again.
To sum up, we have:

1. Negligible impact on Convergence Rate and Communication Guarantees (Theorem 2):  $O(r_{max} h(T))$.
2. Upper Bound on Convergence Rate for Convex and Smooth Optimization (Theorem 3):
\begin{align}
E[l(h_{w_T})-l(h_{w^\star})]
\lesssim \frac{r_{max}b D^2}{\tau R}+\frac{(r_{max}b D^4)^{1/3}}{(\sqrt{\tau}R)^{2/3}}+\frac{D}{\sqrt{K\tau R}}.
\end{align}
3. Lower Bound on Convergence Rate for for Convex and Second-order Differentiable Optimization (Theorem 4):
\begin{align}
E[l(h_{w_T})-l(h_{w^\star})]
\gtrsim \frac{r_{\max}b D^2}{\tau R}+\frac{(r_{max}b D^4)^{1/3}}{(\sqrt{\tau}R)^{2/3}}+\frac{D}{\sqrt{K\tau R}}.
\end{align}
4. Convergence and Communication Bounds for Nonconvex Optimization for  PL with Compression (Theorem 5):
\begin{align}
R\lesssim \Big(\frac{q}{K}+1\Big)\kappa\log\Big(\frac{1}{\epsilon}\Big)\quad and \quad \tau\lesssim\Big(\frac{q+1}{K(q/K+1)\epsilon}\Big).
\end{align}
5. Convergence and Communication Guarantees for Nonconvex Optimization with Adaptive Step-sizes for Non-convex Optimization with Adaptive Step-sizes Assumptions (Theorem 6):
\begin{align}
T\lesssim  \frac{r_{max}}{K\epsilon^3}\quad and \quad R\lesssim\frac{r_{max}}{\epsilon^2}.
\end{align}
6. Oracle Complexity of Proximal Operator for  Composite Optimization with Proximal Operator (Theorem 7):  $O\big(r_{max}\sqrt{\kappa}\log(1/\epsilon)\big)$

7. High-probability Bound for Nonconvex Optimization for Sub-Gaussian Assumption (Theorem 8):

\begin{align}
\frac{1}{T}\sum_{t=1}^T\|\nabla_{w}l(h_{w_t})\|_2^2= O(\sigma \sqrt{\frac{r_m \beta}{T}} + \frac{\sigma^2\log(1/\delta)}{T}).
\end{align}

Reviewer comment (Experiments):
Experimental results are weak. Tacking data heterogeneity is a well-known problem in federated learning (e.g., FedProx [r1], SCAFFOLD [r2], minibatch SGD [r3], etc.). However, it seems that the authors only consider FedAvg and FedBN as baselines. I suggest the authors also compare against these papers, which were also proposed to learn a global model with client data heterogeneity.

Author response:
**Compare against more baselines for tackling data heterogeneity in FL**

Thank you for the comprehensive information on these baseline approaches. All serve as benchmarks in Federated Learning (FL) for bridging the heterogeneity gap. Yet, they differ fundamentally. For instance, FedProx is designed for scenarios where clients possess imbalanced sample sizes, as depicted in Table 1 of [1]. In our case, under the label shift setting, while each client has an equal total number of samples, the class imbalance within each client is more pronounced, akin to what is detailed in Table 7 of our paper. Consequently, straightforward application of these benchmarks might yield less-than-optimal outcomes, despite meticulous tuning.

-------------
|Method (5 clients)| FedProx (0.0001) | FedProx (0.1) | FedProx (0.01) | FedProx (0.001) |
|---------------| ---------------- | ------------- | -------------- | --------------- |
|Overall accuracy (run1)| 0.5514           | 0.0219        | 0.1868         | 0.4898          |
|Overall accuracy (run2)| 0.5685           | 0.0217        | 0.1924         | 0.4935          |
|Overall accuracy (run3)| 0.5620           | 0.0221        | 0.1976         | 0.4727          |
|**Avg. accuracy**| **0.5606 ± 0.0070** |  **0.0219 ± 0.0002** | **0.1923 ± 0.0044** | **0.4853 ± 0.0091** |
----------------

| Method  (5 clients)    | Avg. Accuracy             |
|------------|---------------------------|
| Our IW-ERM | **0.7520 ± 0.0209**       |
| FedAvg     | 0.5472 ± 0.0297           |
| FedBN      | 0.5359 ± 0.0306           |
| FedProx    | 0.5606 ± 0.0070           |
| SCAFFOLD   | 0.5774 ± 0.0036           |
| Upper Bound| 0.8273 ± 0.0041           |
--------

The new tables illustrate that the optimal hyperparameter $\mu$ falls outside the range suggested in [1]. After experimenting with various values, we discovered that setting $\mu$ to 1e-3 offers a slight improvement over FedAvg. Meanwhile, SCAFFOLD achieves a substantially easier gain on MNIST, improving by 3% over FedAvg. Nonetheless, both methods underperform on CIFAR10, with FedProx peaking at 41% accuracy post-tuning — below the FedAvg baseline of 45%. SCAFFOLD's results are further behind.

In summary, these methods are designed to regularize global training and come with the ease of implementation and theoretical assurances. However, they demand either an exhaustive hyperparameter search or struggle with application to new datasets (as SCAFFOLD was only tested on simulated data and EMNIST as per [2]). In contrast, statistical approaches like ours inherently differ. We calculate the importance ratio prior to global training, eliminating the need for any dataset-specific adjustments or adaptations during training.

---------
| CIFAR-10 | Our IW-ERM | FedAvg | FedBN |
|----------|------------|--------|-------|
| Avg. accuracy (100 clients) | **0.5354** | 0.3915 | 0.1537 |
| Avg. accuracy (200 clients) | **0.6216** | 0.5942 | 0.1753 |
----------

Lastly, we increased the number of clients and conducted large-scale experiments. Limited by hardware, we executed these on a single GPU for 200 clients. Here, all methods exhibited improved performance, with our approach maintaining dominance. Intriguingly, we observed enhanced results across the board compared to 100 clients. Our configuration involves duplicating the data on each of the 10 clients 10 or 20 times. We hypothesize this performance boost is due to a broader ensemble of optimizer states from shuffling in the Adam optimizer. It’s important to note that the upper bound in Table 2 is no longer applicable to Table 3, given the introduction of randomness when randomly selecting 5 out of 100/200 clients for training in each round.

[1]. Li, Tian, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith. "Federated optimization in heterogeneous networks." MLSys 2020.

[2]. Karimireddy, Sai Praneeth, Satyen Kale, Mehryar Mohri, Sashank Reddi, Sebastian Stich, and Ananda Theertha Suresh. "SCAFFOLD: Stochastic controlled averaging for federated learning." In ICML2020.

Reviewer comment (Experiments):
The authors did not report the communication round result in the experiments section. It is unclear whether improves over FedAvg or FedBN when there is only limited communication. There is still a huge gap between the IW-ERM and the upper-bound performance (Table 3).

Author response:
**Theorems 1 and 2. Elaborate on the communication rounds theoretically and empirically?**

Theorem 1 primarily provides assurances for the accuracy of ratio estimation on an individual client, establishing that the upper bound of the estimation error is contingent upon the size of the test sample and the calibration quality of the pre-trained predictor. In the context of Federated Learning (FL), this theorem applies to each client by leveraging all available unlabeled test samples. By doing so, we locally estimate the ratio on each client, which allows us to closely approximate the true ratios, thereby ensuring that our estimation is as accurate as feasible. As a result, the implementation of Importance Weighted Empirical Risk Minimization (IWERM) will be as unbiased as possible with respect to the test distribution.

Theorem 2 gives the convergence rates in the context of our ratios with data-dependent constant terms which increase linearly with negligible communication overhead over the ERM-solver baseline without importance weighting. In Appendix F, we establish tight convergence rates and communication guarantee for Eq.(IW-ERM) with a broad range of importance optimization settings.

What’s more, the estimation of ratios and the implementation of IW-ERM introduce negligible additional communication costs during the global training phase. The estimation process is conducted prior to global training, which takes place locally on each client with an off-the-shelf predictor, which can potentially be MLP or other simple networks already existing. Subsequently, the estimated ratios, rather than the raw data, are shared across clients, entailing N^2 exchanges of K-dimensional vectors within one round of communication. Once global training commences, these ratios remain static with iterations T—they are neither updated nor exchanged again.
To sum up, we have:

1. Negligible impact on Convergence Rate and Communication Guarantees (Theorem 2):  $O(r_{max} h(T))$.
2. Upper Bound on Convergence Rate for Convex and Smooth Optimization (Theorem 3):
\begin{align}
E[l(h_{w_T})-l(h_{w^\star})]
\lesssim \frac{r_{max}b D^2}{\tau R}+\frac{(r_{max}b D^4)^{1/3}}{(\sqrt{\tau}R)^{2/3}}+\frac{D}{\sqrt{K\tau R}}.
\end{align}
3. Lower Bound on Convergence Rate for for Convex and Second-order Differentiable Optimization (Theorem 4):
\begin{align}
E[l(h_{w_T})-l(h_{w^\star})]
\gtrsim \frac{r_{\max}b D^2}{\tau R}+\frac{(r_{max}b D^4)^{1/3}}{(\sqrt{\tau}R)^{2/3}}+\frac{D}{\sqrt{K\tau R}}.
\end{align}
4. Convergence and Communication Bounds for Nonconvex Optimization for  PL with Compression (Theorem 5):
\begin{align}
R\lesssim \Big(\frac{q}{K}+1\Big)\kappa\log\Big(\frac{1}{\epsilon}\Big)\quad and \quad \tau\lesssim\Big(\frac{q+1}{K(q/K+1)\epsilon}\Big).
\end{align}
5. Convergence and Communication Guarantees for Nonconvex Optimization with Adaptive Step-sizes for Non-convex Optimization with Adaptive Step-sizes Assumptions (Theorem 6):
\begin{align}
T\lesssim  \frac{r_{max}}{K\epsilon^3}\quad and \quad R\lesssim\frac{r_{max}}{\epsilon^2}.
\end{align}
6. Oracle Complexity of Proximal Operator for  Composite Optimization with Proximal Operator (Theorem 7):  $O\big(r_{max}\sqrt{\kappa}\log(1/\epsilon)\big)$

7. High-probability Bound for Nonconvex Optimization for Sub-Gaussian Assumption (Theorem 8):

\begin{align}
\frac{1}{T}\sum_{t=1}^T\|\nabla_{w}l(h_{w_t})\|_2^2= O(\sigma \sqrt{\frac{r_m \beta}{T}} + \frac{\sigma^2\log(1/\delta)}{T}).
\end{align}

Reviewer comment (Evaluation):
Can you elaborate on the communication round results theoretically and empirically?

Author response:
**Theorems 1 and 2. Elaborate on the communication rounds theoretically and empirically?**

Theorem 1 primarily provides assurances for the accuracy of ratio estimation on an individual client, establishing that the upper bound of the estimation error is contingent upon the size of the test sample and the calibration quality of the pre-trained predictor. In the context of Federated Learning (FL), this theorem applies to each client by leveraging all available unlabeled test samples. By doing so, we locally estimate the ratio on each client, which allows us to closely approximate the true ratios, thereby ensuring that our estimation is as accurate as feasible. As a result, the implementation of Importance Weighted Empirical Risk Minimization (IWERM) will be as unbiased as possible with respect to the test distribution.

Theorem 2 gives the convergence rates in the context of our ratios with data-dependent constant terms which increase linearly with negligible communication overhead over the ERM-solver baseline without importance weighting. In Appendix F, we establish tight convergence rates and communication guarantee for Eq.(IW-ERM) with a broad range of importance optimization settings.

What’s more, the estimation of ratios and the implementation of IW-ERM introduce negligible additional communication costs during the global training phase. The estimation process is conducted prior to global training, which takes place locally on each client with an off-the-shelf predictor, which can potentially be MLP or other simple networks already existing. Subsequently, the estimated ratios, rather than the raw data, are shared across clients, entailing N^2 exchanges of K-dimensional vectors within one round of communication. Once global training commences, these ratios remain static with iterations T—they are neither updated nor exchanged again.
To sum up, we have:

1. Negligible impact on Convergence Rate and Communication Guarantees (Theorem 2):  $O(r_{max} h(T))$.
2. Upper Bound on Convergence Rate for Convex and Smooth Optimization (Theorem 3):
\begin{align}
E[l(h_{w_T})-l(h_{w^\star})]
\lesssim \frac{r_{max}b D^2}{\tau R}+\frac{(r_{max}b D^4)^{1/3}}{(\sqrt{\tau}R)^{2/3}}+\frac{D}{\sqrt{K\tau R}}.
\end{align}
3. Lower Bound on Convergence Rate for for Convex and Second-order Differentiable Optimization (Theorem 4):
\begin{align}
E[l(h_{w_T})-l(h_{w^\star})]
\gtrsim \frac{r_{\max}b D^2}{\tau R}+\frac{(r_{max}b D^4)^{1/3}}{(\sqrt{\tau}R)^{2/3}}+\frac{D}{\sqrt{K\tau R}}.
\end{align}
4. Convergence and Communication Bounds for Nonconvex Optimization for  PL with Compression (Theorem 5):
\begin{align}
R\lesssim \Big(\frac{q}{K}+1\Big)\kappa\log\Big(\frac{1}{\epsilon}\Big)\quad and \quad \tau\lesssim\Big(\frac{q+1}{K(q/K+1)\epsilon}\Big).
\end{align}
5. Convergence and Communication Guarantees for Nonconvex Optimization with Adaptive Step-sizes for Non-convex Optimization with Adaptive Step-sizes Assumptions (Theorem 6):
\begin{align}
T\lesssim  \frac{r_{max}}{K\epsilon^3}\quad and \quad R\lesssim\frac{r_{max}}{\epsilon^2}.
\end{align}
6. Oracle Complexity of Proximal Operator for  Composite Optimization with Proximal Operator (Theorem 7):  $O\big(r_{max}\sqrt{\kappa}\log(1/\epsilon)\big)$

7. High-probability Bound for Nonconvex Optimization for Sub-Gaussian Assumption (Theorem 8):

\begin{align}
\frac{1}{T}\sum_{t=1}^T\|\nabla_{w}l(h_{w_t})\|_2^2= O(\sigma \sqrt{\frac{r_m \beta}{T}} + \frac{\sigma^2\log(1/\delta)}{T}).
\end{align}

Reviewer comment (Experiments):
Can you compare against more baselines for tackling data heterogeneity in federated learning?

Author response:
**Compare against more baselines for tackling data heterogeneity in FL**

Thank you for the comprehensive information on these baseline approaches. All serve as benchmarks in Federated Learning (FL) for bridging the heterogeneity gap. Yet, they differ fundamentally. For instance, FedProx is designed for scenarios where clients possess imbalanced sample sizes, as depicted in Table 1 of [1]. In our case, under the label shift setting, while each client has an equal total number of samples, the class imbalance within each client is more pronounced, akin to what is detailed in Table 7 of our paper. Consequently, straightforward application of these benchmarks might yield less-than-optimal outcomes, despite meticulous tuning.

-------------
|Method (5 clients)| FedProx (0.0001) | FedProx (0.1) | FedProx (0.01) | FedProx (0.001) |
|---------------| ---------------- | ------------- | -------------- | --------------- |
|Overall accuracy (run1)| 0.5514           | 0.0219        | 0.1868         | 0.4898          |
|Overall accuracy (run2)| 0.5685           | 0.0217        | 0.1924         | 0.4935          |
|Overall accuracy (run3)| 0.5620           | 0.0221        | 0.1976         | 0.4727          |
|**Avg. accuracy**| **0.5606 ± 0.0070** |  **0.0219 ± 0.0002** | **0.1923 ± 0.0044** | **0.4853 ± 0.0091** |
----------------

| Method  (5 clients)    | Avg. Accuracy             |
|------------|---------------------------|
| Our IW-ERM | **0.7520 ± 0.0209**       |
| FedAvg     | 0.5472 ± 0.0297           |
| FedBN      | 0.5359 ± 0.0306           |
| FedProx    | 0.5606 ± 0.0070           |
| SCAFFOLD   | 0.5774 ± 0.0036           |
| Upper Bound| 0.8273 ± 0.0041           |
--------

The new tables illustrate that the optimal hyperparameter $\mu$ falls outside the range suggested in [1]. After experimenting with various values, we discovered that setting $\mu$ to 1e-3 offers a slight improvement over FedAvg. Meanwhile, SCAFFOLD achieves a substantially easier gain on MNIST, improving by 3% over FedAvg. Nonetheless, both methods underperform on CIFAR10, with FedProx peaking at 41% accuracy post-tuning — below the FedAvg baseline of 45%. SCAFFOLD's results are further behind.

In summary, these methods are designed to regularize global training and come with the ease of implementation and theoretical assurances. However, they demand either an exhaustive hyperparameter search or struggle with application to new datasets (as SCAFFOLD was only tested on simulated data and EMNIST as per [2]). In contrast, statistical approaches like ours inherently differ. We calculate the importance ratio prior to global training, eliminating the need for any dataset-specific adjustments or adaptations during training.

---------
| CIFAR-10 | Our IW-ERM | FedAvg | FedBN |
|----------|------------|--------|-------|
| Avg. accuracy (100 clients) | **0.5354** | 0.3915 | 0.1537 |
| Avg. accuracy (200 clients) | **0.6216** | 0.5942 | 0.1753 |
----------

Lastly, we increased the number of clients and conducted large-scale experiments. Limited by hardware, we executed these on a single GPU for 200 clients. Here, all methods exhibited improved performance, with our approach maintaining dominance. Intriguingly, we observed enhanced results across the board compared to 100 clients. Our configuration involves duplicating the data on each of the 10 clients 10 or 20 times. We hypothesize this performance boost is due to a broader ensemble of optimizer states from shuffling in the Adam optimizer. It’s important to note that the upper bound in Table 2 is no longer applicable to Table 3, given the introduction of randomness when randomly selecting 5 out of 100/200 clients for training in each round.

[1]. Li, Tian, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith. "Federated optimization in heterogeneous networks." MLSys 2020.

[2]. Karimireddy, Sai Praneeth, Satyen Kale, Mehryar Mohri, Sashank Reddi, Sebastian Stich, and Ananda Theertha Suresh. "SCAFFOLD: Stochastic controlled averaging for federated learning." In ICML2020.