The aim of this paper is to develop a generalized federated domain generalization mechanism that takes into account both intra-client and inter-client domain distributions to collaboratively train a model with strong generalization performance using unshared data from different clients. Assume that there are K clients, each of which collects data independently to form its own local dataset, where the local dataset of the k-th client is denoted as $\left(X_k,Y_k\right)$, $X_k\subset {ℝ}^{N_k\times M_X}$, $Y_k\subset {ℝ}^{N_k\times M_Y}$, and $X_{i,k}$ and $Y_{i,k}$ denote the i-th sample data of the k-th client, each client can train the model independently and can transmit the information with the server. Let $\mathcal{X}\subset {ℝ}^{M_X}$, $\mathcal{Y}\subset {ℝ}^{M_Y}$ and $\mathcal{Z}\subset {ℝ}^{M_Z}$ denote respectively the input space, the output space and the feature space. $F:\mathcal{X}\to \mathcal{Z}$ denotes the feature extractor, $C:\mathcal{Z}\to \mathcal{Y}$ denotes the classifier, $Z_k\subset {ℝ}^{N_k\times M_Z}$ denotes the sample features of the k-th client extracted through the feature extractor $F(\cdot )$, $Z_{k,:,i}$ denotes the i-th feature variable in the feature space of the k-th client and $N_k$ denotes the sample size of the k-th client. Since each client collects data independently, the problem of agnostic distribution bias may exist both intra-client and inter-client, so we would like to give a generalized federated domain generalization mechanism that improves the generalization performance of model on unseen clients while observing federated learning privacy preservation.

FedALN is a network structure for aligning inter-client feature distributions in the federated context proposed in [16] , which is mainly based on Adversarial Learning Networks (ALNs), by generating a dynamic reference distribution and aligning the feature distributions of each client to this reference distribution to reduce the inter-clients’ domain distribution differences and improve the generalization performance of the global model. Its effectiveness has been verified by a large number of experiments. Our proposed method also adopts FedALN to eliminate the impact of inter-client domain bias on model performance and below we provide a brief review of FedALN.

FedALN consists of 3 main components: the feature extractor $F(\cdot )$, the distribution generator $G(\cdot )$ and the discriminator $D(\cdot )$. Among them, the feature extractor $F(\cdot )$ is used to extract real data features $Z$ from the local raw data and the distribution generator $G(\cdot )$ is used to generate the reference distribution and output the generated features $Z^{\prime }$, the discriminator $D(\cdot )$ is employed to differentiate between the real data features $Z$ and the generated features $Z^{\prime }$.

Specifically, each client has a local ALN and the data distribution differences between different clients are converted into feature distribution differences. Firstly, each client aligns the local feature distribution with the generated reference distribution by training the local ALNs, and then the model parameters of the local ALNs of all clients are passed to the server for model parameter aggregation to form a new FedALN to be passed to each client again. Through multiple interactions of the local clients and the server, the feature distributions of different clients will be aligned to the same reference distribution, thus indirectly realizing the alignment of the feature distributions between clients.

Then we will describe specifically the three components of FedALN and the corresponding adversarial loss functions.

Discriminator $D(\cdot )$: The discriminator $D(\cdot )$ is primarily responsible for distinguishing between the true data features $Z$, which are obtained from the feature extractor $F(\cdot )$, and the false features $Z^{\prime }$, generated by the distribution generator $G(\cdot )$. To facilitate this task, one-hot encoded labels are used to ensure that the discriminator $D(\cdot )$ can discriminate between true and false features according to their respective categories.

In the training process, the true data feature $Z$ is regarded as a negative sample and the false feature $Z^{\prime }$ is regarded as a positive sample, input data features (true data feature $Z$, false feature $Z^{\prime }$) and one-hot encoded labels $y$, and output the probability that the feature is judged as a positive sample, i.e., the probability that the feature comes from the distribution generator $G(\cdot )$. The corresponding adversarial loss function is expressed as follows:

${ℒ}_{adv\text{\_}d}=-\left({\mathbb{E}}_{Z~p\left(Z\right)}\left[{\left(1-D\left(Z|y\right)\right)}^2\right]+{\mathbb{E}}_{Z^{\prime }~p\left(Z^{\prime }\right)}\left[D{\left(Z^{\prime }|y\right)}^2\right]\right)\text{,}$ (1)

