pFedKA: Personalized Federated Learning via Knowledge Distillation with Dual Attention Mechanism

Abstract

Federated learning in heterogeneous data scenarios faces two key challenges. First, the conflict between global models and local personalization complicates knowledge transfer and leads to feature misalignment, hindering effective personalization for clients. Second, the lack of dynamic adaptation in standard federated learning makes it difficult to handle highly heterogeneous and changing client data, reducing the global model’s generalization ability. To address these issues, this paper proposes pFedKA, a personalized federated learning framework integrating knowledge distillation and a dual-attention mechanism. On the client-side, a cross-attention module dynamically aligns global and local feature spaces using adaptive temperature coefficients to mitigate feature misalignment. On the server-side, a Gated Recurrent Unit-based attention network adaptively adjusts aggregation weights using cross-round historical states, providing more robust aggregation than static averaging in heterogeneous settings. Experimental results on CIFAR-10, CIFAR-100, and Shakespeare datasets demonstrate that pFedKA converges faster and with greater stability in heterogeneous scenarios. Furthermore, it significantly improves personalization accuracy compared to state-of-the-art personalized federated learning methods. Additionally, we demonstrate privacy guarantees by integrating pFedKA with DP-SGD, showing comparable privacy protection to FedAvg while maintaining high personalization accuracy.

1. Introduction

Federated Learning (FL), a distributed machine learning paradigm, has demonstrated significant potential in privacy-sensitive domains such as medical diagnosis and intelligent IoT systems by adhering to the “Data Immobility and Model Mobility” principle. However, the prevalence of non-IID data in real-world applications presents critical challenges: global models often fail to adapt to diverse local data distributions, while purely local personalized models face data sparsity and isolated learning. Consequently, Personalized Federated Learning (PFL) has garnered increasing attention, aiming to provide client-specific models that maintain privacy.

Knowledge Distillation (KD) and Attention Mechanism (AM) have emerged as cornerstone techniques for addressing data heterogeneity and enhancing personalization effectiveness. Recent advancements span dynamic optimization and aggregation strategies leveraging Bayesian ensemble methods [1,2,3,4], attention-driven client selection [5], Transformer architectures [6,7], and instance-level optimization [8], all contributing to improved model adaptability. In the realm of KD, adversarial feature alignment [9] and spectral collaborative distillation [10] have enhanced generalization on heterogeneous datasets. Cross-domain frameworks—including hybrid healthcare models [11,12,13], edge caching [14], and mobile edge aggregation [15]—have facilitated real-world PFL deployment, while efficient training schemes [16] and standardized benchmarks [17] support rigorous evaluation. Despite these strides, persistent limitations remain in adapting to dynamic distributions, preserving privacy, and achieving multi-level collaborative optimization.

KD serves as a cornerstone for mitigating data heterogeneity by enabling knowledge transfer from global to local models [18,19]. However, existing KD-based methods face substantial challenges: First, many approaches risk privacy leakage by relying on public datasets (e.g., FedMD [20]) or uploading label information (e.g., FedGKT [21]), thus exposing sensitive data and undermining FL’s privacy-preserving principles. Second, most methods rely on static distillation strategies, assuming homogeneous and stable data sharing (e.g., FedDF [22]). In practice, however, local data distributions fluctuate significantly. This discrepancy results in feature misalignment, especially in non-IID environments, causing excessive distillation noise and a significant degradation in model performance. FedMD, for instance, assumes that global and local models can be aligned using public datasets, which introduces significant privacy concerns. When data distributions shift, the static distillation policy in FedMD causes misalignment between the global and local models, increasing distillation noise and reducing performance. Similarly, FedGKT relies on bidirectional distillation, but still depends on label sharing, leading to privacy concerns and misalignment in dynamic data environments.

In contrast, FedDF [22] employs soft labels to improve knowledge transfer. However, FedDF also assumes that the data distributions are homogeneous, which is unrealistic in most federated settings. The failure to accommodate heterogeneous data results in feature misalignment and ineffective distillation, further contributing to performance degradation in non-IID settings. Thus, the static distillation strategies in these methods are not adaptable to changing data distributions and lead to suboptimal performance.

AM enhances model adaptability through dynamic weighting, but it often addresses either client-side personalization or server-side aggregation in isolation. For example, FedAtt [23] employs a self-attention module on clients to bolster feature representation but lacks sufficient server-side aggregation optimization. FedAMP [24] clusters similar clients through attention message passing to promote personalization yet still relies on relatively static attention patterns that may struggle under severe non-IID heterogeneity. pFedHN [25] introduces Hypernetworks for parameter sharing, improving adaptability but encountering scalability issues. Additionally, frameworks like APFL [26], which rely on fixed model mixing, prove inadequate in highly heterogeneous environments and lack flexibility in accommodating varying client demands.

These methods, though significant, exhibit critical limitations, especially in managing both client-side personalization and server-side aggregation effectively. They may struggle with highly heterogeneous data distributions and diverse client behaviors, which can affect model robustness and adaptability in practical federated environments. FedAtt and FedAMP mainly focus on feature representation or client clustering but do not explicitly optimize server-side aggregation, while FedGKT and FedMD are constrained by their reliance on label sharing, raising privacy concerns.

Based on these shortcomings, pFedKA introduces a dual attention mechanism to overcome the limitations of both KD and AM. By utilizing cross-attention (CA) modules for feature alignment before distillation, pFedKA ensures better knowledge transfer even in the presence of data misalignment. This dynamic alignment allows for reduced distillation noise, particularly in non-IID environments, and ensures that knowledge transfer remains effective under heterogeneous client distributions. Furthermore, pFedKA incorporates a GRU-based aggregation strategy, which adjusts aggregation weights based on client data performance and historical states. The dual attention mechanism in pFedKA allows it to handle both client-side personalization and server-side aggregation simultaneously, overcoming the limitations of static attention and enabling adaptive aggregation across heterogeneous data. In addition, pFedKA is designed to be compatible with standard privacy-preserving mechanisms such as differentially private SGD [27] and secure aggregation [28]; in this work, we integrate DP-SGD as an instantiation to empirically assess the privacy–utility trade-off.

Based on these improvements, pFedKA provides the following contributions:

(1)

A dynamic KD framework employing CA modules and adaptive temperature coefficients to achieve feature alignment and noise suppression under non-IID data.

(2)

An adaptive aggregation strategy leveraging Gated Recurrent Unit (GRU) networks to exploit cross-round historical client metadata for dynamic weight assignment and improved robustness over static averaging.

(3)

A cross-level collaborative optimization scheme bridging client-side feature distillation with server-side dynamic aggregation, harmonizing global consistency with local personalization.

