Real World Federated Learning with a Knowledge Distilled Transformer for Cardiac CT Imaging

2024-07-10

AI Summary

Multi-center cardiac CT analysis faces practical challenges including sparse and incomplete annotations across institutions and stringent privacy requirements. Method: We propose the first large-scale (8,104 cases) federated semi-supervised learning framework tailored for clinical deployment. Our approach introduces a novel two-stage federated knowledge distillation scheme: (i) task-specific CNNs generate high-quality pseudo-labels locally; (ii) a Transformer backbone models heterogeneous partial labels across sites, enhanced by a multi-head label adapter and a CNN-Transformer collaborative architecture—enabling, for the first time, federated joint modeling of partial labels. Results: Evaluated on real-world data from eight hospitals, our method significantly outperforms the UNet baseline with substantially improved generalizability. The code and pretrained models are publicly released to advance privacy-preserving, annotation-efficient federated cardiac CT analysis.

Technology Category

• Machine Learning: Distributed Machine Learning & Federated Learning
• Computer Vision: Multi-modal Vision
• Constraint Satisfaction and Optimization: Distributed CSP/Optimization

Application Category

• User Modeling, Personalization and Recommendation: Federated recommendation systems and personalization
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Systems and Infrastructure for Web, Mobile and WoT: Federated Web and WoT systems, including distributed, federated and edge-based data processing

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.

Problem

• Medical Data Annotation
• Cardiac CT Image Analysis
• Patient Privacy Protection

Innovation

• Federated Learning
• Transformer Models
• Cardiac CT Image Analysis