Federated Domain Generalization with Decision Insight Matrix

1 Introduction

Federated Learning (FL) has revolutionized the machine learning landscape by enabling collaborative model training across distributed clients while preserving data privacy [Liao et al., 2025]. In this paradigm, clients train local models on their private data, which are then periodically aggregated by a central server to form a global model, thereby circumventing the need for direct access to raw data [Li et al., 2020a; Li et al., 2024; Chen et al., 2025]. Although FL has demonstrated promising results in scenarios where data is independently and identically distributed (i.i.d.) [McMahan et al., 2017; Shen et al., 2025], real-world applications often present a more complex challenge: clients typically collect data independently, resulting in distinct distributions across different domains. Furthermore, during deployment, models frequently encounter data from previously unseen target domains, leading to a significant distribution shift problem [Huang et al., 2023; Liao et al., 2024; Wan et al., 2024]. This scenario, known as Federated Domain Generalization (FedDG) and illustrated in Figure 1, raises a fundamental challenge beyond traditional FL's data heterogeneity: how can federated models effectively generalize across diverse domains while maintaining privacy constraints?

∗ Corresponding author.

Traditional domain generalization methods focus on learning invariant relationships explicitly from data or representations [Hu et al., 2024; Fu et al., 2025a; Huang et al., 2025]. However, these methods require a centralized setting where data or representations are shared across clients, potentially compromising client privacy. To address this challenge, researchers have begun exploring federated domain generalization [Qi et al., 2024]. To investigate the invariance relationships between clients, various FL methods have been developed, focusing either on the feature level or the logit level [Qiao et al., 2024]. [Zhang et al., 2023] proposed a novel model aggregation method based on locally estimated generalization gaps, but their insight was limited to scenarios where each training domain was treated as a single client. [Huang et al., 2023] proposed a prototype aggregation method, FPL, from the feature perspective. By introducing consistency regularization, it aligns local features with prototypes to explicitly learn invariance in the feature extractor. [Guo et al., 2023] introduced a regularization approach FedIIR based on local empirical risk minimization, which aims to implicitly learn invariance by constraining the parameter space. However, these methods fail to integrate information from both the feature extractor and the classifier [Hu et al., 2023; Zhang et al., 2024; Qi et al., 2025; Fu et al., 2025b].

In response, we suggest a novel approach that emphasizes the decision-making process in the classifier layer of deep neural networks, rather than focusing solely on feature or classifier invariance. In conventional models, final output logits are computed by multiplying the penultimate layer's features with the classifier's weights. A deeper analysis reveals that [Chen et al., 2023] each logit value can be decomposed into the summation of element-wise products between the feature vector and corresponding weight vector . While most FL methods rely on feature prototypes [Wan et al., 2024; Bai et al., 2024] to learn invariance, the intermediate elementwise products (before summation) retain more fine-grained information. Viewing each product term as a contribution to its corresponding logit, we collect these contributions across all classes into a matrix. This matrix, which we term the "insight matrix", encapsulates the model's decision-making process for input classification. By exchanging insight matrices in FL, clients can share cross-domain knowledge representations, thus facilitating invariant learning.

Based on insight matrices decision, we propose a federated domain generalization model, named FedDIM, which provides a new theoretical and practical framework for invariant learning. We assume that different clients have heterogeneous input/output distributions, but a well-generalized model should make decisions based on cues that are consistent across samples and clients. Based on this intuition, we propose a regularization term that promotes similarity between each sample's insight matrix and the mean insight matrix of its corresponding class. Our approach offers two key advantages: First, it explicitly combines the semantic information of the encoder and classifier through fine-grained modeling, thus enhancing the model's ability to adapt to client-domain heterogeneity; Second, the insight matrix, as a natural byproduct of the logit computation process, has minimal computational overhead. Experimental results show that FedDIM improves the model's ability to learn invariant relationships across client domains. In summary, the main contributions of this paper are as follows:

• New perspective: In this paper, considering the privacy of FL, we propose a novel method built upon the concept of category-wise mean insight matrices. Bridging the gap between focusing only on feature invariance or logit invariance in FL, offering new insights into out-ofdistribution (OOD) generalization of FL.
• Simple yet effective algorithm: We introduce an efficient strategy that leverages insight matrices to enhance model robustness. Our method requires minimal modifications to the standard FedAvg algorithm while achieving superior performance. The lightweight implementation adds only a few lines of code while its effectiveness is theoretically guaranteed.
• Superior Performance: We conduct extensive experiments on multiple benchmark datasets. The results demonstrate that FedDIM consistently outperforms existing federated learning methods in Out-of-Distribution generalization.

