⚡ electrical engineering

FedKPer: Tackling Generalization and Personalization in Medical Federated Learning via Knowledge Personalization

The paper introduces FedKPer, a medical federated learning framework that simultaneously enhances generalization and personalization while mitigating catastrophic forgetting by integrating knowledge personalization into local training and employing a selective, label-diverse global aggregation scheme.

Original authors: Zoe Fowler, Ghassan AlRegib

Published 2026-05-04

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Zoe Fowler, Ghassan AlRegib

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine a group of doctors from different hospitals trying to build a single, super-smart AI to help diagnose diseases. They can't share their patient records because of privacy laws, so they use a method called Federated Learning. In this setup, each hospital trains a piece of the AI on its own local data, and then they send just the "brain updates" (not the data) to a central server to combine them into one global model.

The paper, FedKPer, tackles two big problems that happen when these hospitals try to work together:

The "One-Size-Fits-None" Problem (Generalization vs. Personalization):
- The Issue: If the central server just averages everyone's updates, the final AI becomes a "jack of all trades, master of none." It might be okay at everything but great at nothing. It fails to adapt to the specific quirks of Hospital A (which sees mostly elderly patients) while also failing to generalize well to Hospital B (which sees mostly children).
- The Analogy: Imagine a teacher trying to teach a class where some students speak French, some speak Mandarin, and some speak Spanish. If the teacher forces everyone to learn only English, the French speakers struggle, and the Mandarin speakers get confused. If the teacher lets everyone speak only their native language, they never learn to communicate with each other. The goal is to find a balance where they can all understand each other and speak their own language fluently.
The "Forgetful Student" Problem (Forgetting):
- The Issue: As the AI learns new things from a new hospital, it often "forgets" what it learned from the previous ones. It's like a student who studies for a history test, then studies for a math test, and suddenly can't remember the history facts. In medical terms, this means the AI might get worse at diagnosing a disease it used to be good at, just because it saw a new type of patient.

How FedKPer Solves This

The authors propose a new system called FedKPer that acts like a smart, flexible coach for these doctors. It uses two main tricks:

1. The "Trust-But-Verify" Teacher (Knowledge Personalization)

In standard methods, the central server acts like a strict teacher, telling every local hospital exactly what to think. But if the central server is wrong about a specific hospital's patients (because that hospital sees rare diseases), following the teacher blindly hurts the local model.

FedKPer's Fix: Instead of blindly copying the central teacher, each local hospital checks: "Is the teacher actually right about my specific patients?"
The Analogy: Imagine a local chef (the hospital) learning from a famous celebrity chef (the global model). If the celebrity chef suggests a recipe that tastes terrible with the local chef's specific ingredients, FedKPer tells the local chef: "Ignore that part of the recipe. Stick to what works for your ingredients, but keep the good techniques."
The Result: Each hospital keeps its own unique style (Personalization) but still learns useful general techniques (Generalization).

2. The "Fair Voting" System (Smart Aggregation)

When the central server combines the updates from all hospitals, it usually just averages them based on how many patients each hospital has. This is unfair if one hospital only sees one type of disease; their update might dominate and skew the whole group.

FedKPer's Fix: The server looks at two things before voting:
1. Reliability: How well did this hospital's model perform on its own data?
2. Diversity: Does this hospital see a wide variety of diseases, or just one?
The Analogy: Imagine a committee voting on a new policy. Instead of giving everyone one vote, FedKPer gives more voting power to members who are both expert (reliable) and representative of the whole group (diverse). If a member only knows about one tiny topic, their vote counts less so they don't drag the whole committee off-topic.
The Result: The final global model is balanced and doesn't get "hijacked" by hospitals with weird or limited data.

The Results

The paper tested this on medical image datasets (like blood cells and organs). They found that FedKPer:

Did better overall: It improved the balance between being good at everything (generalization) and being good at specific local tasks (personalization).
Forgot less: It didn't lose old knowledge as quickly as other methods.
Was efficient: It achieved these results without taking significantly more time or computing power.

In short, FedKPer is a smarter way for hospitals to collaborate. It ensures the AI learns from everyone without forcing everyone to be the same, and it makes sure the AI remembers what it learned yesterday while learning something new today.

1. Problem Statement

Federated Learning (FL) is a promising paradigm for medical AI, allowing institutions to collaborate without sharing raw patient data. However, medical FL faces a critical challenge: statistical heterogeneity (non-IID data). Patient populations vary significantly across hospitals, leading to two conflicting issues:

Generalization vs. Personalization: Standard FL algorithms (like FedAvg) prioritize a global model that generalizes well across all sites but often fail to adapt to local data distributions, reducing local accuracy. Conversely, Personalized FL (pFL) improves local accuracy but often sacrifices the ability of the global model to generalize to unseen populations.
Catastrophic Forgetting: Under heterogeneity and partial client participation, model updates can overwrite previously learned patterns. This leads to performance degradation on earlier clients or classes, reducing the global model's consistency and reliability.
Ineffective Knowledge Distillation: Existing methods often use the global model as a "teacher" for local clients via knowledge distillation. However, in highly heterogeneous settings, the global model may be a poor teacher for specific local clients, leading to "blind distillation" that misguides local training and harms both personalization and generalization.

