Original: FedSR: A Simple and Effective Domain Generalization Method for Federated Learning
Revised: FedSR: A Simple and Effective Domain Generalization Method for Federated Learning

Original: Anonymous Author(s)
Revised: Anonymous Author(s)

Original: AffiliationAddress email
Revised: AffiliationAddress email

Original: Abstract
Revised: Abstract

Original: Federated Learning (FL) refers to the decentralized and privacy-preserving machinelearning framework in which multiple clients collaborate (with the help of a centralserver) to train a global model without sharing their data.
Revised: Federated Learning (FL) refers to the decentralized and privacy-preserving machinelearning framework in which multiple clients collaborate (with the help of a centralserver) to train a global model without sharing their data.

Original: However, most existing
Revised: However, most existing

Original: FL methods only focus on maximizing the model’s performance on the sourceclients’ data (e.g., mobile users) without considering its generalization ability tounknown target data (e.g., a new user).
Revised: FL methods only focus on maximizing the model’s performance on the sourceclients’ data (e.g., mobile users) without considering its generalization ability tounknown target data (e.g., a new user).

Original: In this paper, we incorporate the problem of
Revised: In this paper, we incorporate the problem of

Original: Domain Generalization (DG) into Federated Learning to tackle the aforementionedissue.
Revised: Domain Generalization (DG) into Federated Learning to tackle the aforementionedissue.

Original: However, virtually all existing DG methods require a centralized settingwhere data is shared across the domains, which violates the principles of decentralized FL and hence not applicable.
Revised: However, virtually all existing DG methods require a centralized settingwhere data is shared across the domains, which violates the principles of decentralized FL and hence not applicable.

Original: To this end, we propose a simple yet novelrepresentation learning framework, namely FedSR, which enables domain generalization while still respecting the decentralized and privacy-preserving natures ofthis FL setting.
Revised: To this end, we propose a simple yet novelrepresentation learning framework, namely FedSR, which enables domain generalization while still respecting the decentralized and privacy-preserving natures ofthis FL setting.

Original: Motivated by classical machine learning algorithms, we aim tolearn a simple representation of the data for better generalization.
Revised: Motivated by classical machine learning algorithms, we aim tolearn a simple representation of the data for better generalization.

Original: In particular,we enforce an L2-norm regularizer on the representation and a conditional mutualinformation (between the representation and the data given the label) regularizerto encourage the model to only learn essential information (while ignoring spurious correlations such as the background).
Revised: In particular,we enforce an L2-norm regularizer on the representation and a conditional mutualinformation (between the representation and the data given the label) regularizerto encourage the model to only learn essential information (while ignoring spurious correlations such as the background).

Original: Furthermore, we provide theoreticalconnections between the above two objectives and representation alignment indomain generalization.
Revised: Furthermore, we provide theoreticalconnections between the above two objectives and representation alignment indomain generalization.

Original: Extensive experimental results suggest that our methodsignificantly outperforms relevant baselines in this particular problem.
Revised: Extensive experimental results suggest that our methodsignificantly outperforms relevant baselines in this particular problem.

Original:  Introduction
Revised:  Introduction

Original: In this paper, we are interested the problem of decentralized Federated Learning (FL), where Kclients collaborate to jointly train a machine learning model with their local (decentralized) data.
Revised: In this paper, we are interested the problem of decentralized Federated Learning (FL), where Kclients collaborate to jointly train a machine learning model with their local (decentralized) data.

Original: Forprivacy reasons, we consider the setting where it is desirable/preferred that the clients do not sharetheir data with each other.
Revised: Forprivacy reasons, we consider the setting where it is desirable/preferred that the clients do not sharetheir data with each other.

Original: For example, [28] brings up an example scenario where mobile devicescollaborate to train a model without sending their data to a central server.
Revised: For example, [28] brings up an example scenario where mobile devicescollaborate to train a model without sending their data to a central server.

Original: Another example is that Khospitals train a diagnosis model together but do not want to share their patients’ data.
Revised: Another example is that Khospitals train a diagnosis model together but do not want to share their patients’ data.

Original: However, a major concern with this machine learning framework is that it often does not account for the possible distribution shift on the target data, a problem commonly referred to as Domain General- ization (DG).
Revised: However, a major concern with this machine learning framework is that it often does not account forthe possible distribution shift on the target data, a problem commonly referred to as Domain Generalization (DG).

Original: Furthermore, and perhaps more problematically, it not straightforward to incorporateexisting DG techniques into this FL setting due to its privacy-preserving nature.
Revised: Furthermore, and perhaps more problematically, it not straightforward to incorporateexisting DG techniques into this FL setting due to its privacy-preserving nature.

Original: Specifically, mostdomain invariant representation learning methods [9, 38, 31, 22, 32] require sharing/comparing therepresentation of the data across domains, while meta-learning based methods [21, 5] require accessto data from a pool of multiple source domains (in a central server).
Revised: Specifically, mostdomain invariant representation learning methods [9, 36, 30, 22, 31] require sharing/comparing therepresentation of the data across domains, while meta-learning based methods [21, 5] require accessto data from a pool of multiple source domains (in a central server).

Original: There have been some early attempts to incorporate domain generalization into the FL framework.
Revised: There have been some early attempts to incorporate domain generalization into the FL framework.

Original: For example, [27] proposes a technique tailored for medical image segmentation which allows the clients to share their data in the frequency space with each other.
Revised: For example, [27] proposes a technique tailored for medical image segmentation which allows the clients to share their data in the frequency space with each other.

Original: This can be viewed as a form of dataleakage and is undesirable in the FL setting we consider.
Revised: This can be viewed as a form of dataleakage and is undesirable in the FL setting we consider.

Original: Meanwhile, [44] aligns the representationdistribution across domains via a reference distribution (from a generative model).
Revised: Meanwhile, [42] aligns the representationdistribution across domains via a reference distribution (from a generative model).

Original: While this doesfully preserve the data privacy, it can be overly complicated to implement in practice, and theperformance is not yet on par with SOTA centralized methods.
Revised: While this doesfully preserve the data privacy, it can be overly complicated to implement in practice, and theperformance is not yet on par with SOTA centralized methods.

Original: Nevertheless, this is one of our mainbaselines in this paper.
Revised: Nevertheless, this is one of our mainbaselines in this paper.

Original: To tackle the above challenges, in this paper, we introduce a simple yet effective method for domaingeneralization in the FL setting.
Revised: To tackle the above challenges, in this paper, we introduce a simple yet effective method for domaingeneralization in the FL setting.

Original: We propose to learn a “simple” representation of the data; specifically,we employ two regularization methods, namely the L2 regularizer on the representation, and theconditional mutual information between the data and the representation given the label.
Revised: We propose to learn a “simple” representation of the data; specifically,we employ two regularization methods, namely the L2 regularizer on the representation, and theconditional mutual information between the data and the representation given the label.

Original: Theseregularizers are easy to implement and do not require the sharing of data between clients, thus fullyrespecting the principles of FL.
Revised: Theseregularizers are easy to implement and do not require the sharing of data between clients, thus fullyrespecting the principles of FL.

Original: In particular, they aim to restrict the amount of information therepresentation can contains, which we hope will help to ignore spurious correlations and lead to bettergeneralization to unseen target domains.
Revised: In particular, they aim to restrict the amount of information therepresentation can contains, which we hope will help to ignore spurious correlations and lead to bettergeneralization to unseen target domains.

Original: With a deeper look, we show that these regularizers act asimplicit forms of marginal and conditional representation alignment , which also further explainsthe performance gain of our method.
Revised: With a deeper look, we show that these regularizers act asimplicit forms of marginal and conditional representation alignment , which also further explainsthe performance gain of our method.

Original: This means that our method attempts to learn a representationwhose marginal distribution and conditional distribution (given the label) are invariant across thedomains without explicitly comparing those distributions between clients.
Revised: This means that our method attempts to learn a representationwhose marginal distribution and conditional distribution (given the label) are invariant across thedomains without explicitly comparing those distributions between clients.

