Dual-Layer Bidirectional Contrastive Federated Learning: Enhancing Feature Extraction and Personalization Simultaneously

Abstract

As a privacy-preserving machine learning paradigm, federated learning enables collaborative training across multiple clients without sharing data. However, when there are discrepancies in the data distributions of individual clients, the training accuracy often significantly decreases. Although existing personalized federated learning approaches aim to preserve each client’s capability on their respective datasets, this is still insufficient to fully address the problem. In this paper, we propose a dual-layer bidirectional contrastive federated learning method (FDCL), which enhances the feature extraction ability of the model while maintaining personalization, thereby improving the overall model performance. Specifically, we divide the model into a feature extraction layer and a classification layer. In the feature extraction layer, we reduce the distance between the client and the global model while enlarging the distance from the previous client version, which enhances the model’s feature extraction capability. In the classification layer, we reduce the distance between the client and the previous version while increasing the distance from the global model, which improves the model’s personalization ability. Extensive experiments demonstrate that our method outperforms other competitive methods in multiple image classification tasks.

1. Introduction

With the rapid development of deep learning technologies, the requirements for model accuracy have continuously increased. High-accuracy models often require large amounts of data for training [1,2]. However, in practical applications, much of the data involve users’ personal privacy, leading to widespread caution among users towards data sharing, resulting in the so-called “data silo” phenomenon [3,4]. This phenomenon not only restricts the flow of data but also affects the model’s performance in diverse data environments. Therefore, how to fully utilize decentralized data resources while protecting user privacy has become a key challenge in improving model performance.

To address this issue, federated learning was introduced [5,6,7,8,9,10,11,12]. The core idea of federated learning is to keep data on the client side and perform model training locally, avoiding the privacy risks of directly accessing raw data. Each client trains a model or gradient on its local dataset, generating local updates, which are then uploaded to the server. The server aggregates the model information uploaded by all clients to generate a global model. This global model attempts to simulate its performance as much as possible under the training conditions of different datasets. When the data distribution across clients is relatively consistent, the global model generated by federated learning can closely resemble a model trained directly on all datasets, ensuring high accuracy [13,14].

However, in reality, there are often significant differences in the data distribution across clients, known as data heterogeneity, which can negatively impact the model’s performance [15]. Traditional federated learning methods use a globally shared model, which often leads to the global model performing well on some clients’ datasets while performing poorly on others, especially when the data distributions differ greatly [16,17,18]. This reduces the model’s generalization ability and adaptability. Therefore, how to address data heterogeneity while ensuring privacy has become an urgent problem in federated learning.

In response, personalized federated learning methods have emerged [19,20,21]. The basic idea of personalized federated learning is to train a personalized model for each client to better adapt to the specific needs and data characteristics of different clients. This approach enables each client to have a dedicated model, improving accuracy on specific tasks. Although personalized learning methods can effectively enhance the performance of individual clients, many existing personalized algorithms tend to focus too much on enhancing personalization capabilities, neglecting the model’s generalization ability, i.e., its universality across different clients. Experimental research has shown that globally shared models typically have stronger feature extraction capabilities. Therefore, if clients can leverage the powerful feature extraction capabilities of the global model while retaining their own personalized features, the overall model’s performance can be effectively enhanced.

Inspired by the above research, this paper proposes a novel personalized federated learning algorithm designed to strengthen the model’s generalization ability while ensuring personalization. We decompose the model structure into two parts: the feature extraction layer and the classification layer. During the training process, we use contrastive learning to bring the client model’s representation in the feature extraction layer closer to that of the global model, allowing the client to learn stronger feature extraction capabilities from the global model. At the same time, we enhance the personalized representation of the client’s classification layer, pushing the classification layer’s representation farther from the global model, thereby maintaining the client’s personalized capabilities. This strategy effectively combines the powerful feature extraction abilities of the global model with the personalized needs of the clients, improving the overall model’s performance and balancing the enhancement of both personalization and generalization capabilities.

Our contributions are summarized as follows:

• We propose a strategy that divides the model into a feature extraction layer and a classification layer, and uses the abstract representations output by these two layers for contrastive training, without relying on data from other clients as negative samples;
• We introduce a novel method that improves the model’s feature extraction capability by minimizing the distance between the client’s feature extraction layer and the global model, while enhancing the model’s personalization capability by maximizing the distance between the client’s classification layer and the global model;
• We conduct extensive experiments on three datasets, and the results demonstrate that our method achieves competitive performance.

2. Related Work

2.1. Personal Federated Learning

Federated learning (FL), as an emerging distributed machine learning approach, aims to train a global model without the need to centralize data. However, due to the heterogeneity in data distribution and communication rates across clients [22,23], along with the limitation of having only a single global model, the federated learning average (FedAVG) method often performs poorly when dealing with non-independent and identically distributed (Non-IID) data. To address this challenge, personalized federated learning (PFL) has been introduced. PFL primarily encompasses two directions: training a global model with good generalization performance and customizing local models for each client. The former focuses on training a global model with broad generalization capability and fine-tuning it on each client after training to adapt quickly to the specific needs of the clients [24,25]. In contrast, the latter assigns different models to each client from the beginning of the training process to better accommodate diverse data distributions [20,26,27,28]. Common approaches for local model customization include model decoupling and model interpolation. Model decoupling achieves a balance between generalization and personalization by dividing the model into distinct functional modules, while model interpolation combines the global and local models by weighted summation, leveraging the strengths of both. In this paper, we propose an innovative method that decouples the model into feature extraction and classification layers, enhancing the performance of these components through contrastive learning, with the goal of simultaneously improving both the generalization ability and the personalization capacity of the model. The difference between personalized federated learning and conventional federated learning is shown in Figure 1.

2.2. Contrastive Learning

In recent years, contrastive learning techniques have been widely applied in unsupervised learning scenarios and have achieved leading performance in the unsupervised training of deep image models and graph models [29,30,31,32,33]. Many studies focus on training encoders to pull embeddings of similar samples closer while pushing dissimilar samples apart, thereby enhancing model performance. Additionally, since contrastive learning enables federated learning to achieve higher accuracy with less data, there have been many works in recent years combining the two approaches. Model-contrastive federated learning (MOON) [34] introduces model contrastive loss, aiming to align the current local model with the global model while pushing the current local model away from the previous round local model. Federated self-supervised learning (FedSSL) [35] processes unlabeled data collected from edge devices through self-supervised contrastive learning. Prototypical contrastive federated learning (FedProc) [36] uses contrastive loss to align local features with global prototypes, thereby reducing the representation gap, where the global class prototypes are distributed by the server. Bias-reduction federated learning (FedBR) [37] uses contrastive learning to align local and global feature spaces by utilizing local data and globally shared proxy data, thereby reducing the bias in local training. However, previous methods fail to simultaneously enhance the model’s feature extraction capability and personalization performance. We propose a new two-way contrastive learning algorithm. While aligning local and global feature spaces, it increases the distance between the classification layers of local and global models to ensure the improvement of personalized performance.

3. Preliminary

In order to demonstrate the validity of our research question and the proposed method, we provide detailed explanations of some necessary issues in this chapter.

3.1. Problem Formulation

Consider a federated learning system with a total of M devices, denoted by the set $\mathcal{M}=\{1,2,\dots ,M\}$. Each device $i\in \mathcal{M}$ holds a local dataset, which can be represented as

$${\mathcal{D}}_i=\left\{(x_{i,d},y_{i,d})\mid d=1,2,\dots ,\left|{\mathcal{D}}_i\right|\right\},$$

(1)

where $x_{i,d}\in {\mathbb{R}}^s$ refers to the d-th input data point, $y_{i,d}\in \mathbb{R}$ is the associated label, and $|{\mathcal{D}}_i|$ indicates the total number of samples on device i. The global dataset, formed by combining the datasets from all devices, is denoted as

$$\mathcal{D}=\bigcup _{i\in \mathcal{M}}{\mathcal{D}}_i,N=\sum _{i\in \mathcal{M}}\left|{\mathcal{D}}_i\right|,$$

(2)

where N is the overall number of pieces of data across all devices.

The objective of personalized federated learning is to collaboratively train a different model for each client using local datasets without exchanging raw data directly. We define the global loss function as follows:

$$\mathcal{L}\left(w\right)={\displaystyle \frac{1}{N}}\sum _{i\in \mathcal{M}}\sum _{(x_{i,d},y_{i,d})\in {\mathcal{D}}_i}𝓁(w;x_{i,d},y_{i,d}),$$

(3)

where $𝓁(·)$ denotes the loss function for a single input. For a particular device i, the local loss function is given by

$${\mathcal{L}}_i\left(w\right)={\displaystyle \frac{1}{|{\mathcal{D}}_i|}}\sum _{(x_{i,d},y_{i,d})\in {\mathcal{D}}_i}𝓁(w;x_{i,d},y_{i,d}).$$

(4)

The personalized federated learning problem is then formulated as the following optimization problem:

$$\underset{w}{min}\mathcal{L}\left(w\right)=\underset{w}{min}{\displaystyle \frac{1}{N}}\sum _{i\in \mathcal{M}}\left|{\mathcal{D}}_i\right|{\mathcal{L}}_i\left(w\right).$$

(5)

3.2. Motivation

In machine learning, most models can be divided into two layers: the feature extraction layer and the classification layer. The last layer of the model is generally defined as the classification layer, while the other layers are defined as feature extraction layers. The feature extraction layer is responsible for extracting abstract features from the input data and projecting them into a high-dimensional space. The classification layer is responsible for dividing and classifying the information in this high-dimensional space.

Our approach is based on the conclusion that, in federated learning, the global model’s feature extraction capability is superior to that of the client models, and the performance bottleneck of a single global model lies in its inability to use different classification heads to address the diverse data distributions of individual clients. In Figure 2, we use t-SNE to demonstrate the feature extraction results under different settings, validating the global model’s superior performance in feature extraction. Among them, the results of Figure 2a,b were obtained using the CIFAR-10 [38] with a Dirichlet distribution ($\alpha =0.5$).

Existing methods typically focus on processing the feature extraction layer or the classification layer individually, which allows the model to only enhance its feature extraction capability or personalization capability separately, but not both simultaneously. As a result, it is difficult to comprehensively improve the overall performance. Therefore, this paper proposes a novel algorithm that strengthens both the model’s feature extraction capability and personalization capability through a bidirectional contrast approach. Specifically, our method brings the feature extraction representations of the client and global models closer together to enhance the client’s feature extraction ability, while distancing the classification representations of the client and global models to enhance the client’s personalized expression ability.

