FedSR: A Simple and Effective Domain Generalization Method for Federated Learning
Authors: A. Tuan Nguyen
Philip Torr
Ser-Nam Lim
Venue: NeurIPS 2022 Accept (2022)
Keywords: domain generalization
representation alignment
federated learning

Abstract: 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. However, most existing 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). In this paper, we incorporate the problem of Domain Generalization (DG) into Federated Learning to tackle the aforementioned issue. 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. 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. Motivated by classical machine learning algorithms, we aim to learn a simple representation of the data for better generalization. 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). Furthermore, we provide theoretical connections between the above two objectives and representation alignment in domain generalization. Extensive experimental results suggest that our method significantly outperforms relevant baselines in this particular problem.

FedSR: A Simple and Effective Domain Generalization Method for Federated Learning
A. Tuan Nguyen University of Oxford tuan@robots.c ox.oox.ccoooo ac.acc..o uk
Philip H. S. Torr University of Oxford philip. torrdeng. ox.ac.
Ser-Nam Lim Meta AI Research sernamlim@fb com
Abstract
Federated Learning (FL.) refers to the decentralized and privacy- prreacc-preseevingg machine learning framework in which multiple clients collaborate (with the help of a central server) to train global model without sharing their data. However, most existing FL methods only focus on maximizing the model's performance on the source clients' data (e.g., mobile users) without considering its generalization ability unknown target data (e.g.. a new user). In this paper, we incorporate the problem of Domain Generalization (DG) into Federated Learning to tackle the aforementioned issue. However, virtually all existing DG methods require a centralized setting where data is shared across the domains, which violates the principles of decen- tralized FL. and hence not applicable. To this end, we propose simple yet novel representation learning framework, namely FedSR, which enables domain general- ization while still respecting the decentralized and privacy-preserving natures of this FL setting. Motivated by classical machine learning algorithms, we aim to learn a simple representation of the data for better generalization. In particular, we enforce an L.2normm 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 spu- rious correlations such as the background). Furthermore, provide theoretical connections between the above two objectives and representation alignment domain generalization. Extensive experimental results suggest that method significantly outperforms relevant baselines in this particular problem.
Introduction
In this paper, we are interested the problem of decentralized Federated Learning (FL., where K clients collaborate to jointly train machine learning model with their local (decentralized) data. For privacy reasons, we consider the setting where is desirableepreferreed that the clients do not share their data with each other. For example, [28] brings up an example scenario where mobile devices collaborate to train a model without sending their data to central server. Another example is that K hospitals train a diagnosis model together but do not want to share their patients' data.
However, a major concern with this machine learning framework is that often does not account for the possible distribution shift the target data, problem commonly referred to Domain General- ization (DG). Furthermore, and perhaps more problematically, it not straightforward to incorporate nature. Specifically, most
existing DG techniques into this FL setting due to its privacy-preservicn domain invariant representation learning methods [9, 38, 31, 1, 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).
There have been some early attempts to incorporate domain generalization into the FL framework. For example, [27] proposes a technique tailored for medical image segmentation which allows the
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

clients to share their data in the frequency space with each other. This can be viewed as form of data leakage and is undesirable in the FL setting we consider. Meanwhile, [44] aligns the representation distribution across domains via reference distribution (from generative model). While this does fully preserve the data privacy, can be overly complicated to implement practice, and the performance is not yet on par with SOTA centralized methods. Nevertheless, this is one of our main baselines in this paper.
To tackle the above challenges, in this paper, we introduce simple yet effective method for domain generalization in the FL setting. We propose to learn "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. These regularizers are easy to implement and do not require the sharing of data between clients, thus fully respecting the principles of FL.. 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. With 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. 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. Our contributions in this work are as follows:
We propose simple method, which aims to learn a simple representation of the data for the prediction task, for DG in the FL setting. To this end, we revisit two common regularization methods, namely the conditional mutual information regularizer and the L2-norm regularizer the representation, to restrict the amount of information it can contain.
Furthermore, we show theoretically and empirically that these two regularizer are well connected with marginal and conditional alignment in the invariant representation learning framework.
We validate the effectiveness of our resulting model with a wide range datasets and model architectures and show that it significantly outperforms existing FL methods the domain generalization task and can even compete with sfate-of-the-ar centralized methods.
2 Related Work
2.1 Federated Learning
One of the first and most common FL algorithms FedAVG [28], which uses several local SGD updates within each communication round to the server to save communication cost. In the context of DG in the FL setting, this method is equivalent to Empirical Risk Minimization setting, and is one of our main baselines. One of the main challenges that FI. needs to tackle is the statistical diversity among the clients (i.e., each client has different local data distribution). great amount of research [23, 47, 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. data of the source clients in the FL setting), Another line of approaches to deal with this problem is to allow each client to have a personalized model FL way (personalized FL) [13, 6, 4, 8, 16]. Although these above
centralized
and jointly train the models 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.
2.2 Domain Generalization
Most of the existing DG methods consider a centralized setting, A predominant and effective approach
to learn a domain-invariant representation [30, 25, 26, 42, 2, 15, 46, 1, 22, 38, 35] (meaning to learn 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. Our method can be viewed as implicit forms of representation alignment, thus more suitable for the FL setting. 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. These methods typically require access to pool of data from source domains (in a centralized server). [45] considers single-domain domain generalization

---

setting, However, it needs to generate a "fake" target distribution from source distribution and then maximize the mutual information between the representations of the two domains (source vs fake target). This might be overly complicated when applied at each client level in Fl. setting and hence less practical. [34] considers federated domain adaptation, where each source domain aligns the representation distribution with the target domain by receiving the representations from this target distribution. 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 s data (including the data representation) for privacy reasons.
Approach
3.1 Preliminaries
We first present the problem of Federated Learning and the importance of Domain Generalization in this setting.
Federated Learning In the conventional FL setting, K clients collaborate with each other to train global model. Each clients i (1 S i K) has its own data distribution (z y), where x E X' is the input and E y is its corresponding label, and a dataset with N, data points: (N.)
D. = -
:y; (. is often assumed that the data distribution px(x, y) changes across the clients (which we will use interchangeably with domains in this paper). Here, the data distributions (x, y) are sampled from a family & of distributions (E can be finite or infinite), i.e., (z( y) ~ E. 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 datapoint (z, y) denoted as l(w; x, y).
(1) Mi ").... (x)
The global objective function to minimize is:
1 f(w) K
Σ fi(w)
where fiw) is the local objective function, se(w)
(p.. (1 y)]
N; E
((u (2(") 3/ ()). Note that the objective in Eq. 1 can also be weighted average among the clients. However, since the analysis and discussion are unchanged, we will use this simple average objective in this paper.
In most Federated Learning algorithms, the optimization process consists of a number of rounds. 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- leaaned parameters from all (active) clients, average them to form the new global parameter, and broadcast back each client at the beginning of next round. By taking multiple steps of the local minimization each round, the algorithm can reduce the communication cost (which is often the bottleneck). For a more detailed description of these FL. methods, we refer the readers [28, 41, 37].
Domain Generalization In the DG problem, we are more interested in training a model to perform well in an unseen target domain pr(x, y) ~ & (which is different from pr y) Vi). Therefore, we hope to minimize not only the expected loss on the source domains/clients (p.(..) [e(w; x. y) but also the loss unseen target domains, in either an average worst-case sense defined below.
In the average case, we want to minimize this loss over the distribution family E:
Eppne (pr(a.)) [[(u; x, v)]
Meanwhile, we might also be interested the worst-case quantity, i.e., the loss of the domaín in which the model performs the worst:
sup
E [ 3E
In conventional machine learning problem, we usually need to observe more datapoints to learn to generalize better to unseen datapoints. Similarly, in this Domain Generalization problem, additionally need to observe more domains in order to generalize better unseen target
3

---

domains. However, we typically have limited number of source domains in these DG problems, so regularization techniques often need to be applied for better generalization. One common technique among those is to learn domain invariant representation, which means to learn a representation of x that has the same marginal and/or conditional distribution across the domains. These methods often help to ignore domain-specific spurious correlations and lead to better generalization. We will discuss this in more detail later the paper.
Generalization Challenge of FL From the objective in Eq. 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. For example, if K clinical institutes in the US and UK collaborate to train 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. Similarly, if we use Federated Learning to train a machine learning model with data collected from multiple mobile phone users, it is crucial that the model can generalize a new user. Thus, domain generalization is very important for Federated Learning.
However, as mentioned earlier, applying existing domain generalization techniques to Federated Learning not straightforward. This 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. For example, the aforementioned domain invariant representation learning approach typically requires sharing and comparing the representation distribution across domains. Towards this end, we propose a simple but effective domain invariant representation learning technique for Federated Learning that fully keeps its decentralized and privacy aspects. 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.
3.2 Problem Setting: Representation Learning for Domain Generalization
In the representation learning framework, we aim to learn a representation z of x with a distribution p(z|z) parameterized by W1, which omit here for notation simplicity. This can be deterministic mapping (i.e., p(z|x)= = 6g.,, (z) where is the dirac delta function) or a probabilistic mapping
(e.g., p(2|z) = NN2;; Aw, (z), (x)) where A is the normal distribution). With the data distribution pa(x, y) of client i, its joint distribution of xr y, is:
(= y, =) = p. (z, v)p(= since = conditionally independent of y given æ
(4)
From the representation we learn a classifter/regresoor that predicts given with the predictive distribution p(ylz) parameterized by W2 (again, we omit we here for notation simplicity). For example, with a classification problem, this is often just linear layer followed by softmax layer to form the predictive distribution.
The predictive distribution of given x of our model is:
p(y|z) Ep(=(=) [Py]=2]
Note that when p(=[æ>) is a deterministic mapping, Eq. 5 simplifies into p(y12 = 9m. (z)).
For both regression and classification, the loss function of a data point xr y is often the negative log predictive, i.e, É(v; xr, v) log [p(c)=))((y(=)) and in this case wr {wi, w2).
Hence, the local objective of client is:
f: (x.)) |- log Ep(=)=) [[v]=)]]
During training, for a probabilistic mapping, often sample single = per from p(2|x) (for deterministic mapping, it is obvious that we only sample one 2), leading to the training objective:
F.(w)) Ep((z.v)
[Ep[ log p(vl2)]
No.
log Poy (= ((())) where e(") a single sample from p(=(sr{")
(8)
N

---

The quantity - log ( is often non -egatiiee for all common predictive distributions. For example, with a categorical predictive distribution 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).
Note that according Jensen's Inequality, (i(w) is upper bound of fi(w), so the training process sound.
The representation learning framework that we discussed so far still has not dealt with the problem of domain generalization. Many existing works DG focus on learning a domain-invariant representa- tion which means to learn a representation which has the same marginal (=)) and/or conditional (p (zly) or p. (y(z)) distributions across the domains. However, this often requires sharing information of the data (the representation) across domains/clients, which is not allowed in Federated Learning.
Motivated by the fact that classical machine learning algorithms make predictions based on simple representation of the data and often generalize better than modern deep learning models, in this paper, we aim learn a Simple Representation of the data (hence the name FedSR). We apply several well-known regularization techniques to restrict the representation's ' complexity and encourage the deep network p(=(ær) to only learn the essential information. Furthermore, we show that these regularization techniques have strong theoretical connections to domain-invarian representation learning and that they can achieve good generalization performance practice.
3.3 FedSR
As mentioned earlier, we employ an L2-norm regularizer (L!-2B)) on the representation and (an upper bound of) the conditional mutual information between the data and the representation given the label (C(:MT) to restrict the amount of information the representation can contain. Moreover, 2!22 also has the effect align the marginal distribution (z) to be centered around 0 as a Gaussian distribution, while gCM1 aligns the conditional distribution pa (zly) to reference distribution (also chosen to be Gaussian the experiment section). These two terms will be explained in detail below. The final local objective function of each client is:
F + a 22 + a C.M CMIICMI
(9)
where 12n and CMI are dyper-parameters (We will also test the two regularizers separately, namely the FedL2R and FedCMI variants.)
3.3.1 L2-norm Regularizer
In optimization and learning theory, 1.2-norm of the weights/parameters is often regarded to as a measure of the model's is complexity. Here, with similar idea, enforce L2-norm Regularizer (L2R) on the representation (not the network parameters) to restrict the complexity of the representa- tion. (Note that some methods like [29] use the L2-norm but in the parameter space to help with the generalization of Federated Learning). Specifically, the regularizer has the form:
RL2R
Ep((x)
(1p= 211
(10)
N:
(m)
12. where 2!th) is a single sample from p(z(a{())
(11)
N
Intuitively, this regularizer will have the effect of encouraging be centered around 0 for all clients, thus helping with the marginal alignment of the representation. This connection can be explained in more detail as below:
Connection to Domain Invariant Representation Learning:
First of all, rewrite 2.2n as:
L2R (p)( [Epe)))) [|]]] Epp((.>) [||=||]] Ex,(,) [||||]]]
(12)
Now, consider "reference" distribution q(2) = (0, o2) (with a small Ø). Then we have:
- log g(2) = ||
(with an additive constant)
(13)
202
5

---

And hence:
e428 = pp)(=) [||=|0]] = 20 |- log [=)) = 20 (D=(t))
(14)
where (z), (()) denotes the cross entropy from g(2) to p;(2).
Note further that:
H(p.(z), (()())))) (LLp())))))
(15)
Hence, if the entropy (p:(z)) does not change significantly during training, minimizing 4 will also minimize (L)g:(=))q(z)} which encourages (z(z) be close q(2), i.e., an implicit alignment of the marginal distribution.
Nevertheless, also according the above equation, if we really want to minimize (=(p((=)lg(2) (for the marginal alignment), it is perhaps more rigorous to minimize the regularizer 2 20 H (pi (=)) instead of g.-2n, However, when a is small enough, this term dominated by 2.22m (meaning that is small, minimizing 2 almost guarantees to minimize (=)()())) Empirically, we found that the variants achieve almost identical performance. We therefore recommend using 6 instead of 8f2R 20 H (z)) for its simplicity.
3.3.2 Conditional Mutual Information regularizer
The mutual information term between the input x and the representation (denoted [z, 2)) is often used to regularize the amount of information 2 can contains [3, 39]. However, this regularizer might be too restrictive in practice its coefficient is not tuned properly, as it encourages the representation to contain no information about the input data. In this paper, we use the conditional mutual information of æ and = given y. as less restricted than the former. We can also later see that this is well-connected with conditional distribution alignment in DG. Recall that this conditional mutual information term is calculated for domain/client as:
(I 2/y) P((Tly)p,
(x, 21y) = Ep.(m.m.)) log
(16)
Intuitively, the terms Je and [x,, =ly) work together to enforce the representation = to contain only the information needed to predict the label y, and no additional information (non-label-related) about
However, unfortunately, this mutual information term tractable because of p4(2ly) (which is hard to compute due to the integration over r). Therefore, we derive an upper bound in order to minimize Specifically:
Let r(2ly) be a(ny) conditional distribution of given y Vy E V. If p(2]z) and r(2ly) have the same support set Vx E X, E V,, we have:
Lemma 1.
[(x, S (p(( (K(p))) = CMMI
(17)
Proof. Provided the Supplementary Material.
Here, the upper bound JCMI can be computed and used as a regularizer when training the represen- tation network. We optimize both p(2 (x) and r(=ly) in order to minimize CCMI Note that r(aly) can be a network that takes y an input and output the distribution r(zly); more simply for classification (when we have a finite number of labels y), we can set r(=ly) = A2; Mty az).. where My: o3 (v I.CC) are the parameters to be optimized, with C' being the number of classes.
Note also that if we use this (upper bound) regularizer eCMI it requires the representation mapping p(2(x)) to be probabilistic (due to the KL term). Hence, in the experimental section, we use probabilistic representation network (p(zkz) (=; Mwn (x), ou (z))) for any models/variants with the @CMr objective, and use deterministic network otherwise. With the empirical dataset D, of client i, we have an empirical estimator of &CMI as below:

---

CCMI =Rp (2,s)) (I.Ip(([c))r<l|))]
N, (3) (r(pp)
kL|p(z|a
(18)
Connection Domain Invariant Representation Learning: We found that this regularizer has a theoretical connection with conditional distribution alignment in domain generalization. We start with the following lemma:
Lemma 2.
C CMMI EE (x.(a
[LIp(c)=))r<<]]]]] Ep.00 [KLpp((
(19)
Proof. Provided in the Supplementary Material.
Therefore, by minimizing CCMI, we also minimize the divergence (=In() Vy, forcing (=(y)) and r(=ly) to be close each other: p=(yy)) R r(zly). Therefore, the model will try to enforce (=ly) R Pj (=] r(=ly)) ) V clients/domains j. i.e., we are doing implicit alignment of the conditional distribution of the representation 2 given the label y, which is a common and effective technique in conventional Domain Generalization. For example, [26, 25, 46] explicitly align this conditional distribution (not suitable for Federated Learning).
3.3.3 Implementation Choices and Optimization
For the federated optimization, we use FedAVG, which is one of the first and perhaps most common FL optimization algorithms (although we can use any federated optimization techniques).
Following common practice, we use hini-batches to approximate Eq. 8, Eq. 18, and Eq. 11 (instead of full-batch training).
For a probabilistic network (when using (CMl), to allow the gradient to backpropagate through the samples z/(") to the network parameters wi, we use the reparameterization trick [18]. With this reparameterizatioon trick, F. and 6.22 can be implemented and back-propagaiae propagated straightforwardly.
Regarding the conditional mutual information regularizer e (C'M we use Gaussian distribution for the variational conditional distribution r(2 r(2ly), hence the KL term in Eq. 18 can be computed analytically (recall that we also use a Gaussian distribution for p(2[z) Here the astute readers might wonder that such simple choice of the variational distribution r(2 (aly) might not be able to approximate Pa (2ly). However, note that since both p(2 (r) and r(2|y) are optimized in order to minimize CCMT it encourages the network to learn simple representation can match (=(y)) (which arguably might be good regularization effect). Furthermore, depending datasets, can, in theory, more complex distributions such as Gaussian mixture as the variational distribution. In such cases, one can use methods such as MC sampling to estimate the KL. term. our experiment section, for simplicity, only use Gaussian distribution for r(2ly).
that the variational distribution r(2ly)
4 Experiments
4.1 Datasets
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.
RotatedMNIST [10]: In this dataset, the MNIST images [19] are rotated counter-clockwise with an angle of 09, 15", 30", 455,, 60% and 75 to form six domains M1.M M30, M Mgu and M75- The task is classification with ten classes (digits 0 to 9). We the same version of the dataset as in [31, 15] (where only 1000 images are rotated to form the domain) for easy comparison. Note that this is slightly different from the version [1] and thus the numbers are not comparable.
PACS [20]: contains 9,991 images from four different domains: art painting, cartoon, photo, sketch.
The task is classification with seven classes.
7

---

Table 1: RotatedMNNIST. Reported numbers from 3 runs
Domains
Models
MD
M15
Mao
Mas
Msa
M7a
Average
HIR [42]
90.34
99.75
99.40
96.17
99.25
91.26
96.03
Centralized
DIVA [155]
93.5
99.3
99.1
99.2
99.3
93.0
97.2
Methods
DGER [46]
90.09
99.24
99.27
99.31
99.45
90.81
96.36
DIRT-GAN [31]
97.2
99.4
99.3
99.3
99.2
97.1
98.6
FedAVG [28]
85.920.6
98. 08.2
98. 98.22.2
98,920.1
98.7=0.2
86.2+0.7
94.5
FedADG [44]
89.9+0.5
98. 98.9=0.2
19.0=0.1
99.1=0.1
98.9=0.2
90.2=0.4
FL methods
FedCMI (ours)
91.5=#03
99.2=0.1
99. 99.3=0.1
99.2=0.1
99.2=0.1
91.0=0.3
96.5
FedL2R (ours)
90.5+0.5
99.1
99.2=0.1
99.1=0.1
99.2=0.1
90.7=0.6
96.3
FedSR (ours)
91.6+0.3
99.3+0.1 99.3+0.1
99.2+0.1
99.3#0.1
91.5=0.3
96.7
Table PACS. Reported numbers are from 3 runs
PACS
Models
Backbone
A
P
S
Average
Centralized Methods
DGER [46]
Resnet18
80.70
76.40
96.65
71.77
81.38
DIRT-GAN [31]
Resnet18
82.56
76.37
95.65
79.89
83.62
FedAVG [28]
Resnet18
77.8=0.5
72.8=0.4
91.9+0.5
78.8+0.3
80.3
FedADG [44]
Resnet18
77.8=0.5
74.7=0.4
92.9+0.3
79.5+0.4
81.2
FL Methods
FedCMI (ours)
Resnet18
80.8+0.4
73.7+0,2
92.8+0.5
79.5+0.2
81.7
FedL2R (ours)
Resnet18
82,2+0.4
75,820.3
92.840.4
81.640.1
83.1
FedSR (ours)
Resnet18
83.2+0.3
76.0+0.3
93.8+0.5
81.9+0.2
83.7
OfficeHome [40]: has 15,500 images of daily objects from four domains: art, clipart, product and real. There are 65 classes in this classification dataset.
DomainNet [33] is large-scale dataset, consisting 586,575 images from 345 classes. These images from 6 domains: clipart, infograph, painting, quickdraw, real, sketch.
4.2 Experimental Setting
For all datasets, we perform "leave-one-domainnour" experiments, where we choose one domain as the target domain, train the model on all remaining domains, and evaluate it on the chosen domain. Each source domain is treated as a client. Following standard practice, we use 90% of available data as training data and 10% as validation data. We train all models with NVIDIA A100 GPUs from our AWS cluster.
For the RotatedMNIST dataset, we use network of two 3x3 convolutional layers and fully connected layer as the representation network go to get representation of 64 dimensions. A single linear layer is then used to map the representation to the ten output classes. This architecture is the same as the network used by [15, 31]. We train our network for 500 epochs with stochastic gradient descent (SGD), using learning rate of 0.001 and minibatch size 64, and report performance on the test domain after the last epoch. Each client performs local optimization iterations within each communication round (F = 5).
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]. We use ResNet50 [14] backbone for OfficeHome and DomainNet since they are more complex datasets. We replace the last fully connected layer of the backbone with linear layer of dimension 512 for ResNet18 and 2048 for ResNet50 to form the representation network. As with the RotatedMNIST experiment, we use a single layer to map from representation to output. 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 5e 5e Similar to the RotatedMNIST experiments, each client performs local optimization iterations within each communication round (E' = 5). However, we also found that method is not sensitive to this

---

Table 3: OfficeHomme. Reported numbers from runs
OfficeHome
Models
Backbone
C
P
R
Average
Centralized
Mixup [43]
Resnet50
64.7
54.7
77.3
79.2
69.0
Methods
CORAI [38]
Resnet50
64.4
55.3
76.7
77.9
68.6
FedAVG [28]
Resnet50
62.2+0.9
55.6+0.9
75.7+0.2
78.2=0.2
67.9
FedADG [44]
Resnet50
63,240.9
57.040.2
76.040.1
77.7=0.5
68.4
FL Methods
FedCMI (ours)
Resnet50
61.8+0.5
55.5+0.9
76.320.1
77.420.1
67.8
FedL2R (ours)
Resnet50
64.5+0.3
56.540.5
76.120.2
77.9+0.2
68.8
FedSR (ours)
Resnet50
65.4+0.5
57.4+0.2
76.2+0.6
78.3+0.3
69.3
Table 4: DomainNet. Reported numbers are from runs
DomainNet
Models
Backbone
C
R
S
AVG
Centralized
MLDG [21]
Resnet50
59.5
19.8
48.3
13.0
59.5
50.4
41.8
Methods
CORAL [38]
Resnet50
58.7
20.9
47.3
13.6
60.2
50.2
41.8
FedAVG [28]
Resnet50
59.3+0.7
16.5+0.9
44.2+0.7
10.8 8
57.240.8
49.8+0.4
39.6
FedADG [44]
Resnet50
60.9+0.6
17.2+0.2
44.3+0.2
12.4+0.2
57.6+0.9
50.3+0.8
40.4
FL Methods
FedCMT (ours)
Resnet50
59.0+0.9
18.0=0.7
44.6+0.5
12.2+0.4
56.2+0.2
50.0=0.4
40.0
FedL2R (ours)
Resnet50
60.2+0.6
18. 18.4
44.940.6
11.040.9
57.840.4
51.540.7
40.6
FedSR (ours)
Resnet50
61.020.6
18.640.4
45.240.5
13.420.6
57.640.2
51.840.3
41.3
hyper-parae prpee-parameter (performance with E 100 is almost identical). For the syper-parameters of our method (3GM7 and BI2R), we follow [11] and use random search to tune their value based only on the validation (no data leakage from the target domain). This protocol is highly recommended by [11] and allows for a fair comparison among the methods. For details of the tuned values, please refer our supplementary material. 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 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 http:://ittub..oo//aaaa http:///..oooouuuuuguuffe/////
4.3 Baselines
For all experiments. we consider the FL algorithms to be our main baselines: FedAVG [28] (a common FL algorithm) and FedADG [44] (an existing FL algorithm with DG ability). For the RotatedMNIST and PACS experiments, for comparison (the numbers are taken directly from the original papers). 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. Therefore, we also include the best performers for each dataset reported in the DomainBed paper. Note that we include these centralized methods for reference only, and they are not direct competitors of our model(s). For this reason, we highlight (in bold) the best performant model for the centralized and FL. settings separately.
also include SOTA centralized domain-invarian methods
4.4 Results
Tables 1, 2, and 4 show that our method (FedSR) outperforms the baseline FedAVG significantly (by margin of 1.5% to 3% on average) and consistently in all experiments. FedADG [44] (a existing method for DG in the FL setting) does outperform FedAVG but not yet on par with centralized methods. Meanwhile, FedSR achieves competitive results when compared to state-of-the- art centralized DG methods (as listed in the first half of the tables). This indicates the effectiveness of our method and the potential for Domain Generalization in the FL setting. As mentioned earlier,

---

our method can be interpreted as an implicit form of representation alignment. Please refer to our supplementary material for the visualization of our representation space, which does show the alignment effect of our method.
5 Conclusion
To conclude, in this paper, we present simple and effective method for implicit alignment of the representation across domains/clients in Federated Learning setting for better domain general- ization ability. We employ two regularization techniques, namely CMI and 1,2R, to learn a simple representation of the data in the hope for a better generalization. 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. Extensive experiments show that our method outperforms relevant FL baselines (FedAVG, FedADG) and can even achieve competitive DG performance when compared to centralized approaches. A potential limitation of our method 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). However, this is common comment for all domain invariant representation methods and not specific to ours. Furthermore, experiments show that method does lead to better representation alignment and prediction performance in practice.
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 Al Fellowship EP/W002981// 1/1. Meta author Ser-Nam Lim is neither supported by the UKRI grant nor has any relationship to the grant.
References
[1] H. Ajakan, P. Germain, H. Larochelle, F. Laviolette, and M. Marchand. Domain Domain-adversarr adversarial neural networks. arXiv preprint arXiv:1412.4446, 2014.
[2] K. Akuzawa, Y. Iwasawa, and Y. Matsuo. Adversarial invariant feature learning with accuracy constraint for domain generalization. In Joint European Conference on Machine Learning and Knowledge Discovery in Databases, pages 315-331. Springer, 2019.
[3] A. A. Alemi, I. Fischer, J. V. Dillon, and K. Murphy. Deep variational information bottleneck. arXiv preprint arXiv: 16/2.00410, 2016.
[4] M. G. Arivazhagan, V. Aggarwal, A. K. Singh, and S. Choudhary. Federated learning with personalization layers. arXiv preprint arXiv:1912.008/8, 2019.
[5] Y. Balaji, S. Sankaranarayanan, and R. Chellappa. Metareg: Towards domain generalization using metta-regularizztion. Advances in neural information processing systems, 31, 2018.
[6] Y. Deng, M. M. Kamani, and M. Mahdavi. Adaptive personalized federated learning, arXiv preprint Xiv:20033.13461 2020.
[7] Y. Du, Xu, H. Xiong, Q. Qiu, X. Zhen, C. G. Snoek, and L. Shao. Learning to learn with variational information bottleneck for domain generalization. In European Conference on Computer Vision, pages 200-216. Springer, 2020.
[8] A. Fallah, A. Mokhtari, and Ozdaglar. Personalized federated learning: A meta-learning approach. arXiv preprint arXiv:2002.077 079448088 2020.
[9] Y. Ganin, E. Ustinova, H. Ajakan, P. Germain, H. Larochelle, F. Laviolette, M. Marchand, and V. Lempitsky. Domaiinadversarial training of neural networks. The journal of machine learning research, 17(1):0066-2330 2016.
[10] M. Ghifary, W. B. Kleijn, M. Zhang, and D. Balduzzi. Domain generalization for object recognition with multi-task autoencoders. In Proceedings of the IEEE International Conference on Computer Vision pages 2551-2559, 2015.
[11] Gulrajani and D. Lopez-Paz. search of lost domain generalization. arXiv preprint drXiv:2007.0143..222 2020.
10

---

[12] F. Haddadpour and M. Mahdavi. On the convergence of local descent methods in federated learning. arXiv preprint arXiv:1910.1422,, 2019.
[13] Hanzely and Richtárik. Federated learning of a mixture of global and local models. arXiv preprint arXiv:2002.05516. 2020.
[14] K. He, X. Zhang, S. Ren, and J. Sun. Deep residual learning for image recognition. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pages 770-778, 2016.
[15] M. Ilse, J. M. Tomezak, C. Louizos, and M. Welling. Diva: Domain invariant variational autoencoders. In Medical Imaging with Deep Learning, pages 322-348. PMLR, 2020.
[16] Y. Jiang, J. Konečný, K. Rush, and S. Kannan. Improving federated learning personalization via model agnostic meta learning. arXiv preprint arXiv: 12488, 2019.
[17] A. Khaled, K. Mishchenko, and P. Richtárik. Tighter theory for local sgd on identical and heterogeneous data. In International Conference on Artificial Intelligence and Statistics, pages 4519-4529. PMLR, 2020.
[18] D. P. Kingma and M. Welling. Auto-encoding variational bayes. arXiv preprint 2013.
[19] Y. LeCun and C. Cortes. MNIST handwritten digit database. 2010. URL httht:/yammm URhhttp://amm..emm lecun com/exdb/ /xdb///nbb////...////////
[20] D. Li, Y. Yang, Y.-Z. Song. and Hospedales. Deeper, broader and artier domain generalization. In International Conference Computer Vision, 2017.
[21] D. Li, Y. Yang, Y.-Z. Song, and T. M. Hospedales. Learning to generalize: Meta-learning for domain generalization. In Thirty-Second AAAI Conference on Artificial Intelligence, 2018.
[22] H. Li, S. J. Pan, S. Wang, and A. C. Kot. Domain generalization with adversarial feature learning. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition. pages 5400-5409, 2018.
[23] T. Li, A. K. Sahu, M. Zaheer, M. Sanjabi, A. Talwalkar, and V. Smith. Federated optimization in heterogencous networks. Proceedings of Machine Learning and Systems, 2:429-450, 2020.
[24] X. Li, K. Huang, W. Yang. S. Wang, and Z. Zhang. On the convergence of fedavg on non-iid data. arXiv preprint arXiv: 1907.02288 189, 2019.
[25] Y. Li, M. Gong, X. Tian, Liu, and D. Tao. Domain generalization via conditional invariant representations. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 32, 2018.
[26] Y. Li, X. Tian, M. Gong, Y. Liu, T. Liu, K. Zhang, and D. Tao. Deep domain generalization via conditional invariant adversarial networks. In Proceedings of the European Conference Computer Vision (ECCV), pages 624,639, 2018
[27] Q Liu, C. Chen, J. Qin, Q Dou, and P.A. Heng. Feddg: Federated domain generalization on medical image segmentation via episodic learning continuous frequency space. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 1013-1023, 2021.
[28] B. McMahan, E. Moore, D. Ramage, S. Hampson, and B. A. y Arcas. Communication-efficcent learning of deep networks from decentralized data. In Artificial intelligence and statistics, pages 1282. PMLR, 2017.
[29] M. Mohri, G. Sivek, and A. T. Suresh. Agnostic federated learning. In International Conference on Machine Learning, pages 4615-4625. PMLR, 2019.
[30] K. Muandet, D. Balduzzi, and B. Schölkopf. Domain generalization via invariant feature representation. In International Conference on Machine Learning, pages 10-18. PMLR, 2013.
11

---

[31] A. T. Nguyen, T. Tran, Y. Gal, and A. G. Baydin. Domain invariant representation learning with domain density transformations. Advances Neural Information Processing Systems, 34, 2021.
[32] A. T. Nguyen, Tran, Y. Gal. P. H. Torr, and A. G. Baydin. Kl guided domain adaptation. arXiv preprint arXiv:2106.07788 2021.
[33] X. Peng, Q. Bai, X. Xia, Z. Huang, K. Saenko, and B. Wang. Moment matching for multi-source domain adaptation. In Proceedings of the IEEE/CVF international conference on computer vision, pages 1406-1415, 2019.
[34] X. Peng, Z. Huang, Y. Zhu, and Saenko. Federated adversarial domain adaptation. arXiv preprint arXiv: 1911 1.02054, 2019.
[35] J. Shen, Y. Qu, W. Zhang, and Y. Yu. Wasserstein distance guided representation learning for domain adaptation. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 32, 2018.
[36] V. Smith, C.-K. Chiang, M. Sanjabi, and A. S. Talwalkar. Federated multi-task learning. Advances in neural information processing systems, 30, 2017.
[37] S. U. Stich. Local sgd converges fast and communicates little. arXiv preprint arXiv:1805.0976) 2018.
[38] B. Sun and K. Saenko. Deep coral: Correlation alignment for deep domain adaptation. In European conference on computer vision, pages 143-450. Springer, 2016.
[39] N. Tishby, F. C. Pereira, and W. Bialek. The information bottleneck method. arXiv preprint physies/0004057, 2000.
[40] H. Venkateswara, J. Eusebio, S. Chakraborty, and S. Panchanathan. Deep hashing network for unsupervised domain adaptation. Proceedings of the IEFE Conference on Computer Vision and Pattern Recognition, pages 5018-5027, 2017.
[41] J. Wang and G. Joshi. Cooperative sgd: A unified framework for the design and analysis of communication-cfiiei.
sgd algorithms. arXiv preprint arXiv:1808.07576, 2018.
[42] Z. Wang, M. Loog, and J. van Gemert. Respecting domain relations: Hypothesis invariance for domain generalization. arXiv preprint arXiv:2010.0759/ 2020.
[43] M. Xu, J. Zhang. B. Ni, Li, C. Wang, Q. Tian, and W. Zhang. Adversarial domain adaptation with domain mixup. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 34, pages 6502-6509, 2020.
[44] L. Zhang, X. Lei, Y. Shi, H. Huang, and C. Chen. Federated learning with domain generalization. arXiv preprint 11 2021.
[45] C. Zhao and W. Shen. Adversarial mutual information-guider single domain generalization network for intelligent fault diagnosis. IEEE Transactions on Industrial Informatics, 2022.
[46] S. Zhao, M. Gong, T. Liu, H. Fu, and D. Tao. Domain generalization via entropy regularization. Advances Neural Information Processing Systems, 33, 2020.
[47] Y. Zhao, M. Li, L. Lai, N. Suda, D. Civin, and V. Chandra. Federated learning with non-iid data. arXiv preprint 1rXiv:1806.om 00582, 2018.
Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope? [Yes]
(b) Did you describe the limitations of your work? [Yes] see Section
(c) Did you discuss any potential negative societal impacts of your work? [N/A]
12

---

(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [Yes]
(b) Did you include complete proofs of all theoretical results? [Yes] Supplementary File
3. 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 URL)? [Yes]
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Yes] see Section
(c) Did you report error bars (e.g., with respect to the random seed after running experi- ments multiple times)? [Yes]
(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)? [Yes] See Section 4
4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators? [Yes]
(b) Did you mention the license of the assets? [N/A]
(c) Did you include any new assets either in the supplemental material or as URL? [Yes]
We include our source code in the Supplementary File
(d) Did you discuss whether and how consent was obtained from people whose data you're 're
using/curating? [N/A]
(e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [N/A]
5. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, applicable? [N/A]
(b) Did you deseribe any potential participant risks, with links to Institutional Review Board (IRB) approvals, applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation? [N/A]
13


Reviews:
[{"review_id": "UHrtC8nXgUA", "reviewer": "Reviewer_BRsg", "summary": "In the submission, a feature representation learning method was proposed to tackle the problem of federated learning based on the out-of-distribution classification setting. Specifically, two regularization terms were introduced, i.e., L2-norm and conditional mutual information regularization term. Experimental results are performed on 4 different benchmark datasets in the community of domain generalization, and some desired performances are reported. ", "questions": "Followed by my first concern, since the main contribution is feature representation learning, therefore, it is necessary to discuss and compare with the DG baseline methods  which does not require domain alignment across domain, and can also be applied based on the FL setting.  ", "limitations": "The limitation of the proposed method has been discussed. ", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strength:\n\n- The problem of federated learning based on OOD setting is interesting and sounding. \n- The proposed method is simple and effective. \n\nWeakness:\n\n- While the proposed method is simple, IMHO, it does not shed light on federated learning. Noted that feature representation regularization is actually not novel in the community of DG (e.g., [A]),  these methods can also be applied to the individual domain followed by FedAvg. \n\n[A] Adversarial Mutual Information-Guided Single Domain Generalization Network for Intelligent Fault Diagnosis\n\n\n- FL based on heterogeneous environment (non iid) has been largely explored. It's interesting to see how these methods perform based on the setting DG. \n\n- I suggest considering more realistic datasets (e.g., medical imaging), which require privacy guarantee for FL setting, for evaluation.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "Qqv-YJRqpN9", "reviewer": "Reviewer_AUa4", "summary": "This work proposes FedSR, a framework for learning (s)imple (r)epresentation in the (fed)erated setting so that the model can be applied to unseen domains/clients (in the sense of domain generalization). The representation is assumed to be one of the hidden layers in the neural network and is learned by minimizing the sum of (a) task prediction loss (2) conditional mutual information (CMI) and (3) L2 regularization (L2R) on the representation (see Eq.(9)). By theoretically justifying CMI and L2R terms, this work shows that together the objective function can lead to good generalization performance as demonstrated in common domain generalization benchmarks.", "questions": "The proposed FedSR provides a solution for domain generalization in the federated learning setting, where each client can only see data on their own data thus comparing data across clients is infeasible. The paper is well-organized and the presentation is clear. Each term in the objective function is well-motivated and their connections to common domain generalization techniques are also interesting. There are a few comments/questions\n\n1. It may be better to explicitly mention that z is assumed to be conditional independent of y given x before Eq.(4)\n\n2. For the experiment, are the hidden dimensions for z independent with separate variances? Or a covariance matrix is learned?\n\n3. The experiments are only using a handful of clients (since each of them corresponds to one domain). However, it is common in federated learning to have tens to hundreds of clients, which the current experiments fail to demonstrate. In this sense, it would be more convincing if the clients are indeed sampled from an underlying distribution as depicted in Eq.(2) and (3). Then (2) and (3) will be important metrics for the baselines and FedSR.\n\n4. It will be interesting to see what the representations of FedSR look like. Currently, the experiments only show the performance advantage in terms of accuracy, but we do not know why we have such an advantage. Maybe looking at the quality of representations can help, such as the CMI and L2R on the hold-out domain, or visualization of the representations.\n\nMinor points\n\n- LHS of (23) misses a p in the log\n- L269 the variational distribution can match (what?)", "limitations": "The authors have sufficiently addressed the limitations.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths\n- The problem is important yet under-explored in the literature\n- The writing and presentation are clear and easy to follow\n- Good empirical performance\n\nWeaknesses\n- The experiment setup is less typical for federated learning", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "OLGeckh6i1", "reviewer": "Reviewer_dm5i", "summary": "Motivated by the need for domain generalization in a federated learning setup (where each client has a different data distribution or domain), this paper proposes two simple regularizers based on the most common L2 norm squared regularizer and (an upper bound of) conditional mutual information (CMI), i.e., $I(x,z|y)$.\nThe L2 norm regularizer is shown to be related to minimizing the KL divergence between a fixed Gaussian distribution and the prior distributions at each client.\nThis can be seen as pushing all the prior distributions to the same *fixed* distribution.\nThus, this encourages marginal alignment of the latent distributions similar to AlignFlow.\nThe CMI regularizer is approximated via a variational approximation similar to the ELBO objective.\nBecause the model is shared, this encourages all client-specific conditional distributions $p_i(z|y)$ to match the global $r(z|y)$---thus making a connection to conditional alignment.\nFinally, the paper gives empirical DG results on multiple datasets showing that both regularizers are helpful.\n", "questions": "- How is (11) related to the ELBO?  I expect this to be very close in terms of derivation using an auxiliary variational distribution r. Is the main difference that you are trying to match by conditioning on diferent things (i.e., $p$ conditions on $x$ but $r$ conditions on $y$)?\n\n- Why discuss CMI first?  I would suggest putting L2R first.\n\n- How does this relate to AlignFlow which aligns distributions by pushing them both towards a fixed standard normal distribution?\n\n[AlignFlow] Grover, A., Chute, C., Shu, R., Cao, Z., & Ermon, S. (2020, April). Alignflow: Cycle consistent learning from multiple domains via normalizing flows. In Proceedings of the AAAI Conference on Artificial Intelligence (Vol. 34, No. 04, pp. 4028-4035).\n\n", "limitations": "- The paper mentions several limitations.\n", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n- Proposes a simple algorithm that merely uses two locally-computable regularizers for domain generalization in a federated learning setting.\n- Demonstrates insightful connections of these regularizers to marginal and conditional distribution alignment.\n- Demonstrates the empirical effectiveness of these simple FL objectives for domain generalization on multiple datasets.\n- Very easy to read paper that answers natural questions throughout the paper (e.g., Line 249-255, which I may suggest putting a little earlier or mentioning it a little earlier).\n\nWeaknesses:\n- It would be nice to evaluate these methods on real-world distribution shifts as in the WILDS benchmark dataset rather than only simulated or pseudo-realistic distribution shifts.\n\n- An exploration of the number of optimization steps before synchronization (i.e., steps per round) would improve the paper by showing how this key FL parameter affects the results.  Currently, the synchronization is set at E = 5 (i.e., I think this means sync every 5 minibatches), which seems quite often unless this means 5 epochs. What is the performance if the models are synchronized every epoch, every 2 epochs or every half epoch?  Or what if E=10, 50, 100, 500?\n\n[WILDS] Koh, P. W., Sagawa, S., Marklund, H., Xie, S. M., Zhang, M., Balsubramani, A., ... & Liang, P. (2021, July). Wilds: A benchmark of in-the-wild distribution shifts. In International Conference on Machine Learning (pp. 5637-5664). PMLR.\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "5jUF2cz1OyD", "reviewer": "Reviewer_FDaC", "summary": "This paper tackles the problem of domain generalization in the federated learning setting. Additional to optimizing the average of the local objective functions, this paper tries to learn the feature representation that is invariant with data distribution by optimizing a conditional mutual information $\\ell^{CMI}$. An $\\ell_2$-norm regularize is added to the loss function for the feature representation to regulate the feature complexity. The proposed algorithm is evaluated on several datasets to show its strength in dealing with data from unseen domains. The algorithm is compared with FedAvg and FedDGand a few centralized methods and has comparable performance as centralized algorithms.", "questions": "Q1: Could the authors include more discussions on existing works in domain generalization for federated learning and compare the proposed algorithm with the existing methods (e.g., [1.1], [1.2])\n\nQ2: Can the author include an algorithm block demonstrating how $\\ell^{CMI}_i$ is updated.\n\nQ3: It is unclear how the data are separated as training/testing data. It is using one domain as testing data and the other domains as training data, or using 90% of the data from all domains as training and 10% of the data from all domains as testing data?", "limitations": "This paper has no negative social impact.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Originality: This paper fails to include several important prior works in federated domain generalization, including [1,1], and [1,2], which are closely related to the proposed method in this paper. In specific, the idea of learning the feature representation with norm regularize has been used in [1.2]. The domain alignment (domain invariant representation) idea has also been used in [1.1]. Therefore, this paper seems to be incremental to the prior works.\n\nQuality: The paper is well organized. It has theoretical justifications for the regularizers and made a clear connection to domain invariant representation learning. However, the analysis is still limited as no generalization or performance bound of the proposed loss is provided. \n\nClarity: The presentation of the paper is clear. The notations and concepts have been clearly defined and explained. Theoretical results are fully discussed and the numerical results are thoroughly explained.\n\nSignificance: This paper tackles an important issue in federated learning of the domain shift between the training and deployment phase. The numerical results show a promising improvement in different datasets.\n\n[1.1] Peng, X., Huang, Z., Zhu, Y., & Saenko, K. (2019, September). Federated Adversarial Domain Adaptation. In International Conference on Learning Representations.\n\n[1.2] Mohri, M., Sivek, G., & Suresh, A. T. (2019, May). Agnostic federated learning. In International Conference on Machine Learning (pp. 4615-4625). PMLR.\n\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept