# FedGen: Generalizable Federated Learning

Praveen Venkateswaran  IBM Research<EMAIL_ADDRESS>,  Vatche
Isahagian  IBM Research<EMAIL_ADDRESS>,  Vinod Muthusamy  IBM Research
<EMAIL_ADDRESS>and  Nalini Venkatasubramanian  UC Irvine
<EMAIL_ADDRESS>
###### Abstract.

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.

##  1\. 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 (McMahan et al.,
2017) 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 (Wang et al., 2020; Kairouz et al., 2019), in
particular its performance when the data across clients are non-iid. While
several improvements have been proposed to handle non-iid data (Li et al.,
2018; Mohri et al., 2019), 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 (Tzeng et al.,
2017; Recht et al., 2019). A classic example was highlighted by Beery et al.
(2018) where 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.(DeGrave et al., 2021) 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 meta-
learning (Bengio et al., 2019), adversarial learning (Albuquerque et al.,
2019), subnetwork selection (Zhang et al., 2021) etc, they are designed for
centralized training and there has not been prior work done for
generalizability in federated settings.

(a) Training Accuracy

(b) Test Accuracy

Figure 1. Comparing the generalizability of FedAvg with and without spurious
correlations in the training data

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 (Choe et al., 2020). 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 real-world domains and demonstrate the significantly
improved generalizability of aggregated models using FedGen.

##  2\. Background

Consider a set of $K$ devices or clients, each with data samples drawn from
$D_{k}\in\mathcal{D}$ distributions on $X\times Y$, where $X$ is the set of
input features and $Y$ is the target variable. We denote $n_{k}$ as the number
of data samples on device $k\in K$, and $N=\sum_{k=1}^{K}n_{k}$ as the total
number of samples across the $K$ devices.

The general federated learning problem involves learning a single global
statistical model from these $K$ 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:

(1)
$\min_{w}F(w)\;,\;\text{where}\;F(w)=\sum_{k=1}^{K}\frac{n_{k}}{N}f_{k}(w_{k})$

where $f_{k}$ is the local objective function for the $k$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:

(2)
$f_{k}(w_{k})=\mathbb{E}_{(x_{k},y_{k})\sim{\mathcal{D}}}[\ell(x_{k},y_{k});w_{k}]$

where $x_{k}\in X$, $y_{k}\in Y$, and $\ell$ is a given loss function. One of
the popular aggregation methods is Federated Averaging (FedAvg) by McMahan et
al.(McMahan et al., 2017) 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
(Vapnik, 1992), 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 (Bonawitz et
al., 2019), and (ii) test distributions differ significantly from training
distributions which is a natural consequence when spurious correlations are
present in the training data (Tzeng et al., 2017).

Several alternative approaches to FedAvg have been proposed to address the
first issue. FedProx proposed by Li et al.(Li et al., 2018) 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 $f_{k}$:

(3) $f_{k}(w_{k})+\frac{\mu}{2}||w_{k}-w^{t}||^{2}$

where the proximal term penalizes large deviations of the client model
parameters, $\mu$ is a hyper-parameter and we see that when $\mu=0$, this
reduces to FedAvg. Agnostic Federated Learning (AFL) by Mohri et al.(Mohri et
al., 2019) is another variant of FedAvg, which optimizes a centralized
distribution that is a mixture of client distributions $D_{k}$ by modeling it
as $D_{\lambda}=\sum\lambda_{k}D_{k}$, where $\lambda$ is the mixture weight.
Their goal is to find a solution that works for any $\lambda$. 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. (2019) tries to find an invariant data representation
$\Phi(X)$ 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:

(4)
$\min_{\Phi}\sum_{D\in\mathcal{D}}\mathcal{L}_{D}(\Phi(X_{D}))+\lambda||\nabla_{w}\mathcal{L}_{D}(w^{\top}\Phi(X_{D}))||^{2}_{2}$

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. (2020) propose a game theoretic approach to IRM, while
Krueger et al. (2021) 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 (Nagarajan et al., 2020; Rosenfeld et al., 2020). Of
particular importance is its reliance on prior knowledge of training data
distributions, absent which, its performance is similar to ERM (Venkateswaran
et al., 2021a). This knowledge is especially hard to obtain 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.

Algorithm 1 Federated Averaging FedAvg

1:Devices $k\in K$, local epochs $E$, learning rate $\eta$, global model $w$
randomly initialized

2:for $t=1\to T$ do

3: Server selects subset $K_{t}$ of $K$ devices at random

4: Server sends $w^{t}$ to devices in $K_{t}$

5: for each client $k\in K_{t}$ in parallel do

6: Update $w^{t}$ for $E$ epochs of SGD to get $w_{k}^{t+1}$

7: Send $w_{k}^{t+1}$ back to the server

8: Server aggregates
$w^{t+1}\leftarrow\sum_{k=1}^{K}\frac{n_{k}}{N}w_{k}^{t+1}$

##  3\. 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 $X$ used by the device models, the set of invariant features
$X^{I}$ across all devices is one where the target prediction probability is
consistent across all data distributions, (i.e.) $p(Y|X_{i}$ $\in
X^{I},\mathcal{D})$ is approximately constant. Conversely, the set of spurious
features $X^{S}$ consists of features whose prediction probabilities differ
across distributions due to the spurious correlations. Hence, it follows that
$X^{I}\cup X^{S}=X$, and $X^{I}\cap X^{S}=\emptyset$, (i.e.) a feature cannot
be both invariant and spurious.

Our approach to determining whether the $i^{\text{th}}$ feature $X_{i}\in X$
is spurious or invariant, is by measuring the stability of its parameter
weight $w_{i}$. Recent work (Javed et al., 2020; Venkateswaran et al., 2021a)
has shown that if $X_{i}$ is a causal or invariant feature, $w_{i}$ converges
to a fixed magnitude, (i.e.) $\mathbb{E}[Y|X_{i}]=c$ for some constant value
$c$, across all training epochs. Whereas if $\mathbb{E}[Y|X_{i}]$ is changing,
$w_{i}$ would keep changing as well, and hence spurious features have
parameter weights that exhibit high variance. This definition is equivalent to
learning features whose correlations with the target variable are stable.

We leverage this intuition for our federated setting and define a set of masks
$M_{k}=\\{m_{1k},...,m_{jk}\\}$ over the input features for each of the $k\in
K$ clients, where $m_{ik}\in\mathbb{R}$ and $j$ is the number of input
features in $X$. 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
$k\in K$ as:

(5) $m_{ik}\leftarrow
m_{ik}+\frac{1}{j}\sum_{i=1}^{j}v(w_{ik})-\alpha(v(w_{ik})),\;\forall
m_{ik}\in M_{k}$

where hyper-parameter $\alpha$ serves as a scaling factor, $v(w_{ik})$ is the
variance of the weights of feature $X_{i}$ on client $k$, and the second term
is the average variance of all feature weights on client $k$. 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 $K$ clients since their variance is lower than
the average (i.e.) $\frac{1}{j}\sum_{i=1}^{j}v(w_{ik})-\alpha(v(w_{ik}))>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 $\sigma$ 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:

(6) $\sigma(m_{ik})\odot X_{i}\rightarrow\begin{cases}X_{i}&\text{if $X_{i}$
is invariant}\\\ 0&\text{if $X_{i}$ is spurious}\end{cases}\;~{}~{},\forall
X_{i}\in X$

where $\odot$ denotes element-wise multiplication. We then define the local
loss on each client in FedGen as:

(7)
$\displaystyle\underset{w_{k}}{\mathrm{min}}\;\underbrace{f_{k}(Z_{k},w_{k})}_{\text{local
loss}}+\underbrace{\lambda\lVert\nabla_{w_{k}}w_{k}^{\top}f_{k}(Z_{k},w_{k})\rVert^{2}_{2}}_{\text{{FedGen}
penalty}}\;,\;\textrm{where}$ (8) $\displaystyle
f_{k}(Z_{k},w_{k})=\mathbb{E}_{(x_{k},y_{k})\sim
D_{k}}[\ell(Z_{k},y_{k});w_{k}]\;,\;\text{s.t.}\;Z_{k}=\sigma(M_{k})\odot
x_{k}$

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:

(9)
$w^{t+1}=\sum_{k=1}^{K}\frac{n_{k}}{N}w_{k}^{t+1},\;\;M^{t+1}\leftarrow\sum_{k=1}^{K}\frac{n_{k}}{N}M_{k}^{t+1}$

Aggregating the masks across the $K$ 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:

(10)
$\min_{w}F(Z,w)\;,\;\text{where}\;F(Z,w)=\sum_{k=1}^{K}\frac{n_{k}}{N}f_{k}(Z_{k},w_{k})$

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.

Algorithm 2 Federated Generalization FedGen

1:Devices $k\in K$, local epochs $E$, learning rate $\eta$, global model $w$
randomly initialized

2:Local masks $M_{k}\in M$, scaling factor $\alpha$, warm up epochs $e^{init}$

3:for $t=1\to T$ do

4: Server selects subset $K_{t}$ of $K$ devices at random

5: Server sends $w^{t}$ to devices in $K_{t}$

6:ClientUpdate:

7:For each client $k\in K$,

8:Initialize masks $m_{ik}=1,\;\forall m_{ik}\in M_{k}$

9:Initialize mean $\textbf{u}_{ik}=0$, variance $\textbf{v}_{ik}=0$

10:for $e=1\to E$ do

11: Sample batch $x_{k},y_{k}$

12: $\ell_{loc}=[\ell(\sigma(M_{k})\odot x_{k},y_{k});w_{k}]$ $\triangleright$
prediction error

13: $\ell_{1}=\lVert w_{k}\rVert_{1}$ $\triangleright$ $L1$ regularization

14:
$\ell_{pen}=\lambda\lVert\nabla_{w_{k}}w_{k}^{\top}f_{k}(Z_{k},w_{k})\rVert^{2}_{2}$
$\triangleright$ FedGen penalty

15: $\mathcal{L}=\ell_{loc}+\ell_{1}+\ell_{pen}$ $\triangleright$ Total loss

16: $w_{k}=w_{k}-\eta\nabla w_{k}\mathcal{L}$

17: $\textbf{u}_{old}=\textbf{u}_{k}$

18: $\textbf{u}_{k}=\beta w_{k}+(1-\beta)\textbf{u}_{old}$ $\triangleright$
Mean estimate

19:
$\textbf{v}_{k}=\delta\textbf{v}_{k}+(1-\delta)(w_{k}-\textbf{u}_{old})^{2}$
$\triangleright$ Variance Estimate

20: if $e>e^{init}$ then $\triangleright$ Update masks

21:
$m_{ik}\mathrel{+}=\frac{1}{j}\sum_{i=1}^{j}\textbf{v}_{k}(w_{ik})-\alpha(\textbf{v}_{k}(w_{ik})),\forall
m_{ik}\in M_{k}$

22:Server aggregates
$w^{t+1}\leftarrow\sum_{k=1}^{K}\frac{n_{k}}{N}w_{k}^{t+1}$

23:Server aggregates
$M^{t+1}\leftarrow\sum_{k=1}^{K}\frac{n_{k}}{N}M_{k}^{t+1}$

######  Theorem 1.

Given training distributions $D^{tr}$ and a test distribution $D^{te}$, the
set of invariant features $X^{I}$ is the saddle point of the following minimax
problem:

$X^{I}=\underset{w}{\textup{min}}\;\underset{X^{I},X^{S}}{\textup{max}}F(Z,w;D^{te})\;\;,\;\textup{where}\;\;Z=\sigma(M)\odot
X\;,$

where $F(Z,w)$ is the FedGen central objective function, $M$ the final
aggregated masks, and $X^{I},X^{S}$ denote the set of invariant and spurious
features respectively such that $X^{I}\cup X^{S}=X$, and $X^{I}\cap
X^{S}=\emptyset$ (i.e.) they are disjoint.

###### Proof.

For every $Z$, we can partition it into invariant variables $Z^{I}$ and non-
invariant variables $Z^{S}$ as:

(11) $Z^{I}=\sigma(M)\odot X^{I}\ ,\;\;\;\;Z^{S}=\sigma(M)\odot X^{S}.$

Consider a test distribution $D^{*te}$, where the set of spurious features
$X^{S*}$ are not predictive of the output $Y$, and only the invariant features
$X^{I}$ are predictive of $Y$, (i.e.)

(12) $p(Y|Z,D^{*te})=p(Y|Z^{I},D^{*te}),\;\;p(Y|Z,D^{tr})=p(Y|Z^{I},D^{tr})$

Therefore, $F(Z,w\;;D^{*te})$

