Overcoming label shift in targeted federated learning

Authors: Edvin Listo Zec, Adam Breitholtz, Fredrik D. Johansson
Primary area: other topics in machine learning (i.e., none of the above)
Keywords: Federated learning, Label shift, distribution shift, non-iid

Abstract:
Federated learning enables multiple actors to collaboratively train models without sharing private data. This unlocks the potential for scaling machine learning to diverse applications. Existing algorithms for this task are well-justified when clients and the intended target domain share the same distribution of features and labels, but this assumption is often violated in real-world scenarios. One common violation is label shift, where the label distributions differ across clients or between clients and the target domain, which can significantly degrade model performance. To address this problem, we propose FedPALS, a novel model aggregation scheme that adapts to label shifts by leveraging knowledge of the target label distribution at the central server. Our approach ensures unbiased updates under stochastic gradient descent, ensuring robust generalization across clients with diverse, label-shifted data. Extensive experiments on image classification demonstrate that FedPALS consistently outperforms standard baselines by aligning model aggregation with the target domain. Our findings reveal that conventional federated learning methods suffer severely in cases of extreme client sparsity, highlighting the critical need for target-aware aggregation. FedPALS offers a principled and practical solution to mitigate label distribution mismatch, ensuring models trained in federated settings can generalize effectively to label-shifted target domains.

Decision: reject
Reviewer scores: 3, 5, 3, 6, 3

Reviews and discussion:
<thread>
creation date:
1729153833959

thread_type:
official_review

<thread_object>
creator: reviewer_p4wj

<content>
summary:
The paper introduces FedPALS, a method to deal with label shift in federated learning. Label shift happens when label distributions change across clients or between clients and target. The authors propose a novel aggregation scheme to optimize a convex combination of client models to ensure that the aggregated model is better suited for the label distribution of the target domain.

strengths:
S1: The problem of label skews is interesting and well positioned.

S2: The paper gives good theoretical proofs on the algorithm.

weaknesses:
W1: The compared baselines are too old. The authors only compare algorithms before 2020 (FedAvg, FedProx, SCAFFOLD). There are several recent works on label shifts in federated learning (e.g., [1, 2, 3]). Lack of these recent baselines cannot convince that the proposed algorithm is SOTA.

W2: The authors work with C=2,3 and beta=0.1 partitions. How about C=1, the most extreme label shift?

W3: Assumes target label distribution known is often not practical. No examples of how to get this info without privacy risk.

[1] Li, Xin-Chun, and De-Chuan Zhan. "Fedrs: Federated learning with restricted softmax for label distribution non-iid data." Proceedings of the 27th ACM SIGKDD conference on knowledge discovery & data mining. 2021.

[2] Zhang, Jie, et al. "Federated learning with label distribution skew via logits calibration." International Conference on Machine Learning. PMLR, 2022.

[3] Diao, Yiqun, Qinbin Li, and Bingsheng He. "Exploiting Label Skews in Federated Learning with Model Concatenation." Proceedings of the AAAI Conference on Artificial Intelligence. Vol. 38. No. 10. 2024.

questions:
See weaknesses. Given the very weak baselines, it is not convincing that the proposed solution is SOTA to reach the bar of acceptance.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730553338030

thread_type:
official_review

<thread_object>
creator: reviewer_otnv

<content>
summary:
This paper addresses the problem of label shift in federated learning, where client data distributions differ from a target domain with a known label distribution, but no target samples are available. The authors propose a novel aggregation scheme called FedPALS that optimizes a convex combination of client models to align with the target label distribution, ensuring that the aggregated model minimizes the target risk. The paper provides theoretical justification for the proposed method and demonstrates its effectiveness through extensive empirical evaluation, showing that it outperforms traditional approaches like FedAvg, FedProx, and SCAFFOLD in scenarios with significant distributional shifts.

strengths:
- The proposed FedPALS method is novel and addresses a specific yet significant problem in federated learning, i.e., label shift. The approach of optimizing a convex combination of client models to align with the target label distribution is well-justified.
- The paper is technically sound, with theoretical analysis and extensive empirical evaluation. The experiments cover a variety of scenarios and datasets, demonstrating the robustness and effectiveness of the proposed method.
- The writing is clear and the presentation is well-organized.