4. Method

4.1. Framework Overview

Our algorithm consists of two parts: In the first part, we divide the model into the feature extraction layer and the classification layer. The last layer of the model is used as the classification layer, and the remaining layers serve as the feature extraction layers. During training, we collect the outputs of each sample from both layers for training in the next part. The second part involves bidirectional contrastive learning. Based on the high-dimensional representations obtained from the different layers in the first part, we use a contrastive learning method to make the feature extraction layer representation of the client close to that of the global model, while pushing the classification layer representation of the client away from that of the global model. This way, we enhance both the feature extraction capability and personalization of the model. The pseudocode of our algorithm is shown in Algorithms 1 and 2, and the schematic diagram of our algorithm is shown in Figure 3.

Algorithm 1 FDCL-Server

Input: global round T, all clients set c, client local steps Q, initial model $w_0$, dataset size $\mathcal{D}$ Output: a trained model $w_T$ 1: for t = 1 to T do 2:    $c_{[1:n]}\leftarrow samplen\mathrm{clients}fromc$ 3:    for i = 1 to n in parallel do 4:      $w_t^i\leftarrow \mathrm{Run}\mathrm{FDCL}\text{-}\mathrm{Client}(w_{t-1},Q)\mathrm{on}{c}_i$ 5:    end for 6:    $w_t\leftarrow {\sum }_{k=1}^n{\displaystyle \frac{|{\mathcal{D}}_i|}{\left|\mathcal{D}\right|}}{w}_t^k$ 7: end for

Algorithm 2 FDCL-Client

Input: all clients c, client learning rate $\eta$, local epoch E, temerature $\tau$, client’s dataset $\mathcal{D}$, local model w, global model $w^g$ Output: a trained model 1: $w^1\leftarrow \mathrm{local}\mathrm{model}w$ 2: $w^g\leftarrow \mathrm{global}\mathrm{model}{w}^g$ 3: for epoch $i=1to E$ do 4:    for each batch $\mathbf{b}=\{x,y\}$ of $\mathcal{D}$ do 5:      ${𝓁}_{\mathrm{task}}\leftarrow \mathrm{TaskLoss}(w^i\left(x\right),y)$ 6:      $P_{el},P_{eg},P_{ep}\leftarrow w_e^i\left(x\right),w_e^g\left(x\right),w_e^{i-1}\left(x\right)$ 7:      $P_{cl},P_{cg},P_{cp}\leftarrow w_c^i\left(x\right),w_c^g\left(x\right),w_c^{i-1}\left(x\right)$ 8:      ${𝓁}_{\mathrm{e}}=-log\left({\displaystyle \frac{exp\left(\mathrm{sim}(P_{el},P_{eg})/\tau \right)}{exp\left(\mathrm{sim}(P_{el},P_{eg})/\tau \right)+exp\left(\mathrm{sim}(P_{el},P_{ep})/\tau \right)}}\right)$ 9:      ${𝓁}_{\mathrm{c}}=-log\left({\displaystyle \frac{exp\left(\mathrm{sim}(P_{cl},P_{cp})/\tau \right)}{exp\left(\mathrm{sim}(P_{cl},P_{cg})/\tau \right)+exp\left(\mathrm{sim}(P_{cl},P_{cp})/\tau \right)}}\right)$ 10:      $𝓁\leftarrow {𝓁}_{\mathrm{task}}+\mu ({𝓁}_e+{𝓁}_c)$ 11:      $w^{i+1}\leftarrow w^i-\eta \nabla 𝓁$ 12:    end for 13: end for 14: Return $w^{E+1}$ to server

4.2. Model Partition

Similar to many existing studies, we divide the model architecture into two main components: the feature extraction layer $M_e$ and the classification layer $M_c$. Specifically, suppose that a model M has a total of N layers, with $N_e$ layers as its encoder; the definition of $M_e$ is as follows:

$$M_e=M_{[1:N_e]}.$$

(6)

Similarly, let us assume that the model starts to be a classification layer from the $N_c$-th layer; then, the definition of $M_c$ is as follows:

$$M_c=M_{[N_c:N]}.$$

(7)

In this structure, $M_e$ is responsible for extracting high-level abstract features from the raw input data. These features often contain the key information and patterns of the data, helping the model to understand the core structure of the data. The classification layer $M_c$, on the other hand, uses the features extracted by $M_e$ to produce the final classification output, mapping these abstract features to specific class labels through a series of operations. We recommend that, in encoder–decoder architectures, the partitioning of $M_c$ and $M_e$ should align with the original architecture’s inherent division. For other architectures, it is advised to designate the final layer as the $M_c$ layer while classifying all preceding layers as $M_e$ layers.

In the following discussion, we denote the output of $M_e$ as $P_e$, representing the intermediate feature representation obtained from the feature extraction layer. Meanwhile, the output of $M_c$ is denoted as $P_c$, which represents the final prediction result obtained through the classification layer.

The advantage of this structure lies in the separation of feature extraction and classification tasks, allowing the model to focus more on learning effective feature representations, thereby improving its performance on complex tasks. Additionally, this architecture provides greater flexibility for future task adaptation and model transfer.

4.3. Bidirectional Contrastive Learning

In this section, we introduce the bidirectional contrastive learning component of our model. To enhance feature extraction capabilities and improve personalization, our model performs contrastive learning not only with the global model but also with the previous version of the local model.