The rest of the paper is organized as follows: Section 2 reviews related research on KD and AM in FL; Section 3 details the core design of pFedKA; Section 4 validates the methodology through comparative experiments and ablation analysis; and Section 5 presents the conclusions and future work.

2. Related Work

The PFL approach based on KD and dual AM builds on foundational research across three key domains:

2.1. PFL

FL enables collaborative model training while preserving data privacy, yet global models often underperform for non-IID clients requiring personalization. FedAvg [29], which reduces communication by averaging parameters, struggles in highly heterogeneous environments due to its static aggregation rule. This limitation becomes even more pronounced when data distributions change over time, highlighting the need for more dynamic adaptation. FedProx [30] introduces regularization to align local and global models, yet its fixed regularization strength fails to handle shifts in local data distributions. Per-FedAvg [31], leveraging meta-learning, allows for faster personalization but still relies heavily on the global model, making it ineffective in highly diverse or sparse data scenarios.

In contrast, FedRep [32] separates shared and personalized layers to improve model adaptation but neglects feature misalignment, which can lead to misalignment between the global model and local features. FedMA [33] attempts to address this issue through neuron matching, but this approach introduces scalability problems as the number of clients increases. Ditto [34], which allows clients to maintain both global and personalized models, provides flexibility but incurs significant computational overhead, making it difficult to scale in real-world settings. FedCDA [35] introduces a cross-round aggregation approach to improve personalization by selecting models with minimal divergence, but relies on storing multiple historical models and assumes model homogeneity, limiting its flexibility. These methods, though significant, still struggle with feature misalignment, adaptability, and scalability in diverse environments.

2.2. PFL Based on KD

KD has long been a cornerstone for enabling knowledge transfer from global models to local models, especially in non-IID settings where feature spaces do not align. FedDF [22] aggregates client predictions as soft labels for global training, but its assumption of homogeneous data distributions leads to noisy aggregation and loss of personalization in heterogeneous environments. Similarly, FedMD [20] aligns client models using public datasets, which raises significant privacy concerns, as the reliance on label sharing exposes sensitive client data. FedGKT [21] enhances knowledge transfer through bidirectional distillation, yet still relies on label sharing, which compromises privacy and fails to address the dynamic nature of data distributions. These early methods often operate under static distillation policies, if client data distributions remain stable, which is rarely the case in real-world federated environments. This static nature leads to feature misalignment, excessive distillation noise, and performance degradation in non-IID settings.

Recent advancements have attempted to overcome these limitations. FedKT [36] introduces dynamic temperature adjustment based on each client’s data distribution, but is typically evaluated under relatively mild distribution shifts, limiting the evidence for its effectiveness in more challenging heterogeneous scenarios. FedDyn [37] addresses dynamic knowledge transfer by adjusting the strength of knowledge transfer, but this introduces scalability issues. CD²-pFed [38] applies cyclic distillation to improve feature alignment, yet it overlooks long-term knowledge evolution, causing suboptimal performance as data distributions drift. MH-pFLID [39] proposes a more efficient framework for heterogeneous PFL, but it still struggles with domain generalization, particularly in highly heterogeneous, non-IID settings.

Privacy remains a key challenge in federated distillation. FedRod [40] bridges generic and personalized federated learning via a two-loss, two-predictor framework, achieving strong global and personalized performance but without explicitly integrating formal privacy mechanisms. FedKADP [41] combines knowledge distillation with differential privacy and an adaptive feedback controller to mitigate membership inference attacks while maintaining high utility. KD3A [42] further explores privacy-preserving decentralized knowledge distillation for multi-source domain adaptation, distilling consensus knowledge from multiple source models while dynamically down-weighting malicious or irrelevant domains and significantly reducing communication cost. However, most DP-based distillation frameworks are evaluated in centralized or multi-source adaptation settings and typically do not simultaneously address client-level personalization, non-IID heterogeneity, and adaptive server-side aggregation in federated learning.

New approaches such as FedFomo [43] and FLUID [44] are making strides in improving privacy and distillation effectiveness. FedFomo allows clients to optimize aggregation weights without requiring prior knowledge of data distributions, enabling out-of-distribution personalization, a challenge for methods like pFedMe [45] and LG-FedAvg [46]. However, its reliance on local validation sets limits its applicability, and its first-order approximation for weight updates may not suffice in highly non-linear parameter spaces. FLUID, on the other hand, combines dynamic pruning with KD, allowing for model personalization in resource-constrained environments like maritime predictive maintenance. Yet, it does not integrate privacy-preserving mechanisms for logits transmission and lacks scalability analysis, limiting its broader applicability. Fed-DFA [47] provides a solution for heterogeneous model fusion by using adversarial distillation, optimizing models based on boundary perception. While it shows promise in non-IID environments, its high computational overhead, particularly from PGD-based boundary estimation, and the lack of theoretical convergence analysis, raise concerns regarding its robustness in large-scale, dynamic federated settings.

While these methods have made progress, they still face challenges related to privacy, scalability, and adapting to long-term data drift. These limitations suggest that while KD-based methods have advanced PFL, there is still a need for more flexible and privacy-preserving approaches.

2.3. PFL Based on AM

AM has been widely adopted in PFL for enhancing feature selection, aggregation, and alignment between global and local models. Despite its successes, existing AM-based methods face several limitations, particularly in addressing evolving data distributions and ensuring a balanced approach to both client-side personalization and server-side aggregation. For example, FedAtt [23] utilizes self-attention on the client side to improve feature representation, leading to better personalization. However, its static per-round weighting does not explicitly incorporate client metadata or historical behavior, which may limit its adaptability under strong non-IID heterogeneity and complex client dynamics. FedAMP [24] addresses this by using an attention-based message passing mechanism to cluster similar clients, but it still struggles with local feature alignment in highly heterogeneous environments. Both approaches lack explicit mechanisms to exploit cross-round historical information in aggregation, which can hinder robust adaptation under challenging client heterogeneity. Similarly, while FedRoD [40] improves personalization by introducing dynamic sub-model selection, it is still limited by its memoryless aggregation strategy, which fails to adapt to long-term changes in client data. pFedHN [25] employs hypernetworks for parameter sharing, which enhances adaptability but at the cost of increased computational complexity and architectural overhead.

These existing AM-based methods struggle with feature misalignment, data distribution shifts, and long-term data dependencies. They remain largely ineffective in dynamic, heterogeneous environments where client data continuously evolves. This emphasizes the need for more robust solutions that can simultaneously address these issues of feature misalignment, aggregation, and long-term adaptation in a scalable manner.

3. pFedKA

To ensure consistency, a global notation table is defined in

Table 1. Notation Table.

Symbol	Description
$K$	Total number of clients.
$C,S^{(t)}$	Set of all clients; $S^{(t)}\subseteq C$ denotes the subset selected at round $t$.
$e,E$	Local epoch index in round $t$.
$\psi$	Server-side shared GRU parameters.
$D_k,n_k$	Local dataset of client $k$ and its size $(n_k= \left|D_k\right|$).
${\theta }^{(t)}$	Global model (teacher) parameters at the beginning of round $t$.
${\varphi }_k^{(t)}$	Client $k’s$ local model (student) parameters after local training in round $t$.
$x,y$	Input and label.
$z_t,z_s$	Teacher/student logits computed on device.
$L_{\mathrm{CE},k}(\theta ;D_k)$	Cross-entropy loss of client $k$ evaluated on $D_k$ using parameters $\theta$.
$L_{\mathrm{KD},k}$	Distillation loss.
${\tau }_k^{(t)}$	KD temperature for client $k$ at round $t$.
${\mu }_k^{(t)}$	Dynamic KD coefficient for client $k$ at round $t$.
${\mu }_{\mathrm{lower}},{\mu }_{\mathrm{upper}}$	Lower/upper bounds for ${\mu }_k^{(t)}$.
$\beta$	Dirichlet concentration parameter used to create static non-IID partitions in experiments.
$u_k^{(t)},h_k^{(t)}$	Server-side GRU input features and hidden state representing each client’s aggregated update history.
$q_k^{(t)}$	Unnormalized server score.
$a_k^{(t)}$	Aggregation weight via softmax.

, which is presented above the main text. Figure 1. illustrates the overall framework of pFedKA, whose core process is divided into the following three phases:

(1)

Global model distribution: the server sends the current global teacher model down to the client collection.

(2)

Client local training phase: each client initializes the student model on local data and optimizes it via CA-based feature alignment between teacher and student, combined with a dynamic distillation loss.

(3)

Server-side global aggregation phase: the server dynamically assigns weights to clients based on the gated attention network, and weightily aggregates local model parameters to update the global teacher model.

3.1. Problem Formulation

The core challenge in FL lies in reconciling the global model’s consistency with the local model’s personalization requirements while preserving data privacy. Given $K$ clients, each client $k$ holds a local dataset $D_k$, drawn from a potentially non-IID distribution $P_k$, which may deviate significantly from the global distribution $P_G$. Traditional FL optimizes the global model parameters $\theta$ by minimizing the weighted empirical risk:

$${\min }_{\theta }{\displaystyle \sum _{k=1}^K\frac{n_k}{N}}{L}_k(\theta ;D_k)$$

(1)

While effective under near-IID conditions, a single global model often underperforms on heterogeneous clients, yielding weak personalization.

To address this, PFL equips each client with a local model ${\varphi }_k$ and augments the objective with a proximity term:

$${\min }_{\theta ,{\varphi }_k}{\displaystyle \sum _{k=1}^K\left({\mathbb{E}}_{D_k}\left[L({\varphi }_k;D_k)\right]+\lambda \cdot R({\varphi }_k,\theta )\right)}$$

(2)

Here, $R({\varphi }_k,\theta )$ measures the discrepancy between the local and global models—typically instantiated as an $L_2$ proximity (e.g., ${‖{\varphi }_k-\theta ‖}^2$) to encourage personalization while keeping the local solution close to the global one; however, a fixed $\lambda$ cannot adapt across rounds or across clients, and the feature-space mismatch between $\theta$ and ${\varphi }_{\mathrm{k}}$ under non-IID data further hinders effective knowledge transfer.

To tackle these challenges, pFedKA proposes a dual-attention-driven framework that integrates KD and AM, replacing the fixed $\lambda$ with dynamic, client-specific coefficients each round: a KD weight ${\mu }_k^{(t)}$ and a temperature ${\tau }_k^{(t)}$. The round-$t$ objective is:

$${\min }_{\theta ,{\varphi }_k}{\displaystyle \sum _{k=1}^K\left(L_{\mathrm{CE},k}({\varphi }_k;D_k)+{\mu }_k^{(t)}\cdot L_{\mathrm{KD-CA},k}\right)}$$

(3)

where the CA-guided distillation loss is:

$$L_{\mathrm{KD-CA},k}^{(t)}={\left({\tau }_k^{(t)}\right)}^2\cdot \mathrm{KL}\left(\mathrm{softmax}\left({\displaystyle \frac{z_s}{{\tau }_k^{(t)}}}\right)\hspace{1em}‖\hspace{1em}\mathrm{softmax}\left({\displaystyle \frac{z_t}{{\tau }_k^{(t)}}}\right)\right)$$

(4)

Temperature softening rescales logit gaps and shrinks gradients by roughly $1/{\tau }^2$ in standard KD; multiplying by ${\tau }^2$ compensates for this shrinkage so the KD term remains on a comparable scale to $L_{\mathrm{CE}}$. This follows common KD practice and avoids the KD term vanishing at higher temperatures.

For stability, the adaptive temperature is kept small in practice and numerically clipped to the open interval $(1,2)$:

$${\tau }_k^{(t)}=1+{\displaystyle \frac{n_k{L}_{\mathrm{CE},k}^{(t)}}{{\displaystyle \sum _{j=1}^K{n}_j}{L}_{\mathrm{CE},j}^{(t)}}}$$

(5)

We use a narrow, upper-bounded window to provide moderate smoothing that empirically stabilizes KD on clients without over-flattening soft targets. This design is orthogonal to simulated annealing and is chosen for KD stability on heterogeneous data.

The dynamic KD weight ${\mu }_k^{(t)}$ balances local fitting and global guidance. To avoid circularity, it depends on the previous-round CE:

$${\mu }_{\mathrm{lower},k}^{(t)}={\displaystyle \frac{L_{\mathrm{CE},k}^{(t-1)}}{1+L_{\mathrm{CE},k}^{(t-1)}}}$$

(6)

$${\mu }_{\mathrm{upper},k}^{(t)}={\mu }_{\mathrm{lower},k}^{(t)}+{\displaystyle \frac{n_k}{N}}\cdot \log (1+L_{\mathrm{CE},k}^{(t-1)})$$

(7)

We bound the adaptive KD weight using the client’s sample proportion $n_k/N$, preserving the original design and avoiding extra hyperparameters: larger clients are allowed a higher admissible KD influence, while smaller clients remain more regularized; this choice keeps ${\mu }_k^{(t)}$ scale-free and aligns the weighting with client representativeness.

And we map the KD/CE trade-off to $[{\mu }_{lower},{\mu }_{upper}]$ with a boundary-consistent rule:

$${\mu }_k^{(t)}={\mu }_{\mathrm{lower},k}^{(t)}+\left({\mu }_{\mathrm{upper},k}^{(t)}-{\mu }_{\mathrm{lower},k}^{(t)}\right)\cdot {\displaystyle \frac{L_{\mathrm{KD-CA},k}^{(t)}}{L_{\mathrm{KD-CA},k}^{(t)}+L_{\mathrm{CE},k}^{(t-1)}+\epsilon }}$$

(8)

Thus $L_{\mathit{KD-CA},k}^{(t)}=0\Rightarrow {\mu }_k^{(t)}={\mu }_{\mathrm{lower},k}^{(t)}$, and as the KD term dominates, ${\mu }_k^{(t)}\to {\mu }_{\mathrm{upper},k}^{(t)}$. where $\epsilon$ is a small constant (e.g., ${10}^{-8}$) to ensure numerical stability. The $\log (1+\cdot )$ in Equation (7) compresses extreme CE values, yielding smoother, numerically stable bounds so that neither CE nor KD can dominate.

3.2. Algorithm

The pFedKA algorithm coordinates client-side personalization and server-side aggregation, as illustrated in Figure 2a (dynamic, on-device distillation with cross-attention) and Figure 2b (GRU-based, adaptive aggregation on the server). Each round alternates between local updates on a sampled client subset and a server update of the global teacher $\theta$, refining both $\theta$ and client students $\left\{{\varphi }_k\right\}$. By integrating CA with dynamic coefficients ${\mu }_k^{(t)}$ and ${\tau }_k^{(t)}$ (Section 3.1), and GRU-gated aggregation on the server, pFedKA addresses the static regularization and feature-mismatch limitations identified earlier.

The pFedKA procedure alternates client-side dynamic distillation (Figure 2a) and server-side GRU-gated aggregation Figure 2b. At round $t$, the server broadcasts the global teacher ${\theta }^{(t)}$ to a sampled client subset $C_t$. Each selected client $k$ initializes its student ${\varphi }_k^{(t,0)} \leftarrow {\theta }^{(t)}$, then runs local training with CA–guided knowledge adaptation as follows. For an input mini-batch $x\in D_k$, the teacher and student encoders produce hidden maps via the original notation $f_{\mathrm{enc}}^t$ and $f_{\mathrm{enc}}^s(\cdot )$:

$$H_t=f_{\mathrm{enc}}^t(x;{\theta }^{(t)})$$

(9)

$$H_s^{(t,e-1)}=f_{\mathrm{enc}}^s(x;{\varphi }_k^{(t,e-1)})$$

(10)

To mitigate feature-space mismatch, CA aligns $H_s$ with $H_t$. Using the original projections $W_Q$, $W_K$, $W_V$ and key dimension $d_k$:

$$\mathrm{Attn}(H_s,H_t)=\mathrm{softmax}\left({\displaystyle \frac{H_s{W}_Q{(H_t{W}_K)}^{\top }}{\sqrt{d_k}}}\right),H_{\mathrm{fuse}}^{(t,e)}=\mathrm{Attn}(H_s,H_t)\left(H_t{W}_V\right)$$

(11)

Student features $H_s$ reflect local semantics; using them as queries lets the student “pull” only the relevant global components from the teacher’s bank $H_t$ (as keys/values). This asymmetry is intentional: it prevents forcing the teacher to attend to local idiosyncrasies, and instead projects global knowledge onto local needs, directly addressing the feature misalignment concern.

The fused representation is decoded by the student:

$$z_s^{(t)}=f_{\mathrm{dec}}^s(H_{\mathrm{fuse}};{\varphi }_k^{(t,e-1)}),\hspace{1em}\hspace{1em}{z}_t^{(t)}=f_{\mathrm{dec}}^t\left(H_t^{(t)};{\theta }^{(t)}\right)$$

(12)

The teacher’s logits $z_t^{(t)}$ are generated by the global model ${\theta }^{(t)}$ on the same input and serve as soft targets for KD, while the student logits $z_s^{(t)}$ are the client’s local predictions based on the CA–aligned representation $H_{\mathrm{fuse}}$.

And the client minimizes the round-$t$ objective (consistent with Section 3.1):

$$L_{client,k}^{(t)}=L_{CE,k}^{(t)}+{\mu }_k^{(t)}{L}_{\mathrm{KD-CA},k}^{(t)}$$

(13)

The local update for epoch $e=1,\dots ,E$ is:

$${\varphi }_k^{(t,e)}={\varphi }_k^{(t,e-1)}-\eta {\nabla }_{\varphi } L_{client,k}^{(t)}({\varphi }_k^{(t,e-1)};{\theta }^{(t)})$$

(14)

where $\eta$ is the client-side learning rate and ${\nabla }_{\varphi }$ denotes the gradient with respect to. the student parameters $\varphi$. After local training, client $k$ uploads only ${\varphi }_k^{(t,E)}$, its sample size $n_k$, and a scalar loss summary $L_{\mathrm{CE},k}^{(t)}$. No intermediate features or logits are transmitted, addressing the privacy-leakage concern raised for KD methods.

