SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Attention score. Consider the  $i$ -th data sample of client  $n$ . Given the feature  $z_{i}^{n}\in \mathbb{R}^{C\times H\times W}$  obtained from the output of the feature extractor, we flatten it to a two-dimensional tensorop  $X_{i}\in \mathbb{R}^{C\times HW}$  with a size of  $(C,HW)$ , where we omit the client index  $n$  for notational simplicity. Now consider another  $X_{j}$ $(j\neq 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\in \mathbb{R}^{HW\times HW}$  between  $X_{i}$  and  $X_{j}$  as follows:
$$
S := \operatorname {S i m} \left(X _ {i}, X _ {j}\right) = \left(\theta_ {q} X _ {j}\right) ^ {T} \left(\theta_ {k} X _ {i}\right), \tag {6}
$$
where  $\theta_q \in \mathbb{R}^{d \times C}$ ,  $\theta_k \in \mathbb{R}^{d \times 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 = \theta_q X_j$  and
keys  $K = \theta_{k}X_{i}$  Then, we reshape  $S\in \mathbb{R}^{HW\times HW}$  to  $S_{r}\in \mathbb{R}^{HW\times H\times W}$ ; the  $(m,n)$ -th spatial feature  $(\in \mathbb{R}^{HW\times 1\times 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} = \mathrm{mean}(S_{r})\in \mathbb{R}^{H\times 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}\gets \mathrm{softmax}(a_{s})\in \mathbb{R}^{H\times W}$ .
Figure 3: Proposed attention-based feature highlighter (Sec. 3.2): Our attention-based learning captures important characteristics within each class (regardless of the domain) for better generalization.
Attention-based weighted averaging. Based on the attention score, we take the weighted average of  $z_{i} \in \mathbb{R}^{C \times H \times W}$  to generate an attention feature  $A(z_{i}) \in \mathbb{R}^{C}$  in which the important parts of the key image, having common characteristics with the query image (in the same class), are emphasized:
$$
A \left(z _ {i}\right) _ {c} = \sum_ {h = 1} ^ {H} \sum_ {w = 1} ^ {W} \left(a _ {s}\right) _ {h, w} \left(z _ {i}\right) _ {c, h, w}. \tag {7}
$$
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  $\theta_q$  and  $\theta_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
$$
\operatorname {S i m} _ {\text {m i x}} \left(X _ {i}, X _ {j}\right) = \left(\frac {\theta_ {q} X _ {j} + \theta_ {q} X _ {i}}{2}\right) ^ {T} \left(\theta_ {k} X _ {i}\right), \tag {8}
$$
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  $\mathrm{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  $\theta_{q}$  and  $\theta_{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.
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  $\alpha$  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 \times 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
Methods	PACS	VLCS
Art	Cartoon	Photo	Sketch	Avg.	Caltech	LabelMe	Pascal	Sun	Avg.
FedAvg [26]	73.67	70.87	90.27	55.70	72.63	93.75	59.30	70.05	69.90	73.25
FedBN [22]	78.42	70.9	90.96	54.07	73.59	94.81	58.59	72.06	70.36	73.96
MixStyle [39]	79.10	76.30	90.10	60.63	76.53	95.20	60.40	72.10	69.93	74.41
DSU [21]	80.43	75.70	92.60	69.87	79.65	96.13	58.77	71.80	71.87	74.64
CCST [2]	71.35	72.40	88.65	64.10	74.13	92.50	61.20	68.20	66.50	72.10
FedDG [25]	71.20	71.40	90.70	59.20	73.13	95.3	57.5	72.8	69.8	73.85
FedSR [28]	76.40	71.25	93.25	60.55	75.36	92.10	60.50	70.75	71.65	73.75
StableFDG (ours)	84.10	78.57	95.40	72.73	82.70	98.13	59.20	73.60	70.27	75.30
(a) PACS and VLCS datasets.
Methods	Office-Home	Digits-DG
Art	Clipart	Product	Real	Avg.	MNIST	MNIST-M	SVHN	SYN	Avg.
FedAvg [26]	57.27	48.23	72.77	74.60	63.22	98.05	70.95	68.95	86.40	81.09
FedBN [22]	57.56	48.13	72.65	74.57	63.23	97.33	72.68	71.77	85.36	81.79
MixStyle [39]	56.05	51.55	70.95	73.25	62.95	97.75	74.25	70.85	85.50	82.09
DSU [21]	58.55	52.60	71.60	73.15	63.98	98.10	75.60	70.47	85.80	82.49
CCST [2]	51.3	51.75	70.2	70.3	60.89	95.10	62.80	56.60	74.90	72.35
FedDG [25]	57.6	48.1	72.55	74.33	63.15	97.97	72.13	71.03	87.87	82.25
FedSR [28]	57.8	48.1	72.1	74.2	63.05	98.00	73.00	68.50	86.70	81.55
StableFDG (ours)	57.57	54.30	72.33	74.97	64.79	97.23	74.53	72.95	85.85	82.64
(b) Office-Home and Digits-DG datasets.
Table 1: Main result 1 (single-domain data distribution): Each client has one source domain in its local data. The proposed StableFDG achieves the best generalization, underscoring its effectiveness.
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.
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.
Methods	Office-Home	VLCS
Art	Clipart	Product	Real	Avg.	Caltech	LabelMe	Pascal	Sun	Avg.
FedAvg [26]	57.70	48.30	72.87	75.33	63.55	93.65	61.10	72.55	65.40	73.18
FedBN [22]	57.07	48.32	72.31	74.57	63.07	94.34	62.61	69.89	69.04	73.97
MixStyle [39]	55.87	52.03	71.10	74.20	63.30	95.20	60.77	73.90	66.73	74.15
DSU [21]	58.60	52.80	71.63	74.00	64.26	96.87	60.23	72.97	68.97	74.76
CCST [2]	52.2	52.2	70.6	72.3	61.83	96.70	60.40	71.40	65.00	73.38
FedDG [25]	57.9	48.6	73.2	75.0	63.68	96.2	60.7	72.4	67.3	74.15
FedSR [28]	58.1	48.2	72.5	75.4	63.55	92.60	60.80	72.15	68.30	73.46
StableFDG (ours)	57.87	54.20	73.10	75.00	65.04	98.50	60.07	74.40	69.43	75.60
(a) Office-Home and VLCS datasets.
Art	Cartoon	Photo	Sketch	Avg.
FedAvg [26]	74.87	74.53	95.30	63.37	77.02
FedBN [22]	77.00	76.83	95.45	67.01	79.07
MixStyle [39]	80.07	77.53	96.23	67.40	80.31
DSU [21]	80.53	76.30	95.37	70.93	80.78
CCST [2]	75.53	75.80	93.53	71.13	79.00
FedDG [25]	75.35	75.85	95.65	61.05	76.98
FedSR [28]	73.83	74.83	95.53	66.03	77.56
StableFDG (ours)	83.97	79.10	96.27	75.67	83.75
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-by-one in a multi-domain data distribution setup using PACS. We compare our results with style-augmentation DG base
(b) PACS dataset.
Table 2: Main result 2 (multi-domain data distribution): Each client has multiple source domains in its local dataset. The results are consistent with the single-domain scenario in Table 1.
Methods	Art	Cartoon	Photo	Sketch	Avg.
MixStyle [39]	80.07	77.53	96.23	67.40	80.31
DSU [21]	80.53	76.30	95.37	70.93	80.78
StableFDG (only style)	82.62	79.01	95.57	74.47	82.92
StableFDG (only attention)	79.98	78.58	95.75	71.35	81.41
StableFDG (both)	83.97	79.10	96.27	75.67	83.75
Table 3: Effect of each component of StableFDG.
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  $\alpha = 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  $\alpha$  is too small, the performance is relatively low since
Figure 4: Effects of exploration level  $\alpha$  (left) and oversampling size (right) in StableFDG.
the model is not able to explore novel styles beyond the client's source domains. If  $\alpha$  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  $\alpha$ , 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
Methods	Art	Cartoon	Photo	Sketch	Avg.
MixStyle [39]	81.61	78.83	96.77	72.29	82.38
DSU [21]	78.84	79.56	95.21	79.39	83.25
StableFDG (ours)	85.02	79.65	96.38	78.35	84.85
(a) Performance in a centralized DG setup.
Methods	Art	Cartoon	Photo	Sketch	Avg.
MixStyle [39]	87.51	81.12	97.48	68.39	83.63
DSU [21]	86.48	81.22	97.21	73.99	84.73
StableFDG (ours)	90.01	83.29	98.02	79.47	87.70
(b) Performance using ResNet-50.
Table 4: The applicability of StableFDG in a centralized DG setup (Table 4a) and performance using a larger model (Table 4b) on the PACS dataset.
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
Figure 5: Visualization of attention score maps of each input image (left: bed, right: radio).
Additional results. Other implementation details, comprehensive ablation studies for each component, discussions on complexity, and results on DomainNet dataset are in Appendix.
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.
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.
[1] 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.
[2] 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.
[35] 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.
[36] 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.
[37] 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.
[38] 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.
[39] 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.
Methods	Clipart	Inforgraph	Painting	Quickdraw	Real	Sketch	Avg.
FedAvg [26]	61.52	24.75	50.83	12.08	60.00	49.92	43.18
FedBN [22]	60.23	24.40	50.43	11.99	59.57	49.66	42.71
MixStyle [39]	61.39	24.33	51.41	13.07	57.90	51.40	43.25
DSU [21]	62.43	24.30	51.87	13.65	58.75	52.40	43.90
FedDG [25]	62.49	23.71	48.42	12.45	61.44	49.11	42.94
FedSR [28]	61.91	25.37	50.54	11.59	62.03	50.13	43.60
StableFDG	62.58	24.12	52.23	14.87	60.60	52.50	44.48
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.
Table 5: Results on DomainNet dataset in a single-domain data distribution scenario.
Methods	Communication load	Computation time	Achievable average accuracy
FedAvg [26]	44.98 MB	4.57 sec	77.02 %
CCST [2]	44.98 MB	9.03 sec	79.00 %
FedDG [25]	44.98 MB	22.44 sec	76.98 %
FedSR [28]	44.98 MB	4.59 sec	77.56 %
StableFDG	45.00 MB	7.39 sec	83.75 %
Table 6: Computation and communication cost comparison.
C Experiments with Three Clients