Towards multi-task medical imaging models: exploring federated learning with vision transformers

Abstract

Medical Imaging models have become commonplace for critical diagnostic tasks such as image segmentation, detection, and classification. They have been proven to perform better than humans and have made diagnostic procedures largely hassle-free with minimum human intervention. However, the training of robust diagnostic models is hindered by two major roadblocks. Firstly, training specialized models for each task requires large amounts of data. Secondly, several privacy laws restrict the sharing of medical data, limiting opportunities for collaborative training. To overcome the first challenge, Multi-Task Learning (MTL) is utilized to perform multiple tasks using a single model. However, while traditional Convolutional Neural Network-based MTL models excel at identifying local features, they struggle to contextualize global features. To address the second challenge, Federated Learning (FL) is used to collaboratively train models by periodically sharing model weights with an aggregation server, avoid-ing direct data communication. However, neural networks are permutation invariant, which means that permuting the nodes in any layer of the network does not affect its prediction outcome. This is a problem for traditional FL methods, as averaging the weight tensors of corresponding layers of multiple separately trained models can result in distorted feature maps.