DualFed: Enjoying both Generalization and Personalization in Federated Learning via Hierachical Representations

In this paper, we consider a standard PFL setting which consists of a central server and $M$ distributed clients. For each client $m\in[M]$, there are totally $N_{m}$ samples $\{\bm{x}_{m}^{i},\bm{y}_{m}^{i}\}_{i=1}^{N_{m}}$ drawn from the distribution $\mathcal{D}_{m}$, where $\bm{x}_{m}^{i}\in\mathcal{X}_{m}\subseteq\mathbb{R}^{n}$ represents the raw input and $\bm{y}_{m}^{i}\in\mathcal{Y}_{m}\subseteq\{0,1\}^{C}$ represents the corresponding label, with $C$ denoting the total number of classes. In Non-IID scenarios within PFL, the data distributions are assumed to be heterogeneous across clients, indicating that $\mathcal{D}_{i}\neq\mathcal{D}_{j},\forall i,j\in\{1,2,\dots M\},i\neq j$.

The goal of a standard PFL setting is to develop a model $\psi_{m}(\cdot)$ parameterized by $\Theta_{m}$ for client $m$. The corresponding optimization objective can be expressed as:

where $\mathcal{L}(\Theta_{1},\ldots,\Theta_{M})$ represents the overall optimization objective for the PFL system, $\mathcal{L}_{m}(\Theta_{m})$ denotes the empirical risk for client $m$. In PFL, directly optimizing $\mathcal{L}(\Theta_{1},\ldots,\Theta_{M})$ is commonly infeasible as the clients cannot access the data on other clients. Therefore, a PFL training procedure typically involves the independently local updating performed on participating clients utilizing their own empirical risk and the model aggregation performed on the server. Specifically, for client $m$, its empirical risk is defined as:

with $\hat{\bm{y}}_{m}^{i}=\psi_{m}(\bm{x}_{m}^{i};\Theta_{m})$ representing the model’s prediction for $\bm{x}_{m}^{i}$, and $\ell:\mathcal{Y}\times\mathcal{Y}\rightarrow\mathbb{R}$ being the loss function that quantifies the prediction error (e.g., cross-entropy loss).

Once the local training on clients is completed, the participating clients upload their updated global parameters within the model to the server. The server then averages the parameters at corresponding positions to generate new global parameters. These global parameters are subsequently distributed to the clients for the next round of local updating. By iteratively performing local training and model aggregation, PFL facilitates collaborative model training without the need to share raw data from the clients.

For the sake of brevity, we occasionally omit the superscript denoting the sample index in subsequent sections of this paper. Additionally, we sometimes denote personalized parameters with the superscrip $p$ (e.g., $\Theta_{m}^{p}$), and global parameters with the superscript $s$ (e.g., $\Theta_{m}^{s}$), to clarify the expressions in the following sections.

As shown Figure 1 (a) and (b), in previous studies of PFL, the model $\Theta_{m}$ is commonly divided into an encoder $f_{m}(\cdot)$ and a classifier $h_{m}(\cdot)$ (Oh et al., 2022; Zhou et al., 2023; Zhu et al., 2022; Arivazhagan et al., 2019; Collins et al., 2021), parameterized by $\theta_{m}^{f}$ and $\theta_{m}^{h}$, respectively. The encoder $f_{m}(\cdot):\mathcal{X}_{m}\rightarrow\mathcal{Z}_{m}$ generally consists of a series of stacked convolutional layers. It maps the raw input $\bm{x}_{m}$ from $\mathcal{X}_{m}\subseteq\mathbb{R}^{n}$ into a representation space $\mathcal{Z}_{m}\subseteq\mathbb{R}^{k}$, which is denoted as $\bm{z}_{m}=f(\bm{x}_{m};\theta_{m}^{f})$. Here, $\bm{z}_{m}\in\mathcal{Z}_{m}$ denotes the representation generated from $\bm{x}_{m}$ utilizing the encoder $f_{m}(\cdot)$. Practically, the dimension of this representation is significantly smaller than that of the raw input, which implies that $k\ll n$. The classifier, $h_{m}(\cdot):\mathcal{Z}_{m}\rightarrow\mathcal{Y}_{m}$, generally includes a fully connected (FC) layer and a softmax layer. It generates the normalized predictions $\hat{\bm{y}}_{m}$ based on the representation $\bm{z}_{m}$, which is indicated as $\hat{\bm{y}}_{m}=h_{m}(\bm{z}_{m};\theta_{m}^{h})$.

Nevertheless, within the encoder-classifier architecture, only the representations after the encoder, referred to as the post-encoder representations, are used for decision-making. This approach can lead to a dilemma in PFL, as generalization and personalization are contradictory objectives, particularly in Non-IID scenarios. More specifically, to enhance model generalization, the post-encoder representations should capture shared information across varying data distributions among clients. On the other hand, enhancing model personalization requires these representations to capture specific information aligned with each client’s local data distribution. When the data distribution varies significantly across clients, these two types of information can be vastly different. Consequently, in this encoder-classifier architecture, ensuring that the post-encoder representations simultaneously meet these two conflicting objectives is a challenging task.

To address the dilemma mentioned earlier, we shift our focus on the process of representation extraction within the deep models. Advanced deep models are typically organized in a hierachical architecture. As shown in previous studies, these models initially extract generalized representations that are transferable across various data distributions (Zeiler and Fergus, 2014; Yosinski et al., 2014; Shwartz-Ziv and Tishby, 2017; Olah et al., 2017; Recanatesi et al., 2019; Bordes et al., 2023; Gui et al., 2023; Wang et al., 2023; Masarczyk et al., 2024). As the model progresses to deeper layers, it gradually discards irrelevant components and retains only information relevant to the specific task. In other words, both the generalized and personalized representations that PFL seeks for are already existed within the model. By leveraging these hidden generalized and personalized representations, we can achieve both high generalization and personalization in PFL. However, directly extracting these specific representations during the representation extraction phase is computationally challenging (Sariyildiz et al., 2023). Therefore, DualFed adopts a simpler strategy by incorporating a personalized projection network, which effectively decouples the generalized and personalized representations.

Figure 2 presents the framework of DualFed. It aligns with the standard training framework of PFL, which includes iterative local training on clients and global model aggregation on the server. The key innovation in DualFed, as compared to previous PFL methods, is the integration of a personalized projection network situated between the encoder and the classifier. We refer to this personalized projection network as $g^{p}_{m}(\cdot)$, with its parameters denoted by $\theta_{m}^{g,p}$. Functionally, this projection network, $g^{p}_{m}(\cdot):\mathcal{Z}_{m}\rightarrow\mathcal{U}_{m}$ is usually a MLP (multi-layer perceptron). By inserting this projection network, the representations produced by the encoder are not directly inputted into the classifier for prediction. Instead, they first pass through the projection network, which remaps them to a personalized representation space $\mathcal{U}_{m}\subseteq\mathbb{R}^{d}$. Formally, we represent this process as $\bm{u}_{m}=g(\bm{z}_{m};\theta_{m}^{g,p})$. For clarity, we term the representation before the projection network (i.e., $\bm{z}_{m}$) as the pre-projection representations, and the representation after the projection network (i.e., $\bm{u}_{m}$) as the post-projection representations.

Drawing on the hierarchical nature of deep model representation extraction, the pre-projection and post-projection representations in our framework exhibit distinct characteristics, aligning with the generalized and personalized objectives of PFL, respectively. Specifically, the pre-projection representations are separated from the final outputs by the projector network, meaning that they are not directly tied to the local tasks on each client. As previously studies have shown, these pre-projection representations are easier transferred across different data distributions (Wang et al., 2022; Sariyildiz et al., 2023). Therefore, in DualFed, the post-projection representations are fed into a global classifier $h_{m}^{s}(\cdot)$, which is parameterized by $\theta^{h,s}_{m}$. Additionally, to encourage the encoder to extract more generalized information, we let the encoder be shared among clients in DualFed. The predictions from this global classifier is expressed as:

Conversely, the post-projection representations are more closely aligned with the final outputs. This implies that these representations are more pertinent to accomplishing tasks related to the local data distribution. In DualFed, to effectively adapt to these local distributions, we utilize a personalized classifier $h^{p}_{m}(\cdot)$ for each client, parameterized by $\theta_{m}^{h,p}$, to adapt to the local data distribution. For a given input $\bm{x}_{m}$, the prediction generated by this personalized classifier can be expressed as follows:

During inference, the final predictions are derived by ensembling the outputs from both the global classifier and the personalized classifier. This process is expressed as follows:

By integrating a personalized projection network between the encoder and the classifier, DualFed effectively separates the contradictory optimization objectives inherent in PFL into distinct stages within the model. This approach resolves the conflict of pursuing contradictory objectives within the representations in the same stage, thereby can achieve a win-win situation between the model generalization and personalization.

In DualFed, each client updates the model for $E$ rounds using its own datasets after revceiving the global models from the sever. In order to fully exploit the hierarchical characteristics of deep model representations and achieve the optimization objectives of PFL, we introduce a stage-wise training procedure for local clients.

At the first stage, we freeze the global classifier and training the main branch of the model. The main branch comprises the global encoder, the personalized projector, and the personalized classifier, with their parameters collectively represented as $\overline{\Theta}_{m}:=\{\theta^{f,s}_{m},\theta^{g,p}_{m},\theta^{h,p}_{m}\}$. This stage allows the model to extract both generalized and personalized representations. To ensure the model’s effectiveness in accomplishing local tasks, we employ cross-entropy loss as the classification loss, as indicated in the following equation:

where $\bm{y}_{m}^{i,c}$ denotes the value at $c^{th}$ class of the one-hot ground-truth label of the $i^{th}$ sample on client $m$, $\hat{\bm{y}}^{p,i,c}_{m}$ represents the normalized prediction probability of $c^{th}$ classes of the $i^{th}$ sample on client $m$ from the personalized classifier.

As the post-projection representations are tailored to adapt to the local data distribution, we further enhance its discrimination by implementing supervised contrastive loss (Khosla et al., 2020), as demonstrated in the following equation:

where $A$ is the full set of samples, $A(i)$ consists of samples in $A$ that belong to the same class as $\bm{x}_{m}^{i}$, $\odot$ is the cosine similarity, and $\tau\in\mathcal{R}^{+}$ is the temperature coefficient.

The optimization objective at this stage is then defined as:

where $\lambda$ denotes the hyperparameter used for balancing these two loss terms, $t$ denotes the local updating epochs.

After updating $\overline{\Theta}_{m}$, we freeze its parameters and train the global classifier using the pre-projection representations to fulfill the local task, as represented in the following equation:

where $\bm{y}_{m}^{i,c}$ denotes the value at $c$ class of the one-hot ground-truth label of the $i_{th}$ sample on client $m$, $\hat{\bm{y}}^{s,i,c}_{m}$ represents the normalized prediction probability of $c$ classes of the $i_{th}$ sample on client $m$ from the global classifier.

The optimization objective in this stage can be expressed as:

In DualFed, both optimization objectives, as described in Eqs. (8) and (10), are optimized using mini-batch stochastic gradient descent (SGD). As evidenced by our experiments, this stage-wise optimization strategy diminishes the impact of local tasks on the pre-projection representations, thereby effectively preserving its generalization.

Once the local updating process is complete, the clients send their global encoder and classifier parameters to the server. The server then aggregates these parameters using the following equation:

Following the model aggregation, the server broadcast the updated model back to the clients to for subsequent local training.

Our experiments are conducted on PACS (Li et al., 2017), DomainNet (Peng et al., 2019), and Office-Home (Venkateswara et al., 2017), each of them consists of several domains. PACS includes $4$ distinct domains: Photo (P), Art Painting (A), Cartoon (C), and Sketch (S), each featuring images from $7$ common categories. DomainNet encompasses $6$ distinct domains: Clipart (C), Infograph (I), Painting (P), Quickdraw (Q), Real (R), and Sketch (S). Initially, each domain comprises $345$ classes, but for our study, we narrow this down to $10$ commonly used classes to create our experimental dataset. Office-Home contains $4$ distinct domains: Art (A), Clipart (C), Product (P), and Real-World (R), each containing $65$ classes. We retain all classes within Office-Home to conduct a comprehensive evaluation of DualFed on a larger-scale dataset.

For these datasets, we select the images from a single domain to form the dataset of an individual client. In both PACS and DomainNet, we choose a subset of $500$ training images per client from the same domain for the training dataset. For Office-Home, we set the number of training samples to $2,000$ for the Clipart, Product, and Real-World domains. In the case of the Art domain, the number is limited to $1,942$, matching the total number of samples available in this domain. All the images from the test dataset are reserved for evaluation for these datasets. We apply random flipping and rotational augmentations to these images during the training.

We perform a comparative analysis against the following methods, including FedAvg(McMahan et al., 2017), FedProx(Li et al., 2020b), FedPer(Arivazhagan et al., 2019), FedRep(Collins et al., 2021), LG-FedAvg(Liang et al., 2020), FedBN(Li et al., 2021), FedProto(Tan et al., 2022a), SphereFed(Dong et al., 2022), Fed-RoD(Chen and Chao, 2022), FedETF(Li et al., 2023). Additionally, the SingleSet method, where separate models are trained and tested for each client using only their private data, is also used for comparison in our experiments.

The adopted encoder is from the one of the ResNet18 model pretrained on the ImageNet dataset (He et al., 2016). It is followed by a projector network, which consists of an FC network with the architecture: [Linear($512$, $256$) - ReLU - BN - Linear($256$, $512$) - BN]. To ensure uniform model capacity, all compared methods employ this Encoder-Projector architecture for representation extraction.

The learning rate is set $0.01$, with a momentum of $0.5$, for all methods except SphereFed . For SphereFed, we set the learning rate to $1.0$ for Office-Home and to $0.1$ for both DomainNet and PACS. The batch size is set to $256$ for all methods. The epoch of local updating is set to $1$ for all methods except FedRep. For FedRep, it has a total of $5$ local epochs, with the initial 4 epochs focusing on classifier optimization and the last epoch on encoder and projector optimization. The total global rounds is set to $300$.

The other hyperparameters for different methods are selected by grid searching. To mitigate cross-domain interference and potential privacy issues related to BN layers, we localize the running-mean and running-var components within these layers for all methods.

To ensure the reliability of our results, each experiment is repeated $5$ times with random seeds: $\{0,1,2,3,4\}$. The subsequent sections will detail the mean and standard deviation of the highest test accuracy achieved during FL training.

Tables 1 - 3 showcase the experimental results of our proposed DualFed alongside other FL methods on the PACS, Office-Home, and DomainNet datasets, respectively. Notably, DualFed presents a significant performance gain in comparison to these SOTA methods.

Interestingly, the SingleSet model stands out as a strong benchmark, despite not collaborating with other clients, particularly in simpler domains, such as the Quickdraw domain in the DomainNet dataset. The underlying reason is that these simpler domains requires less complex semantic information for downstream tasks. In these cases, a personalized encoder’s representations are sufficient, and collaboration for extensive semantic extraction might be unnecessary or even detrimental. This observation is supported by LG-FedAvg’s performance, which, while also utilizing a personalized encoder for representation extraction, outperforms SingleSet by leveraging collaborative training for a global classifier.

However, as the complexity within a domain increases, such as in the Infograph domain of DomainNet, the benefits of sharing the encoder among clients become apparent. This collaborative approach allows the encoder to extract more nuanced semantic information from the raw data, improving overall model performance, as demonstrated by the results of FedAvg and FedProx. FedRep and FedPer, employing a personalized classifier to adapt the representations from the global encoder, often outperform FedAvg and FedProx. However, these methods primarily leverage the global encoder’s representations and do not fully utilize personalized information to cater to the local data distribution on individual clients.

FedProto significantly improves model performance by aligning representations from different clients within a unified representation space. Nonetheless, this alignment can result in a loss of semantic information pertinent to local tasks due to varying data distributions across clients. This issue is even more pronounced in models like SphereFed and FedETF, which employ a predefined classifier for representation alignment and lack specific semantic information about local data.

Fed-RoD adopts an architecture similar to ours, utilizing both global and personalized classifiers to capture generalized and personalized information. However, it attempts to utilize representations at the same stage, posing challenges in simultaneously meeting these two contradictory objectives. In contrast, our proposed method strategically separates these two conflicting objectives into different stages of the model. This division allows us to achieve both generalization and personalization more effectively, ultimately resulting in superior performance across a wider range of scenarios.

Comparison of global and personalized classifiers. To gain a deeper understanding of the behavior of the global and personalized classifiers, we compare their accuracy, individually and in combination, during training. Figure 3 shows the corresponding experimental results on DomainNet. It is evident that personalized classifier significantly surpasses the global one, owing to its better alignment with local data distributions. Nevertheless, the accuracy of the local classifier can be significantly improved by combining its predictions with those from the global classifier. This enhancement is particularly notable in complex domains, such as Infograph. Conversely, in simpler domains like Quickdraw and Sketch, the benefit of combining classifiers becomes less pronounced. This occurs because, in simpler domains, the representations extracted by the personalized projection network are sufficient for each client’s local tasks, thereby reducing the necessity for more diverse representations from the global encoder.