where $Z=F\left(X\right)$ denotes the real data features extracted from the feature extractor, $Z^{\prime }=G\left(\epsilon \right)$ denotes the false features generated from the distribution generator and $\epsilon$ is a random noise sample from a uniform distribution $\left[0,1\right)$, $p\left(Z\right)$ denotes the real feature distribution generated by the raw data $X$ through the feature extractor $F(\cdot )$, and $p\left(Z^{\prime }\right)$ denotes the reference distribution generated by random noise $\epsilon$ through the distribution generator $G(\cdot )$. This adversarial loss function is mainly used to measure the ability of the discriminator to correctly discriminate.

Feature extractor $F(\cdot )$: The feature extractor $F(\cdot )$ is mainly used to extract real data features $Z$ from local raw data, and when it inputs local raw data $X$, it outputs extracted feature information $Z$. During the ALNs training, the feature distribution $p\left(Z\right)$ extracted by the feature extractor $F(\cdot )$ needs to be progressively aligned to the reference distribution $p\left(Z^{\prime }\right)$, i.e., given a trained discriminator $D(\cdot )$ with fixed parameters, when the feature distribution $p\left(Z\right)$ is closer to the reference distribution $p\left(Z^{\prime }\right)$, the more likely the discriminator $D(\cdot )$ fails to correctly discriminate the real data features $Z$ from the feature distribution $p\left(Z\right)$. At this point the discriminator $D(\cdot )$ is used to measure the gap between the feature distribution $p\left(Z\right)$ and the reference distribution $p\left(Z^{\prime }\right)$. The corresponding adversarial loss function of the feature extractor ${ℒ}_{adv\text{\_}f}$ is expressed as follows:

${ℒ}_{adv\text{\_}f}={\mathbb{E}}_{Z~p\left(Z\right)}\left[{\left(1-D\left(Z|y\right)\right)}^2\right]\text{,}$ (2)

where the discriminator $D(\cdot )$ outputs the probability that the features come from the generated distribution $p\left(Z^{\prime }\right)$. This adversarial loss function indicates that the more likely the discriminator fails to correctly discriminate the true features $Z$ from the feature distribution $p\left(Z\right)$, the closer the feature distribution $p\left(Z\right)$ is to the reference distribution $p\left(Z^{\prime }\right)$.

Distribution generator $G(\cdot )$: The distribution generator $G(\cdot )$ is mainly used to generate the dynamic reference distribution $p\left(Z^{\prime }\right)$ and the corresponding generated features $Z^{\prime }$ (also known as spurious features). It inputs random noise $\epsilon$ and one-hot encoded true label $y$, and it outputs false feature $Z^{\prime }$ with the same dimension as the true feature $Z$. In the training process, given a trained discriminator $D(\cdot )$ with fixed parameters, the ${ℒ}_{adv\text{\_}g}$ is primarily employed to adjust the parameters of the distribution generator $G(\cdot )$, and the specific distribution generator adversarial loss function ${ℒ}_{adv\text{\_}g}$ is expressed as follows:

${ℒ}_{adv\text{\_}g}=-{\mathbb{E}}_{Z^{\prime }~p\left(Z^{\prime }\right)}\left[{\left(1-D\left(Z^{\prime }|y\right)\right)}^2\right]\text{.}$ (3)

The above three adversarial loss functions constitute the loss function of FedALN, which indirectly reduces the differences in feature distributions between clients by introducing reference distributions. Training federated prediction models on this basis will be beneficial to avoid the impact of inter-client domain distribution differences on model performance, the FedADG mechanism [16] is proposed based on FedALN. This mechanism sets the classifier $C(\cdot )$ for constructing a complete prediction model. It inputs real data features $Z$, it outputs prediction labels $C\left(Z\right)$, and uses the traditional cross-entropy loss as the loss function of the classifier for model training. The specific FedADG objective function is expressed as follows:

${ℒ}_{\text{FedADG}}={ℒ}_{adv\text{\_}d}+{ℒ}_{adv\text{\_}g}+{\lambda }_0{ℒ}_{adv\text{\_}f}+{\lambda }_1{ℒ}_{err},$ (4)

where ${ℒ}_{err}=\frac{1}{N}{\displaystyle {\sum }_{i=1}^{N}L\left(C\left(F\left(X_i\right)\right),Y_i\right)}$ denotes the classifier loss function, $L\left(\cdot ,\cdot \right)$ denotes the cross-entropy loss, $0<{\lambda }_0,{\lambda }_1<1$ and ${\lambda }_0+{\lambda }_1=1$.

Now consider that there are $K$ clients, each of which collects its local data $\left(X_k,Y_k\right)$ independently, then the objective function of the above FedADG can be reformulated into the following form:

${ℒ}_{\text{FedADG}}={\displaystyle \underset{k=1}{\overset{K}{\sum }}\frac{N_k}{N}}\left({ℒ}_{adv\text{\_}{d}_k}+{ℒ}_{adv\text{\_}{g}_k}+{\lambda }_0{ℒ}_{adv\text{\_}{f}_k}+{\lambda }_1{ℒ}_{er{r}_k}\right),$ (5)

