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)

Original: 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 machine learning framework in which multiple clients collaborate (with the help of a central server) 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 source clients’ data (e.g., mobile users) without considering its generalization ability to unknown 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 aforementioned issue.

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 setting where 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 novel representation learning framework, namely FedSR, which enables domain generalization while still respecting the decentralized and privacy-preserving natures of this 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 to learn 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 mutual information (between the representation and the data given the label) regularizer to 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 theoretical connections between the above two objectives and representation alignment in domain generalization.

Original: Extensive experimental results suggest that our methodsignificantly outperforms relevant baselines in this particular problem.
Revised: Extensive experimental results suggest that our method significantly 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 K clients 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: For privacy reasons, we consider the setting where it is desirable/preferred that the clients do not share their 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 devices collaborate 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 K hospitals 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 for the 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 incorporate existing 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, most domain invariant representation learning methods [9, 38, 31, 22, 32] require sharing/comparing the representation of the data across domains, while meta-learning based methods [21, 5] require access to 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 36th Conference on Neural Information Processing Systems (NeurIPS 2022). 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 data leakage 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, [45] aligns the representation distribution 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 does fully preserve the data privacy, it can be overly complicated to implement in practice, and the performance 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 main baselines 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 domain generalization 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 the conditional 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: These regularizers are easy to implement and do not require the sharing of data between clients, thus fully respecting 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 the representation can contains, which we hope will help to ignore spurious correlations and lead to better generalization 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 as implicit forms of marginal and conditional representation alignment , which also further explains the 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 representation whose marginal distribution and conditional distribution (given the label) are invariant across the domains without explicitly comparing those distributions between clients.

Original: Our contributions in thiswork are as follows:
Revised: Our contributions in this work 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 the prediction 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 regularization methods, namely the conditional mutual information regularizer and the L2-norm regularizer on 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 well connected with marginal and conditional alignment in the invariant representation learning framework.

Original: • We validate the effectiveness of our resulting model with a wide range of datasets and modelarchitectures and show that it significantly outperforms existing FL methods on the domaingeneralization task and can even compete with state-of-the-art centralized methods.
Revised: • We validate the effectiveness of our resulting model with a wide range of datasets and model architectures and show that it significantly outperforms existing FL methods on the domain generalization task and can even compete with state-of-the-art centralized methods.

Original:  Related Work 
Revised:  Related Work 2.1 Federated Learning

Original: 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 SGD updates 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 context of DG in the FL setting, this method is equivalent to Empirical Risk Minimization in a centralized setting, 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 is the 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: A great amount of research [23, 48, 12, 24, 17, 36] has been done to tackle this problem (i.e., how to train 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: Another line of approaches to deal with this problem is to allow each client to have a personalized model and 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 above methods tackle the distribution shift across the source domains/clients, they still have not dealt with the 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 approach is to learn a domain-invariant representation [30, 25, 26, 43, 2, 15, 47, 1, 22, 38, 35] (meaning to learn a representation whose marginal distribution and/or conditional distribution given the label are unchanged across the domains), hoping that this representation generalizes better to an unseen target domain.

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 suitable for 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 to adapt 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 from source domains (in a centralized server).

Original: [45] considers a single-domain domain generalizationsetting.
Revised: [46] considers a single-domain domain generalization setting.

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).
Revised: However, it needs to generate a “fake” target distribution from a source distribution and then maximize 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.
Revised: This might be overly complicated when applied at each client level in a FL setting and hence less 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.
Revised: [34] considers a federated domain adaptation, where each source domain aligns the representation distribution with the target domain by receiving the representations from this target distribution.

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.
Revised: 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:  Approach 3.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 in this 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 to train 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 ) , where x ∈ 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 ) changes across 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 data distributions 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 model will 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: where f i ( w ) is the local objective function,

Original: [Equation]
Revised: [Equation]

