FedeKD: Energy-Based Gating for Robust Federated Knowledge Distillation under Heterogeneous Settings

Abstract: Federated learning (FL) operates in heterogeneous environments, where variations in data distributions and asymmetric model design often result in negative transfer. While federated knowledge distillation (FKD) avoids direct model parameter sharing, existing methods typically rely on public datasets or assume that transferred knowledge is uniformly reliable, which limits their robustness in practice. This paper presents FedeKD, a reliability-aware FKD framework that makes sample-wise trust estimation an explicit component of knowledge transfer, without relying on additional public data. Each client maintains a high-capacity private model for local learning and a lightweight shared proxy model for cross-client knowledge exchange. During training, proxy models are aggregated on the server to form a global proxy, which is then used to guide updates of the private models. At the core of FedeKD is an energy-based gating mechanism that converts task-specific private-proxy disagreement into sample-wise trust weights for backward distillation. This mechanism enables sample-wise weighting of knowledge transfer, where the proxy model contributes more to reliable samples while down-weighting unreliable ones. Extensive experiments on six real-world datasets demonstrate that FedeKD significantly reduces negative transfer under heterogeneous settings while maintaining strong predictive performance.

• The paper introduces a dual-model design that separates private and proxy models to enhance privacy and efficiency.
• It implements an energy-based gating mechanism to scale sample contributions, significantly mitigating negative transfer in heterogeneous environments.
• Empirical results across benchmarks, including CIFAR-10 and various medical imaging datasets, demonstrate improved worst-case performance and robust accuracy.

Energy-Based Gating for Reliable Federated Knowledge Distillation in Heterogeneous Environments

Motivation and Problem Setting

Significant heterogeneity in client data and model architectures is a defining challenge for federated learning (FL), undermining the reliability of transferred knowledge and exacerbating negative transfer effects—where collaborative training leads to degraded local performance relative to pure local learning. Conventional parameter aggregation (FedAvg, FedProx, etc.) assumes uniform trustworthiness of model updates and struggles under distribution and model asymmetry. Existing federated knowledge distillation (FKD) approaches largely overlook sample-level and client-level variability of transferred knowledge reliability, often assuming uniform trust or relying on public data for stabilization, which is not always viable.

The paper "FedeKD: Energy-Based Gating for Robust Federated Knowledge Distillation under Heterogeneous Settings" (2605.05553) presents a robust FKD method tailored to these real-world conditions by explicitly operationalizing sample-wise trust in knowledge transfer. This method eschews reliance on public data and directly addresses the unreliability and potential harm of knowledge transfer under heterogeneity.

FedeKD Framework

FedeKD introduces a dual-model per-client architecture comprising a high-capacity private model (for local learning) and a lightweight proxy model (for federated communication/aggregation). Knowledge transfer is decoupled: full-capacity private models are never shared; only proxy models participate in communication. This separation fundamentally enhances privacy and supports model asymmetry.

Figure 1: The framework of FedeKD, illustrating the split between private and proxy models, with only the latter participating in inter-client communication and aggregation for knowledge distillation.

Global training alternates through the following phases for each communication round:

1. Forward Proxy Distillation: Each client distills its private model's predictions into the proxy model using a task-dependent distillation loss, freezing the private model.
2. Proxy Aggregation: Proxy weights from all clients are transmitted to the server and averaged to form a global proxy model.
3. Energy-Gated Backward Distillation: The global proxy is broadcast to all clients. Each private model incorporates the proxy’s knowledge via a computationally formulated, sample-wise, reliability-aware gating mechanism.

Energy-Based Gating Mechanism

The energy-based gating is the core innovation of FedeKD. For each sample, an energy score is computed to quantify the private–proxy model disagreement:

• Classification tasks: The energy is formulated as a symmetric KL divergence between predicted distributions, normalized by entropy, ensuring scale-invariance and emphasizing sharp, confident contradictions:

$$E(x) = \frac{1}{2} \frac{\mathrm{KL}(p || q) + \mathrm{KL}(q || p)}{H(p) + H(q) + \epsilon}$$

• Regression tasks: The energy is the squared L2 difference:

$E(x) = \frac{1}{2} \|f_k(x) - g(x)\|^2$

Within each minibatch, these energies are batch-normalized and then mapped to trust weights via a logistic (sigmoid) function, controlled by a temperature parameter $\beta$:

$$\tilde{E}_i = \frac{E_i - \mu_B}{s_B + \epsilon_B}, \quad w_i = \rho(-\beta \tilde{E}_i)$$

where $\rho(t) = 1/(1 + e^{-t})$, $\mu_B$ and $s_B$ are the batch mean and standard deviation.

Trust weights scale the contribution of each sample in the backward distillation objective. Critically, gradients do not flow through the gate (stop-gradient), ensuring reliability weights only modulate rather than redirect the distillation process.

This mechanism presents several formal properties:

• Monotonicity: Orderings of energy are preserved in reliability weights.
• Boundedness: Weights are strictly within (0,1), avoiding both hard rejections and overweighting.
• Variational interpretation: The gate constitutes an entropy-regularized soft trust assignment.

Experimental Evaluation

FedeKD is evaluated on six public datasets including both standard benchmarks (FashionMNIST, CIFAR-10) and medical imaging data (OCTMNIST, OrganAMNIST, RetinaMNIST, Diabetic Retinopathy), covering both classification and regression tasks. Heterogeneity is simulated via Dirichlet partitions with varying $\alpha$ and model asymmetry is systematically introduced.

Key metrics:

• Negative transfer ($\Delta$): Difference in performance relative to local-only training, reported at both average and worst-case across clients.
• Robustness metrics: 10th percentile (P10 for classification) and 90th percentile (P90 for regression) transfer, assessing lower and upper tails of client performance.
• Predictive metrics: Client average accuracy/RMSE.

Empirical results show:

• FedeKD achieves minimal or positive average $\Delta$ in both tasks, even under strong heterogeneity ($E(x) = \frac{1}{2} \|f_k(x) - g(x)\|^2$0), significantly outperforming standard parameter aggregation and even strong FKD baselines (FedType) on both average and worst-case negative transfer metrics.
• Particularly, FedeKD is the only framework achieving close-to-zero or negative worst-case $E(x) = \frac{1}{2} \|f_k(x) - g(x)\|^2$1 across all settings, directly contradicting the implicit assumption that adaptive gating is unnecessary in mild heterogeneity.
• Competitive or best-in-class predictive accuracy/RMSE is maintained, indicating no sacrifice of performance when reducing negative transfer.

  Figure 2: Classification accuracy on FashionMNIST and CIFAR-10, and regression RMSE on RetinaMNIST and Diabetic Retinopathy at $E(x) = \frac{1}{2} \|f_k(x) - g(x)\|^2$2.

Ablation studies confirm the necessity of energy-based gating; naive distillation (no gating) suffers pronounced negative transfer. Variants using entropy-based, margin-based, or feature-space energy scores are less effective, both in robustness and variance.

Figure 3: Average $E(x) = \frac{1}{2} \|f_k(x) - g(x)\|^2$3 across varying $E(x) = \frac{1}{2} \|f_k(x) - g(x)\|^2$4 values on CIFAR-10 and OrganAMNIST indicating that moderate gating sharpness ($E(x) = \frac{1}{2} \|f_k(x) - g(x)\|^2$5) optimally balances selectivity and informativeness in gating.

Theoretical Foundations and Analysis

The gating mechanism is analytically grounded. Formal propositions establish monotonicity, boundedness, and the direction-preserving nature of gradients under stop-gradient gating. The variational interpretation frames the mechanism as solving an entropy-regularized assignment, ensuring smooth adaptation without hard thresholds.

Additionally, batch normalization of energy ensures relativity of trust assignment, dynamically adjusting to intra-batch heterogeneity. Stop-gradient design secures that the proxy cannot dominate or reverse the local learning trajectory.

Figure 4: Average $E(x) = \frac{1}{2} \|f_k(x) - g(x)\|^2$6 across $E(x) = \frac{1}{2} \|f_k(x) - g(x)\|^2$7 values for sensitivity analysis, demonstrating robust performance with moderate distillation coefficient selection.

Practical and Theoretical Implications

Privacy and Communication: By restricting aggregation and communication to lightweight proxy models, FedeKD reduces exposure of private parameters and associated data leakage risk. It enables cross-client collaboration without reliance on public data.

Robustness: The batch-wise, energy-based gating framework for trust assignment systematically mitigates negative transfer, particularly in lower-tail and worst-case regimes, without sacrificing average-case accuracy. This demonstrates the value of relative, adaptive trust modeling in FL environments subject to extensive heterogeneity.

Transferability: The generality of the energy-based gating (applicable to both classification and regression, without the need for public datasets) positions the framework as a generic reliability-aware FKD primitive, potentially extensible to broader collaborative or privacy-driven machine learning paradigms.

Limitations and Future Prospects: FedeKD introduces additional computation (dual-model maintenance) and synchronization overhead. Real-world deployment scenarios—especially with unobserved institutional biases and operational fairness constraints—remain open for study. Future work may extend the trust estimation to leverage model calibration, fairness-aware objectives, and communication-efficient distillation regimes.

Conclusion

FedeKD delivers a reliability-aware framework for federated knowledge distillation, operationalizing energy-based sample-wise trust estimation in collaborative learning without public data. It addresses the core challenge of negative transfer under dual heterogeneity—data and model asymmetry—demonstrating strong robustness, theoretical soundness, and adaptability to both classification and regression tasks. The energy-based gating mechanism, grounded in formal analysis and supported by extensive empirical validation, represents an important contribution to robust, privacy-preserving FL in health care and beyond.