Learning Personalized Causally Invariant Representations for Heterogeneous Federated Clients
Authors: Xueyang Tang
Song Guo
Jie ZHANG
Jingcai Guo
Venue: ICLR 2024 poster (2024)
Keywords: Personalized Federated Learning
Invariant Learning
Causality
Out-of-distribution Generalization
Shortcut Learning
Primary area: general machine learning (i.e., none of the above)
TL;DR: We propose a provable shortcut discovery and removal method under personalized federated learning to develop the optimal personalized invariant predictors that can generalize to unseen OOD data for the heterogeneous clients.

Abstract: Personalized federated learning (PFL) has gained great success in tackling the scenarios where target datasets are heterogeneous across the local clients. However, the application of the existing PFL methods to real-world setting is hindered by the common assumption that the test data on each client is in-distribution (IND) with respect to its training data. Due to the bias of training dataset, the modern machine learning model prefers to rely on shortcut which can perform well on the training data but fail to generalize to the unseen test data that is out-of-distribution (OOD). This pervasive phenomenon is called shortcut learning and has attracted plentiful efforts in centralized situations. In PFL, the limited data diversity on federated clients makes mitigating shortcut and meanwhile preserving personalization knowledge rather difficult. In this paper, we analyse this challenging problem by formulating the structural causal models (SCMs) for heterogeneous federated clients. From the proposed SCMs, we derive two significant causal signatures which inspire a provable shortcut discovery and removal method under federated learning, namely FedSDR. Specifically, FedSDR is divided into two steps: 1) utilizing the available training data distributed among local clients to discover all the shortcut features in a collaborative manner. 2) developing the optimal personalized causally invariant predictor for each client by eliminating the discovered shortcut features. We provide theoretical analysis to prove that our method can draw complete shortcut features and produce the optimal personalized invariant predictor that can generalize to unseen OOD data on each client. The experimental results on diverse datasets validate the superiority of FedSDR over the state-of-the-art PFL methods on OOD generalization performance.

Published as a conference paper at ICLR 2024
LEARNING PERSONALIZED CAUSALLY INVARIANT REPRESENTATIONS FOR HETEROGENEOUS FEDER-
ATED CLIENTS
Xueyang Tang', Song Guo"; Jie Zhang & Jingeai Guo
'The Hong Kong Polytechnic University
The Hong Kong University of Science and Technology
ABSTRACT
Personalized federated learning (PFL) has gained great success in tackling the see- narios where target datasets are heterogeneous across the local clients. However, the application of the existing PFL methods to real-world setting is hindered by the common assumption that the test data on each client is in-distribution (IND) with respect to its training data. Due to the bias of training dataset, the modern machine learning model prefers to rely on shortcut which can perform well the training data but fail to generalize to the unseen test data that is out-of-distributio (OOD). This pervasive phenomenon is called shortcut learning and has attracted plentiful efforts in centralized situations. In PFL, the limited data diversity on fed- erated clients makes mitigating shortcut and meanwhile preserving personaliza- tion knowledge rather difficult. In this paper, we analyse this challenging problem by formulating the structural causal models (SCMs) for heterogeneous federated clients. From the proposed SCMs, we derive two significant causal signatures which inspire provable shortcut discovery and removal method under feder- ated learning, namely FedSDR. Specifically, FedSDR is divided into two steps: I) utilizing the available training data distributed among local clients to discover all the shortcut features in a collaborative manner. 2) developing the optimal person- alized causally invariant predictor for each client by eliminating the discovered shortcut features. We provide theoretical analysis to prove that our method can draw complete shortcut features and produce the optimal personalized invariant predictor that can generalize unseen OOD data on each client. The experi- mental results on diverse datasets validate the superiority of FedSDR over the state-of-the fe-theean PFL methods on OOD generalization performance.
INTRODUCTION
Federated learning (FL) allows the participation of massive number of data holders (i.e., clients) that possess limited data to collaboratively train learning models in a datasets across local (McMahan al.] 2017). From the view of the heterogeneity of target clients, we can divide the literature on FL into two branches. 1) Federated learning aims at training global model fit the local data distributions and perform well when the local target datasets are subject to independent and identically distribution (IID). In particular, some works (Deng et al.) 2020 Liu et al. 2021c Nguyen al. 2022) (including robust federated learning and fedcrated do- main generalization) focus on training a global model that can tackle the distribution/domain shift across local training datasets. Unfortunately, the shared global model can diverge from the optimal local solutions when the target datasets are heterogeneous or not IID (i.e., Non-IID) across local clients (Hsieh et al. (2020), since the useful information about personalization dropped. 2) Per- sonalized federated learning develops a personalized model for each client to handle the discrepancy among the local optima when the target datasets across local clients are Non-IID Despite succeeding in handling Non-IID target datasets, all the existing PFL methods neglect the shorteut trap problem which is attracting more and more interest in centralized machine learning.
man-
"oorresooniing author

---

Published as a conference paper at ICLR 2024
Shortcut trap is found pervasive in modern machine learning (Geirhos al. 2020) where models prefer to rely on the shortcut to solve problems due to the bias of training dataset. The utilized shorteut can perform well training data but fails to generalize to unseen test data that is out-of- distribution (OOD) with respect to the training data. For example, there is binary image classifica- tion task where the model needs recognize the pictures of cows and camels (Beery et al. 2018)_ Deep learning model can classify the picture of a cow in a desert background as "camel" at test time, if most of cows appear grass backgrounds and most of camels stand in desert backgrounds in training environments (environments are data subsets that have different data distributions). This dataset bias makes the obtained model choose the background rather than the shape of animals the pictures as the discriminative feature. The similar shortcut trap exists in diverse weal-world scenar- ios (Geirhos et al. 2020). Although many efforts have been attracted to the shortcut trap problem in centralized situations, they focus on mitigating shortcut by extracting environment-invarian (a.k.a. invariant) features. When applying these schemes into PFL, the invariance constraint will climinate all heterogeneous features, including shortcut and personalized features. Therefore, the existing invariant learning schemes can hardly tackle the shortcut trap problem in PFL.
What's 's worse, find the trivial combination of the ex- isting PFL and centralized invariant learning schemes, in- stead of solving the shortcut trap problem in PFL., can even induce worse performance than the better one of them- selves (discussed in the evaluation part). To handle the challenging shortcut trap problem in PFL, we firstly for-
Train-test relation
FL PTL
RFL Fed
0OD
mulate the structural causal models (SCMs) to simulate the
heterogeneous data generating processes on local clients. From the SCMs, we derive causal signature which re- veals that the shorteut is statistical independent with the client/user indicator conditional on label and environment indicator. Inspired by this finding, we design collabora- tive shortcut discovery method which can work well even if there only one available training environment on each client. Then, the personalized causally invariant represen- tations are extracted by utilizing another causal signature that describes the conditional independence between the personalized invariant features and the shortcut features. Finally, the optimal personalized invariant predictors can be elicited from the extracted personalized causally invariant features. The comparison between the coverage of our approach and the related works is illustrated in Figure [ The main contributions of this paper are summarized as follows:
IND
Test-test
IID
Non-IID
relation
Figure 1: The coverage of ours and the re- lated works. a) FL: Federated Learning; b) PFL: Personalized Federated Learning; c) RFL: Robust Federated Learning: d) Fed DG: Federated Domain Generalization. Besides, IND denotes in-distribution.
To the best of our knowledge, we are the first to consider the shortcut trap problem in per- sonalized federated learning and analyse it by formulating the structural causal models for heterogencous clients. Based on the proposed SCMs, we design a provable shortcut dis- covery and removal method to develop the optimal personalized invariant predictor which can generalize to unseen local test distribution for each client.
Theoretically, we demonstrate that the designed shortcut discovery method can draw all the latent shortcut components, then the shortcut removal method can eliminate the discovered shortcut features and produce the optimal personalized invariant predictor for each client.
Empirically, we conduct experiments on several commonly used datasets and the results validate the superiority of our method on ouittf-distribution gener- alization performance, compared with the state-of-the-art competitors.
RELATED WORK
Federated learning. The classic FedAvg (McMahan et al. 2017) performs well if local training datasets are IID. Some methods (Karimireddy et al. 2200): Dieuleveut et al. (2021): Zhang et al. (022 ) mitigate the negative impact of training data heterogeneity on convergence rate, while an- other branch (Deng et al. (2020); Sharma et al. 2202);; ; Sun & Wei (2022) targets at reducing the performance bias of global model on local clients. Besides, few works (Liu et al. (2021c); Nguyen et al. (002)): Guo al. (003ch) investigate the scenarios where the training data heterogeneity ap-
2

---

