Towards Collaborative Fairness in Federated Learning Under Imbalanced Covariate Shift

Abstract

Collaborative fairness is a crucial challenge in federated learning. However, existing approaches often overlook a practical yet complex form of heterogeneity: imbalanced covariate shift. We provide a theoretical analysis of this setting, which motivates the design of FedAKD (Federated Asynchronous Knowledge Distillation) — a simple yet effective approach that balances accurate prediction with collaborative fairness. FedAKD consists of client and server updates. In the client update, we introduce a novel asynchronous knowledge distillation strategy based on our preliminary analysis, which reveals that while correctly predicted samples exhibit similar feature distributions across clients, incorrectly predicted samples show significant variability. This suggests that imbalanced covariate shift primarily arises from misclassified samples. Leveraging this insight, our approach first applies traditional knowledge distillation to update client models while keeping the global model fixed. Next, we select the correctly predicted high-confidence samples and update the global model using these samples, while keeping the client models fixed. The server update simply aggregates all client models. We further provide a theoretical proof of FedAKD ’s convergence. Experimental results on both public datasets (Fashion-MNIST and CIFAR10) and a real-world Electronic Health Records (EHR) dataset demonstrate that FedAKD significantly improves collaborative fairness, enhances predictive accuracy, and fosters client participation, even under highly heterogeneous data distributions.¹

1. Introduction

Federated Learning (FL) has emerged as a promising distributed paradigm that enables multiple participants (or clients) to collaboratively train a global model without sharing their raw local data [6, 9, 15, 24, 31]. However, disparities in data quality, quantity, and distribution among clients make uniform treatment in global model aggregation unfair, particularly for those with higher-quality or larger datasets. To address this, collaborative fairness (CF) has been introduced to ensure that each client’s final reward or benefit is commensurate with its contribution to the global model [13]. In other words, clients with a greater impact on model performance should receive proportionally higher gains.

Several fairness-aware FL approaches, such as CGSV [23], CFFL [13], FedAve [19], and FedSAC [20], have been developed to assign different contribution-based rewards (i.e., weights) during model aggregation. While these methods help mitigate fairness disparities, they still face the following challenges:

• Unrealistic assumptions about data heterogeneity.

Existing CF approaches primarily assume that data heterogeneity arises from imbalanced data sizes [13, 23], imbalanced class distributions [13, 23], or both [1, 28]. However, real-world datasets exhibit greater complexity. Beyond these imbalances, client data often differ significantly in feature distributions, leading to the covariate shift problem [1, 2, 5, 9]. Figure 1(a) presents a preliminary analysis on a real-world electronic health record (EHR) dataset for pancreatic cancer prediction². The x-axis represents client dataset sizes, while the y-axis shows the Kullback–Leibler (KL) divergence between each client’s fitted latent feature distribution and that of the entire dataset. Each circle represents a client, i.e., a U.S. state in the EHR dataset. The preliminary data analysis reveals that not only do clients have varying dataset sizes, but their latent feature distributions also significantly diverge from the global reference distribution. Thus, a more realistic FL heterogeneity setting should account for imbalanced covariate shift rather than just data quantity or class imbalance.

Figure 1:

(a) KL divergence vs. sample size for each client’s local data, revealing both data imbalance and feature covariate shifts. (b)-(e) compare different fairness methods’ contribution definitions to each client’s standalone training accuracy. Their poor correlation highlights the limitations of explicit contribution metrics under extra covariate shifts.

• Weak correlation between client accuracy and assigned contributions.

Beyond unrealistic data distribution assumptions, existing approaches to collaborative fairness typically follow a two-step pipeline: (1) explicitly defining a contribution metric and (2) allocating rewards based on this metric. For example, CGSV [23] estimates client contributions via gradient similarity, CFFL [13] and FedAVE [19] rely on performance improvements measured by a global validation set, and FedSAC [20] bases contributions on standalone local training results. However, in real-world datasets characterized by imbalanced covariate shift, these approaches fail to establish a strong correlation between client accuracy and assigned contributions, contradicting their underlying design assumptions.

Figures 1(b)–(e) illustrate the relationship between client standalone accuracy (x-axis) on the testing set and the learned contribution value (y-axis) either on the training set or validation set under different methods – CFFL, CGSV, FedAVE, and the baseline FedAvg [15] – using the real-world EHR dataset, which is the same as we analyzed in Figure 1(a). Each dot represents an individual client, and in the case of FedAvg, the contribution score is simply the proportion of data owned by the client. The results show that existing CF-based methods exhibit significantly lower Pearson correlation scores compared to the simple FedAvg baseline. These findings highlight the limitations of existing approaches in handling realistic imbalanced covariate shift settings.

Theoretical analysis on imbalanced covariate shift.

This paper aims to develop a novel model to ensure collaborative fairness in federated learning under the realistic yet challenging imbalanced covariate shift setting. To achieve this, in Section 2, we first provide a theoretical analysis in Theorem 1 demonstrating that imbalanced covariate shift – quantified as the KL divergence between each client’s empirical data distribution and the ideal global distribution – is primarily influenced by the perturbation $\delta$. Specifically, if the underlying data distributions of individual clients and the entire dataset follow a multivariate normal distribution $𝒩(\mu ,\mathrm{\Sigma })$, where $\mu$ is the mean vector, we show that $\delta$ and the covariance $\mathrm{\Sigma }$ play key roles in determining the extent of distributional divergence in Theorem 2. The two theorems motivate us to mitigate the imbalanced covariate shift to achieve collaborative fairness by correctly qualifying $\delta$ and $\mathrm{\Sigma }$. However, directly calculating these values is infeasible, as the true data distributions are inherently unknown.

Motivations of model design.

To address this challenge, we conducted a preliminary analysis on the results of FedAvg applied to the entire EHR dataset. For each client k, we categorized the correctly and incorrectly predicted samples as ${ℐ}_k$ and ${𝒟}_k-{ℐ}_k$, respectively, and aggregated these categories across all clients to form the global sets ${\left\{{ℐ}_k\right\}}_{k=1}^K$ and ${\left\{{𝒟}_k-{ℐ}_k\right\}}_{k=1}^K$, where ${𝒟}_k$ denotes the client dataset, and K is the number of clients. Next, we applied principal component analysis (PCA) to project each sample’s latent representation (i.e., the encoder output from each client model) onto a 1-D space and used kernel density estimation (KDE) to estimate the probability density function (PDF) for each set. Figures 2(a) and (b) show a comparison of the global density function with the density functions of two clients (Minnesota: MN and New Hampshire: NH). The x-axis represents the projected PCA 1-D values, and the y-axis represents the estimated PDF values. Similar to Figure 1(a), we also analyze the distribution differences between local and global data in terms of correct and incorrect classifications, as illustrated in Figure 2(c). These results indicate that the primary source of imbalanced covariate shift lies in the “incorrect” samples, as clients generally show agreement on the distribution of “correct” predictions.

Figure 2:

(a) Distribution of locally versus globally correct samples. (b) Distribution of locally versus globally incorrect samples. (c) KL divergence for “right” vs. “wrong” samples. We observe that the feature distributions of correctly classified samples closely resemble the global distribution, whereas those of misclassified samples deviate significantly. This suggests that imbalanced covariate shift primarily arises from incorrectly classified samples.

Our approach.

Building on our theoretical and empirical analysis, we propose FedAKD, a novel framework designed to address the imbalanced covariate shift challenge while ensuring collaborative fairness. FedAKD leverages a new federated asynchronous knowledge distillation approach, comprising client updates and server updates in each communication round t. Specifically, the client update includes three key steps: (1) Global → Local Distillation: Using traditional knowledge distillation [4], we employ the global model ${\mathbf{w}}_g^t$ as a teacher to guide the client model ${\mathbf{w}}_k^t$ in learning from its full training set ${𝒟}_k$. (2) High-confidence Sample Selection: Inspired by our observations in Figure 2, correctly predicted samples positively contribute to the global model update. Thus, we first identify the correctly classified samples from the updated local model, denoted as ${ℐ}_k^t$. (3) Local → Global Distillation: The global model ${\mathbf{w}}_g^t$ is then refined by distilling “high-confidence” client knowledge from ${\mathbf{w}}_k^t$ using each client’s correctly predicted set ${ℐ}_k^t$. This design helps mitigate distortions caused by misclassified data under covariate shift conditions. The updated global model for each client (i.e., ${\mathbf{w}}_{g,k}^t$) is uploaded to the server, where it is aggregated following FedAvg [15] in the server update step. These two updates iterate until FedAKD converges. The theoretical convergence of FedAKD is established in Theorem 3.

Through this two-stage asynchronous distillation process, FedAKD effectively encourages fair collaboration even when clients have highly divergent feature distributions. High-quality participants benefit by sharing more correct samples, whereas lower-quality participants gain from the improved global knowledge, thus collectively promoting collaborative fairness without imposing rigid or impractical contribution measurements.

