⚡ electrical engineering

FedProIn: Mitigating Client Drift for Learnable Prototypes in Federated Medical Imaging

FedProIn is a prototype-based federated learning framework that mitigates client drift in medical imaging by employing feature and prototype divergence losses alongside a normalized influence aggregation strategy, achieving superior performance over existing baselines on heterogeneous datasets.

Original authors: Harsh Kumar, Tarun Kumar Garg, Vaanathi Sundaresan

Published 2026-07-07

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Harsh Kumar, Tarun Kumar Garg, Vaanathi Sundaresan

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 teach a single AI how to recognize skin diseases or blood cells. They can't share their actual patient photos because of privacy laws (like HIPAA). Instead, they use a method called Federated Learning.

Think of this like a group of chefs trying to perfect a single "Master Recipe." Each chef cooks in their own kitchen with their own local ingredients. They send their instructions on how to improve the recipe to a central head chef, who combines them to make a better "Master Recipe" for everyone.

The Problem: "Client Drift"

In the real world, every hospital is different. One might use a specific type of camera, another might see mostly elderly patients, and a third might see mostly children. This is called statistical heterogeneity.

Because of these differences, the local chefs (clients) start to drift away from the Master Recipe.

Feature Drift: One chef starts describing a "red spot" as "bright crimson," while another calls it "deep scarlet." They are looking at the same thing but describing it in totally different ways.
Prototype Drift: The chefs start memorizing their own local quirks. If a chef only sees rare diseases, they might over-focus on those, ignoring common ones.

When the head chef tries to mix these conflicting instructions, the final recipe becomes confused, unstable, and less accurate. This is what the paper calls Client Drift.

The Solution: FedProIn

The authors propose a new system called FedProIn (Federated Learning with Prototype aggregation with Influence). Instead of just averaging the chefs' instructions, they use a smarter approach with three main tricks:

1. Learning "Archetypes" (Learnable Prototypes)

Instead of just averaging the data, the system creates Learnable Prototypes.

The Analogy: Imagine instead of just saying "red spot," the system creates a set of idealized "Archetype Cards" for every disease. There isn't just one card for "Melanoma"; there are several cards representing different types of Melanoma.
These cards are not static; they are "learnable," meaning the AI can actively tweak them to be the perfect representation of a disease, rather than just taking a blurry average of all the photos it has seen.

2. The "Double-Check" System (Mitigating Drift)

To stop the chefs from drifting apart, FedProIn uses two specific "loss functions" (which are just penalty scores for making mistakes):

Feature Divergence Loss (The "Translation" Check): This forces the local chefs to describe things in a way that matches the global "dictionary." If a local chef describes a spot as "crimson," the system gently nudges them to align with the global definition of "red" so everyone is speaking the same language.
Prototype Contrastive Loss (The "Memory" Check): This ensures the local chefs' "Archetype Cards" don't get too weird or biased toward their own kitchen. It pulls their cards closer to the global Master Cards and pushes them away from their own previous, potentially biased versions.

3. The "Smart Vote" (Normalized Influence Aggregation)

This is the most unique part. In standard systems, every chef gets a vote based on how many photos they have. But what if a chef has 1,000 photos of a common disease but zero photos of a rare one? Their vote might drown out the one chef who actually has the rare photos.

FedProIn introduces Normalized Influence Aggregation (NIA).

