Federated Learning under Label Shifts with Guarantees

We thank the reviewer for their thoughtful comments which we address one by one below.

Empirical comparison

For your reference, we upload two plots anonymously here. 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”.

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.

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):

$

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}}.

$

3. Lower Bound on Convergence Rate for for Convex and Second-order Differentiable Optimization (Theorem 4):

$

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}}.

$

4. Convergence and Communication Bounds for Nonconvex Optimization for PL with Compression (Theorem 5):

$

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).

$

5. Convergence and Communication Guarantees for Nonconvex Optimization with Adaptive Step-sizes for Non-convex Optimization with Adaptive Step-sizes Assumptions (Theorem 6):

$

T\lesssim \frac{r_{max}}{K\epsilon^3}\quad and \quad R\lesssim\frac{r_{max}}{\epsilon^2}.

$

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):

$

\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}).

$

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

[2]. Saurabh Garg, Yifan Wu, Sivaraman Balakrishnan, and Zachary C. Lipton. A unified view of label shift estimation. In Advances in neural information processing systems (NeurIPS), 2020.

Dear Reviewer VBZD,

We would like to express our gratitude for your review of our ICLR submission. We have taken your insightful comments into account and have made the necessary adjustments. Should you have any questions or require further information, please feel free to contact us. We look forward to any additional feedback you may have.

This is a friendly reminder that we submitted our responses to your comments and are looking forward to your feedback for further refinement of our paper.

We thank the reviewer for their thoughtful comments which we address one by one below.

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.

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.

For greater clarity and understanding, a detailed derivation of Equation (Reg-Est) within the main body of the paper would be advantageous

We acknowledge that there were certain oversights in the formulation of the equation initially, leading to some variations in the expression. However, despite this, the experimental setup was thoughtfully designed and executed. In the final version, we will add more details to the body since we have an additional page to elaborate further. Most importantly, our study demonstrated impressive performance. The strength of our work lies in the application of ratio estimation methods in practical scenarios, specifically in the context of federated learning. We successfully implemented the model in training sets without any additional information updates, and the performance on test sets nearly matched the theoretical upper bounds, using a fixed ratio for training. Besides, our approach shines, effectively addressing and overcoming challenges related to privacy, convergence rate, and communication costs.

[1]. Zachary C. Lipton. Rlsbench: Domain adaptation under relaxed label shift. In International Conference on Machine Learning (ICML), 2023.

[2]. Antonio Torralba, Rob Fergus, and William T. Freeman. 80 million tiny images: A large data set for nonparametric object and scene recognition. IEEE Transactions on Pattern Analysis and Machine Intelligence, 30(11):1958–1970, 2008.

[3]. Benjamin Recht, Rebecca Roelofs, Ludwig Schmidt, and Vaishaal Shankar. Do CIFAR-10 classifiers generalize to CIFAR-10? arXiv preprint arXiv:1806.00451, 2018.

Dear Reviewer PGzm,

We sincerely appreciate your review of our ICLR submission. Your valuable feedback has been carefully considered and addressed. If you have any further inquiries or require additional details, please do not hesitate to reach out. Your continued input is highly valued.

Just a quick reminder that we have responded to your comments and look forward to your feedback to enhance our work.

We thank the reviewer for their thoughtful comments which we address one by one below.

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):

$

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}}.

$

3. Lower Bound on Convergence Rate for for Convex and Second-order Differentiable Optimization (Theorem 4):

$

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}}.

$

4. Convergence and Communication Bounds for Nonconvex Optimization for PL with Compression (Theorem 5):

$

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).

$

5. Convergence and Communication Guarantees for Nonconvex Optimization with Adaptive Step-sizes for Non-convex Optimization with Adaptive Step-sizes Assumptions (Theorem 6):

$

T\lesssim \frac{r_{max}}{K\epsilon^3}\quad and \quad R\lesssim\frac{r_{max}}{\epsilon^2}.

$

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):

$

\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}).

$

Dear Reviewer HSCY,

Thank you for reviewing our ICLR submission. We have diligently addressed your valuable concerns. If you have any questions or require additional information, please feel free to reach out. We sincerely appreciate your continued feedback.

I am writing to kindly remind you that we submitted our responses, but we have yet to receive your feedback. We highly value your opinions and are eager to further improve our paper based on your insights.