Contributions.

The main contributions of this work include: (1) A new heterogeneity setting. We introduce collaborative fairness under the imbalanced covariate shift, a practical challenge in real-world medical datasets where both sample-size imbalance and feature-distribution mismatch significantly hinder the effectiveness of existing collaborative fairness metrics. (2) A simple yet effective solution. We propose FedAKD, a novel asynchronous distillation framework that first distills knowledge from correctly predicted local samples to improve the global model quality, followed by an inverse distillation step to enhance client learning across the full dataset. (3) Theoretical guarantees.³ We provide a rigorous theoretical analysis of imbalanced covariate shift, expanding the KL divergence parametrically to model real-world heterogeneity. Additionally, we prove the convergence of FedAKD under broad heterogeneity conditions, ensuring both theoretical soundness and improved collaborative fairness. (4) Promising results. We conduct extensive experiments on three datasets, evaluating FedAKD against ten baselines across four heterogeneity settings using three evaluation metrics. The results demonstrate that FedAKD effectively addresses imbalanced covariate shift and outperforms all baselines across diverse heterogeneity scenarios.

2. Imbalanced Covariate Shift Analysis

Imbalanced covariate shift presents a significant and open challenge in federated learning, requiring a deeper mathematical understanding to effectively mitigate disparities in collaborative fairness. To formalize this, we assume the existence of a global data feature distribution $p_{\omega }$. Each client’s data distribution is modeled as a small parametric perturbation, denoted as $p_{\omega +\delta }$, where $\delta \in {ℝ}^N$ represents the perturbation vector and $N$ is the model size. When a client draws an i.i.d. sample set of size $A$ from $p_{\omega +\delta }$, it results in the empirical distribution ${\widehat{p}}_{\omega +\delta }$. Consequently, the imbalanced covariate shift can be quantitatively assessed by measuring the KL divergence between the empirical distribution ${\widehat{p}}_{\omega +\delta }$ and the original global distribution $p_{\omega }$.

Theorem 1 (Ideal Imbalance Covariate Shift Quantification). Let ${\left\{p_{\theta }\right\}}_{\theta \in \mathrm{\Theta }}$ be a smooth parametric family of probability distributions, and let $\omega \in \mathrm{\Theta }$ denote a baseline parameter. Suppose that a perturbed distribution $p_{\omega +\delta }$ is defined by a small perturbation $\delta \in {ℝ}^N$. Given that $p_{\omega +\delta }$ is estimated via an empirical distribution ${\widehat{p}}_{\omega +\delta }$ using $A$ i.i.d. samples and $R$ (free) parameters, we have the following approximation under standard regularity conditions and a large-sample limit:

$$D_{\mathrm{K}\mathrm{L}}\left({\hat{p}}_{\omega +\delta }\Vert p_{\omega }\right)\approx \frac{1}{2}{\delta }^{\mathrm{\top }}I(\omega )\delta +\frac{1}{2}{\delta }^{\mathrm{\top }}\left({\nabla }_{\omega }I(\omega )\delta \right)\delta +\frac{R}{2A},$$

where $I(\omega )$ is the Fisher information matrix at $\omega$, and ${\nabla }_{\omega }I(\omega )$ denotes the gradients with respect to $\omega$.

We assume that the probabilistic distribution follows an $M$-dimensional Gaussian distribution. We then extend Theorem 1 as follows:

Theorem 2 (Imbalance Covariate Shift Quantification Under Gaussian Distribution). Let $p_{\mu ,\mathrm{\Sigma }}(x)=𝒩(x;\mu ,\mathrm{\Sigma })$ be a $M$-dimensional Gaussian distribution, where $\mu \in {ℝ}^M$ is the mean vector and $\mathrm{\Sigma }\in {ℝ}^{M\times M}$ is a symmetric positive-definite covariance matrix. Suppose $\omega =\left({\mu }_0,{\mathrm{\Sigma }}_0\right)$ is a baseline parameter, and consider a small perturbation $\delta =\left({\delta }_{\mu },{\delta }_{\mathrm{\Sigma }}\right),{\theta }^{\prime }=\omega +\delta \approx \left({\mu }_0+{\delta }_{\mu },{\mathrm{\Sigma }}_0+{\delta }_{\mathrm{\Sigma }}\right)$. Under a large-sample limit, the Kullback-Leibler divergence between the empirical distribution ${\widehat{p}}_{{\theta }^{\prime }}$ (fitted from $A$ i.i.d. samples drawn from $p_{{\theta }^{\prime }}$) and the baseline model $p_{\omega }$ can be approximated by:

$$\begin{gathered} D_{\mathrm{K}\mathrm{L}}\left({\widehat{p}}_{{\theta }^{\prime }}\Vert p_{\omega }\right)\approx \frac{1}{4}{‖{\mathrm{\Sigma }}^{-1}{\delta }_{\mathrm{\Sigma }}{\mathrm{\Sigma }}^{-1}‖}_F^2-\frac{1}{2}\mathrm{t}\mathrm{r}\mathrm{a}\mathrm{c}\mathrm{e}\left({\left({\delta }_{\mathrm{\Sigma }}{\mathrm{\Sigma }}^{-1}\right)}^3\right) \\ +\frac{1}{2}{\left({\delta }_{\mu }\right)}^{\mathrm{\top }}\left({\mathrm{\Sigma }}^{-1}\left(I-{\delta }_{\mathrm{\Sigma }}{\mathrm{\Sigma }}^{-1}\right)\right){\delta }_{\mu }+\frac{M(M+3)}{4A}. \end{gathered}$$ (1)

Together, these two theorems provide a unified mathematical framework to model the combined effects of imbalanced sample sizes and covariate shift. This framework provides a principled approach to analyzing how local client distributions diverge from the global baseline, offering deeper insights into collaborative fairness in federated learning.

3. The Proposed FedAKD

While our theoretical analysis in Section 2 provides valuable insights into imbalanced covariate shift, it cannot be directly applied to model design, as the global distribution remains unknown in federated learning. However, these theorems reveal that the imbalanced covariate shift arises due to small perturbations in client distributions. Our preliminary analysis (Figure 2 in Section 1) further suggests that these perturbations predominantly stem from incorrectly classified samples. This observation motivates us to develop an effective collaborative fairness approach named FedAKD that mitigates the imbalanced covariate shift by addressing the impact of misclassified samples via a simple asynchronous knowledge distribution strategy. The algorithm is shown in Algorithm 1.

Similar to existing collaborative fairness approaches in federated learning, FedAKD comprises both client and server updates. However, unlike prior methods that require carefully designing reward weights for each client [13, 19, 20, 23], FedAKD simplifies aggregation by directly following the standard FedAvg [15] in the server update. The novelty of FedAKD lies in the client update, where we introduce a new asynchronous knowledge distillation strategy, inspired by our preliminary analysis. Specifically, the client update consists of three key steps: (1) global → local distillation, (2) high-confidence sample selection, and (3) local → global distillation. Next, we provide the details of these three steps.

3.1. Step 1: Global → Local Distillation

The global model ${\mathbf{w}}_g^t$ contains aggregated knowledge, but forcing all clients to adopt ${\mathbf{w}}_{g,k}^t={\mathbf{w}}_g^t$ directly may harm local performance due to the imbalanced covariate shift. To avoid this issue, we propose global → local distillation, enabling each client to selectively adopt global insights while retaining local specialization by optimizing the following loss:

$$\overrightarrow{{ℒ}_t}=\mathrm{C}\mathrm{E}\left({\mathbf{w}}_k^{t-1};{𝒟}_k\right)+\alpha \mathrm{K}\mathrm{D}\left({\mathbf{w}}_k^{t-1};{\mathbf{w}}_{g,k}^t,{𝒟}_k\right),$$ (2)

where CE denotes the cross-entropy loss, KD is the knowledge distillation loss, and $\alpha$ is the hyperparameter. The gradient update yields:

$${\mathbf{w}}_k^t={\mathbf{w}}_k^{t-1}-\eta \nabla \overrightarrow{{ℒ}_t},$$ (3)

where $\eta$ is the learning rate. In this step, we fix the global model parameters ${\mathbf{w}}_{g,k}^t$ and only update the client model parameters ${\mathbf{w}}_k^{t-1}$ using the full training set ${𝒟}_k$. This procedure allows each client to merge the updated global knowledge with its local parameters, safeguarding performance for distribution-mismatched (yet high-quality) clients. Consequently, no client is penalized for joining the federation, reinforcing the incentives for collaborative fairness under the imbalanced covariate shift.

Algorithm 1.

FedAKD

Require: $K$ clients; local datasets $\left\{{𝒟}_k\right\}$; total rounds $T$; learning rate $\eta$; distillation coefficients $\alpha$ and $\beta$.

1: Initialization:

2:  Generate an initial model ${\mathbf{w}}^0$ (e.g., randomly);