where ${ℒ}_{adv\text{\_}{d}_k}$, ${ℒ}_{adv\text{\_}{g}_k}$, ${ℒ}_{adv\text{\_}{f}_k}$ denote the corresponding local ALNs adversarial losses, which can be obtained by using the client’s local data, and ${ℒ}_{er{r}_k}=\frac{1}{N_k}{\displaystyle {\sum }_{i=1}^{N_k}L\left(C\left(F\left(X_{i,k}\right)\right),Y_{i,k}\right)}$ denotes the local classifier loss function.

FedADG aligns the inter-client feature distributions based on FedALN, eliminating the impact of inter-client feature distribution differences on model performance, but this distribution alignment operation is performed on a client-by-client basis and does not take into account the domain bias and imbalance problems that may also exist within a single client. Since a single client may contain multiple source domain data and there usually is no significant domain labels for intra-client domain distribution alignment, we introduce the idea of SL [21] within a single client, which is used to extract stable features by removing spurious correlations between feature variables.

Specifically, in the k-th client, we use the stable predictive loss function ${ℒ}_{S_k}$ instead of the traditional cross-entropy loss ${ℒ}_{er{r}_k}$ as the loss function of the local classifier $C(\cdot )$, and the specific expression of the stable predictive loss function ${ℒ}_{S_k}$ is as follows

$\begin{array}{l}{ℒ}_{S_k}=\frac{1}{N_k}{\displaystyle \underset{i=1}{\overset{N_k}{\sum }}{w}_{i,k}^{*}}L\left(C\left(F\left(X_{i,k}\right)\right),y_{i,k}\right),\\ w_k^{*}=\underset{w_k\in {\Delta }_{N_k}}{\arg \min }{\displaystyle \underset{1\le m<h\le M_Z}{\sum }{\Vert {\hat{\Sigma }}_{Z_{k,:,m}{Z}_{k,:,h};w_k}\Vert }_F^2},\end{array}$ (6)

where ${\text{Δ}}_{N_k}=\left\{w_k\in {ℝ}_{+}^{N_k}|{\displaystyle {\sum }_{i=1}^{N_k}{w}_{i,k}}=N_k\right\}$, $w_{i,k}\in w_k$ denotes the $i$-th sample

weight of the $k$-th client, $F(\cdot )$ is the feature extractor, $C(\cdot )$ is the classifier, $L\left(\cdot ,\cdot \right)$ is the cross-entropy loss, $Z_{k,:,m}$ denotes the $m$-th feature variable in the feature space of the $k$-th client, and similarly, $Z_{k,:,h}$ denotes the $h$-th feature variable in the feature space of the $k$-th client, and ${}_{\widehat{\text{Σ}}}$ denotes the weighted partial cross-covariance matrix of the feature variables $Z_{k,:,m}$ and $Z_{k,:,h}$, which is used to measure the independence between $Z_{k,:,m}$ and $Z_{k,:,h}$. The specific expression of ${}_{\widehat{\text{Σ}}}$ is as follows:

$\begin{array}{c}{\hat{\Sigma }}_{Z_{k,:,m}{Z}_{k,:,h};w_k}=\frac{1}{N_k}{\displaystyle \underset{i=1}{\overset{N_k}{\sum }}[{\left(w_{i,k}u\left(Z_{k,i,m}\right)-\frac{1}{N_k}{\displaystyle \underset{j=1}{\overset{N_k}{\sum }}{w}_{j,k}}u\left(Z_{k,j,m}\right)\right)}^{\text{T}}}\\ \times \left(w_{i,k}v\left(Z_{k,j,h}\right)-\frac{1}{N_k}{\displaystyle \underset{j=1}{\overset{N_k}{\sum }}{w}_{j,k}}v\left(Z_{k,j,h}\right)\right)],\end{array}$ (7)

where

$u(\cdot )=\left(u_1(\cdot ),u_2(\cdot ),\cdots ,u_{n_A}(\cdot )\right),u_j(\cdot )\in {\mathscr{H}}_{\text{RFF}},\forall j,$

$v(\cdot )=\left(v_1(\cdot ),v_2(\cdot ),\cdots ,v_{n_B}(\cdot )\right),v_j(\cdot )\in {\mathscr{H}}_{\text{RFF}},\forall j.$