(13) $\displaystyle\begin{split}&=H(p(Y|Z,D^{*te});p(Y|Z,D^{tr}))\\\
&\overset{(i)}{=}H(p(Y|Z^{I},D^{*te});p(Y|Z^{I},D^{tr}))\\\
&\overset{(ii)}{=}H(p(Y|\sigma(M)\odot X^{I},D^{*te});p(Y|\sigma(M)\odot
X^{I},D^{tr}))\\\
&\overset{(iii)}{=}H(p(Y|X^{I},D^{*te});p(Y|X^{I},D^{tr}))\\\
&=F(X^{I},w\;;D^{*te})\end{split}$

where $H(\cdot)$ is the cross-entropy loss function. Step $(i)$ is obtained
from applying equation (12). Step $(ii)$ is obtained by applying equation
(11), and step $(iii)$ is due to the property of the masks from equation (6).

Recall that $X^{S*}$ was assumed to be non-predictive of $Y$. However, in most
cases, the spurious feature $X^{S}$ would have some predictive power over $Y$
in the training environment. Hence, from the definition of spurious features
$X^{S}$, their biased influence on the model performance during training will
lead to an increased loss in the worst case test environment:

(14) $\underset{X^{S}}{\textup{max}}\;F(Z,w;D^{te})\geq F(Z,w;D^{*te})$

Recall $X^{I}$ denotes the set of invariant features, thus $p(Y|X^{I},D^{te})$
does not depend on $X^{S}$. Therefore,

(15) $\underset{X^{S}}{\textup{max}}\;F(X^{I},w;D^{te})=F(X^{I},w;D^{*te})$

By combining equations (13), (14), and (15), we have:

(16)
$\underset{X^{S}}{\textup{max}}\;F(Z,w;D^{te})\geq\underset{X^{S}}{\textup{max}}\;F(X^{I},w;D^{te})$

The above formulation holds for all $X^{I}$. Hence, taking the maximum over
$X^{I}$ in equation (16) preserves the inequality,

$\underset{X^{I},X^{S}}{\textup{max}}F(Z,w;D^{te})\geq\underset{X^{I},X^{S}}{\textup{max}}F(X^{I},w;D^{te})$

which in turn implies,

$X^{I}=\underset{w}{\textup{min}}\underset{X^{I},X^{S}}{\textup{max}}F(Z,w;D^{te})$

∎

###  3.1. FedGen Convergence

In this section, we analyze the convergence of FedGen by first introducing
some definitions and assumptions for our convergence analysis.

######  Definition 0.

(Smoothness) The function f is L-smooth with constant L ¿ 0 if $\forall
x_{1},x_{2}$

(17) $f(x_{1})-f(x_{2})\leq\langle\nabla
f(x_{2}),x_{1}-x_{2}\rangle+\frac{L}{2}||x_{1}-x_{2}||^{2}$

######  Definition 0.

(Convexity) The function f is $\mu$-strongly convex with $\mu>0$ if $\forall
x_{1},x_{2}$

(18) $f(x_{1})-f(x_{2})\geq\langle\nabla
f(x_{2}),x_{1}-x_{2}\rangle+\frac{\mu}{2}||x_{1}-x_{2}||^{2}$

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 (Li et al., 2018; Vaswani et
al., 2019; Yin et al., 2018).

######  Definition 0.

(B-local dissimilarity) The local functions $f_{k}$ are B-locally dissimilar
at $Z,w$ if $\mathbb{E}||(\nabla f_{k}(Z,w)||^{2}\leq||\nabla
F(Z,w)||^{2}B^{2}$.

We further define $B(w)=\sqrt{\frac{\mathbb{E}||\nabla
f_{k}(Z,w)||^{2}}{||\nabla F(Z,w)||^{2}}}$ when $||\nabla F(Z,w)||^{2}\neq 0$.
As explained in Li et al.(Li et al., 2018), when all the local functions are
the same (i.e.) distributions of data sampled from the devices are i.i.d, we
have $B(w)\rightarrow 1$ for all $w$. However, in the presence of spurious
correlations, the data distributions are heterogeneous and therefore $B>1$,
and larger the value of $B(w)$, the more the heterogeneity of the local data.
We next make the following assumption on the objective functions.

######  Assumption 1.

Assume that:

  1. (1)

The FedGen central objective function F(Z, w) is bounded from below, (i.e.)
$\min_{w}F(Z,w)>-\infty$

  2. (2)

There exists $\epsilon>0$ such that $\mathbb{E}[\nabla f_{k}(Z,w)]\leq||\nabla
F(Z,w)||$, and $\nabla F(Z,w)^{\top}\mathbb{E}[\nabla
f_{k}(Z,w)]\geq\epsilon||\nabla F(Z,w)||^{2}$ holds for all w.

We note that if $\epsilon=1$, then $\nabla f_{k}(Z,w)$ is an unbiased
estimator of $\nabla F(Z,w)$

######  Theorem 5.

Let Assumption 1 hold. Assume F(Z, w) is L-smooth and $\mu$-strongly convex
and functions $f_{k}$ are B-locally dissimilar with learning rate $\eta$. Then
we can show that for $\rho=-\eta(\epsilon-\frac{L\eta B^{2}}{2})$, at
iteration $t$ of FedGen, we have the following expected decrease in the global
objective:

$\mathbb{E}[F(Z^{t+1},w^{t+1})]-F(Z_{*},w_{*})\leq(1-2\mu\rho)(F(Z^{t},w^{t})-F(Z_{*},w_{*}))$

We present the detailed proof in Appendix A.1

######  Theorem 6.

Let Assumption 1 hold. Assume F(Z, w) is L-smooth and non-convex and functions
$f_{k}$ are B-locally dissimilar with learning rate $\eta$. Also, let
$F(Z^{0},w^{0})-F(Z^{*},w^{*})=\Updelta$, where $w^{*}=\min_{w}F(Z,w)$. Then
after $T$ iterations of FedGen, we have

(19) $\frac{1}{T}\sum_{t=0}^{T-1}\rho\mathbb{E}[||\nabla
F(Z^{t},w^{t})||^{2}]\leq\Updelta$

The proof can be found in Appendix A.2

##  4\. 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 (Arjovsky et al., 2019; Ahuja et al., 2020; Choe et al., 2020;
Krueger et al., 2021), we augment these benchmark datasets with spurious
features $X^{S}$ such that it has a strong spurious correlation with the
target labels $Y$ in the training data distributions (i.e.)
$p(Y|X^{S},D^{tr})\geq\alpha=0.8$, and this spurious relationship does not
hold in the test distribution (i.e.) $p(Y|X^{S},D^{te})\leq\alpha=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.

