Task-Agnostic Federated Learning

TL;DR

The paper tackles the privacy-preserving challenge of learning from multi-institution medical imaging data without exposing task information or raw data. It introduces a task-agnostic SSL-FL framework that pretrains a global Vision Transformer encoder using masked image modeling on decentralized unlabeled data, then fine-tunes downstream tasks with lightweight adapters, enabling transfer to unseen tasks and out-of-network data. Across six fundus datasets, the approach achieves competitive performance with limited centralized labeling, with robust benefits for data-scarce clients and improved generalization to new tasks, though gains over centralized SSL vary by task and data balance. The work demonstrates the potential of federated learning as a multi-task foundation modeling paradigm in privacy-constrained medical imaging, highlighting the importance of data balancing and task-agnostic pretraining for real-world deployment.

Abstract

In the realm of medical imaging, leveraging large-scale datasets from various institutions is crucial for developing precise deep learning models, yet privacy concerns frequently impede data sharing. federated learning (FL) emerges as a prominent solution for preserving privacy while facilitating collaborative learning. However, its application in real-world scenarios faces several obstacles, such as task & data heterogeneity, label scarcity, non-identically distributed (non-IID) data, computational vaiation, etc. In real-world, medical institutions may not want to disclose their tasks to FL server and generalization challenge of out-of-network institutions with un-seen task want to join the on-going federated system. This study address task-agnostic and generalization problem on un-seen tasks by adapting self-supervised FL framework. Utilizing Vision Transformer (ViT) as consensus feature encoder for self-supervised pre-training, no initial labels required, the framework enabling effective representation learning across diverse datasets and tasks. Our extensive evaluations, using various real-world non-IID medical imaging datasets, validate our approach's efficacy, retaining 90\% of F1 accuracy with only 5\% of the training data typically required for centralized approaches and exhibiting superior adaptability to out-of-distribution task. The result indicate that federated learning architecture can be a potential approach toward multi-task foundation modeling.

Figures (3)

• Figure 1: Federated Learning system where clients' tasks are anonymous between each others and hidden from the server. $E_i$ is the $i_{th}$ feature encoder and $D_i$ is the $i_{th}$ decoder.
• Figure 2: Data distribution of two split settings. In (a) clients have different tasks and different datasets (modality) while clients in (b) share the same amount of data from each datasets and do the same tasks for such portion of dataset.
• Figure 3: Example of segmentation results given by SSL-FL framework and local supervision.