Learning Personalized Causally Invariant Representations for Heterogeneous Federated Clients

Dear Reviewer Qvuh,

Thanks for your valuable feedbacks. As the deadline for author-reviewer discussion is approaching, we are concerned you might miss the clarifications that we make for your questions. We would appreciate it very much if you could reassess our work according to the clarifications. Of course, more discussions are welcomed if you still have any concerns or questions for this work.

Thanks for the reviewer's valuable feedbacks and suggestions. The detailed responses to the reviewer's questions and concerns are listed as follows:

1. I believe the main issue with this paper is that the research motivation seems weak. The paper claims to primarily address the scenario where training sets are Non-IID among clients, and where within each client, the training and test sets are also Non-IID. However, I think existing Robust Federated Learning and Federated Domain Generalization methods are capable of handling the aforementioned scenario. Although these two methods primarily focus on the issue of Non-IID training and test sets within each client.

Answer: Unfortunately, both robust federated learning and domain generalization methods cannot handle the scenarios well where target datasets are Non-IID across local clients. Since both Federated Domain Generalization (FedDG) and Robust Federated Learning (RFL) generate a shared global model, they abandon the important personalization information on each client. Therefore, the relation between our method and RFL/FedDG is same as the relation between the existing personalized federated learning (FL) and traditional federated learning (FL).

2. The existing Robust Federated Learning and Federated Domain Generalization methods can also address the shortcut trap problem. Therefore, the contribution of this paper appears to be limited.

Answer: The existing Robust Federated Learning (RFL) and Federated Domain Generalization (FedDG) methods will abandon the important personalization information on each client since they output a shared global model, although they can also mitigate the shortcut features. In many real-world applications (e.g., medical diagnosis and automatic driving), personalization information is of vital importance and cannot be overlooked. The main contribution of this paper is to fully exploit the significant personalization information, and in the meanwhile eliminate the shortcut features. It's rather challenging because the personalized and shortcut features are usually entangled in the real-world scenarios.

3. In Figure 1, I find the classification of the scenarios where RFL&FedDG are applicable not very accurate. In the scenarios where RFL&FedDG are applicable, the Test-test relation can also be IID or Non-IID.

Answer: The existing Robust Federated Learning (RFL) and Federated Domain Generalization (FedDG) methods will abandon the important personalization information on each client since they output a shared global model, when target datasets are Non-IID across local clients. Therefore, we conclude that RFL&FedDG can't cover the scenario well. The relation between our method and RFL/FedDG is same as the relation between the existing personalized federated learning (FL) and traditional federated learning (FL).

4. In the experimental section, the proposed model is only compared to FL and PFL methods. However, I believe it should also be compared to RFL and FedDG methods to provide more convincing experimental results.

Answer: Actually, we included three RFL&FedDG methods as baselines in the experimental section. The considered competitor DRFA is a typical RFA method while FedSR and FedIIR are two state-of-the-art FedDG methods. The overall experimental results of these baseline methods can be found in Table 1.

[1] Yuyang Deng, Mohammad Mahdi Kamani, and Mehrdad Mahdavi. Distributionally robust federated averaging. Advances in Neural Information Processing Systems, 33:15111–15122, 2020.

[2] A Tuan Nguyen, Philip Torr, and Ser-Nam Lim. Fedsr: A simple and effective domain generalization method for federated learning. In Advances in Neural Information Processing Systems, 2022.

[3] Yaming Guo, Kai Guo, Xiaofeng Cao, Tieru Wu, and Yi Chang. Out-of-distribution generalization of federated learning via implicit invariant relationships. In Proceedings of the 40th International Conference on Machine Learning, pp. 11905–11933, 2023.

5. The experimental section lacks detailed information about the configuration of the model used in this paper. More detailed information is needed.

Answer: The detailed information about the adopted models were provided in section 5.1. For CMNIST and CFMNIST, we adopt the deep neural network with one hidden layer of (128, 128) as feature extractor and an subsequent fully-connected layer as classifier. As regard to Waterbird and PACS, ResNet-18 is used as the learning model where the part before the last fully-connected layer works as feature extractor and the last fully connected layer works as classifier.

Thanks a lot for the reviewer's valuable feedbacks and suggestions. The detailed responses to the reviewer's questions and concerns are listed as follows:

1. Can you add a synthetic dataset with a simple model in the experiments? It can be used to verify the theorems by satisfying the assumptions.

