Title: StableFDG: Style and Attention Based Learning for Federated Domain Generalization
Venue: NeurIPS 2023
Accepted: yes
Keywords: Federated Learning

Domain Generalization
Abstract: Traditional federated learning (FL) algorithms operate under the assumption that the data distributions at training (source domains) and testing (target domain) are the same. The fact that domain shifts often occur in practice necessitates equipping FL methods with a domain generalization (DG) capability. However, existing DG algorithms face fundamental challenges in FL setups due to the lack of samples/domains in each client’s local dataset. In this paper, we propose StableFDG, a style and attention based learning strategy for accomplishing federated domain generalization, introducing two key contributions. The first is style-based learning, which enables each client to explore novel styles beyond the original source domains in its local dataset, improving domain diversity based on the proposed style sharing, shifting, and exploration strategies. Our second contribution is an attention-based feature highlighter, which captures the similarities between the features of data samples in the same class, and emphasizes the important/common characteristics to better learn the domain-invariant characteristics of each class in data-poor FL scenarios. Experimental results show that StableFDG outperforms existing baselines on various DG benchmark datasets, demonstrating its efficacy.

StableFDG: Style and Attention Based Learning for Federated Domain Generalization

Abstract

Traditional federated learning (FL) algorithms operate under the assumption that the data distributions at training (source domains) and testing (target domain) are the same. The fact that domain shifts often occur in practice necessitates equipping FL methods with a domain generalization (DG) capability. However, existing DG algorithms face fundamental challenges in FL setups due to the lack of samples/domains in each client's local dataset. In this paper, we propose StableFDG, a style and attention based learning strategy for accomplishing federated domain generalization, introducing two key contributions. The first is style-based learning, which enables each client to explore novel styles beyond the original source domains in its local dataset, improving domain diversity based on the proposed style sharing, shifting, and exploration strategies. Our second contribution is an attention-based feature highlighter, which captures the similarities between the features of data samples in the same class, and emphasizes the important/common characteristics to better learn the domain-invariant characteristics of each class in data-poor FL scenarios. Experimental results show that StableFDG outperforms existing baselines on various DG benchmark datasets, demonstrating its efficacy.

1 Introduction

Federated learning (FL) has now become a key paradigm for training a machine learning model using local data of distributed clients [26, 19, 11]. Without directly sharing each client's data to the third party, FL enables the clients to construct a global model via collaboration. However, although FL has achieved remarkable success, the underlying assumption of previous works is that the data distributions during training and testing are the same. This assumption is not valid in various scenarios with domain shifts; for example, although the FL clients only have data samples that belong to the source domains (e.g., images on sunny days and rainy days), the trained global model should be also able to make reliable predictions for the unseen target domain (e.g., images on snowy days). Therefore, in practice, FL methods have to be equipped with a domain generalization (DG) capability.

Given the source domains in the training phase, DG aims to construct a model that has a generalization capability to predict well on an unseen target domain. Various DG methods have been proposed in a centralized setup [39, 21, 36, 38, 16, 15, 24, 23]. However, directly applying centralized DG schemes to FL can potentially restrict the model performance since each client has limited numbers of data samples and styles in its local dataset. The local models are unable to capture the domain-invariant characteristics due to lack of data/styles in individual clients.

Although several researchers have recently focused on DG for FL [25, 2, 35, 28], they still do not directly handle the fundamental issues that arise from the lack of data and styles (which represent domains) in individual FL clients. These performance limitations become especially prevalent when federating complex DG datasets having large style shifts between domains or having backgrounds unrelated to the prediction of class, as we will see in Sec. 4. Despite the practical significance of federated DG, this field is still in an early stage of research and remains a great challenge.

Contributions. In this paper, we propose StableFDG, a style and attention based learning strategy tailored to federated domain generalization. StableFDG tackles the fundamental challenges in federated DG that arise due to the lack of data/styles in each FL client, with two novel characteristics:

We first propose style-based learning , which exposes the model of each FL client to various styles beyond the source domains in its local dataset. Specifically, we (i) design a style-sharing method that can compensate for the missing styles in each client by sharing the style statistics with other clients; (ii) propose a style-shifting strategy that can select the best styles to be shifted to the new style to balance between the original and new styles; and (iii) develop styleexploration to further expose the model to a wider variety of styles by extrapolating the current styles. Based on these unique characteristics, our style-based learning handles the issue of the lack of styles in each FL client, significantly improving generalization capability.

We also propose an attention-based feature highlighter , which enables the model to focus only on the important/essential parts of the features when making the prediction. Our key contribution here is to utilize an attention module to capture the similarities between the features of data samples in the same class (regardless of the domain), and emphasize the important/common characteristics to better learn the domain-invariant features. Especially in data-poor FL scenarios where models are prone to overfitting to small local datasets, our new strategy provides advantages for complicated DG tasks by removing background noises that are unrelated to class prediction and focusing on the important parts.

The two suggested schemes work in a complementary fashion, each providing one necessary component for federated DG: our style-based learning improves domain diversity, while the attention-based feature highlighter learns domain-invariant characteristics of each class. Experiments on various FL setups using DG benchmarks confirm the advantage of StableFDG over (i) the baselines that directly apply DG methods to FL and (ii) the baselines that are specifically designed for federated DG.

2 Related Works

Federated learning. FL enables multiple clients to train a shared global model or personalized models without directly sharing each client's data with the server or other clients. FedAvg [26] is a well-known early work that sparked interest in FL in the machine learning community. Since then, various FL strategies have been proposed to handle the communication burden issue [31, 9], data heterogeneity issue [20, 12], adversarial attack issue [34, 29, 7], and personalization issue [3, 18, 5]. However, existing FL methods do not have generalization capabilities to predict well on an arbitrary unseen domain. In other words, most prior FL methods are not able to handle the DG problem.

Domain generalization. DGis one of the emerging fields in the AI community due to its significance in practical applications. Existing DG strategies based on domain alignment [17, 24, 23], metalearning [16, 4, 15, 37] and style/data-augmentation [39, 21, 36, 38] have shown great success in a centralized setup where the whole dataset is accessible during training. Recently, style-augmentation methods [39, 21, 36] including MixStyle [39] and DSU [21] are receiving considerable attention due to their high compatibility with various tasks and model architectures. However, although existing DG solutions work well in a centralized setup, they face challenges in FL scenarios; in data-poor FL setups, prior works achieve relatively low performance due to the lack of data samples or domains in each client, resulting in compromised generalization capabilities. The applications of meta-learning or domain alignment methods could be also limited when domain labels are not accessible in each client. Compared to these prior works focusing on a centralized DG setup, we develop a style and attention based DG strategy tailored to FL. The advantages of our methodology against these baselines are confirmed via experiments in Sec. 4.

Federated domain generalization. Only a few recent works [25, 2, 35, 28] have focused on the intersection of FL and DG. Based on the training set distributed across clients, the goal of these works is to construct a global model that is generalizable to an unseen target domain. In [2], the authors