weaknesses:
- The paper lacks a reference to the work "Agnostic Federated Learning" by Mohri et al., presented at ICML 2019. This work also addresses federated learning with unknown target distributions and shares some theoretical similarities with the proposed method in this paper. The authors should include this reference in the related work section and provide a more detailed comparison in both the theoretical and experimental sections to highlight the differences and similarities between the two approaches.
- The paper could be improved by providing a more detailed discussion on the limitations of the proposed method. For example, how does FedPALS perform in scenarios where the label shift assumption does not hold? Are there any cases where FedPALS might fail or underperform?
- The paper discusses the choice of the hyperparameter $\lambda$ but could provide more guidance on how to select this parameter in practice. A more detailed analysis of the impact of $\lambda$ on the performance and robustness of the method would be valuable.

questions:
- Could the authors provide a more detailed comparison with existing methods, particularly those that address distributional shifts in federated learning? How does FedPALS differ from and improve upon these methods?
- Does the proposed method generalize to other types of distributional shifts, such as covariate shift? If so, could the authors provide some initial results or discussion on this topic?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730637511008

thread_type:
official_review

<thread_object>
creator: reviewer_5cxe

<content>
summary:
This paper addresses label shift in federated learning, a challenge where label distributions vary across clients and differ from the target domain, potentially leading to degraded model performance. The authors propose FedPALS, a novel model aggregation approach that aligns client updates with the target label distribution, enabling better generalization under label shift. The method incorporates a weighting strategy that balances the need to match target distribution closely while minimizing variance in updates. Experimental results on datasets with label sparsity demonstrate FedPALS’s effectiveness over baseline federated learning methods, showing improved model accuracy in label-shifted target domains.

strengths:
1. FedPALS introduces a well-justified aggregation technique that adjusts model updates for label shifts, making it highly applicable in non-i.i.d. data settings commonly found in real-world federated learning applications.

2. The paper includes extensive experiments on multiple datasets with varying degrees of label sparsity, demonstrating FedPALS’s superiority over standard methods like FedAvg, FedProx, and SCAFFOLD, thus validating its effectiveness in handling label shifts.

weaknesses:
FedPALS lacks novelty, largely building on existing federated aggregation methods without introducing a fundamentally new approach. It uses outdated baselines; incorporating recent models (2023-2024) would provide better comparisons. The experiments with a limited number of clients (e.g., 3 for PACS, 10 for CIFAR-10) restrict evaluation of FedPALS’s scalability in larger federated settings. Additionally, its optimization for client aggregation weights could become computationally expensive with more clients or labels. The method focuses on label shift alone, though real-world federated learning often involves covariate shifts. Lastly, FedPALS relies on a critical hyperparameter, λ, without clear guidance on selection, which may limit practical deployment.

questions:
1. The novelty of the proposed method (FedPALS) is not strongly highlighted, as it builds upon existing model aggregation schemes in federated learning without introducing a fundamentally new approach. Although it adapts the aggregation scheme to account for label shift, similar issues have been addressed with methods like FedAvg and FedProx, making it essential to clarify the unique contributions FedPALS offers beyond these.

2. The baseline used in this paper is too old. Please use some new baselines from 2023 or 2024 for comparison.

3. The experiments are conducted with a limited number of clients (e.g., 3 clients in PACS and 10 in CIFAR-10 and Fashion-MNIST), which is not representative of typical large-scale federated learning environments. This small client network size limits the evaluation of FedPALS's scalability and robustness, especially under diverse, high-client scenarios that federated learning is intended for.

4. FedPALS introduces an optimization process to determine client aggregation weights, which may become computationally expensive as the number of clients or labels increases. The paper does not address the potential scalability issues of this approach, raising concerns about its applicability in large federated learning settings with numerous heterogeneous clients or labels.

5. The method focuses on label shift, assuming that the input distributions remain consistent across clients and the target domain. However, in real-world federated learning, covariate shift often occurs alongside label shift. Expanding FedPALS to handle both types of shifts would enhance its comprehensiveness and applicability.

6. The method’s performance relies on a crucial hyperparameter, λ, which balances bias and variance in the model. The paper does not offer a clear, practical method for choosing λ, which may pose challenges for real-world deployment where fine-tuning may not be feasible. This lack of guidance could hinder the ease of implementation and generalization of FedPALS.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730676484478

