{'fraction_non_alphanumeric': 0.031236166445329793, 'fraction_numerical': 0.01131584772023019, 'mean_word_length': 4.885849210226348, 'pattern_counts': {'":': 0, '<': 0, '<?xml version=': 0, '>': 3, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 3, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
FedGen: Generalizable Federated Learning

Existing federated learning models that follow the standard risk minimization paradigm of machine learning often fail to generalize in the presence of spurious correlations in the training data. In many real-world distributed settings, spurious correlations exist due to biases and data sampling issues on distributed devices or clients that can erroneously influence models. Current generalization approaches are designed for centralized training and attempt to identify features that have an invariant causal relationship with the target, thereby reducing the effect of spurious features. However, such invariant risk minimization approaches rely on apriori knowledge of training data distributions which is hard to obtain in many applications. In this work, we present a generalizable federated learning framework called FedGen, which allows clients to identify and distinguish between spurious and invariant features in a collaborative manner without prior knowledge of training distributions. We evaluate our approach on real-world datasets from different domains and show that FedGen results in models that achieve significantly better generalization than current federated learning approaches.


INTRODUCTION
The rise in popularity of edge devices (e.g., mobile phones, wearables, sensors) has resulted in the generation of massive amounts of data, exacerbating network bandwidth demands, privacy concerns, and security issues. This has called for distributed model training over a large number of these devices for many machine learning applications. Federated Learning (FL) has emerged as an attractive paradigm, allowing local devices or clients to collaboratively train a shared global model. The typical federated learning paradigm involves two stages -(i) clients train models with their local datasets independently, and (ii) a central server gathers the locally trained models and aggregates them to obtain a shared global model.
A standard approach for federated learning has been FedAvg [28] where parameters of local models are averaged element-wise with weights proportional to sizes of the client datasets. However, FedAvg has been shown to have several shortcomings [19,49], in particular its performance when the data across clients are non-iid. While several improvements have been proposed to handle noniid data [25,29], they rely on empirical risk minimization (ERM) techniques which has been shown by numerous examples to be brittle in generalizing to out-of-distribution (OOD) data. This is partly due to being influenced by spurious correlations and data biases that fail to hold outside training data distributions [35,43]. A classic example was highlighted by Beery et al. [5] where (a) Training Accuracy (b) Test Accuracy Figure 1: Comparing the generalizability of FedAvg with and without spurious correlations in the training data a model, trained to classify images of cows in pastures and camels in the desert, failed when the backgrounds were switched because it was influenced by the spurious correlation (i.e., green pastures with cows and sandy deserts with camels) rather than relying on the invariant features (i.e., the cows and camels themselves). In another example, DeGrave et al. [12] showed how models trained to detect COVID-19 from chest radiographs relied on spurious features (e.g. hospital metadata) rather than medical pathology features, and failed when tested in new hospitals. While there has been an increasing effort to improve the generalization of models to out-of-distribution data using different approaches like metalearning [6], adversarial learning [2], subnetwork selection [52] etc, they are designed for centralized training and there has not been prior work done for generalizability in federated settings.
To demonstrate the sensitivity of FedAvg and associated empirical risk minimization approaches to spurious correlations in the training data, we use the Punctuated SST-2 dataset [10]. It consists of sentences and their binary sentiment labels divided into two training distributions. A punctuation mark, either a '!' or '. ', is introduced as a spurious feature with an 80% and 90% spurious correlation with each of the binary sentiment labels in the two training distributions respectively, and only has a 10% correlation in the test distribution (i.e.) out-of-distribution (OOD) setting. Hence, any model relying on this punctuation feature rather than the sentence while predicting the sentiment would do well during training, but perform poorly on the test set where the spurious correlation no longer holds.
We distribute the training data over ten clients and compare the performance of FedAvg on data with the spurious feature against its performance if the spurious feature was not present. Figure 1 shows the training accuracy and the resulting out-of-distribution test accuracy of FedAvg in the two cases. We observe that when the spurious feature is present, FedAvg is influenced significantly by it, thereby resulting in a higher training accuracy than if it were not present. However, this influence results in poor generalization, and it demonstrates very low test accuracy on the out-of-distribution (OOD) data. The poor generalizability of FedAvg can be clearly seen from the performance of the two cases in Figure 1, since the only difference in the datasets is the presence of a spurious feature.
The example above highlights the drawbacks of existing federated learning approaches. In this work, we address the drawbacks by developing a generalizable federated learning approach called FedGen. Our approach develops a distributed masking function that enables clients to collaboratively identify spurious features or correlations in the data thereby ensuring that the resulting aggregated model relies only on the invariant features. In this work, we focus on classification tasks on time-series and sequential data, such as natural language processing and IoT data where spurious correlations are often present and will extend this to other tasks and domains as part of future work.
Our contributions can be summarized as follows: • We present a generalizable federated learning framework (FedGen) that can handle the presence of spurious correlations or biases in the training data. • FedGen uses a distributed masking function that allows the clients to collaboratively identify and distinguish between spurious and invariant features during the training process even as each client trains on only its local data. • We formally prove the correctness and present results on the convergence of FedGen. • We compare our framework with standard centralized training as well as existing federated learning approaches on a variety of sequential and time-series datasets from realworld domains and demonstrate the significantly improved generalizability of aggregated models using FedGen.

BACKGROUND
Consider a set of devices or clients, each with data samples drawn from ∈ D distributions on × , where is the set of input features and is the target variable. We denote as the number of data samples on device ∈ , and = =1 as the total number of samples across the devices.
The general federated learning problem involves learning a single global statistical model from these devices under the constraint that device-generated data are stored and processed locally, with only intermediate model parameter updates being communicated periodically with a central server, which aggregates them. The goal is typically to minimize the following objective function: where is the local objective function for the th device. The local objective function is often defined by existing approaches using Empirical Risk Minimization (ERM) which minimizes the average training loss in a distribution agnostic manner: where ∈ , ∈ , and ℓ is a given loss function. One of the popular aggregation methods is Federated Averaging (FedAvg) by McMahan et al. [28] which uses ERM and averages the local model parameters of the devices element-wise during aggregation as summarized in Algorithm 1. While ERM has been shown to work well in practice for i.i.d. data [44], it has been shown to fail when (i) training distributions on the devices are non-iid and leading each device to optimize its local objective as opposed to the central objective -potentially hurting convergence or even causing divergence of the aggregated model [7], and (ii) test distributions differ significantly from training distributions which is a natural consequence when spurious correlations are present in the training data [43].
Several alternative approaches to FedAvg have been proposed to address the first issue. FedProx proposed by Li et al. [25] limits the impact of local updates by keeping them close to the global model where they add a proximal term to the local client objective functions : where the proximal term penalizes large deviations of the client model parameters, is a hyper-parameter and we see that when = 0, this reduces to FedAvg. Agnostic Federated Learning (AFL) by Mohri et al. [29] is another variant of FedAvg, which optimizes a centralized distribution that is a mixture of client distributions by modeling it as = , where is the mixture weight. Their goal is to find a solution that works for any . However, these approaches assume that all correlations present in the training distributions are invariant, and therefore cannot handle spurious correlations resulting in aggregated models that are not generalizable as we will demonstrate.
Finding the set of invariant features has emerged as a popular approach for generalization in centralized training settings. Invariant Risk Minimization (IRM) by Arjovsky et al. [4] tries to find an invariant data representation Φ( ) and assumes that the training data comes from multiple distributions, and that features whose distributions vary across the training data are likely to also vary between the training and test datasets and hence should be treated as spurious correlations. They define the following loss or risk function: where the second term serves as a regularizer that balances between predictive power within a distribution (ERM), and the invariance of the predictor across distributions. There have been several extensions to the IRM framework. For instance Ahuja et al. [1] propose a game theoretic approach to IRM, while Krueger et al. [22] introduce the notion of risk extrapolation to encourage strict equality between training risks.
While the goal of identifying the invariant feature set has been shown to be effective for out-of-distribution generalization over a variety of datasets and serves as a base for our approach, IRM based approaches have been shown to have inherent drawbacks [30,37]. Of particular importance is its reliance on prior knowledge of training data distributions, absent which, its performance is similar to ERM [46]. This knowledge is especially hard to obtain

Algorithm 1 Federated Averaging FedAvg
Require: Devices ∈ , local epochs , learning rate , global model randomly initialized 1: for = 1 → do 2: Server selects subset of devices at random 3: Server sends to devices in 4: for each client ∈ in parallel do 5: Update for epochs of SGD to get +1 6: Send +1 back to the server 7: in federated settings, where the local data on each client can itself be composed of multiple distributions and only a subset of clients may participate. In this work, our goal is to bring the benefits of generalization to federated learning, while not requiring any prior knowledge of the distributions in the training data.

FEDGEN: GENERALIZABLE FEDERATED LEARNING
As described in Section 1, generalizable models need to be able to distinguish between spurious and invariant features, thereby ignoring spurious correlations with the target variable present in training data. Given the set of input features used by the device models, the set of invariant features across all devices is one where the target prediction probability is consistent across all data distributions, (i.e.) ( | ∈ , D) is approximately constant. Conversely, the set of spurious features consists of features whose prediction probabilities differ across distributions due to the spurious correlations. Hence, it follows that ∪ = , and ∩ = ∅, (i.e.) a feature cannot be both invariant and spurious. Our approach to determining whether the th feature ∈ is spurious or invariant, is by measuring the stability of its parameter weight . Recent work [17,46]  We leverage this intuition for our federated setting and define a set of masks = { 1 , ..., } over the input features for each of the ∈ clients, where ∈ R and is the number of input features in . We update the masks during training by using the variances in the feature weights to emphasize invariant features and suppress spurious ones. During each training epoch, we update the local masks for each client ∈ as: where hyper-parameter serves as a scaling factor, ( ) is the variance of the weights of feature on client , and the second term is the average variance of all feature weights on client . From the above intuition, we know that the variance of invariant features is low and that of spurious features would be high. Hence, we see that updating the local masks on each client using Equation 5 results in the masks of invariant features gaining in value on each of the clients since their variance is lower than the average (i.e.) 1 =1 ( ) − ( ( )) > 0. Masks of spurious features on the other hand, become progressively negative, since the variance of their weights coupled with the scaling factor is larger than the average which is brought down by invariant features. Since the masks themselves are unbounded, we retain the scale of the original feature values by scaling the masks using the sigmoid function which is bounded between [0, 1]. Hence, using the masks and the update function in Equation 5 on the input features for each local client results in: where ⊙ denotes element-wise multiplication. We then define the local loss on each client in FedGen as: The penalty term here like IRM serves as a regularizer for the local models that penalizes doing too well on one data distribution (i.e., possibly relying on spurious features) and rewards doing well across distributions (i.e., relying on invariant features). FedGen then builds a centralized model by communicating and subsequently aggregating both the local model parameters as well as the local masks of the participating devices or clients. We use the same element-wise averaging for aggregation as used by prior work: Aggregating the masks across the devices or clients allows them to obtain a consensus on invariant and spurious features across the different data distributions that can be present in their local training data. Masks that are positive (i.e. reflecting invariant features) across most clients will continue to be emphasized after aggregation, and likewise for the negative spurious feature masks. The mask aggregation also allows client models relying on the wrong features to obtain the global consensus and ignore spurious features. We define the FedGen central objective function as: We summarize our FedGen approach in Algorithm 2. We next formally prove that FedGen results in an aggregated model that is generalizable in Theorem 1 by formulating a minimax problem and showing that the resulting model minimizes the loss using invariant features, even under the most adverse unseen test distributions.

Theorem 1. Given training distributions
and a test distribution , the set of invariant features is the saddle point of the following minimax problem:
Proof. For every , we can partition it into invariant variables and non-invariant variables as: Consider a test distribution * , where the set of spurious features * are not predictive of the output , and only the invariant features are predictive of , (i.e.) Therefore, ( , ; * ) where (·) is the cross-entropy loss function.
Step ( ) is obtained by applying equation (11), and step ( ) is due to the property of the masks from equation (6).
Recall that * was assumed to be non-predictive of . However, in most cases, the spurious feature would have some predictive power over in the training environment. Hence, from the definition of spurious features , their biased influence on the model performance during training will lead to an increased loss in the worst case test environment: Recall denotes the set of invariant features, thus ( | , ) does not depend on . Therefore, By combining equations (13), (14), and (15), we have: The above formulation holds for all . Hence, taking the maximum over in equation (16)

FedGen Convergence
In this section, we analyze the convergence of FedGen by first introducing some definitions and assumptions for our convergence analysis.
In order to analyze convergence in heterogeneous non-iid settings with constant step-size (as is usually implemented in practice), we also introduce a notion of dissimilarity between the local device functions in the federated network similar to assumptions made by prior work [25,45,51]. We further define ( ) = √︂ E | | ∇ ( , ) | | 2 | | ∇ ( , ) | | 2 when ||∇ ( , )|| 2 ≠ 0. As explained in Li et al. [25], when all the local functions are the same (i.e.) distributions of data sampled from the devices are i.i.d, we have ( ) → 1 for all . However, in the presence of spurious correlations, the data distributions are heterogeneous and therefore > 1, and larger the value of ( ), the more the heterogeneity of the local data. We next make the following assumption on the objective functions.
We present the detailed proof in Appendix A.1 The proof can be found in Appendix A.2

EXPERIMENTS
In this section, we evaluate FedGen on diverse models, clients, and datasets spanning multiple application domains: Natural Language Processing (NLP), Internet-of-Things (IoT), and business process mining (Table 1). Similar to prior work [1,4,10,22], we augment these benchmark datasets with spurious features such that it has a strong spurious correlation with the target labels in the training data distributions (i.e.) ( | , ) ≥ = 0.8, and this spurious relationship does not hold in the test distribution (i.e.) ( | , ) ≤ = 0.1. We assess the quality of generalization of the aggregated model as the classification accuracy obtained on the test environment, disjoint from the set of training environments. Hence, an approach that cannot generalize will do well during training but not on the test set, while a generalizable model will continue to perform well on the test data. We also present results from an ablation study of the masking function as well as additional experiments that highlight the robustness of our approach.

SST-2:
We modify the Stanford Sentiment Treebank (SST-2) [41], a benchmark dataset for binary sentiment analysis with spurious correlations in a similar manner as Choe et al. [10]. We add a spurious feature by pairing each sentiment with a final punctuation mark (. or !) in the corresponding text as follows: where we set as described above. The addition of the punctuation mark as a separate feature, and not as part of the sentence, allows us to isolate the sentence sequence embeddings to mask, while still preserving the structure and meaning of the sentence (Choe et al. [10]). We create two different training distributions with differing strengths of spurious correlations and treat the test data as a third distribution with weak spurious correlation as described above. We distribute the training data among 100 devices or clients.

AG News:
To evaluate FedGen for multi-class NLP predictions, we use the AG News dataset [53], which consists of a corpus of news articles and their titles that have been classified into four categories -(1) World, (2) Sports, (3) Business, and (4) Sci/Tech. There are 30,000 training and 1,900 test examples for each class. We again add a spurious feature that pairs each of the four news categories with a corresponding final punctuation mark in the article from the set P = {. ! @ ∧} as follows: and set as before. We divide the training data into two distributions with strong spurious correlations, and a separate test distribution with a weak correlation. We distribute the training data among 100 devices.

Human Activity Recognition (HAR)
. The HAR dataset [3] consists of smartphone accelerometer and gyroscope readings recorded from 30 individuals performing six activities (walking, standing, sitting, etc.). The model objective is to predict the activity being performed given the sensor readings. The data consists of sequences of 128 time-steps of sensor readings which correspond to any particular activity. We consider the data from the smartphone of each of the individuals to be a client, and divide them into two training and one test distribution such that there is a strong spurious correlation between the specific activity performed and the smartphone model used by individuals in the training test, which does not hold for individuals in the test set as follows: where and refer to the th activity and smartphone model respectively.

Water Quality
Detection. The dataset [14] comprises of water quality data collected from different cities across the United States. It includes sensor readings such as turbidity, conductivity, pH, temperature, DO 2 , in addition to environmental attributes like humidity and precipitation, which correspond to binary labels depicting the occurrence of a contamination event. We divide the training data into two distributions, clustered the humidity into  Table 2: Accuracy achieved by comparison approaches for all datasets low ( ) and high ( ) values and spurious correlate them to the contamination labels (C, ∅) as follows: We distribute the training data among 120 clients corresponding to the different cities.

Business Process Mining:
We use a benchmark process mining dataset (BPIC 2018) 1 which consists of payment applications from German farmers to the European Agricultural Guarantee Fund collected over a period of three years. Each application is processed by one of four departments, and to evaluate our approach in a multidistribution setting, we consider each department as a distribution and use three of them for training and the fourth as our out-ofdistribution (OOD) test data. We consider two business process variants and augment the process trace sequences with a spurious feature denoting the area of the farm. We cluster the area into low ( ) and high ( ) values and spuriously correlate them to the two business process variants (A, B) as follows: and define as before. We distribute the training data among 100 clients.

Benchmarks
We compare FedGen to four benchmark approaches: • ERM: Standard empirical risk minimization approach that minimizes the average loss over the entire training data. We use this as a comparison with centralized training. • FedAvg: Federated averaging approach by McMahan et al. [28] where model parameters are averaged element-wise. • FedProx: Approach by Li et al. [25] to handle heterogeneous non-iid data in federated settings using a proximal term to keep local models close to the global model. • Inv-FedAvg: FedAvg trained on data without any spurious features. This approach reflects the setting where spurious correlations are not present and hence reflects the best performance.

Implementation Details
We implement FedGen, ERM, FedAvg, and Inv-FedAvg using the PyTorch library [31] and implement FedProx based on their publicly 1 https://data.4tu.nl/articles/BPI_Challenge_2018/12688355 available code 2 . For each dataset, we tried both MLP and LSTM classifiers of different architectures and selected the model that performed the best. We used the same model architecture for all comparison approaches to ensure a fair comparison. We used the cross-entropy loss for classification during training and the Adam optimizer. We identify the final models when they have either converged, started to diverge, or run a sufficient number of rounds, whichever comes earlier, similar to definitions by prior work [25,38]. We also assume the central server samples all the local clients to join the training process in every communication round for simplicity. We provide further details on the hyper-parameters, model architectures, and experimental settings in Appendix A.3.

Test Accuracy.
We first measure the accuracy of the final aggregated global model resulting from all the comparison approaches on the out-of-distribution (OOD) test datasets to compare their generalizability. Table 2 summarizes the results, where we see that ERM, FedAvg, and FedProx, are all heavily influenced by the spurious correlations since their training accuracy across the datasets is high. However, this results in low generalization to the test set where the spurious correlations no longer hold. We observe that FedGen on the other hand, ignores the spurious correlations (lower training accuracy), and instead identifies and relies on the invariant features thereby achieving good generalization and high test accuracy across all the datasets. It outperforms ERM, FedAvg, and FedProx, and achieves 24% more accuracy on average across all five datasets. Also importantly, FedGen, even in the presence of spurious correlations, performs nearly as well as Inv-FedAvg (within 2% accuracy on average) which reflects the accuracy when trained on data without any spurious correlations.
In addition, even though FedAvg and Inv-FedAvg rely on the same model and algorithm, the additional spurious features results in a significant increase in FedAvg's training accuracy, which does not translate to generalizability on unseen test distributions, highlighting the importance of building generalizable model training approaches.

Effect of choice of local training epochs.
While a large number of local training epochs can help in reducing the communication costs, previous work has also shown that the value of can impact the performance of FedAvg and can sometimes lead to divergence [8,28,38]. We conduct an experimental study on the effect of the choice of over the different comparison approaches across all five datasets. The candidate local epochs we consider are ∈ {20, 40, 60, 80, 100, 120, 140}. We run each approach till it achieves convergence and report the test accuracy achieved in Figure 2. We observe that the performance of FedAvg and Inv-FedAvg deteriorates with longer local training, thus showcasing its sensitivity to hyperparameter settings, which matches prior observations made in literature. FedProx only partially alleviates this problem and also demonstrates a drop in accuracy for larger number of epochs. On the other hand, we observe that FedGen often benefits from longer training, since it gets to further emphasize invariant features.

Convergence Rate.
In this experiment we compare the global model convergence achieved by all the comparison approaches. We report the convergence rate for the local training epochs that yields the best final aggregated model accuracy over the test set for each of the approaches. Figure 3 shows the convergence rates of all the approaches across the five datasets. We observe that the convergence rate achieved by the approaches is quite similar, and that FedProx converges the quickest, followed by FedGen and FedAvg. This demonstrates that FedGen can result in significantly more generalizable models without a loss in runtime performance compared to existing approaches.

Ablation Study.
We perform an ablation experiment to measure the contribution of the components of our masking function and the FedGen penalty term as shown in Table 3. We first remove the hyperparameter used to scale the masks, and observe that there is a drop in accuracy across all three datasets since the degree of spuriousness may not be fully captured, and hence the local models can get influenced by spurious features. We next remove the masking function (M), essentially reducing the loss function to that used by FedAvg. This is reflected by the accuracy values,   with > 0.8. Table 4 shows that the performance of FedGen and Inv-FedAvg remains consistent even with an increased number of distributions, demonstrating FedGen's robustness, while the other approaches are impacted to a small extent.

RELATED WORK
As edge devices like phones, sensors, and wearables grow in capabilities and popularity, it is increasingly attractive to train models locally using federated learning [24]. To improve the performance of the popular FedAvg approach [28] in statistically heterogeneous settings, FedProx [25] and Agnostic Federated Learning [29] are two approaches that we described in Section 2. Karimireddy et al.
propose SCAFFOLD [20] which uses control variates (variance reduction) to correct for any distribution drift or changes in the local clients. However, the heterogeneity assumed by these approaches does not extend to spurious correlations as we have demonstrated.
Additionally, heuristic approaches have been proposed to tackle the heterogeneity by sharing the local device data or server-side proxy data. However, in addition to placing a burden on network bandwidth, sending local data to the server as in Jeong et al. [18] violates the key privacy assumption of federated learning, and sending globally-shared proxy data to all clients (Zhao et al. [54], Huang et al. [16]) requires effort to carefully generate the data.
There has also been work on improving the communication costs of FedAvg. Wang et al. [49] propose FedMA which matching the neurons of client models before averaging them, thereby reducing the communication needed. CMFL is an approach by Luping et al. [27] where every client checks if its update is relevant for model improvement and reduces the number of irrelevant updates to the server, while FedAsync (Xie et al. [50]) improves training time by focusing on handling stragglers. However, these approaches assume that the training and test datasets come from the same population, which may not happen in many real-world settings, thereby calling for generalizable approaches.
Prior work on machine learning generalization have predominantly looked at computer vision challenges. Data augmentation techniques aim to make the model more robust by training using instances obtained from neighbouring domains hallucinated from the training domains, and thus make the network ready for these neighbouring domains. Shankar et al. [40] augment data using instances perturbed along directions of domain change and use a second classifier to capture this. Volpi et al. [48] apply this to single domain data, while Carlucci et al. [9] apply augmentation to images during training by simultaneously solving an auxiliary unsupervised jigsaw puzzle alongside.
To obtain stable predictions for unseen test distributions, a desirable model relies on only invariant features [39] and causal model discovery [32] aims to find an underlying causal graph to obtain the invariant feature set. Peters et al. [33] adopt statistical tests for mining the invariant set, claiming that the residual of the invariant method should be i.i.d. while Rojas-Carulla et al. [36] perform this for non-linear settings. Decomposition based approaches represent the parameters of the network as the sum of a common parameter and domain-specific parameters during training [11]. Khosla et al. [21] applied decomposition to domain generalization by retaining only the common parameter for inference. Li et al. [23] extended this work to Convolutional Neural Networks (CNNs) where each layer of the network was decomposed into common and specific low-rank components. Piratla et al. [34] recently proposed a more efficient approach that decomposes only the last layer, imposes loss on both the common and domain-specific parameters, and constrains the two parts to be orthogonal.
Another approach is to pose the generalization problem as a meta-learning task which has been studied either in the context of few-shot supervised learning methods which adapt using small amounts of labeled data from the new domain [15], distribution shifts in only test domains [13], or only considering label shifts [26]. There have also been some prior work on generalization for other applications like business process predictions [47] and medical diagnosis using human annotated spurious features [42].

CONCLUSION
We present an approach to improve the generalizability of models trained using federated learning. FedGen overcomes the brittleness of current empirical risk minimization based approaches in the presence of spurious correlations in the training data, while also not needing any prior knowledge of training distributions unlike existing generalization approaches. We develop a masking function for the local clients that allows them to collaboratively identify and distinguish between invariant and spurious features during the training process. The final aggregated model thus relies on the invariant feature set, thereby ensuring its robustness and generalizability to unseen test distributions. We prove that our masking function minimizes loss even under the most adverse test distributions. We show that FedGen achieves good generalization and outperforms existing federated learning approaches on a variety of real-world datasets. We believe that this work motivates the need to improve generalizability in federated learning and as part of future work, we will work on improving the communication costs and extend our approach to other models that satisfy the minimax optimization such as GANs, as well as other prediction tasks.

A APPENDIX A.1 Proof of Theorem 2
Proof. Since ( , ) is L-smooth, we have With Assumption 1 and B-local dissimilarity, we get Since ( , ) is -strongly convex, for some ′ and ′ , we get: Re-arranging,

A.3 Implementation Details
We simulate the federated learning setup (1 server and devices) on a commodity machine with 7 Intel ® Xeon ® E5-2690 v4 CPUs and an NVIDIA ® Tesla P100 GPU. We describe the implementation details of the experiments in the paper as described in Section 4. We searched for hyper-parameters based on the values in Table 5, and chose the configurations with the best performance for each of the approaches.  For both the SST-2 and AG News datasets, we used a 3-layer (50-50-2) MLP with ReLU activations, based on the average of the 50-dimensional GloVe (glove.6B.50d) embedding word vectors as inputs. These are pre-trained word vector embeddings that have been trained on Wikipedia articles. We achieve similar results with the 300-dimensional word vectors as well (glove.6B.300d). For SST-2, we use a learning rate of 0.001 and set the penalty and regularizer weight for FedGen to be 10 4 and 10 −3 respectively with a scaling factor of 10. The regularizer weight for ERM, FedAvg, Inv-FedAvg and FedProx was set to 10 −6 and the proximal term was 10 −3 .
For AG News, we also use a learning rate of 0.001 and set the penalty and regularizer weight for FedGen to be 1 and 10 −6 respectively with a scaling factor of 10. The same regularizer weight was used for the other approaches and the proximal term was 10 −3 .
For the HAR dataset, we use an LSTM with one layer and hidden size of 100, followed by a linear layer of size 100 with ReLU activation. We use a batch size of 2 6 on the clients and a learning rate of 0.01. The regularizer weight for all approaches was set to 10 −3 and FedGen's penalty was 10 with a scaling factor of 10. We used the same proximal term as above.
For the Water Quality dataset, we used a 3-layer MLP (200-200-1) with ReLU activations. The learning rate was 0.07 and the regularizer weight was 10 −6 . We set FedGen's penalty to 1 with a scaling factor of 10.
For BPIC 2018, we use an LSTM with two layers and a hidden size of 120 followed by a linear layer of size 120 with ReLU activation. We use a batch size of 5% on the clients. The learning rate was 0.005 and the regularizer weight was 7 −6 . We set the penalty of FedGen to 1 and the scaling factor to 10.