FedLOE: Federated Domain Generalization via Locally Overfit Ensemble
Authors: Ruqi Bai
David I. Inouye
Venue: Accepted by TMLR ()

Abstract: In federated learning (FL), clients typically access data from just one distribution. Ideally, the learned models would generalize to out-of-distribution (OOD) data, i.e., domain generalization (DG). However, centralized DG methods cannot easily be adapted to the domain separation context, and some existing federated DG methods can be brittle in the large-client, domain-separated regime. To address these challenges, we revisit the classic mixture-of-experts (MoE) idea by viewing each client as an expert on its own dataset. From this perspective, simple federated averaging can be seen as a type of iterative MoE, where the amount of local training determines the strength of each expert. In contrast to the standard FL communication-performance trade-off, we theoretically demonstrate in linear cases and empirically validate in deep models that reducing communication frequency can effectively enhance DG performance, surpassing centralized counterparts (e.g., $+4.34\%$ on PACS). Building on this, we further propose an additional MoE strategy to combine the client-specific classifier heads using standard DG objectives. Our proposed \methodname method can be viewed as an intermediate approach between FedAvg and one-time ensembling. It demonstrates both theoretical soundness and empirical effectiveness. Moreover, \methodname requires fewer communication rounds, highlighting its practical efficiency and scalability.

Published in Transactions on Machine Learning Research (03/2026)
FedLOE: Federated Domain Generalization via Locally Over- fit Ensemble
Ruqi Bai
boil 16
Elmore Family School of Electrical and Computer Engineering Purdue Universily
David I. Inouye
dinouye Dpurrdue. edu
Elmore Family School of Electrical and Computer Engineering Purdue University
Reviewed on OpenReview: https: // openreview. net/ fortmf id= #4T9sK6Ga
Abstract
In federated learning (FL), clients typically access data from just one distribution. Ide- ally, the learned models would generalize to out-of-distribution (OOD) data, iee, domain generalization (DG). However, centralized DG methods cannot easily be adapted to the domain separation context, and some existing federated DG methods can be brittle in the large-client, domain-separated regime. To address these challenges, we revisit the classic mixture-of-experts (MoE) idea by viewing each client as an expert on its own dataset. From this perspective, simple federated averaging can be seen as a type of iterative MoE, where the amount of local training determines the strength of each expert. In contrast to the standard FL commmication--peffomme cases and empirically validate in deep models that reducing communication frequency can effectively enhance DG performance, surpassing centralized counterparts (e.g., +4.34% on PACS). Building on this, we further propose an additional MoE strategy to combine the client-specific classifier heads using standard DG objectives. Our proposed FedLOE method can be viewed as an intermediate approach between FedAvg and one-time ensembling. Tt demonstrates both theoretical soundness and empirical effectiveness. Moreover, FedLOE requires fewer communication rounds, highlighting its practical efficiency and scalability.
trade-off, thcoretically demonstrate in lincar
1 Introduction
In federated learning (FL), each client owns a portion of the overall data. For example, in medical network, each hospital may possess ils own set of chest. X-rays tissue sample images, collected using different imaging devices. Similarly, regional datasets capturing economic, crime, or demographic information are often compiled with varying methodologies. As a result, each client's dataset may originate from distinct domain, leading to what commonly referred to as domain separation (Bai et al., 2024). This setup violates the typical independent and identically distributed (i.i.d.) assumption and presents challenges for model training in FL. Additionally, it raises the question: Con models trained in this setting generalize to clients from previously unseen domains? This is the central concern of domain generalization (DG), which consider in the federated setting where domain separation is pronounced and the number of clients is large. This setting is also known to be challenging cmpirically: the federated DG benchmarking study of Bai et al. (2024) reports that many existing federated DG baselines can substantially degrade under domain separation, especially as the number of clients grows, highlighting a gap between centralized DG assumptions and realistic federated deployment. These findings motivate our focus on mechanisms that remain effective in the large-client, domain- separated regime and do not require sharing data or multi-domain access on any single client.

---