2 Preliminaries

2.1 Problem Setting

In federated domain generalization, the data in each client is sampled from different domains. Let D denote the set of all client domains. We denote the training domain by Dtr = {D1, · · ·DM}, Dtr ⊆ D , where M is the number of training domains (or clients). Let X and Y represent the input space and target space, respectively, the sample contains K classes. Each client c ∈ D holds a local dataset denoted as {(xci , yc i )}nc i=1 , where nc is the number of samples. Let the loss function as L(f(x), y). Then, for each client c, the expected risk as Ec(f) = Exc ,ycL(f(xc ), yc ), and the global expected risk of model f denotes as EDtr (f) = EDtr Ec(f).

The ideal goal of FL training is to minimize the overall loss function on the dataset D. However, in practice, FL typically involves a large number of clients with heterogeneous data distributions, and only a subset of clients participate in the training. This introduces a distribution shift between the participating clients and those not seen during training, leading to the out-of-distribution (OOD) generalization problem [Qi et al., 2024; Xie et al., 2024]. Therefore, instead of optimizing the expected risk over the entire domain, we focus on optimizing the following empirical risk objective:

The federated OOD problem cannot be solved directly since not all potential clients are observed. This is more challenging than ordinary heterogeneous FL. To generalize to nonparticipating clients, the key to our study is how to learn the invariant relationship between inputs and goals.

2.2 Modeling Decision Insight Matrix

In the current deep model, the final output of the decision process involves two main steps: (1) the feature extractor: transforming the input from the original feature space to the feature embedding space, i.e., z = h(φ, x) ∈ RD : X → Z; and (2) the classifier: using the features to compute the final logits, i.e., o = g(w, z) ∈ RK : Z → ˆ Y. Thus the model can be written as f(θ) = g(w) ◦ h(φ), where θ = (φ, w).

In the FedDG scenario, clients primarily focus on samples from their local domains, while the server needs to aggregate inter-domain information from multiple clients to learn invariant relationships. Most existing research has focused on learning invariance by uploading features or logits to regularize the model, but these approaches have certain limitations: ❶ Solely focusing on feature invariance often overlooks the importance of classifier weights across different feature elements, which may lead to biased estimates of feature importance, thereby weakening the model's generalization ability.

• ❷ Although logits implicitly encode the relationship between

the classifier's weights, they only provide rough numerical values and lack fine-grained recognition of cross-domain generalizability, thus lacking deeper insight into the underlying decision-making process.

Since each logit value can be decomposed into the sum of the element-wise products of the feature vector and the corresponding weight vector, we argue that these intermediate product terms retain more granular information. Therefore, assuming the label set contains K classes, the logits can be represented as o = WT z ∈ RK , where W ∈ RD×K is the weight. For simplicity, we ignore the bias term of the classifier. Based on this decomposition, we treat each product term as the contribution of its corresponding logit. Thus, the logic value of class k can be expressed as

We aggregate the contributions of all categories into a matrix, called the "Insight Matrix", which is defined sa following:

The insight matrix I ∈ RD×K encapsulates key information about the model's decision-making process when classifying an input.

3 Methodology

We propose a federated domain generalization solution that utilizes the insight matrix as a key component for exchanging information between servers and clients. The core idea is that the insight matrix of samples from the same class is consistent with its corresponding average value, implying that the insight matrices for the same class across different domains should exhibit similarity. This ensures that the model makes classification decisions based on the same reasoning process. The framework of our method is illustrated in Figure 2.

3.1 Local Mean Class Insight Matrix

According to Eq. (2), the client generates an insight matrix for each sample. We define a local mean insight matrix I(k) c to represent the k-th class. For the c-th client, the average insight matrix is the mean of the insight matrices of the samples belonging to class k.

where ⊙ represents element-wise product, I(k) c ∈ RD×K and nc,k denotes the samples with class k in client c. We calculate the average insight matrix for each class and stack to get the local mean insight matrix Ic = [I(1) c , · · · , I(K) c ] ∈ RK×D×K .

3.2 Global Aggregation and Update

To generalize the global model to unseen clients, it is insufficient to simply aggregate the models of participating clients on the server. Although the clients possess domain information with different distributions, they share the same label space, which enables the participating clients to share a common embedding space. By aggregating the clients' insight matrices based on class information, we can learn invariant relationships in the federated domain generalization scenario. Thus, our aggregated global model and global insight matrix can be expressed as:

where θtc is the model trained by client c in round t, and C is the number of clients sampled per round.

In the FL training process, clients perform random sampling in each round. However, in the FedDG scenario, the

Algorithm 1 FedDIM

Input: total rounds T , local epochs E, total number of clients M, sampled number of clients C, learning rate η, hyperparameter for loss λ

Server executes:

• 1: Initialize global model θ and global insight matrix ¯I
• 2: for each round t = 1 · · · T do
• 3: Server samples subset C of clients
• 4: for each client c ∈ C in parallel do
• 5: {θtc , Ic} ← Clients updates(θt , ¯It )
• 6: end for
• 7: Update global model and calculate the global insight matrix by Eq. (5)
• 8: Update global insight matrix ¯It+1 by Eq. (6)
• 9: end for

Clients updates:

• 1: Initialize local model θtc = θt
• 2: for each local epoch e = 1 · · ·E do
• 3: Sample mini-batch in B:
• 4: Calculate the n-th sample insight matrix In
• 5: Calculate local loss by Eq. (7)
• 6: Update local model: θtc ← θtc - η∇Lc (θtc ; Bc)
• 7: end for
• 8: Calculate the mean class insight matrix Ic by Eq. (4)
• 9: return θtc and Ic

significant differences in client data distributions can lead to considerable variations in the aggregated insight matrix in each round. Therefore, we adopt a momentum to update the global insight matrix in each round. This method reduces the impact of noise in each iteration and balances new information with historical data, making the update process smoother and preventing overreaction to individual training data.

where m is a positive momentum value, and ¯I is initialized from the first iteration to compute the processed insight matrix, and ¯Ig calculated by Eq. (5).

3.3 Local Model Update

Clients update their local models to learn invariant relationships and generate consistent insight matrices across clients. To achieve this, we introduce a regularization term in the local loss, which enables the local model to capture the invariant relationship between data and targets during single-domain learning. Specifically, the loss is defined as follows:

where || · || is the l2 norm, Lce is the cross entropy loss, B is the number of samples in a mini-batch. Ii is the insight matrix for the n-th sample. ¯Ik is the global insight matrix corresponding to the k-th class distributed by the server.

4 Theoretical Analysis

This section presents the theoretical analysis demonstrating how our methods address the distribution shift problem. At first, we provide the following lemma, which is from [BenDavid et al., 2010] to bound the distribution divergence between two different domains.

Lemma 4.1. Let dH∆H(A, B) denotes the domain divergence between two domain distributions A and B. The expected risk gap between A and B is bounded as |EA(θ) - EB(θ)| ≤ 1 2 dH∆H(A, B).

Then, we consider the federated domain generalization setting where the training data follow the distribution D = ⋃|C| c=1 Dc, with Dc denoting the data distribution of client c among |C| total clients [Yan and Guo, 2025]. Each client maintains a training set Dc sampled from Dc with size nc = |Dc|, forming an overall training set D of n = ∑c nc samples and model aggregation weights {pc = nc n }. Let ED(θ) denote the expected risk on D and ˆ ED(θ) denote the empirical risk on D. We define ˆ EDc (θ), ˆ Ic(θ), as the two loss terms in Eq. (7).

Theorem 4.2. Let ˆ θ be the aggregated global model federatedly trained with the proposed overall loss function. Define θ∗T := argminθ ET (θ) and θ∗ c := argminθc ˆ EDc (θc) . Let H be a hypothesis space of VC dimension d. For any δ ∈ (0, 1), with probability at least 1 - δ, the generalization gap of the model ˆ θ on the unseen testing domain T has the following bound,

The complete proof is provided in the supplementary material. In this theorem, dH∆H (Dc, T ) represents the domain divergence between source domain Dc and target domain T , while ∆ represents the optimal model's residual error on both D and T . The theorem establishes that the generalization gap of the global model ˆ θ on target domain T is upper-bounded by two components: a weighted average term and the residual error. The weighted average term incorporates each client's empirical risk, insight matrix, and domain divergence bounds. Our proposed method aims to enhance the global model's generalization ability on T by explicitly minimizing the first two bound terms. The final term, which emerges from converting expected loss to empirical loss, is determined by the dataset size

5 Experiments

5.1 Experimental Setup

Datasets. To evaluate our approach, we conducted experiments in four datasets, RotatedMNIST [Ghifary et al., 2015], is a MNIST dataset of 7000 samples by rotating it at angles of 0°, 15°, 30°, 45°, 60°, and 75°, resulting in six different domains. PACS [Li et al., 2017], has 9991 images consisting of seven object categories in four domains (photo, art, cartoon, and sketch). VLCS [Fang et al., 2013], has 10729 images consisting of five object categories in four domains (Caltech101, LabelMe, SUN09, and VOC2007). OfficeHome [Venkateswara et al., 2017], is an image recognition dataset that includes 15,588 images of 65 classes from four different domains (art, clipart, product, and real-world). These are commonly used in the literature for domain generalization. We adhere to the experimental methodology outlined in FedIIR [Guo et al., 2023]. For all datasets, we perform "leave-one-domain-out" strategy [Gulrajani and LopezPaz, 2020], where we choose one domain as the test domain, train the model on all remaining domains, and evaluate it on the chosen domain. Each source domain is treated as a client. Following standard practice, we use 90% of available data as training data and 10% as validation data.

Considering the FL setting, we explore two scenarios based on the number of clients: the one-domain-one-client scenario and the one-domain-multiple-clients scenario. In the onedomain-one-client scenario, each training domain is treated as an individual client. In the one-domain-multiple-clients scenario, data from each training domain is randomly partitioned into multiple subsets, with each client containing data from one subset of a given training domain. The details of the data partitioning are provided in the Appendix C.1.

Baselines. We consider 2 classic federated methods FedAvg [McMahan et al., 2017], FedProx [Li et al., 2020b], and 4 state-of-the-art federated methods for domain generalization FedADG [Zhang et al., 2021], FedSR [Nguyen et al., 2022], FedIIR [Guo et al., 2023] and FedLGF [Yan and Guo, 2025] as baselines.

Implementation. We design dataset-specific models for each task. For the RotatedMNIST dataset, the feature encoder consists of four convolutional blocks, with ReLU activation, group normalization, and average pooling, followed by a linear classifier. During training, the batch size is 64. For the VLCS and PACS datasets, ResNet-18 is used as the feature encoder, while ResNet-50 is employed for the OfficeHome dataset. The classifiers for these three datasets consist of two fully connected layers. During training, the batch size is 32. For all datasets and scenarios, we set the communication rounds T to 100, with local iteration per round E=1 to accommodate limited local computational resources. Local models are updated using the SGD optimizer with a momentum of 0.9. The best parameters reported in the original paper were selected for the baseline, and the optimal hyperparameters of FedDIM were found by grid search. Each experiment was repeated 3 times and the average value was calculated.

5.2 Experimental Results

We evaluated all methods in two scenarios, where M denotes the total number of clients and C represents the number of sampled clients. Table 1 reports the results for the onedomain-one-client scenario, while Table 2 presents the results for the one-domain-multi-client scenario. For clarity, detailed results for specific domains are provided in the Appendix. C.

The experimental results demonstrate that, under the crossdomain client setting, our proposed method consistently outperforms other state-of-the-art baselines. In terms of average accuracy, we outperform the latest baseline FedLGF by 1.47% across all datasets. These observations validate the effectiveness of our method compared to existing baselines.

As the total number of clients increases, the performance of all methods declines significantly in the one-domainmulti-client scenario. In particular, FedADG and FedSR exhibit the most noticeable performance drops, likely because the increased number of clients makes it difficult to align the distributions across different source domains. To further validate this hypothesis, we extended the number of clients to 100, with the experimental results provided in Table 3 of the Appendix. Under this setting, both FedADG and FedSR perform worse than FedAvg. Additionally, when client data volume is large, the performance of FedIIR also drops significantly, potentially due to the reduced number of samples per category for each client caused by the increased number of clients. In contrast, our method exhibits the smallest performance degradation, highlighting its effectiveness and strong generalization capability in multi-client scenarios.

Loss surface visualization. We visualized the loss surface in a one-domain-one-client scenario using the "art" test domain in the PACS dataset, as shown in Figure 3. In the visualization, we used the global model as the origin and labeled the local models. This approach is consistent with the visualization technique in [Garipov et al., 2018]. It can be observed that compared to FedAvg, our local models all converge in the flat region of the loss surface, and in this way the global model induced is more generalizable. In addition, we find that the gap between the global and local models is much smaller, suggesting that our approach has a clear advantage in maintaining a consistent optimization objective across different domains.