Visualization of generalized and personalized representations. To intuitively understand the generalized and personalized representations, we utilize t-SNE (Van der Maaten and Hinton, 2008) for visualization. Figure 4 illustrates the visualization of both the generalized and personalized representations on DomainNet dataset. In these visualizations, different colors indicate different classes. It is noticeable that the personalized representations are more discriminative than the generalized ones, yet they exhibit lower consistency across clients. This demonstrates that DualFed can effectively separate the representation extraction process into two distinct stages, each characterized by high levels of generalization and personalization, respectively.

Quantitative evaluation of generalized and personalized representations. We employ two metrics to quantitatively evaluate the evolution of generalized and personalized representations during training. To quantify the personaliztion of representations on clients, we adopt the class-wise separation in (Kornblith et al., 2021). Additionally, we adopt the linear centered kernel alignment (CKA) (Kornblith et al., 2019), to measure the generalization ability of representation. Figure 5 presents the varying of $separation$ during the training. The personalized representations can achieve higher class separation compared with the generalized representations. However, as shown in Figure 6, the similarity between clients of generalized representations is significant higher that that of the personalized representations.

Comparison of Training Strategy. DualFed employs a stage-wise training strategy, ensuring that the pre-projection representation remain undisturbed by specific local tasks, thereby maintaining its generalization. Here, we compare this training strategy with the one that training all parameters simultaneously. As shown in Table 4, when $E$ is relatively small (i.e., $E=1$), simultaneous training can, in fact, outperforms stage-wise training. However, as $E$ increases (i.e., $E=20$), simultaneous training lead to a obvious performance drop in PACS and DomainNet. This trend can be attributed to the fact that an increased number of local epochs causes the pre-projection representations to align more closely with the local task, thereby reducing their generalization.

Effect of Projector Architecture. We investigate the impact of the architecture of the projection network in three key aspects: the model depth ($D$), the dimension of hidden layers ($H$), the impact of BN layers. The corresponding results are shown in Table 5. While increasing $D$ can lead to more generalized pre-projection representations, it simultaneously reduces their discriminative power. Therefore, it is advisable to select an optimal $D$ that maintains a balance in the discriminative and generalized ability of the pre-projection representations. Increasing $H$ can enhance the model performance in most times. The importance of BN layers becomes more pronounced as the scale of the dataset increases.

Effect of Position of Global Classifier. In DualFed, we employ a global classifier for generalized representations and a personalized classifier for personalized representations. Here we conduct experiments when placing the global classifier after the projector. In these experiments, we maintained a shared encoder and investigated two configurations: sharing the projection network (DualFed-G) and personalizing it (DualFed-P). As indicated in Table 6, removing the global classifier to the same stage as the personalized classifier results in a significant performance decrease. This underscores the importance of the representations at different stages, as they provide complementary information.

Effect of Personalized Layers. Table 7 presents the model performance with different personalization strategy. The results indicate that combining a global encoder with a personalized projection network significantly enhances model performance, as it integrates both generalized and personalized information.

Communication Costs. We assess the communication costs by using the total number of model parameters transferred to reach a predefined target accuracy during training. For PACS, DomainNet, and Office-Home, the target accuracies are set to 85%, 75%, and 70%, respectively. As illustrated in Table 8, DualFed outperforms other methods by achieving the same target accuracy with lower communication costs, showcasing its practical efficiency.

Effect of Hyper Parameters. We conduct experiments using various hyperparameters, including the temperature coefficient ($\tau$), the loss balance coefficient ($\lambda$), and the number of local epochs ($E$). As depicted in Figure 7, we observe that as $\tau$ increases, its effectiveness in distinguishing between different classes diminishes, thereby losing the advantage of contrastive loss. Figure 8 presents the test accuracy with varying $\lambda$. Setting $\lambda$ to 0 is equivalent to training the model solely with cross-entropy loss. Increasing $\lambda$ enhances the distinctiveness and relevance of personalized representations to the local task, which, in turn, improves model performance. Figure 9 shows the test accuracy with different local epochs, it illustrates that DualFed consistently surpasses FedAvg across various local epochs, demonstrating its robustness to local epochs.