Original: [Figure] [Equation]
Revised: ( 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 average among 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 simple average 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 for several iterations of stochastic gradient descent; then the central server will collect the newly-learned parameters from all (active) clients, average them to form the new global parameter, and broadcast it back 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 minimization each 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 a more detailed description of these FL methods, we refer the readers to [28, 42, 37].

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 perform well 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, we hope 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 )] butalso the loss on unseen target domains, in either an average or worst-case sense as defined below.
Revised: [ ℓ ( w ; x, y )] but also 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 in which 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 learn to 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 target domains .

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, so regularization 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 technique among those is to learn a domain invariant representation, which means to learn a representation z of 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 methods often 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 will discuss 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 FL methods 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 their own 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 phone users, 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 very important 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 centralized training setting with data (or a part of the data, such as a representation) being shared across the domains.

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 approach typically 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: Towards this 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 the problem setting of representation learning in Subsection 3.2, and then present our domain invariant representation 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 distribution p ( 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 deterministic mapping (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 distribution p 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: since z is conditionally independent of y given x [Equation] From the representation z , we learn a classifier/regressor that predicts y given z with the predictive distribution ˆ 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: For example, with a classification problem, this is often just a linear layer followed by a softmax layer to form 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: 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 log predictive, 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 a deterministic mapping, it is obvious that we only sample one z ), leading to the training objective:

Original: [Equation]
Revised: [Equation]

Original:  z ( n ) i is a single sample [Equation]
Revised: N i (cid:88) n = − log ˆ p ( y ( n ) i | z ( n ) i ) , where z ( n ) i is a single sample from p ( z | x ( n ) i )

Original: [Equation]
Revised: [Equation]

Original: 175 The quantity − log ˆ p ( y | z ) is often non-negative for all common predictive distributions.
Revised: 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, this becomes 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 process is 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 of domain 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 information of 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 simple representation 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 several well-known regularization techniques to restrict the representation’s complexity and encourage the deep 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 these regularization techniques have strong theoretical connections to domain-invariant representation learning 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 upper bound 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 i also hasthe effect to align the marginal distribution p
Revised: Moreover, ℓ L 2 R i also has the 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 final local 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.)
Revised: (We will also test the two regularizers separately, namely the FedL2R and FedCMI variants.)

Original: L2- norm Regularizer
Revised: 3.3.1 L2-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.
Revised: In optimization and learning theory, L2-norm of the weights/parameters is often regarded to as a measure 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.
Revised: 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).
Revised: (Note that some methods like [29] use the L2-norm but in the parameter space to help with the generalization of Federated Learning).

Original: Specifically, the regularizer has the form:
Revised: Specifically, the regularizer has the form:

Original: [Equation]
Revised: [Equation]

Original: |

Original: [Equation]

Original: z
Revised: z

Original: [Equation]
Revised: [Equation]

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