Answer: We generate a synthetic dataset using the same strategy as in [2]. Specifically, it is a logistic regression task and the data instance $X$ is generated by $X=g(Z_C^g, Z_C^U, Z_S)$, where the dimensionalities of $Z_C^g$, $Z_C^U$ and $Z_S$ are $d_C^g=3$, $d_C^U=3$ and $d_S=6$ respectively. The linear function $g$ is implemented by one fully-connected layer which has 12 neurons. The latent variables $Z_C^g$, $Z_C^U$ and $Z_S$ are subject to $\mathcal{N}(y\cdot \mu_{c,g}, \sigma_{c,g}^2I)$, $\mathcal{N}(y\cdot \mu_{c,u}, \sigma_{c,u}^2I)$ and $\mathcal{N}(y\cdot \mu_s, \sigma_s^2I)$ respectively. Target variable $y$ is taken from the distribution $\mathbb{P}(y=-1)=\mathbb{P}(y=1)=0.5$. Both $\mu_{c,g}$ and $\mu_{c,u}$ are randomly sampled from $\mathcal{N}(0, 1.5I)$ while $u_s$ is randomly sampled from $\mathcal{N}(0, 0.75I)$. To make the shortcut representation $Z_S$ easier to learn, we choose $\sigma_{c,g}=\sigma_{c,u}=2$ and $\sigma_s=1$ as in [2]. Each fixed value of $\mu_s$ indicates one environment. We generate 10 training environments and 5000 test environments to evaluate the out-of-distribution generalization performance. Each (training/test) environment contains 10000 data samples (X, y) and the training data samples are distributed onto totally 100 clients. The training and test data samples on each client are generated with an identical value of $\mu_{c,u}$. Besides, we choose the client sampling rate as 0.1. The experimental results on this synthetic dataset are shown in the following table:

In particular, when we manually select the causal features $[Z_C^g, Z_C^U]$ as the discriminating features, we find the optimal personalized classifiers achieve an stable accuracy around $97.5$ in different test environments. Therefore, the results shown in Table 1 can demonstrate the effectiveness of our FedSDR on developing the optimal personalized invariant predictors, compared with the state-of-the-art FL and PFL methods.

[2] Rosenfeld, E., Ravikumar, P., & Risteski, A. (2021, July). The Risks of Invariant Risk Minimization. In International Conference on Learning Representations.

2. Can you add experiments without manually adding shortcut features? It is quite important. Currently, I’m not clear what are the real applications of FedSDR.

Answer: Actually, PACS is a real-world dataset that is commonly used in related papers. It's a multi-class classification task and there doesn't exists explicit (or manually adding) shortcut features. Specifically, It consists of 7 classes distributed across 4 environments (or domains). The results on PACS dataset are also shown in Table 1.

Thanks a lot for the reviewer's valuable comments and suggestions. The detailed responses to the reviewer's questions and concerns are listed as follows:

1. The paper’s algorithmic approach (Algorithm 1) deviates from the conventional practice of client selection in Federated Learning (FL). Could the authors elucidate why they've chosen this approach over the typical FL client selection methodology?

Answer: We adopted the same federated learning framework as in pFedMe [1] while the conventional client selection is conducted before local updates. In the adopted framework, the global model is generated by aggregating the updates from the selected clients and the updates from the clients that are not selected are never used. Then on each client, the local model will be initialized by the global model before local update. Therefore, the final output models of these two client selection methods are exactly identical. Our method FedSDR can also be implemented using the conventional client selection methodology and the experimental results will keep unchanged.

[1] T Dinh, C., Tran, N., & Nguyen, J. (2020). Personalized federated learning with moreau envelopes. Advances in Neural Information Processing Systems, 33, 21394-21405.

2. Concerns arise from certain experimental results in Table 1, particularly the unexpected lower accuracies observed for specific baseline methods on CMNIST and CFMNIST. Can the authors provide an explanation to address these concerns?

Answer: As shown in Table 1, our FedSDR achieved around 6.5%, 9%, 3.5% and 2% higher worst-case accuracy than the second best algorithm among the baselines and in the meanwhile reaches the highest average accuracy on CMNIST, CFMNIST, WaterBird and PACS, respectively. So, what do you mean by "the unexpected lower accuracies observed for specific baseline methods on CMNIST and CFMNIST"? In addition, we will make the project codes public soon to enhance the reproducibility of the conducted experiments.

Thanks again for your valuable feedbacks. More discussions are welcomed if you still have any concerns or questions for this work.