We define the components of the models as follows:

• The feature extraction layer of the global model is denoted as $M_{eg}$, and the classification layer is $M_{cg}$;
• The feature extraction layer of the current version of the client model is $M_{el}$, and the classification layer is $M_{cl}$;
• The feature extraction layer of the previous version of the client model is $M_{ep}$, and the classification layer is $M_{cp}$.

The outputs $P_e$ and $P_c$ are named according to the same convention. For example, $P_{eg}$ represents the output of $M_{eg}$, $P_{cg}$ represents the output of $M_{cg}$, and so forth for other pairs.

We first compute the loss ${𝓁}_{\mathrm{e}}$ for the feature extraction layer. Given the powerful feature extraction ability of the global model, the feature extraction layer $M_{el}$ of the current client model should learn the feature extraction capability of the global model as much as possible. Therefore, our goal is to minimize the distance between $M_{el}$ and $M_{eg}$, while maximizing the distance between $M_{el}$ and $M_{ep}$. The loss function for the feature extraction layer is as follows:

$${𝓁}_{\mathrm{e}}=-log\left({\displaystyle \frac{exp\left(\mathrm{sim}(P_{el},P_{eg})/\tau \right)}{exp\left(\mathrm{sim}(P_{el},P_{eg})/\tau \right)+exp\left(\mathrm{sim}(P_{el},P_{ep})/\tau \right)}}\right).$$

(8)

Here, $\tau$ is a hyperparameter. In order to enhance the personalized performance of the client models, we need to increase the distance between the classification head of each client model and the classification head of the global model while reducing the distance between the model’s classification head and the previous version’s classification head to maintain its performance on the local dataset. Our classification layer loss function is as follows:

$${𝓁}_{\mathrm{c}}=-log\left({\displaystyle \frac{exp\left(\mathrm{sim}(P_{cl},P_{cp})/\tau \right)}{exp\left(\mathrm{sim}(P_{cl},P_{cg})/\tau \right)+exp\left(\mathrm{sim}(P_{cl},P_{cp})/\tau \right)}}\right).$$

(9)

Our overall loss function is as follows:

$$𝓁={𝓁}_{task}+\mu ({𝓁}_e+{𝓁}_c).$$

(10)

Here, $\mu$ is a given hyperparameter. In Figure 4, we illustrate the differences between FDCL and other methods. By employing the loss function, we further increase the distance between the client and the server based on PFL, thereby enabling the client’s model to approach the local optimal solution more closely.

5. Experiment

In this section, we demonstrate the excellent performance of the FDCL algorithm by running it on different datasets and comparing it with other algorithms under various experimental settings. Additionally, we investigate the impact of different hyperparameters on the performance of the FDCL algorithm. The Dirichlet distribution in this experiment is automatically partitioned using the FLGO framework. By specifying the number of clients, the $\alpha$ value of the Dirichlet distribution, and the dataset, FLGO can quickly divide the entire dataset into federated datasets ready for experimentation.

5.1. Experimental Setup

The experimental hardware configuration includes Intel Xeon CPU E5-2690, RTX4090, and 128 GB RAM. The operating system is Ubuntu 20.04.4 LTS. We choose to implement our algorithm through the FLGO [39] framework. The FLGO framework is an efficient and easy-to-use federated learning algorithm framework that supports simulating federated learning under various parameters. We simulated virtual scenes of 20 clients on FLGO to simulate federated learning in real-world scenarios.

Datasets. In this chapter, we selected four commonly used image classification datasets, namely, MNIST [40], Fashion-MNIST [41], CIFAR-10, and CIFAR-100 [38], as well as a graph dataset, MUTAG [42], to evaluate the performance of the algorithm. MNIST is a classic handwritten digit dataset, containing 70,000 grayscale images, each with a resolution of 28 × 28 pixels. The dataset is divided into 60,000 training images and 10,000 testing images, covering 10 digit classes. Fashion-MNIST is a dataset similar to MNIST, containing 70,000 grayscale images of fashion products, with the same resolution of 28 × 28 pixels. The dataset is divided into 10 classes, including clothing, footwear, and other fashion items, with the same number of training and testing images as MNIST. The CIFAR-10 dataset contains 60,000 RGB color images, each with a resolution of 32 × 32 pixels, divided into 10 classes, with 6000 images per class. The dataset consists of 50,000 training images and 10,000 testing images. CIFAR-100 is an extended version of CIFAR-10, containing 60,000 color images with a resolution of 32 × 32 pixels, divided into 100 classes, with 600 images per class. The training and testing sets contain 50,000 and 10,000 images, respectively. The MUTAG dataset comprises 188 nitro compounds, with labels indicating whether each compound is aromatic or heteroaromatic. We visualize the data distributions under different $\alpha$ values in Figure 5.

To simulate different types of data distributions, we adopted the Dirichlet distribution, which is commonly used in federated learning, as the data partitioning strategy. The Dirichlet distribution is often used as a prior for data distributions to characterize the data distribution across different clients, where the hyperparameter $\alpha$ controls the level of uniformity in the data distribution. A larger value of $\alpha$ indicates a more uniform distribution, while a smaller value of $\alpha$ reflects more complex data heterogeneity. To better simulate the application scenarios of personalized federated learning, we set $\alpha$ = 0.1 and $\alpha$ = 0.5 to represent cases with severe data heterogeneity and mild data imbalance, respectively.