2. Methodology: FedKPer

The authors propose FedKPer, a framework that simultaneously addresses generalization and personalization through two distinct, coupled components: Adaptive Knowledge Personalization (client-side) and Reliability-Diversity Aggregation (server-side).

A. Knowledge Personalization (Client-Side)

Instead of blindly following the global model, FedKPer introduces an adaptive reliability weight into the local training loss function.

Loss Function: The local loss combines standard Cross-Entropy ( $\ell_{CE}$ ) with a Knowledge Distillation term ( $\ell_{KD}$ ).
$\ell = \ell_{CE}(q_k, y) + \lambda \ell_{KD}(q_g, q_k)$
Adaptive Weight ( $\lambda$ ): The weight $\lambda$ $λ$ is dynamically calculated based on the global model's performance on the local client's data:
$\lambda = \frac{1}{\ell_{CE}(q_g, y)}$
- If the global model is inaccurate on the local data (high loss), $\lambda$ decreases, suppressing the global model's influence.
- If the global model is accurate, $\lambda$ increases, encouraging alignment.
- A cap (e.g., $\lambda \leq 10$ ) is applied to prevent numerical instability.
Goal: This allows clients to retain globally transferable structures while flexibly adapting to their unique label distributions, preventing the "harmful imitation" of a distorted global hypothesis.

B. Generalization via Aggregation (Server-Side)

Standard FedAvg aggregates updates based on dataset size ( $n_k$ ). FedKPer proposes a novel aggregation strategy that prioritizes updates that are both reliable and label-diverse.

Reliability Score: Measured by the local model's accuracy ( $A_k$ ) on its own training set.
Diversity Score: Measured by the entropy of the client's local label histogram ( $d_k$ ). High entropy implies a diverse set of classes; low entropy implies a skewed distribution.
Aggregation Weight ( $\bar{p}_k$ ):
$\bar{p}_k = A_k \cdot (\epsilon + d_k)$
The final aggregation weight is the normalized version of $\bar{p}_k$ .
Goal: This prevents clients with highly skewed data (e.g., only seeing one class) from dominating the global update, ensuring the global model reflects a broader range of label groups and improves generalization.

C. Evaluation Metrics

The paper introduces metrics rooted in forgetting to better evaluate FL algorithms:

Consistency (C): Measures the stability of accuracy after reaching a peak. It calculates the Average Inter-Peak Forgetting Rate (AIPFR), quantifying how much accuracy drops between record-high peaks.
Backward Transfer (BwT): Measures the performance change on previously sampled clients as the global model evolves, indicating how well the model retains prior knowledge.

3. Key Contributions

FedKPer Framework: A unified approach that couples client-side adaptive personalization with server-side diversity-aware aggregation to balance the generalization-personalization trade-off.
Adaptive Distillation: A mechanism that treats the global model as a conditional teacher, only trusting its guidance when it is empirically correct for the specific local client.
Novel Aggregation Strategy: A weighting scheme that favors clients with diverse label distributions and high local reliability, mitigating the bias caused by class imbalance in FL.
Forgetting-Oriented Metrics: The introduction of Consistency and BwT metrics to rigorously evaluate the stability and retention capabilities of FL algorithms, moving beyond simple accuracy reporting.

4. Experimental Results

The authors evaluated FedKPer on three medical imaging datasets (BloodMNIST, OrganCMNIST, OrganSMNIST) with high heterogeneity (Dirichlet distribution $\alpha=0.1$ ).

Performance: FedKPer achieved the highest global accuracy across all three datasets, outperforming the next best method by significant margins (e.g., +9.8% on BloodMNIST).
Balance: It achieved the best Global-Local Balance score, demonstrating superior performance in both generalization and personalization simultaneously.
Retention: FedKPer showed significantly improved Consistency and BwT (less negative values), indicating reduced catastrophic forgetting compared to baselines like FedAvg, FedProx, and FedNTD.
Worst-Client Performance: It achieved the best or second-best "worst-client" accuracy, proving it reduces performance gaps between clients and avoids bias toward specific subsets.
Efficiency: In terms of wall-clock time, FedKPer achieved higher accuracy than FedAvg within the same time budget (e.g., +38.8% improvement at 500s on BloodMNIST).
Ablation Study: Results confirmed that the adaptive distillation term primarily boosts local performance, while the aggregation term primarily boosts global performance; combining both yields the optimal trade-off.

5. Significance

This paper addresses a critical gap in medical AI deployment: the inability of current FL methods to simultaneously adapt to local hospital data while maintaining a robust, generalizable global model.

Clinical Relevance: By reducing forgetting and balancing personalization with generalization, FedKPer ensures that medical models remain reliable across diverse patient populations without sacrificing accuracy for specific institutions.
Methodological Advance: The shift from static distillation to adaptive, reliability-weighted distillation offers a new paradigm for handling heterogeneity in distributed learning.
Evaluation Standard: The proposed forgetting-based metrics provide a more holistic view of FL algorithm behavior, encouraging future research to prioritize stability and retention alongside raw accuracy.