3:  Client side (for each k): set each local model ${\mathbf{w}}_k^0={\mathbf{w}}^0$;

4:  Server side: set the global model ${\mathbf{w}}_g^1={\mathbf{w}}^0$ and distribute

5:  ${\mathbf{w}}_g^1$ to clients;

6: for $t=1,\cdots ,T$ do

7:  // Client Update

8:  for $k=\mathrm{1,\; 2},\cdots ,K$ do

9:   // Step 1: Global → Local Distillation

10:   ${\mathbf{w}}_{g,k}^t\leftarrow {\mathbf{w}}_g^t$;

11:   // Loss computation by fixing ${\mathbf{w}}_{g,k}^t$ using ${𝒟}_k$

12:   $\overrightarrow{{ℒ}_t}=\mathrm{C}\mathrm{E}\left({\mathbf{w}}_k^{t-1};{𝒟}_k\right)+\alpha \mathrm{K}\mathrm{D}\left({\mathbf{w}}_k^{t-1};{\mathbf{w}}_{g,k}^t,{𝒟}_k\right)$;

13:   //update model parameters

14:   ${\mathbf{w}}_k^t\leftarrow {\mathbf{w}}_k^{t-1}-\eta \nabla \overrightarrow{{ℒ}_t}$;

15:   // Step 2: High-confidence Sample Selection

16:   ${ℐ}_k^t=\left\{(\mathbf{x},y)\in {𝒟}_k\mid \mathrm{P}\mathrm{r}\mathrm{e}\mathrm{d}\left({\mathbf{w}}_k^t,\mathbf{x}\right)=y\right\}$;

17:   // Step 3: Local → Global Distillation

18:   // Loss computation by fixing ${\mathbf{w}}_k^t$ using ${ℐ}_k^t$

19:   $\overleftarrow{{ℒ}_t}=\mathrm{C}\mathrm{E}\left({\mathbf{w}}_{g,k}^t;{ℐ}_k^t\right)+\beta \mathrm{K}\mathrm{D}\left({\mathbf{w}}_{g,k}^t;{\mathbf{w}}_k^t,{ℐ}_k^t\right)$;

20:   // Update model parameters

21:   ${\mathbf{w}}_{g,k}^{t+1}\leftarrow {\mathbf{w}}_{g,k}^t-\eta \nabla \overleftarrow{{ℒ}_t}$;

22:   Upload ${\mathbf{w}}_{g,k}^{t+1}$ to the server;

23:  end for

24:  // Server Update

25:  ${\mathbf{w}}_g^{t+1}=\frac{1}{{\sum }_{k=1}^K\left|{𝒟}_k\right|}{\sum }_{k=1}^K\left(\left|{𝒟}_k\right|{\mathbf{w}}_{g,k}^{t+1}\right)$;

26:  Distribute ${\mathbf{w}}_g^{t+1}$ to each client;

27: end for

28: Output: The global model ${\mathbf{w}}_g^T$ and local models $\left\{{\mathbf{w}}_k^T\right\}$.

3.2. Step 2: High-confidence Sample Selection

Our preliminary analysis (Figures 2 in Section 1) suggests that imbalanced covariate shift primarily arises from misclassified samples on each client, whereas correctly classified samples positively contribute to global model learning. To address this, we select high-confidence samples (i.e., correctly classified samples) to update the global model ${\mathbf{w}}_{g,k}^t$, which is denoted as:

$${ℐ}_k^t=\left\{(\mathbf{x},y)\in {𝒟}_k\mid \mathrm{P}\mathrm{r}\mathrm{e}\mathrm{d}\left({\mathbf{w}}_k^t,\mathbf{x}\right)=y\right\}.$$ (4)

3.3. Step 3: Local → Global Distillation

Unlike existing bidirectional knowledge distillation [8,17] that updates two models simultaneously using the same dataset, we propose an asynchronous knowledge distillation approach for this step. Additionally, we leverage only the selected high-confidence samples ${ℐ}_k^t$, enabling the local model ${\mathbf{w}}_k^t$ to guide the learning of the global model ${\mathbf{w}}_{g,k}^t$ by optimizing the following loss:

$$\overleftarrow{{ℒ}_t}=\mathrm{C}\mathrm{E}\left({\mathbf{w}}_{g,k}^t;{ℐ}_k^t\right)+\beta \mathrm{K}\mathrm{D}\left({\mathbf{w}}_{g,k}^t;{\mathbf{w}}_k^t,{ℐ}_k^t\right),$$ (5)

where $\beta$ is the hyperparameter. The gradient update yields:

$${\mathbf{w}}_{g,k}^{t+1}={\mathbf{w}}_{g,k}^t-\eta \nabla \overleftarrow{{ℒ}_t},$$ (6)

The proposed asynchronous knowledge distillation offers three key benefits: (1) Robustness to Noisy Updates. It protects the global model from noisy or erroneous updates by discarding locally misclassified samples. (2) Fair Collaboration. It promotes fairness by allowing high-quality clients—those that classify more samples correctly—to have a stronger influence without explicitly revealing their accuracy or contributions. (3) Adaptability to Imbalanced Covariate Shift. It ensures that even if a client’s data distribution differs significantly from the global average, it can still contribute reliable knowledge.

3.4. Convergence Analysis

3.4.1. Notions and Assumptions.

In this section, we consider a binary classification problem following [16] with input space $X\in {ℝ}^d$ and label space $Y=\{\mathrm{0,\; 1}\}$. We employ a linear classification setting: for each local sample $\mathbf{x}\in X$, the logits are $z={\mathbf{x}}^{\mathrm{\top }}\mathbf{w}$, and the predicted probability is $\hat{y}(\mathbf{x})=\sigma (z)=\frac{1}{1+e^{-z}}$. In our codistillation setup, the distillation temperature is set to $\tau =1$, keeping the standard sigmoid form. We denote the cross-entropy loss by $ℒ(\mathbf{w};𝒟)=\mathrm{C}\mathrm{E}(\mathbf{w};𝒟)=\frac{1}{|𝒟|}{\sum }_{{\mathbf{x}}_i\in 𝒟}\left[-y_i\mathrm{l}\mathrm{o}\mathrm{g}\left(\hat{y}\left({\mathbf{x}}_i\right)\right)-\left(1-y_i\right)\mathrm{l}\mathrm{o}\mathrm{g}\left(1-\hat{y}\left({\mathbf{x}}_i\right)\right)\right]$. Here, $\mathbf{w}$ is the model parameter vector, and $𝒟$ is the training dataset consisting of samples ${\mathbf{x}}_i$ with labels $y_i$. We further introduce the KD loss, denoted by $\mathrm{K}\mathrm{D}\left(\mathbf{w},{\mathbf{w}}_0;𝒟\right)=\frac{1}{|𝒟|}{\sum }_{{\mathbf{x}}_i\in 𝒟}\left[-{\hat{y}}_0\left({\mathbf{x}}_i\right)\mathrm{l}\mathrm{o}\mathrm{g}\left(\hat{y}\left({\mathbf{x}}_i\right)\right)-(1-{\hat{y}}_0\left({\mathbf{x}}_i\right))\mathrm{l}\mathrm{o}\mathrm{g}\left(1-\hat{y}\left({\mathbf{x}}_i\right)\right)\right]$. This can also be written in expectation form as ${𝔼}_{{\mathbf{x}}_i\in 𝒟}\left[-{\hat{y}}_0\left({\mathbf{x}}_i\right)\mathrm{l}\mathrm{o}\mathrm{g}\left(\hat{y}\left({\mathbf{x}}_i\right)\right)-\left(1-{\hat{y}}_0\left({\mathbf{x}}_i\right)\right)\mathrm{l}\mathrm{o}\mathrm{g}\left(1-\hat{y}\left({\mathbf{x}}_i\right)\right)\right]$. Here, ${\hat{y}}_0\left({\mathbf{x}}_i\right)$ is the (fixed) teacher model’s output, ${\hat{y}}_0\left({\mathbf{x}}_i\right)=\sigma \left({\mathbf{x}}_i^{\mathrm{\top }}{\mathbf{w}}_0\right)$, and ${\mathbf{w}}_0$ denotes the teacher’s parameter vector. This KD objective is equivalent (up to a constant) to minimizing the KL divergence from the teacher distribution [10, 16].

We define the notion of a $𝛾$-inexact solution following [10, 16], which quantifies the improvement made by local updates:

Definition 1 (${𝛾}_1$-inexact solution [10, 16]). For a function $\overrightarrow{ℒ}\left(\mathbf{w},{\mathbf{w}}_0;𝒟\right)=ℒ(\mathbf{w};𝒟)+\alpha \mathrm{K}\mathrm{D}\left(\mathbf{w},{\mathbf{w}}_0;𝒟\right)$, and let ${𝛾}_1\in [0,1]$. Suppose ${\mathbf{w}}_0$ is an initial point for ${\mathrm{m}\mathrm{i}\mathrm{n}}_{\mathbf{w}}\overrightarrow{ℒ}\left(\mathbf{w};{\mathbf{w}}_0;𝒟\right)$. We say ${\mathbf{w}}^{\mathrm{*}}$ is a ${𝛾}_1$-inexact solution if $‖\nabla \overrightarrow{ℒ}\left({\mathbf{w}}^{\mathrm{*}},{\mathbf{w}}_0;𝒟\right)‖\le {𝛾}_1‖\nabla \overrightarrow{ℒ}\left({\mathbf{w}}_0,{\mathbf{w}}_0;𝒟\right)‖$.

A smaller ${𝛾}_1$ indicates a greater reduction in the gradient norm relative to the initial point, implying more significant local improvement. Conversely, a larger ${𝛾}_1$ indicates a less complete local optimization. Similarly, for the function $\overleftarrow{ℒ}$ with coefficient $\beta$. We define $\overleftarrow{ℒ}\left(\mathbf{w},{\mathbf{w}}_0;𝒟\right)=ℒ(\mathbf{w};𝒟)+\beta \mathrm{K}\mathrm{D}\left(\mathbf{w},{\mathbf{w}}_0;𝒟\right)$, and let ${𝛾}_2\in [0,1]$. we say ${\mathbf{w}}^{\mathrm{*}}$ is a ${𝛾}_2$-inexact solution with $‖\nabla \overleftarrow{ℒ}\left({\mathbf{w}}^{\mathrm{*}},{\mathbf{w}}_0;𝒟\right)‖\le {𝛾}_2‖\nabla \overleftarrow{ℒ}\left({\mathbf{w}}_0,{\mathbf{w}}_0;𝒟\right)‖$.

Assumption 1 ($L$-smoothness [10, 11]). There exists $L>0$ such that for all $\mathbf{w},{\mathbf{w}}^{\prime },‖\nabla ℒ(\mathbf{w};𝒟)-\nabla ℒ\left({\mathbf{w}}^{\prime };𝒟\right)‖\le L‖\mathbf{w}-{\mathbf{w}}^{\prime }‖$.

Assumption 2 ($\mu$-strong convexity [10, 11]). There exists $\mu >0$ such that for all $\mathbf{w},{\mathbf{w}}^{\prime },ℒ(\mathbf{w};𝒟)\ge ℒ\left({\mathbf{w}}^{\prime };𝒟\right)+\nabla ℒ{\left({\mathbf{w}}^{\prime };𝒟\right)}^{\mathrm{\top }}(\mathbf{w}-{\mathbf{w}}^{\prime })+\frac{\mu }{2}{‖\mathbf{w}-{\mathbf{w}}^{\prime }‖}^2,\overrightarrow{ℒ}\left(\mathbf{w},{\mathbf{w}}_0;𝒟\right)\ge \overrightarrow{ℒ}\left({\mathbf{w}}^{\prime },{\mathbf{w}}_0;𝒟\right)+\nabla \overrightarrow{ℒ}{\left({\mathbf{w}}^{\prime },{\mathbf{w}}_0;𝒟\right)}^{\mathrm{\top }}\left(\mathbf{w}-{\mathbf{w}}^{\prime }\right)+\frac{\mu }{2}{‖\mathbf{w}-{\mathbf{w}}^{\prime }‖}^2$, and $\overleftarrow{ℒ}\left(\mathbf{w},{\mathbf{w}}_0;𝒟\right)\ge \overleftarrow{ℒ}\left({\mathbf{w}}^{\prime },{\mathbf{w}}_0;𝒟\right)+\nabla \overleftarrow{ℒ}{\left({\mathbf{w}}^{\prime },{\mathbf{w}}_0;𝒟\right)}^{\mathrm{\top }}\left(\mathbf{w}-{\mathbf{w}}^{\prime }\right)+\frac{\mu }{2}{‖\mathbf{w}-{\mathbf{w}}^{\prime }‖}^2$.

Assumption 3 (Bounded Gradient Dissimilarity [10, 16]). Let ${𝒟}_g≔{\cup }_k{𝒟}_k$ be the global dataset consisting of all local datasets ${𝒟}_k$. For some $ϵ>0$, define ${𝒮}_c^{ϵ}=\left\{\mathbf{w}\mid {‖\nabla ℒ\left(\mathbf{w};{𝒟}_g\right)‖}^2>ϵ\right\}$. There exists $B_{ϵ}$ such that for all $\mathbf{w}\in {𝒮}_c^{ϵ},B(\mathbf{w})=\sqrt{\frac{{𝔼}_k\left[{‖\nabla ℒ\left(\mathbf{w};{𝒟}_k\right)‖}^2\right]}{{‖\nabla ℒ\left(\mathbf{w};{𝒟}_g\right)‖}^2}}\le B_{ϵ}$.

Here, $B(\mathbf{w})$ measures data heterogeneity across devices. If data are IID and $n_k\to \mathrm{\infty }$, then $B(\mathbf{w})\to 1$. Generally, $B_{ϵ}\ge 1$, and larger values capture more dissimilar data distributions.

Assumption 4 (Bounded Gradient Dissimilarity on Subset). Let $ℐ\subseteq 𝒟$ be a subset. There exists $\theta \ge 0$ such that for all $\mathbf{w}$, $\Vert \nabla ℒ(\mathbf{w};ℐ)-\nabla ℒ(\mathbf{w};𝒟)\Vert \le \theta \Vert \nabla ℒ(\mathbf{w};𝒟)\Vert$.

This condition ensures that the gradient on a chosen subset does not deviate excessively from the gradient on the entire local dataset, thus quantifying the heterogeneity between these two distributions.

3.4.2. Main Results.

Theorem 3 (FedAKD Convergence). Under Assumptions 1–4, assume that ${\mathbf{w}}^t$ is not a stationary solution and the loss function $ℒ$ is $B$-dissimilar, i.e., $B\left({\mathbf{w}}^t\right)\le B$. If $\alpha ,\beta$ and $𝛾≔\mathrm{m}\mathrm{a}\mathrm{x}\left\{{\gamma }_1,{\gamma }_2\right\}$ are chosen such that $r=\left(\frac{4}{\beta ‖{\mathrm{\Omega }}_2‖}+\frac{4}{\alpha ‖{\mathrm{\Omega }}_1‖}\right)B-\frac{L{B}^2}{2}((L(1+𝛾)+\mu )\frac{𝛾(1+\theta )}{\mu }+(1+\theta )+\frac{(1+𝛾)\beta ‖{\mathrm{\Omega }}_2‖}{4\mu }{)}^2-\left(\frac{4}{\beta ‖{\mathrm{\Omega }}_2‖}+\frac{4}{\alpha ‖{\mathrm{\Omega }}_1‖}\right)\left(r_1+r_2\right)B>\mathit{0}$, where $r_1=\left((L(1+𝛾)+\mu )\frac{𝛾(1+\theta )}{\mu }+(1+\theta )+\frac{(1+𝛾)\beta ‖{\mathrm{\Omega }}_2‖}{4\mu }\right)(1+\theta )L+\theta +(L(1+𝛾)+\mu )\frac{𝛾(1+\theta )}{\mu }),r_2=\frac{L(1+𝛾)}{\mu }+𝛾,{\mathrm{\Omega }}_1={𝔼}_k\left[{𝔼}_{{\mathbf{x}}_{k,i}\in {𝒟}_k}\left[{\mathbf{x}}_{k,i}{\mathbf{x}}_{k,i}\right]\right]$ and ${\mathrm{\Omega }}_2={𝔼}_k\left[{𝔼}_{{\mathbf{x}}_{k,i}\in {ℐ}_k^t}\left[{\mathbf{x}}_{k,i}{\mathbf{x}}_{k,i}\right]\right]$, then FedAKD satisfies $ℒ\left({\mathbf{w}}_g^{t+1};{𝒟}_g\right)-ℒ\left({\mathbf{w}}^{\mathrm{*}};{𝒟}_g\right)\le (1-2\mu r)\left[ℒ\left({\mathbf{w}}_g^t;{𝒟}_g\right)-ℒ\left({\mathbf{w}}^{\mathrm{*}};{𝒟}_g\right)\right]$.

4. Simulation Experiments

4.1. Non-IID Settings

Since most real-world federated applications involve non-IID data distributions, particularly in the imbalanced covariate shift setting, we consider the following three non-IID client partitions in our simulation experiments:

• Imbalanced Dataset Sizes (POW) [13, 23]: Each client’s dataset size $\left|{𝒟}_k\right|$ follows a power-law distribution, leading to significant disparities in data quantity across clients. However, the feature distributions remain similar across clients.

• Balanced Covariate Shift (BCS): Clients exhibit covariate shifts in their data feature distributions while maintaining a similar number of samples.

• Imbalanced Covariate Shift (ICS): ICS combines the characteristics of POW and BCS, where client dataset sizes follow a power-law distribution (POW), and different institutions experience significant variations in feature distributions.