Models. For all computer vision datasets, a convolutional neural network (CNN) model was used in this paper for pattern recognition. As a fundamental architecture of neural networks, CNN excel at extracting local features from images, making them highly effective for image classification tasks [43,44]. The core idea of CNN is to utilize convolutional layers to extract local features from the input data and to progressively reduce the spatial dimensions of the features through pooling layers, thereby reducing the computational load and enhancing the model’s generalization ability.

The model used in our experiments consists of a feature extractor and a classifier. The feature extractor consists of two convolutional layers, each followed by a ReLU activation function [45] and a 2 × 2 maximum pooling layer. The first convolutional layer uses 32 “5 × 5” convolutional kernels to extract basic features from the input features, and the second convolutional layer uses 64 “5 × 5” convolutional kernels for feature extraction. After feature extraction, the model maps the extracted features onto a 512-dimensional vector through a fully connected layer. This vector is passed through the ReLU activation function and finally converted into the output of the model by another fully connected layer. The mathematical expression of the ReLU function is ReLU(x) = max(x, 0). By introducing this activation function, we can incorporate nonlinearity while mitigating the issues of gradient explosion and vanishing gradients [46,47].

It is worth noting that, in the experimental section, although we only employ this basic CNN for testing, our method is in fact applicable to all models whose final layer can serve as a classification head, such as a fully connected layer. Our approach solely modifies the loss function without altering the model architecture, which ensures its versatility to a certain extent.

Furthermore, for graph datasets, a default Graph Isomorphism Network (GIN) was used in the experiments in [48].

Hyperparameters. In all our experiments, we carefully selected and tuned the hyperparameters to achieve optimal performance. Specifically, the learning rate was set to 0.01, providing a good balance between fast convergence and stability, with a decay rate of 0.999 to gradually reduce the learning rate during training. The batch size was set to 32, a commonly used value that balances computational efficiency and performance. All shared hyperparameters are kept consistent, and, if there are specific hyperparameters, they are set according to the recommended values from the corresponding papers.

For federated learning, we configured the global communication rounds to 100, ensuring sufficient training iterations, with 5 local epochs per round to allow clients to learn from their local data. The client participation rate was set to 50%, meaning that half of the clients participated in each round, balancing communication costs and model performance. Additionally, the hyperparameter $\beta$ was set to 0.1 to regulate the balance between personalization and global model aggregation. The temperature hyperparameter $\tau$ for knowledge distillation was set to 0.5. Finally, the SGD optimizer was used to update the model parameters during training, ensuring efficient learning.

Baselines. To verify the performance of our method in asynchronous environments, we compared it with several classic and advanced federated learning algorithms across multiple datasets. The algorithms compared include FedAVG [13], DITTO [19], MOON [34], FedCP [49], and FedPer [24]. The introductions of these algorithms are as follows:

• FedAVG: FedAVG is a fundamental federated learning algorithm that performs global model averaging updates by training locally on the client and uploading updates to the server;
• DITTO: DITTO achieves fairness and robustness in heterogeneous environments by introducing a globally regularized multi-task learning framework;
• MOON: MOON uses contrastive learning to enhance the representation ability of global models and conducts personalized learning between clients through contrastive loss functions;
• FedPer: FedPer addresses the issue of non-independent and identically distributed data by learning personalized models for each client, thereby improving personalized performance;
• FedCP: FedCP dynamically learns a conditional policy to disentangle global and personalized information from its features, followed by independent processing through a global processing module and a personalized processing module.

In the experimental section, we kept the hyperparameters common to all methods the same, while setting the hyperparameters specific to each algorithm (if any) to the recommended values in the respective papers.

5.2. Performance Comparison

In this section, we conduct experiments on all baseline methods under fixed global rounds and record the final accuracy achieved by each algorithm, as shown in

Table 1. Final accuracy of different algorithms under various settings after 100 epochs. The bold numbers represent the highest values.

Methods	MNIST		FMNIST		CIFAR-10		CIFAR-100		MUTAG
	$\alpha =\mathbf{0}.\mathbf{1}$	$\alpha =\mathbf{0}.\mathbf{5}$	$\alpha =\mathbf{0}.\mathbf{1}$	$\alpha =\mathbf{0}.\mathbf{5}$	$\alpha =\mathbf{0}.\mathbf{1}$	$\alpha =\mathbf{0}.\mathbf{5}$	$\alpha =\mathbf{0}.\mathbf{1}$	$\alpha =\mathbf{0}.\mathbf{5}$	$\alpha =\mathbf{1}.\mathbf{0}$
FedAVG	92.43	96.74	78.95	79.32	43.36	57.41	18.76	20.17	68.42
DITTO	92.64	96.88	76.19	77.20	36.97	55.88	17.45	19.47	73.68
FedPer	92.98	97.03	78.90	80.19	41.74	48.92	18.67	20.33	63.16
MOON	93.05	97.35	77.57	79.62	42.06	57.06	18.82	20.97	78.95
FedCP	89.85	94.49	72.57	73.21	37.13	39.57	13.49	14.54	69.81
FDCL (Ours)	93.12	97.83	79.27	80.68	43.89	57.43	19.12	20.63	84.21