Original: Our contributions in thiswork are as follows:
Revised: Our contributions in thiswork are as follows:

Original: • We propose a simple method, which aims to learn a simple representation of the data for theprediction task, for DG in the FL setting.
Revised: • We propose a simple method, which aims to learn a simple representation of the data for theprediction task, for DG in the FL setting.

Original: To this end, we revisit two common regularizationmethods, namely the conditional mutual information regularizer and the L2-norm regularizeron the representation, to restrict the amount of information it can contain.
Revised: To this end, we revisit two common regularizationmethods, namely the conditional mutual information regularizer and the L2-norm regularizeron the representation, to restrict the amount of information it can contain.

Original: • Furthermore, we show theoretically and empirically that these two regularizer are wellconnected with marginal and conditional alignment in the invariant representation learningframework.
Revised: • Furthermore, we show theoretically and empirically that these two regularizer are wellconnected with marginal and conditional alignment in the invariant representation learningframework.

Original: • We validate the effectiveness of our resulting model with a wide range of datasets and modelarchitectures and show that it significantly outperforms existing
Revised: • We validate the effectiveness of our resulting model with a wide range of datasets and modelarchitectures and show that it significantly outperforms existing

Original: FL methods on the domaingeneralization task and can even compete with state-of-the-art centralized methods.
Revised: FL methods on the domaingeneralization task and can even compete with state-of-the-art centralized methods.

Original:  Related Work
Revised:  Related Work

Original: 2.1 Federated Learning
Revised: 2.1 Federated Learning

Original: One of the first and most common FL algorithms is FedAVG [28], which uses several local SGDupdates within each communication round to the server to save communication cost.
Revised: One of the first and most common FL algorithms is FedAVG [28], which uses several local SGDupdates within each communication round to the server to save communication cost.

Original: In the contextof DG in the FL setting, this method is equivalent to Empirical Risk Minimization in a centralizedsetting, and is one of our main baselines.
Revised: In the contextof DG in the FL setting, this method is equivalent to Empirical Risk Minimization in a centralizedsetting, and is one of our main baselines.

Original: One of the main challenges that FL needs to tackle isthe statistical diversity among the clients (i.e., each client has a different local data distribution).
Revised: One of the main challenges that FL needs to tackle isthe statistical diversity among the clients (i.e., each client has a different local data distribution).

Original: Agreat amount of research [23, 47, 12, 24, 17, 36] has been done to tackle this problem (i.e., how totrain a shared model to deal with non-i.i.d.
Revised: Agreat amount of research [23, 44, 12, 24, 17, 34] has been done to tackle this problem (i.e., how totrain a shared model to deal with non-i.i.d.

Original: data of the source clients in the FL setting).
Revised: data of the source clients in the FL setting).

Original: Anotherline of approaches to deal with this problem is to allow each client to have a personalized modeland jointly train the models in a FL way (personalized FL) [13, 6, 4, 8, 16].
Revised: Anotherline of approaches to deal with this problem is to allow each client to have a personalized modeland jointly train the models in a FL way (personalized FL) [13, 6, 4, 8, 16].

Original: Although these abovemethods tackle the distribution shift across the source domains/clients, they still have not dealt withthe generalization ability to an unseen target domain – the problem we consider in this paper.
Revised: Although these abovemethods tackle the distribution shift across the source domains/clients, they still have not dealt withthe generalization ability to an unseen target domain – the problem we consider in this paper.

Original: 2.2 Domain Generalization
Revised: 2.2 Domain Generalization

Original: Most of the existing DG methods consider a centralized setting.
Revised: Most of the existing DG methods consider a centralized setting.

Original: A predominant and effective approachis to learn a domain-invariant representation [30, 25, 26, 42, 2, 15, 46, 1, 22, 38, 35] (meaning tolearn a representation whose marginal distribution and/or conditional distribution given the labelare unchanged across the domains), hoping that this representation generalizes better to an unseentarget domain.
Revised: A predominant and effective approachis to learn a domain-invariant representation [29, 25, 26, 40, 2, 15, 43, 1, 22, 36, 33] (meaning tolearn a representation whose marginal distribution and/or conditional distribution given the label areunchanged across the domains), hoping that this representation generalizes better to an unseen targetdomain.

Original: Our method can be viewed as implicit forms of representation alignment, thus moresuitable for the FL setting.
Revised: Our method can be viewed as implicit forms of representation alignment, thus more suitablefor the FL setting.

Original: Another line of approaches for DG is to use the idea of meta-learning[7, 5, 21]; with the logic being that if a model can adapt among the source domains well, it is morelikely to adapt to an unseen target domain.
Revised: Another line of approaches for DG is to use the idea of meta-learning [7, 5, 21];with the logic being that if a model can adapt among the source domains well, it is more likely toadapt to an unseen target domain.

Original: These methods typically require access to a pool of datafrom source domains (in a centralized server).
Revised: These methods typically require access to a pool of data fromsource domains (in a centralized server).

Original: [45] considers a single-domain domain generalizationsetting.

Original: However, it needs to generate a “fake” target distribution from a source distribution and thenmaximize the mutual information between the representations of the two domains (source vs fake target).

Original: This might be overly complicated when applied at each client level in a FL setting and henceless practical.

Original: [34] considers a federated domain adaptation, where each source domain aligns therepresentation distribution with the target domain by receiving the representations from this targetdistribution.

Original: This cannot be applied to our problem setting to align the representation of two clients,since one client cannot have access to another source client’s data (including the data representation)for privacy reasons.

Original:  Approach3.1 Preliminaries
Revised:  Approach3.1 Preliminaries

Original: We first present the problem of Federated Learning and the importance of Domain Generalization inthis setting.
Revised: We first present the problem of Federated Learning and the importance of Domain Generalization inthis setting.

Original: Federated Learning In the conventional FL setting, K clients collaborate with each other totrain a global model.
Revised: Federated Learning In the conventional FL setting, K clients collaborate with each other totrain a global model.

Original: Each clients i ( 1 ≤ i ≤ K ) has its own data distribution p
Revised: Each clients i ( 1 ≤ i ≤ K ) has its own data distribution p

Original: i ( x, y ) , wherex ∈ X is the input and
Revised: i ( x, y ) , wherex ∈ X is the input and

Original: y ∈ Y is its corresponding label, and a dataset with N i data points:
Revised: y ∈ Y is its corresponding label, and a dataset with N i data points:

Original: D i = { ( x (1) i , y (1) i ) , ..., ( x ( N i )
Revised: D i = { ( x (1) i , y (1) i ) , ..., ( x ( N i )

Original: [Equation]

Original:  , y ( N i ) i ) } .
Revised: i , y ( N i ) i ) } .

Original: It is often assumed that the data distribution p i ( x, y ) changesacross the clients (which we will use interchangeably with domains in this paper).
Revised: It is often assumed that the data distribution p i ( x, y ) changesacross the clients (which we will use interchangeably with domains in this paper).

Original: Here, the datadistributions p i ( x, y ) are sampled from a family E of distributions ( E can be finite or infinite), i.e.,p i ( x, y ) ∼ E .
Revised: Here, the datadistributions p i ( x, y ) are sampled from a family E of distributions ( E can be finite or infinite), i.e.,p i ( x, y ) ∼ E .

Original: Our aim is for the clients to jointly train a global model with parameters w (this modelwill be specified later), with the loss function of a datapoint ( x, y ) denoted as ℓ ( w ; x, y ) .
Revised: Our aim is for the clients to jointly train a global model with parameters w (this modelwill be specified later), with the loss function of a datapoint ( x, y ) denoted as ℓ ( w ; x, y ) .

Original: The global objective function to minimize is:
Revised: The global objective function to minimize is:

Original: [Equation]
Revised: [Equation]

Original: 111 where f i ( w ) is the local objective function,
Revised: 103 where f i ( w ) is the local objective function,

Original: [Equation]
Revised: [Equation]

