[Paper Review] Real World Federated Learning with a Knowledge Distilled Transformer for Cardiac CT Imaging

Malte Tölle, Philipp Garthe+23 more|arXiv (Cornell University)|Jul 10, 2024

Advanced X-ray and CT Imaging4 citations

TL;DR

This paper proposes a two-stage semi-supervised federated learning framework that distills knowledge from task-specific CNNs into a single vision transformer for cardiac CT imaging, enabling effective use of partially labeled data across eight hospitals. The method improves model accuracy and generalizability while training a unified transformer to predict multiple cardiac structures, with open-sourced code and weights serving as a foundation model for future research.

ABSTRACT

Federated learning is a renowned technique for utilizing decentralized data while preserving privacy. However, real-world applications often face challenges like partially labeled datasets, where only a few locations have certain expert annotations, leaving large portions of unlabeled data unused. Leveraging these could enhance transformer architectures ability in regimes with small and diversely annotated sets. We conduct the largest federated cardiac CT analysis to date (n=8,104) in a real-world setting across eight hospitals. Our two-step semi-supervised strategy distills knowledge from task-specific CNNs into a transformer. First, CNNs predict on unlabeled data per label type and then the transformer learns from these predictions with label-specific heads. This improves predictive accuracy and enables simultaneous learning of all partial labels across the federation, and outperforms UNet-based models in generalizability on downstream tasks. Code and model weights are made openly available for leveraging future cardiac CT analysis.

Motivation & Objective

• To address the challenge of training robust deep learning models on partially labeled cardiac CT datasets from multiple hospitals due to privacy constraints and label heterogeneity.
• To leverage large amounts of unlabeled data typically unused in clinical settings to improve model performance.
• To develop a federated learning framework that enables simultaneous learning of multiple cardiac structures (e.g., calcification, coronary arteries) across institutions with varying label coverage.
• To create a generalizable, open-source foundation model for cardiac CT imaging that outperforms standard CNNs in downstream tasks.
• To ensure privacy compliance by training models locally and aggregating only model weights, avoiding centralized data storage.

Proposed method

• A two-stage semi-supervised learning strategy is employed: first, task-specific CNNs (for landmark detection and calcification segmentation) are trained on labeled data at each client site.
• Knowledge distillation is applied to transfer knowledge from these CNNs into a single vision transformer model using both labeled and abundant unlabeled data across all clients.
• Federated learning is used to train the transformer model across eight hospital sites, with model weights aggregated centrally after local training while preserving data privacy.
• The final transformer model is fine-tuned on a downstream task—coronary artery segmentation—by only training the last layer, demonstrating improved feature representation.
• The FedBioMed framework ensures secure communication with TLS encryption, client authentication, and IP whitelisting, enabling a unidirectional, secure FL pipeline.
• Custom Docker orchestration and PyTorch-based pipelines with MONAI and nnU-Net are used for preprocessing, training, and validation across sites.

Experimental results

Research questions

• RQ1 Can a single vision transformer model effectively learn multiple cardiac structures (e.g., calcification, coronary arteries) from partially labeled data across multiple hospitals?
• RQ2 Does knowledge distillation from task-specific CNNs into a transformer improve performance on downstream segmentation tasks in a federated setting?
• RQ3 To what extent can unlabeled data from multiple hospitals be leveraged to improve model generalizability in a privacy-preserving federated learning setup?
• RQ4 How does the performance of a federated transformer model compare to a CNN-based baseline in terms of feature quality and segmentation accuracy?
• RQ5 Can a unified foundation model be trained across heterogeneous label distributions without centralized data collection?

Key findings

• The proposed federated knowledge-distilled transformer model achieved superior performance on cardiac CT imaging tasks compared to a baseline UNet, particularly in downstream coronary artery segmentation.
• By leveraging unlabeled data through knowledge distillation, the model improved predictive accuracy and generalizability despite limited labeled data per client.
• The method enabled simultaneous learning of multiple partial labels (e.g., calcification, landmarks, coronary arteries) within a single transformer architecture across the federation.
• Fine-tuning only the last layer of the transformer for coronary artery segmentation yielded more meaningful features than the CNN baseline, indicating better representation learning.
• The final model was successfully trained across eight hospitals with heterogeneous label coverage, demonstrating scalability and robustness in real-world federated settings.
• The code, model weights, and data processing pipelines are openly available under a FAIR-compliant framework, enabling reuse and further development in cardiac imaging research.