Causality Inspired Federated Learning for OOD Generalization
Authors: Jiayuan Zhang
Xuefeng Liu
Jianwei Niu
Shaojie Tang
Haotian Yang
Xinghao Wu
Venue: ICML 2025 poster (2025)
Keywords: federated learning
causality
out-of-distribution generalization
Primary area: optimization->large_scale_parallel_and_distributed

Abstract: The out-of-distribution (OOD) generalization problem in federated learning (FL) has recently attracted significant research interest. A common approach, derived from centralized learning, is to extract causal features which exhibit causal relationships with the label.  However, in FL, the global feature extractor typically captures only invariant causal features shared across clients and thus discards many other causal features that are potentially useful for OOD generalization. To address this problem, we propose FedUni, a simple yet effective architecture trained to extract all possible causal features from any input. FedUni consists of a comprehensive feature extractor, designed to identify a union of all causal feature types in the input, followed by a feature compressor, which discards potential \textit{inactive} causal features.  With this architecture, FedUni can benefit from collaborative training in FL while avoiding the cost of model aggregation (i.e., extracting only invariant features). In addition, to further enhance the feature extractor's ability to capture causal features, FedUni add a causal intervention module on the client side, which employs a counterfactual generator to generate counterfactual examples that simulate distributions shifts. Extensive experiments and theoretical analysis demonstrate that our method significantly improves OOD generalization performance.

Causality Inspired Federated Learning for OOD Generalization
Jiayuan Zhang' Xuefeng Liu 12
Jianwei Niu
12
Shaojie Tang
Haotian Yang'
Xinghao Wu
Abstract
1. Introduction
Federated Learning (FL) has emerged as a privacy- preserving framework for collaborative learning across dis- tributed clients. A practical problem, how enable the FL model to robustly generalize to target clients with un- known distributions, referred to as the out-of-distribuuiin (OOD) generalization problem, has recently attracted sig- nificant research attention. Unfortunately, deep learning models are prone to relying on non-causal relationships (i.e., shortcuts) for prediction (Geirhos et al., 2020), leading to significant performance degradation when the test dis- tributions fall outside the training scope. For instance, in Fig. 1, 'dogs' and 'grass' exhibit high statistical dependence in the training dataset of Client A, which can easily mis- lead the model into making predictions based on the grass background (non-causal features) rather than the shape of dogs (causal features). Recent works centralized learning scenarios have proposed addressing the OOD generalization problem through the lens of causality, focusing learning causal features that maintain an invariant causal relationship to the output label across various data distributions. It is widely recognized that causal features are far more robust to distribution shifts compared to non-causal ones.
The out-of-distribution (OOD) generalization problem in federated learning (FL) has recently attracted significant research interest. common approach, derived from centralized learning, is to extract causal features which exhibit causal re- lationships with the label. However, in FL, the global feature extractor typically captures only invariant causal features shared across clients and thus discards many other causal features that are potentially useful for OOD) generalization. To address this problem, we propose FedUni, sim- ple yet effective architecture trained to extract all possible causal features from any input. FedUni consists of a comprehensive feature extractor, de- signed to identify a union of all causal feature types in the imput, followed by a feature compres- sor, which discards potential inactive causal fea- tures. With this architecture, FedUni can benefit from collaborative training in FL while avoiding the cost of model aggregation (i.c., extracting only invariant features). In addition, to further enhance the feature extractor's ability to capture causal fea- tures, FedUni add a causal intervention module on the client side, which employs counterfac- tual generator to generate counterfactual exam- ples that simulate distributions shifts. Extensive experiments and theoretical analysis demonstrate that our method significantly improves OOD gen- cralization performance.
However, in FL, the global feature extractor is typically trained extract causal features that remain invariant across all clients (Zhang et al., 2021; Nguyen et al., 2022; Tang et al., 2022; Guo et al., 2023b; Liao al., 2024). Although these invariant features are generally robust and reliable, many valuable causal features that are potentially useful for OOD generalization are inevitably discarded. This lim- itation becomes more pronounced in the presence of data heterogeneity, where only small subset of invariant causal features shared across clients. As an illustrative example, consider the task of learning global model for dog classifi- cation in FL, as shown in Fig. 1. In this setup, Client A has real-world photos of dogs, while Client B's dataset consists of dog sketches The global feature extractor, trained to capture invariance across clients, preserves only invariant causal features like shape while discarding client-specific ones such details of the dog's mouth and tail. When the test data distribution matches that of Client A's 'ss training data, discarding these detailed causal features can lead to suboptimal generalization. In contrast, preserving all causal features enables clients to flexibly select those most relevant
Sttte Key Laboratory of Virtual Reality Technology and Sys- tems, School of Computer Science and Engincering. Beihang Uni- versity "Zhongguancun Laboratory, Beijing, China 'Department of Management Science and Systems, School of Management, Center for AI Business Innovation, University Buffalo, Buf- falo, NY, USA. Correspondence to: Jianwei Niu <niujian- wei www.........a....o.a..a. en?. Xuefeng Liu <liu_xuefeng buaa.cdu.cn.
Proceedings of the 42nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Causality inspired Federated Learning for OOD Generalization
into counterfactual examples that simulate distribution shift between source and target domains while preserving the semantic information. The main contributions of this work are outlined as follows:
client A
FL
Background gross
Invariant
Invariant Causal features
x
Train Clients
Causal Features
Target Client
client B
Enpe Teri
Ours
Z
Select
We formulate a structural causal model (SCM) for the OOD generalization problem in FL and propose that the global feature extractor should retain all causal features rather than limiting itself to only the invariant ones shared across clients.
Style:
skotch
Select
Spacific
Target client C
Ta
We introduce an effective method, FedUni, which in- corporates a comprehensive feature extractor to iden- tify the union of all causal feature types in the input. This is followed by a feature compressor to discard potential inactive causal features. To further enhance the ability capture causal features, FedUni employs a counterfactual generator for causal intervention.
Figure 1. Taking dog classification as an example to illustrate Fe- dUni. UndUni preserves all causal features present across clients, rather than only the invariant causal features, thereby achieving better generalization to unseen target clients.
to their target domain, thereby improving generalization performance. For instance, the causal features of samples from the cartoon domain may constitute combination of features found in both real-world photos and sketches
Extensive experiments and theoretical analysis demon- strate that our method significantly improves OOD generalization performance.
To address the above limitation, we propose that the global feature extractor in FL should capture the full set of causal features present across all clients, rather than only the in- variant subset shared by all. However, this objective is chal- lenging because different clients possess varying quantities and types of causal features. Additionally, not all causal fea- tures are necessarily beneficial for generalization to given target client. For example, in Fig. I, the causal features extracted by the global feature extractor contain all causal features from both Client and Client B. However, when deployed target client whose data distribution more similar to Client B's sketch style, the detailed features of dogs become irrelevant to the target domain. Retaining these irrelevant causal features may even degrade generalization performance.
2. Related Work
Causal Feature Learning. One promising approach to ad- dressing OOD generalization is leveraging causal features underlying the observed data (Zhou et al., 2022). A series of works derived from Invariant Risk Minimization (IRM) (Ar- jovsky et al., 2019) formulate invariant causal mechanisms by incorporating information theory and nonlinear predic- tion functions (Ahuja et al., 2021; Krueger et al., 2021; Yang et al., 2023). More closely related to our method, (Chen et al., 2024) reveals that existing methods may still learn fake causal features and introduces conditional mutual information to rectify them. Another line of research gen- erates counterfactual samples and enforces consistency of causal features hetween original and augmented samples, as shown in (Chang et al., 2021; Lv et al., 2022; Noohdani et al., 2024). However, these methods focus on centralized scenarios where all training data is accessible.
To overcome this challenge, we introduce a simple yet effee- tive architecture, named FedUni, which consists of two key components: a comprehensive feature extractor followed by feature compressor. The former is designed extract the union of all causal feature types across clients. Notably, if given input lacks a particular type of causal feature, the corresponding feature is identified as an inactive causal fea- ture. The feature compressor is then employed to effectively discard these inactive causal features, mitigating their poten- tial negative impact on the model's performance. These two components are trained simultaneously, so that the feature extractor can benefit from model aggregation while avoiding the extraction of only invariant features. In addition, to fur- ther enhance the feature extractor's ability to capture causal features and eliminate non-causal ones, we incorporate a causal intervention module on the client side. This mod- ule employs counterfactual generator to transform images
Federated Learning for OOD Generalization. The OOD generalization problem in FL has gained significant research interest. Existing methods extract invariant features via ad- versarial domain alignment (Zhang et al., 2021; Qi et al., 2025; 2023; 2024; Meng et al., 2024), feature regulariza- tion (Nguyen al., 2022), or inter-client gradient alignment (Guo al., 2023b; Wu et al., 2022). However, they primar- ily focus on inter-client invariance, potentially discarding valuable client-specific information beneficial for OOD gen- cralization. Meanwhile, FedSDR (Tang et al., 2023) and FedPIN (Tang et al., 2024) aim to preserve personalized causal features. However, these methods are designed for personalized federated learning (PFL. (Tang et al., 2021;
2

---

Causality inspired Federated Learning for OOD Generalization
2022), where the causal features of each training client are assumed to be known. Our problem is more challenging as the types and quantities of causal features for a target client are unseen during training. Retaining only invariant features may discard valuable ones, whereas preserving all causal features could introduce traning-sppecife causal fea- tures that are irrelevant to the target distribution, potentially degrading OOD generalization performance.
E
U
E
Y
U
Z5
Z
X
X
(a) Existing SCM
(b) Restructured SCM
A more detailed discussion of related work is provided in Appendix A.
Figure 2. SCMs for FL OOD generalization. (The solid circle represents observed variables, the dashed circle represents latent variables, the solid arrow indicates causal dependency, and the dashed arrow represents Markov dependency.)
3. Problem Formulation
Notations. Let X, y and & represent the input, target and environment space respectively. Suppose there are N clients, and the local data De, the e-th client, which contains ne samples, is drawn from the training environ- ment Elr (ee E Etr). For convenience, following (Guo et al., 2023b), we formalize the model as f = w o o where : X' + Z C. R" is the feature extractor that maps the input space X' to the feature representation space Z and w 2 + y is the classifier. The overall model is param- eterized fo(.)
generalization performance. We propose to solve the OOD generalization problem in FL through the lens of causality.
3.2. A Causality Viewpoint to FL OOD Generalization
It is known that OOD generalization is possible under prac- tical causal assumptions (Ahuja al., 2021).
w(((()), where = (w,, <). In the standard FL paradigm, the goal is to learn global model that minimizes the average expected risk over all participat- ing clients. The global expected empirical loss is denoted as R())= = Ee (ewEEE ((fo(z))))) y)), where / the loss function.
Existing SCM. Recently, some works have formulated SCM in FL considering the distribution heterogeneity among clients (Tang et al., 2023; 2024), as shown in Fig. 2a, where is the Environment variables, and E f1....,,N N. is the User/Client indicator. Zg denotes non-causal vari- ables. The causal variables are divided into two compo- nents: the shared/global causal variables Zg and the person- alized/local causal variables Zy The personalized causal variables zy include both Zg and additional invariant in- formation derived from U.
3.1. FL OOD Generalization
The goal of FL OOD generalization to learn a global model fe with D = :D= c [I, }} to enable prediction on samples drawn from arbitrary environments in E. This objective can be formulated as min-max opti- mization problem follows:
Assumption 3.1 (Existing SCM illustrated in Fig. 2a).
Zs == fapu E, Y), 2z == Jivo (U, ZZ), X = Jeen (Zs, Zg. zZ). y w( zZd)
Re(() == max min R (fo)).
(1)
eng
Obviously, the optimization problem in Eq. (1) cannot be solved directly, as only the training environment E,, instead of E is accessible during training. A common assumption in FL OOD generalization is that there exist latent causal variables (features) Ze Rd, which possess the following environment-inarrian property (Arjovsky et al., 2019):
According to Assumption 3.1, existing methods assume that the global model should extract invariant causal variables ze - 2tz= shared across clients. However, as shown by Lemma 3.2, this objective may lead to the ex- tracted invariant features being incomplete.
Lemma 3.2. the data generating process of each client follows the SCM assumption illustrated in Fig. 2a, the global invariance 7g, incomplete from two perspectives:
Zo E = e) =
= e'), Ve, E E. (2)
To enable the global feature extractor f extract Zc, FL methods extract inter-client invariant features from distri- butions (Liu et al., 2021), representations (Nguyen al., 2022), and gradients (Guo et al., 2023b). However, these methods focus only on the invariance of causal features, without considering their completeness. Retaining only an invariant subset of causal features may discard useful in- formation for OOD generalization, leading to suboptimal
Y 1 U Z&: Under Zg, the prediction of Y is still influenced by the client indicator U.
Y 1 U. Y Zg. Given Zg and z Y is con- ditionally independent of the environment variable E and user indicator U.
3

