Personalized Federated Learning for Medical Segmentation using Hypernetworks

Hilit Segev, Gal Chechik

Keywords: HyperNetworks, Personalized Federated Learning (PFL), Image segmentation, Medical Imaging

TL;DR: We develop a personalized federated segmentation approach for medical images based on hypernetworks.

Abstract: In federated learning (FL), several clients jointly train a shared model without sharing their data, maintaining data privacy and reducing communication costs. In personalized federated learning (PFL), each client has their own model, and models are trained jointly. Hypernetworks have been shown to be useful for PFL in classification problems, but it is still not clear how to apply them to problems like segmentation. There, models are very large, and it is not known what parts of models should be personalized, and what parts should be shared across clients. Here, we explore HNs for PFL for solving a problem of image segmentation in the context of medical imaging diagnosis. Using MRI scans for prostate segmentation, we demonstrate that using a hypernetwork to personalize a single convolution layer and the batch-norm layer outperforms local and FL baselines.