Task-Agnostic Federated Learning

2024-06-25

arXiv.org

AI Summary

Federated learning (FL) for medical imaging faces critical challenges including unknown downstream tasks, scarce labeled data, non-independent and identically distributed (Non-IID) data across clients, and poor cross-task generalization. Method: We propose the first task-agnostic self-supervised FL framework, built upon a Vision Transformer-based consensus feature encoder. It integrates unlabeled collaborative representation learning with heterogeneous task adaptation, enabling zero-prior access to unseen tasks. Contribution/Results: Evaluated on real-world Non-IID medical imaging data, our method achieves 90% of the F1 score using only 5% centralized training data—demonstrating robust out-of-distribution task generalization. This work advances FL toward multi-task foundation model paradigms and establishes a new privacy-preserving paradigm for cross-institutional medical AI collaboration.

Technology Category

Machine Learning: Distributed Machine Learning & Federated LearningComputer Vision: Large Vision ModelsMultiagent Systems: Multiagent Learning

Application Category

User Modeling, Personalization and Recommendation: Federated recommendation systems and personalizationSystems and Infrastructure for Web, Mobile and WoT: Federated Web and WoT systems, including distributed, federated and edge-based data processingSearch and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

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.

Problem

• Federated Learning
• Medical Image Analysis
• Privacy Preservation

Innovation

• Adaptive Federated Learning
• Visual Transformer
• Multi-task Capability