Original: [Equation]
Revised: [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.
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.

Original: This connection can be explained inmore detail as below:
Revised: This connection can be explained in more detail as below:

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

Original: [Equation]
Revised: [Equation]

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

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

Original: Then we have:
Revised: Then we have:

Original: [Equation]
Revised: [Equation]

Original: an additive

Original: [Equation]
Revised: [Equation]

Original: (13) 215 And hence:
Revised: And hence:

Original: [Equation]
Revised: [Equation]

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

Original: Note further that:
Revised: Note further that:

Original: [Equation]
Revised: [Equation]

Original: Hence, if the entropy H ( p i ( z )) does not change significantly during training, minimizing ℓ L 2 R i will219 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.
Revised: Hence, if the entropy H ( p i ( z )) does not change significantly during training, minimizing ℓ L 2 R i will also minimize KL [ p i ( z ) | q ( z )] , which encourages p i ( z ) to be close to q ( z ) , i.e., an implicit alignment of 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 .
Revised: Nevertheless, also according to the above equation, if we really want to minimize KL [ p i ( z ) | q ( z )] (for the marginal alignment), it is perhaps more rigorous to minimize the regularizer ℓ L 2 R i [Equation]

Original: However, when σ is small enough, this term is dominated by ℓ L 2 R
Revised: p [Equation] z

Original:  i (meaningthat if σ is small, minimizing ℓ L 2 R i almost guarantees to minimize KL [ p i ( z ) | q ( z )] ).
Revised: [Equation] ) instead of ℓ L 2 R i . However, when σ is small enough, this term is dominated by ℓ L 2 R i (meaning that 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.
Revised: Empirically, we found that the two variants achieve almost identical performance.

Original: We therefore recommend using ℓ L 2 R i instead of ℓ L 2 R i − 2 σ 2 H ( p i ( z )) for its simplicity.
Revised: 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.2 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 often used to regularize the amount of information z can contains [3, 39].

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 might be too restrictive in practice if its coefficient is not tuned properly, as it encourages the representation to 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 mutual information 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 this is 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 mutual information 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 the information 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 is hard 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 to minimize 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 for classification (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 mapping p ( 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 a probabilistic 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 with the ℓ 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 of client 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 regularizer has a theoretical connection with conditional distribution alignment in domain generalization.

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

Original: Lemma 2.
Revised: Lemma 2.

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 , forcing p 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 i ( z | y )
Revised: Therefore, the model will try to enforce p 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 of the conditional distribution of the representation z given the label y , which is a common and effective technique in conventional Domain Generalization.

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

Original: 3.3.3 Implementation Choices and Optimization
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. 18, and Eq.

Original: 11 (insteadof full-batch training).
Revised: 11 (instead of 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 through the 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 this reparameterization 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 the variational 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: 18 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 wonder that such a simple choice of the variational distribution r ( z | y ) might not be able to approximate p 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 , 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).

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 experiment section, for simplicity, we only use a Gaussian distribution for r ( z | y ) .

Original:  Experiments 
Revised:  Experiments 4.1 Datasets

Original: 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 more challenging) 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 an angle 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 [31, 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 that this 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: [Table] [Table] OfficeHome [41]: has 15,500 images of daily objects from four domains: art, clipart, product and real.

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 [33] 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 images are 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 as the 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 data as 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 from our 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 fully connected layer as the representation network g θ to get a representation z of 64 dimensions.

Original: A singlelinear layer is then used to map the representation z to the ten output classes.
Revised: A single 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 the same as the network used by [15, 31].

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 gradient descent (SGD), using a learning rate of 0.001 and minibatch size 64, and report performance on the test 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 each communication 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 representation learning methods, we use the most common choice of backbone network in existing works as the representation 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 the backbone with a linear layer of dimension 512 for ResNet18 and 2048 for ResNet50 to form the representation 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 from the 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 of 5000 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 within each communication round ( E = 5 ).

Revised: However, we also found that our method is not sensitive to this

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: [Table] [Table] hyper-parameter (performance with E = 100 is almost identical). 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 data leakage 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 fair comparison 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 supplementary material.

Original: Data augmentation is also standard practice for real-world computer vision datasets like 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: Data augmentation is also standard practice for real-world computer vision datasets like PACS, OfficeHome and Domainet, and during the training we augment our data as follows: crops of random size and aspect ratio, resizing to 224 × 224 pixels, random horizontal flips, random color jitter, randomly converting the image tile to grayscale with 10% probability, and normalization using the ImageNet channel means and standard deviations. Our code will be released at https://github.com/atuannguyen/FedSR.

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] (a common FL algorithm) and FedADG [45] (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 methods for 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 OfficeHome and DomainNet experiments, since our setting (e.g., backbone network) and evaluation protocol are identical to those of DomainBed [11], our results can be compared directly to the results of the 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 for each 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 for reference 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 (in bold) 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 [45] (a existing method for DG in the FL setting) does outperform FedAVG but is not yet on par with centralized 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 effectiveness of 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 to our supplementary material for the visualization of our representation space, which does show the alignment 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 the representation 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 simple representation 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 these regularizers implicitly align the marginal and conditional distribution of the representation, which are 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 our method 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 method is that the invariance is enforced among the source domains and might fail to generalize to the unseen target domain (since 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 for all 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 show that our method does lead to better representation alignment and prediction performance in practice.

Original: Checklist
Revised: Acknowledgments Author A. Tuan Nguyen acknowledges Meta AI for funding his PhD study. Oxford lab headed by author Philip Torr is supported by the UKRI grant: Turing AI Fellowship EP/W002981/1. Meta AI author Ser-Nam Lim is neither supported by the UKRI grant nor has any relationship to the grant.

Original: For all authors...
Revised: Proof. Recall that the joint distribution of x, y, z for client i is: [Equation]

Original: (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s
Revised: [Equation] We have: [Equation] Notice that:

Original: contributions and scope?
Revised: [Equation] Therefore: [Equation] A.2 Lemma

Original: Yes
Revised: Proof.

Original: (b) Did you describe the limitations of your work?
Revised: We need to prove: [Equation] z

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

Original: (c) Did you discuss any potential negative societal impacts of your work?
Revised: ( since E p i ( x,y,z ) [log r ( z | y )]

Original: [N/A]
Revised: = E p

Original: (d) Have you read the ethics review guidelines and ensured that your paper conforms to
Revised: i ( y,z ) [log r

Original: them?
Revised: z | y )])

Original: [Yes]
Revised: [Equation]

Original: If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?
Revised: [Equation] [Equation] ( x,

Original: [Yes]
Revised: [Equation]

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

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

Original: File
Revised: [Equation] [Equation] [Equation]

Original: If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either in the supplemental material or as a URL)?
Revised: Therefore

Original: [Yes]
Revised: [Equation]

Original: (b) Did you specify all the training details (e
Revised: we only need to prove that E p i ( x,z | y )

Original: g., data splits, hyperparameters,
Revised: [log p ( z | x )] ≥ E p

Original: how they
Revised: i ( z | y ) [log p i

Original: were chosen)?
Revised: z | y )] ∀ y

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

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

Original: [Yes]
Revised: [Equation]

Original: (d) Did you include the total amount of compute and the type of resources used (e g., type
Revised: Notice that the function f ( a ) = a log a, a > 0 is a convex function since f ′′ ( a ) =

Original: of GPUs, internal cluster, or cloud provider)?
Revised: 1 a > 0 ∀ a > 0 .

Revised: Hence, due to Jensen’s Inequality we have:

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

Original: If you are using existing assets (e
Revised: where b = E p i ( x | y ) [ p ( z | x )]

Original: g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?
Revised: = p i ( z | y ) .

Original: [Yes]
Revised: Therefore, E p i ( x | y )

Original: (b) Did you mention the license of the assets? [N/A]
Revised: [ p ( z | x ) log p ( z | x )] ≥ p

Original: (c) Did you include any new assets either in the supplemental material or as a URL?
Revised: i ( z | y ) log p

Original: [Yes]
Revised: i ( z | y ) .

Original: We include our source code in the Supplementary File
Revised: Then, inequality 44 holds, and we conclude our proof.

Original: (d) Did you discuss whether and how consent was obtained from people whose data you’re
Revised: B Details on the ℓ CMIi objective

Original: using/curating?
Revised: As mentioned in the main paper, we use a Gaussian distribution for both p ( z | x ) and r ( z | y ) , so that the KL term can be computed analytically.

Original: [N/A] (e) Did you discuss whether the data you are using/curating contains personally identifiable
Revised: Let w 1 be the parameter of the (probabilistic) representation network.

Original: information or offensive content?
Revised: Given an input x 0 (of a datapoint ( x 0 , y 0 ) ), the network output the mean and standard deviation of the Gaussian distribution p ( z | x 0 ) (both from the final layer of the network), denoted µ w 1 ( x 0 ) and σ w 1 ( x 0 ) . This means that p ( z | x 0 ) =

Original: [N/A]
Revised: N ( z ; µ

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

Original: [Figure]
Revised: [Equation]

Original:  ) Did you include the full text of instructions given to participants and screenshots, if [Figure] applicable?
Revised: 1 ( x 0 ) , σ 2 w 1 ( x 0 )) .

Original: [N/A] [Figure] Did you describe any potential participant risks, with links to Institutional Review [Figure]
Revised: Also as pointed out in the paper, for a classification problem (where there is a finite number of the labels), we set r ( z | y )

Original: Board
Revised: =

Original: approvals, if applicable [N
Revised: N ( z ; µ y , σ 2 y ) , where µ y , σ y ( y = 1 ..C ) are the variational parameters to be optimized.

Original: [Figure]
Revised: Therefore, for the datapoint x 0 , y 0 , we have r ( z | y 0 ) = N ( z ; µ y 0 , σ 2 y 0 ) . Note that when z is a K -dimensional vector, µ w 1 ( x 0 ) , σ w 1 ( x 0 ) , µ y 0 and σ y 0 are all K -dimensional.

Revised: The KL term in ℓ CMI can be computed analytically as:

Revised: [Equation]

Revised: If K > 1 (i.e., z is high dimensional), the calculations in Eq.

Revised: 46 are element-wise, and the result are summed across the dimension K at the end.

Revised: As this computation is deterministic with respect to w 1 and µ y , σ y ( y ∈ 1 ..C ), the gradient w.r.t.

Revised: these parameters can be computed straight-forwardly.

Revised: C Experiments C.1 Hyper-parameters:

Revised: As mentioned in the main text, we use random search to tune the hyper-parameters of our method, namely α CMI and α L 2 R .

Original: Did you include the estimated hourly wage paid to participants and the total amount
Revised: Specifically, the tuned value for those parameters are: RotatedMNIST • FedL2R: α L 2 R = 0 .• FedCMI: α CMI = 0 .•

Revised: FedSR: α L 2 R = 0 .

Original: {+1 , α CMI = 
+}
Revised: {+1 , α CMI = 
+}

Revised: .PACS • FedL2R: α L 2 R = 0 .•

Revised: FedCMI: α CMI = 0 .•

Revised: FedSR: α L 2 R = 0 .

Revised: 01 , α CMI = 0 .

Revised: OfficeHome

Original: spent on participant compensation
Revised: • FedL2R: α L 2 R = 0 .• FedCMI: α CMI = 0 .• FedSR: α L 2 R = 0 . 05 , α CMI = 0 .

Original: [N
Revised: DomainNet • FedL2R: α L 2 R = 0 .• FedCMI: α CMI = 0 .• FedSR: α L 2 R = 0 .

Revised: 01 , α CMI = 0 .

Revised: C.2 Visualization of the Representation space

Original: A
Revised: We also visualize the representation space of our method and observe that it aligns the representation much better than the conventional FL baseline FedAVG (which does not actively attempt to align the representation distributions). Specifically, Figure 1 show the t-SNE [40] visualization of the representation of FedSR and FedAVG in the RotatedMNIST experiment, with M 0 (the target domain) and M 15 (one of the source domains). We can clearly see that both the marginal and the conditional distributions of the representation are aligned better between the two domains for FedSR [Figure]