proposed to share the style statistics across different clients that could be utilized during local updates. However, this method does not utilize the styles beyond the clients' domains and shows limited performance in specific datasets where the data samples are not well separated in the style space. Moreover, it increases the computation and memory costs for generating new images in each client using a pretrained style transfer. In our work, we handle these issues via style exploration and the attention-based feature highlighter to train the model with novel styles while capturing the important knowledge of each class. In [25], the authors proposed to exchange distribution information among clients based on Fourier transform, especially targeting image segmentation tasks for medical data. The authors of [35] proposed a strategy for federated DG based on the domain-invariant feature extractor and an ensemble of domain-specific classifiers. Two regularization losses are developed in [28] aiming to learn a simple representation of data during client-side local updates.

Although these prior works [25, 35, 28] improve DG performance, the authors do not directly handle the issues that arise from limited styles and data in each client. Compared to these works, we take an orthogonal approach based on style and attention based learning to effectively learn style-invariant features while capturing common knowledge of classes. Experimental results in Sec. 4 reveals that our scheme outperforms existing ideas tackling federated DG in practical data distribution setups.

3 Proposed StableFDG Algorithm

Problem formulation. We consider a FL setup with N clients distributed over the network. Let S n = { ( x n i , y n i ) } ρ n i =1 be the local dataset of the n -th client, which consists of ρ n pairs of data sample x and the corresponding label y . Here, each client n can have data samples from either a single or multiple source domains in its local dataset S n . Previous works on FL focus on constructing a global model that predicts well on the overall dataset S = {S 1 , S 2 , . . . , S N } or personalized models that work well on individual local datasets S n . In contrast to the conventional FL setup, given the overall dataset (or source domains) S , the goal of this work is to construct a shared global model w that has a generalization capability to predict well on any unseen target domain T .

Background. Let s ∈ R C × H × W be the feature of a sample which is obtained at a specific layer of the neural network. Here, C , H , W denote the dimensions of channel, height, width, respectively. Given the feature s of a specific data sample, the channel-wise mean µ ( s ) ∈ R C and the channel-wise standard deviation σ ( s ) ∈ R C can be written as

respectively. These variables are known as style statistics as they contain style information of an image in CNNs [10]. Based on these style statistics, various style-augmentation schemes such as MixStyle [39] or DSU [21] have been proposed in the literature targeting a centralized setup.

Overview of approach. Fig. 1 provides an overview of the problem setup and our StableFDG algorithm. As in conventional FL, the training process consists of multiple global rounds, which we index t = 1 , 2 , . . . , T . In the beginning of round t , a selected set of clients download the current global model w t from the server. Before local training begins, each client n computes its own style information Φ n = [ µ n , σ n , Σ n ( µ ) , Σ n ( σ )] using its local dataset according to (2), which will be clarified soon. This information is sent to the server, and the server shares these information with other clients to compensate for the lack of styles or domains in each client. During the local update process, each client selectively shifts the

Unseen target domain

Source domain 1

Source domain 2

Source domain 3

Style sharing

(Sec. 3.1)

Attention-based

weighted averaging

(Sec. 3.2)

Local update

Style shifting and

exploration (Sec. 3.1)

styles of the original data in the mini-batch to the new style (received from the server) via adaptive instance normalization (AdaIN) [10], to improve domain diversity (inner box in Fig. 2b). After this style sharing and shifting process, each client performs style exploration via feature-level oversampling to further expose the model to novel styles beyond the current source domains of each client (outer box in Fig. 2b). Finally, at the output of the feature extractor, we apply our attention-based

: Newly generated styles based on style augmentation

MixStyle

DSU

: Original styles in the client

(a) Existing style-based DG baselines

New styles

(from another client)

Selective style shift

Original styles

(in the client)

Style center

(

𝜇 " #

or

𝜎 " #

)

Style variance

