Overcoming label shift in targeted federated learning

We thank P4Wj for taking the time to review our work. We hope that our answers below address any questions and concerns brought up.

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

A: We are considering the specific setting of label shift between the client aggregate and the target. We also assume the server has access to the label marginals of client and target. This setting is distinct from just label shift between clients and it is unclear whether newer proposed algorithms for the latter, which are not tailored to our setting, would outperform FedProx and SCAFFOLD. Could the reviewer please motivate why the proposed works, which do not handle label shift between client aggregate and target, would be preferable to the baselines we have?

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

A: Thank you for the suggestion. However, $C=1$ represents an extreme case where each client only contains a single label, making the problem trivial at the client level. In such scenarios, local classifiers are unlikely to learn meaningful representations due to the lack of label diversity, and the applicability of the learned model to federated aggregation becomes unclear. Nonetheless, we appreciate the insight.

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

A: If the target label distribution is unknown the approach we propose is not suitable. It may be the case that we have some other information about the target that could be used, but that is another setting that we do not consider in our work.

We thank otNV for their helpul and insightful review. We hope that our answers below address any questions and concerns brought up.

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

A: We thank the reviewer for pointing this out this omission. We have added AFL as a baseline to our experiments.

• 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?

A: We have added some further discussion of the limitations of the method in section 6.

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

A: We have added further discussion and guidance on the choice of $\lambda$ at the end of section 3.

• 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?

A: We have some discussion on why we do not use such approaches in the related work section, in particular on lines 348--356.

• 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?

A: Covariate shift is inherently different from label shift as it typically involves changes in the distributions of high-dimensional continuous variables, such as input features. However, in our setting, we assume that $P(Y∣X)$ and $P(X∣Y)$ are invariant across sources and targets, while $P(Y)$ is changing. This change in $P(Y)$ naturally induces a shift in $P(X)$, meaning that covariate shift is inherently present in all of our experiments.

This relationship is evident from the joint distribution $P(X,Y)=P(X∣Y)P(Y)=P(Y∣X)P(X)$. As $P(Y)$ changes, $P(X)$ is also affected, reflecting aspects of covariate shift within our framework. We will add a clarification in the paper to explain this connection and emphasize how our approach implicitly addresses certain cases where covariate shift arises due to changes in label distribution.

Dear reviewer otNV,

Thanks again for reviewing our work. We believe our revision and response should address your concerns and kindly ask whether you would consider updating your review. If there are any remaining questions we can answer before the deadline tomorrow, please let us know.

Best, Authors

We thank 5CXe for their review. We hope that our answers below address any questions and concerns brought up.

• FedPALS lacks novelty, largely building on existing federated aggregation methods without introducing a fundamentally new approach.

A: We respectfully disagree with the reviewer’s assessment that FedPALS lacks novelty. The core contribution of FedPALS lies in the optimization of aggregation weights using both source and target marginal distributions, a strategy that is specifically tailored to address label shifts in federated learning. This is a fundamentally novel approach, as traditional federated aggregation methods, such as FedAvg, FedProx and SCAFFOLD, do not account for mismatches between the client and target label distributions.

• It uses outdated baselines; incorporating recent models (2023-2024) would provide better comparisons.

A: We are considering the specific setting of label shift between the client aggregate and the target where we assume the server has access to the label marginals of client and target. This setting is novel and it is unclear whether newer proposed algorithms which are not tailored to this setting would outperform FedProx and SCAFFOLD. Would the reviewer please advise which more recent models would provide better comparison than the baselines chosen?

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

A: While some of our experiments, such as PACS, use a smaller number of clients, we also evaluate FedPALS in larger federated settings. Specifically, the iWildCam experiment involves 100 clients (see Figure 3b), and for CIFAR-10, we conduct experiments with varying numbers of clients, including up to 100. The CIFAR-10 results are detailed in the appendix. This demonstrates that FedPALS is scalable and effective across a range of client network sizes.

• Additionally, its optimization for client aggregation weights could become computationally expensive with more clients or labels.

A: We respectfully disagree with the reviewer’s concern regarding the computational expense of optimizing client aggregation weights. This optimization is performed only once and is formulated as a quadratic program over $M$ variables, where $M$ is the number of clients. In all our experiments, this computation takes less than a few seconds to complete, making it computationally negligible compared to the significantly more resource-intensive task of training $M$ deep neural networks, which is the primary focus of federated learning.

• The method focuses on label shift alone, though real-world federated learning often involves covariate shifts.

A: We acknowledge the reviewer’s point; however, the focus of this work is specifically on label shift, not covariate shift. While real-world federated learning may involve both types of shifts, addressing label shift is a distinct and challenging problem that warrants dedicated study.

• Lastly, FedPALS relies on a critical hyperparameter, λ, without clear guidance on selection, which may limit practical deployment.

A: We have added further discussion and guidance on the choice of $\lambda$ at the end of section 3.

Dear reviewer 5CXe,

Thanks again for reviewing our work. We believe our revision and response should address your concerns and kindly ask whether you would consider updating your review. If there are any remaining questions we can answer before the deadline tomorrow, please let us know.

Best, Authors

Thank you for suggesting specific baselines—this substantially helps the discussion!

We have looked into the two papers you suggested and neither tackle the targeted federated learning scenario—i.e., where the target distribution differs from the client aggregate. Moreover, neither [1] or [2] are directly applicable in our setting:

Both FedRC and FedCollab produce sets of models, one for each cluster or coalition. Neither clustering in [1] or [2] is sensitive to a target distribution, and it is not clear how to pick which cluster's model to apply in the target after training. Without data from the target it can't be clustered with the other clients, so there's no obvious choice for which $\theta_k$ in [1] or $\hat{h}_{\alpha_i}$ in [2] to use in the target. Finally, both [1] and [2] optimize the in-cluster client aggregate error which does not conform to our setting, in which we want to optimize the target distribution error.

Between them, the 5 reviewers have now suggested 6 different baselines, and only 1 baseline (AFL) was suggested by two reviewers. We added AFL to the paper since the revision and, in summary, it does not perform much better than other baselines and still substantially worse than our method. None of the suggested methods are designed for targeted FL.

In light of the above, we hope you agree that our setting is new and that there are no recent or old baselines that are specifically suited to it. Based on this, we firmly argue that our selection of baselines is representative and sufficient. Since this was your primary concern, we hope that you consider updating your review.

We will include the suggested baselines in the related work section and contrast their settings with ours. Let us know if we can clarify anything further.

Best, Authors

We thank mzUQ for their thorough and helpful review. We hope that our answers below address your questions and concerns.

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

A: The $\lambda$ parameter is selected such that they represent 25%, 50% and 75% of the ESS compared to the FedAvg solution. This is detailed in Appendix A.3 and A.4. We

• The assumption of this paper is strong. It assumes that the target label distribution 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.

A: We acknowledge that the assumption of the server knowing the target label distribution $T(Y)$ is a strong one. However, this is central to the problem setting we address, where the server can access high-level aggregated label statistics without requiring individual data points from the target domain. This setup is motivated by real-world applications, such as retail or healthcare, where such aggregate statistics are often accessible and privacy-preserving.

Regarding the assumption that the server knows the clients' label distributions, we recognize that this introduces potential privacy concerns. However, as part of the privacy-accuracy tradeoff inherent in federated learning, we believe that the performance gains from utilizing $T(Y)$ and client label information justify exploring this scenario. Notably, sharing aggregate label statistics is less invasive than sharing raw data, making this a reasonable compromise in many practical settings.

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

A: Our conclusions are focused on demonstrating that $T(Y)$ is a valuable signal for federated learning. Even under the assumption of its availability, this does not invalidate the broader significance of leveraging label distribution knowledge for improved performance.

Furthermore, it is unclear how other methods could effectively use $T(Y)$ without substantial modification. To address this, we have combined FedPALS with FedProx in our experiments as it is straightforward to do so, while combining with other baselines like SCAFFOLD is more complex. This ensures a more meaningful comparison of FedPALS’ contribution while recognizing its specialized assumptions.

• The proof of Proposition 2 is confusing.

A: We thank the reviewer for pointing this out. We have changed the proof to be simpler and more straightforward.

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

A: The $\lambda$ parameter is selected such that they represent 25%, 50% and 75% of the ESS compared to the FedAvg solution. This is detailed in appendix A.3 and A.4 for the respective experiments.

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

A: We ackknowledge that the models may not be fully converged here. This experiment is quite large scale and as such takes a lot of compute to perform. We have followed the benchmarking details from Bai et al.

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

A: We thank the reviewer for pointing out paper c1 which we had not mentioned in our related work. We have added this as a baseline to our tasks, see Table 1 and Figure 3 in the revised manuscript. We are currently running the iWildCam experiment and will update if this gets completed before the deadline.

Dear reviewer mzUQ,

Thanks again for reviewing our work. We believe our revision and response should address your concerns and kindly ask whether you would consider updating your review. If there are any remaining questions we can answer before the deadline tomorrow, please let us know.

Best, Authors

We thank VuV9 for their review. As you can see in the revised version of the paper, all of the comments have been addressed.

• (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.

A: We have rephrased this paragraph to reflect better our intent: "A common (implicit) assumption in federated learning is that the learned model will be applied in a target domain $T(X,Y)$ that coincides with the marginal distribution of clients ... This is reflected in trained models being evaluated in terms of their average performance over all clients." The identically-distributed-client setting is a special case of that. We are explicitly challenging and extending this assumption to situations with non-identical distributions.

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

A: We have changed the text to clarify that we refer to client distributions.

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

A: The risk minimization objective is completely general and not specific to any loss function. In our classification experiments we use the cross-entropy loss, for example. We have added a mention of this in section A of the appendix.

• (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.

A: The reviewer is correct that there is a privacy concern with sharing label distributions with the server. However, there is always a trade-off between accuracy and privacy in federated learning since model updates also reveal information about clients' data sets. Our work demonstrates a clear benefit in terms of accuracy (in several cases, more than 10%) if the clients are willing to share their label marginals. It is up to clients to decide whether this tradeoff is worth it for them.

• (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."

A: We have added this suggestion to the revised manuscript.

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

A: It is correct that Proposition 1 does not give guarantees for the multiple-local-iterations case, and we never claim that it does. As the reviewer points out, we already made this remark on ln 285-289. We have moved this remark closer to the proposition itself. Nevertheless, our result is justified in the SGD setting under known label shift, strictly more general than the justification for FedAvg, and we demonstrate empirically that there is a strong benefit to our aggregation scheme also in the non-SGD setting.

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

A: L 161 does not contain a reference to either $\alpha$ or $S$ so it is difficult to know which version of $\alpha$ is intended. Generally, $\alpha$ refers to the weights of a convex combination. It is used for the first time in the definition of the simplex before equation (3). It is then defined for specific combinations in (3), (4), (6), and (7). $S(Y)$ refers to the matrix of all client label marginals. We have clarified this in ``Case 2'' on page 6.

We just wanted to update the reviewers to the fact that the iWildCam experiment has finished and the AFL baseline does not appear to learn effectively at all, only reaching ~0.005 F1 on average. This failure is in line with the results reported in Bai et al. where they report a large reduction in performance in the heterogeneous setting for AFL. See the results for $\lambda=0$ in Table 3 in their paper for reference.