4.2. Federated Data Simulation

In the simulation experiments, we use two image classification datasets: Fashion MNIST [22] and CIFAR10 [7]. Fashion MNIST contains 70,000 grayscale images (28×28) evenly split into 10 classes (e.g., T-shirt/top, trousers). CIFAR10 consists of 60,000 color images (32 × 32) across 10 classes (e.g., airplane, bird). We partition the datasets into training, validation, and testing in a ratio of 7:1:2. In addition, we set the number of clients as K = 10. To simulate the POW partition, we follow a power law with an exponent of 1 to divide the global data into 10 clients. For the k-th client, its data size is $\left|{𝒟}_k\right|=\frac{1}{kZ}\left|{𝒟}_g\right|$, where $Z={\sum }_{k=1}^{10}\frac{1}{k}$.

Simulating the covariate shift setting is nontrivial, and as far as we know, no existing work provides an automated way to generate such federated datasets. To fill this research gap, we design a novel algorithm based on Theorem 2 for covariate shift data generation, as shown in Appendix C Algorithm 2. To stimulate the BCS partition, we set c = 5 and add a perturbation $\delta$ (satisfying ${\delta }^{\mathrm{\top }}{\mathrm{\Sigma }}^{-1}\delta =5$) to the global mean. In such a way, each client receives an equal number of samples (i.e., balanced), thus focusing on the covariate shift while keeping dataset sizes uniform. To stimulate the ICS partition, we run Algorithm 2 again to produce three variants with C = 2, C = 5, and C = 10, representing increasing levels of covariate shift. Meanwhile, the total number of samples is partitioned among the 10 clients according to the power law distribution with the exponent as 1, thus coupling imbalanced dataset sizes with covariate shifts.

4.3. Baselines

We compare our method against two categories of baselines: Collaborative Fairness algorithms designed for non-IID data, and Covariate Shift algorithms focusing on feature-level discrepancies. Specifically, we include CGSV [23], CFFL [13], FedAVE [19], and FedSAC [20], which explicitly address fairness by measuring client contributions or customizing reward allocations. Meanwhile, FedAS [25], FedDC [1], pFedCK [30], and FedMPR [2] focus on mitigating feature-level drift across clients. We also include two traditional baselines: Standalone (each client trains independently without aggregation) and the classic FedAvg [15] for federated averaging. The details of the baselines can be found in Appendix A.

4.4. Implementation

We implement a network consisting of two convolutional layers (each followed by batch normalization and ReLU activation), inter-leaved with max-pooling, and ending with a fully connected output layer for the simulation evaluation. We implement all baselines and our model in PyTorch and train them on an NVIDIA RTX A6000 GPU. All details of parameter setting can be found in Appendix D.

4.5. Evaluation Metrics

Due to the imbalanced covariate shift setting, each client data has a unique distribution. Thus, we conduct local evaluations and then report the average values of all the clients for three runs. Let Acc_p) denote the prediction accuracy of clients after federation. Following [20], we use three metrics:

• Average Client Accuracy, i.e., $\frac{{\sum }_k{\mathbf{A}\mathbf{c}\mathbf{c}}_p[k]}{K}$;

• Maximum Client Accuracy, i.e, $\mathrm{m}\mathrm{a}\mathrm{x}\left({\mathbf{A}\mathbf{c}\mathbf{c}}_p\right)$;

• Collaborative Fairness (CF) Coefficient, which reflects how uniformly performance is distributed across clients. We use the CF defined in [13], i.e, $CF=100\times \rho \left({\mathbf{A}\mathbf{c}\mathbf{c}}_s,{\mathbf{A}\mathbf{c}\mathbf{c}}_p\right)\in [-\mathrm{100,\; 100}]$, where $\rho (\cdot ,\cdot )$ is Pearson’s correlation coefficient, and Acc_s represents the standalone accuracy of clients.

The greater these three metric values, the better the performance.

4.6. Results of Simulation Experiments

Table 1 presents the results on FashionMNIST and CIFAR10 under three types of non-IID partitions. As shown in Table 1(a), our method consistently achieves superior or highly competitive Max Client Accuracy across all partitions for both datasets. Table 1(b) further highlights our approach’s advantage in Avg Client Accuracy, particularly on the more challenging CIFAR10 dataset. The most critical metric, Collaborative Fairness (CF), is reported in Table 1(c), where FedAKD demonstrates significantly higher fairness under varying degrees of non-IID settings. These results collectively validate the effectiveness of FedAKD in improving both accuracy and fairness.

Table 1:

Performance evaluation of simulation experiments on two datasets with three runs.

(a) Max Client Accuracy (%)
Method	FashionMNIST					CIFAR10
	POW	BCS	ICS(2.0)	ICS(5.0)	ICS(10.0)	POW	BCS	ICS(2.0)	ICS(5.0)	ICS(10.0)
Standalone	88.09±0.06	96.06±0.16	96.93±0.16	96.80±0.05	96.94± 0.02	54.86±1.93	54.53±0.66	58.82±1.19	59.95±1.29	56.20±1.71
FedAvg	94.47±0.46	97.50±0.14	94.57±0.32	96.23±0.16	94.32±0.13	67.08±0.17	65.80±0.65	66.08±0.15	64.68±0.60	66.67±0.95
CFFL	94.66±0.42	99.61±0.16	97.41±2.44	98.18±1.59	98.63±1.19	67.76±1.87	70.71±3.34	71.37±0.07	64.34±4.98	69.27±2.03
CGSV	96.33±0.31	98.64±0.11	98.29±0.52	98.51±0.25	97.97±0.09	76.16±2.13	76.31±0.26	76.04±2.28	72.29±2.05	78.40±1.37
FedAVE	92.42±0.93	98.72±0.75	96.44±0.76	95.46±1.83	96.78±0.14	62.39±1.56	58.27±0.09	60.53± 3.10	60.79±0.86	60.96±1.26
FedSAC	96.42±0.56	96.73±0.45	95.33±0.64	97.31±1.34	96.34±0.33	65.27±0.22	63.04±2.57	65.42±0.06	64.13±0.84	65.59±0.08
pFedCK	94.15±0.18	97.52±0.37	97.11±0.15	98.62±0.54	98.25±0.87	80.13±0.87	79.54±0.31	80.88±1.24	79.99±1.46	80.71±0.98
FedDC	91.71±0.08	96.22±0.40	94.13±0.42	94.45±0.39	94.37±0.13	67.10±1.06	64.67±0.25	66.11±0.64	65.63±1.34	65.81±0.65
FedAS	88.20±0.38	98.89±0.21	96.08±0.23	96.76±0.20	96.37±0.33	75.56±1.23	73.13±0.66	74.59±0.79	74.25±1.27	75.58±1.61
FedMPR	94.80±0.23	97.72±0.21	95.20±0.67	95.93±0.33	95.88±0.43	69.06±0.36	66.40±0.57	66.03±0.54	67.84±1.43	69.03±0.80
FedAKD	97.45±0.09	99.67±0.14	99.56±0.04	99.66±0.04	99.56±0.11	83.80±1.12	81.27±0.19	81.93±1.03	82.73±0.38	81.30±1.43
(b) Average Client Accuracy (%)
Method	FashionMNIST					CIFAR10
	POW	BCS	ICS(2.0)	ICS(5.0)	ICS(10.0)	POW	BCS	ICS(2.0)	ICS(5.0)	ICS(10.0)
Standalone	86.26±0.06	92.04±0.03	89.15±0.19	89.50±0.05	90.01±0.13	45.44±0.13	50.03±0.45	47.25±0.38	47.20±0.96	46.33±0.75
FedAvg	91.16±0.24	92.98±0.06	91.60±0.13	91.98±0.08	91.28±0.03	63.95±1.00	62.37±0.66	63.10±0.16	62.03±0.55	62.14±0.69
CFFL	87.79±0.91	93.92±0.04	89.43±0.56	88.39±0.83	87.36±0.36	53.91±3.31	62.88±4.27	46.14±2.92	48.27±3.29	54.18±2.71
CGSV	91.23±0.15	92.21±0.35	90.52±1.35	89.32±0.36	91.14±0.15	66.68±1.50	67.00±2.93	64.32±2.32	63.65±2.11	65.56±1.35
FedAVE	89.44±0.45	92.36±0.19	89.83±0.47	88.99±0.64	89.66±0.04	55.85±1.44	52.19±1.19	55.32±2.09	51.85±2.41	53.42±1.84
FedSAC	89.54±0.63	88.92±0.83	89.75±0.92	88.65±0.54	89.76±0.63	54.63±0.32	60.53±0.81	60.52±0.29	57.52±0.72	60.84±0.64
pFedCK	90.83±0.18	93.63±0.28	91.82±0.28	91.04±0.77	91.98±0.25	67.51±0.52	64.18±0.84	65.19±0.27	65.65±0.71	65.88±0.39
FedDC	87.83±0.46	91.08±1.46	90.06±0.56	90.20±0.24	90.00±0.42	55.34±0.65	53.13±1.31	52.98±0.05	52.30±1.73	55.83±0.73
FedAS	84.36±0.22	91.47±0.23	88.27±0.40	87.84±0.27	88.22±0.24	64.64±0.75	64.68±0.30	63.75±0.60	64.39±0.32	63.81±0.54
FedMPR	91.13±0.04	92.90±0.07	91.74±0.05	91.73±0.13	91.71±0.19	64.95±0.56	62.79±0.48	63.77±0.06	63.20±0.58	64.67±0.07
FedAKD	93.50±0.19	95.68±0.15	92.62±0.20	92.47±0.14	92.92±0.12	70.96±0.58	67.59±0.49	68.96±0.31	68.50±0.42	68.55±1.35
(c) Collaborative Fairness (CF) Coefficient
Method	FashionMNIST					CIFAR10
	POW	BCS	ICS(2.0)	ICS(5.0)	ICS(10.0)	POW	BCS	ICS(2.0)	ICS(5.0)	ICS(10.0)