(

Σ & #(𝜇)

or

Σ & #(𝜎))

: Samples selected via k-means++

Style sharing and shifting

: New styles generated via style exploration (oversampled features)

Style exploration

: Samples shifted to new styles

Overall style

center

,

(b) Proposed style-based learning for federated DG

feature highlighter to extract common/important feature information within each class and emphasize them for better generalization (Fig. 3). When local updates are finished, the server aggregates the client models and proceeds to the next round.

In the following, we first describe our style-based learning in Sec. 3.1, which determines how the style information is shared and utilized in each FL client, and how style exploration is conducted, to overcome the lack of styles in each client. In Sec. 3.2, we show how attention mechanism is utilized to capture the essential characteristics of each class for better generalization in data-poor FL setups.

3.1 Style-based Learning for Federated DG

The main components of our style-based learning are style-sharing, selective style-shifting, and style-exploration, each having its own role to handle the style limitation problem at each client. As illustrated in Fig. 2, this fundamental challenge of federated DG is not directly resolvable using existing style-augmentation DG methods that can only explore limited areas in the style-space.

Step 1: Style information sharing. Based on the data samples in its local dataset S n , at a specific layer of the model, each client n computes the average of channel-wise means, µ n , and the variance of the channel-wise means, Σ n ( µ ) , as follows:

where s n i ∈ R C × H × W is the feature of the i -th sample in the n -th client's local dataset, and the square operation ( · ) 2 works in an element-wise manner. Similarly, the average and the variance of channelwise standard deviations are computed as σ n = 1 ρ n ∑ ρ n i =1 σ ( s n i ) , Σ n ( σ ) = 1 ρ n ∑ ρ n i =1 ( σ n -σ ( s n i )) 2 , respectively. Here, µ n and σ n represent the center of the style statistics of client n , while Σ n ( µ ) and Σ n ( σ ) show how these style statistics of client n are distributed around the center. Now we define

̸

as the style information representing the domain identity of the n -th client in the style-space. Each client n sends Φ n to the server, and the server randomly shares the style information to other clients in a one-to-one manner; client n receives Φ n ′ that belongs to client n ′ ( n = n ′ ) without overlapping with other clients. By compensating for the lack of styles in each FL client, this style-sharing process is the first step that provides an opportunity for each model to get exposed to new styles (blue region in Fig. 2b) beyond the client's original styles (orange region in Fig. 2b).

Step 2: Selective style shifting. Suppose client n received Φ n ′ from the server. Now the question is, how should each client utilize this additional style information during training to improve domain/style diversity? Our idea is to selectively shift the styles of the samples from the original style to the

new style to effectively balance between the original/new source domains. To this end, given a mini-batch with size B , each client runs k -means++ with k = B/ 2 in the style-space for one iteration and selects B/ 2 cluster centers. This enables each client to choose the B/ 2 styles that are similar to the remaining B/ 2 styles, so that the model can get exposed to new styles while not losing the performance on the original styles. The selected B/ 2 samples keep their original styles, while for the remaining B/ 2 samples, we shift the style of their feature s to the new style via AdaIN [10] as f ( s ) = ( σ n ′ + ϵ σ Σ n ′ ( σ )) ( s -µ ( s ) σ ( s ) ) +( µ n ′ + ϵ µ Σ n ′ ( µ )) , where f ( s ) is the new feature shifted from s and ϵ µ ∼ N (0 , Σ n ( µ )) , ϵ σ ∼ N (0 , Σ n ( σ )) are the values sampled from normal distributions. Then, the mini-batch applied with new styles in Φ n ′ is forwarded to the next layer. The inner box in Fig. 2b shows how style shifting is performed in client n based on the new style information Φ n ′ .

Overall, based on the shared style statistics in Step 1, our style shifting in Step 2 balances between the original source domain and the new source domain via k -means++ for better generalization, which cannot be achieved by previous methods in Fig. 2a that rely on interpolation or style-shift near the original style. In the following, we describe our style exploration that can further resolve the style-limitation problem based on feature-level oversampling and extrapolation.

Step 3: Feature-level oversampling. Let s n ∈ R B × C × H × W be a mini-batch of features in client n at a specific layer, obtained after Steps 1 and 2 above. Here, we oversample the features by the minibatch size B in a class-balanced manner; the samples belonging to minority classes are compensated as balanced as possible up to size B . For example, suppose that the number of samples for classes a , b , c are 3, 2, 1, respectively in the mini-batch. Using oversampling size of 6 , we oversample by 1, 2, 3 data points for classes a , b , c , respectively, to balance the mini-batch in terms of classes. Based on this, we obtain a new oversampled mini-batch ˜ s n with size B , and concatenate it with the original mini-batch as follows: ˆ s n = [ s n , ˜ s n ] . This not only mitigates the class-imbalance issue in each client but also provides a good platform for better style exploration; the oversampled features are utilized to explore a wider style-space beyond the current source domains, as we will describe in Step 4. The oversampling size can be adjusted depending on the clients' computation/memory constraints.

Step 4: Style exploration. In order to further expose the model to a wider variety of styles, we transfer the styles of tensors in ˜ s n (i.e., the set of oversampled features) to novel styles beyond the style of each client. Let ˜ s n i be the feature of i -th sample in the oversampled mini-batch ˜ s n . We obtain the new styles by extrapolating the original styles in ˜ s n around the average of channel-wise mean µ n (ˆ s n ) and the average of standard deviations σ n (ˆ s n ) computed on the concatenated mini-batch ˆ s n = [ s n , ˜ s n ] as

where α is the exploration level . We perform AdaIN to shift the style of ˜ s n i to the new style statistics µ new (˜ s n i ) and σ new (˜ s n i ) . If α = 0 , the styles remain unchanged, and as α increases, the styles are shifted farther from the center. The outer box in Fig. 2b describes the concept of our style exploration.

Step 5: Style augmentation. After style exploration, we can apply existing style-augmentation methods during training. In this work, we mix the style statistics of the entire samples in ˆ s n to generate diverse domains as in [39] as µ new (ˆ s n i ) = λ · µ (ˆ s n i ) + (1 -λ ) · µ (ˆ s n j ) and σ new (ˆ s n i ) = λ · σ (ˆ s n i ) + (1 -λ ) · σ (ˆ s n j ) , where ˆ s n i and ˆ s n j are arbitrary two samples in ˆ s n and λ is a mixing parameter sampled from the beta distribution. Below, we wish to highlight two important points.

Remark 1 (Privacy). It is already well-known that there is an inherent clustering of samples based on their domains in the style-space, regardless of their labels [39]. This indicates that label information is not contained in the style statistics, resolving privacy issues. Note that some prior works on federated DG [25, 2] also adopt sharing style information between clients, but in different ways (see Sec. 2).

Remark 2. By enabling the model to explore a wider region in the style space based on the exploration level α , our style exploration in Step 4 is especially beneficial when the target domain is significantly far from the source domains (e.g., Sketch domain in PACS dataset). Existing style-based DG methods (e.g., MixStyle, DSU) are ineffective in this case as they can explore only some limited areas near the original styles in each client (Fig. 2a vs. Fig. 2b), which leads to performance degradation especially in data-poor FL scenarios. It turns out that our scheme has significant performance improvements over these baselines even with a rather arbitrarily chosen α , as we will see in Sec. 4.3.

3.2 Attention-based Feature Highlighter

In this subsection, we describe the second component of our solution, the attention-based feature highlighter, which operates at the output of the feature extractor to tackle the remaining challenge of federated DG: limited generalization in each client due to the lack of data. To handle this issue, we start from our key intuition that the images from the same class have inherently common characteristics regardless of domains to which they belong (e.g., Fig. 3). Based on this insight, we take advantage of attention to find the essential characteristics of images in the same class and emphasize them, which turns out to be very effective in data-poor FL scenarios.

̸

Attention score. Consider the i -th data sample of client n . Given the feature z n i ∈ R C × H × W obtained from the output of the feature extractor, we flatten it to a two-dimensional tensorop X i ∈ R C × HW with a size of ( C, HW ) , where we omit the client index n for notational simplicity. Now consider another X j ( j = i ) in a mini-batch that belongs to the same class as X i , where the domains of X i and X j can be either same or different. Inspired by the attention mechanism [32], we measure the spatial similarity S ∈ R HW × HW between X i and X j as follows:

where θ q ∈ R d × C , θ k ∈ R d × C are the learnable parameters trained to extract important information in each class from the given samples and d is the embedding size of queries Q = θ q X j and keys K = θ k X i . Then, we reshape S ∈ R HW × HW to S r ∈ R HW × H × W ; the ( m,n ) -th spatial feature ( ∈ R HW × 1 × 1 ) of S r represents the similarity between the ( m,n ) -th spatial feature of the key feature X i and the overall spatial feature of the query feature X j . Then by taking the mean of S r along the first dimension, we obtain the attention score a s , which represents how the spatial feature at each location of key X i is similar to the overall features of query X j : a s = mean( S r ) ∈ R H × W . A higher score a s indicates a higher similarity. Finally, we normalize the attention score through softmax function so that the total sum is one: a s ← softmax( a s ) ∈ R H × W .

Attention

Module

Common/important part

from the same class

Classifier

Query sample

𝒋

(Class: alarm clock)

Key sample

𝒊

(Class: alarm clock)

Training loss

Pooling

Important parts

emphasized

Style shifting and

exploration (Sec. 3.1)

𝑧 $

%

𝑎'

𝐴(𝑧 $

% )

Concatenation

Feature

Extractor

Weighted

Averaging

𝑧 +

%

Feature

Extractor

Attention-based weighted averaging. Based on the attention score, we take the weighted average of z i ∈ R C × H × W to generate an attention feature A ( z i ) ∈ R C in which the important parts of the key image , having common characteristics with the query image (in the same class), are emphasized:

As shown in Fig. 3, this attention feature is concatenated with the last embedding feature (e.g., after the global average pooling in ResNet) before the classifier, and it goes through the classifier to compute the loss during training. One important thing to note is that θ q and θ k are trained so that the common features of query and key images become close in the embedding space while unimportant factors such as backgrounds are effectively distanced (i.e., less emphasized).

When implementing attention-based weighted averaging in practice, instead of directly adopting equation (6), we modify the similarity function using the query of its own as

to avoid performance degradation when there is little commonality between key and query images. In other words, StableFDG takes advantage of both cross-attention and self-attention, enabling the model to extract and learn important characteristics across images (via cross-attention), and within the image (via self-attention). A more detailed analysis on (8) can be found in Appendix.

Inference. During testing, given a new test sample, we compute the spatial similarity of the test sample itself as Sim( X i , X i ) based on self-attention, take the weighted average, and concatenate the features to make a prediction.

Remark 3. It is interesting to note that we can adopt self-attention during testing. This is because θ q and θ k are trained to capture the essential characteristics of each class, which enables the attention module to effectively find the important parts of individual samples via self-attention (e.g., see heat map visualization in Fig. 3). Moreover, with only 0.44% of additional model parameters for the attention module, this strategy turns out to be very effective in learning domain-invariant characteristics of classes in data-poor FL setups, handling the key challenge of federated DG.

3.3 StableFDG

Finally, we put together StableFDG. In each FL round, the clients first download the global model from the server and perform style sharing, shifting, and exploration according to 3.1, which are done in the early layers of CNNs where the style information is preserved. Then, at the output of the feature extractor, attention-based weighted averaging is applied according to Sec. 3.2. These two components have their own roles and work in a complementary fashion to handle the challenging DG problem in FL; our style-based strategy is effective in improving the domain diversity, while our attention-based method can directly capture the domain-invariant characteristics of each class. After the local update process, the server aggregates the client models and proceeds to the next round.

Remark 4 (Computational complexity). The computational complexity of StableFDG depends on the oversampling size in Sec. 3.1 and the attention module size in Sec. 3.2, which could be controlled depending on the resource constraints of clients. We show in Sec. 4 that StableFDG achieves the state-of-the-art performance with (i) minimal oversampling size and (ii) negligible cost of attention module. A more detailed discussion on the computational complexity is in Appendix.

4 Experimental Results

4.1 Experimental Setup

Datasets. We consider five datasets commonly adopted in DG literature: PACS [14], VLCS [6], Digits-DG [38], Office-Home [33], and DomainNet [30]. PACS consists of 7 classes from 4 different domains, while VLCS contains 4 domains with 5 classes. Digits-DG is composed of 4 different digit datasets, MNIST [13], MNIST-M [8], SVHN [27], SYN [8], each corresponding to a single domain. Office-Home consists of 65 classes from 4 domains, while DomainNet has 345 classes from 6 domains. The DomainNet results are reported in Appendix.

Data partitioning for FL. When evaluating the model performance, we follow the conventional leave-one-domain-out protocol where one domain is selected as a target and the remaining domains are utilized as sources. Compared to the centralized setup, in FL, the source domains are distributed across the clients. We consider a setup with N = 30 clients and distribute the training set into two different ways: single-domain data distribution and multi-domain data distribution scenarios. In a single-domain setup, we let each client to have training data that belong to a single source domain. Since there are three different source domains during training (except DomainNet), the training data of each domain is distributed across 10 clients uniformly at random. In a multi-domain distribution setup, each client can have multiple domains, but the domain distribution within each client is heterogeneous. For each domain, we sample the heterogeneous proportion from Dirichlet distribution with dimension N = 30 and parameter of 0.5, and distribute the train samples of each domain to individual clients according to the sampled proportion. In Appendix, we also report the results with N = 3 following the settings of prior works in federated DG [2, 28].

Implementation. Following [39], we utilize ResNet-18 pretrained on ImageNet as a backbone while the results on ResNet-50 are reported in Sec. 4.3. The exploration level α is set to 3 for all experiments regardless of datasets. For our attention module, we set the embedding size d of queries Q and keys K to 30, where Q and K matrices are extracted from the output of the last residual block using 1 × 1 convolution such that the channel size of each output is 30. When the attention module is applied, the input dimension of the classifier becomes twice as large since we concatenate the attention feature with the last embedding feature before the classifier. The number of additional model parameters for the attention module is only 0.44% of the entire model. FL is performed for 50 global rounds and we trained the local model for 5 epochs with a mini-batch size of 32. Among a

total of N = 30 clients, 10 clients participate in each global round. All reported results are averaged over three random seeds.

Where to apply style-based modules. Inspired by [39, 21], we apply our style-based modules with a probability of 0.5 at specific layers. In particular, style sharing and shifting are executed at the first residual block of ResNet with a probability 0.5, while the style exploration module is performed at the first or second or third residual blocks independently with probability 0.5 after style sharing/shifting.

Baselines. (i) FL baselines: First, we consider FedAvg [26] and FedBN [22], which are the basic FL baselines not specific to DG. (ii) DG baselines applied to each FL client: We apply MixStyle [39] during the local update process of each client and aggregate the model via FedAvg. Similarly, we also apply DSU [21] at each client and then perform FedAvg to compare with our work. (iii) Federated DG baselines: Among DG schemes tailored to FL, we implement the recently proposed CCST [2], FedDG [25], FedSR [28] and evaluate the performance. For a fair comparison, we reproduced all the baselines in accordance with our experimental setup. The augmentation level in CCST, a hyperparameter that controls the amount of augmented images, is set to 3 as in the original paper [2]. The hyperparameters in FedSR, that control the regularization losses, are tuned to achieve the best performance in our setup. Except for these, we adopted the parameter values in the original papers.

4.2 Main Experimental Results

Single-domain data distribution. Table 1 shows our results in a single-domain data distribution setup. We have the following observations. Compared to the previous results provided in the centralized DG works [39, 21], the performance of each method is generally lower. This is due to the limited numbers of styles and data samples in each FL client, which restricts the generalization performance of individual client models. It can be seen that most of the baselines perform better than FedAvg and FedBN that do not tackle the DG problem. The proposed StableFDG achieves the best average accuracy for all benchmark datasets, where the gain is especially large in PACS having large shifts between domains. In contrast to our scheme, the prior works [2, 25, 28] targeting federated DG show marginal performance gains relative to FedAvg in our practical experimental setup with (i) more clients (which results in less data in each client) and (ii) partial client participations.

Multi-domain data distribution. In Table 2, we report the results in a multi-domain data distribution scenario. Compared to the results in Table 1, most of the schemes achieve improved performance in Table 2. This is because each client has multiple source domains, and thus providing a better platform for each client model to gain generalization ability. The proposed StableFDG still performs the best, demonstrating the effectiveness of our style and attention based learning strategy for federated DG.

(a) Office-Home and VLCS datasets.

4.3 Ablation Studies and Discussions

Effect of each component. To see the effect of each component of StableFDG, in Table 3, we apply our style-based learning and attention-based learning one-byone in a multi-domain data distribution setup using PACS. We compare our results with style-augmentation DG base-

lines, MixStyle [39] and DSU [21]. By applying only our style-based learning, StableFDG already outperforms prior style-augmentation methods. Furthermore, by adopting only one of the proposed components, our scheme performs better than all the baselines in Table 2. Additional ablation studies using other datasets are reported in Appendix.

Effect of hyperparameters. In DG setups, it is generally impractical to tune the hyperparameter using the target domain, because there is no information on the target domain during training. Hence, we used a fixed exploration level α = 3 throughout all experiments without tuning. In Fig. 4, we observe how the hyperparameters affect the target domain performance on PACS. In the first plot of Fig. 4, if α is too small, the performance is relatively low since

0

2

4

6

Exploration level

76

78

80

82

84

Average accuracy

StableFDG

Mixstyle

0

10

20

30

Oversampling size

76

78

80

82

84

Average accuracy

StableFDG

Mixstyle

the model is not able to explore novel styles beyond the client's source domains. If α is too large, the performance could be slightly degraded because the model would explore too many redundant styles. The overall results show that StableFDG still performs better than the baselines with an arbitrarily chosen α , which is a significant advantage of our scheme in the DG setup where hyperparameter tuning is challenging. The second plot of Fig. 4 shows how the oversampling size (introduced in Step 3 of Sec. 3.1) affects the DG performance. StableFDG still outperforms the baseline with minimal oversampling size, indicating that other components of our solution (style sharing/shifting and attention-based components) are already strong enough. The size of oversampling can be determined depending on the clients' computation/memory constraints, with the cost of improved generalization.

Performance in a centralized setup. Although our scheme is tailored to federated DG, the ideas of StableFDG can be also utilized in a centralized setup. In Table 4a, we study the effects of our style and attention based strategies in a centralized DG setting using PACS, while the other settings are

(a) Performance in a centralized DG setup.

(b) Performance using ResNet-50.

the same as in the FL setup. The results demonstrate that the proposed ideas are not only specific to data-poor FL scenarios but also have potentials to be utilized in centralized DG settings.

Performance with ResNet-50. In Table 4b, we also conduct experiments using ResNet-50 on PACS dataset in the multi-domain data distribution scenario. Other settings are exactly the same as in Table 2. The results further confirm the advantage of StableFDG with larger models.

Attention score visualization. To gain an intuitive understanding of the effect of our attention-based learning, in Fig. 5, we visualize the score maps obtained via our attention module at testing. The score maps are interpolated so that it has the same size as the original image. A warmer color indicates a higher value. It can be seen that our attention

module highlights important parts of each class even in the presence of unrelated backgrounds.

Additional results. Other implementation details, comprehensive ablation studies for each component, discussions on complexity, and results on DomainNet dataset are in Appendix.

5 Conclusion

Despite the practical significance, the field of federated domain generalization is still in the early stage of research. In this paper, we proposed StableFDG, a new training strategy tailored to this unexplored area. Our style-based strategy enables the model to get exposed to various novel styles beyond each client's source domains, while our attention-based method captures and emphasizes the important/common characteristics of each class. Extensive experimental results confirmed the advantage of our StableFDG for federated domain generalization with data-poor FL clients.

Limitations and future works. StableFDG requires 0.45 % more communication load compared to FedAvg for sharing the attention module and style statistics, which is the cost for a better DG performance. Further developing our idea to tailor to centralized DG and extending our attention strategy to segmentation/detection DG tasks are also interesting directions for future research.

Acknowledgments

This work was supported by the National Research Foundation of Korea (NRF) grant funded by the Korea government (MSIT) (No. NRF-2019R1I1A2A02061135), by IITP funds from MSIT of Korea (No. 2020-0-00626), by the National Science Foundation (NSF) under grants CNS-2146171 and CPS-2313109, and by the Defense Advanced Research Projects Agency (DARPA) under grant D22AP00168-00.

References

Irwan Bello, Barret Zoph, Ashish Vaswani, Jonathon Shlens, and Quoc V Le. Attention augmented convolutional networks. In Proceedings of the IEEE/CVF international conference on computer vision , pages 3286-3295, 2019.

Junming Chen, Meirui Jiang, Qi Dou, and Qifeng Chen. Federated domain generalization for image recognition via cross-client style transfer. In Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision , pages 361-370, 2023.

Yuyang Deng, Mohammad Mahdi Kamani, and Mehrdad Mahdavi. Adaptive personalized federated learning. arXiv preprint arXiv:2003.13461 , 2020.

Yingjun Du, Jun Xu, Huan Xiong, Qiang Qiu, Xiantong Zhen, Cees GM Snoek, and Ling Shao. Learning to learn with variational information bottleneck for domain generalization. In European Conference on Computer Vision , pages 200-216. Springer, 2020.

Alireza Fallah, Aryan Mokhtari, and Asuman Ozdaglar. Personalized federated learning with theoretical guarantees: A model-agnostic meta-learning approach. Advances in Neural Information Processing Systems , 33:3557-3568, 2020.

Chen Fang, Ye Xu, and Daniel N Rockmore. Unbiased metric learning: On the utilization of multiple datasets and web images for softening bias. In Proceedings of the IEEE International Conference on Computer Vision , pages 1657-1664, 2013.

Minghong Fang, Xiaoyu Cao, Jinyuan Jia, and Neil Gong. Local model poisoning attacks to { Byzantine-Robust } federated learning. In 29th USENIX Security Symposium (USENIX Security 20) , pages 1605-1622, 2020.

Yaroslav Ganin and Victor Lempitsky. Unsupervised domain adaptation by backpropagation. In International conference on machine learning , pages 1180-1189. PMLR, 2015.

Jenny Hamer, Mehryar Mohri, and Ananda Theertha Suresh. Fedboost: A communicationefficient algorithm for federated learning. In International Conference on Machine Learning , pages 3973-3983. PMLR, 2020.

Xun Huang and Serge Belongie. Arbitrary style transfer in real-time with adaptive instance normalization. In Proceedings of the IEEE international conference on computer vision , pages 1501-1510, 2017.

Peter Kairouz, H Brendan McMahan, Brendan Avent, Aurélien Bellet, Mehdi Bennis, Arjun Nitin Bhagoji, Kallista Bonawitz, Zachary Charles, Graham Cormode, Rachel Cummings, et al. Advances and open problems in federated learning. Foundations and Trends® in Machine Learning , 14(1-2):1-210, 2021.

Sai Praneeth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank Reddi, Sebastian Stich, and Ananda Theertha Suresh. Scaffold: Stochastic controlled averaging for federated learning. In International Conference on Machine Learning , pages 5132-5143. PMLR, 2020.

Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient-based learning applied to document recognition. Proceedings of the IEEE , 86(11):2278-2324, 1998.

Da Li, Yongxin Yang, Yi-Zhe Song, and Timothy M Hospedales. Deeper, broader and artier domain generalization. In Proceedings of the IEEE international conference on computer vision , pages 5542-5550, 2017.

Da Li, Yongxin Yang, Yi-Zhe Song, and Timothy M Hospedales. Learning to generalize: Meta-learning for domain generalization. In Thirty-Second AAAI Conference on Artificial Intelligence , 2018.

Da Li, Jianshu Zhang, Yongxin Yang, Cong Liu, Yi-Zhe Song, and Timothy M Hospedales. Episodic training for domain generalization. In Proceedings of the IEEE/CVF International Conference on Computer Vision , pages 1446-1455, 2019.

Haoliang Li, Sinno Jialin Pan, Shiqi Wang, and Alex C Kot. Domain generalization with adversarial feature learning. In Proceedings of the IEEE conference on computer vision and pattern recognition , pages 5400-5409, 2018.

Tian Li, Shengyuan Hu, Ahmad Beirami, and Virginia Smith. Ditto: Fair and robust federated learning through personalization. In International Conference on Machine Learning , pages 6357-6368. PMLR, 2021.

Tian Li, Anit Kumar Sahu, Ameet Talwalkar, and Virginia Smith. Federated learning: Challenges, methods, and future directions. IEEE Signal Processing Magazine , 37(3):50-60, 2020.

Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith. Federated optimization in heterogeneous networks. Proceedings of Machine Learning and Systems , 2:429-450, 2020.

Xiaotong Li, Yongxing Dai, Yixiao Ge, Jun Liu, Ying Shan, and Lingyu Duan. Uncertainty modeling for out-of-distribution generalization. In International Conference on Learning Representations , 2022.

Xiaoxiao Li, Meirui JIANG, Xiaofei Zhang, Michael Kamp, and Qi Dou. Fedbn: Federated learning on non-iid features via local batch normalization. In International Conference on Learning Representations .

Ya Li, Mingming Gong, Xinmei Tian, Tongliang Liu, and Dacheng Tao. Domain generalization via conditional invariant representations. In Proceedings of the AAAI conference on artificial intelligence , volume 32, 2018.

Ya Li, Xinmei Tian, Mingming Gong, Yajing Liu, Tongliang Liu, Kun Zhang, and Dacheng Tao. Deep domain generalization via conditional invariant adversarial networks. In Proceedings of the European Conference on Computer Vision (ECCV) , pages 624-639, 2018.

Quande Liu, Cheng Chen, Jing Qin, Qi Dou, and Pheng-Ann Heng. Feddg: Federated domain generalization on medical image segmentation via episodic learning in continuous frequency space. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 1013-1023, 2021.

Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. Communication-efficient learning of deep networks from decentralized data. In Artificial Intelligence and Statistics , pages 1273-1282, 2017.

Yuval Netzer, Tao Wang, Adam Coates, Alessandro Bissacco, Bo Wu, and Andrew Y Ng. Reading digits in natural images with unsupervised feature learning. 2011.

A Tuan Nguyen, Philip Torr, and Ser-Nam Lim. Fedsr: A simple and effective domain generalization method for federated learning. In Advances in Neural Information Processing Systems , 2022.

Jungwuk Park, Dong-Jun Han, Minseok Choi, and Jaekyun Moon. Sageflow: Robust federated learning against both stragglers and adversaries. Advances in Neural Information Processing Systems , 34:840-851, 2021.

Xingchao Peng, Qinxun Bai, Xide Xia, Zijun Huang, Kate Saenko, and Bo Wang. Moment matching for multi-source domain adaptation. In Proceedings of the IEEE/CVF international conference on computer vision , pages 1406-1415, 2019.

Amirhossein Reisizadeh, Aryan Mokhtari, Hamed Hassani, Ali Jadbabaie, and Ramtin Pedarsani. Fedpaq: A communication-efficient federated learning method with periodic averaging and quantization. In International Conference on Artificial Intelligence and Statistics , pages 20212031. PMLR, 2020.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems , 30, 2017.

Hemanth Venkateswara, Jose Eusebio, Shayok Chakraborty, and Sethuraman Panchanathan. Deep hashing network for unsupervised domain adaptation. In Proceedings of the IEEE conference on computer vision and pattern recognition , pages 5018-5027, 2017.

Hongyi Wang, Kartik Sreenivasan, Shashank Rajput, Harit Vishwakarma, Saurabh Agarwal, Jy-yong Sohn, Kangwook Lee, and Dimitris Papailiopoulos. Attack of the tails: Yes, you really can backdoor federated learning. Advances in Neural Information Processing Systems , 33:16070-16084, 2020.

Guile Wu and Shaogang Gong. Collaborative optimization and aggregation for decentralized domain generalization and adaptation. In Proceedings of the IEEE/CVF International Conference on Computer Vision , pages 6484-6493, 2021.

Yabin Zhang, Minghan Li, Ruihuang Li, Kui Jia, and Lei Zhang. Exact feature distribution matching for arbitrary style transfer and domain generalization. arXiv preprint arXiv:2203.07740 , 2022.

Yuyang Zhao, Zhun Zhong, Fengxiang Yang, Zhiming Luo, Yaojin Lin, Shaozi Li, and Nicu Sebe. Learning to generalize unseen domains via memory-based multi-source meta-learning for person re-identification. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 6277-6286, 2021.

Kaiyang Zhou, Yongxin Yang, Timothy Hospedales, and Tao Xiang. Learning to generate novel domains for domain generalization. In European conference on computer vision , pages 561-578. Springer, 2020.

Kaiyang Zhou, Yongxin Yang, Yu Qiao, and Tao Xiang. Domain generalization with mixstyle. In International Conference on Learning Representations , 2021.

A Results on DomainNet Dataset

To demonstrate the effectiveness of our StableFDG on a larger dataset, we performed experiments on DomainNet dataset in a single-domain data distribution scenario. To this end, we utilize ResNet-50 pretrained on ImageNet. The number of global rounds and mini-batch size are set to 60 and 64, respectively. The remaining settings are the same as those in our main paper. The results in Table 5 show that our StableFDG consistently outperforms not only the centralized DG works but also the prior works on federated DG even with a more complex dataset.

B Discussion on Computational and Communication Costs

Table 6 compares the communication, computation, and average accuracy of different schemes on PACS in a multi-domain data distribution scenario. ResNet-18 is adopted as in our main manuscript. We first compare the uplink communication load of each client in a specific global round. Compared to FedAvg that only transmits the model in each round, our scheme requires additional communication burden for transmitting the style statistics and the attention module, which are negligible. We also compare the computation time by measuring the time required for local update at each client using an GTX 1080 Ti GPU. CCST [2] and FedDG [25] require large computation due to the increased amounts of data samples or multiple backpropagations for meta training. Our scheme requires additional computation caused by style exploration, attention module update, etc., which are the costs for better generalization to the unseen domain.

C Experiments with Three Clients

Different from the prior works [2, 28] for federated DG adopting the usual setting where the number of clients equals the number of source domains, in our main paper, we introduce a more practical experimental setting for federated DG where the source data is distributed to more clients than the number of source domains. For a comparison with them in the same setting, we also provide additional experimental results in the setup with number of clients = number of source domains . Table 7 shows the results on PACS and Office-Home datasets with three clients in a single-domain data distribution scenario. It is confirmed from the results that our StableFDG also achieves better performance compared to the existing works [2, 28] in this simple setting.

(b) Results on Office-Home dataset.

D Effect of Each Component of StableFDG

As mentioned in the main manuscript, we provide further ablation studies on the effect of each component in StableFDG using VLCS dataset. Table 8 shows that each component individually brings performance gain compared to FedAvg. Using both strategies achieves greater performance gains, confirming that the two proposed schemes work in a complementary fashion.

E Ablation Studies on Style-Based Learning

E.1 Randomness in style sharing

In our style based learning, style sharing among clients is performed at random. However, one can think of the strategy where client n receives the style information Φ n ′ that has the largest distance with its own style information Φ n ′ in the style space. Table 9 shows the corresponding result using PACS dataset in a multi-domain data distribution setup. Interestingly, it can be seen that the random selection adopted in this paper performs better, since most of the users generally tend to receive the same style statistics when using the largest distance strategy.

E.2 Effect of the number of shared styles

We also compare the effect of number of styles received at each client in Table 10 using PACS dataset in a multi-domain data distribution setup. The performance increases as the number of received styles increases, with small additional communication load (the vector length of the style information is 128, which is negligible compared to the number of model parameters, which is 11,180,103).

E.3 Effect of k -means++ in style shifting

In Section 3.1, we utilized the k -means++ as a tool for facilitating our key idea in style shifting, which is to effectively balance between the original source domain and the new source domain for better generalization; k-means++ plays a role to select the B/ 2 styles that are similar to the remaining B/ 2 styles in the mini-batch. By doing so, the model can explore new styles while not losing the performance on the original styles. To see this effect, we compare k -means++ vs. random sampling when selecting B/2 samples to be shifted, in Table 11. The results show that strategically selecting the B/2 samples to be shifted achieves better performance especially in the Sketch domain (1.89% gain) that has a large style gap with other domains. We believe that these results motivate and justify our design choice.

E.4 Effect of class-balanced oversampling

In our main paper, we performed the class-balanced oversampling in the feature space to alleviate the class-imbalance issue during style exploration. To confirm the effectiveness of the class-balanced oversampling, we compare it with random oversampling under the same condition where only the style exploration is applied without other components. Table 12 shows the results on Office-Home dataset, where the class distribution is highly imbalanced in a multi-domain data distribution scenario. It can be seen that our class-balanced oversampling achieves higher performance over the simple random sampling, which validates the efficacy of mitigating the class imbalance problem in FL clients.

E.5 Effect of the operation probability in style-based learning

We conduct additional ablation studies on the probability value (defined as p here) utilized to control the operation of the style sharing/shifting and style exploration modules. Larger p means that our scheme is more likely to be activated. In Table 13, we provide results on various probability values in a single-domain data distribution scenario using PACS dataset. From the results, it is confirm that for all p values, the proposed StableFDG outperforms existing baselines, demonstrating that our scheme can work well with an arbitrarily chosen probability p . In detail, when p ranges from 0.3 to 0.7, the high performance is maintained while the performance decreases at both extreme probabilities ( p = 0 . 1 and 0.9). Therefore, it is recommended for practitioners to select the p in an appropriate range, avoiding extreme cases.

E.6 Where to apply the style module

For implementation, style-based learning is applied only in the 1st, 2nd, 3rd blocks among 4 residual blocks in ResNet-18. Note that at the output of the 4th block, label information is dominant rather than style information, which results in degraded performance when style-based schemes are applied. This is confirmed by our new experiments in the table below. It can be seen from the results that if we consider the 4th residual block to apply our style-based learning, the performance gets degraded. This result confirms the intuition that style-based learning should be conducted at the earlier layers where style information is preserved.

F Ablation Studies on Attention-based Feature Highlighter

F.1 Effect of adopting both cross-attention and self-attention

In our main paper, the similarity metric in equation (6) adopts cross-attention, while the metric in equation (8) combines cross-attention and self-attention. When applying only the cross-attentionbased metric in equation (6), we found that the similarity value could become low even when the two samples belong to the same class, in special cases. We handled this issue by adding the self-attention component as in equation (8). Intuitively, by doing this, the attention module is learning to extract the important features across images (via cross-attention), and within the image (via self-attention). Table 15 compares the performance of our StableFDG when using (i) self-attention alone, (ii) cross-attention alone (equation (6)), and (iii) both self and cross attentions at the same time (equation (8)), confirming the advantage of using self-attention and cross-attention together.

F.2 Comparison using the same model size

Our attention module requires 0.44% of additional model parameters to perform the attention-based learning. For a fair comparison to see the effect of our attention-based learning, we consider a

different baseline with the same model size but without attention-based learning. Specifically, the baseline computes the attention score map using only additional convolutional operations and take the weighted average of the feature z i based on the attention score map. Table 16 shows the results using PACS dataset in a single-domain data distribution scenario. The results demonstrate that our attention-based learning achieves performance improvements on all four domains while playing a key role in capturing essential parts of the features.

F.3 Effect of attention in a centralized setup

Now we provide answer to the following question: Instead of the FL setup we focused on, can attention provide benefits in the centralized DG setup? Table 17 shows the results with/without attention module in a centralized setup using PACS dataset. The results show that attention still provides performance improvements in the centralized setup by learning domain-invariant features, although the gain is slightly lower than the gain in the FL setup as shown in Table 3 of the main manuscript. These results indicate that the proposed attention-based learning indeed captures the domain-invariant characteristics of samples, while the scheme provides more benefits in the FL setup where each client is prone to overfitting due to lack of data.

G Other Implementation Details

Our code is built upon the official code of [39] and [1]. During the local update process, we use SGD as an optimizer with a momentum of 0.9 and a weight decay of 5 e -4 . For PACS, Office-Home and VLCS, the learning rate is set to 0.001 and the cosine annealing is used as a scheduler. For Digits-DG, we set the learning rate to 0.02 and the learning rate is decayed by 0.1 every 20 steps. For our attention-based feature highlighter, at least two samples are required to be in the mini-batch of every client. When this condition is not met, additional samples are extracted from the corresponding client's local dataset to facilitate cross-attention.

More detailed description on oversampling: Let s n ∈ R B × C × H × W be a mini-batch of features in client n at a specific layer, obtained after Steps 1 and 2 in the main manuscript. Now given a fixed oversampling size, we oversample the features in the mini-batch to obtain ˜ s n , so that the concatenated mini-batch ˆ s n = [ s n , ˜ s n ] becomes class-balanced as much as possible. Consider a toy example where the number of samples for classes a , b , c are 3, 2, 1, respectively in the mini-batch s n . In this example, if the oversampling size is 3, we randomly choose one data point from class b and two data points from class c (in this case, the same data point is selected for two times with duplication) to obtain ˜ s n , so that the concatenated mini-batch ˆ s n = [ s n , ˜ s n ] becomes class-balanced. If the oversampling size is 1, we oversample one data point in class c to make the concatenated mini-batch to be balanced as much as possible. If the oversampling size is 6, we oversample 1, 2, 3 samples from classes a , b , c , respectively to construct ˜ s n . The concatenated mini-batch ˆ s n = [ s n , ˜ s n ] is utilized for style-based learning and updating the model. This process not only mitigates the class-imbalance issue in each client but also provides a good platform for style exploration by oversampling the features. In our

work, we reported the results with oversampling size of B (which is equal to the mini-batch size), while the effect of the oversampling size is also reported in Fig. 4 of the main manuscript: A larger oversampling size leads to a better performance, and more importantly, our StableFDG outperforms the baseline even without any oversampling.

H Visualization of Attention Score Maps

Finally in Fig. 6, we visualize the attention score maps of different input images during testing. The results confirm the effectiveness of our attention-based feature highlighter to focus on the important parts of each image from the unseen domain.

★

I Pseudo Code Algorithm of Stable FDG

Algorithm 1 summarizes the overall process of our StableFDG.

Peer reviews:
comments: SUMMARY:
This paper proposes a style and attention based learning strategy for federated domain generalization which consists of two stages. The first stage shares style information among clients and then explorates and augments styles, and the second stage highlights the essential characteristics of images with attention mechanism. Experiments  show that the proposed method achieves a new state of the art.

STRENGTHS:
- The style-based learning strategy is novel and intuitive.
- The paper is well written and easy to follow
- The proposed method is technically sound with superior results on the 5 datasets.

WEAKNESSES:
The detail of oversampling is not fully covered in the paper.

QUESTIONS:
- Why can eq.7. avoid performance degradation when there is little commonality between key and query images ?
- How is oversampling implemented. What if the dataset itself is not class balanced ?

LIMITATIONS:
The paper discusses its limitation in sec. 5.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2023-06-23

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors observe that existing domain generalization (DG) algorithms face the issue in the setting of federated learning (FL) because of the lack of samples in each client's local datasets/domains. They propose StableFDG, which utilizes both style-based learning which allows the clients to explore styles different from that of the original source domain, and an attention-based feature highlighter, which captures the similarities and common characteristics of features of samples of the same class. The authors conduct their evaluation on datasets such as Office-Home and Digits-DG.\n\nSTRENGTHS:\n+ The evaluation is thorough, considering most of the popular benchmarks including the more challenging ones such as DomainNet.\n+ The visuals are helpful and illustrative for readers to understand the proposed method\n\nWEAKNESSES:\n- It would seem that the work could greatly benefit from works in the field of style transfer, which also involves extracting domain knowledge. It would be better for the authors to remark on the connections between their style sharing and exploration with style transfer.\n\n- It is not clearly explained what would happen after the style information is obtained by Equation 3. It's evident that this piece of information is used for style-shifting, etc., but how does such an operation relate to the intuition/big picture behind the approach?\n\nQUESTIONS:\n+ Have you considered style transfer approaches to boost your current style-based learning? Style transfer approaches also strive to extract the domain representation (style); would this field help or hinder the performance of StableFDG?\n\nLIMITATIONS:\n- The authors argue that, because the label information is not included in the style parameters, the proposed approach doesn't have privacy issues. However, style statistics/parameters, derived from the domains, still contain domain knowledge. It's not clear and not elaborated why the privacy issue is resolved simply because the labels aren't shared.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn this paper, the authors present StableFDG, where a style-based strategy enables the model to reach a variety of new styles beyond each client's source domain, and an attention-based approach captures and emphasizes the important/common features of each category.\n\nSTRENGTHS:\nThe StableFDG method proposed in the article has achieved good experimental results.\nThe experimental details in the article are particularly well articulated and richly illustrated.\nThe essay writing is well-organized and readable.\n\nWEAKNESSES:\nThe idea of enabling the model to reach a variety of new styles beyond the source domain of each client has been proposed by other method [1], and what the authors have done is more of an incremental improvement.\nIn addition, the authors in a related work evaluate [1] to improve the computational and memory costs, but it itself takes up more memory than [1].\n [1] Junming Chen, Meirui Jiang, Qi Dou, and Qifeng Chen. Federated domain generalization for image recognition via cross-client style transfer. In Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision, pages 361–370, 2023.\n\nThe method proposed in section 3.1 of the article is more like a combination of the methods mentioned in the authors' text.\n\nIn the \"Where to apply style-based modules\" subsection, the authors directly follow the way other methods apply modules, where more detailed ablation experiments could be done or the benefits or reasons for such application could be explained in the article.\n\nThe authors propose a new experimental setup, are the other methods in the paper based on this experimental setup, and what are the advantages of such a new setup over the previous ones?\n\nQUESTIONS:\nSee the weakness.\n\nLIMITATIONS:\nThe authors could have elaborated more fully on their new experimental setup.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper presents a novel domain generation method in the FL framework. The authors make two key contributions to enhance domain generalization in FL. Firstly, they introduce style-based learning, which empowers individual clients to explore new styles by leveraging both their own local statistics and those of other clients. This approach promotes domain diversity. Secondly, the authors propose an attention-based feature highlighter that captures similarities among the features of data samples within the same class. This component enables better learning of domain-invariant characteristics. To validate the efficacy of their proposed method, extensive experiments were conducted. The results demonstrate the superiority of the approach compared to existing methods on various benchmark datasets.\n\nSTRENGTHS:\nThis paper addresses an important practical problem within the FL setup. It provides comprehensive coverage of related works and is well-written for easy understanding. The extensive experiments conducted effectively demonstrate the effectiveness of the proposed method. The supplementary material is valuable.\n\nWEAKNESSES:\nThe adaptation of mixing styles and exploration to FL frameworks is straightforward. Also, the attention-based feature highlighter can be viewed as a variation of cross-attention.\n\nQUESTIONS:\n1) In this paper, statistics are recalculated at each round. Instead, could the running mean and variance of Batch Normalization (BN) be used as a substitute? Also, I am curious about the relationship between style augmentation and BN.\n2) The attention-based feature highlighter combines self-attention and cross-attention for feature refinement. However, it remains unclear why this specific method is tailored to Federated Domain Generalization (DG). Further explanation is needed to clarify this aspect.\n3) During training, both self-attention and cross-attention are used in Eq. (7), but during testing, only self-attention is employed. It would be beneficial to confirm if using self-attention alone during training is not sufficient.\n\nLIMITATIONS:\nDue to page constraints, some answers to certain questions were provided in the supplementary material, particularly in Supplementary Appendix E. It would be beneficial to include relevant content from the supplementary material into the main body of the paper.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}