. Meanwhile, in Figure 6, we also recorded the training curves of all algorithms on different datasets. The experimental results show that, in almost all experimental settings, the final accuracy of our proposed algorithm FDCL is superior to other algorithms. It is noteworthy that, on the CIFAR-100 dataset (with $\alpha =0.5$), our algorithm performs slightly worse than the MOON. Through in-depth analysis, we attribute this phenomenon to the inherent characteristics of our algorithm design: compared to MOON, our method enhances client personalization, which may cause the classification head parameters of individual clients to deviate from the global optimum. In scenarios with weak data heterogeneity (i.e., relatively uniform client data distribution), where the global and local optima exhibit high consistency, this personalization enhancement strategy may lead to a slight performance degradation.

However, it is crucial to emphasize that personalized federated learning is primarily applied in scenarios with significant data heterogeneity. Therefore, FDCL consistently demonstrates superior performance in real-world applications.

5.3. The Impact of Participation Rate

The participation rate refers to the ratio of the number of clients participating in training during each communication round to the total number of clients. In this section, we investigate the performance of our algorithm under varying client participation rates. Extensive experiments were conducted using our algorithm on the CIFAR-10 dataset with $\alpha$ = 0.1 and different participation rates, and the results are presented in

Table 2. Performance comparison under same rounds on CIFAR-10. The bold numbers represent the highest values.

$\mathcal{P}$	20%	30%	50%	70%	100%
FedAVG	37.35	40.49	43.36	43.97	44.38
DITTO	35.79	36.34	36.97	38.96	41.41
FedPer	36.45	39.08	41.74	41.93	43.36
MOON	37.84	40.23	42.06	42.33	42.50
FedCP	34.81	35.52	37.13	38.08	38.57
FDCL (Ours)	38.42	41.54	43.89	44.62	44.92

. The performance of all algorithms improves as the participation rate increases. However, our algorithm, FDCL, consistently demonstrates strong performance across different client participation rates, which highlights its robustness to variations in client participation.

5.4. Effectiveness of the Number of Classification Heads

The number of classification heads refers to the number of layers in $M_c$. To assess the impact of the number of layers in $M_c$ on the performance of the FDCL algorithm, we conducted experiments using the CIFAR-10 dataset with $\alpha$ = 0.1 and evaluated different configurations of $M_c$. The experimental results, as shown in Figure 7, reveal that the accuracy of the FDCL model is highest when $M_c$ is set to 1. In contrast, when the number of layers in $M_c$ is increased, the algorithm’s performance exhibits a noticeable decline, with accuracy decreasing to varying degrees, depending on the number of layers. This trend suggests that increasing the depth of $M_c$ may negatively impact the model’s generalization ability. Based on these observations, we conclude that setting $M_c$ to 1 is optimal for achieving the best performance of the FDCL algorithm.

5.5. The Impact of Hyperparameter $\mu$

The hyperparameter $\mu$ controls the balance between the bidirectional contrastive loss and the total loss. In this section, we conduct extensive experiments to investigate the impact of $\mu$ on algorithm performance and to identify its optimal value. We performed experiments on the FDCL algorithm using different values of $\mu$ on the CIFAR-10 dataset, with the results shown in

Table 3. Performance comparison under the same number rounds on CIFAR-10. The bold numbers represent the highest values.

$\mu$	0.05	0.1	0.3	0.5	0.7	1.0	2.0
Acc	40.52	42.06	39.60	35.73	34.31	36.23	28.55

. In the experiments, the algorithm achieved the highest accuracy when $\mu$ = 0.1. When the value of $\mu$ either decreased or increased, the accuracy of the algorithm decreased to varying degrees. This is because, when $\mu$ is too small, the effect of bidirectional contrastive learning becomes less pronounced, and the performance approaches that of FedAVG. Conversely, when $\mu$ is too large, bidirectional contrastive learning dominates the training process, reducing the impact of task loss, which leads to suboptimal performance on the dataset. Therefore, we recommend setting $\mu$ to 0.1, as this results in better performance of FDCL.

5.6. The Impact of Local Epochs

Local epochs refers to the number of epochs that selected clients train locally during each communication round in federated learning. In this section, we conducted experiments on the FMNIST datasets with different algorithms and varying local epoch settings to investigate the impact of local epochs on algorithm performance. The experimental results are shown in Figure 8. We observed that appropriately increasing the local epochs within the range of 1 to 10 improved the algorithm’s final accuracy. This is because, when local epochs are too small, local training has not stabilized before the model participates in server aggregation, leading to poorer performance. However, we do not recommend setting a very high number of epochs, as this would significantly increase the federated learning training time and introduce the risk of overfitting in the context of personalized federated learning [50].

5.7. Influence of Client Quantity

To investigate the robustness of our method, we also compared its performance with other approaches under varying numbers of clients. In federated learning, increasing the number of clients without expanding the total dataset size leads to a decline in the final model’s accuracy. This occurs because the data volume per client decreases in such scenarios, resulting in a more dispersed distribution of training data, which adversely affects the training process. Our experiments validate this conclusion; the experimental results are shown in the

Table 4. Results comparison across varying numbers of clients on CIFAR-10. The bold numbers represent the highest values.

$\mathit{num}\_\mathit{clients}$	5	20	50
FedAVG	40.76	38.55	35.98
FedPer	41.08	38.73	36.32
MOON	41.22	39.45	37.04
FDCL (Ours)	42.67	40.32	38.73

