Causality Inspired Federated Learning for OOD Generalization

Thank you for your rigorous and responsible review.

In response to the reviewer's concern that “some detected causal features may, in fact, be non-causal,” we clarify that the features described above are indeed causal. The misunderstanding appears to stem from our use of the misleading term “fake causal features” to denote causal features that exist in general but are absent in the current image. Importantly, “fake causal features” do not refer to misidentified non-causal features. Perhaps a more precise term would be “inactive causal features” and we will replace “fake causal features” with “inactive causal features” in the final version.

More specifically, after eliminating the non-causal features $Z_S$ in lines 281–283, the remaining causal features $Z_C^g$ can be further divided into active causal features $Z_C^U$ and inactive causal features $Z_F^U$. The feature compressor then performs adaptive feature selection based on the test data distribution $U$, retaining the active causal features $Z_C^U$ and discarding the inactive causal features $Z_F^U$ that are not present in the current data. For example, a cat paw is a valid causal feature for classifying a cat. However, in an image where only a cat head is visible, the cat paw becomes an inactive causal feature. Inactive causal features may confuse the model, leading to performance degradation, so they should be excluded from the final features.

Thus, the introduction of the counterfactual generator enables the elimination of non-causal features $Z_S$, while the feature compressor eliminates inactive causal features $Z_F^U$ and amplifies active causal features $Z_C^U$.

Thanks for your suggestions, and we will include the citation you pointed out in the final version.

Q1: Experimental Designs

We conduct additional experiments in response to the suggestions about experimental designs.

Add Updated Model Architectures: Following the reviewer's suggestion, we include experiments with ResNet-50 and ViT on PACS dataset, and our method still achieves superior performance. Notably, the selected models in original submission was guided by the experimental settings of existing works to ensure a fair comparison.

Increase Client Number and Add Client Sampling: We added experiments on the CMNIST dataset with 100 clients, adopting the widely used random client sampling strategy by selecting 20 clients per round for training.

Consider Non-iid data: Reviewer suggested considering non-IID data. However, we have already considered non-IID settings. The experiment results presented in Tables 1 and 2 on page 7 of the original submission were all obtained under non-IID conditions, addressing both class imbalance and covariate shift. A detailed description of the non-IID setup can be found in Appendix D.2.

Q2: Broader Scientific Literature and References

We closely follow the reviewer’s suggestion to conduct comparisons, and we will cite these works in the final version. As we have cited the advanced versions from the same team (Tang 2023, 2024), the earlier versions Ref [1,2] were not included in the original submission.

We evaluate the provided methods on the CMNIST dataset with 100 clients. We abbreviate the method in Ref [1] as DRC. Compared to FOOGD, Ditto, and CG-PFL, our method explicitly considers the extraction of causal features. In contrast to DRC, which focuses on preserving causal features specific to individual clients, our approach retains causal features across all clients. The experimental results further demonstrate the effectiveness of our motivation.

Q3: Question about Counterfactual Generator

How semantic information is preserved: We need to clarify that in the original submission we explicitly addressed this issue during model training by incorporating a loss term $L_{REG}$ , which encourages the preservation of semantic information. In particular, we provide a detailed explanation in Section 4.1(2) Preserve semantic information and offer a theoretical proof of its effectiveness in Lemma 4.3 on Page 5.

Generation Sensitivity: We need to clarify that in the original submission we have defined the hyperparameter $\alpha$ to measure sensitivity to generation quality, where $\alpha = 0$ indicates generation without semantic constraints, and larger values of $\alpha$ impose stronger constraints. As shown in Fig. 7(a) on page 8 of the original submission, the accuracy only drops from 90.4% to 88.5% when switching from the best setting to a setting without semantic constraints—a gap of only 1.9%. This demonstrates that FedUni is robust with respect to generation quality and hyperparameters.

Why the simple model architecture works: Our counterfactual generator does not generate images from noise, instead, it adds noise (i.e., distribution shifts) to the original image. Since the original image remains part of the input, a simple model architecture is sufficient to simulate distribution shifts. Furthermore, this architecture is widely used in style transfer methods [1,2,3] and has been shown to efficiently generate distribution shifts while preserving semantic information.

[1] Huang, et al. Arbitrary Style Transfer in Real-time with Adaptive Instance Normalization. 2017.

[2] Guo, et al. Single Domain Generalization via Unsupervised Diversity Probe. 2023.

[3] Yang, et al. Practical Single Domain Generalization via Training-time and Test-time Learning. 2024.