Visualization. To better demonstrate the effectiveness of our method, we performed t-SNE [Van der Maaten and Hinton, 2008] visualization of the features z on the PACS dataset, as shown in Figure 4.

Compared to FedAvg, our method is clearer in clustering classification. On the training domain, FedDIM presents a clear block structure, indicating the effectiveness of the training process. When the model is generalized to the test domain (dark red part), the clustering structure is still obvious. This demonstrates the effectiveness of our method in generalizing to unseen distributions.

5.3 Ablation Study

To confirm the validity of our approach, we designed the following five variants for comparison to assess the independent impact of each component.

• W/λ = 0: The method degenerates to FedAvg.
• W/m = 0: The insight matrix I is generated and fixed for the initial pre-trained model.
• W/m = 1: The insight matrix I is dynamically updated by the current round.
• W/Fea: We replaced 2nd term of Eq. (7) using featureinvariance as Lim = 1B ∑k ∑{n|yi=k} ∥zi - ¯zk∥2 .

feature-invariance: focusing on the consistent expression of the features in the middle layer of the model.

• W/Log: We replaced 2nd term of Eq. (7) using logitinvariance as Lim = 1B ∑k ∑{n|yi=k} ∥oi - ¯ok∥2 .

Table 3 presents the experimental results of our different variants. It can be observed that the performance at W/m = 0 surpasses other variants, indicating that leveraging the insight matrix directly from the pre-trained model facilitates invariant decision-making. While W/m = 1 is outperformed by our momentum updating strategy due to the limited number of samples in a single batch, which cannot adequately consider all samples in the same class. Therefore, we design a scheme for dynamic updating based on historical information.

Moreover, it can be observed that both variants W/Fea and W/Log outperform the case where λ = 0. This is because the invariance constraint we adopt helps the model achieve better generalization and robustness to some extent. However, it is worth noting that this constraint may amplify the influence of irrelevant features, which have large values but correspond to small weights in the decision-making process, thus weakening the overall classification performance. Furthermore, focusing solely on logical invariance does not account for the varying contributions of individual features to the final decision. This can lead to small contributions being boosted to ensure that the summation equals the mean value, causing the model to emphasize irrelevant features and further degrading performance.

logit-invariance: focuing on the stability of the model's predicted probability, which is the behavior of the output layer.

5.4 Visualization of the Convergence Process

We present the convergence behavior of different methods in a one-domain multiple-client scenario. The learning rate for all methods is fixed to the same value. We report the accuracy of each method on both the validation and test sets, as shown in Figure 5. From the figure, it is evident that all methods demonstrate stable convergence on the validation set, with our algorithm achieving relatively superior performance across various datasets, particularly on the VLCS and OfficeHome datasets. Moreover, although the performance of different methods on the test set exhibits varying degrees of fluctuation, overall, the test accuracy of all methods stabilizes in the later stages of training, indicating good convergence behavior of the models on the test set.

5.5 Parameter Sensitivity Analysis

We investigated the effects of momentum coefficient and the loss trade-off parameter in a one-domain-multi-client scenario. We evaluated the sensitivity of the model using four datasets in the range λ ∈ {0.0001, 0.001, 0.01, 0.1, 1} and m ∈ {0.1, 0.3, 0.5, 0.7, 0.9}, as shown in Figure 6. The results show that FedDIM performs consistently when λ ∈ {0.0001, 0.001, 0.01, 0.1}, but performance degrades significantly at λ = 1. Furthermore, models with static momentum (m = 0 or m = 1) perform worse than those with dynamic momentum updates, highlighting the importance of momentum in improving generalization.

6 Conclusion

We study OOD generalization in federated learning through a novel perspective of fine-grained invariant relationship learning, which captures subtle yet crucial patterns across distributed domains. We propose FedDIM, a simple yet effective method that enhances OOD generalization by leveraging insight matrices to distill domain-invariant relationship from heterogeneous client data. Our theoretical analysis shows that FedDIM can effectively generalize to unseen domains by maintaining consistent relationships across different distributions, while empirical results demonstrate stateof-the-art performance on standard federated domain generalization benchmarks, including RotatedMNIST, PACS, VLCS and OfficeHome datasets.

Acknowledgments

The research is supported bythe National Key R&D Program of China (2023YFB2703700), the National Natural Science Foundation of China (62176269).

Contribution Statement

Tianchi Liao and Binghui Xie contribute equally to this work.