FedAvg	45.45±3.60	62.43±1.11	71.76±2.14	87.60±4.47	74.43±2.37	3.62±25.88	26.00±9.58	67.60±6.97	33.97±27.58	74.79±3.52
CFFL	33.85±15.44	79.80±10.74	60.80±20.34	74.22±6.33	67.01±8.69	45.90±31.33	16.94±32.09	55.03±15.49	13.17±7.51	26.82±11.03
CGSV	38.11±14.34	69.03±6.63	52.58±7.38	57.32±11.69	68.46±16.67	36.25±26.04	45.87±19.76	18.50±21.40	−4.15±18.53	5.54±22.85
FedAVE	21.43±23.31	65.10±1.47	78.15±8.32	80.02±7.00	65.79±4.98	16.60±18.27	49.08±15.03	17.07±22.03	24.56±34.43	25.51±27.50
FedSAC	28.10±7.93	79.54±5.13	48.78±14.97	83.30±5.13	71.26±4.76	62.04±12.03	43.61±10.38	56.89±7.30	23.90±12.70	20.71±19.95
pFedCK	23.15±5.26	41.09±13.56	34.15±9.36	36.65±8.61	15.27±15.31	51.14±9.15	70.78±11.15	24.64±21.54	8.76±19.93	36.83±10.66
FedDC	−40.85±18.38	−11.12±22.49	2.39±15.37	18.37±9.41	−15.32±38.39	−4.22±32.08	−9.54±31.25	−21.57±5.97	−9.38±27.68	−0.52±42.33
FedAS	45.09±2.61	75.81±4.81	27.74±4.12	79.61±3.11	70.23±3.12	61.34±6.55	74.28±9.43	70.01±6.07	60.54±4.05	76.54±4.10
FedMPR	31.33±14.55	83.20±2.59	70.43±2.21	76.66±5.21	61.59±3.03	50.93±14.09	40.63±4.25	29.78±18.15	44.96±6.39	53.06±9.69
FedAKD	70.61±5.82	86.88±0.81	78.25±2.02	89.51±2.40	79.72±1.47	88.02±2.09	82.15±2.23	81.25±1.55	82.53±2.42	84.17±2.77

5. Real-World Experiments

5.1. Experimental Settings

The EHR Dataset is extracted from the TriNetX database⁴, which contains patients’ claims data from all 50 states in the USA. This dataset is curated for the early prediction of pancreatic cancer and includes 259,480 de-identified patient records (161,345 negative vs. 98,135 positive). It consists of both static features (sex, zip code) and time-series events (medications, lab tests, vital signs, etc.). Further details on the EHR dataset and experimental settings are provided in Appendix E. In this experiment, we adopt a two-layer bidirectional GRU with attention mechanisms [3, 14, 18, 21, 26, 27, 29] to predict whether a patient eventually develops pancreatic cancer. We use the same baselines and evaluation metrics as the simulation experiments.

5.2. Performance Analysis

Table 2 presents the results on the real-world EHR dataset, which naturally follows an ICS (Imbalanced Covariate Shift) pattern due to varying sample sizes and feature distributions across different medical institutions. As shown in the table, FedAKD significantly outperforms the baselines, particularly in the CF metric, demonstrating a more equitable distribution of benefits among clients. While FedMPR achieve competitive Max Acc, their fairness metrics remain comparatively limited.

Table 2:

Experimental results on the EHR dataset.

Method	CF	Max Acc	Avg. Acc
Standalone	–	76.27±0.13	70.41±0.04
FedAvg	−12.63±5.22	74.75±1.24	70.13±0.21
CFFL	52.63±12.43	73.26±2.49	69.57±0.92
CGSV	46.39±4.24	72.55±1.71	69.31±0.42
FedAVE	38.24±17.24	75.66±1.59	69.24±0.81
FedSAC	70.54±1.24	74.21±0.32	68.33± 0.56
pFedCK	35.28±13.74	76.87±0.18	70.01±0.10
FedDC	17.10±3.21	75.08±0.32	69.61±0.29
FedAS	8.50±3.28	74.89±0.29	69.05±0.25
FedMPR	−5.59±2.31	76.39±0.25	70.10±0.13
FedAKD	78.42±1.09	78.98±1.01	71.23±0.27

5.3. Ablation Studies

To investigate how each component of the proposed FedAKD (Algorithm 1) contributes to its final performance, we conduct three ablation experiments. In each variant, we selectively remove or modify part of the procedure to gauge its impact: (1) Local → Global Distillation Using All Data ${𝒟}_k$. The biggest key finding of this work is to use the correctly predicted subset ${ℐ}_k^t$ to guide the update of ${\mathbf{w}}_{g,k}^t$. This baseline uses all local samples ${𝒟}_k$ for local → global distillation to test whether restricting to accurately labeled data is necessary for effective knowledge distillation. (2) Only Local → Global Distillation. After receiving the aggregated global model, each client simply overwrites its local model with ${\mathbf{w}}_g^t$, instead of performing an additional global → local distillation. This variant helps isolate the effect of double-direction distillation. (3) Correct Sample Count for Aggregation. In the standard procedure, the server weights each client’s update by the total local dataset size. In this ablation, we replace that term with the number of correctly predicted local samples $\left(\left|{ℐ}_k^t\right|\right)$, thus investigating whether “correct sample counts” lead to better aggregation fairness or performance.

The results of these ablation studies are summarized in Table 4. Our findings indicate that the proposed asynchronous knowledge distillation strategy has the most significant impact on both performance and fairness. Additionally, conducting knowledge distillation on the entire dataset negatively affects model performance. However, using the size of either the full dataset or only correctly classified data does not lead to notable performance differences. These results validate the rationale behind our model design.

Table 4:

Ablation study on three variants of FedAKD versus the full method. CF denotes a fairness metric among clients (higher is better), whereas Max Acc and Avg. Acc refer to the maximum and average accuracies (in %), respectively, across global rounds.

Method	CF	Max Acc	Avg. Acc
FedAKD (All-Data)	66.02±2.93	75.20±0.82	70.66±0.42
FedAKD (Single-Dist)	58.49±4.32	73.88±1.61	64.02±2.82
FedAKD (Correct-Agg.)	77.31±1.98	78.22±0.91	71.17±0.66
FedAKD (Full)	78.42±1.09	78.98±1.01	71.23±0.27

5.4. Computational Cost Analysis

In our model design, the client training contains three steps, while baselines also use other techniques to calculate rewards. To validate the efficiency of our model, we show the average per-round running times (in seconds) as depicted in Figure 3a. Although FedAKD shows a slightly higher time consumption per round, primarily due to asynchronous knowledge distillation, its computational overhead is still comparable to other baselines (e.g., CFFL and FedAVE, which involve additional validation or sparsification steps). To further evaluate compute efficiency, we fixed FedAKD to 40 rounds, recorded its total wall-clock time, and then allowed every baseline to train for the same duration. Table 3 shows that—even under this strict budget—FedAKD achieves the largest fairness improvement and the highest predictive performance.

Figure 3:

Efficiency comparison of FedAKD and baselines. (a) Average per-round runtime; (b) Training loss trajectory on the global dataset.

Table 3:

Performance of all methods when constrained to the same wall-clock time as FedAKD (40 rounds).

Method	Round	CF	Max Acc	Avg. Acc
FedAvg	80	−12.63±5.22	74.75±1.24	70.13±0.21
CFFL	43	48.63±11.43	75.25±1.42	67.31±0.95
CGSV	61	33.11±6.31	73.62±0.55	68.84±0.84
FedAVE	39	58.62±9.35	70.35±1.24	66.52±0.99
FedSAC	65	60.22±3.21	74.02±0.44	68.89±0.21
pFedCK	58	33.51±8.51	75.01±0.43	69.11±0.30
FedDC	51	30.10±5.32	73.54±1.35	68.53±0.88
FedAS	59	15.89±4.09	72.88±1.09	68.66±0.11
FedMPR	47	8.22±4.24	74.39±0.21	69.28±0.26
FedAKD	40	71.89±1.82	77.10±0.58	71.02±0.53