Original: [Figure] [Equation]
Revised: 1 N i (cid:80) N i n =1 ℓ ( w ; x ( n ) i , y ( n ) i ) .

Original: Note that the objective in Eq.
Revised: Note that the objective in Eq.

Original: 1 can also be a weighted averageamong the clients.
Revised: 1 can also be a weighted averageamong the clients.

Original: However, since the analysis and discussion are unchanged, we will use this simpleaverage objective in this paper.
Revised: However, since the analysis and discussion are unchanged, we will use this simpleaverage objective in this paper.

Original: In most Federated Learning algorithms, the optimization process consists of a number of rounds.
Revised: In most Federated Learning algorithms, the optimization process consists of a number of rounds.

Original: Within each round, each (active) client runs its own local minimization of the local objective forseveral iterations of stochastic gradient descent; then the central server will collect the newly-learnedparameters from all (active) clients, average them to form the new global parameter, and broadcast itback to each client at the beginning of next round.
Revised: Within each round, each (active) client runs its own local minimization of the local objective forseveral iterations of stochastic gradient descent; then the central server will collect the newly-learnedparameters from all (active) clients, average them to form the new global parameter, and broadcast itback to each client at the beginning of next round.

Original: By taking multiple steps of the local minimizationeach round, the algorithm can reduce the communication cost (which is often the bottleneck).
Revised: By taking multiple steps of the local minimizationeach round, the algorithm can reduce the communication cost (which is often the bottleneck).

Original: For amore detailed description of these FL methods, we refer the readers to [28, 41, 37].
Revised: For amore detailed description of these FL methods, we refer the readers to [28, 39, 35].

Original: Domain Generalization In the DG problem, we are more interested in training a model to performwell in an unseen target domain p T ( x, y ) ∼ E (which is different from p
Revised: Domain Generalization In the DG problem, we are more interested in training a model to performwell in an unseen target domain p T ( x, y ) ∼ E (which is different from p

Original: i ( x, y ) ∀ i ).
Revised: i ( x, y ) ∀ i ).

Original: Therefore, wehope to minimize not only the expected loss on the source domains/clients E p
Revised: Therefore, wehope to minimize not only the expected loss on the source domains/clients E p

Original: i ( x,y )
Revised: i ( x,y )

Original: [ ℓ ( w ; x, y )]
Revised: [ ℓ ( w ; x, y )]

Original: butalso the loss on unseen target domains, in either an average or worst-case sense as defined below.
Revised: butalso the loss on unseen target domains, in either an average or worst-case sense as defined below.

Original: In the average case, we want to minimize this loss over the distribution family E :
Revised: In the average case, we want to minimize this loss over the distribution family E :

Original: [Equation]
Revised: [Equation]

Original: Meanwhile, we might also be interested in the worst-case quantity, i.e., the loss of the domain inwhich the model performs the worst:
Revised: Meanwhile, we might also be interested in the worst-case quantity, i.e., the loss of the domain inwhich the model performs the worst:

Original: [Equation]
Revised: [Equation]

Original: In a conventional machine learning problem, we usually need to observe more datapoints to learnto generalize better to unseen datapoints .
Revised: In a conventional machine learning problem, we usually need to observe more datapoints to learnto generalize better to unseen datapoints .

Original: Similarly, in this Domain Generalization problem,we additionally need to observe more domains in order to generalize better to unseen targetdomains .
Revised: Similarly, in this Domain Generalization problem,we additionally need to observe more domains in order to generalize better to unseen targetdomains .

Original: However, we typically have a limited number of source domains in these DG problems, soregularization techniques often need to be applied for a better generalization.
Revised: However, we typically have a limited number of source domains in these DG problems, soregularization techniques often need to be applied for a better generalization.

Original: One common techniqueamong those is to learn a domain invariant representation, which means to learn a representation zof x that has the same marginal and/or conditional distribution across the domains.
Revised: One common techniqueamong those is to learn a domain invariant representation, which means to learn a representation zof x that has the same marginal and/or conditional distribution across the domains.

Original: These methodsoften help to ignore domain-specific spurious correlations and lead to better generalization.
Revised: These methodsoften help to ignore domain-specific spurious correlations and lead to better generalization.

Original: We willdiscuss this in more detail later in the paper.
Revised: We willdiscuss this in more detail later in the paper.

Original: Generalization Challenge of FL From the objective in Eq.
Revised: Generalization Challenge of FL From the objective in Eq.

Original: 1, it is clear that conventional FLmethods only focus on maximizing the performance (minimizing the loss) on K source clients’ data,which is very problematic in many applications.
Revised: 1, it is clear that conventional FLmethods only focus on maximizing the performance (minimizing the loss) on K source clients’ data,which is very problematic in many applications.

Original: For example, if K clinical institutes in the USand UK collaborate to train a Covid prediction model, the goal is not only to perform well on theirown data distributions but also to generalize well to data from other countries.
Revised: For example, if K clinical institutes in the US and UK collaborate to train a Covid prediction model, the goal is not only to perform well on theirown data distributions but also to generalize well to data from other countries.

Original: Similarly, if we use
Revised: Similarly, if we use

Original: Federated Learning to train a machine learning model with data collected from multiple mobile phoneusers, it is crucial that the model can generalize to a new user.
Revised: Federated Learning to train a machine learning model with data collected from multiple mobile phoneusers, it is crucial that the model can generalize to a new user.

Original: Thus, domain generalization is veryimportant for Federated Learning.
Revised: Thus, domain generalization is veryimportant for Federated Learning.

Original: However, as mentioned earlier, applying existing domain generalization techniques to Federated
Revised: However, as mentioned earlier, applying existing domain generalization techniques to Federated

Original: Learning is not straightforward .
Revised: Learning is not straightforward .

Original: This is because most existing DG methods need a centralizedtraining setting with data (or a part of the data, such as a representation) being shared across thedomains.
Revised: This is because most existing DG methods need a centralizedtraining setting with data (or a part of the data, such as a representation) being shared across thedomains.

Original: For example, the aforementioned domain invariant representation learning approachtypically requires sharing and comparing the representation distribution across domains.
Revised: For example, the aforementioned domain invariant representation learning approachtypically requires sharing and comparing the representation distribution across domains.

Original: Towardsthis end, we propose a simple but effective domain invariant representation learning technique for
Revised: Towardsthis end, we propose a simple but effective domain invariant representation learning technique for

Original: Federated Learning that fully keeps its decentralized and privacy aspects.
Revised: Federated Learning that fully keeps its decentralized and privacy aspects.

Original: We first introduce theproblem setting of representation learning in Subsection 3.2, and then present our domain invariantrepresentation learning methods in Subsection 3.3.
Revised: We first introduce theproblem setting of representation learning in Subsection 3.2, and then present our domain invariantrepresentation learning methods in Subsection 3.3.

Original: 3.2 Problem Setting: Representation Learning for Domain Generalization
Revised: 3.2 Problem Setting: Representation Learning for Domain Generalization

Original: In the representation learning framework, we aim to learn a representation z of x with a distributionp ( z | x ) parameterized by w 1 , which we omit here for notation simplicity.
Revised: In the representation learning framework, we aim to learn a representation z of x with a distributionp ( z | x ) parameterized by w 1 , which we omit here for notation simplicity.

Original: This can be a deterministicmapping (i.e., p ( z | x ) =
Revised: This can be a deterministicmapping (i.e., p ( z | x ) =

Original: δ g w 1 ( x ) ( z ) where δ is the dirac delta function) or a probabilistic mapping(e.g., p ( z | x ) =
Revised: δ g w 1 ( x ) ( z ) where δ is the dirac delta function) or a probabilistic mapping(e.g., p ( z | x ) =

Original: N ( z ; µ w 1 ( x ) , σ 2 w 1 ( x )) where N is the normal distribution).
Revised: N ( z ; µ w 1 ( x ) , σ 2 w 1 ( x )) where N is the normal distribution).

Original: With the data distributionp i ( x, y ) of client i , its joint distribution of x, y, z is:
Revised: With the data distributionp i ( x, y ) of client i , its joint distribution of x, y, z is:

Original: [Equation]
Revised: [Equation]

Original: p ( z | x ) , since z is conditionally independent of y given x (4) 162 From the representation z , we learn a classifier/regressor that predicts y given z with the predictivedistribution ˆ p ( y | z ) parameterized by w 2 (again, we omit w 2 here for notation simplicity).
Revised: 154 From the representation z , we learn a classifier/regressor that predicts y given z with the predictivedistribution ˆ p ( y | z ) parameterized by w 2 (again, we omit w 2 here for notation simplicity).

Original: Forexample, with a classification problem, this is often just a linear layer followed by a softmax layer toform the predictive distribution.
Revised: Forexample, with a classification problem, this is often just a linear layer followed by a softmax layer toform the predictive distribution.

Original: The predictive distribution of y given x of our model is:
Revised: The predictive distribution of y given x of our model is:

Original: [Equation]
Revised: [Equation]

Original: 5) 168 Note that when p ( z | x ) is a deterministic mapping, Eq.
Revised: 160 Note that when p ( z | x ) is a deterministic mapping, Eq.

Original: 5 simplifies into ˆ p ( y | z = g w 1 ( x )) .
Revised: 5 simplifies into ˆ p ( y | z = g w 1 ( x )) .

Original: For both regression and classification, the loss function of a data point x, y is often the negative logpredictive, i.e, ℓ ( w ; x, y )
Revised: For both regression and classification, the loss function of a data point x, y is often the negative logpredictive, i.e, ℓ ( w ; x, y )

Original: = − log E p ( z | x ) [ˆ p ( y | z )] , and in this case w = { w 1 , w 2 } .
Revised: = − log E p ( z | x ) [ˆ p ( y | z )] , and in this case w = { w 1 , w 2 } .

Original: Hence, the local objective of client i is:
Revised: Hence, the local objective of client i is:

Original: [Equation]
Revised: [Equation]

Original: During training, for a probabilistic mapping, we often sample a single z per x from p ( z | x ) (for adeterministic mapping, it is obvious that we only sample one z ), leading to the training objective: 
Revised: During training, for a probabilistic mapping, we often sample a single z per x from p ( z | x ) (for adeterministic mapping, it is obvious that we only sample one z ), leading to the training objective: [Equation]

Original: [Equation]
Revised: [Equation]

Original:  z ( n ) i is a single sample 
Revised: i ) , where z ( n ) i is a single sample from p

Original: [Equation]
Revised: [Equation]

Revised: z | x

Original: [Equation]
Revised: [Equation]

Original: 175 The quantity − log ˆ p ( y | z ) is often non-negative for all common predictive distributions.
Revised: 167 The quantity − log ˆ p ( y | z ) is often non-negative for all common predictive distributions.

Original: For example,with a categorical predictive distribution in a classification problem, this is the cross-entropy loss;while with a Gaussian predictive distribution (with a fixed variance) in a regression problem, thisbecomes the squared error (with an additive constant).
Revised: For example,with a categorical predictive distribution in a classification problem, this is the cross-entropy loss;while with a Gaussian predictive distribution (with a fixed variance) in a regression problem, thisbecomes the squared error (with an additive constant).

Original: Note that according to Jensen’s Inequality, f i ( w ) is an upper bound of f i ( w ) , so the training processis sound.
Revised: Note that according to Jensen’s Inequality, f i ( w ) is an upper bound of f i ( w ) , so the training processis sound.

Original: The representation learning framework that we discussed so far still has not dealt with the problem ofdomain generalization.
Revised: The representation learning framework that we discussed so far still has not dealt with the problem ofdomain generalization.

Original: Many existing works in DG focus on learning a domain-invariant representation z , which means to learn a representation z which has the same marginal ( p i ( z ) ) and/or conditional ( p i ( z | y ) or p i ( y | z ) ) distributions across the domains.
Revised: Many existing works in DG focus on learning a domain-invariant representation z , which means to learn a representation z which has the same marginal ( p i ( z ) ) and/or conditional( p i ( z | y ) or p i ( y | z ) ) distributions across the domains.

Original: However, this often requires sharing informationof the data (the representation) across domains/clients, which is not allowed in Federated Learning.
Revised: However, this often requires sharing informationof the data (the representation) across domains/clients, which is not allowed in Federated Learning.

Original: Motivated by the fact that classical machine learning algorithms make predictions based on a simplerepresentation of the data and often generalize better than modern deep learning models, in this paper,we aim to learn a S imple R epresentation of the data (hence the name FedSR).
Revised: Motivated by the fact that classical machine learning algorithms make predictions based on a simplerepresentation of the data and often generalize better than modern deep learning models, in this paper,we aim to learn a S imple R epresentation of the data (hence the name FedSR).

Original: We apply severalwell-known regularization techniques to restrict the representation’s complexity and encourage thedeep network p ( z | x ) to only learn the essential information.
Revised: We apply severalwell-known regularization techniques to restrict the representation’s complexity and encourage thedeep network p ( z | x ) to only learn the essential information.

Original: Furthermore, we show that theseregularization techniques have strong theoretical connections to domain-invariant representationlearning and that they can achieve good generalization performance in practice.
Revised: Furthermore, we show that theseregularization techniques have strong theoretical connections to domain-invariant representationlearning and that they can achieve good generalization performance in practice.

Original: 3.3 FedSR
Revised: 3.3 FedSR

Original: As mentioned earlier, we employ an L2-norm regularizer ( ℓ L 2 R i ) on the representation and (an upperbound of) the conditional mutual information between the data and the representation given the label( ℓ CMIi ) to restrict the amount of information the representation can contain.
Revised: As mentioned earlier, we employ an L2-norm regularizer ( ℓ L 2 R i ) on the representation and (an upperbound of) the conditional mutual information between the data and the representation given the label( ℓ CMIi ) to restrict the amount of information the representation can contain.

Original: Moreover, ℓ L 2 R
Revised: Moreover, ℓ L 2 R

Original: i also hasthe effect to align the marginal distribution p
Revised: i also hasthe effect to align the marginal distribution p

Original: i ( z ) to be centered around 0 as a Gaussian distribution,while ℓ CMIi aligns
Revised: i ( z ) to be centered around 0 as a Gaussian distribution,while ℓ CMIi aligns

Original: the conditional distribution p i ( z | y ) to a reference distribution (also chosen to be
Revised: the conditional distribution p i ( z | y ) to a reference distribution (also chosen to be

Original: Gaussian in the experiment section).
Revised: Gaussian in the experiment section).

Original: These two terms will be explained in detail below.
Revised: These two terms will be explained in detail below.

Original: The finallocal objective function of each client i is:
Revised: The finallocal objective function of each client i is:

Original: [Equation]
Revised: [Equation]

Original: where α L 2 R and α CMI are hyper-parameters.
Revised: where α L 2 R and α CMI are hyper-parameters.

Original: (We will also test the two regularizers separately,namely the FedL2R and FedCMI variants.)L2-
Revised: (We will also test the two regularizers separately,namely the FedL2R and FedCMI variants.)

Original: norm Regularizer

Original: In optimization and learning theory, L2-norm of the weights/parameters is often regarded to as ameasure of the model’s complexity.

Original: Here, with a similar idea, we enforce an L2-norm Regularizer(L2R) on the representation (not the network parameters) to restrict the complexity of the representation.

Original: (Note that some methods like [29] use the L2-norm but in the parameter space to help with thegeneralization of Federated Learning).

Original: Specifically, the regularizer has the form:

Original: [Equation]

Original: |

Original: [Equation]

Original: z

Original: [Equation]

Original: n ) i || 22 , where z ( n ) i is a single sample from p ( z | x ( n ) i )

Original: [Equation]

Original: Intuitively, this regularizer will have the effect of encouraging z to be centered around 0 for all clients,thus helping with the marginal alignment of the representation.

Original: This connection can be explained inmore detail as below:

Original: Connection to Domain Invariant Representation Learning: First of all, rewrite ℓ L 2 R i as:

