A Personalized Federated Learning Method Based on Knowledge Distillation and Differential Privacy

Abstract

Federated learning allows data to remain decentralized, and various devices work together to train a common machine learning model. This method keeps sensitive data local on devices, protecting privacy. However, privacy protection and non-independent and identically distributed data are significant challenges for many FL techniques currently in use. This paper proposes a personalized federated learning method (FedKADP) that integrates knowledge distillation and differential privacy to address the issues of privacy protection and non-independent and identically distributed data in federated learning. The introduction of a bidirectional feedback mechanism enables the establishment of an interactive tuning loop between knowledge distillation and differential privacy, allowing dynamic tuning and continuous performance optimization while protecting user privacy. By closely monitoring privacy overhead through Rényi differential privacy theory, this approach effectively balances model performance and privacy protection. Experimental results using the MNIST and CIFAR-10 datasets demonstrate that FedKADP performs better than conventional federated learning techniques, particularly when handling non-independent and identically distributed data. It successfully lowers the heterogeneity of the model, accelerates global model convergence, and improves validation accuracy, making it a new approach to federated learning.

1. Introduction

The rise of big data has transformed information processing and knowledge acquisition. Data-driven machine learning techniques offer great potential to improve efficiency across industries, driving innovation in areas such as healthcare and finance. Machine learning has undoubtedly become a key driver of modern technological innovation. However, this wave of innovation also poses privacy and security challenges, especially when it comes to sensitive data such as personal medical records [1] and financial information [2]. Ensuring that personal privacy is not compromised while reaping the benefits of data-driven machine learning has attracted extensive attention from academia, industry, and policy makers.

In this context, federated learning (FL) [3] is a novel machine learning technique that enables multiple participants to work together to jointly train a generic model without exchanging data. This approach optimizes the value of distributed data while protecting privacy and reducing the chance of data leakage. It provides an attractive solution for data-rich but highly privacy-sensitive industries [4].

Although FL avoids direct data exchange, model parameters may still reveal sensitive information. Malicious participants or attackers may be able to infer data features by analyzing gradients, which can be inferred or reconstructed even if the data remain localized [5]. Therefore, further enhancing privacy protection in FL to prevent various potential privacy violations has become a research priority in this area. Abadi et al. [6] were among the first to combine differentially private techniques with deep learning, proposing a differentially private (DP) stochastic gradient descent optimization algorithm. While they found that the loss in model quality when maintaining DP may be acceptable, this loss may still be significant for some applications. Noble et al. [7] introduced a differentially private FL algorithm tailored for highly heterogeneous data and analyzed the tradeoffs between privacy and model utility. However, applying the same level of noise to all clients may not provide the optimal balance between utility and privacy. By introducing noise prior to model aggregation, Kang et al.’s [8] novel differential privacy-based methodology effectively stops information leakage and offers a thorough convergence analysis. However, in real world scenarios, data are often non-independent and identically distributed (non-IID), which may affect the actual performance of the algorithm. Thus, when applying DP to FL, how to balance the effectiveness of model training with privacy protection becomes an urgent problem to solve.

In addition, another major difficulty is the non-IID character of data in FL contexts [9]. In the real world, data originating from different environments and participants lead to significant variations in data distributions, adding additional difficulties to model training [10,11]. Effectively addressing this issue to improve model generalization and overall performance is another important research direction in the field of FL. While privacy protection was taken into consideration, Jeong et al. [12] suggested using joint distillation and augmentation to maximize machine learning training on devices. However, it is still unclear how to ensure the privacy of the created data and quantify privacy protection. Fallah et al. [13] introduced a personalized FL approach using the model-agnostic meta-learning framework, which allows users to fine-tune models according to their local data for personalization. However, meta-learning introduces additional computational steps, and although approximation algorithms have been proposed to reduce the computational burden, they still impose significant computational costs on resource-constrained devices. Tursunboev et al. [14] proposed a new hierarchical federated learning algorithm that was specially used for edge-aided unmanned aerial vehicle networks and demonstrated strong performance in handling non-IID data. However, this method may face challenges in communication overhead and resource management in a wider range of application scenarios. Thus, when dealing with FL in non-IID environments, finding an ideal solution that enhances model performance and generalization, improves communication efficiency, and ensures privacy protection remains an unresolved challenge. This highlights the ongoing need to develop innovative approaches to effectively address these multifaceted issues.

In summary, existing research faces the following major challenges in integrating DP and knowledge distillation into FL. First, although FL avoids direct exchange of data, the exchange of model parameters and updates may still expose sensitive information, posing the risk of inference or reconstruction attacks. Second, the integration of DP and knowledge distillation is difficult to balance between privacy-preserving capabilities and model validity, and the adoption of uniform noise will hinder performance improvement in the later stages of training. Third, the non-ID characteristics of data in FL environments complicate the integration of knowledge distillation, and the inconsistency of data distribution among clients can hinder the knowledge transfer process and affect the generalization and performance of the model. This work proposes a customized FL strategy (FedKADP) based on knowledge refinement and DP to overcome these issues. Its main contributions are as follows:

(1)

A personalized FL framework combining knowledge distillation and DP is proposed to cope with non-IID data environments, incorporating knowledge distillation in the training of local models to enhance the model’s generalization capabilities and applying differential privacy mechanisms to noise perturbation of the output parameters;

(2)

A bidirectional feedback mechanism is proposed that adaptively adjusts the knowledge distillation and DP parameters according to the performance of the model and the user’s privacy requirements, balancing the relationship between privacy protection and model performance to achieve the optimal model performance under the demanded privacy intensity;

(3)

It is experimentally demonstrated that FedKADP improves the robustness and accuracy of data processing in non-IID data environments, maximizes model performance with the same privacy-preserving strength, and significantly reduces the communication cost.

2. Relevant Definitions

2.1. Federated Learning

Federated learning is a distributed machine learning technique that eliminates the requirement to aggregate data from various local data sources into a single central location in order to train models. In FL, training is done locally, and the global model can be updated by aggregating changed model parameters and sending them to a central server using a communication protocol. Using a decentralized learning approach, model training is possible while data are locally kept, reducing the requirement for data transmission and protecting user privacy. Formally, the optimization objective of traditional FL can be expressed as follows:

$${\min }_{w}F(w)={\min }_w{\displaystyle \sum _{k=1}^K\frac{n_k}{n}}{F}_k(w).$$

(1)

Here, $w$ represents the global model parameters; $K$ is the number of clients participating in FL; $n_k$ is the number of data samples of the $k$-th client; $n$ is the sum of the data samples across all clients, which is denoted as $n={\displaystyle {\sum }_{k=1}^K{n}_k}$; and $F_k(w)$ is the local loss function of the $k$-th client.

Each client independently updates its model parameters based on its local dataset and then sends these updates to a central server. The central server aggregates these updates into a new global model. This aggregation process typically uses a weighted average, which can be expressed as follows:

$$w^{t+1}={\displaystyle \sum _{k=1}^K\frac{n_k}{n}}{w}_k^{t+1}.$$

(2)

Here, $w_k^{t+1}$ is the model parameters of the $k$-th client at time step $t+1$, and $w^{t+1}$ represents the updated global model parameters.

FL aims to successfully train a global model using these repeated procedures without explicitly exchanging data. This method preserves privacy while fostering collaborative learning across devices and helps address the data silos issue.

2.2. Knowledge Distillation

A model compression technique called knowledge distillation [15] transfers knowledge from a complicated model to a reduced model in an effort to minimize the size and computing cost of the model. By “distilling” the output data from the complex model into the simplified model, this strategy is put into practice. The simplified model in knowledge distillation is known as the student model, whereas the complex model is generally referred to as the teacher model.

Here, a teacher model is represented by a complex model $T$, and a student model is represented by a simplified model $S$. Given an input sample $x$ and the “hard label” $y$, the soft probability distribution output using the teacher model is $T(x)$, and the soft probability distribution output using the student model is $S(x)$. Thus, the loss function for knowledge distillation can be expressed as follows:

$$L_{KD}=\lambda \cdot H(y,S(x))+(1-\lambda )\cdot KL(T(x),S(x)).$$

(3)

Here, $\lambda$ is a hyperparameter that balances two loss terms, $H$ denotes the cross-entropy loss function, and $KL$ denotes the Kullback–Leibler divergence. The first term $H(y,S(x))$ measures the cross-entropy loss between the “hard labels” and the predictions of the student model, while preserving the ability of the student model to fit the “hard labels”. The second measure $KL(T(x),S(x))$ is the similarity of outputs between the teacher model and the student model, specifically the Kullback–Leibler divergence between the “soft labels” of the teacher model and the predictions of the student model. By minimizing the knowledge distillation loss function, the student model can mimic the behavior of the teacher model as closely as possible without significant loss of accuracy, resulting in model compression and speedup.

2.3. Rényi Differential Privacy

The goal of the privacy protection mechanism known as DP [16] is to safeguard an individual’s privacy by adding noise to output results. The fundamental goal of DP is to secure individual privacy by guaranteeing that the probability distribution of the output results is almost the same regardless of who participates. Given a deep learning mechanism $M$ and a privacy budget $\epsilon$, for any two neighboring datasets $D$ and $D^{\prime }$, and for any output set $S$, if the following condition is satisfied:

$$\mathrm{Pr}[M(D)\in S]\le e^{\epsilon }\cdot \mathrm{Pr}\left[M\left(D^{\prime }\right)\in S\right]+\delta .$$

(4)

then the mechanism $M$ is said to satisfy $(\epsilon ,\delta )-DP$. This satisfies the properties of sequential composition [17] and parallel composition [18] of algorithms. Here, $\delta$ is the probability of failure, and $\epsilon$ is the privacy budget. In addition, $\delta$ is an important parameter in the definition of differential privacy, which allows us to introduce some flexibility in privacy assurance. Specifically, $\delta$ represents the probability that the differential privacy algorithm may “fail” or “violate” the privacy guarantee. When $\delta =0$, the algorithm has pure differential privacy, which means that the guarantee of differential privacy is absolutely established under any circumstances. When $\delta >0$, the algorithm has approximate differential privacy. At this point, the algorithm still satisfies the differential privacy in most cases. However, in rare cases, privacy leakage may occur, and the probability of such privacy leakage will not exceed $\delta$. Therefore, $\delta$ allows a small probability, and the privacy guarantee of the algorithm may not fully meet the expectations.

Rényi differential privacy [19] (RDP) is an extension of DP, and the Gaussian mechanism is a fundamental method for achieving RDP. It introduces Rényi sensitivity, which measures the divergence between two distributions using Rényi entropy. This approach is more flexible than the global sensitivity used in traditional DP, allowing for tailored levels of privacy protection according to specific application contexts and requirements. When applying composition mechanisms, there is no need for special higher-order composition theorems, as the concept of RDP can significantly optimize the computation of privacy overhead.

Definition 1

([19]). Given a training mechanism $M$, $M(D)(Y)$ denotes the probability density of $Y$ obtained by its training on the dataset $D$. If the mechanism $M$ satisfies $(\alpha ,\epsilon )-RDP$ under all $\alpha \in (1,\infty )$, then it also satisfies the following conditions on the neighboring datasets $D$ and $D^{\prime }$:

$${\displaystyle \frac{1}{\alpha -1}}\log E_{\theta ~M\left(D^{\prime }\right)}\left[{\left({\displaystyle \frac{M(D)(Y)}{M\left(D^{\prime }\right)(Y)}}\right)}^{\alpha }\right]\le \epsilon .$$

(5)

Definition 2

([19]). If the mechanism $M$ satisfies $(\alpha ,\epsilon )-RDP$, then for all $\delta \in (0,1)$, $M$ satisfies $(\epsilon +\log (1/\delta )/(\alpha -1),\delta )-DP$.

Definition 3

([19,20]). If there exists a deterministic function that $f:D\to R^d$ with sensitivity $\Delta f=1$, then under the Gaussian mechanism, $f^{\prime }(d)=f(d)+N\left(0,{\sigma }^2\right)$ satisfies $\left(\alpha ,\alpha /2{\sigma }^2\right)-RDP$.

Definition 4

([21]). Given mechanism $Subsampled$ for sampling without replacement from a dataset of size $m$ to obtain a subset of size $n$ (sampling rate $s=n/m$) and mechanism $M$, which is a randomized algorithm taking the subset as input, then, for all $\alpha \in [2,\infty )$, if $M$ satisfies $(\alpha ,\epsilon (\alpha ,\sigma ))-RDP$, the combination of both mechanisms satisfies $\left(\alpha ,{\epsilon }^{\prime }(\alpha ,\sigma )\right)-RDP$, where ${\epsilon }^{\prime }(\alpha ,\sigma )$ satisfies the following condition:

$$\begin{array}{c}{\epsilon }^{\prime }(\alpha ,\sigma )\le {\displaystyle \frac{1}{\alpha -1}}\log \left(1+s^2\left(\begin{array}{l}\alpha \hfill \\ 2\hfill \end{array}\right)\min \left\{4\left(e^{\epsilon (2)}-1\right),e^{\epsilon (2)}\min \left\{2,{\left(e^{\epsilon (\infty )}-1\right)}^2\right\}\right\}\right.\\ \left.+{\displaystyle \sum _{j=3}^{\alpha }{s}^j}\left(\begin{array}{l}\alpha \hfill \\ j\hfill \end{array}\right)e^{(j-1)\epsilon (j)}\min \left\{2,{\left(e^{\epsilon (\infty )}-1\right)}^j\right\}\right).\end{array}$$

(6)

Definition 5

([21]). Given two mechanisms $M_1$ and $M_2$, if $M_1$, which takes dataset $D$ as input, satisfies $\left(\alpha ,{\epsilon }_1\right)-RDP$, and $M_2$, which takes dataset $D$ and the output of $M_1$ as input, satisfies $\left(\alpha ,{\epsilon }_2\right)-RDP$, then the combination of the two mechanisms satisfies $\left(\alpha ,{\epsilon }_1+{\epsilon }_2\right)-RDP$.

3. FedKADP-Specific Implementation

3.1. FedKADP Framework

In order to address model heterogeneity, privacy, and security in federated learning with non-IID data, this research presents a framework. FedKADP, as illustrated in Figure 1, differs from traditional federated learning by using knowledge distillation, leveraging “soft labels” for local model training, integrating DP for enhanced user privacy, and employing a bidirectional feedback mechanism. This mechanism allows dynamic adjustments between knowledge distillation and DP.

3.2. Bidirectional Feedback Mechanism

Protecting the privacy of user data is critical in a FL environment. Concurrently, it is essential to ensure that the model’s ability to generalize, as acquired through knowledge distillation, is not compromised by the implementation of privacy protection measures. The management of knowledge distillation and DP independently could lead to an unnecessary trade-off between privacy protection and model performance. To address this issue, this paper proposes a two-way feedback mechanism that dynamically adjusts the parameters of knowledge distillation and DP to optimize the overall model performance and privacy protection.

The bidirectional feedback mechanism performs a comprehensive evaluation of the model after each round of model training based on the model’s metric in terms of gradient change, loss change, number of communication rounds, and accuracy and adaptively adjusts the application of knowledge distillation and differential privacy based on the metric. The mechanism consists of two main feedback paths:

(1)

From DP to knowledge distillation: the temperature parameter of knowledge distillation is dynamically adjusted based on the actual impact of noise added by DP on model performance;

(2)

From knowledge distillation to DP: the noise parameter of the DP based on the actual impact of knowledge distillation on model performance is adjusted to ensure that model learning is maximized without sacrificing too much privacy.

Specifically, the bidirectional feedback mechanism introduces an evaluation mechanism $Metric$ to assess the stability of the model based on the model’s gradient change, loss change, communication rounds, and validation accuracy, reflecting the impact of knowledge distillation and DP on model performance. The evaluation mechanism $Metric$ can be expressed as follows:

$$Metric=w_g\cdot M_g+w_l\cdot M_l+w_a\cdot M_a+w_t\cdot M_t.$$

(7)

Here, $w_g$, $w_l$, $w_a$, and $w_t$ are the weights of each metric. Adjusting the weights allows the influence of the metric to be flexibly adjusted. $M_g$ is used to assess the changes in gradients, which can be expressed as follows:

$$M_g=\max \left(100\times \left(1-\left|{\displaystyle \frac{g_t-g_{avg}}{g_{avg}}}\right|\right),0\right).$$

(8)

Here, $g_t$ is the gradient $L2$ norm of the current communication round, which is mathematically expressed as $g_t={\left[{\displaystyle {\sum }_{i=1}^n{\left({\theta }_t-{\theta }_{t-1}\right)}^2}\right]}^{\raisebox{1ex}{$1$}\!\left/ \!\raisebox{-1ex}{$2$}\right.}$; ${\theta }_t$ and ${\theta }_{t-1}$ are the current and previous global model parameters, respectively; and $g_{avg}$ is the average of the previous gradient $L2$ norms. $M_l$ is used to assess the changes in loss, which can be expressed as follows:

$$M_l=\left\{\begin{array}{ll}100\times \left(1-r_l\right),\hfill & r_l<0\hfill \\ 100\times \max \left(1-0.5\times r_l,0\right),\hfill & r_l\ge 0 \hfill \end{array}\right..$$

(9)

Here, $r_l$ is the rate of change in loss, which is mathematically expressed as $r_l=L_t-L_{t-1}/\left|L_{t-1}\right|$, and $L_t$ and $L_{t-1}$ are the loss values of the current and previous communication rounds, respectively. $M_a$ is used to assess changes in accuracy, which can be expressed as follows:

$$M_a=100\times {\displaystyle \frac{a_t-a_{t-1}}{a_{t-1}-a_{t-2}}}.$$

(10)

Here, $a_t$, $a_{t-1}$, and $a_{t-2}$ represent the accuracy of the current, previous, and first two communication rounds, respectively. $M_t$ is used to assess the time factor, which can be expressed as follows:

$$M_t={\displaystyle \frac{100}{1+e^{-(\frac{t}{T}-0.5)}}}.$$

(11)

Here, $t$ is the current communication round, and $T$ is the total number of communication rounds.

In summary, the implementation of the bidirectional feedback mechanism in FedKADP is illustrated in Figure 2. This figure depicts an effective bidirectional feedback mechanism that enables the knowledge distillation and DP processes to work in close collaboration, thereby optimizing the overall performance and privacy protection level of the federated learning model.

3.3. Knowledge Distillation Optimization

In traditional knowledge distillation, there is usually only one global teacher model. FedKADP proposes a hierarchical architecture of teacher models. Every local client receives the global model as its teacher model in each communication round. The client uses its local data as input to obtain “soft labels” trained by the local model. Then, the global model is again received as the student model, which participates in the training and updating of the local model. This approach reduces the loss of generalization performance brought on by non-IID data by giving each teacher model the responsibility of training a local client model. It also enables each teacher model to more precisely capture the data characteristics of its assigned client and efficiently direct the learning of the corresponding client model. The traditional federated learning method directly averages the parameters of the client model, which may lead to performance degradation of the global model because the model may be biased toward the specific data distribution of some clients. However, the student model in this method can use “soft labels” to imitate the output of the teacher model and better capture the characteristics of global data distribution, thus alleviating the model deviation caused by uneven data distribution. In addition, the student model can also use a “hard label”-driven approach to more directly fit the real data distribution of the client to ensure high accuracy in the target task. By optimizing the losses related to “hard labels”, the model can better adapt to the specific data characteristics of the client.

Specifically, in each local training iteration, the local data from the client are used as input, and the logits output of the local model (student model) is denoted as ${\widehat{y}}_{student}$. The logits output from the global model (teacher model) is denoted as ${\widehat{y}}_{teacher}$. Both the teacher and student models’ outputs are softened and transformed into probabilistic outputs that can be expressed as follows:

$$p_i=\mathrm{Softmax}({\displaystyle \frac{{\widehat{y}}_{teacher}}{\rho }});$$

(12)

$$q_i=\mathrm{Softmax}({\displaystyle \frac{{\widehat{y}}_{student}}{\rho }}).$$

(13)

Here, $\rho$ is the temperature parameter for softening, which controls the “smoothness” of the Softmax output. A higher temperature $\rho$ makes the probability distribution flatter, which helps the model to focus more on the classes that are difficult to classify during learning. This is especially important when dealing with non-IID data.

Subsequently, the loss function $Los{s}_{KD}$ for knowledge distillation is calculated, and the loss function is updated, which can be expressed as follows:

$$Loss=Los{s}_0+Los{s}_{KD}.$$

(14)

$Los{s}_0$ can be expressed as follows:

$$Los{s}_0=H({\widehat{y}}_{student},y)=-{\displaystyle \frac{1}{N}}{\displaystyle {\sum }_{i=1}^N{\displaystyle {\sum }_j{y}_{ij}}}\log (q_{ij}).$$

(15)

Here, $H(\cdot )$ is the cross-entropy loss, $N$ is the number of samples in the local dataset, $y_{ij}$ is the “hard labels” of the $j$-th class for the $i$-th sample, and $q_{ij}$ is the predicted probability of the student model for the $j$-th class of the $i$-th sample. The model’s fit can be determined by comparing the probability distribution of the student model’s predictions with the distribution of the “hard labels”. The student model can fit the true labels more closely by minimizing the cross-entropy loss, which enhances performance on the original task.

$Los{s}_{KD}$ can be expressed as follows:

$$Los{s}_{KD}=KL(p_i\parallel q_i)\times {\rho }^2={\displaystyle {\sum }_{i=1}^N{\displaystyle {\sum }_j{p}_{ij}}}\log ({\displaystyle \frac{p_{ij}}{q_{ij}}})\times {\rho }^2.$$

(16)

Here, $KL(\cdot )$ is the Kullback–Leibler (KL) divergence, and $p_{ij}$ is the predicted probability of the teacher model for the $j$-th class for the $i$-th sample. The similarity between the soft outputs of the teacher and student models is measured by calculating the difference between them. The rationale for multiplying the calculated KL divergence by the square of the temperature parameter is that when employing the temperature-adjusted Softmax output, the final loss must be scaled in accordance with the chosen temperature parameter to ensure that the magnitude of the loss remains constant. The student model can efficiently learn practical, generalized knowledge from the teacher model while retaining its uniqueness by minimizing the KL divergence.

In order to achieve the adaptive adjustment of temperature parameters and distillation weights based on a bidirectional feedback mechanism, the adjustment for the temperature parameter $\rho$ in FedKADP can be expressed as follows:

$${\rho }_t={\rho }_{\min }+\left({\rho }_{\max }-{\rho }_{\min }\right)\times {\displaystyle \frac{1}{1+e^{-k\left(Metric-Metri{c}_{KD}\right)}}}.$$

(17)

Here, ${\rho }_{\max }$ and ${\rho }_{\min }$ are the maximum and minimum values of the temperature parameter, respectively; $Metri{c}_{KD}$ is the threshold value; and $k$ is a parameter that adjusts the steepness of the curve. The larger the value of $k$ is, the more sensitive the function response. Such an adjustment mechanism allows for the simultaneous achievement of the following goals: increasing the temperature parameter smooths the probability distribution, helping the model learn more general information from the teacher’s model rather than specific noise; and adjusting the distillation weights enhances the influence of the teacher’s model output on the student model, improving the student model’s ability to better extract useful information despite noise and the performance reduction caused by DP.

The process of knowledge distillation optimization in FedKADP is illustrated in Figure 3. Overall, it allows the model to flexibly balance the fit to the “hard labels” and the absorption of knowledge from the teacher model during the learning process. This approach is particularly effective in improving the generalization capabilities of the model when faced with the challenges of non-IID data.

3.4. Adaptive Differential Privacy

In order to prevent user privacy leakage when uploading local model parameters, traditional methods introduce noise equally in each communication round. This could cause the accuracy of the model to deteriorate in later stages. Currently, there are relatively few adaptive noise introduction techniques suitable for FL, especially those aimed at protecting DP at the record level. The FedKADP method employs the Gaussian mechanism to add Gaussian noise to the local model during each iteration of local model training and adaptively adjusts the noise magnitude, while utilizing RDP theory to track privacy overhead and reasonably allocate the privacy budget. The objective is to enhance model accuracy while maintaining the same conditions of privacy budget.

In the context of DP gradient descent algorithms executed on the client side, a cropping technique is frequently employed with the objective of ensuring that $‖w‖\le C$. Here, $w$ represents the model parameters of the client, while $C$ denotes the cropping threshold at the $w$ boundary. Consequently, the client-side model gradient cropping can be expressed as follows:

$$\overline{g}\leftarrow \nabla g\cdot \min \left(1,{\displaystyle \frac{c}{\parallel \nabla g{\parallel }_2}}\right).$$

(18)

Assuming that the batch size $B$ is used to sample the samples for training in the local training, the sensitivity $C$ of the model gradient can be obtained based on the gradient trimming threshold. Therefore, the addition of Gaussian noise to the local model can be expressed as follows:

$$\tilde{g}\leftarrow {\displaystyle \frac{1}{\left|B\right|}}{\displaystyle {\sum }_{b\in \mathcal{B}}\overline{g}}+\mathcal{N}(0,{\sigma }^2{c}^{2}I).$$

(19)

In order to achieve the adaptive adjustment of noise parameters based on a bidirectional feedback mechanism, the adjustment for the noise parameter $\sigma$ in FedKADP can be expressed as follows:

$${\sigma }_t=\left\{\begin{array}{cc}\hfill p{\sigma }_{t-1},\hfill & \hfill Metric\ge Metri{c}_{DP};\hfill \\ \hfill {\sigma }_{t-1},\hfill & \hfill Metric<Metri{c}_{DP}. \hfill \end{array}\right.$$

(20)

Here, $p$ is the attenuation coefficient, and $Metri{c}_{DP}$ is the threshold value. Using such an adjustment mechanism that simultaneously optimizes the noise parameter based on model performance through knowledge distillation, the model can better adapt to data idiosyncrasies and variations, especially in non-IID data. This approach avoids the pitfalls of a one-size-fits-all fixed noise setting, which can often lead to inadequate model performance. With real-time tuning, the addition of noise can be optimized as needed to avoid overprotection at the expense of utility.

3.5. FedKADP-Algorithm Description

In FedKADP, clients conduct training on non-IID data, utilizing knowledge distillation techniques to optimize local model training. This method lessens the detrimental effects of model heterogeneity brought on by non-IID data, minimizes disparities between local models, accelerates global convergence, and enhances model validation accuracy. Additionally, DP is implemented using the Gaussian mechanism, where Gaussian noise is added to local models during each local iteration to create perturbations. Furthermore, a bidirectional feedback mechanism is introduced that allows dynamic adjustments between knowledge distillation and DP. The specific implementations are illustrated in Algorithms 1 and 2.

Algorithm 1. FedKADP Client Algorithm.

Input:$E$ is the number of local training iterations, $w$ represents the model parameters, $\eta$ is the learning rate, $B$ is the batch size, $\sigma$ is the Gaussian noise parameter, $\rho$ is the temperature parameter, and $c$ is the cropping threshold.

1:  $w_{k,t}^{student}\leftarrow w_{t-1}$$, w_{k,t}^{teacher}\leftarrow w_{t-1}$

2:  $\sigma \leftarrow {\sigma }_{t-1},\rho \leftarrow {\rho }_{t-1}$

3:  $\mathcal{B}\leftarrow$$(\mathrm{Split} D_k$ into batches $b$ of size $B$)

4: for $e=0,1,2,\cdots ,E-1$ do

5:        $\mathrm{sample} \mathrm{a} \mathrm{batch} b\in \mathcal{B}$ do

6:                ${\widehat{y}}_{student}\leftarrow f(w_{k,t,e}^{student},x)$$, q_i\leftarrow \mathrm{Softmax}({\displaystyle \frac{{\widehat{y}}_{student}}{\rho }})$

7:                ${\widehat{y}}_{teacher}\leftarrow f(w_{k,t,e}^{teacher},x)$$, p_i\leftarrow \mathrm{Softmax}({\displaystyle \frac{{\widehat{y}}_{teacher}}{\rho }})$

8:                $l\leftarrow Los{s}_0+Los{s}_{KD}$

9:                $\nabla g\leftarrow {\displaystyle \frac{\partial l}{\partial w_{k,t,e}^{student}}}$

10:                $\overline{g}\leftarrow \nabla g\cdot \min \left(1,{\displaystyle \frac{c}{\parallel \nabla g{\parallel }_2}}\right)$

11:                $\tilde{g}\leftarrow {\displaystyle \frac{1}{\left|B\right|}}{\displaystyle {\sum }_{b\in \mathcal{B}}\overline{g}}+\mathcal{N}(0,{\sigma }^2{c}^{2}I)$

12:                $w_{k,t,e+1}^{student}\leftarrow w_{k,t,e}^{student}-\eta \tilde{g}$

13: end for

14: return $w_{k,t}^{student}$

Algorithm 2. FedKADP Server Algorithm.

Input:$T$$\mathrm{is} \mathrm{the} \mathrm{number} \mathrm{of} \mathrm{communication} \mathrm{rounds}, {\epsilon }_T$ is the privacy budget, $K$ is the total number of clients, and $s$ is the client sampling rate.

1:  $w_t\leftarrow w_o$$, {\sigma }_t\leftarrow {\sigma }_0$

2: for $t=0,1,2,\cdots ,T-1$ do

3:          $N\leftarrow \max (K\cdot s,1)$

4:          $S_t\leftarrow$$N$ clients)

5:          for$\mathrm{each} \mathrm{client} k\in S_t$ in parallel do

6:                    Algorithm 1 7:          end for

8:          $w_t\leftarrow {\displaystyle \frac{1}{N}}{\displaystyle {\sum }_{i=1}^N{w}_{k,t}^{student}}$

9:        ${\epsilon }_t\leftarrow$(Privacy overhead calculating)

10:         if ${\epsilon }_t>{\epsilon }_T$

11:              break

12:         end if

13:         ${\sigma }_t,{\rho }_t\leftarrow$$(\mathrm{Adjustment} \mathrm{based} \mathrm{on} Metric$)

14: end for

15: return $w_T$

4. Theoretical Analysis

The following parameters are provided: the client sampling rate $s$, the number of communication rounds $T$, the Gaussian noise parameter ${\sigma }_t$ for each round, the number of clients $n$ participating in each round, the set of participating clients $S_t$ for each round, the number of local training iterations $E$, and the Gaussian mechanism $M^{\prime }(d)=M(d)+N\left(0,{\sigma }_t^2\right)$.

4.1. Privacy Overhead Calculation

In accordance with Definition 3, it can be concluded that when the Gaussian parameter is equal to ${\sigma }_t$, $M^{\prime }\left(d\right)$ satisfies $\left(\alpha ,\epsilon \left(\alpha ,{\sigma }_t\right)\right)-RDP$, where $\epsilon \left(\alpha ,{\sigma }_t\right)\le \alpha /2{\sigma }_t^2$.

In accordance with Definition 4, it can be concluded that after client $k$($k\in S_t$) completes a round of local training, it satisfies $\left(\alpha ,{{\epsilon }_k}^{\prime }\left(\alpha ,{\sigma }_t\right)\right)-RDP$. Here, $\exp (\infty )$ is unbounded in the Gaussian mechanism, where ${{\epsilon }_k}^{\prime }\left(\alpha ,{\sigma }_t\right)$ satisfies the following condition:

$${{\epsilon }_k}^{\prime }\left(\alpha ,{\sigma }_t\right)\le {\displaystyle \frac{1}{\alpha -1}}\log \left(1+s^2\left(\begin{array}{l}\alpha \hfill \\ 2\hfill \end{array}\right)\min \left\{4\left(e^{{\displaystyle \frac{1}{{\sigma }_t^2}}}-1\right),2e^{{\displaystyle \frac{1}{{\sigma }_t^2}}}\right\}+2{\displaystyle \sum _{j=3}^{\alpha }{s}^j}\left(\begin{array}{l}\alpha \hfill \\ j\hfill \end{array}\right)e^{{\displaystyle \frac{j(j-1)}{2{\sigma }_t^2}}}\right).$$

(21)

In accordance with Definition 5, it can be concluded that after client $k$($k\in S_t$) completes $E$ rounds of local training, it satisfies $\left(\alpha ,E{{\epsilon }_k}^{\prime }\left(\alpha ,{\sigma }_t\right)\right)-RDP$, where $E{{\epsilon }_k}^{\prime }\left(\alpha ,{\sigma }_t\right)$ satisfies the following condition:

$$E{{\epsilon }_k}^{\prime }\left(\alpha ,{\sigma }_t\right)\le {\displaystyle \frac{E}{\alpha -1}}\log \left(1+s^2\left(\begin{array}{l}\alpha \hfill \\ 2\hfill \end{array}\right)\min \left\{4\left(e^{{\displaystyle \frac{1}{{\sigma }_t^2}}}-1\right),2e^{{\displaystyle \frac{1}{{\sigma }_t^2}}}\right\}+2{\displaystyle \sum _{j=3}^{\alpha }{s}^j}\left(\begin{array}{l}\alpha \hfill \\ j\hfill \end{array}\right)e^{{\displaystyle \frac{j(j-1)}{2{\sigma }_t^2}}}\right).$$

(22)

In accordance with the theory presented by Noble [7], after the server aggregates the outputs from all clients in $S_t$, it satisfies $\left(\alpha ,{{\epsilon }_t}^{\prime }\left(\alpha ,{\sigma }_t\right)\right)-RDP$, where ${{\epsilon }_t}^{\prime }\left(\alpha ,{\sigma }_t\right)$ satisfies the following condition:

$${{\epsilon }_t}^{\prime }\left(\alpha ,{\sigma }_t\right)\le {\displaystyle \frac{E}{\alpha -1}}\log \left(1+s^2\left(\begin{array}{l}\alpha \hfill \\ 2\hfill \end{array}\right)\min \left\{4\left(e^{{\displaystyle \frac{1}{n{\sigma }_t^2}}}-1\right),2e^{{\displaystyle \frac{1}{n{\sigma }_t^2}}}\right\}+2{\displaystyle \sum _{j=3}^{\alpha }{s}^j}\left(\begin{array}{l}\alpha \hfill \\ j\hfill \end{array}\right)e^{{\displaystyle \frac{j(j-1)}{2n{\sigma }_t^2}}}\right).$$

(23)

4.2. Privacy Analysis

In accordance with the theoretical analysis in Section 4.1 and Definition 5, it can be concluded that FedKADP satisfies $\left(\alpha ,{\displaystyle {\sum }_{t=0}^{T-1}{{\epsilon }_t}^{\prime }\left(\alpha ,{\sigma }_t\right)}\right)-RDP$ after $T$ rounds of communication.

In accordance with Definition 2, it can be concluded that FedKADP satisfies $\left({\displaystyle {\sum }_{t=0}^{T-1}{{\epsilon }_t}^{\prime }\left(\alpha ,{\sigma }_t\right)}+\log \left(1/\delta \right)/(\alpha -1),\delta \right)-DP$.

In conclusion, as long as ${\sigma }_t$ is defined in accordance with the following condition:

$${\displaystyle \sum _{t=0}^{T-1}\frac{\mathrm{E}}{\alpha -1}\log \left(1+s^2\left(\begin{array}{l}\alpha \hfill \\ 2\hfill \end{array}\right)\min \left\{4\left(e^{\frac{1}{n{\sigma }_t^2}}-1\right),2e^{\frac{1}{n{\sigma }_t^2}}\right\}+2{\displaystyle \sum _{j=3}^{\alpha }{s}^j}\left(\begin{array}{l}\alpha \hfill \\ j\hfill \end{array}\right)e^{\frac{j(j-1)}{2n{\sigma }_t^2}}\right)+\frac{\log \left(\frac{1}{\delta }\right)}{\alpha -1}\le \epsilon }.$$

(24)

Then, FedKADP satisfies $\left(\epsilon ,\delta \right)-DP$, and the proof is completed.

5. Experiment

The effectiveness of FedKADP is validated in this section using experiments that assess the impact of parameter changes on the accuracy of FedKADP and compare FedKADP with existing algorithms. Additionally, the privacy overhead is tracked throughout the model training process based on RDP.

The experimental hardware configuration includes an AMD Ryzen 7 6800H CPU, GTX 3060 GPU, and 16 GB RAM, running on Windows 11. The deep learning models and DP noise addition are trained using the PyTorch 2.1.2 framework. The experimental datasets used are the widely utilized MNIST and CIFAR-10 datasets in FL algorithms. Due to limitations in the size of the datasets, the complete datasets are distributed across 100 clients. In a non-IID environment, the data are delivered to various clients after being sorted by labels. In particular, the images of the MNIST dataset are sorted according to their labels after dispersion and normalization. This process ensures that each client receives data that are predominantly label specific, rather than a uniform distribution of all labels. Similarly, the CIFAR-10 dataset is distributed to each client by means of label sorting after transformation and normalization. This approach simulates the real-world diversity and inhomogeneity of data sources from different clients, thus creating a non-IID data environment between clients with different label distributions. For the MNIST dataset, the experiments use a convolutional neural network that consists of two convolutional layers and two fully connected layers. Each convolutional layer is followed by a rectified linear unit activation function and a maximum pooling layer for feature extraction and spatial dimension reduction. The final fully connected layer outputs predictions for ten categories. For the CIFAR-10 dataset, we also conducted experiments using a convolutional neural network consisting of three convolutional layers and two fully connected layers. Each convolutional layer is followed by a rectified linear unit activation function and a maximum pooling layer. The last two fully connected layers map the extracted features to ten categories of predictions. The experiment sets the hyperparameters for the number of communication rounds $T$, the number of local training iterations $E$, the total number of clients $K$, the batch size $B$, the learning rate $\eta$, and the temperature parameters ${\rho }_{\min }$ and ${\rho }_{\max }$ to 100, 20, 100, 32, 0.005, 2, and 3, respectively.

5.1. Visual Analysis of the Bidirectional Feedback Mechanism

This section shows the change in the $Metric$ value throughout the global iteration process in a visual way, with an initial noise parameter $\sigma$ of 5 and a client sampling rate $s$ of 0.1. It analyzes how the bidirectional feedback mechanism dynamically adjusts the knowledge distillation and differential privacy parameters. This helps to more intuitively understand the role of the bidirectional feedback mechanism in optimizing model performance and privacy protection, as illustrated in Figure 4.

From the perspective of differential privacy, this bidirectional feedback mechanism dynamically adjusts the noise parameter of differential privacy according to the value of $Metric$ using Equation (20). When $Metric$ exceeds $Metri{c}_{DP}$, the differential privacy noise parameter ${\sigma }_t$ is proportionally attenuated, i.e., multiplied by the attenuation factor $p$. When $Metric$ is lower than $Metri{c}_{DP}$, the noise parameter remains unchanged. When $Metric$ is low, the performance of the model may be less than optimal, and keeping the noise parameter high at this time can enhance privacy protection and prevent the leakage of sensitive information. When $Metric$ reaches or exceeds $Metri{c}_{DP}$, it means that the model performance is better. At this time, the introduction of noise can be appropriately reduced to further improve the accuracy of the model and reduce the performance loss due to noise. As the training of the model is gradually optimized, the system will adaptively reduce the noise to avoid overprotection, and this mechanism helps to improve the practical application of the model at a later stage, while still meeting the requirements of differential privacy.

From the perspective of knowledge distillation, the bidirectional feedback mechanism dynamically adjusts the temperature parameters of knowledge distillation technology according to the value of $Metric$ using Equation (17). The knowledge distillation temperature parameter ${\rho }_t$ increases as $Metric$ increases. Specifically, when $Metric$ approaches or exceeds $Metri{c}_{KD}$, the temperature parameter tends toward the maximum value ${\rho }_{\max }$. The temperature parameter ${\rho }_t$ controls the student model’s focus on the “soft labels” provided by the teacher model. When the temperature is low, the student model pays more attention to the obvious classification. When the temperature is high, the student model will pay more attention to those samples that are difficult to classify, so as to learn the knowledge of the teacher model better for complex tasks. In the early or middle stage of training, when the $Metric$ value is low, the system will keep the temperature parameter low, which makes the model focus on the basic classification task. As training progresses and metrics improve, the mechanism gradually adjusts the temperature parameters. Thus, the model can learn complex knowledge better, which is helpful to improve the generalization ability of the model in the non-IID data.

5.2. Influence of Relevant Parameters on FedKADP

5.2.1. Client Sampling Rate ($s$)

This section examines the influence of varying client sampling rates $s$ on the validation accuracy of FedKADP, conducting experiments with the MNIST and CIFAR-10 datasets. During the experimental process, the initial noise parameter $\sigma$ is set at 5, and three distinct client sampling rates are employed: 0.01, 0.05, and 0.1. As illustrated in Figure 5. It is evident that the model validation accuracy increases with the number of participating clients. This increase is due to the involvement of more clients, which typically provides more data, diverse data types, additional model parameters, and greater computational resources, all of which collectively contribute to improved model validation accuracy. However, this improvement also raises the computational cost and the risk of privacy leakage. In practical applications, the client sampling rate can be flexibly adjusted based on available computational resources, privacy budgets, and training objectives to optimize training outcomes.

5.2.2. Initial Noise Parameter ($\sigma$)

This section uses tests on the MNIST and CIFAR-10 datasets to investigate how changing the initial noise parameter $\sigma$ affects FedKADP’s validation accuracy. During the experimental process, the client sampling rate $s$ is set at 0.1. Three distinct initial noise magnitudes are employed: 5, 10, and 15. RDP is used to track privacy overhead. As illustrated in Figure 6. The model’s accuracy declines as the noise magnitude increases. This decline is due to the increased addition of noise parameters, which enhance privacy protection and reduce the privacy budget, inevitably compromising some model performance. Specifically, in the MNIST dataset, as the noise level increased from 5 to 10 and 15, the model’s validation accuracy decreased from 86% to 81% and 77.8%, respectively. In the CIFAR-10 dataset, it was observed that as the noise level increased from 5 to 10 and 15, the model’s validation accuracy dropped from 29.5% to 24.9% and 20.1%. These changes directly reflect how the introduction of noise can affect the model’s predictive ability. They also indicate that the model’s performance is highly sensitive to adjustments in noise parameters and the privacy budget. This impact seems to be more pronounced in the CIFAR-10 dataset, which can be attributed to the differences in complexity and features between the two datasets. MNIST images are simple, composed of fixed digit shapes, allowing the model to easily extract key edges and contours, thus making it more tolerant to noise. In contrast, CIFAR-10 images are more complex, containing a variety of colors, textures, and shapes. Noise more easily interferes with the extraction of these complex features, particularly in deeper convolutional layers, leading to a significant reduction in classification performance. In practical applications, the initial noise parameter can be adjusted according to the size of the privacy budget. Every training round’s privacy cost is monitored during the procedure, and the privacy budget is allocated rationally, aiming to balance the effectiveness of training with privacy concerns.

5.3. Membership Inference Attack Experiments

This section perform experiments on membership inference attacks to validate the privacy-preserving ability of our proposed federated learning method based on knowledge distillation and DP. A popular privacy attack is the membership inference attack, which uses the model’s output to infer details about specific training data samples. In this experiment, this work simulated potentially malicious participants and attempted to infer the individual training data associated with them by analyzing the model output.

Table 1. Relevant parameters.

	Epochs	Batch Size	Learning Rate	Learning Rate Decay	Reg
Shadow model	50	128	0.001	0.96	1 × 10−4
Attack model	50	10	0.001	0.96	1 × 10−7

describes the parameters used for both the shadow and attack models in our experiments. The shadow model was trained for 50 epochs with a batch size of 128, a learning rate of 0.001, a learning rate decay of 0.96, and a regularization term of 1 × 10⁻⁴. The attack model was trained for 50 rounds, but with a smaller batch size of 10, a learning rate of 0.001, a learning rate decay of 0.96, and a regularization term of 1 × 10⁻⁷. These settings were chosen in order to adequately simulate a realistic training environment while ensuring that there is sufficient model complexity to perform membership inference attacks.

Table 2. Experimental results.

Method	Label	Precision	Recall	F1-Score	Support
FedKADP	Member	0.48	0.03	0.06	15,000
Fedavg	Member	0.50	0.18	0.26	15,000

shows the results of the membership inference attack experiments. For FedKADP, the precision, recall, and F1 score for recognizing membership samples are 0.48, 0.03, and 0.06, respectively, with 15,000 samples supported. In comparison, traditional FedAvg has a precision of 0.50, a recall of 0.18, and an F1 score of 0.26 with 15,000 samples supported. These results show that while the both methods have similar precision, FedKADP has a significantly lower recall and F1 score, indicating better privacy protection. The lower recall and F1 scores indicate that the attack model is less successful in correctly identifying membership samples, highlighting the effectiveness of FedKADP in reducing membership inference attacks and thus providing stronger privacy protection for the training data.

5.4. Comparison Experiments

The objective of this section is to verify FedKADP’s efficacy using non-IID data under the same privacy budget. To this end, FedKADP is compared with FedKD [22] (with DP), Adap DP-FL [23], and DP-FL [6].

As illustrated in Figure 7, the MNIST dataset is used with an initial noise parameter $\sigma$ of 5 and a client sampling rate $s$ of 0.1. Under the same privacy budget, FedKADP is capable of completing 83 model updates with an accuracy of 85.1%, FedKD (with DP) completes 100 updates with an accuracy of 81.9%, Adap DP-FL completes 80 updates with an accuracy of 81.2%, and DP-FL completes 100 updates with an accuracy of 76.5%. From the perspective of communication costs, FedKD (with DP) reaches an accuracy of 81.9% after 99 rounds of communication, whereas FedKADP achieves the same accuracy in just 42 rounds. Adap DP-FL reaches an accuracy of 81.2% after 75 rounds of communication, whereas FedKADP achieves the same accuracy in just 42 rounds. DP-FL reaches an accuracy of 76.5% after 89 rounds of communication, whereas FedKADP achieves the same accuracy in just 30 rounds.

The CIFAR-10 dataset is used with an initial noise parameter $\sigma$ of 5 and a client sampling rate $s$ of 0.1. Under the same privacy budget, FedKADP is capable of completing 77 model updates with an accuracy of 29.6%, FedKD (with DP) completes 100 updates with an accuracy of 28.5%, Adap DP-FL completes 73 updates with an accuracy of 25.4%, and DP-FL completes 100 updates with an accuracy of 25%. From the perspective of communication costs, FedKD (with DP) reaches an accuracy of 28.5% after 78 rounds of communication, whereas FedKADP achieves the same accuracy in just 64 rounds. Adap DP-FL reaches an accuracy of 25.4% after 73 rounds of communication, whereas FedKADP achieves the same accuracy in just 35 rounds. DP-FL reaches an accuracy of 25% after 68 rounds of communication, whereas FedKADP achieves the same accuracy in just 29 rounds.

6. Conclusions

In conclusion, this work presents FedKADP, a novel federated learning technique that combines differential privacy and knowledge distillation to improve privacy and performance, especially for non-IID data. Employing a bidirectional feedback mechanism, FedKADP dynamically optimizes privacy protection and model performance and reduces communication overhead, promoting faster global model convergence. Experimental results using the MNIST and CIFAR-10 datasets demonstrate that FedKADP outperforms traditional FL methods by effectively reducing model heterogeneity and accelerating convergence, while maintaining robust differential privacy standards. This advancement marks a significant step forward in making FL scalable, efficient, and privacy-preserving, with potential applications in various data-rich but privacy-sensitive domains.

Nevertheless, there are still some limitations to this study. The experiments using this method are performed using the standard datasets MNIST and CIFAR-10 to test the machine learning models. However, these datasets are relatively small and may not fully capture the complexity of non-IID data in real-world scenarios. Future work will try to further optimize this method, and experiments should be performed on larger and more complex datasets. The performance of this method should be compared using different non-IID data to comprehensively evaluate its adaptability and robustness.