Published as a conference paper at ICLR 2024
pears to be domain shift. All the above methods produce a shared global model which can diverge from the local optimal solutions when local target datasets are Non-IID.
Personalized federated learning. Many PFLs Dinh et al. (2020) Hanzely et al. (2020): Fallah
ct al. (2020); et et al. (202T); Tang et al. (2022); Cheng et al. (023)) train the personalized models
with the guidance of global model which embeds in the shared knowledge, DFL (Luo et al.,
2022) disentangles the shared features from the client-specific ones to achieve accurate aggregation
on shared knowledge. Similarly, pFedPara (Hyeon-Woo et al.) 2022) and Factorized-FLL (Jeong &
Hwang 2022) factorize the model parameters into the shared and personalized parts. Another branch ( (Collins et al. (202T): Chen & Chao (2022); Xu et al. (2023)) employs the shared/aligned feature extractor to capture global knowledge and personalized classifiers to encode the personalization information. All of them don't cover the situations where shorteut exists local training datasets.
Shortcut and Invariant learning (IL). Causally invariant predictor is proposed in (Peters et al., 2016), and then applied into deep learning in IRM (Arjovsky et al.] 2019) to mitigate shortcut. Subsequently, Subsequentty_Roseeeeed when training environments are insufficient. IFM (Chen et al., 2022b) lowers the requirement and demands only logarithmic training environments. Some works focus settling IL problem when the environment label is unavailable, e.g., EIIL (Creager et al., 2021), HRM (Liu et al. 2021a b),
(Rosenfeld et al. (202l) prove that IRM and its variants can be still trapped by shortcut
)022). Another branch (Ahuja al. (2021); Chen et al. (2022a): Huh & Baidya (2022)) completes the constraints that IRM misses to improve the performance. The iCaRL (Lu al. 2022) extends II. to non-linear causal representations while
EDNIL (Huang al. (222) and ZIN (Lin et al.
ACTIR (Jiang & Veitch 2022) extends IL to anti-causal scenarios. All these methods are devised for centralized scenarios where all training data is accessed and training environments are sufficient.
PROBLEM FORMULATION
Notations. Let X, y and & denote the input, target and environment space respectively. Data instance is (X, v) E (X.Y,E) V) Suppose there are N clients and the local dataset D,, on client contains M,, samples, u E [N]. The sets of training and test environments on client are denoted by E and E% respectively. We use Ealt as the set of all possible environments the task that client concentrates environment sets are denoted by Egr := U., gu and Eall := U., Edlt. For convenience, we separate the learning model parameterized mapping from X' to y' into two consecutive parts: 1) the feature extractor (& and V denote the invariant and spurious feature extractors respectively) maps from input space X' to latent feature space 2, i.e., (XX) E Z and (X) 2; 2) the classifier outputs a prediction ÿ from latent feature = E Z. For example, the overall model based on the invariant feature extractor is denoted by fa(-) wx)(-)) where fo indicates the function parameterized by 0. We define the expected empirical loss for model so on dataset D as R.o; D) E (ex v)] where is the loss function.
i.e., Etr, Eu C Catty Vu E [N]. In federated learning system, the overall
3.1 INVARIANT LEARNING
Succeeding in mitigating shorteut and solving the OOD generalization problem, invariant learning assumes that there exists some invariant feature & () satisfying the invariance constraint:
(X) = e, = ((X) 2, e,) V2 E e', E Eau-
(1)
Rosenfeld al. 202
proved that IRM (Arjovsky et al. 2019) and its variants need at least dg + (ds is the dimension of shortcut features) training environments to eliminate all shortcut features and elicit the optimal invariant predictor, under linear scenarios.
Definition (Optimal Invariant Predictor) The optimal invariant predictor is elicited based on the complete invariant features which are informative for the target label in the task concerned, i.e., 4* = arg max. [(Y; (X)), where ((::)) denotes the Shannon mutual information between two random variables and t satisfies the above invariance constraint.
3.2 CAUSAL SETUP
In invariant learning (IL). researchers usually formulate a structural causal model to simulate the data generating process in the target task. A valid SCM is depicted by a directed acyclic graph where each
3

---

Published as a conference paper at ICLR 2024
node represents random variable and each edge deseribe directed functional relationship between the corresponding variables (Pearl, 2009). When we study the invariant learning in federated setting, the latent heterogeneity of data generating mechanisms among local clients need to be considered.
E
ZE
Ls
Zc
72
Z
7
ZZ
Z
X
(a) Causal IL
(b) Anti-causal IL
(c) Causal FedSDR
(d) Anti-causal FedSDR
Figure 2: (a) (Arjovsky et al. 2019 Huang et al. 2022 and (b) (Rosenfeld ct al. 2021 Huh & Baidya 2022
give the structural causal models (SCMs) commonly adopted in invariant learning. while (c) and (d) show the SCMs proposed in this paper. Ze and Zs denote the invariant and shortcut features respectively. E is the indicator of shortcut while U is the indicator of user/client. Dotted arrows indicate unstable causal relations can vary in different environments.
Therefore, we propose the SCMs in federated learning by adding the user/client indicator U and deconstructing the invariant features into separate parts: the personalized invariance zg and the chared/global invariance zá. The detailed SCMs are shown in Figure 2 As discussed in the literature on invariant learning, Zs the latent shortcut feature. The functional relation between Zs and label Y can vary across different environments. That is, AZs there always exists some e, e' E Eaul that make P(Y (7sse] e) F P(Y) e') hold. By analogy with the optimal invariant predictor in invariant learning, we provide the definition of the optimal personalized invariant predictor in PFL.
Definition (Optimal Personalized Invariant Predictor) The optimal personalized invariant predictor for client is elicited based the complete invariant features which are informative for target label in the task that client concentrates on. That is, D arg maxq, I(Y: Pu(X)), where 4, satisfies that P(Y P(Y) e, = P(Y| 2, e,e E Z, Ve, c' Edi-
4 METHOD: FEDSDR
When the training environments on each client are insufficient, locally invariant learning can fail discussed before. How about collaborative manner? Unfortunately, personalized invariant features can cause deviation of the invariance constraint as shortcut features do if training environ- ments are collected from different clients. As result, the collaborative invariant learning can clim- inate/preserve both personalized invariant and shortcut features with the same probability. Nonethe- less, how about combining collaborative invariant learning with PFL methods? ? Even though we can get the global invariant features via collaborative IL and conduct local adaptation as in many PFL. schemes (e.g., fine-tuning (Cheng et al. 2023) and L2-regularizzatinn ( Dinh et al. 20200] Hanzely al.) 2020 Li et al.) 2021)), the local adaptation can pick both personalized invariant and shortcut features again since local training environments are insufficient. It turns out the trivial combination can hardly outperform the superior individual one on OOD generalization performance.
FedSDR. In view of the above failure, we turn to the complementary perspective: discovering the shortcut features and removing them instead of straightly constraining invariance. The feasibility of this tack is guaranteed by the causal signatures that we derive from the SCMs in Figure ((c) and ((d)
Lemma If the data generating mechanism of each federated client obeys the causal graph in Figure or the anti-causal graph in Figure 2(d) we can have:
Zs L U | Y.E which means that the shortcut features Zs are conditionally independent of the personalization indicator U given Y and E.
zg T Zs Y and z Zs Y. which means that both the global (Z&) and personalized ( invariant features are conditionally independent of the shortcut features Zg given Y.
Remark 1. The first causal signature Lemma I indicates that we can discover the shortcut fea- tures using training environments across local clients even if the data generating mechanisms
4

---

Published as a conference paper at ICLR 2024
heterogeneous among them. The second causal signature makes it possible develop the optimal personalized invariant predictors with the discovered shortcut features even though there is just one training environment each client, since the relationships between Zg. Ze and 2ş are independent of environment E.. The detailed proof of Lemma] | can be found in the appendix.
In the following sections, will introduce the two-stage implementation of our method in detail.
4.1 THE PROVABLE SHORTCUT DISCOVERY
At the first stage, we need to capture the complete shortcut features in a collaborative manner. Recalling the difference between the definitions of shortcut features and invariant features, we design the following objective to extract the complete shorteut features in collaborative manner:
why V
arg min 1
1EGD D.) R(w(V); D.) &ais (v; Du)h
(2)
N
v. -
1=1
where the first term ((w(()) D.) is adopted to exclude the uninformative features (c.g.. noise). the balaneing weight and the second term lais(V; D.)) is designed for extracting the complete shortcut features. Specifically, we define that
Eaio (V. D.. = EXED. EE Σ (C (Y c(X),.))|Fu (Y ((X). cf
(3)
where (C((P) denotes the Kullback -ulbbeck-Leiblrr divergence between two probability distributions. Pwz (Y V, e,) means that P(Y L ei) is parameterized by the classifier w/ which is trained by:
w = arg min wit
Σ P. (V) Ve, E Etr;
(4)
where p's when client has data samples from environment e, and pt. 0 otherwise.
Since Pu,; (Y V e,):)) is parameterized by the classifier wi to be distribution around wi (V) for any
given To and Cgs adopt simple and effective measure compute the divergence KCL (Y) V, PF Ve, ej E Eait In this way. we can rewrite the overall objective 2 for shortcut discovery as the following bi-level optimization:
ww: (Y MD, e,)) That K.((Pw; w; (Y V [P)l (Y V w,e,)) = _ [llw(()) w/ (
N
Wis V arg min V.w
Σ
(CDD(V D.) (.)( (V; D.))
(5)
N
S.L.
wi arg min -
Σ
pu ((w(((WW)))))) Ves E
(6)
(x) This bi-level opti- mization can be solved by alternatively updating the solutions of the outer and inner objective. Under federated learning, both the outer and inner objective can be divided into N sub- sub-problems that can be settled on local clients respectively. The sever can aggregate the update from local clients to gain the solution V To avoid the outer objective being dominated by maximizing laie (V; D.), we replace ldie with min(a, X Eais (V; D..)) in the practical version, where a is a positive threshold.
where laie D.) = EXED. Ec,ec.. Ee,ee,
Theoretical Analysis. Before continuing to introduce the shortcut removal method, we formally analyse the optimal solution of the Fq. In this theoretical analysis part, we consider the linear data model that Rosenfeld et al. 2021) adopted. Specifically, we focus on the logistic regression problem where label y {-I].
in (Rosenfeld et al.) (0211) we suppose both the invariant features Zo [zg, 20 and shorteut features 755 of sample y drawn from the following Gaussian:
Ze ~ Nyy Me, o.),
~ NN M..
where Me E de and μs E Rdd, The sample of observation is generated by X g(Zc, Zs) where g(-) is a non parameterized function. The parameters Pao a, and function g are independent of environment, while μs and a, varies in different environments.
Assumption The distribution of label Y satisfies the following two conditions: 1) P(Y u) = P(Y), Vu 2) Ply e) = PPl' {e), Vy. v E Y and Wee Eir-
5

---

Published as a conference paper at ICLR 2024
Theorem If the Assumption 7 holds, the function g is linear and the total number of training environments in the federated learning system satisfies [Err] >> dg. then the following two statements equivalent:
(X) depends and only depends on the complete shortcut features Zg. That is, (X) is finction of Zy alone;
TV* is the optima of the Eq. 5 with an appropriately chosen hyper-parameter X.
Remark 2. Theorem guarantees the elaborated Eq. can yield the feature extractor that extracts complete shortcut features and excludes all invariant features. Note that Assumption i about the label distributions in training datasets. Since the shortcut extractor works as an auxiliary model and never part of the optimal personalized invariant predictors, we can sample some data subsets from local training datasets to train the shorteut extractor yr* In this way, the sampled data subsets can Besides, the causal signatures in Lemma [ proof of Theorem T and the complete proof is provided in the appendix.
casily satisfy Assumption [
play critical parts in the
4.2 PERSONALIZED INVARIANT LEARNING WITH SHORTCUT REMOVAL
With the shorteut extractor that depends and only depends on the complete shortcut features Zg, we can extract the most informative invariant features to elicit the optimal personalized invariant predictor for each client. Based on the second causal signature in Lemma I we design the following objective for each client to develop the optimal personalized invariant predictor:
(T("") = arg min (w (b.); Dw) = (R((w((u)) :) Dw) + I(eu; V/ [Y), Vu E [N],
(7)
where (:: denotes the conditional mutual information and is the balancing weight. The optimal personalized invariant predictor is given by fö. := (
Theorem Suppose 1 (X) in the Eq. 7 depends and only depends on the complete shortcut fea- tures Zg. If fō. (Vu [N) the optima of the Eq. with the kyper-parameter ? chosen appropri- ately, then the fø., is the optimal personalized invariant predictor for the client u. Yu [N].
Remark 3. Theorem 2 guarantees that our method can produce the optimal personalized invariant predictor for every chient. Note that I(P, V Y) = is the necessary and sufficient condition for 4., JL V Y. Since 4,, 1. V Y is independent of environment, our method can develop the optimal personalized invariant predictor for every client even though there is only one training environment on each client. The complete proof of Theorem i provided in the appendix.
In the practical implementation, it can be infeasible to compute the exact value of T(P.,; V Y)- Considering the limited computation resources on local clients, we adopt a simple approximating scheme used in (Jiang Veitch 2022) to measure I(P,, V Y). Specifically, we estimate it by
I(Pu: 4** Y) = PxP (but not necessary) condition for E
- E Y1) because I(P., 4+* Y) = is the sufficient
(X) E
E((((X [Y]i 0. With the data samples
on local clients, estimate the conditional mutual information by
Ma
Mr
I(O.; V (Y)
1
Σ 10=1
@x(Xm)
(V (X'(X))
Σ u=1
V (X(.))
M.
2
where (Xm, Vm)) m [Mu] is drawn from dataset Da, 1/f if yn ym and t. = 0 otherwise.
Note that our method can easily cooperate with most of the existing PFL methods to improve their OOD generalization performance by adding /(#.; " Y) into their objectives as a regularization term, since / ( V Y)= 0 can constrain the personalized models to eliminate all shortcut features even though each client has only training environment.
4.3 ALGORITIM DESIGN
In the following contents, we will discuss what the server and local clients need to conduct to develop the optimal personalized invariant predictor få., for each client, Yu [N].
Server Update. Before the algorithm starts, the server initializes the models with random param- eters. At each communication round i, the server firstly selects a fraction of local clients (u E St)

---

Published as a conference paper at ICLR 2024
them. After the selected local clients finish conducting the client update process, the server can receive the local update + and {W +1 i = 1,2, ...., [err] from the selected clients. Then it can update the global solutions by 1,00
and broadeast the current & and {w} | = 1., (E2rl)
Tpi++
|st] Eves: V and w = EuEsi Ceest Py Wht , = 1,2, .... Eerl-
1
P/a
Client Update. Before the algorithm starts, the client initializes the personalized invariant model with random parameters f After receiving the global model y and {wf i = 1, 2, ..., |Err from the server, the local client (V' E (S*) needs to carry on the following two steps: 1) update the personalized invariant model by
S6. -o -nVVE's Du)
for K steps and finally get f6+1 fo. where 1j is the personalized learning rate. 2) The client can conduct R local iterations update the local shortcut extractor. Before it starts, the client initializes the related models Virr= v² and wi = w 1,2, 1,,,,,,,,,, Egyl. During each local iteration
r, the client firstly updates the local shorteut extractor by
whrtt = wh" -V D.)
for one epoch where B denotes the learning rate, and then get the near-optimal environment classi- fiers w/r++1 ; 1, 1,2, ...; I.y. by stochastic gradient descent (on VR(W (VL"") e;)) for T. steps. When completing R local iterations, the client upload the local parameters T = V. and [wftt] = w = 2.., |E to the server for server update.
Due to the space limitation, we place the pseudo-code of the above algorithm and more discussions on the combinations of our method with the prevalent PFL schemes in the appendix.
5 EXPERIMENTS (MORE DETAILS AND RESULTS IN THE APPENDIX
5.1 EXPERIMENTAL SETUP
CoIOTed-MNISTT (CMNIST) (Arjovsky al. 2019) constructed based on MNIST (LeCun al. 1998) via rearranging the images of digit 0-4 into a single class labeled 0 and the images of digit 5-9 into another class labeled 1. Each digit having label 0 is colored green/red with probability p²/1-p" and each digit having label 1 is colored red/green with probability p²/1 - pf, respectively. Thus "color" feature is shortcut in the dataset and the data distribution varies as pr changes. We provide two training environments (pt 0.90 and 0.80) as &er and every local client only has one training environment which is randomly sampled from Etr. To assess the model performance on different test distributions, the test environment each client varies from pf. = 0.00 to 1.00. Considering the heterogeneous data generating process across local clients, the data instances used for constructing the training/test environments each client randomly sampled from only two digit sub-classes labeled O (e.g., digit 1, 2) and two digit sub-classes labeled (e.g., digit 6, 7) without replacement.
Colored Fashion-MNISTT (CFMNIST) (Ahuja al. 2020 is constructed using the same strategy as Colored-MNIS'T, but the original images come from Fashion-MNISTT (Xiao et al.) 2017). Hence, CFMNIST dataset carries more complex feature space than colored-MNIST does.
WaterBird (Sagaaa et al. 2019) considers real-world scenario where the photographs of water- birds usually have water backgrounds while the photographs of landbirds usually have land back- grounds because of the distinct habitats. It makes learning models easily trapped by "background" shortcut when classify "waterbird" and "landbird". In WaterBird, a waterbird is placed onto a wa- ter/land background with probability p²/1-pt and a landbird is placed onto a land/ hand/water background with probability p²/1 pd respectively. We setup two training environments (pf, 0.95 and 0.85) as Ety and each client has only one training environment which is randomly sampled from & The test environment varies from pfc- = 0.00 to 1.00. We notice that the diverse geographic distributions of different bird species naturally accord with the heterogeneity of local data generating process if the federated clients are located in different geographic areas. Considering WaterBird includes 46 waterbird species and 154 landbird species, we distribute 15 (10 separated and 5 overlapped) water- bird species and 51 (34 separated and 17 overlapped) landbird species to each client. The training and test datasets each client contain bird pictures that belong the same bird species.
7

---

Published as a conference paper at ICLR 2024
larger real-world dataset commonly-used in evaluating out-of-distributio (OOD) generalization. It consists of 7 classes distributed across 4 environments (or domains). We adopt the sewee--ne-domainouu strategy to evaluate the OOD generalization performance. Tak- ing personalization into consideration, we split each training domain into two subsets according to classes (i.e., subset consists of dog, elephant and giraffe and another subset consists of guitar, horse, house, and person), and then distribute these two subsets onto two clients respectively. The training and test datasets on each client come from different domains but consist of the same classes.
PACS LL 2017)
Model Selection and Competitors. For CMNIST and CFMNIST, we adopt the deep neural net- work with one hidden layer as feature extractor and an subsequent fully-connected layer as classi- fier. As regard to Waterbird and PACS, ResNet-18 He et al. (2016) is used as the learning model where the part before the last fully-connected layer works as feature extractor and the last fully- connected layer works as classifier. We compare our method (FedSDR) with 10 state-of-the-an algorithms: federated learning methods (FedAvg (McMahan et al., 2017). DRFA (Deng et al.) 2020), FedSR (Nguyen et al. (2022) and FedIIR (Guo al. (023c)), ). and personalized federated learning methods (pFedMe (T Dinh et al. 2020), Ditto (Li al. 202T), I. FTFA (Cheng et al. 2023).
FedRep (Collins al. 2021 FedRoD (Chen & Chao 2022 and FedPAC (Xu et al. 2023]).
Table The overall comparison between performance our method the baselines on four datasets.
Dataset
CMNIST
CFMNIST
WaterBird
PACS
Test accuracy (%)
worst-case
average
worst-case
average
worst-case
average
worst-case
average
FedAvg
3.39
51.03
0.16
50.02
54.13
67.
41.71
47.66
DRFA
21.15
52.81
19.84
53.88
59.75
68.39
42.48
48,95
FedSR
46.93
48.62
47.61
48.90
61.75
71.68
46.76
51.25
FedIRR
47.25
48.39
48.06
49.16
61.24
70.87
47.03
51.58
FTFA
15.42
54.96
11.35
53.52
54.38
69.68
40.89
48.79
pFedMe
21.30
48.53
4.22
51.26
55.63
68.24
45.24
51.33
Ditto
3.02
50.97
0.37
50.12
53.13
68.73
44.95
51.28
FedRep
2.76
50.83
0.11
50.01
52.88
70.23
49.27
53.75
FedRoD
9.09
50.84
1.23
51.57
52.36
70.86
48.16
52.92
FedPAC
1.01
50.05
0.16
50.13
45.08
65.57
49.93
54.20
FedSDR
53.88
55.59
56.92
61.88
65.25
73.20
52.14
56.18
5.2 EXPERIMENTAL RESULTS
Performance Comparison. We summarize the test accuracy of all competitors on different unseen test distributions (11 test distributions in CMNIST, CFMNIST and WaterBird; 4 test distributions in PACS) and figure out the worst-case and average accuracy of each method in Table [ We can find that our method FedSDR consistently outperform the baselines on both worst-case and average test accuracy. In partieular, FedSDR achieves around 6.5%, 9%, 3.5% and 2% higher worst-case accuracy than the second best algorithm and in the meanwhile reaches the highest average accuracy CMNIST, CFMNIST, WaterBird and PACS, respectively.
0.7
0.7
0.33
0.55
1-300
1000
0.59
1900
0.15
022
02. 63 0.6. 05 06 0.7
08 09
022
00
04 02
0.3 0.
35. 06 077
64 09 10
0.30
02
04
04 05
06 01 05
Text anvrrnmmnt (N))
snveanmmntt (85)
"sstenttorment (C))
CMNIST
(b) CEMNIST
WaterBird
Figure 3: The relationship between the test accuracy and the test distribution.
Mitigation of Shorteut Features. Since there exists definite correlation between shortcut features and label in CMNIST, CFMNIST and WaterBird, we can use these three datasets to evaluate how well method can mitigate the shorteut features. The more highly method relies shorteut, the more approximate its test accuracy is to the corresponding pf.. In contrast, a method that eliminates the shorteut can produce consistent test accuracy across different pfc- We evaluate the competitors
8

---

Published as a conference paper at ICLR 2024
under diverse test environment (i.e., Pie) and show the relationships between test accuracy and pie in Figure 3 In particular, "Oracle" represents the scheme where we manually remove the shortcut features (color in CMNIST and CFMNIST; background in WaterBird) from the whole dataset and then train the personalized models using the pre-processed dataset. Hence "Oracle" provides an ideal performance for comparison. We can see that FedSDR can effectively mitigate the shortcut and achieve a more consistent test accuracy than most of the FL and PFL methods. Because FedSDR exploits the personalized invariant features, it consistently achieves higher test accuracy than the federated domain generalization methods which drop the personalization information.
Table 2: Performance comparison between FedSDR and the trivial combination of IL with PFL schemes.
Dataset
CMNIST
CFMNIST
WaterBird
PACS
Test accuracy (%)
worst case
average
worst case
average
wonst-case
average
worst case
averaaee
IRM
46.38
49.14
47.76
49.41
60.38
68.63
46.35
50.83
IRMI
14.32
54.27
11.09
53.48
60.25
69.46
50.04
IRM -L2
45.68
49.04
17.92
19.46
61:25
68.93
18.37
51.98
FedSDR
$3.88
55.59
56.92
61.88
65.25
73.20
52.14
56.18
Necessity of Shortcut Discovery and Removal. At the beginning of Section 4 we analyse that trivial combination of invariant learning scheme with local adaptation (commonly used in PFL) can fail to generate the optimal personalized invariant predictors for local clients. Too validate the superiority of our method on developing the personalized invariant predictors when local training environments are insufficient, we implement two typical personalization skills with the global model being trained by the distributional version of IRM (i.e., IRM in 2 One is L2-norm regularizer
used in PFL (T Dinh et al. 2020 Hanzely & Richtárik 2020; Hanzely et al., 2020 Li et al. 202
and we call this implementation IRM Another one is local Fine- Fine-Tuning Tuning which is proved simple and effective for personalization (Cheng et al. (2023) and we name it IRM
From the results in Table 2 we can find the combinations can hardly improve the OOD general- ization performance. In particular, the local fine-tuning skill can even degrade the performance, compared with baseline IRM The underlying reason is that local adaptation can readily make the personalized model pick up the shortcut features when local training environments By contrast, our shortcut removal method is independent environment and can effectively mitigate the shortcut features even though there is only one training environment each client.
insufficient.
Effect of Hyper--prramete À and y. We evaluate the effects of two significant hyper-parameters in the proposed objective (i.e., À and ?) on model per- formance here. Since the results other datasets present the similar tendency as on WaterBird, we herein focus on WaterBird. The results on other datasets are placed in the appendix. When evaluat- ing the effect of X, we fix 7 = 1.4. When evaluat-
Table 3: Performance of FedSDR on WaterBird with different values of hyper-parameters À and 7.
0.00
0.10
0.50
1.00
10.0
worst-case (%)
61.88
62.51
65.25
61.68
61.74
average (5)
71.34
72.39
73.20
70.61
70.18
0.00
0.10
1.00
1.40
10.0
ing the effect of 7 we fix À 0.5. The results are
worst-case (%)
43.75
44.16
57.64
65.25
48.29
average (%)
66.30
65.73
70.18
73.20
61.86
shown in T When À 0.0, shorteut feature extractor is trained by empirical risk minimization (i.e., ERM). When
0.0, the personalized
models are trained by local ERM. Because models trained by ERM tend to rely on shortcut, the performance of FedSDR is more sensitive to the selection of y than the selection of A.
6 CONCLUSION AND FUTURE WORK
In this paper, we study the challenging shortcut trap problem in PFL. We formulate the SCMs to interpret the heterogeneous data generating mechanisms on federated clients and derive two signif- icant causal signatures which inspire our provable shorteut discovery and removal method. Theo- retical analysis proves the proposed FedSDR can draw all shortcut features and elicit the optimal personalized invariant predictor that can generalize unseen target data for each client. FedSDR can cooperate with most of the existing PFL methods to improve their OOD generalization perfor- mance, which can facilitate the real-world application of PFL. Since the theoretical guarantee is derived under linear cases in this paper, we will extend it to more complex cases in the future work.

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGMENTS
This research was supported by fundings from the Key-Area Research and Development Program of Guangdong Province (No. 302180101400003), Hong Kong RGC Research Impact Fund (No. R5060- 19, No. R5034-18), Areas of Excellence Scheme (AoE/E-601/22-R) General Research Fund (No. 152203/20E, 152244/21E, 152169/22E, 152228/23E).
REFERENCES
Kartik Ahuja, Karthikeyan Shanmugam, Kush Varshney, and Amit Dhurandhar. Invariant risk min- imization games. In International Conference on Machine Learning, pp. 145-155. PMLR, 2020.
Kartik Ahuja, Ethan Caballero, Dinghuai Zhang, -hrrisistt Gagnon-Audet, Yoshua Bengio, Ioannis Mitliagkas, and Irina Rish. Invariance principle meets information bottleneck for out-of- distribution generalization. Advances in Neural Information Processing Systems, 34:3438-233 3450, 2021.
Martin Arjovsky, Léon Bottou, Ishaan Gulrajani, and David Lopez-Paz. Invariant risk minimization. arXiv preprini arXiv: 00000000099.... 2019.
Sara Beery, Grant Van Horn, and Pietro Perona. Recognition in terra incognita. In Proceedings of the European conference on computer vision (ECCV), pp. 456-473, 2018.
Hong- You Chen and Wei-Lun Chao. On bridging generic and personalized federated learning for image classification. International Conference on Learning Representations, 2022.
Yimeng Chen, Ruibin Xiong, Zhi-Ming Ma, and Yanyan Lan. When does group invariant learning survive spurious correlations? Advances in Neural Information Processing Systems, 35:7038- 7051, 2022a.
Yining Chen, Elan Rosenfeld, Mark Sellke, Tengyu Ma, and Andrej Risteski. Iterative feature matching: Toward provable domain generalization with logarithmic environments. Advances in Neural Information Processing Systems, 35:1725-1736, 2022b.
Gary Cheng, Karan Chadha, and John Duchi. Federated asymptotics: a model to compare federated learning algorithms. In International Conference on Artificial Intelligence and Statistics, pp- 10650-10689, PMLR, 2023.
Liam Collins, Hamed Hassani, Aryan Mokhtari, and Sanjay Shakkottai. Exploiting shared repre- sentations for personalized federated learning. In International Conference Machine Learning. pp. 2089-2099. PMLR, 2021.
Elliot Creager, Jörn-Henrik Jacobsen, and Richard Zemel. Environment inference for invariant learning. In International Conference Machine Learning, pp. 2189-2200, PMLR, 2021.
Yuyang Deng, Mohammad Mahdi Kamani, and Mehrdad Mahdavi. Distributionally robust federated averaging, Advances Neural Information Processing Systems, 33:15111-15122 2020.
Aymeric Dieuleveut, Gersende Fort, Eric Moulines, and Geneviève Robin. Federated-em with het- erogeneity mitigation and variance reduction. M. Ranzato, A. Beygelzimer, Dauphin, P.S. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, vol- ume 34, pp. 29553-29553 292636.9966... Curran Associates, Inc., 2021.
Alireza Fallah, Aryan Mokhtari, and Asuman Ozdaglar. Personalized federated learning with the- oretical guarantees: A model-agnostic meta-learning approach. Advances in Neural Information Processing Systems, 33:3557-3568. 2020.
Robert Geirhos, Jömn-Herrik Jacobsen, Claudio Michaelis, Richard Zemel, Wieland Brendel, Matthias Bethge, and Felix A Wichmann. Shortcut learning in deep neural networks. Nature Machine Intelligence, 2(1 (1:1)))6666777 2020.
10

---

Published as a conference paper at ICLR 2024
Tao Guo, Song Guo, and Junxiao Wang. Pfedprompt: Learning personalized prompt for vision- language models in federated learning. In Proceedings of the ACM Web Conference 2023, pp- 1364-1374, 2023a.
Tao Guo, Song Guo, Junxiao Wang, Xueyang Tang, and Wenchao Xu. Promptfi: Let federated participants cooperatively learn prompts instead of models federated learning in age of foundation model. IEEE Transactions on Mobile Computing, 2023b.
Yaming Guo, Kai Guo, Xiaofeng Cao, Tieru Wu, and Yi Chang. Out-of-distribution generalization of federated learning via implicit invariant relationships. In Proceedings of the 40th International Conference on Machine Learning, pp. 19905-11933, 2023c.
Filip Hanzely and Peter Richtárik. Federated learning of mixture of global and local models. arXiv preprint arXiv:2002.055/6 2020.
Filip Hanzely, Slavomír Hanzely, Samuel Horváth, and Peter Richtarik. Lower bounds and opti- mal algorithms for personalized federated learning. Advances in Neural Information Processing Systems, 33, 2020.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp- 770-778, 2016.
Kevin Hsieh, Amar Phamishayee, Onur Mutlu, and Phillip Gibbons. The non-iid data quagmire of decentralized machine learning. International Conference on Machine Learning, pp. 4387- 4398. PMLR, 2020.
Bo-Wei Huang, Keng-Te Liao, Chang-Shene Kao, and Shou-De Lin. Environment diversification with multi-head neural network for invariant learning. Advances in Neural Information Processing Systems, 35:915-922 2022.
Yutao Huang, Lingyang Chu, Zirui Zhou, Lanjun Wang, Jiangchuan Liu, Jian Pei, and Yong Zhang. Personalized cross-silo federated learning on non-iid data. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pp. 7865-7873, 2021.
Dongsung Huh and Avinash Baidya. The missing invariance principle found - the reciprocal twin of invariant risk minimization. In S. Koyejo, Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (eds.), Advances in Neural Information Processing Systems, volume 35, pp. 23023-23035. Curran Associates, Inc., 2022.
Nam Hyeon-Woo Moon Ye-Bin, and Tae-Hyun Oh. Fedpara: Low-rank hadamard product for communication-effiieettt federated learning. In International Conference on Learning Represen- fations, 2022.
Wonyong Jeong and Sung Ju Hwang. Factorized-ff: Personalized federated learning with parameter factorization similarity matching. In Advances Neural Information Processing Systems, 2022.
Yibo Jiang and Victor Veitch. Invariant and transportable representations for anti-causal domain shifts. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (eds.), Advances in Neural Information Processing Systems, volume 35, pp. 20782-20794, Curran Associates, Inc., 2022.
Sai Praneeth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank Reddi, Sebastian Stich, and Ananda Theertha Suresh. SCAFFOLD: Stochastic controlled averaging for federated learning- In Hal Daumé III and Aarti Singh (eds.), Proceedings of the 37th International Conference Machine Learning, volume 119 of Proceedings of Machine Learning Research, pp. 5132-5143. PMLR, 13-18 Jul 2020.
Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient-based learning applied to document recognition. Proceedings of the IEFF, 86(11 (6(11):2778-224- 1998.

---

Published as a conference paper at ICLR 2024
Da Li, Yongxin Yang. Yi-Zhe Song. and Timothy M Hospedales. Deeper, broader and artier domain generalization. Proceedings of the IEEE international conference on computer vision, pp- 5542-5550, 2017. 2017.
Tian Li, Shengyuan Hu, Ahmad Beirami, and Virginia Smith. Ditto: Fair and robust federated learning through personalization. In International Conference on Machine Learning, pp. 6357- 6368. PMLR, 2021.
Yong Lin, Shengyu Zhu, Lu Tan, and Peng Cui. Zin: When and how to learn invariance without environment partition? Advances in Neural Information Processing Systems, 35:24529-22542, 2022.
Jiashuo Liu, Zheyuan Hu, Peng Cui, Bo Li, and Zheyan Shen. Heterogeneous risk minimization. In International Conference on Machine Learning, pp. 6814-6814. PMLR, 2021a.
Jiashuo Liu, Zheyuan Hu, Peng Cui, Bo Li, and Zheyan Shen. Integrated latent heterogeneity and invariance learning in kernel space. Advances in Neural Information Processing Systems, 34: 21720-21733 2021b.
Quande Liu, Cheng Chen, Jing Qin, Qi Dou, and Pheng-Ann Heng. Feddg: Federated domain gener- alization on medical image segmentation via episodic learning in continuous frequency space. In Proceedings of the IEEE/CVF Conference on Computer Vision and Puttern Recognition (CVPR). pp. 1013-1023, June 2021c.
Chaochao Lu, Yuhuai Wu, José Miguel Hemmáneez-Lobato, and Bernhard Schölkopf. Invariant causal representation learning for out-of-distribution generalization. International Conference Learning Representations, 2022.
Zhengquan Luo, Yunlong Wang, Zilei Wang, Zhenan Sun, and Tieniu Tan. Disentangled feder- ated learning for tackling attributes skew via invariant aggregation and diversity transferring. In International Conference on Machine Learning, pp. 14527-1452 PMLR, 2022.
Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. Communication-effiieenn learning of deep networks from decentralized data. In Artificial intelli- gence and statistics, pp. 1273-1282. PMLR, 2017.
A Tuan Nguyen, Philip Torr, and Ser-Nam Lim. Fedsr: A simple and effective domain generalization
method for federated learning. In Advances in Neural Information Processing Systems, 2022.
Judea Pearl. Causality. Cambridge university press, 2009.
Jonas Peters, Peter Bühlmann, and Nicolai Meinshausen. Causal inference by using invariant pre- diction: identification and confidence intervals. Journal of the Royal Statistical Society. Series (Statistical Methodology), pp. 947-1012, 2016.
Elan Rosenfeld, Pradeep Ravikumar, and Andrej Risteski. The risks of invariant risk minimization. In International Conference on Learning Representations, volume 9, 2021.
Shiori Sagawa, Pang Wei Koh, Tatsunori B Hashimoto, and Percy Liang. Distributionally robust neural networks for group shifts: On the importance of regularization for worst-case generaliza- tion. arXiv preprint arXiv: /9/ 1.08311, 2019.
Pranay Sharma, Rohan Panda, Gauri Joshi, and Pramod Varshney. Federated minimax optimiza- tion: Improved convergence analyses and algorithms. In International Conference Machine Learming, pp. 19683-19730. PMLR, 2022.
Virginia Smith, Kai Chiang, Maziar Sanjabi, and Ameet S Talwalkar. Federated multi-task learning, Advances neural information processing systems. 30, 2017.
Zhenyu Sun and Ermin Wei. A communication-effici algorithm with linear convergence for federated minimax learning. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho. and A. Oh (eds.). Advances Neural Information Processing Systems, volume 35, pp. 6060-6073. Curran Associates, Inc., 2022.
12

---

Published as a conference paper at ICLR 2024
Canh T Dinh, Nguyen Tran, and Josh Nguyen. Personalized federated learning with moreau velopes. Advances in Neural Information Processing Systems, 33:21 2911121444405 2020.
Xueyang Tang, Song Guo, and Jingcai Guo. Personalized federated learning with contextualized generalization. In Lud De Raedt (ed.), Proceedings of the Thirty-First International Joint Con- ference on Artificial Intelligence, 1JCAI-22, pp. 2241-2247. International Joint Conferences on Artificial Intelligence Organization, 7 2022. Main Track.
Haoxiang Wang, Haozhe Si, Bo Li, and Han Zhao. Provable domain generalization via invariant- feature subspace recovery. In International Conference on Machine Learning, pp. 0-003323003.. PMLR, 2022.
Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-mnist: a novel image dataset for benchmark- ing machine learning algorithms. arXiv preprint arXiv: 2017
Jian Xu, Xinyi Tong, and Shao-Lun Huang. Personalized federated learning with feature alignment and classifier collaboration. In The Eleventh International Conference on Learning Representa- tions, 2023.
Jie Zhang, Song Guo, Xiaosong Ma, Haozhao Wang, Wenchao Xu, and Feijie Wu. Parameterized knowledge transfer for personalized federated learning. Advances in Neural Information Process- ing Systems, 34: 10094-1104, 2021.
Jie Zhang, Zhiqi Li, Bo Li, Jianghe Xu, Shuang Wu, Shouhong Ding, and Chao Wu. Federated learn- ing with label distribution skew via logits calibration. In International Conference Machine Learning, pp. 26311-26233 PMLR, 2022.
A RELATED WORK
Due to the space limitation, more complete literature review is placed in this section.
Federated learning, The classic FedAvg (McMahan et al. datasets are IID. Some methods (Karimireddy et al. (2020): Dieuleveut et al. (202T): Zhang et al.
2017) performs well if local training
(022); Guo et al. (2023b)) mitigate the negative impact of training data heterogeneity on conver- gence rate, while another branch (Deng et al.] (2020); Sharmma al.) (2022); (un & Wei (2022) targets at reducing the performance bias of global model on local clients. Besides, few works (Liu al. 2021c);; Nguyen al. (2022): Guo et al. (2023c) investigate the scenarios where the training data heterogeneity appears to be domain shift. All the above methods produce a shared global model which can diverge from the local optimal solutions when local target datasets are Non-IID.
Personalized federated learning. Many PFLs [ Dinh et al. (2020); Hanzely al. (2020);
lah et al. 22020); Li et al. (202T): Tang et al. 2022); Cheng et al. (2023); Guo et al. (2023a) train the personalized models with the guidance of global model which embeds in the shared knowl- edge. Some researchers study the parameterized knowledge transfer between similar clients, c.g.,
MOCHA (Smith et al. 2017). FedAMP (Huang et al., 2021) and KT-pFI. (Zhang al., 2021 DFL (Luo et al. 2022) disentangles the shared features from the client-specific ones to achieve shared knowledge. Similarly, pFedPara (Hyeon-Woo et al.) 2022) and Factorized-FL (Jeong & Hwang 2022) factorizes the model parameters into the shared and per- sonalized parts. Another branch (Collins et al. (202T); Chen & Chao (002);; Xu et al. (2023) employs the shared/aligned feature extractor to capture global knowledge and personalized classi- fiers to encode the personalization information. All of them don't on'' cover the situations where there exists shortcut in local training datasets.
accurate aggregation
Shortcut and Invariant learning (IL.) Causally invariant predictor is proposed (Peters et al., 2016), and then applied into deep learning in IRM (Arjovsky et al. 2019) to mitigate shortcut. Subsequently, (Rosenfel et al. (202I) prove that IRM and its variants can be still trapped by shortcut when training environments are insufficient. IFM (Chen et al. 2022b) lowers the requirement and demands only logarithmic training environments. Some works focus on settling IL problem when the environment label is unavailable, e.g., EIIL (Creager al. 2021), HRM (Liu et al.) 202 b), EDNIL (Huang et al. (2022) and ZIN (Lin et al. 2022). Another branch (Ahuja et al. (2021): Chen al. (2022a): Huh & Baidya (022)) completes the constraints that IRM misses to improve the
13

---

Published as a conference paper at ICLR 2024
performance. The iCaRL (Lu al. 2022
extends IL to non-linear causal representations while ACTIR (Jiang & Veitch 2022) extends IL to anti-causal scenarios. All these methods are devised for centralized scenarios where all training data is accessed and training environments are sufficient.
B COMPLETE THEORETICAL PROOFS
In this section, we will provide the complete proofs of the theorems stated in the main text.
B. PROOF LEMMA
U
U
E
ZZ
Z
7
(a) Causal FedSDR
(b) Anti-causal FedSDR
Figure 4: The structural causal models (SCMs) considered in this paper.
From the above structural causal models, we derive the following two useful causal properties.
Lemma 1
If the data generating mechanism of each federated client obeys the causal graph in
Figure ((a) or the anti-causal graph in Figure 4(b) we can have:
Zs L U 111 Y E which means that the shortcut features Zs are conditionally independent of the personalization indicator U given Y and E.
Zz 1 Zg Y and Zy Zs Y, which means that both the global (Z&) and personalized (Z() invariant features are conditionally independent of the shortcut features Zg given
Proof. According to the causal Markov condition (Theorem 1.4.1) proved in (Pearl 2009), we know that the variable Zs is independent of all its nondescendants, given its parents the (Markov) causal graph. Since Y and E are the parent variables of Zs and U is nondescendant of Zs. the first causal signature Lemma i is guaranteed. Besides, based on the d-separation criterion (Pearl) 2009) can find the variable Y d-separates Zg: from Zs and d-separates Zy from Zs in the SCMs. Therefore, we get the second causal signature in Lemma
B.2 PROOF OF THEOREM 1
If the Assumption holds, the function g is linear and the total number of training environments in the federated learning system satisfies EEy, >> de then the following two statements are equivalent:
Theorem
(X) depends and only depends the complete shorteut features Zs;
1²* is the optima of the 5q.5 with an appropriately chosen hyper-parameter X.
Proof. We write the linear feature extractors that can recover the latent features ([Zc, Zsl) from the observation X & (X) (( Zc, Zs)) AZc + BZs, where A and B are fixed transformation matrices. This formulation is also adopted in the theoretical analysis in (Rosenfeld et al. 2021) and /Wang al. 2022). For the concerned logistic regression,
get closed form
14

---

Published as a conference paper at ICLR 2024
for the distribution P(Y| v) as:
P(Y V := (Y AZC + BZs)
PO(AZc + BZs (Y)
IPT (AZC + BBss)
P (-AZc- + BBs (Y)
Σy Pe (Y - -y)P=(AZe + BZs Y y)
Since Assumption holds, we have PY 1/) (V=(v=))),
= =V), E EY. We can derive that
(AAc + BZg 18)
Ply | AZc + BZs) = Ey Pa -Bss Y yzy)
IP IPa + B7g 11) P'AA + BZs Y y
y) + P (AZc + BZg Y = =-yy)
Vy E f+1}. BIxl0=BBslY=-2 1+ P'Ace BZs
Because we have Zc 1 Zg Y from Theorem 1, we can get the probability density of AZC + B7s as follows:
AZc BZs y NNy M2 Mx),
(8)
where μ2= = Atto + Bfg and Σ, = AAT 02 BBB o3. Thus, we can get the P(Y | V as:
pe (y v
Pc AAZ BZ
1 + F (AZc+ B2x|
E
1+ exp(-x y. 2N 14)
where E-1 represents the generalized inverse of Eg, i.e., 1= I.
According to Lemma F.2. proved in the appendix of (Rosenfeld et al. 2021), the optimal clas- sifier based on the feature extractor (XX AZc + BZg is sufficiently and necessarily given
by 2(AATo? + BBT d?) (Aya + Bys). That is, we have Pu: v)) 1+0xp(=9/27T5; m.))
Vy E {=1}, if and only if Wf E arg min En ((; Ve: E Er.
Therefore, we can calculate the KL-divergence between Put (Y V e;) and (Y v,,e) by
KC((P (Y V. e:) |F4 (Y (.
Pwt (v] | v e,) log Pw; (y) V vz Pug (y V e,e))
Σ =E(+1)
=
exp( y- yy1--y--777 pp)
Σ ve(1l}
log
1 + exp(p-1 24' Ei kc)
exeeep(X (ppp--y-- y 24/ Dz. w2)
1 exp( 220 x.g1 (
log
1 exp(-22p -207 Σz, (1)
1 + exp( -2V/ kc)
+ (xpp2v/T -
+
log
1 + exp(2V''' Dーe k!}
+ exp 22 m$n
1
+ exp(21) Tx mg)
cxp(2NV 1=)
log
+ log
+ 1xpp-22/700 2 Nー. k$)
2xp(222 2i, mو)
exp دوح vz)
1 + exp(2v/ 2 23 " (u{) + exp (2 (xp T3, mp1
log
1 + exp(24 TE m
1 + 1eep(22V Σラ (1d
24 (2 //= Σi (0!
log
+
1 cxxp22 m1)
1 + exp( 25 m1)
15

---

Published as a conference paper at ICLR 2024
Similarly, get that
(L (Y ee (Y v,e)))
Puy v,eg)
E Puy (y U ej) log NE(=1)
Puy (y) es)
exp E 1b) 1 exp (2 Σi uz)
24 ( mr Σー kC)
= log
1 1 exp( -2V exx((-2 p²2)
Combining the above results, we can get that (F (Y (,,ea) (Y e,,))) KL (Y V ( (V ei)))
1
{2 (5 μ En {m}
1 1+ cxp( - k$) + e k1)
T2
Ti
2 0, Vei,e E Eau, VV E H.
Since the absolute value of term Tj (i.e., (T)) monotonically increases with term [Tp] increasing, the objective maxy (CC(P (Y) W, [.e.e) (Y V.e,)) (FC (Y V.e,) e,) PP: (Y V, e4)) is equivalent to maxy |2 uh - اg wz)ll Therefore, the second term & EN laie (V; Dw) Eq.5 can be written as
Lai ( v; Dw) Ev
Σ
Σ
12 -i, m3)l
1=1
ecEEe EfEer
4l[(Ant Bb(() - (Ape + B E
Σ
Σ
(AAT o? oB)
e EfEler
B - p1)
Σ
Σ
Ey V 2
(AATo? + BB o?)²
According to the mentioned AZc + BZs ~~ NNy M1 E.), we can get the density
P(VV)
Σ P(Y PP Y y) DEY
With the Assumption holding, can get the mean E T BBT o3. Therefore, we have
0 and the variance D[V] = AAT o² +
N
1
Lase (I; Da) = Ev
Σ
|2 (Σ - Σm
-
N
B|| 22) 12
Σ
(D (E(
(AA' 02 + BB o3
e,EE EfEEsr
B =2)
Σ
Σ
AAT 0² + BBT
E.E EFEEtr
ΣΣ e,EEL CyEE
ll -1212
AA BET
0Z + 0z
We can find that maximizing the above objective will make A = 0 and BB" + 0. Moreover, when |Eell d,, maximizing Ee,ee.. Ee, E&,. ||, ugll will make rank (B) d,. In the meanwhile, satisfying A 0 and rank(B) d, will in turn maximize the objective | N Eais (v; D.)-
In Eq.5, we utilize a Lagrangian multiplier to solve the constrained optimization and the balancing weight is X. Therefore, Theorem 2 gets proved.
16

---

Published as a conference paper at ICLR 2024
B.3 PROOF OF THEOREM? 2
Theorem 2 Suppose V (X) ) in the E depends and only depends on the complete shortcut fea- tures Zz. If Sä. (Vu E [N] is the optima of the Eq.7 7 with the hyper-purameter ? chosen appropri- ately, then the fo. the optimal personalized invariant predictor for the client u, 'u [N).
Proof. We know that minimizing (P,) Du) is the sufficient condition of maximizing (x))), and ((. V (Y = 0 is equivalent to Py 1 V Y. According to the prop- erty of Lagrangian multiplier, the objective in E is equivalent to the constrained optimization where the constrain is u y Y) = O, with the appropriately chosen 7. Combining with the second causal signature in Theorem 1, Theorem 3 gets proved.
C MORE DETAILS OF EXPERIMENTS
In this section, we will include more detailed setups and discussions on the evaluation part. Code is available at
C.I NON-IID DATA PARTITION
For CMNIST and CFMNIST datasets, we provide two training environments (pirain = 0.90 and 0.80) Etrain and every local client only has one training environment which randomly sam- pled from the training environment set Etrain- To assess the model performance on different test distributions, the test environment on each client varies across Piest 0.00, 0.10, .... 0.90, 1.00. Considering the heterogeneous data generating process across local clients, the data instances used for constructing the training/test environments on each client are randomly sampled from only two digit sub-classes (1 separated and 1 overlapped) labeled 0 and two digit sub-classes (1 separated and overlapped) labeled without replacement. Specifically, we totally simulate eight local clients and one server in the federated learning system. For example, the data instances on client I are randomly sampled from digit O, 1, 5, 6; the data instances on client 2 are randomly sampled from digit 1, 2, 6, 7; the data instances client 3 randomly sampled from digit 2, 3, 7, 8; and the data instances on client 8 are randomly sampled from digit 3, 4, 8, 9.
As regard to WaterBird, distribute 15 (10 separated and overlapped) waterbird species and 51 (34 separated and 17 overlapped) landbird species to each local client. Both the training and test data instances are constructed using bird photographs randomly sampled from the corresponding bird species in the bird dataset and background photographs randomly selected from the background dataset without replacement. Similarly, we totally simulate eight local clients and one server in the federated learning system.
PACS consists of classes (i.e., dog, elephant, giraffe, guitar, horse, house, and person) distributed (i.e., Art Painting, Cartoon, Photo and Sketch). We adopt the "leave-
across 4 domains/envirommenn one-domain-oute strategy to evaluate the example, when we evaluate the performance on Art Painting domain, we use the remaining three domains (i.e., Cartoon, Photo and Sketch) as training environments. Taking personalization into consideration, split each training domain into two subsets according to classes (i.e., one subset consists of dog, elephant and giraffe and another subset consists of guitar, horse, house, and person), and then distribute these two subsets onto two clients respectively. The training and test datasets on each client come from different domains but consist of the same classes.
(OOD) generalization performance. For
C.2 HYPER-PARAMETER
The hyper- papee--arramterssccom of the competitors and our algorithm are tuned to make the accuracy on the validation environment (i.c., pf..t = 0.10) as high as possible. Specifically, the mainly used hyper- parameters in the evaluation part are listed as follows: Global communication round: T = 600,
Local iterations: R = 10, Personalized epochs to update the personalized invariant predictors:
K 10, Local batch size: B 50. Global learning rate: 3 0.0001, Personalized learning
rate: 17 = 0.0001, Diserepancy threshold: a = 1.0, Balancing weight: = 0.5, Balancing weight:
7 1.4, Optimizer: Adam.
17

---

Published as a conference paper at ICLR 2024
C.3 IMPLEMENTATION
Besides, the experiments are implemented in PyTorch. We simulate set of clients and a centralized server one deep learning workstation (Intel(R) Core(TM) 19-12900K CPU @ 3.20GHz with one NVIDIA GeForce RTX 3090 GPU).
C.4 ALGORITHMS
Algorithm FedSDR: Federated Learning with Shortcut Discovery and Removal
Input: T, R. K, 3, 77, a, y.
1: Initialize the models Vo, {weli E [|,,, 1b. (8, Ju [N]]
2: for = to T- do
3: Server sends the global models (w² fwili E [[&,)]]] to the participating local clients.
4: for local device
N in parallel do
5: Initialization: 1p.00 <- Tp', [wf, + wili E [[)
6: for k
K do Update the personalized invariant model: ft,b+1 = fr. nnVeR
7:
C.., D.).
f8. f1+1.0 + f
8:
9:
for t to do
10:
Update the shortcut extractor: 11r = v.r BVEDD (Vh;; D.).
11:
Update the environment classifiers for K epochs with VR(W (wt"))
12:
Randomly select a subset (S') of the users to upload the local approximation:
13:
uett1 + "r. and {w]t+ + w = 1, 1,2, 1 l.
14:
Global aggregation: shortcut extractor T [ST Evesi Th++1 and environment classi-
P'_ fiers wi+1 = Eves: Exest Lia i = 1, 1,2,... (Eer)-
return the personalized invariant models ( lu [N]].
the evaluation part, the data distributions among the local clients only overlap slightly. Therefore, the personalized invariant models are trained and updated locally and do not participate in the global aggregation. The corresponding algorithm is shown in Algorithm [
Algorithm FedSDR (+FedAvg): Federated Learning with Shortcut Discovery and Removal
Input: T, R. K, 3, 7, a, 7,
1: Initialize the models 1/, f and (wfli [[Etrl]].
2: for to T 1 do
3: Server sends the global models (V' ff and (wili E [[E))]]]) to the participating local clients.
4: for local device u
N in parallel do
5: Initialization: T &- v.. fo. t fg and fwf, < wili [8
6: for r = to R do
7:
Update personalized invariant model: pt,++1 = 1/1 - V D.). Update the shorteut extractor: T. = Th!." - (Thr; Da).
8:
9:
Update the environment classifiers for K epochs with VR(W (R( (s))
10:
Randomly select subset (S') of the users to upload the local approximation:
11:
wi+1 + ".. ff + st.R and (w/+1 11 + w...t i== 1, [8
Global aggregation: the shorteut extractor 4e+1 1 Enes: V++1 the environment
pl 1+1 classifiers w/+1 Enese Evest t
1,2,.., |Erl and the global invariant model
st+1 1 Enest s++1
return the global and personalized invariant models f/, UTT." [N]].
As mentioned in the main text, the proposed shortcut discovery and removal method can easily cooperate with most of the existing federated and personalized federated learning method to improve
18

---

Published as a conference paper at ICLR 2024
the out-of-distribution generalization performance via adding shorteut discovery and removal regularization term. We provide an example combination FedSDR with FedAvg (MMMahan et al. 2017) and PFedMe (T Dinh et al. 2020) in Algorithm 2 and Algorithm 3 respectively. Of course, the combinations with more other federated and personalized federated learning methods can be explored in the future.
Algorithm 3 FedSDR (+pFiedMe): Federated Learning with Shortcut Discovery and Removal
Input: T, R, K,,,K,,,,, 77, at 7,
1: Initialize the models V S8 and (w]li [[
2: for I =0 T- do
3: Server sends the global models (V' fgz and {wilit E [[&r)]]] to the participating local clients. 1 to N in parallel do
4: for local device
5:
Initialization: T0.0 +- wp fo fg and (wf, < wfli [8%1]. for r = 0 to do
6:
7:
for k = 0 to K 1 do
8:
Update the personalized invariant model:
9:
f0 pn pr,k
Sn kr Du) + 7 pt.ck
f6()
10:
Update the global invariant model: f gt.r+1 = fo
BT (fh pr.Kk
11:
Update the shortcut extractor: 1pr,r+1 wi."" WeED (Vrr Da). Update the environment classifiers for K epochs with VR(w (wtr))
12:
13:
Randomly select subset (S*) of the users to upload the local approximation:
14:
pt.R and (witt) ve(+1 +- V f +- Se.
R.
1,2, ..., E
Global aggregation: the shortcut extractor di+11 |st| Evesi D the environment
classifiers w Enes' Evest m wha 1,2,...,/Err and the global invariant model
oi
1241 1ST Eves: S
16: return the global and personalized invariant models fJ URRK K lu E [N]).
C.5 MORE EXPERIMENTAL RESULTS
Experiment on synthetic dataset To verify the provided theoretical guarantees under linear cases, we generate synthetic dataset using the same strategy as in (Rosenfeld et al. (021). Specifically, it is a logistic regression task and the data instance X is generated by X g( Zi Zú Zs). where the dimensionalities of Zg. z and Zş are di 3, de function g is implemented by one fully-connected layer which has 12 neurons. The latent variables zg. zy and Zs are subject N(y He.g. 22.g() (y Me, 2 I) and and NN (4 o22) respectively. Target variable y is taken from the distribution P(y -1) PP 1) 0.5. Both He.g and He. are randomly sampled from (N 1.57) while μs is randomly sampled from A/ 0.75I. To make the shorteut representation Zs easier to learn, we choose oe.g Te.u = and o, as in (Rosenfeld et al. 202T). Each fixed value of иı indicates one specified environment. We generate 10 training environments and 5000 test environments generalization performance. Each (training/test) environment contains 10000 data samples (X, v) and the training data samples are distributed onto totally 100 clients. The training and test data samples each client are generated with an identical value of Me.a. Besides, we choose the client sampling rate as 0.1. The experimental results on this synthetic dataset are shown in Table In
and dg 6 respectively. The linear
evaluate the out-of-distributioo
Table 4: The performance FedSDR and the competitors on the synthetic dataset.
Algorithm
FedAvg DRFA FedSR FeduR FTFA
PFedMe
Ditto
FedRep FedRoD FedPAC FedSDR
worst-case(%)
3.06
62.41
63.09
67.39
1.32
10.58
7.76
2.57
21.53
8.98
92.49
average(%)
85.56
69.64
70.53
70.75
96.26
95.72
96.50
97.80
97.24
98.77
96.07
particular, when we manually select the causal features (8. as the discriminating features, find the optimal personalized classifiers achieve stable accuracy around 97.5 in different test
19

---

Published as a conference paper at ICLR 2024
environments. Therefore, the results shown in Table can demonstrate the effectiveness of our FedSDR on developing the optimal personalized invariant predictors, compared with the state-of- the-art FL and PFL methods.
Scalability of FedSDR In the evaluation part of the main text, we simulate clients in the ex- periments on CMNIST, CFMNIST and WaterBird. The experiments on PACS are conducted on clients. To further evaluate the scalability of FedSDR, we firstly partition CMNIST dateset into subsets using the same strategy adopted to simulate 8 clients. And then, we randomly distribute each subset onto clients. In this way, we totally construct 80 clients for CMNIST dataset. Similarly, we construct 80, 80, 60 clients for CFMNIST, WaterBird and PACS respectively. When evaluating the model performance on these four datasets, we adopt a client sampling rate of 0.1. The experimental results shown in Table:
Table 5: The overall comparison between the performance of our method FedSDR and the baselines with large number of clients.
Dataset
CMNIST
CFMNIST
WaterBird
PACS
Test accuracy (%)
worst-case
average
worst-case
average
worst-case
average
worst-case
average
FedAvg
1.74
46.82
0.77
45.62
48.65
61.57
33.75
40.18
DRFA
14.94
47.24
15.51
47.14
52.34
60.43
36.17
41.75
FedSR
40.29
43.64
41.16
43.27
55.63
64.32
39.03
43.40
FedIIR
41.18
42.93
41.80
43.58
54.31
64.60
40.15
44.37
FTFA
11.51
49.28
7.20
47.57
50.25
63.39
34.65
42.19
PFedMc
17.28
44.13
2.42
47.95
50.01
61.97
41.06
45.84
Ditto
1.98
45.84
1.80
45.71
49.08
63.38
10.18
46.30
FedRep
1.56
46.20
16.14
48.12
64.52
42.16
17.58
FedRoD
6.53
16.86
1.60
47.43
49.56
65.49
42.68
16.61
FedPAC
0.38
45.64
0.23
44.88
42.61
63.81
44.19
49.71
FedSDR 50.41
51.85
52.81
57.14
59.96
68.09
48.07
51,55
The results show that FedSDR can still outperform the competitors when there large number of clients in the federated learning system, which can validate the scalability of the proposed FedSDR.
More invariant learning methods We consider combinations of the personalized federated learn- ing schemes with more sate-of-thheaat invariant learning methods, including IRM-IB (Ahuja al. 2021) and MRI (Huh & Baidya 2022). The experimental results shown in Table 6 We can find that our algorithm FedSDR can outperform these combinations on both worst-case and average-case performances.
Table 6: The comparison between the performance of FedSDR the centralized invariant learning.
Dataset
CMNIST
CFMNIST
WaterBird
Test case
worst-case (%)
average (%)
worst-case (%)
average (%)
worst-case (%)
average (%)
IRM
46.38(+0.61)
49.14(10.89)
47.76(=0.68)
60.38(+1.54) 49.41(#81)
68.63(+1.91)
IRM-IB
49.78(#0.72) 48.15(+0.34)
72)
49.04(=0.52)
50.54(+66)
60.75(+1143)
69.82(+1.64)
MRI
46.81(10.49)
48.92(10.84)
48.84(10.61)
49.33(193)
61.64(11.02)
69.47(11.35)
TRM'-FT
54.27(+0074) 14.32(+0.44)
11.09(=0.81)
58.48(+0.74)
60.25(#0.78)
69.46(+1.25)
IRM -1.2
45.68(+0.58)
49.04(+0.83)
47.92(=0.42)
49.46(+0.92)
61.25(=0.86)
68.93(+1..6)
IRM-IB FT
17.21()0.23) 54.35(10.56)
13.47(10.38)
53.14(10.52)
60.84(10.64)
70.31(+1.06)
IRM-IB IB -12
47.79 50.13(100.7)
49.73(=0.51)
51.65(11.75)
61.38(10.86)
69.08(+1.13)
MRI -FT
15.52(10.61) 53.98(10.92))
12.13(#0.79)
54.21(10.88)
61.14(20.93)
70.84(+1.64)
MRI -12
(+.21)) 52.37(++0.8))
18.73( 48.73(+0.85)
50.67(+0.87)
62.18
69.16(+1.71)
FedSDR
53.88 (+-0.24)
55.59 (+0.65)
56.92 (+0.73) 61.88 (+ 1.02)
65.25 (+-0.97)
73.20 (+1.28)
20

---

Published as a conference paper at ICLR 2024
Compatibility of FedSDR with typical PFL schemes We implement Algorithm 2 and Algo- rithm 3 mentioned above on three datasets (i.e., Colored-MNIST. Colored-FMNIST and WaterBird). The values of the generic hyper- parameters are set same FedSDR. In particular, = 0.5 and y = 1.0 * 10 for Colored-MNIST and Colored-FMNIST. As to WaterBird dataset, we choose À = 0.5 and
choose
1.4. The experimental results are shown in the following Table Note that the performances of FedSDR+FedAvg and Fed are evaluated with the personalized invariant models output by Algorithm 2 and Algorithm 3 respectively. The results show that the model performance can be further improved when we combine the proposed shortcut discovery and removal method with the prevalent federated learning algorithms.
Table 7: The combinations of our method with other federated learning schemes.
Dataset
CMNIST
CFMNIST
WaterBird
Test case
worst-case (%) average (%)
worst-case (%) average (%)
worst-case (%)
average (%)
FedSDR
53.88(+0.24) 55.59(++0.65)
56.92(+0.73)
61 61.88(+1.)) 18(01.002) .2) 65.25(++0.97)
73.20(+1.29)
FedSDR+FodAvg 51.76(+0.28 28) 56.01
+
11.56(+0.99
66.73 +1.133
74.27(+1.44
FedSDR+PFedMe 53.54(=0.19) 55.93(11.20) 56.69 00668) 63.22(10.70
67.32 10.81)
74.38( 74.38(
The effect of pyper-parameters As discussed in the main text, there are two significant hyper- parameters in the proposed objective 5 and 6 i.e., and y. We will evaluate the effects of these two prper-parameters on the performance of our method FedSDR. When we evaluate the effect of the hyper-parameter X, we fix 7 = 1.0 10 for CMNIST and CFMNIST datasets, and 7 1.4 for WaterBird dataset. When we evaluate the effect of the hyper-parameter % we fix À = 0.5 for CMNIST and CFMNIST datasets, and = 0.5 for WaterBird dataset. The results on CMNIST. CFMNIST and WaterBird are given in Table 8 Table 9 and Table 10 respectively.
Table 8: The performance of FedSDR CMNIST dataset with different hyper-parameters À and 7.
0.10
0.20
0.30
0.40
0.50
0.60
0.70
0.80
0.90
1.00
10.0
worst-case (%)
53.73
53.62
53.18
53.93
53.88
53.51
53.83
53.66
53.79
53.82
52.35
average (%)
55.08
55.37
55.72
55.68
55.59
54.97
55.21
55.83
55.52
55.19
54.20
7( (x110*)
0.01
0.10
0.50
0.60
0.80
1.00
1.50
2.00
5.00
10.0
100.0
worst-case (%)
1.21
1.79
43.85
49.73
53.69
53.88
52.87
52.91
51.64
50.72
49.86
average (%)
50.61
50.78
52.21
52.83
53.49
55.59
54.17
54.93
53.31
51.35
50.17
Table 9: The performance of FedSDR on CFMNIST dataset with different hyper-parameters À and 7.
0.10
0.20
0.30
0.40
0.50
0.60
0.70
0.80
0.90
1.00
10.0
worst-case (%)
56.31
56.17
56.89
56.57
56.92
56.18
56.36
57.15
56.57
55.36
54.62
average (%)
61.62
62.07
61.78
61.57
61.88
61.83
61.75
61.53
61.25
61.03
60.77
7(x10*)
0.01
0.10
0.50
0.60
0.80
1.00
1.50
2.00
5.00
10.0
100.0
worst-case (%)
0.19
0.27
52.68
55.57
56.28
56.92
54.79
53.15
53.05
51.68
49.59
average (%)
50.19
50.73
57.63
60.54
61.72
61.88
61.06
59.53
58.74
55.24
51.38
C.6 COMPLETE EXPERIMENTAL RESULTS
The test accuracy shown the main text are calculated based on 3 trials with different random seeds. Due to space limitation, only give the average value in the evaluation part of the main text. Here, the complete results on CMNIST, CFMNIST and WaterBird are provided in Table [I] and the detailed results PACS are shown in Table 12 Besides the detailed results for evaluating the effect of hyper-parameter À and 7 on WaterBird dataset given in Table 10
21

---

Published as a conference paper at ICLR 2024
Table 10: The performance of FedSDR on WaterBird dataset with different hyper-parameters and
0.10
0.20
0.30
0.40
0.50
0.60
0.70
0.80
0.90
1.00
10.0
worst-case (%)
62.51
64.13
63.87
64.91
65.25
64.16
63.24
63.67
62.15
61.68
61.74
average (%)
72.39
71.87
71.56
72.78
73.20
72.15
71.76
72.09
71.14
70.61
70.18
0.10
1.00
1.10
1.20
1.30
1.40
1.50
1.60
1.80
2.00
10.0
worst-case (%)
44.16
57.64
60.33
61.58
63.75
65.25
64.52
63.26
63.58
60.95
48.29
average (%)
65.73
70.18
71.78
71.13
72.73
73.20
72.14
70.54
71.19
68.73
64.86
Table 11: The complete results on CMNIST, CFMNIST and WaterBird datascts.
Dataset
CMNIST
CFMNIST
WaterBird
Test case
worst-case (%)
average (%)
worst-case (%)
average (%)
worst-case (%)
average (%)
FedAvg
3.39
51.03(+0.67)
0.16(+0.36)
50.02(+10.42)
67.95(#1.03) 54.13(110.75)
DRFA
21.15(10.21
19.84(#1.21) (2..1(+:0:45
03.88(11.80)
59.75(+1.41)
68.39(+1.50)
FedSR
46.93 +0.12)
48.62(+00.3)
47.61 (+0.39))
48.90(+00.33)
61.75(+0.81)
71.68(+1.17)
FedIIR
47.25 (-0.26)
48.39(+0.42)
48.06(+0.50)
49.16(+0.51)
61.24(+0.76)
70.87(+0.94)
FTFA
15.42(10:28)
54.96(110.1)
11.35(10.81)
53.52(10..5)
54.38(11..0))
69.68(11.36)
pFedMe
21.30(+0010) 48.53(+0024)
4.22(++0.4)
51.26(+0051)
55.63(+0.64)
68.24(+00.83)
Ditto
3.02(=0.32) 50.97(+0.59)
0.37(+0.42)
50.12(+0.63)
53.13(++0.3)
68.73(+0.72)
FedRep
2.76(+0.34) 50.83(#0.68)
0.11(+0.37)
50.01(#0:49)
52.88(+0.56)
70.23(+0.92)
FedRoD
9.09(10.32)
50.84(10.21)
1.23(10.15)
52.36(10.78) 51.57(110.8)
70.86(10.88)
FedPAC
(.0
50.05(10.54)
0.16(40.11)
50.13(10.41)
15.08(10.84)
65.57(+1.12)
FedSDR
53.88 (+:0.41) 55.59 (4.0655
56.92 (4.0.33
61.88 (+1.02)
65.25 (+.0.97) 73.20 (+1.28
Table 12: The detailed results PACS dataset.
Algorithm
Art Painting (%)
Cartoon (%)
Photo (%)
Sketch (%)
average (%)
FedAvg
41.71(=0.62)
51.21(=0.16)
52.92(=0.44) 44.79(+0..27)
47.66
DRFA
52.84(=0.33) 40.48(
53.37(=0.64) 47.12(40.52)
48.95
FedSR
46.76(+0.34)
54.15(-0.24)
48.67 +0.43)
55.43( (5.71)
51.25
FedIIR
47.03(+00.0)
54.66(=0.34)
49.51
55.12(=00.5)
51.58
FTFA
40.89(+0.51)
52.56(=0.28)
47.38(+0.76)
54.31(=0.37)
48.79
PFedMe
45.24(+0.43)
54.57(=0.52)
49.19(+0.56)
56.33(=0.76)
51.33
Ditto
(4.95
53.48(=0.21)
50.71(+0.65)
55.98(=0.57)
51.28
FedRep
49.27(+0.26)
56.08(-0.35)
51.48(41.73)
58.15(10.82)
53.75
FedRoD
48.16(+10.8)
55.13(=0.41)
51.02(40.81)
57.38(=0..9)
52.92
FedPAC
(9.93(+0.31)
55.86(=0.53)
52.35(+0.58)
58.67(=0.58)
54.20
IRMt
46.35(+0.81)
53.71(=0.79)
48.52(+0.86)
54.73(
50.83
IRMi-FT
43.18( L0.56)
53.23(=0.45)
48.27 48.27(40.66)
(5.73)
50.04
IRM1-L2
(00.66
54.08(=00.27)
49.35(+0.51)
55.91(=0.42)
51.98
FedSDR
52.14(+0.75)
58.23( 0.84)
55.49(40.67)
58.84(10.91)
56.18
FedSDR+FedAvg
53.26( +0.63)
58.91 (..96(=0.76)
36.37(40.52
59.37
56.98
FedSDR+PFedMe
52.58(+0.522
58.63(=0.57)
55.85(+0.49)
58.72(=0.60)
56.45
22


Reviews:
[{"review_id": "z0hHEHo2HC", "reviewer": "Reviewer_zjFS", "summary": "The paper focuses on addressing the \"shortcut trap\" issue within personalized federated learning (pFL) by introducing FedSDR. This solution aims to identify and eliminate shortcut features, leading to enhanced performance for individual clients within pFL on their respective local datasets. The authors support their approach with theoretical proofs and comprehensive experiments, demonstrating the effectiveness of FedSDR across multiple scenarios.", "questions": "1. The paper’s algorithmic approach (Algorithm 1) deviates from the conventional practice of client selection in Federated Learning (FL). Could the authors elucidate why they've chosen this approach over the typical FL client selection methodology?\n\n2. Concerns arise from certain experimental results in Table 1, particularly the unexpected lower accuracies observed for specific baseline methods on CMNIST and CFMNIST. Can the authors provide an explanation to address these concerns?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Originality: The paper pioneers the exploration of the shortcut trap problem in pFL.\n\n2. Solid Theoretical Support: The inclusion of strong theoretical foundations.\n\n3. Sufficient Experiments: The comprehensive set of experiments strengthens the paper's contributions.", "weaknesses": "1. Notation Conciseness: The notation could be more concise, potentially enhancing the paper's readability.\n\n2. Confusing Training Process: A section of the training process appears confusing and requires clarification.\n\n3. Inconclusive Experimental Results: Certain aspects of the experimental outcomes lack conviction and need further clarification for robustness.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "igBLdFK3G2", "reviewer": "Reviewer_2uxd", "summary": "The paper proposes a new personalized federated learning approach named FedSDR. Considering generalization on unseen test data, the paper utilizes invariant learning in the federated setting. Specifically, FedSDR first extracts shortcut features that are irrelevant to the task and them remove it to extract the most informative personalized invariant features by carefully designing the objectives. Experiments show that FedSDR outperforms the other baselines in the settings where the test data distribution shifts.", "questions": "1. Can you add a synthetic dataset with a simple model in the experiments? It can be used to verify the theorems by satisfying the assumptions.\n\n2. Can you add experiments without manually adding shortcut features? It is quite important. Currently, I’m not clear what are the real applications of FedSDR.\n\n3. Can you add experiments that increase the number of clients and adopt client sampling?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Applying invariant learning in the federated setting is interesting and promising.\n\n2. The organization of the paper is clear.\n\n3. FedSDR significantly outperforms the other baselines.", "weaknesses": "1. My main concern is about the experimental setting. Experiments are based on the simulated setting where the authors manually add shortcut features and change the test data distributions. Based on the motivation, the shortcut features should naturally exist in datasets. Experiments on real-world natural datasets are necessary. Otherwise, the application of FedSDR may be very limited.\n\n2. The theoretical analysis has strong assumptions, e.g., logistic regression, linear function, etc. The analysis may not be applicable in the experimental settings.\n\n3. The number of clients used in the experiments is missing. Experiments to evaluate the scalability of FedSDR are not provided.\n\n4. Typos: 1. Page 4: “Theorem 1” -> “Lemma 1”; 2. Page 6: “guarantee” -> “guarantees”", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "3k6Q6iCEsx", "reviewer": "Reviewer_Qvuh", "summary": "The paper proposes a novel approach to mitigating shortcut learning in personalized federated learning, which is a challenging problem in real-world settings. The proposed method, FedSDR, utilizes structural causal models to discover and remove shortcuts while preserving personalization knowledge.", "questions": "Please see above", "limitations": "", "rating": 3, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. This paper provides extensive background information, offering readers significant convenience in understanding the topic.\n\n2. The authors conducted extensive experiments on some real-world datasets, validating the excellent performance of the proposed method.", "weaknesses": "1. I believe the main issue with this paper is that the research motivation seems weak. The paper claims to primarily address the scenario where training sets are Non-IID among clients, and where within each client, the training and test sets are also Non-IID. However, I think existing Robust Federated Learning and Federated Domain Generalization methods are capable of handling the aforementioned scenario. Although these two methods primarily focus on the issue of Non-IID training and test sets within each client.\n\n2. The paper states, \"To the best of our knowledge, we are the first to consider the shortcut trap problem in personalized federated learning and analyze it by formulating the structural causal models for heterogeneous clients.\" While it's true that this paper is indeed the first to study the use of PFL+SCM to address the shortcut trap problem, existing Robust Federated Learning and Federated Domain Generalization methods can also address the shortcut trap problem. Therefore, the contribution of this paper appears to be limited.\n\n3. In Figure 1, I find the classification of the scenarios where RFL&FedDG are applicable not very accurate. In the scenarios where RFL&FedDG are applicable, the Test-test relation can also be IID or Non-IID.\n\n4. In the experimental section, the proposed model is only compared to FL and PFL methods. However, I believe it should also be compared to RFL and FedDG methods to provide more convincing experimental results.\n\n5. The experimental section lacks detailed information about the configuration of the model used in this paper. More detailed information is needed.\n\n6. The paper lacks a thorough discussion of the limitations of the proposed FedSDR method. While the authors mention some potential limitations in passing, a more detailed discussion of the assumptions and constraints of the method would be helpful for readers to better understand its applicability in different scenarios. \n\n7. Although the authors provide some high-level explanations of how FedSDR works, a more detailed discussion of the causal modeling framework and how it is used to address shortcut learning would be helpful for readers who are not familiar with this area of research.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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:

Decision: Accept (poster)