5.5. Convergence Analysis

We next analyze the convergence behavior of these methods by evaluating the global parameter ${\mathbf{w}}_g^t$ on the global EHR dataset, i.e., computing $ℒ\left({\mathbf{w}}_g^t;{𝒟}_g\right)$ at each round. Figure 3b shows two representative convergence curves, comparing FedAvg and FedAKD. We observe that FedAKD converges slightly more slowly in the early stages; however, it remains stable and ultimately achieves a low global loss. This result empirically verifies Theorem 3, demonstrating that the global model in FedAKD converges effectively to a desirable minimum on ${𝒟}_g$

6. Related Work

This work mainly focuses on collaborative fairness (CF) in federated learning, which regards the global model as the core reward and seeks to ensure that the final performance of each client reflects its actual contribution. For instance, CGSV [23] computes a cosine gradient Shapley value to measure how closely each client’s local gradient aligns with the global gradient, and allocates model updates based on this similarity. CFFL [13] relies on a public validation set to evaluate each client’s data diversity and local-model performance, then allocates rewards accordingly. FedAVE [19] computes reputation by examining each client’s local model performance and data distribution, offering better adaptability to various distributional scenarios. FedSAC [20] avoids the need for a global validation set by distributing varying submodels to high-contribution clients. However, it depends on pre-known standalone training results, which may be unrealistic in real-world deployments. Moreover, submodel-based pruning alone cannot fully address feature-distribution mismatch, as high-quality but distribution-mismatched data may still be underrepresented.

7. Conclusion

This paper investigates a practical yet challenging form of heterogeneity that impacts collaborative fairness: imbalanced covariate shift. To address this issue, we propose a novel approach, FedAKD (Federated Asynchronous Knowledge Distillation), which mitigates the effects of imbalanced covariate shift by excluding incorrectly predicted samples from the global model update—an insight derived from our preliminary findings. Experimental results on three datasets compared against ten baselines demonstrate the effectiveness and fairness of FedAKD across various heterogeneity settings in federated learning.

CCS Concepts.

• Information systems → Information systems applications.

A. Baselines

In this appendix, we present the details of the baseline methods used in our experiments. We classify them into two main categories and also include two additional traditional baselines: Standalone training and the classic FedAvg.

Standalone and FedAvg.

The first traditional baseline is Standalone, in which each client trains independently without any model aggregation. This approach ignores the potential benefits of federated collaboration, providing a lower-bound performance reference. Next, we employ FedAvg [15], the standard federated averaging method. FedAvg updates the global model by performing weighted averaging of the locally trained models from all clients, thereby enabling knowledge sharing while minimizing data exchange.

Collaborative Fairness Baselines.

We next consider a set of baselines specifically aimed at improving collaborative fairness. CFFL [13] focuses on distributing rewards proportionally by measuring each client’s contribution differences. In practice, it evaluates client updates on a held-out validation set (or an equivalent performance benchmark) before deciding how to compensate high- versus low-contribution clients. On the other hand, CGSV [23] elies on gradient-based importance metrics—specifically using cosine similarities of gradients—so it does not require a separate validation set. Each client’s contribution is gauged by comparing its gradient direction against the aggregated global gradient. Additionally, FedAVE [19] incorporates explicit fairness by maintaining a global validation set at the server. It periodically tests each client’s model update on this validation set to compute a “reputation” score, which then guides how the final global model is aggregated. FedSAC [20] similarly uses a global validation set for evaluating each client’s local update. Clients receive a proportionate “reward gradient” based on their evaluated performance, ensuring that clients with higher impact on the validation set obtain a larger share of the global update.

Personalized FL (Covariate Shift) Baselines.

Since our approach also falls under the umbrella of Personalized Federated Learning (FL), we include baselines originally designed to tackle feature-level (covariate) shifts, even though they do not explicitly target collaborative fairness. FedDC [1] tackles non-IID data by introducing a drift variable that aligns local models more closely with the global model. FedAS [25] alleviates intra- and inter-client inconsistencies through federated parameter alignment and client synchronization. pFedCK [30] clusters clients by update similarity and performs mutual knowledge distillation between interactive and personalized models to enhance robustness under data heterogeneity. Lastly, FedMPR [2] combines iterative magnitude pruning with regularization techniques to improve robustness under highly heterogeneous client data distributions.

By comparing these diverse baselines, we can comprehensively evaluate our proposed method from both collaborative fairness and feature-level drift perspectives.

B. Hyperparameter Tuning

In Figure 4, we vary both $\alpha$ (global-to-local distillation) and $\beta$ (local-to-global distillation) from 0.2 to 1.0 (25 grid points in total), and observe how they affect three metrics: collaborative fairness (CF), max accuracy, and mean accuracy. In subfigure 4(a), the CF score reaches its peak (around 78) when $\alpha =\beta =0.8$. Moreover, when both $\alpha$ and $\beta$ exceed 0.8, the fairness scores generally remain high, suggesting that stronger global-local knowledge exchange can boost collaborative fairness. By contrast, if $\beta$ becomes too small, the local models fail to effectively transfer information back to the global model, thus degrading overall fairness. On the other hand, subfigures 4(b) and 4(c) indicate that $\alpha$ and $\beta$ have a relatively minor impact on max and mean accuracy. Hence, to achieve a good trade-off between fairness and accuracy, choosing values of $\alpha$ and $\beta$ near 1.0 is preferable in practice.

C. Covariate Shift Data Generation

In this section, we detail how to construct a covariate-shifted dataset from a baseline distribution with Algorithm 2. We assume the underlying global dataset ${𝒟}_g$ approximately follows a single Gaussian distribution $𝒩(\mu ,\mathrm{\Sigma })$, from which we draw subsets for different clients. For each client, we shift the features by a fixed Mahalanobis distance.

As discussed in Theorem 2, under a large-sample limit, the Kullback-Leibler (KL) divergence can be used to measure how much the client-specific distribution $p_{{\theta }^{\prime }}$ deviates from the baseline model $p_{\omega }$. By fixing ${‖{\delta }_k‖}_{{\mathrm{\Sigma }}^{-1}}^2={\delta }_k^{\mathrm{\top }}{\mathrm{\Sigma }}^{-1}{\delta }_k=C$ for each client, we ensure a controlled covariate shift in the mean. The covariance Σ remains the same, and the degree of shift is comparable across all clients. This construction thus provides a systematic way to generate imbalanced covariate shifts, enabling direct measurement and comparison of fairness or performance in federated learning experiments.

D. Implementation Details

D.1. Common Hyperparameters of All Algorithms Used in Simulation

Number of clients (NUM_CLIENTS) is set to 10. This indicates how many total clients are simulated or trained independently in the scenario. Global rounds (NUM_GLOBAL_ROUNDS) is

Algorithm 2.

Covariate Shift Data Generation

Require: Global labeled dataset ${𝒟}_g={\left\{\left(x_i,y_i\right)\right\}}_{i=1}^{\left|{𝒟}_g\right|}$; constant $C>0$; number of clients $K$; desired sample counts ${\left\{n_k\right\}}_{k=1}^K$ (such that ${\sum }_{k=1}^K{n}_k=⌊\left|{𝒟}_g\right|/2⌋$); feature dimension $d$.

Ensure: Covariate-shifted client datasets ${\left\{{𝒟}_k\right\}}_{k=1}^K$, where each ${𝒟}_k$ retains the original labels but has shifted features (via importance sampling with respect to a Gaussian distribution with mean ${\mu }_k$ and covariance $\mathrm{\Sigma }$).

1: Vectorize Features: Convert each feature $x_i$ in ${𝒟}_g$ into a vector in ${ℝ}^d$. Let X be the resulting $\left|{𝒟}_g\right|\times d$ matrix.

2: Estimate Baseline Gaussian: $$\mu =\frac{1}{\left|{𝒟}_g\right|}{\sum }_{i=1}^{\left|{𝒟}_g\right|}{x}_i,\mathrm{\Sigma }=\mathrm{C}\mathrm{o}\mathrm{v}(\mathbf{X}).$$

3: Half-sampling setup:

4: Let $M=⌊\left|{𝒟}_g\right|/2⌋$. We only use half of the global dataset for experiments. Hence, set ${\sum }_{k=1}^K{n}_k=M$. Adjust $\left\{n_k\right\}$ if needed.

5: for $k=1,\dots ,K$ do

6:   Compute mean shift ${\delta }_k$:

Sample a vector ${\delta }_k\in {ℝ}^d$ such that ${\delta }_k^{\mathrm{\top }}{\mathrm{\Sigma }}^{-1}{\delta }_k=C$.