Dataset | Devices/Clients | Train Seq. | Test Seq. | #Spurious
---|---|---|---|---
SST-2 | 100 | 67,350 | 873 | 1
AG News | 100 | 120,000 | 7600 | 1
HAR | 30 | 7352 | 2947 | 2
Water Quality | 120 | 132,227 | 14,237 | 1
BPIC 2018 | 100 | 306,615 | 63,692 | 1

Table 1. Summary of Federated Datasets

Algorithm | SST-2 | AG News | HAR | Water Quality | BPIC 2018
---|---|---|---|---|---
| Train | Test | Train | Test | Train | Test |  |  |  |
ERM | 85.6 $\pm$ 0.1 | 9.6 $\pm$ 0.3 | 92.3 $\pm$ 0.1 | 63.4 $\pm$ 0.5 | 98.1 $\pm$ 0.3 | 73.2 $\pm$ 0.5 | 88.7 $\pm$ 0.3 | 65.3 $\pm$ 0.3 | 77.1 $\pm$ 0.7 | 73.7 $\pm$ 0.4
FedAvg | 83.3 $\pm$ 0.1 | 9.6 $\pm$ 0.1 | 91.5 $\pm$ 0.6 | 61.3 $\pm$ 0.2 | 98.0 $\pm$ 0.5 | 70.1 $\pm$ 1.6 | 89.3 $\pm$ 0.7 | 63.8 $\pm$ 1.6 | 77.1 $\pm$ 2.1 | 71.7 $\pm$ 1.3
FedProx | 75.0 $\pm$ 0.1 | 9.6 $\pm$ 0.1 | 91.3 $\pm$ 0.8 | 61.1 $\pm$ 0.3 | 96.7 $\pm$ 0.5 | 71.5 $\pm$ 0.6 | 87.7 $\pm$ 1.1 | 66.1 $\pm$ 0.7 | 73.9 $\pm$ 1.1 | 71.2 $\pm$ 0.5
FedGen (ours) | 58.3 $\pm$ 0.2 | 60.9 $\pm$ 0.3 | 81.8 $\pm$ 0.3 | 82.3 $\pm$ 0.1 | 94.2 $\pm$ 1.2 | 87.5 $\pm$ 0.3 | 64.3 $\pm$ 0.2 | 85.6 $\pm$ 1.1 | 74.8 $\pm$ 0.8 | 86.1 $\pm$ 0.4
Inv-FedAvg | 66.6 $\pm$ 0.2 | 62.5 $\pm$ 0.2 | 82.1 $\pm$ 0.4 | 82.8 $\pm$ 0.1 | 93.0 $\pm$ 0.4 | 89.9 $\pm$ 0.5 | 74.9 $\pm$ 1.1 | 88.1 $\pm$ 0.6 | 78.0 $\pm$ 0.3 | 88.7 $\pm$ 0.3

Table 2. Accuracy achieved by comparison approaches for all datasets

###  4.1. Datasets

####  4.1.1. SST-2:

We modify the Stanford Sentiment Treebank (SST-2) (Socher et al., 2013), a
benchmark dataset for binary sentiment analysis with spurious correlations in
a similar manner as Choe et al. (2020). We add a spurious feature $X^{S}$ by
pairing each sentiment with a final punctuation mark (. or !) in the
corresponding text as follows:

$\displaystyle p(Y=0\;|\;X^{S}=.\;,D)=p(Y=1\;|\;X^{S}=\;!\;,D)$
$\displaystyle=\alpha$ $\displaystyle
p(Y=0\;|\;X^{S}=\;!\;,D)=p(Y=1\;|\;X^{S}=.\;,D)$ $\displaystyle=1-\alpha$

where we set $\alpha$ 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. (2020)). 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.

####  4.1.2. AG News:

To evaluate FedGen for multi-class NLP predictions, we use the AG News dataset
(Zhang et al., 2015), 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 $\mathcal{P}=\\{.\;!\;$@$\;\wedge\\}$ as follows:

$\displaystyle p(Y=0\;|\;X^{S}=.\;,D)=p(Y=1\;|\;X^{S}=!\;,D)$
$\displaystyle=\alpha$ $\displaystyle
p(Y=2\;|\;X^{S}=@\;,D)=p(Y=3\;|\;X^{S}=\wedge\;,D)$ $\displaystyle=\alpha$
$\displaystyle
p(Y=y_{j}\;|\;X^{S}=\text{rand}(\mathcal{P}\setminus\mathcal{P}_{j})\;,D)$
$\displaystyle=1-\alpha$

and set $\alpha$ 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.

####  4.1.3. Human Activity Recognition (HAR)

The HAR dataset (Anguita et al., 2013) 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:

$\displaystyle p(Y=y_{j}\;|\;X^{S}=S_{j}\;,D)$ $\displaystyle=\alpha$
$\displaystyle p(Y=y_{j}\;|\;X^{S}=\text{rand}(S\setminus S_{j})\;,D)$
$\displaystyle=1-\alpha$

where $y_{j}$ and $S_{j}$ refer to the $j$th activity and smartphone model
respectively.

####  4.1.4. Water Quality Detection

The dataset (Environmental Protection Agency, 2021) comprises of water quality
data collected from different cities across the United States. It includes
sensor readings such as turbidity, conductivity, pH, temperature, DO2, 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 low ($L$) and high ($H$) values and spurious correlate them to the
contamination labels ($\mathcal{C},\emptyset$) as follows:

$\displaystyle p(Y=\mathcal{C}\;|\;X^{S}\in H\;,D)=p(Y=\emptyset\;|\;X^{S}\in
L\;,D)$ $\displaystyle=\alpha$ $\displaystyle p(Y=\mathcal{C}\;|\;X^{S}\in
L\;,D)=p(Y=\emptyset\;|\;X^{S}\in H\;,D)$ $\displaystyle=1-\alpha$

We distribute the training data among 120 clients corresponding to the
different cities.

####  4.1.5. Business Process Mining:

We use a benchmark process mining dataset (BPIC
2018)111https://data.4tu.nl/articles/BPI_Challenge_2018/12688355 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 multi-distribution setting, we consider each department as a
distribution and use three of them for training and the fourth as our out-of-
distribution (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 ($L$) and high ($H$) values and
spuriously correlate them to the two business process variants
($\mathcal{A},\mathcal{B}$) as follows:

$\displaystyle p(Y=\mathcal{A}\;|\;X^{S}\in
H\;,D)=p(Y=\mathcal{B}\;|\;X^{S}\in L\;,D)$ $\displaystyle=\alpha$
$\displaystyle p(Y=\mathcal{A}\;|\;X^{S}\in
L\;,D)=p(Y=\mathcal{B}\;|\;X^{S}\in H\;,D)$ $\displaystyle=1-\alpha$

and define $\alpha$ as before. We distribute the training data among 100
clients.

###  4.2. 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. (McMahan et al., 2017)
where model parameters are averaged element-wise.

  * •

FedProx: Approach by Li et al. (Li et al., 2018) 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.

###  4.3. Implementation Details

We implement FedGen, ERM, FedAvg, and Inv-FedAvg using the PyTorch library
(Paszke et al., 2019) and implement FedProx based on their publicly available
code222https://github.com/litian96/FedProx. 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 (Li et al., 2018; Sahu et al., 2018). 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.

###  4.4. Results

(a) SST-2

(b) AG News

(c) HAR

(d) Water Quality

(e) BPIC 2018

Figure 2. Effect of number of local training epochs on the different
approaches across all datasets

(a) SST-2

(b) AG News

(c) HAR

(d) Water Quality

(e) BPIC 2018

Figure 3. Convergence rates of the different approaches across all datasets

####  4.4.1. 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.

####  4.4.2. Effect of choice of local training epochs

While a large number of local training epochs $E$ can help in reducing the
communication costs, previous work has also shown that the value of $E$ can
impact the performance of FedAvg and can sometimes lead to divergence (McMahan
et al., 2017; Caldas et al., 2018; Sahu et al., 2018). We conduct an
experimental study on the effect of the choice of $E$ over the different
comparison approaches across all five datasets. The candidate local epochs we
consider are $E\in\\{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.

####  4.4.3. 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 $E$ 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.

####  4.4.4. 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 $\alpha$ 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 $\sigma(\textbf{M})$, essentially reducing the loss function to that
used by FedAvg. This is reflected by the accuracy values, which are similar to
those achieved by FedAvg, thereby demonstrating that our proposed masking
function is the primary reason for model generalization. We finally remove the
penalty term while optimizing model weights (Equation 7) and again observe a
drop in accuracy across all datasets since the regularization is important to
ensure that the model does not get influenced by spurious correlations in any
one distribution.

Algorithm | Test Accuracy (%)
---|---
| SST-2 | AG News | HAR | Water Quality | BPIC 2018
FedGen | 60.9 | 82.3 | 87.5 | 85.6 | 86.1
$-$ scaling ($\alpha$) | 9.6 | 55.6 | 74.9 | 65.5 | 71.2
$-$ $\sigma(\textbf{M})$ | 9.6 | 61.1 | 68.1 | 63.4 | 70.0
$-$ penalty | 9.6 | 59.7 | 70.4 | 61.3 | 69.3

Table 3. Ablation Study

Algorithm | Number of Distributions
---|---
| 2 | 4 | 6 | 8
ERM | 63.4 | 63.0 | 62.6 | 62.0
FedAvg | 61.3 | 61.0 | 60.3 | 60.0
FedProx | 61.1 | 60.7 | 60.6 | 60.3
FedGen (ours) | 82.3 | 82.1 | 82.1 | 82.2
Inv-FedAvg | 82.8 | 82.8 | 82.6 | 82.8

Table 4. Varying Number of Training Distributions

####  4.4.5. Varying Number of Training Distributions

We measure the robustness of FedGen to varying numbers of distributions in the
training data using the AG News dataset. We augment the dataset with more
distributions with differing levels of spurious correlations with
$\alpha>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.

##  5\. 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 (Li et al., 2020). To improve the performance of the
popular FedAvg approach (McMahan et al., 2017) in statistically heterogeneous
settings, FedProx (Li et al., 2018) and Agnostic Federated Learning (Mohri et
al., 2019) are two approaches that we described in Section 2. Karimireddy et
al. propose SCAFFOLD (Karimireddy et al., 2020) 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.(Jeong et al., 2018) violates the key
privacy assumption of federated learning, and sending globally-shared proxy
data to all clients (Zhao et al.(Zhao et al., 2018), Huang et al.(Huang et
al., 2018)) requires effort to carefully generate the data.

There has also been work on improving the communication costs of FedAvg. Wang
et al.(Wang et al., 2020) 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.(Luping et al., 2019) 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.(Xie et al.,
2019)) 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. (Shankar et al., 2018) augment data
using instances perturbed along directions of domain change and use a second
classifier to capture this. Volpi et al. (Volpi et al., 2018) apply this to
single domain data, while Carlucci et al. (Carlucci et al., 2019) 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 (Schölkopf et al., 2012) and causal model
discovery (Pearl, 2009) aims to find an underlying causal graph to obtain the
invariant feature set. Peters et al.(Peters et al., 2016) 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.(Rojas-Carulla et
al., 2018) 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 (Daumé III, 2009).
Khosla et al. (Khosla et al., 2012) applied decomposition to domain
generalization by retaining only the common parameter for inference. Li et al.
(Li et al., 2017) 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. (Piratla et al., 2020) 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
(Finn et al., 2017), distribution shifts in only test domains (Dou et al.,
2019), or only considering label shifts (Lipton et al., 2018). There have also
been some prior work on generalization for other applications like business
process predictions (Venkateswaran et al., 2021b) and medical diagnosis using
human annotated spurious features (Srivastava et al., 2020).

##  6\. 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.

## References

  * (1)
  * Ahuja et al. (2020) Kartik Ahuja, Karthikeyan Shanmugam, Kush Varshney, and Amit Dhurandhar. 2020.  Invariant risk minimization games. In _International Conference on Machine Learning_. PMLR, 145–155.
  * Albuquerque et al. (2019) Isabela Albuquerque, João Monteiro, Tiago H Falk, and Ioannis Mitliagkas. 2019.  Adversarial target-invariant representation learning for domain generalization.  _arXiv preprint arXiv:1911.00804_ (2019).
  * Anguita et al. (2013) Davide Anguita, Alessandro Ghio, Luca Oneto, Xavier Parra, Jorge Luis Reyes-Ortiz, et al. 2013\.  A public domain dataset for human activity recognition using smartphones.. In _Esann_ , Vol. 3. 3.
  * Arjovsky et al. (2019) Martin Arjovsky, Léon Bottou, Ishaan Gulrajani, and David Lopez-Paz. 2019.  Invariant risk minimization.  _arXiv preprint arXiv:1907.02893_ (2019).
  * Beery et al. (2018) Sara Beery, Grant Van Horn, and Pietro Perona. 2018\.  Recognition in terra incognita. In _Proceedings of the European Conference on Computer Vision (ECCV)_. 456–473.
  * Bengio et al. (2019) Yoshua Bengio, Tristan Deleu, Nasim Rahaman, Rosemary Ke, Sébastien Lachapelle, Olexa Bilaniuk, Anirudh Goyal, and Christopher Pal. 2019\.  A meta-transfer objective for learning to disentangle causal mechanisms.  _arXiv preprint arXiv:1901.10912_ (2019).
  * Bonawitz et al. (2019) Keith Bonawitz, Hubert Eichner, Wolfgang Grieskamp, Dzmitry Huba, Alex Ingerman, Vladimir Ivanov, Chloe Kiddon, Jakub Konečnỳ, Stefano Mazzocchi, H Brendan McMahan, et al. 2019\.  Towards federated learning at scale: System design.  _arXiv preprint arXiv:1902.01046_ (2019).
  * Caldas et al. (2018) Sebastian Caldas, Sai Meher Karthik Duddu, Peter Wu, Tian Li, Jakub Konečnỳ, H Brendan McMahan, Virginia Smith, and Ameet Talwalkar. 2018.  Leaf: A benchmark for federated settings.  _arXiv preprint arXiv:1812.01097_ (2018).
  * Carlucci et al. (2019) Fabio M Carlucci, Antonio D’Innocente, Silvia Bucci, Barbara Caputo, and Tatiana Tommasi. 2019\.  Domain generalization by solving jigsaw puzzles. In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_. 2229–2238.
  * Choe et al. (2020) Yo Joong Choe, Jiyeon Ham, and Kyubyong Park. 2020\.  An empirical study of invariant risk minimization.  _arXiv preprint arXiv:2004.05007_ (2020).
  * Daumé III (2009) Hal Daumé III. 2009\.  Frustratingly easy domain adaptation.  _arXiv preprint arXiv:0907.1815_ (2009).
  * DeGrave et al. (2021) Alex J DeGrave, Joseph D Janizek, and Su-In Lee. 2021\.  AI for radiographic COVID-19 detection selects shortcuts over signal.  _Nature Machine Intelligence_ 3, 7 (2021), 610–619.
  * Dou et al. (2019) Qi Dou, Daniel Coelho de Castro, Konstantinos Kamnitsas, and Ben Glocker. 2019.  Domain generalization via model-agnostic learning of semantic features. In _Advances in Neural Information Processing Systems_. 6450–6461.
  * Environmental Protection Agency (2021) Environmental Protection Agency. 2021\.  Open Data US EPA.  https://www.epa.gov/data
  * Finn et al. (2017) Chelsea Finn, Pieter Abbeel, and Sergey Levine. 2017\.  Model-agnostic meta-learning for fast adaptation of deep networks. In _International conference on machine learning_. PMLR, 1126–1135.
  * Huang et al. (2018) Li Huang, Yifeng Yin, Zeng Fu, Shifa Zhang, Hao Deng, and Dianbo Liu. 2018\.  Loadaboost: Loss-based adaboost federated machine learning on medical data.  _arXiv preprint arXiv:1811.12629_ (2018).
  * Javed et al. (2020) Khurram Javed, Martha White, and Yoshua Bengio. 2020\.  Learning causal models online.  _arXiv preprint arXiv:2006.07461_ (2020).
  * Jeong et al. (2018) Eunjeong Jeong, Seungeun Oh, Hyesung Kim, Jihong Park, Mehdi Bennis, and Seong-Lyun Kim. 2018\.  Communication-efficient on-device machine learning: Federated distillation and augmentation under non-iid private data.  _arXiv preprint arXiv:1811.11479_ (2018).
  * Kairouz et al. (2019) Peter Kairouz, H Brendan McMahan, Brendan Avent, Aurélien Bellet, Mehdi Bennis, Arjun Nitin Bhagoji, Kallista Bonawitz, Zachary Charles, Graham Cormode, Rachel Cummings, et al. 2019\.  Advances and open problems in federated learning.  _arXiv preprint arXiv:1912.04977_ (2019).
  * Karimireddy et al. (2020) Sai Praneeth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank Reddi, Sebastian Stich, and Ananda Theertha Suresh. 2020.  Scaffold: Stochastic controlled averaging for federated learning. In _International Conference on Machine Learning_. PMLR, 5132–5143.
  * Khosla et al. (2012) Aditya Khosla, Tinghui Zhou, Tomasz Malisiewicz, Alexei A Efros, and Antonio Torralba. 2012\.  Undoing the damage of dataset bias. In _European Conference on Computer Vision_. Springer, 158–171.
  * Krueger et al. (2021) David Krueger, Ethan Caballero, Joern-Henrik Jacobsen, Amy Zhang, Jonathan Binas, Dinghuai Zhang, Remi Le Priol, and Aaron Courville. 2021\.  Out-of-distribution generalization via risk extrapolation (rex). In _International Conference on Machine Learning_. PMLR, 5815–5826.
  * Li et al. (2017) Da Li, Yongxin Yang, Yi-Zhe Song, and Timothy M Hospedales. 2017\.  Deeper, broader and artier domain generalization. In _Proceedings of the IEEE international conference on computer vision_. 5542–5550.
  * Li et al. (2020) Tian Li, Anit Kumar Sahu, Ameet Talwalkar, and Virginia Smith. 2020\.  Federated learning: Challenges, methods, and future directions.  _IEEE Signal Processing Magazine_ 37, 3 (2020), 50–60.
  * Li et al. (2018) Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith. 2018\.  Federated optimization in heterogeneous networks.  _arXiv preprint arXiv:1812.06127_ (2018).
  * Lipton et al. (2018) Zachary Lipton, Yu-Xiang Wang, and Alexander Smola. 2018\.  Detecting and correcting for label shift with black box predictors. In _International conference on machine learning_. PMLR, 3122–3130.
  * Luping et al. (2019) Wang Luping, Wang Wei, and LI Bo. 2019.  CMFL: Mitigating communication overhead for federated learning. In _2019 IEEE 39th international conference on distributed computing systems (ICDCS)_. IEEE, 954–964.
  * McMahan et al. (2017) Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. 2017\.  Communication-efficient learning of deep networks from decentralized data. In _Artificial intelligence and statistics_. PMLR, 1273–1282.
  * Mohri et al. (2019) Mehryar Mohri, Gary Sivek, and Ananda Theertha Suresh. 2019\.  Agnostic federated learning. In _International Conference on Machine Learning_. PMLR, 4615–4625.
  * Nagarajan et al. (2020) Vaishnavh Nagarajan, Anders Andreassen, and Behnam Neyshabur. 2020.  Understanding the failure modes of out-of-distribution generalization.  _arXiv preprint arXiv:2010.15775_ (2020).
  * Paszke et al. (2019) Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. 2019\.  Pytorch: An imperative style, high-performance deep learning library.  _Advances in neural information processing systems_ 32 (2019), 8026–8037.
  * Pearl (2009) Judea Pearl. 2009\.  _Causality_.  Cambridge university press.
  * Peters et al. (2016) Jonas Peters, Peter Bühlmann, and Nicolai Meinshausen. 2016\.  Causal inference by using invariant prediction: identification and confidence intervals.  _Journal of the Royal Statistical Society: Series B (Statistical Methodology)_ 78, 5 (2016), 947–1012.
  * Piratla et al. (2020) Vihari Piratla, Praneeth Netrapalli, and Sunita Sarawagi. 2020\.  Efficient domain generalization via common-specific low-rank decomposition. In _International Conference on Machine Learning_. PMLR, 7728–7738.
  * Recht et al. (2019) Benjamin Recht, Rebecca Roelofs, Ludwig Schmidt, and Vaishaal Shankar. 2019.  Do imagenet classifiers generalize to imagenet?. In _International Conference on Machine Learning_. PMLR, 5389–5400.
  * Rojas-Carulla et al. (2018) Mateo Rojas-Carulla, Bernhard Schölkopf, Richard Turner, and Jonas Peters. 2018.  Invariant models for causal transfer learning.  _The Journal of Machine Learning Research_ 19, 1 (2018), 1309–1342.
  * Rosenfeld et al. (2020) Elan Rosenfeld, Pradeep Ravikumar, and Andrej Risteski. 2020\.  The risks of invariant risk minimization.  _arXiv preprint arXiv:2010.05761_ (2020).
  * Sahu et al. (2018) Anit Kumar Sahu, Tian Li, Maziar Sanjabi, Manzil Zaheer, Ameet Talwalkar, and Virginia Smith. 2018\.  On the convergence of federated optimization in heterogeneous networks.  _arXiv preprint arXiv:1812.06127_ 3 (2018), 3.
  * Schölkopf et al. (2012) Bernhard Schölkopf, Dominik Janzing, Jonas Peters, Eleni Sgouritsa, Kun Zhang, and Joris Mooij. 2012.  On causal and anticausal learning.  _arXiv preprint arXiv:1206.6471_ (2012).
  * Shankar et al. (2018) Shiv Shankar, Vihari Piratla, Soumen Chakrabarti, Siddhartha Chaudhuri, Preethi Jyothi, and Sunita Sarawagi. 2018.  Generalizing across domains via cross-gradient training.  _arXiv preprint arXiv:1804.10745_ (2018).
  * Socher et al. (2013) Richard Socher, Alex Perelygin, Jean Wu, Jason Chuang, Christopher D Manning, Andrew Y Ng, and Christopher Potts. 2013.  Recursive deep models for semantic compositionality over a sentiment treebank. In _Proceedings of the 2013 conference on empirical methods in natural language processing_. 1631–1642.
  * Srivastava et al. (2020) Megha Srivastava, Tatsunori Hashimoto, and Percy Liang. 2020\.  Robustness to spurious correlations via human annotations. In _International Conference on Machine Learning_. PMLR, 9109–9119.
  * Tzeng et al. (2017) Eric Tzeng, Judy Hoffman, Kate Saenko, and Trevor Darrell. 2017\.  Adversarial discriminative domain adaptation. In _Proceedings of the IEEE conference on computer vision and pattern recognition_. 7167–7176.
  * Vapnik (1992) Vladimir Vapnik. 1992\.  Principles of risk minimization for learning theory. In _Advances in neural information processing systems_. 831–838.
  * Vaswani et al. (2019) Sharan Vaswani, Francis Bach, and Mark Schmidt. 2019\.  Fast and faster convergence of sgd for over-parameterized models and an accelerated perceptron. In _The 22nd International Conference on Artificial Intelligence and Statistics_. PMLR, 1195–1204.
  * Venkateswaran et al. (2021a) Praveen Venkateswaran, Vinod Muthusamy, Vatche Isahagian, and Nalini Venkatasubramanian. 2021a.  Environment Agnostic Invariant Risk Minimization for Classification of Sequential Datasets. In _Proceedings of the 27th ACM SIGKDD Conference on Knowledge Discovery & Data Mining_. 1615–1624.
  * Venkateswaran et al. (2021b) Praveen Venkateswaran, Vinod Muthusamy, Vatche Isahagian, and Nalini Venkatasubramanian. 2021b.  Robust and Generalizable Predictive Models for Business Processes. In _International Conference on Business Process Management_. Springer, 105–122.
  * Volpi et al. (2018) Riccardo Volpi, Hongseok Namkoong, Ozan Sener, John C Duchi, Vittorio Murino, and Silvio Savarese. 2018\.  Generalizing to unseen domains via adversarial data augmentation.  _Advances in neural information processing systems_ 31 (2018).
  * Wang et al. (2020) Hongyi Wang, Mikhail Yurochkin, Yuekai Sun, Dimitris Papailiopoulos, and Yasaman Khazaeni. 2020\.  Federated learning with matched averaging.  _arXiv preprint arXiv:2002.06440_ (2020).
  * Xie et al. (2019) Cong Xie, Sanmi Koyejo, and Indranil Gupta. 2019.  Asynchronous federated optimization.  _arXiv preprint arXiv:1903.03934_ (2019).
  * Yin et al. (2018) Dong Yin, Ashwin Pananjady, Max Lam, Dimitris Papailiopoulos, Kannan Ramchandran, and Peter Bartlett. 2018.  Gradient diversity: a key ingredient for scalable distributed learning. In _International Conference on Artificial Intelligence and Statistics_. PMLR, 1998–2007.
  * Zhang et al. (2021) Dinghuai Zhang, Kartik Ahuja, Yilun Xu, Yisen Wang, and Aaron Courville. 2021.  Can Subnetwork Structure be the Key to Out-of-Distribution Generalization?. In _International Conference on Machine Learning_. PMLR, 12356–12367.
  * Zhang et al. (2015) Xiang Zhang, Junbo Zhao, and Yann LeCun. 2015.  Character-level convolutional networks for text classification.  _Advances in neural information processing systems_ 28 (2015), 649–657.
  * Zhao et al. (2018) Yue Zhao, Meng Li, Liangzhen Lai, Naveen Suda, Damon Civin, and Vikas Chandra. 2018\.  Federated learning with non-iid data.  _arXiv preprint arXiv:1806.00582_ (2018).

##  Appendix A Appendix

###  A.1. Proof of Theorem 2

###### Proof.

Since $F(Z,w)$ is L-smooth, we have

(20) $\displaystyle\begin{split}&F(Z^{t+1},w^{t+1})-F(Z^{t},w^{t})\\\
&\leq\langle\nabla
F(Z^{t},w^{t}),w^{t+1}-w^{t}\rangle+\frac{L}{2}||w^{t+1}-w^{t}||^{2}\\\
&\leq-\nabla F(Z^{t},w^{t})^{\top}\eta\nabla
f_{k}(Z^{t},w^{t})+\frac{L}{2}||\eta\nabla
f_{k}(Z^{t},w^{t})||^{2}\end{split}$

With Assumption 1 and B-local dissimilarity, we get

(21)
$\displaystyle\begin{split}&\mathbb{E}[F(Z^{t+1},w^{t+1})]-F(Z^{t},w^{t})\\\
&\leq-\eta\nabla F(Z^{t},w^{t})^{\top}\mathbb{E}[\nabla
f_{k}(Z^{t},w^{t})]+\frac{L\eta^{2}}{2}\mathbb{E}[||\nabla
f_{k}(Z^{t},w^{t})||^{2}]\\\ &\leq-\eta\epsilon||\nabla
F(Z^{t},w^{t})||^{2}+\frac{L\eta^{2}B^{2}}{2}||\nabla F(Z^{t},w^{t})||^{2}\\\
&\leq-\eta(\epsilon-\frac{L\eta B^{2}}{2})||\nabla
F(Z^{t},w^{t})||^{2}\end{split}$

Setting $\rho=-\eta(\epsilon-\frac{L\eta B^{2}}{2})$ we get,

(22) $\mathbb{E}[F(Z^{t+1},w^{t+1})]-F(Z^{t},w^{t})\leq-\rho||\nabla
F(Z^{t},w^{t})||^{2}$

Since $F(Z,w)$ is $\mu$-strongly convex, for some $Z^{\prime}$ and
$w^{\prime}$, we get:

(23) $F(Z^{\prime},w^{\prime})-F(Z^{t},w^{t})\geq\nabla
F(Z^{t},w^{t})^{\top}(w^{\prime}-w^{t})+\frac{\mu}{2}||w^{\prime}-w^{t}||^{2}$

Since $F(Z^{\prime},w^{\prime})$ is quadratic in $w^{\prime}$, its minimal
value is obtained when $\nabla F(Z^{\prime},w^{\prime})=\nabla
F(Z^{t},w^{t})+\mu(w^{\prime}-w^{t})=0$. This occurs when
$w^{\prime}=w^{t}-\frac{\nabla F(Z^{t},w^{t})}{\mu}$. Substituting for
$w^{\prime}$ in Equation 23 we get,

(24) $F(Z_{*},w_{*})\geq F(Z^{\prime},w^{\prime})\geq
F(Z^{t},w^{t})-\frac{||\nabla F(Z^{t},w^{t})||^{2}}{2\mu}$

(25) $2\mu(F(Z^{t},w^{t})-F(Z_{*},w_{*}))\leq||\nabla F(Z^{t},w^{t})||^{2}$

Substituting into Equation 22 we get,

(26)
$\mathbb{E}[F(Z^{t+1},w^{t+1})]-F(Z^{t},w^{t})\leq-2\mu\rho(F(Z^{t},w^{t})-F(Z_{*},w_{*}))$

Re-arranging and subtracting $F(Z_{*},w_{*})$ from both sides,

(27)
$\mathbb{E}[F(Z^{t+1},w^{t+1})]-F(Z_{*},w_{*})\leq(1-2\mu\rho)(F(Z^{t},w^{t})-F(Z_{*},w_{*}))$

∎

###  A.2. Proof of Theorem 3

###### Proof.

By taking the full expectation of Equation 21, we get

(28)
$\mathbb{E}[F(Z^{t+1},w^{t+1})]\leq\mathbb{E}[F(Z^{t},w^{t})]-\rho\mathbb{E}[||\nabla
F(Z^{t},w^{t})||^{2}]$

Then summing up Equation 28 over global iteration T, we obtain

(29) $\mathbb{E}[F(Z^{t+1},w^{t+1})]\leq
F(Z^{0},w^{0})-\sum_{t=0}^{T-1}\rho\mathbb{E}[||\nabla F(Z^{t},w^{t})||^{2}]$

From Assumption 1, we see that
$F(Z^{*},w^{*})\leq\mathbb{E}[F(Z^{t+1},w^{t+1})]$, using which we can write

(30) $F(Z^{*},w^{*})\leq
F(Z^{0},w^{0})-\frac{1}{T}\sum_{t=0}^{T-1}\eta\rho\mathbb{E}[||\nabla
F(Z^{t},w^{t})||^{2}]$

Re-arranging,

(31) $\frac{1}{T}\sum_{t=0}^{T-1}\rho\mathbb{E}[||\nabla
F(Z^{t},w^{t})||^{2}]\leq F(Z^{0},w^{0})-F(Z^{*},w^{*})$

∎

###  A.3. Implementation Details

We simulate the federated learning setup (1 server and $N$ 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.

Hyper-parameter | Range
---|---
Learning rate | [$10^{-5},10^{-1}$]
Rounds | [$11,201$]
Epochs | [$5,150$]
Regularization weight | [$10^{-6},10^{-2}$]
Penalty | [$10,10^{4}$]
Scaling factor | [$10^{-2},10$]
Proximal term (FedProx) | [$10^{-5},10^{-1}$]

Table 5. Hyper-parameter ranges tried for all comparison 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 $7e^{-6}$. We set the penalty of FedGen to $1$ and the
scaling factor to $10$.