Original: [Equation]

Original: Now, consider a “reference” distribution q ( z ) =

Original: N (0 , σ 2 I ) (with a small σ ).

Original: Then we have:

Original: [Equation]

Original: an additive

Original: [Equation]

Original: (13) 215 And hence:

Original: [Equation]

Original: 14) where H ( p i ( z ) , q ( z )) denotes the cross entropy from q ( z ) to p i ( z ) .

Original: Note further that:

Original: [Equation]

Original: Hence, if the entropy H ( p i ( z )) does not change significantly during training, minimizing ℓ L 2 R i will 219 also minimize KL [ p i ( z ) | q ( z )] , which encourages p i ( z ) to be close to q ( z ) , i.e., an implicit alignmentof the marginal distribution.

Original: Nevertheless, also according to the above equation, if we really want to minimize KL [ p i ( z ) | q ( z )] (forthe marginal alignment), it is perhaps more rigorous to minimize the regularizer ℓ L 2 R i − 2 σ 2 H ( p i ( z ))instead of ℓ L 2 R i .

Original: However, when σ is small enough, this term is dominated by ℓ L 2 R

Original: i (meaningthat if σ is small, minimizing ℓ L 2 R i almost guarantees to minimize KL [ p i ( z ) | q ( z )] ).

Original: Empirically,we found that the two variants achieve almost identical performance; furthermore, even withoutregularizing − 2 σ 2 H ( p i ( z )) , the entropy term H ( p i ( z )) does not change significantly during training.

Original: We therefore recommend using ℓ L 2 R i instead of ℓ L 2 R i − 2 σ 2 H ( p i ( z )) for its simplicity.

Original: 3.3.2 Conditional Mutual Information regularizer
Revised: 3.3.1 Conditional Mutual Information regularizer

Original: The mutual information term between the input x and the representation z (denoted I ( x, z ) ) is oftenused to regularize the amount of information z can contains [3, 39].
Revised: The mutual information term between the input x and the representation z (denoted I ( x, z ) ) is oftenused to regularize the amount of information z can contains [3, 37].

Original: However, this regularizer mightbe too restrictive in practice if its coefficient is not tuned properly, as it encourages the representationto contain no information about the input data.
Revised: However, this regularizer mightbe too restrictive in practice if its coefficient is not tuned properly, as it encourages the representationto contain no information about the input data.

Original: In this paper, we use the conditional mutualinformation of x and z given y , as it is less restricted than the former.
Revised: In this paper, we use the conditional mutualinformation of x and z given y , as it is less restricted than the former.

Original: We can also later see that thisis well-connected with conditional distribution alignment in DG.
Revised: We can also later see that thisis well-connected with conditional distribution alignment in DG.

Original: Recall that this conditional mutualinformation term is calculated for domain/client i as:
Revised: Recall that this conditional mutualinformation term is calculated for domain/client i as:

Original: [Equation]
Revised: [Equation]

Original: [Equation] Intuitively, the terms f i and I i ( x, z | y ) work together to enforce the representation z to contain only theinformation needed to predict the label y , and no additional information (non-label-related) about x .
Revised: Intuitively, the terms f i and I i ( x, z | y ) work together to enforce the representation z to contain only theinformation needed to predict the label y , and no additional information (non-label-related) about x .

Original: However, unfortunately, this mutual information term is not tractable because of p i ( z | y )
Revised: However, unfortunately, this mutual information term is not tractable because of p i ( z | y )

Original: (which ishard to compute due to the integration over x ).
Revised: (which ishard to compute due to the integration over x ).

Original: Therefore, we derive an upper bound in order tominimize it.
Revised: Therefore, we derive an upper bound in order tominimize it.

Original: Specifically:
Revised: Specifically:

Original: Lemma 1.
Revised: Lemma 1.

Original: Let r ( z | y ) be a(ny) conditional distribution of z given y ∀ y ∈ Y .
Revised: Let r ( z | y ) be a(ny) conditional distribution of z given y ∀ y ∈ Y .

Original: If p ( z | x ) and r ( z | y )have the same support set ∀ x ∈ X , y ∈ Y , we have:
Revised: If p ( z | x ) and r ( z | y )have the same support set ∀ x ∈ X , y ∈ Y , we have:

Original: [Equation]
Revised: [Equation]

Original: Proof.
Revised: Proof.

Original: Provided in the Supplementary Material.
Revised: Provided in the Supplementary Material.

Original: Here, the upper bound ℓ CMIi can be computed and used as a regularizer when training the representation network.
Revised: Here, the upper bound ℓ CMIi can be computed and used as a regularizer when training the representation network.

Original: We optimize both p ( z | x ) and r ( z | y ) in order to minimize ℓ CMIi .
Revised: We optimize both p ( z | x ) and r ( z | y ) in order to minimize ℓ CMIi .

Original: Note that r ( z | y )can be a network that takes y as an input and output the distribution r ( z | y ) ; or more simply forclassification (when we have a finite number of labels y ), we can set r ( z | y )
Revised: Note that r ( z | y )can be a network that takes y as an input and output the distribution r ( z | y ) ; or more simply forclassification (when we have a finite number of labels y ), we can set r ( z | y )

Original: =
Revised: =

Original: N ( z ; µ y , σ 2 y ) , whereµ y , σ 2 y ( y = 1 ..C ) are the parameters to be optimized, with C being the number of classes.
Revised: N ( z ; µ y , σ 2 y ) , whereµ y , σ 2 y ( y = 1 ..C ) are the parameters to be optimized, with C being the number of classes.

Original: Note also that if we use this (upper bound) regularizer ℓ CMIi , it requires the representation mappingp ( z | x ) to be probabilistic (due to the KL term).
Revised: Note also that if we use this (upper bound) regularizer ℓ CMIi , it requires the representation mappingp ( z | x ) to be probabilistic (due to the KL term).

Original: Hence, in the experimental section, we use aprobabilistic representation network ( p ( z | x ) =
Revised: Hence, in the experimental section, we use aprobabilistic representation network ( p ( z | x ) =

Original: N ( z ; µ w 1 ( x ) , σ 2 w 1 ( x )) ) for any models/variants withthe ℓ CMI i objective, and use a deterministic network otherwise.
Revised: N ( z ; µ w 1 ( x ) , σ 2 w 1 ( x )) ) for any models/variants withthe ℓ CMIi objective, and use a deterministic network otherwise.

Original: With the empirical dataset D i ofclient i , we have an empirical estimator of ℓ CMIi as below:
Revised: With the empirical dataset D i ofclient i , we have an empirical estimator of ℓ CMIi as below:

Original: [Equation]
Revised: [Equation]

Original: Connection to Domain Invariant Representation Learning:
Revised: Connection to Domain Invariant Representation Learning:

Original: We found out that this regularizerhas a theoretical connection with conditional distribution alignment in domain generalization.
Revised: We found out that this regularizerhas a theoretical connection with conditional distribution alignment in domain generalization.

Original: Westart with the following lemma:
Revised: Westart with the following lemma:

Original: Lemma 2.
Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: Proof.
Revised: Proof.

Original: Provided in the Supplementary Material.
Revised: Provided in the Supplementary Material.

Original: Therefore, by minimizing ℓ CMIi , we also minimize the divergence KL [ p i ( z | y ) | r ( z | y )] ∀ y , forcingp i ( z | y ) and r ( z | y ) to be close to each other:
Revised: Therefore, by minimizing ℓ CMIi , we also minimize the divergence KL [ p i ( z | y ) | r ( z | y )] ∀ y , forcingp i ( z | y ) and r ( z | y ) to be close to each other:

Original: p i ( z | y )
Revised: p i ( z | y )

Original: ≈
Revised: ≈

Original: r
Revised: r

Original: ( z | y ) .
Revised: ( z | y ) .

Original: Therefore, the model will try toenforce p
Revised: Therefore, the model will try toenforce p

Original: i ( z | y )
Revised: i ( z | y )