7:   ${\mu }_k\leftarrow \mu +{\delta }_k▹$ Shifted mean for client $k$.

8:   Construct Shifted Dataset for client $k$:

(1) Compute importance weights

For each $\left(x_i,y_i\right)$ in ${𝒟}_g$, compute: $$w_i=\mathrm{e}\mathrm{x}\mathrm{p}\left(-\frac{1}{2}{\left(x_i-{\mu }_k\right)}^{\mathrm{\top }}{\mathrm{\Sigma }}^{-1}\left(x_i-{\mu }_k\right)\right).$$

Then normalize: $${\tilde{w}}_i=\frac{w_i}{{\sum }_{j=1}^{\left|{𝒟}_g\right|}{w}_j}.$$

(2) Sample from ${𝒟}_g$ by ${\tilde{w}}_i$

Sample $n_k$ pairs $\left(x_i,y_i\right)$ from ${𝒟}_g$ with replacement according to probabilities ${\tilde{w}}_i$.

Store these sampled pairs in ${𝒟}_k$.

9: end for

10: return ${\left\{{𝒟}_k\right\}}_{k=1}^K$.

commonly set to 20. This is the total number of federated communication rounds. Local epochs (LOCAL_EPOCHS) is often set to 1. It denotes how many epochs each client trains on its local data per global round. Batch size (BATCH_SIZE) is often set to 32. Learning rates (learning_rates) are set to 0.001 for FashionMNIST and 0.005 for CIFAR10. These control the local step size for optimizing SGD.

D.2. Algorithm-Specific Hyperparameters

Table 5 details each algorithm’s unique or particularly important hyperparameters, along with their default values (or ranges) and a brief explanation.

Figure 4:

Results of hyperparameter tuning on the EHR dataset.

Table 5:

Algorithm-Specific Hyperparameters (Condensed)

Algorithm	Special Hyperparameters	Value
CFFL	1) THETA_U (grad upload) 2) CLIP_NORM (clip thr.) 3) C_TH (rep. thr.) 4) ALPHA (rep. update)	(1) 0.5 (2) 5.0 (3) 0.05 (4) 1.0
CGSV	1) ALPHA_R (mov. avg) 2) BETA (sim. scaling) 3) SPARSITY 4) ALTRUISM	(1) 0.9 (2) 2.0 (3) True (4) 1.0
FedAVE	1) UPLOAD_FRAC 2) DOWNLOAD_FRAC_BASE 3) ALPHA / BETA	(1) 0.5 (2) 0.3 (3) 0.9 / 1.0
FedAvg	—	—
FedDC	1) ALPHA (penalty) 2) drift_vars	(1) 1.0 (2) init=0
FedMPR	PRUNE_PERCENT	0.1
FedProx	MU (prox. coeff.)	0.01
FedSAC	1) BETA (c_i mapping) 2) mid_round	(1) 2.0 (2) 15
SCAFFOLD	1) ${\eta }_g$ (global LR) 2) ${\eta }_l$ (local LR) 3) K (local steps) 4) c_global, c_local	(1) 0.005 (2) 0.1 (3) 1 (4) init=0
FedAKD	1) Distill $\alpha$ 2) Distill $\beta$ 3) Temp T	(1) 1.0 (2) 1.0 (3) 1.0

E. The EHR Dataset

We begin with a real-world healthcare dataset composed of 17 tables. To simplify preprocessing and retain the most relevant information, we keep only: Patient Demographic Table, Diagnosis Table, Procedure Table, Medication Drug Table, Lab Result Table, and Vital Sign Table.

Data Processing.

We merge these tables by patient_ID, ensuring each patient record contains both static features (e.g., zipcode, sex) and a series of events (medical codes plus numerical attributes, and the patient’s age at each event). Because multiple medical coding

Table 6:

State Data Statistics

State	Total	Pos.	Neg.	State	Total	Pos.	Neg.
AK	558	196	362	MT	636	221	415
AL	3,410	1,292	2,118	NC	7,263	2,222	5,041
AR	2,341	842	1,499	ND	605	179	426
AZ	5,521	2,347	3,174	NE	1,429	424	1,005
CA	20,040	7,116	12,924	NH	900	344	556
CO	4,164	1,362	2,802	NJ	7,347	3,762	3,585
CT	2,595	1,180	1,415	NM	1,203	416	787
DE	805	342	463	NV	1,994	792	1,202
FL	18,898	8,158	10,740	NY	18,268	8,134	10,134
GA	7,901	2,461	5,440	OH	11,634	4,339	7,295
HI	1,060	433	627	OK	2,575	865	1,710
IA	2,951	1,093	1,858	OR	3,637	1,322	2,315
ID	1,021	361	660	PA	11,817	4,658	7,159
IL	8,932	3,497	5,435	RI	423	158	265
IN	4,388	1,523	2,865	SC	3,386	1,194	2,192
KS	1,857	524	1,333	SD	651	266	385
KY	4,110	1,396	2,714	TN	5,671	2,038	3,633
LA	3,316	1,150	2,166	TX	15,785	5,152	1,0633
MA	3,492	1,453	2,039	UT	2,117	495	1,622
MD	5,145	1,959	3,186	VA	6,057	1,924	4,133
ME	1,183	473	710	VT	445	150	295
MI	9,744	3,629	6,115	WA	6,247	2,009	4,238
MN	3,504	1,371	2,133	WI	2,893	1,281	1,612
MO	4,341	1,572	2,769	WV	1,712	617	1,095
MS	2,246	710	1,536	WY	358	113	245

systems are used across the U.S. healthcare spectrum, we unify these codes into a single standardized terminology.

Next, a board-certified medical oncologist provides a set of pancreatic cancer diagnostic codes, which are used to extract data labels. In ICD-10, the codes for malignant neoplasm of the pancreas include:

• C25 (Malignant neoplasm of pancreas, general),

• C25.0 (Malignant neoplasm of head of pancreas),

• C25.1 (Malignant neoplasm of body of pancreas),

• C25.2 (Malignant neoplasm of tail of pancreas),

• C25.3 (Malignant neoplasm of pancreatic duct),

• C25.4 (Malignant neoplasm of endocrine pancreas),

• C25.7 (Malignant neoplasm of other specified parts of pancreas),

• C25.8 (Malignant neoplasm of overlapping lesions of pancreas),

• C25.9 (Malignant neoplasm of pancreas, unspecified).

In ICD-9, the corresponding codes are:

• 157 (Malignant neoplasm of pancreas, general),

• 157.0 (Malignant neoplasm of head of pancreas),

• 157.1 (Malignant neoplasm of body of pancreas),

• 157.2 (Malignant neoplasm of tail of pancreas),

• 157.3 (Malignant neoplasm of pancreatic duct),

• 157.4 (Malignant neoplasm of islets of Langerhans),

• 157.8 (Malignant neoplasm of other specified sites of pancreas),

• 157.9 (Malignant neoplasm of pancreas, unspecified).

For each patient, if any of these codes appear in the longitudinal record, we set the label y = 1; furthermore, we remove any events that occur at or after the time of pancreatic cancer diagnosis to prevent data leakage. If none of the pancreatic cancer codes appear for a patient, we set that patient’s label to y = 0.

After these preprocessing steps, we obtain 265,085 de-identified patient samples spanning 50 U.S. states.⁵ We randomly sample 10% (26,509) of this dataset as a global validation set. The remaining 90% of the data is returned to each state, where each state splits its local data into a local training set and a local test set (e.g. 7:2). The binary label y indicates whether the patient eventually develops pancreatic cancer. Table 6 summarizes, for each of the 50 states, the number of local samples, how many are used for training/testing, and the distribution of positive/negative labels. Globally, we observe 99,604 positive and 165,481 negative samples from local total data(train and test).

Model and Training Setup.

To handle this longitudinal EHR data, we embed the event codes, concatenate them with numerical attributes (e.g. age, value), and feed the sequence into a two-layer bidirectional GRU, which captures both forward and backward dependencies. We then apply an attention mechanism over the GRU outputs, computing importance weights for each time step and aggregating them into a context vector. This vector is concatenated with the embedded static features (sex, postal_code), and a linear classifier predicts whether the patient will develop pancreatic cancer.

We adopt a focal loss [12] to mitigate label imbalance and train with an Adam optimizer (default PyTorch settings). The mini-batch size is set to B = 64, GRU hidden dimension to 256, dropout probability to 0.3, and learning rate to lr = 0.00001. We embed sex and postal_code into vectors of dimensions 8 and 16, respectively, zero-padding time-series inputs (with sequence lengths tracked via pack_padded_sequence). Specifically, the two-layer GRU produces a 2 × 256-dim vector per time step, mapped by the attention layer to a 256-dim energy vector and finally to scalar scores for softmax weighting. Ultimately, we train and validate our model in this federated setting, treating each state as one client (K = 50).