here, $Z_{k,i,m}$ denotes the $i$-th sample of the $m$-th feature variable from the $k$-th client. The vectors $u$ and $v$ are the Random Fourier Features mapping functions from the Random Fourier Feature function space ${\mathscr{H}}_{\text{RFF}}$, which is employed to capture the nonlinear dependence among the complex features, and the Random Fourier Feature function space ${\mathscr{H}}_{\text{RFF}}$ is expressed as follows:

${\mathscr{H}}_{\text{RFF}}=\left\{h:x\to \sqrt{2}\cos \left(\omega x+\varphi \right)|\omega ~N\left(0,1\right),\varphi ~\text{Uniform}\left(0,2\pi \right)\right\}.$ (8)

The core of SL is to remove the dependency between features by sample reweighting, i.e., learning the sample weights $w_k$ to make the weighted partial cross-covariance matrix ${}_{\widehat{\text{Σ}}}$ converge to 0. And the Frobenius norm is used to construct the independence test statistic, i.e., when

$I_{Z_{k,:,m}{Z}_{k,:,h};w_k}={\Vert {}_{\widehat{\text{Σ}}}\Vert }_F^2$ tends to 0, the feature variables $Z_{k,:,m}$ and $Z_{k,:,h}$

tend to be independent. Within a single client, for all feature variables, we learn the optimal $w_k^{\text{*}}$ by equation (0) to minimize the linear and nonlinear dependencies of all features within the client, and learn the stable features to eliminate the impact of feature distribution bias of intra-client on the model performance.

In summary, our proposed Stable-FedADG approach consists of two main aspects. First, within a single client, we use the stable prediction loss function ${ℒ}_{S_k}$ as the loss function of the local classifier $C(\cdot )$ to learn stable features, and second, we base our method on FedALN to adversarially align inter-client domain distributions. The method simultaneously eliminates the effects caused by intra-client and inter-client domain distribution bias to learn domain-invariant features in the federated context and further improves the generalization performance of the model. The specific Stable-FedADG objective function is expressed as follows:

${ℒ}_{\text{Stable-FedADG}}={\displaystyle \underset{k=1}{\overset{K}{\sum }}\frac{N_k}{N}}\left({ℒ}_{adv\text{\_}{d}_k}+{ℒ}_{adv\text{\_}{g}_k}+{\lambda }_0{ℒ}_{adv\text{\_}{f}_k}+{\lambda }_1{ℒ}_{S_k}\right).$ (9)

The detailed training process of the Stable-FedADG method is shown in Algorithm 1. The training process of Stable-FedADG consists of two main parts: server aggregation and client update.

Each of the three settings above simulates a different federated scenario. Following [16] , 70% of each client’s data is randomly chosen as the training set and 30% as the validation set. The network architecture in this paper is the same as that in FedADG [16] . For the PACS dataset, the feature extractor is constructed using ImageNet pretrained AlexNet and the OfficeHome dataset used ResNet50, and both without its last layer. Additionally, the classifier, distribution generator, and discriminator are all composed of two fully-connected layers. The output layer of the discriminator is configured to have a size of 1. In all experiments, the number of local iterations is set as $E_0=3$ and $E_1=5$ with a batch size of 64. The communication occurs 30 times. The local model is updated via the SGD method, with ${\lambda }_0=0.85$ and ${\lambda }_1=0.15$. The specific initial learning rates are shown in Table 1 , where $l{r}_f$ represents the learning rate of the feature extractor and the classifier, $l{r}_g$ and $l{r}_d$ are for the distribution generator and discriminator respectively, and $l{r}_w$ is for learning sample weights within the client.

In this paper, we present the experimental results of different algorithms in three settings, including the classic FL algorithms FedAvg [4] , FedProx [7] and SCAFFOLD [8] , and the existing FDG algorithms FedIIR [22] , FedSR [23] and FedADG [16] . Among them, FedADG is a method that improves the generalization performance of the model only by aligning the feature distributions among clients.

Table 2 and Table 3 display the test accuracies of different models on the PACS dataset and OfficeHome dataset, respectively. Each result is the average of three repeats. Analysing Table 2 and Table 3 , it can be seen that the test accuracies of Stable-FedADG and FedADG outperform FedAvg significantly in both the LODO setting and the mixed setting, indicating that when inter-client feature distributions vary greatly, adversarial aligning inter-client feature distributions are conducive to the learning of domain-invariant features and boost model generalization. Meanwhile, Stable-FedADG outperforms FedAvg and FedADG in all three settings, indicating that intra-client stable feature learning can further improve the FL global model’s generalization. Besides, Stable-FedADG also outperforms other FL and FDG algorithms in all three settings, indicating that considering simultaneously intra-client and inter-client domain distribution differences facilitates further improving the generalization performance of the model in various federated scenarios.

Continued