. Notably, the proposed framework, as a personalized federated learning method, performs well under uneven data distribution and outperforms other personalized federated learning approaches. This advantage stems from our division of the model into two parts, which simultaneously enhances both feature extraction capability and personalization, thereby improving the model’s overall performance.

5.8. Ablation Study

We conducted an ablation study on the CIFAR-10 and CIFAR-100 datasets using a Dirichlet distribution with $\alpha$ = 0.01 to evaluate the impact of different modules. Since a more imbalanced data distribution can better demonstrate the advantages of our proposed method, we adopted this $\alpha$ value. The main components of our algorithm are ${𝓁}_c$ and ${𝓁}_e$. Therefore, in this section, we conduct an ablation experiment on these two components. The experimental results are shown in

Table 5. Ablation experiment results.

${𝓁}_{\mathit{e}}$	${𝓁}_{\mathit{c}}$	CIFAR-10	CIFAR-100
✗	✗	36.63	16.06
✗	✓	38.87	18.22
✓	✗	39.45	18.31
✓	✓	40.32	18.48

. Our experiments reveal that incorporating ${𝓁}_c$ improves the model’s accuracy on CIFAR-10 by 2.24%, while ${𝓁}_e$ yields a more significant enhancement (36.63→39.45), indicating that ${𝓁}_e$ plays a more crucial role in feature learning. A similar trend is observed on CIFAR-100. When both ${𝓁}_e$ and ${𝓁}_c$ are employed simultaneously, the model achieves optimal performance (CIFAR-10: 40.32%, CIFAR-100: 18.48%), demonstrating their synergistic effect. This is because ${𝓁}_e$ optimizes feature representation, whereas ${𝓁}_c$ further refines the classification boundaries.

6. Analysis

6.1. Communication Overhead

Our method only transmits the model itself during the client–server communication process, with no additional information involved in the transmission, thus introducing no extra communication costs.

6.2. Scalability Under Low-Resource Conditions

Since introducing contrastive learning on the client side increases computational costs, we analyzed the additional overhead that it incurs to evaluate its scalability under low-resource conditions.

We measured the execution time per round for different algorithms under identical conditions, where each round involved 100 clients participating in training, with each client running 5 local epochs. As shown in

Table 6. Time required for different algorithms to run one round.

Method	CIFAR-10	CIFAR-100	MNIST
FedAVG	53.76 s	87.37 s	56.47 s
FedPer	56.41 s	92.40 s	54.75 s
Ditto	89.83 s	116.09 s	82.01 s
MOON	66.77 s	95.26 s	64.24 s
FDCL (Ours)	68.32 s	96.95 s	67.47 s

, the results demonstrate that, although the introduction of contrastive learning incurs additional computational overhead, the impact remains limited. Thus, our method maintains satisfactory operational efficiency, even in resource-constrained scenarios.

6.3. Convergence Analysis

We provide a detailed convergence analysis for the FDCL algorithm under standard federated learning assumptions.

Consider a federated learning system with K clients, each with a local dataset ${\mathcal{D}}_k$. The global objective is

$$F\left(w\right)=\sum _{k=1}^K{p}_k{𝓁}_{\mathrm{task},k}\left(w_k\right),$$

(11)

where $p_k=|{\mathcal{D}}_k|/|\mathcal{D}|$ and $w_k=(w_{e,k},w_{c,k})$ are the feature extraction and classification layer parameters.

In FDCL, each client optimizes

$${𝓁}_k\left(w_k\right)={𝓁}_{\mathrm{task},k}\left(w_k\right)+\mu \left({𝓁}_{e,k}(w_k,w^g,w_k^{\mathrm{prev}})+{𝓁}_{c,k}(w_k,w^g,w_k^{\mathrm{prev}})\right),$$

(12)

where $\mu >0$ balances task and contrastive losses, $w^g$ is the global model, and $w_k^{\mathrm{prev}}$ is the previous local model.

6.3.1. Assumptions

• Smoothness: ${𝓁}_{\mathrm{task},k}$, ${𝓁}_{e,k}$, and ${𝓁}_{c,k}$ are L-smooth;
• Strong Convexity: ${𝓁}_{\mathrm{task},k}$ is ${\mu }_0$-strongly convex and ${𝓁}_k$ is ${\mu }_0$-strongly convex if ${𝓁}_{e,k}$ and ${𝓁}_{c,k}$ are convex;
• Bounded Heterogeneity: $\mathbb{E}\parallel \nabla {𝓁}_{\mathrm{task},k}{\left(w\right)-\nabla F\left(w\right)\parallel }^2\le {\sigma }^2$;
• Bounded Contrastive Gradient: $\parallel \nabla ({𝓁}_{e,k}+{𝓁}_{c,k})\parallel \le G$.

6.3.2. Local Update

For client k at round t, starting from $w_k^t=w_t^g$, perform Q SGD steps with step size $\eta$:

$$w_k^{t,q+1}=w_k^{t,q}-\eta \nabla {𝓁}_k\left(w_k^{t,q}\right).$$

(13)

Using $L^{\prime }$-smoothness,

$${𝓁}_k\left(w_k^{t,q+1}\right)\le {𝓁}_k\left(w_k^{t,q}\right)-\eta \left(1-{\displaystyle \frac{L^{\prime }\eta }{2}}\right){\parallel \nabla {𝓁}_k\left(w_k^{t,q}\right)\parallel }^2.$$