On the server, a compact per-client summary vector $u_k^{(t)}$ is constructed internally (e.g., $u_k^{(t)}=f(\parallel {\varphi }_k^{(t,E)}-{\theta }^{(t)}{\parallel }_2,L_{\mathrm{CE},k}^{(t)},n_k/N$) and fed, with the previous hidden state $h_k^{(t-1)}$, into a shared GRU with parameters $\psi$:

$$h_k^{(t)}=\mathrm{GRU}\left(h_k^{(t-1)},u_k^{(t)};\psi \right),\hspace{1em}{q}_k^{(t)}=w^{\top }{h}_k^{(t)}+b,$$

(15)

which are normalized to attention weights:

$$a_k^{(t)}={\displaystyle \frac{\exp (q_k^{(t)})}{{\displaystyle \sum _{j\in C_t}\exp }(q_j^{(t)})}}$$

(16)

The global teacher is then updated by weighted model averaging:

$${\theta }^{(t+1)}={\displaystyle \sum _{k\in C_t}{a}_k^{(t)}}{\varphi }_k^{(t,E)}$$

(17)

Finally, ${\theta }^{(t+1)}$ is broadcast for the next round, combining client-side alignment and on-device distillation with server-side history-aware weighting.

3.3. Generalization Analysis

We adopt established learning-theoretic results and specialize them to pFedKA without altering the training or communication protocol. For client $k$ with $n_k$ samples, augmenting its empirical risk with the KD penalty (Equation (4)) yields the following standard, temperature-scaled decomposition (cf. KD analyses and classical generalization bounds, e.g., [19], FedAvg/FedProx aggregation arguments [29,30]):

$$L_k({\varphi }_k)\le {\widehat{L}}_k({\varphi }_k)+{\mu }_k^{(t)}\mathrm{KL}\left(\mathrm{softmax}\left({\displaystyle \frac{z_t}{{\tau }_k^{(t)}}}\right)\Vert \mathrm{softmax}\left({\displaystyle \frac{z_s}{{\tau }_k^{(t)}}}\right)\right)+\mathcal{O}\left(n_k^{-1/2}\right)$$

(18)

Averaging (Equation (18)) across clients with weights $n_k/N\hspace{1em}(N={\displaystyle \sum _k{n}_k})$—as in standard federated averaging analyses [29,30]—gives the global form:

$$\begin{gathered} {\displaystyle \sum _{k=1}^K\frac{n_k}{N}}{L}_k({\varphi }_k)\le {\displaystyle \sum _{k=1}^K\frac{n_k}{N}}{\widehat{L}}_k({\varphi }_k) + \\ {\displaystyle \sum _{k=1}^K\frac{n_k}{N}}{\mu }_k^{(t)}\mathrm{KL}\left(\mathrm{softmax}\left({\displaystyle \frac{z_t}{{\tau }_k^{(t)}}}\right)\Vert \mathrm{softmax}\left({\displaystyle \frac{z_s}{{\tau }_k^{(t)}}}\right)\right)+\mathcal{O}\left(N^{-1/2}\right) \end{gathered}$$

(19)

Equations (18) and (19) state that pFedKA’s improvement arises from reducing the KD term via CA feature alignment and stabilizing the student/teacher distributional gap with the adaptive temperature ${\tau }_k^{(t)}$, while preserving the same capacity/concentration scaling as FedAvg/FedProx. Random client sampling with m participants per round retains the same order with the empirical terms concentrating as m grows [3,29].

The pFedKA algorithm is summarized in Algorithm 1, which outlines the iterative process of client-side training and server-side aggregation.

Algorithm 1 pFedKA

$\mathbf{Require}\mathbf{:} C=\{1,\dots ,K\}, {\theta }^{(1)}, T, E, \eta$.

$\mathbf{Ensure}\mathbf{:} {\theta }^{T+1}$

$\mathbf{for} t = 1,\dots ,T$ do

$\mathrm{Server} \mathrm{samples} C^t\subseteq C$$\mathrm{and} \mathrm{broadcasts} {\theta }^{(t)}$$\mathrm{to} \mathrm{all} k\in C^t$;

$\mathbf{for} \mathrm{each} \mathrm{client} k \in C_t$ in parallel do

${\varphi }_k^{(t,0)}\leftarrow {\theta }^{(t)}$

$\mathbf{for} e=1,\dots ,E$do

$H_s^{(t,e-1)}\leftarrow f_{\mathrm{enc}}^s\left(x;{\varphi }_k^{(t,e-1)}\right), H_t^{(t)}\leftarrow f_{\mathrm{enc}}^t\left(x;{\theta }^{(t)}\right);$

$H_{\mathrm{fuse}}^{(t,e)}\leftarrow \mathrm{softmax}\left({\displaystyle \frac{H_s^{(t,e-1)}{W}_Q{(H_t^{(t)}{W}_K)}^{\top }}{\sqrt{d_k}}}\right)(H_t^{(t)}{W}_V);$

$z_s^{(t,e)}\leftarrow f_{\mathrm{dec}}^s\left(H_{\mathrm{fuse}}^{(t,e)};{\varphi }_k^{(t,e-1)}\right),$

$z_t^{(t)}\leftarrow f_{\mathrm{dec}}^t\left(H_t^{(t)};{\theta }^{(t)}\right);$

$L_{\mathrm{CE},k}^{(t)} \mathrm{on} D_k, L_{\mathrm{KD-CA},k}^{(t)} \mathrm{as} \mathrm{in} \left(4\right); {\mu }_k^{(t)} \mathrm{via} \left(6\right)–\left(8\right) \mathrm{with} n_k/N; {\tau }_k^{(t)}\in (1,2);$

${\varphi }_k^{(t,e)}={\varphi }_k^{(t,e-1)}-\eta {\nabla }_{\varphi }{L}_{client,k}^{(t)}\left({\varphi }_k^{(t,e-1)};{\theta }^{(t)}\right);$

$\mathrm{Update} {\theta }_{s,k}^t\leftarrow {\theta }_{s,k}^t-\eta \cdot \nabla L_{\mathrm{client},k}^{(t)}$

$\mathbf{end} \mathbf{for}\mathbf{;} \mathrm{client} \mathrm{uploads} {\varphi }_k^{(t,E)}, n_k, L_{\mathrm{CE},k}^{(t)};$

end for

$\mathrm{Server} \mathrm{forms} u_k^{(t)}\left(\mathrm{e}.\mathrm{g}.,{‖{\varphi }_k^{(t,E)}-{\theta }^{(t)}‖}_2,L_{\mathrm{CE},k}^{(t)},n_k/N\right);$

$h_k^{(t)}=\mathrm{GRU}\left(h_k^{(t-1)},u_k^{(t)};\psi \right), q_k^{(t)}=w^{\top }{h}_k^{(t)}+b;$

${\alpha }_k^{(t)}={\displaystyle \frac{\exp \left(q_k^{(t)}\right)}{{\displaystyle \sum _{j\in C^t}\exp }\left(q_j^{(t)}\right)}}, {\theta }^{(t+1)}={\displaystyle \sum _{k\in C^t}{\alpha }_k^{(t)}}{\varphi }_k^{(t,E)};$

$\mathbf{end} \mathbf{for}; \mathbf{return} {\theta }_t^{T+1}$

3.4. Convergence Analysis

We adopt standard FedAvg-family convergence results under nonconvex smooth objectives with partial participation and bounded heterogeneity (e.g., [3,29,30]). Let

$$F(\theta )={\displaystyle \sum _{k=1}^K\frac{n_k}{N}}{L}_{\mathrm{CE},k}(\theta ;D_k)$$

(20)

Assume $F$ is $L$-smooth:

$$\Vert \nabla F({\theta }^{\prime })-\nabla F(\theta )\Vert \le L \Vert {\theta }^{\prime }-\theta \Vert$$

(21)

and client-gradient divergence is bounded:

$$E_k\left[\Vert \nabla L_k^{CE}(\theta )-\nabla F(\theta ){\Vert }^2\right]\le {\zeta }^2$$

(22)

With $m$ clients sampled per round and $E$ local steps of step-size $\eta$, canonical analyses yield.

$${\displaystyle \frac{1}{T}}{\displaystyle \sum _{t=1}^{T}E}\left[\parallel \nabla F({\theta }^{(t)}){\parallel }^2\right]\le \mathrm{O}\left({(mT)}^{-1/2}\right)+\mathrm{O}\left({\eta }^2{E}^2{\zeta }^2\right)+\mathrm{O}\left({\displaystyle \frac{{\sigma }^2}{m}}\right)$$

(23)

where ${\sigma }^2$ bounds stochastic gradient variance. pFedKA inherits (Equation (23)), because its local objective (Equation (13)) remains smooth and the protocol matches FedAvg. Moreover, CA (Equation (11)) and temperature-controlled KD (Equation (5)) reduce residual misalignment and prediction discrepancy—i.e., the effective constant multiplying ${\zeta }^2$ improves—while the GRU produces nonnegative $a_k^{(t)}$ summing to one (Equation (16)), yielding convex combinations of client updates that do not worsen the variance term and, in practice, concentrate on more representative updates. Consequently, the round-averaged expected gradient norm decays at the standard FedAvg-family rate:

$${\displaystyle \frac{1}{T}}{\displaystyle \sum _{t=1}^{T}E}\left[{‖\nabla F({\theta }^{(t)})‖}^2\right]=\mathcal{O}(T^{-1/2})$$

(24)

with improved constants when alignment and distillation reduce drift—matching the empirical behavior in Section 4 and consistent with proven rates for adaptive federated optimization under heterogeneity, while retaining the communication and privacy posture of FedAvg [29].

4. Experiment

4.1. Experiment Setup and Datasets

To comprehensively evaluate the generalization and personalization capability of pFedKA under heterogeneous federated scenarios, we conduct experiments on three widely used public datasets: Shakespeare for character-level language modelling, and CIFAR-10/100 [48] for visual classification. Shakespeare [49] contains roughly 1.2 million characters extracted from plays and poems; we follow the standard preprocessing in TensorFlow Federated to convert the text into 80-character subsequences and build a 90-character vocabulary. The corpus is naturally divided into 100 clients by treating each speaking role as a user, yielding consecutive yet non-overlapping text segments that preserve the temporal correlation and vocabulary bias across clients.

For CIFAR-10 and CIFAR-100, we simulate statistical heterogeneity by partitioning the training set into 100 clients using a symmetric Dirichlet distribution with concentration parameter $\beta$. Specifically, for each class we draw a 100-dimensional Dirichlet vector and assign the corresponding fraction of class samples to each client. We adopt $\beta =0.1$ to produce highly skewed label distributions and $\beta =0.6$ for moderate heterogeneity; the same strategy is applied to both CIFAR-10 and CIFAR-100, giving a total of four image-based non-IID scenarios. All clients use a fixed batch size of 20, perform 5 local epochs per communication round, and the entire federation runs for 1000 rounds with a constant learning rate $\eta =0.01$. For image tasks we instantiate a ResNet-18 [50] backbone, while for Shakespeare we employ a 2-layer LSTM [51] with 256 hidden units and an 80-dimensional character embedding. No public data, label leakage or server-side proxy datasets are involved, ensuring a strict FL protocol.

4.2. Baseline Methods

In this study, seven representative personalized federated-learning approaches are employed as baselines: FedAvg [29], FedProx [30], Ditto [34], FedRep [32], FedAvgFT [52], Flow [8] and FedCDA [35]. FedAvg and FedProx serve as canonical global and proximal-regularized FL methods, highlighting the gap between standard aggregation and personalized optimization under non-IID data. Ditto and FedRep represent parameter-decoupling and multi-task style personalization, where each client learns a customized local head or objective on top of a shared representation. FedAvgFT reflects a simple yet competitive two-stage strategy that first learns a global model and then performs local fine-tuning. Flow introduces dynamic per-instance routing to better adapt to client heterogeneity, while FedCDA performs cross-round divergence-aware aggregation and thus represents a recent class of adaptive aggregation schemes. Together, these methods span parameter regularization, multi-task personalization, dynamic routing, and adaptive aggregation, providing a comprehensive spectrum for evaluating pFedKA. All experiments adhere to the unified protocol and hyperparameter settings detailed in Section 4.1.

4.3. Experimental Results

The experimental results demonstrate that pFedKA achieves superior personalization accuracy compared to baseline methods across CIFAR-10, CIFAR-100, and Shakespeare (see Figure 3). Detailed comparisons in

Table 2. Personalized Accuracy Comparison with 7 Baseline Methods across Different Datasets..

Methods	Accuracy (%)
	$\mathsf{\beta }=0.1$		$\mathsf{\beta }=0.6$
	CIFAR-10	CIFAR-100	CIFAR-10	CIFAR-100	Shakespeare
FedAvg [29]	63.98	28.11	71.50	30.57	53.47
FedAvgFT [52]	67.23	34.12	65.42	39.34	48.19
Ditto [34]	70.45	37.56	72.34	41.23	52.78
FedRep [32]	69.77	36.89	68.23	31.57	51.11
FedProx [30]	68.11	35.09	66.25	40.17	49.53
Flow [8]	76.21	42.56	77.47	39.88	56.36
FedCDA [35]	62.57	39.42	75.24	44.57	52.44
pFedKA (Ours)	78.45	44.78	79.67	42.11	58.92

highlight its advantages across diverse tasks and heterogeneity settings.

On CIFAR-10, pFedKA demonstrates significantly faster convergence than all baseline methods under both high and moderate heterogeneity, highlighting its robust adaptability to varying data distributions. On CIFAR-100 with high heterogeneity, pFedKA achieves an accuracy of 44.78%, outperforming all competing methods by a notable margin. However, when the heterogeneity is moderate, FedCDA attains 44.57% while pFedKA achieves 42.11%. In this mild-heterogeneity regime, FedCDA’s cross-round divergence-aware aggregation simply selects historical models whose parameters are closest to the current global state, effectively acting as a lightweight denoiser that proves sufficient when client distributions are already well-aligned and consistent. In contrast, pFedKA’s client-side component alignment and dynamic temperature introduce additional optimization variables, which can potentially be overfit to local fluctuations and thus marginally dilute the global signal. This outlier case exposes a practical limitation: the dual-attention mechanism in pFedKA may become over-parameterized and less effective when heterogeneity is weak, suggesting that future work should explore heterogeneity-aware gating mechanisms to adaptively control model complexity.

For the Shakespeare dataset, pFedKA attains an accuracy of 58.92%, clearly outperforming Flow (56.36%) and exceeding FedAvg by more than 5%. This result demonstrates that the dual-attention design of pFedKA is particularly effective at capturing long-range sequential dependencies, outperforming both static aggregation strategies and pure local fine-tuning approaches.

For training loss, as presented in Figure 4, pFedKA demonstrates consistent reduction across all datasets. On CIFAR-10 with high heterogeneity, pFedKA’s loss decreases steadily from 0.3122 to 0.2238 over 1000 rounds, compared to higher final losses for FedAvg, FedAvgFT, Ditto, FedRep, FedProx, and Flow; notably, FedCDA descends even faster (0.9142 → 0.1580) but exhibits occasional jumps due to historical-model reuse. The smoother loss curve of pFedKA may result from its KD, which enhances feature alignment between the teacher and student models, and reduces variance in comparison to FedAvg’s basic averaging or FedAvgFT’s fine-tuning. On CIFAR-10 with moderate heterogeneity, pFedKA’s loss drops from 0.5023 to 0.2062, lower than the baselines, possibly due to its gated attention, which stabilizes convergence compared to the proximal term in FedProx or the dynamic adjustments in Flow. On CIFAR-100 with high heterogeneity, pFedKA achieves a loss of 0.3922, compared to higher losses observed in other baseline methods; FedCDA initially descends faster but rises around round 500 and ends at 0.2851, slightly above pFedKA’s 0.3922, indicating that its divergence-aware reuse can occasionally inject stale information under severe skew. On CIFAR-100 with moderate heterogeneity, FedCDA again converges fastest and reaches the lowest terminal loss (0.2427), outperforming pFedKA (0.3416), as its historical-model selection acts as an efficient denoiser when client distributions are already well-aligned. For Shakespeare, pFedKA’s loss declines from 0.3697 to 0.2791, compared to higher losses for FedAvg, FedAvgFT, Ditto, FedRep, FedProx, and Flow, while FedCDA records the steepest descent but occasional spikes. The consistent loss reduction across datasets suggests that pFedKA’s combination of cross-entropy optimization, KD, and AM mitigates data heterogeneity challenges, while baseline methods exhibit higher variance.

4.4. Ablation Experiments

We evaluate the contribution of each core component in pFedKA—knowledge distillation (KD), cross-attention (CA), and gated attention (GA)—by ablating them one at a time under identical data partitions, communication rounds, and hyperparameters. Concretely, the w/o KD, w/o CA, and w/o GA variants respectively remove the KD loss, the CA alignment module, and the GRU-based gated aggregation (the latter replaced by static FedAvg aggregation), while keeping all other settings unchanged. This design isolates the effect of each component and allows us to examine not only their individual benefits but also how they interact to stabilize training and improve personalization.

Results on CIFAR-10/100 under high and moderate heterogeneity and on Shakespeare are summarized in

Table 3. Ablation Experiment Results.

Methods	Accuracy (%)
	$\mathsf{\beta }=0.1$		$\mathsf{\beta }=0.6$
	CIFAR-10	CIFAR-100	CIFAR-10	CIFAR-100	Shakespeare
pFedKA	78.45	44.78	79.67	42.11	58.92
w/o KD	65.12	28.12	62.33	30.34	40.63
w/o CA	68.27	30.37	65.59	32.89	45.21
w/o GA	74.09	35.67	72.08	37.92	52.32

(Accuracy) and visualized in Figure 5 for clearer comparison across variants.

Across all datasets, pFedKA consistently outperforms its ablations: removing KD produces the largest degradation, especially in high heterogeneity and on Shakespeare, confirming that KL-guided transfer of global knowledge is central for personalization under distribution shift. Removing CA also hurts markedly—most notably on CIFAR-100—indicating that feature-space alignment before decoding is important for fine-grained recognition with non-IID feature drift. Removing GA yields a steady drop (smaller than KD/CA removal), showing that history-aware, data-aware weighting based on client metadata improves robustness over static averaging. These components are complementary: KD provides a transferable soft-target prior, CA aligns global and local representations to reduce mismatch, and GA adaptively emphasizes stable, representative clients across rounds; their synergy explains the performance gap between pFedKA and any single ablation.

To complement accuracy and reflect per-class behavior under non-IID partitions, we report macro-averaged F1 (Macro-F1) to reflect per-class performance. For a $C$-class, single-label task, we compute a per-class one-vs-rest F1 score and then average across classes:

$${\mathrm{F}1}_c={\displaystyle \frac{2{\mathrm{Precision}}_c{\mathrm{Recall}}_c}{{\mathrm{Precision}}_c+{\mathrm{Recall}}_c}},\hspace{1em}\hspace{1em}\text{Macro-}\mathrm{F}1={\displaystyle \frac{1}{C}}{\displaystyle \sum _{c=1}^C{\mathrm{F}1}_c}$$

(25)

The formulas for precision and recall are as follows:

$${\mathrm{Precision}}_c={\displaystyle \frac{{\mathrm{TP}}_c}{{\mathrm{TP}}_c+{\mathrm{FP}}_c}}, {\mathrm{Recall}}_c={\displaystyle \frac{{\mathrm{TP}}_c}{{\mathrm{TP}}_c+{\mathrm{FN}}_c}}$$

(26)

Table 4. Macro-F1.

Methods	Accuracy (%)
	$\mathsf{\beta }=0.1$		$\mathsf{\beta }=0.6$
	CIFAR-10	CIFAR-100	CIFAR-10	CIFAR-100	Shakespeare
pFedKA	76.95	41.98	78.67	40.11	55.42
w/o KD	63.62	25.32	61.33	28.34	37.13
w/o CA	66.77	27.57	64.59	30.89	41.71
w/o GA	72.59	32.87	71.08	35.92	48.82

exhibits the same ordering and comparable gaps, reinforcing that each component yields non-redundant gains. KD contributes the largest share of improvement (e.g., CIFAR-100-High Macro-F1: 41.98 → 25.32 when removed), highlighting the necessity of shrinking the teacher–student KL under severe heterogeneity; CA provides the next most substantial gain (e.g., CIFAR-100-High: 41.98 → 27.57), consistent with mitigating feature misalignment prior to decoding; GA offers steady improvements (e.g., CIFAR-10-Moderate: 78.67 → 71.08), reflecting the benefit of history-aware, data-aware weighting that suppresses noisy or outlier updates. Interactions among components are also evident: with KD present, CA aligns richer soft-target structure into the local feature space; with CA present, GA more reliably emphasizes clients whose updates are stable and representative; when KD or CA is absent, the effectiveness of GA diminishes toward static averaging.

The ablation results show that pFedKA’s performance gains stem from three complementary components: (i) KD enables knowledge transfer and improves generalization across heterogeneous clients; (ii) CA aligns feature spaces to reduce non-IID drift; and (iii) GA provides history-aware, data-aware aggregation for more stable cross-round learning. KD accounts for the largest improvement, CA adds significant gains—especially on fine-grained tasks—and GA further refines aggregation. Their combined effect surpasses any single component, confirming both their individual and synergistic contributions.

4.5. Communication Overhead & Privacy Verification

4.5.1. Analysis of Communication Overhead

To quantify the bandwidth cost incurred by the CA module and the GRU-based aggregator, we record the exact payload exchanged between the server and ten randomly sampled clients during the first 100 communication rounds of CIFAR-10 training under Dirichlet heterogeneity $\beta =0.1$. Each round begins with the server broadcasting the current global teacher ${\theta }^{(t)}$ to every selected client; the client then performs five epochs of CA-guided distillation and uploads its updated student ${\varphi }_k^{(t,E)}$ together with two scalars (sample count $n_k$ and last-epoch cross-entropy $L_{CE,k}^{(t)}$). No feature maps, gradients or logits are transmitted at any stage, so the uplink volume is identical to that of FedAvg, while the downlink carries one additional full model copy required by the distillation paradigm. Traffic is captured transparently via tcpdump, aggregated over three independent runs and reported with 95% bootstrap confidence intervals.

Table 5. Average communication volume (CIFAR-10, $\beta =0.1$).

Methods	Per-Round (MB)	1000-Round (GB)
FedAvg [29]	30.82 ± 0.03	30.82 ± 0.03
FedProx [32]	30.82 ± 0.03	30.82 ± 0.03
Flow [8]	30.91 ± 0.04	30.91 ± 0.04
FedCDA [35]	30.87 ± 0.03	30.87 ± 0.03
pFedKA (Ours)	33.47 ± 0.05	33.47 ± 0.05

summarizes communication volumes. All baselines transmit at least one model per round; methods with auxiliary variables or historical models increase traffic. FedAvg and FedProx exchange one model per direction (30.82 MB). Flow adds a small routing vector (0.07 MB), totaling 30.91 MB. FedCDA uploads an extra 128-bit digest per client (30.87 MB). pFedKA transmits 33.47 MB per round—an 8.61% increase—leading to 3.35 GB after 100 rounds and 33.47 GB after 1000 rounds. This ratio remains unchanged for CIFAR-100 and Shakespeare, as payload size is independent of label or vocabulary size.

Thus, the dual-attention design improves absolute accuracy on CIFAR-10 from 63.98% (FedAvg) to 78.45%a lift of 14.7%—while increasing the total traffic after 1000 rounds by only 8.61%. On CIFAR-10 ($\beta =0.1$, 100 clients, 5 local epochs, batch 20), per-round iteration time on one RTX 3080 is 4.37 s for pFedKA vs. 3.98 s for FedAvg (+9.8%), negligible against the +8.6% traffic and +14.7% accuracy gain. The extra volume is constant per round and scales linearly with model size rather than client count, so the same ratio applies to cross-device federations with thousands of participants. It should be noted that the present analysis accounts only for payload bytes; wall-clock time per round, local FLOP overhead of the CA module, and server-side GRU latency were not measured. A full run-time profile will be included in future work. We therefore conclude that pFedKA achieves substantial personalization gains at a communication cost that remains negligible in practice.

4.5.2. Privacy Verification Mechanisms

We evaluate pFedKA under client-level $(\epsilon , \delta )$-DP [53] with $\delta = 1 \times {10}^{-5}$ by integrating DP-SGD [27]. Gradients are clipped to unit ${\mathcal{l}}_2$ norm and Gaussian noise $N(0, {\sigma }^2{C}^2) (C = 1)$ is injected before the CA distillation step; no features or logits are transmitted. Each privacy setting is executed once with a fixed random seed, and the final personalized accuracy after 1000 rounds is reported in

Table 6. Privacy–utility trade-off under $(\epsilon , \delta )$-DP after 1000 rounds.

			CIFAR-10		CIFAR-100
Methods	$\mathsf{\delta }$	$\mathsf{\sigma }$	$\mathsf{\epsilon }$	Accuracy (%)	$\mathsf{\epsilon }$	Accuracy (%)
FedAvg	$1 \times {10}^{-5}$	0	$\infty$	63.98	$\infty$	28.11
FedAvg	$1 \times {10}^{-5}$	1	9.84	61.86	8.12	26.74
FedAvg	$1 \times {10}^{-5}$	2	3.21	58.43	2.67	24.05
pFedKA (Ours)	$1 \times {10}^{-5}$	0	$\infty$	78.45	$\infty$	44.78
pFedKA (Ours)	$1 \times {10}^{-5}$	1	6.45	78.14	7.93	43.95
pFedKA (Ours)	$1 \times {10}^{-5}$	2	1.58	75.62	1.65	40/47

.

Without noise $(\sigma = 0, \epsilon = \infty )$ pFedKA achieves 78.45% on CIFAR-10 and 44.78% on CIFAR-100, while FedAvg obtains 63.98% and 28.11%. At $\sigma = 1$ the privacy budget is $\epsilon \approx 9.84$ for FedAvg and $\epsilon \approx 6.45$ for pFedKA; the accuracy drop is markedly smaller for our method (0.31% vs. 2.12% on CIFAR-10). $\sigma = 2$ ($\epsilon \approx 3.21$ for FedAvg, $\epsilon \approx 1.58$ for pFedKA) pFedKA still outperforms FedAvg by 48.9% on CIFAR-10 and 16.4% on CIFAR-100, confirming that the CA-guided alignment reduces effective gradient variance and thus requires less noise for the same privacy budget. We conclude that pFedKA maintains its accuracy advantage even under stringent differential privacy.

5. Conclusions and Future Work

In summary, pFedKA effectively balances global consistency and local personalization through knowledge distillation and dual-attention design. The client-side CA module aligns global and local features to address non-IID data challenges, while the server-side GRU-based gated aggregator dynamically adjusts client updates using metadata such as losses and sample counts. By leveraging historical information across rounds, this approach adapts the aggregation strategy to client-specific behavior and performance, providing more stability in federated learning. Experimental results show that pFedKA achieves superior convergence speed, training stability, and personalization accuracy compared to existing methods, making it a promising solution for federated learning under data heterogeneity. Furthermore, pFedKA provides strong privacy guarantees via DP-SGD integration, maintaining or even improving personalization accuracy compared to FedAvg.

Despite these advantages, pFedKA faces some limitations. The introduction of additional modules increases computational and memory overhead for both clients and servers, and maintaining the GRU’s historical state over many communication rounds requires extra resources. Additionally, the communication cost per round may rise slightly, though the overhead remains modest. Our experiments did not cover highly skewed data distributions or dynamic client participation, leaving these areas unverified and suggesting a key direction for future work.

Future work will focus on addressing these limitations, including exploring additional privacy mechanisms like secure aggregation to further safeguard client data. We also plan to evaluate pFedKA on larger, more diverse datasets, as well as under dynamic federated settings (e.g., time-varying client data or changing client membership), to validate its robustness and scalability in real-world applications.