{'fraction_non_alphanumeric': 0.038393725680933855, 'fraction_numerical': 0.01091318093385214, 'mean_word_length': 4.851476342938456, 'pattern_counts': {'":': 0, '<': 2, '<?xml version=': 0, '>': 1, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
Mitigating Data Heterogeneity in Federated Learning with Data Augmentation

Federated Learning (FL) is a prominent framework that enables training a centralized model while securing user privacy by fusing local, decentralized models. In this setting, one major obstacle is data heterogeneity, i.e., each client having non-identically and independently distributed (non-IID) data. This is analogous to the context of Domain Generalization (DG), where each client can be treated as a different domain. However, while many approaches in DG tackle data heterogeneity from the algorithmic perspective, recent evidence suggests that data augmentation can induce equal or greater performance. Motivated by this connection, we present federated versions of popular DG algorithms, and show that by applying appropriate data augmentation, we can mitigate data heterogeneity in the federated setting, and obtain higher accuracy on unseen clients. Equipped with data augmentation, we can achieve state-of-the-art performance using even the most basic Federated Averaging algorithm, with much sparser communication.


Introduction
With the increasing computation power of edge devices, Federated Learning (FL) [1] provides a solution to data privacy constraints using distributed computation. In this paradigm, each user (or client) has a dedicated model that is trained locally and solely on the said user's data collection. During this process, a server is responsible for collecting these local models. By combining their parameters in a principled way, Federated Learning aims to generate a powerful centralized model.
One important aspect of this distributed setting is that each client's dataset has a unique manifestation, which generally reflects a fraction of data distribution of all clients. This implies that the samples of each client are not drawn from the combined distribution in an independently and identically distributed (IID) fashion (i.e. non-IID). For instance, a user's image collection may have certain subjects of interest, backgrounds, lighting conditions and viewpoints. With this considerable limitation, local models will vary substantially, and when aggregated, will render poor generalization [2].
The compartmental view of data share similarities with the popular field of Domain Generalization (DG) [3]. In this framework, instead of having observations drawn from a single distribution, we assume that data is originated from multiple domains. With this setting, the goal of DG is to find a robust classifier that generalizes beyond the domains seen during training. Different strategies in Domain Generalization have successfully improved out-of-distribution performance by modifying the objective function accounting for the different domains [4,5,6]. Drawing a connection to the federated setting, each domain may correspond to one or more clients, and thus many DG algorithms can be transferred to the FL setting (See § 3), though not always easily.
However, recent evidence [7] suggests that, with all conditions equal, when equipped with data augmentation, no algorithm outperforms the elementary Empirical Risk Minimization (ERM) algorithm. An intuitive explanation is that data augmentation acts as an invariance mechanism, discarding potential spurious correlations that are commonly encountered in each domain.
Motivated by the aforementioned results, and considering the data heterogeneity present in the federated setting, in this work, we show that data augmentation also poses a straightforward yet effective solution to mitigate these differences in FL. By selecting transformations that are pertinent to the overall distribution, we can apply data augmentation to each client data, generating more well-distributed and similar data distributions, while maintaining privacy.
Contributions. We study Federated Domain Generalization as shown in Figure 1. Each client owns a proprietary dataset, and the goal is to learn a model through federated learning, such that it generalizes both in-domain (ID) and out-of-domain (OOD). 2 The in-domain datasets are the private ones of the participating clients during training, and a out-of-domain dataset resides in a client that has not participated in the federated training process. This setting bridges the fields of Domain Generalization and Federated Learning. Motivated by the success of data augmentation in DG, we show that, from a causal perspective, data augmentation can alleviate the non-IIDness in FL, improving both generalization and convergence: • Generalization: with proper data augmentation, Federated Learning gives better OOD performance, and even the basic FedAvg algorithm compares favorably against state-of-the-art alternatives. • Convergence: Since the non-IIDness is mitigated with data augmentation, one can perform much sparser communication to achieve comparable performance as the centralized setting, with as few as two communication rounds.

The non-IIDness of Federated Learning
One of the main challenges of FL is data heterogeneity, i.e. the dissimilarity between each client's data distribution. Suppose D i is the data distribution for client i, then the non-IIDness means that D i and D j could differ significantly for i = j, while in conventional IID distributed learning, D i = D j for any i = j. The goal of federated learning is to find a global minimizer θ * of the total loss, Here (x, y) ∼ D i is an (input, label) pair drawn from D i , and is the loss function (e.g. cross entropy) that measures the discrepancy between the model predictionŷ θ (x) and the true label y. The model θ, usually parametrized as deep neural networks, takes an input x and output a probability vector for prediction. In practice, we draw samples from each distribution D i to estimate the expectation.
In the IID case, the gradient ∇f i is also an unbiased estimate of ∇F , implying that the minimizer θ * of (1) is also (approximately) a minimizer of the individual f i 's. To save communication cost, we can compute the one-shot average Table 1 summarizes the connection between DG and FL. As another example, GroupDRO [5] and AFL [13] follow the same principle of minimizing the worst-case client (or source domain): where ∆ := {λ ∈ R n : 1 λ = 1, λ ≥ 0} is the probability simplex. The probability simplex can be generalized to ∆(λ min ) := {λ ∈ R n : 1 λ = 1, λ ≥ λ min 1}, from which the DG algorithm called MM-REx [6] has been designed. We can similarly write down the federated version, that we call Generalized-AFL (Gen-AFL). Moreover, we use this connection to devise another algorithm, named Variance Minimization (VM), from [6]. Essentially, by replacing the gradient ∇f i with the pseudo-gradient, i.e., the sum of local gradient steps for client i, one can obtain the corresponding FL algorithm from a DG algorithm. However, not every DG algorithm can be easily transferable. For example, the popular IRM [4] algorithm suffers from gradient sensitivity after we move it to the FL setting. Full details of Gen-AFL, VM, and IRM can be found in Appendix A.

Data augmentation in domain generalization
The similarity between domain generalization and federated learning motivates us to bring insights from one field to the other. As explained in [7], with proper data augmentation, ERM has the state-of-the-art performance, compared to a large variety of DG algorithms. Does the same conclusion hold for federated learning, with the private source domains and sparse communication? To answer this question, we study the effect of data augmentation in § 4.
Due to the similarity between DG and FL, we will use the word "client" and "environment" interchangeably when not causing confusion.

Data Augmentation for Federated Learning
Inspired by the success of data augmentation in domain generalization, in this section, we propose to mitigate the non-IIDness in federated learning from the data perspective. This is orthogonal to the algorithmic approach to federated learning, e.g. [13,12].
Our key observation is that data augmentation can reduce the heterogeneity of client distributions. As discussed in § 2, federated learning is quite challenging in the presence of non-IID data. By applying proper data augmentation to each local client, the client distributions are more similar and FL will suffer less from the non-IID issue.
Data augmentation is a widely used strategy in machine learning for increasing the diversity of samples. It often allows us to obtain more robust models with better performance effectively. In this work, we define data augmentation as transformations in the input space. More specifically, denote X as the input space, a data augmentation function T is defined as T : X → X . In supervised learning, data augmentation is typically constructed to preserve ground-truth labels.

Data augmentation with different environments: a causal approach
To better understand how data augmentation can improve similarity among different clients (environments), we propose a causal model that explains this approach, as illustrated in Figure 2. The key observation we make in the causal model is that all environments are generated from a common cause Z, which is a random variable. This random variable is sometimes called the semantic content [16], which is generated from the true label Y . For example, we could generate the digit images (Z) from digit labels (Y ). This data generating principle is also known as anti-causal [17].
However, individuals write digits differently, e.g. handwriting can vary. Besides, even the same person can write the same digit in different ways. Therefore, for a client i, we use a random variable ε i to account for these variations. The causal mechanism allows us to generate samples for each client, represented by random variable X i : This is known as the Structural causal Model (SCM) [18]. Note that Z and ε i are independent, and the causal mechanism g remains invariant across environments. For example, in Rotated MNIST [19] Z represents the original digit images, and each ε i represents a fixed angle. The function g is the rotation operation of Z with the angle information given by ε i .
So far, we have talked about the data generation process from the causal perspective. Now we discuss the importance of data augmentation in federated domain generalization, which is two-fold. First, we use data augmentation to homogenize different domains that contain the same semantic meaning but various environmental factors (e.g. the rotation angle in Rotated MNIST, the style of painting the same object). Essentially, if we have a general sense of the underlying SCM, then we can take advantage of it to generate new samples, by applying proper data augmentation. In the ideal case, the data augmentation T for each X i := g(Z, ε i ) induces a transformation T for ε i , without modifying Z. Eventually, we want all the transformed environmental variables T (ε i ) to match an environment-agnostic variable ε. With clients distributions more aligned, we have less data heterogeneity and therefore, better convergence.
Augmentation also has the role of promoting generalization. The data augmentation T can expand the training distributions, potentially incorporating unseen domains, thus entailing better OOD generalization. Figure 2: A causal model for data augmentation. Y : the random variable for the label; Z: the random variable for the semantic meaning; X i : the random variable for the samples in environment i; ε i : the random variable for the environmental factor in environment i; X i = g(Z, ε i ) denotes the structural causal equation for the data generating process; See § 4.2 for a concrete example. Choosing the "right" data augmentation to the "right" extent is vital for our approach, after all, data augmentation needs prior knowledge of the data generating process g(Z, ε), which may not be possible. In such cases, data transformations may not be capable of mapping all environmental factors to an agnostic representation. Nevertheless, even with augmentations marginally related to the causal structure, we can still expand training distributions, increasing their overlap and favoring generalization. We further demonstrate this in the following case study, as well as in the experimental results of § 5.

Rotated MNIST: a case study
Formulation. We give a concrete example for the causal model with the Rotated MNIST dataset [19]. In this example, we have: Y : the label of the digits; Z : the original MNIST digit images; (9) ε i : the rotation angle for the digits; X i : the rotated digit images.
Specifically, each ε i follows a Dirac distribution δ ε=ti , with {t 1 , t 2 , . . . , t 5 } = {0 • , 15 • , . . . , 60 • }. The data generating function g now is: The data augmentation we consider is: Let us now study how the data augmentation encourages different client distributions to be similar. We can see that after the data augmentation, we have: This equation tells us that applying the data augmentation T α onto g(Z, ε) is equivalent to applying a related transformation T α directly on the environmental variable ε. In general, based on our (perhaps incomplete) prior knowledge of the data generation g, we can devise appropriate data augmentation T that solely induces operation on ε i . However, if the data augmentation T is not related to the data generation process g, such as Gaussian blur, then applying T onto g(Z, ε) may reduce data heterogeneity, but will also inevitably modify the semantic variable Z, and degrade the generalization (see also § 5).
Proper data augmentation can homogenize client domains. We now show that the domains for clients are indeed more similar after proper data augmentation. Since the data generation function g remains the same throughout, it suffices to measure the distance between the joint distributions of (Z, ε i ). We use total variation as the measure. Suppose we have two distributions with density functions p 1 and p 2 , and the union of their supports to be supp. Then the total variation (TV) distance between the two distributions is: From the independence assumption of our SCM model, the joint distribution of (Z, ε) is the product of marginals. Hence, for (Z, ε 1 ) ∼ p Z × p 1 and (Z, ε 2 ) ∼ p Z × p 2 , the TV distance is: Therefore, it suffices to study the TV distance for the distributions of the environmental variable ε. Before the data augmentation, the TV distance is between two Dirac distributions δ ε=t1 and δ ε=t2 : We use U (a, b) to represent a uniform distribution on the interval [a, b]. After the data augmentation T α , the TV distance is between two uniform distributions, U (t 1 − α, t 1 + α) and U (t 2 − α, t 2 + α). WLOG, we assume that t 1 < t 2 . With the definition of uniform distribution, we find that: and This means that with increasing data augmentation (larger α), different domains would look more similar.

Experiments
In this section, we examine the generalization performance of FL algorithms and show how data augmentation addresses the non-IIDness in federated learning, and thus the communication issue.

Experimental setup
Since our task is federated domain generalization, we borrow datasets from domain generalization to set up our federated learning experiments. We assume full participation throughout.
Rotated MNIST [19] is a domain generalization dataset adapted from MNIST [20]. In this dataset, each domain has its digits rotated by a fixed angle, from In-distribution performance is evaluated over a validation set, 10% of the training domain, while the domain of 75 • is set aside for OOD. For this task, we use the convolutional network as reported in [21], trained with batch size of 64 and learning rate of 1e-3 with the Adam optimizer [22]. For each client, we take the local steps E = 200 [1].
PACS [23] is an image dataset categorized into 10 classes that are scattered across four different domains, each having a distinct trait: photograph (P), art (A), cartoon (C) and sketch (S). For this experiment, the sketch domain (S) is set aside for testing.
For the federated experiments in PACS, the training domains are further divided into four subsets, with 12 clients in total. The splitting method is based on [24], using a symmetric Dirichlet distribution (with parameter 200) to assign examples from each class in non-IID fashion. For this task we use ResNet-18 [25], pretrained on ImageNet, with batch size of 32 and learning rate of 5e-5. We use the Adam optimizer [22] and for each client, we take the local step E = 200 [1].
OfficeHome [26] is analogous to PACS, having four distinct image domains: Art (A), ClipArt (C), Product (P) and Real-World (R). However, in OfficeHome the number of classes is much greater, with 65 categories in total. In this experiment, we use the same settings as reported in PACS, using the ClipArt domain for testing. For further discussion on the datasets and experiments, see App. B.1.

The importance of an appropriate augmentation
Augmentation methods used. As reported in § 4.2, akin to the causal mechanism that generates different domains, in Rotated MNIST, we apply increasing levels of random rotation in (12), with α ∈ {0 • , 15 • , 30 • , 45 • , 60 • }. We also measure the effect of applying uncorrelated transformations, such as Gaussian blur. Since the causal mechanism is not as patent in PACS and OfficeHome, we measure the effect of stacking multiple transformations that are pertinent in distinguishing the different domains and do not compromise the semantic content. We first apply positional transformations, such as random cropping and horizontal flipping, here denoted as weak augmentation. In moderated augmentation, we further include color transformations, such as random conversion to grayscale and color jittering.
In general, out-of-domain accuracies increase with stronger levels of suitable data augmentation (see Figure 3). This verifies that suitable data augmentation can significantly improve out-of-domain generalization, even in the federated Figure 3: Out of distribution accuracy of federated algorithms across different data augmentation schemes. Augmentations are further described in § 5.2.
learning setting. As we explained in § 4.2, this is because our data augmentation alleviates the data heterogeneity of training and test clients.
However, for transformations that are irrelevant to the task, such as Gaussian blur in Rotated MNIST, the performance is even worse than no data augmentation. This shows us the importance of choosing the right data augmentation that is relevant to the data generating process, as otherwise data augmentation can even degrade generalization. From the causal perspective (Figure 2), irrelevant augmentation also changes the semantic variable Z.

Data augmentation and data heterogeneity
As discussed in § 2, federated algorithms may not converge to the optimal solution under non-IIDness. However, one effect we observe is that with appropriate augmentation, federated learning can achieve comparable performance to the centralized setting (see Figure 5). This is possible because, in such cases, data augmentation increases the similarity of local distributions, shrinking the data heterogeneity (non-IIDness) among clients. To support this claim, we adopt (5) from [11] for measuring data heterogeneity. We use the global model obtained by FedAvg in each augmentation scheme as a proxy of the optimal model θ * . We estimate the gradient norm of the global model across all in-domain and out-of-domain test sets and take their average. With increasing levels of proper augmentation, we observe that the averaged gradient norm σ 2 decreases, indicating that the augmented domains are more similar. However, if the augmentation is not relevant to the environmental variable (see Figure 2), such as Gaussian blur for Rotated MNIST, then the augmentation does not help to alleviate the non-IIDness.

Data augmentation for communication efficiency
The reduction of data heterogeneity has an additional advantage in the federated setting. By augmenting the clients' data, and thus generating more similar data distributions, local models can be trained in isolation for longer without suffering from weight divergence, as reported in [2]. To validate this conjecture, we propose the following experiment using the Rotated MNIST dataset and Federated Averaging. We test the effect of increasing local steps, and decreasing communication rounds, on out-of-domain generalization, with varying degrees of augmentation, as shown in Figure 5.
Our results demonstrate that, with the right augmentation, we can obtain comparatively good performance with as few as two rounds of communication, as it is shown in E = 128000 in Figure 5. In the case where each local gradient step is followed by a communication round, i.e., with E = 1 local step, federated learning reduces to centralized training. However, under suitable augmentations, we show that the performance gap between centralized and federated learning decreases. This further motivates the search for effective transformations in more practical settings.
In-domain vs. out-of-domain. Last but not least, in both centralized and federated cases, we find that in-domain test accuracy is greater than out-of-domain ( Figure 5), a gap of which is further enlarged in federated training. This emphasizes the challenges of OOD generalization in federated learning. As mentioned before, one potential approach is proper data augmentation, with which we can show that the gap between in-domain and out-of-domain accuracy diminishes.

Related work
Federated Learning: In the seminal work [1], despite claiming that Federated Averaging (FedAvg) is able to cope with non-IID data distributions, extensive empirical evidence suggests a degradation in performance under such settings [27,2]. This is further supported by [28] where they demonstrate the convergence of FedAvg to be linearly affected by the data heterogeneity. This is due to the divergence of various client models when trained with data heterogeneity [2]. Since then, one prolific line of work has been proposing changes to the aggregation scheme of local models [29,30,31] or the objective function [12] to account for these differences. Other techniques consist of dynamically increasing the communication frequency [27] or introducing a subset of data shared across devices to mitigate model divergence [2], despite being less viable in practice. Our formulation of Federated Domain Generalization also relates to the fields of Federated Domain Adaptation [32,33] and Federated Transfer Learning [34,35], where the concept of different domains is also present.
Domain Generalization: One of the primary assumptions of Domain Generalization is that data is not IID across domains. However, contrary to Federated Learning, this is typically deemed beneficial, as it can be exploited to design more robust predictors. DG techniques typically fall into one of three categories: data manipulation, representation learning, and learning strategy [36]. The first aims to learn a general representation across domains by manipulating the inputs either via data augmentation [37,38], or data generation [39,40]. In representation learning, the goal is to learn a predictor while guaranteeing some stable property across domains, achieved through various techniques, such kernel-based methods [41], domain adversarial learning [42], explicit feature alignment [43], or invariant predictors [44,45,4,46,47,48,47]. The latter also has formulations connected to FL [49]. The category of learning strategy gathers techniques that employ different strategies to the learning task, such as using meta-learning [50], or distributionally robust optimization [5,6], among others.
Data Augmentation: Examples abound on the employment of data augmentation for generalization [51] and selfsupervised learning [52,53,54]. Its most elementary use consists of applying handcrafted augmentations, whose success depends on the relationship with the task in hand. For this reason, the choice of augmentation is usually done based on prior knowledge, although some works [55,56] have attempted to devise a transformation strategy during training. In more recent work, the goal is to learn such task-related augmentations. Among different strategies, these may depend on generative models [57,58] or adversarial training [38,59]. In terms of theory, data augmentation can be linked to regularization theory [60,61], group theory [62], as well as to the notion of causality [47,63,48].
Causal modelling for domain generation: several approaches attempt to formulate the data generating process of different domains using a causal model. Examples abound in learning tasks in both causal and anti-causal directions [44,64,65,45,4,17]. Similar to our work, in [64,66,67,47] the input is causally related to the semantic content and an environmental factor (also denoted as style). Data augmentation is used in [48,67,66] as a perturbation over the style -also seen as an intervention in [47]-to promote model invariance. Further, [68] uses group theory to formalize the notion of data augmentation, devising a general principle to learn generative causal models.

Conclusion
In this paper, we study the out-of-domain generalization of federated learning, and propose to use data augmentation for solving the non-IIDness. This approach is inspired from the connection to domain generalization, from which we can borrow insights for designing FL algorithms as well. Our key observation is that, with proper data augmentation, the client domains are more similar and therefore the FL training suffers less from non-IIDness and OOD generalization can be improved. We demonstrate this point through extensive experiments. Our approach is distinct from the popular line of FL research on better optimization and aggregation, and focuses more on the data side. We hope this data-driven perspective can inspire more exploration into designing proper data augmentation for federated learning and OOD generalization in general.

A Connecting Federated Learning and Domain Generalization Algorithms
In this appendix we provide more details about the FL algorithms that we transfer from the field of domain generalization, including Generalized AFL (Gen-AFL), Variance Minimization (VM) and Federated IRM (Fed-IRM). We summarize the connections drawn in this paper between Federated Learning and Domain Generalization algorithms, as well as their objective function formulation in Table 2.

A.1 Generalized AFL
In domain generalization, one extension of GroupDRO is Minimax Risk Extrapolation (MM-REx, [6]), where one replaces the simplex ∆ in (7) with ∆(λ min ) := {λ ∈ R n : 1 λ = 1, λ ≥ λ min 1}. MM-REx can be treated as a generalization of GroupDRO and ERM: if λ min = 0 then we have GroupDRO (AFL), where the highest weight is put on the worst-case client; if λ min > 0, then we want our algorithm to be more uniform (specifically if λ min = 1/n, then we have ERM (FedAvg)). The extrapolation occurs when we have λ min < 0: we enforce larger weights on the worst-case clients and negative weights on the best clients. In other words, we are extrapolating beyond the convex hull of the client distributions.
Proof. The projection on the left hand side of (18) can be written as: proj Note that ∆(λ min ) can also be rewritten as: Therefore, we can rewrite: With this transformation we can obtain (18).

A.2 Variance Minimization
Another extrapolation algorithm proposed in [6] is called V-REx, which aims to minimize a linear combination of the variance and the mean: 8 Output: global model θ T the predictor is composed of a feature encoder g and a classifier h on top. The formulation of IRM is: In [4], the authors propose to treat h as a fixed scalar classifier, and thus the optimization problem becomes: Following a similar procedure of FedAvg, we propose the federated version of IRM in Algorithm 3. It is important to mention that other federated versions of IRM have been proposed in [49]. However, these differ from our approach, as they comprise some form of data communication, either through intermediate representations of the feature encoder (as in CausalFed), or through a subset shared with all clients and server (as in CausalFedGSD). For this reason, we did not include these methods in our analysis. Nevertheless, our Fed-IRM formulation resembles CausalFedGSD, with the distinction that Fed-IRM does not share any data between clients nor the server.
In our formulation of the federated version, when E = 1 we have the exact agreement to the DG formulation. However, once local models are trained under IRM with less frequent communication rounds, both in-domain and out-of-domain performance start to degrade, as we further discuss in the results of Rotated MNIST, in section B.1.

B Experiments
We present the experimental details in this section. Part of the code was based on the DomainBed 3 test suite, whose license can be found at https://github.com/facebookresearch/DomainBed/blob/main/LICENSE. We performed the experiments in a cluster of GPUs, including NVIDIA V100 and P100.

B.1 Benchmark experiments
In this subsection, we report details over the experiments concerning federated and domain generalization algorithms, as reported in Table 2, also including FedProx, which does not possess a centralized counterpart. In these experiments, we trained all models for 16000 gradient steps, which assured convergence in all feasible cases. In federated learning, these gradient steps are further divided into 80 communication rounds, totalizing 200 local steps per round. Apart from FedProx, which implements the proximal operator with SGD, we train all other algorithms with Adam using the default parameters as reported in the PyTorch library (β 1 = 0.9, β 2 = 0.999), except for the learning rate, which was adapted for each task (see § 5).
The hyperparameters of the DG algorithms were chosen according to the DomainBed implementation. For the FL algorithms, hyperparameter tuning is used for Gen-AFL as well as VM. For Gen-AFL, λ min = −1 is used for all algorithms, while for VM, β = 10 is used in Rotated MNIST and β = 0.5 is used for PACS and OfficeHome.
The results for Fed-IRM are not reported in Table 3 due to the challenges in tuning the model locally. One important distinction from the centralized version is that, in local training, the penalty term comprises only the gradient on the  ID  OOD  ID  OOD  ID  OOD  ID  OOD  ID  OOD  ID  OOD  -  Generally, more frequent communication is beneficial in federated learning, whereas in the limit of local step E = 1 we have the equivalent to centralized training. Surprisingly, the results in PACS indicate that the federated regime, despite having slightly lower in-domain performance, significantly outperformed centralized training in terms of out-of-domain accuracy. We conjecture that, in this case, alternating between local training and recurrent model aggregation may act as a regularizer (similar to early stopping and dropout) that is favorable for the test domain at hand. This interesting phenomenon may require further investigation. As for OfficeHome, the task is much harder than in PACS. Despite having slightly more samples, in OfficeHome, these are distributed along a much larger set of labels than in PACS. In federated learning, this challenge becomes even harder.
After the dataset split, the resulting clients possess a small fraction of examples for each class, which results in poor generalization for both in and out-of-domain.
The results of OfficeHome and PACS differ from the ones found on [7] for multiple reasons. The first distinction is model capacity, as in [7], experiments were conducted with ResNet50 networks instead of ResNet18. Furthermore, [7] further performs an extensive hyperparameter search as well as model selection based on different criteria.

B.2 Data augmentation and local steps
In this experiment, the objective is to validate whether data augmentation could aid the homogenization of clients, thus enabling the federated training with sparser communication rounds. For this, we trained FedAvg with varying local steps, communication rounds and data transformations over Rotated MNIST, whose results are reported in Table 6.  ID  OOD  ID  OOD  ID  OOD  ID  OOD  ID  OOD  ID  OOD  2000  5

B.3 Data augmentation and data heterogeneity
In this analysis, we use the models trained in the experiment aforementioned and compute the gradient norm with respect to each domain. As previously mentioned in § 5.3, in the training domains, the gradient is obtained over the validation set of each domain, while the full set is used for the test domain. The results for Rotated MNIST, PACS and OfficeHome are reported in Tables 7, 8, and 9, respectively.