---

Causality inspired Federated Learning for OOD Generalization
The first claim in Lemma 3.2 suggests that the global model may exhibit inconsistent performance across differ- ent ID/OOD data distributions. The second claim indicates that combining Zú and zğ can further reduce the empirical risk, and that z should be retained for better inference performance. This indicates that existing SCMs are not well suited for the OOD generalization problem in FL, where the global model is expected to exhibit fairness and robust adaptability across heterogeneous data distributions.
4. Methodology
Overview. Inspired by the restructured SCM discussed above, we propose a novel FL paradigm named FedUni, consisting of three components: (1) a causal intervention module the client side to distinguish causal features Zi from non-causal ones Zs (Section 4.1), (2) an inac- tive causal feature compressor that discards inactive causal features z from all causal features Zó specific to the data distribution (Section 4.2), and (3) A comprehensive feature extractor designed to capture the union of all causal fea- ture types present in a given input, achieved by training jointly with the other two modules. The comprehensive feature extractor and inactive causal feature compressor are uploaded for global aggregation and transferred to the tar- get clients, while the causal intervention module remains local to the training clients. The theoretical analysis of the generalization boundary is presented in Appendix C.
Restructured SCM. To better address the FL OOD gen- eralization problem, we propose reconstructed SCM, as shown in Fig. 2b. A key distinction in our formulation lies in the definition of the global causal variable Zg. Unlike ex- isting approaches that define zg as the intersection of client- specific causal variables, we define -g- {U=M= =
as the union of zy across all clients for any given input. In addition, since the global causal variable may contain information that is not required for the current task, we introduce z to represent these Inactive causal features, which are causal features specific to certain distributions but non-causal for the target task. Formally, we define this as z = Z
4.1. Causal Intervention Module
We expect the FL model to extract causal features z& while eliminating non-causal ones Zg. Unfortunately, directly re- constructing the causal features Zü impractical. Nonethe- less, inspired by (Lv et al., 2022), we propose two properties that these features should satisfy.
Assumption 3.3 (Our Restructured SCM illustrated in Fig. 2b).
Zs = fapu ( Y), zg= = (Zn) t = fino ZE), x= = sgen Zz. Lz, ) Y = w(Z))
Proposition 4.1. Causal variables Lg should satisfy the following two properties:
Lemma 3.4. If the data generating process of each client follows the SCM assumption illustrated in Fig. 2b, then the following three statements hold:
Invariance under causal intervention: Based on the first claim in Lemma 3.4, we have Zg: 1 E Y. Thus, performing an intervention upon E does not make changes to Z2g-
Zs 1 12g. 2 Y: Given Y, both the 7& and z4p are conditionally independent of Zg.
Sufficiency of semantic information: The causal vari- able ZiC should be causally sufficient the label pre- diction, i.e. maximizing T(7g. Y).
Y 1 IU. Zg: Considering that 2 S Zg, com- ponents in (U, F] are conditionally independent of the target Y given 2& and zŞ.
Existing methods typically perform predefined data augmen- tations on samples serve causal intervention module (Lyv et al., 2022; Chen et al., 2023). However, these ap- proaches lack diversity and are insufficient to simulate the distribution shifts in the real world. Inspired by recent single domain generalization (Guo et al., 2023a; Yang et al., 2024), we employ a trainable counterfactual generator g XX' < X", parameterized by 1/, to generate more diverse counterfactual samples while preserving semantic information. Below, we propose the optimization objective for g.
Y y 1z. 24. Y is not conditionally independent of Z& given Zs and z%
Lemma 3.4 highlights the potential for designing better FL method for OOD generalization. Based on the first claim, the non-causal features can be directly excluded from the causal ones, and the second claim indicates that the global invariant features zg can be extracted via collaborative learning among clients. Moreover, the third claim indicates that we should reconsider the Markov dependency between Zg. and ZV. and these variables need to be distinguished based on their differing roles in label prediction when con- ditioned on each other. All proofs related to this section are provided in Appendix B.
(1) Diversify non-causal variables. To increase the diver- sity of generated counterfactual samples, the generated im- age x' should have the minimal correlation with the source image x. Mutual information (MI), denoted as I(x; z'), provides natural measure of this correlation (Wang et al.,

---

Causality inspired Federated Learning for OOD Generalization
2021):
4.2. Inactive Causal Feature Compressor
p( [x] p(x')=
We now proceed discuss how to compress Z from zg. Notably, the features extracted by the global feature extrac- tor consist of the true causal features Zy given the distribu- tion U, the inactive causal features z and the non-causal parts Zs, which can be denoted as:
1(x;a") = (xx(..))) [log
H(x')- (1'1x, (3)
where (x') is the entropy of x' and H (x' (r)) is the condi- tional entropy of x' given I. To prevent (a') = O, during optimization, we employ normalized variant (x;x² of the mutual information, namely the coefficient of constraint (Press, 2007):
x(x) zy zy U Zs, where z = Zz/ZZ.
(8)
To preserve semantic information, we need to maximize the conditional mutual information (CMI) to strengthen the causal path zg + Y:
min (x;n)( min [(x;z") max + (xx))
H(zl(z) & (xx))
(4)
where Nxx;; x") if and only x' and æ are independent. For simplicity, we use parameter-free approximation to estimate the solution.
max [[; Z Z Zs].
(9)
2g
Similarly, observe an inactive causal dependency path between Z and Y (2% - 42 - -Y). As the label pre- diction Y is expected to be conditionally independent of zy. we have Y 1 z ng. Zs. To compress the inactive causal features, we propose minimizing the CMI between Y and ZP. conditioned on 2ó and Zs:
Lemma 4.2. H (x'lx) can be approximated by the condi- tional probability p(z' (xr) as C Li-Laplacians, as follows:
(x' (xx) = - (log (()) v [|]]==)((())))i)
Similarly, (xz') can be approximated by the marginal dis- tribution of z':
min [x; 2 25 Es}.
(10)
H(x') - P(x')) x B (1bx))ll
For simplicity, we assume that after the causal intervention steps, the non-causal features Zs have been removed, i.e. s" (X) zy U z and Z &0 *x(x) Then, Eq. (9)
The detailed proof of the above lemma can be found in Appendix B.4. Based on this approximation, we define the following information-theoreti loss:
and Eq. (10) can be simplified to:
(<xlla) + (x²)
sO max (W(x) xll (1)()
DDV = max
(5)
mingy AI(Y: ZZ |
- I -
z3]. (11)
Lemma 4.4. The CMI admits the following decompositions:
(2) Preserve semantic information. To ensure that the semantic information of x' and does not change drastically, we impose constraints on the unbounded generation domain. Lemma 4.3. Maximizing the diversity loss is approximately equivalent to maximizing the lower bounds of LDIV, which satisfy the following inequality:
T (X)/22)
(x)/z zY} = -1 (X) (Y))
The proof of this lemma is provided in Appendix B.5. Ac- cording to Lemma 4.4, Eq. (11) can be decomposed as:
max LDVV 2 max E.y $
(1
|
min (Y 60 (X) (X)/Z) (H(1)Z)). 2
(12)
[xd]
The objective implies that we only need to compress fea- tures from ox()) to strengthen the causal dependency path ny + Y and discourage the inactive causal path z + zg + Y. We propose an inactive causal feature compressor 8: 2 - Z', which is a simple two-layer archi- tecture. The compressor s(-) receives the features extracted by the global feature extractor ox-) and performs soft feature selection to compress Z
It can be observed that generating the counterfactual image x' essentially involves maximizing ||||=/||, ||a]|| llr-
The detailed proof of this lemma can be found in Appendix B.4. Following Lemma 4.3, we constrain the distance be- tween the generated image and the original image by adding a penalty on the -norm of the samples:
LRGG = min E.e || 1|1 - ||
(6)
Optimization. The optimization objective of the inactive causal feature compressor in Eq. (11) can be expressed as follows:
Optimization. The optimization objective for the generator is defined as follows, where a is hyper-parameter used to adjust the strength of the regularization term:
C. min E. o v(x)).v)] m.
(13)
Cg min - LDIV + aLREG. &
[E, (=.[((((1s(((x)))) s(((x))) @ o(())
5

---

Causality inspired Federated Learning for OOD Generalization
where À is a hyper-parameter and @ denotes the element- wise product and (1 (((x))) o ((x)) represents the com- pressed part Zy-
This encourages the feature extractor to capture richer causal features. Meanwhile, this can prevent certain feature loca- tions from being consistently ignored, which could other- wise lead to the inclusion of non non-informative features.
Algorithm FedUni
Optimization. The overall optimization objective can be written in the following form, where used to control the proportion of the causal intervention loss in the training objective:
1: Input: T, K, N, 7) 4g ( 2
2: for ( 0 - do
3: Randomly select a client subset St from N clients.
4: Broadcast w ($)), = to selected clients.
5: for each client E S do
b min LSEM + LADV + BLcr.
(17)
02
6: for local steps to K- do Update counterfactual generator:
The detailed optimization procedure is outlined in Algo- rithm 1.
7:
8:
3/ = v/ Lg
9:
Update compressive global model:
10:
whe & = w/, $ - V
5. Experiments
11:
Update feature compressor:
12:
sy wh = syy w/, L
5.1. Experimental Settings
13:
end for
Datasets. To evaluate the effectiveness of our method, we conduct experiments on both the spurious correlation dataset, i.e. Waterbirds (Sagawa et al., 2019), Colored- MNIST/FMNIST (Arjovsky al., 2019; Ahuja et al., 2020) and the cross-domain datasets, i..., Digits, PACS (Li et al., 2017).
14:
end for
15: Model aggregation:
16:
wy (g,)),, Sg |.T E. (Z [527 st
17: end for
18: Return: global model wy ( (Bg),
Configurations. For Waterbirds and PACS, we use the ResNet-18 (He et al., 2016) pretrained on ImageNet (Deng et al., 2009) as backbone. With regard to CMNIST and CFMNIST, an MLP with one hidden layer serves as the model. The AlexNet (Deng et al., 2009) is selected for Digits dataset. Unless otherwise mentioned, the local up- date step is 5 and the mini-batch size is 64. The learning rate Ir = 1.41 x 10-4 In our setting, a, B, À are set to 0.1, 1, 0.01 respectively. The test results are based on the model that performs best on the validation set sampled from the training data with sampling ratio of 0.2. For more detailed experimental settings, please refer to Appendix D.
4.3. Comprehensive Feature Extractor
The goal of the comprehensive feature extractor is extract all causal features from any input across all clients, which can be achieved by training it simultaneously with the two aforementioned modules.
Firstly, with the causal intervention module, Zs can eliminated by minimizing the causal effect of environmental changes. We use the contrastive loss (Chen al., 2020) the objective to enforce the consistency of features between the original image and the counterfactual image z':
Comparison. We compare our method with 9 state-of-the- art algorithms: FedAvg (McMahan et al., 2017) which is a classic FL method; three methods for addressing hetero- geneity (FedProx (Li et al., 2020), Scaffold (Karimireddy et al., 2020), MOON (Li et al., 2021b)); and four FL general- ization methods (FedSR (Nguyen et al., 2022), FedlRR (Guo et al., 2023b), FedDG-GA (Zhang al., 2023), FedSDR (Tang al., 2023)). It is worth noting that FedSDR is de- signed for Personalized Federated Learning (PFL.); in our experiments, we use a variant that combines FedSDR with FedAvg.
Lcr = Eg log ((()))) Ex`EB ((X)))
(14) where the original and counterfactual scenes of the same instance treated a positive pair, while the features drawn from other counterfactual samples are considered negative pairs. Moreover, the causal features should be sufficient for the label prediction. We use the ERM loss (Vapnik, 2013) to optimize the model to learn the semantic information related to target> Y:
LSEM = R(f) + E,
(15)
5.2. Overall Performance
Secondly, with the inactive causal compressor module, we require that unselected features 2% still retain some seman- tically relevant information:
Effect of spurious correlations. To validate the effective- ness of the proposed method in eliminating the impact of spurious correlations and thereby enhancing OOD general- ization performance, we evaluate the test accuracy of the
LADV min E, (1 s(((x))) d(x)), v)]. (16) G.w
6

---

Causality inspired Federated Learning for OOD Generalization
global model on a range of diverse test data distributions on Waterbirds, CMNIST, and CFMNIST. The worst-case ('ID-worst' and 'OOD-worst) and average-case ('ID-avg' and 'OOD-avg") ) accuracy are summarized in Table 1. Due to the preservation of richer causal information during train- ing, our method demonstrates superior performance on both ID and OOD data distributions compared to other SOTA methods.
Dupa Card ZAS5 Cacal
Fisher Soor
MT
MM
SV
UP
P
Figure 4. Fisher Scores of causal and inactive causal features on Digits and PACS datasets.
Effect of domain shift. Compared with spurious correla- tion datasets, the non-causal features (i.e. style) are tightly coupled with the causal features of images, making them more difficult to remove. Moreover, the data distributions across different domains exhibit stronger heterogeneity. As shown in Table 2, our method demonstrates significantly better generalization performance across almost all domains, with reduced variance in cross-domain generalization.
distributions, the inactive causal feature compressor gener- ates mask s(x(z)) for each given data distribution, and the mean of these masks is used as the overall causal features types for the current data distribution. Fig. 5a shows the cosine similarity of feature selection masks among training clients, grouped by dataset origins: MNISTM (1-2), SVHN (3-4), SYN (5-6), and USPS (7-10). It can be observed that clients from the same domain exhibit a higher overlap in the selection of causal features. We further analyze the similarity in causal feature selection between test data and the training clients, as shown in Fig. 5b. It is found that domains present in the training data retain feature selection consistency within their respective groups. For the unseen test domain MNIST, its feature selection overlaps more with MNISTM, which shares similar shape, and USPS, which shares the same black and white color scheme. These find- ings further validate the motivation behind FedUni, which is designed to select task-relevant causal features for differ- ent data distributions. Furthermore, domains with higher similarity tend to exhibit more consistent causal feature selection.
5.3. Verification
Analysis of causal and inactive causal features. FedUni decouples the features extracted by the comprehensive fea- ture extractor into two parts, causal features zú and inac- tive causal features z Fig. 3 presents the T-SNE (Van der Maaten & Hinton, 2008) visualization of causal and inac- rive causal features on the Digits dataset, where the first row represents causal features and the second row represents inactive causal features. The visualization demonstrates that causal features exhibit higher discriminative characteristics, as they have stronger causal relationship with the label compared to inactive causal features. To provide a quanti- tative evaluation, we further compute the Fisher Score (Gu et al., 2012) in Fig. 4, which measures feature discriminabil- ity by assessing the ratio of inter-class variance to intra-class variance, formulated as F- = E, ni(u: - / n,o?? where nia Ma. and of represent the number of samples, mean, and variance of the i-th class, respectively. In addition, it is worth noting that the gap in Fisher Scores between causal and inactive causal features increases as the dataset com- plexity rises. This indicates that more comprehensive causal features are extracted from complex datasets, which contain a higher proportion of inactive causal features.
clexhon
(a) Training clients.
(b) Test data and training clients
MNIST
MNISTM
SVIN
Figure 5. Overlapped ratios of causal features.
Analysis of causal features redundancy. We now inves- tigate how different components of our setup affect the learned features. In this section, we compare different methods, which are: (1) Base: FedAvg using ERM update locally; (2) +CI: Incorporating a causal intervention module to the base model; (3) +SL. Incorporating inactive causal feature compressor to the base model; (4) +CI+SL; Incor- porating both of the aforementioned modules and using the features extracted by the feature extractor for prediction; (5) +S)+ Based on (4), using the features after the
Figure 3. T-SNE visualization of causal and inactive causal fea- tures the Digits dataset.
Analysis of overlapped ratio of causal features. To ana- lyze the overlap of selected causal features among different
7

---

Causality inspired Federated Learning for OOD Generalization
Datasets
Waterbirds
CMNIST
CEMNIST
Test Acc (%6)
ID-avg
IDwwrrt
OOD-avg
ODD-worst
ID-avg
ID-oorst
OOD-avg
OOD-worst
ID-avg
ID-worst
OOD-avg
WOD-worst
FedAvg
87.11
86.30
79.09
71.26
94,11
85.52
68.53
27.63
92.55
64.41
87.48
44.52
FedProx
87.31
86.41
78.80
70.42
94.32
85,42
68.23
26.10
94.43
74,43
86.40
36.82
Scaffold
87.41
86.20
78.41
70.95
93.56
84.54
68.47
29.16
94.01
71.41
86.70
10
MOON
85.42
84.69
79.25
73.27
93.77
83.15
73.92
39.62
93.60
69.20
86.64
39.71
FedSR
87.27
86.55
79.86
74.08
95.19
85,70
72.69
37.41
92.87
67.22
86.60
43.12
FedIIR
87.12
86.14
78.67
70.56
93.99
84,91
68.99
28.69
92.92
65.72
87.08
43.21
FedDG-GA
86.50
85.64
79.21
72.16
94.68
88,33
69.02
26.90
92.58
64.32
87.22
44.13
FedSDR
86,93
86.22
80.60
24.59
94.44
83.71
68.72
27.81
94.53
74.10
86.23
40.52
Ours
88,92
88.38
83.29
78.05
95.87
89.53
88,03
56.41
94.43
74.54
88.70
46.25
Table 1. Performance comparison of different methods on Waterbirds, CMNIST and CFMNIST datasets. (Bolded values represent the best performance the experiment setting, while underlined values indicate the second-best performance)
Datasets
Digits
PACS
Test (9)
MNIST
MNISTM
SVHN
USPS
Avg.
Cartoon
Photo
Sketch
Avg.
FedAvg
1..38
34.0020.69
22.62
41. 20.07
80.51
52.25
66.27
61.07a..6
86.76 86.76.0.70
47.69
65.44.0.21
FedProx
81.54_
21. 20.77
38 18.000.07
77.2
500.62... 17.77
70.6910s3
64.43
84 84.861100
19.51 10.41
67.
Scaffold
85.9220.04
22.772 35.09
to
30
79.82
52:933g0:41
28.4
61. 18..55
87.4221.00
49.36
20.09
MOON
82.85
31881077
22.62 10.42
41 41.811077
78.12
52.05_a.56
71.07 jass
68.51
87.9811.4
13.11
67.02
FedSR
88.91.0.31
36 aa.s
25.23 25.23..o.3
28. 28.09.1..
87,30
57.21 ta.1
20.53
61.49
88. 18 18.18
51.61 12.30
68.71
FedlIR
26.27=0079
22.61 20,90
11.03=120
81.43+3.10
84.04
70.71=0a0
29..2
87.3320.2
48.48
07.35g0.44
FedDG-GA GA
85.54.0.57
23.55 34.40aass
40.12g0.7
79.54
52.65
67 14.84
64.961a.s
g7.51
48.11
67.11
RedSDR
60.72
33. 34.744
22.98
44.73=127
89.45
35.92=0.40
70.71
66. 06.76+0.72
90.40
40.58
20.30
89.8 89.85.0.97
g1811
35.42g2.45
66. 66.29 20.05
90.48_0a3
66.8520054
71.08
78.7611.1
89.07
76.04.114
76.24.0.00
Table 2. Performance comparison of different methods on Digits and PACS datasets.
inactive causal feature compressor for prediction.
Redundancy
Basc
We calculate the redundancy of features on the Digits dataset
0.8
+CI
using the formula: R = # Σ; E; Ip(X,, X;)l, where ) is the Pearson correlation between pair of fea- ture dimensions i and j. As shown in Fig. 6, both the compressor and the causal intervention module can improve test accuracy as well as feature redundancy. The observed low test accuracy and feature redundancy in the base model suggest that it has extracted non-causal features (i.e., short- cuts) for prediction, while discarding diverse causal features. The incorporation of the compressor retains client-specific causal features, while the causal intervention module climi- nates spurious correlations. Together, they better preserve the underlying causal features. However, not all causal features contribute meaningfully the test lask-some are inactive. Filtering features through the compressor allows the model to discard redundant or inactive features, thereby improving both redundancy metrics and predictive perfor- mance.
X,)
0.6
+CH+SL
+CH+SL (s)
60
80 80
85
90
Mcan Acc. (%)
OOD Ace. (%)
Figure 6. The relationship between feature redundancy and ID/OOD test accuracy. Gray and Orange arrows denote changes due to adding feature compressor and causal intervention module, respectively.
demonstrating the effectiveness of integrating both compo- nents.
Components
Test Accuracy (%)
Epiv
1. CRea
MNIST
MNISTM
SVHN
SYN
USPS
Avg.
82.39
34.69
80.51
52.26
x
(1.31)
(0.59)
(0.16)
(0.71)
(1.16)
(0.46)
5.4. A Ablation Study
/ (0.47)
(2.13)
49.20
(0.72)
(1.04)
(0 (7)37)
87.89
(0.67)
63.37
x 88.19
36.81
23.90
48.67
87.44
57.00
Ablation Study on Model Components. We conduct abla- tion studies to evaluate the individual and combined contri- butions of FedUni components, as shown in Table 3. Both the generator and compressor individually enhance perfor- mance compared to their absence. Notably, their combi- nation achieves the best performance across all domains,
(0.31)
(0.24)
(0.38)
(0.32)
(0.08)
(0.14)
88.63
$2.40
35.42
66.29
90.48
66.65
v
(0.97)
(1.81)
(0.98)
(0:95)
(0.64)
Table 3. Ablation study on the generator and filter module on Dig- its dataset. (The gray background represents the variance re- peated experiments.)
8

---

Causality inspired Federated Learning for OOD Generalization
References
Ahuja, K., Shanmugam, K., Varshney, K., and Dhurandhar, A. Invariant risk minimization games. In International Conference on Machine Learning, pp. 145-155. PMLR, 2020.
24
84
(a) a
(b) /
(c) X
Ahuja, K., Caballero, E.,, Zhang, D., Gagnon-Audet, J.-C., Bengio, Y., Mitliagkas, I,, and Rish, I. Invariance princi- ple meets information bottleneck for out-of-distribution generalization. Advances in Neural Information Process- ing Systems, 34:3438 3450, 2021.
Figure? 7. The ablation of
CMNIST.
Arjovsky, M., Bottou, L., Gulrajani. I., and Lopez- Paz, D. Invariant risk minimization. arXiv preprint arXiv:
Hyper-parameter Sensitivity. The proposed method in- volves three hyper-parameters o, 3, and X, where a rep- resents the degree of the regularization term in Eq. (7), 3 corresponds to the proportion of the causal intervention loss in the training objective Eq. (17), and denotes the selec- tion ratio of the soft-feature filter in Eq. (12). We conduct experiments on the CMNIST dataset and performed 5 re- peated runs. We vary one hyper-parameter while keeping the other two fixed. The results are shown in Fig. 7. It can be observed that our approach is robust to hyper-parameter changes, with a maximum fluctuation of 2.57%. This indi- cates that FedUni does not rely heavily on careful hyper- parameter tuning and can maintain stable performance under wide range of configurations.
2019.
Carlini, N., Liu, C., Erlingsson, U., Kos, J., and Song, D. The secret sharer: Evaluating and testing unintended memorization in neural networks. In 28th USENIX se- curity symposium (USENIX security 19), pp. 267-284, 2019.
Chang, C.-H., Adam, G. A., and Goldenberg, A. Towards robust classification model by counterfactual and invari- ant data generation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 15212-15221, 2021.
Chen, J., Gao, Z., Wu, X., and Luo, Meta-causa learning for single domain generalization. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 7683-7692, 2023.
6. Conclusion
In this work, we address the OOD generalization challenge in FL by proposing FedUni, a novel framework that re- tains all possible causal features across clients rather than only an invariant subset. Our approach leverages a com- prehensive feature extractor and feature compressor to extract the union of causal features while suppressing in- active causal features. Additionally, we introduce a causal intervention module that employs counterfactual generation to further enhance OOD generalization. Through extensive theoretical analysis and empirical evaluation, we demon- strate that FedUni significantly improves OOD generaliza- tion. These findings highlight the importance of preserving diverse causal features in FL.
Chen, T., Korbllihh, S., Norouzi, M., and Hinton, G. A simple framework for contrastive learning of visual rep- resentations. In International conference on machine learning, pp. 1597-1607. PMLR, 2020.
Chen, Z., Zheng, Y., Lai, Z.-R., Guan, Q. and Lin, L. Diag- nosing and rectifying fake ood invariance: A restructured causal approach. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 38, pp. 11471-11479, 2024.
Deng, J., Dong, W., Socher, R., Li, L..J., Li, K., and Fei-Fei, L. Imagenet: A large-scale hierarchical image database. In 2009 IEEE conference computer vision and pattern recognition, pp. 248 448-255, leee, 2009.
Acknowledgment
This work was supported by the National Natural Science Foundation of China under Grants 62372028 and 62372027.
Denker, J., Gardner, W., Graf, H., Henderson, D., Howard, R., Hubbard, W., Jackel, I.. D., Baird, H., and Guyon, I. Neural network recognizer for hand-written zip code dig- its. Advances in neural information processing systems, 1, 1988.
Impact Statement
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none of which we feel must be specifically highlighted here.
Duan, J.-H., Li, W., and Lu, S. Feddna: Federated learning with decoupled normalization layer aggregation for non- iid data. In Machine Learning and Knowledge Discovery

---

Causality inspired Federated Learning for OOD Generalization
in Dalabases. Research Track: European Conference, ECMI. PKDD 2021, Bilbao, Spain, September 13-17, 2021, Proceedings, Part I 21, pp. 722-737. Springer, 2021.
Krueger, D., Caballero, E., Jacobsen, J.-H., Zhang, A., Bi- nas, J., Zhang. D., Le Priol, R., and Courville, A. Out-of- distribution generalization via risk extrapolation (rex). In International conference on machine learning, pp. 5815- 5826. PMLR, 2021.
Ganin, Y. and Lempitsky, V. Unsupervised domain adapta- tion by backpropagation. In International conference on machine learning, pp. 1180-1189, PMLR, 2015.
Langley, P. Crafting papers on machine learning. In Langley, P. (ed.), Proceedings of the 17th International Conference on Machine Learning (ICMI. 2000), pp. 1207-1216, Stan- ford, CA, 2000. Morgan Kaufmann.
Gatys, L. A., Eeker, A. S., and Bethge, M. Image style trans- fer using convolutional neural networks. In Proceedings of the IEEE conference on computer vision and paltern recognition, pp. 2423-2423, 2016.
LeCun, Y., Bottou, I.., Bengio, Y., and Haffner, P. Gradient- based learning applied to document recognition. Proceed- ings of the IEEE, 86(1 ( 1):2278 2324, 1998.
Geirhos, R., Jacobsen, J.-H., Michaelis, C., Zemel, R., Bren- del, W., Bethge, M., and Wichmann, A. Shortcut learn- ing in deep neural networks. Nature Machine Intelligence, 2(1 ((11))665-440) 673, 2020.
Li, B., Shen, Y., Wang, Y., Zhu, w., Li, D., Keutzer, K., and Zhao, H. Invariant information bottleneck for domain generalization. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 36, pp. 7399-7407, 2022.
Gu, Q., Li, Z., and Han, J. Generalized fisher score for feature selection. arXiv preprint arXiv: arXiv:/202.377 120.37725, 2012.
Li, D., Yang, Y., Song, Y.-Z., and Hospedales, T. M. Deeper, broader and artier domain generalization. In Proceedings of the IEEE international conference on computer vision, pp. 5542-5550, 2017.
Guo, K., Ding, R., Qiu, T., Zhu, X., Wu, Z., Wang, L., and Fang, H. Single domain generalization via unsuper- vised diversity probe. In Proceedings of the 31st ACM International Conference on Multimedia, pp. 2101-2111, 2023a.
Li, L.,, Gao, K., Cao, J., Huang, Z., Weng, Y., Mi, X., Yu, Z., Li, X., and Xia, B. Progressive domain expansion network for single domain generalization. In Proceedings of the IEEE/CVE Conference on Computer Vision and Pattern Recognition, pp. 224-233, 2021a.
Guo, Y., Guo, K. Cao. X., Wu, T., and Chang, Y. Out- of-distribution generalization of federated learning via implicit invariant relationships. International Confer- ence on Machine Learning, pp. 11905-11933, PMLR, 2023b.
Li, Q. He, B., and Song, D. Modtl-contrastive federated learning. In Proceedings of the IEEE/CVF conference computer vision and pattern recognition, pp. 10713- 10722, 2021b.
Hao, W., EI-Khamy, M., Lee, J., Zhang, J., Liang, K. J., Chen, C., and Duke, L. C. Towards fair federated learning with zero-shot data augmentation. In Proceedings of the IEEF/CVF conference on computer vision and pattern recognition, pp. 3310-3319 2021.
Li, T., Sahu, A. K., Zaheer, M., Sanjabi, M., Talwalkar, A., and Smith, V. Federated optimization in heterogeneous networks. Proceedings of Machine learning and systems, 2:429 450, 2020.
He, K., Zhang, X., Ren, S., and Sun, J. Deep residual learn- ing for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770-778, 2016.
Liao, X., Liu, W., Zhou, P., Yu, F., Xu, J., Wang. J., Wang, W., Chen, C., and Zheng, X. Foogd: Federated collab- oration for both autoof-distribution generalization and detection. Advances in Neural Information Processing Systems, 37:132908-133 132945, 2024.
Hsu, T.-M. H., Qi, H., and Brown, M. Measuring the effects of non-identical data distribution for federated visual clas- sification. arXiv preprint arXiv: 1et::999900000 0039006355 2019.
Liu, Q. Chen, C., Qin, J., Dou, Q., and Heng, P.-A. Feddg: Federated domain generalization on medical image seg- mentation via episodic learning in continuous frequency space. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 1013-1023, 2021.
Huang, X. and Belongie, S. Arbitrary style transfer in real- time with adaptive instance normalization. In Proceed- ings of the IEEE international conference on computer vision, pp. 1501-1510, 2017.
Lv, F., Liang. J., Li, S., Zang, B., Liu, C. H., Wang, 7., and Liu, D. Causality inspired representation learning for domain generalization. In Proceedings of the IEFF/CVF conference on computer vision and pattern recognition, pp. 8046-8056, 2022.
Karimireddy, S. P., Kale, S., Mohri, M., Reddi, S., Stich, S., and Suresh, A. T. Scaffold: Stochastic controlled averag- ing for federated learning. In International conference on machine learning, pp. 5132-5143. PMLR, 2020.
10

---

Causality inspired Federated Learning for OOD Generalization
Massart, P. About the constants in talagrand's concentra- tion inequalities for empirical processes. The Annals of Probability, 28(2)863--84, 2000.
Qiao, F., Zhao, L., and Peng. X. Learning to learn single domain generalization. In Proceedings of the IEEF/CVF conference computer vision and pattern recognition. 2020.
pp.
McMahan, B., Moore, E., Ramage, D., Hampson, S., and y Arcas, B. A. coammunicatioo--ffciiee learning of deep networks from decentralized data. In Artificial intelli- gence and statistics, pp. 1273-1282. PMLR, 2017.
Reichenbach, H. The theory of probability. Univ of Califor- nia Press, 1971.
Sagawa, S., Koh, P. W., Hashimoto, T. B., and Liang, P. Distributionally robust neural networks for group shifts: On the importance of regularization for worst-case gener- alization. arXiv preprint
Mendieta, M., Yang, T., Wang, P., Lee, M., Ding, 7.,, and Chen, C. Local learning matters: Rethinking data het- erogeneity in federated learning. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 8397-84000 2022.
2019.
Sauer, N. Combinatorial Theory, Series A, 13(1):145-147, 1972. ISSN 0097-3165.
On the density of families of sets. Journal of
Meng, L.., Qi, Z.., Wu, L.., Du, X., Li, Z.., Cui, L.,, and Meng, X. Improving global generalization and local personalization for federated learning. IEEE Transactions on Neural Networks and Learning Systems, 36, 2024.
Selvaraju, R. R., Cogswell, M., Das, A., Vedantam, R., Parikh, D., and Batra, D. Grad-cam: Visual explana- tions from deep networks via gradient- localization. In Proceedings of the IEEE international conference on computer vision, pp. 118-626,770 2017.
Mohri, M. Foundations of machine learning. 2018
Netzer, Y. Wang, T., Coates, A., Bissacco, A., Wu, B., Ng, A. Y., et al. Reading digits in natural images with unsupervised feature learning. In NIPS workshop on deep learning and unsupervised learning, volume 2011, pp. Granada, 2011.
Shalev-Shwartz S. and Ben-David, S. Understanding ma- chine learning: From theory to algorithms. Cambridge university press, 2014.
Shannon, C. E. A mathematical theory of communication.
Nguyen, A. T., Torr, P., and Lim. S. N. Fedsr: A simple and effective domain generalization method for federated learning. Advances in Neural Information Processing Systems, 35:38831-3884. 2022.
The Bell system technical journal, 27(3):379 423, 1948.
Sharma, P., Panda, R., Joshi, G., and Varshney, P. Federated minimax optimization: Improved convergence analyses and algorithms. In International Conference on Machine Learning, pp. 19683-19730, PMLR, 2022.
Noohdani, F. H., Hosseini, P., Parast, A. Y., Araghi, H. Y., A com-
and Baghshah, M. S. aecompose-and-comppee positional approach to mitigating spurious correlation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Putlern Recognition, pp. 27662-27671, 2024.
Tang, X., Guo, S., and Guo, J. Personalized federated learn- ing with contextualized generalization. arXiv preprint arXiv:2106.14 /3044, 2021.
Press, W. H. Numerical recipes 3rd edition: The art of scientific computing. Cambridge university press, 2007.
Tang, X., Guo, S., and Zhang, Exploiting personalized invariance for better nut-of-distribution generalization in federated learning. arXiv preprint arXiv:2211.112433 2022.
Qi. Z., Meng, L., Chen, Z., Hu, H., Lin, H., and Meng, X. Cross-silo prototypical calibration for federated learning with non-iid data. In Proceedings of the 31st ACM In- ternational Conference on Multimedia, pp. 4099-3107, 2023.
Tang, X., Guo, S., Zhang, J., and Guo, J. Learning personal- ized causally invariant representations for heterogeneous federated clients. The Twelfth International Conference on Learning Representations, 2023.
Qi, 7., He, W., Meng, X., and Meng, I.. Attentive modeling and distillation for -ottof-distribution generalization of federated learning. In 2024 IEEE International Confer- ence on Multimedia and Expo (ICME), pp. 1-6. IEFF, 2024.
Tang, X., Guo, S., Guo, J., Zhang, J., and Yu, Causally motivated personalized federated invariant learning with shortcut-averse intommation-thooree regularization. In Forty-first International Conference on Machine Learn- ing. 2024.
Qi. Z., Meng, L., and et al. Cross-silo feature space align- ment for federated learning on clients with imbalanced data. In The 39th Annual AAAI Conference on Artificial Intelligence (AAAI-25), pp. 19986-19994 2025.
Van der Maaten, L. and Hinton, G. Visualizing data using t-sne. Journal of machine learning research, 9(11), 2008

---

Causality inspired Federated Learning for OOD Generalization
Vapnik, V. The nature of statistical learning theory. Springer science & business media, 2013.
Proceedings of the 30th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, pp. 3794-3805, 2024.
Volpi, R., Namkoong, H.,, Sener, O., Duchi, J. C., Murino, V., and Savarese, S. Generalizing to unseen domains via adversarial data augmentation. Advances in neural information processing systems, 31, 2018.
Yoon, T., Shin, S., Hwang. S. J., and Yang. E. Fedmix: Ap- proximation of mixup under mean augmented federated learning. arXiv preprint arXiv: Stti::100700233, 2021.
Wang. J., Liu, Q. Liang. H., Joshi, G., and Poor, H. V. Tack- ling the objective inconsistency problem in heterogeneous federated optimization. Advances neural information processing systems, 1-7623, 2020.
Zhang, J., Li, Z., Li, B., Xu, J., Wu, S., Ding, S., and
Wu, C. Federated learning with label distribution skew via logits calibration. In International Conference on Machine Learning, pp. 2631263112222 1-26329. PMLR, 2022.
Wang. Z., Luo, Y., Qiu, R., Huang, Z., and Baktashmotlagh, M. Learning to diversify for single domain generalization. In Proceedings of the IEEE/CVF International Confer- ence on Computer Vision, pp. 834-843, 2021.
Zhang, J., Liu, X., Zhang, Y., Zhu, G., Niu, J., and Tang, S. Enabling collaborative test-time adaptation in dynamic environment via federated learning. In Proceedings of the 30th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, pp. 1191-4 2024.
Welinder, P., Branson, S., Mita, T., Wah, C., Schroff, F., Belongie, S., and Perona, P. Caltech-ucsd birds 200. 2010.
Zhang, L.,, Lei, X., Shi, Y., Huang, H., and Chen, C. Feder- ated learning with domain generalization. arXiv preprint arXiv:2111.1043: 2021.
Wu, X., Niu, J., Liu, X., Ren, T., Huang, Z., and Li, Z. pfedgf: Enabling personalized federated learning via gra- dient fusion. In 2022 IEEE International Parallel and Dis- tributed Processing Symposium (IPDPS), pp. 639-649, 2022.
Zhang, R., Xu, Q. Yao, J., Zhang, Y., Tian, Q. and Wang, Y. Federated domain generalization with generalization adjustment. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 3954- 3963, 2023.
Wu, X., Liu, X., Niu, J., Zhu, G., and Tang, S. Bold but cau- tious: Unlocking the potential of personalized federated learning through cautiously aggressive collaboration. In Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV), pp. 19375-19384, October 2023.
Zhou, B., Lapedriza, A., Xiao, J., Torralba, A., and Oliva, A. Learning deep features for scene recognition using places database. Advances in neural information processing systems, 27, 2014.
Zhou, K., Liu, Z., Qiao, Y., Xiang, T., and Loy, C. C. Do- main generalization: A survey. IEEE Transactions on Pattern Analysis and Machine Intelligence, 45(4):4396 4415, 2022.
Wu, X., Liu, X., Niu, J., Wang, H., Tang. S., Zhu, G., and Su, H. Decoupling general and personalized knowledge in federated learning via additive and low-rank decom- position. In Proceedings of the 32nd ACM International Conference Multimedia, MM '24, pp. 7172-7181, New York, NY, USA, 2024.
Zhu, G., Liu, X., Tang, S., and Niu, J. Aligning before ag- gregating: Enabling cross-domain federated learning via consistent feature extraction. In 2022 IEEE 42nd Inter- national Conference on Distributed Computing Systems (ICDCS), pp. 809-80 109-881,,22 2022.
Wu, X., Niu, J., Liu, X., Zhu, G., Zhang, J., and Tang. S. Enhancing visual representation with textual semantics: Textual semantics-powered prototypes for heterogeneous federated learning, 2025.
Zhu, G., Liu, X., Niu, J., Tang, S., Wu, X., and Zhang, J. Du- alfed: enjoying both generalization and personalization in federated learning via hierachical representations, In Proceedings of the 32nd ACM International Conference on Multimedia, pp. 1 11600-1 1069, 2024a.
Xiao, H., Rasul, K. and Vollgraf, R. Pashion-mnist: a novel image dataset for benchmarking machine learning 2017.
algorithms. arXiv preprint
Yang, M., Fang, 7.,, Zhang, Y., Du, Y., Liu, F., Ton, J.-F., Wang, J., and Wang, Invariant learning via probability of sufficient and necessary causes. Advances in Neural Information Processing Systems, 36:79832-79883 2023.
Zhu, G., Liu, X., Niu, J., Wei, Y.. Tang, S., and Zhang, J. Learning by imitating the classics: Mitigating class imbalance in federated learning via simulated centralized learning. Expert Systems with Applications, 255:1 2024b.
Yang, S., Zhang, Z., and Gu, L. Practical single domain generalization via training-time and test-time learning. In
12

---

Causality inspired Federated Learning for OOD Generalization
A. Related Work
Causal Feature Learning. An important research direction to tackle OOD generalization is to exploit the causal features behind the observed data (Zhou et al., 2022). A series of works formulate invariant causal mechanisms in terms of features. Invariant risk minimization (IRM) (Arjovsky et al., 2019) provides a solution for learning causal features, but it dramatically deteriorates when applied to an over- ver-parameterized model. From this perspective, some works (Ahuja et al., 2021; Krueger et al., 2021) further extend the IRM by incorporating information theory and nonlinear prediction functions. IIB (Li et al., 2022) minimizes invariant risks for nonlinear classifiers through a tractable mutual information- based loss function. CaSN (Yang et al., 2023) introduces the classical concept of the Probability of Necessity and Sufficiency (PNS) to simultaneously capture information on both sufficient and necessary causes and proposes PNS risk to optimize OOD generalization capability, Another series of papers achieves this goal by generating counterfactual images and constraining the consistency between the features before and after generation (Chang al., 2021). CIRL (Lv al., 2022) introduce Fourier transformation as a causal intervention module to enforce features to be separable from non-causal factors, jointly independent, and causally sufficient for classification. DaC (Noohdani et al., 2024) leverages the model's Class Activation Map (CAM) to identify causal components in images, then intervenes by composing different elements to generate counterfactual samples and retrains the model for improved generalization. However, these works focus on centralized scenarios where all training data is accessible.
Single Domain Generalization is more challenging yet practical domain generalization task, where the training data comes from only a single distribution. A typical line of methods synthesize diverse samples from hypothetical domains and ensure consistency between features extracted from the source and synthetic samples. Among the earliest attempts, ADA (Volpi et al., 2018) and M-ADA (Qiao et al., 2020) introduce an adversarial data augmentation strategy to tackle the worst-case domain generalization problem. Recent studies inspired by style transfer methods (Gatys et al., 2016) focus on generating samples with novel styles. To name a few, PDEN (Li et al., 2021 a) proposes a progressive domain expansion framework. L2D (Wang et al., 2021) presents a style-complement module designed to synthesize images through the optimization of the mutual information's 's upper bound. UDP (Guo al., 2023a) directly augments images using in the image space, decoupling the training objectives between the generator and the discriminator, thereby avoiding mutual interference. Despite their remarkable effectiveness, these methods often introduce additional computational burdens. In this paper, we take into account the limited computational capabilities of clients in FL and devises more lightweight algorithm.
General Federated Learning (GFL). The classic FedAvg (McMahan et al., 2017) algorithm performs well on IID datasets, but its performance significantly drops under the presence of data heterogeneity. One line of methods has been proposed to mitigate the impact of data heterogeneity during the client-side training process (Zhang et al., 2022; Mendieta et al., 2022; Zhu et al., 2024b; Wu et al., 2024; Zhu et al., 2022; Wu et al., 2023; 2025; Zhu et al., 2024b). For example, FedProx (Li et al., 2020) introduces an L2-norm term between the global model parameters and the local parameters. SCAFFOLD (Karimireddy et al., 2020) adds correction terms for local gradients to ensure that the local update moves towards the true optimum. MOON (Li et al., 2021b) propose a model-level contrastive learning approach inspired by conventional data-level contrastive frameworks. Another line of studies focus on enhancing the server-side aggregation process. FedNova (Wang et al., 2020) enhances model aggregation by addressing heterogencous local steps caused by varying data quantities or training epochs. FedDNA (Duan et al., 2021) selectively apply importance-weigheee adversarial learning to batch normalization layers while maintaining standard FedAvg for gradient- parameters. In the context of server-side optimization (Sharma et al., 2022). FedAvgM (Hsu et al., 2019) implements momentum--cccelrratee SGD. While these methods comprehensively address distribution shifts among participating clients, it critically overlooks the generalization requirements of non pon-paaticiiaaiigg clients. Unlike these works, our method attempts to learn latent causal relationships that work equally well for both the participating and non- non-participating clients.
Federated Learning for OOD Generalization. The duttr--distribution (OOD) generalization problem in FL has recently attracted significant research interest (Zhang ct al., 2024; Zhu et al., 2024a). Data augmentation- based research involves applying various transformations to the existing training data (Yoon et al., 2021; Hao et al., 2021). Several approaches enhance OOD generalization by optimizing the weighting strategies for global model aggregation. FedDG-GA (Zhang et al., 2023) dynamically calibrates aggregation weights through generalization adjustment to achieve tighter theoretical bounds. More similar to our approach, several studies concentrate invariant (causal) feature extraction through adversarial domain alignment or causal representation learning. FedADG (Zhang et al., 2021) introduces a federated adversarial domain adaptation framework to align data distributions across all clients. FedSR (Nguyen et al., 2022) imposes dual regularization mechanisms on latent representations and conditional mutual information, suppressing spurious correlations while preserving invariant features. FL GAME (Zhang et al., 2021) establishes a multi-agent game-theoretic architecture
13

---

Causality inspired Federated Learning for OOD Generalization
enabling decentralized parallel computation through Nash equilibrium-driven optimization, which systematically extracts client-invariant causal features via information bottleneck constraints. FedlIR (Guo al., 2023b) implicitly learns causal relationships for OOD generalization by quantifying disagreement across clients and reducing through inter-client gradient alignment in parameter space. However, these methods focus on mining the invariant features across clients, which inevitably leads to the loss of some valuable information.
Causal Feature with FL OOD Generalization. Recently, inspired by causal feature learning in centralized learning scenarios, a growing number of methods aim to extract causal features to improve the OOD generalization ability of FL model (Tang et al., 2021; 2022). FedSDR(Tang et al., 2023) constructs structural causal model (SCM) for heterogeneous clients to collaboratively identify non-causal features and learn the optimal personalized causally invariant predictor. FedPIN (Tang et al., 2024) distinguishes personalized causal features from non-causal ones using global invariant features as anchors, formulating an information-theooe
constraint to enable shorteut-averse personalized invariant learning, thereby improving OOD generalization in PFL. However, these methods are especially designed for PFL, where the causal features of each client's s data are observable during training. In contrast, in the GFL setting we focus on, the causal features present in the unknown test data are unseen during training.
B. Theoretical Guarantees
B.1. The detailed explanation of Assumption 3.1 and the proof of Lemma 3.2
Assumption 3.1 (Existing SCM illustrated in Fig. 2a)
Zs = (E,Y). 2ze == [un Zz), = Jeen (Zs, Zz, zZ), Y == wZZ,
Explanation. Here, Zg represents spurious non-causal features. It is function of the environment E. and target Y. zo is the local causal variables, which is determined by the client indicator U and the global causal features Zg through the invariable causal function fine. The data X is generated by the generation function fyen. which takes into account the spurious variables Zs. the global causal variables Zg. and the local causal variables zy. Finally, the target Y is determined by the classifier w based on the global causal variables z& and the local causal variables zy-
Lemma 3.2. If the data generating process of each client follows the SCM assumption illustrated in Fig. 2a, the global invariance Zg incomplete from two perspectives:
Y L U Zg: Under zg. the prediction of Y is still influenced by the client indicator U.
Y 1 [U, F] Z8, ze: Given 7g. (i.e. global/shared invariant features) and ze (i.e. local/personalized invariant features), Y is conditionally independent of the environment variable E and user indicator U.
Proof. To prove the non-independence Y y U Zg. we were the d-separation criterion. Specifically, there is an unblocked causal path from U to Y (U + zg + Y) that does not pass through zg. This implies that even when conditioned on zg. the prediction of Y is still influenced by the client indicator U.
To prove the conditional independence Y [U, E] 2g. zE. we observe that the variables Zg and Zy d-separate Y from both the environment indicator E and the user/client indicator U.
B.2. The detailed explanation of Assumption 3.3 and the proof of Lemma 3.4
Assumption 3.3 (Our Restructured SCM illustrated in Fig. 2b)
Zs = fayu fagu(E, Y), zg. = f.no 22 = feno 28),
X = fgen (z. 2g, zF), Y == w(Z)).
Explanation. Here, Zg represents spurious non-causal variable and is a function of the environment E and the target Y. Different from Assumption 3.1, zg is now the only cause of the target Y, which means encompasses the local causal
14

---

Causality inspired Federated Learning for OOD Generalization
knowledge carried by zg in Assumption 3.1. We have new variable z (where F stands for False). Z is only affected by the client indicator U and influences Y through affecting Z 7 local causal variable, similar z in Assumption 3.1, representing some causal information that holds locally at the client but my not be transferred well to all environments.
Lemma 3.4. If the data generating process of each client follows the SCM assumption illustrated in Fig. 2b, then the following three statements hold:
Vs 1 1 1 Y: Given Y, both the zz and Z are conditionally independent of the non-causal variable Zg. Y 1 [U, " D3] Zz: Components in [U, E]] are conditionally independent of the target Y given Zÿ and zY. Considering that z7 _ L this expression a simplified notation.
Y 1 Ls, zp: Y is not conditionally independent of Z2 given Zs and z
Proof. We first prove that Zs is conditionally independent of both Zg and Zy given Y. According to the d-separation criterion, Zs and Te zy share no common parent variables, and the only path from Zg. z to Zs is blocked by Y, leading to the conditionally independence between Zs and Zo. ZT.
To prove the second statement, we observe the conditional independence Y E] zg. Since zy is the only cause variable of the target Y, blocks the conditional correlation between U, and Therefore, the conditional independence holds.
Finally, we turn to the third statement. We argue that Y is not conditionally independent of Za given Zg and Zp, as there exists a direct causal path from Zg to Y that cannot be blocked by Zs and Z This implies that z and Zg play different roles in the prediction process and also justifies this non-independence.
B.3. The proof of Proposition 4.1
Proposition 4.1. Causal variables Zg should satisfy the following two properties:
Invariance under causal intervention: Based on the first claim in Lemma 3.4, we have Zg 1 E Y. Thus, performing intervention upon E does not make changes to Zg.
Sufficiency of semantic information: The causal variable zg should be causally sufficient to the label prediction, i.e. maximizing T YZ)
Proof. To demonstrate invariance under causal intervention property, we first examine the causal relationships among the variables involved. There no direct or indirect casual path from E. to Zg that is not blocked by Y. Specifically, the generation of Zg is mainly through z which has no direct causal link from E. Therefore, Zg and E are d-separated conditioned on Y. This means that intervening on E does not affect Zg. thus satisfying the invariance property under causal intervention.
To demonstrate the sufficiency of semantic information property, we refer to the SCM model in Assumption 3.3, where Y = =Z())) This implies that Y is determined entirely by Zg. From the perspective of information, the mutual information is given by Y) = HY) - HY Zz). Thus, maximizing / [Zg; Y) is to minimize (Y zz). According to the causal relationships among variables, Z% is the only cause of the target and there no other variable that can provide additional information for predicting Y other than Zg. which should minimize (Y 2g)), thus satisfying the sufficiency of semantic information property.
B.4. The proof of Lemma 4.2 and Lemma 4.3
Lemma 4.2. H (z)(xz can be approximated by the conditional probability p(='(x)) as Li-Laplacians, follows:
H(s'(x) (log ((<(l)) Eg[l (x)
Similarly, Hxx') can be approximated by the marginal distribution of x':
H(x') = E.. P(x')) [=))))))
15

---

Causality inspired Federated Learning for OOD Generalization
Proof. The L1-Laplace distribution is often used to approximate conditional entropy, which effectively captures deviations under noise-perturhed conditions (Carlini et al., 2019; Guo et al., 2023a). In the first approximation, the conditional probability is modeled Li-Laplace distribution centered at (xx) with identity covariance.
P(a1 = (()), & exp(-IIa (x))))
(18)
By using this approximation above, the conclusion in Lemma 4.2 holds. Similarly, in the subsequent approximation, we model the distribution using an L1 Lapplee distribution centered at zero,
- = -(x';,11 exp(-Ilv(=))l))
(19)
Lemma 4.3. Maximizing the diversity loss is approximately equivalent to maximizing the lower bounds of Lprv, which satisfy the following inequality:
- 1 x
max LDIv
max E. 1
W
be observed that generating the counterfactual image xf essentially involves maximizing ||||=||| ||>1|||1p-
Proof. Since ||>||| and ||'|| are positive, the left-hand side of the above inequality can be rewritten follows:
| 2
|||="
|2]|1]
(20)
[x'
Notice that for most practical distributions, especially when (bx) (or x²) is close to || =/I1 tends to be greater than or equal to |||= - || cause changes in the L1 -nom are generally larger compared to element wise differences. Therefore, it's reasonable to conclude:
(b(())
x/11 [b(a)]1
2 E. (1
||0/|1 |
)
Eg
(21)
Combining the analysis above, we can derive:
||0||a
max LDNV 2 max Ex 1
(22)
|x'||
B.5. The proof Lemma 4.4
Lemma 4.4. The CMI admits the following decompositions:
-
(x)=" 1 (X))
Proof. Under our Restructured SCM in Assumption 3.3, our goal is to distinguish Z from Z effectively adapt to test data distributions. It can be observed that we can distinguish them based on their differing behaviors when conditioned on each other in predicting Y. In line with (Chen et al., 2024), we adopt CMI as principled metric to analyze the dependencies among different feature variables. However, (Chen et al., 2024) focuses on extracting causal features and computing the mutual information between causal features and label Y, whereas our focus is on eliminating inactive causal features I from global causal features zg.
16

---

Causality inspired Federated Learning for OOD Generalization
The CMI between X and Y given Z can be decomposed into entropy terms as follows:
[[; H( + H( (|Z))
(23)
Note that when (XX is well trained, the models encourage o² (X( z U zg. therefore, Zy (X()/ZZ
N Z = HY;
2z1Zz).
(24)
By the definition of conditional entropy (Shannon, 1948), obtain:
,
-H(Z),,
(25)
where H(Zp, 2g) is the joint entropy of z and z Hz)) is the entropy of z and H(Y, Zp z is the joint entropy
of 2 2 and Y. From the above decomposition in Eq. (23), the following equation can be derived:
Z(Z) = [H H (Z)) - [H 42) - H 2
22H z8) ((Z)) H(Y, HYY, (g)
= =2H z5) - H(Z)) - (7 28) + HHZ),
= (1Z4).2)) (2 - H
-Z(Z)/Z)) HY 24 .
(26)
Then, substituting the above expression into Eq. (24) and simplifying yields:
+ -1 - 10 (X)) + ((Y1"((X)/2)
(27)
In the same way, we can prove the second equation in the lemma:
T[: (x)z/
(28)
C. Theoretical analysis of generalization boundary
To analyze the generalization boundary of FedUni, we follow (Mohri, 2018) and use Sopnik-Chervonenkis (VC) dimension theory (Reichenbach, 1971). VC dimension is an important concept in statistical learning theory used to measure the complexity and expressiveness of a hypothesis class, which is closely related to model's generalization ability. Generally, models with higher VC dimension indicate that the hypothesis class can represent more complex functions and may perform well on training data but might overfit when applied to testing data.
As mentioned in Section 3, each client in our model consists of a feature extractor and a classifier, and can be formulated as (-) = (.))).)). where 0. = (wer dz). Considering only the stages following feature extraction, the global optimization objective is given Eq. (17).
noteworthy that all components of the optimization objective are upper-bounded. Without loss of generality, we assume the loss function is bounded above by 1. Let H. be a family of functions that take values of the one-hot vectors of length i for k-classification problem, where the VC-dimension is d. Then, for any 8 > with probability at least 1 - d, for any hr, E H. the following holds:
Eo d B[R(h,,)) < (Rtss 2/ 2d
log 1 2m
(29)
where R(h,,) and (h,,)) are the generalization error and empirical error on sample respectively. Since the empirical error in our method is upper- upper-bounded, the generalization error has upper bound as well. To establish this result, we
17

---

Causality inspired Federated Learning for OOD Generalization
build upon the analysis in (Mohri, 2018), generalizing it from the centralized binary classification setting multi-client, multi-class scenario. Firstly, let FF he family of functions mapping from 2 to [0.1], in which 2 = X ) is the input space.
Then, for any d > O, with probability at least 1 - å over i..dd sample S of size m, the following holds for each f E F:
1 B[f())) -
m
log 2m
m E f(z+) 29 (F)
(30)
where 9t is the Rademacher complexity (Shalev-Shwartzz & Ben-David, 2014). Our focus shifts to proving that the right side of Eq. (30) is bounded, with the key challenge lying in demonstrating that Rm, (F) is bounded.
Let H. be family of functions that take values in the one-hot vectors of length & for a k-classification problem, and F be the family of loss functions associated with H. for the zero-one loss, i.e., F = {(x, y) + 1x()) h E H). For any sample
= { yr....., (Im, 9m)) of elements in (X²,") let Sx denote its projection onto X. Then, the following relation holds between the empirical Rademacher complexities of and H:
Rs(F) = -ssx (H).
(31)
Using Eq. (30) and Eq. (31), we have
log g
R(h) S R2(h) + 2m (Ht +
(32)
2m
where R(h) is the generalization error and (h) is the empirical error. According to Massart's s Lemma (Massart, 2000) and Sauer's Lemma (Sauer, 1972), which together establish an explicit relationship among the Rademacher complexity, the growth function, and the VC-dimension, we can derive the following result based on Eq. (32): for any 5 >> 0, with probability at least 1 - 8, for every function h E H. the following upper bound on the generalization error R(h) holds:
2 log (m)
log 2m
2d em d
log 2m
Rhh) < Rh(h) + +
< (h) + 22
(33)
where (m) is the growth function of H and d is the VC-dimension of H. Since Eq. (33) holds for all clients in our method, we now focus on the expected generalization error for each client. In the multi-client scenario, the result is shown in the following equation, therehy proving the theoretical generalization bound of FedUni:
em 2d log
log 3
El,( < E[Rs (he, 2/
d
(34)
2m
D. More Details about Experiments
D.1. Description of the dataset
Waterbirds is synthetic dataset that combines bird photographs from the Caltech-UCSI 201s22000-0011 (CUB) dataset (Welinder et al., 2010) with backgrounds from the Places dataset (Zhou et al., 2014). This dataset provides binary classification task for 'landbird' and 'waterbird', and the background is spuriously correlated with the class. We use p, to quantify the degree of spurious correlation, where landbirds have a probability of ps with a land background and (1 - -Pps) with a water background.
CoLOred-MNIST (CMNIST) is variant of the original MNIST (LeCun al., 1998) dataset via rearranging the images of digits 0-4 into one class (class 0) and the images of digits 5-9 into another class (class 1). The color is spuriously correlated with the class, where each digit with label 0 is colored green/red with probability ps/1 Ps, and this probability is reversed for label 1.
Golored-Fashion MNIST (CFMNIST) is constructed from Fashion-MNIS" (Xiao et al., 2017) using the same strategy as CMNIST.
18

---

Causality inspired Federated Learning for OOD Generalization
Digits consists of 5 different datasets: MNIST (LeCun al., 1998), SVHN (Netzer al., 2011), MNISTM (Ganin & Lempitsky, 2015), SYN (Ganin & Lempitsky, 2015) and USPS (Denker al., 1988). Each dataset is considered unique domain which may be different from the rest of the domains in font style, background, and color.
PACS consists of a total of 9,991 samples, each with dimension of (3, 224, 224). It includes 7 classes and covers four domains: art, cartoon, photo, sketch.
D.2. Implementation Details
For Waterbirds, CMNIST, and CFMNIST datasets, the degree of spurious correlation ps for each source client is randomly sampled from [0.85: : 1.0], while for the test data ps varies from 0. to 1.0. Considering the heterogeneity across clients, the training data on each client are sampled from a subset of the original classes. Specifically, we conduct experiments on 10 clients. We distribute 10 waterbird species (4 separated and 6 overlapped) and 19 landbird species (15 separated and overlapped) to each client on the Waterbirds dataset. And the CMNIISTCFMNNITT dataset, we re-label the data with labels {0, 1, 2, 3, 4} class 0, and the data with labels {5, 6, 7, 8, 9} as class 1. Furthermore, we randomly select two different digit subclasses from each of class 0 and class 1. Since the test distribution in the real world is unknown, the model may be deployed to source clients, which share the same distribution as the training data, or to unseen clients, whose data distribution shifts from the training process. We aim to enhance OOD generalization performance while maintaining the discriminative ability on the ID data distribution as much as possible. We refer to the results for p, E {1.0, 0.9, 0.8} as the in-distribution (ID) case, since this degree close to the training data distribution, while the remaining values correspond to the OOD case.
For Digits and PACS datasets, we follow the 'leave-oneoout rule, where we choose domain target domain, train the model on all remaining domains, and evaluate on the target domain. We conduct experiments on 10 clients, the data on each client are sampled from the same domain, and the domains on different clients can be repeated. For example, in the Digits dataset, there are 5 domains in total, with training clients covering 4 of these domains. The number of clients in each training domain is 2, 2, 2, and 4, respectively, and each domain contains 1000 training samples. In the PACS dataset, there are 4 domains in total, with training clients covering 3 of these domains. The number of clients in each training domain 3, 3, and 4, respectively, and each domain contains 750 training samples.
E. Additional Method Details
E.1. Structure of the counterfactual generator
Our generator (z;; 0p) aims to synthesize novel and diverse samples to extend the distribution of the source domain. We chose a model architecture that is widely used in the field of style transfer (Guo et al., 2023a; Yang et al., 2024), due to its proven effectiveness in capturing and transferring domain information. consists of two convolution layers and Adal N noise- The output synthesized by (x;;; (v) is defined as follows:
bb:;; )+) Conv( Aatoe(Ada/noostii Nnoiee (Conv())));
(35)
where Cone represents the convolution layers and AdaINnoise is variant of adaptive instance normalization (Huang & Belongie, 2017). To achieve counterfactual generation, AdaINnoiet replaces affine transformation parameters with style mean j E R and style variance â E TRS while simultaneously incorporating the variable noise v) E hRxxxx into the semantic content, with h, w and e indicating the height, width, and channel, respectively.
F. Additional Experimental Results.
F.1. Model attention visualization.
To demonstrate that our method relies on causal features rather than non-causal ones, we generate visual explanations using Grad-CAN (Selvaraju al., 2017). The models trained on the Waterbirds dataset using various FL methods. Grad-CAM, an extension of CAM, is a tool that provides visual insights into CNN-based image classifiers by highlighting image regions that contribute most to the model's predictions, thus allowing us to assess the features the model uses as the basis for its classification decisions.
The detailed results can be found in Fig. 8, which shows that the model trained with FedUni focuses more on the causal feature regions relevant to classification, such as the body parts and feather textures of birds. In contrast, the comparative
19

---

Causality inspired Federated Learning for OOD Generalization
models rely spurious features, such as the background. These findings provide strong evidence that method can effectively extract invariant causal features while avoiding spurious associations.
F.2. Mitigation of spurious correlations.
We further evaluate the extent of spurious associations by training models using different FL methods three datasets: CMNIST, CFMNIST, and Waterbirds. The degree to which each model relies spurious features, such as background and color, is quantified by the probability Ps.
In Table 4, 5, and 6, we demonstrate the relationship between test accuracy and for each model. Our proposed model outperforms the SOTA methods by exhibiting lower dependence on spurious associations across all three datasets. This indicates that our model effectively mitigates reliance on spurious associations.
p
1
0.9
0.8
0.7
0.6
0.5
0.4
0.3
0.2
0.1
0
FedAvg
87.92
86.3
85.07
82.84
82.26
80.35
78.99
77.40
75,59
74.26
70.18
-worst
85.01
82.67
82.66
78.81
79.18
77.44
75.71
72.86
71.79
70.18
67.5
FedProx
88.2
86.41
85.55
82.91
81.8
80.36
78.06
76.69
75.02
73.69
71.88
-worst
84.3
83.07
80.78
79.6
77.43
76.77
74.43
72.32
70.72
69.49
67.86
Scaffold
87.89
86.2
84.97
82.49
82.26
80.19
78.85
77.18
75.79
74.02
72.54
-worst
84.33
83,25
81.83
79.78
78.66
78.16
76.19
74.21
73.02
70.22
68.95
Moon
86.15
84.69
83.93
82.53
81.64
80.56
78.9
77.99
77.01
75.67
74.64
-worst
80.78
79.37
79.89
78.84
76.54
76.42
75.31
73.93
72.84
72.13
70.89
FedSR
87.99
86.55
85.49
83.45
82.52
81.2
79.41
78.28
77.06
75.74
74.26
-worst
83.25
82.89
81.31
80.69
77.62
77.08
74.91
74.64
73.57
72.92
68.93
FedlIR
88.1
86.14
85.61
82.7
81.88
80.27
78.79
76.54
75.25
73.52
72.01
-worst
85.01
83.03
82.19
79.54
78.16
77.44
74.82
72.32
71.66
70.04
68.39
FedDG-GA
87.35
85.64
84.79
82.86
82.38
79.95
78.63
77.31
75.84
74.49
73.13
-worst
82.54
82.36
81.13
80.07
79.4
76.17
74.19
73.11
70.94
69.49
67.51
FedSDR
87.64
86.22
85.74
84
83.24
81.86
80.78
79.4
78.69
77.7
76.67
-worst
84.52
82,46
82.54
81.72
79.72
78.67
77.86
76.43
75.54
73.65
71.84
Ours
89.45
88.38
87.81
86.33
85.64
84.25
83.08
82.24
81.4
80.61
79.41
-worst
86.07
85.26
84.3
84.13
83.25
81.39
80.78
79.11
78.21
78.16
76.61
Table 4. The relationship hetween average and worst test accuracy and test distribution specified by ps on Waterbirds dataset. (Rows without background color show average test accuracy results.)
F.3. Convergence behavior.
To provide a clear illustration of the convergence behavior of our method, we plot the test accuracy over global communication rounds on the Digits dataset, as shown in Fig. In the early stages of training, the test accuracy shows some fluctuations because the model has yet fully captured the causal features. However, the number of communication rounds increases, the model gradually learns more robust feature representations, enhancing its ability to handle diverse data distributions. Eventually, the curve becomes stable. This trend demonstrates that our method steadily converges throughout the training process and maintains stable performance later stages, indicating its effectiveness and robustness heterogeneous data scenarios.
G. Discussion of Limitations and Future Direction.
In this section, we discuss the limitations of our method and the potential solutions.
Despite the promising results our proposed method has achieved in reducing the reliance on spurious associations, it still suffers from some limitations.
Firstly, in terms computation, compared with the FedAvg method, the CMI filter and counterfactual generator included in our proposed method bring relatively large computational burden. This restricts the scale of the model and the amount
20

---

Causality inspired Federated Learning for OOD Generalization
p
0.9
0.8
0.7
0.6
0.5
0.4
0.3
0.2
0.1
0
FedAvg
96.72
91.92
87.1
82.46
77.67
72.73
68.53
63.28
59.21
53.92
49.18
-worst
90.5
85.4
82.3
74.9
67
54.6
48.2
39.7
32.7
26.1
FedProx
96.4
91.81
87.31
82.28
77.7
73.05
68.81
63.6
59.55
54.56
49.97
-worst
90.3
85.5
81.2
75.9
68.2
63
56.1
40.8
34.2
27.6
Scaffold
95.8
91.32
86.81
81.92
77.58
72.83
68.68
63.64
59.71
54.77
50.27
-worst
88.7
84.5
80.9
75.9
68.8
62.4
56.8
48.8
42.4
35.6
29.1
Moon
94.31
92.89
91.61
90.54
88.93
88.04
86.49
85.41
84.36
82.86
81.53
-worst
72.3
69.2
64.5
61.6
58.7
55
50.4
50.5
46.8
42.8
39.7
FedSR
97.18
93.2
89.39
84.74
80.96
76.46
72.52
68.69
64.97
60.18
56.34
-worst
90.3
85.7
82.3
77.7
72.45
66.7
61.9
56.2
49
41.8
37.4
FedIIR
96.22
91.77
87.25
82.4
77.93
73.39
69.3
64. 64.14
60.38
55.33
50.79
89.3
84.9
80.2
74.9
67.5
62.5
55.4
49.4
41.5
35.1
28.6
FedDG-GA
96.92
92.44
87.53
82.68
78.2
73.38
69.4
64.1
60.16
55.22
50.55
-worst
92.7
88.3
81.8
76
69.2
62
55.7
48
42
34.2
26.9
FedSDR
96.73
92.15
87.5
82.47
78.07
73.32
68.83
63.77
59.77
54.71
50.01
-worst
88.9
83.7
80.5
75.3
69.6
61.2
56.7
48.2
42.7
35.1
27.8
Ours
96.57
95.16
93.8
92.37
90.92
89.46
87.88
86.46
85.39
83.75
82.28
-worst
92
89.5
86.1
82.1
78
74.7
69.7
68.4
64.2
60
56.4
Table 5. The relationship between average and worst test accuracy and test distribution specified by P. CMNIST dataset. (Rows without background color show average test accuracy results.)
p
0.9
0.8
0.7
0.6
0.5
0.4
0.3
0.2
0.1
0
FedAvg
93.13
91.96
91.14
90.4
89.31
88.38
87.21
86.61
85.79
84.72
83.71
-worst
66.7
64.4
61.2
59.9
58.4
52
52.2
48.8
46.9
44.5
FedProx
95.23
93.62
92.08
91.1
89.03
88.07
86.28
84.92
83.57
82.05
80.54
-worst
78.3
74.4
68.1
66
60.4
56.8
50.4
50
45.6
40.7
36.8
Scaffold
94.73
93.29
91.93
90.81
89.16
88.12 18.12
86.56
85.35
84.29
82.73
81.36
-worst
74.4
71.4
66.3
63.2
60.2
56.5
52.2
51.9
48.2
44.3
41.1
Moon
95.53
92.01
88.38
84.62
80.82
77.44
74.01
70.16
67.18
63.19
59.5
-worst
86.1
83.1
80.2
77
72.5
67.3
61.7
55.8
51.3
45.4
39.6
FedSR
93.4
92.34
91.14
90.19
88.62
87.89
86.18
85.68
84.48
83.27
81.99
-worst
69.5
67.2
62.7
61.2
58.5
55.9
51.2
50.6
48.4
45.9
43
FedIIR
93.53
92.3
91.26
90.43
89.14
88.22
86.86
86.04
85.07
83.96
82.78
-worst
68.3
65.7
62.1
60.3
58.6
55.1
51.7
51.8
48.2
45.6
43.2
FedDG-GA
93.2
91.95
91.07
90.31
89.19
88.13 13
86.97
86.32
85.41
84.29
83.26
-Worst
66.8
64.3
61.1
59.7
58.3
54.6
51.5
52
48.5
46.4
44.1
FedSDR
95.34
93.71
92.13
90.95
89,04
87.82
86.21
84.68
83.32
81.77
80.14
-worst
77.5
74.1
68.9
65.7
61.5
57.8
53.6
51.6
47.8
44.4
40.5
Ours
94.95
93.91
92.83
91.9
90.61
89.79
88.73
87.68
86.67
85.58
84.51
-worst
76.5
74.5
69.4
66.6
63.5
60.7
57.2
55.6
$2.5
49.3
46.2
Table 6. The relationship between average and worst test accuracy and test distribution specified by on CFMNIST dataset. (Rows without background color show average test accuracy results.)
21

---

Causality inspired Federated Learning for OOD Generalization
Origin
FedAvg
FedDG-GA
FedlIR
FedProx
FedSDR
MOON
Scaffold
FedSR
Ours
Figure 8. Visualization of the effects of some federated learning methods (including ours) on the Waterbirds dataset. The red areas are important for prediction the correct class, and blue areas are vice versa.
Communcation Kounds
Comminsaton Rounds
Communicanol Rounds
Comsuniation Rounds
Communicaiio Rounds
(a) SYN
(b) MNIST
(c) USPS
(d) SVHN
(c) MNISTM
Figure 9, Accuracy curves with global communication rounds increase.
22

---

Causality inspired Federated Learning for OOD Generalization
of data that be processed. As result, our method fits better in cross-silo FL scenarios where devices have sufficient computing power.
Secondly, device heterogeneity has not been considered. In FL., different clients may have diverse hardware conditions, including network environments and computing capabilities. Some clients can handle complex causal feature extraction and selection, while others may struggle to complete these tasks due to limitations in computing capabilities, network, or memory. On the other hand, the datasets on different clients may also vary. Some clients may have very few or low-quality datasets, which can affect the local learning effect and, in turn, the overall performance. It is important to note that the current framework proposed in this paper assumes uniform computational resources across all clients. This assumption simplifies the experimental setup but does not fully capture the complexity of real-world scenarios involving heterogeneous devices and resource constraints.
To address the aforementioned limitations, several promising research directions can be explored. A key area is optimizing the computational efficiency of our method. By exploring more efficient algorithms and techniques, we can reduce the computational burden and make the method applicable to a wider range of model scales and data amounts. Additionally, designing customized lightweight approximation schemes for clients with low computing power and unstable network conditions could further enhance the adaptability of our method.
23


Reviews:
[{"review_id": "pzmNb1EPft", "reviewer": "Reviewer_FyJy", "summary": "The paper proposes FedUni, a framework for federated learning. The framework extracts causal features from different clients in training and flexibly selects the features applicable for the target client. It differs from the existing methods in terms of not using a fixed sharing feature pool. The paper provides experiments to show the improvement brought by the novel design.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your valuable suggestions, and we will enlarge the text font in Fig 5, 7 in the final version.\n\n**Q1: Computational Cost** \n\nWe 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.\n\nFirstly, 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.\n\n|Run Time (s)|FedAvg|FedProx|Scaffold|Moon|FedSR|FedIIR|FedDG-GA|FedSDR|Ours|\n|-|-|-|-|-|-|-|-|-|-|\n|NVIDIA GeForce RTX 2080|1.33|1.87|1.94|2.04|2.12|1.95|3.05|5.02|3.26|\n|NVIDIA GeForce RTX 3090|0.77|0.92|0.91|1.13|1.53|0.86|1.40|2.60|1.57|\n|**Test Acc (%) （OOD-avg）**|68.53|68.23|68.47|73.92|72.69|68.99|69.02|68.72|**88.03**|\n\nIn addition, we presented a brief proof for the convergence rate, showing that our convergence speed is also comparable to that of other methods. \n\nWe 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})$.\n\n**Q2: When to choose FedUni** \n\nThe 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**.\n\n**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. \n\n**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. \n\n**Q3: Performance Gains in Homogeneous Setting** \n\nReviewer 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.\n\n|Test Acc (%)|FedAvg|FedProx|Scaffold|Moon|FedSR|FedIIR|FedDG-GA|FedSDR|Ours|\n|-|-|-|-|-|-|-|-|-|-|\n|ID|97.08| 96.68| 97.06| 97.08|97.42|96.07|96.92|95.18|**98.16**|\n|OOD| 85.27| 85.46| 87.83|86.20|86.18|90.25|88.24|89.81|**97.30**|\n\n**Q4: More rigorous argument for Lemma 4.2** \n\nThanks for your review! We will revise Lemma 4.2 into a more rigorous form in the final version.\n\nInspired 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. \nFirst, 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]$.\n\n[1] Pedro Savarese,et al. Information-Theoretic Segmentation by Inpainting Error Maximization. 2021.", "claims_and_evidence": "The theoretical claims are well-supported by proofs and explanations. I found them easy to follow and understand. The novel specification of SCM looks natural to me, and the advantages are well presented in the experiments.", "methods_and_evaluation_criteria": "The proposed method and evaluation makes sense to the specific problem.", "theoretical_claims": "I checked Lemma 3.2, 3.4, 4.2. For Lemma 3.2 and 3.4, the techniques are pretty standard and look correct to me. \n\n\nI think the $\\approx$ notion looks hand-wavy in Lemma 4.2, thus I suggest making a more rigorous argument.", "experimental_designs_or_analyses": "The design (Sec 5.1) and analysis (Sec 5.2) both look solid to me.", "relation_to_broader_scientific_literature": "The design of fake causal features is a smart and novel idea. It may be useful for the federated learning tasks where the target client requires domain knowledge (not just common knowledge) from a specific client in training. Yet for the tasks that the training clients are very similar to each other, the method may not make a large difference despite requiring more computational cost.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "I like the presentation of the paper. Though it is mathematically dense and has a lot of content both theoretically and empirically, it does not cause a burden to read. \n\n\nIt could be helpful to make the text font in some figures (Fig 5, 7) larger.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "1. What is the computational cost of the current method? I see it is briefly mentioned in Appendix G. I’m wondering if the authors could provide details from either experiments (actual run time) or theories (time complexity/asymptotic convergence rate).\n\n2. How should practitioners know when to choose FedUni over the prior methods?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "CJzAxQBKT5", "reviewer": "Reviewer_KHjr", "summary": "This paper addresses the challenge of OOD generalization in FL by rethinking how causal features are extracted across clients. The authors argue that instead of limiting the global feature extractor to invariant features (like the traditional FL methods), it should capture the union of all causal features present in the data, thereby preserving richer and more diverse information that can enhance OOD performance by leveraging both collaborative training and targeted causal interventions. The proposed FedUni comprising three core components. 1) A comprehensive feature extractor is designed to identify all potential causal features from any input across clients. 2) A fake causal feature compressor is employed to filter out client-specific or spurious features that do not contribute to the target task. 3) A causal intervention module on the client side uses a trainable counterfactual generator to create modified examples that simulate distribution shifts while preserving semantic content. Experimental results on different datasets and theoretical analysis are provided in the paper.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your suggestions, and we will include the citation you pointed out in the final version.\n\n**Q1: Experimental Designs** \n\nWe conduct additional experiments in response to the suggestions about experimental designs.\n\n**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.\n\n|ResNet50|Art|Cartoon|Photo|Sketch|Avg|\n|-|-|-|-|-|-|\n|FedAvg|75.07|75.47|92.80|76.20|79.88|\n|FedSR|74.27|71.60|92.13|67.87|76.47|\n|FedGD-GA|77.47|72.80|**94.27**|69.73|78.57|\n|FedIIR|74.80|71.87|91.33|71.87|77.47|\n|Ours|**80.27**|**75.87**|93.27|**79.60**|**82.25**|\n\n|ViT-Base-32|Art|Cartoon|Photo|Sketch|Avg|\n|-|-|-|-|-|-|\n|FedAvg|80.01|74.53|93.07|69.73|79.33|\n|FedSR|73.60|68.13|87.33|65.07|73.53|\n|FedGD-GA|80.34|71.27|93.07|65.87|77.63|\n|FedIIR|78.93|72.67|87.87|68.93|77.10|\n|Ours|**80.40**|**74.60**|**94.93**|**78.13**|**82.02**|\n\n**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. \n\n|Test Acc (%)|FedAvg|FedProx|Scaffold|Moon|FedSR|FedIIR|FedDG-GA|FedSDR|Ours|\n|-|-|-|-|-|-|-|-|-|-|\n|ID|97.65|97.86|98.25|97.66|98.45|96.04|94.5|98.19|**99.61**|\n|OOD|76.48|76.59|83.51|76.79|80.08|72.69|73.86|76.77|**93.87**|\n\n**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.\n\n**Q2: Broader Scientific Literature and References**\n\nWe 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. \n\nWe 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.\n\n|Test Acc (%)|DRC [ref 1]|CG-PFL [ref 2]|Ditto [ref 3]|Ours|\n|-|-|-|-|-|\n|ID|97.17|98.04|96.15|**99.61**|\n|OOD|76.75|79.94|65.23|**93.87**|\n\n| | Test Acc (%)|\n|-|-|\n|FOOGD|79.61|\n|Ours|**90.26**|\n\n**Q3: Question about Counterfactual Generator** \n\n**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. \n\n**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.\n\n**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.\n\n\n[1] Huang, et al. Arbitrary Style Transfer in Real-time with Adaptive Instance Normalization. 2017. \n\n[2] Guo, et al. Single Domain Generalization via Unsupervised Diversity Probe. 2023.\n\n[3] Yang, et al. Practical Single Domain Generalization via Training-time and Test-time Learning. 2024.\n\n**Q4: Computation Cost** \n\nPlease refer to our response to Reviewer FyJy 4.  FedUni achieves significantly higher accuracy, while maintaining comparable computational cost to many state-of-art approaches.", "claims_and_evidence": "The process of capturing causal features while rejecting fake ones relies heavily on the performance of a counterfactual generator. This separation could be in practice challenging and dependent on the model type, architecture, size, hyperparameters, etc. In this paper, a very simplistic model of two convolution layers along with some normalization was considered; but, how we can justify using such simplistic model for the purpose, especially for practical size images (not to small pixel size)?", "methods_and_evaluation_criteria": "Yes.", "theoretical_claims": "Seem to be correct.", "experimental_designs_or_analyses": "1.\tThe experimental evaluation in this paper is undermined by the use of outdated and overly simplistic models, namely ResNet-18, a one-hidden layer MLP, and AlexNet. These choices do not reflect the current SOTA architectures (ResNet-50/101, ViT, etc) that are standard in top-tier research. These modern architectures have fundamentally different inductive biases, scalability, and computational characteristics. Therefore, the results currently presented in this paper may not reliably generalize to realistic settings. \n2.\tThe number of clients seems to be small, N=10, to capture real-world scenarios. In many recent papers on FL, much more clients are often considered, e.g., 100.\n3.\tKey practical issues of FL such as non-iid data and client sampling have not been considered.", "relation_to_broader_scientific_literature": "The problem of OOD generalization can be also addressed through personalized FL (PFL). It seems that the authors claim that the PFL approach is limited, by citing a couple of works such as (Tang 2023, 2024). But, the literature of PFL is very rich, and some of them explicitly tackle the issue of OOD generalization while numerous others tackle the issue implicitly (but, possibly effectively). Few examples are:\n\n[Ref 1] Exploiting Personalized Invariance for Better Out-of-distribution Generalization in Federated Learning\n\n[Ref 2] Personalized Federated Learning with Contextualized Generalization\n\n[Ref 3] Ditto: Fair and Robust Personalized Federated Learning\n\nI am hoping to see fair and proper comparison against those approaches of PFL.", "essential_references_not_discussed": "See above comment. Relevant references for PFL for OOD generalization should be cited (and if possible, the performance should be compared). Also, FOOGD: Federated Collaboration for Both Out-of-distribution Generalization and Detection (NeurIPS 2024) should be considered (and if possible, the performance should be compared).", "other_strengths_and_weaknesses": "Another weakness: In the paper, the intervention module is a key component, which generates counterfactual examples that simulate distributions shifts. However, it is not clear how the authors can guarantee that these generated examples accurately simulate distribution shifts without altering the underlying semantics and how sensitive the method is to potential failures of proper generation? In Appendix E, the authors clarify that the generator is a simplistic model composed of two convolution layers and some adaptive normalization. But, I am not convinced that this overly simplistic network can generate high-quality realistic samples that simulate distributions shifts, without changing the underlying semantics.", "other_comments_or_suggestions": "The proposed method can be computationally demanding on the client side due to the extra modules, which raises a question of scalability. In the current experimental setup, the model architecture and number of clients are all small. More experiments on large-scale or highly heterogeneous federated settings would help assess the scalability of the proposed scheme.", "questions_for_authors": "My concerns and questions have been raised in the above comments.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "vqVTh7H8Ku", "reviewer": "Reviewer_Lsn9", "summary": "The authors identify a limitation in federated learning, arguing that existing methods primarily capture invariant causal features across clients/environments while failing to account for variant causal features. To address this, they propose a method designed to capture both invariant and variant causal features—both of which are direct causes of the target variable, y. Their proposed architecture, FedUni, is trained to extract all possible causal features from input data, acknowledging that some extracted features may be non-causal. To mitigate the risk of incorporating spurious correlations, they introduce a feature compressor. Additionally, they incorporate a causal intervention module on the client side, leveraging a counterfactual generator to create counterfactual examples.\n\nThe authors provide extensive experiments and theoretical analysis, claiming that their method significantly enhances out-of-distribution (OOD) generalization.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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. \n\nMore 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. \n\nThus, 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$.", "claims_and_evidence": "While the approach appears well-motivated, the effectiveness of capturing variant causal features—especially under real-world federated settings—remains an open question. A deeper examination of the robustness of their counterfactual generation and feature compression mechanisms would be valuable in assessing the practical impact of this work.", "methods_and_evaluation_criteria": "The motivation and proposed method are well-founded, and the choice of benchmark datasets is appropriate.", "theoretical_claims": "No concern.", "experimental_designs_or_analyses": "Experiments are reasonable.", "relation_to_broader_scientific_literature": "It is relevant to the broader machine learning community.", "essential_references_not_discussed": "Not that I am aware of.", "other_strengths_and_weaknesses": "The paper has a good motivation. However, it lacks clarity to explain and support each component. \n\nIn lines 281–283, the authors state, \"Firstly, with the causal intervention module, Z_S can be eliminated by minimizing the causal effect of environmental changes.\" However, this claim seems questionable, as the authors themselves acknowledge that some detected causal features may, in fact, be non-causal.\n\nIt remains unclear how misidentified non-causal features are effectively reduced while true causal features are amplified. Further clarification on this mechanism would be beneficial.", "other_comments_or_suggestions": "Already mentioned in above.", "questions_for_authors": "Please address my comment in \"Other Strengths And Weaknesses\".", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "GEfzU9Zusu", "reviewer": "Reviewer_MJM2", "summary": "I am not familiar with the field of this submission. I suggest the AC find another reviewer and disregard my comments. Apologies for the inconvenience.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your rigorous and responsible review.", "claims_and_evidence": "I am not familiar with the field of this submission. I suggest the AC find another reviewer and disregard my comments. Apologies for the inconvenience.", "methods_and_evaluation_criteria": "I am not familiar with the field of this submission. I suggest the AC find another reviewer and disregard my comments. Apologies for the inconvenience.", "theoretical_claims": "I am not familiar with the field of this submission. I suggest the AC find another reviewer and disregard my comments. Apologies for the inconvenience.", "experimental_designs_or_analyses": "I am not familiar with the field of this submission. I suggest the AC find another reviewer and disregard my comments. Apologies for the inconvenience.", "relation_to_broader_scientific_literature": "I am not familiar with the field of this submission. I suggest the AC find another reviewer and disregard my comments. Apologies for the inconvenience.", "essential_references_not_discussed": "I am not familiar with the field of this submission. I suggest the AC find another reviewer and disregard my comments. Apologies for the inconvenience.", "other_strengths_and_weaknesses": "No", "other_comments_or_suggestions": "No", "questions_for_authors": "No", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: The reviewers generally agree that this paper merits acceptance; therefore, I recommend acceptance.

*Strengths.* The paper proposes a novel method that addresses the performance cost of only relying on invariant features when there might be other potentially casual features for some clients. The expected standard theoretical results are provided, and the empirical results cover reasonable datasets and baselines.

*Weaknesses.* The paper is quite dense. While the regularizers provided can be considered related to the desired properties, e.g., maintaining semantics, there is no provided guarantee or test of it.

The authors provided additional experiments to improve the empirical results during rebuttal and clarified reviewer questions. Most reviewers were positive before the rebuttal and remained positive after. I recommend acceptance conditioned on following through on the clarifications and discussions from the review process. Additionally, more discussion on privacy and the potential impacts of it in federated learning and in the context of the proposed method is neccesary.
Decision: Accept (poster)