Q4: Computation Cost

Please refer to our response to Reviewer FyJy 4. FedUni achieves significantly higher accuracy, while maintaining comparable computational cost to many state-of-art approaches.

Thanks for your valuable suggestions, and we will enlarge the text font in Fig 5, 7 in the final version.

Q1: Computational Cost

We provided the run time per communication round and a theoretical proof for the convergence rate. Due to space limitations, we provide a brief proof here, and the full proof will be included in the appendix of the final version.

Firstly, we measured the runtime per communication round (20 clients with 5 local training epochs) on two commonly used GPUs. As illustrated in the table below, FedUni achieves approximately 19% higher accuracy compared to the second-best baseline, while maintaining comparable computation cost to many state-of-art approaches.

In addition, we presented a brief proof for the convergence rate, showing that our convergence speed is also comparable to that of other methods.

We assume that the objective functions $(F_{1}, \cdots, F_{N})$ satisfy L-smoothness and μ-strong convexity. The variances of the stochastic gradients are bounded, and the expected squared norms are also bounded. Meanwhile, $(\Gamma=F-\sum_{k=1}^{N} p_{k} F_{k})$ is used to quantify the degree of non-iid data. In the case of full device participation, by choosing appropriate parameters $\kappa=\frac{L}{\mu}$, $\gamma = (\max{8\kappa, E})$ and the learning rate $\eta_{t}=\frac{2}{\mu(\gamma + t)}$. Our proposed algorithm satisfies $\mathbb{E}[F(w_{T})]-F^* \leq \frac{\kappa}{\gamma+T-1}(\frac{2B}{\mu}+\frac{\mu \gamma}{2} \mathbb{E}\| w_{1}-w^*| ^{2})$, where $B=\sum_{k=1}^{N} p_{k}^{2} \sigma_{k}^{2}+6 L \Gamma+8(E-1)^{2} G^{2}$, $\sigma_{k}^{2}$ represents the upper bound of the variance of the stochastic gradient in the k-th device, E represents the number of local updates performed by each device between two communications. This indicates that when all devices participate and specific conditions are met, our proposed method can converge to the vicinity of the global optimal solution at a rate of $O(\frac{1}{T})$.

Q2: When to choose FedUni

The main advantage of FedUni is its ability to extract and retain all causal features from any input, making it well-suited for scenarios with unknown test data distributions and client heterogeneity.

Unknown test data distributions: Our method is inherently distribution-agnostic, performing robustly in both ID and OOD settings. In contrast, existing approaches typically focus on either ID or OOD settings. ID-based methods inevitably exploit spurious correlations in the training data, undermining OOD generalization, while OOD-based methods usually discard client-specific knowledge, reducing ID discrimination ability. FedUni addresses both issues concurrently by incorporating a counterfactual generator to enhance OOD generalization and a feature compressor to preserve valuable client-specific information.

Client heterogeneity: Our method is particularly effective in settings with data heterogeneity among clients, where existing approaches may discard valuable client-specific knowledge. Our method can preserve a broader set of causal features, leading to improved model performance.

Q3: Performance Gains in Homogeneous Setting

Reviewer pointed out that when training clients are similar, FedUni may not yield significant improvements. We conducted experiments in homogeneous setting (on the CMNIST dataset). FedUni still can outperforms other approaches. Whether to use FedUni in homogeneous setting depends on the trade-off between performance gain and computational cost.

Q4: More rigorous argument for Lemma 4.2

Thanks for your review! We will revise Lemma 4.2 into a more rigorous form in the final version.

Inspired by ref [1], Lemma 4.2 relies on adopting specific $l$-norms to approximate the entropy terms in the coefficient of constraint. In particular,we rely on two main approximations. First, for the conditional entropy $H(x'|x)=-\mathbb{E}_X(\log P(x'|x)) \approx \mathbb{E}_X[\|\|x-\psi(x)\|\|_1]$, where the approximation adopted in the last step amounts to assigning a $l_1$-Laplace distribution with identity covariance to the conditional probabilities:$P(x'|x)= \mathcal{L}(x';\mu(x), I)\propto \exp (-\|\|x'-x\|\|_1)$. Similiarily, for $H(x')$, we can derive $H(x')\approx\mathbb{E}_X[\|\|\psi(x)\|\|_1]$.

[1] Pedro Savarese,et al. Information-Theoretic Segmentation by Inpainting Error Maximization. 2021.