SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Jungwuk Park*
savertm@kaist.ac.kr
Dong-Jun Han*
han762@purdue.edu
Jinho Kim
SK Hynix
jinho123.kim@sk.com
Shiqiang Wang
wangshiq@us.ibm.com
Christopher G. Brinton
Jaekyun Moon
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.
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/style 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 style-exploration 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.
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. DG is 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], meta-learning [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  $\mathcal{S}_n = \{(x_i^n,y_i^n)\}_{i = 1}^{\rho_n}$  be the local dataset of the  $n$ -th client, which consists of  $\rho_{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  $\mathcal{S}_n$ . Previous works on FL focus on constructing a global model that predicts well on the overall dataset  $\mathcal{S} = \{\mathcal{S}_1,\mathcal{S}_2,\dots ,\mathcal{S}_N\}$  or personalized models that work well on individual local datasets  $\mathcal{S}_n$ . In contrast to the conventional FL setup, given the overall dataset (or source domains)  $\mathcal{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  $\mathcal{T}$ .
Background. Let  $s \in \mathbb{R}^{C \times H \times 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  $\mu(s) \in \mathbb{R}^C$  and the channel-wise standard deviation  $\sigma(s) \in \mathbb{R}^C$  can be written as
$$
\mu (s) _ {c} = \frac {1}{H W} \sum_ {h = 1} ^ {H} \sum_ {w = 1} ^ {W} s _ {c, h, w}, \quad \sigma^ {2} (s) _ {c} = \frac {1}{H W} \sum_ {h = 1} ^ {H} \sum_ {w = 1} ^ {W} \left(s _ {c, h, w} - \mu (s) _ {c}\right) ^ {2}, \tag {1}
$$
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,\dots ,T$ . In the beginning of round  $t$ , a selected set of clients download the current global model  $\mathbf{w}_t$  from the server. Before local training begins, each client  $n$  computes its own style information  $\Phi_n = [\mu_n,\sigma_n,\Sigma_n(\mu),\Sigma_n(\sigma)]$  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
Figure 1: Overview of our problem setup and algorithm for federated domain generalization. Each client can have a single source domain as described in Fig. 1, or even multiple source domains in its local dataset.
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
(a) Existing style-based DG baselines
(b) Proposed style-based learning for federated DG
Figure 2: Proposed style-based learning strategy (Sec. 3.1): Compared to existing style-based DG methods that rely on interpolation or style shift near the original style of each sample, our style-based learning (i) effectively utilizes other FL clients' styles based on style sharing and shifting, and (ii) enables the model to further explore a wider style space via feature-level oversampling and extrapolation, handling the issue of domain-limited FL settings.
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,  $\mu_{n}$ , and the variance of the channel-wise means,  $\Sigma_{n}(\mu)$ , as follows:
$$
\mu_ {n} = \frac {1}{\rho_ {n}} \sum_ {i = 1} ^ {\rho_ {n}} \mu \left(s _ {i} ^ {n}\right), \quad \Sigma_ {n} (\mu) = \frac {1}{\rho_ {n}} \sum_ {i = 1} ^ {\rho_ {n}} \left(\mu_ {n} - \mu \left(s _ {i} ^ {n}\right)\right) ^ {2}, \tag {2}
$$
where  $s_i^n \in \mathbb{R}^{C \times H \times W}$  is the feature of the  $i$ -th sample in the  $n$ -th client's local dataset, and the square operation  $(\cdot)^2$  works in an element-wise manner. Similarly, the average and the variance of channelwise standard deviations are computed as  $\sigma_n = \frac{1}{\rho_n} \sum_{i=1}^{\rho_n} \sigma(s_i^n)$ ,  $\Sigma_n(\sigma) = \frac{1}{\rho_n} \sum_{i=1}^{\rho_n} (\sigma_n - \sigma(s_i^n))^2$ , respectively. Here,  $\mu_n$  and  $\sigma_n$  represent the center of the style statistics of client  $n$ , while  $\Sigma_n(\mu)$  and  $\Sigma_n(\sigma)$  show how these style statistics of client  $n$  are distributed around the center. Now we define
$$
\Phi_ {n} = \left[ \mu_ {n}, \sigma_ {n}, \Sigma_ {n} (\mu), \Sigma_ {n} (\sigma) \right] \tag {3}
$$
as the style information representing the domain identity of the  $n$ -th client in the style-space. Each client  $n$  sends  $\Phi_n$  to the server, and the server randomly shares the style information to other clients in a one-to-one manner; client  $n$  receives  $\Phi_{n'}$  that belongs to client  $n'$  ( $n \neq 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  $\Phi_{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) = (\sigma_{n'} + \epsilon_{\sigma}\Sigma_{n'}(\sigma))\left(\frac{s - \mu(s)}{\sigma(s)}\right) + (\mu_{n'} + \epsilon_{\mu}\Sigma_{n'}(\mu))$ , where  $f(s)$  is the new feature shifted from  $s$  and  $\epsilon_{\mu} \sim \mathcal{N}(0,\Sigma_n(\mu))$ ,  $\epsilon_{\sigma} \sim \mathcal{N}(0,\Sigma_n(\sigma))$  are the values sampled from normal distributions. Then, the mini-batch applied with new styles in  $\Phi_{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  $\Phi_{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 \in \mathbb{R}^{B \times C \times H \times 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 mini-batch 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  $\tilde{s}^n$  with size  $B$ , and concatenate it with the original mini-batch as follows:  $\hat{s}^n = [s^n, \tilde{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  $\tilde{s}^n$  (i.e., the set of oversampled features) to novel styles beyond the style of each client. Let  $\tilde{s}_i^n$  be the feature of  $i$ -th sample in the oversampled mini-batch  $\tilde{s}^n$ . We obtain the new styles by extrapolating the original styles in  $\tilde{s}^n$  around the average of channel-wise mean  $\mu_n(\hat{s}^n)$  and the average of standard deviations  $\sigma_n(\hat{s}^n)$  computed on the concatenated mini-batch  $\hat{s}^n = [s^n, \tilde{s}^n]$  as
$$
\mu_ {n e w} \left(\tilde {s} _ {i} ^ {n}\right) = \mu \left(\tilde {s} _ {i} ^ {n}\right) + \alpha \cdot \left(\mu \left(\tilde {s} _ {i} ^ {n}\right) - \mu_ {n} \left(s ^ {n}\right)\right), \tag {4}
$$
$$
\sigma_ {\text {n e w}} \left(\tilde {s} _ {i} ^ {n}\right) = \sigma \left(\tilde {s} _ {i} ^ {n}\right) + \alpha \cdot \left(\sigma \left(\tilde {s} _ {i} ^ {n}\right) - \sigma_ {n} \left(s ^ {n}\right)\right), \tag {5}
$$
where  $\alpha$  is the exploration level. We perform AdaIN to shift the style of  $\tilde{s}_i^n$  to the new style statistics  $\mu_{new}(\tilde{s}_i^n)$  and  $\sigma_{new}(\tilde{s}_i^n)$ . If  $\alpha = 0$ , the styles remain unchanged, and as  $\alpha$  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  $\hat{s}^n$  to generate diverse domains as in [39] as  $\mu_{new}(\hat{s}_i^n) = \lambda \cdot \mu (\hat{s}_i^n) + (1 - \lambda)\cdot \mu (\hat{s}_j^n)$  and  $\sigma_{new}(\hat{s}_i^n) = \lambda \cdot \sigma (\hat{s}_i^n) + (1 - \lambda)\cdot \sigma (\hat{s}_j^n)$ , where  $\hat{s}_i^n$  and  $\hat{s}_j^n$  are arbitrary two samples in  $\hat{s}^n$  and  $\lambda$  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  $\alpha$ , 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  $\alpha$ , 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.