Original: ≈ p j ( z | y )
Revised: ≈ p j ( z | y )

Original: ( ≈ r ( z | y ))
Revised: ( ≈ r ( z | y ))

Original: ∀ clients/domains i, j , i.e., we are doing implicit alignment ofthe conditional distribution of the representation z given the label y , which is a common and effectivetechnique in conventional Domain Generalization.
Revised: ∀ clients/domains i, j , i.e., we are doing implicit alignment ofthe conditional distribution of the representation z given the label y , which is a common and effectivetechnique in conventional Domain Generalization.

Original: For example, [26, 25, 46] explicitly align thisconditional distribution (not suitable for Federated Learning).
Revised: For example, [26, 25, 43] explicitly align thisconditional distribution (not suitable for Federated Learning).

Original: 3.3.3 Implementation Choices and Optimization
Revised: 3.3.2 L2-norm Regularizer

Revised: In optimization and learning theory, L2-norm of the weights/parameters is often regarded to as a measure of the model’s complexity. Here, with a similar idea, we enforce an L2-norm Regularizer (L2R)on the representation (not the network parameters) to restrict the complexity of the representation. Specifically, the regularizer has the form:

Revised: [Equation]

Revised: |

Revised: [Equation]

Revised: z

Revised: [Equation]

Revised: i || 22 , where z ( n ) i is a single sample from p ( z | x ( n ) i )

Revised: [Equation]

Revised: [Equation]

Revised: Intuitively, this regularizer will have the effect of encouraging z to be centered around 0 for all clients,thus helping with the marginal alignment of the representation.

Revised: This connection can be explained inmore detail as below:

Revised: Connection to Domain Invariant Representation Learning: First of all, rewrite ℓ L 2 R i as:

Revised: [Equation]

Revised: Now, consider a “reference” distribution q ( z ) =

Revised: N (0 , σ 2 I ) (with a small σ ).

Revised: Then we have:

Revised: [Equation]

Revised: with an additive constant) (17) 242 And hence:

Revised: [Equation]

Revised: [Equation]

Revised: where H ( p i ( z ) , q ( z )) denotes the cross entropy from q ( z ) to p i ( z ) .

Revised: Note further that:

Revised: [Equation]

Revised: Hence, if the entropy H ( p i ( z )) does not change significantly during training, minimizing ℓ L 2 R i willalso minimize KL [ p i ( z ) | q ( z )] , which encourages p i ( z ) to be close to q ( z ) , i.e., an implicit alignmentof the marginal distribution.

Revised: Nevertheless, also according to the above equation, if we really want to minimize KL [ p i ( z ) | q ( z )] (forthe marginal alignment), it is perhaps more rigorous to minimize the regularizer ℓ L 2 R i − 2 σ 2 H ( p i ( z ))instead of ℓ L 2 R i .

Revised: However, when σ is small enough, this term is dominated by ℓ L 2 R

Revised: i (meaningthat if σ is small, minimizing ℓ L 2 R i almost guarantees to minimize KL [ p i ( z ) | q ( z )] ).

Revised: Empirically,we found that the two variants achieve almost identical performance; furthermore, even withoutregularizing − 2 σ 2 H ( p i ( z )) , the entropy term H ( p i ( z )) does not change significantly during training.

Revised: We therefore recommend using ℓ L 2 R i instead of ℓ L 2 R i − 2 σ 2 H ( p i ( z )) for its simplicity.

Revised: 3.3.3 Implementation Choices and Optimization

Original: For the federated optimization, we use FedAVG, which is one of the first and perhaps most common
Revised: For the federated optimization, we use FedAVG, which is one of the first and perhaps most common

Original: FL optimization algorithms (although we can use any federated optimization techniques).
Revised: FL optimization algorithms (although we can use any federated optimization techniques).

Original: Following common practice, we use mini-batches to approximate Eq. 8, Eq. 18, and Eq.
Revised: Following common practice, we use mini-batches to approximate Eq. 8, Eq. 12, and Eq.

Original: 11 (insteadof full-batch training).
Revised: 15 (insteadof full-batch training).

Original: For a probabilistic network (when using ℓ CMI ), to allow the gradient to backpropagate throughthe samples z ( n ) i to the network parameters w 1 , we use the reparameterization trick [18].
Revised: For a probabilistic network (when using ℓ CMI ), to allow the gradient to backpropagate throughthe samples z ( n ) i to the network parameters w 1 , we use the reparameterization trick [18].

Original: With thisreparameterization trick, f i and ℓ L 2 R i can be implemented and back-propagated straightforwardly.
Revised: With thisreparameterization trick, f i and ℓ L 2 R i can be implemented and back-propagated straightforwardly.

Original: Regarding the conditional mutual information regularizer ℓ CMIi , we use a Gaussian distribution for thevariational conditional distribution r ( z | y ) , hence the KL term in Eq.
Revised: Regarding the conditional mutual information regularizer ℓ CMIi , we use a Gaussian distribution for thevariational conditional distribution r ( z | y ) , hence the KL term in Eq.

Original: 18 can be computed analytically(recall that we also use a Gaussian distribution for p ( z | x ) ).
Revised: 12 can be computed analytically(recall that we also use a Gaussian distribution for p ( z | x ) ).

Original: Here the astute readers might wonderthat such a simple choice of the variational distribution r ( z | y ) might not be able to approximatep i ( z | y ) .
Revised: Here the astute readers might wonderthat such a simple choice of the variational distribution r ( z | y ) might not be able to approximatep i ( z | y ) .

Original: However, note that since both p ( z | x ) and r ( z | y ) are optimized in order to minimize ℓ CMIi ,it encourages the network to learn a simple representation so that the variational distribution r ( z | y )can match p i ( z | y ) (which arguably might be a good regularization effect).
Revised: However, note that since both p ( z | x ) and r ( z | y ) are optimized in order to minimize ℓ CMIi , itencourages the network to learn a simple representation so that the variational distribution can match(which arguably might be a good regularization effect).

Original: Furthermore, dependingon datasets, we can, in theory, use more complex distributions such as a Gaussian mixture as thevariational distribution.
Revised: Furthermore, depending on datasets, we can,in theory, use more complex distributions such as a Gaussian mixture as the variational distribution.

Original: In such cases, one can use methods such as MC sampling to estimate the KLterm.
Revised: In such cases, one can use methods such as MC sampling to estimate the KL term.

Original: In our experiment section, for simplicity, we only use a Gaussian distribution for r ( z | y ) .
Revised: In our experimentsection, for simplicity, we only use a Gaussian distribution for r ( z | y ) .

Original:  Experiments
Revised:  Experiments

Original: 4.1 Datasets
Revised: 4.1 Datasets

Original: To evaluate our method, we perform experiments in four datasets (ranging from easy to morechallenging) that are commonly used in the literature for domain generalization.
Revised: To evaluate our method, we perform experiments in four datasets (ranging from easy to morechallenging) that are commonly used in the literature for domain generalization.

Original: RotatedMNIST [10]: In this dataset, the MNIST images [19] are rotated counter-clockwise with anangle of 0 ◦ , 15 ◦ , 30 ◦ , 45 ◦ , 60 ◦ and 75 ◦ to form six domains M 0 , M 15 , M 30 , M 45 , M 60 and M 75 .
Revised: RotatedMNIST [10]: In this dataset, the MNIST images [19] are rotated counter-clockwise with anangle of 0 ◦ , 15 ◦ , 30 ◦ , 45 ◦ , 60 ◦ and 75 ◦ to form six domains M 0 , M 15 , M 30 , M 45 , M 60 and M 75 .

Original: The task is classification with ten classes (digits 0 to 9).
Revised: The task is classification with ten classes (digits 0 to 9).

Original: We use the same version of the dataset as in[31, 15] (where only 1000 images are rotated to form the domain) for easy comparison.
Revised: We use the same version of the dataset as in[30, 15] (where only 1000 images are rotated to form the domain) for easy comparison.

Original: Note thatthis is slightly different from the version in [11] and thus the numbers are not comparable.
Revised: Note thatthis is slightly different from the version in [11] and thus the numbers are not comparable.