(14)

With $\eta \le 1/L^{\prime }$,

$${𝓁}_k\left(w_k^{t,q+1}\right)\le {𝓁}_k\left(w_k^{t,q}\right)-{\displaystyle \frac{\eta }{2}}{\parallel \nabla {𝓁}_k\left(w_k^{t,q}\right)\parallel }^2.$$

(15)

By ${\mu }_0$-strong convexity,

$${𝓁}_k\left(w_k^{t,q+1}\right)-{𝓁}_k\left(w_k^{*}\right)\le (1-\eta {\mu }_0)({𝓁}_k\left(w_k^{t,q}\right)-{𝓁}_k\left(w_k^{*}\right)).$$

(16)

After Q steps,

$${𝓁}_k\left(w_k^{t+1}\right)-{𝓁}_k\left(w_k^{*}\right)\le {(1-\eta {\mu }_0)}^Q({𝓁}_k\left(w_k^t\right)-{𝓁}_k\left(w_k^{*}\right)).$$

(17)

6.3.3. Local Deviation

Since $\nabla {𝓁}_k\left(w_k^{*}\right)=0$,

$$\nabla {𝓁}_{\mathrm{task},k}\left(w_k^{*}\right)=-\mu \nabla ({𝓁}_{e,k}+{𝓁}_{c,k}),\parallel \nabla {𝓁}_{\mathrm{task},k}\left(w_k^{*}\right)\parallel \le \mu G.$$

(18)

By strong convexity,

$$\parallel w_k^{*}-w_{\mathrm{task},k}^{*}\parallel \le {\displaystyle \frac{\mu G}{{\mu }_0}}.$$

(19)

6.3.4. Global Aggregation

The global update is

$$w_{t+1}^g=\sum _{k=1}^K{p}_k{w}_k^{t+1}.$$

(20)

For $Q=1$,

$$\mathbb{E}\parallel w_{t+1}^g-w^{*}{\parallel }^2={\parallel w_t^g-w^{*}\parallel }^2-2\eta \langle w_t^g-w^{*},\sum _{k=1}^K{p}_k\nabla {𝓁}_k\left(w_k^t\right)\rangle +{\eta }^2\mathbb{E}{∥\sum _{k=1}^K{p}_k\nabla {𝓁}_k\left(w_k^t\right)∥}^2.$$

(21)

Compute

$$\mathbb{E}\sum _{k=1}^K{p}_k\nabla {𝓁}_k\left(w_k^t\right)=\nabla F\left(w_t^g\right)+\mu \sum _{k=1}^K{p}_k\nabla ({𝓁}_{e,k}+{𝓁}_{c,k}).$$

(22)

$$\begin{array}{cc}\hfill \mathbb{E}\parallel w_{t+1}^g-w^{*}{\parallel }^2& \le \parallel w_t^g-w^{*}{\parallel }^2-2\eta \left(F\left(w_t^g\right)-F\left(w^{*}\right)+{\displaystyle \frac{{\mu }_0}{2}}{\parallel w_t^g-w^{*}\parallel }^2\right)\hfill \\ \hfill & +2\eta \mu G\parallel w_t^g-w^{*}\parallel \hfill \\ \hfill & +2{\eta }^2(\parallel \nabla F\left(w_t^g\right){\parallel }^2+{\mu }^2{G}^2+{\sigma }^2).\hfill \end{array}$$

(23)

6.3.5. Convergence Rate

Let $e_t=F\left(w_t^g\right)-F\left(w^{*}\right)$;

$$\mathbb{E}{e}_{t+1}\le (1-\eta {\mu }_0)e_t+\eta \mu G+{\eta }^2({\mu }^2{G}^2+{\sigma }^2).$$

(24)

Solve

$$\mathbb{E}{e}_T\le {(1-\eta {\mu }_0)}^T{e}_0+{\displaystyle \frac{\eta \mu G+{\eta }^2({\mu }^2{G}^2+{\sigma }^2)}{\eta {\mu }_0}}.$$

(25)

With $\eta =1/\left(L^{\prime }T\right)$,

$$\mathbb{E}[F\left(w_T^g\right)-F\left(w^{*}\right)]\le \mathcal{O}\left({\displaystyle \frac{1}{T}}\right)+\mathcal{O}\left(\mu \right).$$

(26)

Therefore, it can be proven that FDCL is convergent.

7. Conclusions

In this paper, we propose a novel dual-layer bidirectional contrastive federated learning method to address the challenges of data heterogeneity and model personalization in federated learning. By decomposing the model into two layers—the feature extraction layer and the classification layer—and applying bidirectional contrastive learning, our method enhances the feature extraction capabilities of the global model while maintaining strong personalization for individual clients. Experimental results across multiple datasets demonstrate that FDCL outperforms existing federated learning algorithms, such as FedAVG, DITTO, and MOON, in terms of classification accuracy and robustness, particularly under conditions of significant data heterogeneity.

The proposed approach effectively balances the benefits of global model generalization and client-specific personalization, making it a promising solution for real-world federated learning applications. In future work, we will aim to refine our model’s scalability and explore its performance in additional domains, while further optimizing key hyperparameters to enhance both efficiency and accuracy.