Federated Prototype Learning

Instead of transmitting full model weights or gradients, clients share only class prototypes—compact vectors representing the central tendency of each class in the latent space. A prototype is typically computed as the mean embedding of all samples belonging to a class on that client. This reduces communication overhead from megabytes to kilobytes per round, making it ideal for bandwidth-constrained hospital networks.

Because prototypes are class-conditional statistics, they are inherently robust to label distribution skew. A hospital specializing in rare diseases can contribute prototypes for those classes without requiring balanced local datasets. The global model aggregates prototypes per class across clients, effectively reconstructing a balanced representation of all classes even when no single client has a complete label distribution.

Local training is guided by a prototype-based regularization loss that pulls local embeddings toward their corresponding global prototypes while pushing them away from other classes. This dual objective—standard cross-entropy plus prototype distance minimization—ensures that local feature extractors learn representations aligned across institutions without sharing raw data or violating privacy constraints.

Unlike weight-averaging methods such as FedAvg, prototype learning decouples model architecture from aggregation. Clients can use different neural network backbones as long as they project inputs into a shared embedding space of the same dimensionality. This flexibility allows resource-constrained clinics to deploy lightweight models while research hospitals use deeper architectures, all participating in the same federated system.

Prototypes are aggregate statistics rather than individual data points, providing an inherent layer of privacy. A single prototype represents the mean of potentially hundreds of patient records, making membership inference attacks significantly harder. When combined with differential privacy by adding calibrated noise to prototypes before transmission, the approach achieves formal privacy guarantees suitable for HIPAA and GDPR compliance.

The prototype framework naturally supports few-shot learning scenarios. When a new client joins with only a handful of labeled examples, its local prototypes can be compared against the global prototype set for nearest-neighbor classification. This enables rapid onboarding of new hospitals without requiring extensive local training, a critical advantage for clinical networks expanding to new sites.

Instead of sharing raw data or full model weights, clients compute and transmit class-representative vectors (prototypes). A prototype is a compact, averaged embedding for each class label present in a client's local dataset. This naturally abstracts away the specifics of individual patient records, providing an inherent layer of privacy by design. The global server aggregates these prototypes to form a shared, discriminative embedding space, which is then redistributed to clients for local inference. This method is highly effective under label distribution skew, where different hospitals have vastly different proportions of disease cases.

Federated Prototype Learning is specifically architected to combat statistical heterogeneity. By sharing abstract class representations rather than gradients, the process is less susceptible to the client drift that plagues traditional Federated Averaging (FedAvg) on non-IID data. The global objective shifts from finding a single consensus weight vector to learning a well-separated embedding space where local data can be accurately classified. This makes it robust against both feature distribution skew and label distribution skew, as the prototypes capture the local geometry of each class without forcing a one-size-fits-all model.

A primary advantage is the drastic reduction in communication overhead. Transmitting a few prototype vectors per class is orders of magnitude smaller than sharing a full deep learning model with millions of parameters. This makes the technique ideal for cross-device federated learning scenarios involving edge devices or hospitals with limited bandwidth. The communication cost scales with the number of classes, not the model size, enabling the use of large, expressive neural networks on the client side without a corresponding increase in network traffic.

Because clients only share the output of their penultimate layer (the prototypes), there is no requirement for all clients to use the same model architecture. A research hospital with powerful GPUs can train a large Vision Transformer, while a rural clinic can train a lightweight MobileNet. As long as they agree on the dimensionality of the prototype embedding space, the global server can aggregate their knowledge. This model-agnostic property is a significant step beyond standard federated learning, allowing institutions to leverage their specific computational resources without constraint.

Federated Prototype Learning shares a philosophical lineage with Federated Knowledge Distillation (FKD). Both methods transmit compact, abstracted knowledge rather than model parameters. In FKD, clients share soft label predictions (logits) on a public dataset. In prototype learning, they share class centroids. Both techniques inherently support heterogeneous model architectures and provide strong privacy guarantees. Prototype learning can be seen as a more data-efficient variant when a suitable public dataset for distillation is unavailable, as it relies solely on local data statistics.