Post-Deployment Adaptation with Access to Source Data via Federated Learning and Source-Target Remote Gradient Alignment

TL;DR

This paper presents FedPDA, a novel framework that enables post-deployment adaptation of neural networks by leveraging remote source data through federated learning techniques, improving performance on medical imaging tasks under distribution shifts.

Contribution

It introduces FedPDA, which allows source data to inform model adaptation via remote gradient exchange, and proposes StarAlign, a new gradient alignment method for better domain-specific learning.

Findings

FedPDA improves model adaptation in medical imaging tasks.

StarAlign effectively aligns gradients for target-specific learning.

The method outperforms previous adaptation techniques.

Abstract

Deployment of Deep Neural Networks in medical imaging is hindered by distribution shift between training data and data processed after deployment, causing performance degradation. Post-Deployment Adaptation (PDA) addresses this by tailoring a pre-trained, deployed model to the target data distribution using limited labelled or entirely unlabelled target data, while assuming no access to source training data as they cannot be deployed with the model due to privacy concerns and their large size. This makes reliable adaptation challenging due to limited learning signal. This paper challenges this assumption and introduces FedPDA, a novel adaptation framework that brings the utility of learning from remote data from Federated Learning into PDA. FedPDA enables a deployed model to obtain information from source data via remote gradient exchange, while aiming to optimize the model specifically for the target domain. Tailored for FedPDA, we introduce a novel optimization method StarAlign (Source-Target Remote Gradient Alignment) that aligns gradients between source-target domain pairs by maximizing their inner product, to facilitate learning a target-specific model. We demonstrate the method's effectiveness using multi-center databases for the tasks of cancer metastases detection and skin lesion classification, where our method compares favourably to previous work. Code is available at: https://github.com/FelixWag/StarAlign