Published in Transactions on Machine Learning Research (03/2026)
A straightforward approach might be to adapt centralized DG objectives to each client. However, many popular DG methods assume access to multiple domains during training (Arjovsky et al., 2019; Shi et al., 2022; Sun and Saenko, 2016; Li et al., 2018b; Sagawa et al., 2020; Ganin et al., 2016; Zhou et al., 2021; Li et al., 2018a) 8a), which does not hold in the domain separation sctting of FL. Existing methods that address DG in FL her do not scale well with many clients (Zhang et al., 2021; Nguyen et al., 2022) or may compromise privacy by sharing the frequency spectra of images (Liu et al., 2021).
To address these challenges, revisit the mixturre-of-experts perspective, treating each client as expert its local domain. Rather than being a drawback, domain-specific data may enable clients to specialize, potentially offering complementary information across the network. The key challenge then becomes: How can these expert models be effectively combined to support out-of-domain generalization?
We consider two combination strategies. First, we explore an implicit combination through the standard FedAvg algorithm and find that communication frequency plays an important role in DG performance (scc Section 3). Surprisingly, in some cases, FedAvg with infrequent communication achieves more stable and improved DG performance compared to centralized ERM (which corresponds to communicating at every iteration). We hypothesize that frequent communication may lead models to overfit to common spurious correlations, while infrequent communication allows them to retain diverse representations. Our theoretical analysis supports this behavior in a simplified linear setup (Section 3).
Second, we propose an explicit ensembling step for combining locally trained classifier heads. This step is motivated by observations in the DG literature: while DG methods can be effective in simple linear settings (Arjovsky et al., 2019; Sun and Saenko, 2016), they often underperform on overparameterized deep models (Gulrajaní and Lopez-Paz, 2021; Koh et al., 2021). Morcover, recent studics suggest that the deep feature extractor may already encode useful invariances, but the classifier head can emphasize spurious features (Wald et al., 2022; Rosenfeld et al., 2022). Drawing on these insights, we adopt a two-stage training process that overfits client-specific classifiers and then combines them using a standard DG objective. Our contributions are as follows:
1. We explore a vixtuur--ff-xxpett interpretation of federated DG and observe that FedAvg with reduced communication frequency can lead to improved out-of-domain generalization under domain separation. 2. We provide a theoretical analysis in simplified linear setting that sheds light on the conditions under which our approach may outperform centralized FRM.
3. We introduce a two-stage algorithm, FedLOE, that first performs local overfitting followed by aggregation using parameter averaging and a DG objective, respectively.
4. We present empirical results on several real-world datasets, providing evidence that our approach performs comparably or favorably to existing baselines and is robust to varying numbers of clients.
2 Background
Consider a featurizer gx(x)) RDD - Rdd with parameters and a linear classifier head denoted as hy(z) Rd Rm with parameters v where z = go (x).. Let E(s,)) and L(0, v) = Ep(za) [ete.)]((hss (go be the per-sample and expected loss respectively. Let [A] represent the set of integers up to A: [A] 2 (1,2. A). C denotes the number of clients, and K denotes the number of training domains.
2.1 Domain generalization (DG)
In domain generalization, we assume there is a set of K domain-specific training distributions where the k-th joint distribution is denoted as p(x: y|k) and the marginal probability of cach domain is p(k). Additionally, there are one or more domain-specific unseen test distributions denoted by p(æ, w/kk,, where h > K. The goal of DG is to perform well on these unseen test distributions, which can be formalized as minimizing the expected loss over the set of unseen test distributions, i.c.,
min Σ (0)) 0,
(1)
K
2

---

Published in Transactions on Machine Learning Research (03/2026)
Clearly, because the test domains are unknown, practical proxy objective to use standard empirical risk minimization (ERM) by minimizing the expected loss (i.e., risk) of the training domain distributions:
where (g(0,)) - E Epta. (go y)].
min
Σ p())Cx(0,0)) 6,45
where La are the expected losses of the training distributions. Despite its simplicity, simple ERM has been difficult to beat for the DG task (Gulrajani and Lopez-Paz, 2021), particularly when training deep non-lincar models. Several common approaches add type of DG regularization term to the ERM objective.
Common approaches to DG include representation learning, such as domain-invariant representation learning via kernel methods (Muandet et al., 2013; Ghifary et al., 2016), invariant risk minimization (Arjovsky al., 2019; Krueger al., 2021), and domain adversarial neural networks (Sun and Saenko, 2016; Li et al., 2018b). Bosides domain-invariant learning, several methods add regularizations to the gradient computations (Shi et al., 2022; Rame el al., 2022). Other approaches include distributionally robust optimization (Sagawa el al., 2020), which optimizes for the worst-case distribution scenario across training domains; and meta- wheaallerrinng (Fimn et al., 2017; Li et al., 2018a), which uses paaninggto-learn mechanism to acquire general knowledge by constructing meta-learning tasks that simulate domain shifts.
Most of these methods can be viewed as adding regularization penalty r(0, v) with tuning parameter À the simple ERM objective:
min 0,0
((,) Ar(
(3)
We select three representative examples that we will use in our experiments. In the IRMv1 method from Arjovsky et al. (2019), the regularization term encourages gradients to be zero:
Firm (1, (0. v)) & Σ Epp Vo, Vo 2)
Another popular choice is to align the gradients as in Fish (Shi et al., 2022) by computing the inner product of the domain-specific loss functions:
(0, v 2
EEpero
E
Epc.pe (Vo.o v), Vo.
The regularization in REx (Krueger et al,, 2021) seeks to reduce the loss variance across domains;
FRE (0, xb) A Var (0, ()., ... Lx(0,v)))
2.2 Federated domain generalization
In the federated DG setup, we assume C' clients, each possessing data from only a single training domain a specific form of client heterogeneity known as domain separation (Bai et al., 2024). Formally, client e will have samples from one domain distribution, i.c., p(æ,, yle) = p(æ,, y/k) for some domain k. In our experiments, we assume the number of clients is greater than or equal to the number of training domains (C' K), specifically where C' is fairly large in contrast to K, which may only be 4 domains (e.g., in PACS). Importantly, this type of client heterogeneity is distinct from label imbalance, Le., 3k * k k(((() * p(ylk',, which is the most common heterogeneity considered in the FL literature. Rather, we assume that the joint distribution could be different, i.e. 3k + k', p(æ,, y/k) + p(æ, vlk'), but we do not assunee any particular lype of shift between distributions. This more general case of client heterogeneity than class imbalance. Crucially, it important to explicitly distinguish this federated DG setting from standard FL client heterogeneity. Traditional FL generally focuses on mitigating label skew or covariate shift to optimize a global model for the participating clients, or to personalize models for their specific local distributions (iee., in-domain
3

---

Published in Transactions on Machine Learning Research (03/2026)
performance). In contrast, federated DG under domain separation explicitly targets out-of-distribution (OOD) generalization. The primary objective is not to maximize performance on the K domains present during training, but to extract robust, invariant representations that transfer to completely unseen target domains. Consequently, the core challenge shifts: rather than merely aggregating local knowledge or aligning client updates to satisfy current clients, the system must prevent models from overfitting to common spurious correlations that would degrade performance novel environments. Like ERM the centralized case, the standard FL learning algorithm the FedAvg algorithm (McMahan et al., 2017), which minimizes the expected loss on each client and then averages the parameters from all clients. Prior empirical studies suggest that this federated DG regime is particularly brittle for many existing approaches. For example, Bai et al. (2024) systematically evaluates federated DG methods under domain separation and reports that several baselines may struggle to maintain robust. OOD performance in this setting. This motivates our emphasis on (i) communication--fficcent mechanisms that scale to many clients and (ii) procedures that preserve client-specific specialization rather than forcing frequent synchronization.
Using our notation, FedAvg can be compactly formalized performing sequence of compulational steps denoted by e [T'] where T is the maximum number of computational epochs. In each computational step, FedAvg updates the model and sometimes also synchronizes and averages all parameters across clients: for all c C [C],
Ue 1+
(lt1) = (V - VVCol
We 1 (
ift TT . otherwise, Σ Wo (
(Wttl,00 ( 0ct1)
(4)
where a is the step size, and T S [T] denotes a set of synchronization indices. If T = [T] then the synchronization of the sequences is performed every epoch, which corresponds to using mini-batch SGD with mini-batch size C' to solve ERM. If T = {T}, then (4) amounts to one-shot averaging. A proper choice of the synchronization set T allows us to obtain an expressive model while keeping the influence of common spurious features relatively low. T is Lask-dependent and can be determined experimentally. Furthermore, the averaging weights p(c') are chosen based on the training size.
Formulated in this way, we notice that FodAvg can be interpreted as a type of local SGD (Stich, 2019), which was originally aimed at improving standard supervised learning (i.e., in-domain accuracy). In Section 3, we show that unlike standard supervised learning teikss-wheree Stich (2019) show that more frequent commu- nication improves in-domain accuracy the federated DG accuracy is not monotonic w.r.t. communication frequency.
To the best of our knowledge, only a few works seek to solve federated DG. In particular, Nguyen et al. (2022) proposed FedSR, enabling domain generalization while still respecting the distributed and privacy- preserving nature of the FT. context by enforcing an l norm regularizer and a conditional mutual information regularizer on the representation. Liu et al. (2021) propose FedDG, a federated learning paradigm specifically designed for medical image classification. The proposed method requires sharing the amplitude spectrum of images among local clients, which violates standard privacy protocols. Zhang et al. (2021) applies generative adversarial networks (GANs) (Goodfellow et al., 2020) in the FL context, where it first trains the featurizer and classifier by minimizing the empirical loss, then trains the generator and the discriminator using a GAN- based (Goodfellow et al., 2020) approach. Additionally, FedGMA (Tenison et al., 2023) proposes applying a mask on the gradient on the server side to align updates across domains.
3 Unraveling the benefits of expert parameter ensembling for domain generalization
In this section, we provide motivation through a simple experiment showing that frequent communication can lcad to out-of-distributoon (i.c., DG) overfitting. We then show theoretically that infrequent communication leads to better DG accuracy on a linear structural causal model. The linear setting is chosen for closed-form population solutions that disentangle the mechanism of interest. from aptimization/fiitteee /fnite-samttle effects. It.
4

---

Published in Transactions on Machine Learning Research (03/2026)
is intended as explanatory and intuitive toy model; experiments demonstrate the same qualitative behavior in deep networks.
Motivating observation: infrequent communication leads to better DG accuracy We start with a surprising phenomenon related to the effect of communication frequency on DG accuracy when using FedAvg. To demonstrate this phenomenon, conducted an experiment PACS with 20 clients using the same model and a fixed computational budget of 256 local epochs, while varying the number of communication rounds (i.e,, how many times model parameters are averaged). We follow the standard leave-one-domain-out DG protocol: we designate one PACS domain as the target domain that is completely unseen during training, and train only on the remaining source domains. We report two lest-lime metrics: in-domain accuracy is computed on the held-out test splits of the source domains, while DG (out-of-domain) accuracy is computed on the held-out test split of the target domain. To form 20 clients in the federated setting, we partition the source-domain training data across 20 clients, such that each client contains samples from a single source domain, and run FedAAg under different communication frequencies with the same total local-epoch budget.
As seen in Figure| 1, while in-domain accuracy steadily increases and stays nearly constant, the optimal communication frequency for DG is actually a very small mumber of communications. performance, FedAvg with 4 communications attains the best performance and stability as seen by the low variance. In contrast, its centralized counterpart, mini-batch SGD with total of 256 communications, performs poorly. These results indicate that increasing communication frequency in federated learning can potentially hinder DG performance.
terms of DG
Invariant feature
inv
fredioe
0.9
Stage 2
A Aerm
Accuracy
Spurious feature
0.8
Stage
0.7
domain generalization
0.6
in domain
More communication
0o
100 150
200
250
Less communication
Total communications
While more communicatinn improves in- domain test accuracy (evaluatod on held-out test splits of the source/training domains), more commu- nication may harm out-of-distribution domain gener- alization (DG) test accuracy (evaluated the held- out test split of an unseen target domain). The left- most points represent one-shot. averaging (iee., only one communication), while the rightmost point cor- responds to ERM (i.e., communicate every batch).
Figure 1:
Figure 2: This figure illustrates how the infreeuent communication of FedLOE's first stage (green) may implicitly regularize the solution away from the ERM solution 6ern but closer to the robust invariant solu- tion 0mn, and FedLOE's second stage (red) explicitly moves toward the invariant solution via DG objec- lives, while standard Fed with frequent commu- nication (blue) converges to the ERM solution.
Thcorctical analysis of expert parameter ensembling for a linear structural equation model We now theoretically validate our empirical finding by analyzing a linear DG problem to provide further theoretical insight into expert parameter ensembling. Our analysis sheds light on the potential for parameter averaging of locally overfit models (i.c., experts) to indeed improve DG performance. We expect that this result could be extended to more complex linear settings (and possibly non-linear ones) but focus on simple clear example here. Concretely, we extend the linear model example from Arjovsky et al. (2019) to multiple
5

---

Published in Transactions on Machine Learning Research (03/2026)
invariant and spurious features. Consider the following structural equation model:
xI L1, y= + Eo. 2s = yas + E3+
(5)
where fi, E 21, are exogenous noise variables, R is the regression target, ær R is the invariant feature vector, Is E IR is the spurious feature vector, and az E R and as E R are the invariant and spurious parameters, respectively. Note that y is causally generated by zy while as are spurious features because they are generated from (i.c., they are causal descendants).
To frame the DG problem, we assume that each domain or environment shares the above causal model (including the same az and as) but has domain-specific noise distributions for & and & Additionally, for the FL context, assume that each client has access to data from single domain. We formalize these assumplions next.
Assumption 3.1 (Fimite-moment domain-dependent exogenous noises). We assume the exogenous noises in (5) satisfy:
1. The components of Éι, Ê2, and &a have finite second moments.
2. &Ι and E2 are zero mean with depended variances:
E[G]] = O,
[[GIC]] = (A) 1.
[[[2] = O,
B = o2(..
3. E3 is zero mean with dodeene-nneeennden variance:
El O,
El = = J.
A Assumption 3.2 (Domain separation). Each client c e [C] has data from unique, non-overlapping domain k [K] where C = K in this example.
Given this problem context, we consider the OLS estimator of the coefficient a == [az, as) for predicting from the concatenated features x = [ ass) E Rdd:
â E arg min E(nyy) [(y- x aa)
(6)
&cE
The domain-invariane i.e., robust solution, is to simply ignore the spurious features æs and only use the linear coefficients of the invariant features, i.e., a == [az, 0 This solution also corresponds to finding the true causal mechanism that generates y. We compare two distinct cases in this context: (1) the clients communicate every epoch so that the solution converges to the ERM solution (Stich, 2019), or (2) the clients solve their local problems and only communicate once, i.e., one-shot averaging of the model parameters.
Theorem 3.3. Given the least squares problem (6) under Assumplion 3.1 and Assumplion 3.2, one-shol averaging will put relatively more weight on the invariant features, i.e.,
ERM
@OSA COSA
ERM S
and OSA will be closer than ERM the invariant solution a - Jazz T,, i.e,,
|| a²1 >
aa'l.
(8)
Remark 3.4 "Population-rise risk setting vs. finite-sample extensions). Theorem 3.3 is stated and proved for the population least-squares objective in (6). Accordingly, the proof relies only on the first and second moments in Assumption 3.1 and does not require tail conditions such as sub-Gaussianity. A finite-sample analogue (replacing population moments by empirical estimates on each client/doman concentration terms depending on the number of clients per domain and the number of samples per client; establishing such bounds typically benefits from stronger distributional assumptions (c.g., sub- sub-Gaussian bounded-moment conditions), which we leave as future work.
would introduce additional
Corollary 3.5. Under the above assumption, the DG risk of ERM is higher than that of OSA. i.e.,
E [v x a @ERM)2] > E [1V x. a 0
(9)
where (xr, D rest-
6

---

Published in Transactions on Machine Learning Research (03/2026)
While the full proofs are provided in Appendix A and Appendix B, respectively, the key to the proof relies on the fact. that the harmonic mean (the ERM solution) is less than the arithmetic mean (the OSA solution). Theorem 3.3 shows that simple one-shot ensembling of local estimators can be more robust than ERM. Furthermore, Appendix B shows that OSA could achieve lower DG risk. This validates our empirical findings and suggests that frequent communication may over-emphasize spurious features. Further theoretical analysis would be required to analyze cases in between these two extremes that are more likely to have better DG performance than either extreme as seen empirically in Figure 1. Given our empirical results, we expect this type of behavior to hold beyond this specific example.
4 FedLOE: locally overfit ensemble algorithm for federated DG
In this section, we establish framework of new algorithms by creating experts via client-specific training followed by expert ensembling via linear parameter combination. Our method aims to reduce overfitting to common spurious features, induce robust features, and control data privacy loss. Our algorithm is designed with two stages that share the same structure: (1) construct expert models by overfitting locally and (2) combine experts via linear combination as outlined in Table 1 and illustrated in Figure 2.
Table 1: FedLOE: An iterative two stage framework for federated DG that locally overfits and ensembles first all model parameters (0, v)) and then only the linear classifier head parameters sb.
Stage 1: Learning shared features
Stage 2: Robustifying classifier head
Step 1: Locally overfit
SGD on 0 and vb Simple average of parameters
SGD on &b
Step 2: Ensemble parameters
Linearly combine via DG objective
In stage one, we draw from the key observation in Section 3 and reinterpret FedAvg with infrequent com- munication as locally overfitting and then averaging the local expert parameters. Infrequent communication helps to avoid common spurious features while still learning an expressive model. Additionally, infrequent communication better for creases DG robustness compared to synchronizing every mini-batch (which is equivalent to centralized SGD), which requires the most frequent communication. In stage two, we freeze the feature extractor parameters but allow each client to become an expert by training its own lincar classifier head locally on its own dataset (note this requires no communication). Then we ensemble the expert classifier heads using standard DG objectives to make the final ensemble classifier more robust.
FL on edge devices. Even by itself, stage one in-
4.1 Stage 1: Learning shared features via FedAvg with infrequent communication
Motivated by the observation in Section 3 that frequent averaging may amplify common spurious features, we use FodAvg with an infrequent commuication schedule (4) where T = Tı. This can be viewed as overfitting each client model on its own data and ensembling these models via parameter averaging where the communication frequency determines how specialized each expert becomes. Ti is tuned experimentally and generally is a much smaller set than all possible time points. For example, in Figure the optimal choice is Ti 4. By the end of this stage, all clients share common featurizer 96 and linear classifier vift, An ablation study in Section 5.1 also confirms that if we replace stage 1 with the most frequent communications, i.e. mini-batch SGD, the DG accuracy would decrease significantly. This verifies our hypothesis from the exploratory experiments in Figure 1 that frequent communication induces common spurious features. At the end of the first stage, the featurizer parameters are fixed and we will focus on the last. linear layer because prior work has shown that the deep featurizer may be robust enough and only the last classifier layer noods to be made robust. (Wald et al., 2022).
4.2 Stage 2: Robustifying the linear classifier head
In Stage 2, we follow similar overfitting and ensembling strategy as in Stage but freeze the featurizer and update only the linear classifier head. The key difference is that we compute specialized linear combination
7

---

Published in Transactions on Machine Learning Research (03/2026)
of the classifier heads via DG objectives rather than simple average as in FedAvg. To do so, each client randomly holds out small local luming set (disjoint from its local training data); this tuning set is used only in Stage 2.2 to learn the ensemble weights.
Stage 2.1: Local overfitting of classifier heads. In this substage, the clients become experts on their own local datascts by training only their classifier heads the E [C] on their local datasets {aftaann (fraiu) ) Jigttra. We assunne each client e partitions its local labeled data into two disjoint subsets: a local training set Dirain used in Stage 2.1 to fit tee, and small local tuning set Diune that is held out from training and used only in Stage 2.2 to learn the ensemble weights on the server. Unless otherwise stated, De is formed by uniform random split of the client's local data and is fixed throughout training for a given run or seed.
While keeping the featurizer 97 from stage 1 fixed, specifically, for each c E [C], we initialize at the current classifier head and use SGD to solve
Ve _ arg min Cc,0, vbe)
(10)
tee
Thus, by the end of substage 2.1, we obtain C distinct linear classifiers e1,-..,,0c Ve, where each ve E mxd
Stage 2.2: Robust ensemble of classifier heads via DG methods. Stage (infrequent-coomuictttto FedAvg) is typically the primary driver of DG gains, while Stage 2 provides an optional, low-commuiiatioon refinement of the classifier head. In this stage, we learn a robust linear combination of the expert classifier heads using DG objectives. For instance, IRM (Arjovsky et al., 2019) or Fish (Shi et al., 2022). The additional improvement from Stage 2 can be objeetive-dependent and detaset-depenndet (e.g., whether IRM//Fish/llR /REx are effective in that regime). To do so, the server collects the expert heads v1....,44c from Stage 2.1 and estimates ensemble weights using a small tuning set from each client. This fine-tuming can be accomplished in two equivalent ways depending on the trust context; additionally, this stage can be omitted privacy is paramount over performance gains.
Protocol A viscr-visible heads). The server broadcasts all expert lincar classifiers ,i...., VC to cach client. Each client e evaluates every expert d' E [C] on its local luning examples (ær, v) E Diuue and sends only the resulting predictions and labels to the server. Concretely, for each client and each expert d', let Te,(()) A Ae (gg(z)) denote the logit prediction on input æ. Each client sends these predictions to the server, and the server then estimates the ensemble weights @ by minimizing (11).
Protocol B (trusted server). Alternatively, to avoid broadcasting client heads to other clients, each client sends only the frozen features 9z(x) and labels for æ E Diune to the server. The server then computes Te,e( (x) locally for all experts e' and solves the same objective (11). Both protocols optimize the same ensemble objective but rely on different trust assumptions.
C r(e) E(
C [ Se ( (o)
Te
w) ] + x
= arg min &ERC
(1I)
where r()) denotes a DG rogularization term computed from the tuning-set predictions (x)) and labels across clients (see Appendix C.1 for the specific forms used in our experiments), and À is a regularization parameter.
From another perspoctive, because both the classifiers and this meta-combination are lincar, the resulting ensemble classifier can be simplified to single linear classifier by taking combination of the client classifiers, i.c., Vensemble E Babe. This second perspective allows us to write both steps of this stage in similar manner to FedAvg with the important exception that the linear combination weights @ are learned via the DG objective in Equation (11), i.c.,
--
re Co Te-1) 26-1),
ift $ T2
Vc, 32
(12)
3-1
if E Tz
8

---

Published in Transactions on Machine Learning Research (03/2026)
In summary, our FedLOE provides general framework for domain generalization methods in the federated context, which can incorporate newly designed centralized DG methods in the second stage to provide better DG accuracy. See Section 5 for experiment results. The framework naturally fits into the FL context with a low communication burden. The key idea of the framework is to "freee the clients by letting them overfit in-domain to avoid common spurious features, further personalize their classifiers, and encourage invariant predictions by solving linear classification problem with DG regularization. See Algorithm 1 in the appendix for more details. Our ablation study in Section 5.1 shows that without Stage 2, or when training Stage 2 with SGD, the DG accuracy would stay the same as FedAvg.
5 Experiments
In this section, we evaluate FedLOE on three real-world datasets: PACS (Li et al., 2017), OfficeHome (Venkateswara et al., 2017), and TWildCam-Wilds (Koh et al., 2021). We choose popular IRM (Arjovsky et al., 2019), Fish (Shi et al., 2022), and REx (Krueger et al., 2021) as the DG regularizers for Stage 2, and we compare them against FodAvg (McMahan et al., 2017). FedDG (Liu al., 2021), FedSR (Nguyen et al., 2022), and FedGMA (Tenison et al., 2023), which were originally designed solve DG in the FL regime. Our experimental setup follows the domain scparation protocol detailed in Bai et al. (2024): each client contains data from single domain, and we evaluate OOD generalization to unseen target domain under the standard leave-one-domain-orr DG protocol. Note that our focus is federated domain generalization under domain separation, where each client observes data from a single domain and the goal is OOD performance on unseen domains. We therefore compare primarily to methods proposed for federated DG and to FedAvg variants that control communication. We include centralized FRM as contextual upper/lower reference points, but many centralized DG objectives require multi-domain access to compute cross-domain penalties, which is not directly available in the domain scparation sctting. In Section 5.1, explore the impact of the number of clients and the elfect of Stage 2. For all these experiments, except. for the ablation study, we report average performance as well as the standard error over 5 different runs. We explain the most important experimental settings here, but refer the reader to Appendix D for more details. For FedLOE, each client holds out a fixed fraction its local data as Drune, which is disjoint from Dirain for Stage 2.2; all reported in-domain and DG results are computed on the standard held-out test splits.
Table 2: PACS: We show that our methods outperform centralized SGD. Our methods, as well as FedAvg with infrequent communication, are more stable.
DG accuracy (by domain)
Test Domain
Photo
Art
Cartoon
Sketch
Average
Mini-batch SGD
93.37 + 0.52
81.98 + 1.71
77.72 + 0.67
75.98 + 2.92
82.26
FedDG
96.67 = 0.28
84.40 : 0.92
75.77 + 0.85
74.97 + 1.74
82.95
FedSR
92.63 0.81
79.73 + 1.19
73.90 3.08
69.93 + 0.10
79.05
FedGMA
97.82 + 0.75
88.17 + 1.01
77.40 + 0.94
79.30 + 0.32
85.67
Our FedAvg
97.70 = 0.00
88.98 = 0.00
78.63 + 0.01
78.61 + 0.02
85.98
FedLOE- IRM
97.67 0.77
86.04 + 0.72
80.01 1.01
79.01 + 0.98
85.68
FedLOE- Fish
98.25 1 0.25
88.62 1 0.59
78.53 1 0.72
81.01
0.69
86.60
FedLOE-REs
97.75 0.94
86.66 1 0.71
78.87 0.55
79.44 1 0.70
85.68
5.1 Main results
We present the experiment results of our methods as well as the baseline methods PACS in Table 2, OfficeHome in Table 3, and IWildCam in Table 4. We observe that FedAvg outperforms centralized ERM on all three datasets. This behavior is consistent with prior observations that domain separation is a difficult regime for existing federated DG baselines (Bai et al., 2024), and it highlights the importance of mechanisms such as controlled communication that preserve useful client/ /omaii diversity. This validates our thoory that frequent communication across different domains leads to overfitting to spurious features, and our FodLOE method with both stages is comparable to or better than Fed on both the PACS and OfficeHome datasets.
9

---

Published in Transactions on Machine Learning Research (03/2026)
Table 3: OfficeHome: We show that our methods outperform centralized SGD, Our methods, as well as FedAvg with infrequent communication, are more stable.
DG accuracy (by domain)
Test Domain
Art.
Clipart
Product
Real-world
Average
Mini-batch SGD
60.06 2.01
50.32
1.25
73.20 +: 1.49
75.63 1.85
64.80
FedDG
58.58 0.74
50.92 0.89
75.41 1:
75.97 0.98
65.22
FedSR
1.54 1 0.01
1.53 1 0.01
1.52 1 0.01
1.53 1 0.01
1.53
FedGMA
58.19 + 2.74
3.65 42.33
68.90 + 1.64
71.70 2.07
60.28
Our FedAvg
63.70 0.01
50.71 + 0.02
76.74 0.00
78.74 0.01
67.47
FedLOE-IRM
65.55 + 0.83
51.23 0.45
76.10
78.95 1.03
67.95
FedLOE-Fish
64.79 + 0.59
50.79 0.39
76.93 + 0.44
78.4
67.72
FexEOE-REx
64.01 = 0.29
50.52 + 0.38
77.32 + 0.34
78.10 0.42
67.49
This suggests that in some cases, our second stage of robust ensembling can improve performance. Further- more, we achieve the highest performance on the leaderboard of the WILDS Benchmark using ResNet-50 without targeted augmentation in this client heterogemeity setting using infrequent-commuiccie It is also worth noting that our MaLLOE-IRM, FedLOE-Fish, and FedLOE- REx methods exhibit less ef- feetiveness compared FedAvg. This could be attributed to the fact that IRM, Fish, and REx did not demonstrate significant effectiveness in the cantralized setting, as reflected in the WILDS benchmark leader- board (Koh et al., 2021).
FedAvg.
0.81 PACS with cartoon as test domain
0.66
Office Home with Art as test doman.
0.65
6.95
0.44
0.77
0.52
0.2
0.00
0.0
25
Number of Shenis 150
20
102
175
Hurmorr of Cinrs
Conmoncation
Commuricaton
Figure 3: OOD accuracy with changing num- Figure 4: DG accuracy per communication the ber of clients C. The second panel is zoomed-in PACS and OfficeHome datasets; the second stage of view of the first figure, showing the details of cach our method starts at 12 and 24 respectively (bluc). method.
The Influence of the Number of Clients C: In this study, we investigate the performance of various algorithms in the FL context as the number of clients C increases on PACS. As shown in Figure 3, OOD accuracy is plotted against the mum- ber of clients, where each client only contains samples from single domain. The results indicate that: 1) the performance of FedSR degrades the number of clients increases. 2) In con- trast, proposed method, FedLOE, not only achieves higher accuracy but also exhibits greater robustness to larger client number settings. We emphasize that our client-scaling plot in Figure 3 is intended to highlight this observed brittleness in a reprosentative baseline (FedSR); we do not claim that all prior foderated DG methods necessarily degrade with increasing C under every implementation and tuning.
Table 4: IWildCam: Our methods outper- form centralized SGD. Our FedAvg with infrequent communication performs better and is more stable.
Methods
Fl score
Mini-barch SGD
31.61 = 1.88
FedDG
31.15 + 1.72
FedSR
0.01 + 0.00
FedGMA
26.05 2.24
Our FedAvg
$3 03.02 1 0.61
FedLOE-IRM
31.02 = 0.65
FedLOE-Fise
31.14 0.91
FedLOE-RExx
31.96 0.76
10

---

Published in Transactions on Machine Learning Research (03/2026)
In Figure 4, the DG accuracy per commmication is plotted, illustrating the performance boost over Fed on both the PACS and OfficeHome datasets. Specifically, the results demon- strate that FedAvg achieves the same accuracy with the same number of communications and computations as our FodLOE-IRM, indicating that FodAvg converges.
The Effect of Stage 2:
6 Related work
In this section, we discuss some related approaches.
6.1 Federated domain generalization
Federated domain generalization (Federated DG) has emerged as natural extension of centralized domain generalization, especially suited for real-world applications where data is inherently distributed across mul- tiple domains. In such settings, the need for robust generalization to unseen domains is both intuilive and essential. However, Federated DG introduces two key challenges. First, the presence of non-iid (heteroge- ncous) data distributions across clients complicates model convergence, even before considering discrepancies between training and testing domains. Second, strict privacy requirements and communication limitations often preclude direct data sharing across domains, thereby rendering many traditional, centralized domain generalization approaches infeasible in federated settings. While standard federated learning algorithms have made significant strides in handling client heterogeneity, their primary objectives fundamentally differ from those of federated domain generalization. Well-established methods such as FedProx (Li et al., 2020) and SCAFFOLD (Karimireddy et al., 2020) are specifically designed to improve optimization stability, handle client drift, and ensure convergence when local data distributions diverge (e.g., under label skew). However, these methods optimize primarily for in-domain performance across the participating training clients; they do not directly encode domaín generalization objectives or explicitly encourage representation invariance across domains. Consequently, they are not direct competitors for evaluating out-of-distribution generalization to completely unseen target domains. For this reason, our empirical evaluations focus primarily on methods explicitly designed to tackle the federated DG setting, such as FedDG (Liu et al., 2021), FedSR (Nguyen et al., 2022), and FedGMA (Tenison et al., 2023), which actively attempt to bridge the domain gap under privacy raistsaints. We view standard heterogeneous FL optimization techniques as complementary to our goals, as they could potentially be integrated alongside DG objectives to further stabilize local training.
6.2 Model ensemble
Model ensemble methods aim to improve generalization by combining predictions from multiple models, often trained under varying conditions or on different subsets of data. In the context of Federated DG, ensemble methods have been leveraged to account for domain shifts without requiring centralized data shar- ing. For instance, FedEM (Marfoq et al., 2021) ensembles personalized models from different clients to form more generalizable global predictor. FedCE (Cai et al., 2023) continues this line of work by introduc- ing a contrastive ensembling framework that encourages consistency across local models while maintaining domain-specific diversity. These approaches highlight the strength of ensemble- based strategies in mitigating the effects of domain heterogeneity and enhancing robustness to unseen target domains. However, ensem- bling can incur increased computational and communication overhead, particularly in resource-constraneel federated environments, making scalability a key concern.
6.3 Personalization and fine-tuning in FL
Personalization aims to improve per-client in-domain performance by combining a shared global model with client-specific adaptation. A common and strong bascline is to run FedAvg and then finc-tune locally on each client, which has been widely observed to perform well in heterogeneous FL. Representative personal- ization approaches include meta-learning-bre representation methods such as FedRep (Collins et al., 2021). Recent theory also helps explain why FodAvg followed by local fine-tuning can learn transferable representations across tasks/ /clinnss (Collins et al., 2022). Our setting is complementary: we study domain genernlization to unseen domains under domain separation,
methods like Per-FedAvg (Fallah et al., 2020) and shared-
11

---

Published in Transactions on Machine Learning Research (03/2026)
without assuming access to target-domain data for adaptation at test time. Nevertheless, personalization and DG are closely related, and combining FedT.OF with post-hoc local fine-tuning is an interesting direction for future work.
6.4 Benefits and limits of local steps
The role of multiple local steps in FedAvg//lcal SGD has been studied extensively from an optimization perspective, including conditions under which local updates help or hurt convergence under heterogeneity (Woodworth et al., 2020; Patel et al., 2024). Beyond optimization, recent works suggest that local updates can improve representation learning and transfer, helping explain the empirical strength of FedAvg with fine-tuning (Bao al., 2024; Collins et al., 2022). These insights are related but do not directly imply our main observation for federated DG: in our conain-separation DG setting, OOD accuracy is not monotone in communiiaaton frequency and can peak at an intermediate number of communication rounds (Figure 1). Moreover, our linear SCM analysis explicitly models invariant vs. spurious features and shows how averaging locally overfit experts can downweight spurious directions, yielding improved DG risk relative to centralized ERM, see Theorem 3.3.
7 Conclusion and discussion
Our paper introduces novel mixture-of-experts perspective FL DG in the domain separation context, where each client is viewed as an expert. Our method, while preserving privacy, still shares model parameters with the server. We discuss privacy and potential leakage risks in the Broader Impact section, and note that standard defenses such as secure aggregation and DP can be incorporated. We could further apply differential privacy (Dwork et al., 2006) to strengthen privacy protections.
The key observation is that locally overfitting and then refitting can actually avoid collusion among experts that they do not find common set of spurious features noted in several recent works. Given our perspective and these observations, we design a novel federated DG method, FedLOE, that combines two stages. The first stage uses the natural and implicit combination strategy of federated averaging to learn a robust deep featurizer. The second stage focuses on making the linear classifier head robust by iteratively locally overfitting the heads and then optimizing the best linear combination of these heads using standard DG objectives. Given that the second stage employs standard DG objectives as subroutines, it offers flexibil- ity that allows us to integrate any future methods aimed at enhancing DG accuracy. More broadly, this work suggests that client heterogeneity in FT. may actually be beneficial for OOD generalization, in contrast to the usual assumption that data heterogeneity impairs FL methods. Finally, while our linear structural equation model provides clear theoretical intuition, and our empirical results demonstrale solid gains, extending this theoretical framework to less restrictive models remains an exciting avenue for future work. Additionally, broadening the empirical comparisons to explore how standard heterogencous FL optimization tochniques (e.g., FedProx, SCAFFOLD) might complement or interact with explicit DG objectives could inspire useful follow-up research in robust federated systems.
Broader impact
This work studies federated domain generalization under domain separation, aiming to improve out-of- distribution robustness without centralizing raw client data. Potential positive impacts include better gen- cralization in safety-critical and high-stakes deployments (e.g., medical imaging across hospitals or sensing across regions), where distribution shift common.
Privacy and security considerations. Our methods do not require sharing raw examples. However, like most FL approaches, they still involve exchanging model information that may leak sensitive attributes. In Stage 1, clients commuicate model updates/parament parameters in FodAvg, which can be vulnerable to inference attacks (e.g., gradient inversion membership inference) in adversarial settings.
12

---

Published in Transactions on Machine Learning Research (03/2026)
Stage 2 introduces additional communication that warrants explicit discussion. In Stage 2.1, clients may send their classifier heads to the server; while these heads are smaller than full models, they may still encode information about local data, In Stage 2.2, our two protocols make different trust assumptions: (i) Protocol A broadcasts client heads to peers and collects tuning-set predictions ploiiiss /logits and labels; sharing heads and predictions can increase exposure to inference attacks or reveal domain-speeific information. (ii) Protocol B assumes a trusted server and avoids broadcasting heads to other clients, but may require clients to send tuning-set features and labels the server, which may also leak private information. In practice, these risks can be mitigated by combining our approach with standard FL privacy mechanisms, such as secure aggre- gation, encrypted communication, differential privacy (DP) communicated updates/ predictions/ features, or restricting the tuning set size and granularity (e.g., sharing only aggregated statistics rather than per- example logits). However, we do not make any explicit formal claims of privacy in this paper and urge that caution be used for any privacy-sensitiv applications.
Acknowledgement
R.B. and D.I. acknowledge support from NSF (II--2212097), ARL (W911NF-2020-221) and ONR (N00014- 23 23-C-1016). Any opinions, findings, and conclusions or recommendations expressed in this material are those of the authors and do not necessarily reflect the views of the sponsor.
References
Martin Arjovsky, Lóon Bottou, Ishaan Gulrajani, and David Lopez-Paz. Invariant risk minimization. arXiv preprinl Xrv: 028993, 2019.
Ruqi Bai, Saurabh Bagchi, and David I. Inouye. Benchmarking algorithms for federated domain generaliza- tion. In International Conference on Learning Representations, 2024.
Yajie Bao, Michael Crawshaw, and Mingrui Liu. Provable benefits of local steps in heterogeneous federated learning for neural networks: A feature learning perspective. In Proceedings of the International Conference on Machine Learning, 2024.
Luxin Cai, Naiyue Chen, Yuanzhouhan Cao, Jiahuan He, and Yidong Li. Fedce: Personalized federated learning method based on clustering ensembles. In Proceedings of the ACM International Conference on Multimedia, pages 1625 1633, 2023.
Liam Collins, Hamed Hassani, Aryan Mokhtari, and Sanjay Shakkottai. Exploiting shared representations for personalized federated learning. In International Conference on Machine Learning, pages 2089 2099, 2021.
Liam Collins, Hamed Hassani, Aryan Mokhtari, and Sanjay Shakkottai. Fedavg with fine tuning: Local updates lead to representation learning. In Advances in Neural Information Processing Systems, volume 35, pages 10572-10586, 2022.
Cynthia Dwork, Frank McSherry, Kobbi Nissim, and Adam Smith. Calibrating noise to sensitivity in private data analysis. In Theory of Cryptography Conference, pages 265-284, 2006.
Alireza Fallah, Aryan Mokhtari, and Asuman Ozdaglar. Personalized federated learning with theoretical guarantees: A modcl-agnostic meta-learning approach. In Advances in Neural Information Processing Systems, volume 33, pages 3557-3568, 2020.
Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adaptation of deep networks. In International Conference on Machine Learning, pages 1126-1135, 2017.
Yaroslav Ganin, Evgeniya Ustinova, Hana Ajakan, Pascal Germain, Hugo Larochelle, François Laviolette, Mario March, and Victor Lempitsky. Donaannadversarial training of neural networks. Journal of Machine Learning Research, 17(5)): 1135, 2016.
13

---

Published in Transactions on Machine Learning Research (03/2026)
Muhammad Ghifary, David Balduzzi, W. Bastiaan Kleijn, and Mengjie Zhang. Scatter component analysis: A unified framework for domain adaptation and domain generalization. IEEE Transactions Pottern Analysis and Machine Intelligence, 39(7):1414 (100:444--130,, 2016.
lan Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial networks. Commumications of the ACM, 63(1)): 139-144, 2020.
Ishaan Gulrajani and David Lopez-Pase Paz. In search of lost domain generalization. In International Conference on Learning Representations, 2021.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pages 770 778, 2016.
Sai Praneeth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank Reddi, Sebastian U. Stich, and Ananda Theertha Suresh. Scaffold: Stochastic controlled averaging for federated learning, In International Conference Machine Learning, pages 1132-51433 2020.
Pang Wei Koh, Shiori Sagawa, Henrik Marklund, Sang Michael Xie, Marvin Zhang, Akshay Balsubramani,
Weihua Hu, Michihiro Yasunaga, Richard Lanas Phillips, Trena Gao, et al. Wilds: A benchmark of in-the-wild distribution shifts. In International Conference on Machine Learning, pages 5637 5664, 2021.
David Krueger, Ethan Caballero, Joern-Henrik Jacobsen, Amy Zhang, Jonathan Binas, Dinghuai Zhang, Remi Le Priol, and Aaron Courville. Out-of-distribution generalization via risk extrapolation (rex). In International Conference on Machine Learning, pages 9815-58201 2021.
Da Li, Yongxin Yang. Yi-Zhe Song, and Timothy M. Hospedales. Deeper, broader and artier domain generalization. In Proceedings of the IEEE International Conference on Compuler Vision, pages 5542- 5550, 2017.
Da Li, Yongxin Yang, Yi-Zhe Song, and Timothy M. Hospedales. Learning to generalize: Meta-lcarning for domain generalization. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 32, 2018a.
Haoliang Li, Sinno Jialin Pan, Shiqi Wang, and Alex C. Kot. Domain generalization with adversarial feature learning. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pages 540005409, 2018b.
Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet S. Talwalkar, and Virginia Smith. Feder- ated optimization in heterogeneous networks. In Proceedings of Machine Learning and Systems, volume 2, pages 2020.
Quande Liu, Cheng Chen, Jing Qin, Qi Dou, and Pheng-Am Aoeng-Ann Heng. Feddg: Federated domain generalization medical image segmentation via episodic learning in continuous frequency space. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 1013-1023, 2021.
Othmanc Marfoq, Giovanni Neglia, Aurélien Bellet, Lactitia Kameni, and Richard Vidal. Foderated multi- (ask learning under mixture of distributions. In Advances in Neural Information Processing Systems, volume 34, pages 15434-15447, 2021.
H. Brendan McMahan, Fider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. Commuicctiion-ffiitee learning of deep networks from decentralized data. In Artificial Intelligence and Statisties, pages 1273-1282, 2017.
Krikamol Muandet, David Balduzzi, and Bernhard Schölkopf. Domain generalization via invariant fcature representation. In International Conference on Machine Learning, pages 10-18, 2013.
A. Tuan Nguyen, Ser-Nam Lim, and Philip H. S. Torr. Fedsr: Simple and effective domain generalization method for federated learning. In International Conference Learning Representations, 2022.
14

---

Published in Transactions on Machine Learning Research (03/2026)
Kumar Kshitij Patel, Brian Bullins, Ohad Shamir, and Nati Srebro. The limits and potentials of local SGD for distributed learning with heterogeneous data. In Proceedings of the 37th Conference on Learning Theory (COLT), 2024.
Alexandre Rame, Corentin Dancette, and Matthieu Cord. Fishr: Invariant gradient variances for out-of- distribution generalization. In International Conference Machine Learning, pages 18347-18377, 2022.
Elan Rosenfeld, Pradeep Kumar Ravikumar, and Andrej Risteski. Domain-adjusted regression or: ERM
may already learn features sufficient for distri-dist generalization. In NeurIPS Workshop on Distribulion Shifls: Connecling Melhods and Applications, 2022.
Shiori Sagawa, Pang Wei Koh, Tatsunori B. Hashimolo, and Percy Liang. Distributionally robust neural networks. In International Conference on Learning Representations, 2020.
Yuge Shi, Jeffrey Seely, Philip H. S. Torr, Siddharth N, Awni Hannun, Nicolas Usunier, and Gabriel Synnaeve. Gradient matching for domain generalization. In International Conference on Learning Representations, 2022.
Schastian U. Stich. Local SGD converges fast and communicates little. In International Conference Learning Representations, 2019.
Baochen Sun and Kate Saenko. Deep coral: Correlation alignment for deep domain adaptation. In Computer Vision-ECCV Workshops, pages 443 4830 2016.
Irene Tenison, Sai Aravind Sreeramadas, Vaikkunth Mugunthan, Edouard Oyallon, Irina Rish, and Eugene Belilovsky. Gradient masked averaging for federated learning. Transactions on Machine Learning Research, 2023.
Hemanth Venkateswara, Jose Eusebio, Shayok Chakraborty, and Sethuraman Panchanathan. Deep hashing network for unsupervised domain adaptation. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pages 5018-500 5027, 2017.
Yoav Wald, Gal Yona, Uri Shalit, and Yair Carmon. Malign overfitting: Interpolation and invariance fun- damentally at odds. In NewrIPS Workshop on Distribution. Shifts: Connecting Methods and Applications, 2022.
Blake Woodworth, Kumar Kshitij Patel, Sebastian U. Stich, Zheng Dai, Brian Bullins, H. Brendan McMahan, Ohad Shamir, and Nati Srebro. Is local SGD better than minibatch SGD? In Proceedings of the 37th International Conference on Machine Learning (ICML), 2020.
Liling Zhang, Xinyu Lei, Yichun Shi, Hongyu Huang, and Chao Chen. Federated learning with domain generalization. arXio preprint 1rXiv:2111.1044 2021.
Kaiyang Zhou, Yongxin Yang, Yu Qiao, and Tao Xiang. Domain generalization with mixstyle. In Interna- tional Conference on Learning Representations, 2021.
15

---

Published in Transactions on Machine Learning Research (03/2026)
Appendix
A Proof of Theorem 3.3
Proof. At the solution &, we have
0L
0,
00g
85.
0.
80g
0g
For all invariant features, we have
0L
2E
xr x α y)]
d0z
2E [ (x Ta-aaa atas - E2]
2B
[x â [[]a
2E
[ (āz ar) + a5ās) 0.
(13)
therefore, holds that which is equivalent to For all spurious features, we have
E [raa (ar- āz)l E ás
(14)
E E (az-āz)) āz)]
E E aga as)
(15)
0L
2TE [ (r: â - y)]
Öās
as
2E [ (a â - 2I arT E2]
2EE
[ ( (σiz az) + [s)]s]] - 2E [s] arastz + Eas + E3t2]
(b)
2E [xs(x) (āz az) +zJāas] 2E [5
2 [yas + (5T (ōzz - ax) + (yos + Ea)) [ss] 2F [ ass
21E [I aras ( (²z -x)) + 2E [[ az + aasa EGEs ās} 2E [2as]]
0,
where (b) follows from the independence of &1, & and Es. Therefore, we have
E T aras { (āz az))+++2 (ET az Ex) ass + EsE as) =E =E
(16)
Combining (15) and (16), obtain
E E (az âz araāās) 00
(17)
E [ET or [&T (āz az)] 3 as] + + [[<T ar Pasa + E3E as] 0.
Solving this system yields
{
ar ār + arajās.
E ((T az)² ajās + 2 as = E [[<TT oz) &a)asa + E3E ās)
(18)
Using the assumption E[G,T]] Eg,
16

---

Published in Transactions on Machine Learning Research (03/2026)
{
ag ã + azazās, (o) ة as Ezãs.
(19)
(19) suggests that our assumption on the noise is mild. As long as either ε2 or Es depends on the environment, the regressed spurious feature &g will also depend on the environment. In particular, given that &3 centered i..d. sub-Gaussian noise, we have
{
at āz + aragās.
(20)
(o] ag a as = ozās.
which implics
02 as 1
agās
(21)
02 as 2 + 03
Combining (21) with (20), we obtain
o σo
10,
as
(22)
of os ãs
P+o
as
Therefore, we have
0
a
aas
|as
Regressing solely on the invariant correlations
2
yiclds
JEI
à arg min E [|y atal? oiERiiI
(23)
aL = x ( y)1
da
E [ x az E2]
[ (a- - aτ) = 0.
(24)
where (a) follows from the independence of &1 and &2. and the fact that both E and centered.
Therefore, â = [azz O)T.
Consider using ERM to solve (6); i.e., by using all the data (x;, vi)ee(kk. we obtain
of = L Ey ((m.)) (2())) n = : 1
k
(1)
(25)
K
where we assume 17 satisfies Pn i) 1 for all i [K]. Note that this is a reasonable assumption given that the amount of data from different distributions is the same. If the training data sizes differ, the distribution of can be adjusted accordingly. Similarly, we have
K
03
(1).
(26)
K
Therefore, for ERM, we obtain
K
o0(
FRM
+ E ビm= o(
(27)
> o
2
FRMM
Σ i-1 o3(1)]] as 1| Σ df(( 0s.
17

---

Published in Transactions on Machine Learning Research (03/2026)
If each client locally uses their own domain data to solve (6), and we then perform one-shot averaging the solutions, we obtain
(j00)sss P+o az,
(28)
OSA
2 ff P+o as.
Comparing the solutions of ERM and one-shot averaging, when d3 d3, we have
Σ Σー of Σ dz())|ass 2 +
K
Σ
03
Σ (o2(1)
i=
K ass 1 P/O +1)
K ==
(o)
as 12/03 +
1
(29)
K
as 2 + o3
where the incquality follows from the fact that the harmonic mcan is less than or equal to the arithmetic mean. Therefore, it holds that
a
ERM
< Q7
COA
(30)
Furthermore, we have
Σ K
(i)
K
2i=1 03
K Σ
o2(1) as| 2 + Σ i=
k
03
Li=1
K
(i)sss 1 + 2i=1 o.
aas 1
K
K
02(1) as 2 + +o²
|as|
1 K
(1
(31)
K
02
as
+ +3
which implics
FRM OSA
as
(32)
Given that (2(1) for i E [K] are distinct, conclude that the above incquality is strict; i.c.,
ERM
QOSA and a
ERM
a
(33)
Recall that a² = [[zz, TT then (33) implics
| EIM az 2| > |
OSA
arl 2 and
BRM
0s > ||
OSA
0s
(34)
Therefore, we have
ERM
2212
-I@aRM
az 12 + a
BRM
0s 1
||aOSA
09
+
||a
OSA
0s
1SA
a'll ||.
(35)
addition, we have
K
1 o
K
o
Σ
2 Σ
K 03
o3
1=1
18

---

Published in Transactions on Machine Learning Research (03/2026)
where 9i for i E [K] is defined as
A >
9:
As
where
A A
03(1)/03
as
+1
K
k=
Therefore, the robust-to-spurion ratio is given by
a
FRM
a3
az
FRM
||
K 1 (i
K
a
S
|
|
K ビ
|az||
K (
aas
4
1 Σ 1==
4
|x
K
A
[as
K
1
ag
D K
ass
K 1=
A
|
(36)
| |@OSA
B Additional proof that OSA will have better DG accuracy than ERM from linear example
Corollary B.1. Under the assumptions of Theorem 3.3, then the DG risk of ERM is higher than for OSA:
(( [(V 1u--a x T@ > E(Da))Dan 11y- - x x 2
(37)
Proof. We first derive the covariances for a single domain:
Σx oîf
(38)
Es.s Ei( E + ta)(f &2) (
(39)
aTB(GiET + EkK]asa + E[
(40)
(llaz 1 o +
(41)
Σr,
E|rra3) E(G((( + falas (
(42)
Eigif oro + E16 + 6.4
(43)
=ofara..
(44)
19

---

Published in Transactions on Machine Learning Research (03/2026)
We now decompose the test-domain risk where the domain index is not in the training dataset, i..., & [C].
as follows:
E (maa))DDe (yy x [i)²
[ + (2(1) a]
-((ey) Dtont
(0 a + (2( x xTa)
[ (a* [] + ož(1)
((e.y))Denn
(a*" -a--a)
@) xx (a* (a- -a) + +1)
(a* @) E (aan))Denn xD [1 - a) - (1)
(45)
(ax -a) Emnoo (ox* - ② + (1)
(46)
(1 wrraa
Σxz((1) Ei,((i)
(1-0))o 12)
+ +
(47)
Es,((i) Es,s(i)
u
(1 wr)az
of(iTzz
(
(1-dz)ar
WSQs + +
(48)
osas
of(i)asaz
(Iaz]|20f(1)
oZ asas ošIs
iz) oz [1) + (2(())+20())) oz)1 0ss (B)(()
(49)
(11aa (1) + o1))) (as0)) + + as o1)
Note that the algorithm specific terms are only via uI and ús because all other terms are constant w.r.t. the learned parameters ã. Finally, from the previous theorem proof ((33)), we know that
0 Σ (1 - 07 OSA < (1 W7
ERM <
(50)
0 < wsg
OSA
< u 1
ERM
(51)
Combining these facts with Eqn. 49,
easily arrive at the result:
E E(m.g)~D.nn
(V x (w-=TaasAy) TO
(52)
= (1 17 OSA | + 2:0 OSA w7 (0SA | as (1))(1)
(53)
+ (w) OSA 1 ((1(ar) (of(1) + ( 12 + o :) + o2(1)
< (1 a = | |10001 2m (1
(1- ||
(54)
+ (WERM) )²(( ((||ar] (0)(( - - as)sllss +ogl ( + +z(
((rgg))-Duns [( T
(55)
C FedLOE algorithm pseudo-code
We present our training method FedLOE in Algorithm
C.1 Penalty choices in the FL setting
For the invariant manager training, in every iteration t E T2, we recall (11) for convenience:
e ( Se (oeetee 5² ) + A.r
( LoeWe.v
arg min dERO
ΣN(r Pl vice)
where Btune p( denotes an empirical average over the tuning set Drune on client c, À is the penalty parameter, and r(.) is penalty term to induce invariance, which be chosen using different methods in (3). For
20

---

Published in Transactions on Machine Learning Research (03/2026)
Algorithm FedLOE: Federated DG via Local Overfitting and Ensembling
Input: Training datasets { Me
train 1 Ic=1 and fine- tinetuning datasets fat Ve Ic=1 on each of the C
clients; iterations for stage 1 and 2: Ti,12: T2; communication indices Ti and T2 for stage and 2; local computations E1 and Eg.
# Slage 1: Few shot Federaled Averaging to Learn 0: (4)
# Stage 2: Robustifying the Linear Classifier Head
for t E {Ti + +1...., T } do
if $ T2, # Stage 2.1: Locally overfit classifier heads, distribute to all clients
[Client] Ve, vol Mini-batch SGD (0, Wt E2)
[Client Server) Vc, de
{Server Client}
Elsc if t E T2, # Stage 2.2: Robust training # Slep 1 Server collect predictions on fine lunving dalasel
[Client] Ve, { Vé.e - gg(a!!me) ( C'ECI
{Client + Server) Ve, 106 Jere(c]
# Step 2 Train invariant manager. {Server] solves Equation (11) to obtain Z end for Return: (0, (0th, TT
[Server + Client} = =. D
example, to encourage Ee dyDe, to be simultaneously optimal for different client domains e [C], we choose Fdg as Jirmv: which enforces the gradients be zero (Arjovsky et al., 2019):
tune "
Tirm (0, 10)
Σ Epp
Vol==1of
Σ WesDDass ve
Alternatively, we can set Tdg to Ffish;, which aligns the gradients as in Fish (Shi et al., 2022):
yo tune ) Vol
c
))
Fiss (0, +b)
Σ
Epons
Vol
Σ
ss Vo. tume
exte
Furthermore, we can set Fdg to FREx: which reduces the loss variance 0:
]
FREx FREx(0,
ΣV Vp.
Σ oj Yasi Ve
tune
D Additional experimental details
D.1 Dataset
PACS Gulrajani and Lopez-Paz (2021) is an image dataset designed for domain generalization classification, comprising a total of 4 distinct domains: Photo (1, 670 images), Art Painting (2, 048 images), Cartoon (2, 344 images), and Sketch (3, 929 images). Within each domain, there are categories present. For this dataset, we use each domain as the test domain in each run, and evaluate the average classification accuracy.
OfficeHome Gulrajani and Lopez-Paz (2021) is also an image dataset designed for domain generalization classification containing 4 domains namely Art (2, 427 images), Clipart (4, 365 images), Product (4, 439 images), and Real-World (4, 357 images) images. Within each domain, there are 65 categories present, which makes it a harder setting than PACS. Similar to PACS, for this dataset, we use each domain as the test. domain each run, and evaluate the average classification accuracy.
IWildCam Koh et al. (2021) consists of diverse collection of wild animals captured by 343 camera traps situated across various natural habitats worldwide. This dataset. offers multi-class classification task,
21

---

Published in Transactions on Machine Learning Research (03/2026)
comprising total of 203, 029 data samples representing 182 distinct animal species. Our primary objective is to achieve high classification accuracy, particularly for rare animal species. Due to this goal, we employ the macro-Fl score as our chosen metric.
D.2 Data partitioning
For our experiments on the PACS and OfficeHomee datasets, adopt the following evaluation protocol. We select one domain as the test domain for each experiment, while the remaining 3 domains serve as the training domains. First, we allocate 5% of the training dataset for validation and another 5% for in-domain testing.
Subsequently, we divide the remaining 90% of the dataset among the clients, ensuring that each client exclusively contains data from a single domain. For each client, we utilize 90% of its data for training (wtrain, y yrrrin), while the remaining 10% is designated for tuning as (w ytume). This partitioning strategy enables us to train and evaluate our models on distinct domains while maintaining a balanced and controlled experimental setup.
On the IWildCam dataset, we adhere to the official splits method provided by the WILDS benchmark (Koh al., 2021) for consistency and comparability. Furthermore, we partition the training data among the clients to ensure that each training domain is present in only one client. This deliberate distribution strategy guarantoes maximum domain heterogencity within the client sct.
D.3 Neural network structure
In all our experiments, we use ResNet-50 (He al., 2016) as our featurizer, excluding the last fully connected layer. The leaturizer produces an output of size 2048. Subsequently, each client maintains its own linear classifier. In stage 1 of our approach, the linear classifiers belonging to each client are averaged using the default weights of FedAvg. In stage 2, we advance this by averaging the linear classifiers using learned weights.
D.4 Model selection, early stopping and other hyperparameters
In this section, report the hyperparameters used in the experiments. Please refer to Table 5.
Table 5: Hyperparameters for our infrequent FedAvg and our FedLOF
Dataset
PACS
Officellome
IWildCam
Hyperparameter searching
Best performance held-out validation datasct
Stopping criterion
Fix iteration
Best performance on held-out validation dataset
Batch size
128
128
128
Number of iterations at stage
256
256
5120
Stage total communication
8
8
80
Stage 1 optimizer
ADAM
ADAM
ADAM
Stage learning rate
0.00005
0.00003
0.00005
Stage 2 mmm of iterations (FedLOE only)
8
8
320
Stage 2 total communicaiion (FedLOE only)
5
Stage 2 optimizer (FedLOE only)
SGD
SGD
SGD
Stuge 2 learning rate [FedLOE only)
0.05
0.05
0.05
Stage 2 weight docay (FedLOE only)
0.0005
0.0005
0.0005
IRM è (FedLOF. IRM only)
100
100
100
Fish meta Ir (Fe.ILOF.Fiie Fish only)
0.05
0.05
0.01
REx (FedLOE-Ree Rex only)
100
100
100
E
Code repository for reproducibility
Please see the code here:
22


Reviews:
[{"review_id": "ndBPYxXqXT", "reviewer": "Reviewer_Jrjg", "summary": "", "questions": "", "limitations": "", "rating": null, "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": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper identifies and studies a phenomenon where federated-style optimization can outperform centralized optimization methods when evaluated on domain-generalization capabilities.  While surprising, the authors offer theoretical insights into this behavior through a rigorous analysis on specially-constructed linear regression problem.  They go on to propose a new class of algorithms which they coin \"FedLOE-* (Locally  Overfit Ensemble)\".  These algorithms consist of two stages, and two steps within each stage, with the variants differing primarily in the specific domain-generalization objective used to linearly combine classification heads.  The empirical results demonstrate that both FedAvg and the proposed FedLOE methods achieve better DG scores consistently across 4 tasks to further support their theory on linear models.\n\nStrengths:\n1. The paper identifies an interesting and surprising new phenomenon that should be of interest to the broader FL community.\n2. The paper conclusively demonstrates that their finding is not merely a coincidence, through theoretical analysis on linear models and empirical analysis with non-linear models on a benchmark suite of 4 datasets.\n3.  The paper is organized well and easy to follow and read.\n\nWeaknesses:\n1. The paper would benefit from explicitly discussing the connections between domain generalization and over fitting as it's traditionally understood, potentially doing more experiments.\n2. FL is a rather indirect way to solve the DG problem and the algorithm likely doesn't get at the root of the problem. While it's interesting that it outperforms the SGD baseline, it's hard to imagine that it would be better than a centralized algorithm with an appropriately regularized objective.\n3. The paper is missing key details in important places that make it impossible to fully understand some of the results presented without additional context (see below).", "claims_and_evidence_yn": "Yes", "claims_explanation": "The authors have done a good job to study this problem on a toy problem theoretically, and on real problems empirically. I am convinced by the evidence provided, even though I am someone surprised by the phenomenon itself.", "audience_yn": "Yes", "audience_explanation": "This research should be of interest to the broader federated learning community. The authors identified a phenomenon that could be the basis for future research.", "requested_changes": "* I was confused by the parenthetical on page 2 \"which corresponds to communication at every epoch\".  Should this be every \"step\" or is there an implicit assumption that the ERM problem is being solved with a full-batch method?\n\n* The motivating observation in section 3 seems to be missing some details in the experimental setup.  Is the in-domain = train dataset and dg = test dataset, or are they both different test datasets and if so how are they constructed?\n\n* Stage 2.2 needs some clarifications:\n- What is the tuning dataset?  Does each client randomly split their data into (train, tune) subsets?\n- FL algorithms are largely motivated by their nice privacy properties: updates from each user are aggregated immediately on the server and incorporating into the model before broadcasting back to clients. The mechanism you describe of \"broadcasting all the linear classifiers to each client\" seems to break this.  Since this is fundamental to the proposed algorithm variants, it is somewhat concerning.", "broader_impact_concerns": "No concerns.", "additional_comments": "* Stage 2.2 \"broadcast the all the\" typo.  It also seems to be a weird table placement in the middle of the paragraph.", "extra_scores": {}, "extra_text": {}}, {"review_id": "LnaTPuxT0c", "reviewer": "Reviewer_QuVr", "summary": "", "questions": "", "limitations": "", "rating": null, "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": "", "claims_and_evidence": "No", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The authors consider the problem of federated domain generalization. In domain generalization, the goal is to train a model using data available from a specific set of domains, such that it performs well on data from unseen domains. To formalize this in the federated setting, they consider each client containing data from a specific domain, and the goal is to perform well on unseen clients.\n\nThe most popular approach in centralised settings is to perform empirical risk minimisation on the available clients.\n\nFor the federated settings, the authors first show theoretically for linear models with spurious and invariant features, and each client having subgaussian features with different variances that empirical risk minimization, which corresponds to FedAvg with number of local steps $=1$ performs worse than running local steps on each machine and then performing a single step of averaging (OSA), see Theorem 3.3.\nThey also verify this fact through experiments (see Fig 1a), where FedAvg with $>1$ local steps, performs better than FedAvg with $1$ local step, keeping the computational budget fixed to 256 epochs.\n\nThen, they formulate a family of domain generalization algorithms, called FedLOE, described in Table 1. The key steps for this are to locally overfit parameters on clients followed by ensembling parameters from different clients. Their instantiations of FedLOE first train all parameters via FedAvg on all clients, then combine the classifier heads via a Domain Generalization objective used in previous works to obtain \"IRM, Fish\" and \"REx\" variants each of which uses an existing regularization.  Additionally, FedAvg is also an instance of their FedLOE family.\n\nThey compare FedAvg, 3 instances of FedLOE algorithms and existing baselines on Federated Domain Generalization (FedDG, FedSR, FedGMA) on 3 domain adaptation datasets -- PACS, OfficeHome and IWildCam, where the training domains are separate from test domains. Their algorithms (including FedAvg) obtain the best performance across baselines, although the difference between their new FedLOE variants and FedAvg is not large, and often FedAvg performs the best overall. \n\n\n## Strengths --\n1. **Novel Theoretical Result**: Comparison of One-shot averaging to simple ERM is a novel contribution. I don't think one-shot averaging has been previously analyzed in this theoretical setup exactly. Their theoretical setup, although adapted from existing works, is pretty general. The proof uses the key AM-GM-HM inequality that harmonic mean $\\leq$ arithmetic mean, where OSA yields HM and FedAvg with local steps $=1$ yields AM.\n2. **Well-formulated Insight about benefit of local steps**: If local steps are beneficial or not is an active open question in federated learning, and it has still not been fully resolved for all aspects like optimization and generalization. Atleast for the case of domain generalization, the authors have concrete evidence of it being useful. This insight might be used in generalization analysis of federated learning, not just for federated domain generalization.\n3. **FedLOE algorithms outperform baselines**: FedLOE algorithms (including FedAvg) outperform all existing Federated DG baselines. They also provided the full code and implementation details. \n\n## Weaknesses --\n1. **Theoretical Model**: The authors don't really use the fact that their features are sub-Gaussian which they state in point 1. of  Assumption 3.1. From the proof, it looks like they only need to define the first and second moments of their variables. Further, they minimize the population risks directly, not the empirical risks, which would have resulted in conditions on sample complexity for the number of clients from a domain and the number of local datapoints on a client. Also, the only difference between different domains is that their features are from a subGaussian distribution with mean $0$ and different variance for different domains. Due to this they can use the AM-HM inequality, as they can compute the minimizer of population risk and OSA explicitly, and they only depend on noise and feature variance scales. While the same model can be extended to different domains with different means, which would strengthen the result.\n2. **FedLOE doesn't seem to improve on FedAvg by a lot**: The difference between the other FedLOE variants and FedAvg seems very small. So, I'm not sure what the benefit of introducing other FedLOE variants is. If they could have a setting where other FedLOE variants are strictly better than FedAvg, it would improve the paper. They could start with even a synthetic data setting. As for the baselines, I think more baselines especially those using personalization should be compared as they might even fall into FedLOE (for instance the meta learning approach in (Fallah et al 2020).  \n3. **Related Works Not Discussed**: Several related works are relevant to this problem but have not been discussed.\n     - **Personalization**: The topic of personalization in federated learning, where the goal is to perform well on each client locally while also using the data on other clients to learn a good global model is extremely relevant to this problem. Further, in practice, fine-tuning after FedAvg (Wu et al 2022) seems to be one of the best performing personalization approaches, and I suspect that it might perform better for domain generalization problem even here. \n     - **Existing work on benefit of local steps**: Note that precise conditions for benefit of local steps in FedAvg from an optimization perspective have been well studied (Woodworth et al 2020,  Patel et al 2024), but not discussed here. Further, even from a generalization (not necessarily domain generalization) perspective, (Bao et al 2024) and (Collins et al 2020) show that local steps actually help learning features. (Collins et al 2022) is very closely related to their theoretical model: they use a linear model $B^\\star w_j^\\star$ where $B^\\star$ is shared across clients (in this case across domains), while $w_j^\\star$ is specific to clients. I think the authors should provide a thorough comparison to these cases, and why their insight for domain generalization cannot be inferred from these papers, especially (Collins et al 2022). \n4. **Main Message of the Paper**: This is more of a personal opinion. When reading the paper, it seemed like the authors were trying to convey why local steps are necessary for federated domain generalization until the theoretical result. Then, they move on to introducing the FedLOE framework, which seems completely detached from the original insight and motivation of local steps being necessary. This makes it very hard to find the main message of the paper. If the main message is that local steps are necessary, then there should be more discussion and theoretical/experimental analysis of this claim. If the main message of the paper is to introduce the FedLOE family of algorithms, then the authors should show settings either theoretically/via experiments, where FedLOE algorithms outperform FedAvg significantly.  \n\n**References**\n- (Wu et al 2022) Motley: Benchmarking Heterogeneity and Personalization in Federated Learning.\n- (Fallah et al 2020) Personalized Federated Learning with Theoretical Guarantees: A Model-Agnostic Meta-Learning Approach. NeurIPS.\n- (Woodworth et al 2020) Is Local SGD Better than Minibatch SGD? ICML.\n- (Patel et al 2024) The Limits and Potentials of Local SGD for Distributed Heterogeneous Learning with Intermittent Communication. COLT.\n- (Bao et al 2024) Provable Benefits of Local Steps in Heterogeneous Federated Learning for Neural Networks: A Feature Learning Perspective. ICML.\n- (Collins et al 2022) FedAvg with Fine Tuning: Local Updates Lead to Representation Learning. NeurIPS.", "claims_and_evidence_yn": "No", "claims_explanation": "Overall, I think the paper has some contributions that are  well motivated like the local step insight and formulation of FedLOE family. However, the paper does not focus on either one of them concretely, and for each of them there are significant weaknesses : for local steps, sufficient comparison to existing works not provided, and for FedLOE algorithms, I'm not sure how they improve upon FedAvg.", "audience_yn": "Yes", "audience_explanation": "Yes, the topic of federated domain generalization is still an active research area in federated learning.", "requested_changes": "Please address the Weaknesses in the Summary section.\n\n- Apart from this, I think the authors could use some existing results from the benchmark paper on federated domain generalization which they cite (Bai et al 2023), where all these federated domain generalization baselines perform poorly under the domain separation assumption. These could be used to strengthen their argument.\n\n- ``Prior federated DG methods perform poorly when the number of clients is large\": This line is in the abstract and something similar has been repeated in the introduction, but the authors don't provide an explanation for this. Only FedSR performs poorly when we increase number of clients which is seen in Fig 2. What about other baselines?\n- I suspect Fig 1b is a hypothesis of what is happening. Can the authors somehow show that this hypothesis is true, especially the part about  $\\hat{\\theta}_{Fedloe}$ ? Does this figure follow directly from Theorem 3.3, for the linear model where FedLOE corresponds to OSA?  \n\n- The authors state that some of the baselines require sending image spectrum information to all clients, so are not private. But, their method especially Stage 2.2 (Page 7) also sends a small portion of predictions of each client to the server. Is this also not private? Either the authors should not dismiss baselines because of privacy, if their own methods are not private, or clearly specify how their methods incur less privacy loss. Stage 2.2 states that ``Notice that we only send m dimensional prediction which are much less than the original D dimensional\ndata\". Under this logic, the baselines which send image spectrum also don't send the entire dataset.\n- Please fix these Typos:\n  - Figure 1 has no caption. It's subfigures have captions. \n  - In the Theory section and it's proof in appendix, the authors use a notation $ \\alpha_1 \\leq \\alpha_2 $, for two vectors $\\alpha_1$ and $\\alpha_2$. Please specify that this notation implies that each coordinate of $\\alpha_1$ is $\\leq$ each coordinate of $\\alpha_2$, as there is no comparison operation between two vectors.\n  - There is an \"f\" before Appendix B.", "broader_impact_concerns": "There are no additional concerns about the broader impact of this paper other than those for any paper on machine learning or distributed learning.", "additional_comments": "- What is the difference between \"Our FedAvg\" and the vanilla FedAvg algorithm?", "extra_scores": {}, "extra_text": {}}, {"review_id": "ZWh5cgejjq", "reviewer": "Reviewer_E6f2", "summary": "", "questions": "", "limitations": "", "rating": null, "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": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper tackles federated domain generalization (DG), where multiple clients hold heterogeneous source-domain data distributions, while the goal is to train a global model that generalizes to unseen target domains. \nMethod-wise, the paper proposes FedLOE, a two-stage (“locally overfit, then ensemble”) framework that explicitly decomposes the model into a shared feature extractor and a linear prediction head:\n- Stage 1: run local SGD with infrequent communication (periodic averaging), intentionally allowing clients to overfit locally before averaging. According to the authors’ analysis, this yields feature extractors less tied to common spurious cues than fully synchronized training.\n\n- Stage 2: freeze the feature extractors, let each client overfit a local linear head, and then learn global ensemble weights for these heads by optimizing a DG objective (IRM/Fish/REx) using the prediction vectors on client fine-tuning sets, producing a single global ensemble head.\n\nThe main insight is to reframe infrequent communication as a DG benefit, supported by analysis in a linear causal model\n\nEmpirically, the framework delivers practical performance gains over standard FL DG baselines.\n\n**Strengths**:\n- The paper is generally clearly presented and  relatively easy to follow.\n- It offers a relatively new insight by reframing local overfitting (via infrequent communication) as beneficial for DG, supported by a linear causal analysis.\n- It presents a simple, modular design with a clear two-stage recipe that plugs into standard FL loops and DG losses\n- The proposed method shows empirical gains, generally outperforming FL-DG baselines.\n\n**Weaknesses**:\n- The novelty appears incremental, as the components (local SGD, head-only training, ensembling, DG regularizers) are established and the main contribution is their integration.\n- The theoretical analysis is confined to a simple linear setting, so its practical implications for deep models may be limited.\n- The discussion of FL relevance could be strengthened\n- The comparisons and ablations could be more comprehensive, with additional baselines and sensitivity analyses", "claims_and_evidence_yn": "Yes", "claims_explanation": "The paper’s central claim that \"local overfitting with infrequent communication, followed by a prediction-level ensemble of client heads, improves DG\" is relatively well supported. \n\nThe linear causal analysis aligns with the intuition, yet its scope is limited and doesn’t guarantee the effect for deep models. Empirically, the method often outperforms FL-DG baselines, which is encouraging, but the evidence would be stronger with (1) broader baselines (e.g., standard FL methods such as FedProx/SCAFFOLD, plus centralized DG baselines), (2) more comprehensive ablations/sensitivity (T₁/T₂, fintuning-set size, number of clients per domain), (3) statistical rigor (multiple seeds with means and standard deviations or confidence intervals).", "audience_yn": "Yes", "audience_explanation": "While the principled gap to prior work is modest, the practical combination of infrequent local-SGD with prediction-level DG-guided head ensembling may be interesting to parts of the FL (DG) community.", "requested_changes": "- While the experiments convey the main message that local overfitting plus ensembling can improve OOD performance under typical FDG evaluations, the baseline suite seems somewhat limited; consider adding a few key comparisons or briefly justifying why certain baselines (e.g., standard FL methods such as FedProx/SCAFFOLD and centralized DG baselines) are out of scope or inapplicable.\n\n- The connection to and distinction from general FL (with inherently non-IID client data) could be strengthened. In the design considerations, clarify which choices are driven by the domain generalization objective and what gaps remain in standard FL techniques (and potentially benchmark against stronger FL methods beyond FedAvg, reporting their performance in the DG setting) to better contextualize the contribution.\n\n- It would also help to more clearly articulate the distinction between this federated DG formulation and standard (centralized) DG, making explicit which design choices and analyses are driven by federated constraints (e.g., communication, aggregation protocol, privacy assumptions).", "broader_impact_concerns": "No \"Broader Impact\" section is included. While I don’t see major risks beyond the usual concerns in federated learning, the privacy implications may be stated more explicitly, e.g., discuss potential leakage from sharing heads/predictions. Currently this appears only briefly in the Conclusion/Discussion.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Decision: Accept (poster)