The Analogy: Instead of counting heads, the system counts who actually used the Archetype Cards.
If a specific "Rare Disease Card" was used frequently by a specific chef to classify their patients, that chef gets a strong vote for that specific card.
If a chef has a card that they never actually used (maybe they didn't have any patients with that condition), their vote for that card is silenced. This prevents "noise" from drowning out important, rare information.

The Results

The authors tested this system on two real medical datasets:

HAM10000: Skin lesion images.
Matek-19: Blood cell images.

They simulated a scenario where the data was messy and uneven (Non-IID), which is how real life works.

The Outcome: FedProIn outperformed all other methods. It was better at keeping the "Master Recipe" accurate even when the local kitchens were very different.
The Numbers: On the skin dataset, it achieved 81.1% accuracy in the messy, real-world scenario (compared to lower scores for other methods). On the blood cell dataset, it hit 95.8%.

Summary

In short, FedProIn is a smarter way for AI models to learn together without sharing private data. It stops them from getting confused by local differences by:

Creating flexible "Archetype Cards" for diseases.
Forcing everyone to speak the same "feature language."
Giving a louder voice to the chefs who actually know the rare diseases, rather than just the ones with the most photos.

The paper concludes that this method makes the global AI model much more robust and accurate in the messy, diverse reality of medical imaging.

Technical Summary: FedProIn – Mitigating Client Drift for Learnable Prototypes in Federated Medical Imaging

Problem Statement

Federated Learning (FL) in medical imaging is significantly hindered by statistical heterogeneity arising from variations in scanner protocols, acquisition methods, and diverse patient populations across different medical centers. This non-IID (non-Independent and Identically Distributed) data distribution induces client drift, where local models diverge from the global model during optimization. Specifically, identical instances are mapped to distinct regions in the feature space by different clients, and local representations become biased toward client-specific characteristics.

Existing approaches face limitations:

Constraint-based regularization (e.g., FedProx, MOON) struggles to align semantic representations in high-dimensional spaces.
Standard prototype-based methods often assume a single global prototype per class (via mean feature vectors), which is inadequate for heterogeneous medical data.
Multi-prototype methods often introduce high computational complexity or rely on unimodal representations that fail to optimize decision boundaries effectively.

Methodology: FedProIn

The authors propose FedProIn (Federated Learning with Prototype aggregation with Influence), a framework that utilizes Learnable Prototypes (LPs) as optimizable model parameters rather than static mean vectors. The global model consists of a feature encoder $f_\theta$ and a set of prototypes $P \in \mathbb{R}^{C \times M \times d}$ , where $M$ is the number of prototypes per class.

1. Decomposition of Client Drift

FedProIn explicitly addresses client drift by decomposing it into two components and applying specific regularization losses:

Feature Drift: Addressed via Feature Divergence Loss (FDL). Clients maintain a frozen copy of the global feature extractor ( $\theta^{t-1}$ ). The local encoder is regularized to minimize the mean square error between local feature embeddings and the reference global features, ensuring semantic alignment.
Prototype Drift: Addressed via Prototype Contrastive Loss (PCL). This loss pulls updated local prototypes closer to the current global prototypes while pushing them away from their previous local states. This prevents prototypes from becoming overly biased toward local data shifts.

The total local objective combines the primary classification loss ( $\ell_{Task}$ ) with these regularization terms:
$\ell_{Total} = \ell_{Task} + \lambda_{FDL}\ell_{FDL} + \lambda_{PCL}\ell_{PCL}$

To handle class imbalance in long-tailed medical datasets, the method employs class-balanced sampling and adaptive augmentation to ensure rare disease subtypes are adequately represented during prototype updates.

2. Normalized Influence Aggregation (NIA)

Standard aggregation (e.g., FedAvg) is insufficient for multiple prototypes because it weights updates solely by dataset size, ignoring prototype utilization. This can amplify noise from underutilized prototypes or suppress specialized representations for rare cases.

FedProIn introduces Normalized Influence Aggregation (NIA):

Influence Matrix ( $I_i$ ): During local training, each client tracks an influence score for every prototype, counting how many times a specific prototype served as the nearest neighbor for local samples.
Aggregation: Global prototypes are aggregated by weighting them proportionally to their usage scores ( $I_{k,c}^i$ ) rather than just the client's dataset size. This down-weights noisy updates from unused prototypes and preserves specialized prototypes relevant to specific client subsets.

Key Contributions

Learnable Prototypes Formulation: Adopting prototypes as optimizable parameters (jointly with the feature extractor) rather than static mean vectors to maximize class separation margins.
Dual Regularization Strategy: Explicitly mitigating client drift by decomposing it into feature and prototype drift, addressed via FDL and PCL respectively.
Normalized Influence Aggregation (NIA): A novel aggregation mechanism that adaptively weights client prototypes based on their local utility, addressing issues of heterogeneous utilization and sparse activation in multi-prototype settings.
Empirical Validation: Comprehensive evaluation on public medical imaging datasets demonstrating robustness against statistical heterogeneity.

Experimental Results

The framework was evaluated on two public long-tailed medical imaging datasets: HAM10000 (dermatoscopic images) and Matek-19 (peripheral blood cell images). Experiments were conducted under both IID and Non-IID (Dirichlet $\alpha=0.5$ ) settings with $K=10$ clients.

Performance Highlights:

HAM10000: FedProIn achieved 83.5% accuracy (IID) and 81.1% accuracy (Non-IID), outperforming all baselines including FedAvg, FedProx, MOON, and other prototype-based methods (FedProto, FedTGP, FedPLVM).
Matek-19: FedProIn achieved 96.2% accuracy (IID) and 95.8% accuracy (Non-IID).
Robustness: Under Non-IID conditions, existing prototype methods suffered significant performance drops (e.g., 6–7% accuracy loss on HAM10000), whereas FedProIn restricted relative losses to merely 2.4% and 0.4% on HAM10000 and Matek-19, respectively.
Ablation Study: Results confirmed that both FDL and PCL are necessary for optimal performance. Removing FDL caused the sharpest drop in performance, indicating feature drift is the dominant failure mode. The full model significantly outperformed single-component variants and the degenerate case (equivalent to FedAvg).

Significance and Claims

The paper claims that FedProIn effectively mitigates client drift in federated medical imaging by leveraging learnable prototypes and influence-aware aggregation. The authors assert that their approach:

Achieves superior generalization compared to existing FL baselines and personalized methods, particularly in highly heterogeneous (Non-IID) environments.
Successfully addresses the limitations of single-prototype assumptions and unweighted aggregation in multi-prototype settings.
Provides a robust solution for medical imaging where data heterogeneity and class imbalance are prevalent.

The authors note that future work will explore adapting FedProIn for personalized FL (removing the global model dependency) and extending the framework to heterogeneous FL settings with diverse model architectures. The code is made available to support reproducibility.