Technical Summary: FedeKD

Problem Statement

Federated Learning (FL) enables collaborative model training across decentralized clients without sharing raw data. However, real-world FL deployments face significant challenges due to data heterogeneity (non-IID distributions) and model asymmetry. These conditions often lead to negative transfer, where knowledge shared between clients degrades local performance rather than improving it.

Existing Federated Knowledge Distillation (FKD) methods attempt to mitigate this by exchanging model outputs (logits or representations) instead of parameters. However, current approaches suffer from two primary limitations:

1. Uniform Reliability Assumption: They typically assume that transferred knowledge is uniformly reliable across all samples and clients, ignoring the variability in knowledge quality caused by distribution shifts or architectural differences.
2. Dependency on Public Data: Many methods rely on additional public datasets or shared data representations to stabilize training, which is often unavailable or raises privacy concerns in practical scenarios.

Consequently, blindly transferring knowledge from a "teacher" (global proxy) to a "student" (local private model) can introduce misleading signals, causing severe performance degradation compared to local training.

Methodology: FedeKD

The paper proposes FedeKD (Energy-Based Gating for Robust Federated Knowledge Distillation), a framework designed to perform sample-wise trust estimation during knowledge transfer without requiring public data.

Architecture

FedeKD employs an asymmetric private-proxy architecture at each client:

• Private Model ($f_k$): A high-capacity network responsible for the primary local learning task.
• Proxy Model ($g_k$): A lightweight network used exclusively for cross-client knowledge exchange. All proxy models share a common architecture.

Training Process

The framework operates in sequential stages per communication round:

1. Forward Proxy Distillation: Each client freezes its private model and trains the local proxy to mimic the private model's predictions on local data.
2. Proxy Aggregation: Local proxies are uploaded to the server, aggregated into a global proxy, and broadcast back to all clients.
3. Energy-Gated Backward Distillation: The global proxy guides updates to the local private model. Crucially, this step utilizes a reliability-aware gating mechanism to determine how much trust to place in the proxy's guidance for each specific sample.

The Energy-Based Gating Mechanism

The core innovation is the conversion of private-proxy disagreement into sample-wise trust weights.

• Energy Score ($E(x)$): Measures the disagreement between the private model and the global proxy.
  • For Classification: Uses an entropy-normalized symmetric Kullback-Leibler (KL) divergence. This normalizes the disagreement by the total predictive entropy, ensuring that high-confidence contradictions are penalized more heavily than diffuse, high-entropy disagreements.
  • For Regression: Uses the squared error between predictions.
• Trust Weight ($w_i$): The energy scores are batch-normalized and passed through a logistic gate to produce continuous trust weights:
  $$w_i = \rho(-\beta \tilde{E}_i)$$
  where $\tilde{E}_i$ is the batch-normalized energy and $\beta$ controls the sharpness.
• Update Rule: The private model is updated using a weighted combination of the supervised loss and the distillation loss. Samples with high disagreement (high energy) receive lower trust weights, effectively down-weighting unreliable knowledge, while samples with low disagreement receive higher weights.

The gating mechanism is batch-relative, meaning trust is assigned based on how a sample's disagreement compares to others in the same minibatch, rather than a fixed global threshold. The weights are detached during backpropagation to ensure they act as reliability scalars without introducing new optimization paths.

Key Contributions

1. Reliability-Aware Backward Distillation: The paper introduces a novel objective that shifts the focus from how to exchange knowledge to how much each transferred signal should be trusted. This is achieved without public data by leveraging task-specific private-proxy disagreement.
2. Asymmetric Private-Proxy Framework: The design separates local learning (high-capacity private model) from communication (lightweight proxy), reducing the exposure of sensitive parameters while enabling robust knowledge exchange.
3. Energy-Based Gating: A continuous, differentiable gating mechanism that maps disagreement into trust weights. It handles both classification (via entropy-calibrated KL divergence) and regression (via squared error) within a unified framework.
4. Theoretical Properties: The authors establish that the gating mechanism preserves the direction of distillation gradients (preventing reversal of the learning signal) and provides bounded, non-degenerate trust weights.

Experimental Results

FedeKD was evaluated on six real-world datasets (four classification: FashionMNIST, CIFAR-10, OCTMNIST, OrganAMNIST; two regression: RetinaMNIST, Diabetic Retinopathy) under varying levels of heterogeneity (controlled by Dirichlet concentration parameter $\alpha$).

• Negative Transfer Mitigation: FedeKD consistently achieved the lowest average-case ($\Delta_{avg}$) and worst-case ($\Delta_{worst}$) negative transfer across all datasets and heterogeneity levels. Under severe heterogeneity ($\alpha=0.1$), it significantly outperformed baselines like FedAvg, FedProx, FedDyn, and the recent FedType.
• Robustness: The method demonstrated superior robustness in the lower tail of the client performance distribution (P10 $\Delta$ for classification) and the upper tail of the error distribution (P90 $\Delta$ for regression), indicating it protects the most vulnerable clients and mitigates extreme errors.
• Predictive Performance: Despite prioritizing robustness, FedeKD maintained competitive or superior mean client accuracy and RMSE compared to baselines, proving that reducing negative transfer does not come at the cost of overall predictive performance.
• Ablation Studies:
  • Removing the gating mechanism ("No Gating") resulted in significantly higher negative transfer, confirming the necessity of reliability-aware weighting.
  • Alternative energy functions (entropy, margin, feature-distance) were less consistent than the proposed KL-based energy.
  • The gating sharpness parameter $\beta$ and distillation weight $\lambda_{kd}$ showed optimal performance in moderate ranges, balancing the trade-off between utilizing shared knowledge and preserving local learning.

Significance and Claims

The paper claims that FedeKD addresses a fundamental gap in heterogeneous FL: the lack of mechanisms to dynamically assess the reliability of transferred knowledge. By treating knowledge reliability as a first-class object, FedeKD mitigates negative transfer without the privacy risks or data requirements of existing methods.

The authors position FedeKD as a step toward more robust and scalable federated learning, particularly for privacy-sensitive domains like healthcare where data heterogeneity is prevalent and the cost of negative transfer (e.g., misdiagnosis) is high. The work emphasizes that adaptive, sample-wise trust estimation is essential for effective collaboration in non-IID environments, moving beyond the assumption that all transferred signals are equally beneficial.

The paper acknowledges limitations, noting that experiments were conducted on public benchmarks with simulated partitions and that the regression energy measures functional disagreement rather than calibrated uncertainty. Future work is suggested to focus on real-world multi-site deployments, fairness-aware evaluation, and adaptive trust estimation.