Original: PACS [20]: contains 9,991 images from four different domains: art painting, cartoon, photo, sketch.
Revised: PACS [20]: contains 9,991 images from four different domains: art painting, cartoon, photo, sketch.

Original: The task is classification with seven classes.
Revised: The task is classification with seven classes.

Original: OfficeHome [40]: has 15,500 images of daily objects from four domains: art, clipart, product andreal.
Revised: OfficeHome [38]: has 15,500 images of daily objects from four domains: art, clipart, product andreal.

Original: There are 65 classes in this classification dataset.
Revised: There are 65 classes in this classification dataset.

Original: DomainNet [33] is a large-scale dataset, consisting 586,575 images from 345 classes.
Revised: DomainNet [32] is a large-scale dataset, consisting 586,575 images from 345 classes.

Original: These imagesare from 6 domains: clipart, infograph, painting, quickdraw, real, sketch.
Revised: These imagesare from 6 domains: clipart, infograph, painting, quickdraw, real, sketch.

Original: 4.2 Experimental Setting
Revised: 4.2 Experimental Setting

Original: For all datasets, we perform “leave-one-domain-out” experiments, where we choose one domain asthe target domain, train the model on all remaining domains, and evaluate it on the chosen domain.
Revised: For all datasets, we perform “leave-one-domain-out” experiments, where we choose one domain asthe target domain, train the model on all remaining domains, and evaluate it on the chosen domain.

Original: Each source domain is treated as a client.
Revised: Each source domain is treated as a client.

Original: Following standard practice, we use 90% of available dataas training data and 10% as validation data.
Revised: Following standard practice, we use 90% of available dataas training data and 10% as validation data.

Original: We train all our models with NVIDIA A100 GPUs fromour AWS cluster.
Revised: We train all our models with NVIDIA A100 GPUs fromour AWS cluster.

Original: For the RotatedMNIST dataset, we use a network of two 3x3 convolutional layers and a fullyconnected layer as the representation network g θ to get a representation z of 64 dimensions. 
Revised: For the RotatedMNIST dataset, we use a network of two 3x3 convolutional layers and a fullyconnected layer as the representation network g θ to get a representation z of 64 dimensions. A single [Table]

Original: A singlelinear layer is then used to map the representation z to the ten output classes.
Revised: [Table] linear layer is then used to map the representation z to the ten output classes.

Original: This architecture is thesame as the network used by [15, 31].
Revised: This architecture is thesame as the network used by [15, 30].

Original: We train our network for 500 epochs with stochastic gradient [Table] [Table] descent (SGD), using a learning rate of 0.001 and minibatch size 64, and report performance on thetest domain after the last epoch.
Revised: We train our network for 500 epochs with stochastic gradientdescent (SGD), using a learning rate of 0.001 and minibatch size 64, and report performance on thetest domain after the last epoch.

Original: Each client performs 5 local optimization iterations within eachcommunication round ( E = 5 ).
Revised: Each client performs 5 local optimization iterations within eachcommunication round ( E = 5 ).

Original: For the PACS datasets, for easy comparison with existing centralized domain invariant representationlearning methods, we use the most common choice of backbone network in existing works as therepresentation networks, i.e., Resnet18 [14].
Revised: For the PACS datasets, for easy comparison with existing centralized domain invariant representationlearning methods, we use the most common choice of backbone network in existing works as therepresentation networks, i.e., Resnet18 [14].

Original: We use a ResNet50 [14] backbone for OfficeHome and
Revised: We use a ResNet50 [14] backbone for OfficeHome and

Original: DomainNet since they are more complex datasets.
Revised: DomainNet since they are more complex datasets.

Original: We replace the last fully connected layer of thebackbone with a linear layer of dimension 512 for ResNet18 and 2048 for ResNet50 to form therepresentation network.
Revised: We replace the last fully connected layer of thebackbone with a linear layer of dimension 512 for ResNet18 and 2048 for ResNet50 to form therepresentation network.

Original: As with the RotatedMNIST experiment, we use a single layer to map fromthe representation z to the output.
Revised: As with the RotatedMNIST experiment, we use a single layer to map fromthe representation z to the output.

Original: Each local client uses stochastic gradient descent (SGD) (a total of5000 iterations) with learning rate 0.01, momentum 0.9, minibatch size 64, and weight decay 5 e − 4 .
Revised: Each local client uses stochastic gradient descent (SGD) (a total of5000 iterations) with learning rate 0.01, momentum 0.9, minibatch size 64, and weight decay 5 e − 4 .

Original: Similar to the RotatedMNIST experiments, each client performs 5 local optimization iterations withineach communication round ( E = 5 ).
Revised: Similar to the RotatedMNIST experiments, each client performs 5 local optimization iterations withineach communication round ( E = 5 ).

Original: For the hyper-parameters of our method ( β CMI and β L 2 R ),we follow [11] and use random search to tune their value based only on the validation set (no dataleakage from the target domain).
Revised: For the hyper-parameters of our method ( β CMI and β L 2 R ),we follow [11] and use random search to tune their value based only on the validation set (no dataleakage from the target domain).

Original: This protocol is highly recommended by [11] and allows for a faircomparison among the methods.
Revised: This protocol is highly recommended by [11] and allows for a faircomparison among the methods.

Original: For details of the tuned values, please refer to our supplementarymaterial.
Revised: For details of the tuned values, please refer to our supplementarymaterial.

Original: Data augmentation is also standard practice for real-world computer vision datasets like
Revised: Data augmentation is also standard practice for real-world computer vision datasets like

Original: PACS, OfficeHome and Domainet, and during the training we augment our data as follows: crops ofrandom size and aspect ratio, resizing to 224 × 224 pixels, random horizontal flips, random colorjitter, randomly converting the image tile to grayscale with 10% probability, and normalization usingthe ImageNet channel means and standard deviations.
Revised: PACS, OfficeHome and Domainet, and during the training we augment our data as follows: crops ofrandom size and aspect ratio, resizing to 224 × 224 pixels, random horizontal flips, random colorjitter, randomly converting the image tile to grayscale with 10% probability, and normalization usingthe ImageNet channel means and standard deviations.

Original: 4.3 Baselines
Revised: 4.3 Baselines

Original: For all experiments, we consider the FL algorithms to be our main baselines: FedAVG [28] (acommon FL algorithm) and FedADG [44] (an existing FL algorithm with DG ability).
Revised: For all experiments, we consider the FL algorithms to be our main baselines: FedAVG [28] (acommon FL algorithm) and FedADG [42] (an existing FL algorithm with DG ability).

Original: For the
Revised: For the

Original: RotatedMNIST and PACS experiments, we also include SOTA centralized domain-invariant methodsfor comparison (the numbers are taken directly from the original papers).
Revised: RotatedMNIST and PACS experiments, we also include SOTA centralized domain-invariant methodsfor comparison (the numbers are taken directly from the original papers).

Original: For the OfficeHomeand DomainNet experiments, since our setting (e.g., backbone network) and evaluation protocolare identical to those of DomainBed [11], our results can be compared directly to the results ofthe centralized methods recorded in that paper.
Revised: For the OfficeHomeand DomainNet experiments, since our setting (e.g., backbone network) and evaluation protocol [Table] [Table] are identical to those of DomainBed [11], our results can be compared directly to the results ofthe centralized methods recorded in that paper.

Original: Therefore, we also include the best performers foreach dataset reported in the DomainBed paper.
Revised: Therefore, we also include the best performers foreach dataset reported in the DomainBed paper.

Original: Note that we include these centralized methods for [Table] [Table] reference only, and they are not direct competitors of our model(s).
Revised: Note that we include these centralized methods forreference only, and they are not direct competitors of our model(s).

Original: For this reason, we highlight (inbold) the best performant model for the centralized and FL settings separately.
Revised: For this reason, we highlight (inbold) the best performant model for the centralized and FL settings separately.

Original: 4.4 Results
Revised: 4.4 Results

Original: Tables 1, 2, 3 and 4 show that our method (FedSR) outperforms the baseline FedAVG significantly(by a margin of 1.5% to 3% on average) and consistently in all experiments.
Revised: Tables 1, 2, 3 and 4 show that our method (FedSR) outperforms the baseline FedAVG significantly(by a margin of 1.5% to 3% on average) and consistently in all experiments.

Original: FedADG [44] (aexisting method for DG in the FL setting) does outperform FedAVG but is not yet on par withcentralized methods.
Revised: FedADG [42] (aexisting method for DG in the FL setting) does outperform FedAVG but is not yet on par withcentralized methods.

Original: Meanwhile, FedSR achieves competitive results when compared to state-of-theart centralized DG methods (as listed in the first half of the tables).
Revised: Meanwhile, FedSR achieves competitive results when compared to state-of-theart centralized DG methods (as listed in the first half of the tables).

Original: This indicates the effectivenessof our method and the potential for Domain Generalization in the FL setting.
Revised: This indicates the effectivenessof our method and the potential for Domain Generalization in the FL setting.

Original: As mentioned earlier,our method can be interpreted as an implicit form of representation alignment.
Revised: As mentioned earlier,our method can be interpreted as an implicit form of representation alignment.

Original: Please refer toour supplementary material for the visualization of our representation space, which does show thealignment effect of our method.
Revised: Please refer toour supplementary material for the visualization of our representation space, which does show thealignment effect of our method.

Original:  Conclusion
Revised:  Conclusion

Original: To conclude, in this paper, we present a simple and effective method for implicit alignment of therepresentation across domains/clients in a Federated Learning setting for a better domain generalization ability.
Revised: To conclude, in this paper, we present a simple and effective method for implicit alignment of therepresentation across domains/clients in a Federated Learning setting for a better domain generalization ability.

Original: We employ two regularization techniques, namely CMI and L2R, to learn a simplerepresentation of the data in the hope for a better generalization.
Revised: We employ two regularization techniques, namely CMI and L2R, to learn a simplerepresentation of the data in the hope for a better generalization.

Original: Furthermore, we also show that theseregularizers implicitly align the marginal and conditional distribution of the representation, whichare shown to be effective in the domain generalization problem.
Revised: Furthermore, we also show that theseregularizers implicitly align the marginal and conditional distribution of the representation, whichare shown to be effective in the domain generalization problem.

Original: Extensive experiments show that ourmethod outperforms relevant FL baselines (FedAVG, FedADG) and can even achieve competitive
Revised: Extensive experiments show that ourmethod outperforms relevant FL baselines (FedAVG, FedADG) and can even achieve competitive

Original: DG performance when compared to centralized approaches.
Revised: DG performance when compared to centralized approaches.

Original: A potential limitation of our methodis that the invariance is enforced among the source domains and might fail to generalize to the unseentarget domain (since the target domain can be arbitrary).
Revised: A potential limitation of our methodis that the invariance is enforced among the source domains and might fail to generalize to the unseentarget domain (seen the target domain can be arbitrary).

Original: However, this is a common comment forall domain invariant representation methods and not specific to ours.
Revised: However, this is a common comment forall domain invariant representation methods and not specific to ours.

Original: Furthermore, experiments showthat our method does lead to better representation alignment and prediction performance in practice.
Revised: Furthermore, experiments showthat our method does lead to better representation alignment and prediction performance in practice.

Original: Checklist
Revised: Checklist

Original: For all authors...
Revised: For all authors...

Original: (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s
Revised: (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s

Original: contributions and scope?
Revised: contributions and scope?

Original: Yes
Revised: [Yes]

Original: (b) Did you describe the limitations of your work?
Revised: (b) Did you describe the limitations of your work?

Original: [Yes] see Section 
Revised: [Yes] see Section 5(c)

Original: (c) Did you discuss any potential negative societal impacts of your work?
Revised: Did you discuss any potential negative societal impacts of your work?

Original: [N/A]
Revised: [N/A]

Original: (d) Have you read the ethics review guidelines and ensured that your paper conforms to
Revised: (d) Have you read the ethics review guidelines and ensured that your paper conforms to

Original: them?
Revised: them?

Original: [Yes]
Revised: [Yes]

Original: If you are including theoretical results...
Revised: If you are including theoretical results...

Original: (a) Did you state the full set of assumptions of all theoretical results?
Revised: (a) Did you state the full set of assumptions of all theoretical results?

Original: [Yes]
Revised: [Yes]

Original: (b) Did you include complete proofs of all theoretical results?
Revised: (b) Did you include complete proofs of all theoretical results?

Original: [Yes] see our Supplementary
Revised: [Yes] see our Supplementary

Original: File
Revised: File

Original: If you ran experiments...
Revised: If you ran experiments...

Original: (a) Did you include the code, data, and instructions needed to reproduce the main experi-
Revised: (a) Did you include the code, data, and instructions needed to reproduce the main experi-

Original: mental results (either in the supplemental material or as a URL)?
Revised: mental results (either in the supplemental material or as a URL)?

Original: [Yes]
Revised: [Yes]

Original: (b) Did you specify all the training details (e
Revised: (b) Did you specify all the training details (e

Original: g., data splits, hyperparameters, how they were chosen)?
Revised: g., data splits, hyperparameters, how theywere chosen)?

Original: [Yes] see Section
Revised: [Yes] see Section

Original: (c) Did you report error bars (e g., with respect to the random seed after running experi- ments multiple times)?
Revised: (c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?

Original: [Yes]
Revised: [Yes]

Original: (d) Did you include the total amount of compute and the type of resources used (e g., type of GPUs, internal cluster, or cloud provider)?
Revised: (d) Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)?

Original: [Yes] See Section
Revised: [Yes] See Section

Original: If you are using existing assets (e
Revised: If you are using existing assets (e

Original: g., code, data, models) or curating/releasing new assets...
Revised: g., code, data, models) or curating/releasing new assets...

Original: (a) If your work uses existing assets, did you cite the creators?
Revised: (a) If your work uses existing assets, did you cite the creators?

Original: [Yes]
Revised: [Yes]

Original: (b) Did you mention the license of the assets?
Revised: (b) Did you mention the license of the assets?

Original: [N/A]
Revised: [N/A]

Original: (c) Did you include any new assets either in the supplemental material or as a URL?
Revised: (c) Did you include any new assets either in the supplemental material or as a URL?

Original: [Yes]
Revised: [Yes]

Original: We include our source code in the Supplementary File
Revised: We include our source code in the Supplementary File

Original: (d) Did you discuss whether and how consent was obtained from people whose data you’re
Revised: (d) Did you discuss whether and how consent was obtained from people whose data you’re

Original: using/curating?
Revised: using/curating?

Original: [N/A]
Revised: [N/A]

Original: (e) Did you discuss whether the data you are using/curating contains personally identifiable
Revised: (e) Did you discuss whether the data you are using/curating contains personally identifiable

Original: information or offensive content?
Revised: information or offensive content?

Original: [N/A]
Revised: [N/A]

Original: If you used crowdsourcing or conducted research with human subjects...
Revised: If you used crowdsourcing or conducted research with human subjects...

Original: [Figure] ) Did you include the full text of instructions given to participants and screenshots, if[Figure]
Revised: (a) Did you include the full text of instructions given to participants and screenshots, if

Original: applicable?
Revised: applicable?

Original: [N/A]
Revised: [N/A]

Original: [Figure] Did you describe any potential participant risks, with links to Institutional Review[Figure]
Revised: (b) Did you describe any potential participant risks, with links to Institutional Review

Original: Board  approvals, if applicable
Revised: Board (IRB) approvals, if applicable?

Original: [N [Figure]
Revised: [N/A]

Original:  Did you include the estimated hourly wage paid to participants and the total amount
Revised: (c) Did you include the estimated hourly wage paid to participants and the total amount

Original: spent on participant compensation
Revised: spent on participant compensation?

Original: [N A
Revised: [N/A]