thread_type:
official_review

<thread_object>
creator: reviewer_mzuq

<content>
summary:
This paper studied a new FL problem under label shifts where the server knows the different label distributions of clients and target domains are known to the central server but unknown to the clients. To solve this problem, this paper optimized the aggregation weights using the label distributions. Experimental results demonstrated the effectiveness of the proposed methods over several baselines.

strengths:
(S1) A novel FL problem with label shifts is formulated by assuming that the server knows the different label distributions of clients and target domains are known to the central server but unknown to the clients

(S2) A novel parameter aggregation strategy is proposed based on the label distributions. It also balances the effective sample size and the alignment with the target label distribution.

(S3) Experimental results on several data sets demonstrate the effectiveness of the proposed method.

weaknesses:
(1) The training procedures of the proposed FedPALS method based on equation (7) are confusing. The parameter $\lambda$ largely affects the model performance in the experiments. Though section 3.2 provides several options, it is unclear how $\lambda$ is selected in the experiments, e.g., Figure 3(a)(b).

(2) The assumption of this paper is strong. 
- It assumes that the target label distribution $T (Y )$ is known to the server, but no training examples are available for the target client. It would be better to provide more realistic FL scenarios to illustrate the importance of this special problem settings.
- It assumes that the server knows the different label distributions of clients. This can also result in the privacy concern.

(3) The prior information regarding $T (Y )$ is only available for the proposed FedPALS approach. It might be unfair for performance comparison between FedPALS and baselines.

questions:
(1) The proof of Proposition 2 is confusing. When $\mu = 2\sum_i \frac{\alpha_i^*}{n_i}$, how can it derive $\mu = \frac{2}{sum_i n_i}$ using the the primal feasibility condition?

(2) How is the parameter $\lambda$ selected in Figure 3(a)(b)?

(3) It seems that the approaches in Figure 3(b) do not converge yet after 80 communication rounds, as their target F1-scores keep increasing.

(4) The paper [c1] also studied the FL problem under unknown target client. It can assume no prior information regarding the target client. Thus, the developed AFL can be one of the strong baselines for FedPALS in the experiments.

[c1] Mohri, Mehryar, Gary Sivek, and Ananda Theertha Suresh. "Agnostic federated learning." In International conference on machine learning, pp. 4615-4625. PMLR, 2019.
</content>
</thread_object>

</thread>
<thread>
creation date:
1731097758794

thread_type:
official_review

<thread_object>
creator: reviewer_vuv9

<content>
summary:
This work addresses the realistic scenario of generating a global model for an arbitrary target distribution that may differ from the distribution of the aggregated client training datasets. To tackle this challenge, the authors assume that the central server has knowledge of the different label distributions of the clients and target domains, while this information remains unknown to the clients. They introduce FedPALS, a method that effectively balances the trade-off between bias (related to the target distribution relative to the aggregated training dataset) and variance stemming from the number of each client dataset.

strengths:
This paper is the first to tackle the realistic federated learning scenario of generating a global model for an arbitrary target distribution that may differ from the aggregated client training datasets.

weaknesses:
**Weaknesses**

The presentation of this paper is notably poor, as it contains numerous inaccuracies and unclear statements throughout. For instance:

1. (L87-88 and L136)  "A common assumption in federated learning is that all client distributions are identical." 
-> This is not right. Federated learning generally allows for differences in client distributions, with data heterogeneity being one of the main challenges.

2. (L 97-98)-> the phrase "Given a set of clients S_1, ..., S_M" -> these are the set of client's marginal distribution.

3. (L98)There is a lack of clarity regarding the loss function; the specific loss being used is not mentioned.

4. (L108) "While the server has access to all marginal label distributions," reflects an assumption made in the paper. Given the emphasis on data privacy in federated learning, this assumption requires a much deeper justification. The authors only reference a single paper (Ramakrishna & Dan, 2022) in (L113-115) to support this assumption, which is insufficient. More extensive discussion is needed to validate this critical assumption.

5. (L161) The phrase "As we see in Table 1"  is unclear, as Table 1 appears several pages later. A more helpful reference would be "As we see in Table 1 in Section 5."

6. Proposition 1 discusses the importance of the aggregated training dataset's distribution being equal to that of the target dataset. However, establishing this equivalence alone does not address the challenges of federated learning. The issue of data heterogeneity among clients must be resolved when the aggregated training dataset's distribution matches the target dataset's distribution, particularly when local iterations are performed only once. In practice, due to communication issues, multiple local iterations are typically needed, and the implications of unbiasedness discussed in Proposition 1 are not sufficiently emphasized until lines 285-289.

7. (L 161) There is no definition provided for vector alpha and S  which creates confusion.

8. (L 300) The validation set's specific dataset type is not defined. If it is identical to the target dataset, this would constitute a dangerous approach that could be seen as cheating.

9. The inclusion of the hyperparameter lambda in (L 302) is a significant weakness, as it necessitates searching for optimal lambda values in arbitrary settings, which may vary across different settings.

10. The term "this benchmark" in (L 359) lacks clarity. It should specify what is being referred to. Additionally, the experimental setup appears weak, as modern federated learning research typically uses at least CIFAR-100, with ImageNet or Tiny ImageNet as benchmarks. This study seems to rely on a dataset with fewer than 10 classes, which lacks challenge. Moreover, with only 10 clients and an alpha value of 0.1 for data heterogeneity, it is unclear how performance will change as data heterogeneity increases.

11. In (L 444-445), the phrase "varying the number of clients C across 10 clients" is confusing, as C does not accurately convey the number of clients.

questions:
I am curious about the performance comparison of the proposed method with baseline federated learning algorithms (such as SCAFFOLD) in scenarios where the label distribution of the aggregated training dataset matches that of the target domain.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732205090264

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General response

comment:
We thank the reviewers for their time and effort in evaluating our work and for their helpful comments which we have used to improve our paper. We have responded to each reviewer and are confident that our rebuttal responds to the concerns which were brought up. Below we provide a summarised list of main points of our rebuttal and changes we have made to the revised manuscript. In the manuscript, changes are marked with blue text for your convenience.

* Added Baseline: hybrid FedProx+FedPals  (Figures 2c, 3a,b)
* Added Baseline: Agnostic Federated Learning (AFL), suggested by Reviewers mzUQ, otNV, to tasks other than iWildCam, (Table 1, Figures 2c-3a)
* Added discussion regarding the label distributions being known by the server (Reviewer VuV9, mzUQ)
* Added discussion on covariate shift (Reviewer )5CXe, otNV)
* Added discussion on the privacy-accuracy tradeoff involved in clients sharing label distribution (Reviewer mzUQ, P4Wj)
* Extended guidelines for the choice of hyperparameter $\lambda$ (Reviewer VuV9, mzUQ, 5CXe)
* Clarified proof of Proposition 2 (Reviewer mzUQ)
* Incorporated all proposed changes and clarifications to the text by Reviewer VuV9
* MOVED the synthetic experiment to the appendix to accommodate these changes.
* Explained why the computational cost of computing the weighting $\alpha$ is insignificant compared to training the model
</content>
</thread_object>

</thread>


Metareview:
metareview:
**Summary:** The paper proposes a novel method to address label shifts in federated learning by optimizing model aggregation weights based on known target label distributions. The authors provide theoretical justifications for their approach and conduct empirical evaluations across several datasets, demonstrating improvements over standard baselines like FedAvg and FedProx. The proposed method aims to balance bias and variance while aligning the aggregated model with the target distribution, addressing challenges associated with label shifts in non-i.i.d. data settings.

**Decision:** Despite addressing an important problem and proposing a reasonable approach, the paper has significant shortcomings that limit its potential impact. Several reviewers raised concerns about the method's assumptions, including the requirement for knowledge of the target label distribution and client label marginals, which is often impractical in real-world scenarios. The experimental setup was also criticized for relying on outdated baselines and limited scalability evaluations. More importantly, the presentation of this paper could also be further improved. Because of this, the novelty of the novelty of this paper is challenging to recognize; problem setups, terminology definition, and some experimental details (hyper-parameter tuning) are unclear and confusing. While the authors addressed some concerns during the rebuttal, such as adding new baselines and clarifying theoretical contributions, the responses did not fully alleviate reviewers' critiques.

These concerns outweigh the contributions, leading to the decision to reject. During the reviewer-AC discussion period, no objections were raised to this decision.